Fix: Ensure automatic rendering on .step if render was called initially (new gym spec)

Test Suite Improvements

.github/ISSUE_TEMPLATE/bug_report.md

@@ -0,0 +1,32 @@
+---
+name: Bug report
+about: Create a report to help us improve
+title: ''
+labels: ''
+assignees: ''
+
+---
+
+**Describe the bug**
+A clear and concise description of what the bug is.
+
+**To Reproduce**
+Steps to reproduce the behavior:
+1. Go to '...'
+2. Select variables '....'
+3. Execute '....'
+4. See error
+
+**Expected behavior**
+A clear and concise description of what you expected to happen.
+
+**Screenshots**
+If applicable, add screenshots to help explain your problem.
+
+**Desktop (please complete the following information):**
+ - OS: [e.g. Ubunutu]
+ - Python version [e.g. 3.8]
+ - Version [e.g. 1.0]
+
+**Additional context**
+Add any other context about the problem here.

.github/ISSUE_TEMPLATE/feature_request.md

@@ -0,0 +1,20 @@
+---
+name: Feature request
+about: Suggest an idea for this project
+title: ''
+labels: ''
+assignees: ''
+
+---
+
+**Is your feature request related to a problem? Please describe.**
+A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]
+
+**Describe the solution you'd like**
+A clear and concise description of what you want to happen.
+
+**Describe alternatives you've considered**
+A clear and concise description of any alternative solutions or features you've considered.
+
+**Additional context**
+Add any other context or screenshots about the feature request here.

.github/workflows/publish-to-pypi.yml

