prodmp box pushing works

fancy_gym/black_box/black_box_wrapper.py

@@ -9,6 +9,7 @@ from fancy_gym.black_box.controller.base_controller import BaseController
 from fancy_gym.black_box.raw_interface_wrapper import RawInterfaceWrapper
 from fancy_gym.utils.utils import get_numpy
 
+import torch
 
 class BlackBoxWrapper(gym.ObservationWrapper):
 
@@ -80,7 +81,9 @@ class BlackBoxWrapper(gym.ObservationWrapper):
         bc_time = np.array(0 if not self.do_replanning else self.current_traj_steps * self.dt)
         # TODO we could think about initializing with the previous desired value in order to have a smooth transition
         #  at least from the planning point of view.
-        self.traj_gen.set_boundary_conditions(bc_time, self.current_pos, self.current_vel)
+        self.traj_gen.set_boundary_conditions(torch.as_tensor(bc_time),
+                                              torch.as_tensor(self.current_pos),
+                                              torch.as_tensor(self.current_vel))
         duration = None if self.learn_sub_trajectories else self.duration
         self.traj_gen.set_duration(duration, self.dt)
         # traj_dict = self.traj_gen.get_trajs(get_pos=True, get_vel=True)

fancy_gym/black_box/factory/basis_generator_factory.py

@@ -1,4 +1,4 @@
-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 +10,8 @@ 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":
+        return ProDMPBasisGenerator(phase_generator, **kwargs)
     elif basis_generator_type == "rhythmic":
         raise NotImplementedError()
         # return RhythmicBasisGenerator(phase_generator, **kwargs)