Repository: pytorch/rl

Issue:
[BUG] `ToTensorImage` doesn't update the observation spec properly if `unsqueeze=True`
## Describe the bug

If you use `ToTensorImage` with the unsqueeze parameter set, the observations are unsqueezed correctly. However, the observation spec isn't updated.

## To Reproduce

You can just use any `TransformedEnv` that has a `ToTensorImage` transformation with `unsqueeze=True`. Then, try to create a fake tensordict with `fake_tensordict()`. The dimensions of the observation will be wrong as they don't include the unsqueezed dimension.

```python
import torchrl

base_env = ...
env = TransformedEnv(
    base_env,
    Compose(
        ToTensorImage(in_keys="observation", unsqueeze=True),
        StepCounter(),
    ),
)
print(env.fake_tensordict())
```

## Expected behavior

The fake tensordict should have an extra dimension to account for the unsqueezed dimension.

## System info
`torchrl` was installed through Python Poetry and is running on Python 3.10.8.

```python
import torchrl, numpy, sys
print(torchrl.__version__, numpy.__version__, sys.version, sys.platform)
None 1.24.3 3.10.8 (tags/v3.10.8:aaaf517, Oct 11 2022, 16:50:30) [MSC v.1933 64 bit (AMD64)] win32
```

## Additional context

N/A

## Reason and Possible fixes

I will submit a PR that outlines a possible fix.

## Checklist

- [x] I have checked that there is no similar issue in the repo (**required**)
- [x] I have read the [documentation](https://github.com/pytorch/rl/tree/main/docs/) (**required**)
- [x] I have provided a minimal working example to reproduce the bug (**required**)



Discussion:

Resolving patch:
diff --git a/torchrl/envs/transforms/transforms.py b/torchrl/envs/transforms/transforms.py
--- a/torchrl/envs/transforms/transforms.py
+++ b/torchrl/envs/transforms/transforms.py
@@ -951,15 +951,17 @@ def _apply_transform(self, observation: torch.FloatTensor) -> torch.Tensor:
             *list(range(observation.ndimension() - 3)), -1, -3, -2
         )
         observation = observation.div(255).to(self.dtype)
-        if observation.ndimension() == 3 and self.unsqueeze:
+        if self._should_unsqueeze(observation):
             observation = observation.unsqueeze(0)
         return observation
 
     @_apply_to_composite
     def transform_observation_spec(self, observation_spec: TensorSpec) -> TensorSpec:
         observation_spec = self._pixel_observation(observation_spec)
+        unsqueeze_dim = [1] if self._should_unsqueeze(observation_spec) else []
         observation_spec.shape = torch.Size(
             [
+                *unsqueeze_dim,
                 *observation_spec.shape[:-3],
                 observation_spec.shape[-1],
                 observation_spec.shape[-3],
@@ -969,6 +971,14 @@ def transform_observation_spec(self, observation_spec: TensorSpec) -> TensorSpec
         observation_spec.dtype = self.dtype
         return observation_spec
 
+    def _should_unsqueeze(self, observation_like: torch.FloatTensor | TensorSpec):
+        has_3_dimensions = False
+        if isinstance(observation_like, torch.FloatTensor):
+            has_3_dimensions = observation_like.ndimension() == 3
+        else:
+            has_3_dimensions = len(observation_like.shape) == 3
+        return has_3_dimensions and self.unsqueeze
+
     def _pixel_observation(self, spec: TensorSpec) -> None:
         if isinstance(spec.space, ContinuousBox):
             spec.space.maximum = self._apply_transform(spec.space.maximum)


Test patch: