Merge remote-tracking branch 'origin/Add-ProDMP-envs' into mujoco_binding

# Conflicts:
#	fancy_gym/black_box/factory/basis_generator_factory.py

fancy_gym/black_box/controller/base_controller.py

@@ -2,3 +2,6 @@ class BaseController:
 
     def get_action(self, des_pos, des_vel, c_pos, c_vel):
         raise NotImplementedError
+
+    def __call__(self, des_pos, des_vel, c_pos, c_vel):
+        return self.get_action(des_pos, des_vel, c_pos, c_vel)

fancy_gym/black_box/controller/meta_world_controller.py

@@ -18,7 +18,8 @@ class MetaWorldController(BaseController):
         cur_pos = c_pos[:-1]
         xyz_pos = des_pos[:-1]
 
-        assert xyz_pos.shape == cur_pos.shape, \
-            f"Mismatch in dimension between desired position {xyz_pos.shape} and current position {cur_pos.shape}"
+        if xyz_pos.shape != cur_pos.shape:
+            raise ValueError(f"Mismatch in dimension between desired position"
+                             f" {xyz_pos.shape} and current position {cur_pos.shape}")
         trq = np.hstack([(xyz_pos - cur_pos), gripper_pos])
         return trq

fancy_gym/black_box/controller/pd_controller.py

@@ -8,7 +8,6 @@ class PDController(BaseController):
     A PD-Controller. Using position and velocity information from a provided environment,
     the tracking_controller calculates a response based on the desired position and velocity
 
