[Test] Baseline benchmarks for async collectors and DreamerV3

.github/async-benchmark-requirements.txt

@@ -1,6 +1,7 @@
 # Pinned async trend v1 dependencies: CPython 3.10, Linux x86-64, CUDA 12.6.
 # Use stable wheels so a cache eviction cannot depend on expired nightly wheels.
 # TensorDict source is pinned in benchmarks.yml; bump the trend version with this lock.
+antlr4-python3-runtime==4.9.3
 cloudpickle==3.1.2
 cmake==4.4.0
 cuda-bindings==12.9.7
@@ -10,6 +11,7 @@ exceptiongroup==1.3.1
 filelock==3.32.5
 fsspec==2026.7.0
 hoptorch==0.1.4
+hydra-core==1.3.6
 importlib-metadata==9.0.0
 iniconfig==2.3.0
 jinja2==3.1.6
@@ -33,6 +35,7 @@ nvidia-nccl-cu12==2.29.3
 nvidia-nvjitlink-cu12==12.9.86
 nvidia-nvshmem-cu12==3.4.5
 nvidia-nvtx-cu12==12.6.77
+omegaconf==2.3.1
 orjson==3.11.9
 packaging==26.2
 pluggy==1.6.0
@@ -45,6 +48,7 @@ pytest==9.1.1
 pytest-benchmark==5.2.3
 pytest-timeout==2.4.0
 pyvers==0.2.3
+pyyaml==6.0.3
 setuptools==83.0.0
 setuptools-scm==10.2.0
 sympy==1.14.0

.github/benchmark-requirements.txt

@@ -5,6 +5,7 @@ absl-py==2.5.0
 accelerate==1.14.0
 ale-py==0.8.1
 annotated-doc==0.0.4
+antlr4-python3-runtime==4.9.3
 anyio==4.14.2
 attrs==26.1.0
 autorom==0.4.2
@@ -36,6 +37,7 @@ hoptorch==0.1.4
 httpcore==1.0.9
 httpx==0.28.1
 huggingface-hub==1.23.0
+hydra-core==1.3.6
 idna==3.18
 importlib-metadata==9.0.0
 importlib-resources==7.1.0
@@ -71,6 +73,7 @@ nvidia-nccl-cu12==2.29.3
 nvidia-nvjitlink-cu12==12.9.86
 nvidia-nvshmem-cu12==3.4.5
 nvidia-nvtx-cu12==12.6.77
+omegaconf==2.3.1
 orjson==3.11.9
 packaging==26.2
 pandas==2.3.3

.github/workflows/benchmarks.yml

@@ -11,6 +11,7 @@ on:
     paths:
       - 'torchrl/**'
       - 'benchmarks/**'
+      - 'sota-implementations/dreamer_v3/**'
       - '.github/workflows/benchmarks.yml'
       - '.github/async-benchmark-requirements.txt'
       - '.github/scripts/summarize_async_benchmarks.py'
@@ -108,9 +109,22 @@ jobs:
           deb [snapshot=20260907T000000Z] http://security.ubuntu.com/ubuntu/ jammy-security main universe
           APT
             APT_OPTIONS=(-o Dir::Etc::sourcelist=/tmp/benchmark-sources.list -o Dir::Etc::sourceparts=-)
-            apt-get "${APT_OPTIONS[@]}" update
-            apt-get "${APT_OPTIONS[@]}" install -y --no-install-recommends \
-              git gcc g++ make cmake libpython3.10-dev
+            # The snapshot service answers 502/503 intermittently. A failed
+            # install left git missing and failed every later step of the
+            # push-triggered runs, so retry before giving up.
+            for attempt in 1 2 3 4 5 6; do
+              if apt-get "${APT_OPTIONS[@]}" update \
+                && apt-get "${APT_OPTIONS[@]}" install -y --no-install-recommends \
+                  git gcc g++ make cmake libpython3.10-dev; then
+                break
+              fi
+              if [ "$attempt" = 6 ]; then
+                echo "Installing system packages from the snapshot failed after $attempt attempts" >&2
+                exit 1
+              fi
+              sleep 30
+            done
+            command -v git
           else
             apt-get update -y
             apt-get install software-properties-common cmake -y
@@ -177,7 +191,7 @@ jobs:
             python -m pip install -c /tmp/torch-constraints.txt "pybind11[global]"
             python -m pip install -c /tmp/torch-constraints.txt cloudpickle packaging importlib_metadata numpy orjson "pyvers>=0.2.3,<0.3.0"
             python -m pip install --no-deps git+https://github.com/pytorch/tensordict
-            python -m pip install -c /tmp/torch-constraints.txt safetensors tqdm pandas numpy matplotlib ray psutil
+            python -m pip install -c /tmp/torch-constraints.txt safetensors tqdm pandas numpy matplotlib ray psutil hydra-core omegaconf
 
             bash .github/unittest/helpers/assert_torch_version.sh nightly
             bash .github/unittest/helpers/assert_torch_tensordict_versions.sh nightly
