diff --git a/shimmy/atari_env.py b/shimmy/atari_env.py index a7ddef51..0ccf6d9d 100644 --- a/shimmy/atari_env.py +++ b/shimmy/atari_env.py @@ -21,7 +21,7 @@ import numpy as np from gymnasium.error import Error from gymnasium.spaces import Box, Discrete -from gymnasium.utils.ezpickle import EzPickle +from gymnasium.utils import EzPickle if sys.version_info < (3, 11): from typing_extensions import NotRequired, TypedDict diff --git a/shimmy/bsuite_compatibility.py b/shimmy/bsuite_compatibility.py index 9f3e8928..c5e2d291 100644 --- a/shimmy/bsuite_compatibility.py +++ b/shimmy/bsuite_compatibility.py @@ -8,6 +8,7 @@ from bsuite.environments import Environment from gymnasium.core import ObsType from gymnasium.error import UnsupportedMode +from gymnasium.utils import EzPickle from shimmy.utils.dm_env import dm_env_step2gym_step, dm_spec2gym_space @@ -17,7 +18,7 @@ np.int = int # pyright: ignore[reportGeneralTypeIssues] -class BSuiteCompatibilityV0(gymnasium.Env[ObsType, np.ndarray]): +class BSuiteCompatibilityV0(gymnasium.Env[ObsType, np.ndarray], EzPickle): """A compatibility wrapper that converts a BSuite environment into a gymnasium environment. Note: @@ -33,6 +34,7 @@ def __init__( render_mode: str | None = None, ): """Initialises the environment with a render mode along with render information.""" + EzPickle.__init__(self, env, render_mode) self._env = env self.observation_space = dm_spec2gym_space(env.observation_spec()) diff --git a/shimmy/dm_control_compatibility.py b/shimmy/dm_control_compatibility.py index 314a1d26..e0f7068b 100644 --- a/shimmy/dm_control_compatibility.py +++ b/shimmy/dm_control_compatibility.py @@ -16,6 +16,7 @@ from dm_control.rl import control from gymnasium.core import ObsType from gymnasium.envs.mujoco.mujoco_rendering import MujocoRenderer +from gymnasium.utils import EzPickle from shimmy.utils.dm_env import dm_env_step2gym_step, dm_spec2gym_space @@ -27,7 +28,7 @@ class EnvType(Enum): RL_CONTROL = 1 -class DmControlCompatibilityV0(gymnasium.Env[ObsType, np.ndarray]): +class DmControlCompatibilityV0(gymnasium.Env[ObsType, np.ndarray], EzPickle): """This compatibility wrapper converts a dm-control environment into a gymnasium environment. Dm-control is DeepMind's software stack for physics-based simulation and Reinforcement Learning environments, using MuJoCo physics. @@ -57,6 +58,14 @@ def __init__( camera_id: int = 0, ): """Initialises the environment with a render mode along with render information.""" + EzPickle.__init__( + self, + env, + render_mode, + render_height, + render_width, + camera_id, + ) self._env = env self.env_type = self._find_env_type(env) diff --git a/shimmy/dm_control_multiagent_compatibility.py b/shimmy/dm_control_multiagent_compatibility.py index 30ce7ae2..3641afce 100644 --- a/shimmy/dm_control_multiagent_compatibility.py +++ b/shimmy/dm_control_multiagent_compatibility.py @@ -10,6 +10,7 @@ import gymnasium import numpy as np from gymnasium.envs.mujoco.mujoco_rendering import MujocoRenderer +from gymnasium.utils import EzPickle from pettingzoo.utils.env import ActionDict, AgentID, ObsDict, ParallelEnv from shimmy.utils.dm_env import dm_obs2gym_obs, dm_spec2gym_space @@ -62,7 +63,7 @@ def _unravel_ma_timestep( ) -class DmControlMultiAgentCompatibilityV0(ParallelEnv): +class DmControlMultiAgentCompatibilityV0(ParallelEnv, EzPickle): """This compatibility wrapper converts multi-agent dm-control environments, primarily soccer, into a Pettingzoo environment. Dm-control is DeepMind's software stack for physics-based simulation and Reinforcement Learning environments, @@ -84,6 +85,7 @@ def __init__( env (dm_env.Environment): dm control multi-agent environment render_mode (Optional[str]): render_mode """ + EzPickle.__init__(self, env, render_mode) super().__init__() self._env = env self.render_mode = render_mode diff --git a/shimmy/dm_lab_compatibility.py b/shimmy/dm_lab_compatibility.py index ef44f0f0..0cddcd8a 100644 --- a/shimmy/dm_lab_compatibility.py +++ b/shimmy/dm_lab_compatibility.py @@ -6,11 +6,12 @@ import gymnasium as gym import numpy as np from gymnasium.core import ObsType +from gymnasium.utils import EzPickle from shimmy.utils.dm_lab import dm_lab_obs2gym_obs_space, dm_lab_spec2gym_space -class DmLabCompatibilityV0(gym.Env[ObsType, Dict[str, np.ndarray]]): +class DmLabCompatibilityV0(gym.Env[ObsType, Dict[str, np.ndarray]], EzPickle): """This compatibility wrapper converts a dm_lab-control environment into a gymnasium environment. DeepMind Lab is a 3D learning environment based on id Software's Quake III Arena via ioquake3 and @@ -27,6 +28,7 @@ def __init__( render_mode: None = None, ): """Initialises the environment with a render mode along with render information.""" + EzPickle.__init__(self, env, render_mode) self._env = env # need to do this to figure out what observation spec the user used diff --git a/shimmy/meltingpot_compatibility.py b/shimmy/meltingpot_compatibility.py index 6b6a930e..f4b7033e 100644 --- a/shimmy/meltingpot_compatibility.py +++ b/shimmy/meltingpot_compatibility.py @@ -15,7 +15,7 @@ import meltingpot.python import numpy as np import pygame -from gymnasium.utils.ezpickle import EzPickle +from gymnasium.utils import EzPickle from ml_collections import config_dict from pettingzoo.utils.env import ActionDict, AgentID, ObsDict, ParallelEnv diff --git a/shimmy/openai_gym_compatibility.py b/shimmy/openai_gym_compatibility.py index 33a6b980..d73f8ce9 100644 --- a/shimmy/openai_gym_compatibility.py +++ b/shimmy/openai_gym_compatibility.py @@ -21,6 +21,7 @@ Text, Tuple, ) +from gymnasium.utils import EzPickle from gymnasium.utils.step_api_compatibility import ( convert_to_terminated_truncated_step_api, ) @@ -40,7 +41,7 @@ GYM_IMPORT_ERROR = None -class GymV26CompatibilityV0(gymnasium.Env[ObsType, ActType]): +class GymV26CompatibilityV0(gymnasium.Env[ObsType, ActType], EzPickle): """This compatibility layer converts a Gym v26 environment to a Gymnasium environment. Gym is the original open source Python library for developing and comparing reinforcement learning algorithms @@ -65,6 +66,7 @@ def __init__( make_kwargs: Additional keyword arguments for make env: An gym environment to wrap. """ + EzPickle.__init__(self, env_id, make_kwargs, env) if GYM_IMPORT_ERROR is not None: raise error.DependencyNotInstalled( f"{GYM_IMPORT_ERROR} (Hint: You need to install gym with `pip install gym` to use gym environments" diff --git a/shimmy/openspiel_compatibility.py b/shimmy/openspiel_compatibility.py index 839bc7b0..e79288c3 100644 --- a/shimmy/openspiel_compatibility.py +++ b/shimmy/openspiel_compatibility.py @@ -5,14 +5,14 @@ from typing import Any, Dict, Optional import numpy as np -import pettingzoo as pz import pyspiel from gymnasium import spaces -from gymnasium.utils import seeding +from gymnasium.utils import EzPickle, seeding +from pettingzoo import AECEnv from pettingzoo.utils.env import AgentID, ObsType -class OpenspielCompatibilityV0(pz.AECEnv): +class OpenspielCompatibilityV0(AECEnv, EzPickle): """This compatibility wrapper converts an openspiel environment into a pettingzoo environment. OpenSpiel is a collection of environments and algorithms for research in general reinforcement learning @@ -35,6 +35,7 @@ def __init__( game (pyspiel.Game): game render_mode (Optional[str]): render_mode """ + EzPickle.__init__(self, game, render_mode) super().__init__() self.game = game self.possible_agents = [