Merge branch 'dmc_integration' into dmc_integration

README.md

@@ -1,105 +1,212 @@
 ## ALR Robotics Control Environments
 
-This repository collects custom Robotics environments not included in benchmark suites like OpenAI gym, rllab, etc. 
-Creating a custom (Mujoco) gym environment can be done according to [this guide](https://github.com/openai/gym/blob/master/docs/creating-environments.md).
-For stochastic search problems with gym interface use the `Rosenbrock-v0` reference implementation.
-We also support to solve environments with Dynamic Movement Primitives (DMPs) and Probabilistic Movement Primitives (DetPMP, we only consider the mean usually). 
-
-## Step-based Environments
-Currently we have the following environments: 
-
-### Mujoco
-
-|Name| Description|Horizon|Action Dimension|Observation Dimension
-|---|---|---|---|---|
-|`ALRReacher-v0`|Modified (5 links) Mujoco gym's `Reacher-v2` (2 links)| 200 | 5 | 21
-|`ALRReacherSparse-v0`|Same as `ALRReacher-v0`, but the distance penalty is only provided in the last time step.| 200 | 5 | 21
-|`ALRReacherSparseBalanced-v0`|Same as `ALRReacherSparse-v0`, but the end-effector has to remain upright.| 200 | 5 | 21
-|`ALRLongReacher-v0`|Modified (7 links) Mujoco gym's `Reacher-v2` (2 links)| 200 | 7 | 27
-|`ALRLongReacherSparse-v0`|Same as `ALRLongReacher-v0`, but the distance penalty is only provided in the last time step.| 200 | 7 | 27
-|`ALRLongReacherSparseBalanced-v0`|Same as `ALRLongReacherSparse-v0`, but the end-effector has to remain upright.| 200 | 7 | 27
-|`ALRBallInACupSimple-v0`| Ball-in-a-cup task where a robot needs to catch a ball attached to a cup at its end-effector. | 4000 | 3 | wip
-|`ALRBallInACup-v0`| Ball-in-a-cup task where a robot needs to catch a ball attached to a cup at its end-effector | 4000 | 7 | wip
-|`ALRBallInACupGoal-v0`| Similiar to `ALRBallInACupSimple-v0` but the ball needs to be caught at a specified goal position | 4000 | 7 | wip
-    
-### Classic Control
-
-|Name| Description|Horizon|Action Dimension|Observation Dimension
-|---|---|---|---|---|
-|`SimpleReacher-v0`| Simple reaching task (2 links) without any physics simulation. Provides no reward until 150 time steps. This allows the agent to explore the space, but requires precise actions towards the end of the trajectory.| 200 | 2 | 9
-|`LongSimpleReacher-v0`| Simple reaching task (5 links) without any physics simulation. Provides no reward until 150 time steps. This allows the agent to explore the space, but requires precise actions towards the end of the trajectory.| 200 | 5 | 18
-|`ViaPointReacher-v0`| Simple reaching task leveraging a via point, which supports self collision detection. Provides a reward only at 100 and 199 for reaching the viapoint and goal point, respectively.| 200 | 5 | 18 
-|`HoleReacher-v0`| 5 link reaching task where the end-effector needs to reach into a narrow hole without collding with itself or walls | 200 | 5 | 18
+This project offers a large verity of reinforcement learning environments under a unifying interface base on OpenAI gym.
+Besides, some custom environments we also provide support for the benchmark suites
+[OpenAI gym](https://gym.openai.com/),
+[DeepMind Control](https://deepmind.com/research/publications/2020/dm-control-Software-and-Tasks-for-Continuous-Control)
+(DMC), and [Metaworld](https://meta-world.github.io/). Custom (Mujoco) gym environment can be created according
+to [this guide](https://github.com/openai/gym/blob/master/docs/creating-environments.md). Unlike existing libraries, we
+further support to control agents with Dynamic Movement Primitives (DMPs) and Probabilistic Movement Primitives (DetPMP,
+we only consider the mean usually).
 
 ## Motion Primitive Environments (Episodic environments)
-Unlike step-based environments, these motion primitive (MP) environments are closer to stochastic search and what can be found in robotics. They always execute a full trajectory, which is computed by a Dynamic Motion Primitive (DMP) or Probabilitic Motion Primitive (DetPMP) and translated into individual actions with a controller, e.g. a PD controller. The actual Controller, however, depends on the type of environment, i.e. position, velocity, or torque controlled.
-The goal is to learn the parametrization of the motion primitives in order to generate a suitable trajectory. 
-MP This can also be done in a contextual setting, where all changing elements of the task are exposed once in the beginning. This requires to find a new parametrization for each trajectory.
-All environments provide the full cumulative episode reward and additional information about early terminations, e.g. due to collisions. 
 
-### Classic Control
-|Name| Description|Horizon|Action Dimension|Context Dimension
-|---|---|---|---|---|
-|`ViaPointReacherDMP-v0`| A DMP provides a trajectory for the `ViaPointReacher-v0` task. | 200 | 25
-|`HoleReacherFixedGoalDMP-v0`| A DMP provides a trajectory for the `HoleReacher-v0` task with a fixed goal attractor. | 200 | 25
-|`HoleReacherDMP-v0`| A DMP provides a trajectory for the `HoleReacher-v0` task. The goal attractor needs to be learned. | 200 | 30 
-|`ALRBallInACupSimpleDMP-v0`| A DMP provides a trajectory for the `ALRBallInACupSimple-v0` task where only 3 joints are actuated. | 4000 | 15
-|`ALRBallInACupDMP-v0`| A DMP provides a trajectory for the `ALRBallInACup-v0` task. | 4000 | 35
-|`ALRBallInACupGoalDMP-v0`| A DMP provides a trajectory for the `ALRBallInACupGoal-v0` task. | 4000 | 35 | 3 
+Unlike step-based environments, motion primitive (MP) environments are closer related to stochastic search, black box
+optimization and methods that often used in robotics. MP environments are trajectory-based and always execute a full
+trajectory, which is generated by a Dynamic Motion Primitive (DMP) or a Probabilistic Motion Primitive (DetPMP). The
+generated trajectory is translated into individual step-wise actions by a controller. The exact choice of controller is,
+however, dependent on the type of environment. We currently support position, velocity, and PD-Controllers for position,
+velocity and torque control, respectively. The goal of all MP environments is still to learn a policy. Yet, an action
+represents the parametrization of the motion primitives to generate a suitable trajectory. Additionally, in this
+framework we support the above setting for the contextual setting, for which we expose all changing substates of the
+task as a single observation in the beginning. This requires to predict a new action/MP parametrization for each
+trajectory. All environments provide the next to the cumulative episode reward also all collected information from each
+step as part of the info dictionary. This information should, however, mainly be used for debugging and logging.
 
-[//]:  |`HoleReacherDetPMP-v0`|
+|Key| Description|
+|---|---|
+`trajectory`| Generated trajectory from MP
+`step_actions`| Step-wise executed action based on controller output
+`step_observations`| Step-wise intermediate observations
+`step_rewards`| Step-wise rewards
+`trajectory_length`| Total number of environment interactions
+`other`| All other information from the underlying environment are returned as a list with length `trajectory_length` maintaining the original key. In case some information are not provided every time step, the missing values are filled with `None`.
 
-### OpenAI gym Environments
-These environments are wrapped-versions of their OpenAI-gym counterparts.
+## Installation
 
-|Name| Description|Trajectory Horizon|Action Dimension|Context Dimension
-|---|---|---|---|---|
-|`ContinuousMountainCarDetPMP-v0`| A DetPmP wrapped version of the ContinuousMountainCar-v0 environment. | 100 | 1
-|`ReacherDetPMP-v2`| A DetPmP wrapped version of the Reacher-v2 environment. | 50 | 2
-|`FetchSlideDenseDetPMP-v1`| A DetPmP wrapped version of the FetchSlideDense-v1 environment. | 50 | 4 
-|`FetchReachDenseDetPMP-v1`| A DetPmP wrapped version of the FetchReachDense-v1 environment. | 50 | 4
-
-### Deep Mind Control Suite Environments
-These environments are wrapped-versions of their Deep Mind Control Suite (DMC) counterparts.
-Given most task can be solved in shorter horizon lengths than the original 1000 steps, we often shorten the episodes for those task. 
-
-|Name| Description|Trajectory Horizon|Action Dimension|Context Dimension
-|---|---|---|---|---|
-|`dmc_ball_in_cup-catch_detpmp-v0`| A DetPmP wrapped version of the "catch" task for the "ball_in_cup" environment. | 50 | 10 | 2
-|`dmc_ball_in_cup-catch_dmp-v0`| A DMP wrapped version of the "catch" task for the "ball_in_cup" environment. | 50| 10 | 2
-|`dmc_reacher-easy_detpmp-v0`| A DetPmP wrapped version of the "easy" task for the "reacher" environment. | 1000 | 10 | 4
-|`dmc_reacher-easy_dmp-v0`| A DMP wrapped version of the "easy" task for the "reacher" environment. | 1000| 10 | 4
-|`dmc_reacher-hard_detpmp-v0`| A DetPmP wrapped version of the "hard" task for the "reacher" environment.| 1000 | 10 | 4
-|`dmc_reacher-hard_dmp-v0`| A DMP wrapped version of the "hard" task for the "reacher" environment. | 1000 | 10 | 4
-
-## Install
 1. Clone the repository
+
 ```bash 
 git clone git@github.com:ALRhub/alr_envs.git
 ```
+
 2. Go to the folder
+
 ```bash 
 cd alr_envs
 ```
+
 3. Install with
+
 ```bash 
 pip install -e . 
 ```
-4. Use (see [example.py](alr_envs/examples/examples_general.py)): 
-```python
-import gym
 
-env = gym.make('alr_envs:SimpleReacher-v0')
+## Using the framework
+
+We prepared [multiple examples](alr_envs/examples/), please have a look there for more specific examples.
+
+### Step-wise environments
+
+```python
+import alr_envs
+
+env = alr_envs.make('HoleReacher-v0', seed=1)
 state = env.reset()
 
-for i in range(10000):
+for i in range(1000):
     state, reward, done, info = env.step(env.action_space.sample())
     if i % 5 == 0:
         env.render()
 
     if done:
         state = env.reset()
+``` 
+
+For Deepmind control tasks we expect the `env_id` to be specified as `domain_name-task_name` or for manipulation tasks
+as `manipulation-environment_name`. All other environments can be created based on their original name.
+
+Existing MP tasks can be created the same way as above. Just keep in mind, calling `step()` always executs a full
+trajectory.
+
+```python
+import alr_envs
+
+env = alr_envs.make('HoleReacherDetPMP-v0', seed=1)
+# render() can be called once in the beginning with all necessary arguments. To turn it of again just call render(None). 
+env.render()
+
+state = env.reset()
+
+for i in range(5):
+    state, reward, done, info = env.step(env.action_space.sample())
+
+    # Not really necessary as the environments resets itself after each trajectory anyway.
+    state = env.reset()
+```
+
+To show all available environments, we provide some additional convenience. Each value will return a dictionary with two
+keys `DMP` and `DetPMP` that store a list of available environment names.
+
+```python
+import alr_envs
+
+print("Custom MP tasks:")
+print(alr_envs.ALL_ALR_MOTION_PRIMITIVE_ENVIRONMENTS)
+
+print("OpenAI Gym MP tasks:")
+print(alr_envs.ALL_GYM_MOTION_PRIMITIVE_ENVIRONMENTS)
+
+print("Deepmind Control MP tasks:")
+print(alr_envs.ALL_DEEPMIND_MOTION_PRIMITIVE_ENVIRONMENTS)
+
+print("MetaWorld MP tasks:")
+print(alr_envs.ALL_METAWORLD_MOTION_PRIMITIVE_ENVIRONMENTS)
+```
+
+### How to create a new MP task
+
+In case a required task is not supported yet in the MP framework, it can be created relatively easy. For the task at
+hand, the following interface needs to be implemented.
+
+```python
+import numpy as np
+from mp_env_api import MPEnvWrapper
+
+
+class MPWrapper(MPEnvWrapper):
+
+    @property
+    def active_obs(self):
+        """
+            Returns boolean mask for each substate in the full observation.
+            It determines whether the observation is returned for the contextual case or not.
+            This effectively allows to filter unwanted or unnecessary observations from the full step-based case.
+            E.g. Velocities starting at 0 are only changing after the first action. Given we only receive the first  
+            observation, the velocities are not necessary in the observation for the MP task.
+        """
+        return np.ones(self.observation_space.shape, dtype=bool)
+
+    @property
+    def current_vel(self):
+        """
+            Returns the current velocity of the action/control dimension. 
+            The dimensionality has to match the action/control dimension.
+            This is not required when exclusively using position control, 
+            it should, however, be implemented regardless.
+            E.g. The joint velocities that are directly or indirectly controlled by the action.
+        """
+        raise NotImplementedError()
+
+    @property
+    def current_pos(self):
+        """
+            Returns the current position of the action/control dimension. 
+            The dimensionality has to match the action/control dimension.
+            This is not required when exclusively using velocity control, 
+            it should, however, be implemented regardless.
+            E.g. The joint positions that are directly or indirectly controlled by the action.
+        """
+        raise NotImplementedError()
+
+    @property
+    def goal_pos(self):
+        """
+            Returns a predefined final position of the action/control dimension.
+            This is only required for the DMP and is most of the time learned instead.
+        """
+        raise NotImplementedError()
+
+    @property
+    def dt(self):
+        """
+            Returns the time between two simulated steps of the environment
+        """
+        raise NotImplementedError()
 
 ```
 
-For an example using a DMP wrapped env and asynchronous sampling look at [mp_env_async_sampler.py](./alr_envs/utils/mp_env_async_sampler.py)
+If you created a new task wrapper, feel free to open a PR, so we can integrate it for others to use as well. 
+Without the integration the task can still be used. A rough outline can be shown here, for more details we recommend 
+having a look at the [examples](alr_envs/examples/).
+
+```python
+import alr_envs
+
+# Base environment name, according to structure of above example
+base_env_id = "ball_in_cup-catch"
+
+# Replace this wrapper with the custom wrapper for your environment by inheriting from the MPEnvWrapper.
+# You can also add other gym.Wrappers in case they are needed, 
+# e.g. gym.wrappers.FlattenObservation for dict observations
+wrappers = [alr_envs.dmc.suite.ball_in_cup.MPWrapper]
+mp_kwargs = {...}
+kwargs = {...}
+env = alr_envs.make_dmp_env(base_env_id, wrappers=wrappers, seed=1, mp_kwargs=mp_kwargs, **kwargs)
+# OR for a deterministic ProMP (other mp_kwargs are required):
+# env = alr_envs.make_detpmp_env(base_env, wrappers=wrappers, seed=seed, mp_kwargs=mp_args)
+
+rewards = 0
+obs = env.reset()
+
+# number of samples/full trajectories (multiple environment steps)
+for i in range(5):
+    ac = env.action_space.sample()
+    obs, reward, done, info = env.step(ac)
+    rewards += reward
+
+    if done:
+        print(base_env_id, rewards)
+        rewards = 0
+        obs = env.reset()
+```

alr_envs/__init__.py

@@ -1,901 +1,15 @@
-from gym.envs.registration import register
-from gym.wrappers import FlattenObservation
+from alr_envs import dmc, meta, open_ai
+from alr_envs.utils.make_env_helpers import make, make_detpmp_env, make_dmp_env, make_rank
+from alr_envs.utils import make_dmc
 
-from alr_envs import classic_control, dmc, open_ai
+# Convenience function for all MP environments
+from .alr import ALL_ALR_MOTION_PRIMITIVE_ENVIRONMENTS
+from .dmc import ALL_DEEPMIND_MOTION_PRIMITIVE_ENVIRONMENTS
+from .meta import ALL_METAWORLD_MOTION_PRIMITIVE_ENVIRONMENTS
+from .open_ai import ALL_GYM_MOTION_PRIMITIVE_ENVIRONMENTS
 
-from alr_envs.utils.make_env_helpers import make_dmp_env
-from alr_envs.utils.make_env_helpers import make_detpmp_env
-from alr_envs.utils.make_env_helpers import make_env
-from alr_envs.utils.make_env_helpers import make_env_rank
-
-# Mujoco
-
-## Reacher
-register(
-    id='ALRReacher-v0',
-    entry_point='alr_envs.mujoco:ALRReacherEnv',
-    max_episode_steps=200,
-    kwargs={
-        "steps_before_reward": 0,
-        "n_links": 5,
-        "balance": False,
-    }
-)
-
-register(
-    id='ALRReacherSparse-v0',
-    entry_point='alr_envs.mujoco:ALRReacherEnv',
-    max_episode_steps=200,
-    kwargs={
-        "steps_before_reward": 200,
-        "n_links": 5,
-        "balance": False,
-    }
-)
-
-register(
-    id='ALRReacherSparseBalanced-v0',
-    entry_point='alr_envs.mujoco:ALRReacherEnv',
-    max_episode_steps=200,
-    kwargs={
-        "steps_before_reward": 200,
-        "n_links": 5,
-        "balance": True,
-    }
-)
-
-register(
-    id='ALRLongReacher-v0',
-    entry_point='alr_envs.mujoco:ALRReacherEnv',
-    max_episode_steps=200,
-    kwargs={
-        "steps_before_reward": 0,
-        "n_links": 7,
-        "balance": False,
-    }
-)
-
-register(
-    id='ALRLongReacherSparse-v0',
-    entry_point='alr_envs.mujoco:ALRReacherEnv',
-    max_episode_steps=200,
-    kwargs={
-        "steps_before_reward": 200,
-        "n_links": 7,
-        "balance": False,
-    }
-)
-
-register(
-    id='ALRLongReacherSparseBalanced-v0',
-    entry_point='alr_envs.mujoco:ALRReacherEnv',
-    max_episode_steps=200,
-    kwargs={
-        "steps_before_reward": 200,
-        "n_links": 7,
-        "balance": True,
-    }
-)
-
-## Balancing Reacher
-
-register(
-    id='Balancing-v0',
-    entry_point='alr_envs.mujoco:BalancingEnv',
-    max_episode_steps=200,
-    kwargs={
-        "n_links": 5,
-    }
-)
-
-# Classic control
-
-## Simple Reacher
-register(
-    id='SimpleReacher-v0',
-    entry_point='alr_envs.classic_control:SimpleReacherEnv',
-    max_episode_steps=200,
-    kwargs={
-        "n_links": 2,
-    }
-)
-
-register(
-    id='SimpleReacher-v1',
-    entry_point='alr_envs.classic_control:SimpleReacherEnv',
-    max_episode_steps=200,
-    kwargs={
-        "n_links": 2,
-        "random_start": False
-    }
-)
-
-register(
-    id='LongSimpleReacher-v0',
-    entry_point='alr_envs.classic_control:SimpleReacherEnv',
-    max_episode_steps=200,
-    kwargs={
-        "n_links": 5,
-    }
-)
-
-register(
-    id='LongSimpleReacher-v1',
-    entry_point='alr_envs.classic_control:SimpleReacherEnv',
-    max_episode_steps=200,
-    kwargs={
-        "n_links": 5,
-        "random_start": False
-    }
-)
-
-## Viapoint Reacher
-
-register(
-    id='ViaPointReacher-v0',
-    entry_point='alr_envs.classic_control:ViaPointReacher',
-    max_episode_steps=200,
-    kwargs={
-        "n_links": 5,
-        "allow_self_collision": False,
-        "collision_penalty": 1000
-    }
-)
-
-## Hole Reacher
-register(
-    id='HoleReacher-v0',
-    entry_point='alr_envs.classic_control:HoleReacherEnv',
-    max_episode_steps=200,
-    kwargs={
-        "n_links": 5,
-        "random_start": True,
-        "allow_self_collision": False,
-        "allow_wall_collision": False,
-        "hole_width": None,
-        "hole_depth": 1,
-        "hole_x": None,
-        "collision_penalty": 100,
-    }
-)
-
-register(
-    id='HoleReacher-v1',
-    entry_point='alr_envs.classic_control:HoleReacherEnv',
-    max_episode_steps=200,
-    kwargs={
-        "n_links": 5,
-        "random_start": False,
-        "allow_self_collision": False,
-        "allow_wall_collision": False,
-        "hole_width": 0.25,
-        "hole_depth": 1,
-        "hole_x": None,
-        "collision_penalty": 100,
-    }
-)
-
-register(
-    id='HoleReacher-v2',
-    entry_point='alr_envs.classic_control:HoleReacherEnv',
-    max_episode_steps=200,
-    kwargs={
-        "n_links": 5,
-        "random_start": False,
-        "allow_self_collision": False,
-        "allow_wall_collision": False,
-        "hole_width": 0.25,
-        "hole_depth": 1,
-        "hole_x": 2,
-        "collision_penalty": 100,
-    }
-)
-
-# Short Continuous Mountain Car
-register(
-    id="MountainCarContinuous-v1",
-    entry_point="gym.envs.classic_control:Continuous_MountainCarEnv",
-    max_episode_steps=100,
-    reward_threshold=90.0,
-)
-
-# Motion Primitive Environments
-
-## Simple Reacher
-versions = ["SimpleReacher-v0", "SimpleReacher-v1", "LongSimpleReacher-v0", "LongSimpleReacher-v1"]
-for v in versions:
-    name = v.split("-")
-    register(
-        id=f'{name[0]}DMP-{name[1]}',
-        entry_point='alr_envs.utils.make_env_helpers:make_dmp_env_helper',
-        # max_episode_steps=1,
-        kwargs={
-            "name": f"alr_envs:{v}",
-            "wrappers": [classic_control.simple_reacher.MPWrapper],
-            "mp_kwargs": {
-                "num_dof": 2 if "long" not in v.lower() else 5,
-                "num_basis": 5,
-                "duration": 20,
-                "alpha_phase": 2,
-                "learn_goal": True,
-                "policy_type": "velocity",
-                "weights_scale": 50,
-            }
-        }
-    )
-
-register(
-    id='ViaPointReacherDMP-v0',
-    entry_point='alr_envs.utils.make_env_helpers:make_dmp_env_helper',
-    # max_episode_steps=1,
-    kwargs={
-        "name": "alr_envs:ViaPointReacher-v0",
-        "wrappers": [classic_control.viapoint_reacher.MPWrapper],
-        "mp_kwargs": {
-            "num_dof": 5,
-            "num_basis": 5,
-            "duration": 2,
-            "learn_goal": True,
-            "alpha_phase": 2,
-            "policy_type": "velocity",
-            "weights_scale": 50,
-        }
-    }
-)
-
-register(
-    id='ViaPointReacherDetPMP-v0',
-    entry_point='alr_envs.utils.make_env_helpers:make_detpmp_env_helper',
-    # max_episode_steps=1,
-    kwargs={
-        "name": "alr_envs:ViaPointReacher-v0",
-        "wrappers": [classic_control.viapoint_reacher.MPWrapper],
-        "mp_kwargs": {
-            "num_dof": 5,
-            "num_basis": 5,
-            "duration": 2,
-            "width": 0.025,
-            "policy_type": "velocity",
-            "weights_scale": 0.2,
-            "zero_start": True
-        }
-    }
-)
-
-## Hole Reacher
-versions = ["v0", "v1", "v2"]
-for v in versions:
-    register(
-        id=f'HoleReacherDMP-{v}',
-        entry_point='alr_envs.utils.make_env_helpers:make_dmp_env_helper',
-        # max_episode_steps=1,
-        kwargs={
-            "name": f"alr_envs:HoleReacher-{v}",
-            "wrappers": [classic_control.hole_reacher.MPWrapper],
-            "mp_kwargs": {
-                "num_dof": 5,
-                "num_basis": 5,
-                "duration": 2,
-                "learn_goal": True,
-                "alpha_phase": 2,
-                "bandwidth_factor": 2,
-                "policy_type": "velocity",
-                "weights_scale": 50,
-                "goal_scale": 0.1
-            }
-        }
-    )
-
-    register(
-        id=f'HoleReacherDetPMP-{v}',
-        entry_point='alr_envs.utils.make_env_helpers:make_detpmp_env_helper',
-        kwargs={
-            "name": f"alr_envs:HoleReacher-{v}",
-            "wrappers": [classic_control.hole_reacher.MPWrapper],
-            "mp_kwargs": {
-                "num_dof": 5,
-                "num_basis": 5,
-                "duration": 2,
-                "width": 0.025,
-                "policy_type": "velocity",
-                "weights_scale": 0.2,
-                "zero_start": True
-            }
-        }
-    )
-
-## Deep Mind Control Suite (DMC)
-### Suite
-
-register(
-    id=f'dmc_ball_in_cup-catch_dmp-v0',
-    entry_point='alr_envs.utils.make_env_helpers:make_dmp_env_helper',
-    # max_episode_steps=1,
-    kwargs={
-        "name": f"ball_in_cup-catch",
-        "time_limit": 1,
-        "episode_length": 50,
-        "wrappers": [dmc.suite.ball_in_cup.MPWrapper],
-        "mp_kwargs": {
-            "num_dof": 2,
-            "num_basis": 5,
-            "duration": 1,
-            "learn_goal": True,
-            "alpha_phase": 2,
-            "bandwidth_factor": 2,
-            "policy_type": "motor",
-            "weights_scale": 50,
-            "goal_scale": 0.1,
-            "policy_kwargs": {
-                "p_gains": 50,
-                "d_gains": 1
-            }
-        }
-    }
-)
-
-register(
-    id=f'dmc_ball_in_cup-catch_detpmp-v0',
-    entry_point='alr_envs.utils.make_env_helpers:make_detpmp_env_helper',
-    kwargs={
-        "name": f"ball_in_cup-catch",
-        "time_limit": 1,
-        "episode_length": 50,
-        "wrappers": [dmc.suite.ball_in_cup.MPWrapper],
-        "mp_kwargs": {
-            "num_dof": 2,
-            "num_basis": 5,
-            "duration": 1,
-            "width": 0.025,
-            "policy_type": "motor",
-            "weights_scale": 0.2,
-            "zero_start": True,
-            "policy_kwargs": {
-                "p_gains": 50,
-                "d_gains": 1
-            }
-        }
-    }
-)
-
-# TODO tune episode length for all below
-register(
-    id=f'dmc_reacher-easy_dmp-v0',
-    entry_point='alr_envs.utils.make_env_helpers:make_dmp_env_helper',
-    # max_episode_steps=1,
-    kwargs={
-        "name": f"reacher-easy",
-        "time_limit": 1,
-        "episode_length": 50,
-        "wrappers": [dmc.suite.reacher.MPWrapper],
-        "mp_kwargs": {
-            "num_dof": 2,
-            "num_basis": 5,
-            "duration": 1,
-            "learn_goal": True,
-            "alpha_phase": 2,
-            "bandwidth_factor": 2,
-            "policy_type": "motor",
-            "weights_scale": 50,
-            "goal_scale": 0.1,
-            "policy_kwargs": {
-                "p_gains": 50,
-                "d_gains": 1
-            }
-        }
-    }
-)
-
-register(
-    id=f'dmc_reacher-easy_detpmp-v0',
-    entry_point='alr_envs.utils.make_env_helpers:make_detpmp_env_helper',
-    kwargs={
-        "name": f"reacher-easy",
-        "time_limit": 1,
-        "episode_length": 50,
-        "wrappers": [dmc.suite.reacher.MPWrapper],
-        "mp_kwargs": {
-            "num_dof": 2,
-            "num_basis": 5,
-            "duration": 1,
-            "width": 0.025,
-            "policy_type": "motor",
-            "weights_scale": 0.2,
-            "zero_start": True,
-            "policy_kwargs": {
-                "p_gains": 50,
-                "d_gains": 1
-            }
-        }
-    }
-)
-
-register(
-    id=f'dmc_reacher-hard_dmp-v0',
-    entry_point='alr_envs.utils.make_env_helpers:make_dmp_env_helper',
-    # max_episode_steps=1,
-    kwargs={
-        "name": f"reacher-hard",
-        "time_limit": 1,
-        "episode_length": 50,
-        "wrappers": [dmc.suite.reacher.MPWrapper],
-        "mp_kwargs": {
-            "num_dof": 2,
-            "num_basis": 5,
-            "duration": 1,
-            "learn_goal": True,
-            "alpha_phase": 2,
-            "bandwidth_factor": 2,
-            "policy_type": "motor",
-            "weights_scale": 50,
-            "goal_scale": 0.1,
-            "policy_kwargs": {
-                "p_gains": 50,
-                "d_gains": 1
-            }
-        }
-    }
-)
-
-register(
-    id=f'dmc_reacher-hard_detpmp-v0',
-    entry_point='alr_envs.utils.make_env_helpers:make_detpmp_env_helper',
-    kwargs={
-        "name": f"reacher-hard",
-        "time_limit": 1,
-        "episode_length": 50,
-        "wrappers": [dmc.suite.reacher.MPWrapper],
-        "mp_kwargs": {
-            "num_dof": 2,
-            "num_basis": 5,
-            "duration": 1,
-            "width": 0.025,
-            "policy_type": "motor",
-            "weights_scale": 0.2,
-            "zero_start": True,
-            "policy_kwargs": {
-                "p_gains": 50,
-                "d_gains": 1
-            }
-        }
-    }
-)
-register(
-    id=f'dmc_cartpole-balance_dmp-v0',
-    entry_point='alr_envs.utils.make_env_helpers:make_dmp_env_helper',
-    # max_episode_steps=1,
-    kwargs={
-        "name": f"cartpole-balance",
-        # "time_limit": 1,
-        "camera_id": 0,
-        "episode_length": 1000,
-        "wrappers": [dmc.suite.cartpole.MPWrapper],
-        "mp_kwargs": {
-            "num_dof": 1,
-            "num_basis": 5,
-            "duration": 10,
-            "learn_goal": True,
-            "alpha_phase": 2,
-            "bandwidth_factor": 2,
-            "policy_type": "motor",
-            "weights_scale": 50,
-            "goal_scale": 0.1,
-            "policy_kwargs": {
-                "p_gains": 10,
-                "d_gains": 10
-            }
-        }
-    }
-)
-
-register(
-    id=f'dmc_cartpole-balance_detpmp-v0',
-    entry_point='alr_envs.utils.make_env_helpers:make_detpmp_env_helper',
-    kwargs={
-        "name": f"cartpole-balance",
-        # "time_limit": 1,
-        "camera_id": 0,
-        "episode_length": 1000,
-        "wrappers": [dmc.suite.cartpole.MPWrapper],
-        "mp_kwargs": {
-            "num_dof": 1,
-            "num_basis": 5,
-            "duration": 10,
-            "width": 0.025,
-            "policy_type": "motor",
-            "weights_scale": 0.2,
-            "zero_start": True,
-            "policy_kwargs": {
-                "p_gains": 10,
-                "d_gains": 10
-            }
-        }
-    }
-)
-register(
-    id=f'dmc_cartpole-balance_sparse_dmp-v0',
-    entry_point='alr_envs.utils.make_env_helpers:make_dmp_env_helper',
-    # max_episode_steps=1,
-    kwargs={
-        "name": f"cartpole-balance_sparse",
-        # "time_limit": 1,
-        "camera_id": 0,
-        "episode_length": 1000,
-        "wrappers": [dmc.suite.cartpole.MPWrapper],
-        "mp_kwargs": {
-            "num_dof": 1,
-            "num_basis": 5,
-            "duration": 10,
-            "learn_goal": True,
-            "alpha_phase": 2,
-            "bandwidth_factor": 2,
-            "policy_type": "motor",
-            "weights_scale": 50,
-            "goal_scale": 0.1,
-            "policy_kwargs": {
-                "p_gains": 10,
-                "d_gains": 10
-            }
-        }
-    }
-)
-
-register(
-    id=f'dmc_cartpole-balance_sparse_detpmp-v0',
-    entry_point='alr_envs.utils.make_env_helpers:make_detpmp_env_helper',
-    kwargs={
-        "name": f"cartpole-balance_sparse",
-        # "time_limit": 1,
-        "camera_id": 0,
-        "episode_length": 1000,
-        "wrappers": [dmc.suite.cartpole.MPWrapper],
-        "mp_kwargs": {
-            "num_dof": 1,
-            "num_basis": 5,
-            "duration": 10,
-            "width": 0.025,
-            "policy_type": "motor",
-            "weights_scale": 0.2,
-            "zero_start": True,
-            "policy_kwargs": {
-                "p_gains": 10,
-                "d_gains": 10
-            }
-        }
-    }
-)
-
-register(
-    id=f'dmc_cartpole-swingup_dmp-v0',
-    entry_point='alr_envs.utils.make_env_helpers:make_dmp_env_helper',
-    # max_episode_steps=1,
-    kwargs={
-        "name": f"cartpole-swingup",
-        # "time_limit": 1,
-        "camera_id": 0,
-        "episode_length": 1000,
-        "wrappers": [dmc.suite.cartpole.MPWrapper],
-        "mp_kwargs": {
-            "num_dof": 1,
-            "num_basis": 5,
-            "duration": 10,
-            "learn_goal": True,
-            "alpha_phase": 2,
-            "bandwidth_factor": 2,
-            "policy_type": "motor",
-            "weights_scale": 50,
-            "goal_scale": 0.1,
-            "policy_kwargs": {
-                "p_gains": 10,
-                "d_gains": 10
-            }
-        }
-    }
-)
-
-register(
-    id=f'dmc_cartpole-swingup_detpmp-v0',
-    entry_point='alr_envs.utils.make_env_helpers:make_detpmp_env_helper',
-    kwargs={
-        "name": f"cartpole-swingup",
-        # "time_limit": 1,
-        "camera_id": 0,
-        "episode_length": 1000,
-        "wrappers": [dmc.suite.cartpole.MPWrapper],
-        "mp_kwargs": {
-            "num_dof": 1,
-            "num_basis": 5,
-            "duration": 10,
-            "width": 0.025,
-            "policy_type": "motor",
-            "weights_scale": 0.2,
-            "zero_start": True,
-            "policy_kwargs": {
-                "p_gains": 10,
-                "d_gains": 10
-            }
-        }
-    }
-)
-register(
-    id=f'dmc_cartpole-swingup_sparse_dmp-v0',
-    entry_point='alr_envs.utils.make_env_helpers:make_dmp_env_helper',
-    # max_episode_steps=1,
-    kwargs={
-        "name": f"cartpole-swingup_sparse",
-        # "time_limit": 1,
-        "camera_id": 0,
-        "episode_length": 1000,
-        "wrappers": [dmc.suite.cartpole.MPWrapper],
-        "mp_kwargs": {
-            "num_dof": 1,
-            "num_basis": 5,
-            "duration": 10,
-            "learn_goal": True,
-            "alpha_phase": 2,
-            "bandwidth_factor": 2,
-            "policy_type": "motor",
-            "weights_scale": 50,
-            "goal_scale": 0.1,
-            "policy_kwargs": {
-                "p_gains": 10,
-                "d_gains": 10
-            }
-        }
-    }
-)
-
-register(
-    id=f'dmc_cartpole-swingup_sparse_detpmp-v0',
-    entry_point='alr_envs.utils.make_env_helpers:make_detpmp_env_helper',
-    kwargs={
-        "name": f"cartpole-swingup_sparse",
-        # "time_limit": 1,
-        "camera_id": 0,
-        "episode_length": 1000,
-        "wrappers": [dmc.suite.cartpole.MPWrapper],
-        "mp_kwargs": {
-            "num_dof": 1,
-            "num_basis": 5,
-            "duration": 10,
-            "width": 0.025,
-            "policy_type": "motor",
-            "weights_scale": 0.2,
-            "zero_start": True,
-            "policy_kwargs": {
-                "p_gains": 10,
-                "d_gains": 10
-            }
-        }
-    }
-)
-register(
-    id=f'dmc_cartpole-two_poles_dmp-v0',
-    entry_point='alr_envs.utils.make_env_helpers:make_dmp_env_helper',
-    # max_episode_steps=1,
-    kwargs={
-        "name": f"cartpole-two_poles",
-        # "time_limit": 1,
-        "camera_id": 0,
-        "episode_length": 1000,
-        # "wrappers": [partial(DMCCartpoleMPWrapper, n_poles=2)],
-        "wrappers": [dmc.suite.cartpole.TwoPolesMPWrapper],
-        "mp_kwargs": {
-            "num_dof": 1,
-            "num_basis": 5,
-            "duration": 10,
-            "learn_goal": True,
-            "alpha_phase": 2,
-            "bandwidth_factor": 2,
-            "policy_type": "motor",
-            "weights_scale": 50,
-            "goal_scale": 0.1,
-            "policy_kwargs": {
-                "p_gains": 10,
-                "d_gains": 10
-            }
-        }
-    }
-)
-
-register(
-    id=f'dmc_cartpole-two_poles_detpmp-v0',
-    entry_point='alr_envs.utils.make_env_helpers:make_detpmp_env_helper',
-    kwargs={
-        "name": f"cartpole-two_poles",
-        # "time_limit": 1,
-        "camera_id": 0,
-        "episode_length": 1000,
-        # "wrappers": [partial(DMCCartpoleMPWrapper, n_poles=2)],
-        "wrappers": [dmc.suite.cartpole.TwoPolesMPWrapper],
-        "mp_kwargs": {
-            "num_dof": 1,
-            "num_basis": 5,
-            "duration": 10,
-            "width": 0.025,
-            "policy_type": "motor",
-            "weights_scale": 0.2,
-            "zero_start": True,
-            "policy_kwargs": {
-                "p_gains": 10,
-                "d_gains": 10
-            }
-        }
-    }
-)
-register(
-    id=f'dmc_cartpole-three_poles_dmp-v0',
-    entry_point='alr_envs.utils.make_env_helpers:make_dmp_env_helper',
-    # max_episode_steps=1,
-    kwargs={
-        "name": f"cartpole-three_poles",
-        # "time_limit": 1,
-        "camera_id": 0,
-        "episode_length": 1000,
-        # "wrappers": [partial(DMCCartpoleMPWrapper, n_poles=3)],
-        "wrappers": [dmc.suite.cartpole.ThreePolesMPWrapper],
-        "mp_kwargs": {
-            "num_dof": 1,
-            "num_basis": 5,
-            "duration": 10,
-            "learn_goal": True,
-            "alpha_phase": 2,
-            "bandwidth_factor": 2,
-            "policy_type": "motor",
-            "weights_scale": 50,
-            "goal_scale": 0.1,
-            "policy_kwargs": {
-                "p_gains": 10,
-                "d_gains": 10
-            }
-        }
-    }
-)
-
-register(
-    id=f'dmc_cartpole-three_poles_detpmp-v0',
-    entry_point='alr_envs.utils.make_env_helpers:make_detpmp_env_helper',
-    kwargs={
-        "name": f"cartpole-three_poles",
-        # "time_limit": 1,
-        "camera_id": 0,
-        "episode_length": 1000,
-        # "wrappers": [partial(DMCCartpoleMPWrapper, n_poles=3)],
-        "wrappers": [dmc.suite.cartpole.ThreePolesMPWrapper],
-        "mp_kwargs": {
-            "num_dof": 1,
-            "num_basis": 5,
-            "duration": 10,
-            "width": 0.025,
-            "policy_type": "motor",
-            "weights_scale": 0.2,
-            "zero_start": True,
-            "policy_kwargs": {
-                "p_gains": 10,
-                "d_gains": 10
-            }
-        }
-    }
-)
-
-### Manipulation
-
-register(
-    id=f'dmc_manipulation-reach_site_dmp-v0',
-    entry_point='alr_envs.utils.make_env_helpers:make_dmp_env_helper',
-    # max_episode_steps=1,
-    kwargs={
-        "name": f"manipulation-reach_site_features",
-        # "time_limit": 1,
-        "episode_length": 250,
-        "wrappers": [dmc.manipulation.reach.MPWrapper],
-        "mp_kwargs": {
-            "num_dof": 9,
-            "num_basis": 5,
-            "duration": 10,
-            "learn_goal": True,
-            "alpha_phase": 2,
-            "bandwidth_factor": 2,
-            "policy_type": "velocity",
-            "weights_scale": 50,
-            "goal_scale": 0.1,
-        }
-    }
-)
-
-register(
-    id=f'dmc_manipulation-reach_site_detpmp-v0',
-    entry_point='alr_envs.utils.make_env_helpers:make_detpmp_env_helper',
-    kwargs={
-        "name": f"manipulation-reach_site_features",
-        # "time_limit": 1,
-        "episode_length": 250,
-        "wrappers": [dmc.manipulation.reach.MPWrapper],
-        "mp_kwargs": {
-            "num_dof": 9,
-            "num_basis": 5,
-            "duration": 10,
-            "width": 0.025,
-            "policy_type": "velocity",
-            "weights_scale": 0.2,
-            "zero_start": True,
-        }
-    }
-)
-
-## Open AI
-register(
-    id='ContinuousMountainCarDetPMP-v1',
-    entry_point='alr_envs.utils.make_env_helpers:make_detpmp_env_helper',
-    kwargs={
-        "name": "alr_envs:MountainCarContinuous-v1",
-        "wrappers": [open_ai.classic_control.continuous_mountain_car.MPWrapper],
-        "mp_kwargs": {
-            "num_dof": 1,
-            "num_basis": 4,
-            "duration": 2,
-            "post_traj_time": 0,
-            "width": 0.02,
-            "policy_type": "motor",
-            "policy_kwargs": {
-                "p_gains": 1.,
-                "d_gains": 1.
-            }
-        }
-    }
-)
-
-register(
-    id='ReacherDetPMP-v2',
-    entry_point='alr_envs.utils.make_env_helpers:make_detpmp_env_helper',
-    kwargs={
-        "name": "gym.envs.mujoco:Reacher-v2",
-        "wrappers": [open_ai.mujoco.reacher_v2.MPWrapper],
-        "mp_kwargs": {
-            "num_dof": 2,
-            "num_basis": 6,
-            "duration": 1,
-            "post_traj_time": 0,
-            "width": 0.02,
-            "policy_type": "motor",
-            "policy_kwargs": {
-                "p_gains": .6,
-                "d_gains": .075
-            }
-        }
-    }
-)
-
-register(
-    id='FetchSlideDenseDetPMP-v1',
-    entry_point='alr_envs.utils.make_env_helpers:make_detpmp_env_helper',
-    kwargs={
-        "name": "gym.envs.robotics:FetchSlideDense-v1",
-        "wrappers": [FlattenObservation, open_ai.robotics.fetch.MPWrapper],
-        "mp_kwargs": {
-            "num_dof": 4,
-            "num_basis": 5,
-            "duration": 2,
-            "post_traj_time": 0,
-            "width": 0.02,
-            "policy_type": "position"
-        }
-    }
-)
-
-register(
-    id='FetchReachDenseDetPMP-v1',
-    entry_point='alr_envs.utils.make_env_helpers:make_detpmp_env_helper',
-    kwargs={
-        "name": "gym.envs.robotics:FetchReachDense-v1",
-        "wrappers": [FlattenObservation, open_ai.robotics.fetch.MPWrapper],
-        "mp_kwargs": {
-            "num_dof": 4,
-            "num_basis": 5,
-            "duration": 2,
-            "post_traj_time": 0,
-            "width": 0.02,
-            "policy_type": "position"
-        }
-    }
-)
+ALL_MOTION_PRIMITIVE_ENVIRONMENTS = {
+    key: value + ALL_DEEPMIND_MOTION_PRIMITIVE_ENVIRONMENTS[key] +
+         ALL_GYM_MOTION_PRIMITIVE_ENVIRONMENTS[key] +
+         ALL_METAWORLD_MOTION_PRIMITIVE_ENVIRONMENTS[key]
+    for key, value in ALL_ALR_MOTION_PRIMITIVE_ENVIRONMENTS.items()}

alr_envs/alr/__init__.py

@@ -0,0 +1,329 @@
+from gym import register
+
+from . import classic_control, mujoco
+from .classic_control.hole_reacher.hole_reacher import HoleReacherEnv
+from .classic_control.simple_reacher.simple_reacher import SimpleReacherEnv
+from .classic_control.viapoint_reacher.viapoint_reacher import ViaPointReacherEnv
+from .mujoco.ball_in_a_cup.ball_in_a_cup import ALRBallInACupEnv
+from .mujoco.ball_in_a_cup.biac_pd import ALRBallInACupPDEnv
+from .mujoco.reacher.alr_reacher import ALRReacherEnv
+from .mujoco.reacher.balancing import BalancingEnv
+
+ALL_ALR_MOTION_PRIMITIVE_ENVIRONMENTS = {"DMP": [], "DetPMP": []}
+
+# Classic Control
+## Simple Reacher
+register(
+    id='SimpleReacher-v0',
+    entry_point='alr_envs.alr.classic_control:SimpleReacherEnv',
+    max_episode_steps=200,
+    kwargs={
+        "n_links": 2,
+    }
+)
+
+register(
+    id='SimpleReacher-v1',
+    entry_point='alr_envs.alr.classic_control:SimpleReacherEnv',
+    max_episode_steps=200,
+    kwargs={
+        "n_links": 2,
+        "random_start": False
+    }
+)
+
+register(
+    id='LongSimpleReacher-v0',
+    entry_point='alr_envs.alr.classic_control:SimpleReacherEnv',
+    max_episode_steps=200,
+    kwargs={
+        "n_links": 5,
+    }
+)
+
+register(
+    id='LongSimpleReacher-v1',
+    entry_point='alr_envs.alr.classic_control:SimpleReacherEnv',
+    max_episode_steps=200,
+    kwargs={
+        "n_links": 5,
+        "random_start": False
+    }
+)
+
+## Viapoint Reacher
+
+register(
+    id='ViaPointReacher-v0',
+    entry_point='alr_envs.alr.classic_control:ViaPointReacherEnv',
+    max_episode_steps=200,
+    kwargs={
+        "n_links": 5,
+        "allow_self_collision": False,
+        "collision_penalty": 1000
+    }
+)
+
+## Hole Reacher
+register(
+    id='HoleReacher-v0',
+    entry_point='alr_envs.alr.classic_control:HoleReacherEnv',
+    max_episode_steps=200,
+    kwargs={
+        "n_links": 5,
+        "random_start": True,
+        "allow_self_collision": False,
+        "allow_wall_collision": False,
+        "hole_width": None,
+        "hole_depth": 1,
+        "hole_x": None,
+        "collision_penalty": 100,
+    }
+)
+
+register(
+    id='HoleReacher-v1',
+    entry_point='alr_envs.alr.classic_control:HoleReacherEnv',
+    max_episode_steps=200,
+    kwargs={
+        "n_links": 5,
+        "random_start": False,
+        "allow_self_collision": False,
+        "allow_wall_collision": False,
+        "hole_width": 0.25,
+        "hole_depth": 1,
+        "hole_x": None,
+        "collision_penalty": 100,
+    }
+)
+
+register(
+    id='HoleReacher-v2',
+    entry_point='alr_envs.alr.classic_control:HoleReacherEnv',
+    max_episode_steps=200,
+    kwargs={
+        "n_links": 5,
+        "random_start": False,
+        "allow_self_collision": False,
+        "allow_wall_collision": False,
+        "hole_width": 0.25,
+        "hole_depth": 1,
+        "hole_x": 2,
+        "collision_penalty": 100,
+    }
+)
+
+# Mujoco
+
+## Reacher
+register(
+    id='ALRReacher-v0',
+    entry_point='alr_envs.alr.mujoco:ALRReacherEnv',
+    max_episode_steps=200,
+    kwargs={
+        "steps_before_reward": 0,
+        "n_links": 5,
+        "balance": False,
+    }
+)
+
+register(
+    id='ALRReacherSparse-v0',
+    entry_point='alr_envs.alr.mujoco:ALRReacherEnv',
+    max_episode_steps=200,
+    kwargs={
+        "steps_before_reward": 200,
+        "n_links": 5,
+        "balance": False,
+    }
+)
+
+register(
+    id='ALRReacherSparseBalanced-v0',
+    entry_point='alr_envs.alr.mujoco:ALRReacherEnv',
+    max_episode_steps=200,
+    kwargs={
+        "steps_before_reward": 200,
+        "n_links": 5,
+        "balance": True,
+    }
+)
+
+register(
+    id='ALRLongReacher-v0',
+    entry_point='alr_envs.alr.mujoco:ALRReacherEnv',
+    max_episode_steps=200,
+    kwargs={
+        "steps_before_reward": 0,
+        "n_links": 7,
+        "balance": False,
+    }
+)
+
+register(
+    id='ALRLongReacherSparse-v0',
+    entry_point='alr_envs.alr.mujoco:ALRReacherEnv',
+    max_episode_steps=200,
+    kwargs={
+        "steps_before_reward": 200,
+        "n_links": 7,
+        "balance": False,
+    }
+)
+
+register(
+    id='ALRLongReacherSparseBalanced-v0',
+    entry_point='alr_envs.alr.mujoco:ALRReacherEnv',
+    max_episode_steps=200,
+    kwargs={
+        "steps_before_reward": 200,
+        "n_links": 7,
+        "balance": True,
+    }
+)
+
+## Balancing Reacher
+
+register(
+    id='Balancing-v0',
+    entry_point='alr_envs.alr.mujoco:BalancingEnv',
+    max_episode_steps=200,
+    kwargs={
+        "n_links": 5,
+    }
+)
+
+# Motion Primitive Environments
+
+## Simple Reacher
+_versions = ["SimpleReacher-v0", "SimpleReacher-v1", "LongSimpleReacher-v0", "LongSimpleReacher-v1"]
+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": [classic_control.simple_reacher.MPWrapper],
+            "mp_kwargs": {
+                "num_dof": 2 if "long" not in _v.lower() else 5,
+                "num_basis": 5,
+                "duration": 20,
+                "alpha_phase": 2,
+                "learn_goal": True,
+                "policy_type": "velocity",
+                "weights_scale": 50,
+            }
+        }
+    )
+    ALL_ALR_MOTION_PRIMITIVE_ENVIRONMENTS["DMP"].append(_env_id)
+
+    _env_id = f'{_name[0]}DetPMP-{_name[1]}'
+    register(
+        id=_env_id,
+        entry_point='alr_envs.utils.make_env_helpers:make_detpmp_env_helper',
+        # max_episode_steps=1,
+        kwargs={
+            "name": f"alr_envs:{_v}",
+            "wrappers": [classic_control.simple_reacher.MPWrapper],
+            "mp_kwargs": {
+                "num_dof": 2 if "long" not in _v.lower() else 5,
+                "num_basis": 5,
+                "duration": 20,
+                "width": 0.025,
+                "policy_type": "velocity",
+                "weights_scale": 0.2,
+                "zero_start": True
+            }
+        }
+    )
+    ALL_ALR_MOTION_PRIMITIVE_ENVIRONMENTS["DetPMP"].append(_env_id)
+
+# Viapoint reacher
+register(
+    id='ViaPointReacherDMP-v0',
+    entry_point='alr_envs.utils.make_env_helpers:make_dmp_env_helper',
+    # max_episode_steps=1,
+    kwargs={
+        "name": "alr_envs:ViaPointReacher-v0",
+        "wrappers": [classic_control.viapoint_reacher.MPWrapper],
+        "mp_kwargs": {
+            "num_dof": 5,
+            "num_basis": 5,
+            "duration": 2,
+            "learn_goal": True,
+            "alpha_phase": 2,
+            "policy_type": "velocity",
+            "weights_scale": 50,
+        }
+    }
+)
+ALL_ALR_MOTION_PRIMITIVE_ENVIRONMENTS["DMP"].append("ViaPointReacherDMP-v0")
+
+register(
+    id='ViaPointReacherDetPMP-v0',
+    entry_point='alr_envs.utils.make_env_helpers:make_detpmp_env_helper',
+    # max_episode_steps=1,
+    kwargs={
+        "name": "alr_envs:ViaPointReacher-v0",
+        "wrappers": [classic_control.viapoint_reacher.MPWrapper],
+        "mp_kwargs": {
+            "num_dof": 5,
+            "num_basis": 5,
+            "duration": 2,
+            "width": 0.025,
+            "policy_type": "velocity",
+            "weights_scale": 0.2,
+            "zero_start": True
+        }
+    }
+)
+ALL_ALR_MOTION_PRIMITIVE_ENVIRONMENTS["DetPMP"].append("ViaPointReacherDetPMP-v0")
+
+## Hole Reacher
+_versions = ["v0", "v1", "v2"]
+for _v in _versions:
+    _env_id = f'HoleReacherDMP-{_v}'
+    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:HoleReacher-{_v}",
+            "wrappers": [classic_control.hole_reacher.MPWrapper],
+            "mp_kwargs": {
+                "num_dof": 5,
+                "num_basis": 5,
+                "duration": 2,
+                "learn_goal": True,
+                "alpha_phase": 2,
+                "bandwidth_factor": 2,
+                "policy_type": "velocity",
+                "weights_scale": 50,
+                "goal_scale": 0.1
+            }
+        }
+    )
+    ALL_ALR_MOTION_PRIMITIVE_ENVIRONMENTS["DMP"].append(_env_id)
+
+    _env_id = f'HoleReacherDetPMP-{_v}'
+    register(
+        id=_env_id,
+        entry_point='alr_envs.utils.make_env_helpers:make_detpmp_env_helper',
+        kwargs={
+            "name": f"alr_envs:HoleReacher-{_v}",
+            "wrappers": [classic_control.hole_reacher.MPWrapper],
+            "mp_kwargs": {
+                "num_dof": 5,
+                "num_basis": 5,
+                "duration": 2,
+                "width": 0.025,
+                "policy_type": "velocity",
+                "weights_scale": 0.2,
+                "zero_start": True
+            }
+        }
+    )
+    ALL_ALR_MOTION_PRIMITIVE_ENVIRONMENTS["DetPMP"].append(_env_id)

alr_envs/alr/classic_control/README.MD

@@ -0,0 +1,21 @@
+### Classic Control
+
+## Step-based Environments
+|Name| Description|Horizon|Action Dimension|Observation Dimension
+|---|---|---|---|---|
+|`SimpleReacher-v0`| Simple reaching task (2 links) without any physics simulation. Provides no reward until 150 time steps. This allows the agent to explore the space, but requires precise actions towards the end of the trajectory.| 200 | 2 | 9
+|`LongSimpleReacher-v0`| Simple reaching task (5 links) without any physics simulation. Provides no reward until 150 time steps. This allows the agent to explore the space, but requires precise actions towards the end of the trajectory.| 200 | 5 | 18
+|`ViaPointReacher-v0`| Simple reaching task leveraging a via point, which supports self collision detection. Provides a reward only at 100 and 199 for reaching the viapoint and goal point, respectively.| 200 | 5 | 18 
+|`HoleReacher-v0`| 5 link reaching task where the end-effector needs to reach into a narrow hole without collding with itself or walls | 200 | 5 | 18
+
+## MP Environments
+|Name| Description|Horizon|Action Dimension|Context Dimension
+|---|---|---|---|---|
+|`ViaPointReacherDMP-v0`| A DMP provides a trajectory for the `ViaPointReacher-v0` task. | 200 | 25
+|`HoleReacherFixedGoalDMP-v0`| A DMP provides a trajectory for the `HoleReacher-v0` task with a fixed goal attractor. | 200 | 25
+|`HoleReacherDMP-v0`| A DMP provides a trajectory for the `HoleReacher-v0` task. The goal attractor needs to be learned. | 200 | 30 
+|`ALRBallInACupSimpleDMP-v0`| A DMP provides a trajectory for the `ALRBallInACupSimple-v0` task where only 3 joints are actuated. | 4000 | 15
+|`ALRBallInACupDMP-v0`| A DMP provides a trajectory for the `ALRBallInACup-v0` task. | 4000 | 35
+|`ALRBallInACupGoalDMP-v0`| A DMP provides a trajectory for the `ALRBallInACupGoal-v0` task. | 4000 | 35 | 3
+
+[//]:  |`HoleReacherDetPMP-v0`|

alr_envs/alr/classic_control/__init__.py

@@ -0,0 +1,3 @@
+from .hole_reacher.hole_reacher import HoleReacherEnv
+from .simple_reacher.simple_reacher import SimpleReacherEnv
+from .viapoint_reacher.viapoint_reacher import ViaPointReacherEnv

alr_envs/alr/classic_control/hole_reacher/hole_reacher.py

@@ -6,7 +6,7 @@ import numpy as np
 from gym.utils import seeding
 from matplotlib import patches
 
-from alr_envs.classic_control.utils import check_self_collision
+from alr_envs.alr.classic_control.utils import check_self_collision
 
 
 class HoleReacherEnv(gym.Env):

alr_envs/alr/classic_control/viapoint_reacher/viapoint_reacher.py

@@ -5,10 +5,10 @@ import matplotlib.pyplot as plt
 import numpy as np
 from gym.utils import seeding
 
-from alr_envs.classic_control.utils import check_self_collision
+from alr_envs.alr.classic_control.utils import check_self_collision
 
 
-class ViaPointReacher(gym.Env):
+class ViaPointReacherEnv(gym.Env):
 
     def __init__(self, n_links, random_start: bool = False, via_target: Union[None, Iterable] = None,
                  target: Union[None, Iterable] = None, allow_self_collision=False, collision_penalty=1000):

alr_envs/alr/mujoco/README.MD

@@ -0,0 +1,15 @@
+# Custom Mujoco tasks
+
+## Step-based Environments
+|Name| Description|Horizon|Action Dimension|Observation Dimension
+|---|---|---|---|---|
+|`ALRReacher-v0`|Modified (5 links) Mujoco gym's `Reacher-v2` (2 links)| 200 | 5 | 21
+|`ALRReacherSparse-v0`|Same as `ALRReacher-v0`, but the distance penalty is only provided in the last time step.| 200 | 5 | 21
+|`ALRReacherSparseBalanced-v0`|Same as `ALRReacherSparse-v0`, but the end-effector has to remain upright.| 200 | 5 | 21
+|`ALRLongReacher-v0`|Modified (7 links) Mujoco gym's `Reacher-v2` (2 links)| 200 | 7 | 27
+|`ALRLongReacherSparse-v0`|Same as `ALRLongReacher-v0`, but the distance penalty is only provided in the last time step.| 200 | 7 | 27
+|`ALRLongReacherSparseBalanced-v0`|Same as `ALRLongReacherSparse-v0`, but the end-effector has to remain upright.| 200 | 7 | 27
+|`ALRBallInACupSimple-v0`| Ball-in-a-cup task where a robot needs to catch a ball attached to a cup at its end-effector. | 4000 | 3 | wip
+|`ALRBallInACup-v0`| Ball-in-a-cup task where a robot needs to catch a ball attached to a cup at its end-effector | 4000 | 7 | wip
+|`ALRBallInACupGoal-v0`| Similar to `ALRBallInACupSimple-v0` but the ball needs to be caught at a specified goal position | 4000 | 7 | wip
+    

alr_envs/alr/mujoco/__init__.py

@@ -0,0 +1,4 @@
+from .reacher.alr_reacher import ALRReacherEnv
+from .reacher.balancing import BalancingEnv
+from .ball_in_a_cup.ball_in_a_cup import ALRBallInACupEnv
+from .ball_in_a_cup.biac_pd import ALRBallInACupPDEnv

alr_envs/alr/mujoco/ball_in_a_cup/ball_in_a_cup.py

@@ -35,10 +35,10 @@ class ALRBallInACupEnv(MujocoEnv, utils.EzPickle):
         self.sim_time = 8  # seconds
         self.sim_steps = int(self.sim_time / self.dt)
         if reward_type == "no_context":
-            from alr_envs.mujoco.ball_in_a_cup.ball_in_a_cup_reward_simple import BallInACupReward
+            from alr_envs.alr.mujoco.ball_in_a_cup.ball_in_a_cup_reward_simple import BallInACupReward
             reward_function = BallInACupReward
         elif reward_type == "contextual_goal":
-            from alr_envs.mujoco.ball_in_a_cup.ball_in_a_cup_reward import BallInACupReward
+            from alr_envs.alr.mujoco.ball_in_a_cup.ball_in_a_cup_reward import BallInACupReward
             reward_function = BallInACupReward
         else:
             raise ValueError("Unknown reward type: {}".format(reward_type))

alr_envs/alr/mujoco/ball_in_a_cup/ball_in_a_cup_reward.py

@@ -1,5 +1,5 @@
 import numpy as np
-from alr_envs.mujoco import alr_reward_fct
+from alr_envs.alr.mujoco import alr_reward_fct
 
 
 class BallInACupReward(alr_reward_fct.AlrReward):

alr_envs/alr/mujoco/ball_in_a_cup/ball_in_a_cup_reward_simple.py

@@ -1,5 +1,5 @@
 import numpy as np
-from alr_envs.mujoco import alr_reward_fct
+from alr_envs.alr.mujoco import alr_reward_fct
 
 
 class BallInACupReward(alr_reward_fct.AlrReward):

alr_envs/alr/mujoco/ball_in_a_cup/biac_pd.py

@@ -42,10 +42,10 @@ class ALRBallInACupPDEnv(mujoco_env.MujocoEnv, utils.EzPickle):
         self._dt = 0.02
         self.ep_length = 4000  # based on 8 seconds with dt = 0.02 int(self.sim_time / self.dt)
         if reward_type == "no_context":
-            from alr_envs.mujoco.ball_in_a_cup.ball_in_a_cup_reward_simple import BallInACupReward
+            from alr_envs.alr.mujoco.ball_in_a_cup.ball_in_a_cup_reward_simple import BallInACupReward
             reward_function = BallInACupReward
         elif reward_type == "contextual_goal":
-            from alr_envs.mujoco.ball_in_a_cup.ball_in_a_cup_reward import BallInACupReward
+            from alr_envs.alr.mujoco.ball_in_a_cup.ball_in_a_cup_reward import BallInACupReward
             reward_function = BallInACupReward
         else:
             raise ValueError("Unknown reward type: {}".format(reward_type))

alr_envs/alr/mujoco/ball_in_a_cup/utils.py

@@ -1,4 +1,4 @@
-from alr_envs.mujoco.ball_in_a_cup.ball_in_a_cup import ALRBallInACupEnv
+from alr_envs.alr.mujoco.ball_in_a_cup.ball_in_a_cup import ALRBallInACupEnv
 from mp_env_api.mp_wrappers.detpmp_wrapper import DetPMPWrapper
 from mp_env_api.mp_wrappers.dmp_wrapper import DmpWrapper
 

alr_envs/alr/mujoco/beerpong/beerpong.py

@@ -37,7 +37,7 @@ class ALRBeerpongEnv(MujocoEnv, utils.EzPickle):
         self.sim_time = 8  # seconds
         self.sim_steps = int(self.sim_time / self.dt)
         if reward_function is None:
-            from alr_envs.mujoco.beerpong.beerpong_reward import BeerpongReward
+            from alr_envs.alr.mujoco.beerpong.beerpong_reward import BeerpongReward
             reward_function = BeerpongReward
         self.reward_function = reward_function(self.sim, self.sim_steps)
         self.cup_robot_id = self.sim.model._site_name2id["cup_robot_final"]

alr_envs/alr/mujoco/beerpong/beerpong_reward.py

@@ -1,5 +1,5 @@
 import numpy as np
-from alr_envs.mujoco import alr_reward_fct
+from alr_envs.alr.mujoco import alr_reward_fct
 
 
 class BeerpongReward(alr_reward_fct.AlrReward):

alr_envs/alr/mujoco/beerpong/beerpong_reward_simple.py

@@ -1,5 +1,5 @@
 import numpy as np
-from alr_envs.mujoco import alr_reward_fct
+from alr_envs.alr.mujoco import alr_reward_fct
 
 
 class BeerpongReward(alr_reward_fct.AlrReward):

alr_envs/alr/mujoco/beerpong/beerpong_simple.py

@@ -38,7 +38,7 @@ class ALRBeerpongEnv(MujocoEnv, utils.EzPickle):
         self.sim_time = 8  # seconds
         self.sim_steps = int(self.sim_time / self.dt)
         if reward_function is None:
-            from alr_envs.mujoco.beerpong.beerpong_reward_simple import BeerpongReward
+            from alr_envs.alr.mujoco.beerpong.beerpong_reward_simple import BeerpongReward
             reward_function = BeerpongReward
         self.reward_function = reward_function(self.sim, self.sim_steps)
         self.cup_robot_id = self.sim.model._site_name2id["cup_robot_final"]

alr_envs/alr/mujoco/beerpong/utils.py

@@ -1,6 +1,6 @@
 from alr_envs.utils.mps.detpmp_wrapper import DetPMPWrapper
-from alr_envs.mujoco.beerpong.beerpong import ALRBeerpongEnv
-from alr_envs.mujoco.beerpong.beerpong_simple import ALRBeerpongEnv as ALRBeerpongEnvSimple
+from alr_envs.alr.mujoco.beerpong.beerpong import ALRBeerpongEnv
+from alr_envs.alr.mujoco.beerpong.beerpong_simple import ALRBeerpongEnv as ALRBeerpongEnvSimple
 
 
 def make_contextual_env(rank, seed=0):

alr_envs/alr/mujoco/gym_table_tennis/envs/table_tennis_env.py

@@ -2,9 +2,9 @@ import numpy as np
 from gym import spaces
 from gym.envs.robotics import robot_env, utils
 # import xml.etree.ElementTree as ET
-from alr_envs.mujoco.gym_table_tennis.utils.rewards.hierarchical_reward import HierarchicalRewardTableTennis
+from alr_envs.alr.mujoco.gym_table_tennis.utils.rewards.hierarchical_reward import HierarchicalRewardTableTennis
 import glfw
-from alr_envs.mujoco.gym_table_tennis.utils.experiment import ball_initialize
+from alr_envs.alr.mujoco.gym_table_tennis.utils.experiment import ball_initialize
 from pathlib import Path
 import os
 

alr_envs/alr/mujoco/gym_table_tennis/utils/experiment.py

@@ -1,6 +1,6 @@
 import numpy as np
 from gym.utils import seeding
-from alr_envs.mujoco.gym_table_tennis.utils.util import read_yaml, read_json
+from alr_envs.alr.mujoco.gym_table_tennis.utils.util import read_yaml, read_json
 from pathlib import Path