fixed OpenAI fetch tasks; added nicer imports

alr_envs/open_ai/__init__.py

@@ -0,0 +1,3 @@
+from alr_envs.open_ai.mujoco import reacher_v2
+from alr_envs.open_ai.robotics import fetch
+from alr_envs.open_ai.classic_control import continuous_mountain_car

alr_envs/open_ai/classic_control/continuous_mountain_car/__init__.py

@@ -0,0 +1 @@
+from .mp_wrapper import MPWrapper

alr_envs/open_ai/classic_control/continuous_mountain_car/mp_wrapper.py

@@ -1,7 +1,7 @@
 from typing import Union
 
 import numpy as np
-from mp_env_api.interface_wrappers.mp_env_wrapper import MPEnvWrapper
+from mp_env_api import MPEnvWrapper
 
 
 class MPWrapper(MPEnvWrapper):

alr_envs/open_ai/continuous_mountain_car/__init__.py

@@ -1 +0,0 @@
-from alr_envs.open_ai.continuous_mountain_car.mp_wrapper import MPWrapper

alr_envs/open_ai/fetch/__init__.py

@@ -1 +0,0 @@
-from alr_envs.open_ai.fetch.mp_wrapper import MPWrapper

alr_envs/open_ai/fetch/mp_wrapper.py

@@ -1,22 +0,0 @@
-from typing import Union
-
-import numpy as np
-from mp_env_api.interface_wrappers.mp_env_wrapper import MPEnvWrapper
-
-
-class MPWrapper(MPEnvWrapper):
-    @property
-    def current_vel(self) -> Union[float, int, np.ndarray]:
-        return self.unwrapped._get_obs()["observation"][-5:-1]
-
-    @property
-    def current_pos(self) -> Union[float, int, np.ndarray]:
-        return self.unwrapped._get_obs()["observation"][:4]
-
-    @property
-    def goal_pos(self):
-        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/open_ai/mujoco/reacher_v2/__init__.py

@@ -0,0 +1 @@
+from .mp_wrapper import MPWrapper

alr_envs/open_ai/mujoco/reacher_v2/mp_wrapper.py

@@ -1,7 +1,7 @@
 from typing import Union
 
 import numpy as np
-from mp_env_api.interface_wrappers.mp_env_wrapper import MPEnvWrapper
+from mp_env_api import MPEnvWrapper
 
 
 class MPWrapper(MPEnvWrapper):

alr_envs/open_ai/reacher_v2/__init__.py

@@ -1 +0,0 @@
-from alr_envs.open_ai.reacher_v2.mp_wrapper import MPWrapper

alr_envs/open_ai/robotics/fetch/__init__.py

@@ -0,0 +1 @@
+from .mp_wrapper import MPWrapper

alr_envs/open_ai/robotics/fetch/mp_wrapper.py

@@ -0,0 +1,49 @@
+from typing import Union
+
+import numpy as np
+
+from mp_env_api import MPEnvWrapper
+
+
+class MPWrapper(MPEnvWrapper):
+
+    @property
+    def active_obs(self):
+        return np.hstack([
+            [False] * 3,  # achieved goal
+            [True] * 3,  # desired/true goal
+            [False] * 3,  # grip pos
+            [True, True, False] * int(self.has_object),  # object position
+            [True, True, False] * int(self.has_object),  # object relative position
+            [False] * 2,  # gripper state
+            [False] * 3 * int(self.has_object),  # object rotation
+            [False] * 3 * int(self.has_object),  # object velocity position
+            [False] * 3 * int(self.has_object),  # object velocity rotation
+            [False] * 3,  # grip velocity position
+            [False] * 2,  # gripper velocity
+        ]).astype(bool)
+
+    @property
+    def current_vel(self) -> Union[float, int, np.ndarray]:
+        dt = self.sim.nsubsteps * self.sim.model.opt.timestep
+        grip_velp = self.sim.data.get_site_xvelp("robot0:grip") * dt
+        # gripper state should be symmetric for left and right.
+        # They are controlled with only one action for both gripper joints
+        gripper_state = self.sim.data.get_joint_qvel('robot0:r_gripper_finger_joint') * dt
+        return np.hstack([grip_velp, gripper_state])
+
+    @property
+    def current_pos(self) -> Union[float, int, np.ndarray]:
+        grip_pos = self.sim.data.get_site_xpos("robot0:grip")
+        # gripper state should be symmetric for left and right.
+        # They are controlled with only one action for both gripper joints
+        gripper_state = self.sim.data.get_joint_qpos('robot0:r_gripper_finger_joint')
+        return np.hstack([grip_pos, gripper_state])
+
+    @property
+    def goal_pos(self):
+        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