Repository: kevinzakka/dm_env_wrappers (language: python)

Issue:
Support gym and gymnasium converter.
Hey!

I was wondering if you would be willing to put a gym and gymnasium wrapper in this repo. These are wrappers that I use quite regularly from Acme and it would be quite helpful for them to be put here so that we can use them outside the Acme package.

More concretely, to avoid having these as a required dependency, I think we can [`lazy-import`](https://github.com/google-deepmind/launchpad/blob/master/launchpad/lazy_loader.py) these. This is what I do from one of my project in progress.

```python
from corax.utils import lazy_loader

with lazy_loader.LazyImports(__name__, False):
    # This would expose these symbols as top-level classes while still not depending on gym and gymnasium.
    from corax.wrappers.gym_wrapper import GymWrapper
    from corax.wrappers.gymnasium_wrapper import GymnasiumWrapper
del lazy_loader

```


Discussion / hints:
Hi @ethanluoyc, yeah I think it makes a lot of sense to have them in here. Would you be willing to make a PR?

Fix patch:
diff --git a/dm_env_wrappers/__init__.py b/dm_env_wrappers/__init__.py
index 11b2e2e..d3b11e3 100644
--- a/dm_env_wrappers/__init__.py
+++ b/dm_env_wrappers/__init__.py
@@ -21,6 +21,14 @@
 from dm_env_wrappers._src.validate_spec import ValidateActionSpecWrapper
 from dm_env_wrappers._src.video import VideoWrapper
 
+from dm_env_wrappers._src import lazy_loader
+
+with lazy_loader.LazyImports(__name__, False):
+    from dm_env_wrappers._src.gymnasium_wrapper import GymnasiumWrapper
+    from dm_env_wrappers._src.gym_wrapper import GymWrapper
+
+del lazy_loader  # lazy_loader should not be exported
+
 __version__ = "0.0.12"
 
 __all__ = (
@@ -35,6 +43,8 @@
     "EpisodeStatisticsWrapper",
     "ExpandScalarObservationShapesWrapper",
     "FrameStackingWrapper",
+    "GymnasiumWrapper",
+    "GymWrapper",
     "ObservationActionRewardWrapper",
     "SinglePrecisionWrapper",
     "StepLimitWrapper",
diff --git a/dm_env_wrappers/_src/gym_wrapper.py b/dm_env_wrappers/_src/gym_wrapper.py
new file mode 100644
index 0000000..da71fd7
--- /dev/null
+++ b/dm_env_wrappers/_src/gym_wrapper.py
@@ -0,0 +1,161 @@
+# Copyright 2018 DeepMind Technologies Limited. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Wraps an OpenAI Gym environment to be used as a dm_env environment."""
+
+from typing import Any, Dict, Optional
+
+import dm_env
+import gym
+import numpy as np
+import tree
+from gym import spaces
+from dm_env import specs
+
+
+class GymWrapper(dm_env.Environment):
+    """Environment wrapper for OpenAI Gym environments.
+
+    NOTE: This currently only supports the old-style step API.
+    """
+
+    # Note: we don't inherit from base.EnvironmentWrapper because that class
+    # assumes that the wrapped environment is a dm_env.Environment.
+
+    def __init__(self, environment: gym.Env):
+        self._environment = environment
+        self._reset_next_step = True
+        self._last_info: Optional[Dict[str, Any]] = None
+
+        # Convert action and observation specs.
+        obs_space = self._environment.observation_space
+        act_space = self._environment.action_space
+        self._observation_spec = _convert_to_spec(obs_space, name="observation")
+        self._action_spec = _convert_to_spec(act_space, name="action")
+
+    def reset(self) -> dm_env.TimeStep:
+        """Resets the episode."""
+        self._reset_next_step = False
+        observation = self._environment.reset()
+        # Reset the diagnostic information.
+        self._last_info = None
+        return dm_env.restart(observation)
+
+    def step(self, action) -> dm_env.TimeStep:
+        """Steps the environment."""
+        if self._reset_next_step:
+            return self.reset()
+
+        observation, reward, done, info = self._environment.step(action)
+        self._reset_next_step = done
+        self._last_info = info
+
+        # Convert the type of the reward based on the spec, respecting the scalar or
+        # array property.
+        reward = tree.map_structure(
+            lambda x, t: (  # pylint: disable=g-long-lambda
+                t.dtype.type(x) if np.isscalar(x) else np.asarray(x, dtype=t.dtype)
+            ),
+            reward,
+            self.reward_spec(),
+        )
+
+        if done:
+            truncated = info.get("TimeLimit.truncated", False)
+            if truncated:
+                return dm_env.truncation(reward, observation)
+            return dm_env.termination(reward, observation)
+        return dm_env.transition(reward, observation)
+
+    def observation_spec(self):
+        return self._observation_spec
+
+    def action_spec(self):
+        return self._action_spec
+
+    def get_info(self) -> Optional[Dict[str, Any]]:
+        """Returns the last info returned from env.step(action).
+
+        Returns:
+          info: dictionary of diagnostic information from the last environment step
+        """
+        return self._last_info
+
+    @property
+    def environment(self) -> gym.Env:
+        """Returns the wrapped environment."""
+        return self._environment
+
+    def __getattr__(self, name: str):
+        if name.startswith("__"):
+            raise AttributeError(
+                "attempted to get missing private attribute '{}'".format(name)
+            )
+        return getattr(self._environment, name)
+
+    def close(self):
+        self._environment.close()
+
+
+def _convert_to_spec(space: gym.Space, name: Optional[str] = None):
+    """Converts an OpenAI Gym space to a dm_env spec or nested structure of specs.
+
+    Box, MultiBinary and MultiDiscrete Gym spaces are converted to BoundedArray
+    specs. Discrete OpenAI spaces are converted to DiscreteArray specs. Tuple and
+    Dict spaces are recursively converted to tuples and dictionaries of specs.
+
+    Args:
+      space: The Gym space to convert.
+      name: Optional name to apply to all return spec(s).
+
+    Returns:
+      A dm_env spec or nested structure of specs, corresponding to the input
+      space.
+    """
+    if isinstance(space, spaces.Discrete):
+        return specs.DiscreteArray(num_values=space.n, dtype=space.dtype, name=name)  # type: ignore
+
+    elif isinstance(space, spaces.Box):
+        return specs.BoundedArray(
+            shape=space.shape,
+            dtype=space.dtype,
+            minimum=space.low,
+            maximum=space.high,
+            name=name,
+        )
+
+    elif isinstance(space, spaces.MultiBinary):
+        return specs.BoundedArray(
+            shape=space.shape, dtype=space.dtype, minimum=0.0, maximum=1.0, name=name
+        )
+
+    elif isinstance(space, spaces.MultiDiscrete):
+        return specs.BoundedArray(
+            shape=space.shape,
+            dtype=space.dtype,
+            minimum=np.zeros(space.shape),
+            maximum=space.nvec - 1,
+            name=name,
+        )
+
+    elif isinstance(space, spaces.Tuple):
+        return tuple(_convert_to_spec(s, name) for s in space.spaces)
+
+    elif isinstance(space, spaces.Dict):
+        return {
+            key: _convert_to_spec(value, key) for key, value in space.spaces.items()
+        }
+
+    else:
+        raise ValueError("Unexpected gym space: {}".format(space))
diff --git a/dm_env_wrappers/_src/gymnasium_wrapper.py b/dm_env_wrappers/_src/gymnasium_wrapper.py
new file mode 100644
index 0000000..636dbc1
--- /dev/null
+++ b/dm_env_wrappers/_src/gymnasium_wrapper.py
@@ -0,0 +1,146 @@
+"""Wraps an Farama Gymnasium environment to be used as a dm_env environment."""
+
+from typing import Any, Dict, Optional
+
+import dm_env
+import gymnasium as gym
+import numpy as np
+import tree
+from gymnasium import spaces
+from dm_env import specs
+
+
+class GymnasiumWrapper(dm_env.Environment):
+    """Environment wrapper for OpenAI Gym environments."""
+
+    # Note: we don't inherit from base.EnvironmentWrapper because that class
+    # assumes that the wrapped environment is a dm_env.Environment.
+
+    def __init__(self, environment: gym.Env):
+        self._environment = environment
+        self._reset_next_step = True
+        self._last_info: Optional[Dict[str, Any]] = None
+
+        # Convert action and observation specs.
+        obs_space = self._environment.observation_space
+        act_space = self._environment.action_space
+        self._observation_spec = _convert_to_spec(obs_space, name="observation")
+        self._action_spec = _convert_to_spec(act_space, name="action")
+
+    def reset(self) -> dm_env.TimeStep:
+        """Resets the episode."""
+        self._reset_next_step = False
+        observation, info = self._environment.reset()
+        # Reset the diagnostic information.
+        self._last_info = info
+        return dm_env.restart(observation)
+
+    def step(self, action) -> dm_env.TimeStep:
+        """Steps the environment."""
+        if self._reset_next_step:
+            return self.reset()
+
+        observation, reward, terminated, truncated, info = self._environment.step(
+            action
+        )
+
+        self._reset_next_step = terminated or truncated
+        self._last_info = info
+
+        # Convert the type of the reward based on the spec, respecting the scalar or
+        # array property.
+        reward = tree.map_structure(
+            lambda x, t: (  # pylint: disable=g-long-lambda
+                t.dtype.type(x) if np.isscalar(x) else np.asarray(x, dtype=t.dtype)
+            ),
+            reward,
+            self.reward_spec(),
+        )
+
+        if terminated or truncated:
+            if truncated:
+                return dm_env.truncation(reward, observation)
+            return dm_env.termination(reward, observation)
+        return dm_env.transition(reward, observation)
+
+    def observation_spec(self):
+        return self._observation_spec
+
+    def action_spec(self):
+        return self._action_spec
+
+    def get_info(self) -> Optional[Dict[str, Any]]:
+        """Returns the last info returned from env.step(action).
+
+        Returns:
+          info: dictionary of diagnostic information from the last environment step
+        """
+        return self._last_info
+
+    @property
+    def environment(self) -> gym.Env:
+        """Returns the wrapped environment."""
+        return self._environment
+
+    def __getattr__(self, name: str):
+        if name.startswith("__"):
+            raise AttributeError(
+                "attempted to get missing private attribute '{}'".format(name)
+            )
+        return getattr(self._environment, name)
+
+    def close(self):
+        self._environment.close()
+
+
+def _convert_to_spec(space: gym.Space, name: Optional[str] = None):
+    """Converts an OpenAI Gym space to a dm_env spec or nested structure of specs.
+
+    Box, MultiBinary and MultiDiscrete Gym spaces are converted to BoundedArray
+    specs. Discrete OpenAI spaces are converted to DiscreteArray specs. Tuple and
+    Dict spaces are recursively converted to tuples and dictionaries of specs.
+
+    Args:
+      space: The Gym space to convert.
+      name: Optional name to apply to all return spec(s).
+
+    Returns:
+      A dm_env spec or nested structure of specs, corresponding to the input
+      space.
+    """
+    if isinstance(space, spaces.Discrete):
+        return specs.DiscreteArray(num_values=space.n, dtype=space.dtype, name=name)  # type: ignore
+
+    elif isinstance(space, spaces.Box):
+        return specs.BoundedArray(
+            shape=space.shape,
+            dtype=space.dtype,
+            minimum=space.low,
+            maximum=space.high,
+            name=name,
+        )
+
+    elif isinstance(space, spaces.MultiBinary):
+        return specs.BoundedArray(
+            shape=space.shape, dtype=space.dtype, minimum=0.0, maximum=1.0, name=name
+        )
+
+    elif isinstance(space, spaces.MultiDiscrete):
+        return specs.BoundedArray(
+            shape=space.shape,
+            dtype=space.dtype,
+            minimum=np.zeros(space.shape),
+            maximum=space.nvec - 1,
+            name=name,
+        )
+
+    elif isinstance(space, spaces.Tuple):
+        return tuple(_convert_to_spec(s, name) for s in space.spaces)
+
+    elif isinstance(space, spaces.Dict):
+        return {
+            key: _convert_to_spec(value, key) for key, value in space.spaces.items()
+        }
+
+    else:
+        raise ValueError("Unexpected gym space: {}".format(space))
diff --git a/dm_env_wrappers/_src/lazy_loader.py b/dm_env_wrappers/_src/lazy_loader.py
new file mode 100644
index 0000000..37294c2
--- /dev/null
+++ b/dm_env_wrappers/_src/lazy_loader.py
@@ -0,0 +1,139 @@
+# Copyright 2020 DeepMind Technologies Limited. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""LazyLoader providing support for loading symbols upon first reference."""
+
+import importlib
+import sys
+import typing
+
+# Set to true in your local workspace to work around issues with lazy imports.
+# Please report any cases of usage to stanczyk@, so we can address them.
+_DISABLE_LAZY_IMPORTS = False
+
+
+class LazyImports:
+    """Turns all imports performed in its scope into lazy.
+
+    It is used to avoid pulling in large dependencies.
+    """
+
+    def __init__(self, parent, add_to_globals=True):
+        """Initializes LazyImports.
+
+        Args:
+          parent: Name of the parent module doing the imports.
+          add_to_globals: Whether imported symbols should be added to importing
+              module globals. Rule of thumb is to not add for __init__ modules,
+              as it is important for these modules to resolve Lazy symbol upon its
+              first access (not when accessing members of the Lazy import).
+        """
+        self._parent = sys.modules[parent]
+        self._globals = vars(self._parent)
+        self._add_to_globals = add_to_globals
+        self._symbols = dict()
+
+    def parent(self):
+        return self._parent
+
+    def _find_and_load(self, name, import_):
+        del import_
+        return _ModuleToLoad(name, self)
+
+    def _handle_fromlist(self, module, fromlist, import_):
+        del import_
+        assert len(fromlist) == 1
+        if not isinstance(module, _ModuleToLoad):
+            return _ModuleToLoad(module.__name__, self)
+        return module
+
+    def __enter__(self):
+        if typing.TYPE_CHECKING or _DISABLE_LAZY_IMPORTS:
+            return
+        # Start capturing import statements.
+        self._org_find_and_load = importlib._bootstrap._find_and_load
+        self._org_handle_fromlist = importlib._bootstrap._handle_fromlist
+        importlib._bootstrap._find_and_load = self._find_and_load
+        importlib._bootstrap._handle_fromlist = self._handle_fromlist
+
+    def __exit__(self, type_, value, traceback):
+        if typing.TYPE_CHECKING or _DISABLE_LAZY_IMPORTS:
+            return
+        # Stop capturing import statements.
+        importlib._bootstrap._find_and_load = self._org_find_and_load
+        importlib._bootstrap._handle_fromlist = self._org_handle_fromlist
+
+        # If symbols are not added to globals then __getattr__ has to be patched.
+        if not self._add_to_globals:
+            self._parent.__getattr__ = self.__getattr__
+
+        # Make sure lazy symbols know their local name within the importing module.
+        for name in list(self._globals.keys()):
+            member = self._globals[name]
+            if isinstance(member, _MemberToLoad):
+                member.set_export_name(name)
+                if not self._add_to_globals:
+                    del self._globals[name]
+                    self._symbols[name] = member
+
+    def __getattr__(self, name):
+        if name in self._symbols:
+            res = self._symbols[name].load()
+            self._parent.__dict__[name] = res
+            return res
+        raise AttributeError(f"module {self._parent!r} has no attribute {name!r}")
+
+
+class _MemberToLoad:
+    """Module member to load in a lazy way."""
+
+    def __init__(self, name, parent):
+        self._name = name
+        self._export_name = name
+        self._parent = parent
+
+    def __getattr__(self, item):
+        return getattr(self.load(), item)
+
+    def __call__(self, *args, **kwargs):
+        return self.load()(*args, **kwargs)  # type: ignore
+
+    def load(self):
+        try:
+            res = importlib.import_module(f"{self._parent.name()}.{self._name}")
+        except ModuleNotFoundError:
+            res = importlib.import_module(self._parent.name()).__dict__[self._name]
+
+        self._parent.parent().parent().__dict__[self._export_name] = res
+        return res
+
+    def set_export_name(self, name):
+        self._export_name = name
+
+
+class _ModuleToLoad:
+    """Module to load in a lazy way."""
+
+    def __init__(self, name, parent):
+        self._name = name
+        self._parent = parent
+
+    def name(self):
+        return self._name
+
+    def parent(self):
+        return self._parent
+
+    def __getattr__(self, item):
+        return _MemberToLoad(item, self)
diff --git a/setup.py b/setup.py
index 383e389..83d933d 100644
--- a/setup.py
+++ b/setup.py
@@ -27,11 +27,18 @@
     "scipy",
 ]
 
-test_requirements = [
-    "absl-py",
-    "pytest-xdist",
-    "dm_control >= 1.0.7",
-]
+gym_requirements = ["gym<0.24.0"]
+gymnasium_requirements = ["gymnasium"]
+
+test_requirements = (
+    [
+        "absl-py",
+        "pytest-xdist",
+        "dm_control >= 1.0.7",
+    ]
+    + gym_requirements
+    + gymnasium_requirements
+)
 
 dev_requirements = [
     "black",
@@ -82,5 +89,7 @@
     extras_require={
         "test": test_requirements,
         "dev": dev_requirements,
+        "gym": gym_requirements,
+        "gymnasium": gymnasium_requirements,
     },
 )


Test patch:
diff --git a/dm_env_wrappers/_src/gym_wrapper_test .py b/dm_env_wrappers/_src/gym_wrapper_test .py
new file mode 100644
index 0000000..6dc41bd
--- /dev/null
+++ b/dm_env_wrappers/_src/gym_wrapper_test .py
@@ -0,0 +1,69 @@
+# type: ignore
+"""Tests for gymnasium_wrapper."""
+
+import gym
+import numpy as np
+from absl.testing import absltest
+from dm_env import specs
+
+from dm_env_wrappers._src import gym_wrapper
+
+
+class GymWrapperTest(absltest.TestCase):
+    def test_gym_cartpole(self):
+        env = gym.GymWrapper(gym.make("CartPole-v0"))
+
+        # Test converted observation spec.
+        observation_spec: specs.BoundedArray = env.observation_spec()
+        self.assertEqual(type(observation_spec), specs.BoundedArray)
+        self.assertEqual(observation_spec.shape, (4,))
+        self.assertEqual(observation_spec.minimum.shape, (4,))
+        self.assertEqual(observation_spec.maximum.shape, (4,))
+        self.assertEqual(observation_spec.dtype, np.dtype("float32"))
+
+        # Test converted action spec.
+        action_spec: specs.BoundedArray = env.action_spec()
+        self.assertEqual(type(action_spec), specs.DiscreteArray)
+        self.assertEqual(action_spec.shape, ())
+        self.assertEqual(action_spec.minimum, 0)
+        self.assertEqual(action_spec.maximum, 1)
+        self.assertEqual(action_spec.num_values, 2)
+        self.assertEqual(action_spec.dtype, np.dtype("int64"))
+
+        # Test step.
+        timestep = env.reset()
+        self.assertTrue(timestep.first())
+        timestep = env.step(1)
+        self.assertEqual(timestep.reward, 1.0)
+        self.assertTrue(np.isscalar(timestep.reward))
+        self.assertEqual(timestep.observation.shape, (4,))
+        env.close()
+
+    def test_early_truncation(self):
+        # Pendulum has no early termination condition. Recent versions of gym force
+        # to use v1. We try both in case an earlier version is installed.
+        try:
+            gym_env = gym.make("Pendulum-v1")
+        except:  # noqa
+            gym_env = gym.make("Pendulum-v0")
+        env = gym_wrapper.GymWrapper(gym_env)
+        ts = env.reset()
+        while not ts.last():
+            ts = env.step(env.action_spec().generate_value())
+        self.assertEqual(ts.discount, 1.0)
+        self.assertTrue(np.isscalar(ts.reward))
+        env.close()
+
+    def test_multi_discrete(self):
+        space = gym.spaces.MultiDiscrete([2, 3])
+        spec = gym_wrapper._convert_to_spec(space)
+
+        spec.validate([0, 0])
+        spec.validate([1, 2])
+
+        self.assertRaises(ValueError, spec.validate, [2, 2])
+        self.assertRaises(ValueError, spec.validate, [1, 3])
+
+
+if __name__ == "__main__":
+    absltest.main()
diff --git a/dm_env_wrappers/_src/gymnasium_wrapper_test.py b/dm_env_wrappers/_src/gymnasium_wrapper_test.py
new file mode 100644
index 0000000..a8a8e2c
--- /dev/null
+++ b/dm_env_wrappers/_src/gymnasium_wrapper_test.py
@@ -0,0 +1,68 @@
+# type: ignore
+"""Tests for gymnasium_wrapper."""
+
+import numpy as np
+from absl.testing import absltest
+from dm_env import specs
+import gymnasium
+from dm_env_wrappers._src import gymnasium_wrapper
+
+
+class GymWrapperTest(absltest.TestCase):
+    def test_gym_cartpole(self):
+        env = gymnasium_wrapper.GymnasiumWrapper(gymnasium.make("CartPole-v0"))
+
+        # Test converted observation spec.
+        observation_spec: specs.BoundedArray = env.observation_spec()
+        self.assertEqual(type(observation_spec), specs.BoundedArray)
+        self.assertEqual(observation_spec.shape, (4,))
+        self.assertEqual(observation_spec.minimum.shape, (4,))
+        self.assertEqual(observation_spec.maximum.shape, (4,))
+        self.assertEqual(observation_spec.dtype, np.dtype("float32"))
+
+        # Test converted action spec.
+        action_spec: specs.BoundedArray = env.action_spec()
+        self.assertEqual(type(action_spec), specs.DiscreteArray)
+        self.assertEqual(action_spec.shape, ())
+        self.assertEqual(action_spec.minimum, 0)
+        self.assertEqual(action_spec.maximum, 1)
+        self.assertEqual(action_spec.num_values, 2)
+        self.assertEqual(action_spec.dtype, np.dtype("int64"))
+
+        # Test step.
+        timestep = env.reset()
+        self.assertTrue(timestep.first())
+        timestep = env.step(1)
+        self.assertEqual(timestep.reward, 1.0)
+        self.assertTrue(np.isscalar(timestep.reward))
+        self.assertEqual(timestep.observation.shape, (4,))
+        env.close()
+
+    def test_early_truncation(self):
+        # Pendulum has no early termination condition. Recent versions of gym force
+        # to use v1. We try both in case an earlier version is installed.
+        try:
+            gym_env = gymnasium.make("Pendulum-v1")
+        except:  # noqa
+            gym_env = gymnasium.make("Pendulum-v0")
+        env = gymnasium_wrapper.GymnasiumWrapper(gym_env)
+        ts = env.reset()
+        while not ts.last():
+            ts = env.step(env.action_spec().generate_value())
+        self.assertEqual(ts.discount, 1.0)
+        self.assertTrue(np.isscalar(ts.reward))
+        env.close()
+
+    def test_multi_discrete(self):
+        space = gymnasium.spaces.MultiDiscrete([2, 3])
+        spec = gymnasium_wrapper._convert_to_spec(space)
+
+        spec.validate([0, 0])
+        spec.validate([1, 2])
+
+        self.assertRaises(ValueError, spec.validate, [2, 2])
+        self.assertRaises(ValueError, spec.validate, [1, 3])
+
+
+if __name__ == "__main__":
+    absltest.main()