Repository: pytorch/rl

Issue:
[BUG] `fake_tensordict()` in `EnvBase` is incompatible with `OneHotDiscreteTensorSpec`
## Describe the bug

When using collectors like `SyncDataCollector`, the constructor will try to call the `env.fake_tensordict()` (`env` is an argument of the constructor). Unfortunately, the implementation of `env.fake_tensordict()` will generate a fake observation tensordict using something like `observation_spec.zero()`. But for an observation spec like `OneHotDiscreteTensorSpec`, an all-zero tensor is illegal, which will lead to unexpected behaviors for the downstream models.

## To Reproduce

An example code to reproduce the behavior:

```python
from typing import Optional

import torch
from tensordict import TensorDict, TensorDictBase
from tensordict.nn import TensorDictModule
from torchrl.envs import EnvBase
from torchrl.data import (
    CompositeSpec,
    UnboundedContinuousTensorSpec,
    OneHotDiscreteTensorSpec,
)
from torchrl.collectors import SyncDataCollector


class DemoEnv(EnvBase):
    batch_locked = True

    def __init__(self):
        super().__init__()

        # For demo purpose, params are hard coded.
        self.node_num = 3
        self.batch_size = (2,)

        self.set_seed(torch.empty((), dtype=torch.int64).random_().item())
        self._make_spec()

    def _make_spec(self):
        """Set the observation, action, and reward specs."""
        self.observation_spec = CompositeSpec(
            current_node=OneHotDiscreteTensorSpec(
                n=self.node_num,
                shape=(*self.batch_size, self.node_num),
                dtype=torch.bool,
            ),
            shape=self.batch_size,
        )
        self.state_spec = self.observation_spec.clone()
        self.action_spec = UnboundedContinuousTensorSpec(shape=self.batch_size)
        self.reward_spec = UnboundedContinuousTensorSpec(shape=self.batch_size)

    def _set_seed(self, seed: Optional[int]):
        rng = torch.manual_seed(seed)
        self.rng = rng

    def _step(self, td: TensorDictBase) -> TensorDict:
        # Trivial step function for demo purpose.
        out = TensorDict(
            {
                "next": {
                    "current_node": td["current_node"],
                    "reward": td["action"],
                    "done": td["done"],
                }
            },
            td.shape,
        )
        return out

    def _reset(self, td) -> TensorDict:
        current_node = torch.zeros(*self.batch_size, self.node_num).bool()
        current_node[:, 0] = True

        out = TensorDict(
            {
                "current_node": current_node,
            },
            self.batch_size,
        )

        return out


def trival_policy(x):
    out = 1 / torch.sum(x, dim=-1)
    assert not torch.isinf(out).any(), "Oops! Inf value!"
    return out


env = DemoEnv()
policy = TensorDictModule(
    module=trival_policy,
    in_keys=["current_node"],
    out_keys=["action"],
)

collector = SyncDataCollector(
    env,
    policy,
    frames_per_batch=10,
    total_frames=100,
)
```
You can always see the assertion exception raised here.

## Expected behavior

The example code should not raise an assertion exception, which means a one-hot style tensor should be delivered to the policy module!

## System info

Though I'm using the stable version of TorchRL, this bug may also exist in the nightly version. 

```python
import torchrl, numpy, sys
print(torchrl.__version__, numpy.__version__, sys.version, sys.platform)
```
```
None 1.24.3 3.10.11 | packaged by conda-forge | (main, May 10 2023, 18:58:44) [GCC 11.3.0] linux
```

## Reason and Possible fixes

As I said in the description section, the implementation of `env.fake_tensordict()` will generate a fake observation tensordict using something like `observation_spec.zero()`, which leads to this bug.

current main branch code:
https://github.com/pytorch/rl/blob/95b1bfec3411b154870163a777537ee0657801c3/torchrl/envs/common.py#L1674

stable branch code:
https://github.com/pytorch/rl/blob/6d030c9ec1db72c252a23dff04dcad2a553678d2/torchrl/envs/common.py#L895

Note that in other cases, the same problem may also exist for the other spec (like action spec).

## 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/collectors/collectors.py b/torchrl/collectors/collectors.py
--- a/torchrl/collectors/collectors.py
+++ b/torchrl/collectors/collectors.py
@@ -621,64 +621,45 @@ def __init__(
         )
 
         with torch.no_grad():
