Added more tests
This commit is contained in:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user