-    :param env: A position environment
     :param p_gains: Factors for the proportional gains
     :param d_gains: Factors for the differential gains
     """
@@ -20,9 +19,11 @@ class PDController(BaseController):
         self.d_gains = d_gains
 
     def get_action(self, des_pos, des_vel, c_pos, c_vel):
-        assert des_pos.shape == c_pos.shape, \
-            f"Mismatch in dimension between desired position {des_pos.shape} and current position {c_pos.shape}"
-        assert des_vel.shape == c_vel.shape, \
-            f"Mismatch in dimension between desired velocity {des_vel.shape} and current velocity {c_vel.shape}"
+        if des_pos.shape != c_pos.shape:
+            raise ValueError(f"Mismatch in dimension between desired position "
+                             f"{des_pos.shape} and current position {c_pos.shape}")
+        if des_vel.shape != c_vel.shape:
+            raise ValueError(f"Mismatch in dimension between desired velocity"
+                             f" {des_vel.shape} and current velocity {c_vel.shape}")
         trq = self.p_gains * (des_pos - c_pos) + self.d_gains * (des_vel - c_vel)
         return trq

fancy_gym/black_box/factory/basis_generator_factory.py

@@ -1,4 +1,5 @@
-from mp_pytorch.basis_gn import NormalizedRBFBasisGenerator, ZeroPaddingNormalizedRBFBasisGenerator, ProDMPBasisGenerator
+from mp_pytorch.basis_gn import NormalizedRBFBasisGenerator, ZeroPaddingNormalizedRBFBasisGenerator, \
+    ProDMPBasisGenerator
 from mp_pytorch.phase_gn import PhaseGenerator
 
 ALL_TYPES = ["rbf", "zero_rbf", "rhythmic"]
@@ -11,6 +12,8 @@ def get_basis_generator(basis_generator_type: str, phase_generator: PhaseGenerat
     elif basis_generator_type == "zero_rbf":
         return ZeroPaddingNormalizedRBFBasisGenerator(phase_generator, **kwargs)
     elif basis_generator_type == "prodmp":
+        from mp_pytorch.phase_gn import ExpDecayPhaseGenerator
+        assert isinstance(phase_generator, ExpDecayPhaseGenerator)
         return ProDMPBasisGenerator(phase_generator, **kwargs)
     elif basis_generator_type == "rhythmic":
         raise NotImplementedError()

fancy_gym/envs/__init__.py

@@ -248,7 +248,6 @@ register(
     max_episode_steps=FIXED_RELEASE_STEP,
 )
 
-
 # movement Primitive Environments
 
 ## Simple Reacher

fancy_gym/meta/README.MD

@@ -1,6 +1,6 @@
 # MetaWorld Wrappers
 
-These are the Environment Wrappers for selected [Metaworld](https://meta-world.github.io/) environments in order to use our Motion Primitive gym interface with them. 
+These are the Environment Wrappers for selected [Metaworld](https://meta-world.github.io/) environments in order to use our Movement Primitive gym interface with them. 
 All Metaworld environments have a 39 dimensional observation space with the same structure. The tasks differ only in the objective and the initial observations that are randomized. 
 Unused observations are zeroed out. E.g. for `Button-Press-v2` the observation mask looks the following:
 ```python

fancy_gym/meta/__init__.py

@@ -28,11 +28,31 @@ DEFAULT_BB_DICT_ProMP = {
     }
 }
 
+DEFAULT_BB_DICT_ProDMP = {
+    "name": 'EnvName',
+    "wrappers": [],
+    "trajectory_generator_kwargs": {
+        'trajectory_generator_type': 'prodmp'
+    },
+    "phase_generator_kwargs": {
+        'phase_generator_type': 'exp'
+    },
+    "controller_kwargs": {
+        'controller_type': 'metaworld',
+    },
+    "basis_generator_kwargs": {
+        'basis_generator_type': 'prodmp',
+        'num_basis': 5
+    }
+}
+
 _goal_change_envs = ["assembly-v2", "pick-out-of-hole-v2", "plate-slide-v2", "plate-slide-back-v2",
                      "plate-slide-side-v2", "plate-slide-back-side-v2"]
 for _task in _goal_change_envs:
     task_id_split = _task.split("-")
     name = "".join([s.capitalize() for s in task_id_split[:-1]])
+
+    # ProMP
     _env_id = f'{name}ProMP-{task_id_split[-1]}'
     kwargs_dict_goal_change_promp = deepcopy(DEFAULT_BB_DICT_ProMP)
     kwargs_dict_goal_change_promp['wrappers'].append(goal_change_mp_wrapper.MPWrapper)
@@ -45,10 +65,25 @@ for _task in _goal_change_envs:
     )
     ALL_METAWORLD_MOVEMENT_PRIMITIVE_ENVIRONMENTS["ProMP"].append(_env_id)
 
+    # ProDMP
+    _env_id = f'{name}ProDMP-{task_id_split[-1]}'
+    kwargs_dict_goal_change_prodmp = deepcopy(DEFAULT_BB_DICT_ProDMP)
+    kwargs_dict_goal_change_prodmp['wrappers'].append(goal_change_mp_wrapper.MPWrapper)
+    kwargs_dict_goal_change_prodmp['name'] = f'metaworld:{_task}'
+
+    register(
+        id=_env_id,
+        entry_point='fancy_gym.utils.make_env_helpers:make_bb_env_helper',
+        kwargs=kwargs_dict_goal_change_prodmp
+    )
+    ALL_METAWORLD_MOVEMENT_PRIMITIVE_ENVIRONMENTS["ProDMP"].append(_env_id)
+
 _object_change_envs = ["bin-picking-v2", "hammer-v2", "sweep-into-v2"]
 for _task in _object_change_envs:
     task_id_split = _task.split("-")
     name = "".join([s.capitalize() for s in task_id_split[:-1]])
+
+    # ProMP
     _env_id = f'{name}ProMP-{task_id_split[-1]}'
     kwargs_dict_object_change_promp = deepcopy(DEFAULT_BB_DICT_ProMP)
     kwargs_dict_object_change_promp['wrappers'].append(object_change_mp_wrapper.MPWrapper)
@@ -60,6 +95,18 @@ for _task in _object_change_envs:
     )
     ALL_METAWORLD_MOVEMENT_PRIMITIVE_ENVIRONMENTS["ProMP"].append(_env_id)
 
+    # ProDMP
+    _env_id = f'{name}ProDMP-{task_id_split[-1]}'
+    kwargs_dict_object_change_prodmp = deepcopy(DEFAULT_BB_DICT_ProDMP)
+    kwargs_dict_object_change_prodmp['wrappers'].append(object_change_mp_wrapper.MPWrapper)
+    kwargs_dict_object_change_prodmp['name'] = f'metaworld:{_task}'
+    register(
+        id=_env_id,
+        entry_point='fancy_gym.utils.make_env_helpers:make_bb_env_helper',
+        kwargs=kwargs_dict_object_change_prodmp
+    )
+    ALL_METAWORLD_MOVEMENT_PRIMITIVE_ENVIRONMENTS["ProDMP"].append(_env_id)
+
 _goal_and_object_change_envs = ["box-close-v2", "button-press-v2", "button-press-wall-v2", "button-press-topdown-v2",
                                 "button-press-topdown-wall-v2", "coffee-button-v2", "coffee-pull-v2",
                                 "coffee-push-v2", "dial-turn-v2", "disassemble-v2", "door-close-v2",
@@ -74,6 +121,8 @@ _goal_and_object_change_envs = ["box-close-v2", "button-press-v2", "button-press
 for _task in _goal_and_object_change_envs:
     task_id_split = _task.split("-")
     name = "".join([s.capitalize() for s in task_id_split[:-1]])
+
+    # ProMP
     _env_id = f'{name}ProMP-{task_id_split[-1]}'
     kwargs_dict_goal_and_object_change_promp = deepcopy(DEFAULT_BB_DICT_ProMP)
     kwargs_dict_goal_and_object_change_promp['wrappers'].append(goal_object_change_mp_wrapper.MPWrapper)
@@ -86,10 +135,26 @@ for _task in _goal_and_object_change_envs:
     )
     ALL_METAWORLD_MOVEMENT_PRIMITIVE_ENVIRONMENTS["ProMP"].append(_env_id)
 
+    # ProDMP
+    _env_id = f'{name}ProDMP-{task_id_split[-1]}'
+    kwargs_dict_goal_and_object_change_prodmp = deepcopy(DEFAULT_BB_DICT_ProDMP)
+    kwargs_dict_goal_and_object_change_prodmp['wrappers'].append(goal_object_change_mp_wrapper.MPWrapper)
+    kwargs_dict_goal_and_object_change_prodmp['name'] = f'metaworld:{_task}'
+
+    register(
+        id=_env_id,
+        entry_point='fancy_gym.utils.make_env_helpers:make_bb_env_helper',
+        kwargs=kwargs_dict_goal_and_object_change_prodmp
+    )
+    ALL_METAWORLD_MOVEMENT_PRIMITIVE_ENVIRONMENTS["ProDMP"].append(_env_id)
+
+
 _goal_and_endeffector_change_envs = ["basketball-v2"]
 for _task in _goal_and_endeffector_change_envs:
     task_id_split = _task.split("-")
     name = "".join([s.capitalize() for s in task_id_split[:-1]])
+
+    # ProMP
     _env_id = f'{name}ProMP-{task_id_split[-1]}'
     kwargs_dict_goal_and_endeffector_change_promp = deepcopy(DEFAULT_BB_DICT_ProMP)
     kwargs_dict_goal_and_endeffector_change_promp['wrappers'].append(goal_endeffector_change_mp_wrapper.MPWrapper)
@@ -101,3 +166,16 @@ for _task in _goal_and_endeffector_change_envs:
         kwargs=kwargs_dict_goal_and_endeffector_change_promp
     )
     ALL_METAWORLD_MOVEMENT_PRIMITIVE_ENVIRONMENTS["ProMP"].append(_env_id)
+
+    # ProDMP
+    _env_id = f'{name}ProDMP-{task_id_split[-1]}'
+    kwargs_dict_goal_and_endeffector_change_prodmp = deepcopy(DEFAULT_BB_DICT_ProDMP)
+    kwargs_dict_goal_and_endeffector_change_prodmp['wrappers'].append(goal_endeffector_change_mp_wrapper.MPWrapper)
+    kwargs_dict_goal_and_endeffector_change_prodmp['name'] = f'metaworld:{_task}'
+
+    register(
+        id=_env_id,
+        entry_point='fancy_gym.utils.make_env_helpers:make_bb_env_helper',
+        kwargs=kwargs_dict_goal_and_endeffector_change_prodmp
+    )
+    ALL_METAWORLD_MOVEMENT_PRIMITIVE_ENVIRONMENTS["ProDMP"].append(_env_id)

fancy_gym/utils/make_env_helpers.py

@@ -141,7 +141,7 @@ def make_bb(
     Returns: DMP wrapped gym env
 
     """
-    _verify_time_limit(traj_gen_kwargs.get("duration", None), kwargs.get("time_limit", None))
+    _verify_time_limit(traj_gen_kwargs.get("duration"), kwargs.get("time_limit"))
 
     learn_sub_trajs = black_box_kwargs.get('learn_sub_trajectories')
     do_replanning = black_box_kwargs.get('replanning_schedule')