From b40869d0e433f94b2f25e5f16ab499c0e80ca75e Mon Sep 17 00:00:00 2001
From: gehring <clement.gehring@gmail.com>
Date: Thu, 26 Dec 2019 16:22:17 -0500
Subject: [PATCH] Wrapper for the dm_env interface (#6468)

---
 rllib/env/__init__.py             |  3 +-
 rllib/env/dm_env_wrapper.py       | 98 +++++++++++++++++++++++++++++++
 rllib/examples/dmlab_watermaze.py | 33 +++++++++++
 3 files changed, 133 insertions(+), 1 deletion(-)
 create mode 100644 rllib/env/dm_env_wrapper.py
 create mode 100644 rllib/examples/dmlab_watermaze.py

diff --git a/rllib/env/__init__.py b/rllib/env/__init__.py
index 87eb931df..a454f0dd9 100644
--- a/rllib/env/__init__.py
+++ b/rllib/env/__init__.py
@@ -1,4 +1,5 @@
 from ray.rllib.env.base_env import BaseEnv
+from ray.rllib.env.dm_env_wrapper import DMEnv
 from ray.rllib.env.multi_agent_env import MultiAgentEnv
 from ray.rllib.env.external_env import ExternalEnv
 from ray.rllib.env.serving_env import ServingEnv
@@ -7,5 +8,5 @@ from ray.rllib.env.env_context import EnvContext
 
 __all__ = [
     "BaseEnv", "MultiAgentEnv", "ExternalEnv", "VectorEnv", "ServingEnv",
-    "EnvContext"
+    "EnvContext", "DMEnv"
 ]
diff --git a/rllib/env/dm_env_wrapper.py b/rllib/env/dm_env_wrapper.py
new file mode 100644
index 000000000..37c9df7dc
--- /dev/null
+++ b/rllib/env/dm_env_wrapper.py
@@ -0,0 +1,98 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import gym
+from gym import spaces
+
+import numpy as np
+
+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")
diff --git a/rllib/examples/dmlab_watermaze.py b/rllib/examples/dmlab_watermaze.py
new file mode 100644
index 000000000..5a23357b9
--- /dev/null
+++ b/rllib/examples/dmlab_watermaze.py
@@ -0,0 +1,33 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from deepmind_lab import dmenv_module
+
+from ray.rllib import env
+
+
+class Watermaze(env.DMEnv):
+    def __init__(self, env_config):
+        lab = dmenv_module.Lab(
+            "contributed/dmlab30/rooms_watermaze",
+            ["RGBD"],
+            config=env_config,
+        )
+        super(Watermaze, self).__init__(lab)
+
+
+env = Watermaze({"width": "320", "height": "160"})
+print(env.action_space)
+
+for i in range(2):
+    print(
+        env.step({
+            "CROUCH": 0.,
+            "FIRE": 0.,
+            "JUMP": 0.,
+            "LOOK_DOWN_UP_PIXELS_PER_FRAME": 0.,
+            "LOOK_LEFT_RIGHT_PIXELS_PER_FRAME": 0.,
+            "MOVE_BACK_FORWARD": 0.,
+            "STRAFE_LEFT_RIGHT": 0.
+        }))