current state

alr_envs/black_box/factory/basis_generator_factory.py

@@ -0,0 +1,17 @@
+from mp_pytorch import PhaseGenerator, NormalizedRBFBasisGenerator, ZeroStartNormalizedRBFBasisGenerator
+from mp_pytorch.basis_gn.rhytmic_basis import RhythmicBasisGenerator
+
+ALL_TYPES = ["rbf", "zero_rbf", "rhythmic"]
+
+
+def get_basis_generator(basis_generator_type: str, phase_generator: PhaseGenerator, **kwargs):
+    basis_generator_type = basis_generator_type.lower()
+    if basis_generator_type == "rbf":
+        return NormalizedRBFBasisGenerator(phase_generator, **kwargs)
+    elif basis_generator_type == "zero_rbf":
+        return ZeroStartNormalizedRBFBasisGenerator(phase_generator, **kwargs)
+    elif basis_generator_type == "rhythmic":
+        return RhythmicBasisGenerator(phase_generator, **kwargs)
+    else:
+        raise ValueError(f"Specified basis generator type {basis_generator_type} not supported, "
+                         f"please choose one of {ALL_TYPES}.")

alr_envs/black_box/factory/phase_generator_factory.py

@@ -0,0 +1,20 @@
+from mp_pytorch import LinearPhaseGenerator, ExpDecayPhaseGenerator
+from mp_pytorch.phase_gn.rhythmic_phase_generator import RhythmicPhaseGenerator
+from mp_pytorch.phase_gn.smooth_phase_generator import SmoothPhaseGenerator
+
+ALL_TYPES = ["linear", "exp", "rhythmic", "smooth"]
+
+
+def get_phase_generator(phase_generator_type, **kwargs):
+    phase_generator_type = phase_generator_type.lower()
+    if phase_generator_type == "linear":
+        return LinearPhaseGenerator(**kwargs)
+    elif phase_generator_type == "exp":
+        return ExpDecayPhaseGenerator(**kwargs)
+    elif phase_generator_type == "rhythmic":
+        return RhythmicPhaseGenerator(**kwargs)
+    elif phase_generator_type == "smooth":
+        return SmoothPhaseGenerator(**kwargs)
+    else:
+        raise ValueError(f"Specified phase generator type {phase_generator_type} not supported, "
+                         f"please choose one of {ALL_TYPES}.")

alr_envs/black_box/factory/trajectory_generator_factory.py

@@ -0,0 +1,22 @@
+from mp_pytorch.mp.dmp import DMP
+from mp_pytorch.mp.promp import ProMP
+from mp_pytorch.mp.idmp import IDMP
+
+from mp_pytorch.basis_gn.basis_generator import BasisGenerator
+
+ALL_TYPES = ["promp", "dmp", "idmp"]
+
+
+def get_trajectory_generator(
+        trajectory_generator_type: str, action_dim: int, basis_generator: BasisGenerator, **kwargs
+):
+    trajectory_generator_type = trajectory_generator_type.lower()
+    if trajectory_generator_type == "promp":
+        return ProMP(basis_generator, action_dim, **kwargs)
+    elif trajectory_generator_type == "dmp":
+        return DMP(basis_generator, action_dim, **kwargs)
+    elif trajectory_generator_type == 'idmp':
+        return IDMP(basis_generator, action_dim, **kwargs)
+    else:
+        raise ValueError(f"Specified movement primitive type {trajectory_generator_type} not supported, "
+                         f"please choose one of {ALL_TYPES}.")