@@ -0,0 +1,52 @@
+name: Publish Python package to PyPI
+
+on:
+  push:
+    tags:
+      - '*'
+
+jobs:
+  publish:
+    name: Publish to PyPI
+    runs-on: ubuntu-latest
+    if: false && startsWith(github.ref, 'refs/tags/')  # Only run on tagged commits
+
+    steps:
+    - name: Check out code
+      uses: actions/checkout@v4
+      with:
+        fetch-depth: 0  # This fetches all history for all branches and tags
+
+    - name: Verify tag is on master branch
+      run: |
+        TAG_IS_ON_MASTER=$(git branch -r --contains ${{ github.ref }} | grep 'origin/master')
+        if [ -z "$TAG_IS_ON_MASTER" ]; then
+          echo "Tag is not on the master branch. Cancelling the workflow."
+          exit 1
+        fi
+        echo "Tag is on the master branch. Proceeding with the workflow."
+
+    - name: Set up Python
+      uses: actions/setup-python@v4
+      with:
+        python-version: "3.x"
+
+    - name: Install pypa/build/setuptools/twine
+      run: >-
+        python3 -m
+        pip install
+        build setuptools twine
+        --user
+
+    - name: Prevent fallback onto setup.py
+      run: rm setup.py
+
+    - name: Build a binary wheel and a source tarball
+      run: python3 -m build
+
+    - name: Publish to PyPI
+      env:
+        TWINE_USERNAME: __token__
+        TWINE_PASSWORD: ${{ secrets.PYPI_TOKEN }}
+      run: twine upload dist/*
+

.github/workflows/publish-to-test-pypi.yml

@@ -0,0 +1,52 @@
+name: Publish Python package to TestPyPI
+
+on:
+  push:
+    tags:
+      - '*'
+
+jobs:
+  publish:
+    name: Publish to TestPyPI
+    runs-on: ubuntu-latest
+    if: startsWith(github.ref, 'refs/tags/')  # Only run on tagged commits
+
+    steps:
+    - name: Check out code
+      uses: actions/checkout@v4
+      with:
+        fetch-depth: 0  # This fetches all history for all branches and tags
+
+    - name: Verify tag is on master branch
+      run: |
+        TAG_IS_ON_MASTER=$(git branch -r --contains ${{ github.ref }} | grep 'origin/master')
+        if [ -z "$TAG_IS_ON_MASTER" ]; then
+          echo "Tag is not on the master branch. Cancelling the workflow."
+          exit 1
+        fi
+        echo "Tag is on the master branch. Proceeding with the workflow."
+
+    - name: Set up Python
+      uses: actions/setup-python@v4
+      with:
+        python-version: "3.x"
+
+    - name: Install pypa/build/setuptools/twine
+      run: >-
+        python3 -m
+        pip install
+        build setuptools twine
+        --user
+
+    - name: Prevent fallback onto setup.py
+      run: rm setup.py
+
+    - name: Build a binary wheel and a source tarball
+      run: python3 -m build
+
+    - name: Publish to TestPyPI
+      env:
+        TWINE_USERNAME: __token__
+        TWINE_PASSWORD: ${{ secrets.TEST_PYPI_TOKEN }}
+      run: twine upload --repository-url https://test.pypi.org/legacy/ dist/*
+

CODE_OF_CONDUCT.md

@@ -0,0 +1,128 @@
+# Contributor Covenant Code of Conduct
+
+## Our Pledge
+
+We as members, contributors, and leaders pledge to make participation in our
+community a harassment-free experience for everyone, regardless of age, body
+size, visible or invisible disability, ethnicity, sex characteristics, gender
+identity and expression, level of experience, education, socio-economic status,
+nationality, personal appearance, race, religion, or sexual identity
+and orientation.
+
+We pledge to act and interact in ways that contribute to an open, welcoming,
+diverse, inclusive, and healthy community.
+
+## Our Standards
+
+Examples of behavior that contributes to a positive environment for our
+community include:
+
+* Demonstrating empathy and kindness toward other people
+* Being respectful of differing opinions, viewpoints, and experiences
+* Giving and gracefully accepting constructive feedback
+* Accepting responsibility and apologizing to those affected by our mistakes,
+  and learning from the experience
+* Focusing on what is best not just for us as individuals, but for the
+  overall community
+
+Examples of unacceptable behavior include:
+
+* The use of sexualized language or imagery, and sexual attention or
+  advances of any kind
+* Trolling, insulting or derogatory comments, and personal or political attacks
+* Public or private harassment
+* Publishing others' private information, such as a physical or email
+  address, without their explicit permission
+* Other conduct which could reasonably be considered inappropriate in a
+  professional setting
+
+## Enforcement Responsibilities
+
+Community leaders are responsible for clarifying and enforcing our standards of
+acceptable behavior and will take appropriate and fair corrective action in
+response to any behavior that they deem inappropriate, threatening, offensive,
+or harmful.
+
+Community leaders have the right and responsibility to remove, edit, or reject
+comments, commits, code, wiki edits, issues, and other contributions that are
+not aligned to this Code of Conduct, and will communicate reasons for moderation
+decisions when appropriate.
+
+## Scope
+
+This Code of Conduct applies within all community spaces, and also applies when
+an individual is officially representing the community in public spaces.
+Examples of representing our community include using an official e-mail address,
+posting via an official social media account, or acting as an appointed
+representative at an online or offline event.
+
+## Enforcement
+
+Instances of abusive, harassing, or otherwise unacceptable behavior may be
+reported to the community leaders responsible for enforcement at
+their respective email addresses.
+All complaints will be reviewed and investigated promptly and fairly.
+
+All community leaders are obligated to respect the privacy and security of the
+reporter of any incident.
+
+## Enforcement Guidelines
+
+Community leaders will follow these Community Impact Guidelines in determining
+the consequences for any action they deem in violation of this Code of Conduct:
+
+### 1. Correction
+
+**Community Impact**: Use of inappropriate language or other behavior deemed
+unprofessional or unwelcome in the community.
+
+**Consequence**: A private, written warning from community leaders, providing
+clarity around the nature of the violation and an explanation of why the
+behavior was inappropriate. A public apology may be requested.
+
+### 2. Warning
+
+**Community Impact**: A violation through a single incident or series
+of actions.
+
+**Consequence**: A warning with consequences for continued behavior. No
+interaction with the people involved, including unsolicited interaction with
+those enforcing the Code of Conduct, for a specified period of time. This
+includes avoiding interactions in community spaces as well as external channels
+like social media. Violating these terms may lead to a temporary or
+permanent ban.
+
+### 3. Temporary Ban
+
+**Community Impact**: A serious violation of community standards, including
+sustained inappropriate behavior.
+
+**Consequence**: A temporary ban from any sort of interaction or public
+communication with the community for a specified period of time. No public or
+private interaction with the people involved, including unsolicited interaction
+with those enforcing the Code of Conduct, is allowed during this period.
+Violating these terms may lead to a permanent ban.
+
+### 4. Permanent Ban
+
+**Community Impact**: Demonstrating a pattern of violation of community
+standards, including sustained inappropriate behavior,  harassment of an
+individual, or aggression toward or disparagement of classes of individuals.
+
+**Consequence**: A permanent ban from any sort of public interaction within
+the community.
+
+## Attribution
+
+This Code of Conduct is adapted from the [Contributor Covenant][homepage],
+version 2.0, available at
+https://www.contributor-covenant.org/version/2/0/code_of_conduct.html.
+
+Community Impact Guidelines were inspired by [Mozilla's code of conduct
+enforcement ladder](https://github.com/mozilla/diversity).
+
+[homepage]: https://www.contributor-covenant.org
+
+For answers to common questions about this code of conduct, see the FAQ at
+https://www.contributor-covenant.org/faq. Translations are available at
+https://www.contributor-covenant.org/translations.

CONTRIBUTING.md

@@ -0,0 +1,28 @@
+# Contribution Guidelines
+
+We welcome and appreciate contributions to this repository from the community. To ensure a positive and productive collaboration, we have established the following guidelines for contributing to this project:
+
+**Code of Conduct**: Please abide by our Code of Conduct, which sets the standards for respectful and inclusive behavior within the community.
+
+**Issues and Feature Requests**: Feel free to open new issues to report bugs or suggest new features. Before submitting, please check if a similar issue or request already exists.
+
+**Pull Requests**: We encourage pull requests for bug fixes, new features, and improvements. Please follow these steps when submitting a pull request:
+- Fork this repository and create a new branch for your changes.
+- Ensure your code aligns with our coding standards.
+- Provide a clear and concise description of your changes.
+- Test your changes thoroughly.
+- Make sure your code is well-documented.
+- Code Style: Follow our coding style guidelines, which include formatting, naming conventions, and other coding standards. Adhering to these guidelines will help streamline the review process.
+
+**Collaboration**: Be prepared for constructive feedback during the review process. We aim to maintain the quality and consistency of the codebase.
+
+**License**: Ensure your contributions comply with the project's existing open-source license. By contributing, you grant us the right to distribute your code under the project's license.
+
+**Ownership**: When you submit a contribution, you confirm that you have the right to license your code to us and that your work does not violate any existing patents, trademarks, or intellectual property rights.
+
+**Maintainers**: The repository maintainers have the final authority on merging or rejecting contributions. They will work collaboratively with contributors to ensure a smooth process.
+
+**Community Support**: As part of our community, please consider helping others, answering questions, and providing support to fellow contributors.
+
+We appreciate your interest in contributing to our project. Your involvement helps make this community vibrant and successful.
+Thank you for being part of this project!

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 Autonomous Learning Robots Lab @ KIT
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

MANIFEST.in

@@ -0,0 +1,14 @@
+# Include the README
+include README.md
+
+# Include the license, Code of Conduct and Contributing guidelines
+include LICENSE
+include CODE_OF_CONDUCT.md
+include CONTRIBUTING.md
+
+# Include stl and xml files from the fancy_gym/envs/mujoco directory
+recursive-include fancy_gym/envs/mujoco *.stl
+recursive-include fancy_gym/envs/mujoco *.xml
+
+# Also shipping the most important part of fancy gym
+include icon.svg

README.md

@@ -1,101 +1,141 @@
-# Fancy Gym
+<h1 align="center">
+  <br>
+  <img src='https://raw.githubusercontent.com/ALRhub/fancy_gym/master/icon.svg' width="250px">
+  <br><br>
+  <b>Fancy Gym</b>
+  <br><br>
+</h1>
 
-`fancy_gym` offers a large variety of reinforcement learning environments under the unifying interface
-of [OpenAI gym](https://gymlibrary.dev/). We provide support (under the OpenAI gym interface) for the benchmark suites
-[DeepMind Control](https://deepmind.com/research/publications/2020/dm-control-Software-and-Tasks-for-Continuous-Control)
-(DMC) and [Metaworld](https://meta-world.github.io/). If those are not sufficient and you want to create your own custom
-gym environments, use [this guide](https://www.gymlibrary.dev/content/environment_creation/). We highly appreciate it, if
-you would then submit a PR for this environment to become part of `fancy_gym`.  
-In comparison to existing libraries, we additionally support to control agents with movement primitives, such as Dynamic
-Movement Primitives (DMPs) and Probabilistic Movement Primitives (ProMP).
+|  &#x2757; Fancy Gym has recently received a major refactor, which also updated many of the used dependencies to current versions. The update has brought some breaking changes. If you want to access the old version, check out the [legacy branch](https://github.com/ALRhub/fancy_gym/tree/legacy). Find out more about what changed [here](https://github.com/ALRhub/fancy_gym/pull/75). |
+| --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+
+Built upon the foundation of [Gymnasium](https://gymnasium.farama.org/) (a maintained fork of OpenAI’s renowned Gym library) `fancy_gym` offers a comprehensive collection of reinforcement learning environments.
+
+**Key Features**:
+
+- **New Challenging Environments**: `fancy_gym` includes several new environments (Panda Box Pushing, Table Tennis, etc.) that present a higher degree of difficulty, pushing the boundaries of reinforcement learning research.
+- **Support for Movement Primitives**: `fancy_gym` supports a range of movement primitives (MPs), including Dynamic Movement Primitives (DMPs), Probabilistic Movement Primitives (ProMP), and Probabilistic Dynamic Movement Primitives (ProDMP).
+- **Upgrade to Movement Primitives**: With our framework, it's straightforward to transform standard Gymnasium environments into environments that support movement primitives.
+- **Benchmark Suite Compatibility**: `fancy_gym` makes it easy to access renowned benchmark suites such as [DeepMind Control](https://deepmind.com/research/publications/2020/dm-control-Software-and-Tasks-for-Continuous-Control) and [Metaworld](https://meta-world.github.io/), whether you want to use them in the regular step-based setting or using MPs.
+- **Contribute Your Own Environments**: If you're inspired to create custom gym environments, both step-based and with movement primitives, this [guide](https://gymnasium.farama.org/tutorials/gymnasium_basics/environment_creation/) will assist you. We encourage and highly appreciate submissions via PRs to integrate these environments into `fancy_gym`.
 
 ## Movement Primitive Environments (Episode-Based/Black-Box Environments)
 
-Unlike step-based environments, movement primitive (MP) environments are closer related to stochastic search, black-box
-optimization, and methods that are often used in traditional robotics and control. MP environments are typically
-episode-based and execute a full trajectory, which is generated by a trajectory generator, such as a Dynamic Movement
-Primitive (DMP) or a Probabilistic Movement Primitive (ProMP). The generated trajectory is translated into individual
-step-wise actions by a trajectory tracking 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 as well as a special controller for the MetaWorld control suite.  
-The goal of all MP environments is still to learn an optimal policy. Yet, an action represents the parametrization of
-the motion primitives to generate a suitable trajectory. Additionally, in this framework we support all of this also for
-the contextual setting, i.e. we expose the context space - a subset of the observation space - in the beginning of the
-episode. This requires to predict a new action/MP parametrization for each context.
+<p align="justify">
+Movement primitive (MP) environments differ from traditional step-based environments. They align more with concepts from stochastic search, black-box optimization, and methods commonly found in classical robotics and control. Instead of individual steps, MP environments operate on an episode basis, executing complete trajectories. These trajectories are produced by trajectory generators like Dynamic Movement Primitives (DMP), Probabilistic Movement Primitives (ProMP) or Probabilistic Dynamic Movement Primitives (ProDMP).
+</p>
+<p align="justify">
+Once generated, these trajectories are converted into step-by-step actions using a trajectory tracking controller. The specific controller chosen depends on the environment's requirements. Currently, we support position, velocity, and PD-Controllers tailored for position, velocity, and torque control. Additionally, we have a specialized controller designed for the MetaWorld control suite.
+</p>
+<p align="justify">
+While the overarching objective of MP environments remains the learning of an optimal policy, the actions here represent the parametrization of motion primitives to craft the right trajectory. Our framework further enhances this by accommodating a contextual setting. At the episode's onset, we present the context space—a subset of the observation space. This demands the prediction of a new action or MP parametrization for every unique context.
+</p>
 
 ## Installation
 
-1. Clone the repository
+We recommend installing `fancy_gym` into a virtual environment as provided by [venv](https://docs.python.org/3/library/venv.html). 3rd party alternatives to venv like [Poetry](https://python-poetry.org/) or [Conda](https://docs.conda.io/en/latest/) can also be used.
 
+### Installation from PyPI (recommended)
+
+Install `fancy_gym` via
+```bash
+pip install fancy_gym
+```
+
+We have a few optional dependencies. If you also want to install those use
+
+```bash
+# to install all optional dependencies
+pip install 'fancy_gym[all]'
+
+# or choose only those you want
+pip install 'fancy_gym[dmc,box2d,mujoco-legacy,jax,testing]'
+```
+
+Pip can not automatically install up-to-date versions of metaworld, since they are not avaible on PyPI yet.
+Install metaworld via
+
+```bash
+pip install metaworld@git+https://github.com/Farama-Foundation/Metaworld.git@d155d0051630bb365ea6a824e02c66c068947439#egg=metaworld
+```
+
+### Installation from master
+
+1. Clone the repository
 ```bash
 git clone git@github.com:ALRhub/fancy_gym.git
 ```
 
 2. Go to the folder
-
 ```bash
 cd fancy_gym
 ```
 
 3. Install with
-
 ```bash
 pip install -e .
 ```
 
-In case you want to use dm_control oder metaworld, you can install them by specifying extras
-
+We have a few optional dependencies. If you also want to install those use
 ```bash
-pip install -e .[dmc,metaworld]
+# to install all optional dependencies
+pip install -e '.[all]'
+
+# or choose only those you want
+pip install -e '.[dmc,box2d,mujoco-legacy,jax,testing]'
 ```
 
-> **Note:**   
-> While our library already fully supports the new mujoco bindings, metaworld still relies on
-> [mujoco_py](https://github.com/openai/mujoco-py), hence make sure to have mujoco 2.1 installed beforehand.
+Metaworld has to be installed manually with
+```bash
+pip install metaworld@git+https://github.com/Farama-Foundation/Metaworld.git@d155d0051630bb365ea6a824e02c66c068947439#egg=metaworld
+```
 
 ## How to use Fancy Gym
 
-We will only show the basics here and prepared [multiple examples](fancy_gym/examples/) for a more detailed look.
+We will only show the basics here and prepared [multiple examples](https://github.com/ALRhub/fancy_gym/tree/master/fancy_gym/examples/) for a more detailed look.
 
-### Step-wise Environments
+### Step-Based Environments
+
+Regular step based environments added by Fancy Gym are added into the `fancy/` namespace.
+
+|  &#x2757; Legacy versions of Fancy Gym used `fancy_gym.make(...)`. This is no longer supported and will raise an Exception on new versions. |
+| ----------------------------------------------------------------------------------------------------------------------------------------------- |
 
 ```python
+import gymnasium as gym
 import fancy_gym
 
-env = fancy_gym.make('Reacher5d-v0', seed=1)
-obs = env.reset()
+env = gym.make('fancy/Reacher5d-v0', render_mode='human')
+# or env = gym.make('metaworld/reach-v2', render_mode='human') # fancy_gym allows access to all metaworld ML1 tasks via the metaworld/ NS
+# or env = gym.make('dm_control/ball_in_cup-catch-v0', render_mode='human')
+# or env = gym.make('Reacher-v2', render_mode='human')
+observation = env.reset(seed=1)
+env.render()
 
 for i in range(1000):
     action = env.action_space.sample()
-    obs, reward, done, info = env.step(action)
-    if i % 5 == 0:
-        env.render()
+    observation, reward, terminated, truncated, info = env.step(action)
 
-    if done:
-        obs = env.reset()
+    if terminated or truncated:
+        observation, info = env.reset()
 ```
 
-When using `dm_control` tasks we expect the `env_id` to be specified as `dmc:domain_name-task_name` or for manipulation
-tasks as `dmc:manipulation-environment_name`. For `metaworld` tasks, we require the structure `metaworld:env_id-v2`, our
-custom tasks and standard gym environments can be created without prefixes.
-
 ### Black-box Environments
 
-All environments provide by default the cumulative episode reward, this can however be changed if necessary. Optionally,
-each environment returns all collected information from each step as part of the infos. This information is, however,
-mainly meant for debugging as well as logging and not for training.
+All environments provide by default the cumulative episode reward, this can however be changed if necessary. Optionally, each environment returns all collected information from each step as part of the infos. This information is, however, mainly meant for debugging as well as logging and not for training.
 
-|Key| Description|Type
-|---|---|---|
-`positions`| Generated trajectory from MP | Optional
-`velocities`| Generated trajectory from MP | Optional
-`step_actions`| Step-wise executed action based on controller output | Optional
-`step_observations`| Step-wise intermediate observations | Optional
-`step_rewards`| Step-wise rewards | Optional
-`trajectory_length`| Total number of environment interactions | Always
-`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`. | Always
+| Key                 | Description                                                                                                                                                                                                                                      | Type     |
+| ------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | -------- |
+| `positions`         | Generated trajectory from MP                                                                                                                                                                                                                     | Optional |
+| `velocities`        | Generated trajectory from MP                                                                                                                                                                                                                     | Optional |
+| `step_actions`      | Step-wise executed action based on controller output                                                                                                                                                                                             | Optional |
+| `step_observations` | Step-wise intermediate observations                                                                                                                                                                                                              | Optional |
+| `step_rewards`      | Step-wise rewards                                                                                                                                                                                                                                | Optional |
+| `trajectory_length` | Total number of environment interactions                                                                                                                                                                                                         | Always   |
+| `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`. | Always   |
 
-Existing MP tasks can be created the same way as above. Just keep in mind, calling `step()` executes a full trajectory.
+Existing MP tasks can be created the same way as above. The namespace of a MP-variant of an environment is given by `<original namespace>_<MP name>/`.
+Just keep in mind, calling `step()` executes a full trajectory.
 
 > **Note:**  
 > Currently, we are also in the process of enabling replanning as well as learning of sub-trajectories.
@@ -105,30 +145,35 @@ Existing MP tasks can be created the same way as above. Just keep in mind, calli
 > Feel free to try it and open an issue with any problems that occur.
 
 ```python
+import gymnasium as gym
 import fancy_gym
 
-env = fancy_gym.make('Reacher5dProMP-v0', seed=1)
-# render() can be called once in the beginning with all necessary arguments.
-# To turn it of again just call render() without any arguments. 
-env.render(mode='human')
+env = gym.make('fancy_ProMP/Reacher5d-v0', render_mode="human")
+# or env = gym.make('metaworld_ProDMP/reach-v2', render_mode="human")
+# or env = gym.make('dm_control_DMP/ball_in_cup-catch-v0', render_mode="human")
+# or env = gym.make('gym_ProMP/Reacher-v2', render_mode="human") # mp versions of envs added directly by gymnasium are in the gym_<MP-type> NS
 
 # This returns the context information, not the full state observation
-obs = env.reset()
+observation, info = env.reset(seed=1)
+env.render()
 
 for i in range(5):
     action = env.action_space.sample()
-    obs, reward, done, info = env.step(action)
+    observation, reward, terminated, truncated, info = env.step(action)
 
-    # Done is always True as we are working on the episode level, hence we always reset()
-    obs = env.reset()
+    # terminated or truncated is always True as we are working on the episode level, hence we always reset()
+    observation, info = env.reset()
 ```
 
 To show all available environments, we provide some additional convenience variables. All of them return a dictionary
-with two keys `DMP` and `ProMP` that store a list of available environment ids.
+with the keys `DMP`, `ProMP`, `ProDMP` and `all` that store a list of available environment ids.
 
 ```python
 import fancy_gym
 
+print("All Black-box tasks:")
+print(fancy_gym.ALL_MOVEMENT_PRIMITIVE_ENVIRONMENTS)
+
 print("Fancy Black-box tasks:")
 print(fancy_gym.ALL_FANCY_MOVEMENT_PRIMITIVE_ENVIRONMENTS)
 
@@ -140,22 +185,30 @@ print(fancy_gym.ALL_DMC_MOVEMENT_PRIMITIVE_ENVIRONMENTS)
 
 print("MetaWorld Black-box tasks:")
 print(fancy_gym.ALL_METAWORLD_MOVEMENT_PRIMITIVE_ENVIRONMENTS)
+
+print("If you add custom envs, their mp versions will be found in:")
+print(fancy_gym.MOVEMENT_PRIMITIVE_ENVIRONMENTS_FOR_NS['<my_custom_namespace>'])
 ```
 
 ### 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](fancy_gym/black_box/raw_interface_wrapper.py) needs to be implemented.
+hand, the following [interface](https://github.com/ALRhub/fancy_gym/tree/master/fancy_gym/black_box/raw_interface_wrapper.py) needs to be implemented.
 
 ```python
 from abc import abstractmethod
 from typing import Union, Tuple
 
-import gym
+import gymnasium as gym
 import numpy as np
 
 
 class RawInterfaceWrapper(gym.Wrapper):
+    mp_config = {
+        'ProMP': {},
+        'DMP': {},
+        'ProDMP': {},
+    }
 
     @property
     def context_mask(self) -> np.ndarray:
@@ -167,7 +220,6 @@ class RawInterfaceWrapper(gym.Wrapper):
             context/part of the first observation, the velocities are not necessary in the observation for the task.
             Returns:
                 bool array representing the indices of the observations
-
         """
         return np.ones(self.env.observation_space.shape[0], dtype=bool)
 
@@ -197,34 +249,91 @@ class RawInterfaceWrapper(gym.Wrapper):
 
 ```
 
-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](fancy_gym/examples/).
+Default configurations for MPs can be overitten by defining attributes in mp_config.
+Available parameters are documented in the [MP_PyTorch Userguide](https://github.com/ALRhub/MP_PyTorch/blob/main/doc/README.md).
 
 ```python
-import fancy_gym
+class RawInterfaceWrapper(gym.Wrapper):
+    mp_config = {
+        'ProMP': {
+            'phase_generator_kwargs': {
+                'phase_generator_type': 'linear'
+                # When selecting another generator type, the default configuration will not be merged for the attribute.
+            },
+            'controller_kwargs': {
+                'p_gains': 0.5 * np.array([1.0, 4.0, 2.0, 4.0, 1.0, 4.0, 1.0]),
+                'd_gains': 0.5 * np.array([0.1, 0.4, 0.2, 0.4, 0.1, 0.4, 0.1]),
+            },
+            'basis_generator_kwargs': {
+                'num_basis': 3,
+                'num_basis_zero_start': 1,
+                'num_basis_zero_goal': 1,
+            },
+        },
+        'DMP': {},
+        'ProDMP': {}.
+    }
 
-# Base environment name, according to structure of above example
-base_env_id = "dmc:ball_in_cup-catch"
+    [...]
+```
 
-# Replace this wrapper with the custom wrapper for your environment by inheriting from the RawInferfaceWrapper.
-# You can also add other gym.Wrappers in case they are needed, 
-# e.g. gym.wrappers.FlattenObservation for dict observations
-wrappers = [fancy_gym.dmc.suite.ball_in_cup.MPWrapper]
-kwargs = {...}
-env = fancy_gym.make_bb(base_env_id, wrappers=wrappers, seed=0, **kwargs)
+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](https://github.com/ALRhub/fancy_gym/tree/master/fancy_gym/examples/).
+
+If the step-based is already registered with gym, you can simply do the following:
+
+```python
+fancy_gym.upgrade(
+    id='custom/cool_new_env-v0',
+    mp_wrapper=my_custom_MPWrapper
+)
+```
+
+If the step-based is not yet registered with gym we can add both the step-based and MP-versions via
+
+```python
+fancy_gym.register(
+    id='custom/cool_new_env-v0',
+    entry_point=my_custom_env,
+    mp_wrapper=my_custom_MPWrapper
+)
+```
+
+From this point on, you can access MP-version of your environments via
+
+```python
+env = gym.make('custom_ProDMP/cool_new_env-v0')
 
 rewards = 0
-obs = env.reset()
+observation, info = 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)
+    observation, reward, terminated, truncated, info = env.step(ac)
     rewards += reward
 
-    if done:
-        print(base_env_id, rewards)
+    if terminated or truncated:
+        print(rewards)
         rewards = 0
-        obs = env.reset()
+        observation, info = env.reset()
 ```
+
+## Citing the Project
+
+To cite this repository in publications:
+
+```bibtex
+@software{fancy_gym,
+	title = {Fancy Gym},
+	author = {Otto, Fabian and Celik, Onur and Roth, Dominik and Zhou, Hongyi},
+	abstract = {Fancy Gym: Unifying interface for various RL benchmarks with support for Black Box approaches.},
+	url = {https://github.com/ALRhub/fancy_gym},
+	organization = {Autonomous Learning Robots Lab (ALR) at KIT},
+}
+```
+
+## Icon Attribution
+
+The icon is based on the [Gymnasium](https://github.com/Farama-Foundation/Gymnasium) icon as can be found [here](https://gymnasium.farama.org/_static/img/gymnasium_black.svg).

fancy_gym/__init__.py

@@ -1,13 +1,20 @@
 from fancy_gym import dmc, meta, open_ai
-from fancy_gym.utils.make_env_helpers import make, make_bb, make_rank
-from .dmc import ALL_DMC_MOVEMENT_PRIMITIVE_ENVIRONMENTS
-# Convenience function for all MP environments
-from .envs import ALL_FANCY_MOVEMENT_PRIMITIVE_ENVIRONMENTS
-from .meta import ALL_METAWORLD_MOVEMENT_PRIMITIVE_ENVIRONMENTS
-from .open_ai import ALL_GYM_MOVEMENT_PRIMITIVE_ENVIRONMENTS
+from fancy_gym import envs as fancy
+from fancy_gym.utils.make_env_helpers import make_bb
+from .envs.registry import register, upgrade
+from .envs.registry import ALL_MOVEMENT_PRIMITIVE_ENVIRONMENTS, MOVEMENT_PRIMITIVE_ENVIRONMENTS_FOR_NS
 
-ALL_MOVEMENT_PRIMITIVE_ENVIRONMENTS = {
-    key: value + ALL_DMC_MOVEMENT_PRIMITIVE_ENVIRONMENTS[key] +
-         ALL_GYM_MOVEMENT_PRIMITIVE_ENVIRONMENTS[key] +
-         ALL_METAWORLD_MOVEMENT_PRIMITIVE_ENVIRONMENTS[key]
-    for key, value in ALL_FANCY_MOVEMENT_PRIMITIVE_ENVIRONMENTS.items()}
+ALL_DMC_MOVEMENT_PRIMITIVE_ENVIRONMENTS = MOVEMENT_PRIMITIVE_ENVIRONMENTS_FOR_NS['dm_control']
+ALL_FANCY_MOVEMENT_PRIMITIVE_ENVIRONMENTS = MOVEMENT_PRIMITIVE_ENVIRONMENTS_FOR_NS['fancy']
+if 'metaworld' in MOVEMENT_PRIMITIVE_ENVIRONMENTS_FOR_NS:
+    ALL_METAWORLD_MOVEMENT_PRIMITIVE_ENVIRONMENTS = MOVEMENT_PRIMITIVE_ENVIRONMENTS_FOR_NS['metaworld']
+else:
+    ALL_METAWORLD_MOVEMENT_PRIMITIVE_ENVIRONMENTS = 'Metaworld is not installed.'
+ALL_GYM_MOVEMENT_PRIMITIVE_ENVIRONMENTS = MOVEMENT_PRIMITIVE_ENVIRONMENTS_FOR_NS['gym']
+
+
+def make(*args, **kwargs):
+    """
+    As part of the refactor of Fancy Gym and upgrade to gymnasium the use of fancy_gym.make has been discontinued. Regular gym.make should be used instead. For more details check out the github README. If your codebase was build for older versions of Fancy Gym and relies on the old behavior and dependency versions, please check out the legacy branch.
+    """
+    raise Exception('As part of the refactor of Fancy Gym and upgrade to gymnasium the use of fancy_gym.make has been discontinued. Regular gym.make should be used instead. For more details check out the github README. If your codebase was build for older versions of Fancy Gym and relies on the old behavior and dependency versions, please check out the legacy branch.')

fancy_gym/black_box/black_box_wrapper.py

@@ -1,8 +1,9 @@
-from typing import Tuple, Optional, Callable
+from typing import Tuple, Optional, Callable, Dict, Any
 
-import gym
+import gymnasium as gym
 import numpy as np
-from gym import spaces
+from gymnasium import spaces
+from gymnasium.core import ObsType
 from mp_pytorch.mp.mp_interfaces import MPInterface
 
 from fancy_gym.black_box.controller.base_controller import BaseController
@@ -67,13 +68,14 @@ class BlackBoxWrapper(gym.ObservationWrapper):
         self.reward_aggregation = reward_aggregation
 
         # spaces
-        self.return_context_observation = not (learn_sub_trajectories or self.do_replanning)
+        self.return_context_observation = not (
+            learn_sub_trajectories or self.do_replanning)
         self.traj_gen_action_space = self._get_traj_gen_action_space()
         self.action_space = self._get_action_space()
         self.observation_space = self._get_observation_space()
 
         # rendering
-        self.render_kwargs = {}
+        self.do_render = False
         self.verbose = verbose
 
         # condition value
@@ -99,14 +101,17 @@ class BlackBoxWrapper(gym.ObservationWrapper):
             # If we do not do this, the traj_gen assumes we are continuing the trajectory.
             self.traj_gen.reset()
 
-        clipped_params = np.clip(action, self.traj_gen_action_space.low, self.traj_gen_action_space.high)
+        clipped_params = np.clip(
+            action, self.traj_gen_action_space.low, self.traj_gen_action_space.high)
         self.traj_gen.set_params(clipped_params)
-        init_time = np.array(0 if not self.do_replanning else self.current_traj_steps * self.dt)
+        init_time = np.array(
+            0 if not self.do_replanning else self.current_traj_steps * self.dt)
 
-        condition_pos = self.condition_pos if self.condition_pos is not None else self.current_pos
-        condition_vel = self.condition_vel if self.condition_vel is not None else self.current_vel
+        condition_pos = self.condition_pos if self.condition_pos is not None else self.env.get_wrapper_attr('current_pos')
+        condition_vel = self.condition_vel if self.condition_vel is not None else self.env.get_wrapper_attr('current_vel')
 
-        self.traj_gen.set_initial_conditions(init_time, condition_pos, condition_vel)
+        self.traj_gen.set_initial_conditions(
+            init_time, condition_pos, condition_vel)
         self.traj_gen.set_duration(duration, self.dt)
 
         position = get_numpy(self.traj_gen.get_traj_pos())
@@ -153,24 +158,27 @@ class BlackBoxWrapper(gym.ObservationWrapper):
         trajectory_length = len(position)
         rewards = np.zeros(shape=(trajectory_length,))
         if self.verbose >= 2:
-            actions = np.zeros(shape=(trajectory_length,) + self.env.action_space.shape)
+            actions = np.zeros(shape=(trajectory_length,) +
+                               self.env.action_space.shape)
             observations = np.zeros(shape=(trajectory_length,) + self.env.observation_space.shape,
                                     dtype=self.env.observation_space.dtype)
 
         infos = dict()
-        done = False
+        terminated, truncated = False, False
 
         if not traj_is_valid:
-            obs, trajectory_return, done, infos = self.env.invalid_traj_callback(action, position, velocity,
-                                                                                 self.return_context_observation,
-                                                                                 self.tau_bound, self.delay_bound)
-            return self.observation(obs), trajectory_return, done, infos
+            obs, trajectory_return, terminated, truncated, infos = self.env.invalid_traj_callback(action, position, velocity,
+                                                                                                  self.return_context_observation, self.tau_bound, self.delay_bound)
+            return self.observation(obs), trajectory_return, terminated, truncated, infos
 
         self.plan_steps += 1
         for t, (pos, vel) in enumerate(zip(position, velocity)):
-            step_action = self.tracking_controller.get_action(pos, vel, self.current_pos, self.current_vel)
-            c_action = np.clip(step_action, self.env.action_space.low, self.env.action_space.high)
-            obs, c_reward, done, info = self.env.step(c_action)
+            step_action = self.tracking_controller.get_action(
+                pos, vel, self.env.get_wrapper_attr('current_pos'), self.env.get_wrapper_attr('current_vel'))
+            c_action = np.clip(
+                step_action, self.env.action_space.low, self.env.action_space.high)
+            obs, c_reward, terminated, truncated, info = self.env.step(
+                c_action)
             rewards[t] = c_reward
 
             if self.verbose >= 2:
@@ -182,12 +190,11 @@ class BlackBoxWrapper(gym.ObservationWrapper):
                 elems[t] = v
                 infos[k] = elems
 
-            if self.render_kwargs:
-                self.env.render(**self.render_kwargs)
+            if self.do_render:
+                self.env.render()
 
-            if done or (self.replanning_schedule(self.current_pos, self.current_vel, obs, c_action,
-                                                t + 1 + self.current_traj_steps)
-                    and self.plan_steps < self.max_planning_times):
+
+            if terminated or truncated or (self.replanning_schedule(self.env.get_wrapper_attr('current_pos'), self.env.get_wrapper_attr('current_vel'), obs, c_action, t + 1 + self.current_traj_steps) and self.plan_steps < self.max_planning_times):
 
                 if self.condition_on_desired:
                     self.condition_pos = pos
@@ -207,17 +214,16 @@ class BlackBoxWrapper(gym.ObservationWrapper):
 
         infos['trajectory_length'] = t + 1
         trajectory_return = self.reward_aggregation(rewards[:t + 1])
-        return self.observation(obs), trajectory_return, done, infos
+        return self.observation(obs), trajectory_return, terminated, truncated, infos
 
-    def render(self, **kwargs):
-        """Only set render options here, such that they can be used during the rollout.
-        This only needs to be called once"""
-        self.render_kwargs = kwargs
+    def render(self):
+        self.do_render = True
 
-    def reset(self, *, seed: Optional[int] = None, return_info: bool = False, options: Optional[dict] = None):
+    def reset(self, *, seed: Optional[int] = None, options: Optional[Dict[str, Any]] = None) \
+            -> Tuple[ObsType, Dict[str, Any]]:
         self.current_traj_steps = 0
         self.plan_steps = 0
         self.traj_gen.reset()
         self.condition_pos = None
         self.condition_vel = None
-        return super(BlackBoxWrapper, self).reset()
+        return super(BlackBoxWrapper, self).reset(seed=seed, options=options)

fancy_gym/black_box/factory/controller_factory.py

@@ -11,11 +11,11 @@ def get_controller(controller_type: str, **kwargs):
     if controller_type == "motor":
         return PDController(**kwargs)
     elif controller_type == "velocity":
-        return VelController()
+        return VelController(**kwargs)
     elif controller_type == "position":
-        return PosController()
+        return PosController(**kwargs)
     elif controller_type == "metaworld":
-        return MetaWorldController()
+        return MetaWorldController(**kwargs)
     else:
         raise ValueError(f"Specified controller type {controller_type} not supported, "
                          f"please choose one of {ALL_TYPES}.")

fancy_gym/black_box/raw_interface_wrapper.py

@@ -1,6 +1,6 @@
 from typing import Union, Tuple
 
-import gym
+import gymnasium as gym
 import numpy as np
 from mp_pytorch.mp.mp_interfaces import MPInterface
 
@@ -114,7 +114,8 @@ class RawInterfaceWrapper(gym.Wrapper):
         Returns:
             obs: artificial observation if the trajectory is invalid, by default a zero vector
             reward: artificial reward if the trajectory is invalid, by default 0
-            done: artificial done if the trajectory is invalid, by default True
+            terminated: artificial terminated if the trajectory is invalid, by default True
+            truncated: artificial truncated if the trajectory is invalid, by default False
             info: artificial info if the trajectory is invalid, by default empty dict
         """
-        return np.zeros(1), 0, True, {}
+        return np.zeros(1), 0, True, False, {}

fancy_gym/dmc/README.MD

@@ -9,11 +9,11 @@ environments in order to use our Motion Primitive gym interface with them.
 [//]: <> (These environments are wrapped-versions of their Deep Mind Control Suite &#40;DMC&#41; 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_promp-v0`| A ProMP wrapped version of the "catch" task for the "ball_in_cup" environment. | 1000 | 10 | 2
-|`dmc_ball_in_cup-catch_dmp-v0`| A DMP wrapped version of the "catch" task for the "ball_in_cup" environment. | 1000| 10 | 2
-|`dmc_reacher-easy_promp-v0`| A ProMP 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_promp-v0`| A ProMP 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
+| Name                                     | Description                                                                    | Trajectory Horizon | Action Dimension | Context Dimension |
+| ---------------------------------------- | ------------------------------------------------------------------------------ | ------------------ | ---------------- | ----------------- |
+| `dm_control_ProDMP/ball_in_cup-catch-v0` | A ProMP wrapped version of the "catch" task for the "ball_in_cup" environment. | 1000               | 10               | 2                 |
+| `dm_control_DMP/ball_in_cup-catch-v0`    | A DMP wrapped version of the "catch" task for the "ball_in_cup" environment.   | 1000               | 10               | 2                 |
+| `dm_control_ProDMP/reacher-easy-v0`      | A ProMP wrapped version of the "easy" task for the "reacher" environment.      | 1000               | 10               | 4                 |
+| `dm_control_DMP/reacher-easy-v0`         | A DMP wrapped version of the "easy" task for the "reacher" environment.        | 1000               | 10               | 4                 |
+| `dm_control_ProDMP/reacher-hard-v0`      | A ProMP wrapped version of the "hard" task for the "reacher" environment.      | 1000               | 10               | 4                 |
+| `dm_control_DMP/reacher-hard-v0`         | A DMP wrapped version of the "hard" task for the "reacher" environment.        | 1000               | 10               | 4                 |

fancy_gym/dmc/__init__.py

@@ -1,245 +1,61 @@
-from gym.envs.registration import register
 from copy import deepcopy
 
+from gymnasium.wrappers import FlattenObservation
+from gymnasium.envs.registration import register
+
+from ..envs.registry import register
+
 from . import manipulation, suite
 
-ALL_DMC_MOVEMENT_PRIMITIVE_ENVIRONMENTS = {"DMP": [], "ProMP": [], "ProDMP": []}
-
-
-DEFAULT_BB_DICT_ProMP = {
-    "name": 'EnvName',
-    "wrappers": [],
-    "trajectory_generator_kwargs": {
-        'trajectory_generator_type': 'promp'
-    },
-    "phase_generator_kwargs": {
-        'phase_generator_type': 'linear'
-    },
-    "controller_kwargs": {
-        'controller_type': 'motor',
-        "p_gains": 50.,
-        "d_gains": 1.,
-    },
-    "basis_generator_kwargs": {
-        'basis_generator_type': 'zero_rbf',
-        'num_basis': 5,
-        'num_basis_zero_start': 1
-    }
-}
-
-DEFAULT_BB_DICT_DMP = {
-    "name": 'EnvName',
-    "wrappers": [],
-    "trajectory_generator_kwargs": {
-        'trajectory_generator_type': 'dmp'
-    },
-    "phase_generator_kwargs": {
-        'phase_generator_type': 'exp'
-    },
-    "controller_kwargs": {
-        'controller_type': 'motor',
-        "p_gains": 50.,
-        "d_gains": 1.,
-    },
-    "basis_generator_kwargs": {
-        'basis_generator_type': 'rbf',
-        'num_basis': 5
-    }
-}
-
 # DeepMind Control Suite (DMC)
-kwargs_dict_bic_dmp = deepcopy(DEFAULT_BB_DICT_DMP)
-kwargs_dict_bic_dmp['name'] = f"dmc:ball_in_cup-catch"
-kwargs_dict_bic_dmp['wrappers'].append(suite.ball_in_cup.MPWrapper)
-# bandwidth_factor=2
-kwargs_dict_bic_dmp['phase_generator_kwargs']['alpha_phase'] = 2
-kwargs_dict_bic_dmp['trajectory_generator_kwargs']['weight_scale'] = 10  # TODO: weight scale 1, but goal scale 0.1
 register(
-    id=f'dmc_ball_in_cup-catch_dmp-v0',
-    entry_point='fancy_gym.utils.make_env_helpers:make_bb_env_helper',
-    kwargs=kwargs_dict_bic_dmp
+    id=f"dm_control/ball_in_cup-catch-v0",
+    register_step_based=False,
+    mp_wrapper=suite.ball_in_cup.MPWrapper,
+    add_mp_types=['DMP', 'ProMP'],
 )
-ALL_DMC_MOVEMENT_PRIMITIVE_ENVIRONMENTS["DMP"].append("dmc_ball_in_cup-catch_dmp-v0")
 
-kwargs_dict_bic_promp = deepcopy(DEFAULT_BB_DICT_ProMP)
-kwargs_dict_bic_promp['name'] = f"dmc:ball_in_cup-catch"
-kwargs_dict_bic_promp['wrappers'].append(suite.ball_in_cup.MPWrapper)
 register(
-    id=f'dmc_ball_in_cup-catch_promp-v0',
-    entry_point='fancy_gym.utils.make_env_helpers:make_bb_env_helper',
-    kwargs=kwargs_dict_bic_promp
+    id=f"dm_control/reacher-easy-v0",
+    register_step_based=False,
+    mp_wrapper=suite.reacher.MPWrapper,
+    add_mp_types=['DMP', 'ProMP'],
 )
-ALL_DMC_MOVEMENT_PRIMITIVE_ENVIRONMENTS["ProMP"].append("dmc_ball_in_cup-catch_promp-v0")
 
-kwargs_dict_reacher_easy_dmp = deepcopy(DEFAULT_BB_DICT_DMP)
-kwargs_dict_reacher_easy_dmp['name'] = f"dmc:reacher-easy"
-kwargs_dict_reacher_easy_dmp['wrappers'].append(suite.reacher.MPWrapper)
-# bandwidth_factor=2
-kwargs_dict_reacher_easy_dmp['phase_generator_kwargs']['alpha_phase'] = 2
-# TODO: weight scale 50, but goal scale 0.1
-kwargs_dict_reacher_easy_dmp['trajectory_generator_kwargs']['weight_scale'] = 500
 register(
-    id=f'dmc_reacher-easy_dmp-v0',
-    entry_point='fancy_gym.utils.make_env_helpers:make_bb_env_helper',
-    kwargs=kwargs_dict_bic_dmp
+    id=f"dm_control/reacher-hard-v0",
+    register_step_based=False,
+    mp_wrapper=suite.reacher.MPWrapper,
+    add_mp_types=['DMP', 'ProMP'],
 )
-ALL_DMC_MOVEMENT_PRIMITIVE_ENVIRONMENTS["DMP"].append("dmc_reacher-easy_dmp-v0")
-
-kwargs_dict_reacher_easy_promp = deepcopy(DEFAULT_BB_DICT_ProMP)
-kwargs_dict_reacher_easy_promp['name'] = f"dmc:reacher-easy"
-kwargs_dict_reacher_easy_promp['wrappers'].append(suite.reacher.MPWrapper)
-kwargs_dict_reacher_easy_promp['trajectory_generator_kwargs']['weight_scale'] = 0.2
-register(
-    id=f'dmc_reacher-easy_promp-v0',
-    entry_point='fancy_gym.utils.make_env_helpers:make_bb_env_helper',
-    kwargs=kwargs_dict_reacher_easy_promp
-)
-ALL_DMC_MOVEMENT_PRIMITIVE_ENVIRONMENTS["ProMP"].append("dmc_reacher-easy_promp-v0")
-
-kwargs_dict_reacher_hard_dmp = deepcopy(DEFAULT_BB_DICT_DMP)
-kwargs_dict_reacher_hard_dmp['name'] = f"dmc:reacher-hard"
-kwargs_dict_reacher_hard_dmp['wrappers'].append(suite.reacher.MPWrapper)
-# bandwidth_factor = 2
-kwargs_dict_reacher_hard_dmp['phase_generator_kwargs']['alpha_phase'] = 2
-# TODO: weight scale 50, but goal scale 0.1
-kwargs_dict_reacher_hard_dmp['trajectory_generator_kwargs']['weight_scale'] = 500
-register(
-    id=f'dmc_reacher-hard_dmp-v0',
-    entry_point='fancy_gym.utils.make_env_helpers:make_bb_env_helper',
-    kwargs=kwargs_dict_reacher_hard_dmp
-)
-ALL_DMC_MOVEMENT_PRIMITIVE_ENVIRONMENTS["DMP"].append("dmc_reacher-hard_dmp-v0")
-
-kwargs_dict_reacher_hard_promp = deepcopy(DEFAULT_BB_DICT_ProMP)
-kwargs_dict_reacher_hard_promp['name'] = f"dmc:reacher-hard"
-kwargs_dict_reacher_hard_promp['wrappers'].append(suite.reacher.MPWrapper)
-kwargs_dict_reacher_hard_promp['trajectory_generator_kwargs']['weight_scale'] = 0.2
-register(
-    id=f'dmc_reacher-hard_promp-v0',
-    entry_point='fancy_gym.utils.make_env_helpers:make_bb_env_helper',
-    kwargs=kwargs_dict_reacher_hard_promp
-)
-ALL_DMC_MOVEMENT_PRIMITIVE_ENVIRONMENTS["ProMP"].append("dmc_reacher-hard_promp-v0")
 
 _dmc_cartpole_tasks = ["balance", "balance_sparse", "swingup", "swingup_sparse"]
-
 for _task in _dmc_cartpole_tasks:
-    _env_id = f'dmc_cartpole-{_task}_dmp-v0'
-    kwargs_dict_cartpole_dmp = deepcopy(DEFAULT_BB_DICT_DMP)
-    kwargs_dict_cartpole_dmp['name'] = f"dmc:cartpole-{_task}"
-    kwargs_dict_cartpole_dmp['wrappers'].append(suite.cartpole.MPWrapper)
-    # bandwidth_factor = 2
-    kwargs_dict_cartpole_dmp['phase_generator_kwargs']['alpha_phase'] = 2
-    # TODO: weight scale 50, but goal scale 0.1
-    kwargs_dict_cartpole_dmp['trajectory_generator_kwargs']['weight_scale'] = 500
-    kwargs_dict_cartpole_dmp['controller_kwargs']['p_gains'] = 10
-    kwargs_dict_cartpole_dmp['controller_kwargs']['d_gains'] = 10
     register(
-        id=_env_id,
-        entry_point='fancy_gym.utils.make_env_helpers:make_bb_env_helper',
-        kwargs=kwargs_dict_cartpole_dmp
+        id=f'dm_control/cartpole-{_task}-v0',
+        register_step_based=False,
+        mp_wrapper=suite.cartpole.MPWrapper,
+        add_mp_types=['DMP', 'ProMP'],
     )
-    ALL_DMC_MOVEMENT_PRIMITIVE_ENVIRONMENTS["DMP"].append(_env_id)
 
-    _env_id = f'dmc_cartpole-{_task}_promp-v0'
-    kwargs_dict_cartpole_promp = deepcopy(DEFAULT_BB_DICT_ProMP)
-    kwargs_dict_cartpole_promp['name'] = f"dmc:cartpole-{_task}"
-    kwargs_dict_cartpole_promp['wrappers'].append(suite.cartpole.MPWrapper)
-    kwargs_dict_cartpole_promp['controller_kwargs']['p_gains'] = 10
-    kwargs_dict_cartpole_promp['controller_kwargs']['d_gains'] = 10
-    kwargs_dict_cartpole_promp['trajectory_generator_kwargs']['weight_scale'] = 0.2
 register(
-        id=_env_id,
-        entry_point='fancy_gym.utils.make_env_helpers:make_bb_env_helper',
-        kwargs=kwargs_dict_cartpole_promp
+    id=f"dm_control/cartpole-two_poles-v0",
+    register_step_based=False,
+    mp_wrapper=suite.cartpole.TwoPolesMPWrapper,
+    add_mp_types=['DMP', 'ProMP'],
 )
-    ALL_DMC_MOVEMENT_PRIMITIVE_ENVIRONMENTS["ProMP"].append(_env_id)
 
-kwargs_dict_cartpole2poles_dmp = deepcopy(DEFAULT_BB_DICT_DMP)
-kwargs_dict_cartpole2poles_dmp['name'] = f"dmc:cartpole-two_poles"
-kwargs_dict_cartpole2poles_dmp['wrappers'].append(suite.cartpole.TwoPolesMPWrapper)
-# bandwidth_factor = 2
-kwargs_dict_cartpole2poles_dmp['phase_generator_kwargs']['alpha_phase'] = 2
-# TODO: weight scale 50, but goal scale 0.1
-kwargs_dict_cartpole2poles_dmp['trajectory_generator_kwargs']['weight_scale'] = 500
-kwargs_dict_cartpole2poles_dmp['controller_kwargs']['p_gains'] = 10
-kwargs_dict_cartpole2poles_dmp['controller_kwargs']['d_gains'] = 10
-_env_id = f'dmc_cartpole-two_poles_dmp-v0'
 register(
-    id=_env_id,
-    entry_point='fancy_gym.utils.make_env_helpers:make_bb_env_helper',
-    kwargs=kwargs_dict_cartpole2poles_dmp
+    id=f"dm_control/cartpole-three_poles-v0",
+    register_step_based=False,
+    mp_wrapper=suite.cartpole.ThreePolesMPWrapper,
+    add_mp_types=['DMP', 'ProMP'],
 )
-ALL_DMC_MOVEMENT_PRIMITIVE_ENVIRONMENTS["DMP"].append(_env_id)
-
-kwargs_dict_cartpole2poles_promp = deepcopy(DEFAULT_BB_DICT_ProMP)
-kwargs_dict_cartpole2poles_promp['name'] = f"dmc:cartpole-two_poles"
-kwargs_dict_cartpole2poles_promp['wrappers'].append(suite.cartpole.TwoPolesMPWrapper)
-kwargs_dict_cartpole2poles_promp['controller_kwargs']['p_gains'] = 10
-kwargs_dict_cartpole2poles_promp['controller_kwargs']['d_gains'] = 10
-kwargs_dict_cartpole2poles_promp['trajectory_generator_kwargs']['weight_scale'] = 0.2
-_env_id = f'dmc_cartpole-two_poles_promp-v0'
-register(
-    id=_env_id,
-    entry_point='fancy_gym.utils.make_env_helpers:make_bb_env_helper',
-    kwargs=kwargs_dict_cartpole2poles_promp
-)
-ALL_DMC_MOVEMENT_PRIMITIVE_ENVIRONMENTS["ProMP"].append(_env_id)
-
-kwargs_dict_cartpole3poles_dmp = deepcopy(DEFAULT_BB_DICT_DMP)
-kwargs_dict_cartpole3poles_dmp['name'] = f"dmc:cartpole-three_poles"
-kwargs_dict_cartpole3poles_dmp['wrappers'].append(suite.cartpole.ThreePolesMPWrapper)
-# bandwidth_factor = 2
-kwargs_dict_cartpole3poles_dmp['phase_generator_kwargs']['alpha_phase'] = 2
-# TODO: weight scale 50, but goal scale 0.1
-kwargs_dict_cartpole3poles_dmp['trajectory_generator_kwargs']['weight_scale'] = 500
-kwargs_dict_cartpole3poles_dmp['controller_kwargs']['p_gains'] = 10
-kwargs_dict_cartpole3poles_dmp['controller_kwargs']['d_gains'] = 10
-_env_id = f'dmc_cartpole-three_poles_dmp-v0'
-register(
-    id=_env_id,
-    entry_point='fancy_gym.utils.make_env_helpers:make_bb_env_helper',
-    kwargs=kwargs_dict_cartpole3poles_dmp
-)
-ALL_DMC_MOVEMENT_PRIMITIVE_ENVIRONMENTS["DMP"].append(_env_id)
-
-kwargs_dict_cartpole3poles_promp = deepcopy(DEFAULT_BB_DICT_ProMP)
-kwargs_dict_cartpole3poles_promp['name'] = f"dmc:cartpole-three_poles"
-kwargs_dict_cartpole3poles_promp['wrappers'].append(suite.cartpole.ThreePolesMPWrapper)
-kwargs_dict_cartpole3poles_promp['controller_kwargs']['p_gains'] = 10
-kwargs_dict_cartpole3poles_promp['controller_kwargs']['d_gains'] = 10
-kwargs_dict_cartpole3poles_promp['trajectory_generator_kwargs']['weight_scale'] = 0.2
-_env_id = f'dmc_cartpole-three_poles_promp-v0'
-register(
-    id=_env_id,
-    entry_point='fancy_gym.utils.make_env_helpers:make_bb_env_helper',
-    kwargs=kwargs_dict_cartpole3poles_promp
-)
-ALL_DMC_MOVEMENT_PRIMITIVE_ENVIRONMENTS["ProMP"].append(_env_id)
 
 # DeepMind Manipulation
-kwargs_dict_mani_reach_site_features_dmp = deepcopy(DEFAULT_BB_DICT_DMP)
-kwargs_dict_mani_reach_site_features_dmp['name'] = f"dmc:manipulation-reach_site_features"
-kwargs_dict_mani_reach_site_features_dmp['wrappers'].append(manipulation.reach_site.MPWrapper)
-kwargs_dict_mani_reach_site_features_dmp['phase_generator_kwargs']['alpha_phase'] = 2
-# TODO: weight scale 50, but goal scale 0.1
-kwargs_dict_mani_reach_site_features_dmp['trajectory_generator_kwargs']['weight_scale'] = 500
-kwargs_dict_mani_reach_site_features_dmp['controller_kwargs']['controller_type'] = 'velocity'
 register(
-    id=f'dmc_manipulation-reach_site_dmp-v0',
-    entry_point='fancy_gym.utils.make_env_helpers:make_bb_env_helper',
-    kwargs=kwargs_dict_mani_reach_site_features_dmp
+    id=f"dm_control/reach_site_features-v0",
+    register_step_based=False,
+    mp_wrapper=manipulation.reach_site.MPWrapper,
+    add_mp_types=['DMP', 'ProMP'],
 )
-ALL_DMC_MOVEMENT_PRIMITIVE_ENVIRONMENTS["DMP"].append("dmc_manipulation-reach_site_dmp-v0")
-
-kwargs_dict_mani_reach_site_features_promp = deepcopy(DEFAULT_BB_DICT_ProMP)
-kwargs_dict_mani_reach_site_features_promp['name'] = f"dmc:manipulation-reach_site_features"
-kwargs_dict_mani_reach_site_features_promp['wrappers'].append(manipulation.reach_site.MPWrapper)
-kwargs_dict_mani_reach_site_features_promp['trajectory_generator_kwargs']['weight_scale'] = 0.2
-kwargs_dict_mani_reach_site_features_promp['controller_kwargs']['controller_type'] = 'velocity'
-register(
-    id=f'dmc_manipulation-reach_site_promp-v0',
-    entry_point='fancy_gym.utils.make_env_helpers:make_bb_env_helper',
-    kwargs=kwargs_dict_mani_reach_site_features_promp
-)
-ALL_DMC_MOVEMENT_PRIMITIVE_ENVIRONMENTS["ProMP"].append("dmc_manipulation-reach_site_promp-v0")

fancy_gym/dmc/dmc_wrapper.py

@@ -1,186 +0,0 @@
-# Adopted from: https://github.com/denisyarats/dmc2gym/blob/master/dmc2gym/wrappers.py
-# License: MIT
-# Copyright (c) 2020 Denis Yarats
-import collections
-from collections.abc import MutableMapping
-from typing import Any, Dict, Tuple, Optional, Union, Callable
-
-import gym
-import numpy as np
-from dm_control import composer
-from dm_control.rl import control
-from dm_env import specs
-from gym import spaces
-from gym.core import ObsType
-
-
-def _spec_to_box(spec):
-    def extract_min_max(s):
-        assert s.dtype == np.float64 or s.dtype == np.float32, \
-            f"Only float64 and float32 types are allowed, instead {s.dtype} was found"
-        dim = int(np.prod(s.shape))
-        if type(s) == specs.Array:
-            bound = np.inf * np.ones(dim, dtype=s.dtype)
-            return -bound, bound
-        elif type(s) == specs.BoundedArray:
-            zeros = np.zeros(dim, dtype=s.dtype)
-            return s.minimum + zeros, s.maximum + zeros
-
-    mins, maxs = [], []
-    for s in spec:
-        mn, mx = extract_min_max(s)
-        mins.append(mn)
-        maxs.append(mx)
-    low = np.concatenate(mins, axis=0)
-    high = np.concatenate(maxs, axis=0)
-    assert low.shape == high.shape
-    return spaces.Box(low, high, dtype=s.dtype)
-
-
-def _flatten_obs(obs: MutableMapping):
-    """
-    Flattens an observation of type MutableMapping, e.g. a dict to a 1D array.
-    Args:
-        obs: observation to flatten
-
-    Returns: 1D array of observation
-
-    """
-
-    if not isinstance(obs, MutableMapping):
-        raise ValueError(f'Requires dict-like observations structure. {type(obs)} found.')
-
-    # Keep key order consistent for non OrderedDicts
-    keys = obs.keys() if isinstance(obs, collections.OrderedDict) else sorted(obs.keys())
-
-    obs_vals = [np.array([obs[key]]) if np.isscalar(obs[key]) else obs[key].ravel() for key in keys]
-    return np.concatenate(obs_vals)
-
-
-class DMCWrapper(gym.Env):
-    def __init__(self,
-                 env: Callable[[], Union[composer.Environment, control.Environment]],
-                 ):
-
-        # TODO: Currently this is required to be a function because dmc does not allow to copy composers environments
-        self._env = env()
-
-        # action and observation space
-        self._action_space = _spec_to_box([self._env.action_spec()])
-        self._observation_space = _spec_to_box(self._env.observation_spec().values())
-
-        self._window = None
-        self.id = 'dmc'
-
-    def __getattr__(self, item):
-        """Propagate only non-existent properties to wrapped env."""
-        if item.startswith('_'):
-            raise AttributeError("attempted to get missing private attribute '{}'".format(item))
-        if item in self.__dict__:
-            return getattr(self, item)
-        return getattr(self._env, item)
-
-    def _get_obs(self, time_step):
-        obs = _flatten_obs(time_step.observation).astype(self.observation_space.dtype)
-        return obs
-
-    @property
-    def observation_space(self):
-        return self._observation_space
-
-    @property
-    def action_space(self):
-        return self._action_space
-
-    @property
-    def dt(self):
-        return self._env.control_timestep()
-
-    def seed(self, seed=None):
-        self._action_space.seed(seed)
-        self._observation_space.seed(seed)
-
-    def step(self, action) -> Tuple[np.ndarray, float, bool, Dict[str, Any]]:
-        assert self._action_space.contains(action)
-        extra = {'internal_state': self._env.physics.get_state().copy()}
-
-        time_step = self._env.step(action)
-        reward = time_step.reward or 0.
-        done = time_step.last()
-        obs = self._get_obs(time_step)
-        extra['discount'] = time_step.discount
-
-        return obs, reward, done, extra
-
-    def reset(self, *, seed: Optional[int] = None, return_info: bool = False,
-              options: Optional[dict] = None, ) -> Union[ObsType, Tuple[ObsType, dict]]:
-        time_step = self._env.reset()
-        obs = self._get_obs(time_step)
-        return obs
-
-    def render(self, mode='rgb_array', height=240, width=320, camera_id=-1, overlays=(), depth=False,
-               segmentation=False, scene_option=None, render_flag_overrides=None):
-
-        # assert mode == 'rgb_array', 'only support rgb_array mode, given %s' % mode
-        if mode == "rgb_array":
-            return self._env.physics.render(height=height, width=width, camera_id=camera_id, overlays=overlays,
-                                            depth=depth, segmentation=segmentation, scene_option=scene_option,
-                                            render_flag_overrides=render_flag_overrides)
-
-        # Render max available buffer size. Larger is only possible by altering the XML.
-        img = self._env.physics.render(height=self._env.physics.model.vis.global_.offheight,
-                                       width=self._env.physics.model.vis.global_.offwidth,
-                                       camera_id=camera_id, overlays=overlays, depth=depth, segmentation=segmentation,
-                                       scene_option=scene_option, render_flag_overrides=render_flag_overrides)
-
-        if depth:
-            img = np.dstack([img.astype(np.uint8)] * 3)
-
-        if mode == 'human':
-            try:
-                import cv2
-                if self._window is None:
-                    self._window = cv2.namedWindow(self.id, cv2.WINDOW_AUTOSIZE)
-                cv2.imshow(self.id, img[..., ::-1])  # Image in BGR
-                cv2.waitKey(1)
-            except ImportError:
-                raise gym.error.DependencyNotInstalled("Rendering requires opencv. Run `pip install opencv-python`")
-            # PYGAME seems to destroy some global rendering configs from the physics render
-            # except ImportError:
-            #     import pygame
-            #     img_copy = img.copy().transpose((1, 0, 2))
-            #     if self._window is None:
-            #         pygame.init()
-            #         pygame.display.init()
-            #         self._window = pygame.display.set_mode(img_copy.shape[:2])
-            #         self.clock = pygame.time.Clock()
-            #
-            #     surf = pygame.surfarray.make_surface(img_copy)
-            #     self._window.blit(surf, (0, 0))
-            #     pygame.event.pump()
-            #     self.clock.tick(30)
-            #     pygame.display.flip()
-
-    def close(self):
-        super().close()
-        if self._window is not None:
-            try:
-                import cv2
-                cv2.destroyWindow(self.id)
-            except ImportError:
-                import pygame
-
-                pygame.display.quit()
-                pygame.quit()
-
-    @property
-    def reward_range(self) -> Tuple[float, float]:
-        reward_spec = self._env.reward_spec()
-        if isinstance(reward_spec, specs.BoundedArray):
-            return reward_spec.minimum, reward_spec.maximum
-        return -float('inf'), float('inf')
-
-    @property
-    def metadata(self):
-        return {'render.modes': ['human', 'rgb_array'],
-                'video.frames_per_second': round(1.0 / self._env.control_timestep())}

fancy_gym/dmc/manipulation/reach_site/mp_wrapper.py

@@ -6,6 +6,28 @@ from fancy_gym.black_box.raw_interface_wrapper import RawInterfaceWrapper
 
 
 class MPWrapper(RawInterfaceWrapper):
+    mp_config = {
+        'ProMP': {
+            'controller_kwargs': {
+                'p_gains': 50.0,
+            },
+            'trajectory_generator_kwargs': {
+                'weights_scale': 0.2,
+            },
+        },
+        'DMP': {
+            'controller_kwargs': {
+                'p_gains': 50.0,
+            },
+            'phase_generator': {
+                'alpha_phase': 2,
+            },
+            'trajectory_generator_kwargs': {
+                'weights_scale': 500,
+            },
+        },
+        'ProDMP': {},
+    }
 
     @property
     def context_mask(self) -> np.ndarray:
@@ -35,4 +57,4 @@ class MPWrapper(RawInterfaceWrapper):
 
     @property
     def dt(self) -> Union[float, int]:
-        return self.env.dt
+        return self.env.control_timestep()

fancy_gym/dmc/suite/ball_in_cup/mp_wrapper.py

@@ -6,6 +6,25 @@ from fancy_gym.black_box.raw_interface_wrapper import RawInterfaceWrapper
 
 
 class MPWrapper(RawInterfaceWrapper):
+    mp_config = {
+        'ProMP': {
+            'controller_kwargs': {
+                'p_gains': 50.0,
+            },
+        },
+        'DMP': {
+            'controller_kwargs': {
+                'p_gains': 50.0,
+            },
+            'phase_generator': {
+                'alpha_phase': 2,
+            },
+            'trajectory_generator_kwargs': {
+                'weights_scale': 10
+            },
+        },
+        'ProDMP': {},
+    }
 
     @property
     def context_mask(self) -> np.ndarray:
@@ -31,4 +50,4 @@ class MPWrapper(RawInterfaceWrapper):
 
     @property
     def dt(self) -> Union[float, int]:
-        return self.env.dt
+        return self.env.control_timestep()

fancy_gym/dmc/suite/cartpole/mp_wrapper.py

@@ -6,6 +6,30 @@ from fancy_gym.black_box.raw_interface_wrapper import RawInterfaceWrapper
 
 
 class MPWrapper(RawInterfaceWrapper):
+    mp_config = {
+        'ProMP': {
+            'controller_kwargs': {
+                'p_gains': 10,
+                'd_gains': 10,
+            },
+            'trajectory_generator_kwargs': {
+                'weights_scale': 0.2,
+            },
+        },
+        'DMP': {
+            'controller_kwargs': {
+                'p_gains': 10,
+                'd_gains': 10,
+            },
+            'phase_generator': {
+                'alpha_phase': 2,
+            },
+            'trajectory_generator_kwargs': {
+                'weights_scale': 500,
+            },
+        },
+        'ProDMP': {},
+    }
 
     def __init__(self, env, n_poles: int = 1):
         self.n_poles = n_poles
@@ -35,7 +59,7 @@ class MPWrapper(RawInterfaceWrapper):
 
     @property
     def dt(self) -> Union[float, int]:
-        return self.env.dt
+        return self.env.control_timestep()
 
 
 class TwoPolesMPWrapper(MPWrapper):

fancy_gym/dmc/suite/reacher/mp_wrapper.py

@@ -6,6 +6,30 @@ from fancy_gym.black_box.raw_interface_wrapper import RawInterfaceWrapper
 
 
 class MPWrapper(RawInterfaceWrapper):
+    mp_config = {
+        'ProMP': {
+            'controller_kwargs': {
+                'p_gains': 50.0,
+                'd_gains': 1.0,
+            },
+            'trajectory_generator_kwargs': {
+                'weights_scale': 0.2,
+            },
+        },
+        'DMP': {
+            'controller_kwargs': {
+                'p_gains': 50.0,
+                'd_gains': 1.0,
+            },
+            'phase_generator': {
+                'alpha_phase': 2,
+            },
+            'trajectory_generator_kwargs': {
+                'weights_scale': 500,
+            },
+        },
+        'ProDMP': {},
+    }
 
     @property
     def context_mask(self) -> np.ndarray:
@@ -30,4 +54,4 @@ class MPWrapper(RawInterfaceWrapper):
 
     @property
     def dt(self) -> Union[float, int]:
-        return self.env.dt
+        return self.env.control_timestep()

fancy_gym/envs/classic_control/README.MD

@@ -1,18 +1,20 @@
 ### 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
+
+| Name                         | Description                                                                                                                                                                                                         | Horizon | Action Dimension | Observation Dimension |
+| ---------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------- | ---------------- | --------------------- |
+| `fancy/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                     |
+| `fancy/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                    |
+| `fancy/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                    |
+| `fancy/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
 
-[//]:  |`HoleReacherProMPP-v0`|
+| Name                                | Description                                                                                              | Horizon | Action Dimension | Context Dimension |
+| ----------------------------------- | -------------------------------------------------------------------------------------------------------- | ------- | ---------------- | ----------------- |
+| `fancy_DMP/ViaPointReacher-v0`      | A DMP provides a trajectory for the `fancy/ViaPointReacher-v0` task.                                     | 200     | 25               |
+| `fancy_DMP/HoleReacherFixedGoal-v0` | A DMP provides a trajectory for the `fancy/HoleReacher-v0` task with a fixed goal attractor.             | 200     | 25               |
+| `fancy_DMP/HoleReacher-v0`          | A DMP provides a trajectory for the `fancy/HoleReacher-v0` task. The goal attractor needs to be learned. | 200     | 30               |
+
+[//]: |`fancy/HoleReacherProMPP-v0`|

fancy_gym/envs/classic_control/base_reacher/base_reacher.py

@@ -1,10 +1,10 @@
-from typing import Union, Tuple, Optional
+from typing import Union, Tuple, Optional, Any, Dict
 
-import gym
+import gymnasium as gym
 import numpy as np
-from gym import spaces
-from gym.core import ObsType
-from gym.utils import seeding
+from gymnasium import spaces
+from gymnasium.core import ObsType
+from gymnasium.utils import seeding
 
 from fancy_gym.envs.classic_control.utils import intersect
 
@@ -14,12 +14,14 @@ class BaseReacherEnv(gym.Env):
     Base class for all reaching environments.
     """
 
-    def __init__(self, n_links: int, random_start: bool = True, allow_self_collision: bool = False):
+    def __init__(self, n_links: int, random_start: bool = True, allow_self_collision: bool = False, render_mode: str = None):
         super().__init__()
         self.link_lengths = np.ones(n_links)
         self.n_links = n_links
         self._dt = 0.01
 
+        self.render_mode = render_mode
+
         self.random_start = random_start
 
         self.allow_self_collision = allow_self_collision
@@ -55,7 +57,6 @@ class BaseReacherEnv(gym.Env):
         self.fig = None
 
         self._steps = 0
-        self.seed()
 
     @property
     def dt(self) -> Union[float, int]:
@@ -69,10 +70,15 @@ class BaseReacherEnv(gym.Env):
     def current_vel(self):
         return self._angle_velocity.copy()
 
-    def reset(self, *, seed: Optional[int] = None, return_info: bool = False,
-              options: Optional[dict] = None, ) -> Union[ObsType, Tuple[ObsType, dict]]:
+    def reset(self, *, seed: Optional[int] = None, options: Optional[Dict[str, Any]] = None) \
+            -> Tuple[ObsType, Dict[str, Any]]:
         # Sample only orientation of first link, i.e. the arm is always straight.
-        if self.random_start:
+        super(BaseReacherEnv, self).reset(seed=seed, options=options)
+        try:
+            random_start = options.get('random_start', self.random_start)
+        except AttributeError:
+            random_start = self.random_start
+        if random_start:
             first_joint = self.np_random.uniform(np.pi / 4, 3 * np.pi / 4)
             self._joint_angles = np.hstack([[first_joint], np.zeros(self.n_links - 1)])
             self._start_pos = self._joint_angles.copy()
@@ -84,7 +90,7 @@ class BaseReacherEnv(gym.Env):
         self._update_joints()
         self._steps = 0
 
-        return self._get_obs().copy()
+        return self._get_obs().copy(), {}
 
     def _update_joints(self):
         """
@@ -124,10 +130,6 @@ class BaseReacherEnv(gym.Env):
     def _terminate(self, info) -> bool:
         raise NotImplementedError
 
-    def seed(self, seed=None):
-        self.np_random, seed = seeding.np_random(seed)
-        return [seed]
-
     def close(self):
         super(BaseReacherEnv, self).close()
         del self.fig

fancy_gym/envs/classic_control/base_reacher/base_reacher_direct.py

@@ -1,5 +1,5 @@
 import numpy as np
-from gym import spaces
+from gymnasium import spaces
 
 from fancy_gym.envs.classic_control.base_reacher.base_reacher import BaseReacherEnv
 
@@ -10,8 +10,8 @@ class BaseReacherDirectEnv(BaseReacherEnv):
     """
 
     def __init__(self, n_links: int, random_start: bool = True,
-                 allow_self_collision: bool = False):
-        super().__init__(n_links, random_start, allow_self_collision)
+                 allow_self_collision: bool = False, **kwargs):
+        super().__init__(n_links, random_start, allow_self_collision, **kwargs)
 
         self.max_vel = 2 * np.pi
         action_bound = np.ones((self.n_links,)) * self.max_vel
@@ -32,6 +32,7 @@ class BaseReacherDirectEnv(BaseReacherEnv):
         reward, info = self._get_reward(action)
 
         self._steps += 1
-        done = self._terminate(info)
+        terminated = self._terminate(info)
+        truncated = False
 
-        return self._get_obs().copy(), reward, done, info
+        return self._get_obs().copy(), reward, terminated, truncated, info

fancy_gym/envs/classic_control/base_reacher/base_reacher_torque.py

@@ -1,5 +1,5 @@
 import numpy as np
-from gym import spaces
+from gymnasium import spaces
 
 from fancy_gym.envs.classic_control.base_reacher.base_reacher import BaseReacherEnv
 
@@ -10,8 +10,8 @@ class BaseReacherTorqueEnv(BaseReacherEnv):
     """
 
     def __init__(self, n_links: int, random_start: bool = True,
-                 allow_self_collision: bool = False):
-        super().__init__(n_links, random_start, allow_self_collision)
+                 allow_self_collision: bool = False, **kwargs):
+        super().__init__(n_links, random_start, allow_self_collision, **kwargs)
 
         self.max_torque = 1000
         action_bound = np.ones((self.n_links,)) * self.max_torque
@@ -31,6 +31,7 @@ class BaseReacherTorqueEnv(BaseReacherEnv):
         reward, info = self._get_reward(action)
 
         self._steps += 1
-        done = False
+        terminated = False
+        truncated = False
 
-        return self._get_obs().copy(), reward, done, info
+        return self._get_obs().copy(), reward, terminated, truncated, info

fancy_gym/envs/classic_control/hole_reacher/hole_reacher.py

@@ -1,22 +1,25 @@
-from typing import Union, Optional, Tuple
+from typing import Union, Optional, Tuple, Any, Dict
 
-import gym
+import gymnasium as gym
 import matplotlib.pyplot as plt
 import numpy as np
-from gym.core import ObsType
+from gymnasium import spaces
+from gymnasium.core import ObsType
 from matplotlib import patches
 
 from fancy_gym.envs.classic_control.base_reacher.base_reacher_direct import BaseReacherDirectEnv
+from . import MPWrapper
 
 MAX_EPISODE_STEPS_HOLEREACHER = 200
 
 
 class HoleReacherEnv(BaseReacherDirectEnv):
+
     def __init__(self, n_links: int, hole_x: Union[None, float] = None, hole_depth: Union[None, float] = None,
                  hole_width: float = 1., random_start: bool = False, allow_self_collision: bool = False,
-                 allow_wall_collision: bool = False, collision_penalty: float = 1000, rew_fct: str = "simple"):
+                 allow_wall_collision: bool = False, collision_penalty: float = 1000, rew_fct: str = "simple", **kwargs):
 
-        super().__init__(n_links, random_start, allow_self_collision)
+        super().__init__(n_links, random_start, allow_self_collision, **kwargs)
 
         # provided initial parameters
         self.initial_x = hole_x  # x-position of center of hole
@@ -40,7 +43,7 @@ class HoleReacherEnv(BaseReacherDirectEnv):
             [np.inf]  # env steps, because reward start after n steps TODO: Maybe
         ])
         # self.action_space = gym.spaces.Box(low=-action_bound, high=action_bound, shape=action_bound.shape)
-        self.observation_space = gym.spaces.Box(low=-state_bound, high=state_bound, shape=state_bound.shape)
+        self.observation_space = spaces.Box(low=-state_bound, high=state_bound, shape=state_bound.shape)
 
         if rew_fct == "simple":
             from fancy_gym.envs.classic_control.hole_reacher.hr_simple_reward import HolereacherReward
@@ -54,13 +57,18 @@ class HoleReacherEnv(BaseReacherDirectEnv):
         else:
             raise ValueError("Unknown reward function {}".format(rew_fct))
 
-    def reset(self, *, seed: Optional[int] = None, return_info: bool = False,
-              options: Optional[dict] = None, ) -> Union[ObsType, Tuple[ObsType, dict]]:
+    def reset(self, *, seed: Optional[int] = None, options: Optional[Dict[str, Any]] = None) \
+            -> Tuple[ObsType, Dict[str, Any]]:
+
+        # initialize seed here as the random goal needs to be generated before the super reset()
+        gym.Env.reset(self, seed=seed, options=options)
+
         self._generate_hole()
         self._set_patches()
         self.reward_function.reset()
 
-        return super().reset()
+        # do not provide seed to avoid setting it twice
+        return super(HoleReacherEnv, self).reset(options=options)
 
     def _get_reward(self, action: np.ndarray) -> (float, dict):
         return self.reward_function.get_reward(self)
@@ -170,7 +178,7 @@ class HoleReacherEnv(BaseReacherDirectEnv):
 
         return False
 
-    def render(self, mode='human'):
+    def render(self):
         if self.fig is None:
             # Create base figure once on the beginning. Afterwards only update
             plt.ion()
@@ -189,7 +197,7 @@ class HoleReacherEnv(BaseReacherDirectEnv):
         self.fig.gca().set_title(
             f"Iteration: {self._steps}, distance: {np.linalg.norm(self.end_effector - self._goal) ** 2}")
 
-        if mode == "human":
+        if self.render_mode == "human":
 
             # arm
             self.line.set_data(self._joints[:, 0], self._joints[:, 1])
@@ -197,7 +205,7 @@ class HoleReacherEnv(BaseReacherDirectEnv):
             self.fig.canvas.draw()
             self.fig.canvas.flush_events()
 
-        elif mode == "partial":
+        elif self.render_mode == "partial":
             if self._steps % 20 == 0 or self._steps in [1, 199] or self._is_collided:
                 # Arm
                 plt.plot(self._joints[:, 0], self._joints[:, 1], 'ro-', markerfacecolor='k',
@@ -223,16 +231,3 @@ class HoleReacherEnv(BaseReacherDirectEnv):
             self.fig.gca().add_patch(left_block)
             self.fig.gca().add_patch(right_block)
             self.fig.gca().add_patch(hole_floor)
-
-
-if __name__ == "__main__":
-
-    env = HoleReacherEnv(5)
-    env.reset()
-
-    for i in range(10000):
-        ac = env.action_space.sample()
-        obs, rew, done, info = env.step(ac)
-        env.render()
-        if done:
-            env.reset()

fancy_gym/envs/classic_control/hole_reacher/mp_wrapper.py

@@ -7,6 +7,30 @@ from fancy_gym.black_box.raw_interface_wrapper import RawInterfaceWrapper
 
 class MPWrapper(RawInterfaceWrapper):
 
+    mp_config = {
+        'ProMP': {
+            'controller_kwargs': {
+                'controller_type': 'velocity',
+            },
+            'trajectory_generator_kwargs': {
+                'weights_scale': 2,
+            },
+        },
+        'DMP': {
+            'controller_kwargs': {
+                'controller_type': 'velocity',
+            },
+            'trajectory_generator_kwargs': {
+                # TODO: Before it was weight scale 50 and goal scale 0.1. We now only have weight scale and thus set it to 500. Check
+                'weights_scale': 500,
+            },
+            'phase_generator_kwargs': {
+                'alpha_phase': 2.5,
+            },
+        },
+        'ProDMP': {},
+    }
+
     @property
     def context_mask(self):
         return np.hstack([

fancy_gym/envs/classic_control/simple_reacher/mp_wrapper.py

@@ -7,6 +7,28 @@ from fancy_gym.black_box.raw_interface_wrapper import RawInterfaceWrapper
 
 class MPWrapper(RawInterfaceWrapper):
 
+    mp_config = {
+        'ProMP': {
+            'controller_kwargs': {
+                'p_gains': 0.6,
+                'd_gains': 0.075,
+            },
+        },
+        'DMP': {
+            'controller_kwargs': {
+                'p_gains': 0.6,
+                'd_gains': 0.075,
+            },
+            'trajectory_generator_kwargs': {
+                'weights_scale': 50,
+            },
+            'phase_generator_kwargs': {
+                'alpha_phase': 2,
+            },
+        },
+        'ProDMP': {},
+    }
+
     @property
     def context_mask(self):
         return np.hstack([

fancy_gym/envs/classic_control/simple_reacher/simple_reacher.py

@@ -1,11 +1,12 @@
-from typing import Iterable, Union, Optional, Tuple
+from typing import Iterable, Union, Optional, Tuple, Any, Dict
 
 import matplotlib.pyplot as plt
 import numpy as np
-from gym import spaces
-from gym.core import ObsType
+from gymnasium import spaces
+from gymnasium.core import ObsType
 
 from fancy_gym.envs.classic_control.base_reacher.base_reacher_torque import BaseReacherTorqueEnv
+from . import MPWrapper
 
 
 class SimpleReacherEnv(BaseReacherTorqueEnv):
@@ -16,8 +17,8 @@ class SimpleReacherEnv(BaseReacherTorqueEnv):
     """
 
     def __init__(self, n_links: int, target: Union[None, Iterable] = None, random_start: bool = True,
-                 allow_self_collision: bool = False, ):
-        super().__init__(n_links, random_start, allow_self_collision)
+                 allow_self_collision: bool = False, **kwargs):
+        super().__init__(n_links, random_start, allow_self_collision, **kwargs)
 
         # provided initial parameters
         self.inital_target = target
@@ -42,11 +43,15 @@ class SimpleReacherEnv(BaseReacherTorqueEnv):
     # def start_pos(self):
     #     return self._start_pos
 
-    def reset(self, *, seed: Optional[int] = None, return_info: bool = False,
-              options: Optional[dict] = None, ) -> Union[ObsType, Tuple[ObsType, dict]]:
+    def reset(self, *, seed: Optional[int] = None, options: Optional[Dict[str, Any]] = None) \
+            -> Tuple[ObsType, Dict[str, Any]]:
+        # Reset twice to ensure we return obs after generating goal and generating goal after executing seeded reset.
+        # (Env will not behave deterministic otherwise)
+        # Yes, there is probably a more elegant solution to this problem...
         self._generate_goal()
-
-        return super().reset()
+        super().reset(seed=seed, options=options)
+        self._generate_goal()
+        return super().reset(seed=seed, options=options)
 
     def _get_reward(self, action: np.ndarray):
         diff = self.end_effector - self._goal
@@ -93,7 +98,7 @@ class SimpleReacherEnv(BaseReacherTorqueEnv):
     def _check_collisions(self) -> bool:
         return self._check_self_collision()
 
-    def render(self, mode='human'):  # pragma: no cover
+    def render(self):  # pragma: no cover
         if self.fig is None:
             # Create base figure once on the beginning. Afterwards only update
             plt.ion()
@@ -127,15 +132,3 @@ class SimpleReacherEnv(BaseReacherTorqueEnv):
 
         self.fig.canvas.draw()
         self.fig.canvas.flush_events()
-
-
-if __name__ == "__main__":
-    env = SimpleReacherEnv(5)
-    env.reset()
-    for i in range(200):
-        ac = env.action_space.sample()
-        obs, rew, done, info = env.step(ac)
-
-        env.render()
-        if done:
-            break

fancy_gym/envs/classic_control/viapoint_reacher/mp_wrapper.py

@@ -7,6 +7,26 @@ from fancy_gym.black_box.raw_interface_wrapper import RawInterfaceWrapper
 
 class MPWrapper(RawInterfaceWrapper):
 
+    mp_config = {
+        'ProMP': {
+            'controller_kwargs': {
+                'controller_type': 'velocity',
+            },
+        },
+        'DMP': {
+            'controller_kwargs': {
+                'controller_type': 'velocity',
+            },
+            'trajectory_generator_kwargs': {
+                'weights_scale': 50,
+            },
+            'phase_generator_kwargs': {
+                'alpha_phase': 2,
+            },
+        },
+        'ProDMP': {},
+    }
+
     @property
     def context_mask(self):
         return np.hstack([

fancy_gym/envs/classic_control/viapoint_reacher/viapoint_reacher.py

@@ -1,19 +1,21 @@
-from typing import Iterable, Union, Tuple, Optional
+from typing import Iterable, Union, Tuple, Optional, Any, Dict
 
-import gym
+import gymnasium as gym
 import matplotlib.pyplot as plt
 import numpy as np
-from gym.core import ObsType
+from gymnasium import spaces
+from gymnasium.core import ObsType
 
 from fancy_gym.envs.classic_control.base_reacher.base_reacher_direct import BaseReacherDirectEnv
+from . import MPWrapper
 
 
 class ViaPointReacherEnv(BaseReacherDirectEnv):
 
     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):
+                 target: Union[None, Iterable] = None, allow_self_collision=False, collision_penalty=1000, **kwargs):
 
-        super().__init__(n_links, random_start, allow_self_collision)
+        super().__init__(n_links, random_start, allow_self_collision, **kwargs)
 
         # provided initial parameters
         self.intitial_target = target  # provided target value
@@ -34,16 +36,21 @@ class ViaPointReacherEnv(BaseReacherDirectEnv):
             [np.inf] * 2,  # x-y coordinates of target distance
             [np.inf]  # env steps, because reward start after n steps
         ])
-        self.observation_space = gym.spaces.Box(low=-state_bound, high=state_bound, shape=state_bound.shape)
+        self.observation_space = spaces.Box(low=-state_bound, high=state_bound, shape=state_bound.shape)
 
     # @property
     # def start_pos(self):
     #     return self._start_pos
 
-    def reset(self, *, seed: Optional[int] = None, return_info: bool = False,
-              options: Optional[dict] = None, ) -> Union[ObsType, Tuple[ObsType, dict]]:
+    def reset(self, *, seed: Optional[int] = None, options: Optional[Dict[str, Any]] = None) \
+            -> Tuple[ObsType, Dict[str, Any]]:
+        # Reset twice to ensure we return obs after generating goal and generating goal after executing seeded reset.
+        # (Env will not behave deterministic otherwise)
+        # Yes, there is probably a more elegant solution to this problem...
         self._generate_goal()
-        return super().reset()
+        super().reset(seed=seed, options=options)
+        self._generate_goal()
+        return super().reset(seed=seed, options=options)
 
     def _generate_goal(self):
         # TODO: Maybe improve this later, this can yield quite a lot of invalid settings
@@ -116,7 +123,7 @@ class ViaPointReacherEnv(BaseReacherDirectEnv):
     def _check_collisions(self) -> bool:
         return self._check_self_collision()
 
-    def render(self, mode='human'):
+    def render(self):
         goal_pos = self._goal.T
         via_pos = self._via_point.T
 
@@ -139,7 +146,7 @@ class ViaPointReacherEnv(BaseReacherDirectEnv):
 
         self.fig.gca().set_title(f"Iteration: {self._steps}, distance: {self.end_effector - self._goal}")
 
-        if mode == "human":
+        if self.render_mode == "human":
             # goal
             if self._steps == 1:
                 self.goal_point_plot.set_data(goal_pos[0], goal_pos[1])
@@ -151,7 +158,7 @@ class ViaPointReacherEnv(BaseReacherDirectEnv):
             self.fig.canvas.draw()
             self.fig.canvas.flush_events()
 
-        elif mode == "partial":
+        elif self.render_mode == "partial":
             if self._steps == 1:
                 # fig, ax = plt.subplots()
                 # Add the patch to the Axes
@@ -171,7 +178,7 @@ class ViaPointReacherEnv(BaseReacherDirectEnv):
                 plt.ylim([-1.1, lim])
                 plt.pause(0.01)
 
-        elif mode == "final":
+        elif self.render_mode == "final":
             if self._steps == 199 or self._is_collided:
                 # fig, ax = plt.subplots()
 
@@ -183,16 +190,3 @@ class ViaPointReacherEnv(BaseReacherDirectEnv):
                 plt.plot(self._joints[:, 0], self._joints[:, 1], 'ro-', markerfacecolor='k')
 
                 plt.pause(0.01)
-
-
-if __name__ == "__main__":
-
-    env = ViaPointReacherEnv(5)
-    env.reset()
-
-    for i in range(10000):
-        ac = env.action_space.sample()
-        obs, rew, done, info = env.step(ac)
-        env.render()
-        if done:
-            env.reset()

fancy_gym/envs/mujoco/README.MD

@@ -1,15 +1,48 @@
 # 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
 
+| Name                                       | Description                                                                                        | Horizon | Action Dimension | Observation Dimension |
+| ------------------------------------------ | -------------------------------------------------------------------------------------------------- | ------- | ---------------- | --------------------- |
+| `fancy/Reacher-v0`                         | Modified (5 links) gymnasiums's mujoco `Reacher-v2` (2 links)                                      | 200     | 5                | 21                    |
+| `fancy/ReacherSparse-v0`                   | Same as `fancy/Reacher-v0`, but the distance penalty is only provided in the last time step.       | 200     | 5                | 21                    |
+| `fancy/ReacherSparseBalanced-v0`           | Same as `fancy/ReacherSparse-v0`, but the end-effector has to remain upright.                      | 200     | 5                | 21                    |
+| `fancy/LongReacher-v0`                     | Modified (7 links) gymnasiums's mujoco `Reacher-v2` (2 links)                                      | 200     | 7                | 27                    |
+| `fancy/LongReacherSparse-v0`               | Same as `fancy/LongReacher-v0`, but the distance penalty is only provided in the last time step.   | 200     | 7                | 27                    |
+| `fancy/LongReacherSparseBalanced-v0`       | Same as `fancy/LongReacherSparse-v0`, but the end-effector has to remain upright.                  | 200     | 7                | 27                    |
+| `fancy/Reacher5d-v0`                       | Reacher task with 5 links, based on Gymnasium's `gym.envs.mujoco.ReacherEnv`                       | 200     | 5                | 20                    |
+| `fancy/Reacher5dSparse-v0`                 | Sparse Reacher task with 5 links, based on Gymnasium's `gym.envs.mujoco.ReacherEnv`                | 200     | 5                | 20                    |
+| `fancy/Reacher7d-v0`                       | Reacher task with 7 links, based on Gymnasium's `gym.envs.mujoco.ReacherEnv`                       | 200     | 7                | 22                    |
+| `fancy/Reacher7dSparse-v0`                 | Sparse Reacher task with 7 links, based on Gymnasium's `gym.envs.mujoco.ReacherEnv`                | 200     | 7                | 22                    |
+| `fancy/HopperJumpSparse-v0`                | Hopper Jump task with sparse rewards, based on Gymnasium's `gym.envs.mujoco.Hopper`                | 250     | 3                | 15 / 16\*             |
+| `fancy/HopperJump-v0`                      | Hopper Jump task with continuous rewards, based on Gymnasium's `gym.envs.mujoco.Hopper`            | 250     | 3                | 15 / 16\*             |
+| `fancy/AntJump-v0`                         | Ant Jump task, based on Gymnasium's `gym.envs.mujoco.Ant`                                          | 200     | 8                | 119                   |
+| `fancy/HalfCheetahJump-v0`                 | HalfCheetah Jump task, based on Gymnasium's `gym.envs.mujoco.HalfCheetah`                          | 100     | 6                | 112                   |
+| `fancy/HopperJumpOnBox-v0`                 | Hopper Jump on Box task, based on Gymnasium's `gym.envs.mujoco.Hopper`                             | 250     | 4                | 16 / 100\*            |
+| `fancy/HopperThrow-v0`                     | Hopper Throw task, based on Gymnasium's `gym.envs.mujoco.Hopper`                                   | 250     | 3                | 18 / 100\*            |
+| `fancy/HopperThrowInBasket-v0`             | Hopper Throw in Basket task, based on Gymnasium's `gym.envs.mujoco.Hopper`                         | 250     | 3                | 18 / 100\*            |
+| `fancy/Walker2DJump-v0`                    | Walker 2D Jump task, based on Gymnasium's `gym.envs.mujoco.Walker2d`                               | 300     | 6                | 18 / 19\*             |
+| `fancy/BeerPong-v0`                        | Beer Pong task, based on a custom environment with multiple task variations                        | 300     | 3                | 29                    |
+| `fancy/BeerPongStepBased-v0`               | Step-based Beer Pong task, based on a custom environment with episodic rewards                     | 300     | 3                | 29                    |
+| `fancy/BeerPongFixedRelease-v0`            | Beer Pong with fixed release, based on a custom environment with episodic rewards                  | 300     | 3                | 29                    |
+| `fancy/BoxPushingDense-v0`                 | Custom Box-pushing task with dense rewards                                                         | 100     | 3                | 13                    |
+| `fancy/BoxPushingTemporalSparse-v0`        | Custom Box-pushing task with temporally sparse rewards                                             | 100     | 3                | 13                    |
+| `fancy/BoxPushingTemporalSpatialSparse-v0` | Custom Box-pushing task with temporally and spatially sparse rewards                               | 100     | 3                | 13                    |
+| `fancy/TableTennis2D-v0`                   | Table Tennis task with 2D context, based on a custom environment for table tennis                  | 350     | 7                | 19                    |
+| `fancy/TableTennis2DReplan-v0`             | Table Tennis task with 2D context and replanning, based on a custom environment for table tennis   | 350     | 7                | 19                    |
+| `fancy/TableTennis4D-v0`                   | Table Tennis task with 4D context, based on a custom environment for table tennis                  | 350     | 7                | 22                    |
+| `fancy/TableTennis4DReplan-v0`             | Table Tennis task with 4D context and replanning, based on a custom environment for table tennis   | 350     | 7                | 22                    |
+| `fancy/TableTennisWind-v0`                 | Table Tennis task with wind effects, based on a custom environment for table tennis                | 350     | 7                | 19                    |
+| `fancy/TableTennisGoalSwitching-v0`        | Table Tennis task with goal switching, based on a custom environment for table tennis              | 350     | 7                | 19                    |
+| `fancy/TableTennisWindReplan-v0`           | Table Tennis task with wind effects and replanning, based on a custom environment for table tennis | 350     | 7                | 19                    |
+
+\*Observation dimensions depend on configuration.
+
+<!--
+No longer used?
+| Name                        | Description                                                                                         | Horizon | Action Dimension | Observation Dimension |
+| --------------------------- | --------------------------------------------------------------------------------------------------- | ------- | ---------------- | --------------------- |
+| `fancy/BallInACupSimple-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                   |
+| `fancy/BallInACup-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                   |
+| `fancy/BallInACupGoal-v0`   | Similar to `fancy/BallInACupSimple-v0` but the ball needs to be caught at a specified goal position | 4000    | 7                | wip                   |
+-->

fancy_gym/envs/mujoco/__init__.py

@@ -9,3 +9,8 @@ from .reacher.reacher import ReacherEnv
 from .walker_2d_jump.walker_2d_jump import Walker2dJumpEnv
 from .box_pushing.box_pushing_env import BoxPushingDense, BoxPushingTemporalSparse, BoxPushingTemporalSpatialSparse
 from .table_tennis.table_tennis_env import TableTennisEnv, TableTennisWind, TableTennisGoalSwitching
+
+try:
+    from .air_hockey.air_hockey_env_wrapper import AirHockeyEnv
+except ModuleNotFoundError:
+    print("[FANCY GYM] Air Hockey not available (depends on mushroom-rl, dmc, mujoco)")

fancy_gym/envs/mujoco/air_hockey/LICENSE

@@ -0,0 +1,26 @@
+MIT License
+
+Copyright (c) 2022 Puze Liu, Jonas Guenster, Davide Tateo.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS," WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+---
+
+This project is a derivative of [AirHockeyChallenge](https://github.com/AirHockeyChallenge/air_hockey_challenge).
+The changes are mostly focused on adapting the provided environments to fancy_gym.

fancy_gym/envs/mujoco/air_hockey/air_hockey_env_wrapper.py

@@ -0,0 +1,202 @@
+from copy import deepcopy
+import numpy as np
+from gymnasium import spaces
+
+import fancy_gym.envs.mujoco.air_hockey.constraints as constraints
+from fancy_gym.envs.mujoco.air_hockey import position_control_wrapper as position
+from fancy_gym.envs.mujoco.air_hockey.utils import robot_to_world
+from mushroom_rl.core import Environment
+
+class AirHockeyEnv(Environment):
+    metadata = {"render_modes": ["human", "rgb_array"], "render_fps": 50}
+
+    def __init__(self, env_mode=None, interpolation_order=3, render_mode=None, width=1920, height=1080, **kwargs):
+        """
+        Environment Constructor
+
+        Args:
+            env [string]:
+                The string to specify the running environments. Available environments: [3dof-hit, 3dof-defend, 7dof-hit, 7dof-defend, tournament].
+            interpolation_order (int, 3): Type of interpolation used, has to correspond to action shape. Order 1-5 are
+                    polynomial interpolation of the degree. Order -1 is linear interpolation of position and velocity.
+                    Set Order to None in order to turn off interpolation. In this case the action has to be a trajectory
+                    of position, velocity and acceleration of the shape (20, 3, n_joints)
+        """
+
+        env_dict = {
+            "tournament": position.IiwaPositionTournament,
+
+            "7dof-hit": position.IiwaPositionHit,
+            "7dof-defend": position.IiwaPositionDefend,
+
+            "3dof-hit": position.PlanarPositionHit,
+            "3dof-defend": position.PlanarPositionDefend,
+
+            "7dof-hit-airhockit2023": position.IiwaPositionHitAirhocKIT2023,
+            "7dof-defend-airhockit2023": position.IiwaPositionDefendAirhocKIT2023,
+        }
+
+        if env_mode not in env_dict:
+            raise Exception(f"Please specify one of the environments in {list(env_dict.keys())} for env_mode parameter!")
+
+        if env_mode == "tournament" and type(interpolation_order) != tuple:
+            interpolation_order = (interpolation_order, interpolation_order)
+
+        self.render_mode = render_mode
+        self.render_human_active = False
+
+        # Determine headless mode based on render_mode
+        headless = self.render_mode == 'rgb_array'
+        
+        # Prepare viewer_params
+        viewer_params = kwargs.get('viewer_params', {})
+        viewer_params.update({'headless': headless, 'width': width, 'height': height})
+        kwargs['viewer_params'] = viewer_params
+
+        self.base_env = env_dict[env_mode](interpolation_order=interpolation_order, **kwargs)
+        self.env_name = env_mode
+        self.env_info = self.base_env.env_info
+
+        if hasattr(self.base_env, "wrapper_obs_space") and hasattr(self.base_env, "wrapper_act_space"):
+            self.observation_space = self.base_env.wrapper_obs_space
+            self.action_space = self.base_env.wrapper_act_space
+        else:
+            single_robot_obs_size = len(self.base_env.info.observation_space.low)
+            if env_mode == "tournament":
+                self.observation_space = spaces.Box(low=-np.inf, high=np.inf, shape=(2,single_robot_obs_size), dtype=np.float64)
+            else:
+                self.observation_space = spaces.Box(low=-np.inf, high=np.inf, shape=(single_robot_obs_size,), dtype=np.float64)
+            robot_info = self.env_info["robot"]
+
+            if env_mode != "tournament":
+                if interpolation_order in [1, 2]:
+                    self.action_space = spaces.Box(low=robot_info["joint_pos_limit"][0], high=robot_info["joint_pos_limit"][1])
+                if interpolation_order in [3, 4, -1]:
+                    self.action_space = spaces.Box(low=np.vstack([robot_info["joint_pos_limit"][0], robot_info["joint_vel_limit"][0]]), 
+                                                high=np.vstack([robot_info["joint_pos_limit"][1], robot_info["joint_vel_limit"][1]]))
+                if interpolation_order in [5]:
+                    self.action_space = spaces.Box(low=np.vstack([robot_info["joint_pos_limit"][0], robot_info["joint_vel_limit"][0], robot_info["joint_acc_limit"][0]]), 
+                                                high=np.vstack([robot_info["joint_pos_limit"][1], robot_info["joint_vel_limit"][1], robot_info["joint_acc_limit"][1]]))
+            else:
+                acts = [None, None]
+                for i in range(2):
+                    if interpolation_order[i] in [1, 2]:
+                        acts[i] = spaces.Box(low=robot_info["joint_pos_limit"][0], high=robot_info["joint_pos_limit"][1])
+                    if interpolation_order[i] in [3, 4, -1]:
+                        acts[i] = spaces.Box(low=np.vstack([robot_info["joint_pos_limit"][0], robot_info["joint_vel_limit"][0]]), 
+                                                high=np.vstack([robot_info["joint_pos_limit"][1], robot_info["joint_vel_limit"][1]]))
+                    if interpolation_order[i] in [5]:
+                        acts[i] = spaces.Box(low=np.vstack([robot_info["joint_pos_limit"][0], robot_info["joint_vel_limit"][0], robot_info["joint_acc_limit"][0]]), 
+                                                high=np.vstack([robot_info["joint_pos_limit"][1], robot_info["joint_vel_limit"][1], robot_info["joint_acc_limit"][1]]))
+                self.action_space = spaces.Tuple((acts[0], acts[1]))
+
+        constraint_list = constraints.ConstraintList()
+        constraint_list.add(constraints.JointPositionConstraint(self.env_info))
+        constraint_list.add(constraints.JointVelocityConstraint(self.env_info))
+        constraint_list.add(constraints.EndEffectorConstraint(self.env_info))
+        if "7dof" in self.env_name or self.env_name == "tournament":
+            constraint_list.add(constraints.LinkConstraint(self.env_info))
+
+        self.env_info['constraints'] = constraint_list
+        self.env_info['env_name'] = self.env_name
+
+        super().__init__(self.base_env.info)
+
+    def step(self, action):
+        obs, reward, done, info = self.base_env.step(action)
+
+        if "tournament" in self.env_name:
+            info["constraints_value"] = list()
+            info["jerk"] = list()
+            for i in range(2):
+                obs_agent = obs[i * int(len(obs) / 2): (i + 1) * int(len(obs) / 2)]
+                info["constraints_value"].append(deepcopy(self.env_info['constraints'].fun(
+                    obs_agent[self.env_info['joint_pos_ids']], obs_agent[self.env_info['joint_vel_ids']])))
+                info["jerk"].append(
+                    self.base_env.jerk[i * self.env_info['robot']['n_joints']:(i + 1) * self.env_info['robot'][
+                        'n_joints']])
+
+            info["score"] = self.base_env.score
+            info["faults"] = self.base_env.faults
+
+        else:
+            info["constraints_value"] = deepcopy(self.env_info['constraints'].fun(obs[self.env_info['joint_pos_ids']],
+                                                                                  obs[self.env_info['joint_vel_ids']]))
+            info["jerk"] = self.base_env.jerk
+            info["success"] = self.check_success(obs)
+
+        if self.env_info['env_name'] == "tournament":
+            obs = np.array(np.split(obs, 2))
+
+        if self.render_human_active:
+            self.base_env.render()
+
+        return obs, reward, done, False, info
+
+    def render(self):
+        if self.render_mode == 'rgb_array':
+            return self.base_env.render(record = True)
+        elif self.render_mode == 'human':
+            self.render_human_active = True
+            self.base_env.render()
+        else:
+            raise ValueError(f"Unsupported render mode: '{self.render_mode}'")
+            
+    def reset(self, seed=None, options={}):
+        self.base_env.seed(seed)
+        obs = self.base_env.reset()
+        if self.env_info['env_name'] == "tournament":
+            obs = np.array(np.split(obs, 2))
+        return obs, {}
+
+    def check_success(self, obs):
+        puck_pos, puck_vel = self.base_env.get_puck(obs)
+
+        puck_pos, _ = robot_to_world(self.base_env.env_info["robot"]["base_frame"][0], translation=puck_pos)
+        success = 0
+
+        if "hit" in self.env_name:
+            if puck_pos[0] - self.base_env.env_info['table']['length'] / 2 > 0 and \
+                    np.abs(puck_pos[1]) - self.base_env.env_info['table']['goal_width'] / 2 < 0:
+                success = 1
+
+        elif "defend" in self.env_name:
+            if -0.8 < puck_pos[0] <= -0.2 and puck_vel[0] < 0.1:
+                success = 1
+
+        elif "prepare" in self.env_name:
+            if -0.8 < puck_pos[0] <= -0.2 and np.abs(puck_pos[1]) < 0.39105 and puck_vel[0] < 0.1:
+                success = 1
+        return success
+
+    @property
+    def unwrapped(self):
+        return self
+    
+    def close(self):
+        self.base_env.stop()
+
+
+if __name__ == "__main__":
+    env = AirHockeyEnv(env_mode="7dof-hit")
+    env.reset()
+
+    R = 0.
+    J = 0.
+    gamma = 1.
+    steps = 0
+    while True:
+        action = np.random.uniform(-1, 1, (2, env.env_info['robot']['n_joints'])) * 3
+        observation, reward, done, info = env.step(action)
+        env.render()
+        gamma *= env.info.gamma
+        J += gamma * reward
+        R += reward
+        steps += 1
+        if done or steps > env.info.horizon:
+            print("J: ", J, " R: ", R)
+            R = 0.
+            J = 0.
+            gamma = 1.
+            steps = 0
+            env.reset()

fancy_gym/envs/mujoco/air_hockey/constraints/__init__.py

@@ -0,0 +1 @@
+from .constraints import *

fancy_gym/envs/mujoco/air_hockey/constraints/constraints.py

@@ -0,0 +1,212 @@
+import copy
+
+import numpy as np
+
+from fancy_gym.envs.mujoco.air_hockey.utils.kinematics import forward_kinematics, jacobian
+
+
+class Constraint:
+    def __init__(self, env_info, output_dim, **kwargs):
+        """
+        Constructor
+
+        Args
+        ----
+        env_info: dict
+            A dictionary contains information about the environment;
+        output_dim: int
+            The output dimension of the constraints.
+        **kwargs: dict
+            A dictionary contains agent related information.
+        """
+        self._env_info = env_info
+        self._name = None
+
+        self.output_dim = output_dim
+
+        self._fun_value = np.zeros(self.output_dim)
+        self._jac_value = np.zeros((self.output_dim, 2 * env_info["robot"]["n_joints"]))
+        self._q_prev = None
+        self._dq_prev = None
+
+    @property
+    def name(self):
+        """
+        The name of the constraints
+
+        """
+        return self._name
+
+    def fun(self, q, dq):
+        """
+        The function of the constraint.
+
+        Args
+        ----
+        q: numpy.ndarray, (num_joints,)
+            The joint position of the robot
+        dq: numpy.ndarray, (num_joints,)
+            The joint velocity of the robot
+
+        Returns
+        -------
+        numpy.ndarray, (out_dim,):
+            The value computed by the constraints function.
+        """
+        if np.equal(q, self._q_prev).all() and np.equal(dq, self._dq_prev):
+            return self._fun_value
+        else:
+            self._jacobian(q, dq)
+            return self._fun(q, dq)
+
+    def jacobian(self, q, dq):
+        """
+        Jacobian is the derivative of the constraint function w.r.t the robot joint position and velocity.
+
+        Args
+        ----
+        q: ndarray, (num_joints,)
+            The joint position of the robot
+        dq: ndarray, (num_joints,)
+            The joint velocity of the robot
+
+        Returns
+        -------
+        numpy.ndarray, (dim_output, num_joints * 2):
+            The flattened jacobian of the constraint function J = [dc / dq, dc / dq_dot]
+
+        """
+        if np.equal(q, self._q_prev).all() and np.equal(dq, self._dq_prev):
+            return self._fun_value
+        else:
+            self._fun(q, dq)
+            return self._jacobian(q, dq)
+
+    def _fun(self, q, dq):
+        raise NotImplementedError
+
+    def _jacobian(self, q, dq):
+        raise NotImplementedError
+
+
+class ConstraintList:
+    def __init__(self):
+        self.constraints = dict()
+
+    def keys(self):
+        return self.constraints.keys()
+
+    def get(self, key):
+        return self.constraints.get(key)
+
+    def add(self, c):
+        self.constraints.update({c.name: c})
+
+    def delete(self, name):
+        del self.constraints[name]
+
+    def fun(self, q, dq):
+        return {key: self.constraints[key].fun(q, dq) for key in self.constraints}
+
+    def jacobian(self, q, dq):
+        return {key: self.constraints[key].jacobian(q, dq) for key in self.constraints}
+
+
+class JointPositionConstraint(Constraint):
+    def __init__(self, env_info, **kwargs):
+        super().__init__(env_info, output_dim=2 * env_info["robot"]["n_joints"], **kwargs)
+        self.joint_limits = self._env_info['robot']['joint_pos_limit'] * 0.95
+        self._name = 'joint_pos_constr'
+
+    def _fun(self, q, dq):
+        self._fun_value[:int(self.output_dim / 2)] = q - self.joint_limits[1]
+        self._fun_value[int(self.output_dim / 2):] = self.joint_limits[0] - q
+        return self._fun_value
+
+    def _jacobian(self, q, dq):
+        self._jac_value[:int(self.output_dim / 2), :int(self.output_dim / 2)] = np.eye(
+            self._env_info['robot']['n_joints'])
+        self._jac_value[int(self.output_dim / 2):, :int(self.output_dim / 2)] = -np.eye(
+            self._env_info['robot']['n_joints'])
+        return self._jac_value
+
+
+class JointVelocityConstraint(Constraint):
+    def __init__(self, env_info, **kwargs):
+        super().__init__(env_info, output_dim=2 * env_info["robot"]["n_joints"], **kwargs)
+        self.joint_limits = self._env_info['robot']['joint_vel_limit'] * 0.95
+        self._name = 'joint_vel_constr'
+
+    def _fun(self, q, dq):
+        self._fun_value[:int(self.output_dim / 2)] = dq - self.joint_limits[1]
+        self._fun_value[int(self.output_dim / 2):] = self.joint_limits[0] - dq
+        return self._fun_value
+
+    def _jacobian(self, q, dq):
+        self._jac_value[:int(self.output_dim / 2), int(self.output_dim / 2):] = np.eye(
+            self._env_info['robot']['n_joints'])
+        self._jac_value[int(self.output_dim / 2):, int(self.output_dim / 2):] = -np.eye(
+            self._env_info['robot']['n_joints'])
+        return self._jac_value
+
+
+class EndEffectorConstraint(Constraint):
+    def __init__(self, env_info, **kwargs):
+        # 1 Dimension on x direction: x > x_lb
+        # 2 Dimension on y direction: y > y_lb, y < y_ub
+        # 2 Dimension on z direction: z > z_lb, z < z_ub
+        super().__init__(env_info, output_dim=5, **kwargs)
+        self._name = "ee_constr"
+        tolerance = 0.02
+
+        self.robot_model = copy.deepcopy(self._env_info['robot']['robot_model'])
+        self.robot_data = copy.deepcopy(self._env_info['robot']['robot_data'])
+
+        self.x_lb = - self._env_info['robot']['base_frame'][0][0, 3] - (
+                self._env_info['table']['length'] / 2 - self._env_info['mallet']['radius'])
+        self.y_lb = - (self._env_info['table']['width'] / 2 - self._env_info['mallet']['radius'])
+        self.y_ub = (self._env_info['table']['width'] / 2 - self._env_info['mallet']['radius'])
+        self.z_lb = self._env_info['robot']['ee_desired_height'] - tolerance
+        self.z_ub = self._env_info['robot']['ee_desired_height'] + tolerance
+
+    def _fun(self, q, dq):
+        ee_pos, _ = forward_kinematics(self.robot_model, self.robot_data, q)
+        self._fun_value = np.array([-ee_pos[0] + self.x_lb,
+                                    -ee_pos[1] + self.y_lb, ee_pos[1] - self.y_ub,
+                                    -ee_pos[2] + self.z_lb, ee_pos[2] - self.z_ub])
+        return self._fun_value
+
+    def _jacobian(self, q, dq):
+        jac = jacobian(self.robot_model, self.robot_data, q)
+        dc_dx = np.array([[-1, 0., 0.], [0., -1., 0.], [0., 1., 0.], [0., 0., -1.], [0., 0., 1.]])
+        self._jac_value[:, :self._env_info['robot']['n_joints']] = dc_dx @ jac[:3, :self._env_info['robot']['n_joints']]
+        return self._jac_value
+
+
+class LinkConstraint(Constraint):
+    def __init__(self, env_info, **kwargs):
+        # 1 Dimension: wrist_z > minimum_height
+        # 2 Dimension: elbow_z > minimum_height
+        super().__init__(env_info, output_dim=2, **kwargs)
+        self._name = "link_constr"
+
+        self.robot_model = copy.deepcopy(self._env_info['robot']['robot_model'])
+        self.robot_data = copy.deepcopy(self._env_info['robot']['robot_data'])
+
+        self.z_lb = 0.25
+
+    def _fun(self, q, dq):
+        wrist_pos, _ = forward_kinematics(self.robot_model, self.robot_data, q, link="7")
+        elbow_pos, _ = forward_kinematics(self.robot_model, self.robot_data, q, link="4")
+        self._fun_value = np.array([-wrist_pos[2] + self.z_lb,
+                                    -elbow_pos[2] + self.z_lb])
+        return self._fun_value
+
+    def _jacobian(self, q, dq):
+        jac_wrist = jacobian(self.robot_model, self.robot_data, q, link="7")
+        jac_elbow = jacobian(self.robot_model, self.robot_data, q, link="4")
+        self._jac_value[:, :self._env_info['robot']['n_joints']] = np.vstack([
+            -jac_wrist[2, :self._env_info['robot']['n_joints']],
+            -jac_elbow[2, :self._env_info['robot']['n_joints']],
+        ])
+        return self._jac_value

fancy_gym/envs/mujoco/air_hockey/data/iiwas/double.xml

@@ -0,0 +1,58 @@
+<mujoco model="AirHockeySingle">
+    <include file="iiwa1.xml"/>
+
+    <include file="iiwa2.xml"/>
+
+    <include file="../table.xml"/>
+
+    <contact>
+        <exclude body1="table_surface" body2="iiwa_1/base"/>
+        <exclude body1="table_surface" body2="iiwa_1/link_1"/>
+        <exclude body1="table_surface" body2="iiwa_1/link_2"/>
+        <exclude body1="table_surface" body2="iiwa_1/link_3"/>
+        <exclude body1="table_surface" body2="iiwa_1/link_4"/>
+        <exclude body1="table_surface" body2="iiwa_1/link_5"/>
+        <exclude body1="table_surface" body2="iiwa_1/link_6"/>
+        <exclude body1="table_surface" body2="iiwa_1/link_7"/>
+        <exclude body1="table_surface" body2="iiwa_1/striker_joint_link"/>
+        <exclude body1="table_surface" body2="iiwa_1/striker_mallet"/>
+
+        <exclude body1="rim" body2="iiwa_1/base"/>
+        <exclude body1="rim" body2="iiwa_1/link_1"/>
+        <exclude body1="rim" body2="iiwa_1/link_2"/>
+        <exclude body1="rim" body2="iiwa_1/link_3"/>
+        <exclude body1="rim" body2="iiwa_1/link_4"/>
+        <exclude body1="rim" body2="iiwa_1/link_5"/>
+        <exclude body1="rim" body2="iiwa_1/link_6"/>
+        <exclude body1="rim" body2="iiwa_1/link_7"/>
+        <exclude body1="rim" body2="iiwa_1/striker_joint_link"/>
+        <exclude body1="rim" body2="iiwa_1/striker_mallet"/>
+
+        <exclude body1="world" body2="iiwa_1/striker_mallet"/>
+
+        <exclude body1="table_surface" body2="iiwa_2/base"/>
+        <exclude body1="table_surface" body2="iiwa_2/link_1"/>
+        <exclude body1="table_surface" body2="iiwa_2/link_2"/>
+        <exclude body1="table_surface" body2="iiwa_2/link_3"/>
+        <exclude body1="table_surface" body2="iiwa_2/link_4"/>
+        <exclude body1="table_surface" body2="iiwa_2/link_5"/>
+        <exclude body1="table_surface" body2="iiwa_2/link_6"/>
+        <exclude body1="table_surface" body2="iiwa_2/link_7"/>
+        <exclude body1="table_surface" body2="iiwa_2/striker_joint_link"/>
+        <exclude body1="table_surface" body2="iiwa_2/striker_mallet"/>
+
+        <exclude body1="rim" body2="iiwa_2/base"/>
+        <exclude body1="rim" body2="iiwa_2/link_1"/>
+        <exclude body1="rim" body2="iiwa_2/link_2"/>
+        <exclude body1="rim" body2="iiwa_2/link_3"/>
+        <exclude body1="rim" body2="iiwa_2/link_4"/>
+        <exclude body1="rim" body2="iiwa_2/link_5"/>
+        <exclude body1="rim" body2="iiwa_2/link_6"/>
+        <exclude body1="rim" body2="iiwa_2/link_7"/>
+        <exclude body1="rim" body2="iiwa_2/striker_joint_link"/>
+        <exclude body1="rim" body2="iiwa_2/striker_mallet"/>
+
+        <exclude body1="world" body2="iiwa_2/striker_mallet"/>
+    </contact>
+
+</mujoco>

fancy_gym/envs/mujoco/air_hockey/data/iiwas/iiwa1.xml

@@ -0,0 +1,108 @@
+<mujoco model="iiwa_1">
+    <compiler angle="radian" autolimits="true" meshdir="assets"/>
+
+    <default>
+        <default class="vis">
+            <geom contype="0" conaffinity="0"/>
+        </default>
+        <default class="robot">
+            <geom condim="4" solref="0.02 0.3" priority="2"/>
+        </default>
+    </default>
+    <asset>
+        <mesh name="link_0" file="link_0.stl"/>
+        <mesh name="link_1" file="link_1.stl"/>
+        <mesh name="link_2" file="link_2.stl"/>
+        <mesh name="link_3" file="link_3.stl"/>
+        <mesh name="link_4" file="link_4.stl"/>
+        <mesh name="link_5" file="link_5.stl"/>
+        <mesh name="link_6" file="link_6.stl"/>
+        <mesh name="link_7" file="link_7.stl"/>
+        <mesh name="EE_arm" file="EE_arm.stl"/>
+        <mesh name="EE_mallet_foam" file="EE_mallet_foam.stl"/>
+    </asset>
+    <worldbody>
+        <body name="iiwa_1/base" pos="-1.51 0 -0.1">
+            <geom type="mesh" rgba="0.4 0.4 0.4 1" mesh="link_0" class="vis"/>
+            <body name="iiwa_1/link_1" pos="0 0 0.1575">
+                <inertial pos="4.007709e-06 -0.033936 0.122467" mass="8.240527"
+                          fullinertia="0.021981 0.022182 0.008234 -2.897243e-07 6.3165236e-07 0.003285"/>
+                <joint name="iiwa_1/joint_1" pos="0 0 0" axis="0 0 1" range="-2.96706 2.96706" damping="0.33032"
+                       frictionloss="0.384477"/>
+                <geom type="mesh" rgba="1 0.423529 0.0392157 1" mesh="link_1" class="vis"/>
+                <body name="iiwa_1/link_2" pos="0 0 0.2025" quat="0 0 0.707107 0.707107">
+                    <inertial pos="0.003402 0.034792 0.046725" mass="6.357896"
+                              fullinertia="0.015565 0.005180 0.015484 -4.147301e-06 1.192255e-05 0.002538"/>
+                    <joint name="iiwa_1/joint_2" pos="0 0 0" axis="0 0 1" range="-2.0944 2.0944" damping="0.21216"
+                           frictionloss="0.496333"/>
+                    <geom type="mesh" rgba="1 0.423529 0.0392157 1" mesh="link_2" class="vis"/>
+                    <body name="iiwa_1/link_3" pos="0 0.2045 0" quat="0 0 0.707107 0.707107">
+                        <inertial pos="-0.001452 0.031526 0.133584" mass="4.042756"
+                                  fullinertia="0.010914 0.010381 0.003139 -3.540575e-06 -9.059062e-06 -0.002128"/>
+                        <joint name="iiwa_1/joint_3" pos="0 0 0" axis="0 0 1" range="-2.96706 2.96706" damping="0.1"
+                               frictionloss="0.173951"/>
+                        <geom type="mesh" rgba="1 0.423529 0.0392157 1" mesh="link_3" class="vis"/>
+                        <body name="iiwa_1/link_4" pos="0 0 0.2155" quat="0.707107 0.707107 0 0">
+                            <inertial pos="-0.002527 0.053508 0.037205" mass="3.642249"
+                                      fullinertia="0.007536 0.002538 0.007206 -5.707028e-06 2.781894e-06 0.001256"/>
+                            <joint name="iiwa_1/joint_4" pos="0 0 0" axis="0 0 1" range="-2.0944 2.0944"
+                                   damping="0.219041" frictionloss="0.3751"/>
+                            <geom type="mesh" rgba="1 0.423529 0.0392157 1" mesh="link_4" class="vis"/>
+                            <body name="iiwa_1/link_5" pos="0 0.1845 0" quat="0 0 0.707107 0.707107">
+                                <inertial pos="0.001855 0.024573 0.080131" mass="2.580896"
+                                          fullinertia="0.005201 0.004488 0.002242 1.089316e-07 9.035623e-07 -0.001613"/>
+                                <joint name="iiwa_1/joint_5" pos="0 0 0" axis="0 0 1" range="-2.96706 2.96706"
+                                       damping="0.185923" frictionloss="0.481099"/>
+                                <geom type="mesh" rgba="1 0.423529 0.0392157 1" mesh="link_5" class="vis"/>
+                                <body name="iiwa_1/link_6" pos="0 0 0.2155" quat="0.707107 0.707107 0 0">
+                                    <inertial pos="-0.001739 -0.001973 -0.002502" mass="2.760564"
+                                              fullinertia="0.002534 0.001821 0.002393 -1.311766e-06 9.508242e-07 0.000134"/>
+                                    <joint name="iiwa_1/joint_6" pos="0 0 0" axis="0 0 1" range="-2.0944 2.0944"
+                                           damping="0.1" frictionloss="0.196149"/>
+                                    <geom type="mesh" rgba="1 0.423529 0.0392157 1" mesh="link_6" class="vis"/>
+                                    <body name="iiwa_1/link_7" pos="0 0.081 0" quat="0 0 0.707107 0.707107">
+                                        <inertial pos="0.000735 0.000387 0.026460" mass="1.285417"
+                                                  fullinertia="0.000151 0.000150 0.000187 -7.223100e-08 2.038333e-06 -3.396830e-07"/>
+                                        <joint name="iiwa_1/joint_7" pos="0 0 0" axis="0 0 1" range="-3.05433 3.05433"
+                                               damping="0.1" frictionloss="0.299238" armature="0.01"/>
+                                        <geom type="mesh" rgba="0.4 0.4 0.4 1" mesh="link_7" class="vis"/>
+                                        <geom pos="0 0 0.07" type="mesh" rgba="0.3 0.3 0.3 1" mesh="EE_arm"
+                                              class="vis"/>
+                                        <body name="iiwa_1/striker_joint_link" pos="0 0 0.585">
+                                            <inertial pos="0 0 0" mass="0.1" diaginertia="0.001 0.001 0.001"/>
+                                            <body name="iiwa_1/striker_mallet" pos="0 0 0">
+                                                <inertial pos="0 0 0.0682827" mass="0.283"
+                                                          diaginertia="0.005 0.005 0.005"/>
+                                                <joint name="iiwa_1/striker_joint_1" pos="0 0 0" axis="0 1 0"
+                                                       range="-1.5708 1.5708" damping="0.0"/>
+                                                <joint name="iiwa_1/striker_joint_2" pos="0 0 0" axis="1 0 0"
+                                                       range="-1.5708 1.5708" damping="0.0"/>
+                                                <geom type="mesh" rgba="0.3 0.3 0.3 1" mesh="EE_mallet_foam"
+                                                      class="vis"/>
+                                                <geom name="iiwa_1/ee" type="cylinder" rgba="0.3 0.3 0.3 0.1"
+                                                      size="0.04815 0.03" pos="0 0 0.0505" friction="0 0 0"/>
+                                            </body>
+                                        </body>
+                                    </body>
+                                </body>
+                            </body>
+                        </body>
+                    </body>
+                </body>
+            </body>
+        </body>
+    </worldbody>
+
+
+    <actuator>
+        <motor name="iiwa_1/joint_1" joint="iiwa_1/joint_1" ctrlrange="-320 320"/>
+        <motor name="iiwa_1/joint_2" joint="iiwa_1/joint_2" ctrlrange="-320 320"/>
+        <motor name="iiwa_1/joint_3" joint="iiwa_1/joint_3" ctrlrange="-176 176"/>
+        <motor name="iiwa_1/joint_4" joint="iiwa_1/joint_4" ctrlrange="-176 176"/>
+        <motor name="iiwa_1/joint_5" joint="iiwa_1/joint_5" ctrlrange="-110 110"/>
+        <motor name="iiwa_1/joint_6" joint="iiwa_1/joint_6" ctrlrange="-40 40"/>
+        <motor name="iiwa_1/joint_7" joint="iiwa_1/joint_7" ctrlrange="-40 40"/>
+        <motor name="iiwa_1/striker_joint_1" joint="iiwa_1/striker_joint_1" ctrlrange="-10 10"/>
+        <motor name="iiwa_1/striker_joint_2" joint="iiwa_1/striker_joint_2" ctrlrange="-10 10"/>
+    </actuator>
+</mujoco>

fancy_gym/envs/mujoco/air_hockey/data/iiwas/iiwa2.xml

@@ -0,0 +1,89 @@
+<mujoco model="iiwa_2">
+    <compiler angle="radian" autolimits="true" meshdir="assets"/>
+
+    <worldbody>
+        <body name="iiwa_2/base" pos="1.51 0 -0.1" quat="0 0 0 1">
+            <geom type="mesh" rgba="0.4 0.4 0.4 1" mesh="link_0" class="vis"/>
+            <body name="iiwa_2/link_1" pos="0 0 0.1575">
+                <inertial pos="4.007709e-06 -0.033936 0.122467" mass="8.240527"
+                          fullinertia="0.021981 0.022182 0.008234 -2.897243e-07 6.3165236e-07 0.003285"/>
+                <joint name="iiwa_2/joint_1" pos="0 0 0" axis="0 0 1" range="-2.96706 2.96706" damping="0.33032"
+                       frictionloss="0.384477"/>
+                <geom type="mesh" rgba="1 0.423529 0.0392157 1" mesh="link_1" class="vis"/>
+                <body name="iiwa_2/link_2" pos="0 0 0.2025" quat="0 0 0.707107 0.707107">
+                    <inertial pos="0.003402 0.034792 0.046725" mass="6.357896"
+                              fullinertia="0.015565 0.005180 0.015484 -4.147301e-06 1.192255e-05 0.002538"/>
+                    <joint name="iiwa_2/joint_2" pos="0 0 0" axis="0 0 1" range="-2.0944 2.0944" damping="0.21216"
+                           frictionloss="0.496333"/>
+                    <geom type="mesh" rgba="1 0.423529 0.0392157 1" mesh="link_2" class="vis"/>
+                    <body name="iiwa_2/link_3" pos="0 0.2045 0" quat="0 0 0.707107 0.707107">
+                        <inertial pos="-0.001452 0.031526 0.133584" mass="4.042756"
+                                  fullinertia="0.010914 0.010381 0.003139 -3.540575e-06 -9.059062e-06 -0.002128"/>
+                        <joint name="iiwa_2/joint_3" pos="0 0 0" axis="0 0 1" range="-2.96706 2.96706" damping="0.1"
+                               frictionloss="0.173951"/>
+                        <geom type="mesh" rgba="1 0.423529 0.0392157 1" mesh="link_3" class="vis"/>
+                        <body name="iiwa_2/link_4" pos="0 0 0.2155" quat="0.707107 0.707107 0 0">
+                            <inertial pos="-0.002527 0.053508 0.037205" mass="3.642249"
+                                      fullinertia="0.007536 0.002538 0.007206 -5.707028e-06 2.781894e-06 0.001256"/>
+                            <joint name="iiwa_2/joint_4" pos="0 0 0" axis="0 0 1" range="-2.0944 2.0944"
+                                   damping="0.219041" frictionloss="0.3751"/>
+                            <geom type="mesh" rgba="1 0.423529 0.0392157 1" mesh="link_4" class="vis"/>
+                            <body name="iiwa_2/link_5" pos="0 0.1845 0" quat="0 0 0.707107 0.707107">
+                                <inertial pos="0.001855 0.024573 0.080131" mass="2.580896"
+                                          fullinertia="0.005201 0.004488 0.002242 1.089316e-07 9.035623e-07 -0.001613"/>
+                                <joint name="iiwa_2/joint_5" pos="0 0 0" axis="0 0 1" range="-2.96706 2.96706"
+                                       damping="0.185923" frictionloss="0.481099"/>
+                                <geom type="mesh" rgba="1 0.423529 0.0392157 1" mesh="link_5" class="vis"/>
+                                <body name="iiwa_2/link_6" pos="0 0 0.2155" quat="0.707107 0.707107 0 0">
+                                    <inertial pos="-0.001739 -0.001973 -0.002502" mass="2.760564"
+                                              fullinertia="0.002534 0.001821 0.002393 -1.311766e-06 9.508242e-07 0.000134"/>
+                                    <joint name="iiwa_2/joint_6" pos="0 0 0" axis="0 0 1" range="-2.0944 2.0944"
+                                           damping="0.1" frictionloss="0.196149"/>
+                                    <geom type="mesh" rgba="1 0.423529 0.0392157 1" mesh="link_6" class="vis"/>
+                                    <body name="iiwa_2/link_7" pos="0 0.081 0" quat="0 0 0.707107 0.707107">
+                                        <inertial pos="0.000735 0.000387 0.026460" mass="1.285417"
+                                                  fullinertia="0.000151 0.000150 0.000187 -7.223100e-08 2.038333e-06 -3.396830e-07"/>
+                                        <joint name="iiwa_2/joint_7" pos="0 0 0" axis="0 0 1" range="-3.05433 3.05433"
+                                               damping="0.1" frictionloss="0.299238" armature="0.01"/>
+                                        <geom type="mesh" rgba="0.4 0.4 0.4 1" mesh="link_7" class="vis"/>
+                                        <geom pos="0 0 0.07" type="mesh" rgba="0.3 0.3 0.3 1" mesh="EE_arm"
+                                              class="vis"/>
+                                        <body name="iiwa_2/striker_joint_link" pos="0 0 0.585">
+                                            <inertial pos="0 0 0" mass="0.1" diaginertia="0.001 0.001 0.001"/>
+                                            <body name="iiwa_2/striker_mallet" pos="0 0 0">
+                                                <inertial pos="0 0 0.0682827" mass="0.283"
+                                                          diaginertia="0.005 0.005 0.005"/>
+                                                <joint name="iiwa_2/striker_joint_1" pos="0 0 0" axis="0 1 0"
+                                                       range="-1.5708 1.5708" damping="0.0"/>
+                                                <joint name="iiwa_2/striker_joint_2" pos="0 0 0" axis="1 0 0"
+                                                       range="-1.5708 1.5708" damping="0.0"/>
+                                                <geom type="mesh" rgba="0.3 0.3 0.3 1" mesh="EE_mallet_foam"
+                                                      class="vis"/>
+                                                <geom name="iiwa_2/ee" type="cylinder" rgba="0.3 0.3 0.3 0.1"
+                                                      size="0.04815 0.03" pos="0 0 0.0505" friction="0 0 0"/>
+                                            </body>
+                                        </body>
+                                    </body>
+                                </body>
+                            </body>
+                        </body>
+                    </body>
+                </body>
+            </body>
+        </body>
+    </worldbody>
+
+
+    <actuator>
+        <motor name="iiwa_2/joint_1" joint="iiwa_2/joint_1" ctrlrange="-320 320"/>
+        <motor name="iiwa_2/joint_2" joint="iiwa_2/joint_2" ctrlrange="-320 320"/>
+        <motor name="iiwa_2/joint_3" joint="iiwa_2/joint_3" ctrlrange="-176 176"/>
+        <motor name="iiwa_2/joint_4" joint="iiwa_2/joint_4" ctrlrange="-176 176"/>
+        <motor name="iiwa_2/joint_5" joint="iiwa_2/joint_5" ctrlrange="-110 110"/>
+        <motor name="iiwa_2/joint_6" joint="iiwa_2/joint_6" ctrlrange="-40 40"/>
+        <motor name="iiwa_2/joint_7" joint="iiwa_2/joint_7" ctrlrange="-40 40"/>
+        <motor name="iiwa_2/striker_joint_1" joint="iiwa_2/striker_joint_1" ctrlrange="-10 10"/>
+        <motor name="iiwa_2/striker_joint_2" joint="iiwa_2/striker_joint_2" ctrlrange="-10 10"/>
+    </actuator>
+
+</mujoco>

fancy_gym/envs/mujoco/air_hockey/data/iiwas/iiwa_only.xml

@@ -0,0 +1,94 @@
+<mujoco model="iiwa_1">
+    <compiler angle="radian" autolimits="true" meshdir="assets"/>
+
+    <default>
+        <default class="vis">
+            <geom contype="0" conaffinity="0"/>
+        </default>
+        <default class="robot">
+            <geom condim="4" solref="0.02 0.3" priority="2"/>
+        </default>
+    </default>
+    <asset>
+        <mesh name="link_0" file="link_0.stl"/>
+        <mesh name="link_1" file="link_1.stl"/>
+        <mesh name="link_2" file="link_2.stl"/>
+        <mesh name="link_3" file="link_3.stl"/>
+        <mesh name="link_4" file="link_4.stl"/>
+        <mesh name="link_5" file="link_5.stl"/>
+        <mesh name="link_6" file="link_6.stl"/>
+        <mesh name="link_7" file="link_7.stl"/>
+        <mesh name="EE_arm" file="EE_arm.stl"/>
+        <mesh name="EE_mallet_foam" file="EE_mallet_foam.stl"/>
+    </asset>
+    <worldbody>
+        <body name="iiwa_1/base" pos="0.0 0 0.0">
+            <geom type="mesh" rgba="0.4 0.4 0.4 1" mesh="link_0" class="vis"/>
+            <body name="iiwa_1/link_1" pos="0 0 0.1575">
+                <inertial pos="4.007709e-06 -0.033936 0.122467" mass="8.240527"
+                          fullinertia="0.021981 0.022182 0.008234 -2.897243e-07 6.3165236e-07 0.003285"/>
+                <joint name="iiwa_1/joint_1" pos="0 0 0" axis="0 0 1" range="-2.96706 2.96706" damping="0.33032"
+                       frictionloss="0.384477"/>
+                <geom type="mesh" rgba="1 0.423529 0.0392157 1" mesh="link_1" class="vis"/>
+                <body name="iiwa_1/link_2" pos="0 0 0.2025" quat="0 0 0.707107 0.707107">
+                    <inertial pos="0.003402 0.034792 0.046725" mass="6.357896"
+                              fullinertia="0.015565 0.005180 0.015484 -4.147301e-06 1.192255e-05 0.002538"/>
+                    <joint name="iiwa_1/joint_2" pos="0 0 0" axis="0 0 1" range="-2.0944 2.0944" damping="0.21216"
+                           frictionloss="0.496333"/>
+                    <geom type="mesh" rgba="1 0.423529 0.0392157 1" mesh="link_2" class="vis"/>
+                    <body name="iiwa_1/link_3" pos="0 0.2045 0" quat="0 0 0.707107 0.707107">
+                        <inertial pos="-0.001452 0.031526 0.133584" mass="4.042756"
+                                  fullinertia="0.010914 0.010381 0.003139 -3.540575e-06 -9.059062e-06 -0.002128"/>
+                        <joint name="iiwa_1/joint_3" pos="0 0 0" axis="0 0 1" range="-2.96706 2.96706" damping="0.1"
+                               frictionloss="0.173951"/>
+                        <geom type="mesh" rgba="1 0.423529 0.0392157 1" mesh="link_3" class="vis"/>
+                        <body name="iiwa_1/link_4" pos="0 0 0.2155" quat="0.707107 0.707107 0 0">
+                            <inertial pos="-0.002527 0.053508 0.037205" mass="3.642249"
+                                      fullinertia="0.007536 0.002538 0.007206 -5.707028e-06 2.781894e-06 0.001256"/>
+                            <joint name="iiwa_1/joint_4" pos="0 0 0" axis="0 0 1" range="-2.0944 2.0944"
+                                   damping="0.219041" frictionloss="0.3751"/>
+                            <geom type="mesh" rgba="1 0.423529 0.0392157 1" mesh="link_4" class="vis"/>
+                            <body name="iiwa_1/link_5" pos="0 0.1845 0" quat="0 0 0.707107 0.707107">
+                                <inertial pos="0.001855 0.024573 0.080131" mass="2.580896"
+                                          fullinertia="0.005201 0.004488 0.002242 1.089316e-07 9.035623e-07 -0.001613"/>
+                                <joint name="iiwa_1/joint_5" pos="0 0 0" axis="0 0 1" range="-2.96706 2.96706"
+                                       damping="0.185923" frictionloss="0.481099"/>
+                                <geom type="mesh" rgba="1 0.423529 0.0392157 1" mesh="link_5" class="vis"/>
+                                <body name="iiwa_1/link_6" pos="0 0 0.2155" quat="0.707107 0.707107 0 0">
+                                    <inertial pos="-0.001739 -0.001973 -0.002502" mass="2.760564"
+                                              fullinertia="0.002534 0.001821 0.002393 -1.311766e-06 9.508242e-07 0.000134"/>
+                                    <joint name="iiwa_1/joint_6" pos="0 0 0" axis="0 0 1" range="-2.0944 2.0944"
+                                           damping="0.1" frictionloss="0.196149"/>
+                                    <geom type="mesh" rgba="1 0.423529 0.0392157 1" mesh="link_6" class="vis"/>
+                                    <body name="iiwa_1/link_7" pos="0 0.081 0" quat="0 0 0.707107 0.707107">
+                                        <inertial pos="0.000735 0.000387 0.026460" mass="1.285417"
+                                                  fullinertia="0.000151 0.000150 0.000187 -7.223100e-08 2.038333e-06 -3.396830e-07"/>
+                                        <joint name="iiwa_1/joint_7" pos="0 0 0" axis="0 0 1" range="-3.05433 3.05433"
+                                               damping="0.1" frictionloss="0.299238" armature="0.01"/>
+                                        <geom type="mesh" rgba="0.4 0.4 0.4 1" mesh="link_7" class="vis"/>
+                                        <geom pos="0 0 0.07" type="mesh" rgba="0.3 0.3 0.3 1" mesh="EE_arm"
+                                              class="vis"/>
+                                        <body name="iiwa_1/striker_joint_link" pos="0 0 0.585">
+                                            <inertial pos="0 0 0" mass="0.1" diaginertia="0.001 0.001 0.001"/>
+                                        </body>
+                                    </body>
+                                </body>
+                            </body>
+                        </body>
+                    </body>
+                </body>
+            </body>
+        </body>
+    </worldbody>
+
+
+    <actuator>
+        <motor name="iiwa_1/joint_1" joint="iiwa_1/joint_1" ctrlrange="-320 320"/>
+        <motor name="iiwa_1/joint_2" joint="iiwa_1/joint_2" ctrlrange="-320 320"/>
+        <motor name="iiwa_1/joint_3" joint="iiwa_1/joint_3" ctrlrange="-176 176"/>
+        <motor name="iiwa_1/joint_4" joint="iiwa_1/joint_4" ctrlrange="-176 176"/>
+        <motor name="iiwa_1/joint_5" joint="iiwa_1/joint_5" ctrlrange="-110 110"/>
+        <motor name="iiwa_1/joint_6" joint="iiwa_1/joint_6" ctrlrange="-40 40"/>
+        <motor name="iiwa_1/joint_7" joint="iiwa_1/joint_7" ctrlrange="-40 40"/>
+    </actuator>
+</mujoco>

fancy_gym/envs/mujoco/air_hockey/data/iiwas/single.xml

@@ -0,0 +1,32 @@
+<mujoco model="AirHockeySingle">
+
+    <include file="iiwa1.xml"/>
+    <include file="../table.xml"/>
+
+    <contact>
+        <exclude body1="table_surface" body2="iiwa_1/base"/>
+        <exclude body1="table_surface" body2="iiwa_1/link_1"/>
+        <exclude body1="table_surface" body2="iiwa_1/link_2"/>
+        <exclude body1="table_surface" body2="iiwa_1/link_3"/>
+        <exclude body1="table_surface" body2="iiwa_1/link_4"/>
+        <exclude body1="table_surface" body2="iiwa_1/link_5"/>
+        <exclude body1="table_surface" body2="iiwa_1/link_6"/>
+        <exclude body1="table_surface" body2="iiwa_1/link_7"/>
+        <exclude body1="table_surface" body2="iiwa_1/striker_joint_link"/>
+        <exclude body1="table_surface" body2="iiwa_1/striker_mallet"/>
+
+        <exclude body1="rim" body2="iiwa_1/base"/>
+        <exclude body1="rim" body2="iiwa_1/link_1"/>
+        <exclude body1="rim" body2="iiwa_1/link_2"/>
+        <exclude body1="rim" body2="iiwa_1/link_3"/>
+        <exclude body1="rim" body2="iiwa_1/link_4"/>
+        <exclude body1="rim" body2="iiwa_1/link_5"/>
+        <exclude body1="rim" body2="iiwa_1/link_6"/>
+        <exclude body1="rim" body2="iiwa_1/link_7"/>
+        <exclude body1="rim" body2="iiwa_1/striker_joint_link"/>
+        <exclude body1="rim" body2="iiwa_1/striker_mallet"/>
+
+        <exclude body1="world" body2="iiwa_1/striker_mallet"/>
+    </contact>
+
+</mujoco>

fancy_gym/envs/mujoco/air_hockey/data/planar/double.xml

@@ -0,0 +1,37 @@
+<mujoco model="AirHockeySingle">
+
+    <include file="planar_robot_1.xml"/>
+
+    <include file="planar_robot_2.xml"/>
+
+    <include file="../table.xml"/>
+
+    <contact>
+        <exclude body1="planar_robot_1/body_ee" body2="table_surface"/>
+        <exclude body1="planar_robot_1/body_hand" body2="table_surface"/>
+        <exclude body1="planar_robot_1/body_3" body2="table_surface"/>
+        <exclude body1="planar_robot_1/body_2" body2="table_surface"/>
+        <exclude body1="planar_robot_1/body_1" body2="table_surface"/>
+        <exclude body1="planar_robot_1/base" body2="table_surface"/>
+        <exclude body1="planar_robot_1/body_ee" body2="rim"/>
+        <exclude body1="planar_robot_1/body_hand" body2="rim"/>
+        <exclude body1="planar_robot_1/body_3" body2="rim"/>
+        <exclude body1="planar_robot_1/body_2" body2="rim"/>
+        <exclude body1="planar_robot_1/body_1" body2="rim"/>
+        <exclude body1="planar_robot_1/base" body2="rim"/>
+
+        <exclude body1="planar_robot_2/body_ee" body2="table_surface"/>
+        <exclude body1="planar_robot_2/body_hand" body2="table_surface"/>
+        <exclude body1="planar_robot_2/body_3" body2="table_surface"/>
+        <exclude body1="planar_robot_2/body_2" body2="table_surface"/>
+        <exclude body1="planar_robot_2/body_1" body2="table_surface"/>
+        <exclude body1="planar_robot_2/base" body2="table_surface"/>
+        <exclude body1="planar_robot_2/body_ee" body2="rim"/>
+        <exclude body1="planar_robot_2/body_hand" body2="rim"/>
+        <exclude body1="planar_robot_2/body_3" body2="rim"/>
+        <exclude body1="planar_robot_2/body_2" body2="rim"/>
+        <exclude body1="planar_robot_2/body_1" body2="rim"/>
+        <exclude body1="planar_robot_2/base" body2="rim"/>
+    </contact>
+
+</mujoco>

fancy_gym/envs/mujoco/air_hockey/data/planar/planar_robot_1.xml

@@ -0,0 +1,61 @@
+<mujoco model="planar_robot_1">
+    <compiler autolimits="true" angle="radian"/>
+    <asset>
+        <material name="dark_red" rgba="0.58 0.03 0.25 1"/>
+        <material name="black" rgba="0.1 0.1 0.11 1"/>
+    </asset>
+
+    <default>
+        <default class="visual">
+            <geom contype="0" conaffinity="0"/>
+        </default>
+
+        <default class="robot">
+            <geom condim="4" solref="0.02 0.3" priority="2"/>
+        </default>
+    </default>
+
+    <worldbody>
+        <body name="planar_robot_1/base" pos="-1.51 0 -0.1">
+            <geom type="cylinder" material="black" size="0.075 0.2" pos="0 0 0.2" class="visual"/>
+            <inertial pos="0 0 0" mass="1" diaginertia="0.05 0.018 0.044"/>
+            <body name="planar_robot_1/body_1" pos="0 0 0.25">
+                <joint name="planar_robot_1/joint_1" axis="0 0 1" range="-2.9670597283903604 2.9670597283903604"/>
+                <inertial pos="0.25 0 0" mass="3" diaginertia="0.064375 0.064375 0.00375"/>
+                <geom class="robot" type="cylinder" material="dark_red" size="0.04 0.275" euler="0 1.57079632679 0"
+                      pos="0.275 0 0"/>
+                <body name="planar_robot_1/body_2" pos="0.55 0 0">
+                    <joint name="planar_robot_1/joint_2" axis="0 0 1" range="-1.8 1.8"/>
+                    <inertial pos="0.2 0 0" mass="2" diaginertia="0.0335 0.0335 0.003"/>
+                    <geom class="robot" type="cylinder" material="dark_red" size="0.04 0.22" euler="0 1.57079632679 0"
+                          pos="0.22 0 0"/>
+                    <geom class="robot" type="sphere" material="dark_red" size="0.05"/>
+                    <body name="planar_robot_1/body_3" pos="0.44 0 0">
+                        <joint name="planar_robot_1/joint_3" axis="0 0 1"
+                               range="-2.0943951023931953 2.0943951023931953"/>
+                        <inertial pos="0.2 0 0" mass="2" diaginertia="0.0335 0.0335 0.003"/>
+                        <geom class="robot" type="cylinder" material="dark_red" size="0.04 0.22"
+                              euler="0 1.57079632679 0" pos="0.22 0 0"/>
+                        <geom class="robot" type="sphere" material="dark_red" size="0.05"/>
+                        <body name="planar_robot_1/body_hand" pos="0.44 0 0">
+                            <inertial pos="0 0 0" mass="0.1" diaginertia="0.0008 0.0023 0.0023"/>
+                            <geom class="robot" type="sphere" material="dark_red" size="0.05"/>
+                            <geom class="robot" type="cylinder" material="black" size="0.01 0.075" pos="0 0 -0.075"/>
+                            <body name="planar_robot_1/body_ee" pos="0 0 -0.15">
+                                <inertial pos="0 0 0" mass="0.1" diaginertia="0.0008 0.0023 0.0023"/>
+                                <geom class="robot" name="planar_robot_1/ee" type="cylinder" material="black"
+                                      size="0.04815 0.01" pos="0 0 0.01" friction="0 0 0"/>
+                            </body>
+                        </body>
+                    </body>
+                </body>
+            </body>
+        </body>
+    </worldbody>
+
+    <actuator>
+        <motor name="planar_robot_1/joint_1" joint="planar_robot_1/joint_1" ctrlrange="-100 100"/>
+        <motor name="planar_robot_1/joint_2" joint="planar_robot_1/joint_2" ctrlrange="-50 50"/>
+        <motor name="planar_robot_1/joint_3" joint="planar_robot_1/joint_3" ctrlrange="-30 30"/>
+    </actuator>
+</mujoco>

fancy_gym/envs/mujoco/air_hockey/data/planar/planar_robot_2.xml

@@ -0,0 +1,46 @@
+<mujoco model="planar_robot_2">
+    <compiler autolimits="true" angle="radian"/>
+
+    <worldbody>
+        <body name="planar_robot_2/base" pos="1.51 0 -0.1" quat="0 0 0 1">
+            <geom type="cylinder" material="black" size="0.075 0.2" pos="0 0 0.2" class="visual"/>
+            <inertial pos="0 0 0" mass="1" diaginertia="0.05 0.018 0.044"/>
+            <body name="planar_robot_2/body_1" pos="0 0 0.25">
+                <joint name="planar_robot_2/joint_1" axis="0 0 1" range="-2.9670597283903604 2.9670597283903604"/>
+                <inertial pos="0.25 0 0" mass="3" diaginertia="0.064375 0.064375 0.00375"/>
+                <geom type="cylinder" material="dark_red" size="0.04 0.275" euler="0 1.57079632679 0" pos="0.275 0 0"/>
+                <body name="planar_robot_2/body_2" pos="0.55 0 0">
+                    <joint name="planar_robot_2/joint_2" axis="0 0 1" range="-1.8 1.8"/>
+                    <inertial pos="0.2 0 0" mass="2" diaginertia="0.0335 0.0335 0.003"/>
+                    <geom type="cylinder" material="dark_red" size="0.04 0.22" euler="0 1.57079632679 0"
+                          pos="0.22 0 0"/>
+                    <geom type="sphere" material="dark_red" size="0.05"/>
+                    <body name="planar_robot_2/body_3" pos="0.44 0 0">
+                        <joint name="planar_robot_2/joint_3" axis="0 0 1"
+                               range="-2.0943951023931953 2.0943951023931953"/>
+                        <inertial pos="0.2 0 0" mass="2" diaginertia="0.0335 0.0335 0.003"/>
+                        <geom type="cylinder" material="dark_red" size="0.04 0.22" euler="0 1.57079632679 0"
+                              pos="0.22 0 0"/>
+                        <geom type="sphere" material="dark_red" size="0.05"/>
+                        <body name="planar_robot_2/body_hand" pos="0.44 0 0">
+                            <inertial pos="0 0 0" mass="0.1" diaginertia="0.0008 0.0023 0.0023"/>
+                            <geom type="sphere" material="dark_red" size="0.05"/>
+                            <geom type="cylinder" material="black" size="0.01 0.075" pos="0 0 -0.075"/>
+                            <body name="planar_robot_2/body_ee" pos="0 0 -0.15">
+                                <inertial pos="0 0 0" mass="0.1" diaginertia="0.0008 0.0023 0.0023"/>
+                                <geom name="planar_robot_2/ee" type="cylinder" material="black" size="0.04815 0.01"
+                                      pos="0 0 0.01" friction="0 0 0"/>
+                            </body>
+                        </body>
+                    </body>
+                </body>
+            </body>
+        </body>
+    </worldbody>
+
+    <actuator>
+        <motor name="planar_robot_2/joint_1" joint="planar_robot_2/joint_1" ctrlrange="-100 100"/>
+        <motor name="planar_robot_2/joint_2" joint="planar_robot_2/joint_2" ctrlrange="-50 50"/>
+        <motor name="planar_robot_2/joint_3" joint="planar_robot_2/joint_3" ctrlrange="-30 30"/>
+    </actuator>
+</mujoco>

fancy_gym/envs/mujoco/air_hockey/data/planar/single.xml

@@ -0,0 +1,24 @@
+<mujoco model="AirHockeySingle">
+
+
+    <include file="planar_robot_1.xml"/>
+
+    <include file="../table.xml"/>
+
+
+    <contact>
+        <exclude body1="planar_robot_1/body_ee" body2="table_surface"/>
+        <exclude body1="planar_robot_1/body_hand" body2="table_surface"/>
+        <exclude body1="planar_robot_1/body_3" body2="table_surface"/>
+        <exclude body1="planar_robot_1/body_2" body2="table_surface"/>
+        <exclude body1="planar_robot_1/body_1" body2="table_surface"/>
+        <exclude body1="planar_robot_1/base" body2="table_surface"/>
+        <exclude body1="planar_robot_1/body_ee" body2="rim"/>
+        <exclude body1="planar_robot_1/body_hand" body2="rim"/>
+        <exclude body1="planar_robot_1/body_3" body2="rim"/>
+        <exclude body1="planar_robot_1/body_2" body2="rim"/>
+        <exclude body1="planar_robot_1/body_1" body2="rim"/>
+        <exclude body1="planar_robot_1/base" body2="rim"/>
+    </contact>
+
+</mujoco>

fancy_gym/envs/mujoco/air_hockey/data/table.xml

@@ -0,0 +1,89 @@
+<mujoco model="table">
+
+    <option timestep="0.001" cone="elliptic" impratio="1"/>
+
+    <asset>
+        <material name="grey" specular="0.5" shininess="0.25" rgba="0.8 0.8 0.8 1"/>
+        <material name="white" specular="0.5" shininess="0.25" rgba="1.0 1.0 1.0 1"/>
+        <material name="red" specular="0.5" shininess="0.25" rgba="1.0 0.0 0.0 1"/>
+        <material name="blue" specular="0.5" shininess="0.25" rgba="0.0 0.0 1.0 1"/>
+        <material name="transparent" specular="0.5" shininess="0.25" rgba="0.0 0.0 1.0 0"/>
+
+        <texture type="skybox" builtin="gradient" rgb1="0.3 0.5 0.7" rgb2="0 0 0" width="512"
+                 height="3072"/>
+        <texture type="2d" name="groundplane" builtin="checker" mark="edge" rgb1="0.2 0.3 0.4"
+                 rgb2="0.1 0.2 0.3" markrgb="0.8 0.8 0.8" width="300" height="300"/>
+        <material name="groundplane" texture="groundplane" texuniform="true" texrepeat="5 5"
+                  reflectance="0.2"/>
+    </asset>
+
+    <default>
+        <geom condim="4" solref="0.02 0.3"/>
+        <default class="rim">
+            <geom type="box" material="grey" condim="6" friction="10000 0.0 0.0" priority="1"
+                  solref="-2000000 -250" solimp="0.99 0.999 0.001 0.5 2"/>
+        </default>
+    </default>
+
+    <worldbody>
+        <body name="table">
+            <body name="table_surface">
+                <geom name="surface" type="box" material="white" size="1.064 0.609 0.0505" pos="0 0 -0.0505"
+                      condim="4" friction="0.08 0.08 0.0" priority="1"
+                      solref="0.0125 0.5" solimp="0.9 0.999 0.001 0.5 2"/>
+            </body>
+            <body name="rim">
+<!--                <geom name="rim_home_l" class="rim" size="0.045 0.197 0.005" pos="-1.019 0.322 0.005"/>-->
+<!--                <geom name="rim_home_r" class="rim" size="0.045 0.197 0.005" pos="-1.019 -0.322 0.005"/>-->
+
+                <geom name="rim_home_l" class="rim" size="0.045 0.1945 0.005" pos="-1.019 0.3245 0.005"/>
+                <geom name="rim_home_r" class="rim" size="0.045 0.1945 0.005" pos="-1.019 -0.3245 0.005"/>
+                <geom name="rim_home_bound_l" class="rim" type="cylinder" size="0.005 0.005" pos="-0.979 0.13 0.005"/>
+                <geom name="rim_home_bound_l_tail" class="rim" size="0.0425 0.005 0.005" pos="-1.0215 0.13 0.005"/>
+                <geom name="rim_home_bound_r" class="rim" type="cylinder" size="0.005 0.005" pos="-0.979 -0.13 0.005"/>
+                <geom name="rim_home_bound_r_tail" class="rim" size="0.0425 0.005 0.005" pos="-1.0215 -0.13 0.005"/>
+
+                <geom name="rim_home_top" class="rim" size="0.045 0.519 0.01" pos="-1.019 0 0.02"/>
+
+                <geom name="rim_left" class="rim" size="1.064 0.045 0.015" pos="0 0.564 0.015"/>
+                <geom name="rim_right" class="rim" size="1.064 0.045 0.015" pos="0 -0.564 0.015"/>
+
+<!--                <geom name="rim_away_l" class="rim" size="0.045 0.197 0.005" pos="1.019 0.322 0.005"/>-->
+<!--                <geom name="rim_away_r" class="rim" size="0.045 0.197 0.005" pos="1.019 -0.322 0.005"/>-->
+
+                <geom name="rim_away_l" class="rim" size="0.045 0.1945 0.005" pos="1.019 0.3245 0.005"/>
+                <geom name="rim_away_r" class="rim" size="0.045 0.1945 0.005" pos="1.019 -0.3245 0.005"/>
+                <geom name="rim_away_bound_l" class="rim" type="cylinder" size="0.005 0.005" pos="0.979 0.13 0.005"/>
+                <geom name="rim_away_bound_l_tail" class="rim" size="0.0425 0.005 0.005" pos="1.0215 0.13 0.005"/>
+                <geom name="rim_away_bound_r" class="rim" type="cylinder" size="0.005 0.005" pos="0.979 -0.13 0.005"/>
+                <geom name="rim_away_bound_r_tail" class="rim" size="0.0425 0.005 0.005" pos="1.0215 -0.13 0.005"/>
+                <geom name="rim_away_top" class="rim" size="0.045 0.519 0.01" pos="1.019 0 0.02"/>
+            </body>
+        </body>
+
+        <body name="base" pos="0 0 0">
+            <joint name="puck_x" type="slide" axis="1 0 0" damping="0.005" limited="false"/>
+            <joint name="puck_y" type="slide" axis="0 1 0" damping="0.005" limited="false"/>
+            <joint name="puck_yaw" type="hinge" axis="0 0 1" damping="2e-6" limited="false"/>
+            <body name="puck">
+                <geom pos="0 0 0" name="puck" type="cylinder" material="red" size="0.03165 0.003"
+                      condim="4" priority="0"/>
+                <geom pos="0.02 0 0" type="cylinder" material="blue" size="0.01 0.0031"
+                      condim="4" contype="0" conaffinity="0"/>
+                <inertial pos="0 0 0" mass="0.01" diaginertia="2.5e-6 2.5e-6 5e-6"/>
+            </body>
+        </body>
+
+        <site name="puck_vis" type="ellipsoid" size="0.03165 0.03165 0.004" pos="0 0 -0.05"/>
+        <site name="puck_vis_rot" type="cylinder" size="0.0045 0.006" rgba="1 0 0 1" pos="0 0 -0.05"/>
+    </worldbody>
+
+    <contact>
+        <exclude body1="puck" body2="table_surface"/>
+    </contact>
+
+    <worldbody>
+        <light pos="0 0 3" dir="0 0 -1" directional="true"/>
+        <geom pos="0 0 -0.1" name="floor" size="0 0 0.05" type="plane" material="groundplane"/>
+    </worldbody>
+</mujoco>

fancy_gym/envs/mujoco/air_hockey/position_control_wrapper.py

@@ -0,0 +1,274 @@
+from collections import deque
+
+import mujoco
+import numpy as np
+import scipy
+
+from fancy_gym.envs.mujoco.air_hockey import seven_dof
+from fancy_gym.envs.mujoco.air_hockey import three_dof
+
+
+class PositionControl:
+    def __init__(self, p_gain, d_gain, i_gain, interpolation_order=3, debug=False, *args, **kwargs):
+        """
+            Mixin that adds position controller to mujoco environments.
+
+            Args:
+                p_gain (float): Proportional controller gain
+                d_gain (float): Differential controller gain
+                i_gain (float): Integral controller gain
+                interpolation_order (int, 3): Type of interpolation used, has to correspond to action shape. Order 1-5 are
+                    polynomial interpolation of the degree. Order -1 is linear interpolation of position and velocity.
+                    Set Order to None in order to turn off interpolation. In this case the action has to be a trajectory
+                    of position, velocity and acceleration of the shape (20, 3, n_joints)
+                    In the case of 2 agents it is a tuple, which describes the interpolation order for each agent
+                debug (bool, True): If true it logs the controller performance into controller_record queue. The order of the
+                    entries is desired_pos, current_pos, desired_vel, current_vel, desired_acc, jerk.
+            """
+
+        self.debug = debug
+
+        super(PositionControl, self).__init__(*args, **kwargs)
+
+        self.robot_model = self.env_info['robot']['robot_model']
+        self.robot_data = self.env_info['robot']['robot_data']
+
+        self.p_gain = np.array(p_gain * self.n_agents)
+        self.d_gain = np.array(d_gain * self.n_agents)
+        self.i_gain = np.array(i_gain * self.n_agents)
+
+        self.prev_pos = np.zeros(len(self.actuator_joint_ids))
+        self.prev_vel = np.zeros(len(self.actuator_joint_ids))
+        self.prev_acc = np.zeros(len(self.actuator_joint_ids))
+        self.i_error = np.zeros(len(self.actuator_joint_ids))
+        self.prev_controller_cmd_pos = np.zeros(len(self.actuator_joint_ids))
+
+        self.interp_order = interpolation_order if type(interpolation_order) is tuple else (interpolation_order,)
+
+        self._num_env_joints = len(self.actuator_joint_ids)
+        self.n_robot_joints = self.env_info['robot']["n_joints"]
+
+        self.action_shape = [None] * self.n_agents
+
+        for i in range(self.n_agents):
+            if self.interp_order[i] is None:
+                self.action_shape[i] = (int(self.dt / self._timestep), 3, self.n_robot_joints)
+            elif self.interp_order[i] in [1, 2]:
+                self.action_shape[i] = (self.n_robot_joints,)
+            elif self.interp_order[i] in [3, 4, -1]:
+                self.action_shape[i] = (2, self.n_robot_joints)
+            elif self.interp_order[i] == 5:
+                self.action_shape[i] = (3, self.n_robot_joints)
+
+        self.traj = None
+
+        self.jerk = np.zeros(self._num_env_joints)
+
+        if self.debug:
+            self.controller_record = deque(maxlen=self.info.horizon * self._n_intermediate_steps)
+
+    def _enforce_safety_limits(self, desired_pos, desired_vel):
+        # ROS safe controller
+        pos = self.prev_controller_cmd_pos
+        k = 20
+
+        joint_pos_lim = np.tile(self.env_info['robot']['joint_pos_limit'], (1, self.n_agents))
+        joint_vel_lim = np.tile(self.env_info['robot']['joint_vel_limit'], (1, self.n_agents))
+
+        min_vel = np.minimum(np.maximum(-k * (pos - joint_pos_lim[0]), joint_vel_lim[0]), joint_vel_lim[1])
+
+        max_vel = np.minimum(np.maximum(-k * (pos - joint_pos_lim[1]), joint_vel_lim[0]), joint_vel_lim[1])
+
+        clipped_vel = np.minimum(np.maximum(desired_vel, min_vel), max_vel)
+
+        min_pos = pos + min_vel * self._timestep
+        max_pos = pos + max_vel * self._timestep
+
+        clipped_pos = np.minimum(np.maximum(desired_pos, min_pos), max_pos)
+        self.prev_controller_cmd_pos = clipped_pos.copy()
+
+        return clipped_pos, clipped_vel
+
+    def _controller(self, desired_pos, desired_vel, desired_acc, current_pos, current_vel):
+        clipped_pos, clipped_vel = self._enforce_safety_limits(desired_pos, desired_vel)
+
+        error = (clipped_pos - current_pos)
+
+        self.i_error += self.i_gain * error * self._timestep
+        torque = self.p_gain * error + self.d_gain * (clipped_vel - current_vel) + self.i_error
+
+        # Acceleration FeedForward
+        tau_ff = np.zeros(self.robot_model.nv)
+        for i in range(self.n_agents):
+            robot_joint_ids = np.arange(self.n_robot_joints) + self.n_robot_joints * i
+            self.robot_data.qpos = current_pos[robot_joint_ids]
+            self.robot_data.qvel = current_vel[robot_joint_ids]
+            acc_ff = desired_acc[robot_joint_ids]
+            mujoco.mj_forward(self.robot_model, self.robot_data)
+
+            mujoco.mj_mulM(self.robot_model, self.robot_data, tau_ff, acc_ff)
+            torque[robot_joint_ids] += tau_ff
+
+            # Gravity Compensation and Coriolis and Centrifugal force
+            torque[robot_joint_ids] += self.robot_data.qfrc_bias
+
+            torque[robot_joint_ids] = np.minimum(np.maximum(torque[robot_joint_ids],
+                                                            self.robot_model.actuator_ctrlrange[:, 0]),
+                                                 self.robot_model.actuator_ctrlrange[:, 1])
+
+        if self.debug:
+            self.controller_record.append(
+                np.concatenate([desired_pos, current_pos, desired_vel, current_vel, desired_acc, self.jerk]))
+
+        return torque
+
+    def _interpolate_trajectory(self, interp_order, action, i=0):
+        tf = self.dt
+        prev_pos = self.prev_pos[i*self.n_robot_joints:(i+1)*self.n_robot_joints]
+        prev_vel = self.prev_vel[i*self.n_robot_joints:(i+1)*self.n_robot_joints]
+        prev_acc = self.prev_acc[i*self.n_robot_joints:(i+1)*self.n_robot_joints]
+        if interp_order == 1 and action.ndim == 1:
+            coef = np.array([[1, 0], [1, tf]])
+            results = np.vstack([prev_pos, action])
+        elif interp_order == 2 and action.ndim == 1:
+            coef = np.array([[1, 0, 0], [1, tf, tf ** 2], [0, 1, 0]])
+            if np.linalg.norm(action - prev_pos) < 1e-3:
+                prev_vel = np.zeros_like(prev_vel)
+            results = np.vstack([prev_pos, action, prev_vel])
+        elif interp_order == 3 and action.shape[0] == 2:
+            coef = np.array([[1, 0, 0, 0], [1, tf, tf ** 2, tf ** 3], [0, 1, 0, 0], [0, 1, 2 * tf, 3 * tf ** 2]])
+            results = np.vstack([prev_pos, action[0], prev_vel, action[1]])
+        elif interp_order == 4 and action.shape[0] == 2:
+            coef = np.array([[1, 0, 0, 0, 0], [1, tf, tf ** 2, tf ** 3, tf ** 4],
+                             [0, 1, 0, 0, 0], [0, 1, 2 * tf, 3 * tf ** 2, 4 * tf ** 3],
+                             [0, 0, 2, 0, 0]])
+            results = np.vstack([prev_pos, action[0], prev_vel, action[1], prev_acc])
+        elif interp_order == 5 and action.shape[0] == 3:
+            coef = np.array([[1, 0, 0, 0, 0, 0], [1, tf, tf ** 2, tf ** 3, tf ** 4, tf ** 5],
+                             [0, 1, 0, 0, 0, 0], [0, 1, 2 * tf, 3 * tf ** 2, 4 * tf ** 3, 5 * tf ** 4],
+                             [0, 0, 2, 0, 0, 0], [0, 0, 2, 6 * tf, 12 * tf ** 2, 20 * tf ** 3]])
+            results = np.vstack([prev_pos, action[0], prev_vel, action[1], prev_acc, action[2]])
+        elif interp_order == -1:
+            # Interpolate position and velocity linearly
+            pass
+        else:
+            raise ValueError("Undefined interpolator order or the action dimension does not match!")
+
+        if interp_order > 0:
+            A = scipy.linalg.block_diag(*[coef] * self.n_robot_joints)
+            y = results.reshape(-2, order='F')
+            weights = np.linalg.solve(A, y).reshape(self.n_robot_joints, interp_order + 1)
+            weights_d = np.polynomial.polynomial.polyder(weights, axis=1)
+            weights_dd = np.polynomial.polynomial.polyder(weights_d, axis=1)
+        elif interp_order == -1:
+            weights = np.vstack([prev_pos, (action[0] - prev_pos) / self.dt]).T
+            weights_d = np.vstack([prev_vel, (action[1] - prev_vel) / self.dt]).T
+            weights_dd = np.polynomial.polynomial.polyder(weights_d, axis=1)
+
+        if interp_order in [3, 4, 5]:
+            self.jerk[i*self.n_robot_joints:(i+1)*self.n_robot_joints] = np.abs(weights_dd[:, 1]) + np.abs(weights_dd[:, 0] - prev_acc) / self._timestep
+        else:
+            self.jerk[i*self.n_robot_joints:(i+1)*self.n_robot_joints] = np.ones_like(prev_acc) * np.inf
+
+        self.prev_pos[i*self.n_robot_joints:(i+1)*self.n_robot_joints] = np.polynomial.polynomial.polyval(tf, weights.T)
+        self.prev_vel[i*self.n_robot_joints:(i+1)*self.n_robot_joints] = np.polynomial.polynomial.polyval(tf, weights_d.T)
+        self.prev_acc[i*self.n_robot_joints:(i+1)*self.n_robot_joints] = np.polynomial.polynomial.polyval(tf, weights_dd.T)
+
+        for t in np.linspace(self._timestep, self.dt, self._n_intermediate_steps):
+            q = np.polynomial.polynomial.polyval(t, weights.T)
+            qd = np.polynomial.polynomial.polyval(t, weights_d.T)
+            qdd = np.polynomial.polynomial.polyval(t, weights_dd.T)
+            yield q, qd, qdd
+
+    def reset(self, obs=None):
+        obs = super(PositionControl, self).reset(obs)
+        self.prev_pos = self._data.qpos[self.actuator_joint_ids]
+        self.prev_vel = self._data.qvel[self.actuator_joint_ids]
+        self.prev_acc = np.zeros(len(self.actuator_joint_ids))
+        self.i_error = np.zeros(len(self.actuator_joint_ids))
+        self.prev_controller_cmd_pos = self._data.qpos[self.actuator_joint_ids]
+
+        if self.debug:
+            self.controller_record = deque(maxlen=self.info.horizon * self._n_intermediate_steps)
+        return obs
+
+    def _step_init(self, obs, action):
+        super(PositionControl, self)._step_init(obs, action)
+
+        if self.n_agents == 1:
+            self.traj = self._create_traj(self.interp_order[0], action)
+        else:
+            def _traj():
+                traj_1 = self._create_traj(self.interp_order[0], action[0], 0)
+                traj_2 = self._create_traj(self.interp_order[1], action[1], 1)
+
+                for a1, a2 in zip(traj_1, traj_2):
+                    yield np.hstack([a1, a2])
+
+            self.traj = _traj()
+
+    def _create_traj(self, interp_order, action, i=0):
+        if interp_order is None:
+            return iter(action)
+        return self._interpolate_trajectory(interp_order, action, i)
+
+    def _compute_action(self, obs, action):
+        cur_pos, cur_vel = self.get_joints(obs)
+
+        desired_pos, desired_vel, desired_acc = next(self.traj)
+
+        return self._controller(desired_pos, desired_vel, desired_acc, cur_pos, cur_vel)
+
+    def _preprocess_action(self, action):
+        action = super(PositionControl, self)._preprocess_action(action)
+
+        if self.n_agents == 1:
+            assert action.shape == self.action_shape[0], f"Unexpected action shape. Expected {self.action_shape[0]} but got" \
+                                                      f" {action.shape}"
+        else:
+            for i in range(self.n_agents):
+                assert action[i].shape == self.action_shape[i], f"Unexpected action shape. Expected {self.action_shape[i]} but got" \
+                                                          f" {action[i].shape}"
+
+        return action
+
+
+class PositionControlIIWA(PositionControl):
+    def __init__(self, *args, **kwargs):
+        p_gain = [1500., 1500., 1200., 1200., 1000., 1000., 500.]
+        d_gain = [60, 80, 60, 30, 10, 1, 0.5]
+        i_gain = [0, 0, 0, 0, 0, 0, 0]
+
+        super(PositionControlIIWA, self).__init__(p_gain=p_gain, d_gain=d_gain, i_gain=i_gain, *args, **kwargs)
+
+
+class PositionControlPlanar(PositionControl):
+    def __init__(self, *args, **kwargs):
+        p_gain = [960, 480, 240]
+        d_gain = [60, 20, 4]
+        i_gain = [0, 0, 0]
+        super(PositionControlPlanar, self).__init__(p_gain=p_gain, d_gain=d_gain, i_gain=i_gain, *args, **kwargs)
+
+
+class PlanarPositionHit(PositionControlPlanar, three_dof.AirHockeyHit):
+    pass
+
+
+class PlanarPositionDefend(PositionControlPlanar, three_dof.AirHockeyDefend):
+    pass
+
+
+class IiwaPositionHit(PositionControlIIWA, seven_dof.AirHockeyHit):
+    pass
+
+class IiwaPositionHitAirhocKIT2023(PositionControlIIWA, seven_dof.AirHockeyHitAirhocKIT2023):
+    pass
+
+class IiwaPositionDefend(PositionControlIIWA, seven_dof.AirHockeyDefend):
+    pass
+
+class IiwaPositionDefendAirhocKIT2023(PositionControlIIWA, seven_dof.AirHockeyDefendAirhocKIT2023):
+    pass
+
+class IiwaPositionTournament(PositionControlIIWA, seven_dof.AirHockeyTournament):
+    pass

fancy_gym/envs/mujoco/air_hockey/seven_dof/__init__.py

@@ -0,0 +1,4 @@
+from .env_base import AirHockeyBase
+from .tournament import AirHockeyTournament
+from .hit import AirHockeyHit, AirHockeyHitAirhocKIT2023
+from .defend import AirHockeyDefend, AirHockeyDefendAirhocKIT2023

fancy_gym/envs/mujoco/air_hockey/seven_dof/airhockit_base_env.py

@@ -0,0 +1,114 @@
+import numpy as np
+from gymnasium import spaces
+
+from fancy_gym.envs.mujoco.air_hockey.seven_dof.env_single import AirHockeySingle
+from fancy_gym.envs.mujoco.air_hockey.utils import inverse_kinematics, forward_kinematics, jacobian
+
+class AirhocKIT2023BaseEnv(AirHockeySingle):
+    def __init__(self, noise=False, **kwargs):
+        super().__init__(**kwargs)
+        obs_low = np.hstack([[-np.inf] * 37])
+        obs_high = np.hstack([[np.inf] * 37])
+        self.wrapper_obs_space = spaces.Box(low=obs_low, high=obs_high, dtype=np.float64)
+        self.wrapper_act_space = spaces.Box(low=np.repeat(-100., 6), high=np.repeat(100., 6))
+        self.noise = noise
+
+    # We don't need puck yaw observations
+    def filter_obs(self, obs):
+        obs = np.hstack([obs[0:2], obs[3:5], obs[6:12], obs[13:19], obs[20:]])
+        return obs
+
+    def add_noise(self, obs):
+        if not self.noise:
+            return
+        obs[self.env_info["puck_pos_ids"]] += np.random.normal(0, 0.001, 3)
+        obs[self.env_info["puck_vel_ids"]] += np.random.normal(0, 0.1, 3)
+    
+    def reset(self):
+        self.last_acceleration = np.repeat(0., 6)
+        obs = super().reset()
+        self.add_noise(obs)
+        self.interp_pos = obs[self.env_info["joint_pos_ids"]][:-1]
+        self.interp_vel = obs[self.env_info["joint_vel_ids"]][:-1]
+
+        self.last_planned_world_pos = self._fk(self.interp_pos)
+        obs = np.hstack([
+            obs, self.interp_pos, self.interp_vel, self.last_acceleration, self.last_planned_world_pos
+        ])
+        return self.filter_obs(obs)
+
+    def step(self, action):
+        action /= 10
+
+        new_vel = self.interp_vel + action
+        
+        jerk = 2 * (new_vel - self.interp_vel - self.last_acceleration * 0.02) / (0.02 ** 2)
+        new_pos = self.interp_pos + self.interp_vel * 0.02 + (1/2) * self.last_acceleration * (0.02 ** 2) + (1/6) * jerk * (0.02 ** 3)
+        abs_action = np.vstack([np.hstack([new_pos, 0]), np.hstack([new_vel, 0])])
+
+        self.interp_pos = new_pos
+        self.interp_vel = new_vel
+        self.last_acceleration += jerk * 0.02
+
+        obs, rew, done, info = super().step(abs_action)
+        self.add_noise(obs)
+        self.last_planned_world_pos = self._fk(self.interp_pos)
+        obs = np.hstack([
+            obs, self.interp_pos, self.interp_vel, self.last_acceleration, self.last_planned_world_pos
+        ])
+        
+        fatal_rew = self.check_fatal(obs)
+        if fatal_rew != 0:
+            return self.filter_obs(obs), fatal_rew, True, info
+        
+        return self.filter_obs(obs), rew, done, info
+    
+    def check_constraints(self, constraint_values):
+        fatal_rew = 0
+
+        j_pos_constr = constraint_values["joint_pos_constr"]
+        if j_pos_constr.max() > 0:
+            fatal_rew += j_pos_constr.max()
+
+        j_vel_constr = constraint_values["joint_vel_constr"]
+        if j_vel_constr.max() > 0:
+            fatal_rew += j_vel_constr.max()
+
+        ee_constr = constraint_values["ee_constr"]
+        if ee_constr.max() > 0:
+            fatal_rew += ee_constr.max()
+
+        link_constr = constraint_values["link_constr"]
+        if link_constr.max() > 0:
+            fatal_rew += link_constr.max()
+        
+        return -fatal_rew
+
+    def check_fatal(self, obs):
+        fatal_rew = 0
+        
+        q = obs[self.env_info["joint_pos_ids"]]
+        qd = obs[self.env_info["joint_vel_ids"]]
+        constraint_values_obs = self.env_info["constraints"].fun(q, qd)
+        fatal_rew += self.check_constraints(constraint_values_obs)
+
+        return -fatal_rew
+
+    def _fk(self, pos):
+        res, _ = forward_kinematics(self.env_info["robot"]["robot_model"],
+                                    self.env_info["robot"]["robot_data"], pos)
+        return res.astype(np.float32)
+
+    def _ik(self, world_pos, init_q=None):
+        success, pos = inverse_kinematics(self.env_info["robot"]["robot_model"],
+                                          self.env_info["robot"]["robot_data"],
+                                          world_pos,
+                                          initial_q=init_q)
+        pos = pos.astype(np.float32)
+        assert success
+        return pos
+
+    def _jacobian(self, pos):
+        return jacobian(self.env_info["robot"]["robot_model"],
+                        self.env_info["robot"]["robot_data"],
+                        pos).astype(np.float32)

fancy_gym/envs/mujoco/air_hockey/seven_dof/defend.py

@@ -0,0 +1,166 @@
+import numpy as np
+
+from fancy_gym.envs.mujoco.air_hockey.seven_dof.env_single import AirHockeySingle
+from fancy_gym.envs.mujoco.air_hockey.seven_dof.airhockit_base_env import AirhocKIT2023BaseEnv
+
+
+class AirHockeyDefend(AirHockeySingle):
+    """
+        Class for the air hockey defending task.
+        The agent should stop the puck at the line x=-0.6.
+    """
+    def __init__(self, gamma=0.99, horizon=500, viewer_params={}):
+        self.init_velocity_range = (1, 3)
+        self.start_range = np.array([[0.29, 0.65], [-0.4, 0.4]])  # Table Frame
+        super().__init__(gamma=gamma, horizon=horizon, viewer_params=viewer_params)
+
+    def setup(self, obs):
+        puck_pos = np.random.rand(2) * (self.start_range[:, 1] - self.start_range[:, 0]) + self.start_range[:, 0]
+
+        lin_vel = np.random.uniform(self.init_velocity_range[0], self.init_velocity_range[1])
+        angle = np.random.uniform(-0.5, 0.5)
+
+        puck_vel = np.zeros(3)
+        puck_vel[0] = -np.cos(angle) * lin_vel
+        puck_vel[1] = np.sin(angle) * lin_vel
+        puck_vel[2] = np.random.uniform(-10, 10)
+
+        self._write_data("puck_x_pos", puck_pos[0])
+        self._write_data("puck_y_pos", puck_pos[1])
+        self._write_data("puck_x_vel", puck_vel[0])
+        self._write_data("puck_y_vel", puck_vel[1])
+        self._write_data("puck_yaw_vel", puck_vel[2])
+
+        super().setup(obs)
+
+    def reward(self, state, action, next_state, absorbing):
+        return 0
+
+    def is_absorbing(self, state):
+        puck_pos, puck_vel = self.get_puck(state)
+        # If puck is over the middle line and moving towards opponent
+        if puck_pos[0] > 0 and puck_vel[0] > 0:
+            return True
+        if np.linalg.norm(puck_vel[:2]) < 0.1:
+            return True
+        return super().is_absorbing(state)
+
+class AirHockeyDefendAirhocKIT2023(AirhocKIT2023BaseEnv):
+    def __init__(self, gamma=0.99, horizon=200, viewer_params={}, **kwargs):
+        super().__init__(gamma=gamma, horizon=horizon, viewer_params=viewer_params, **kwargs)
+        self.init_velocity_range = (1, 3)
+        self.start_range = np.array([[0.4, 0.75], [-0.4, 0.4]])  # Table Frame
+        self._setup_metrics()
+        
+    def setup(self, obs):
+        self._setup_metrics()
+        puck_pos = np.random.rand(2) * (self.start_range[:, 1] - self.start_range[:, 0]) + self.start_range[:, 0]
+
+        lin_vel = np.random.uniform(self.init_velocity_range[0], self.init_velocity_range[1])
+        angle = np.random.uniform(-0.5, 0.5)
+
+        puck_vel = np.zeros(3)
+        puck_vel[0] = -np.cos(angle) * lin_vel
+        puck_vel[1] = np.sin(angle) * lin_vel
+        puck_vel[2] = np.random.uniform(-10, 10)
+
+        self._write_data("puck_x_pos", puck_pos[0])
+        self._write_data("puck_y_pos", puck_pos[1])
+        self._write_data("puck_x_vel", puck_vel[0])
+        self._write_data("puck_y_vel", puck_vel[1])
+        self._write_data("puck_yaw_vel", puck_vel[2])
+
+        super().setup(obs)
+
+    def reset(self, *args):
+        obs = super().reset()
+        self.hit_step_flag = False
+        self.hit_step = False
+        self.received_hit_reward = False
+        self.give_reward_next = False
+        return obs
+
+    def _setup_metrics(self):
+        self.episode_steps = 0
+        self.has_hit = False
+
+    def _simulation_post_step(self):
+        if not self.has_hit:
+            self.has_hit = self._check_collision("puck", "robot_1/ee")
+
+        super()._simulation_post_step()
+
+    def _step_finalize(self):
+        self.episode_steps += 1
+        return super()._step_finalize()
+    
+    def reward(self, state, action, next_state, absorbing):
+        puck_pos, puck_vel = self.get_puck(next_state)
+        ee_pos, _ = self.get_ee()
+        rew = 0.01
+        if -0.7 < puck_pos[0] <= -0.2 and np.linalg.norm(puck_vel[:2]) < 0.1:
+            assert absorbing
+            rew += 70
+
+        if self.has_hit and not self.hit_step_flag:
+            self.hit_step_flag = True
+            self.hit_step = True
+        else:
+            self.hit_step = False
+
+        f = lambda puck_vel: 30 + 100 * (100 ** (-0.25 * np.linalg.norm(puck_vel[:2])))
+        if not self.give_reward_next and not self.received_hit_reward and self.hit_step and ee_pos[0] < puck_pos[0]:
+            self.hit_this_step = True
+            if np.linalg.norm(puck_vel[:2]) < 0.1:
+                return rew + f(puck_vel)
+            self.give_reward_next = True
+            return rew
+
+        if not self.received_hit_reward and self.give_reward_next:
+            self.received_hit_reward = True
+            if puck_vel[0] >= -0.2:
+                return rew + f(puck_vel)
+            return rew
+        else:
+            return rew
+
+    def is_absorbing(self, obs):
+        puck_pos, puck_vel = self.get_puck(obs)
+        # If puck is over the middle line and moving towards opponent
+        if puck_pos[0] > 0 and puck_vel[0] > 0:
+            return True
+
+        if self.episode_steps == self._mdp_info.horizon:
+            return True
+            
+        if np.linalg.norm(puck_vel[:2]) < 0.1:
+            return True
+        return super().is_absorbing(obs)
+
+
+if __name__ == '__main__':
+    env = AirHockeyDefend()
+
+    R = 0.
+    J = 0.
+    gamma = 1.
+    steps = 0
+    env.reset()
+    env.render()
+    while True:
+        # action = np.random.uniform(-1, 1, env.info.action_space.low.shape) * 8
+        action = np.zeros(7)
+        observation, reward, done, info = env.step(action)
+        env.render()
+        print(observation)
+        gamma *= env.info.gamma
+        J += gamma * reward
+        R += reward
+        steps += 1
+        if done or steps > env.info.horizon:
+            print("J: ", J, " R: ", R)
+            R = 0.
+            J = 0.
+            gamma = 1.
+            steps = 0
+            env.reset()

fancy_gym/envs/mujoco/air_hockey/seven_dof/env_base.py

@@ -0,0 +1,258 @@
+import os
+
+import mujoco
+import numpy as np
+from scipy.spatial.transform import Rotation as R
+
+from fancy_gym.envs.mujoco.air_hockey.data.iiwas import __file__ as env_path
+from fancy_gym.envs.mujoco.air_hockey.utils.universal_joint_plugin import UniversalJointPlugin
+from mushroom_rl.environments.mujoco import MuJoCo, ObservationType
+from mushroom_rl.utils.spaces import Box
+
+
+"""
+    Abstract class for all AirHockey Environments.
+
+"""
+class AirHockeyBase(MuJoCo):
+    def __init__(self, gamma=0.99, horizon=500, timestep=1 / 1000., n_intermediate_steps=20, n_substeps=1,
+                 n_agents=1, viewer_params={}):
+
+        """
+            Constructor.
+
+            Args:
+                n_agents (int, 1): number of agent to be used in the environment (one or two)
+        """
+
+        self.n_agents = n_agents
+
+        action_spec = []
+        observation_spec = [("puck_x_pos", "puck_x", ObservationType.JOINT_POS),
+                            ("puck_y_pos", "puck_y", ObservationType.JOINT_POS),
+                            ("puck_yaw_pos", "puck_yaw", ObservationType.JOINT_POS),
+                            ("puck_x_vel", "puck_x", ObservationType.JOINT_VEL),
+                            ("puck_y_vel", "puck_y", ObservationType.JOINT_VEL),
+                            ("puck_yaw_vel", "puck_yaw", ObservationType.JOINT_VEL)]
+
+        additional_data = [("puck_x_pos", "puck_x", ObservationType.JOINT_POS),
+                           ("puck_y_pos", "puck_y", ObservationType.JOINT_POS),
+                           ("puck_yaw_pos", "puck_yaw", ObservationType.JOINT_POS),
+                           ("puck_x_vel", "puck_x", ObservationType.JOINT_VEL),
+                           ("puck_y_vel", "puck_y", ObservationType.JOINT_VEL),
+                           ("puck_yaw_vel", "puck_yaw", ObservationType.JOINT_VEL)]
+
+        collision_spec = [("puck", ["puck"]),
+                          ("rim", ["rim_home_l", "rim_home_r", "rim_away_l", "rim_away_r", "rim_left", "rim_right"]),
+                          ("rim_short_sides", ["rim_home_l", "rim_home_r", "rim_away_l", "rim_away_r"])]
+
+        if 1 <= self.n_agents <= 2:
+            scene = os.path.join(os.path.dirname(os.path.abspath(env_path)), "single.xml")
+
+            action_spec += ["iiwa_1/joint_1", "iiwa_1/joint_2", "iiwa_1/joint_3", "iiwa_1/joint_4", "iiwa_1/joint_5",
+                            "iiwa_1/joint_6", "iiwa_1/joint_7"]
+
+            observation_spec += [("robot_1/joint_1_pos", "iiwa_1/joint_1", ObservationType.JOINT_POS),
+                                 ("robot_1/joint_2_pos", "iiwa_1/joint_2", ObservationType.JOINT_POS),
+                                 ("robot_1/joint_3_pos", "iiwa_1/joint_3", ObservationType.JOINT_POS),
+                                 ("robot_1/joint_4_pos", "iiwa_1/joint_4", ObservationType.JOINT_POS),
+                                 ("robot_1/joint_5_pos", "iiwa_1/joint_5", ObservationType.JOINT_POS),
+                                 ("robot_1/joint_6_pos", "iiwa_1/joint_6", ObservationType.JOINT_POS),
+                                 ("robot_1/joint_7_pos", "iiwa_1/joint_7", ObservationType.JOINT_POS),
+                                 ("robot_1/joint_1_vel", "iiwa_1/joint_1", ObservationType.JOINT_VEL),
+                                 ("robot_1/joint_2_vel", "iiwa_1/joint_2", ObservationType.JOINT_VEL),
+                                 ("robot_1/joint_3_vel", "iiwa_1/joint_3", ObservationType.JOINT_VEL),
+                                 ("robot_1/joint_4_vel", "iiwa_1/joint_4", ObservationType.JOINT_VEL),
+                                 ("robot_1/joint_5_vel", "iiwa_1/joint_5", ObservationType.JOINT_VEL),
+                                 ("robot_1/joint_6_vel", "iiwa_1/joint_6", ObservationType.JOINT_VEL),
+                                 ("robot_1/joint_7_vel", "iiwa_1/joint_7", ObservationType.JOINT_VEL)]
+
+            additional_data += [("robot_1/joint_8_pos", "iiwa_1/striker_joint_1", ObservationType.JOINT_POS),
+                                ("robot_1/joint_9_pos", "iiwa_1/striker_joint_2", ObservationType.JOINT_POS),
+                                ("robot_1/joint_8_vel", "iiwa_1/striker_joint_1", ObservationType.JOINT_VEL),
+                                ("robot_1/joint_9_vel", "iiwa_1/striker_joint_2", ObservationType.JOINT_VEL),
+                                ("robot_1/ee_pos", "iiwa_1/striker_mallet", ObservationType.BODY_POS),
+                                ("robot_1/ee_vel", "iiwa_1/striker_mallet", ObservationType.BODY_VEL),
+                                ("robot_1/rod_rot", "iiwa_1/striker_joint_link", ObservationType.BODY_ROT)]
+
+            collision_spec += [("robot_1/ee", ["iiwa_1/ee"])]
+
+            if self.n_agents == 2:
+                scene = os.path.join(os.path.dirname(os.path.abspath(env_path)), "double.xml")
+
+                observation_spec += [("robot_1/opponent_ee_pos", "iiwa_2/striker_joint_link", ObservationType.BODY_POS)]
+
+                action_spec += ["iiwa_2/joint_1", "iiwa_2/joint_2", "iiwa_2/joint_3", "iiwa_2/joint_4",
+                                "iiwa_2/joint_5",
+                                "iiwa_2/joint_6", "iiwa_2/joint_7"]
+
+                observation_spec += [("robot_2/puck_x_pos", "puck_x", ObservationType.JOINT_POS),
+                                     ("robot_2/puck_y_pos", "puck_y", ObservationType.JOINT_POS),
+                                     ("robot_2/puck_yaw_pos", "puck_yaw", ObservationType.JOINT_POS),
+                                     ("robot_2/puck_x_vel", "puck_x", ObservationType.JOINT_VEL),
+                                     ("robot_2/puck_y_vel", "puck_y", ObservationType.JOINT_VEL),
+                                     ("robot_2/puck_yaw_vel", "puck_yaw", ObservationType.JOINT_VEL),
+                                     ("robot_2/joint_1_pos", "iiwa_2/joint_1", ObservationType.JOINT_POS),
+                                     ("robot_2/joint_2_pos", "iiwa_2/joint_2", ObservationType.JOINT_POS),
+                                     ("robot_2/joint_3_pos", "iiwa_2/joint_3", ObservationType.JOINT_POS),
+                                     ("robot_2/joint_4_pos", "iiwa_2/joint_4", ObservationType.JOINT_POS),
+                                     ("robot_2/joint_5_pos", "iiwa_2/joint_5", ObservationType.JOINT_POS),
+                                     ("robot_2/joint_6_pos", "iiwa_2/joint_6", ObservationType.JOINT_POS),
+                                     ("robot_2/joint_7_pos", "iiwa_2/joint_7", ObservationType.JOINT_POS),
+                                     ("robot_2/joint_1_vel", "iiwa_2/joint_1", ObservationType.JOINT_VEL),
+                                     ("robot_2/joint_2_vel", "iiwa_2/joint_2", ObservationType.JOINT_VEL),
+                                     ("robot_2/joint_3_vel", "iiwa_2/joint_3", ObservationType.JOINT_VEL),
+                                     ("robot_2/joint_4_vel", "iiwa_2/joint_4", ObservationType.JOINT_VEL),
+                                     ("robot_2/joint_5_vel", "iiwa_2/joint_5", ObservationType.JOINT_VEL),
+                                     ("robot_2/joint_6_vel", "iiwa_2/joint_6", ObservationType.JOINT_VEL),
+                                     ("robot_2/joint_7_vel", "iiwa_2/joint_7", ObservationType.JOINT_VEL)]
+
+                observation_spec += [("robot_2/opponent_ee_pos", "iiwa_1/striker_joint_link", ObservationType.BODY_POS)]
+
+                additional_data += [("robot_2/joint_8_pos", "iiwa_2/striker_joint_1", ObservationType.JOINT_POS),
+                                    ("robot_2/joint_9_pos", "iiwa_2/striker_joint_2", ObservationType.JOINT_POS),
+                                    ("robot_2/joint_8_vel", "iiwa_2/striker_joint_1", ObservationType.JOINT_VEL),
+                                    ("robot_2/joint_9_vel", "iiwa_2/striker_joint_2", ObservationType.JOINT_VEL),
+                                    ("robot_2/ee_pos", "iiwa_2/striker_mallet", ObservationType.BODY_POS),
+                                    ("robot_2/ee_vel", "iiwa_2/striker_mallet", ObservationType.BODY_VEL),
+                                    ("robot_2/rod_rot", "iiwa_2/striker_joint_link", ObservationType.BODY_ROT)]
+
+                collision_spec += [("robot_2/ee", ["iiwa_2/ee"])]
+        else:
+            raise ValueError('n_agents should be 1 or 2')
+
+        self.env_info = dict()
+        self.env_info['table'] = {"length": 1.948, "width": 1.038, "goal_width": 0.25}
+        self.env_info['puck'] = {"radius": 0.03165}
+        self.env_info['mallet'] = {"radius": 0.04815}
+        self.env_info['n_agents'] = self.n_agents
+        self.env_info['robot'] = {
+            "n_joints": 7,
+            "ee_desired_height": 0.1645,
+            "joint_vel_limit": np.array([[-85, -85, -100, -75, -130, -135, -135],
+                                         [85, 85, 100, 75, 130, 135, 135]]) / 180. * np.pi,
+            "joint_acc_limit": np.array([[-85, -85, -100, -75, -130, -135, -135],
+                                         [85, 85, 100, 75, 130, 135, 135]]) / 180. * np.pi * 10,
+            "base_frame": [],
+            "universal_height": 0.0645,
+            "control_frequency": 50,
+        }
+
+        self.env_info['puck_pos_ids'] = [0, 1, 2]
+        self.env_info['puck_vel_ids'] = [3, 4, 5]
+        self.env_info['joint_pos_ids'] = [6, 7, 8, 9, 10, 11, 12]
+        self.env_info['joint_vel_ids'] = [13, 14, 15, 16, 17, 18, 19]
+        if self.n_agents == 2:
+            self.env_info['opponent_ee_ids'] = [20, 21, 22]
+        else:
+            self.env_info['opponent_ee_ids'] = []
+
+        max_joint_vel = ([np.inf] * 3 + list(self.env_info["robot"]["joint_vel_limit"][1, :7])) * self.n_agents
+
+        super().__init__(scene, action_spec, observation_spec, gamma, horizon, timestep, n_substeps,
+                         n_intermediate_steps, additional_data, collision_spec, max_joint_vel, **viewer_params)
+
+        # Construct the mujoco model at origin
+        robot_model = mujoco.MjModel.from_xml_path(
+            os.path.join(os.path.dirname(os.path.abspath(env_path)), "iiwa_only.xml"))
+        robot_model.body('iiwa_1/base').pos = np.zeros(3)
+        robot_data = mujoco.MjData(robot_model)
+
+        # Add env_info that requires mujoco models
+        self.env_info['dt'] = self.dt
+        self.env_info["robot"]["joint_pos_limit"] = np.array(
+            [self._model.joint(f"iiwa_1/joint_{i + 1}").range for i in range(7)]).T
+        self.env_info["robot"]["robot_model"] = robot_model
+        self.env_info["robot"]["robot_data"] = robot_data
+        self.env_info["rl_info"] = self.info
+
+        frame_T = np.eye(4)
+        temp = np.zeros((9, 1))
+        mujoco.mju_quat2Mat(temp, self._model.body("iiwa_1/base").quat)
+        frame_T[:3, :3] = temp.reshape(3, 3)
+        frame_T[:3, 3] = self._model.body("iiwa_1/base").pos
+        self.env_info['robot']['base_frame'].append(frame_T.copy())
+
+        if self.n_agents == 2:
+            mujoco.mju_quat2Mat(temp, self._model.body("iiwa_2/base").quat)
+            frame_T[:3, :3] = temp.reshape(3, 3)
+            frame_T[:3, 3] = self._model.body("iiwa_2/base").pos
+            self.env_info['robot']['base_frame'].append(frame_T.copy())
+
+        # Ids of the joint, which are controller by the action space
+        self.actuator_joint_ids = [self._model.joint(name).id for name in action_spec]
+
+        self.universal_joint_plugin = UniversalJointPlugin(self._model, self._data, self.env_info)
+
+    def _modify_mdp_info(self, mdp_info):
+        obs_low = np.array([0, -1, -np.pi, -20., -20., -100,
+                            *np.array([self._model.joint(f"iiwa_1/joint_{i + 1}").range[0]
+                                       for i in range(self.env_info['robot']['n_joints'])]),
+                            *self.env_info['robot']['joint_vel_limit'][0]])
+        obs_high = np.array([3.02, 1, np.pi, 20., 20., 100,
+                             *np.array([self._model.joint(f"iiwa_1/joint_{i + 1}").range[1]
+                                        for i in range(self.env_info['robot']['n_joints'])]),
+                             *self.env_info['robot']['joint_vel_limit'][1]])
+        if self.n_agents == 2:
+            obs_low = np.concatenate([obs_low, [1.5, -1.5, -1.5]])
+            obs_high = np.concatenate([obs_high, [4.5, 1.5, 1.5]])
+        mdp_info.observation_space = Box(obs_low, obs_high)
+        return mdp_info
+
+    def _simulation_pre_step(self):
+        self.universal_joint_plugin.update()
+
+    def is_absorbing(self, obs):
+        boundary = np.array([self.env_info['table']['length'], self.env_info['table']['width']]) / 2
+        puck_pos, puck_vel = self.get_puck(obs)
+
+        if np.any(np.abs(puck_pos[:2]) > boundary) or np.linalg.norm(puck_vel) > 100:
+            return True
+        return False
+
+    @staticmethod
+    def _puck_2d_in_robot_frame(puck_in, robot_frame, type='pose'):
+        if type == 'pose':
+            puck_w = np.eye(4)
+            puck_w[:2, 3] = puck_in[:2]
+            puck_w[:3, :3] = R.from_euler("xyz", [0., 0., puck_in[2]]).as_matrix()
+
+            puck_r = np.linalg.inv(robot_frame) @ puck_w
+            puck_out = np.concatenate([puck_r[:2, 3],
+                                       R.from_matrix(puck_r[:3, :3]).as_euler('xyz')[2:3]])
+
+        if type == 'vel':
+            rot_mat = robot_frame[:3, :3]
+
+            vel_lin = np.array([*puck_in[:2], 0])
+            vel_ang = np.array([0., 0., puck_in[2]])
+
+            vel_lin_r = rot_mat.T @ vel_lin
+            vel_ang_r = rot_mat.T @ vel_ang
+
+            puck_out = np.concatenate([vel_lin_r[:2], vel_ang_r[2:3]])
+        return puck_out
+
+    def get_puck(self, obs):
+        """
+        Getting the puck properties from the observations
+        Args:
+            obs: The current observation
+
+        Returns:
+            ([pos_x, pos_y, yaw], [lin_vel_x, lin_vel_y, yaw_vel])
+
+        """
+        puck_pos = np.concatenate([self.obs_helper.get_from_obs(obs, "puck_x_pos"),
+                                   self.obs_helper.get_from_obs(obs, "puck_y_pos"),
+                                   self.obs_helper.get_from_obs(obs, "puck_yaw_pos")])
+        puck_vel = np.concatenate([self.obs_helper.get_from_obs(obs, "puck_x_vel"),
+                                   self.obs_helper.get_from_obs(obs, "puck_y_vel"),
+                                   self.obs_helper.get_from_obs(obs, "puck_yaw_vel")])
+        return puck_pos, puck_vel
+
+    def get_ee(self):
+        raise NotImplementedError
+
+    def get_joints(self, obs):
+        raise NotImplementedError

fancy_gym/envs/mujoco/air_hockey/seven_dof/env_double.py

@@ -0,0 +1,171 @@
+import mujoco
+import numpy as np
+from scipy.spatial.transform import Rotation as R
+
+from fancy_gym.envs.mujoco.air_hockey.seven_dof.env_base import AirHockeyBase
+from fancy_gym.envs.mujoco.air_hockey.utils import inverse_kinematics
+
+class AirHockeyDouble(AirHockeyBase):
+    """
+    Base class for two agents air hockey tasks.
+    """
+
+    def __init__(self, gamma=0.99, horizon=500, viewer_params={}):
+
+        super().__init__(gamma=gamma, horizon=horizon, n_agents=2, viewer_params=viewer_params)
+
+        self._compute_init_state()
+
+        self.filter_ratio = 0.274
+        self.q_pos_prev = np.zeros(self.env_info["robot"]["n_joints"] * self.env_info["n_agents"])
+        self.q_vel_prev = np.zeros(self.env_info["robot"]["n_joints"] * self.env_info["n_agents"])
+
+    def _compute_init_state(self):
+        init_state = np.array([0., -0.1961, 0., -1.8436, 0., 0.9704, 0.])
+
+        success, self.init_state = inverse_kinematics(self.env_info['robot']['robot_model'],
+                                                      self.env_info['robot']['robot_data'],
+                                                      np.array([0.65, 0., 0.1645]),
+                                                      R.from_euler('xyz', [0, 5 / 6 * np.pi, 0]).as_matrix(),
+                                                      initial_q=init_state)
+
+        assert success is True
+
+    def get_ee(self, robot=1):
+        """
+        Getting the ee properties from the current internal state the selected robot. Can also be obtained via forward kinematics
+        on the current joint position, this function exists to avoid redundant computations.
+        Args:
+            robot: ID of robot, either 1 or 2
+
+        Returns: ([pos_x, pos_y, pos_z], [ang_vel_x, ang_vel_y, ang_vel_z, lin_vel_x, lin_vel_y, lin_vel_z])
+        """
+        ee_pos = self._read_data("robot_" + str(robot) + "/ee_pos")
+
+        ee_vel = self._read_data("robot_" + str(robot) + "/ee_vel")
+
+        return ee_pos, ee_vel
+
+    def get_joints(self, obs, agent=None):
+        """
+        Get joint position and velocity of the robots
+        Can choose the robot with agent = 1 / 2. If agent is None both are returned
+        """
+        if agent:
+            q_pos = np.zeros(7)
+            q_vel = np.zeros(7)
+            for i in range(7):
+                q_pos[i] = self.obs_helper.get_from_obs(obs, "robot_" + str(agent) + "/joint_" + str(i + 1) + "_pos")[0]
+                q_vel[i] = self.obs_helper.get_from_obs(obs, "robot_" + str(agent) + "/joint_" + str(i + 1) + "_vel")[0]
+        else:
+            q_pos = np.zeros(14)
+            q_vel = np.zeros(14)
+            for i in range(7):
+                q_pos[i] = self.obs_helper.get_from_obs(obs, "robot_1/joint_" + str(i + 1) + "_pos")[0]
+                q_vel[i] = self.obs_helper.get_from_obs(obs, "robot_1/joint_" + str(i + 1) + "_vel")[0]
+
+                q_pos[i + 7] = self.obs_helper.get_from_obs(obs, "robot_2/joint_" + str(i + 1) + "_pos")[0]
+                q_vel[i + 7] = self.obs_helper.get_from_obs(obs, "robot_2/joint_" + str(i + 1) + "_vel")[0]
+
+        return q_pos, q_vel
+
+    def _create_observation(self, obs):
+        # Filter the joint velocity
+        q_pos, q_vel = self.get_joints(obs)
+        q_vel_filter = self.filter_ratio * q_vel + (1 - self.filter_ratio) * self.q_vel_prev
+        self.q_pos_prev = q_pos
+        self.q_vel_prev = q_vel_filter
+
+        for i in range(7):
+            self.obs_helper.get_from_obs(obs, "robot_1/joint_" + str(i + 1) + "_vel")[:] = q_vel_filter[i]
+            self.obs_helper.get_from_obs(obs, "robot_2/joint_" + str(i + 1) + "_vel")[:] = q_vel_filter[i + 7]
+
+        # Wrap puck's rotation angle to [-pi, pi)
+        yaw_angle = self.obs_helper.get_from_obs(obs, "puck_yaw_pos")
+        self.obs_helper.get_from_obs(obs, "puck_yaw_pos")[:] = (yaw_angle + np.pi) % (2 * np.pi) - np.pi
+        return obs
+
+    def _modify_observation(self, obs):
+        new_obs = obs.copy()
+
+        puck_pos, puck_vel = self.get_puck(new_obs)
+
+        puck_pos_1 = self._puck_2d_in_robot_frame(puck_pos, self.env_info['robot']['base_frame'][0])
+        puck_vel_1 = self._puck_2d_in_robot_frame(puck_vel, self.env_info['robot']['base_frame'][0], type='vel')
+
+        self.obs_helper.get_from_obs(new_obs, "puck_x_pos")[:] = puck_pos_1[0]
+        self.obs_helper.get_from_obs(new_obs, "puck_y_pos")[:] = puck_pos_1[1]
+        self.obs_helper.get_from_obs(new_obs, "puck_yaw_pos")[:] = puck_pos_1[2]
+
+        self.obs_helper.get_from_obs(new_obs, "puck_x_vel")[:] = puck_vel_1[0]
+        self.obs_helper.get_from_obs(new_obs, "puck_y_vel")[:] = puck_vel_1[1]
+        self.obs_helper.get_from_obs(new_obs, "puck_yaw_vel")[:] = puck_vel_1[2]
+
+        opponent_pos_1 = self.obs_helper.get_from_obs(new_obs, 'robot_1/opponent_ee_pos')
+        self.obs_helper.get_from_obs(new_obs, 'robot_1/opponent_ee_pos')[:] = \
+            (np.linalg.inv(self.env_info['robot']['base_frame'][0]) @ np.concatenate([opponent_pos_1, [1]]))[:3]
+
+        puck_pos_2 = self._puck_2d_in_robot_frame(puck_pos, self.env_info['robot']['base_frame'][1])
+        puck_vel_2 = self._puck_2d_in_robot_frame(puck_vel, self.env_info['robot']['base_frame'][1], type='vel')
+
+        self.obs_helper.get_from_obs(new_obs, "robot_2/puck_x_pos")[:] = puck_pos_2[0]
+        self.obs_helper.get_from_obs(new_obs, "robot_2/puck_y_pos")[:] = puck_pos_2[1]
+        self.obs_helper.get_from_obs(new_obs, "robot_2/puck_yaw_pos")[:] = puck_pos_2[2]
+
+        self.obs_helper.get_from_obs(new_obs, "robot_2/puck_x_vel")[:] = puck_vel_2[0]
+        self.obs_helper.get_from_obs(new_obs, "robot_2/puck_y_vel")[:] = puck_vel_2[1]
+        self.obs_helper.get_from_obs(new_obs, "robot_2/puck_yaw_vel")[:] = puck_vel_2[2]
+
+        opponent_pos_2 = self.obs_helper.get_from_obs(new_obs, 'robot_2/opponent_ee_pos')
+        self.obs_helper.get_from_obs(new_obs, 'robot_2/opponent_ee_pos')[:] = \
+            (np.linalg.inv(self.env_info['robot']['base_frame'][1]) @ np.concatenate([opponent_pos_2, [1]]))[:3]
+
+        return new_obs
+
+    def setup(self, obs):
+        for i in range(7):
+            self._data.joint("iiwa_1/joint_" + str(i + 1)).qpos = self.init_state[i]
+            self._data.joint("iiwa_2/joint_" + str(i + 1)).qpos = self.init_state[i]
+
+            self.q_pos_prev[i] = self.init_state[i]
+            self.q_pos_prev[i + 7] = self.init_state[i]
+            self.q_vel_prev[i] = self._data.joint("iiwa_1/joint_" + str(i + 1)).qvel[0]
+            self.q_vel_prev[i + 7] = self._data.joint("iiwa_2/joint_" + str(i + 1)).qvel[0]
+
+        self.universal_joint_plugin.reset()
+
+        super().setup(obs)
+        # Update body positions, needed for _compute_universal_joint
+        mujoco.mj_fwdPosition(self._model, self._data)
+
+    def reward(self, state, action, next_state, absorbing):
+        return 0
+
+
+def main():
+    env = AirHockeyDouble(viewer_params={'start_paused': True})
+    env.reset()
+    env.render()
+    R = 0.
+    J = 0.
+    gamma = 1.
+    steps = 0
+    while True:
+        action = np.zeros(14)
+        observation, reward, done, info = env.step(action)
+        env.render()
+        gamma *= env.info.gamma
+        J += gamma * reward
+        R += reward
+        steps += 1
+        if done or steps > env.info.horizon:
+            print("J: ", J, " R: ", R)
+            R = 0.
+            J = 0.
+            gamma = 1.
+            steps = 0
+            env.reset()
+
+
+if __name__ == '__main__':
+    main()

fancy_gym/envs/mujoco/air_hockey/seven_dof/env_single.py

@@ -0,0 +1,104 @@
+import mujoco
+import numpy as np
+from scipy.spatial.transform import Rotation as R
+
+from fancy_gym.envs.mujoco.air_hockey.seven_dof.env_base import AirHockeyBase
+from fancy_gym.envs.mujoco.air_hockey.utils import inverse_kinematics
+
+class AirHockeySingle(AirHockeyBase):
+    """
+    Base class for single agent air hockey tasks.
+
+    """
+
+    def __init__(self, gamma=0.99, horizon=500, viewer_params={}):
+
+        super().__init__(gamma=gamma, horizon=horizon, n_agents=1, viewer_params=viewer_params)
+
+        self._compute_init_state()
+
+        self.filter_ratio = 0.274
+        self.q_pos_prev = np.zeros(self.env_info["robot"]["n_joints"])
+        self.q_vel_prev = np.zeros(self.env_info["robot"]["n_joints"])
+
+    def _compute_init_state(self):
+        init_state = np.array([0., -0.1961, 0., -1.8436, 0., 0.9704, 0.])
+
+        success, self.init_state = inverse_kinematics(self.env_info['robot']['robot_model'],
+                                                      self.env_info['robot']['robot_data'],
+                                                      np.array([0.65, 0., 0.1645]),
+                                                      R.from_euler('xyz', [0, 5 / 6 * np.pi, 0]).as_matrix(),
+                                                      initial_q=init_state)
+
+        assert success is True
+
+    def get_ee(self):
+        """
+        Getting the ee properties from the current internal state. Can also be obtained via forward kinematics
+        on the current joint position, this function exists to avoid redundant computations.
+
+        Returns:
+            ([pos_x, pos_y, pos_z], [ang_vel_x, ang_vel_y, ang_vel_z, lin_vel_x, lin_vel_y, lin_vel_z])
+        """
+        ee_pos = self._read_data("robot_1/ee_pos")
+
+        ee_vel = self._read_data("robot_1/ee_vel")
+
+        return ee_pos, ee_vel
+
+    def get_joints(self, obs):
+        """
+        Get joint position and velocity of the robot
+        """
+        q_pos = np.zeros(7)
+        q_vel = np.zeros(7)
+        for i in range(7):
+            q_pos[i] = self.obs_helper.get_from_obs(obs, "robot_1/joint_" + str(i + 1) + "_pos")[0]
+            q_vel[i] = self.obs_helper.get_from_obs(obs, "robot_1/joint_" + str(i + 1) + "_vel")[0]
+
+        return q_pos, q_vel
+
+    def _create_observation(self, obs):
+        # Filter the joint velocity
+        q_pos, q_vel = self.get_joints(obs)
+        q_vel_filter = self.filter_ratio * q_vel + (1 - self.filter_ratio) * self.q_vel_prev
+        self.q_pos_prev = q_pos
+        self.q_vel_prev = q_vel_filter
+
+        for i in range(7):
+            self.obs_helper.get_from_obs(obs, "robot_1/joint_" + str(i + 1) + "_vel")[:] = q_vel_filter[i]
+
+        yaw_angle = self.obs_helper.get_from_obs(obs, "puck_yaw_pos")
+        self.obs_helper.get_from_obs(obs, "puck_yaw_pos")[:] = (yaw_angle + np.pi) % (2 * np.pi) - np.pi
+        return obs
+
+    def _modify_observation(self, obs):
+        new_obs = obs.copy()
+        puck_pos, puck_vel = self.get_puck(new_obs)
+
+        puck_pos = self._puck_2d_in_robot_frame(puck_pos, self.env_info['robot']['base_frame'][0])
+
+        puck_vel = self._puck_2d_in_robot_frame(puck_vel, self.env_info['robot']['base_frame'][0], type='vel')
+
+        self.obs_helper.get_from_obs(new_obs, "puck_x_pos")[:] = puck_pos[0]
+        self.obs_helper.get_from_obs(new_obs, "puck_y_pos")[:] = puck_pos[1]
+        self.obs_helper.get_from_obs(new_obs, "puck_yaw_pos")[:] = puck_pos[2]
+
+        self.obs_helper.get_from_obs(new_obs, "puck_x_vel")[:] = puck_vel[0]
+        self.obs_helper.get_from_obs(new_obs, "puck_y_vel")[:] = puck_vel[1]
+        self.obs_helper.get_from_obs(new_obs, "puck_yaw_vel")[:] = puck_vel[2]
+
+        return new_obs
+
+    def setup(self, obs):
+        for i in range(7):
+            self._data.joint("iiwa_1/joint_" + str(i + 1)).qpos = self.init_state[i]
+            self.q_pos_prev[i] = self.init_state[i]
+            self.q_vel_prev[i] = self._data.joint("iiwa_1/joint_" + str(i + 1)).qvel[0]
+
+        self.universal_joint_plugin.reset()
+
+        super().setup(obs)
+
+        # Update body positions, needed for _compute_universal_joint
+        mujoco.mj_fwdPosition(self._model, self._data)

fancy_gym/envs/mujoco/air_hockey/seven_dof/hit.py

@@ -0,0 +1,159 @@
+import numpy as np
+
+from fancy_gym.envs.mujoco.air_hockey.seven_dof.env_single import AirHockeySingle
+from fancy_gym.envs.mujoco.air_hockey.seven_dof.airhockit_base_env import AirhocKIT2023BaseEnv
+
+
+class AirHockeyHit(AirHockeySingle):
+    """
+        Class for the air hockey hitting task.
+    """
+    def __init__(self, gamma=0.99, horizon=500, moving_init=True, viewer_params={}):
+        """
+            Constructor
+            Args:
+                opponent_agent(Agent, None): Agent which controls the opponent
+                moving_init(bool, False): If true, initialize the puck with inital velocity.
+        """
+        super().__init__(gamma=gamma, horizon=horizon, viewer_params=viewer_params)
+
+        self.moving_init = moving_init
+        hit_width = self.env_info['table']['width'] / 2 - self.env_info['puck']['radius'] - \
+                    self.env_info['mallet']['radius'] * 2
+        self.hit_range = np.array([[-0.7, -0.2], [-hit_width, hit_width]])  # Table Frame
+        self.init_velocity_range = (0, 0.5)  # Table Frame
+        self.init_ee_range = np.array([[0.60, 1.25], [-0.4, 0.4]])  # Robot Frame
+
+
+    def setup(self, obs):
+        # Initial position of the puck
+        puck_pos = np.random.rand(2) * (self.hit_range[:, 1] - self.hit_range[:, 0]) + self.hit_range[:, 0]
+
+        self._write_data("puck_x_pos", puck_pos[0])
+        self._write_data("puck_y_pos", puck_pos[1])
+
+        if self.moving_init:
+            lin_vel = np.random.uniform(self.init_velocity_range[0], self.init_velocity_range[1])
+            angle = np.random.uniform(-np.pi / 2 - 0.1, np.pi / 2 + 0.1)
+            puck_vel = np.zeros(3)
+            puck_vel[0] = -np.cos(angle) * lin_vel
+            puck_vel[1] = np.sin(angle) * lin_vel
+            puck_vel[2] = np.random.uniform(-2, 2)
+
+            self._write_data("puck_x_vel", puck_vel[0])
+            self._write_data("puck_y_vel", puck_vel[1])
+            self._write_data("puck_yaw_vel", puck_vel[2])
+
+        super(AirHockeyHit, self).setup(obs)
+
+    def reward(self, state, action, next_state, absorbing):
+        return 0
+
+    def is_absorbing(self, obs):
+        puck_pos, puck_vel = self.get_puck(obs)
+        # Stop if the puck bounces back on the opponents wall
+        if puck_pos[0] > 0 and puck_vel[0] < 0:
+            return True
+        return super(AirHockeyHit, self).is_absorbing(obs)
+
+class AirHockeyHitAirhocKIT2023(AirhocKIT2023BaseEnv):
+    def __init__(self, gamma=0.99, horizon=500, moving_init=True, viewer_params={}, **kwargs):
+        super().__init__(gamma=gamma, horizon=horizon, viewer_params=viewer_params, **kwargs)
+
+        self.moving_init = moving_init
+        hit_width = self.env_info['table']['width'] / 2 - self.env_info['puck']['radius'] - \
+                    self.env_info['mallet']['radius'] * 2
+        self.hit_range = np.array([[-0.7, -0.2], [-hit_width, hit_width]])  # Table Frame
+        self.init_velocity_range = (0, 0.5)  # Table Frame
+        self.init_ee_range = np.array([[0.60, 1.25], [-0.4, 0.4]])  # Robot Frame
+        self._setup_metrics()
+
+    def reset(self, *args):
+        obs = super().reset()
+        self.last_ee_pos = self.last_planned_world_pos.copy()
+        self.last_ee_pos[0] -= 1.51
+        return obs
+
+    def setup(self, obs):
+        self._setup_metrics()
+        puck_pos = np.random.rand(2) * (self.hit_range[:, 1] - self.hit_range[:, 0]) + self.hit_range[:, 0]
+
+        self._write_data("puck_x_pos", puck_pos[0])
+        self._write_data("puck_y_pos", puck_pos[1])
+
+        if self.moving_init:
+            lin_vel = np.random.uniform(self.init_velocity_range[0], self.init_velocity_range[1])
+            angle = np.random.uniform(-np.pi / 2 - 0.1, np.pi / 2 + 0.1)
+            puck_vel = np.zeros(3)
+            puck_vel[0] = -np.cos(angle) * lin_vel
+            puck_vel[1] = np.sin(angle) * lin_vel
+            puck_vel[2] = np.random.uniform(-2, 2)
+
+            self._write_data("puck_x_vel", puck_vel[0])
+            self._write_data("puck_y_vel", puck_vel[1])
+            self._write_data("puck_yaw_vel", puck_vel[2])
+
+        super().setup(obs)
+
+    def _setup_metrics(self):
+        self.episode_steps = 0
+        self.has_scored = False
+
+    def _step_finalize(self):
+        cur_obs = self._create_observation(self.obs_helper._build_obs(self._data))
+        puck_pos, _ = self.get_puck(cur_obs)  # world frame [x, y, z] and [x', y', z']
+
+        if not self.has_scored:
+            boundary = np.array([self.env_info['table']['length'], self.env_info['table']['width']]) / 2
+            self.has_scored = np.any(np.abs(puck_pos[:2]) > boundary) and puck_pos[0] > 0
+
+        self.episode_steps += 1
+        return super()._step_finalize()
+    
+    def reward(self, state, action, next_state, absorbing):
+        rew = 0
+        puck_pos, puck_vel = self.get_puck(next_state)
+        ee_pos, _ = self.get_ee()
+        ee_vel = (ee_pos - self.last_ee_pos) / 0.02
+        self.last_ee_pos = ee_pos
+
+        if puck_vel[0] < 0.25 and puck_pos[0] < 0:
+            ee_puck_dir = (puck_pos - ee_pos)[:2]
+            ee_puck_dir = ee_puck_dir / np.linalg.norm(ee_puck_dir)
+            rew += 1 * max(0, np.dot(ee_puck_dir, ee_vel[:2]))
+        else:
+            rew += 10 * np.linalg.norm(puck_vel[:2])
+
+        if self.has_scored:
+            rew += 2000 + 5000 * np.linalg.norm(puck_vel[:2])
+        
+        return rew
+
+    def is_absorbing(self, obs):
+        puck_pos, puck_vel = self.get_puck(obs)
+        # Stop if the puck bounces back on the opponents wall
+        if puck_pos[0] > 0 and puck_vel[0] < 0:
+            return True
+
+        if self.has_scored:
+            return True
+
+        if self.episode_steps == self._mdp_info.horizon:
+            return True
+
+        return super().is_absorbing(obs)
+
+
+if __name__ == '__main__':
+    env = AirHockeyHit(moving_init=True)
+    env.reset()
+
+    steps = 0
+    while True:
+        action = np.zeros(7)
+
+        observation, reward, done, info = env.step(action)
+        env.render()
+        if done or steps > env.info.horizon:
+            steps = 0
+            env.reset()

fancy_gym/envs/mujoco/air_hockey/seven_dof/tournament.py

@@ -0,0 +1,111 @@
+import mujoco
+import numpy as np
+
+from fancy_gym.envs.mujoco.air_hockey.seven_dof.env_double import AirHockeyDouble
+
+
+class AirHockeyTournament(AirHockeyDouble):
+    """
+        Class for the air hockey tournament. Consists of 2 robots which should play against each other.
+        When the puck is on one side for more than 15 seconds the puck is reset and the player gets a penalty.
+        If a player accumulates 3 penalties his score is reduced by 1.
+    """
+    def __init__(self, gamma=0.99, horizon=15000, viewer_params={}, agent_name="Agent", opponent_name="Opponent"):
+        self.agent_name = agent_name
+        self.opponent_name = opponent_name
+
+        self.score = [0, 0]
+        self.faults = [0, 0]
+        self.start_side = None
+
+        self.timer = 0
+
+        def custom_render_callback(viewport, context):
+            names = f"Agents \nScores \nFaults "
+            data = f"{self.agent_name} - {self.opponent_name}\n "
+            data += f"{self.score[0]} - {self.score[1]}\n "
+            data += f"{self.faults[0]} - {self.faults[1]}"
+            mujoco.mjr_overlay(mujoco.mjtFont.mjFONT_BIG, mujoco.mjtGridPos.mjGRID_TOPLEFT, viewport, names, data, context)
+
+        viewer_params["custom_render_callback"] = custom_render_callback
+        super().__init__(gamma=gamma, horizon=horizon, viewer_params=viewer_params)
+
+        hit_width = self.env_info['table']['width'] / 2 - self.env_info['puck']['radius'] - \
+                    self.env_info['mallet']['radius'] * 2
+
+        self.hit_range = np.array([[-0.7, -0.2], [-hit_width, hit_width]])  # Table Frame
+
+    def setup(self, obs):
+        if self.start_side == None:
+            self.start_side = np.random.choice([1, -1])
+            self.prev_side = self.start_side
+
+        # Initial position of the puck
+        puck_pos = np.random.rand(2) * (self.hit_range[:, 1] - self.hit_range[:, 0]) + self.hit_range[:, 0]
+
+        self._write_data("puck_x_pos", puck_pos[0] * self.start_side)
+        self._write_data("puck_y_pos", puck_pos[1])
+
+        self.prev_side = self.start_side
+        self.timer = 0
+
+        super(AirHockeyTournament, self).setup(obs)
+
+    def reward(self, state, action, next_state, absorbing):
+        return 0
+
+    def is_absorbing(self, obs):
+        puck_pos, puck_vel = self.get_puck(obs)
+
+        # Puck stuck on one side for more than 15s
+        if np.sign(puck_pos[0]) == self.prev_side:
+            self.timer += self.dt
+        else:
+            self.prev_side *= -1
+            self.timer = 0
+
+        if self.timer > 15.0 and np.abs(puck_pos[0]) >= 0.15:
+            if self.prev_side == -1:
+                self.faults[0] += 1
+                self.start_side = -1
+                if self.faults[0] % 3 == 0:
+                    self.score[1] += 1
+            else:
+                self.faults[1] += 1
+                self.start_side = 1
+                if self.faults[1] % 3 == 0:
+                    self.score[0] += 1
+
+            return True
+
+        # Puck in Goal
+        if (np.abs(puck_pos[1]) - self.env_info['table']['goal_width'] / 2) <= 0:
+            if puck_pos[0] > self.env_info['table']['length'] / 2:
+                self.score[0] += 1
+                self.start_side = -1
+                return True
+
+            if puck_pos[0] < -self.env_info['table']['length'] / 2:
+                self.score[1] += 1
+                self.start_side = 1
+                return True
+
+        # Puck stuck in the middle
+        if np.abs(puck_pos[0]) < 0.15 and np.linalg.norm(puck_vel[0]) < 0.025:
+            return True
+        return super(AirHockeyTournament, self).is_absorbing(obs)
+
+
+if __name__ == '__main__':
+    env = AirHockeyTournament()
+    env.reset()
+
+    steps = 0
+    while True:
+        action = np.zeros(14)
+        steps += 1
+        observation, reward, done, info = env.step(action)
+        env.render()
+        if done or steps > env.info.horizon:
+            steps = 0
+            env.reset()

fancy_gym/envs/mujoco/air_hockey/three_dof/__init__.py

@@ -0,0 +1,3 @@
+from .env_base import AirHockeyBase
+from .defend import AirHockeyDefend
+from .hit import AirHockeyHit

fancy_gym/envs/mujoco/air_hockey/three_dof/defend.py

@@ -0,0 +1,76 @@
+import numpy as np
+
+from fancy_gym.envs.mujoco.air_hockey.three_dof.env_single import AirHockeySingle
+
+
+class AirHockeyDefend(AirHockeySingle):
+    """
+    Class for the air hockey defending task.
+    The agent should stop the puck at the line x=-0.6.
+    """
+
+    def __init__(self, gamma=0.99, horizon=500, viewer_params={}):
+
+        self.init_velocity_range = (1, 3)
+
+        self.start_range = np.array([[0.29, 0.65], [-0.4, 0.4]])  # Table Frame
+        self.init_ee_range = np.array([[0.60, 1.25], [-0.4, 0.4]])  # Robot Frame
+        super().__init__(gamma=gamma, horizon=horizon, viewer_params=viewer_params)
+
+    def setup(self, state=None):
+        # Set initial puck parameters
+        puck_pos = np.random.rand(2) * (self.start_range[:, 1] - self.start_range[:, 0]) + self.start_range[:, 0]
+
+        lin_vel = np.random.uniform(self.init_velocity_range[0], self.init_velocity_range[1])
+        angle = np.random.uniform(-0.5, 0.5)
+
+        puck_vel = np.zeros(3)
+        puck_vel[0] = -np.cos(angle) * lin_vel
+        puck_vel[1] = np.sin(angle) * lin_vel
+        puck_vel[2] = np.random.uniform(-10, 10)
+
+        self._write_data("puck_x_pos", puck_pos[0])
+        self._write_data("puck_y_pos", puck_pos[1])
+        self._write_data("puck_x_vel", puck_vel[0])
+        self._write_data("puck_y_vel", puck_vel[1])
+        self._write_data("puck_yaw_vel", puck_vel[2])
+
+        super(AirHockeyDefend, self).setup(state)
+
+    def reward(self, state, action, next_state, absorbing):
+        return 0
+
+    def is_absorbing(self, state):
+        puck_pos, puck_vel = self.get_puck(state)
+        # If puck is over the middle line and moving towards opponent
+        if puck_pos[0] > 0 and puck_vel[0] > 0:
+            return True
+        if np.linalg.norm(puck_vel[:2]) < 0.1:
+            return True
+        return super().is_absorbing(state)
+
+
+if __name__ == '__main__':
+    env = AirHockeyDefend()
+
+    R = 0.
+    J = 0.
+    gamma = 1.
+    steps = 0
+    env.reset()
+    while True:
+        action = np.zeros(3)
+        observation, reward, done, info = env.step(action)
+        env.render()
+        gamma *= env.info.gamma
+        J += gamma * reward
+        R += reward
+        steps += 1
+
+        if done or steps > env.info.horizon:
+            print("J: ", J, " R: ", R)
+            R = 0.
+            J = 0.
+            gamma = 1.
+            steps = 0
+            env.reset()

fancy_gym/envs/mujoco/air_hockey/three_dof/env_base.py

@@ -0,0 +1,224 @@
+import os
+
+import mujoco
+import numpy as np
+from scipy.spatial.transform import Rotation as R
+
+from fancy_gym.envs.mujoco.air_hockey.data.planar import __file__ as env_path
+from mushroom_rl.environments.mujoco import MuJoCo, ObservationType
+from mushroom_rl.utils.spaces import Box
+
+
+class AirHockeyBase(MuJoCo):
+    """
+    Abstract class for all AirHockey Environments.
+
+    """
+
+    def __init__(self, gamma=0.99, horizon=500, timestep=1 / 1000., n_intermediate_steps=20, n_substeps=1,
+                 n_agents=1, viewer_params={}):
+        """
+        Constructor.
+
+        Args:
+            n_agents (int, 1): number of agent to be used in the environment (one or two)
+        """
+
+        self.n_agents = n_agents
+
+        action_spec = []
+
+        observation_spec = [("puck_x_pos", "puck_x", ObservationType.JOINT_POS),
+                            ("puck_y_pos", "puck_y", ObservationType.JOINT_POS),
+                            ("puck_yaw_pos", "puck_yaw", ObservationType.JOINT_POS),
+                            ("puck_x_vel", "puck_x", ObservationType.JOINT_VEL),
+                            ("puck_y_vel", "puck_y", ObservationType.JOINT_VEL),
+                            ("puck_yaw_vel", "puck_yaw", ObservationType.JOINT_VEL)]
+
+        additional_data = [("puck_x_pos", "puck_x", ObservationType.JOINT_POS),
+                           ("puck_y_pos", "puck_y", ObservationType.JOINT_POS),
+                           ("puck_yaw_pos", "puck_yaw", ObservationType.JOINT_POS),
+                           ("puck_x_vel", "puck_x", ObservationType.JOINT_VEL),
+                           ("puck_y_vel", "puck_y", ObservationType.JOINT_VEL),
+                           ("puck_yaw_vel", "puck_yaw", ObservationType.JOINT_VEL)]
+
+        collision_spec = [("puck", ["puck"]),
+                          ("rim", ["rim_home_l", "rim_home_r", "rim_away_l", "rim_away_r", "rim_left", "rim_right"]),
+                          ("rim_short_sides", ["rim_home_l", "rim_home_r", "rim_away_l", "rim_away_r"])]
+
+        if 1 <= self.n_agents <= 2:
+            scene = os.path.join(os.path.dirname(os.path.abspath(env_path)), "single.xml")
+
+            action_spec += ["planar_robot_1/joint_1", "planar_robot_1/joint_2", "planar_robot_1/joint_3"]
+
+            observation_spec += [("robot_1/joint_1_pos", "planar_robot_1/joint_1", ObservationType.JOINT_POS),
+                                 ("robot_1/joint_2_pos", "planar_robot_1/joint_2", ObservationType.JOINT_POS),
+                                 ("robot_1/joint_3_pos", "planar_robot_1/joint_3", ObservationType.JOINT_POS),
+                                 ("robot_1/joint_1_vel", "planar_robot_1/joint_1", ObservationType.JOINT_VEL),
+                                 ("robot_1/joint_2_vel", "planar_robot_1/joint_2", ObservationType.JOINT_VEL),
+                                 ("robot_1/joint_3_vel", "planar_robot_1/joint_3", ObservationType.JOINT_VEL)]
+
+            additional_data += [("robot_1/ee_pos", "planar_robot_1/body_ee", ObservationType.BODY_POS),
+                                ("robot_1/ee_vel", "planar_robot_1/body_ee", ObservationType.BODY_VEL)]
+
+            collision_spec += [("robot_1/ee", ["planar_robot_1/ee"])]
+
+            if self.n_agents == 2:
+                scene = os.path.join(os.path.dirname(os.path.abspath(env_path)), "double.xml")
+
+                observation_spec += [("robot_1/opponent_ee_pos", "planar_robot_2/body_ee", ObservationType.BODY_POS)]
+
+                action_spec += ["planar_robot_2/joint_1", "planar_robot_2/joint_2", "planar_robot_2/joint_3"]
+                # Add puck pos/vel again to transform into second agents frame
+                observation_spec += [("robot_2/puck_x_pos", "puck_x", ObservationType.JOINT_POS),
+                                     ("robot_2/puck_y_pos", "puck_y", ObservationType.JOINT_POS),
+                                     ("robot_2/puck_yaw_pos", "puck_yaw", ObservationType.JOINT_POS),
+                                     ("robot_2/puck_x_vel", "puck_x", ObservationType.JOINT_VEL),
+                                     ("robot_2/puck_y_vel", "puck_y", ObservationType.JOINT_VEL),
+                                     ("robot_2/puck_yaw_vel", "puck_yaw", ObservationType.JOINT_VEL),
+                                     ("robot_2/joint_1_pos", "planar_robot_2/joint_1", ObservationType.JOINT_POS),
+                                     ("robot_2/joint_2_pos", "planar_robot_2/joint_2", ObservationType.JOINT_POS),
+                                     ("robot_2/joint_3_pos", "planar_robot_2/joint_3", ObservationType.JOINT_POS),
+                                     ("robot_2/joint_1_vel", "planar_robot_2/joint_1", ObservationType.JOINT_VEL),
+                                     ("robot_2/joint_2_vel", "planar_robot_2/joint_2", ObservationType.JOINT_VEL),
+                                     ("robot_2/joint_3_vel", "planar_robot_2/joint_3", ObservationType.JOINT_VEL)]
+
+                observation_spec += [("robot_2/opponent_ee_pos", "planar_robot_1/body_ee", ObservationType.BODY_POS)]
+
+                additional_data += [("robot_2/ee_pos", "planar_robot_2/body_ee", ObservationType.BODY_POS),
+                                    ("robot_2/ee_vel", "planar_robot_2/body_ee", ObservationType.BODY_VEL)]
+
+                collision_spec += [("robot_2/ee", ["planar_robot_2/ee"])]
+        else:
+            raise ValueError('n_agents should be 1 or 2')
+
+        self.env_info = dict()
+        self.env_info['table'] = {"length": 1.948, "width": 1.038, "goal_width": 0.25}
+        self.env_info['puck'] = {"radius": 0.03165}
+        self.env_info['mallet'] = {"radius": 0.04815}
+        self.env_info['n_agents'] = self.n_agents
+        self.env_info['robot'] = {
+            "n_joints": 3,
+            "ee_desired_height": 0.1,
+            "joint_vel_limit": np.array([[-np.pi / 2, -np.pi / 2, -np.pi * 2 / 3],
+                                         [np.pi / 2, np.pi / 2, np.pi * 2 / 3]]),
+
+            "joint_acc_limit": np.array([[-2 * np.pi, -2 * np.pi, -2 * 4 / 3 * np.pi],
+                                         [2 * np.pi, 2 * np.pi, 2 * 4 / 3 * np.pi]]),
+            "base_frame": [],
+            "control_frequency": 50,
+        }
+
+        self.env_info['puck_pos_ids'] = [0, 1, 2]
+        self.env_info['puck_vel_ids'] = [3, 4, 5]
+        self.env_info['joint_pos_ids'] = [6, 7, 8]
+        self.env_info['joint_vel_ids'] = [9, 10, 11]
+        if self.n_agents == 2:
+            self.env_info['opponent_ee_ids'] = [13, 14, 15]
+        else:
+            self.env_info['opponent_ee_ids'] = []
+
+        max_joint_vel = ([np.inf] * 3 + list(self.env_info["robot"]["joint_vel_limit"][1, :3])) * self.n_agents
+
+        super().__init__(scene, action_spec, observation_spec, gamma, horizon, timestep, n_substeps,
+                         n_intermediate_steps, additional_data, collision_spec, max_joint_vel, **viewer_params)
+
+        # Construct the mujoco model at origin
+        robot_model = mujoco.MjModel.from_xml_path(
+            os.path.join(os.path.dirname(os.path.abspath(env_path)), "planar_robot_1.xml"))
+        robot_model.body('planar_robot_1/base').pos = np.zeros(3)
+        robot_data = mujoco.MjData(robot_model)
+
+        # Add env_info that requires mujoco models
+        self.env_info['dt'] = self.dt
+        self.env_info["robot"]["joint_pos_limit"] = np.array(
+            [self._model.joint(f"planar_robot_1/joint_{i + 1}").range for i in range(3)]).T
+        self.env_info["robot"]["robot_model"] = robot_model
+        self.env_info["robot"]["robot_data"] = robot_data
+        self.env_info["rl_info"] = self.info
+
+        frame_T = np.eye(4)
+        temp = np.zeros((9, 1))
+        mujoco.mju_quat2Mat(temp, self._model.body("planar_robot_1/base").quat)
+        frame_T[:3, :3] = temp.reshape(3, 3)
+        frame_T[:3, 3] = self._model.body("planar_robot_1/base").pos
+        self.env_info['robot']['base_frame'].append(frame_T.copy())
+
+        if self.n_agents == 2:
+            mujoco.mju_quat2Mat(temp, self._model.body("planar_robot_2/base").quat)
+            frame_T[:3, :3] = temp.reshape(3, 3)
+            frame_T[:3, 3] = self._model.body("planar_robot_2/base").pos
+            self.env_info['robot']['base_frame'].append(frame_T.copy())
+
+        # Ids of the joint, which are controller by the action space
+        self.actuator_joint_ids = [self._model.joint(name).id for name in action_spec]
+
+    def _modify_mdp_info(self, mdp_info):
+        obs_low = np.array([0, -1, -np.pi, -20., -20., -100,
+                            *np.array([self._model.joint(f"planar_robot_1/joint_{i + 1}").range[0]
+                                       for i in range(self.env_info['robot']['n_joints'])]),
+                            *self.env_info['robot']['joint_vel_limit'][0]])
+        obs_high = np.array([3.02, 1, np.pi, 20., 20., 100,
+                             *np.array([self._model.joint(f"planar_robot_1/joint_{i + 1}").range[1]
+                                        for i in range(self.env_info['robot']['n_joints'])]),
+                             *self.env_info['robot']['joint_vel_limit'][1]])
+        if self.n_agents == 2:
+            obs_low = np.concatenate([obs_low, [1.5, -1.5, -1.5]])
+            obs_high = np.concatenate([obs_high, [4.5, 1.5, 1.5]])
+        mdp_info.observation_space = Box(obs_low, obs_high)
+        return mdp_info
+
+    def is_absorbing(self, obs):
+        boundary = np.array([self.env_info['table']['length'], self.env_info['table']['width']]) / 2
+        puck_pos, puck_vel = self.get_puck(obs)
+
+        if np.any(np.abs(puck_pos[:2]) > boundary) or np.linalg.norm(puck_vel) > 100:
+            return True
+        return False
+
+    @staticmethod
+    def _puck_2d_in_robot_frame(puck_in, robot_frame, type='pose'):
+        if type == 'pose':
+            puck_w = np.eye(4)
+            puck_w[:2, 3] = puck_in[:2]
+            puck_w[:3, :3] = R.from_euler("xyz", [0., 0., puck_in[2]]).as_matrix()
+
+            puck_r = np.linalg.inv(robot_frame) @ puck_w
+            puck_out = np.concatenate([puck_r[:2, 3],
+                                       R.from_matrix(puck_r[:3, :3]).as_euler('xyz')[2:3]])
+
+        if type == 'vel':
+            rot_mat = robot_frame[:3, :3]
+
+            vel_lin = np.array([*puck_in[:2], 0])
+            vel_ang = np.array([0., 0., puck_in[2]])
+
+            vel_lin_r = rot_mat.T @ vel_lin
+            vel_ang_r = rot_mat.T @ vel_ang
+
+            puck_out = np.concatenate([vel_lin_r[:2], vel_ang_r[2:3]])
+        return puck_out
+
+    def get_puck(self, obs):
+        """
+        Getting the puck properties from the observations
+        Args:
+            obs: The current observation
+
+        Returns:
+            ([pos_x, pos_y, yaw], [lin_vel_x, lin_vel_y, yaw_vel])
+
+        """
+        puck_pos = np.concatenate([self.obs_helper.get_from_obs(obs, "puck_x_pos"),
+                                   self.obs_helper.get_from_obs(obs, "puck_y_pos"),
+                                   self.obs_helper.get_from_obs(obs, "puck_yaw_pos")])
+        puck_vel = np.concatenate([self.obs_helper.get_from_obs(obs, "puck_x_vel"),
+                                   self.obs_helper.get_from_obs(obs, "puck_y_vel"),
+                                   self.obs_helper.get_from_obs(obs, "puck_yaw_vel")])
+        return puck_pos, puck_vel
+
+    def get_ee(self):
+        raise NotImplementedError
+
+    def get_joints(self, obs):
+        raise NotImplementedError

fancy_gym/envs/mujoco/air_hockey/three_dof/env_single.py

@@ -0,0 +1,91 @@
+import mujoco
+import numpy as np
+
+from fancy_gym.envs.mujoco.air_hockey.three_dof.env_base import AirHockeyBase
+
+
+class AirHockeySingle(AirHockeyBase):
+    """
+    Base class for single agent air hockey tasks.
+
+    """
+
+    def __init__(self, gamma=0.99, horizon=500, viewer_params={}):
+
+        """
+        Constructor.
+
+        """
+        self.init_state = np.array([-1.15570723, 1.30024401, 1.44280414])
+        super().__init__(gamma=gamma, horizon=horizon, n_agents=1, viewer_params=viewer_params)
+
+        self.filter_ratio = 0.274
+        self.q_pos_prev = np.zeros(self.env_info["robot"]["n_joints"])
+        self.q_vel_prev = np.zeros(self.env_info["robot"]["n_joints"])
+
+    def get_ee(self):
+        """
+        Getting the ee properties from the current internal state. Can also be obtained via forward kinematics
+        on the current joint position, this function exists to avoid redundant computations.
+
+        Returns:
+            ([pos_x, pos_y, pos_z], [ang_vel_x, ang_vel_y, ang_vel_z, lin_vel_x, lin_vel_y, lin_vel_z])
+        """
+        ee_pos = self._read_data("robot_1/ee_pos")
+
+        ee_vel = self._read_data("robot_1/ee_vel")
+
+        return ee_pos, ee_vel
+
+    def get_joints(self, obs):
+        """
+        Get joint position and velocity of the robot
+        """
+        q_pos = np.zeros(3)
+        q_vel = np.zeros(3)
+        for i in range(3):
+            q_pos[i] = self.obs_helper.get_from_obs(obs, "robot_1/joint_" + str(i + 1) + "_pos")[0]
+            q_vel[i] = self.obs_helper.get_from_obs(obs, "robot_1/joint_" + str(i + 1) + "_vel")[0]
+
+        return q_pos, q_vel
+
+    def _modify_observation(self, obs):
+        new_obs = obs.copy()
+        puck_pos, puck_vel = self.get_puck(obs)
+
+        puck_pos = self._puck_2d_in_robot_frame(puck_pos, self.env_info['robot']['base_frame'][0])
+
+        puck_vel = self._puck_2d_in_robot_frame(puck_vel, self.env_info['robot']['base_frame'][0], type='vel')
+
+        self.obs_helper.get_from_obs(new_obs, "puck_x_pos")[:] = puck_pos[0]
+        self.obs_helper.get_from_obs(new_obs, "puck_y_pos")[:] = puck_pos[1]
+        self.obs_helper.get_from_obs(new_obs, "puck_yaw_pos")[:] = puck_pos[2]
+
+        self.obs_helper.get_from_obs(new_obs, "puck_x_vel")[:] = puck_vel[0]
+        self.obs_helper.get_from_obs(new_obs, "puck_y_vel")[:] = puck_vel[1]
+        self.obs_helper.get_from_obs(new_obs, "puck_yaw_vel")[:] = puck_vel[2]
+
+        return new_obs
+
+    def setup(self, state=None):
+        for i in range(3):
+            self._data.joint("planar_robot_1/joint_" + str(i + 1)).qpos = self.init_state[i]
+            self.q_pos_prev[i] = self.init_state[i]
+            self.q_vel_prev[i] = self._data.joint("planar_robot_1/joint_" + str(i + 1)).qvel[0]
+
+        mujoco.mj_fwdPosition(self._model, self._data)
+        super().setup(state)
+
+    def _create_observation(self, obs):
+        # Filter the joint velocity
+        q_pos, q_vel = self.get_joints(obs)
+        q_vel_filter = self.filter_ratio * q_vel + (1 - self.filter_ratio) * self.q_vel_prev
+        self.q_pos_prev = q_pos
+        self.q_vel_prev = q_vel_filter
+
+        for i in range(3):
+            self.obs_helper.get_from_obs(obs, "robot_1/joint_" + str(i + 1) + "_vel")[:] = q_vel_filter[i]
+
+        yaw_angle = self.obs_helper.get_from_obs(obs, "puck_yaw_pos")
+        self.obs_helper.get_from_obs(obs, "puck_yaw_pos")[:] = (yaw_angle + np.pi) % (2 * np.pi) - np.pi
+        return obs

fancy_gym/envs/mujoco/air_hockey/three_dof/hit.py

@@ -0,0 +1,84 @@
+import numpy as np
+
+from fancy_gym.envs.mujoco.air_hockey.three_dof.env_single import AirHockeySingle
+
+class AirHockeyHit(AirHockeySingle):
+    """
+    Class for the air hockey hitting task.
+    """
+
+    def __init__(self, gamma=0.99, horizon=500, moving_init=False, viewer_params={}):
+        """
+        Constructor
+        Args:
+            moving_init(bool, False): If true, initialize the puck with inital velocity.
+        """
+        super().__init__(gamma=gamma, horizon=horizon, viewer_params=viewer_params)
+
+        self.moving_init = moving_init
+        hit_width = self.env_info['table']['width'] / 2 - self.env_info['puck']['radius'] - \
+                    self.env_info['mallet']['radius'] * 2
+        self.hit_range = np.array([[-0.7, -0.2], [-hit_width, hit_width]])  # Table Frame
+        self.init_velocity_range = (0, 0.5)  # Table Frame
+        self.init_ee_range = np.array([[0.60, 1.25], [-0.4, 0.4]])  # Robot Frame
+
+    def setup(self, state=None):
+        # Initial position of the puck
+        puck_pos = np.random.rand(2) * (self.hit_range[:, 1] - self.hit_range[:, 0]) + self.hit_range[:, 0]
+
+        # self.init_state = np.array([-0.9273, 0.9273, np.pi / 2])
+
+        self._write_data("puck_x_pos", puck_pos[0])
+        self._write_data("puck_y_pos", puck_pos[1])
+
+        if self.moving_init:
+            lin_vel = np.random.uniform(self.init_velocity_range[0], self.init_velocity_range[1])
+            angle = np.random.uniform(-np.pi / 2 - 0.1, np.pi / 2 + 0.1)
+            puck_vel = np.zeros(3)
+            puck_vel[0] = -np.cos(angle) * lin_vel
+            puck_vel[1] = np.sin(angle) * lin_vel
+            puck_vel[2] = np.random.uniform(-2, 2, 1)
+
+            self._write_data("puck_x_vel", puck_vel[0])
+            self._write_data("puck_y_vel", puck_vel[1])
+            self._write_data("puck_yaw_vel", puck_vel[2])
+
+        super(AirHockeyHit, self).setup(state)
+
+    def reward(self, state, action, next_state, absorbing):
+        return 0
+
+    def is_absorbing(self, obs):
+        puck_pos, puck_vel = self.get_puck(obs)
+        # Stop if the puck bounces back on the opponents wall
+        if puck_pos[0] > 0 and puck_vel[0] < 0:
+            return True
+        return super(AirHockeyHit, self).is_absorbing(obs)
+
+
+if __name__ == '__main__':
+    env = AirHockeyHit(moving_init=False)
+
+    env.reset()
+    env.render()
+    R = 0.
+    J = 0.
+    gamma = 1.
+    steps = 0
+    while True:
+        action = np.zeros(3)
+
+        observation, reward, done, info = env.step(action)
+        env.render()
+
+        gamma *= env.info.gamma
+        J += gamma * reward
+        R += reward
+        steps += 1
+        if done or steps > env.info.horizon:
+            print("J: ", J, " R: ", R)
+            R = 0.
+            J = 0.
+            gamma = 1.
+            steps = 0
+            env.reset()

fancy_gym/envs/mujoco/air_hockey/utils/__init__.py

@@ -0,0 +1,2 @@
+from .kinematics import inverse_kinematics, forward_kinematics, jacobian
+from .transformations import robot_to_world, world_to_robot