From 7e537447e1765d102af55695025e3eee18c78354 Mon Sep 17 00:00:00 2001
From: Mustafa Enes Batur <mustafa.enes99@gmail.com>
Date: Sun, 19 Nov 2023 22:45:48 +0100
Subject: [PATCH] Register environments separately

---
 fancy_gym/envs/__init__.py | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

diff --git a/fancy_gym/envs/__init__.py b/fancy_gym/envs/__init__.py
index 438f322..b3fcb8f 100644
--- a/fancy_gym/envs/__init__.py
+++ b/fancy_gym/envs/__init__.py
@@ -290,8 +290,13 @@ register(
     }
 )
 
-register(
-    id='fancy/AirHockey-v0',
-    entry_point='fancy_gym.envs.mujoco:AirHockeyEnv',
-    max_episode_steps=45000
-)
\ No newline at end of file
+# Air Hockey environments
+for env_mode in ["tournament", "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,
+        kwargs={
+            'env_mode': env_mode
+        }
+    )