diff --git a/fancy_gym/envs/__init__.py b/fancy_gym/envs/__init__.py index c004349..8c0fa23 100644 --- a/fancy_gym/envs/__init__.py +++ b/fancy_gym/envs/__init__.py @@ -291,13 +291,23 @@ register( ) # Air Hockey environments -for env_mode in ["tournament", "7dof-hit", "7dof-defend", "3dof-hit", "3dof-defend"]: +for env_mode in ["7dof-hit", "7dof-defend", "3dof-hit", "3dof-defend"]: register( id=f'fancy/AirHockey-{env_mode}-v0', entry_point='fancy_gym.envs.mujoco:AirHockeyEnv', - max_episode_steps=45000, + max_episode_steps=500, add_mp_types=[], kwargs={ 'env_mode': env_mode } ) + +register( + id=f'fancy/AirHockey-tournament-v0', + entry_point='fancy_gym.envs.mujoco:AirHockeyEnv', + max_episode_steps=15000, + add_mp_types=[], + kwargs={ + 'env_mode': 'tournament' + } +)