[Refactor] Rename TensorDictSequence to TensorDictSequential (#440)

The TensorDictSequence class behaves like the nn.Sequential module, hence those names should match.

Co-authored-by: Ronert Obst <ronert@fb.com>

Author: ronert

README.md

@@ -112,19 +112,19 @@ algorithms. For instance, here's how to code a rollout in TorchRL:
     + out = tensordict["out"]
     ```
 
-    The `TensorDictSequence` class allows to branch sequences of `nn.Module` instances in a highly modular way.
+    The `TensorDictSequential` class allows to branch sequences of `nn.Module` instances in a highly modular way.
     For instance, here is an implementation of a transformer using the encoder and decoder blocks:
     ```python
     encoder_module = TransformerEncoder(...)
     encoder = TensorDictModule(encoder_module, in_keys=["src", "src_mask"], out_keys=["memory"])
     decoder_module = TransformerDecoder(...)
     decoder = TensorDictModule(decoder_module, in_keys=["tgt", "memory"], out_keys=["output"])
-    transformer = TensorDictSequence(encoder, decoder)
+    transformer = TensorDictSequential(encoder, decoder)
     assert transformer.in_keys == ["src", "src_mask", "tgt"]
     assert transformer.out_keys == ["memory", "output"]
     ```
 
-    `TensorDictSequence` allows to isolate subgraphs by querying a set of desired input / output keys:
+    `TensorDictSequential` allows to isolate subgraphs by querying a set of desired input / output keys:
     ```python
     transformer.select_subsequence(out_keys=["memory"])  # returns the encoder
     transformer.select_subsequence(in_keys=["tgt", "memory"])  # returns the decoder

test/test_exploration.py

@@ -14,7 +14,7 @@
 from torchrl.data.tensordict.tensordict import TensorDict
 from torchrl.envs.transforms.transforms import gSDENoise
 from torchrl.envs.utils import set_exploration_mode
-from torchrl.modules import TensorDictModule, TensorDictSequence
+from torchrl.modules import TensorDictModule, TensorDictSequential
 from torchrl.modules.distributions import TanhNormal
 from torchrl.modules.distributions.continuous import (
     IndependentNormal,
@@ -227,7 +227,7 @@ def test_gsde(
     if gSDE:
         model = torch.nn.LazyLinear(action_dim)
         in_keys = ["observation"]
-        module = TensorDictSequence(
+        module = TensorDictSequential(
             TensorDictModule(model, in_keys=in_keys, out_keys=["action"]),
             TensorDictModule(
                 LazygSDEModule(),

test/test_functorch.py

@@ -5,7 +5,7 @@
 from functorch import vmap
 from torch import nn
 from torchrl.data import TensorDict
-from torchrl.modules import TensorDictModule, TensorDictSequence
+from torchrl.modules import TensorDictModule, TensorDictSequential
 from torchrl.modules.functional_modules import (
     FunctionalModuleWithBuffers,
     FunctionalModule,
@@ -158,7 +158,7 @@ def test_vmap_tdsequence(moduletype, batch_params):
         tdmodule1 = TensorDictModule(fmodule1, in_keys=["x"], out_keys=["y"])
         tdmodule2 = TensorDictModule(fmodule2, in_keys=["y"], out_keys=["z"])
         params = TensorDict({"0": params1, "1": params2}, [])
-        tdmodule = TensorDictSequence(tdmodule1, tdmodule2)
+        tdmodule = TensorDictSequential(tdmodule1, tdmodule2)
         assert {"0", "1"} == set(params.keys())
         x = torch.randn(10, 1, 3)
         td = TensorDict({"x": x}, [10])
@@ -174,7 +174,7 @@ def test_vmap_tdsequence(moduletype, batch_params):
         tdmodule2 = TensorDictModule(fmodule2, in_keys=["y"], out_keys=["z"])
         params = TensorDict({"0": params1, "1": params2}, [])
         buffers = TensorDict({"0": buffers1, "1": buffers2}, [])
-        tdmodule = TensorDictSequence(tdmodule1, tdmodule2)
+        tdmodule = TensorDictSequential(tdmodule1, tdmodule2)
         assert {"0", "1"} == set(params.keys())
         assert {"0", "1"} == set(buffers.keys())
         x = torch.randn(10, 2, 3)
@@ -209,7 +209,7 @@ def test_vmap_tdsequence_nativebuilt(moduletype, batch_params):
     if moduletype == "linear":
         tdmodule1 = TensorDictModule(module1, in_keys=["x"], out_keys=["y"])
         tdmodule2 = TensorDictModule(module2, in_keys=["y"], out_keys=["z"])
-        tdmodule = TensorDictSequence(tdmodule1, tdmodule2)
+        tdmodule = TensorDictSequential(tdmodule1, tdmodule2)
         tdmodule, (params, buffers) = tdmodule.make_functional_with_buffers(native=True)
         assert {"0", "1"} == set(params.keys())
         x = torch.randn(10, 1, 3)
@@ -225,7 +225,7 @@ def test_vmap_tdsequence_nativebuilt(moduletype, batch_params):
     elif moduletype == "bn1":
         tdmodule1 = TensorDictModule(module1, in_keys=["x"], out_keys=["y"])
         tdmodule2 = TensorDictModule(module2, in_keys=["y"], out_keys=["z"])
-        tdmodule = TensorDictSequence(tdmodule1, tdmodule2)
+        tdmodule = TensorDictSequential(tdmodule1, tdmodule2)
         tdmodule, (params, buffers) = tdmodule.make_functional_with_buffers(native=True)
         assert {"0", "1"} == set(params.keys())
         assert {"0", "1"} == set(buffers.keys())

test/test_tensordictmodules.py

@@ -24,7 +24,7 @@
 from torchrl.modules.tensordict_module.probabilistic import (
     ProbabilisticTensorDictModule,
 )
-from torchrl.modules.tensordict_module.sequence import TensorDictSequence
+from torchrl.modules.tensordict_module.sequence import TensorDictSequential
 
 
 class TestTDModule:
@@ -785,7 +785,7 @@ def test_key_exclusion(self):
         module3 = TensorDictModule(
             nn.Linear(3, 4), in_keys=["foo1", "key3"], out_keys=["key2"]
         )
-        seq = TensorDictSequence(module1, module2, module3)
+        seq = TensorDictSequential(module1, module2, module3)
         assert set(seq.in_keys) == {"key1", "key2", "key3"}
         assert set(seq.out_keys) == {"foo1", "key1", "key2"}
 
@@ -838,7 +838,7 @@ def test_stateful(self, safe, spec_type, lazy):
                 safe=False,
                 **kwargs
             )
-            tdmodule = TensorDictSequence(tdmodule1, dummy_tdmodule, tdmodule2)
+            tdmodule = TensorDictSequential(tdmodule1, dummy_tdmodule, tdmodule2)
 
         assert hasattr(tdmodule, "__setitem__")
         assert len(tdmodule) == 3
@@ -926,7 +926,7 @@ def test_stateful_probabilistic(self, safe, spec_type, lazy):
                 safe=False,
                 **kwargs
             )
-            tdmodule = TensorDictSequence(tdmodule1, dummy_tdmodule, tdmodule2)
+            tdmodule = TensorDictSequential(tdmodule1, dummy_tdmodule, tdmodule2)
 
         assert hasattr(tdmodule, "__setitem__")
         assert len(tdmodule) == 3
@@ -998,7 +998,7 @@ def test_functional(self, safe, spec_type):
                 out_keys=["out"],
                 safe=safe,
             )
-            tdmodule = TensorDictSequence(tdmodule1, dummy_tdmodule, tdmodule2)
+            tdmodule = TensorDictSequential(tdmodule1, dummy_tdmodule, tdmodule2)
 
         assert hasattr(tdmodule, "__setitem__")
         assert len(tdmodule) == 3
@@ -1082,7 +1082,7 @@ def test_functional_probabilistic(self, safe, spec_type):
                 safe=safe,
                 **kwargs
             )
-            tdmodule = TensorDictSequence(tdmodule1, dummy_tdmodule, tdmodule2)
+            tdmodule = TensorDictSequential(tdmodule1, dummy_tdmodule, tdmodule2)
 
         assert hasattr(tdmodule, "__setitem__")
         assert len(tdmodule) == 3
@@ -1162,7 +1162,7 @@ def test_functional_with_buffer(
                 out_keys=["out"],
                 safe=safe,
             )
-            tdmodule = TensorDictSequence(tdmodule1, dummy_tdmodule, tdmodule2)
+            tdmodule = TensorDictSequential(tdmodule1, dummy_tdmodule, tdmodule2)
 
         assert hasattr(tdmodule, "__setitem__")
         assert len(tdmodule) == 3
@@ -1255,7 +1255,7 @@ def test_functional_with_buffer_probabilistic(
                 safe=safe,
                 **kwargs
             )
-            tdmodule = TensorDictSequence(tdmodule1, dummy_tdmodule, tdmodule2)
+            tdmodule = TensorDictSequential(tdmodule1, dummy_tdmodule, tdmodule2)
 
         assert hasattr(tdmodule, "__setitem__")
         assert len(tdmodule) == 3
@@ -1331,7 +1331,7 @@ def test_functional_with_buffer_probabilistic_laterconstruct(
                 safe=safe,
                 **kwargs
             )
-            tdmodule = TensorDictSequence(tdmodule1, tdmodule2)
+            tdmodule = TensorDictSequential(tdmodule1, tdmodule2)
 
         tdmodule, (params, buffers) = tdmodule.make_functional_with_buffers()
 
@@ -1396,7 +1396,7 @@ def test_vmap(self, safe, spec_type):
                 out_keys=["out"],
                 safe=safe,
             )
-            tdmodule = TensorDictSequence(tdmodule1, dummy_tdmodule, tdmodule2)
+            tdmodule = TensorDictSequential(tdmodule1, dummy_tdmodule, tdmodule2)
 
         assert hasattr(tdmodule, "__setitem__")
         assert len(tdmodule) == 3
@@ -1496,7 +1496,7 @@ def test_vmap_probabilistic(self, safe, spec_type):
                 safe=safe,
                 **kwargs
             )
-            tdmodule = TensorDictSequence(tdmodule1, tdmodule2)
+            tdmodule = TensorDictSequential(tdmodule1, tdmodule2)
 
         # vmap = True
         params = [p.repeat(10, *[1 for _ in p.shape]) for p in params]
@@ -1546,7 +1546,7 @@ def test_submodule_sequence(self, functional):
             in_keys=["hidden"],
             out_keys=["out"],
         )
