Repository: pytorch/rl

Issue:
[Feature Request] Add info dict key-spec pairs to `env.observation_spec` if they're provided to a gym environment through `set_info_dict_reader`
## Motivation

gym's environments can pass some more info through a dictionary `info`:
`obs, reward, done, terminated = env.step(action)`.
This dictionary can contain additional information about the state of the env: time spent doing the task, whether the task is solved etc.
We wrap the obs, reward and done in the output `TensorDict`. We can also place the info if they're numerical values (see #234):
```python
        >>> from torchrl.envs.libs.gym import GymWrapper
        >>> from torchrl.envs import default_info_dict_reader
        >>> reader = default_info_dict_reader(["my_info_key"])
        >>> # assuming "some_env-v0" returns a dict with a key "my_info_key"
        >>> env = GymWrapper(gym.make("some_env-v0"))
        >>> env.set_info_dict_reader(info_dict_reader=reader)
        >>> tensordict = env.reset()
        >>> tensordict = env.rand_step(tensordict)
        >>> assert "my_info_key" in tensordict.keys()
```

Problem is that we also register what are the expected observations (domain, dtype, device etc) in an `observation_spec` attribute. It is important for us to know what to expect as output from the env. However the info is not yet registered in `observation_spec`. This is a problem for parallel environments as we would wish to pre-allocate the tensors that we place in share memory to pass information from one process to another: for that, we use the `observation_spec` which is cheaper than resetting the env.

## Solution

We should be able to provide the specs corresponding to the info keys in the `default_info_dict_reader`, and the `set_info_dict_reader` should be able to read them. If no spec is provided, a default unidimensional, floating-point unbounded spec should be assumed.


Discussion:

Resolving patch:
diff --git a/torchrl/envs/gym_like.py b/torchrl/envs/gym_like.py
--- a/torchrl/envs/gym_like.py
+++ b/torchrl/envs/gym_like.py
@@ -5,20 +5,38 @@
 
 from __future__ import annotations
 
+import abc
 import warnings
-from typing import Optional, Union, Tuple, Any, Dict
+from typing import List, Optional, Sequence, Union, Tuple, Any, Dict
 
 import numpy as np
 import torch
 
 from torchrl.data import TensorDict
+from torchrl.data.tensor_specs import TensorSpec, UnboundedContinuousTensorSpec
 from torchrl.data.tensordict.tensordict import TensorDictBase
 from torchrl.envs.common import _EnvWrapper
 
 __all__ = ["GymLikeEnv", "default_info_dict_reader"]
 
 
-class default_info_dict_reader:
+class BaseInfoDictReader(metaclass=abc.ABCMeta):
+    """
+    Base class for info-readers.
+    """
+
+    @abc.abstractmethod
+    def __call__(
+        self, info_dict: Dict[str, Any], tensordict: TensorDictBase
+    ) -> TensorDictBase:
+        raise NotImplementedError
+
+    @abc.abstractproperty
+    def info_spec(self) -> Dict[str, TensorSpec]:
+        raise NotImplementedError
+
+
+class default_info_dict_reader(BaseInfoDictReader):
     """
     Default info-key reader.
 
@@ -39,11 +57,30 @@ class default_info_dict_reader:
 
     """
 
-    def __init__(self, keys=None):
+    def __init__(
+        self,
+        keys: List[str] = None,
+        spec: Union[Sequence[TensorSpec], Dict[str, TensorSpec]] = None,
+    ):
         if keys is None:
             keys = []
         self.keys = keys
 
+        if isinstance(spec, Sequence):
+            if len(spec) != len(self.keys):
+                raise ValueError(
+                    "If specifying specs for info keys with a sequence, the "
+                    "length of the sequence must match the number of keys"
+                )
+            self._info_spec = dict(zip(self.keys, spec))
+        else:
+            if spec is None:
+                spec = {}
+
+            self._info_spec = {
+                key: spec.get(key, UnboundedContinuousTensorSpec()) for key in self.keys
+            }
+
     def __call__(
         self, info_dict: Dict[str, Any], tensordict: TensorDictBase
     ) -> TensorDictBase:
@@ -57,9 +94,13 @@ def __call__(
                 tensordict[key] = info_dict[key]
         return tensordict
 
+    @property
+    def info_spec(self) -> Dict[str, TensorSpec]:
+        return self._info_spec
+
 
 class GymLikeEnv(_EnvWrapper):
-    _info_dict_reader: callable
+    _info_dict_reader: BaseInfoDictReader
 
     """
     A gym-like env is an environment whose behaviour is similar to gym environments in what
@@ -216,7 +257,7 @@ def _output_transform(self, step_outputs_tuple: Tuple) -> Tuple:
             )
         return step_outputs_tuple
 
-    def set_info_dict_reader(self, info_dict_reader: callable) -> GymLikeEnv:
+    def set_info_dict_reader(self, info_dict_reader: BaseInfoDictReader) -> GymLikeEnv:
         """
         Sets an info_dict_reader function. This function should take as input an
         info_dict dictionary and the tensordict returned by the step function, and
@@ -240,6 +281,8 @@ def set_info_dict_reader(self, info_dict_reader: callable) -> GymLikeEnv:
 
         """
         self.info_dict_reader = info_dict_reader
+        for info_key, spec in info_dict_reader.info_spec.items():
+            self.observation_spec[info_key] = spec
         return self
 
     def __repr__(self) -> str:


Test patch:
diff --git a/test/test_env.py b/test/test_env.py
--- a/test/test_env.py
+++ b/test/test_env.py
@@ -28,11 +28,13 @@
     MultOneHotDiscreteTensorSpec,
     BoundedTensorSpec,
     NdBoundedTensorSpec,
+    UnboundedContinuousTensorSpec,
 )
 from torchrl.data.tensordict.tensordict import assert_allclose_td, TensorDict
 from torchrl.envs import EnvCreator, ObservationNorm, CatTensors, DoubleToFloat
+from torchrl.envs.gym_like import default_info_dict_reader
 from torchrl.envs.libs.dm_control import _has_dmc, DMControlEnv
-from torchrl.envs.libs.gym import _has_gym, GymEnv
+from torchrl.envs.libs.gym import GymEnv, GymWrapper, _has_gym
 from torchrl.envs.transforms import (
     TransformedEnv,
     Compose,
@@ -1121,6 +1123,34 @@ def test_batch_unlocked_with_batch_size(device):
         env.step(td_expanded)
 
 
+@pytest.mark.skipif(not _has_gym, reason="no gym")
+def test_info_dict_reader(seed=0):
+    import gym
+
+    env = GymWrapper(gym.make("HalfCheetah-v4"))
+    env.set_info_dict_reader(default_info_dict_reader(["x_position"]))
+
+    assert "x_position" in env.observation_spec.keys()
+    assert isinstance(env.observation_spec["x_position"], UnboundedContinuousTensorSpec)
+
+    tensordict = env.reset()
+    tensordict = env.rand_step(tensordict)
+
+    assert env.observation_spec["x_position"].is_in(tensordict["x_position"])
+
+    env2 = GymWrapper(gym.make("HalfCheetah-v4"))
+    env2.set_info_dict_reader(
+        default_info_dict_reader(
+            ["x_position"], spec={"x_position": OneHotDiscreteTensorSpec(5)}
+        )
+    )
+
+    tensordict2 = env2.reset()
+    tensordict2 = env2.rand_step(tensordict2)
+
+    assert not env2.observation_spec["x_position"].is_in(tensordict2["x_position"])
+
+
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()
     pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)