@@ -221,7 +235,8 @@ jobs:
               python -m pytest -v --rank 0 --timeout=300 --benchmark-only \
                 --benchmark-save-data --benchmark-json "async-${repeat}.json" \
                 test_envs_benchmark.py test_collectors_benchmark.py \
-                -k 'async_env_pool or async_collection_pixels'
+                test_dreamer_v3_benchmark.py \
+                -k 'async_env_pool or async_collection_pixels or dreamer_v3_async'
             done
           else
             python -m pytest -vvv --rank 0 --timeout=240 --benchmark-only \

benchmarks/ASYNC_BENCHMARKS.md

@@ -1,8 +1,9 @@
 # Continuous async environment benchmarks
 
 The [Continuous Benchmark workflow](https://github.com/pytorch/rl/actions/workflows/benchmarks.yml)
-runs the short async suite after each main-branch merge touching TorchRL or
-benchmarks. The full suite still runs through the nightly orchestrator.
+runs the short async suite after each main-branch merge touching TorchRL,
+benchmarks or the DreamerV3 example. The full suite still runs through the nightly
+orchestrator.
 The [async trend dashboard](https://pytorch.org/rl/dev/bench/async/) appears after
 the first successful main run. Each point links to its source commit. Branch
 runs produce the same summary and downloadable artifacts without publishing to
@@ -51,6 +52,16 @@ becomes public in #4272. They keep one environment per worker because that
 transport rejects grouped workers. The same workload and batching limits apply;
 these remain separate series alongside the grouped integrated curve.
 
+`async-process-slots-integrated` is the cumulative curve of the direct-transport
+pipeline, the counterpart of `async-shm-integrated`. It keeps one environment per
+worker (the transport rejects grouped workers), enables static inference batches
+on GPU when that option is available, and sends 64 transitions per worker message
+once `AsyncBatchedCollector` accepts `transition_chunk_size`; before that it
+matches `async-process-slots`. Each point records its configuration in the
+`execution` field, so a step in this series can be attributed to the option that
+changed. Use the fixed `async-process-slots*` series to separate the individual
+changes.
+
 `test_async_collection_pixels_64_envs[mode]` adds a separate CUDA acceptance
 comparison for `async-shm` and `async-process-slots`: 64 environments, one per
 worker, with the same CNN plus two 1024-wide layers, one-millisecond environment
@@ -68,6 +79,44 @@ in its tooltip, summary and raw JSON. These explicit configuration transitions
 show the combined pipeline as features land; use the fixed modes to separate
 individual changes and detect regressions.
 
+## DreamerV3 training series
+
+`test_dreamer_v3_async_training[backend-ratio]` (`test_dreamer_v3_benchmark.py`)
+measures the `sota-implementations/dreamer_v3` example end to end. The example
+runs unmodified in a child process using the fixed `dreamer_v3.yaml` benchmark
+configuration, independent of the example defaults. Changes to its training loop,
+replay write-back and inference wiring show up here. The workload is
+`bench_dreamer_v3_env.py`: eight environment processes producing 3 x 64 x 64 uint8
+pixels, an 8-float vector and three boolean
+milestones, with one-millisecond steps, six discrete actions and episode lengths
+staggered by environment index. Collection is asynchronous with shared-memory
+exchange and one environment per worker; the learner runs eagerly on the GPU
+(`optimization.compile=off`) with a small network configuration, replay batches of
+16 sequences of 32 records and replay context write-back enabled. Inference is
+limited to one request per batch: the current thread backend can mix exploration
+contexts while the learner runs, causing larger batches to fail. Keep this limit
+fixed after the bug is resolved; the collection-only series measures batched
+inference.
+
+| Series | Learner updates per collected batch | Emphasis |
+| --- | --- | --- |
+| `ratio2` | `train_ratio=2`: one update per 256-transition batch | collection path, driver overhead |
+| `ratio16` | `train_ratio=16`: eight updates per batch | learner step, replay sampling and write-back |
+
+`thread` serves the acting policy from the training process; `process` serves it
+from a dedicated inference process and skips until the example exposes
+`collector.inference_backend`. The series are CUDA-only.
+
+Throughput is read from the example's own metrics log (`logger.metrics_jsonl`,
+one `train` record per 256-transition batch): after the first learner update,
+two warm-up rounds of 1,024 environment steps pass unmeasured, then five
+measured rounds each wait for the next 1,024 steps to be logged. Setup, process
+startup, replay warm-up and the final shutdown stay outside the measurement.
+The summary reports frames per second like the other series; the raw JSON adds
+learner updates per second, total measured updates and the process-tree RSS. The
+pinned lock includes `hydra-core` and `omegaconf` for the example. CUDA-graph
+inference batches stay disabled in this workload.
+
 ## Reading results
 
 The dashboard plots **transitions per second**, higher is better. Each point is
@@ -114,26 +163,27 @@ To rerun the complete nightly workload, choose `suite: full`.
 Merging a pull request that carries the `benchmarks/trigger` label dispatches
 the full suite on main immediately (`benchmarks_post_merge.yml`), so the merge
 gets a trend point without waiting for the nightly sample. The short async
-suite still runs on every main merge touching `torchrl/` or `benchmarks/`.
+suite also runs on main merges touching `torchrl/`, `benchmarks/` or
+`sota-implementations/dreamer_v3/`.
 
 Locally, install `benchmarks/requirements.txt` and run from the repository:
 
 ```sh
 OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 PYTHONHASHSEED=0 python -m pytest \
   benchmarks/test_envs_benchmark.py benchmarks/test_collectors_benchmark.py \
-  -k 'async_env_pool or async_collection_pixels' --timeout=300 \
+  benchmarks/test_dreamer_v3_benchmark.py \
+  -k 'async_env_pool or async_collection_pixels or dreamer_v3_async' --timeout=300 \
   --benchmark-only --benchmark-save-data --benchmark-json=async-1.json
 ```
 
 Repeat in three separate processes with filenames `async-1.json` through
 `async-3.json`, then run
 `python .github/scripts/summarize_async_benchmarks.py <results-directory> --summary summary.md`.
 
-Merge this benchmark/CI change before #4269 to record the current eager baseline.
-Then review #4269, #4271 and #4272 in order. Their graph, grouping and coordination
-changes appear against the continuing eager series. Replay integration follows
-#4273, #4274 and #4275; it must retain these collection workloads and introduce
-separately named replay-inclusive series if the measured operation changes.
+Merge the benchmark additions and record a successful main-branch run before
+merging #4304, #4305, #4306, #4307 or #4308. Keep the workload fixed while those
+changes land; the process-inference series starts when #4304 exposes its config
+option, alongside the continuing thread-inference baseline.
 
 ## CI health and publication
 

benchmarks/bench_dreamer_v3_env.py

@@ -0,0 +1,97 @@
+# 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.
+"""Fixed fake image-and-vector workload for the DreamerV3 training benchmark.
+
+The environment is importable by name from environment worker processes and by
+the example's ``env.factory`` option (``bench_dreamer_v3_env:FakePixelEnv``). Keep
+its observation layout and step cost stable: the continuous benchmark compares
+revisions on this exact workload (see ASYNC_BENCHMARKS.md).
+"""
+from __future__ import annotations
+
+import time
+
+import torch
+from tensordict import TensorDict
+
+from torchrl.data import Bounded, Composite, OneHot, Unbounded
+from torchrl.envs import EnvBase
+
+PIXEL_SHAPE = (3, 64, 64)
+VECTOR_DIM = 8
+NUM_MILESTONES = 3
+
+
+class FakePixelEnv(EnvBase):
+    """Uint8 pixels, a float vector and boolean milestones with a fixed step cost.
+
+    Episode lengths are staggered by ``env_index`` so resets are not
+    synchronized across environments. Rewards depend on the chosen action, so
+    the learner sees a non-degenerate target.
+    """
+
+    def __init__(
+        self,
+        *,
+        seed: int | None = None,
+        env_index: int = 0,
+        num_envs: int = 1,
+        episode_length: int = 200,
+        num_actions: int = 6,
+        step_latency_s: float = 0.001,
+    ):
+        super().__init__(device="cpu", batch_size=torch.Size([]))
+        self.observation_spec = Composite(
+            pixels=Bounded(0, 255, PIXEL_SHAPE, dtype=torch.uint8),
+            vector=Unbounded((VECTOR_DIM,)),
+            obtained=Unbounded((NUM_MILESTONES,), dtype=torch.bool),
+            shape=(),
+        )
+        self.action_spec = OneHot(num_actions, dtype=torch.float32)
+        self.reward_spec = Unbounded((1,))
+        self.episode_length = episode_length + 7 * (env_index % max(num_envs, 1))
+        self.step_latency_s = step_latency_s
+        self._t = 0
+        self._set_seed(seed)
+
+    def _set_seed(self, seed: int | None):
+        self.rng = torch.Generator().manual_seed(0 if seed is None else int(seed))
+
+    def _observation(self) -> TensorDict:
+        return TensorDict(
+            {
+                "pixels": torch.randint(
+                    0, 256, PIXEL_SHAPE, generator=self.rng, dtype=torch.uint8
+                ),
+                "vector": torch.randn(VECTOR_DIM, generator=self.rng),
+                "obtained": torch.tensor(
+                    [
+                        self._t > self.episode_length // 4,
+                        self._t > self.episode_length // 2,
+                        self._t > 3 * self.episode_length // 4,
+                    ]
+                ),
+            },
+            [],
+        )
+
+    def _reset(self, tensordict=None, **kwargs) -> TensorDict:
+        self._t = 0
+        observation = self._observation()
+        observation.set("done", torch.zeros(1, dtype=torch.bool))
+        observation.set("terminated", torch.zeros(1, dtype=torch.bool))
+        return observation
+
+    def _step(self, tensordict: TensorDict) -> TensorDict:
+        if self.step_latency_s > 0:
+            time.sleep(self.step_latency_s)
+        self._t += 1
+        observation = self._observation()
+        action = tensordict["action"].float().argmax(-1, keepdim=True)
+        observation.set("reward", 0.1 * action.float())
+        done = torch.tensor([self._t >= self.episode_length])
+        observation.set("done", done)
+        observation.set("terminated", done)
+        return observation