-        td_module = TensorDictSequence(td_module_1, td_module_2)
+        td_module = TensorDictSequential(td_module_1, td_module_2)
 
         if functional:
             td_1 = TensorDict({"in": torch.randn(5, 3)}, [5])
@@ -1640,7 +1640,7 @@ def test_sequential_partial(self, stack, functional):
             safe=True,
             **kwargs
         )
-        tdmodule = TensorDictSequence(
+        tdmodule = TensorDictSequential(
             tdmodule1, tdmodule2, tdmodule3, partial_tolerant=True
         )
 
@@ -1686,7 +1686,7 @@ def test_subsequence_weight_update(self):
             in_keys=["hidden"],
             out_keys=["out"],
         )
-        td_module = TensorDictSequence(td_module_1, td_module_2)
+        td_module = TensorDictSequential(td_module_1, td_module_2)
 
         td_1 = TensorDict({"in": torch.randn(5, 3)}, [5])
         sub_seq_1 = td_module.select_subsequence(out_keys=["hidden"])

torchrl/modules/models/exploration.py

@@ -262,7 +262,7 @@ class gSDEModule(nn.Module):
             to the sampled action.
 
     Examples:
-        >>> from torchrl.modules import TensorDictModule, TensorDictSequence, ProbabilisticActor, TanhNormal
+        >>> from torchrl.modules import TensorDictModule, TensorDictSequential, ProbabilisticActor, TanhNormal
         >>> from torchrl.data import TensorDict
         >>> batch, state_dim, action_dim = 3, 7, 5
         >>> model = nn.Linear(state_dim, action_dim)
