Repository: pytorch/rl

Issue:
[BUG] Discrete SAC vmap runtime error when using MLP + LSTM architecture
## Describe the bug

Using Discrete SAC with an only-LSTM model works fine. However, using an MLP + LSTM architecture results in the following error when  the objective module calling vmap on the Q networks:

> RuntimeError: vmap: index_put_(self, *extra_args) is not possible because there exists a Tensor `other` in extra_args that has more elements than `self`. This happened due to `other` being vmapped over but `self` not being vmapped over in a vmap. Please try to use out-of-place operators instead of index_put_. If said operator is being called inside the PyTorch framework, please file a bug report instead.

## To Reproduce

The following script reproduces the error

```python

import tqdm
import torch
from tensordict.nn import TensorDictModule, TensorDictSequential
from torchrl.envs.libs.gym import GymEnv
from torchrl.data import LazyMemmapStorage, TensorDictReplayBuffer
from torchrl.modules.distributions import OneHotCategorical
from torchrl.modules import ProbabilisticActor, GRUModule, MLP
from torchrl.collectors import SyncDataCollector
from torchrl.objectives import DiscreteSACLoss
from torchrl.envs import (
    ParallelEnv,
    TransformedEnv,
    InitTracker,
    StepCounter,
    RewardSum,
)


def create_model(input_size, output_size, hidden_size=256, num_layers=3, out_key="logits"):

    embedding_module = TensorDictModule(
        in_keys=["observation"],
        out_keys=["embed"],
        # module=torch.nn.Identity(),  # this works
        # module=MLP(in_features=input_size, out_features=input_size, num_cells=[]), # this raises RuntimeError
        module=torch.nn.Linear(input_size, input_size), # this raises RuntimeError
    )
    
    # RuntimeError: vmap: index_put_(self, *extra_args) is not possible because there exists
    # a Tensor `other` in extra_args that has more elements than `self`. This happened due to
    # `other` being vmapped over but `self` not being vmapped over in a vmap. Please try to use
    # out-of-place operators instead of index_put_. If said operator is being called inside the
    # PyTorch framework, please file a bug report instead.

    lstm_module = GRUModule(
        input_size=input_size,
        hidden_size=hidden_size,
        num_layers=num_layers,
        in_key="embed",
        out_key="features",
        python_based=True,
    )
    mlp = TensorDictModule(
        MLP(
            in_features=hidden_size,
            out_features=output_size,
            num_cells=[],
        ),
        in_keys=["features"],
        out_keys=[out_key],
    )

    inference_model = TensorDictSequential(embedding_module, lstm_module, mlp)
    training_model = TensorDictSequential(embedding_module, lstm_module.set_recurrent_mode(), mlp)

    return inference_model, training_model


def create_rhs_transform(input_size, hidden_size=256, num_layers=3):
    lstm_module = GRUModule(
        input_size=input_size,
        hidden_size=hidden_size,
        num_layers=num_layers,
        in_key="observation",
        out_key="features",
    )
    return lstm_module.make_tensordict_primer()


def main():

    device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
    test_env = GymEnv("CartPole-v1", device=device, categorical_action_encoding=True)
    action_spec = test_env.action_spec.space
    observation_spec = test_env.observation_spec["observation"]

    def create_env_fn():
        env = GymEnv("CartPole-v1", device=device)
        env = TransformedEnv(env)
        env.append_transform(create_rhs_transform(input_size=observation_spec.shape[-1]))
        env.append_transform(InitTracker())
        return env

    # Models
    ##################

    inference_actor, training_actor = create_model(input_size=observation_spec.shape[-1], output_size=action_spec.n)
    inference_actor = ProbabilisticActor(
        module=inference_actor,
        in_keys=["logits"],
        out_keys=["action"],
        distribution_class=OneHotCategorical,
        return_log_prob=True,
    )
    training_actor = ProbabilisticActor(
        module=training_actor,
        in_keys=["logits"],
        out_keys=["action"],
        distribution_class=OneHotCategorical,
        return_log_prob=True,
    )
    inference_actor = inference_actor.to(device)
    training_actor = training_actor.to(device)
    _, training_critic = create_model(input_size=observation_spec.shape[-1], output_size=action_spec.n, out_key="action_value")
    training_critic = training_critic.to(device)

    # Collector
    ##################

    collector = SyncDataCollector(
        create_env_fn=create_env_fn,
        policy=inference_actor,
        frames_per_batch=100,
        total_frames=1000,
        device=device,
        storing_device=device,
        max_frames_per_traj=-1,
        split_trajs=False,
    )

    # Buffer
    ##################

    buffer = TensorDictReplayBuffer(
        storage=LazyMemmapStorage(100),
        batch_size=8,
    )

    # Loss
    ##################

    loss_module = DiscreteSACLoss(
        actor_network=training_actor,
        qvalue_network=training_critic,
        num_actions=action_spec.n,
        num_qvalue_nets=2,
        loss_function="smooth_l1",
    )
    loss_module.make_value_estimator(gamma=0.99)

    # Collection loop
    ##################

    for data in tqdm.tqdm(collector):
        buffer.extend(data.reshape(-1).cpu())
        batch = buffer.sample()
        loss = loss_module(batch.cuda())

    collector.shutdown()
    print("Success!")


if __name__ == "__main__":
    main()

```

## Checklist

- [ ] I have checked that there is no similar issue in the repo (**required**)
- [ ] I have read the [documentation](https://github.com/pytorch/rl/tree/main/docs/) (**required**)
- [ ] I have provided a minimal working example to reproduce the bug (**required**)



Discussion:

Resolving patch:
diff --git a/torchrl/objectives/value/utils.py b/torchrl/objectives/value/utils.py
--- a/torchrl/objectives/value/utils.py
+++ b/torchrl/objectives/value/utils.py
@@ -8,6 +8,7 @@
 import torch
 
 from tensordict import TensorDictBase
+from tensordict.utils import expand_right
 
 
 def _custom_conv1d(tensor: torch.Tensor, filter: torch.Tensor):
@@ -294,8 +295,8 @@ def _fill_tensor(tensor):
             dtype=tensor.dtype,
             device=tensor.device,
         )
-        empty_tensor[mask] = tensor
-        return empty_tensor
+        mask_expand = expand_right(mask, (*mask.shape, *tensor.shape[1:]))
+        return torch.masked_scatter(empty_tensor, mask_expand, tensor.reshape(-1))
 
     if isinstance(tensor, TensorDictBase):
         tensor = tensor.apply(_fill_tensor, batch_size=[*shape])


Test patch:
diff --git a/test/test_tensordictmodules.py b/test/test_tensordictmodules.py
--- a/test/test_tensordictmodules.py
+++ b/test/test_tensordictmodules.py
@@ -1836,6 +1836,62 @@ def create_transformed_env():
             assert (data.get(("next", "recurrent_state_c")) != 0.0).all()
             assert (data.get("recurrent_state_c") != 0.0).any()
 
+    def test_lstm_vmap_complex_model(self):
+        # Tests that all ops in GRU are compatible with VMAP (when build using
+        # the PT backend).
+        # This used to fail when splitting the input based on the is_init mask.
+        # This test is intended not only as a non-regression test but also
+        # to make sure that any change provided to RNNs is compliant with vmap
+        torch.manual_seed(0)
+        input_size = 4
+        hidden_size = 5
+        num_layers = 1
+        output_size = 3
+        out_key = "out"
+
+        embedding_module = TensorDictModule(
+            in_keys=["observation"],
+            out_keys=["embed"],
+            module=torch.nn.Linear(input_size, input_size),
+        )
+
+        lstm_module = LSTMModule(
+            input_size=input_size,
+            hidden_size=hidden_size,
+            num_layers=num_layers,
+            in_key="embed",
+            out_key="features",
+            python_based=True,
+        )
+        mlp = TensorDictModule(
+            MLP(
+                in_features=hidden_size,
+                out_features=output_size,
+                num_cells=[],
+            ),
+            in_keys=["features"],
+            out_keys=[out_key],
+        )
+        training_model = TensorDictSequential(
+            embedding_module, lstm_module.set_recurrent_mode(), mlp
+        )
+        is_init = torch.zeros(50, 11, 1, dtype=torch.bool).bernoulli_(0.1)
+        data = TensorDict(
+            {"observation": torch.randn(50, 11, input_size), "is_init": is_init},
+            [50, 11],
+        )
+        training_model(data)
+        params = TensorDict.from_module(training_model)
+        params = params.expand(2)
+
+        def call(data, params):
+            with params.to_module(training_model):
+                return training_model(data)
+
+        assert torch.vmap(call, (None, 0))(data, params).shape == torch.Size(
+            (2, 50, 11)
+        )
+
 
 class TestGRUModule:
     def test_errs(self):
@@ -2106,6 +2162,62 @@ def create_transformed_env():
             assert (data.get("recurrent_state") != 0.0).any()
             assert (data.get(("next", "recurrent_state")) != 0.0).all()
 
+    def test_gru_vmap_complex_model(self):
+        # Tests that all ops in GRU are compatible with VMAP (when build using
+        # the PT backend).
+        # This used to fail when splitting the input based on the is_init mask.
+        # This test is intended not only as a non-regression test but also
+        # to make sure that any change provided to RNNs is compliant with vmap
+        torch.manual_seed(0)
+        input_size = 4
+        hidden_size = 5
+        num_layers = 1
+        output_size = 3
+        out_key = "out"
+
+        embedding_module = TensorDictModule(
+            in_keys=["observation"],
+            out_keys=["embed"],
+            module=torch.nn.Linear(input_size, input_size),
+        )
+
+        lstm_module = GRUModule(
+            input_size=input_size,
+            hidden_size=hidden_size,
+            num_layers=num_layers,
+            in_key="embed",
+            out_key="features",
+            python_based=True,
+        )
+        mlp = TensorDictModule(
+            MLP(
+                in_features=hidden_size,
+                out_features=output_size,
+                num_cells=[],
+            ),
+            in_keys=["features"],
+            out_keys=[out_key],
+        )
+        training_model = TensorDictSequential(
+            embedding_module, lstm_module.set_recurrent_mode(), mlp
+        )
+        is_init = torch.zeros(50, 11, 1, dtype=torch.bool).bernoulli_(0.1)
+        data = TensorDict(
+            {"observation": torch.randn(50, 11, input_size), "is_init": is_init},
+            [50, 11],
+        )
+        training_model(data)
+        params = TensorDict.from_module(training_model)
+        params = params.expand(2)
+
+        def call(data, params):
+            with params.to_module(training_model):
+                return training_model(data)
+
+        assert torch.vmap(call, (None, 0))(data, params).shape == torch.Size(
+            (2, 50, 11)
+        )
+
 
 def test_safe_specs():