Merge pull request #87 from kayendns/airhockey-off-screen-rendering

Implemented render_mode 'rgb_array' for AirHockeyEnv

.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/*
+

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,12 +1,12 @@
 <h1 align="center">
   <br>
-  <img src='./icon.svg' width="250px">
+  <img src='https://raw.githubusercontent.com/ALRhub/fancy_gym/master/icon.svg' width="250px">
   <br><br>
   <b>Fancy Gym</b>
   <br><br>
 </h1>
 
-| :exclamation: 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). |
+|  &#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.
@@ -33,57 +33,88 @@ While the overarching objective of MP environments remains the learning of an op
 
 ## 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
-git clone git@github.com:ALRhub/fancy_gym.git
+pip install fancy_gym
-```
-
-2. Go to the folder
-
-```bash
-cd fancy_gym
-```
-
-3. Install with
-
-```bash
-pip install -e .
 ```
 
 We have a few optional dependencies. If you also want to install those use
 
 ```bash
-pip install -e '.[all]' # to install all optional dependencies
+# to install all optional dependencies
-pip install -e '.[dmc,metaworld,box2d,mujoco,mujoco-legacy,jax,testing]' # or choose only those you want
+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 .
+```
+
+We have a few optional dependencies. If you also want to install those use
+```bash
+# to install all optional dependencies
+pip install -e '.[all]'
+
+# or choose only those you want
+pip install -e '.[dmc,box2d,mujoco-legacy,jax,testing]'
+```
+
+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-Based Environments
 
 Regular step based environments added by Fancy Gym are added into the `fancy/` namespace.
 
-| :exclamation: Legacy versions of Fancy Gym used `fancy_gym.make(...)`. This is no longer supported and will raise an Exception on new versions. |
+|  &#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 = gym.make('fancy/Reacher5d-v0')
+env = gym.make('fancy/Reacher5d-v0', render_mode='human')
-# or env = gym.make('metaworld/reach-v2') # fancy_gym allows access to all metaworld ML1 tasks via the metaworld/ NS
+# 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')
+# or env = gym.make('dm_control/ball_in_cup-catch-v0', render_mode='human')
-# or env = gym.make('Reacher-v2')
+# 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()
     observation, reward, terminated, truncated, info = env.step(action)
-    if i % 5 == 0:
-        env.render()
 
     if terminated or truncated:
         observation, info = env.reset()
@@ -117,17 +148,14 @@ Just keep in mind, calling `step()` executes a full trajectory.
 import gymnasium as gym
 import fancy_gym
 
-env = gym.make('fancy_ProMP/Reacher5d-v0')
+env = gym.make('fancy_ProMP/Reacher5d-v0', render_mode="human")
-# or env = gym.make('metaworld_ProDMP/reach-v2')
+# or env = gym.make('metaworld_ProDMP/reach-v2', render_mode="human")
-# or env = gym.make('dm_control_DMP/ball_in_cup-catch-v0')
+# or env = gym.make('dm_control_DMP/ball_in_cup-catch-v0', render_mode="human")
-# or env = gym.make('gym_ProMP/Reacher-v2') # mp versions of envs added directly by gymnasium are in the gym_<MP-type> NS
+# 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
-
-# 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')
 
 # This returns the context information, not the full state observation
 observation, info = env.reset(seed=1)
+env.render()
 
 for i in range(5):
     action = env.action_space.sample()
@@ -165,7 +193,7 @@ 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
@@ -251,7 +279,7 @@ 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/).
+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:
 

fancy_gym/__init__.py

@@ -6,7 +6,10 @@ from .envs.registry import ALL_MOVEMENT_PRIMITIVE_ENVIRONMENTS, MOVEMENT_PRIMITI
 
 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']
 
 

fancy_gym/envs/__init__.py

@@ -291,7 +291,7 @@ register(
 )
 
 # Air Hockey environments
-for env_mode in ["7dof-hit", "7dof-defend", "3dof-hit", "3dof-defend"]:
+for env_mode in ["7dof-hit", "7dof-defend", "3dof-hit", "3dof-defend", "7dof-hit-airhockit2023", "7dof-defend-airhockit2023"]:
     register(
         id=f'fancy/AirHockey-{env_mode}-v0',
         entry_point='fancy_gym.envs.mujoco:AirHockeyEnv',

fancy_gym/envs/mujoco/air_hockey/air_hockey_env_wrapper.py

@@ -8,9 +8,9 @@ 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"], "render_fps": 50}
+    metadata = {"render_modes": ["human", "rgb_array"], "render_fps": 50}
 
-    def __init__(self, env_mode=None, interpolation_order=3, render_mode=None, **kwargs):
+    def __init__(self, env_mode=None, interpolation_order=3, render_mode=None, width=1920, height=1080, **kwargs):
         """
         Environment Constructor
 
@@ -30,7 +30,10 @@ class AirHockeyEnv(Environment):
             "7dof-defend": position.IiwaPositionDefend,
 
             "3dof-hit": position.PlanarPositionHit,
-            "3dof-defend": position.PlanarPositionDefend
+            "3dof-defend": position.PlanarPositionDefend,
+
+            "7dof-hit-airhockit2023": position.IiwaPositionHitAirhocKIT2023,
+            "7dof-defend-airhockit2023": position.IiwaPositionDefendAirhocKIT2023,
         }
 
         if env_mode not in env_dict:
@@ -39,9 +42,25 @@ class AirHockeyEnv(Environment):
         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)
@@ -81,9 +100,6 @@ class AirHockeyEnv(Environment):
         self.env_info['constraints'] = constraint_list
         self.env_info['env_name'] = self.env_name
 
-        self.render_mode = render_mode
-        self.render_human_active = False
-
         super().__init__(self.base_env.info)
 
     def step(self, action):
@@ -118,7 +134,13 @@ class AirHockeyEnv(Environment):
         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)

fancy_gym/envs/mujoco/air_hockey/position_control_wrapper.py

@@ -261,10 +261,14 @@ class PlanarPositionDefend(PositionControlPlanar, three_dof.AirHockeyDefend):
 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

@@ -1,4 +1,4 @@
 from .env_base import AirHockeyBase
 from .tournament import AirHockeyTournament
-from .hit import AirHockeyHit
+from .hit import AirHockeyHit, AirHockeyHitAirhocKIT2023
-from .defend import AirHockeyDefend
+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

@@ -1,6 +1,7 @@
 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):
@@ -10,9 +11,7 @@ class AirHockeyDefend(AirHockeySingle):
     """
     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, obs):
@@ -32,7 +31,7 @@ class AirHockeyDefend(AirHockeySingle):
         self._write_data("puck_y_vel", puck_vel[1])
         self._write_data("puck_yaw_vel", puck_vel[2])
 
-        super(AirHockeyDefend, self).setup(obs)
+        super().setup(obs)
 
     def reward(self, state, action, next_state, absorbing):
         return 0
@@ -46,6 +45,98 @@ class AirHockeyDefend(AirHockeySingle):
             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()

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

@@ -1,6 +1,7 @@
 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):
@@ -14,9 +15,6 @@ class AirHockeyHit(AirHockeySingle):
                 opponent_agent(Agent, None): Agent which controls the opponent
                 moving_init(bool, False): If true, initialize the puck with inital velocity.
         """
-        self.hit_range = np.array([[-0.65, -0.25], [-0.4, 0.4]])  # Table Frame
-        self.init_velocity_range = (0, 0.5)  # Table Frame
-
         super().__init__(gamma=gamma, horizon=horizon, viewer_params=viewer_params)
 
         self.moving_init = moving_init
@@ -58,6 +56,93 @@ class AirHockeyHit(AirHockeySingle):
             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)

fancy_gym/meta/__init__.py

@@ -7,6 +7,13 @@ from ..envs.registry import register
 from . import goal_object_change_mp_wrapper, goal_change_mp_wrapper, goal_endeffector_change_mp_wrapper, \
     object_change_mp_wrapper
 
+try:
+    import metaworld
+except ModuleNotFoundError:
+    print('[FANCY GYM] Metaworld not avaible.')
+else:
+    # Will only get executed, if import succeeds
+
     from . import metaworld_adapter
 
     metaworld_adapter.register_all_ML1()

fancy_gym/meta/metaworld_adapter.py

@@ -11,11 +11,7 @@ import numpy as np
 
 from fancy_gym.utils.env_compatibility import EnvCompatibility
 
-try:
 import metaworld
-except Exception:
-    print('[FANCY GYM] Metaworld not avaible')
-
 
 class FixMetaworldHasIncorrectObsSpaceWrapper(gym.Wrapper, gym.utils.RecordConstructorArgs):
     def __init__(self, env: gym.Env):

pyproject.toml

@@ -0,0 +1,79 @@
+[project]
+name = "fancy_gym"
+version = "0.1.4"
+description = "Fancy Gym: Unifying interface for various RL benchmarks with support for Black Box approaches."
+readme = "README.md"
+authors = [
+    {name = "Fabian Otto", email = "fabian.otto@uni-tuebingen.de"},
+    {name = "Onur Celik", email = "celik@kit.edu"},
+    {name = "Dominik Roth", email = "fancy_gym@dominik-roth.eu"},
+    {name = "Hongyi Zhou", email = "hongyi.zhou@kit.edu"}
+]
+license = { text = "MIT License" }
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: MIT License",
+    "Natural Language :: English",
+    "Operating System :: OS Independent",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.7",
+    "Programming Language :: Python :: 3.8",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11"
+]
+
+dependencies = [
+    "mp_pytorch<=0.1.3",
+    "mujoco==2.3.3",
+    "gymnasium[mujoco]>=0.26.0"
+]
+
+requires-python = ">=3.7"
+
+[project.urls]
+"Homepage" = "https://github.com/ALRhub/fancy_gym/"
+#"Documentation" = "https://github.com/ALRhub/fancy_gym/"
+"Bug Tracker" = "https://github.com/ALRhub/fancy_gym/issues"
+#"Repository" = "https://github.com/ALRhub/fancy_gym/"
+
+[build-system]
+requires = ["setuptools>=61.0.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project.optional-dependencies]
+dmc = ["shimmy[dm-control]", "Shimmy==1.0.0"]
+# PyPi does not allow external dependencies. Metaworld will have to be installed manually until Farama publishes up-to-date version of metaworld on PyPi.
+#metaworld = ["metaworld @ git+https://github.com/Farama-Foundation/Metaworld.git@d155d0051630bb365ea6a824e02c66c068947439#egg=metaworld"]
+box2d = ["gymnasium[box2d]>=0.26.0"]
+mujoco-legacy = ["mujoco-py>=2.1,<2.2", "cython<3"]
+jax = ["jax>=0.4.0", "jaxlib>=0.4.0"]
+
+all = [
+    # include all the optional dependencies
+    "shimmy[dm-control]",
+    "Shimmy==1.0.0",
+    #"metaworld @ git+https://github.com/Farama-Foundation/Metaworld.git@d155d0051630bb365ea6a824e02c66c068947439#egg=metaworld",
+    "mujoco==2.3.3",
+    "gymnasium[box2d,mujoco]>=0.26.0",
+    "mujoco-py>=2.1,<2.2",
+    "cython<3",
+    "jax>=0.4.0",
+    "jaxlib>=0.4.0"
+]
+
+testing = [
+    "pytest",
+    # include all the optional dependencies as well
+    "shimmy[dm-control]",
+    "Shimmy==1.0.0",
+    #"metaworld @ git+https://github.com/Farama-Foundation/Metaworld.git@d155d0051630bb365ea6a824e02c66c068947439#egg=metaworld",
+    "mujoco==2.3.3",
+    "gymnasium[box2d,mujoco]>=0.26.0",
+    "mujoco-py>=2.1,<2.2",
+    "cython<3",
+    "jax>=0.4.0",
+    "jaxlib>=0.4.0"
+]

setup.py

@@ -1,15 +1,17 @@
+# We still provide a setup.py for backwards compatability.
+# But the pyproject.toml should be prefered.
 import itertools
 from pathlib import Path
 from typing import List
 
 from setuptools import setup, find_packages
 
+print('[!] You are currently installing/building fancy_gym via setup.py. This is only provided for backwards-compatability. Please use the pyproject.toml instead.')
+
 # Environment-specific dependencies for dmc and metaworld
 extras = {
     'dmc': ['shimmy[dm-control]', 'Shimmy==1.0.0'],
-    'metaworld': ['mujoco==2.3.3', 'metaworld @ git+https://github.com/Farama-Foundation/Metaworld.git@d155d0051630bb365ea6a824e02c66c068947439#egg=metaworld'],
     'box2d': ['gymnasium[box2d]>=0.26.0'],
-    'mujoco': ['mujoco==2.3.3', 'gymnasium[mujoco]>0.26.0'],
     'mujoco-legacy': ['mujoco-py >=2.1,<2.2', 'cython<3'],
     'jax': ["jax >=0.4.0", "jaxlib >=0.4.0"],
     'mushroom-rl': ['mushroom-rl'],
@@ -36,7 +38,7 @@ def find_package_data(extensions_to_include: List[str]) -> List[str]:
 setup(
     author='Fabian Otto, Onur Celik, Dominik Roth, Hongyi Zhou',
     name='fancy_gym',
-    version='1.0',
+    version='0.1.0',
     classifiers=[
         'Development Status :: 4 - Beta',
         'Intended Audience :: Science/Research',
@@ -49,11 +51,13 @@ setup(
         'Programming Language :: Python :: 3.8',
         'Programming Language :: Python :: 3.9',
         'Programming Language :: Python :: 3.10',
+        'Programming Language :: Python :: 3.11',
     ],
     extras_require=extras,
     install_requires=[
-        'gymnasium>=0.26.0',
+        'mp_pytorch<=0.1.3',
-        'mp_pytorch<=0.1.3'
+        'mujoco==2.3.3',
+        'gymnasium[mujoco]>=0.26.0'
     ],
     packages=[package for package in find_packages(
     ) if package.startswith("fancy_gym")],

test/test_all_gym_builtin_envs.py

@@ -12,7 +12,9 @@ GYM_IDS = [spec.id for spec in gym.envs.registry.values() if
            not isinstance(spec.entry_point, Callable) and
            "fancy_gym" not in spec.entry_point and 'make_bb_env_helper' not in spec.entry_point
            and 'jax' not in spec.id.lower()
-           and 'jax' not in spec.id.lower()
+           and 'shimmy' not in spec.id.lower()
+           and 'ale_py' not in spec.id.lower()
+           and 'tabular' not in spec.id.lower()
            and not re.match(r'GymV2.Environment', spec.id)
            ]
 GYM_MP_IDS = fancy_gym.ALL_DMC_MOVEMENT_PRIMITIVE_ENVIRONMENTS['all']

test/test_black_box.py

@@ -103,7 +103,7 @@ def test_verbosity(mp_type: str, env_wrap: Tuple[str, Type[RawInterfaceWrapper]]
 
     env_step = make(env_id)
     env_step.reset()
-    _obs, _reward, _terminated, _truncated, info = env.step(env.action_space.sample())
+    _obs, _reward, _terminated, _truncated, info = env_step.step(env_step.action_space.sample())
     info_keys_step = info.keys()
 
     assert all(e in info_keys for e in info_keys_step)