@@ -274,7 +274,7 @@ class gSDEModule(nn.Module):
         >>> stochatstic_part = ProbabilisticActor(stochatstic_part,
         ...      dist_param_keys=["loc", "scale"],
         ...      distribution_class=TanhNormal)
-        >>> stochatstic_policy = TensorDictSequence(deterministic_policy, stochatstic_part)
+        >>> stochatstic_policy = TensorDictSequential(deterministic_policy, stochatstic_part)
         >>> tensordict = TensorDict({'obs': torch.randn(state_dim), '_epx_gSDE': torch.zeros(1)}, [])
         >>> _ = stochatstic_policy(tensordict)
         >>> print(tensordict)

torchrl/modules/tensordict_module/actors.py

@@ -16,7 +16,7 @@
 from torchrl.modules.tensordict_module.probabilistic import (
     ProbabilisticTensorDictModule,
 )
-from torchrl.modules.tensordict_module.sequence import TensorDictSequence
+from torchrl.modules.tensordict_module.sequence import TensorDictSequential
 
 __all__ = [
     "Actor",
@@ -498,7 +498,7 @@ def __init__(
         )
 
 
-class ActorValueOperator(TensorDictSequence):
+class ActorValueOperator(TensorDictSequential):
     """
     Actor-value operator.
 
@@ -615,21 +615,21 @@ def __init__(
             value_operator,
         )
 
-    def get_policy_operator(self) -> TensorDictSequence:
+    def get_policy_operator(self) -> TensorDictSequential:
         """
 
         Returns a stand-alone policy operator that maps an observation to an action.
 
         """
-        return TensorDictSequence(self.module[0], self.module[1])
+        return TensorDictSequential(self.module[0], self.module[1])
 
-    def get_value_operator(self) -> TensorDictSequence:
+    def get_value_operator(self) -> TensorDictSequential:
         """
 
         Returns a stand-alone value network operator that maps an observation to a value estimate.
 
         """
-        return TensorDictSequence(self.module[0], self.module[2])
+        return TensorDictSequential(self.module[0], self.module[2])
 
 
 class ActorCriticOperator(ActorValueOperator):
@@ -772,7 +772,7 @@ def get_value_operator(self) -> TensorDictModuleWrapper:
         )
 
 
-class ActorCriticWrapper(TensorDictSequence):
+class ActorCriticWrapper(TensorDictSequential):
     """
     Actor-value operator without common module.
 
@@ -863,15 +863,15 @@ def __init__(
             value_operator,
         )
 
-    def get_policy_operator(self) -> TensorDictSequence:
+    def get_policy_operator(self) -> TensorDictSequential:
         """
 
         Returns a stand-alone policy operator that maps an observation to an action.
 
         """
         return self.module[0]
 
-    def get_value_operator(self) -> TensorDictSequence:
+    def get_value_operator(self) -> TensorDictSequential:
         """
 
         Returns a stand-alone value network operator that maps an observation to a value estimate.