new reacher mp wrapper for Philipp

alr_envs/alr/__init__.py

@@ -668,21 +668,33 @@ for _v in _versions:
     _env_id = f'{_name[0]}ProMP-{_name[1]}'
     register(
         id=_env_id,
-        entry_point='alr_envs.utils.make_env_helpers:make_promp_env_helper',
+        entry_point='alr_envs.utils.make_env_helpers:make_mp_env_helper',
         kwargs={
             "name": f"alr_envs:{_v}",
-            "wrappers": [mujoco.reacher.MPWrapper],
+            "wrappers": [mujoco.reacher.NewMPWrapper],
-            "mp_kwargs": {
+            "ep_wrapper_kwargs": {
-                "num_dof": 5 if "long" not in _v.lower() else 7,
+                "weight_scale": 1
-                "num_basis": 2,
+            },
-                "duration": 4,
+            "movement_primitives_kwargs": {
-                "policy_type": "motor",
+                'movement_primitives_type': 'promp',
-                "weights_scale": 5,
+                'action_dim': 5 if "long" not in _v.lower() else 7
-                "zero_start": True,
+            },
-                "policy_kwargs": {
+            "phase_generator_kwargs": {
-                    "p_gains": 1,
+                'phase_generator_type': 'linear',
-                    "d_gains": 0.1
+                'delay': 0,
-                }
+                'tau': 4,  # initial value
+                'learn_tau': False,
+                'learn_delay': False
+            },
+            "controller_kwargs": {
+                'controller_type': 'motor',
+                "p_gains": 1,
+                "d_gains": 0.1
+            },
+            "basis_generator_kwargs": {
+                'basis_generator_type': 'zero_rbf',
+                'num_basis': 2,
+                'num_basis_zero_start': 1
             }
         }
     )
@@ -1275,4 +1287,4 @@ for _v in _versions:
             }
         }
     )
-    ALL_ALR_MOTION_PRIMITIVE_ENVIRONMENTS["ProMP"].append(_env_id)
+    ALL_ALR_MOTION_PRIMITIVE_ENVIRONMENTS["ProMP"].append(_env_id)

alr_envs/alr/mujoco/reacher/__init__.py

@@ -1 +1,2 @@
-from .mp_wrapper import MPWrapper
+from .mp_wrapper import MPWrapper
+from .new_mp_wrapper import NewMPWrapper

alr_envs/alr/mujoco/reacher/alr_reacher.py

@@ -149,4 +149,4 @@ if __name__ == '__main__':
         if d:
             env.reset()
 
-    env.close()
+    env.close()

alr_envs/alr/mujoco/reacher/new_mp_wrapper.py

@@ -3,7 +3,7 @@ from typing import Union, Tuple
 import numpy as np
 
 
-class MPWrapper(EpisodicWrapper):
+class NewMPWrapper(EpisodicWrapper):
 
     @property
     def current_pos(self) -> Union[float, int, np.ndarray, Tuple]:
@@ -21,4 +21,4 @@ class MPWrapper(EpisodicWrapper):
             [False] * 3,  # goal distance
             # self.get_body_com("target"),  # only return target to make problem harder
             [False],  # step
-        ])
+        ])