[RLlib] Env directory cleanup and tests. (#13082)

rllib/env/dm_env_wrapper.py

@@ -1,94 +1,10 @@
-import gym
-from gym import spaces
+from ray.rllib.env.wrappers.dm_env_wrapper import DMEnv as DE
+from ray.rllib.utils.deprecation import deprecation_warning
 
-import numpy as np
+deprecation_warning(
+    old="ray.rllib.env.dm_env_wrapper.DMEnv",
+    new="ray.rllib.env.wrappers.dm_env_wrapper.DMEnv",
+    error=False,
+)
 
-try:
-    from dm_env import specs
-except ImportError:
-    specs = None
-
-
-def _convert_spec_to_space(spec):
-    if isinstance(spec, dict):
-        return spaces.Dict(
-            {k: _convert_spec_to_space(v)
-             for k, v in spec.items()})
-    if isinstance(spec, specs.DiscreteArray):
-        return spaces.Discrete(spec.num_values)
-    elif isinstance(spec, specs.BoundedArray):
-        return spaces.Box(
-            low=np.asscalar(spec.minimum),
-            high=np.asscalar(spec.maximum),
-            shape=spec.shape,
-            dtype=spec.dtype)
-    elif isinstance(spec, specs.Array):
-        return spaces.Box(
-            low=-float("inf"),
-            high=float("inf"),
-            shape=spec.shape,
-            dtype=spec.dtype)
-
-    raise NotImplementedError(
-        ("Could not convert `Array` spec of type {} to Gym space. "
-         "Attempted to convert: {}").format(type(spec), spec))
-
-
-class DMEnv(gym.Env):
-    """A `gym.Env` wrapper for the `dm_env` API.
-    """
-
-    metadata = {"render.modes": ["rgb_array"]}
-
-    def __init__(self, dm_env):
-        super(DMEnv, self).__init__()
-        self._env = dm_env
-        self._prev_obs = None
-
-        if specs is None:
-            raise RuntimeError((
-                "The `specs` module from `dm_env` was not imported. Make sure "
-                "`dm_env` is installed and visible in the current python "
-                "environment."))
-
-    def step(self, action):
-        ts = self._env.step(action)
-
-        reward = ts.reward
-        if reward is None:
-            reward = 0.
-
-        return ts.observation, reward, ts.last(), {"discount": ts.discount}
-
-    def reset(self):
-        ts = self._env.reset()
-        return ts.observation
-
-    def render(self, mode="rgb_array"):
-        if self._prev_obs is None:
-            raise ValueError(
-                "Environment not started. Make sure to reset before rendering."
-            )
-
-        if mode == "rgb_array":
-            return self._prev_obs
-        else:
-            raise NotImplementedError(
-                "Render mode '{}' is not supported.".format(mode))
-
-    @property
-    def action_space(self):
-        spec = self._env.action_spec()
-        return _convert_spec_to_space(spec)
-
-    @property
-    def observation_space(self):
-        spec = self._env.observation_spec()
-        return _convert_spec_to_space(spec)
-
-    @property
-    def reward_range(self):
-        spec = self._env.reward_spec()
-        if isinstance(spec, specs.BoundedArray):
-            return spec.minimum, spec.maximum
-        return -float("inf"), float("inf")
+DMEnv = DE