bugfixes

alr_envs/alr/mujoco/beerpong/mp_wrapper.py

@@ -7,7 +7,8 @@ from alr_envs.black_box.raw_interface_wrapper import RawInterfaceWrapper
 
 class MPWrapper(RawInterfaceWrapper):
 
-    def get_context_mask(self):
+    @property
+    def context_mask(self) -> np.ndarray:
         return np.hstack([
             [False] * 7,  # cos
             [False] * 7,  # sin
@@ -15,16 +16,16 @@ class MPWrapper(RawInterfaceWrapper):
             [False] * 3,  # cup_goal_diff_final
             [False] * 3,  # cup_goal_diff_top
             [True] * 2,  # xy position of cup
-            [False]  # env steps
+            # [False]  # env steps
         ])
 
     @property
     def current_pos(self) -> Union[float, int, np.ndarray, Tuple]:
-        return self.env.sim.data.qpos[0:7].copy()
+        return self.env.data.qpos[0:7].copy()
 
     @property
     def current_vel(self) -> Union[float, int, np.ndarray, Tuple]:
-        return self.env.sim.data.qvel[0:7].copy()
+        return self.env.data.qvel[0:7].copy()
 
     # TODO: Fix this
     def _episode_callback(self, action: np.ndarray, mp) -> Tuple[np.ndarray, Union[np.ndarray, None]]: