Added ALRReacherProMP

alr_envs/alr/__init__.py

@@ -10,7 +10,7 @@ from .mujoco.ball_in_a_cup.biac_pd import ALRBallInACupPDEnv
 from .mujoco.reacher.alr_reacher import ALRReacherEnv
 from .mujoco.reacher.balancing import BalancingEnv
 
-from alr_envs.alr.mujoco.table_tennis.tt_gym import MAX_EPISODE_STEPS
+from .mujoco.table_tennis.tt_gym import MAX_EPISODE_STEPS
 
 ALL_ALR_MOTION_PRIMITIVE_ENVIRONMENTS = {"DMP": [], "ProMP": []}
 
@@ -363,6 +363,58 @@ for _v in _versions:
         }
     )
     ALL_ALR_MOTION_PRIMITIVE_ENVIRONMENTS["ProMP"].append(_env_id)
+    
+## ALRReacher
+_versions = ["ALRReacher-v0", "ALRLongReacher-v0", "ALRReacherSparse-v0", "ALRLongReacherSparse-v0"]
+for _v in _versions:
+    _name = _v.split("-")
+    _env_id = f'{_name[0]}DMP-{_name[1]}'
+    register(
+        id=_env_id,
+        entry_point='alr_envs.utils.make_env_helpers:make_dmp_env_helper',
+        # max_episode_steps=1,
+        kwargs={
+            "name": f"alr_envs:{_v}",
+            "wrappers": [mujoco.reacher.MPWrapper],
+            "mp_kwargs": {
+                "num_dof": 5 if "long" not in _v.lower() else 7,
+                "num_basis": 5,
+                "duration": 4,
+                "alpha_phase": 2,
+                "learn_goal": True,
+                "policy_type": "motor",
+                "weights_scale": 1,
+                "policy_kwargs": {
+                    "p_gains": 1,
+                    "d_gains": 0.1
+                }
+            }
+        }
+    )
+    ALL_ALR_MOTION_PRIMITIVE_ENVIRONMENTS["DMP"].append(_env_id)
+
+    _env_id = f'{_name[0]}ProMP-{_name[1]}'
+    register(
+        id=_env_id,
+        entry_point='alr_envs.utils.make_env_helpers:make_promp_env_helper',
+        kwargs={
+            "name": f"alr_envs:{_v}",
+            "wrappers": [mujoco.reacher.MPWrapper],
+            "mp_kwargs": {
+                "num_dof": 5 if "long" not in _v.lower() else 7,
+                "num_basis": 5,
+                "duration": 4,
+                "policy_type": "motor",
+                "weights_scale": 1,
+                "zero_start": True,
+                "policy_kwargs": {
+                    "p_gains": 1,
+                    "d_gains": 0.1
+                }
+            }
+        }
+    )
+    ALL_ALR_MOTION_PRIMITIVE_ENVIRONMENTS["ProMP"].append(_env_id)
 
 ## Beerpong
 _versions = ["v0", "v1", "v2", "v3"]

alr_envs/alr/mujoco/reacher/__init__.py

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

alr_envs/alr/mujoco/reacher/alr_reacher.py

@@ -42,7 +42,10 @@ class ALRReacherEnv(MujocoEnv, utils.EzPickle):
         if self._steps >= self.steps_before_reward:
             vec = self.get_body_com("fingertip") - self.get_body_com("target")
             reward_dist -= self.reward_weight * np.linalg.norm(vec)
-            angular_vel -= np.linalg.norm(self.sim.data.qvel.flat[:self.n_links])
+            if self.steps_before_reward > 0:
+                # avoid giving this penalty for normal step based case
+                angular_vel -= np.linalg.norm(self.sim.data.qvel.flat[:self.n_links])
+                # angular_vel -= np.square(self.sim.data.qvel.flat[:self.n_links]).sum()
         reward_ctrl = - np.square(a).sum()
 
         if self.balance:
@@ -61,14 +64,29 @@ class ALRReacherEnv(MujocoEnv, utils.EzPickle):
     def viewer_setup(self):
         self.viewer.cam.trackbodyid = 0
 
+    # def reset_model(self):
+    #     qpos = self.np_random.uniform(low=-0.1, high=0.1, size=self.model.nq) + self.init_qpos
+    #     while True:
+    #         self.goal = self.np_random.uniform(low=-self.n_links / 10, high=self.n_links / 10, size=2)
+    #         if np.linalg.norm(self.goal) < self.n_links / 10:
+    #             break
+    #     qpos[-2:] = self.goal
+    #     qvel = self.init_qvel + self.np_random.uniform(low=-.005, high=.005, size=self.model.nv)
+    #     qvel[-2:] = 0
+    #     self.set_state(qpos, qvel)
+    #     self._steps = 0
+    #
+    #     return self._get_obs()
+
     def reset_model(self):
-        qpos = self.np_random.uniform(low=-0.1, high=0.1, size=self.model.nq) + self.init_qpos
-        while True:
-            self.goal = self.np_random.uniform(low=-self.n_links / 10, high=self.n_links / 10, size=2)
-            if np.linalg.norm(self.goal) < self.n_links / 10:
-                break
+        qpos = self.init_qpos
+        if not hasattr(self, "goal"):
+            while True:
+                self.goal = self.np_random.uniform(low=-self.n_links / 10, high=self.n_links / 10, size=2)
+                if np.linalg.norm(self.goal) < self.n_links / 10:
+                    break
         qpos[-2:] = self.goal
-        qvel = self.init_qvel + self.np_random.uniform(low=-.005, high=.005, size=self.model.nv)
+        qvel = self.init_qvel
         qvel[-2:] = 0
         self.set_state(qpos, qvel)
         self._steps = 0

alr_envs/alr/mujoco/reacher/mp_wrapper.py

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