Added more tests

This commit is contained in:
Fabian
2022-10-14 14:02:12 +02:00
parent 915ffbe928
commit 00f622e913
6 changed files with 89 additions and 7 deletions
@@ -1,4 +1,5 @@
from mp_pytorch.basis_gn import NormalizedRBFBasisGenerator, ZeroPaddingNormalizedRBFBasisGenerator
from mp_pytorch.basis_gn import NormalizedRBFBasisGenerator, ZeroPaddingNormalizedRBFBasisGenerator, \
ProDMPBasisGenerator
from mp_pytorch.phase_gn import PhaseGenerator
ALL_TYPES = ["rbf", "zero_rbf", "rhythmic"]
@@ -10,6 +11,10 @@ def get_basis_generator(basis_generator_type: str, phase_generator: PhaseGenerat
return NormalizedRBFBasisGenerator(phase_generator, **kwargs)
elif basis_generator_type == "zero_rbf":
return ZeroPaddingNormalizedRBFBasisGenerator(phase_generator, **kwargs)
elif basis_generator_type == "prodmp":
from mp_pytorch.phase_gn import ExpDecayPhaseGenerator
assert isinstance(phase_generator, ExpDecayPhaseGenerator)
return ProDMPBasisGenerator(phase_generator, **kwargs)
elif basis_generator_type == "rhythmic":
raise NotImplementedError()
# return RhythmicBasisGenerator(phase_generator, **kwargs)