start integrating mp_pytorch lib

alr_envs/alr/mujoco/beerpong/new_mp_wrapper.py

@@ -0,0 +1,20 @@
+from mp_wrapper import BaseMPWrapper
+from typing import Union, Tuple
+import numpy as np
+
+
+class MPWrapper(BaseMPWrapper):
+
+    def current_pos(self) -> Union[float, int, np.ndarray, Tuple]:
+        return self.env.sim.data.qpos[0:7].copy()
+
+    def current_vel(self) -> Union[float, int, np.ndarray, Tuple]:
+        return self.env.sim.data.qvel[0:7].copy()
+
+    def set_active_obs(self):
+        return np.hstack([
+            [False] * 7,  # cos
+            [False] * 7,  # sin
+            [True] * 2,  # xy position of cup
+            [False]  # env steps
+        ])

alr_envs/alr/mujoco/reacher/new_mp_wrapper.py

@@ -0,0 +1,24 @@
+from mp_wrapper import BaseMPWrapper
+from typing import Union, Tuple
+import numpy as np
+
+
+class MPWrapper(BaseMPWrapper):
+
+    @property
+    def current_pos(self) -> Union[float, int, np.ndarray, Tuple]:
+        return self.env.sim.data.qpos.flat[:self.env.n_links]
+    @property
+    def current_vel(self) -> Union[float, int, np.ndarray, Tuple]:
+        return self.env.sim.data.qvel.flat[:self.env.n_links]
+
+    def set_active_obs(self):
+        return np.concatenate([
+            [False] * self.env.n_links,  # cos
+            [False] * self.env.n_links,  # sin
+            [True] * 2,  # goal position
+            [False] * self.env.n_links,  # angular velocity
+            [False] * 3,  # goal distance
+            # self.get_body_com("target"),  # only return target to make problem harder
+            [False],  # step
+        ])