[Algorithm] GTrXL PPO with compact window replay and transformer tutorial

.github/unittest/linux_sota/scripts/run_all.sh

@@ -136,7 +136,8 @@ python -c "import functorch"
 
 # install tensordict
 if [[ "$RELEASE" == 0 ]]; then
-  uv pip install --no-deps git+https://github.com/pytorch/tensordict.git
+  # Draft dependency: TensorDict #1789; remove after it merges.
+  uv pip install --no-deps git+https://github.com/pytorch/tensordict.git@2d252bfae380f150a6136930d7d7a5833d68ecf1
 else
   uv pip install --no-deps tensordict
 fi

.github/unittest/linux_sota/scripts/test_sota.py

@@ -17,6 +17,16 @@
 ), "Composite LP must be set to False. Run this test with COMPOSITE_LP_AGGREGATE=0"
 
 commands = {
+    "gtrxl_ppo": """python sota-implementations/gtrxl/gtrxl_ppo.py \
+  env.num_envs=2 \
+  collector.steps_per_batch=9 \
+  collector.total_frames=36 \
+  replay.window_length=4 \
+  replay.batch_size=3 \
+  replay.storage=memmap \
+  loss.epochs=2 \
+  logger.eval_envs=2
+""",
     "vla_grpo": """python sota-implementations/vla_grpo/vla-grpo.py \
   collector.groups_per_iter=2 \
   collector.group_size=2 \

.github/unittest/tutorials/scripts/run_all.sh

@@ -176,7 +176,8 @@ uv pip install --no-progress \
 
 # install tensordict
 if [[ "$RELEASE" == 0 ]]; then
-  uv pip install --no-progress --no-deps git+https://github.com/pytorch/tensordict.git
+  # Draft dependency: TensorDict #1789; remove after it merges.
+  uv pip install --no-progress --no-deps git+https://github.com/pytorch/tensordict.git@2d252bfae380f150a6136930d7d7a5833d68ecf1
 else
   uv pip install --no-progress --no-deps tensordict
 fi

.github/workflows/docs.yml

@@ -73,7 +73,8 @@ jobs:
           # Use nightly PyTorch for main, nightly, and PRs
           python -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu -U --quiet --root-user-action=ignore
           python -m pip install cloudpickle packaging importlib_metadata numpy orjson "pyvers>=0.2.3,<0.3.0" --quiet --root-user-action=ignore
-          python -m pip install git+https://github.com/pytorch/tensordict.git --no-deps --quiet --root-user-action=ignore
+          # Draft dependency: TensorDict #1789; remove after it merges.
+          python -m pip install git+https://github.com/pytorch/tensordict.git@2d252bfae380f150a6136930d7d7a5833d68ecf1 --no-deps --quiet --root-user-action=ignore
           # torchcodec nightly wheels, matching the torch nightly above.
           python -m pip install --pre torchcodec --index-url https://download.pytorch.org/whl/nightly/cpu --quiet --root-user-action=ignore
         fi

benchmarks/ad_hoc/bench_gtrxl_windows.py

@@ -0,0 +1,99 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+"""Compare per-step-state and parallel compact-window GTrXL training on CPU."""
+from __future__ import annotations
+
+import argparse
+import json
+import platform
+from pathlib import Path
+
+import torch
+from tensordict import TensorDict
+from torch.utils.benchmark import Timer
+from torchrl.modules import GTrXL, set_recurrent_mode, TransformerModule
+
+
+def train_window(module, data):
+    module.zero_grad(set_to_none=True)
+    with set_recurrent_mode(True):
+        module(data.clone())["features"].square().mean().backward()
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--output", type=Path, required=True)
+    args = parser.parse_args()
+    torch.set_num_threads(1)
+    results = []
+    for batch, length, memory_len, width in [
+        (1, 8, 8, 16),
+        (32, 32, 16, 32),
+        (32, 64, 64, 64),
+    ]:
+        torch.manual_seed(0)
+        module = TransformerModule(
+            transformer=GTrXL(7, width, 2, memory_len=memory_len),
+            in_keys=["observation", "state"],
+            out_keys=["features", ("next", "state")],
+        )
+        state = module.transformer.state_spec.zero([batch])
+        state["memory"].normal_()
+        state["valid"].fill_(True)
+        observation = torch.randn(batch, length, 7)
+        is_init = torch.zeros(batch, length, 1, dtype=torch.bool)
+        compact = TensorDict(
+            {"observation": observation, "is_init": is_init, "state": state}, [batch]
+        )
+        dense = TensorDict(
+            {
+                "observation": observation,
+                "is_init": is_init,
+                "state": state.unsqueeze(-1).expand(batch, length),
+            },
+            [batch, length],
+        )
+        with set_recurrent_mode(True):
+            expected = module(dense.clone())["features"]
+            actual = module(compact.clone())["features"]
+        torch.testing.assert_close(actual, expected, atol=2e-5, rtol=2e-5)
+        for layout, data in (("per_step", dense), ("compact", compact)):
+            measurement = Timer(
+                stmt="train_window(module, data)",
+                globals={"train_window": train_window, "module": module, "data": data},
+                num_threads=1,
+            ).blocked_autorange(min_run_time=1)
+            memory_bytes = batch * 2 * memory_len * width * observation.element_size()
+            results.append(
+                {
+                    "layout": layout,
+                    "batch": batch,
+                    "window_length": length,
+                    "memory_len": memory_len,
+                    "hidden_size": width,
+                    "latency_ms": measurement.median * 1000,
+                    "transitions_per_second": batch * length / measurement.median,
+                    "carry_payload_bytes": memory_bytes
+                    * (length if layout == "per_step" else 1),
+                }
+            )
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    args.output.write_text(
+        json.dumps(
+            {
+                "platform": platform.platform(),
+                "torch_version": torch.__version__,
+                "num_threads": 1,
+                "description": "Clone input, forward and backward with identical weights/tensors; no optimizer or collection. Plain TensorDict in both paths.",
+                "results": results,
+            },
+            indent=2,
+        )
+        + "\n"
+    )
+
+
+if __name__ == "__main__":
+    main()