Repository: pytorch/rl

Issue:
[Feature Request] Allow data collectors to receive policies that read and write regular tensors
## Motivation

As a modular library, we'd like our components to be as orthogonal to each other as possible.
As such, we should aim at minimising the inter-dependency of the components (at least on what is exposed to the users).
An example of such a restriction is the fact that our data collectors expect policies that are `TensorDictModule` instances.
It should be possible to automatically build a TensorDictModule just by reading the signature of a policy.

Example:
```
policy = MyRegularModule()
collector = datacollector(make_env, policy)
assert isinstance(collector.policy, TensorDictModule) 
assert collector.policy.out_keys == ["action"] 
assert collector.policy.module is policy 
```

We would assume that the first output of the policy is an action (and has the "action" key). 


### Checks when passing the policy
- If the policy is not an nn.Module, there's nothing we can do with it.
- The policy signature must match the environment `observation_spec` (without the `"next_"` prefixes), eg the following code snippet would only work if the env has keys "next_observation" and "next_stuff" in its observation_spec:
    ```
    class MyPolicy(nn.Module):
        def forward(self, observation, stuff):
            return foo()
    ```


Open questions: 
- how can we support policies that have multiple outputs? (e.g. a log_probability, a hidden state etc)?
  We could name these "output1" ... "outputN"?



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
@@ -4,6 +4,7 @@
 # LICENSE file in the root directory of this source tree.
 
 import abc
+import inspect
 import os
 import queue
 import time
@@ -15,12 +16,13 @@
 
 import numpy as np
 import torch
+import torch.nn as nn
 from torch import multiprocessing as mp
 from torch.utils.data import IterableDataset
 
 from torchrl.envs.utils import set_exploration_mode, step_mdp
 from .._utils import _check_for_faulty_process, prod
-from ..modules.tensordict_module import ProbabilisticTensorDictModule
+from ..modules.tensordict_module import ProbabilisticTensorDictModule, TensorDictModule
 from .utils import split_trajectories
 
 __all__ = [
@@ -81,6 +83,36 @@ def recursive_map_to_cpu(dictionary: OrderedDict) -> OrderedDict:
     )
 
 
+def _policy_is_tensordict_compatible(policy: nn.Module):
+    sig = inspect.signature(policy.forward)
+
+    if isinstance(policy, TensorDictModule) or (
+        len(sig.parameters) == 1
+        and hasattr(policy, "in_keys")
+        and hasattr(policy, "out_keys")
+    ):
+        # if the policy is a TensorDictModule or takes a single argument and defines
+        # in_keys and out_keys then we assume it can already deal with TensorDict input
+        # to forward and we return True
+        return True
+    elif not hasattr(policy, "in_keys") and not hasattr(policy, "out_keys"):
+        # if it's not a TensorDictModule, and in_keys and out_keys are not defined then
+        # we assume no TensorDict compatibility and will try to wrap it.
+        return False
+
+    # if in_keys or out_keys were defined but policy is not a TensorDictModule or
+    # accepts multiple arguments then it's likely the user is trying to do something
+    # that will have undetermined behaviour, we raise an error
+    raise TypeError(
+        "Received a policy that defines in_keys or out_keys and also expects multiple "
+        "arguments to policy.forward. If the policy is compatible with TensorDict, it "
+        "should take a single argument of type TensorDict to policy.forward and define "
+        "both in_keys and out_keys. Alternatively, policy.forward can accept "
+        "arbitrarily many tensor inputs and leave in_keys and out_keys undefined and "
+        "TorchRL will attempt to automatically wrap the policy with a TensorDictModule."
+    )
+
+
 class _DataCollector(IterableDataset, metaclass=abc.ABCMeta):
     def _get_policy_and_device(
         self,
@@ -91,6 +123,7 @@ def _get_policy_and_device(
             ]
         ] = None,
         device: Optional[DEVICE_TYPING] = None,
+        observation_spec: TensorSpec = None,
     ) -> Tuple[
         ProbabilisticTensorDictModule, torch.device, Union[None, Callable[[], dict]]
     ]:
