mp wrapper fixes

alr_envs/alr/mujoco/half_cheetah_jump/__init__.py

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

alr_envs/alr/mujoco/half_cheetah_jump/half_cheetah_jump.py

@@ -1,4 +1,7 @@
 import os
+from typing import Tuple, Union, Optional
+
+from gym.core import ObsType
 from gym.envs.mujoco.half_cheetah_v3 import HalfCheetahEnv
 import numpy as np
 
@@ -20,7 +23,7 @@ class ALRHalfCheetahJumpEnv(HalfCheetahEnv):
                  max_episode_steps=100):
         self.current_step = 0
         self.max_height = 0
-        self.max_episode_steps = max_episode_steps
+        # self.max_episode_steps = max_episode_steps
         self.goal = 0
         self.context = context
         xml_file = os.path.join(os.path.dirname(__file__), "assets", xml_file)
@@ -37,15 +40,15 @@ class ALRHalfCheetahJumpEnv(HalfCheetahEnv):
 
         ## Didnt use fell_over, because base env also has no done condition - Paul and Marc
         # fell_over = abs(self.sim.data.qpos[2]) > 2.5  # how to figure out if the cheetah fell over? -> 2.5 oke?
-        # TODO: Should a fall over be checked herE?
+        # TODO: Should a fall over be checked here?
         done = False
 
         ctrl_cost = self.control_cost(action)
         costs = ctrl_cost
 
-        if self.current_step == self.max_episode_steps:
-            height_goal_distance = -10*np.linalg.norm(self.max_height - self.goal) + 1e-8 if self.context \
-                                                                                       else self.max_height
+        if self.current_step == MAX_EPISODE_STEPS_HALFCHEETAHJUMP:
+            height_goal_distance = -10 * np.linalg.norm(self.max_height - self.goal) + 1e-8 if self.context \
+                else self.max_height
             rewards = self._forward_reward_weight * height_goal_distance
         else:
             rewards = 0
@@ -62,7 +65,8 @@ class ALRHalfCheetahJumpEnv(HalfCheetahEnv):
     def _get_obs(self):
         return np.append(super()._get_obs(), self.goal)
 
-    def reset(self):
+    def reset(self, *, seed: Optional[int] = None, return_info: bool = False,
+              options: Optional[dict] = None, ) -> Union[ObsType, Tuple[ObsType, dict]]:
         self.max_height = 0
         self.current_step = 0
         self.goal = np.random.uniform(1.1, 1.6, 1)  # 1.1 1.6
@@ -80,21 +84,3 @@ class ALRHalfCheetahJumpEnv(HalfCheetahEnv):
 
         observation = self._get_obs()
         return observation
-
-if __name__ == '__main__':
-    render_mode = "human"  # "human" or "partial" or "final"
-    env = ALRHalfCheetahJumpEnv()
-    obs = env.reset()
-
-    for i in range(2000):
-        # objective.load_result("/tmp/cma")
-        # test with random actions
-        ac = env.action_space.sample()
-        obs, rew, d, info = env.step(ac)
-        if i % 10 == 0:
-            env.render(mode=render_mode)
-        if d:
-            print('After ', i, ' steps, done: ', d)
-            env.reset()
-
-    env.close()

alr_envs/alr/mujoco/half_cheetah_jump/mp_wrapper.py

@@ -10,7 +10,7 @@ class MPWrapper(RawInterfaceWrapper):
     def context_mask(self) -> np.ndarray:
         return np.hstack([
             [False] * 17,
-            [True] # goal height
+            [True]  # goal height
         ])
 
     @property
@@ -20,11 +20,3 @@ class MPWrapper(RawInterfaceWrapper):
     @property
     def current_vel(self) -> Union[float, int, np.ndarray, Tuple]:
         return self.env.sim.data.qvel[3:9].copy()
-
-    @property
-    def goal_pos(self) -> Union[float, int, np.ndarray, Tuple]:
-        raise ValueError("Goal position is not available and has to be learnt based on the environment.")
-
-    @property
-    def dt(self) -> Union[float, int]:
-        return self.env.dt

alr_envs/alr/mujoco/half_cheetah_jump/new_mp_wrapper.py

@@ -1,22 +0,0 @@
-from typing import Tuple, Union
-
-import numpy as np
-
-from alr_envs.black_box.raw_interface_wrapper import RawInterfaceWrapper
-
-
-class MPWrapper(RawInterfaceWrapper):
-    def context_mask(self):
-        return np.hstack([
-            [False] * 17,
-            [True] # goal height
-        ])
-
-    @property
-    def current_pos(self) -> Union[float, int, np.ndarray]:
-        return self.env.sim.data.qpos[3:9].copy()
-
-    @property
-    def current_vel(self) -> Union[float, int, np.ndarray, Tuple]:
-        return self.env.sim.data.qvel[3:9].copy()
-