diff --git a/fancy_gym/envs/__init__.py b/fancy_gym/envs/__init__.py
index a970b51..d38194b 100644
--- a/fancy_gym/envs/__init__.py
+++ b/fancy_gym/envs/__init__.py
@@ -496,8 +496,12 @@ for _v in _versions:
     kwargs_dict_box_pushing_prodmp['name'] = _v
     kwargs_dict_box_pushing_prodmp['controller_kwargs']['p_gains'] = 0.01 * np.array([120., 120., 120., 120., 50., 30., 10.])
     kwargs_dict_box_pushing_prodmp['controller_kwargs']['d_gains'] = 0.01 * np.array([10., 10., 10., 10., 6., 5., 3.])
-    kwargs_dict_box_pushing_prodmp['trajectory_generator_kwargs']['weights_scale'] = np.array([3.4944e+01, 4.3734e+01, 9.6711e+01, 2.4429e+02, 5.8272e+02])
-    kwargs_dict_box_pushing_prodmp['trajectory_generator_kwargs']['goal_scale'] = 3.1264e-01
+    # kwargs_dict_box_pushing_prodmp['trajectory_generator_kwargs']['weights_scale'] = np.array([3.4944e+01, 4.3734e+01, 9.6711e+01, 2.4429e+02, 5.8272e+02])
+    # kwargs_dict_box_pushing_prodmp['trajectory_generator_kwargs']['goal_scale'] = 3.1264e-01
+    kwargs_dict_box_pushing_prodmp['trajectory_generator_kwargs']['weights_scale'] = np.array([100., 166., 500., 1000.])
+    kwargs_dict_box_pushing_prodmp['trajectory_generator_kwargs']['goal_scale'] = 1.
+    kwargs_dict_box_pushing_prodmp['basis_generator_kwargs']['num_basis'] = 4
+    kwargs_dict_box_pushing_prodmp['basis_generator_kwargs']['alpha'] = 10.
     kwargs_dict_box_pushing_prodmp['black_box_kwargs']['replanning_schedule'] = lambda pos, vel, obs, action, t: t % 25 == 0
     register(
         id=_env_id,
diff --git a/fancy_gym/envs/mujoco/box_pushing/mp_wrapper.py b/fancy_gym/envs/mujoco/box_pushing/mp_wrapper.py
index c321dc1..52f6e13 100644
--- a/fancy_gym/envs/mujoco/box_pushing/mp_wrapper.py
+++ b/fancy_gym/envs/mujoco/box_pushing/mp_wrapper.py
@@ -11,8 +11,8 @@ class MPWrapper(RawInterfaceWrapper):
     @property
     def context_mask(self):
         return np.hstack([
-            [False] * 7,  # joints position
-            [False] * 7,  # joints velocity
+            [True] * 7,  # joints position
+            [True] * 7,  # joints velocity
             [False] * 7,  # joints gravity compensation
             [False] * 3,  # position of rod tip
             [False] * 4,  # orientation of rod