wrappers updated

alr_envs/alr/mujoco/hopper_jump/hopper_jump.py

@@ -54,7 +54,8 @@ class ALRHopperJumpEnv(HopperEnv):
         self.current_step += 1
         self.do_simulation(action, self.frame_skip)
         height_after = self.get_body_com("torso")[2]
-        site_pos_after = self.sim.data.site_xpos[self.model.site_name2id('foot_site')].copy()
+        # site_pos_after = self.sim.data.site_xpos[self.model.site_name2id('foot_site')].copy()
+        site_pos_after = self.get_body_com('foot_site')
         self.max_height = max(height_after, self.max_height)
 
         ctrl_cost = self.control_cost(action)

alr_envs/alr/mujoco/hopper_jump/mp_wrapper.py

@@ -2,12 +2,12 @@ from typing import Tuple, Union
 
 import numpy as np
 
-from mp_env_api import MPEnvWrapper
+from alr_envs.mp.raw_interface_wrapper import RawInterfaceWrapper
 
 
-class MPWrapper(MPEnvWrapper):
+class MPWrapper(RawInterfaceWrapper):
     @property
-    def active_obs(self):
+    def context_mask(self) -> np.ndarray:
         return np.hstack([
             [False] * (5 + int(not self.exclude_current_positions_from_observation)),  # position
             [False] * 6,  # velocity