mp wrapper fixes

alr_envs/alr/mujoco/beerpong/__init__.py

@@ -1 +1 @@
-from .new_mp_wrapper import MPWrapper
+from .mp_wrapper import MPWrapper

alr_envs/alr/mujoco/beerpong/beerpong.py

@@ -1,12 +1,11 @@
 import os
+from typing import Optional
 
 import mujoco_py.builder
 import numpy as np
 from gym import utils
 from gym.envs.mujoco import MujocoEnv
 
-from alr_envs.alr.mujoco.beerpong.deprecated.beerpong_reward_staged import BeerPongReward
-
 # XML Variables
 ROBOT_COLLISION_OBJ = ["wrist_palm_link_convex_geom",
                        "wrist_pitch_link_convex_decomposition_p1_geom",
@@ -76,7 +75,7 @@ class BeerPongEnv(MujocoEnv, utils.EzPickle):
     def start_vel(self):
         return self._start_vel
 
-    def reset(self):
+    def reset(self, *, seed: Optional[int] = None, return_info: bool = False, options: Optional[dict] = None):
         self.dists = []
         self.dists_final = []
         self.action_costs = []

alr_envs/alr/mujoco/beerpong/mp_wrapper.py

@@ -1,4 +1,4 @@
-from typing import Tuple, Union
+from typing import Union, Tuple
 
 import numpy as np
 
@@ -7,34 +7,36 @@ from alr_envs.black_box.raw_interface_wrapper import RawInterfaceWrapper
 
 class MPWrapper(RawInterfaceWrapper):
 
-    @property
-    def context_mask(self) -> np.ndarray:
+    def get_context_mask(self):
         return np.hstack([
             [False] * 7,  # cos
             [False] * 7,  # sin
-            [False] * 7, # joint velocities
-            [False] * 3, # cup_goal_diff_final
-            [False] * 3, # cup_goal_diff_top
+            [False] * 7,  # joint velocities
+            [False] * 3,  # cup_goal_diff_final
+            [False] * 3,  # cup_goal_diff_top
             [True] * 2,  # xy position of cup
             [False]  # env steps
         ])
 
-    @property
-    def start_pos(self):
-        return self._start_pos
-
     @property
     def current_pos(self) -> Union[float, int, np.ndarray, Tuple]:
-        return self.sim.data.qpos[0:7].copy()
+        return self.env.sim.data.qpos[0:7].copy()
 
     @property
     def current_vel(self) -> Union[float, int, np.ndarray, Tuple]:
-        return self.sim.data.qvel[0:7].copy()
+        return self.env.sim.data.qvel[0:7].copy()
 
-    @property
-    def goal_pos(self):
-        raise ValueError("Goal position is not available and has to be learnt based on the environment.")
+    # TODO: Fix this
+    def _episode_callback(self, action: np.ndarray, mp) -> Tuple[np.ndarray, Union[np.ndarray, None]]:
+        if mp.learn_tau:
+            self.env.env.release_step = action[0] / self.env.dt  # Tau value
+            return action, None
+        else:
+            return action, None
 
-    @property
-    def dt(self) -> Union[float, int]:
-        return self.env.dt
+    def set_context(self, context):
+        xyz = np.zeros(3)
+        xyz[:2] = context
+        xyz[-1] = 0.840
+        self.env.env.model.body_pos[self.env.env.cup_table_id] = xyz
+        return self.get_observation_from_step(self.env.env._get_obs())

alr_envs/alr/mujoco/beerpong/new_mp_wrapper.py

@@ -1,41 +0,0 @@
-from typing import Union, Tuple
-
-import numpy as np
-
-from alr_envs.black_box.raw_interface_wrapper import RawInterfaceWrapper
-
-
-class MPWrapper(RawInterfaceWrapper):
-    @property
-    def current_pos(self) -> Union[float, int, np.ndarray, Tuple]:
-        return self.env.sim.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()
-
-    def get_context_mask(self):
-        return np.hstack([
-            [False] * 7,  # cos
-            [False] * 7,  # sin
-            [False] * 7,  # joint velocities
-            [False] * 3,  # cup_goal_diff_final
-            [False] * 3,  # cup_goal_diff_top
-            [True] * 2,  # xy position of cup
-            [False]  # env steps
-        ])
-
-    # TODO: Fix this
-    def _episode_callback(self, action: np.ndarray, mp) -> Tuple[np.ndarray, Union[np.ndarray, None]]:
-        if mp.learn_tau:
-            self.env.env.release_step = action[0] / self.env.dt  # Tau value
-            return action, None
-        else:
-            return action, None
-
-    def set_context(self, context):
-        xyz = np.zeros(3)
-        xyz[:2] = context
-        xyz[-1] = 0.840
-        self.env.env.model.body_pos[self.env.env.cup_table_id] = xyz
-        return self.get_observation_from_step(self.env.env._get_obs())