-            self._tensordict_out = env.fake_tensordict()
+            self._tensordict_out = self.env.fake_tensordict()
+        # If the policy has a valid spec, we use it
         if (
-            hasattr(self.policy, "spec")
-            and self.policy.spec is not None
-            and all(
-                v is not None for v in self.policy.spec.values(True, True)
-            )  # if a spec is None, we don't know anything about it
-            # and set(self.policy.spec.keys(True, True)) == set(self.policy.out_keys)
-            and any(
-                key not in self._tensordict_out.keys(isinstance(key, tuple))
-                for key in self.policy.spec.keys(True, True)
-            )
-        ):
-            # if policy spec is non-empty, all the values are not None and the keys
-            # match the out_keys we assume the user has given all relevant information
-            # the policy could have more keys than the env:
-            policy_spec = self.policy.spec
-            if policy_spec.ndim < self._tensordict_out.ndim:
-                policy_spec = policy_spec.expand(self._tensordict_out.shape)
-            for key, spec in policy_spec.items(True, True):
-                if key in self._tensordict_out.keys(isinstance(key, tuple)):
-                    continue
-                self._tensordict_out.set(key, spec.zero())
-            self._tensordict_out = (
-                self._tensordict_out.unsqueeze(-1)
-                .expand(*env.batch_size, self.frames_per_batch)
-                .clone()
-            )
-        elif (
             hasattr(self.policy, "spec")
             and self.policy.spec is not None
             and all(v is not None for v in self.policy.spec.values(True, True))
-            and all(
-                key in self._tensordict_out.keys(isinstance(key, tuple))
-                for key in self.policy.spec.keys(True, True)
-            )
         ):
-            # reach this if the policy has specs and they match with the fake tensordict
-            self._tensordict_out = (
-                self._tensordict_out.unsqueeze(-1)
-                .expand(*env.batch_size, self.frames_per_batch)
-                .clone()
-            )
+            if any(
+                key not in self._tensordict_out.keys(isinstance(key, tuple))
+                for key in self.policy.spec.keys(True, True)
+            ):
+                # if policy spec is non-empty, all the values are not None and the keys
+                # match the out_keys we assume the user has given all relevant information
+                # the policy could have more keys than the env:
+                policy_spec = self.policy.spec
+                if policy_spec.ndim < self._tensordict_out.ndim:
+                    policy_spec = policy_spec.expand(self._tensordict_out.shape)
+                for key, spec in policy_spec.items(True, True):
+                    if key in self._tensordict_out.keys(isinstance(key, tuple)):
+                        continue
+                    self._tensordict_out.set(key, spec.zero())
+
         else:
             # otherwise, we perform a small number of steps with the policy to
             # determine the relevant keys with which to pre-populate _tensordict_out.
             # This is the safest thing to do if the spec has None fields or if there is
             # no spec at all.
             # See #505 for additional context.
+            self._tensordict_out.update(self._tensordict)
             with torch.no_grad():
-                self._tensordict_out = self._tensordict_out.to(self.device)
-                self._tensordict_out = self.policy(self._tensordict_out).unsqueeze(-1)
-            self._tensordict_out = (
-                self._tensordict_out.expand(*env.batch_size, self.frames_per_batch)
-                .clone()
-                .zero_()
-            )
-        # in addition to outputs of the policy, we add traj_ids and step_count to
+                self._tensordict_out = self.policy(self._tensordict_out.to(self.device))
+
+        self._tensordict_out = (
+            self._tensordict_out.unsqueeze(-1)
+            .expand(*env.batch_size, self.frames_per_batch)
+            .clone()
+            .zero_()
+        )
+        # in addition to outputs of the policy, we add traj_ids to
         # _tensordict_out which will be collected during rollout
         self._tensordict_out = self._tensordict_out.to(self.storing_device)
         self._tensordict_out.set(


Test patch:
diff --git a/test/test_collector.py b/test/test_collector.py
--- a/test/test_collector.py
+++ b/test/test_collector.py
@@ -1397,6 +1397,26 @@ def test_reset_heterogeneous_envs():
     ).all()
 
 
+def test_policy_with_mask():
+    env = CountingBatchedEnv(start_val=torch.tensor(10), max_steps=torch.tensor(1e5))
+
+    def policy(td):
+        obs = td.get("observation")
+        # This policy cannot work with obs all 0s
+        if not obs.any():
+            raise AssertionError
+        action = obs.clone()
+        td.set("action", action)
+        return td
+
+    collector = SyncDataCollector(
+        env, policy=policy, frames_per_batch=10, total_frames=20
+    )
+    for _ in collector:
+        break
+    collector.shutdown()
+
+
 class TestNestedEnvsCollector:
     def test_multi_collector_nested_env_consistency(self, seed=1):
         env = NestedCountingEnv()