Repository navigation
[BugFix] Fix AsyncEnvPool, weight-sync and multi-collector replay storage correctness bugs - #4701
Merged
theap06 merged 9 commits intoNov 7, 2026
Conversation
AsyncEnvPool.step, _step and step_and_maybe_reset wrote "env_index" into their input as torch.arange, while reset, _reset and the registered NonTensor spec use a NonTensorStack. A Collector keeps that input as its carrier and stacks it into a rollout buffer built from the spec, so Collector(AsyncEnvPool(...)) failed on its first batch with every backend. All five entry points now share _set_env_index, which writes the NonTensorStack expanded to the input batch size.
…, rewards and reset masks
With exchange="shm", wrapping the pool in a TransformedEnv or a Collector
did not work:
- write_input rejected every input entry outside the child schema, such
as "step_count" and "episode_reward" from outer transforms or
("collector", "traj_ids") from a collector. These entries now stay in
the parent process. For the public reset, step and step_and_maybe_reset
they are merged back into the results, which is what the queue
exchange and SerialEnv return. The private _reset and _step drop them,
as their outputs never carry their inputs.
- The worker published the output of the private _step (reward, next
observation, done flags) into the root of the result slot, which has
no reward entry. It now publishes under the slot's "next" entry, which
has exactly that schema, and _async_private_step_recv reads it from
there.
- The slot schema had no "_reset" entry, so partial resets requested by
an outer TransformedEnv could not reach the children. The exchange
template now adds one boolean entry per reset key.
Non-tensor values at schema keys still raise.
…queue ProcessorAsyncEnvPool._drain_result_queues treated every EOFError or OSError from get_nowait as "item consumed, keep draining". That holds when a removed item's tensors cannot be rebuilt, but a queue whose pipe is closed raises OSError on every call without consuming anything, so the loop never ended. This is what hung test_recv_timeout_bounds_whole_call[threading-queue] for 120 s in the optdeps CI run of 2026-11-06 (job 127602086906). The pytest-timeout traceback shows the main thread inside AsyncEnvPool.__del__ -> shutdown -> _drain_result_queues -> get_nowait -> poll of a multiprocessing pool. That pool came from the preceding multiprocessing variant of the same test, where pytest.raises keeps it in a reference cycle. The cyclic garbage collector reclaimed it during the threading variant and finalized the queue connections before the pool's __del__ ran. The threading recv timeout logic itself is fine: that variant takes 2 s when no stale pool is collected during it. After a failed get, the drain now checks whether the queue can still be polled and moves on to the next queue if it cannot.
…the weights MultiProcessWeightSyncScheme is documented as sending weight copies through queues, but on CPU the workers ended up reading the trainer's own parameters: pickling a tensor through a torch.multiprocessing queue moves its storage to shared memory in place, the sender queued the live parameters, and the worker bound the received tensors as its parameters (apply_weights(inplace=False)). Every optimizer step on the trainer was therefore visible to the workers without update_policy_weights_(), and an update landing during a policy call produced a forward that read two weight versions (168 torn forwards out of 3000 steps in a 2-worker repro with an 8M-element parameter). The sender now queues a shared-memory copy of the weights taken at send() time (one per unique entry for the initial sync), the worker clones the initial weights into private parameters, and every later update is copied into those parameters while holding the collector's weight-update lock, the same lock SharedMemWeightSyncScheme(atomic=True) uses. receive() reads the queue before taking the lock so a manual receive does not block policy calls while waiting for weights. On CUDA the workers already received copies whenever the policy was not on the worker device, so this aligns the CPU behavior with the documented contract rather than changing it. test_update_weights now expects MultiProcessWeightSyncScheme workers not to see unsent trainer weights on CPU, and the tear regression test is parametrized over both schemes.
…anging
With weight_sync_schemes={"policy": ..., "postproc.adv_module": ...} (as in
sota-implementations/ppo-async with advantage_on=worker), the receiver applied
the weights to the worker's module and then cascaded the update into
Collector.update_policy_weights_(model_id="postproc.adv_module"). The leaf
Collector only knew how to update its policy, so _maybe_fallback_update raised
KeyError("Collector has no local weight target ..."). The exception was
swallowed by the background receiver thread, no acknowledgment was sent, and
the sender waited forever (MultiProcessWeightSyncScheme) or for
WEIGHT_SYNC_TIMEOUT per worker (SharedMemWeightSyncScheme(atomic=True)).
Collector._maybe_fallback_update now resolves any other model_id as an
attribute path on the collector and copies the weights into that module in
place, so the cascade succeeds and single-process users can update such a
module the same way. A non-policy update no longer bumps the policy version
or notifies the policy of a weight update.
The receiver loops of SharedMemWeightSyncScheme and
MultiProcessWeightSyncScheme no longer leave the sender waiting when applying
the weights fails: the error is logged on the worker, a failure
acknowledgment is sent, and the sender raises a RuntimeError naming the
worker and the error once every targeted worker has answered, so later sends
stay aligned with their acknowledgments.
Tests: a MultiSyncCollector whose GAE postproc value network is synced by its
own scheme receives the trainer's weights and update_policy_weights_ returns
promptly, for both schemes; and a MultiProcessWeightSyncScheme update the
workers cannot apply raises instead of hanging, after which a valid update
still goes through.
…me length A multiprocess collector initializes an empty LazyTensorStorage in the main process by extending it with a fake batch of shape (*env_batch_size, 1), then the workers write batches of shape (*env_batch_size, T) with T = frames_per_batch_worker // env_batch_size.numel(). When the storage has fewer dims than the batch (the default ndim=1 with a batched env), each stored row was laid out with a time length of 1, so the first worker write with T > 1 failed with "indexed destination TensorDict batch size is [4, 1] ... differs from the source batch size [4, 4]" and killed the worker. A single-process Collector writing the same batches to the same storage works, since the storage is laid out from the first real write. The fake batch now has the same time length as the workers' batches, and requested_frames_per_batch is set before the replay buffer is initialized so that frames_per_batch_worker() can be used there. Unbatched envs and T == 1 get the same layout as before. examples/satellite/sac_per_fully_async.py hit this in its CI configuration (4 envs, 4 frames per env), but exited 0 because nothing in the main process noticed the dead worker: the round-robin writer had already advanced the buffer length before the failed write, so the trainer sampled the storage's zero-initialized rows. The example now raises after shutting the collector down if its worker exited with an error.
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/4701
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit db67e25 with merge base c5a5e3a ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
github-actions
Bot
added
BugFix
Documentation
Writing env_index as a NonTensorStack made AsyncEnvPool.step 11-15% slower than main: tensordict.set moved every freshly built stack to the input's device, and each child env then carried the NonTensorData through its step. The stack is now built with its items already on the device (still fresh on every call, so callers own the entry), and env_index is stripped from the data dispatched to the children, which never read it (the workers write their own index on the way back). Queue and threading are back to main's speed; shm stays about 0.1 ms per step slower with mock envs. Adds a full-batch rollout series to the AsyncEnvPool benchmarks.
A send copied every weight tensor into its own shared-memory segment, which each worker then opened, so the cost grew with the number of tensors. A tensordict of CPU weights is now consolidated into a single shared buffer. For a 130-tensor policy a send takes 10.3 ms instead of 21.9 ms (8.3 ms on main, where workers alias the trainer's weights). Adds test_multiprocess_weight_update to the collector benchmarks.
With exchange="shm", extra inputs under "next" were merged into the next root returned by step_and_maybe_reset. The queue exchange and SerialEnv drop them there, as step_mdp does.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Fixes crashes, hangs and stale weights in
AsyncEnvPool, weight synchronization and multi-collector replay storage. Each fix has a regression test that fails or hangs on main.Collector(AsyncEnvPool(...))no longer crashes:env_indexis written as aNonTensorStackon every call.AsyncEnvPoolshared-memory exchange now carries keys added by outer transforms, rewards on_step, and partial-reset masks.AsyncEnvPoolshutdown no longer spins forever on a closed result queue.MultiProcessWeightSyncSchemeworkers keep their own copy of the weights instead of reading the trainer's live tensors (which let policies act on half-written weights). Weights are sent in one shared buffer. Adds a weight-sync benchmark.Breaking changes
MultiProcessWeightSyncSchemeworkers no longer see trainer weight changes untilupdate_policy_weights_()is called. All in-tree users already call it. Each worker holds one extra copy of the weights, and an 8M-parameter send takes ~12 ms instead of ~1 ms.update_policy_weights_()raise instead of hanging.backend="threading"andstack="lazy",AsyncEnvPool.step(td)no longer writes"next"into the caller'std; use the returned tensordict.KeyError.Motivation and Context
These crash a collector, hang the trainer, or let a policy act on half-written weights.
Types of changes
Checklist