Add pickling tests, adapt all envs to be picklable (#928)

Co-authored-by: Ariel Kwiatkowski <[email protected]>
Farama-Foundation · Apr 21, 2023 · 6aca84c · 6aca84c
1 parent c6e19ad
commit 6aca84c
Show file tree

Hide file tree

Showing 50 changed files with 369 additions and 165 deletions.
diff --git a/pettingzoo/atari/base_atari_env.py b/pettingzoo/atari/base_atari_env.py
@@ -49,16 +49,16 @@ def __init__(
         """
         EzPickle.__init__(
             self,
-            game,
-            num_players,
-            mode_num,
-            seed,
-            obs_type,
-            full_action_space,
-            env_name,
-            max_cycles,
-            render_mode,
-            auto_rom_install_path,
+            game=game,
+            num_players=num_players,
+            mode_num=mode_num,
+            seed=seed,
+            obs_type=obs_type,
+            full_action_space=full_action_space,
+            env_name=env_name,
+            max_cycles=max_cycles,
+            render_mode=render_mode,
+            auto_rom_install_path=auto_rom_install_path,
         )
 
         assert obs_type in (

diff --git a/pettingzoo/atari/basketball_pong/basketball_pong.py b/pettingzoo/atari/basketball_pong/basketball_pong.py
@@ -68,7 +68,11 @@
 import os
 from glob import glob
 
-from ..base_atari_env import BaseAtariEnv, base_env_wrapper_fn, parallel_wrapper_fn
+from pettingzoo.atari.base_atari_env import (
+    BaseAtariEnv,
+    base_env_wrapper_fn,
+    parallel_wrapper_fn,
+)
 
 
 def raw_env(num_players=2, **kwargs):

diff --git a/pettingzoo/atari/boxing/boxing.py b/pettingzoo/atari/boxing/boxing.py
@@ -78,7 +78,11 @@
 import os
 from glob import glob
 
-from ..base_atari_env import BaseAtariEnv, base_env_wrapper_fn, parallel_wrapper_fn
+from pettingzoo.atari.base_atari_env import (
+    BaseAtariEnv,
+    base_env_wrapper_fn,
+    parallel_wrapper_fn,
+)
 
 
 def raw_env(**kwargs):

diff --git a/pettingzoo/atari/combat_plane/combat_plane.py b/pettingzoo/atari/combat_plane/combat_plane.py
@@ -88,7 +88,11 @@
 import os
 from glob import glob
 
-from ..base_atari_env import BaseAtariEnv, base_env_wrapper_fn, parallel_wrapper_fn
+from pettingzoo.atari.base_atari_env import (
+    BaseAtariEnv,
+    base_env_wrapper_fn,
+    parallel_wrapper_fn,
+)
 
 avaliable_versions = {
     "bi-plane": 15,

diff --git a/pettingzoo/atari/combat_tank/combat_tank.py b/pettingzoo/atari/combat_tank/combat_tank.py
@@ -86,7 +86,11 @@
 import warnings
 from glob import glob
 
-from ..base_atari_env import BaseAtariEnv, base_env_wrapper_fn, parallel_wrapper_fn
+from pettingzoo.atari.base_atari_env import (
+    BaseAtariEnv,
+    base_env_wrapper_fn,
+    parallel_wrapper_fn,
+)
 
 
 def raw_env(has_maze=True, is_invisible=False, billiard_hit=True, **kwargs):

diff --git a/pettingzoo/atari/double_dunk/double_dunk.py b/pettingzoo/atari/double_dunk/double_dunk.py
@@ -78,7 +78,11 @@
 import os
 from glob import glob
 
-from ..base_atari_env import BaseAtariEnv, base_env_wrapper_fn, parallel_wrapper_fn
+from pettingzoo.atari.base_atari_env import (
+    BaseAtariEnv,
+    base_env_wrapper_fn,
+    parallel_wrapper_fn,
+)
 
 
 def raw_env(**kwargs):

diff --git a/pettingzoo/atari/entombed_competitive/entombed_competitive.py b/pettingzoo/atari/entombed_competitive/entombed_competitive.py
@@ -75,7 +75,11 @@
 import os
 from glob import glob
 
-from ..base_atari_env import BaseAtariEnv, base_env_wrapper_fn, parallel_wrapper_fn
+from pettingzoo.atari.base_atari_env import (
+    BaseAtariEnv,
+    base_env_wrapper_fn,
+    parallel_wrapper_fn,
+)
 
 
 def raw_env(**kwargs):

diff --git a/pettingzoo/atari/entombed_cooperative/entombed_cooperative.py b/pettingzoo/atari/entombed_cooperative/entombed_cooperative.py
@@ -83,7 +83,11 @@
 import os
 from glob import glob
 
-from ..base_atari_env import BaseAtariEnv, base_env_wrapper_fn, parallel_wrapper_fn
+from pettingzoo.atari.base_atari_env import (
+    BaseAtariEnv,
+    base_env_wrapper_fn,
+    parallel_wrapper_fn,
+)
 
 
 def raw_env(**kwargs):

diff --git a/pettingzoo/atari/flag_capture/flag_capture.py b/pettingzoo/atari/flag_capture/flag_capture.py
@@ -72,7 +72,11 @@
 import os
 from glob import glob
 
-from ..base_atari_env import BaseAtariEnv, base_env_wrapper_fn, parallel_wrapper_fn
+from pettingzoo.atari.base_atari_env import (
+    BaseAtariEnv,
+    base_env_wrapper_fn,
+    parallel_wrapper_fn,
+)
 
 
 def raw_env(**kwargs):

diff --git a/pettingzoo/atari/foozpong/foozpong.py b/pettingzoo/atari/foozpong/foozpong.py
@@ -73,7 +73,11 @@
 import os
 from glob import glob
 
-from ..base_atari_env import BaseAtariEnv, base_env_wrapper_fn, parallel_wrapper_fn
+from pettingzoo.atari.base_atari_env import (
+    BaseAtariEnv,
+    base_env_wrapper_fn,
+    parallel_wrapper_fn,
+)
 
 
 def raw_env(num_players=4, **kwargs):

diff --git a/pettingzoo/atari/ice_hockey/ice_hockey.py b/pettingzoo/atari/ice_hockey/ice_hockey.py
@@ -71,7 +71,11 @@
 import os
 from glob import glob
 
-from ..base_atari_env import BaseAtariEnv, base_env_wrapper_fn, parallel_wrapper_fn
+from pettingzoo.atari.base_atari_env import (
+    BaseAtariEnv,
+    base_env_wrapper_fn,
+    parallel_wrapper_fn,
+)
 
 
 def raw_env(**kwargs):

diff --git a/pettingzoo/atari/joust/joust.py b/pettingzoo/atari/joust/joust.py
@@ -75,7 +75,11 @@
 import os
 from glob import glob
 
-from ..base_atari_env import BaseAtariEnv, base_env_wrapper_fn, parallel_wrapper_fn
+from pettingzoo.atari.base_atari_env import (
+    BaseAtariEnv,
+    base_env_wrapper_fn,
+    parallel_wrapper_fn,
+)
 
 
 def raw_env(**kwargs):

diff --git a/pettingzoo/atari/mario_bros/mario_bros.py b/pettingzoo/atari/mario_bros/mario_bros.py
@@ -79,7 +79,11 @@
 import os
 from glob import glob
 
-from ..base_atari_env import BaseAtariEnv, base_env_wrapper_fn, parallel_wrapper_fn
+from pettingzoo.atari.base_atari_env import (
+    BaseAtariEnv,
+    base_env_wrapper_fn,
+    parallel_wrapper_fn,
+)
 
 
 def raw_env(**kwargs):

diff --git a/pettingzoo/atari/maze_craze/maze_craze.py b/pettingzoo/atari/maze_craze/maze_craze.py
@@ -88,7 +88,11 @@
 import warnings
 from glob import glob
 
-from ..base_atari_env import BaseAtariEnv, base_env_wrapper_fn, parallel_wrapper_fn
+from pettingzoo.atari.base_atari_env import (
+    BaseAtariEnv,
+    base_env_wrapper_fn,
+    parallel_wrapper_fn,
+)
 
 avaliable_versions = {
     "robbers": 2,

diff --git a/pettingzoo/atari/othello/othello.py b/pettingzoo/atari/othello/othello.py
@@ -75,7 +75,11 @@
 import os
 from glob import glob
 
-from ..base_atari_env import BaseAtariEnv, base_env_wrapper_fn, parallel_wrapper_fn
+from pettingzoo.atari.base_atari_env import (
+    BaseAtariEnv,
+    base_env_wrapper_fn,
+    parallel_wrapper_fn,
+)
 
 
 def raw_env(**kwargs):

diff --git a/pettingzoo/atari/pong/pong.py b/pettingzoo/atari/pong/pong.py
@@ -70,7 +70,11 @@
 import os
 from glob import glob
 
-from ..base_atari_env import BaseAtariEnv, base_env_wrapper_fn, parallel_wrapper_fn
+from pettingzoo.atari.base_atari_env import (
+    BaseAtariEnv,
+    base_env_wrapper_fn,
+    parallel_wrapper_fn,
+)
 
 avaliable_2p_versions = {
     "classic": 4,

diff --git a/pettingzoo/atari/quadrapong/quadrapong.py b/pettingzoo/atari/quadrapong/quadrapong.py
@@ -65,7 +65,11 @@
 import os
 from glob import glob
 
-from ..base_atari_env import BaseAtariEnv, base_env_wrapper_fn, parallel_wrapper_fn
+from pettingzoo.atari.base_atari_env import (
+    BaseAtariEnv,
+    base_env_wrapper_fn,
+    parallel_wrapper_fn,
+)
 
 
 def raw_env(**kwargs):

diff --git a/pettingzoo/atari/space_invaders/space_invaders.py b/pettingzoo/atari/space_invaders/space_invaders.py
@@ -78,7 +78,11 @@
 import os
 from glob import glob
 
-from ..base_atari_env import BaseAtariEnv, base_env_wrapper_fn, parallel_wrapper_fn
+from pettingzoo.atari.base_atari_env import (
+    BaseAtariEnv,
+    base_env_wrapper_fn,
+    parallel_wrapper_fn,
+)
 
 
 def raw_env(

diff --git a/pettingzoo/atari/space_war/space_war.py b/pettingzoo/atari/space_war/space_war.py
@@ -72,7 +72,11 @@
 import os
 from glob import glob
 
-from ..base_atari_env import BaseAtariEnv, base_env_wrapper_fn, parallel_wrapper_fn
+from pettingzoo.atari.base_atari_env import (
+    BaseAtariEnv,
+    base_env_wrapper_fn,
+    parallel_wrapper_fn,
+)
 
 
 def raw_env(**kwargs):

diff --git a/pettingzoo/atari/surround/surround.py b/pettingzoo/atari/surround/surround.py
@@ -60,7 +60,11 @@
 import os
 from glob import glob
 
-from ..base_atari_env import BaseAtariEnv, base_env_wrapper_fn, parallel_wrapper_fn
+from pettingzoo.atari.base_atari_env import (
+    BaseAtariEnv,
+    base_env_wrapper_fn,
+    parallel_wrapper_fn,
+)
 
 
 def raw_env(**kwargs):

diff --git a/pettingzoo/atari/tennis/tennis.py b/pettingzoo/atari/tennis/tennis.py
@@ -74,7 +74,11 @@
 import os
 from glob import glob
 
-from ..base_atari_env import BaseAtariEnv, base_env_wrapper_fn, parallel_wrapper_fn
+from pettingzoo.atari.base_atari_env import (
+    BaseAtariEnv,
+    base_env_wrapper_fn,
+    parallel_wrapper_fn,
+)
 
 
 def raw_env(**kwargs):

diff --git a/pettingzoo/atari/video_checkers/video_checkers.py b/pettingzoo/atari/video_checkers/video_checkers.py
@@ -64,7 +64,11 @@
 import os
 from glob import glob
 
-from ..base_atari_env import BaseAtariEnv, base_env_wrapper_fn, parallel_wrapper_fn
+from pettingzoo.atari.base_atari_env import (
+    BaseAtariEnv,
+    base_env_wrapper_fn,
+    parallel_wrapper_fn,
+)
 
 
 def raw_env(**kwargs):

diff --git a/pettingzoo/atari/volleyball_pong/volleyball_pong.py b/pettingzoo/atari/volleyball_pong/volleyball_pong.py
@@ -73,7 +73,11 @@
 import os
 from glob import glob
 
-from ..base_atari_env import BaseAtariEnv, base_env_wrapper_fn, parallel_wrapper_fn
+from pettingzoo.atari.base_atari_env import (
+    BaseAtariEnv,
+    base_env_wrapper_fn,
+    parallel_wrapper_fn,
+)
 
 
 def raw_env(num_players=4, **kwargs):

diff --git a/pettingzoo/atari/warlords/warlords.py b/pettingzoo/atari/warlords/warlords.py
@@ -60,7 +60,11 @@
 import os
 from glob import glob
 
-from ..base_atari_env import BaseAtariEnv, base_env_wrapper_fn, parallel_wrapper_fn
+from pettingzoo.atari.base_atari_env import (
+    BaseAtariEnv,
+    base_env_wrapper_fn,
+    parallel_wrapper_fn,
+)
 
 
 def raw_env(**kwargs):

diff --git a/pettingzoo/atari/wizard_of_wor/wizard_of_wor.py b/pettingzoo/atari/wizard_of_wor/wizard_of_wor.py
@@ -66,7 +66,11 @@
 import os
 from glob import glob
 
-from ..base_atari_env import BaseAtariEnv, base_env_wrapper_fn, parallel_wrapper_fn
+from pettingzoo.atari.base_atari_env import (
+    BaseAtariEnv,
+    base_env_wrapper_fn,
+    parallel_wrapper_fn,
+)
 
 
 def raw_env(**kwargs):

diff --git a/pettingzoo/butterfly/knights_archers_zombies/knights_archers_zombies.py b/pettingzoo/butterfly/knights_archers_zombies/knights_archers_zombies.py
@@ -238,20 +238,20 @@ def __init__(
     ):
         EzPickle.__init__(
             self,
-            spawn_rate,
-            num_archers,
-            num_knights,
-            max_zombies,
-            max_arrows,
-            killable_knights,
-            killable_archers,
-            pad_observation,
-            line_death,
-            max_cycles,
-            vector_state,
-            use_typemasks,
-            sequence_space,
-            render_mode,
+            spawn_rate=spawn_rate,
+            num_archers=num_archers,
+            num_knights=num_knights,
+            max_zombies=max_zombies,
+            max_arrows=max_arrows,
+            killable_knights=killable_knights,
+            killable_archers=killable_archers,
+            pad_observation=pad_observation,
+            line_death=line_death,
+            max_cycles=max_cycles,
+            vector_state=vector_state,
+            use_typemasks=use_typemasks,
+            sequence_space=sequence_space,
+            render_mode=render_mode,
         )
         # variable state space
         self.sequence_space = sequence_space