Repository: pytorch/rl

Issue:
Not knowing Which key is causing shape mismatch
Hi there. I'm having problems with my environment rollout. Sometimes my environment shapes are changing somehow. If in the except block I'm going to mention, you don't set any type for the exception it would be much better. Because alongside the exception about shape mismatch another exception would raise saying which key caused this problem.
here is the line I'm suggesting to change:
https://github.com/pytorch/rl/blob/3787a9edeea03cb9e98b928824072b939c327776/torchrl/data/tensor_specs.py#L3978C20-L3978C28
and here is the line a raising shape mismatch:
https://github.com/pytorch/rl/blob/3787a9edeea03cb9e98b928824072b939c327776/torchrl/data/tensor_specs.py#L587

note that if the shapes of observations are different currently just the shape mismatch exception would be raised due to different exception types. 


Discussion:
> Sometimes my environment shapes are changing somehow

Is this expected? We have allowed dynamic shapes in nightlies, it will be part of the next major release. You can just set `-1` to the shapes that should change.

RE your suggestion:
what about capturing this `RuntimeError` separately, and tell user which key is the culprit of the shape mismatch? 
> Is this expected?

No. The changes of shapes aren't desired. My codes are having a bug. 


> what about capturing this `RuntimeError` separately

I agree. That would be nicer.
Ok I guess my initial understanding was that you were expecting your shapes to change and wanted to work around the exceptions, I now understand that the problem is that the exception are not informative enough
I can patch that quickly! 
If you can think of anything else that would help you debug this sort of situation I'd be happy to help

Resolving patch:
diff --git a/torchrl/data/tensor_specs.py b/torchrl/data/tensor_specs.py
--- a/torchrl/data/tensor_specs.py
+++ b/torchrl/data/tensor_specs.py
@@ -3979,6 +3979,10 @@ def encode(
                 raise KeyError(
                     f"The CompositeSpec instance with keys {self.keys()} does not have a '{key}' key."
                 )
+            except RuntimeError as err:
+                raise RuntimeError(
+                    f"Encoding key {key} raised a RuntimeError. Scroll up to know more."
+                ) from err
         return out
 
     def __repr__(self) -> str:


Test patch:
diff --git a/test/test_specs.py b/test/test_specs.py
--- a/test/test_specs.py
+++ b/test/test_specs.py
@@ -1907,6 +1907,22 @@ def test_unboundeddiscrete(
         assert spec == torch.stack(spec.unbind(0), 0)
         assert spec == torch.stack(spec.unbind(-1), -1)
 
+    def test_composite_encode_err(self):
+        c = CompositeSpec(
+            a=UnboundedContinuousTensorSpec(
+                1,
+            ),
+            b=UnboundedContinuousTensorSpec(
+                2,
+            ),
+        )
+        with pytest.raises(KeyError, match="The CompositeSpec instance with keys"):
+            c.encode({"c": 0})
+        with pytest.raises(
+            RuntimeError, match="raised a RuntimeError. Scroll up to know more"
+        ):
+            c.encode({"a": 0, "b": 0})
+
 
 @pytest.mark.parametrize(
     "device",