Add mappings from env name to env module (e.g., mpe_environments) i…

…n addition to `all_environments` (#1155)
Farama-Foundation · Jan 18, 2024 · a8224bd · a8224bd
1 parent ada1691
commit a8224bd
Show file tree

Hide file tree

Showing 6 changed files with 130 additions and 104 deletions.
diff --git a/pettingzoo/atari/all_modules.py b/pettingzoo/atari/all_modules.py
@@ -0,0 +1,53 @@
+from pettingzoo.atari import (
+    basketball_pong_v3,
+    boxing_v2,
+    combat_plane_v2,
+    combat_tank_v2,
+    double_dunk_v3,
+    entombed_competitive_v3,
+    entombed_cooperative_v3,
+    flag_capture_v2,
+    foozpong_v3,
+    ice_hockey_v2,
+    joust_v3,
+    mario_bros_v3,
+    maze_craze_v3,
+    othello_v3,
+    pong_v3,
+    quadrapong_v4,
+    space_invaders_v2,
+    space_war_v2,
+    surround_v2,
+    tennis_v3,
+    video_checkers_v4,
+    volleyball_pong_v3,
+    warlords_v3,
+    wizard_of_wor_v3,
+)
+
+atari_environments = {
+    "atari/basketball_pong_v3": basketball_pong_v3,
+    "atari/boxing_v2": boxing_v2,
+    "atari/combat_tank_v2": combat_tank_v2,
+    "atari/combat_plane_v2": combat_plane_v2,
+    "atari/double_dunk_v3": double_dunk_v3,
+    "atari/entombed_competitive_v3": entombed_competitive_v3,
+    "atari/entombed_cooperative_v3": entombed_cooperative_v3,
+    "atari/flag_capture_v2": flag_capture_v2,
+    "atari/foozpong_v3": foozpong_v3,
+    "atari/joust_v3": joust_v3,
+    "atari/ice_hockey_v2": ice_hockey_v2,
+    "atari/maze_craze_v3": maze_craze_v3,
+    "atari/mario_bros_v3": mario_bros_v3,
+    "atari/othello_v3": othello_v3,
+    "atari/pong_v3": pong_v3,
+    "atari/quadrapong_v4": quadrapong_v4,
+    "atari/space_invaders_v2": space_invaders_v2,
+    "atari/space_war_v2": space_war_v2,
+    "atari/surround_v2": surround_v2,
+    "atari/tennis_v3": tennis_v3,
+    "atari/video_checkers_v4": video_checkers_v4,
+    "atari/volleyball_pong_v3": volleyball_pong_v3,
+    "atari/wizard_of_wor_v3": wizard_of_wor_v3,
+    "atari/warlords_v3": warlords_v3,
+}
diff --git a/pettingzoo/butterfly/all_modules.py b/pettingzoo/butterfly/all_modules.py
@@ -0,0 +1,11 @@
+from pettingzoo.butterfly import (
+    cooperative_pong_v5,
+    knights_archers_zombies_v10,
+    pistonball_v6,
+)
+
+butterfly_environments = {
+    "butterfly/knights_archers_zombies_v10": knights_archers_zombies_v10,
+    "butterfly/pistonball_v6": pistonball_v6,
+    "butterfly/cooperative_pong_v5": cooperative_pong_v5,
+}
diff --git a/pettingzoo/classic/all_modules.py b/pettingzoo/classic/all_modules.py
@@ -0,0 +1,25 @@
+from pettingzoo.classic import (
+    chess_v6,
+    connect_four_v3,
+    gin_rummy_v4,
+    go_v5,
+    hanabi_v5,
+    leduc_holdem_v4,
+    rps_v2,
+    texas_holdem_no_limit_v6,
+    texas_holdem_v4,
+    tictactoe_v3,
+)
+
+classic_environments = {
+    "classic/chess_v6": chess_v6,
+    "classic/rps_v2": rps_v2,
+    "classic/connect_four_v3": connect_four_v3,
+    "classic/tictactoe_v3": tictactoe_v3,
+    "classic/leduc_holdem_v4": leduc_holdem_v4,
+    "classic/texas_holdem_v4": texas_holdem_v4,
+    "classic/texas_holdem_no_limit_v6": texas_holdem_no_limit_v6,
+    "classic/gin_rummy_v4": gin_rummy_v4,
+    "classic/go_v5": go_v5,
+    "classic/hanabi_v5": hanabi_v5,
+}
diff --git a/pettingzoo/mpe/all_modules.py b/pettingzoo/mpe/all_modules.py
@@ -0,0 +1,23 @@
+from pettingzoo.mpe import (
+    simple_adversary_v3,
+    simple_crypto_v3,
+    simple_push_v3,
+    simple_reference_v3,
+    simple_speaker_listener_v4,
+    simple_spread_v3,
+    simple_tag_v3,
+    simple_v3,
+    simple_world_comm_v3,
+)
+
+mpe_environments = {
+    "mpe/simple_adversary_v3": simple_adversary_v3,
+    "mpe/simple_crypto_v3": simple_crypto_v3,
+    "mpe/simple_push_v3": simple_push_v3,
+    "mpe/simple_reference_v3": simple_reference_v3,
+    "mpe/simple_speaker_listener_v4": simple_speaker_listener_v4,
+    "mpe/simple_spread_v3": simple_spread_v3,
+    "mpe/simple_tag_v3": simple_tag_v3,
+    "mpe/simple_world_comm_v3": simple_world_comm_v3,
+    "mpe/simple_v3": simple_v3,
+}
diff --git a/pettingzoo/sisl/all_modules.py b/pettingzoo/sisl/all_modules.py
@@ -0,0 +1,7 @@
+from pettingzoo.sisl import multiwalker_v9, pursuit_v4, waterworld_v4
+
+sisl_environments = {
+    "sisl/multiwalker_v9": multiwalker_v9,
+    "sisl/waterworld_v4": waterworld_v4,
+    "sisl/pursuit_v4": pursuit_v4,
+}
diff --git a/pettingzoo/utils/all_modules.py b/pettingzoo/utils/all_modules.py
@@ -1,61 +1,12 @@
-from pettingzoo.atari import (
-    basketball_pong_v3,
-    boxing_v2,
-    combat_plane_v2,
-    combat_tank_v2,
-    double_dunk_v3,
-    entombed_competitive_v3,
-    entombed_cooperative_v3,
-    flag_capture_v2,
-    foozpong_v3,
-    ice_hockey_v2,
-    joust_v3,
-    mario_bros_v3,
-    maze_craze_v3,
-    othello_v3,
-    pong_v3,
-    quadrapong_v4,
-    space_invaders_v2,
-    space_war_v2,
-    surround_v2,
-    tennis_v3,
-    video_checkers_v4,
-    volleyball_pong_v3,
-    warlords_v3,
-    wizard_of_wor_v3,
-)
-from pettingzoo.butterfly import (
-    cooperative_pong_v5,
-    knights_archers_zombies_v10,
-    pistonball_v6,
-)
-from pettingzoo.classic import (
-    chess_v6,
-    connect_four_v3,
-    gin_rummy_v4,
-    go_v5,
-    hanabi_v5,
-    leduc_holdem_v4,
-    rps_v2,
-    texas_holdem_no_limit_v6,
-    texas_holdem_v4,
-    tictactoe_v3,
-)
-from pettingzoo.mpe import (
-    simple_adversary_v3,
-    simple_crypto_v3,
-    simple_push_v3,
-    simple_reference_v3,
-    simple_speaker_listener_v4,
-    simple_spread_v3,
-    simple_tag_v3,
-    simple_v3,
-    simple_world_comm_v3,
-)
-from pettingzoo.sisl import multiwalker_v9, pursuit_v4, waterworld_v4
+from pettingzoo.atari.all_modules import atari_environments
+from pettingzoo.butterfly.all_modules import butterfly_environments
+from pettingzoo.classic.all_modules import classic_environments
+from pettingzoo.mpe.all_modules import mpe_environments
+from pettingzoo.sisl.all_modules import sisl_environments
 
 all_prefixes = ["atari", "classic", "butterfly", "mpe", "sisl"]
 
+# environments which have manual policy scripts, allowing interactive play
 manual_environments = {
     "butterfly/knights_archers_zombies",
     "butterfly/pistonball",
@@ -64,53 +15,9 @@
 }
 
 all_environments = {
-    "atari/basketball_pong_v3": basketball_pong_v3,
-    "atari/boxing_v2": boxing_v2,
-    "atari/combat_tank_v2": combat_tank_v2,
-    "atari/combat_plane_v2": combat_plane_v2,
-    "atari/double_dunk_v3": double_dunk_v3,
-    "atari/entombed_competitive_v3": entombed_competitive_v3,
-    "atari/entombed_cooperative_v3": entombed_cooperative_v3,
-    "atari/flag_capture_v2": flag_capture_v2,
-    "atari/foozpong_v3": foozpong_v3,
-    "atari/joust_v3": joust_v3,
-    "atari/ice_hockey_v2": ice_hockey_v2,
-    "atari/maze_craze_v3": maze_craze_v3,
-    "atari/mario_bros_v3": mario_bros_v3,
-    "atari/othello_v3": othello_v3,
-    "atari/pong_v3": pong_v3,
-    "atari/quadrapong_v4": quadrapong_v4,
-    "atari/space_invaders_v2": space_invaders_v2,
-    "atari/space_war_v2": space_war_v2,
-    "atari/surround_v2": surround_v2,
-    "atari/tennis_v3": tennis_v3,
-    "atari/video_checkers_v4": video_checkers_v4,
-    "atari/volleyball_pong_v3": volleyball_pong_v3,
-    "atari/wizard_of_wor_v3": wizard_of_wor_v3,
-    "atari/warlords_v3": warlords_v3,
-    "classic/chess_v6": chess_v6,
-    "classic/rps_v2": rps_v2,
-    "classic/connect_four_v3": connect_four_v3,
-    "classic/tictactoe_v3": tictactoe_v3,
-    "classic/leduc_holdem_v4": leduc_holdem_v4,
-    "classic/texas_holdem_v4": texas_holdem_v4,
-    "classic/texas_holdem_no_limit_v6": texas_holdem_no_limit_v6,
-    "classic/gin_rummy_v4": gin_rummy_v4,
-    "classic/go_v5": go_v5,
-    "classic/hanabi_v5": hanabi_v5,
-    "butterfly/knights_archers_zombies_v10": knights_archers_zombies_v10,
-    "butterfly/pistonball_v6": pistonball_v6,
-    "butterfly/cooperative_pong_v5": cooperative_pong_v5,
-    "mpe/simple_adversary_v3": simple_adversary_v3,
-    "mpe/simple_crypto_v3": simple_crypto_v3,
-    "mpe/simple_push_v3": simple_push_v3,
-    "mpe/simple_reference_v3": simple_reference_v3,
-    "mpe/simple_speaker_listener_v4": simple_speaker_listener_v4,
-    "mpe/simple_spread_v3": simple_spread_v3,
-    "mpe/simple_tag_v3": simple_tag_v3,
-    "mpe/simple_world_comm_v3": simple_world_comm_v3,
-    "mpe/simple_v3": simple_v3,
-    "sisl/multiwalker_v9": multiwalker_v9,
-    "sisl/waterworld_v4": waterworld_v4,
-    "sisl/pursuit_v4": pursuit_v4,
+    **atari_environments,
+    **butterfly_environments,
+    **classic_environments,
+    **mpe_environments,
+    **sisl_environments,
 }