[Feature] Capture wrong spec transforms (1/N)

ghstack-source-id: f2d938b
Pull Request resolved: #2805

Author: Vincent Moens

test/test_env.py

@@ -55,12 +55,14 @@
 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, gym_backend, GymEnv, GymWrapper
-from torchrl.envs.transforms import Compose, StepCounter, TransformedEnv
-from torchrl.envs.transforms.transforms import (
+from torchrl.envs.transforms import (
     AutoResetEnv,
     AutoResetTransform,
+    Compose,
+    StepCounter,
     Tokenizer,
     Transform,
+    TransformedEnv,
     UnsqueezeTransform,
 )
 from torchrl.envs.utils import (
@@ -3770,6 +3772,28 @@ def test_str2str_rb_slicesampler(self):
         else:
             raise RuntimeError("Failed to sample both trajs")
 
+    def test_env_with_str_append(self):
+        class StrAppender(Transform):
+            def transform_observation_spec(self, observation_spec):
+                return observation_spec.set("str", NonTensor(example_data="a string"))
+
+            def _step(self, td, next_td):
+                s = td["str"]
+
+                s += "-" + str(int(s.split("-")[-1]) + 1)
+                next_td["str"] = s
+                return next_td
+
+            def _reset(self, td, reset_td):
+                return reset_td.set("str", "0")
+
+        env = TransformedEnv(CountingEnv(), StrAppender())
+        r = env.rollout(10)
+        r_unbind = r.unbind(0)
+        for ep_prev, ep_next in zip(r_unbind[:-1], r_unbind[1:]):
+            assert ep_prev["next", "str"].startswith(ep_prev["str"])
+            assert ep_next["str"] == ep_prev["next", "str"]
+
     def test_env_with_tensorclass(self):
         env = EnvWithTensorClass()
         env.check_env_specs()

torchrl/envs/transforms/rlhf.py

@@ -187,9 +187,6 @@ def _step(
     forward = _call
 
     def transform_output_spec(self, output_spec: Composite) -> Composite:
-        output_spec = super().transform_output_spec(output_spec)
-        # todo: here we'll need to use the reward_key once it's implemented
-        # parent = self.parent
         in_key = unravel_key(self.in_keys[0])
         out_key = unravel_key(self.out_keys[0])
 
@@ -205,18 +202,13 @@ def transform_output_spec(self, output_spec: Composite) -> Composite:
             )
         elif in_key == "reward":
             parent = self.parent
-            reward_spec = Unbounded(
-                device=output_spec.device,
-                shape=output_spec["full_reward_spec"][parent.reward_key].shape,
-            )
+            reward_spec = output_spec["full_reward_spec"][parent.reward_key].clone()
             # then we need to populate the output keys
             observation_spec = output_spec["full_observation_spec"]
             observation_spec[out_key] = reward_spec
         else:
             observation_spec = output_spec["full_observation_spec"]
-            reward_spec = Unbounded(
-                device=output_spec.device, shape=observation_spec[in_key].shape
-            )
+            reward_spec = observation_spec[in_key].clone()
             # then we need to populate the output keys
             observation_spec[out_key] = reward_spec
         return output_spec

torchrl/envs/transforms/transforms.py

@@ -560,6 +560,21 @@ def transform_output_spec(self, output_spec: Composite) -> Composite:
             output_spec["full_done_spec"] = self.transform_done_spec(
                 output_spec["full_done_spec"]
             )
+        output_spec_keys = [
+            unravel_key(k[1:]) for k in output_spec.keys(True) if isinstance(k, tuple)
+        ]
+        out_keys = {unravel_key(k) for k in self.out_keys}
+        in_keys = {unravel_key(k) for k in self.in_keys}
+        for key in out_keys - in_keys:
+            if unravel_key(key) not in output_spec_keys:
+                warnings.warn(
+                    f"The key '{key}' is unaccounted for by the transform (expected keys {output_spec_keys}). "
+                    f"Every new entry in the tensordict resulting from a call to a transform must be "
+                    f"registered in the specs for torchrl rollouts to be consistently built. "
+                    f"Make sure transform_output_spec/transform_observation_spec/... is coded correctly. "
+                    "This warning will trigger a KeyError in v0.9, make sure to adapt your code accordingly.",
+                    category=FutureWarning,
+                )
         return output_spec
 
     def transform_input_spec(self, input_spec: TensorSpec) -> TensorSpec:
@@ -1468,33 +1483,57 @@ def transform_input_spec(self, input_spec: TensorSpec) -> TensorSpec:
         # the action spec from the env, map it using t0 then t1 (going from in to out).
         for t in self.transforms:
             input_spec = t.transform_input_spec(input_spec)
+            if not isinstance(input_spec, Composite):
+                raise TypeError(
+                    f"Expected Compose but got {type(input_spec)} with transform {t}"
+                )
         return input_spec
 
     def transform_action_spec(self, action_spec: TensorSpec) -> TensorSpec:
         # To understand why we don't invert, look up at transform_input_spec
         for t in self.transforms:
             action_spec = t.transform_action_spec(action_spec)
+            if not isinstance(action_spec, TensorSpec):
+                raise TypeError(
+                    f"Expected TensorSpec but got {type(action_spec)} with transform {t}"
+                )
         return action_spec
 
     def transform_state_spec(self, state_spec: TensorSpec) -> TensorSpec:
         # To understand why we don't invert, look up at transform_input_spec
         for t in self.transforms:
             state_spec = t.transform_state_spec(state_spec)
+            if not isinstance(state_spec, Composite):
+                raise TypeError(
+                    f"Expected Compose but got {type(state_spec)} with transform {t}"
+                )
         return state_spec
 
     def transform_observation_spec(self, observation_spec: TensorSpec) -> TensorSpec:
         for t in self.transforms:
             observation_spec = t.transform_observation_spec(observation_spec)
+            if not isinstance(observation_spec, TensorSpec):
+                raise TypeError(
+                    f"Expected TensorSpec but got {type(observation_spec)} with transform {t}"
+                )
         return observation_spec
 
     def transform_output_spec(self, output_spec: TensorSpec) -> TensorSpec:
         for t in self.transforms:
             output_spec = t.transform_output_spec(output_spec)
+            if not isinstance(output_spec, Composite):
+                raise TypeError(
+                    f"Expected Compose but got {type(output_spec)} with transform {t}"
+                )
         return output_spec
 
     def transform_reward_spec(self, reward_spec: TensorSpec) -> TensorSpec:
         for t in self.transforms:
             reward_spec = t.transform_reward_spec(reward_spec)
+            if not isinstance(reward_spec, TensorSpec):
+                raise TypeError(
+                    f"Expected TensorSpec but got {type(reward_spec)} with transform {t}"
+                )
         return reward_spec
 
     def __getitem__(self, item: Union[int, slice, List]) -> Union: