wrappers updated

alr_envs/dmc/suite/ball_in_cup/mp_wrapper.py

@@ -2,13 +2,13 @@ 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:
         # Besides the ball position, the environment is always set to 0.
         return np.hstack([
             [False] * 2,  # cup position

alr_envs/dmc/suite/cartpole/mp_wrapper.py

@@ -2,18 +2,17 @@ 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):
 
     def __init__(self, env, n_poles: int = 1):
         self.n_poles = n_poles
         super().__init__(env)
 
-
     @property
-    def active_obs(self):
+    def context_mask(self) -> np.ndarray:
         # Besides the ball position, the environment is always set to 0.
         return np.hstack([
             [True],  # slider position

alr_envs/dmc/suite/reacher/mp_wrapper.py

@@ -2,13 +2,13 @@ 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:
         # Joint and target positions are randomized, velocities are always set to 0.
         return np.hstack([
             [True] * 2,  # joint position