@@ -105,6 +138,7 @@ def _get_policy_and_device(
             policy (ProbabilisticTensorDictModule, optional): a policy to be used
             device (int, str or torch.device, optional): device where to place
                 the policy
+            observation_spec (TensorSpec, optional): spec of the observations
 
         """
         # if create_env_fn is not None:
@@ -124,6 +158,46 @@ def _get_policy_and_device(
                     "env must be provided to _get_policy_and_device if policy is None"
                 )
             policy = RandomPolicy(self.env.action_spec)
+        elif isinstance(policy, nn.Module):
+            # TODO: revisit these checks when we have determined whether arbitrary
+            # callables should be supported as policies.
+            if not _policy_is_tensordict_compatible(policy):
+                # policy is a nn.Module that doesn't operate on tensordicts directly
+                # so we attempt to auto-wrap policy with TensorDictModule
+                if observation_spec is None:
+                    raise ValueError(
+                        "Unable to read observation_spec from the environment. This is "
+                        "required to check compatibility of the environment and policy "
+                        "since the policy is a nn.Module that operates on tensors "
+                        "rather than a TensorDictModule or a nn.Module that accepts a "
+                        "TensorDict as input and defines in_keys and out_keys."
+                    )
+                sig = inspect.signature(policy.forward)
+                next_observation = {
+                    key[5:]: value
+                    for key, value in observation_spec.rand().items()
+                    if key.startswith("next_")
+                }
+                if set(sig.parameters) == set(next_observation):
+                    out_keys = ["action"]
+                    output = policy(**next_observation)
+
+                    if isinstance(output, tuple):
+                        out_keys.extend(f"output{i+1}" for i in range(len(output) - 1))
+
+                    policy = TensorDictModule(
+                        policy, in_keys=list(sig.parameters), out_keys=out_keys
+                    )
+                else:
+                    raise TypeError(
+                        "Arguments to policy.forward are incompatible with entries in "
+                        "env.observation_spec. If you want TorchRL to automatically "
+                        "wrap your policy with a TensorDictModule then the arguments "
+                        "to policy.forward must correspond one-to-one with entries in "
+                        "env.observation_spec that are prefixed with 'next_'. For more "
+                        "complex behaviour and more control you can consider writing "
+                        "your own TensorDictModule."
+                    )
 
         try:
             policy_device = next(policy.parameters()).device
@@ -298,6 +372,7 @@ def __init__(
         (self.policy, self.device, self.get_weights_fn,) = self._get_policy_and_device(
             policy=policy,
             device=device,
+            observation_spec=self.env.observation_spec,
         )
 
         self.env_device = env.device
@@ -324,10 +399,10 @@ def __init__(
         )
 
         if (
-            hasattr(policy, "spec")
-            and policy.spec is not None
-            and all(v is not None for v in policy.spec.values())
-            and set(policy.spec.keys()) == set(policy.out_keys)
+            hasattr(self.policy, "spec")
+            and self.policy.spec is not None
+            and all(v is not None for v in self.policy.spec.values())
+            and set(self.policy.spec.keys()) == set(self.policy.out_keys)
         ):
             # 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
@@ -338,7 +413,7 @@ def __init__(
                     "done": torch.zeros(
                         env.batch_size, dtype=torch.bool, device=env.device
                     ),
-                    **policy.spec.zero(env.batch_size),
+                    **self.policy.spec.zero(env.batch_size),
                 },
                 env.batch_size,
                 device=env.device,
@@ -793,12 +868,24 @@ def device_err_msg(device_name, devices_list):
             )
         self._policy_dict = {}
         self._get_weights_fn_dict = {}
-        for i, _device in enumerate(devices):
+
+        for i, (_device, create_env, kwargs) in enumerate(
+            zip(devices, self.create_env_fn, self.create_env_kwargs)
+        ):
             if _device in self._policy_dict:
                 devices[i] = _device
                 continue
+
+            if hasattr(create_env, "observation_spec"):
+                observation_spec = create_env.observation_spec
+            else:
+                try:
+                    observation_spec = create_env(**kwargs).observation_spec
+                except:  # noqa
+                    observation_spec = None
+
             _policy, _device, _get_weight_fn = self._get_policy_and_device(
-                policy=policy, device=_device
+                policy=policy, device=_device, observation_spec=observation_spec
             )
             self._policy_dict[_device] = _policy
             self._get_weights_fn_dict[_device] = _get_weight_fn
diff --git a/torchrl/modules/models/models.py b/torchrl/modules/models/models.py
--- a/torchrl/modules/models/models.py
+++ b/torchrl/modules/models/models.py
@@ -168,7 +168,7 @@ def __init__(
         device: Optional[DEVICE_TYPING] = None,
     ):
         if out_features is None:
-            raise ValueError("out_feature must be specified for MLP.")
+            raise ValueError("out_features must be specified for MLP.")
 
         default_num_cells = 32
         if num_cells is None:


Test patch:
diff --git a/test/test_collector.py b/test/test_collector.py
--- a/test/test_collector.py
+++ b/test/test_collector.py
@@ -28,6 +28,7 @@
 from torchrl.data import (
     CompositeSpec,
     NdUnboundedContinuousTensorSpec,
+    TensorDict,
     UnboundedContinuousTensorSpec,
 )
 from torchrl.data.tensordict.tensordict import assert_allclose_td
@@ -44,6 +45,42 @@
 # torch.set_default_dtype(torch.double)
 
 
+class WrappablePolicy(nn.Module):
+    def __init__(self, out_features: int, multiple_outputs: bool = False):
+        super().__init__()
+        self.multiple_outputs = multiple_outputs
+        self.linear = nn.LazyLinear(out_features)
+
+    def forward(self, observation):
+        output = self.linear(observation)
+        if self.multiple_outputs:
+            return output, output.sum(), output.min(), output.max()
+        return self.linear(observation)
+
+
+class TensorDictCompatiblePolicy(nn.Module):
+    def __init__(self, out_features: int):
+        super().__init__()
+        self.in_keys = ["observation"]
+        self.out_keys = ["action"]
+        self.linear = nn.LazyLinear(out_features)
+
+    def forward(self, tensordict):
+        return TensorDict(
+            {self.out_keys[0]: self.linear(tensordict.get(self.in_keys[0]))},
+            [],
+        )
+
+
+class UnwrappablePolicy(nn.Module):
+    def __init__(self, out_features: int):
+        super().__init__()
+        self.linear = nn.LazyLinear(out_features)
+
+    def forward(self, observation, other_stuff):
+        return self.linear(observation), other_stuff.sum()
+
+
 class ParametricPolicyNet(nn.Module):
     def __init__(self):
         super().__init__()
@@ -869,6 +906,97 @@ def test_collector_output_keys(collector_class, init_random_frames, explicit_spe
     del collector
 
 
+@pytest.mark.skipif(not _has_gym, reason="test designed with GymEnv")
+@pytest.mark.parametrize(
+    "collector_class",
+    [
+        SyncDataCollector,
+        MultiaSyncDataCollector,
+        MultiSyncDataCollector,
+    ],
+)
+class TestAutoWrap:
+    num_envs = 3
+
+    @pytest.fixture
+    def env_maker(self):
+        from torchrl.envs.libs.gym import GymEnv
+
+        return lambda: GymEnv("Pendulum-v1")
+
+    def _create_collector_kwargs(self, env_maker, collector_class, policy):
+        collector_kwargs = {"create_env_fn": env_maker, "policy": policy}
+
+        if collector_class is not SyncDataCollector:
+            collector_kwargs["create_env_fn"] = [
+                collector_kwargs["create_env_fn"] for _ in range(self.num_envs)
+            ]
+
+        return collector_kwargs
+
+    @pytest.mark.parametrize("multiple_outputs", [False, True])
+    def test_auto_wrap_modules(self, collector_class, multiple_outputs, env_maker):
+        policy = WrappablePolicy(
+            out_features=env_maker().action_spec.shape[-1],
+            multiple_outputs=multiple_outputs,
+        )
+        collector = collector_class(
+            **self._create_collector_kwargs(env_maker, collector_class, policy)
+        )
+
+        out_keys = ["action"]
+        if multiple_outputs:
+            out_keys.extend(f"output{i}" for i in range(1, 4))
+
+        if collector_class is not SyncDataCollector:
+            assert all(
+                isinstance(p, TensorDictModule) for p in collector._policy_dict.values()
+            )
+            assert all(p.out_keys == out_keys for p in collector._policy_dict.values())
+            assert all(p.module is policy for p in collector._policy_dict.values())
+        else:
+            assert isinstance(collector.policy, TensorDictModule)
+            assert collector.policy.out_keys == out_keys
+            assert collector.policy.module is policy
+
+    def test_no_wrap_compatible_module(self, collector_class, env_maker):
+        policy = TensorDictCompatiblePolicy(
+            out_features=env_maker().action_spec.shape[-1]
+        )
+
+        collector = collector_class(
+            **self._create_collector_kwargs(env_maker, collector_class, policy)
+        )
+
+        if collector_class is not SyncDataCollector:
+            assert all(
+                isinstance(p, TensorDictCompatiblePolicy)
+                for p in collector._policy_dict.values()
+            )
+            assert all(
+                p.out_keys == ["action"] for p in collector._policy_dict.values()
+            )
+            assert all(p is policy for p in collector._policy_dict.values())
+        else:
+            assert isinstance(collector.policy, TensorDictCompatiblePolicy)
+            assert collector.policy.out_keys == ["action"]
+            assert collector.policy is policy
+
+    def test_auto_wrap_error(self, collector_class, env_maker):
+        policy = UnwrappablePolicy(out_features=env_maker().action_spec.shape[-1])
+
+        with pytest.raises(
+            TypeError,
+            match=(
+                "Arguments to policy.forward are incompatible with entries in "
+                "env.observation_spec."
+            ),
+        ):
+            collector_class(
+                **self._create_collector_kwargs(env_maker, collector_class, policy)
+            )
+
+
 def weight_reset(m):
     if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
         m.reset_parameters()