init

pytorch · vmoens · Jan 19, 2024 · Jan 13, 2024 · Jan 14, 2024 · Jan 15, 2024
commit c4c46f90b160fbea3f2015cea2329e0dc683d81e
diff --git a/torchrl/data/tensor_specs.py b/torchrl/data/tensor_specs.py
@@ -1569,6 +1569,14 @@ def __init__(
             shape, ContinuousBox(low, high, device=device), device, dtype, "continuous"
         )
 
+    @property
+    def low(self):
+        return self.space.low
+
+    @property
+    def high(self):
+        return self.space.high
+
     def expand(self, *shape):
         if len(shape) == 1 and isinstance(shape[0], (tuple, list, torch.Size)):
             shape = shape[0]
@@ -2373,6 +2381,10 @@ def __init__(
         super().__init__(shape, space, device, dtype, domain="discrete")
         self.update_mask(mask)
 
+    @property
+    def n(self):
+        return self.space.n
+
     def update_mask(self, mask):
         if mask is not None:
             try:

diff --git a/torchrl/envs/libs/_gym_utils.py b/torchrl/envs/libs/_gym_utils.py
@@ -0,0 +1,120 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+from __future__ import annotations
+
+from torch.utils._pytree import tree_map
+
+from torchrl._utils import implement_for
+from torchrl.envs.libs.gym import _torchrl_to_gym_spec_transform
+
+
+class TorchRLGymWrapper(gymnasium.Env):
+    def __init__(self, env_cls, to_numpy=False, **kwargs):
+        self.torchrl_env = env_cls(**kwargs)
+        super().__init__()
+        self.action_space = _torchrl_to_gym_spec_transform(
+            self.torchrl_env.action_spec,
+            categorical_action_encoding=self.torchrl_env.categorical_action_encoding,
+        )
+        self.observation_space = _torchrl_to_gym_spec_transform(
+            self.torchrl_env.observation_spec,
+            categorical_action_encoding=self.torchrl_env.categorical_action_encoding,
+        )
+        self.to_numpy = to_numpy
+
+    @implement_for("gymnasium")
+    def step(self, action):
+        self._tensordict.set("action", action)
+        self.torchrl_env.step(self._tensordict)
+        _tensordict = step_mdp(self._tensordict)
+        keys = list(self.torchrl_env.observation_spec.keys())
+        observation = (
+            self._tensordict.get("next")
+            .select(*self.torchrl_env.observation_spec.keys())
+            .to_dict()
+        )
+        reward = self._tensordict.get(("next", "reward"))
+        terminated = self._tensordict.get(("next", "terminated"))
+        truncated = self._tensordict.get(("next", "truncated"))
+        info = {}
+        self._tensordict = _tensordict
+        out = (observation, reward, terminated, truncated, info)
+        if self.to_numpy():
+            out = tree_map(lambda x: x.detach().cpu().numpy(), out)
+        return out
+
+    @implement_for("gym", "0.26", None)
+    def step(self, action):
+        self._tensordict.set("action", action)
+        self.torchrl_env.step(self._tensordict)
+        _tensordict = step_mdp(self._tensordict)
+        keys = list(self.torchrl_env.observation_spec.keys())
+        observation = (
+            self._tensordict.get("next")
+            .select(*self.torchrl_env.observation_spec.keys())
+            .to_dict()
+        )
+        reward = self._tensordict.get(("next", "reward"))
+        terminated = self._tensordict.get(("next", "terminated"))
+        truncated = self._tensordict.get(("next", "truncated"))
+        info = {}
+        self._tensordict = _tensordict
+        out = (observation, reward, terminated, truncated, info)
+        if self.to_numpy():
+            out = tree_map(lambda x: x.detach().cpu().numpy(), out)
+        return out
+
+    @implement_for("gym", None, "0.26")
+    def step(self, action):
+        self._tensordict.set("action", action)
+        self.torchrl_env.step(self._tensordict)
+        _tensordict = step_mdp(self._tensordict)
+        keys = list(self.torchrl_env.observation_spec.keys())
+        observation = (
+            self._tensordict.get("next")
+            .select(*self.torchrl_env.observation_spec.keys())
+            .to_dict()
+        )
+        reward = self._tensordict.get(("next", "reward"))
+        done = self._tensordict.get(("next", "done"))
+        info = {}
+        self._tensordict = _tensordict
+        out = (observation, reward, done, info)
+        if self.to_numpy():
+            out = tree_map(lambda x: x.detach().cpu().numpy(), out)
+        return out
+
+    @implement_for("gymnasium")
+    def reset(self):
+        self._tensordict = self.torchrl_env.reset()
+        observation = self._tensordict.select(
+            *self.torchrl_env.observation_spec.keys()
+        ).to_dict()
+        out = observation, {}
+        if self.to_numpy():
+            out = tree_map(lambda x: x.detach().cpu().numpy(), out)
+        return out
+
+    @implement_for("gym", None, "0.26")
+    def reset(self):
+        self._tensordict = self.torchrl_env.reset()
+        observation = self._tensordict.select(
+            *self.torchrl_env.observation_spec.keys()
+        ).to_dict()
+        out = observation
+        if self.to_numpy():
+            out = tree_map(lambda x: x.detach().cpu().numpy(), out)
+        return out
+
+    @implement_for("gym", "0.26", None)
+    def reset(self):
+        self._tensordict = self.torchrl_env.reset()
+        observation = self._tensordict.select(
+            *self.torchrl_env.observation_spec.keys()
+        ).to_dict()
+        out = observation, {}
+        if self.to_numpy():
+            out = tree_map(lambda x: x.detach().cpu().numpy(), out)
+        return out
diff --git a/torchrl/envs/libs/gym.py b/torchrl/envs/libs/gym.py
@@ -30,7 +30,7 @@
     TensorSpec,
     UnboundedContinuousTensorSpec,
 )
