[Performance] Chunk process-slot worker results in AsyncBatchedCollector

[vmoens]

Stack: #4305 → #4306 → #4307 → #4308 (base).

Batch worker transitions to reduce queue, unpickling, and replay-write overhead.

collector = AsyncBatchedCollector(
    ...,
    transition_chunk_size=64,  # ProcessSlotTransport only; default: 1
)

Chunks preserve stream order and frame budgets. Larger chunks increase latency and in-flight memory; batches become dense with tensor env_index.

Coverage: ordering/budgets, replay routing, trajectories, pause/shutdown. Reported CPU benchmark: ~6× less driver CPU/frame.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/4305

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 6108be6 with merge base b7f348e ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[vmoens]

Measured locally:     0.226 -> 0.037 ms/frame (~6x less driver CPU)
Projected collection: ~1.5x; ceiling ~370 -> ~1,000 steps/s

The collection estimates are unmeasured GPU projections. Reaching the ceiling requires sufficient learner throughput (#4307).

[vmoens]

Measured on the target hardware (64-environment DreamerV3 run on a pixel task, one 4-GPU node, process-slot inference server with a static batch of 64, compiled and CUDA-graph learner step, bf16, train ratio 16), comparing the full stack #4305 + #4306 + #4307 + #4308 (with transition_chunk_size=64) against the same configuration without these four PRs, over the same elapsed window of the two runs (1,000 s to 4,428 s after training start):

	environment steps/s	learner updates/s (32x256)
without the stack	361	0.70
with the stack	1,249	2.44
speed-up	3.46x	3.46x

The unthrottled collection phase before training reached 1,254 steps/s (298 before), the learner GPU now shows sustained full-utilization bursts instead of about a quarter duty cycle, and the run had no errors over the measured 1.3 hours. The two rates move together because the train ratio couples them, so the stack lifted both the collection ceiling and the learner cadence past the previous limits. This measures the combined effect of the four PRs; per-PR attribution has not been measured.

[vmoens]

Learning check (09 Sep): on the 64-environment task the stack keeps its throughput without CUDA-graph inference batches (1,234 env steps/s and 2.41 updates/s with inference_static_batch_size=null, against 1,249 with them) and the learning curve then matches the previous example (episode score 3.9 / 4.2 / 4.4 at 3-4.5M environment steps). The collapse seen with static batches is explained and guarded in #4310 (the acting policy rewrites its state/belief input keys, which a captured CUDA graph cannot follow).