diff --git a/fancy_gym/black_box/black_box_wrapper.py b/fancy_gym/black_box/black_box_wrapper.py
index fb2bbbf..df2206a 100644
--- a/fancy_gym/black_box/black_box_wrapper.py
+++ b/fancy_gym/black_box/black_box_wrapper.py
@@ -50,7 +50,7 @@ class BlackBoxWrapper(gym.ObservationWrapper):
         self.tracking_controller = tracking_controller
         # self.time_steps = np.linspace(0, self.duration, self.traj_steps)
         # self.traj_gen.set_mp_times(self.time_steps)
-        self.traj_gen.set_duration(np.array([self.duration]), np.array([self.dt]))
+        self.traj_gen.set_duration(self.duration - self.dt, self.dt)
 
         # reward computation
         self.reward_aggregation = reward_aggregation
@@ -78,8 +78,8 @@ class BlackBoxWrapper(gym.ObservationWrapper):
         self.traj_gen.set_boundary_conditions(
             bc_time=np.array(0) if not self.do_replanning else np.array([self.current_traj_steps * self.dt]),
             bc_pos=self.current_pos, bc_vel=self.current_vel)
-        self.traj_gen.set_duration(None if self.learn_sub_trajectories else np.array([self.duration]),
-                                   np.array([self.dt]))
+        # TODO remove the - self.dt after Bruces fix.
+        self.traj_gen.set_duration(None if self.learn_sub_trajectories else self.duration - self.dt, self.dt)
         traj_dict = self.traj_gen.get_trajs(get_pos=True, get_vel=True)
         trajectory_tensor, velocity_tensor = traj_dict['pos'], traj_dict['vel']
 
@@ -87,7 +87,7 @@ class BlackBoxWrapper(gym.ObservationWrapper):
 
     def _get_traj_gen_action_space(self):
         """This function can be used to set up an individual space for the parameters of the traj_gen."""
-        min_action_bounds, max_action_bounds = self.traj_gen.get_param_bounds()
+        min_action_bounds, max_action_bounds = self.traj_gen.get_params_bounds().t()
         action_space = gym.spaces.Box(low=min_action_bounds.numpy(), high=max_action_bounds.numpy(),
                                       dtype=self.env.action_space.dtype)
         return action_space
diff --git a/fancy_gym/black_box/factory/trajectory_generator_factory.py b/fancy_gym/black_box/factory/trajectory_generator_factory.py
index 2b93a6c..5a36fdd 100644
--- a/fancy_gym/black_box/factory/trajectory_generator_factory.py
+++ b/fancy_gym/black_box/factory/trajectory_generator_factory.py
@@ -12,7 +12,9 @@ def get_trajectory_generator(
         return ProMP(basis_generator, action_dim, **kwargs)
     elif trajectory_generator_type == "dmp":
         return DMP(basis_generator, action_dim, **kwargs)
-    elif trajectory_generator_type == 'idmp':
+    elif trajectory_generator_type == 'prodmp':
+        from mp_pytorch.basis_gn import ProDMPBasisGenerator
+        assert isinstance(basis_generator, ProDMPBasisGenerator)
         return ProDMP(basis_generator, action_dim, **kwargs)
     else:
         raise ValueError(f"Specified movement primitive type {trajectory_generator_type} not supported, "
diff --git a/fancy_gym/envs/__init__.py b/fancy_gym/envs/__init__.py
index 9f0299e..ddd9c68 100644
--- a/fancy_gym/envs/__init__.py
+++ b/fancy_gym/envs/__init__.py
@@ -126,7 +126,7 @@ for _dims in [5, 7]:
     register(
         id=f'Reacher{_dims}dSparse-v0',
         entry_point='fancy_gym.envs.mujoco:ReacherEnv',
-        max_episode_steps=MAX_EPISODE_STEPS_REACHER,
+        max_episode_steps=5,
         kwargs={
             "sparse": True,
             'reward_weight': 200,