[BugFix] Bound direct off-policy replay collection

[vmoens]

Summary

• reserve asynchronous Ray collector batches against total_frames before launching actor calls, preventing direct replay writes after the collection target
• track asynchronous Trainer progress with the collector's transition count instead of replay write_count, preserving multidimensional replay items
• connect lazy distributed replay clients to explicitly pre-bound extend and sample endpoints before falling back to payload bootstrap
• retain explicit [B, T] -> [B * T] collector postprocessing for standard feed-forward off-policy workloads without imposing that layout on general replay usage

Root causes

Frame accounting versus replay item accounting

A vectorized collector can emit a rollout with batch shape [num_envs, time]. A replay buffer may intentionally store that rollout as num_envs sequence items, or a workload may flatten it into num_envs * time transition items before insertion. Both layouts are valid.

Replay write_count counts stored items, so it is not necessarily a transition counter. Using it to drive Trainer progress made collection accounting depend on replay layout: grouped sequence storage advanced write_count by num_envs while the collector had produced num_envs * time transitions.

The Trainer now uses collector.collected_frames for asynchronous progress and termination. Replay write_count retains its normal item-count semantics. Regression coverage runs the same Ray Trainer with both grouped [N, T] samples and flat [N] samples and verifies that each completes at the same collected-frame target.

DQN, SAC, DDPG, and TD3 in the validation suite still explicitly flatten collector rollouts before direct insertion because those particular workloads use transition replay. That policy remains in the workload configuration rather than the generic Trainer.

In-flight collection

The asynchronous iterator previously launched one replacement task after every completed batch. Once the target was reached, it waited for the remaining tasks to finish. Direct-replay collectors write inside those actor calls, so every in-flight call completed an extra replay insertion even though its result was never yielded.

The iterator now includes in-flight batches when deciding whether another actor call can be scheduled. This bounds both collected frames and direct replay writes without constraining the shape of each stored replay item.

Pre-bound distributed transport

When extend_spec and sample_spec are supplied, the replay owner establishes the distributed payload endpoints during startup. The lazy client nevertheless used the first extend payload and first sampled result to call the Ray bootstrap methods. The actor did not need those payloads once the schemas were bound, but Ray had already serialized and transferred them.

The lazy client now asks the owner for an existing endpoint first. It retains first-use payload bootstrap as the fallback when no explicit schema was supplied. This keeps dynamic schemas working while avoiding a full Ray transfer of the initial large collector and learner batches.

Tests

• python -m pytest -q test/rb/test_rb_distributed.py::TestRayRB::test_prebound_distributed_replay_skips_payload_bootstrap test/rb/test_rb_distributed.py::TestRayRB::test_ray_replay_with_prebound_gloo_transport test/rb/test_rb_distributed.py::TestRayRB::test_ray_replay_with_gloo_transport test/test_distributed.py::TestRayCollector::test_async_ray_collector_does_not_write_unreserved_batches test/test_distributed.py::TestRayCollector::test_async_trainer_tracks_frames_independently_of_replay_shape test/test_distributed.py::TestRayCollector::test_dqn_trainer_ray_backend test/test_distributed.py::TestRayCollector::test_offpolicy_trainer_ray_backend test/test_distributed.py::TestRayCollector::test_ray_collector_pause_drains_and_resumes test/test_distributed.py::TestRayCollector::test_distributed_collector_sync (13 passed)
• pre-commit run --files torchrl/collectors/distributed/ray.py torchrl/data/replay_buffers/ray_buffer.py torchrl/data/replay_buffers/replay_buffers.py torchrl/trainers/trainers.py test/rb/test_rb_distributed.py test/test_distributed.py
• 100K-environment SAC validation completed 1,000,000 exact replay writes with flat transition samples after the pre-bound endpoint change

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ You can merge normally! (1 Unrelated Failure)

As of commit 85023be with merge base 702709c ():

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Libs Tests on Linux / unittests-gym (3.10, 12.8) / linux-job (gh) (trunk failure)
  ##[error]The operation was canceled.

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