-from torchrl.data.utils import numpy_to_torch_dtype_dict
+from torchrl.data.utils import numpy_to_torch_dtype_dict, torch_to_numpy_dtype_dict
 from torchrl.envs.batched_envs import CloudpickleWrapper
 from torchrl.envs.common import _EnvPostInit
 
@@ -332,6 +332,63 @@ def _gym_to_torchrl_spec_transform(
         )
 
 
+def _torchrl_to_gym_spec_transform(
+    spec,
+    batch_size=torch.Size([]),
+    categorical_action_encoding=False,
+    remap_state_to_observation: bool = True,
+) -> TensorSpec:
+    """Maps TorchRL specs to gym spaces.
+
+    Args:
+        spec: the torchrl spec to transform.
+        batch_size (torch.Size): batch-size of the input specs.
+        categorical_action_encoding: whether discrete spaces should be mapped to categorical or one-hot.
+            Defaults to one-hot.
+        remap_state_to_observation: whether to rename the 'state' key of Dict specs to "observation". Default is true.
+
+    """
+    gym_spaces = gym_backend("spaces")
+    shape = spec.shape[len(batch_size) :]
+    if isinstance(spec, MultiDiscreteTensorSpec):
+        return gym_spaces.multi_discrete.MultiDiscrete(
+            spec.nvec, dtype=torch_to_numpy_dtype_dict[spec.dtype]
+        )
+    if isinstance(spec, MultiOneHotDiscreteTensorSpec):
+        return gym_spaces.multi_discrete.MultiDiscrete(spec.nvec)
+    if isinstance(spec, DiscreteTensorSpec):
+        return gym_spaces.discrete.Discrete(
+            spec.n
+        )  # dtype=torch_to_numpy_dtype_dict[spec.dtype])
+    if isinstance(spec, OneHotDiscreteTensorSpec):
+        return gym_spaces.discrete.Discrete(spec.n)
+    if isinstance(spec, BinaryDiscreteTensorSpec):
+        return gym_spaces.multi_binary.MultiBinary(
+            spec.n, dtype=torch_to_numpy_dtype_dict[spec.dtype]
+        )
+    if isinstance(spec, UnboundedContinuousTensorSpec):
+        return gym_spaces.Box(-float("inf"), float("inf"), shape)
+    if isinstance(spec, BoundedTensorSpec):
+        return gym_spaces.Box(
+            spec.low.detach().cpu().numpy(), spec.high.detach().cpu().numpy(), shape
+        )
+    if isinstance(spec, CompositeSpec):
+        return dict(
+            **{
+                key: _torchrl_to_gym_spec_transform(
+                    val,
+                    batch_size=spec.shape,
+                    categorical_action_encoding=categorical_action_encoding,
+                )
+                for key, val in spec.items()
+            }
+        )
+    else:
+        raise NotImplementedError(
+            f"spec of type {type(spec).__name__} is currently unaccounted for"
+        )
+
+
 def _get_envs(to_dict=False) -> List:
     if not _has_gym:
         raise ImportError("Gym(nasium) could not be found in your virtual environment.")