Repository: pytorch/rl

Issue:
[Feature Request] Gym spec conversion modifying key names
In the gym to torchrl spec transform https://github.com/pytorch/rl/blob/main/torchrl/envs/libs/gym.py#L181
the key "state" of dict spaces is overwritten with a new name. 

Is it necessary to do it in this location? users (me) might want to transform their gym specs to torchrl without this side effect.


Discussion:
Yep it isn't perfect. We can change that but it'll be bc breaking so I would be happy with any solution that does not break things before the 0.3 release (leaving users with 0.2 with a warning)
One simple one is adding an argument like parse_state_key with default true?

Resolving patch:
diff --git a/torchrl/envs/libs/gym.py b/torchrl/envs/libs/gym.py
--- a/torchrl/envs/libs/gym.py
+++ b/torchrl/envs/libs/gym.py
@@ -179,12 +179,21 @@ def gym_backend(submodule=None):
 
 
 def _gym_to_torchrl_spec_transform(
-    spec, dtype=None, device="cpu", categorical_action_encoding=False
+    spec,
+    dtype=None,
+    device="cpu",
+    categorical_action_encoding=False,
+    remap_state_to_observation: bool = True,
 ) -> TensorSpec:
     """Maps the gym specs to the TorchRL specs.
 
-    By convention, 'state' keys of Dict specs will be renamed "observation" to match the
-    default TorchRL keys.
+    Args:
+        spec: the gym space to transform
+        dtype: a dtype to use for the spec. Defaults to`spec.dtype`.
+        device: the device for the spec. Defaults to "cpu".
+        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 = gym_backend()
@@ -241,7 +250,11 @@ def _gym_to_torchrl_spec_transform(
         spec_out = {}
         for k in spec.keys():
             key = k
-            if k == "state" and "observation" not in spec.keys():
+            if (
+                remap_state_to_observation
+                and k == "state"
+                and "observation" not in spec.keys()
+            ):
                 # we rename "state" in "observation" as "observation" is the conventional name
                 # for single observation in torchrl.
                 # naming it 'state' will result in envs that have a different name for the state vector
@@ -251,6 +264,7 @@ def _gym_to_torchrl_spec_transform(
                 spec[k],
                 device=device,
                 categorical_action_encoding=categorical_action_encoding,
+                remap_state_to_observation=remap_state_to_observation,
             )
         return CompositeSpec(**spec_out)
     elif isinstance(spec, gym.spaces.dict.Dict):
@@ -258,6 +272,7 @@ def _gym_to_torchrl_spec_transform(
             spec.spaces,
             device=device,
             categorical_action_encoding=categorical_action_encoding,
+            remap_state_to_observation=remap_state_to_observation,
         )
     else:
         raise NotImplementedError(
diff --git a/torchrl/envs/libs/vmas.py b/torchrl/envs/libs/vmas.py
--- a/torchrl/envs/libs/vmas.py
+++ b/torchrl/envs/libs/vmas.py
@@ -164,12 +164,14 @@ def _make_specs(
                     self.action_space[agent_index],
                     categorical_action_encoding=True,
                     device=self.device,
+                    remap_state_to_observation=False,
                 )
             )  # shape = (n_actions_per_agent,)
             observation_specs.append(
                 _gym_to_torchrl_spec_transform(
                     self.observation_space[agent_index],
                     device=self.device,
+                    remap_state_to_observation=False,
                 )
             )  # shape = (n_obs_per_agent,)
             reward_specs.append(


Test patch: