Repository: pytorch/rl

Issue:
Make `TDSequence` support indexing
A `TDSequence` is the equivalent of `nn.Sequence` for TDModules. 
As of now, it does not support indexing. A short method mapping `TDSequence.__getitem__(index)` to `TDSequence.module.__getitem__(index)` should be implemented.

Tests should be written in `test/test_tdmodules.py` to check that it behaves as expected.


Discussion:

Resolving patch:
diff --git a/torchrl/modules/td_module/common.py b/torchrl/modules/td_module/common.py
--- a/torchrl/modules/td_module/common.py
+++ b/torchrl/modules/td_module/common.py
@@ -791,6 +791,8 @@ class TDSequence(TDModule):
 
     """
 
+    module: nn.ModuleList
+
     def __init__(
         self,
         *modules: TDModule,
@@ -896,6 +898,15 @@ def forward(self, tensordict: _TensorDict, **kwargs) -> _TensorDict:
     def __len__(self):
         return len(self.module)
 
+    def __getitem__(self, index: Union[int, slice]) -> TDModule:
+        return self.module.__getitem__(index)
+
+    def __setitem__(self, index: int, tdmodule: TDModule) -> None:
+        return self.module.__setitem__(idx=index, module=tdmodule)
+
+    def __delitem__(self, index: Union[int, slice]) -> None:
+        self.module.__delitem__(idx=index)
+
     @property
     def spec(self):
         kwargs = {}


Test patch:
diff --git a/test/test_tdmodules.py b/test/test_tdmodules.py
--- a/test/test_tdmodules.py
+++ b/test/test_tdmodules.py
@@ -314,9 +314,11 @@ def test_stateful(self, safe, spec_type, probabilistic, lazy):
         param_multiplier = 2 if probabilistic else 1
         if lazy:
             net1 = nn.LazyLinear(4)
+            dummy_net = nn.LazyLinear(4)
             net2 = nn.LazyLinear(4 * param_multiplier)
         else:
             net1 = nn.Linear(3, 4)
+            dummy_net = nn.Linear(4, 4)
             net2 = nn.Linear(4, 4 * param_multiplier)
         if probabilistic:
             net2 = NormalParamWrapper(net2)
@@ -345,6 +347,13 @@ def test_stateful(self, safe, spec_type, probabilistic, lazy):
                 out_keys=["hidden"],
                 safe=False,
             )
+            dummy_tdmodule = TDModule(
+                dummy_net,
+                None,
+                in_keys=["hidden"],
+                out_keys=["hidden"],
+                safe=False,
+            )
             tdmodule2 = tdclass(
                 spec=spec,
                 module=net2,
@@ -353,7 +362,21 @@ def test_stateful(self, safe, spec_type, probabilistic, lazy):
                 safe=False,
                 **kwargs
             )
-            tdmodule = TDSequence(tdmodule1, tdmodule2)
+            tdmodule = TDSequence(tdmodule1, dummy_tdmodule, tdmodule2)
+
+        assert hasattr(tdmodule, "__setitem__")
+        assert len(tdmodule) == 3
+        tdmodule[1] = tdmodule2
+        assert len(tdmodule) == 3
+
+        assert hasattr(tdmodule, "__delitem__")
+        assert len(tdmodule) == 3
+        del tdmodule[2]
+        assert len(tdmodule) == 2
+
+        assert hasattr(tdmodule, "__getitem__")
+        assert tdmodule[0] is tdmodule1
+        assert tdmodule[1] is tdmodule2
 
         td = TensorDict({"in": torch.randn(3, 3)}, [3])
         tdmodule(td)
@@ -383,11 +406,13 @@ def test_functional(self, safe, spec_type, probabilistic):
         param_multiplier = 2 if probabilistic else 1
 
         net1 = nn.Linear(3, 4)
+        dummy_net = nn.Linear(4, 4)
         net2 = nn.Linear(4, 4 * param_multiplier)
         if probabilistic:
             net2 = NormalParamWrapper(net2)
 
         fnet1, params1 = make_functional(net1)
+        fdummy_net, _ = make_functional(dummy_net)
         fnet2, params2 = make_functional(net2)
         params = list(params1) + list(params2)
 
@@ -409,16 +434,29 @@ def test_functional(self, safe, spec_type, probabilistic):
             pytest.skip("safe and spec is None is checked elsewhere")
         else:
             tdmodule1 = TDModule(
-                fnet1,
-                None,
-                in_keys=["in"],
-                out_keys=["hidden"],
-                safe=False,
+                fnet1, None, in_keys=["in"], out_keys=["hidden"], safe=False
+            )
+            dummy_tdmodule = TDModule(
+                fdummy_net, None, in_keys=["hidden"], out_keys=["hidden"], safe=False
             )
             tdmodule2 = tdclass(
                 fnet2, spec, in_keys=["hidden"], out_keys=["out"], safe=safe, **kwargs
             )
-            tdmodule = TDSequence(tdmodule1, tdmodule2)
+            tdmodule = TDSequence(tdmodule1, dummy_tdmodule, tdmodule2)
+
+        assert hasattr(tdmodule, "__setitem__")
+        assert len(tdmodule) == 3
+        tdmodule[1] = tdmodule2
+        assert len(tdmodule) == 3
+
+        assert hasattr(tdmodule, "__delitem__")
+        assert len(tdmodule) == 3
+        del tdmodule[2]
+        assert len(tdmodule) == 2
+
+        assert hasattr(tdmodule, "__getitem__")
+        assert tdmodule[0] is tdmodule1
+        assert tdmodule[1] is tdmodule2
 
         td = TensorDict({"in": torch.randn(3, 3)}, [3])
         tdmodule(td, params=params)
@@ -448,6 +486,7 @@ def test_functional_with_buffer(self, safe, spec_type, probabilistic):
         param_multiplier = 2 if probabilistic else 1
 
         net1 = nn.Sequential(nn.Linear(7, 7), nn.BatchNorm1d(7))
+        dummy_net = nn.Sequential(nn.Linear(7, 7), nn.BatchNorm1d(7))
         net2 = nn.Sequential(
             nn.Linear(7, 7 * param_multiplier), nn.BatchNorm1d(7 * param_multiplier)
         )
@@ -455,6 +494,7 @@ def test_functional_with_buffer(self, safe, spec_type, probabilistic):
             net2 = NormalParamWrapper(net2)
 
         fnet1, params1, buffers1 = make_functional_with_buffers(net1)
+        fdummy_net, _, _ = make_functional_with_buffers(dummy_net)
         fnet2, params2, buffers2 = make_functional_with_buffers(net2)
 
         params = list(params1) + list(params2)
@@ -478,16 +518,29 @@ def test_functional_with_buffer(self, safe, spec_type, probabilistic):
             pytest.skip("safe and spec is None is checked elsewhere")
         else:
             tdmodule1 = TDModule(
-                fnet1,
-                None,
-                in_keys=["in"],
-                out_keys=["hidden"],
-                safe=False,
+                fnet1, None, in_keys=["in"], out_keys=["hidden"], safe=False
+            )
+            dummy_tdmodule = TDModule(
+                fdummy_net, None, in_keys=["hidden"], out_keys=["hidden"], safe=False
             )
             tdmodule2 = tdclass(
                 fnet2, spec, in_keys=["hidden"], out_keys=["out"], safe=safe, **kwargs
             )
-            tdmodule = TDSequence(tdmodule1, tdmodule2)
+            tdmodule = TDSequence(tdmodule1, dummy_tdmodule, tdmodule2)
+
+        assert hasattr(tdmodule, "__setitem__")
+        assert len(tdmodule) == 3
+        tdmodule[1] = tdmodule2
+        assert len(tdmodule) == 3
+
+        assert hasattr(tdmodule, "__delitem__")
+        assert len(tdmodule) == 3
+        del tdmodule[2]
+        assert len(tdmodule) == 2
+
+        assert hasattr(tdmodule, "__getitem__")
+        assert tdmodule[0] is tdmodule1
+        assert tdmodule[1] is tdmodule2
 
         td = TensorDict({"in": torch.randn(3, 7)}, [3])
         tdmodule(td, params=params, buffers=buffers)
@@ -518,11 +571,13 @@ def test_vmap(self, safe, spec_type, probabilistic):
         param_multiplier = 2 if probabilistic else 1
 
         net1 = nn.Linear(3, 4)
+        dummy_net = nn.Linear(4, 4)
         net2 = nn.Linear(4, 4 * param_multiplier)
         if probabilistic:
             net2 = NormalParamWrapper(net2)
 
         fnet1, params1 = make_functional(net1)
+        fdummy_net, _ = make_functional(dummy_net)
         fnet2, params2 = make_functional(net2)
         params = params1 + params2
 
@@ -550,10 +605,31 @@ def test_vmap(self, safe, spec_type, probabilistic):
                 out_keys=["hidden"],
                 safe=False,
             )
+            dummy_tdmodule = TDModule(
+                fdummy_net,
+                None,
+                in_keys=["hidden"],
+                out_keys=["hidden"],
+                safe=False,
+            )
             tdmodule2 = tdclass(
                 fnet2, spec, in_keys=["hidden"], out_keys=["out"], safe=safe, **kwargs
             )
-            tdmodule = TDSequence(tdmodule1, tdmodule2)
+            tdmodule = TDSequence(tdmodule1, dummy_tdmodule, tdmodule2)
+
+        assert hasattr(tdmodule, "__setitem__")
+        assert len(tdmodule) == 3
+        tdmodule[1] = tdmodule2
+        assert len(tdmodule) == 3
+
+        assert hasattr(tdmodule, "__delitem__")
+        assert len(tdmodule) == 3
+        del tdmodule[2]
+        assert len(tdmodule) == 2
+
+        assert hasattr(tdmodule, "__getitem__")
+        assert tdmodule[0] is tdmodule1
+        assert tdmodule[1] is tdmodule2
 
         # vmap = True
         params = [p.repeat(10, *[1 for _ in p.shape]) for p in params]