From 0a211a0b2af62437a981408308513a6ceec6a763 Mon Sep 17 00:00:00 2001 From: Achintya P Date: Thu, 5 Nov 2026 20:12:31 -0800 Subject: [PATCH 1/9] [BugFix] Write AsyncEnvPool env_index as a NonTensorStack on every call 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. --- test/envs/test_special.py | 46 +++++++++++++++++++++++++++++++++++++- torchrl/envs/async_envs.py | 30 ++++++++++++++++--------- 2 files changed, 64 insertions(+), 12 deletions(-) diff --git a/test/envs/test_special.py b/test/envs/test_special.py index 08ba09c7902..4f9b9a194ee 100644 --- a/test/envs/test_special.py +++ b/test/envs/test_special.py @@ -42,7 +42,7 @@ _to_device_mps_safe, ) from torchrl.envs.env_creator import get_env_metadata -from torchrl.envs.transforms import StepCounter, TransformedEnv +from torchrl.envs.transforms import RewardScaling, StepCounter, TransformedEnv from torchrl.envs.transforms.transforms import Tokenizer from torchrl.envs.utils import check_env_specs from torchrl.modules import RandomPolicy @@ -53,6 +53,7 @@ ContinuousActionVecMockEnv, CountingBatchedEnv, CountingEnv, + CountingEnvCountPolicy, CountingEnvWithString, DiscreteActionConvMockEnv, DiscreteActionConvMockEnvNumpy, @@ -84,6 +85,14 @@ def _step(self, tensordict): return super()._step(tensordict) +def _rewarded_counting_env(index: int) -> TransformedEnv: + """A CountingEnv whose reward is ``index + 1`` at every step and which is done after 3 steps.""" + return TransformedEnv( + CountingEnv(max_steps=3, start_val=index), + RewardScaling(loc=float(index + 1), scale=1.0), + ) + + class _DifferentSpecCountingEnv(CountingEnv): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -1145,6 +1154,41 @@ def test_async_transformed(self, backend, make_envs): finally: base_env._maybe_shutdown() + @pytest.mark.parametrize( + "backend,exchange", + [ + ("threading", None), + ("multiprocessing", "queue"), + ], + ) + @set_capture_non_tensor_stack(False) + def test_collector_matches_serial_env(self, backend, exchange): + """A Collector over an AsyncEnvPool yields the batches of a SerialEnv. + + The pool wrote ``"env_index"`` as a tensor on step but as a + ``NonTensorStack`` on reset, so the collector could not stack its + rollout. + """ + makers = [partial(_rewarded_counting_env, index) for index in range(3)] + kwargs = {} if exchange is None else {"exchange": exchange} + pool = AsyncEnvPool(makers, backend=backend, **kwargs) + serial = SerialEnv(3, makers) + policy = CountingEnvCountPolicy(serial.action_spec) + try: + pool_batches = list( + Collector(pool, policy, frames_per_batch=15, total_frames=30) + ) + serial_batches = list( + Collector(serial, policy, frames_per_batch=15, total_frames=30) + ) + assert len(pool_batches) == len(serial_batches) == 2 + for pool_batch, serial_batch in zip(pool_batches, serial_batches): + assert pool_batch["env_index"] == [[index] * 5 for index in range(3)] + assert_allclose_td(pool_batch.exclude("env_index"), serial_batch) + finally: + pool.close(raise_if_closed=False) + serial.close(raise_if_closed=False) + @set_capture_non_tensor_stack(False) def test_shared_memory_exchange(self, make_envs): env = AsyncEnvPool(make_envs, backend="multiprocessing", exchange="shm") diff --git a/torchrl/envs/async_envs.py b/torchrl/envs/async_envs.py index 17e4f2c5879..e5e20a86291 100644 --- a/torchrl/envs/async_envs.py +++ b/torchrl/envs/async_envs.py @@ -481,10 +481,7 @@ def _reset( tensordict = TensorDict( batch_size=(self.num_envs,) + self.env_batch_sizes[0] ) - env_idx_nt = NonTensorStack(*range(tensordict.shape[0])) - while env_idx_nt.batch_dims < tensordict.batch_dims: - env_idx_nt = expand_as_right(env_idx_nt, tensordict) - tensordict[self._env_idx_key] = env_idx_nt + self._set_env_index(tensordict) self._async_private_reset_send(tensordict) tensordict = self._async_private_reset_recv(min_get=self.num_envs) return tensordict @@ -492,7 +489,7 @@ def _reset( def _step(self, tensordict: TensorDictBase) -> TensorDictBase: if self._current_step > 0: raise RuntimeError("Some envs are still processing a step.") - tensordict.set(self._env_idx_key, torch.arange(tensordict.shape[0])) + self._set_env_index(tensordict) self._async_private_step_send(tensordict) tensordict = self._async_private_step_recv(min_get=self.num_envs) # Using pop instead of del to account for tensorclasses @@ -504,7 +501,7 @@ def step_and_maybe_reset( ) -> tuple[TensorDictBase, TensorDictBase]: if self._current_step_reset > 0: raise RuntimeError("Some envs are still processing a step.") - tensordict.set(self._env_idx_key, torch.arange(tensordict.shape[0])) + self._set_env_index(tensordict) self.async_step_and_maybe_reset_send(tensordict) tensordict, tensordict_ = self.async_step_and_maybe_reset_recv( min_get=self.num_envs @@ -514,7 +511,7 @@ def step_and_maybe_reset( def step(self, tensordict: TensorDictBase) -> TensorDictBase: if self._current_step > 0: raise RuntimeError("Some envs are still processing a step.") - tensordict.set(self._env_idx_key, torch.arange(tensordict.shape[0])) + self._set_env_index(tensordict) self.async_step_send(tensordict) tensordict = self.async_step_recv(min_get=self.num_envs) return tensordict @@ -543,14 +540,25 @@ def reset( tensordict = TensorDict( batch_size=(self.num_envs,) + self.env_batch_sizes[0] ) - indices = NonTensorStack(*range(tensordict.shape[0])) - if indices.shape != tensordict.shape: - indices = expand_as_right(indices, tensordict) - tensordict[self._env_idx_key] = indices + self._set_env_index(tensordict) self.async_reset_send(tensordict) tensordict = self.async_reset_recv(min_get=self.num_envs) return tensordict + def _set_env_index(self, tensordict: TensorDictBase) -> None: + """Writes the pool-wide ``"env_index"`` entry into ``tensordict`` in place. + + The entry is a :class:`~tensordict.NonTensorStack` matching the + ``NonTensor`` spec registered in :attr:`full_state_spec`, so that the + data written by :meth:`reset`, :meth:`step` and + :meth:`step_and_maybe_reset` all share one type and can be stacked + together (e.g. by a collector writing into a preallocated rollout). + """ + env_idx_nt = NonTensorStack(*range(tensordict.shape[0])) + if env_idx_nt.shape != tensordict.shape: + env_idx_nt = expand_as_right(env_idx_nt, tensordict) + tensordict.set(self._env_idx_key, env_idx_nt) + def _sort_results(self, results, *other_results): # Extract env indices from results. When child envs have a batch dimension # (e.g., batch_size=(1,)), r[self._env_idx_key] may be a 1D sequence From b812d8d8024db7ac627c02ef834eba1d5016ea23 Mon Sep 17 00:00:00 2001 From: Achintya P Date: Thu, 5 Nov 2026 20:14:15 -0800 Subject: [PATCH 2/9] [BugFix] Let the AsyncEnvPool shared-memory exchange carry outer keys, 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. --- docs/source/reference/envs_vectorized.rst | 6 +- test/envs/test_special.py | 62 ++++++++++-- torchrl/envs/_async_exchange.py | 110 ++++++++++++++++------ torchrl/envs/async_envs.py | 77 +++++++++++++-- 4 files changed, 211 insertions(+), 44 deletions(-) diff --git a/docs/source/reference/envs_vectorized.rst b/docs/source/reference/envs_vectorized.rst index b102fe60221..38ee5265f54 100644 --- a/docs/source/reference/envs_vectorized.rst +++ b/docs/source/reference/envs_vectorized.rst @@ -346,8 +346,10 @@ TensorDict payloads through queues: ) This mode requires identical child batch sizes, TensorDict keys, tensor shapes, -and CPU devices. Queue exchange remains the default for dynamic or non-tensor -data. Receive methods also accept ``max_get`` and ``timeout``. ``min_get`` is a +and CPU devices. Input entries outside the child schema, such as keys written by +transforms wrapping the pool or by a collector, stay in the parent process and +are merged back into the results, as with the queue exchange. Queue exchange +remains the default for dynamic or non-tensor data. Receive methods also accept ``max_get`` and ``timeout``. ``min_get`` is a hard lower bound. The timeout starts when the first result arrives and does not limit the wait for ``min_get`` results. After ``min_get`` is reached, the call returns when that deadline expires or ``max_get`` results have arrived. Exchange diff --git a/test/envs/test_special.py b/test/envs/test_special.py index 4f9b9a194ee..8bfd47dca37 100644 --- a/test/envs/test_special.py +++ b/test/envs/test_special.py @@ -42,7 +42,13 @@ _to_device_mps_safe, ) from torchrl.envs.env_creator import get_env_metadata -from torchrl.envs.transforms import RewardScaling, StepCounter, TransformedEnv +from torchrl.envs.transforms import ( + Compose, + RewardScaling, + RewardSum, + StepCounter, + TransformedEnv, +) from torchrl.envs.transforms.transforms import Tokenizer from torchrl.envs.utils import check_env_specs from torchrl.modules import RandomPolicy @@ -1159,6 +1165,7 @@ def test_async_transformed(self, backend, make_envs): [ ("threading", None), ("multiprocessing", "queue"), + ("multiprocessing", "shm"), ], ) @set_capture_non_tensor_stack(False) @@ -1167,7 +1174,8 @@ def test_collector_matches_serial_env(self, backend, exchange): The pool wrote ``"env_index"`` as a tensor on step but as a ``NonTensorStack`` on reset, so the collector could not stack its - rollout. + rollout. With ``exchange="shm"`` the collector's ``("collector", + "traj_ids")`` input was also rejected instead of being passed through. """ makers = [partial(_rewarded_counting_env, index) for index in range(3)] kwargs = {} if exchange is None else {"exchange": exchange} @@ -1189,6 +1197,43 @@ def test_collector_matches_serial_env(self, backend, exchange): pool.close(raise_if_closed=False) serial.close(raise_if_closed=False) + @pytest.mark.parametrize("exchange", ["queue", "shm"]) + @set_capture_non_tensor_stack(False) + def test_outer_transforms_match_serial_env(self, exchange): + """Transforms wrapping the pool produce the rollout they produce around a SerialEnv. + + With ``exchange="shm"`` the pool rejected the keys written by the + outer transforms (``"step_count"``, ``"episode_reward"``), lost the + reward of its private ``_step`` and could not forward the partial + ``"_reset"`` mask. Each child is done after a different number of + steps, so the rollout goes through partial resets. + """ + makers = [partial(_rewarded_counting_env, index) for index in range(3)] + transform = Compose(StepCounter(), RewardSum()) + pool_env = TransformedEnv( + AsyncEnvPool(makers, backend="multiprocessing", exchange=exchange), + transform.clone(), + ) + serial_env = TransformedEnv(SerialEnv(3, makers), transform.clone()) + policy = CountingEnvCountPolicy(serial_env.action_spec) + try: + rollout = pool_env.rollout(6, policy, break_when_any_done=False) + expected = serial_env.rollout(6, policy, break_when_any_done=False) + assert_allclose_td(rollout.exclude("env_index"), expected) + torch.testing.assert_close( + expected["next", "episode_reward"].squeeze(-1), + torch.tensor( + [ + [1.0, 2.0, 3.0, 4.0, 1.0, 2.0], + [2.0, 4.0, 6.0, 2.0, 4.0, 6.0], + [3.0, 6.0, 3.0, 6.0, 3.0, 6.0], + ] + ), + ) + finally: + pool_env.close(raise_if_closed=False) + serial_env.close(raise_if_closed=False) + @set_capture_non_tensor_stack(False) def test_shared_memory_exchange(self, make_envs): env = AsyncEnvPool(make_envs, backend="multiprocessing", exchange="shm") @@ -1208,15 +1253,20 @@ def test_shared_memory_exchange(self, make_envs): assert stats["avg_batch_to_action_ms"] > 0 assert stats["consumer_busy_fraction"] > 0 - invalid = next_step.clone() - invalid.set("unknown_input", torch.zeros(invalid.shape)) + extra = next_step.clone() + extra.set("unknown_input", torch.arange(env.num_envs, dtype=torch.float32)) + env.async_step_send(extra) + echoed = env.async_step_recv(min_get=env.num_envs) + torch.testing.assert_close(echoed["unknown_input"], extra["unknown_input"]) + + invalid = echoed.exclude("next") + invalid.set("observation", NonTensorStack(*["a"] * env.num_envs)) with pytest.raises(KeyError) as exc_info: env.async_step_send(invalid) message = str(exc_info.value) - assert "unknown_input" in message + assert "observation" in message assert "fixed exchange schema" in message assert "action" in message - assert "observation" in message finally: env._maybe_shutdown() diff --git a/torchrl/envs/_async_exchange.py b/torchrl/envs/_async_exchange.py index 7563db44d38..be61be3c0b7 100644 --- a/torchrl/envs/_async_exchange.py +++ b/torchrl/envs/_async_exchange.py @@ -10,9 +10,9 @@ from typing import Any import torch -from tensordict import TensorDictBase +from tensordict import is_leaf_nontensor, TensorDictBase from tensordict.tensorclass import NonTensorData, NonTensorStack -from tensordict.utils import expand_as_right +from tensordict.utils import expand_as_right, NestedKey from torchrl._utils import timeit @@ -108,6 +108,7 @@ def __init__(self, fake_tensordicts: list[TensorDictBase]) -> None: self.result_slots = self.result_buffer.unbind(0) self.next_slots = self.next_buffer.unbind(0) self._input_keys = set(self.input_buffer.keys(True, True)) + self._extra_inputs: dict[int, TensorDictBase] = {} self._lock = threading.Lock() self._clock = timeit("async_env_shared_exchange").start() self._leased_at: dict[int, float] = {} @@ -187,30 +188,71 @@ def write_input( tensordict: TensorDictBase, *, record_action: bool = True, + keep_extra_inputs: bool = False, ) -> tuple: + """Writes the schema entries of ``tensordict`` into the input slot of ``env_index``. + + Entries outside the exchange schema (for instance keys written by a + transform wrapping the pool, or by a collector) cannot be consumed by + the child env. When ``keep_extra_inputs`` is ``True`` they are cloned + and kept in the parent process, then merged back into the next result + read for ``env_index``, matching what the child's public ``reset``, + ``step`` and ``step_and_maybe_reset`` return with the queue exchange. + Otherwise (private ``_reset`` / ``_step`` calls, whose outputs do not + carry their inputs) they are ignored. + + Returns: + the tuple of schema keys written to the slot, for the worker to + select. + """ tensor_keys = [] + extra_keys = [] unsupported = [] - for key, value in tensordict.items(True, True): + for key, value in tensordict.items(True, True, is_leaf=is_leaf_nontensor): if key == "env_index": continue - if not isinstance(value, torch.Tensor) or key not in self._input_keys: + if key not in self._input_keys: + extra_keys.append(key) + elif not isinstance(value, torch.Tensor): unsupported.append(key) else: tensor_keys.append(key) if unsupported: expected_keys = sorted(self._input_keys, key=repr) raise KeyError( - f"Shared slot exchange received unsupported input at keys {unsupported}. " + f"Shared slot exchange received non-tensor input at keys {unsupported}. " f"The fixed exchange schema expects tensor keys {expected_keys}. " "Use exchange='queue' for dynamic data." ) self.input_slots[env_index].update_( tensordict.select(*tensor_keys, strict=True) ) + if keep_extra_inputs and extra_keys: + self._extra_inputs[env_index] = tensordict.select( + *extra_keys, strict=True + ).clone() + else: + self._extra_inputs.pop(env_index, None) if record_action: self.record_action(env_index) return tuple(tensor_keys) + def _with_extra_inputs( + self, env_index: int, *results: TensorDictBase + ) -> tuple[TensorDictBase, ...]: + """Merges the extra inputs kept for ``env_index`` into ``results``. + + Each result gets its own copy so that the root and ``"next"`` outputs + of ``step_and_maybe_reset`` never alias each other. + """ + extra = self._extra_inputs.pop(env_index, None) + if extra is None: + return results + merged = [results[0].update(extra)] + for result in results[1:]: + merged.append(result.update(extra.clone())) + return tuple(merged) + @staticmethod def publish( slot: TensorDictBase, tensordict: TensorDictBase, clock: timeit @@ -292,12 +334,30 @@ def _record_received( self._ready_dwell_s += now - ready_s self._ready_dwell_count += 1 - def read(self, descriptors: list[tuple], stack_func) -> TensorDictBase: + def read( + self, + descriptors: list[tuple], + stack_func, + *, + slot_key: NestedKey | None = None, + ) -> TensorDictBase: + """Stacks the results described by ``descriptors``. + + ``slot_key`` selects a sub-tensordict of each result slot; the private + ``_step`` publishes its output under ``"next"`` since that is where the + step output keys (e.g. the reward) live in the exchange schema. + """ results = [] indices = [] for env_index, keys, _ in descriptors: indices.append(env_index) - results.append(self.result_slots[env_index].select(*keys, strict=True)) + slot = self.result_slots[env_index] + if slot_key is not None: + slot = slot.get(slot_key) + (result,) = self._with_extra_inputs( + env_index, slot.select(*keys, strict=True) + ) + results.append(result) result = stack_func(results) indices_data = NonTensorStack(*indices) while indices_data.batch_dims < result.batch_dims: @@ -307,12 +367,10 @@ def read(self, descriptors: list[tuple], stack_func) -> TensorDictBase: def read_one(self, descriptor: tuple) -> TensorDictBase: env_index, keys, _ = descriptor - return ( - self.result_slots[env_index] - .select(*keys, strict=True) - .clone() - .set("env_index", NonTensorData(data=env_index)) + (result,) = self._with_extra_inputs( + env_index, self.result_slots[env_index].select(*keys, strict=True).clone() ) + return result.set("env_index", NonTensorData(data=env_index)) def read_pair( self, descriptors: list[tuple], stack_func @@ -322,12 +380,13 @@ def read_pair( indices = [] for env_index, result_keys, next_keys, _ in descriptors: indices.append(env_index) - results.append( - self.result_slots[env_index].select(*result_keys, strict=True) - ) - next_results.append( - self.next_slots[env_index].select(*next_keys, strict=True) + result, next_result = self._with_extra_inputs( + env_index, + self.result_slots[env_index].select(*result_keys, strict=True), + self.next_slots[env_index].select(*next_keys, strict=True), ) + results.append(result) + next_results.append(next_result) result = stack_func(results) next_result = stack_func(next_results) indices_data = NonTensorStack(*indices) @@ -340,18 +399,13 @@ def read_pair( def read_pair_one(self, descriptor: tuple) -> tuple[TensorDictBase, TensorDictBase]: env_index, result_keys, next_keys, _ = descriptor index_data = NonTensorData(data=env_index) - result = ( - self.result_slots[env_index] - .select(*result_keys, strict=True) - .clone() - .set("env_index", index_data) - ) - next_result = ( - self.next_slots[env_index] - .select(*next_keys, strict=True) - .clone() - .set("env_index", index_data.clone()) + result, next_result = self._with_extra_inputs( + env_index, + self.result_slots[env_index].select(*result_keys, strict=True).clone(), + self.next_slots[env_index].select(*next_keys, strict=True).clone(), ) + result.set("env_index", index_data) + next_result.set("env_index", index_data.clone()) return result, next_result def record_action(self, env_index: int) -> None: diff --git a/torchrl/envs/async_envs.py b/torchrl/envs/async_envs.py index e5e20a86291..4f70a8c30b1 100644 --- a/torchrl/envs/async_envs.py +++ b/torchrl/envs/async_envs.py @@ -28,7 +28,7 @@ from tensordict.tensorclass import NonTensorData, NonTensorStack from tensordict.utils import expand_as_right, NestedKey -from torchrl._utils import logger as torchrl_logger, timeit +from torchrl._utils import _replace_last, logger as torchrl_logger, timeit from torchrl.data.tensor_specs import Composite, NonTensor, TensorSpec from torchrl.envs._async_exchange import _receive_batch, _SharedSlotExchange from torchrl.envs.common import _EnvPostInit, EnvBase @@ -121,7 +121,10 @@ class AsyncEnvPool(EnvBase, metaclass=_AsyncEnvMeta): memory so retaining results does not retain one mapping per tensor. ``"shm"`` stores fixed-shape tensor data in shared slots and sends only readiness descriptors through queues; it requires identical, - fixed-shape, CPU, tensor-only schemas across workers. ``"auto"`` + fixed-shape, CPU, tensor-only schemas across workers. Input + entries outside that schema (e.g. written by transforms wrapping + the pool or by a collector) stay in the parent process and are + merged back into the results, as the queue exchange does. ``"auto"`` selects ``"shm"`` when the env schema supports it and falls back to ``"queue"`` otherwise (the resolution is reported by :attr:`resolved_exchange` and logged on fallback). Defaults to @@ -1023,11 +1026,15 @@ def _send_worker_batches( *, per_env: bool, record_action: bool, + keep_extra_inputs: bool, ) -> None: requests = [[] for _ in range(self.num_workers)] for env_index, local_td in zip(env_idx, local_tds, strict=True): data = self._prepare_worker_data( - env_index, local_td, record_action=record_action + env_index, + local_td, + record_action=record_action, + keep_extra_inputs=keep_extra_inputs, ) requests[self._env_to_worker[env_index]].append((env_index, data)) for worker_index, worker_requests in enumerate(requests): @@ -1040,11 +1047,15 @@ def _prepare_worker_data( tensordict: TensorDictBase, *, record_action: bool, + keep_extra_inputs: bool, ): if self._slot_exchange is None: return tensordict return self._slot_exchange.write_input( - env_index, tensordict, record_action=record_action + env_index, + tensordict, + record_action=record_action, + keep_extra_inputs=keep_extra_inputs, ) def _check_worker_errors(self) -> None: @@ -1121,6 +1132,7 @@ def async_step_send( tensordict.unbind(0), per_env=_per_env, record_action=True, + keep_extra_inputs=True, ) if not _per_env: self._current_step = self._current_step + len(env_idx) @@ -1148,6 +1160,21 @@ def async_step_recv( check_worker_errors=self._check_worker_errors, ) return self._slot_exchange.read_one(descriptor) + return self._recv_step_batch(min_get, max_get, timeout, slot_key=None) + + def _recv_step_batch( + self, + min_get: int | None, + max_get: int | None, + timeout: float | None, + *, + slot_key: NestedKey | None, + ) -> TensorDictBase: + """Receives a batch of step results from the shared step queue. + + ``slot_key`` is forwarded to the shared-memory exchange reader and + names the sub-tensordict of the result slots the workers wrote to. + """ if min_get is None: min_get = self.min_get if min_get > self._current_step: @@ -1164,7 +1191,7 @@ def async_step_recv( self._current_step = self._current_step - len(r) if self._slot_exchange is not None: idx = [item[0] for item in r] - result = self._slot_exchange.read(r, self._stack_func) + result = self._slot_exchange.read(r, self._stack_func, slot_key=slot_key) self._busy.difference_update(idx) return result r, idx = self._sort_results(r) @@ -1188,10 +1215,24 @@ def _async_private_step_send( tensordict.unbind(0), per_env=False, record_action=True, + keep_extra_inputs=False, ) self._current_step = self._current_step + len(env_idx) - _async_private_step_recv = async_step_recv + def _async_private_step_recv( + self, + min_get: int = 1, + *, + max_get: int | None = None, + timeout: float | None = None, + ) -> TensorDictBase: + """Receives the outputs of private ``_step`` calls. + + Workers publish a ``_step`` output under the ``"next"`` entry of their + shared result slot, which is where the exchange schema stores the step + outputs (reward included). + """ + return self._recv_step_batch(min_get, max_get, timeout, slot_key="next") def async_step_and_maybe_reset_send( self, tensordict: TensorDictBase, env_index: int | list[int] | None = None @@ -1213,6 +1254,7 @@ def async_step_and_maybe_reset_send( tensordict.unbind(0), per_env=_per_env, record_action=True, + keep_extra_inputs=True, ) def async_step_and_maybe_reset_recv( @@ -1285,6 +1327,7 @@ def async_reset_send( tensordict.unbind(0), per_env=_per_env, record_action=False, + keep_extra_inputs=True, ) def async_reset_recv( @@ -1352,6 +1395,7 @@ def _async_private_reset_send( tensordict.unbind(0), per_env=False, record_action=False, + keep_extra_inputs=False, ) _async_private_reset_recv = async_reset_recv @@ -1507,6 +1551,23 @@ def _worker_exec( for env_index, data in requests: local_input_queues[env_index].put((msg, data, per_env)) + @staticmethod + def _exchange_template(env: EnvBase) -> TensorDictBase: + """Returns the schema of the shared-memory exchange slots for ``env``. + + This is the env's fake tensordict plus one boolean entry per reset key + (e.g. ``"_reset"``), shaped like the matching done entry, so that the + partial-reset masks written by transforms wrapping the pool reach the + child env's reset. + """ + template = env.fake_tensordict() + for done_key, *_ in env.done_keys_groups: + template.set( + _replace_last(done_key, "_reset"), + torch.zeros_like(template.get(done_key)), + ) + return template + @classmethod def _env_exec( cls, @@ -1560,7 +1621,7 @@ def _env_exec( if msg == "get_specs": output_queue.put((i, env.specs)) elif msg == "get_fake_tensordict": - output_queue.put((i, env.fake_tensordict())) + output_queue.put((i, cls._exchange_template(env))) elif msg == "init_shm": shared_slots = data output_queue.put(True) @@ -1633,7 +1694,7 @@ def _env_exec( step_queue.put(data) else: keys, ready_s = _SharedSlotExchange.publish( - shared_slots[1], data, shared_slots[3] + shared_slots[1].get("next"), data, shared_slots[3] ) step_queue.put((i, keys, ready_s)) elif msg == "shutdown": From de1a960f50e031473183d1c92fbfb5591a50a97c Mon Sep 17 00:00:00 2001 From: Achintya P Date: Thu, 5 Nov 2026 20:15:30 -0800 Subject: [PATCH 3/9] [BugFix] Stop AsyncEnvPool shutdown from spinning on a closed result 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. --- test/envs/test_special.py | 20 ++++++++++++++++++++ torchrl/envs/async_envs.py | 25 ++++++++++++++++++++----- 2 files changed, 40 insertions(+), 5 deletions(-) diff --git a/test/envs/test_special.py b/test/envs/test_special.py index 8bfd47dca37..d6aa8c04485 100644 --- a/test/envs/test_special.py +++ b/test/envs/test_special.py @@ -1461,6 +1461,26 @@ def test_recv_timeout_bounds_whole_call(self, backend, exchange): finally: env._maybe_shutdown() + def test_shutdown_returns_when_result_queue_pipe_closed(self): + """Shutdown terminates when a result queue's pipe is already closed. + + When the garbage collector reclaims a pool caught in a reference + cycle, it may finalize the result queues' connections before + ``AsyncEnvPool.__del__`` shuts the pool down. Draining a closed queue + fails on every ``get`` without consuming anything, which made the + drain loop spin forever and hang whatever test ran next. + """ + env = AsyncEnvPool( + [CountingEnv, CountingEnv], backend="multiprocessing", exchange="queue" + ) + env.reset() + env.close() + env.step_queue._reader.close() + closer = threading.Thread(target=env.close, daemon=True) + closer.start() + closer.join(timeout=30) + assert not closer.is_alive() + @set_capture_non_tensor_stack(False) def test_queue_shutdown_with_unread_results(self): """Shutdown must not deadlock when unread results are still buffered. diff --git a/torchrl/envs/async_envs.py b/torchrl/envs/async_envs.py index 4f70a8c30b1..9c2c4c702ec 100644 --- a/torchrl/envs/async_envs.py +++ b/torchrl/envs/async_envs.py @@ -1417,11 +1417,26 @@ def _drain_result_queues(self) -> None: except queue.Empty: break except (EOFError, OSError): - # The item has already been removed from the queue, but - # rebuilding a discarded tensor can fail once its worker's - # resource sharer has exited. Keep draining the remaining - # items so other workers can finish flushing their queues. - continue + if not self._is_readable(result_queue): + break + + @staticmethod + def _is_readable(result_queue) -> bool: + """Whether ``result_queue`` can still be polled after a failed ``get``. + + A ``get`` that fails while draining either removed an item whose + tensors could not be rebuilt (its worker's resource sharer has exited), + in which case the queue stays readable and draining continues so other + workers can flush their queues, or it hit a closed pipe. The latter + happens when the garbage collector finalizes the queue connections + before :meth:`AsyncEnvPool.__del__` runs: every subsequent ``get`` + fails without consuming anything, so draining must stop. + """ + try: + result_queue.empty() + except OSError: + return False + return True def shutdown(self): deadline = timeit("async_env_shutdown_deadline").start() From 937ce68fb468dbb6dba629e5ceca5f8b30acb2db Mon Sep 17 00:00:00 2001 From: Achintya P Date: Thu, 5 Nov 2026 20:19:43 -0800 Subject: [PATCH 4/9] [BugFix] Give MultiProcessWeightSyncScheme workers their own copy of 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. --- test/test_collectors.py | 22 ++++-- torchrl/weight_update/_mp.py | 130 +++++++++++++++++++++++++++-------- 2 files changed, 119 insertions(+), 33 deletions(-) diff --git a/test/test_collectors.py b/test/test_collectors.py index 47e0a330904..02d00259005 100644 --- a/test/test_collectors.py +++ b/test/test_collectors.py @@ -2744,8 +2744,10 @@ def create_env(): state_dict = collector.state_dict() policy_state_dict = policy.state_dict() # check they don't match - workers_see_unsent_weights = not torch.cuda.is_available() and not ( - scheme is not None and scheme.atomic + workers_see_unsent_weights = ( + not torch.cuda.is_available() + and not isinstance(scheme, MultiProcessWeightSyncScheme) + and not (scheme is not None and scheme.atomic) ) for worker in range(3): for k in state_dict[f"worker{worker}"]["policy_state_dict"]: @@ -4803,7 +4805,19 @@ def forward(self, state): second = self.second.detach().clone() return torch.zeros_like(state), first, second - def test_shared_mem_update_never_tears_running_policy(self): + @pytest.mark.parametrize( + "make_scheme", + [ + pytest.param( + functools.partial(SharedMemWeightSyncScheme, atomic=True), + id="SharedMemWeightSyncScheme-atomic", + ), + pytest.param( + MultiProcessWeightSyncScheme, id="MultiProcessWeightSyncScheme" + ), + ], + ) + def test_update_never_tears_running_policy(self, make_scheme): policy = TensorDictModule( TestUpdateParams.SlowTwoReadModule(), in_keys=["state"], @@ -4815,7 +4829,7 @@ def test_shared_mem_update_never_tears_running_policy(self): policy, frames_per_batch=4, total_frames=-1, - weight_sync_schemes={"policy": SharedMemWeightSyncScheme(atomic=True)}, + weight_sync_schemes={"policy": make_scheme()}, ) num_updates = 15 torn_calls = 0 diff --git a/torchrl/weight_update/_mp.py b/torchrl/weight_update/_mp.py index 63483d25e59..45574446ba0 100644 --- a/torchrl/weight_update/_mp.py +++ b/torchrl/weight_update/_mp.py @@ -1,11 +1,13 @@ from __future__ import annotations +import contextlib from collections.abc import Callable from typing import Any import torch from tensordict import TensorDictBase from torch import multiprocessing as mp, nn +from torchrl._utils import logger as torchrl_logger from torchrl.weight_update._shared import SharedMemWeightSyncScheme from torchrl.weight_update.utils import _resolve_model from torchrl.weight_update.weight_sync_schemes import ( @@ -14,6 +16,32 @@ ) +def _copy_weights( + weights: TensorDictBase | dict | None, share: bool = False +) -> TensorDictBase | dict | None: + """Return a copy of ``weights`` that shares no storage with them. + + With ``share=True``, CPU tensors of the copy are moved to shared memory, so + that pickling them through a queue to several workers neither moves the + original tensors to shared memory nor races to move the copy. + """ + if isinstance(weights, TensorDictBase): + copy = weights.detach().clone() + if share: + copy.share_memory_() + return copy + if isinstance(weights, dict): + copy = {} + for key, value in weights.items(): + if isinstance(value, torch.Tensor): + value = value.detach().clone() + if share: + value.share_memory_() + copy[key] = value + return copy + return weights + + @register_weight_sync_backend("process") class MultiProcessWeightSyncScheme(SharedMemWeightSyncScheme): """Weight synchronization for multiprocess operations using queues. @@ -22,9 +50,15 @@ class MultiProcessWeightSyncScheme(SharedMemWeightSyncScheme): Unlike the parent SharedMemWeightSyncScheme which uses shared memory for in-place updates, this scheme sends actual weight copies through queues to workers. + Each worker keeps its own copy of the weights: changes made to the sender's + model (for instance an optimizer step) only reach the workers through + ``send()``, which copies the weights at the time it is called. + A background thread on the receiver side listens for "receive" instructions from the sender. When an instruction arrives, the thread receives the weights - from the weight queue and applies them to the model. + from the weight queue and copies them into the model while holding the lock + that the receiving collector holds around each policy call, so a policy + call always reads a single, complete version of the weights. It follows the same two-phase pattern as SharedMemWeightSyncScheme: @@ -299,7 +333,9 @@ def send( - list[int]: Send to specific workers Note: If sync=True (default), this is a blocking call that ensures - specified workers are updated before returning. + specified workers are updated before returning. The weights are copied + before they are queued, so changes made to them after ``send()`` + returns never reach the workers. """ from torchrl.collectors.utils import _stage_unshareable_weights_on_cpu @@ -318,7 +354,9 @@ def send( strategy=self._strategy, context=context, ) - prepared_weights = _stage_unshareable_weights_on_cpu(prepared_weights) + prepared_weights = _copy_weights( + _stage_unshareable_weights_on_cpu(prepared_weights), share=True + ) transports = list(self._iterate_transports(worker_ids)) @@ -393,10 +431,14 @@ def _setup_connection_and_weights_on_sender_impl( # Send to workers sequentially via queues (no ACK - workers are still initializing) # This allows GC to clean up each worker's weights before creating the next + copies = {} for i, transport in enumerate(self._iterate_transports()): if worker_idx is not None and i != worker_idx: continue - worker_weights = params_map[i] + source = params_map[i] + if id(source) not in copies: + copies[id(source)] = _copy_weights(source, share=True) + worker_weights = copies[id(source)] if hasattr(transport, "send_weights_async"): transport.send_weights_async(worker_weights, model_id=self._model_id) else: @@ -446,23 +488,71 @@ def _setup_connection_and_weights_on_receiver_impl( # Apply weights to model if weights is not None and self.model is not None: - self._strategy.apply_weights(self.model, weights, inplace=False) + self._strategy.apply_weights( + self.model, _copy_weights(weights), inplace=False + ) # Start background receiver thread self._start_background_receiver() + def _weight_update_lock(self) -> contextlib.AbstractContextManager: + """Return the lock that the receiving context holds around policy calls. + + Weights are always copied into the worker's own parameters, so the + copy is made while holding this lock whatever the value of ``atomic``. + """ + return getattr(self.context, "_weight_update_lock", contextlib.nullcontext()) + + def _apply_received_weights(self, weights: Any) -> None: + """Copy received weights into the model and cascade the update. + + The cascade through the context's ``update_policy_weights_`` reads the + model's freshly copied parameters, so a context whose policy is the + model itself does not copy the weights a second time. + """ + model = self.model + if weights is not None and model is not None: + self._strategy.apply_weights(model, weights, inplace=True) + weights = self._strategy.extract_weights(model) + if weights is None: + return + model_id = self._model_id or "policy" + if self.context is not None and hasattr(self.context, "update_policy_weights_"): + self.context.update_policy_weights_( + model_id=model_id, policy_or_weights=weights + ) + + def receive(self, timeout: float | None = None) -> Any | None: + """Receive the next weights from the queue and apply them to the model. + + The queue is read without holding the weight-update lock, so a policy + call is only blocked while the weights are copied. + """ + if not self.initialized_on_receiver: + raise RuntimeError( + "Must be initialized on receiver before receiving weights" + ) + if not self.synchronized_on_receiver: + raise RuntimeError( + "Must be synchronized on receiver before receiving weights" + ) + if self._receiver_transport is None: + return None + weights = self._receiver_transport.receive_weights(timeout=timeout) + with self._weight_update_lock(): + self._apply_received_weights(weights) + return weights + def _background_receive_loop(self): """Background thread loop that waits for instructions and receives weights. This loop: 1. Waits for a "receive" instruction from the sender 2. Receives weights from the weight queue - 3. Applies them to the model + 3. Copies them into the model while holding the weight-update lock 4. Sends an acknowledgment back to the sender 5. Repeats until stop event is set or "stop" instruction received """ - from torchrl._utils import logger as torchrl_logger - while not self._stop_event.is_set(): try: instruction = self._wait_for_instruction() @@ -470,28 +560,10 @@ def _background_receive_loop(self): # Stop event was set or timeout continue if instruction == "receive": - # Receive weights from transport (blocking) if self._receiver_transport is not None: - weights = self._receiver_transport.receive_weights( - model=self.model, - strategy=self._strategy, - ) - - if weights is not None: - # Cascade weight update to sub-collectors if context supports it. - # When the context is a leaf Collector, its - # update_policy_weights_ also bumps the local - # PolicyVersion transform — so we don't need a - # separate increment_version() call here. - model_id = self._model_id or "policy" - if self.context is not None and hasattr( - self.context, "update_policy_weights_" - ): - self.context.update_policy_weights_( - model_id=model_id, policy_or_weights=weights - ) - - # Send acknowledgment + weights = self._receiver_transport.receive_weights() + with self._weight_update_lock(): + self._apply_received_weights(weights) self._send_ack("updated") elif instruction == "stop": From b68543af1cab9a2957334c5da7c694255e027dfe Mon Sep 17 00:00:00 2001 From: Achintya P Date: Thu, 5 Nov 2026 20:20:30 -0800 Subject: [PATCH 5/9] [BugFix] Sync non-policy models through weight sync schemes without hanging 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. --- test/test_collectors.py | 95 ++++++++++++++++++++++++++++++++ torchrl/collectors/_single.py | 47 ++++++++++++++-- torchrl/weight_update/_mp.py | 31 ++++++++--- torchrl/weight_update/_shared.py | 63 ++++++++++++++++++--- 4 files changed, 215 insertions(+), 21 deletions(-) diff --git a/test/test_collectors.py b/test/test_collectors.py index 02d00259005..09eb89d015e 100644 --- a/test/test_collectors.py +++ b/test/test_collectors.py @@ -117,6 +117,7 @@ SafeModule, TanhNormal, ) +from torchrl.objectives.value import GAE from torchrl.testing import ( CARTPOLE_VERSIONED, check_rollout_consistency_multikey_env, @@ -4847,6 +4848,100 @@ def test_update_never_tears_running_policy(self, make_scheme): assert torn_calls == 0 assert (data["first"] > 0).all() + @pytest.mark.parametrize( + "make_scheme", + [ + pytest.param( + functools.partial(SharedMemWeightSyncScheme, atomic=True), + id="SharedMemWeightSyncScheme-atomic", + ), + pytest.param( + MultiProcessWeightSyncScheme, id="MultiProcessWeightSyncScheme" + ), + ], + ) + def test_update_syncs_postproc_module(self, make_scheme): + """A module of the workers' postproc synced by its own scheme gets the trainer's weights.""" + env = ContinuousActionVecMockEnv() + obs_dim = env.observation_spec["observation"].shape[-1] + action_dim = env.action_spec.shape[-1] + policy = TensorDictModule( + nn.Linear(obs_dim, action_dim), in_keys=["observation"], out_keys=["action"] + ) + value_net = TensorDictModule( + nn.Linear(obs_dim, 1), in_keys=["observation"], out_keys=["state_value"] + ) + col = MultiSyncCollector( + [ContinuousActionVecMockEnv] * 2, + policy, + frames_per_batch=20, + total_frames=-1, + postproc=GAE(gamma=0.9, lmbda=0.9, value_network=value_net), + replay_buffer=ReplayBuffer(storage=LazyTensorStorage(100), shared=True), + weight_sync_schemes={ + "policy": make_scheme(), + "postproc.value_network": make_scheme(), + }, + ) + try: + next(iter(col)) + with torch.no_grad(): + policy.module.weight.fill_(0.25) + value_net.module.weight.fill_(0.5) + value_net.module.bias.fill_(-2.0) + start = time.monotonic() + col.update_policy_weights_( + weights_dict={"policy": None, "postproc.value_network": None} + ) + assert time.monotonic() - start < 30 + policy_weights = col.get_distant_attr("policy.module.weight") + value_weights = col.get_distant_attr("postproc.value_network.module.weight") + value_biases = col.get_distant_attr("postproc.value_network.module.bias") + finally: + col.shutdown() + for policy_weight, value_weight, value_bias in zip( + policy_weights, value_weights, value_biases + ): + torch.testing.assert_close( + policy_weight.detach(), torch.full((action_dim, obs_dim), 0.25) + ) + torch.testing.assert_close( + value_weight.detach(), torch.full((1, obs_dim), 0.5) + ) + torch.testing.assert_close(value_bias.detach(), torch.full((1,), -2.0)) + + def test_multiprocess_update_raises_when_worker_cannot_apply(self): + """A worker that fails to apply an update makes send() raise instead of hang.""" + env = ContinuousActionVecMockEnv() + obs_dim = env.observation_spec["observation"].shape[-1] + action_dim = env.action_spec.shape[-1] + policy = TensorDictModule( + nn.Linear(obs_dim, action_dim), in_keys=["observation"], out_keys=["action"] + ) + col = MultiSyncCollector( + [ContinuousActionVecMockEnv] * 2, + policy, + frames_per_batch=20, + total_frames=-1, + weight_sync_schemes={"policy": MultiProcessWeightSyncScheme()}, + ) + try: + wrong_shape = TensorDict.from_module(policy).apply( + lambda t: torch.zeros(*t.shape, 2) + ) + with pytest.raises(RuntimeError, match="failed on the receiver side"): + col.update_policy_weights_(wrong_shape) + with torch.no_grad(): + policy.module.weight.fill_(0.25) + col.update_policy_weights_() + policy_weights = col.get_distant_attr("policy.module.weight") + finally: + col.shutdown() + for policy_weight in policy_weights: + torch.testing.assert_close( + policy_weight.detach(), torch.full((action_dim, obs_dim), 0.25) + ) + @pytest.mark.parametrize("atomic", [None, False, True]) def test_shared_mem_atomic_default_warns(self, atomic): scheme = SharedMemWeightSyncScheme(atomic=atomic) diff --git a/torchrl/collectors/_single.py b/torchrl/collectors/_single.py index 53695d432f6..faaa65edd61 100644 --- a/torchrl/collectors/_single.py +++ b/torchrl/collectors/_single.py @@ -60,7 +60,7 @@ from torchrl.modules.inference_server._config import _resolve_device_config from torchrl.modules.utils.utils import _maybe_append_env_transforms_from_module from torchrl.weight_update.utils import _resolve_model -from torchrl.weight_update.weight_sync_schemes import WeightSyncScheme +from torchrl.weight_update.weight_sync_schemes import WeightStrategy, WeightSyncScheme def _cuda_sync_if_initialized(): @@ -1712,6 +1712,8 @@ def update_policy_weights_( super().update_policy_weights_( policy_or_weights=policy_or_weights, worker_ids=worker_ids, **kwargs ) + if kwargs.get("model_id") not in (None, "policy"): + return mark_weight_update(getattr(self, "_policy_w_state_dict", None)) # Bump the local PolicyVersion transform (if track_policy_version is on). @@ -1730,11 +1732,16 @@ def _maybe_fallback_update( *, model_id: str | None = None, ) -> None: - """Copy weights from original policy to internal policy when no scheme configured.""" + """Copy weights from original policy to internal policy when no scheme configured. + + A ``model_id`` other than ``"policy"`` names another module of the + collector by its attribute path (for instance + ``"postproc.adv_module"``), whose parameters are overwritten in place + with ``policy_or_weights``. + """ if model_id is not None and model_id != "policy": - raise KeyError( - f"Collector has no local weight target for model_id={model_id!r}." - ) + self._update_local_model(model_id, policy_or_weights) + return # Get source weights - either from argument or from original policy if policy_or_weights is not None: @@ -1756,6 +1763,36 @@ def _maybe_fallback_update( return raise RuntimeError("Collector has no mutable local policy weight target.") + def _update_local_model( + self, + model_id: str, + policy_or_weights: TensorDictBase | nn.Module | dict | None, + ) -> None: + """Copy ``policy_or_weights`` into the module found at ``model_id``. + + A multiprocess worker reaches this when a weight sync scheme registered + for a module other than the policy cascades its update to the worker's + collector. + """ + try: + target = _resolve_model(self, model_id) + except (AttributeError, KeyError, IndexError): + target = None + if not isinstance(target, nn.Module): + raise KeyError( + f"Collector has no local weight target for model_id={model_id!r}." + ) + if policy_or_weights is None: + raise RuntimeError( + f"Collector cannot update model_id={model_id!r} without explicit " + "weights." + ) + strategy = WeightStrategy(extract_as="tensordict") + if isinstance(policy_or_weights, nn.Module): + policy_or_weights = strategy.extract_weights(policy_or_weights) + strategy.apply_weights(target, policy_or_weights, inplace=True) + mark_weight_update(target) + def set_seed(self, seed: int, static_seed: bool = False) -> int: """Sets the seeds of the environments stored in the DataCollector. diff --git a/torchrl/weight_update/_mp.py b/torchrl/weight_update/_mp.py index 45574446ba0..7c3fb6926b8 100644 --- a/torchrl/weight_update/_mp.py +++ b/torchrl/weight_update/_mp.py @@ -8,7 +8,11 @@ from tensordict import TensorDictBase from torch import multiprocessing as mp, nn from torchrl._utils import logger as torchrl_logger -from torchrl.weight_update._shared import SharedMemWeightSyncScheme +from torchrl.weight_update._shared import ( + _ACK_UPDATED, + _failed_ack, + SharedMemWeightSyncScheme, +) from torchrl.weight_update.utils import _resolve_model from torchrl.weight_update.weight_sync_schemes import ( register_weight_sync_backend, @@ -58,7 +62,9 @@ class MultiProcessWeightSyncScheme(SharedMemWeightSyncScheme): from the sender. When an instruction arrives, the thread receives the weights from the weight queue and copies them into the model while holding the lock that the receiving collector holds around each policy call, so a policy - call always reads a single, complete version of the weights. + call always reads a single, complete version of the weights. If the weights + cannot be applied, the worker reports it and ``send()`` raises a + ``RuntimeError`` instead of waiting for an acknowledgment that never comes. It follows the same two-phase pattern as SharedMemWeightSyncScheme: @@ -550,7 +556,8 @@ def _background_receive_loop(self): 1. Waits for a "receive" instruction from the sender 2. Receives weights from the weight queue 3. Copies them into the model while holding the weight-update lock - 4. Sends an acknowledgment back to the sender + 4. Sends an acknowledgment back to the sender, which reports a failure + if the weights could not be received or applied 5. Repeats until stop event is set or "stop" instruction received """ while not self._stop_event.is_set(): @@ -560,11 +567,19 @@ def _background_receive_loop(self): # Stop event was set or timeout continue if instruction == "receive": - if self._receiver_transport is not None: - weights = self._receiver_transport.receive_weights() - with self._weight_update_lock(): - self._apply_received_weights(weights) - self._send_ack("updated") + try: + if self._receiver_transport is not None: + weights = self._receiver_transport.receive_weights() + with self._weight_update_lock(): + self._apply_received_weights(weights) + except Exception as error: + torchrl_logger.exception( + "MultiProcessWeightSyncScheme: could not apply the " + f"weights of model_id={self._model_id!r}." + ) + self._send_ack(_failed_ack(error)) + else: + self._send_ack(_ACK_UPDATED) elif instruction == "stop": break diff --git a/torchrl/weight_update/_shared.py b/torchrl/weight_update/_shared.py index ca6514274fc..79592738fe9 100644 --- a/torchrl/weight_update/_shared.py +++ b/torchrl/weight_update/_shared.py @@ -54,6 +54,28 @@ def _synchronize_cuda(weights: TensorDictBase) -> None: torch.cuda.synchronize(device) +_ACK_UPDATED = "updated" +_ACK_FAILED = "failed" + + +def _failed_ack(error: BaseException) -> str: + """Return the acknowledgment a receiver sends when it could not apply weights.""" + return f"{_ACK_FAILED}: {type(error).__name__}: {error}" + + +def _raise_on_failed_acks(model_id: str | None, failures: dict[int, str]) -> None: + """Raise if any worker acknowledged a weight update it could not apply.""" + if not failures: + return + details = "; ".join( + f"worker {worker_idx}: {ack}" for worker_idx, ack in sorted(failures.items()) + ) + raise RuntimeError( + f"Weight update for model_id={model_id!r} failed on the receiver side " + f"({details})." + ) + + def _log_weight_sync( *, model_id: str | None, @@ -915,6 +937,10 @@ def _wait_for_ack( Args: worker_ids: Which workers to wait for (None = all workers). timeout: Maximum time to wait (seconds). None means block indefinitely. + + Raises: + RuntimeError: If a worker reports that it could not apply the + weights, once every targeted worker has answered. """ if not self._ack_queues: return # No ack queues, nothing to wait for @@ -926,19 +952,24 @@ def _wait_for_ack( else: target_workers = list(worker_ids) + failures = {} for worker_idx in target_workers: if worker_idx not in self._ack_queues: raise ValueError(f"Worker {worker_idx} not registered") try: ack = self._ack_queues[worker_idx].get(timeout=timeout) - if ack != "updated": - torchrl_logger.warning( - f"Unexpected ack from worker {worker_idx}: {ack}" - ) except Exception as e: torchrl_logger.warning( f"Timeout waiting for ack from worker {worker_idx}: {e}" ) + continue + if isinstance(ack, str) and ack.startswith(_ACK_FAILED): + failures[worker_idx] = ack + elif ack != _ACK_UPDATED: + torchrl_logger.warning( + f"Unexpected ack from worker {worker_idx}: {ack}" + ) + _raise_on_failed_acks(self._model_id, failures) def create_transport(self, **kwargs) -> TransportBackend: """Create shared memory transport. @@ -1105,12 +1136,17 @@ def _wait_for_unacknowledged_workers(self) -> None: sends only collect whatever it has acknowledged in the meantime, without waiting, until it catches up. A dead worker therefore does not block the sender. + + Raises: + RuntimeError: If a worker reports that it could not apply an + update, once every pending acknowledgment has been collected. """ + failures = {} for worker_idx, pending in list(self._pending_acks.items()): stalled = worker_idx in self._stalled_workers while pending: try: - self._ack_queues[worker_idx].get( + ack = self._ack_queues[worker_idx].get( timeout=0 if stalled else WEIGHT_SYNC_TIMEOUT ) except Empty: @@ -1123,11 +1159,14 @@ def _wait_for_unacknowledged_workers(self) -> None: self._stalled_workers.add(worker_idx) break pending -= 1 + if isinstance(ack, str) and ack.startswith(_ACK_FAILED): + failures[worker_idx] = ack if pending: self._pending_acks[worker_idx] = pending else: del self._pending_acks[worker_idx] self._stalled_workers.discard(worker_idx) + _raise_on_failed_acks(self._model_id, failures) @property def weights(self) -> Any | None: @@ -1220,9 +1259,17 @@ def _background_receive_loop(self): # Stop event was set or timeout continue if instruction == "receive": - with self._weight_update_lock(): - self._apply_shared_weights() - self._send_ack("updated") + try: + with self._weight_update_lock(): + self._apply_shared_weights() + except Exception as error: + torchrl_logger.exception( + "SharedMemWeightSyncScheme: could not apply the weights " + f"of model_id={self._model_id!r}." + ) + self._send_ack(_failed_ack(error)) + else: + self._send_ack(_ACK_UPDATED) elif instruction == "stop": break else: From 1a92ed3a88d449d68e97263f11b71b011bebcfca Mon Sep 17 00:00:00 2001 From: Achintya P Date: Thu, 5 Nov 2026 20:20:45 -0800 Subject: [PATCH 6/9] [BugFix] Lay out multi-collector replay storages with the workers' time 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. --- examples/satellite/sac_per_fully_async.py | 4 ++++ test/test_collectors.py | 27 +++++++++++++++++++++++ torchrl/collectors/_multi_base.py | 16 ++++++++++++-- 3 files changed, 45 insertions(+), 2 deletions(-) diff --git a/examples/satellite/sac_per_fully_async.py b/examples/satellite/sac_per_fully_async.py index 50e8f4cf1de..b33bc15ad3e 100644 --- a/examples/satellite/sac_per_fully_async.py +++ b/examples/satellite/sac_per_fully_async.py @@ -672,6 +672,10 @@ def _on_eval_result(result): collector.async_shutdown(timeout=30.0) except Exception as e: torchrl_logger.warning(f"collector.async_shutdown failed: {e}") + if any(proc.exitcode is not None and proc.exitcode > 0 for proc in collector.procs): + raise RuntimeError( + "The collector worker exited with an error; its traceback is in the log above." + ) if evaluator is not None: if evaluator.pending: try: diff --git a/test/test_collectors.py b/test/test_collectors.py index 09eb89d015e..83dfa604537 100644 --- a/test/test_collectors.py +++ b/test/test_collectors.py @@ -7137,6 +7137,33 @@ def test_collector_rb_multiasync( ).all(), steps_counts assert (idsdiff >= 0).all() + def test_multi_collector_rb_batched_env_writes_time_rows(self): + """A batched env writes rows of frames_per_batch // num_envs steps into a 1-D storage.""" + make_env = functools.partial(CountingEnv, max_steps=100, batch_size=[2]) + probe = make_env() + try: + policy = CountingEnvCountPolicy(probe.action_spec) + finally: + probe.close(raise_if_closed=False) + rb = ReplayBuffer(storage=LazyTensorStorage(16), shared=True) + collector = MultiSyncCollector( + [make_env], + policy, + replay_buffer=rb, + frames_per_batch=8, + total_frames=16, + ) + try: + for data in collector: + assert data is None + finally: + collector.shutdown() + observations = rb[: len(rb)]["observation"].squeeze(-1) + expected = torch.tensor( + [[0, 1, 2, 3], [0, 1, 2, 3], [4, 5, 6, 7], [4, 5, 6, 7]] + ) + torch.testing.assert_close(observations, expected.to(observations.dtype)) + def test_multi_async_replay_buffer_device_metadata(self): probe = CountingEnv() try: diff --git a/torchrl/collectors/_multi_base.py b/torchrl/collectors/_multi_base.py index 088a054eeef..a9931754968 100644 --- a/torchrl/collectors/_multi_base.py +++ b/torchrl/collectors/_multi_base.py @@ -518,6 +518,7 @@ def __init__( create_env_fn, total_frames_per_batch = self._setup_workers_and_env_fns( create_env_fn, num_workers, frames_per_batch ) + self.requested_frames_per_batch = total_frames_per_batch # Set up basic configuration self.set_truncated = set_truncated @@ -604,7 +605,6 @@ def __init__( ) # Set up split trajectories - self.requested_frames_per_batch = total_frames_per_batch self.reset_when_done = reset_when_done self._setup_split_trajs(split_trajs, reset_when_done) @@ -1071,6 +1071,13 @@ def _validate_cat_results(self, cat_results: str | int | None) -> None: self.cat_results = cat_results def _check_replay_buffer_init(self): + """Initialise an empty replay-buffer storage before the workers write to it. + + The storage is laid out from a fake batch of the same shape as the + batches the workers write, ``(*env_batch_size, time)`` with ``time`` + the number of steps each worker collects per env, so that a storage + with fewer dimensions than that batch stores rows of the right length. + """ if self.replay_buffer is None: return @@ -1154,8 +1161,13 @@ def _check_replay_buffer_init(self): fake_td["collector", "traj_ids"] = torch.zeros( fake_td.shape, dtype=torch.long ) + time_steps = -( + -self.frames_per_batch_worker(worker_idx=0) + // max(fake_td.numel(), 1) + ) + fake_td = fake_td.unsqueeze(-1).expand(*fake_td.shape, time_steps) # Use extend to avoid time-related transforms to fail - self.replay_buffer.extend(fake_td.unsqueeze(-1)) + self.replay_buffer.extend(fake_td) self.replay_buffer.empty() def _should_init_replay_buffer_from_worker(self, storage): From 7bfffd7e8d5b56e198c90f714b9a30d4bade9092 Mon Sep 17 00:00:00 2001 From: Achintya P Date: Fri, 6 Nov 2026 01:11:32 -0800 Subject: [PATCH 7/9] Keep the AsyncEnvPool env_index stack off the step hot path 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. --- benchmarks/test_envs_benchmark.py | 23 +++++++++++++++++++++++ torchrl/envs/async_envs.py | 15 ++++++++++++++- 2 files changed, 37 insertions(+), 1 deletion(-) diff --git a/benchmarks/test_envs_benchmark.py b/benchmarks/test_envs_benchmark.py index 96f6ee259bc..0c15d680564 100644 --- a/benchmarks/test_envs_benchmark.py +++ b/benchmarks/test_envs_benchmark.py @@ -428,6 +428,29 @@ def test_async_env_pool_per_env_dispatch(benchmark, exchange): pool._maybe_shutdown() +@pytest.mark.parametrize("exchange", ["queue", "shm"]) +def test_async_env_pool_rollout(benchmark, exchange): + """Full-batch step cost of AsyncEnvPool through env.rollout with free envs.""" + with set_capture_non_tensor_stack(False): + pool = AsyncEnvPool( + [ContinuousActionVecMockEnv] * ASYNC_POOL_DISPATCH_ENVS, + backend="multiprocessing", + exchange=exchange, + ) + try: + benchmark.extra_info["num_envs"] = ASYNC_POOL_DISPATCH_ENVS + benchmark.pedantic( + pool.rollout, + args=(100,), + kwargs={"break_when_any_done": False}, + rounds=5, + warmup_rounds=1, + iterations=1, + ) + finally: + pool._maybe_shutdown() + + # Note: the queue exchange is deliberately not parametrized here. At 32 # workers the queue path hangs during setup/reset with the spawn start method # (macOS-confirmed; tracked for a fix in the AsyncEnvPool v2 sequence). The diff --git a/torchrl/envs/async_envs.py b/torchrl/envs/async_envs.py index 9c2c4c702ec..76a0ec3d38f 100644 --- a/torchrl/envs/async_envs.py +++ b/torchrl/envs/async_envs.py @@ -556,8 +556,19 @@ def _set_env_index(self, tensordict: TensorDictBase) -> None: data written by :meth:`reset`, :meth:`step` and :meth:`step_and_maybe_reset` all share one type and can be stacked together (e.g. by a collector writing into a preallocated rollout). + + The stack is built fresh on each call, so the caller owns the entry + and may modify it, with its items created directly on the + tensordict's device: moving a stack to a device while ``set`` + validates it costs more than the step of a light environment. """ - env_idx_nt = NonTensorStack(*range(tensordict.shape[0])) + device = tensordict.device + env_idx_nt = NonTensorStack( + *( + NonTensorData(index, device=device) + for index in range(tensordict.shape[0]) + ) + ) if env_idx_nt.shape != tensordict.shape: env_idx_nt = expand_as_right(env_idx_nt, tensordict) tensordict.set(self._env_idx_key, env_idx_nt) @@ -629,6 +640,8 @@ def _maybe_make_tensordict(self, tensordict, env_index, make_if_none): ) tensordict[self._env_idx_key] = NonTensorStack(*env_index) env_idx = env_index + if tensordict is not None and self._env_idx_key in tensordict.keys(): + tensordict = tensordict.exclude(self._env_idx_key) return tensordict, env_idx @abc.abstractmethod From ce2db24f5dca9e1ef0cfd52f2d45b02feacda21e Mon Sep 17 00:00:00 2001 From: Achintya P Date: Fri, 6 Nov 2026 16:04:30 -0800 Subject: [PATCH 8/9] Send MultiProcessWeightSyncScheme weights in one shared buffer 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. --- benchmarks/test_collectors_benchmark.py | 32 +++++++++++++++++++++++++ torchrl/weight_update/_mp.py | 16 ++++++++++++- 2 files changed, 47 insertions(+), 1 deletion(-) diff --git a/benchmarks/test_collectors_benchmark.py b/benchmarks/test_collectors_benchmark.py index fed37604d4e..6390c87d7c8 100644 --- a/benchmarks/test_collectors_benchmark.py +++ b/benchmarks/test_collectors_benchmark.py @@ -49,6 +49,7 @@ ContinuousActionVecMockEnv, MockBatchedLockedEnv, ) +from torchrl.weight_update import MultiProcessWeightSyncScheme class _ResetLatencyPixelEnv(PixelMockEnv): @@ -681,6 +682,37 @@ def test_async_weight_update(benchmark): collector.shutdown() +def test_multiprocess_weight_update(benchmark): + """Benchmark pushing a 130-tensor policy through ``MultiProcessWeightSyncScheme``. + + Every send copies the weights into fresh shared memory and each worker + opens that memory, so the cost depends on the number of tensors as well as + on their size. + """ + policy = TensorDictModule( + MLP(in_features=7, out_features=7, num_cells=[64] * 64), + in_keys=["observation"], + out_keys=["action"], + ) + collector = MultiSyncCollector( + [ContinuousActionVecMockEnv, ContinuousActionVecMockEnv], + policy, + total_frames=-1, + frames_per_batch=20, + weight_sync_schemes={"policy": MultiProcessWeightSyncScheme()}, + ) + try: + next(iter(collector)) + benchmark.pedantic( + collector.update_policy_weights_, + iterations=10, + rounds=5, + warmup_rounds=1, + ) + finally: + collector.shutdown() + + @pytest.mark.parametrize("payload_size", [65536, 262144], ids=["256KiB", "1MiB"]) def test_parallel_env_compact_obs(benchmark, payload_size): """Benchmark copies for a collector over a large-payload ParallelEnv.""" diff --git a/torchrl/weight_update/_mp.py b/torchrl/weight_update/_mp.py index 7c3fb6926b8..9f3fd58f764 100644 --- a/torchrl/weight_update/_mp.py +++ b/torchrl/weight_update/_mp.py @@ -20,6 +20,15 @@ ) +def _all_cpu_tensors(weights: TensorDictBase) -> bool: + found = False + for leaf in weights.values(include_nested=True, leaves_only=True): + if not isinstance(leaf, torch.Tensor) or leaf.device.type != "cpu": + return False + found = True + return found + + def _copy_weights( weights: TensorDictBase | dict | None, share: bool = False ) -> TensorDictBase | dict | None: @@ -27,9 +36,14 @@ def _copy_weights( With ``share=True``, CPU tensors of the copy are moved to shared memory, so that pickling them through a queue to several workers neither moves the - original tensors to shared memory nor races to move the copy. + original tensors to shared memory nor races to move the copy. A tensordict + whose leaves are all CPU tensors is copied into a single consolidated + shared-memory buffer, so that each send creates and each worker opens one + shared segment rather than one per tensor. """ if isinstance(weights, TensorDictBase): + if share and _all_cpu_tensors(weights): + return weights.detach().consolidate(share_memory=True) copy = weights.detach().clone() if share: copy.share_memory_() From db67e25883922535ee0b9f2eb264de6da90f6075 Mon Sep 17 00:00:00 2001 From: Achintya P Date: Fri, 6 Nov 2026 16:04:33 -0800 Subject: [PATCH 9/9] Keep extra next entries out of the shm step_and_maybe_reset next root 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. --- test/envs/test_special.py | 28 ++++++++++++++++++++++++++++ torchrl/envs/_async_exchange.py | 7 +++++-- 2 files changed, 33 insertions(+), 2 deletions(-) diff --git a/test/envs/test_special.py b/test/envs/test_special.py index d6aa8c04485..b5a6e55f5b3 100644 --- a/test/envs/test_special.py +++ b/test/envs/test_special.py @@ -1234,6 +1234,34 @@ def test_outer_transforms_match_serial_env(self, exchange): pool_env.close(raise_if_closed=False) serial_env.close(raise_if_closed=False) + @pytest.mark.parametrize("exchange", ["queue", "shm"]) + @set_capture_non_tensor_stack(False) + def test_step_and_maybe_reset_extra_inputs_match_serial_env(self, exchange): + """Entries outside the env specs come back where a SerialEnv puts them. + + The root output keeps the extra root and ``"next"`` entries of the + input, and the next root keeps only the extra root entries. + """ + makers = [CountingEnv, CountingEnv] + pool = AsyncEnvPool(makers, backend="multiprocessing", exchange=exchange) + serial = SerialEnv(2, makers) + try: + outputs = [] + for env in (pool, serial): + td = env.reset() + td["action"] = env.full_action_spec["action"].one() + td["tag"] = torch.tensor([3.0, 4.0]) + td["next", "tag"] = torch.tensor([7.0, 8.0]) + outputs.append(env.step_and_maybe_reset(td)) + (pool_root, pool_next), (serial_root, serial_next) = outputs + assert_allclose_td(pool_root.exclude("env_index"), serial_root) + assert_allclose_td(pool_next.exclude("env_index"), serial_next) + assert "next" not in serial_next.keys() + torch.testing.assert_close(serial_next["tag"], torch.tensor([3.0, 4.0])) + finally: + pool.close(raise_if_closed=False) + serial.close(raise_if_closed=False) + @set_capture_non_tensor_stack(False) def test_shared_memory_exchange(self, make_envs): env = AsyncEnvPool(make_envs, backend="multiprocessing", exchange="shm") diff --git a/torchrl/envs/_async_exchange.py b/torchrl/envs/_async_exchange.py index be61be3c0b7..e12cb469774 100644 --- a/torchrl/envs/_async_exchange.py +++ b/torchrl/envs/_async_exchange.py @@ -243,14 +243,17 @@ def _with_extra_inputs( """Merges the extra inputs kept for ``env_index`` into ``results``. Each result gets its own copy so that the root and ``"next"`` outputs - of ``step_and_maybe_reset`` never alias each other. + of ``step_and_maybe_reset`` never alias each other. The next root of + ``step_and_maybe_reset`` does not receive extra ``"next"`` entries, + as :func:`~torchrl.envs.utils.step_mdp` drops them with the queue + exchange. """ extra = self._extra_inputs.pop(env_index, None) if extra is None: return results merged = [results[0].update(extra)] for result in results[1:]: - merged.append(result.update(extra.clone())) + merged.append(result.update(extra.exclude("next").clone())) return tuple(merged) @staticmethod