diff --git a/docs/sphinx/source/en/2-user_guide/1-training/0-index.md b/docs/sphinx/source/en/2-user_guide/1-training/0-index.md index 7fd29e864..7c6a248fd 100644 --- a/docs/sphinx/source/en/2-user_guide/1-training/0-index.md +++ b/docs/sphinx/source/en/2-user_guide/1-training/0-index.md @@ -50,6 +50,12 @@ Run UniLab inside the checked-in Linux NVIDIA image workflow. Reproduce and operate the scoped single-GPU G1 FlashSAC/MJWarp long soak. ::: +:::{grid-item-card} Sync-free autoreset +:link: 8-sync_free_autoreset +:link-type: doc +Opt-in fixed-shape autoreset, reset CUDA graph, and `ResetSelection` terms. +::: + :::: ## When to Drop to `scripts/train_*.py` @@ -78,4 +84,5 @@ Day-to-day runs should use the unified CLI. Reach for the low-level 5-resume_and_checkpoints 6-docker 7-tensor_runtime_production +8-sync_free_autoreset ``` diff --git a/docs/sphinx/source/en/2-user_guide/1-training/8-sync_free_autoreset.md b/docs/sphinx/source/en/2-user_guide/1-training/8-sync_free_autoreset.md new file mode 100644 index 000000000..6f8948a74 --- /dev/null +++ b/docs/sphinx/source/en/2-user_guide/1-training/8-sync_free_autoreset.md @@ -0,0 +1,184 @@ +# Sync-Free Autoreset and Reset CUDA Graph + +By default, `TorchEnv` resets done rows through row indices: it reads the done +mask on the host (`any()`), builds `env_ids` with `nonzero`, and launches reset +kernels whose shape depends on the number of reset rows. Each of these steps +synchronizes the host with the GPU. + +An opt-in fixed-shape path removes these synchronizations, and on CUDA it can +replay the whole reset segment as one CUDA graph. Every switch defaults to +`False`, so tasks that do not opt in keep the selected-row behavior unchanged. + +Implementation: `src/unilab/base/torch_env.py`, +`src/unilab/base/reset_selection.py`, `src/unilab/base/autoreset_graph.py`, +`src/unilab/base/host_transitions.py`. Tests: +`tests/envs/test_fixed_shape_autoreset.py`, +`tests/tasks/test_g1_motion_fixed_shape.py`. + +## Requirements + +- A Manager-based task whose every reset participant accepts a + `ResetSelection` (see "Opting terms into `ResetSelection`" below). + The G1 motion-tracking terms (`MotionCommand`, `JointPositionAction`, the + motion body rewards and the anchor/body height terminations) already opt in. +- `training.defer_finite_validation=true`, the default in every shipped + algorithm config. +- For the CUDA graph: a CUDA env device, and a backend that declares the + `external_graph_capture` capability and implements + `torch_graph_capture_scope`. MJWarp does both, from the UniSim release that + ships `unisim.backend.graph_capture`. + +## Configuration flags + +All flags are fields of `ManagerBasedRlEnvCfg` (or of the base `EnvCfg` for the +backend flag). They live in the `env` section of the owner YAML. They are not +present in the shipped owner YAMLs, so set them with `++` from the command line: + +```bash +uv run train --algo flashsac --task g1_motion_tracking --sim mjwarp \ + ++env.masked_autoreset=true \ + ++env.persistent_step_buffers=true \ + ++env.autoreset_cuda_graph=true \ + ++env.mjwarp_defer_tensor_validation=true +``` + +| Flag | Default | Effect | +| --- | --- | --- | +| `masked_autoreset` | `False` | Runs the autoreset on every step as a full-width `reset(ResetSelection)` under the device done mask, with no `any()` gate and no `nonzero`. Requires `defer_finite_validation`. | +| `persistent_step_buffers` | `False` | Allocates the observation, reward, terminated and truncated carriers once and updates them with `copy_`. Consumers receive one of `published_buffer_depth` persistent publication slots. | +| `published_buffer_depth` | `2` | Size of the publication ring. A published state stays valid for `published_buffer_depth - 1` further steps. `0` publishes the working carriers themselves, which the next step overwrites. | +| `autoreset_cuda_graph` | `False` | Captures the masked autoreset segment, including the backend reset, into one `torch.cuda.graph` and replays it afterwards. Requires `masked_autoreset`, `persistent_step_buffers`, a CUDA device and a backend with external graph capture. | +| `autoreset_graph_warmup` | `2` | Eager executions before capture, so that lazily allocated buffers exist outside the graph. | +| `autoreset_graph_capture_mode` | `"thread_local"` | `torch.cuda.graph` capture error mode: `"global"`, `"thread_local"` or `"relaxed"`. | +| `mjwarp_defer_tensor_validation` | `False` | Latches MJWarp's fail-closed operand and reset-mask checks into device verdicts instead of reading each one on the host. The env publishes them at its deferred-validation poll. Inside a graph capture MJWarp always defers. The flag also removes the host reads from the eager steps outside the captured segment. UniLab forwards it only when it is `True`. | + +Invalid combinations fail when the config is validated: for example, +`autoreset_cuda_graph` without `masked_autoreset` and `persistent_step_buffers`, +or `masked_autoreset` without `defer_finite_validation`. + +## Behavior notes + +**RNG draw order.** With `masked_autoreset` the reset lifecycle runs on every +step; an all-false mask leaves the state unchanged. Full-width samplers, such as +the `MotionCommand` resample or the reset-event interval times, therefore draw +from the Manager RNG on steps where no row is done. With the same seed, +trajectories differ from the selected-row path. Compare the two modes +statistically, not value by value; MJWarp is not bitwise deterministic in any +case. + +**`final_observation`.** It is published on every step and is meaningful only +on done rows. + +**Deferred validation polling.** Every fail-closed check in the reset segment is +latched on the device. This covers finite checks, `TorchEnv.require` verdicts, +reset-mask containment and MJWarp/packed-reset verdicts. +`poll_finite_validation` reads all of them together every +`training.finite_validation_interval` steps (32 in the shipped configs). A +failure is therefore raised at the next poll, up to `interval - 1` steps after +the faulty step. Packed-reset failures name the poll step +(`first detected at poll boundary, step N`). To localize a fault, rerun with +`training.finite_validation_interval=1`; the masked path still requires +`training.defer_finite_validation=true`. + +**Stale diagnostic counters.** In masked mode the env has no host-side reset row +count or per-phase split. Consequences: + +- The per-reset detail keys in `info["timing"]` + (`RESET_DONE_DETAIL_TIMING_KEYS` in `src/unilab/base/backend_timing.py`, for + example `reset_done_count`) are written as `0` on every step. Only + `reset_done_reset_call_ms` is filled. A `0` there does not mean that no row + reset. +- Episode logs (`Episode_Termination/*`, `Episode_Metrics/*` and command reset + metrics) accumulate on the device. They are published into `info["log"]` + only at the deferred-validation poll, and each value covers every episode + that completed since the previous poll. Read them as interval aggregates, + not per-step values. + +**Host state under graph replay.** A CUDA graph replays device work only. Python +state that the reset segment changes, such as a pending-metrics flag or a host +lower bound, must be declared with +`unilab.base.host_transitions.record_host_transition`. The graph owner +re-applies these transitions after every replay. Undeclared host state keeps the +value it had at capture time. + +## Opting terms into `ResetSelection` + +A `ResetSelection` (`src/unilab/base/reset_selection.py`) is a persistent +full-width bool `mask` plus `row_ids = arange(num_envs)`. It is passed through +the existing `env_ids` reset arguments. A participant opts in by declaring +support and handling the selection as a masked full-width write. + +- **Class-based terms** (commands, actions, class reward/termination/metric + terms, class observation terms, stateful event terms) set + `supports_reset_selection = True` and accept a `ResetSelection` in `reset`. +- **Function-style terms** (for example reset-mode event functions) use the + `@reset_selection_capable` decorator, and their body must handle a + `ResetSelection` argument. +- Write reset state with the helpers instead of indexing: + - `fill_rows_(tensor, env_ids, value)` writes a constant; + - `assign_rows_(tensor, env_ids, values)` writes a full-width tensor, or + calls a factory `values(n)` that returns `n` rows (`num_envs` rows under a + selection). + + Both accept row ids, `None`, a slice or a selection. +- Do not read the mask on the host (`bool(mask.any())`, `.item()`, + `nonzero`) inside `reset`; that reintroduces the synchronization. + +```python +from typing import ClassVar + +import torch + +from unilab.base.reset_selection import ResetSelection, fill_rows_ +from unilab.envs.mdp.actions.actions import BaseAction + + +class MyAction(BaseAction): + # Re-declare on every subclass that overrides reset(). + supports_reset_selection: ClassVar[bool] = True + + def reset(self, env_ids: torch.Tensor | ResetSelection | None = None) -> None: + super().reset(env_ids) + fill_rows_(self._my_state, env_ids, 0.0) +``` + +A declaration counts only at or below the class that defines the effective +`reset`. A subclass that overrides `reset` does not inherit its parent's opt-in +and is reported as a blocker until it declares support itself. A +`ResetSelection` is not a tensor, so unconverted code that receives one by +mistake fails at the first indexing operation instead of resetting wrong rows. + +Some managers impose extra rules: + +- `CurriculumManager` terms cannot run in a sync-free reset, because they + publish host curriculum state. +- The reward, termination, observation and metrics managers check only their + class-based terms. Stateless function terms need no declaration. + +### Debugging `reset_selection_blockers` + +With `masked_autoreset`, `init_state` calls +`ManagerBasedRlEnv.reset_selection_blockers()`. It collects one entry per +unsupported participant from every manager and raises: + +```text +NotImplementedError: masked_autoreset needs every reset participant to accept a +ResetSelection; unsupported: EventManager: (); ... +``` + +Each entry has the form `: ()`. A +manager that cannot take a selection for a structural reason reports a sentence +instead. One example is `TerminationManager: episode counts are read on the +host unless defer_finite_validation is enabled`. + +To list blockers without starting a training run, construct the env with the +flag off and call the method: + +```python +env.reset_selection_blockers() # [] when the task can enable masked_autoreset +``` + +Then convert, replace or remove each named term. Reset-randomization events are +common blockers; the G1 motion-tracking FlashSAC owner +(`src/unilab/conf/flashsac/task/g1_motion_tracking/base.yaml`) is a DR-free +baseline for this reason. diff --git a/src/unilab/base/autoreset_graph.py b/src/unilab/base/autoreset_graph.py new file mode 100644 index 000000000..c381bca9f --- /dev/null +++ b/src/unilab/base/autoreset_graph.py @@ -0,0 +1,132 @@ +"""CUDA graph capture of the fixed-shape autoreset segment. + +``TorchEnv`` with ``masked_autoreset`` and ``persistent_step_buffers`` runs its +autoreset as one fixed-shape body (``TorchEnv._masked_autoreset_body``): every +tensor it touches has a fixed address, nothing in it reads the device from the +host, and host bookkeeping is declared through ``record_host_transition``. +:class:`AutoresetGraph` captures that body once and replays it afterwards. + +Capture protocol (all public, no object-graph inspection): + +1. ``warmup`` eager executions allocate every lazily created buffer, scratch + and verdict flag outside the graph. +2. The body runs inside ``torch.cuda.graph`` on a side stream. Within it, the + backend's ``torch_graph_capture_scope`` routes simulator kernels onto the + capture stream and returns an ``ExternalGraphCapture`` handle, which the env + binds so that env/manager host transitions record on the same handle as the + backend's own. +3. Capture records device work without executing it, so the graph is replayed + once immediately; the host transitions already ran during capture. +4. Every later call checks the backend's replay guards, replays the graph and + re-applies the recorded host transitions in recording order. +""" + +from __future__ import annotations + +from collections.abc import Callable, Sequence +from typing import Any + +import torch +from unisim.backend.graph_capture import ExternalGraphCapture, GraphCaptureError + + +class AutoresetGraph: + """Capture-once, replay-always runner for a fixed-shape autoreset body.""" + + def __init__( + self, + *, + backend: Any, + device: torch.device, + body: Callable[[], None], + bind_capture: Callable[[ExternalGraphCapture | None], None], + generators: Sequence[torch.Generator] = (), + warmup: int = 2, + capture_error_mode: str = "thread_local", + ) -> None: + if device.type != "cuda": + raise ValueError(f"AutoresetGraph requires a CUDA device, got {device}") + if isinstance(warmup, bool) or not isinstance(warmup, int) or warmup < 1: + raise ValueError(f"AutoresetGraph warmup must be an integer >= 1, got {warmup!r}") + scope = getattr(backend, "torch_graph_capture_scope", None) + if not callable(scope): + raise NotImplementedError( + f"{type(backend).__name__} does not implement torch_graph_capture_scope" + ) + for generator in generators: + if generator.device.type != "cuda": + raise ValueError( + f"AutoresetGraph generators must be CUDA generators, got {generator.device}" + ) + self._backend = backend + self._device = device + self._body = body + self._bind_capture = bind_capture + self._generators = tuple(generators) + self._warmup = warmup + self._capture_error_mode = capture_error_mode + self._eager_calls = 0 + self._graph: torch.cuda.CUDAGraph | None = None + self._capture: ExternalGraphCapture | None = None + self.replays = 0 + + @property + def captured(self) -> bool: + return self._graph is not None + + @property + def host_transition_names(self) -> tuple[str, ...]: + """Names of the host transitions replayed after every graph replay.""" + return () if self._capture is None else self._capture.host_transition_names + + def run(self) -> None: + """Execute the body: eagerly during warmup, then by graph replay.""" + if self._graph is not None: + capture = self._capture + assert capture is not None + capture.check_replay_guards() + self._graph.replay() + capture.replay_host_transitions() + self.replays += 1 + return + if self._eager_calls < self._warmup: + self._eager_calls += 1 + self._body() + return + self._capture_and_run() + + def _capture_and_run(self) -> None: + current = torch.cuda.current_stream(self._device) + stream = torch.cuda.Stream(device=self._device) + stream.wait_stream(current) + graph = torch.cuda.CUDAGraph() + for generator in self._generators: + graph.register_generator_state(generator) + capture: ExternalGraphCapture | None = None + try: + with torch.cuda.graph( + graph, stream=stream, capture_error_mode=self._capture_error_mode + ): + with self._backend.torch_graph_capture_scope( + stream, capture_error_mode=self._capture_error_mode + ) as capture: + self._bind_capture(capture) + try: + self._body() + finally: + self._bind_capture(None) + except GraphCaptureError: + raise + except Exception as exc: + raise GraphCaptureError(f"autoreset graph capture failed: {exc}") from exc + current.wait_stream(stream) + assert capture is not None + self._graph = graph + self._capture = capture + # Capture recorded the device work without executing it; the host + # transitions already ran once during capture. + graph.replay() + self.replays += 1 + + +__all__ = ["AutoresetGraph"] diff --git a/src/unilab/base/backend_factory.py b/src/unilab/base/backend_factory.py index 630d088ff..8be30b825 100644 --- a/src/unilab/base/backend_factory.py +++ b/src/unilab/base/backend_factory.py @@ -270,6 +270,8 @@ def env_backend_kwargs(cfg: "EnvCfg") -> dict[str, Any]: # Always explicit: never fall back to UniSim's blocking MJWarp default. # Other backends ignore this option. result["mjwarp_validate_tensor_finite"] = cfg.mjwarp_validate_tensor_finite + if cfg.mjwarp_defer_tensor_validation: + result["mjwarp_defer_tensor_validation"] = True return result diff --git a/src/unilab/base/base.py b/src/unilab/base/base.py index 3e03b3917..7eb144e75 100644 --- a/src/unilab/base/base.py +++ b/src/unilab/base/base.py @@ -80,6 +80,12 @@ class EnvCfg: # explicitly and disables the duplicate backend check by default. Set it to # True only to debug a non-finite operand at the exact backend call. mjwarp_validate_tensor_finite: bool = False + # Latch MJWarp's fail-closed operand/bounds checks into device verdicts + # instead of reading them on the host; TorchEnv publishes them with its + # deferred finite validation. Required for a host-sync-free step when + # mjwarp_validate_tensor_finite is True. Forwarded only when True, so + # UniSim releases without the option keep working. + mjwarp_defer_tensor_validation: bool = False # ``newton`` owns MuJoCo-Warp 3.11 capacity and requires explicit CUDA # placement; the process binder supplies the rank-local device when unset. newton_device: Optional[str] = None @@ -185,6 +191,11 @@ def validate(self): "mjwarp_validate_tensor_finite must be bool, " f"got {self.mjwarp_validate_tensor_finite!r}" ) + if not isinstance(self.mjwarp_defer_tensor_validation, bool): + raise ValueError( + "mjwarp_defer_tensor_validation must be bool, " + f"got {self.mjwarp_defer_tensor_validation!r}" + ) if self.newton_device is not None and ( not isinstance(self.newton_device, str) or not self.newton_device.strip() ): diff --git a/src/unilab/base/host_transitions.py b/src/unilab/base/host_transitions.py new file mode 100644 index 000000000..2827856d2 --- /dev/null +++ b/src/unilab/base/host_transitions.py @@ -0,0 +1,39 @@ +"""Explicit host-state transitions for CUDA-graph-captured env segments. + +A CUDA graph replays device work only. Python state that a captured segment +updates (a "pending metrics" flag, an invalidated cache, a rebinding of a +published carrier) would otherwise keep the value it had after the capture +pass. Code that runs inside a capturable segment therefore declares each such +update as a named transition through :func:`record_host_transition`: + +- outside a capture the transition simply runs; +- during a capture it runs once and is recorded on the active + ``unisim.backend.graph_capture.ExternalGraphCapture``, and the graph owner + re-applies it after every replay. + +Only state that other code changes between two executions of the segment needs +a transition. Assignments whose value is the same on every execution (a flag +set and cleared inside the segment, a binding to a persistent buffer) are +idempotent and replay-correct without one. +""" + +from __future__ import annotations + +from collections.abc import Callable + + +def record_host_transition(owner: object, name: str, transition: Callable[[], None]) -> None: + """Run ``transition`` now and replay it after graph replays of the segment. + + ``owner`` is the environment (or anything exposing + ``record_host_transition(name, transition)``); owners without it run the + transition eagerly, which is the correct behavior outside any capture. + """ + recorder = getattr(owner, "record_host_transition", None) + if recorder is None: + transition() + return + recorder(name, transition) + + +__all__ = ["record_host_transition"] diff --git a/src/unilab/base/reset_selection.py b/src/unilab/base/reset_selection.py new file mode 100644 index 000000000..6829853ff --- /dev/null +++ b/src/unilab/base/reset_selection.py @@ -0,0 +1,192 @@ +"""Fixed-shape selected-row reset (``ResetSelection``). + +The public reset contract speaks IsaacLab/mjlab ``env_ids`` row indices. Row +indices have a data-dependent shape: building them from a device done mask +needs ``nonzero`` (a host synchronization), and every consumer then launches +kernels whose shape depends on how many rows reset. Neither is compatible with +a host-sync-free step or with CUDA graph capture. + +``ResetSelection`` is the fixed-shape alternative. It carries a persistent +full-width bool ``mask`` that the environment rewrites in place before each +reset; every participant applies its reset as a masked full-width update, so +shapes never depend on the selected row count and no host read decides them. + +The selection is opt-in per participant. A manager, term or buffer declares +support with the ``supports_reset_selection`` attribute (``True``) and accepts +a ``ResetSelection`` wherever it accepts ``env_ids``. Participants that do not +declare support never receive one: the environment negotiates support before +it enables the masked autoreset and reports every blocker by name. Code that +receives a selection by mistake fails closed, because a ``ResetSelection`` is +not a tensor and cannot be used as an index. +""" + +from __future__ import annotations + +from collections.abc import Callable, Iterable +from dataclasses import dataclass +from typing import Any, TypeVar + +import torch + +SUPPORTS_RESET_SELECTION = "supports_reset_selection" + +_T = TypeVar("_T") + + +@dataclass(frozen=True, eq=False) +class ResetSelection: + """Full-width reset selection: row ``i`` resets when ``mask[i]`` is true. + + ``mask`` is owned by the environment and has a fixed address for the + environment's lifetime; its contents change between resets. ``row_ids`` is + the persistent ``arange(num_envs)``; it is *not* the set of reset rows and + exists for full-width gathers. + """ + + mask: torch.Tensor + row_ids: torch.Tensor + + def __post_init__(self) -> None: + if not isinstance(self.mask, torch.Tensor) or self.mask.dtype != torch.bool: + raise TypeError("ResetSelection.mask must be a bool torch.Tensor") + if self.mask.ndim != 1: + raise ValueError(f"ResetSelection.mask must be 1-D; got shape {tuple(self.mask.shape)}") + if ( + not isinstance(self.row_ids, torch.Tensor) + or self.row_ids.dtype != torch.int64 + or tuple(self.row_ids.shape) != tuple(self.mask.shape) + or self.row_ids.device != self.mask.device + ): + raise ValueError("ResetSelection.row_ids must be int64 arange(num_envs) on mask.device") + + @classmethod + def allocate(cls, num_envs: int, device: torch.device | str) -> ResetSelection: + """Allocate a persistent selection with an all-false mask.""" + device = torch.device(device) + return cls( + mask=torch.zeros(num_envs, dtype=torch.bool, device=device), + row_ids=torch.arange(num_envs, dtype=torch.int64, device=device), + ) + + @property + def num_envs(self) -> int: + return int(self.mask.shape[0]) + + @property + def device(self) -> torch.device: + return self.mask.device + + def row_mask(self, ndim: int, *, row_dim: int = 0) -> torch.Tensor: + """Return ``mask`` shaped to broadcast over a ``ndim`` tensor's row dim.""" + if not 0 <= row_dim < ndim: + raise ValueError(f"row_dim {row_dim} outside a {ndim}-D tensor") + shape = [1] * ndim + shape[row_dim] = self.num_envs + return self.mask.view(shape) + + +def is_reset_selection(selector: object) -> bool: + return isinstance(selector, ResetSelection) + + +def supports_reset_selection(participant: object) -> bool: + """Whether ``participant`` declared that it accepts a ``ResetSelection``. + + For class-based participants the declaration must sit at or below the + class that defines the effective ``reset``: a subclass that overrides + ``reset`` does not inherit its parent's declaration and has to opt in + again, so an unconverted override is reported instead of failing later. + """ + if getattr(participant, SUPPORTS_RESET_SELECTION, False) is not True: + return False + cls = participant if isinstance(participant, type) else type(participant) + if not isinstance(participant, type) and SUPPORTS_RESET_SELECTION in vars(participant): + return True # declared on the instance itself (function terms) + mro = cls.__mro__ + declared_by = next((c for c in mro if SUPPORTS_RESET_SELECTION in vars(c)), None) + reset_by = next((c for c in mro if "reset" in vars(c)), None) + if declared_by is None or reset_by is None: + return True + return issubclass(declared_by, reset_by) + + +def reset_selection_capable(participant: _T) -> _T: + """Mark a function-style term (for example a reset event) as selection-capable. + + Use only for callables whose body handles a ``ResetSelection`` argument; + class-based participants set ``supports_reset_selection = True`` instead. + """ + setattr(participant, SUPPORTS_RESET_SELECTION, True) + return participant + + +def selection_blockers(owner: str, participants: Iterable[tuple[str, object]]) -> list[str]: + """Name every participant that has not declared ``ResetSelection`` support.""" + return [ + f"{owner}: {name} ({_describe(participant)})" + for name, participant in participants + if not supports_reset_selection(participant) + ] + + +def _describe(participant: object) -> str: + # Functions and classes carry __qualname__; instances report their type. + return getattr(participant, "__qualname__", None) or type(participant).__qualname__ + + +def fill_rows_( + tensor: torch.Tensor, + rows: torch.Tensor | slice | ResetSelection | None, + value: float | int | bool, + *, + row_dim: int = 0, +) -> None: + """``tensor[rows] = value`` along ``row_dim`` for any reset selector.""" + if isinstance(rows, ResetSelection): + tensor.masked_fill_(rows.row_mask(tensor.ndim, row_dim=row_dim), value) + return + index: tuple[Any, ...] = (slice(None),) * row_dim + (slice(None) if rows is None else rows,) + tensor[index] = value + + +def assign_rows_( + tensor: torch.Tensor, + rows: torch.Tensor | slice | ResetSelection | None, + values: torch.Tensor | Callable[[int], torch.Tensor], + *, + row_dim: int = 0, +) -> None: + """Write new values into the selected rows. + + ``values`` is either a full-width tensor (row ``i`` of ``values`` lands in + row ``i`` of ``tensor``) or a factory ``values(n)`` producing ``n`` rows. + Under a ``ResetSelection`` the factory is called with ``num_envs`` and the + result is merged through the mask, so its shape stays fixed. + """ + if isinstance(rows, ResetSelection): + full = values(rows.num_envs) if callable(values) else values + tensor.copy_(torch.where(rows.row_mask(tensor.ndim, row_dim=row_dim), full, tensor)) + return + index: tuple[Any, ...] = (slice(None),) * row_dim + (slice(None) if rows is None else rows,) + if callable(values): + if rows is None: + count = tensor.shape[row_dim] + elif isinstance(rows, slice): + count = len(range(*rows.indices(tensor.shape[row_dim]))) + else: + count = int(rows.numel()) + tensor[index] = values(count) + return + tensor[index] = values[index] + + +__all__ = [ + "SUPPORTS_RESET_SELECTION", + "ResetSelection", + "assign_rows_", + "fill_rows_", + "is_reset_selection", + "reset_selection_capable", + "selection_blockers", + "supports_reset_selection", +] diff --git a/src/unilab/base/reset_selectors.py b/src/unilab/base/reset_selectors.py index 582898814..54ecb64bb 100644 --- a/src/unilab/base/reset_selectors.py +++ b/src/unilab/base/reset_selectors.py @@ -14,8 +14,12 @@ from __future__ import annotations +from typing import overload + import torch +from unilab.base.reset_selection import ResetSelection + class ResetMaskBuffer: """Reusable fixed-width bool reset mask owned by one conversion site. @@ -40,14 +44,23 @@ def num_envs(self) -> int: def device(self) -> torch.device: return self._mask.device - def scatter(self, env_ids: torch.Tensor) -> torch.Tensor: - """Rewrite the buffer from row indices and return the aliased mask.""" + def scatter(self, env_ids: torch.Tensor | ResetSelection) -> torch.Tensor: + """Rewrite the buffer from row indices and return the aliased mask. + + A ``ResetSelection`` is copied without scattering (fixed shape, no + host read). + """ + if isinstance(env_ids, ResetSelection): + self._mask.copy_(env_ids.mask) + return self._mask self._mask.zero_() if env_ids.numel(): self._mask[env_ids.to(device=self._mask.device, dtype=torch.int64)] = True return self._mask - def resolve(self, selector: torch.Tensor | slice | None, *, owner: str) -> torch.Tensor: + def resolve( + self, selector: torch.Tensor | slice | ResetSelection | None, *, owner: str + ) -> torch.Tensor: """Resolve a public reset selector like ``resolve_reset_mask``. Every selector is normalized to ``env_ids`` and scattered into this @@ -59,19 +72,50 @@ def resolve(self, selector: torch.Tensor | slice | None, *, owner: str) -> torch return self.scatter(env_ids) +@overload +def resolve_reset_env_ids( + selector: ResetSelection, + *, + num_envs: int, + device: torch.device, + owner: str, +) -> ResetSelection: ... + + +@overload def resolve_reset_env_ids( selector: torch.Tensor | slice | None, *, num_envs: int, device: torch.device, owner: str, -) -> torch.Tensor: +) -> torch.Tensor: ... + + +def resolve_reset_env_ids( + selector: torch.Tensor | slice | ResetSelection | None, + *, + num_envs: int, + device: torch.device, + owner: str, +) -> torch.Tensor | ResetSelection: """Normalize a public reset selector to an int64 row-index tensor. ``None`` selects every row; a slice selects ``arange(num_envs)[slice]``; an integer tensor is validated and moved onto ``device``. Fixed-width bool masks are rejected: the v1.4 mask form was removed in v1.5. + + A ``ResetSelection`` is returned unchanged after a host-only structural + check: it is the fixed-shape selector of selection-capable participants + (see ``unilab.base.reset_selection``), valid by construction. """ + if isinstance(selector, ResetSelection): + if selector.num_envs != num_envs or not _same_device(selector.device, device): + raise ValueError( + f"{owner} ResetSelection covers {selector.num_envs} rows on " + f"{selector.device}; expected {num_envs} rows on {device}" + ) + return selector if selector is None: return torch.arange(num_envs, dtype=torch.int64, device=device) if isinstance(selector, slice): @@ -98,6 +142,13 @@ def resolve_reset_env_ids( return ids +def _same_device(left: torch.device, right: torch.device | str) -> bool: + right = torch.device(right) + return left.type == right.type and ( + left.index is None or right.index is None or left.index == right.index + ) + + def reset_env_ids_mask( env_ids: torch.Tensor, *, @@ -109,7 +160,7 @@ def reset_env_ids_mask( def resolve_reset_mask( - selector: torch.Tensor | slice | None, + selector: torch.Tensor | slice | ResetSelection | None, *, num_envs: int, device: torch.device, @@ -118,9 +169,11 @@ def resolve_reset_mask( """Resolve a public reset selector to the internal fixed-width bool mask. The selector is normalized to ``env_ids`` and scattered into a freshly - allocated mask. + allocated mask; a ``ResetSelection`` contributes a copy of its mask. """ env_ids = resolve_reset_env_ids(selector, num_envs=num_envs, device=device, owner=owner) + if isinstance(env_ids, ResetSelection): + return env_ids.mask.clone() return reset_env_ids_mask(env_ids, num_envs=num_envs, device=device) diff --git a/src/unilab/base/reset_state.py b/src/unilab/base/reset_state.py index 41e8fb0ab..4b91e0537 100644 --- a/src/unilab/base/reset_state.py +++ b/src/unilab/base/reset_state.py @@ -10,7 +10,7 @@ from __future__ import annotations import time -from collections.abc import Iterator +from collections.abc import Callable, Iterator from contextlib import contextmanager from typing import TYPE_CHECKING, Any, cast @@ -99,6 +99,7 @@ def __init__( ) -> None: self._backend = backend self._num_envs = backend.num_envs + self._deferred_mask_require: Callable[[torch.Tensor, str, str], None] | None = None self._selected_default_qpos = default_qpos self.scene_layout = scene_layout if scene_layout is not None and any( @@ -170,6 +171,21 @@ def active(self) -> bool: """Whether a reset lifecycle currently owns the transaction.""" return self._active or self._tensor_active + def set_deferred_mask_checks( + self, require: Callable[[torch.Tensor, str, str], None] | None + ) -> None: + """Route fixed-width mask consistency checks through ``require``. + + ``require(ok, name, message)`` latches a 0-d bool device verdict that + the owning environment publishes at its single host boundary. While + installed, the transaction never reads a mask on the host: equality and + containment checks become deferred verdicts, and the empty-mask + short-circuit is skipped (a masked commit of zero rows is a no-op). + The environment installs it only for fixed-shape (``ResetSelection``) + autoresets; ``None`` restores the eager checks. + """ + self._deferred_mask_require = require + def set_finite_verdict_sink(self, verdict: torch.Tensor | None) -> None: """Accumulate reset payload finite verdicts into the env's carrier. @@ -1121,7 +1137,11 @@ def reset_to_default(self, reset_mask: torch.Tensor, *, term_name: str) -> None: "the active mask" ) self._record_tensor_write_mask(reset_mask) - elif not torch.equal(active_mask, reset_mask): + elif not self._same_mask( + active_mask, + reset_mask, + f"EventManager term '{term_name}' attempted default reset outside the active mask", + ): raise ValueError( f"EventManager term '{term_name}' attempted default reset outside the active mask" ) @@ -1221,7 +1241,12 @@ def write_joint_state_mask( f"EventManager term '{term_name}' attempted joint-state mutation outside " "the active mask" ) - elif not torch.equal(active_mask, reset_mask): + elif not self._same_mask( + active_mask, + reset_mask, + f"EventManager term '{term_name}' attempted joint-state mutation outside " + "the active mask", + ): raise ValueError( f"EventManager term '{term_name}' attempted joint-state mutation outside " "the active mask" @@ -1401,7 +1426,11 @@ def _require_tensor_mask_write( f"ManagerBased tensor mask {capability} attempted a write outside " "the active mask" ) - elif not torch.equal(reset_mask, self._tensor_active_mask): + elif not self._same_mask( + self._tensor_active_mask, + reset_mask, + f"ManagerBased tensor mask {capability} must use the transaction mask", + ): raise ValueError(f"ManagerBased tensor mask {capability} must use the transaction mask") self._requesting_terms.add(capability) return self._tensor_active_mask @@ -1416,12 +1445,38 @@ def _tensor_mask_scope(self, reset_mask: torch.Tensor, *, capability: str) -> to return reset_mask & active_mask return reset_mask + def _same_mask(self, active_mask: torch.Tensor, reset_mask: torch.Tensor, message: str) -> bool: + """Whether ``reset_mask`` equals the transaction mask (deferred if installed).""" + if reset_mask is active_mask: + return True + require = self._deferred_mask_require + if require is not None: + require( + torch.eq(reset_mask, active_mask).all(), + "reset_state.transaction_mask", + message, + ) + return True + return bool(torch.equal(active_mask, reset_mask)) + def _tensor_mask_has_outside_rows( self, reset_mask: torch.Tensor, active_mask: torch.Tensor ) -> bool: + require = self._deferred_mask_require + if require is not None: + require( + ~(reset_mask & ~active_mask).any(), + "reset_state.outside_rows", + "ManagerBased tensor mask write outside the active mask", + ) + return False return bool((reset_mask & ~active_mask).any()) def _mask_is_empty(self, reset_mask: torch.Tensor) -> bool: + if self._deferred_mask_require is not None: + # Fixed-shape mode: masked writes and the masked backend commit are + # no-ops for unselected rows, so emptiness is never decided on host. + return False self._mask_count.copy_(reset_mask.sum()) return bool(self._mask_count.item() == 0) @@ -2458,7 +2513,11 @@ def _validate_tensor_values( ) if not values.is_floating_point(): raise TypeError(f"EventManager term '{term_name}' {capability} must be floating") - if not bool(torch.isfinite(values).all()): + sink = self._finite_verdict_sink + if sink is not None: + # Deferred mode: latch on the env carrier; its poll boundary raises. + sink.logical_and_(torch.isfinite(values).all()) + elif not bool(torch.isfinite(values).all()): raise ValueError(f"EventManager term '{term_name}' {capability} contains NaN or Inf") return values.contiguous() diff --git a/src/unilab/base/torch_env.py b/src/unilab/base/torch_env.py index 9518e9625..a7037df5a 100644 --- a/src/unilab/base/torch_env.py +++ b/src/unilab/base/torch_env.py @@ -30,12 +30,15 @@ TensorExecution, torch_device_matches, ) +from unisim.backend.graph_capture import DeviceVerdicts, ExternalGraphCapture +from unilab.base.autoreset_graph import AutoresetGraph from unilab.base.backend_timing import ( RESET_DONE_DETAIL_TIMING_KEYS, ) from unilab.base.base import ABEnv, EnvCfg, EnvPlayCapabilities from unilab.base.cpu_runtime import apply_cpu_confinement +from unilab.base.reset_selection import ResetSelection from unilab.base.scene import SceneCfg from unilab.dtype_config import get_global_dtype @@ -84,6 +87,17 @@ class TorchEnv(ABEnv): rewards fail closed. """ + # Fixed-shape step defaults; ``_configure_fixed_shape_step`` overrides them + # in the constructor. Declared on the class so constructor-free instances + # (unit-test shells built with ``__new__``) take the ordinary step path. + _masked_autoreset: bool = False + _persistent_step_buffers: bool = False + _autoreset_reset_active: bool = False + _reset_selection: ResetSelection | None = None + _active_graph_capture: ExternalGraphCapture | None = None + _autoreset_graph: AutoresetGraph | None = None + _published_state: TorchEnvState | None = None + def __init__( self, cfg: EnvCfg, @@ -125,6 +139,78 @@ def __init__( raise ValueError("finite_validation_interval must be a positive integer") self._finite_validation_interval = interval self._finite_verdict = torch.ones((), dtype=torch.bool, device=self._device) + # Named deferred fail-closed checks (``require``); published together + # with ``_finite_verdict`` by the single read in poll_finite_validation. + self._verdicts = DeviceVerdicts() + self._configure_fixed_shape_step(cfg) + + def _configure_fixed_shape_step(self, cfg: EnvCfg) -> None: + """Resolve the opt-in sync-free / graph-capturable step switches. + + - ``masked_autoreset``: autoreset runs every step as a fixed-shape + ``reset(ResetSelection)`` under the device done mask (no ``any()`` + gate, no ``nonzero``); requires deferred finite validation. + - ``persistent_step_buffers``: the working step carriers (observations, + reward, terminated, truncated) are allocated once and updated with + ``copy_``; consumers receive a ring of ``published_buffer_depth`` + persistent slots (``0`` publishes the working carriers themselves). + - ``autoreset_cuda_graph``: the masked autoreset segment is captured into + one CUDA graph after ``autoreset_graph_warmup`` eager executions and + replayed afterwards; requires both switches above, a CUDA device and + a backend declaring ``external_graph_capture``. + """ + + def flag(name: str) -> bool: + value = getattr(cfg, name, False) + if not isinstance(value, bool): + raise TypeError(f"{name} must be bool, got {type(value).__name__}") + return value + + def bounded_int(name: str, default: int, minimum: int) -> int: + value = getattr(cfg, name, default) + if isinstance(value, bool) or not isinstance(value, int) or value < minimum: + raise ValueError(f"{name} must be an integer >= {minimum}, got {value!r}") + return value + + self._masked_autoreset = flag("masked_autoreset") + self._persistent_step_buffers = flag("persistent_step_buffers") + self._autoreset_cuda_graph = flag("autoreset_cuda_graph") + self._published_buffer_depth = bounded_int("published_buffer_depth", 2, 0) + self._autoreset_graph_warmup = bounded_int("autoreset_graph_warmup", 2, 1) + capture_mode = getattr(cfg, "autoreset_graph_capture_mode", "thread_local") + if capture_mode not in ("global", "thread_local", "relaxed"): + raise ValueError( + "autoreset_graph_capture_mode must be 'global', 'thread_local' or 'relaxed'; " + f"got {capture_mode!r}" + ) + self._autoreset_graph_capture_mode = capture_mode + if self._masked_autoreset and not self._defer_finite_validation: + raise ValueError( + "masked_autoreset requires defer_finite_validation: its fail-closed checks " + "are published at the deferred validation boundary" + ) + if self._autoreset_cuda_graph: + if not (self._masked_autoreset and self._persistent_step_buffers): + raise ValueError( + "autoreset_cuda_graph requires masked_autoreset and persistent_step_buffers" + ) + if self._device.type != "cuda": + raise ValueError( + f"autoreset_cuda_graph requires a CUDA env device, got {self._device}" + ) + self._reset_selection: ResetSelection | None = ( + ResetSelection.allocate(self._num_envs, self._device) + if self._masked_autoreset + else None + ) + # Device mirror of ``step_counter`` for reset participants that record + # the trigger step inside a fixed-shape (and possibly captured) reset. + self._step_counter_tensor = torch.zeros((), dtype=torch.int64, device=self._device) + self._active_graph_capture: ExternalGraphCapture | None = None + self._autoreset_graph: AutoresetGraph | None = None + self._publication_ring: list[TorchEnvState] = [] + self._publication_cursor = 0 + self._published_state: TorchEnvState | None = None @property def device(self) -> torch.device: @@ -181,8 +267,64 @@ def num_envs(self) -> int: @property def state(self) -> Optional[TorchEnvState]: + """The latest published transition. + + With ``persistent_step_buffers`` this is a publication slot whose + tensors stay unchanged for ``published_buffer_depth - 1`` further + steps (or, at depth 0, the working carriers that the next step + overwrites in place). + """ + if self._published_state is not None: + return self._published_state return self._state + @property + def masked_autoreset(self) -> bool: + return self._masked_autoreset + + @property + def persistent_step_buffers(self) -> bool: + return self._persistent_step_buffers + + @property + def step_counter_tensor(self) -> torch.Tensor: + """0-d int64 device mirror of ``step_counter`` (fixed address).""" + return self._step_counter_tensor + + @property + def autoreset_graph(self) -> AutoresetGraph | None: + return self._autoreset_graph + + def require(self, ok: torch.Tensor, name: str, message: str) -> None: + """Fail-closed check of a 0-d bool device verdict. + + With deferred finite validation the verdict is AND-ed into the named + persistent flag ``name`` (no host read, CUDA-graph safe) and published + at the next ``poll_finite_validation`` boundary, which raises + ``ValueError(message)``. Otherwise it is read now and raises at once. + """ + if self._defer_finite_validation: + self._verdicts.record(name, ok, message) + return + if not bool(ok): + raise ValueError(message) + + def record_host_transition(self, name: str, transition: Callable[[], None]) -> None: + """Run a host-state update and replay it after autoreset graph replays. + + See ``unilab.base.host_transitions``. Outside a capture this simply + runs ``transition``. + """ + capture = self._active_graph_capture + if capture is None: + transition() + return + capture.record_host_transition(f"env.{name}", transition) + + def reset_selection_blockers(self) -> list[str]: + """Reset participants that cannot take a ``ResetSelection``.""" + return [f"{type(self).__qualname__} does not implement a ResetSelection reset"] + @property def obs_groups_spec(self) -> dict[str, int]: raise NotImplementedError( @@ -224,11 +366,26 @@ def init_state(self) -> TorchEnvState: info={"steps": initial_steps}, ) self._last_initial_episode_steps = initial_steps.clone() + if self._masked_autoreset: + blockers = self.reset_selection_blockers() + if blockers: + raise NotImplementedError( + "masked_autoreset needs every reset participant to accept a " + "ResetSelection; unsupported: " + "; ".join(blockers) + ) # terminated is statically all-ones here; the bootstrap reset must not - # re-read the done mask from the device. + # re-read the done mask from the device. The bootstrap is a cold path + # and always uses explicit row ids. self._reset_done_envs(done_active=True) self._clear_step_final_observation() - return self._state + if self._autoreset_cuda_graph and self._autoreset_graph is None: + self._autoreset_graph = self._build_autoreset_graph() + if self._persistent_step_buffers: + self._allocate_publication_ring() + self._publish_step_state() + state = self.state + assert state is not None + return state def _initial_episode_steps(self) -> torch.Tensor: """Allocate the public int64 episode counter tensor before reset.""" @@ -333,6 +490,8 @@ def step(self, actions: torch.Tensor) -> TorchEnvState: self._state.info["steps"] += 1 self.step_counter += 1 + if self._masked_autoreset: + self._step_counter_tensor.fill_(self.step_counter) truncated = self._compute_truncated(self._state) self._state.truncated.logical_or_(truncated) @@ -344,21 +503,31 @@ def step(self, actions: torch.Tensor) -> TorchEnvState: # synchronization point after an event-backed asynchronous backend # step, so its wait belongs to the backend step phase rather than to # reset bookkeeping. - did_reset = self._autoreset and bool((self._state.terminated | self._state.truncated).any()) + # The fixed-shape autoreset has no gate: it runs every step under the + # device done mask, so nothing here reads the device. + did_reset = self._autoreset and ( + self._masked_autoreset or bool((self._state.terminated | self._state.truncated).any()) + ) step_core_ms += (time.perf_counter() - phase) * 1000.0 step_core_cpu_ms += (_cpu_time() - phase_cpu) * 1000.0 phase = time.perf_counter() phase_cpu = _cpu_time() if did_reset: - self._reset_done_envs(done_active=True) + if self._masked_autoreset: + self._run_masked_autoreset() + else: + self._reset_done_envs(done_active=True) reset_done_ms = (time.perf_counter() - phase) * 1000.0 reset_done_cpu_ms = (_cpu_time() - phase_cpu) * 1000.0 self.poll_finite_validation() timing = self._state.info.setdefault("timing", {}) - if not did_reset: + if not did_reset or self._masked_autoreset: + # The fixed-shape path has no host row count or per-phase split. self._clear_reset_done_detail_timing(timing) + if did_reset: + timing["reset_done_reset_call_ms"] = reset_done_ms timing["env_step_total_ms"] = (time.perf_counter() - started) * 1000.0 timing["action_validate_ms"] = action_validate_ms timing["apply_action_ms"] = apply_action_ms @@ -383,7 +552,11 @@ def step(self, actions: torch.Tensor) -> TorchEnvState: completion = backend_result.get("completion") if completion is not None: self._state.info["completion"] = completion - return self._state + if self._persistent_step_buffers: + self._publish_step_state() + state = self.state + assert state is not None + return state @property def requires_substep_state_feedback(self) -> bool: @@ -764,6 +937,111 @@ def _reset_done_envs(self, *, done_active: bool = False) -> None: self._clear_reset_done_detail_timing(timing) timing.update(detail_timing) + def _run_masked_autoreset(self) -> None: + """Execute the fixed-shape autoreset, eagerly or by graph replay.""" + graph = self._autoreset_graph + if graph is not None: + graph.run() + else: + self._masked_autoreset_body() + + def _masked_autoreset_body(self) -> None: + """Fixed-shape autoreset: device work plus declared host transitions. + + Every tensor written here has a fixed address (persistent working + carriers, the persistent ``ResetSelection`` mask, final-observation + scratch), nothing reads the device from the host, and host state that + the rest of the step changes is updated only through + ``record_host_transition``. That makes the body capturable as one CUDA + graph (``autoreset_cuda_graph``) with no eager epilogue. + """ + assert self._state is not None + selection = self._reset_selection + assert selection is not None + state = self._state + done = selection.mask + torch.logical_or(state.terminated, state.truncated, out=done) + self._reset_done_count.copy_(done.sum()) + state.info["steps"].masked_fill_(done, 0) + final_observation = self._ensure_final_observation_scratch() + for name, values in state.obs.items(): + final_observation[name].copy_(values) + # step() clears final_observation before update_state on every step. + self.record_host_transition("final_observation", self._publish_final_observation) + self._autoreset_reset_active = True + self._autoreset_scattered_reset_obs = False + try: + new_obs, reset_info = self.reset(selection) # type: ignore[arg-type] + finally: + self._autoreset_reset_active = False + if not self._autoreset_scattered_reset_obs: + self._validate_reset_observation_mask(new_obs) + for name, values in new_obs.items(): + state.obs[name].copy_(torch.where(done[:, None], values, state.obs[name])) + self._scatter_reset_info_mask(reset_info, done) + + def _publish_final_observation(self) -> None: + assert self._state is not None + self._state.final_observation = self._final_observation_scratch + + def _build_autoreset_graph(self) -> AutoresetGraph: + capabilities = self._backend.get_tensor_capabilities() + if not getattr(capabilities, "external_graph_capture", False): + raise NotImplementedError( + f"autoreset_cuda_graph requires backend '{self._backend.backend_type}' to " + "declare the external_graph_capture tensor capability" + ) + return AutoresetGraph( + backend=self._backend, + device=self._device, + body=self._masked_autoreset_body, + bind_capture=self._bind_graph_capture, + generators=self._graph_generators(), + warmup=self._autoreset_graph_warmup, + capture_error_mode=self._autoreset_graph_capture_mode, + ) + + def _bind_graph_capture(self, capture: ExternalGraphCapture | None) -> None: + self._active_graph_capture = capture + + def _graph_generators(self) -> list[torch.Generator]: + """CUDA generators consumed by the autoreset segment (besides the default).""" + return [] + + def _allocate_publication_ring(self) -> None: + assert self._state is not None + if self._publication_ring or self._published_buffer_depth == 0: + return + state = self._state + self._publication_ring = [ + TorchEnvState( + obs={name: torch.empty_like(value) for name, value in state.obs.items()}, + reward=torch.empty_like(state.reward), + terminated=torch.empty_like(state.terminated), + truncated=torch.empty_like(state.truncated), + info=state.info, + ) + for _ in range(self._published_buffer_depth) + ] + + def _publish_step_state(self) -> None: + """Publish the working carriers into the next persistent ring slot.""" + state = self._state + assert state is not None + if self._published_buffer_depth == 0: + self._published_state = state + return + slot = self._publication_ring[self._publication_cursor] + self._publication_cursor = (self._publication_cursor + 1) % self._published_buffer_depth + for name, value in state.obs.items(): + slot.obs[name].copy_(value) + slot.reward.copy_(state.reward) + slot.terminated.copy_(state.terminated) + slot.truncated.copy_(state.truncated) + slot.info = state.info + slot.final_observation = state.final_observation + self._published_state = slot + def _validate_reset_observation_mask(self, obs: Mapping[str, torch.Tensor]) -> None: if not isinstance(obs, dict) or set(obs) != set(self.obs_groups_spec): raise ValueError("TorchEnv mask reset observations do not match obs_groups_spec") @@ -1073,10 +1351,15 @@ def poll_finite_validation(self, *, force: bool = False) -> bool: if isinstance(bad_verdict, torch.Tensor): guard_bad = bad_verdict verdicts.append(torch.logical_not(bad_verdict)) + named_sets = self._named_verdict_sets() + for named in named_sets: + verdicts.extend(named.device_flags()) valid = bool(torch.stack(verdicts).all().detach().cpu().item()) if valid: self._finite_verdict.fill_(True) + for named in named_sets: + named.clear() for manager in managers: reset = getattr(manager, "reset_finite_verdict", None) if callable(reset): @@ -1088,6 +1371,9 @@ def poll_finite_validation(self, *, force: bool = False) -> bool: return True # Detailed scans and diagnostics are reserved for the abnormal path. + for named in named_sets: + # Raises ValueError naming every failed fail-closed check. + named.flush() base_valid = bool(self._finite_verdict.detach().cpu().item()) self._finite_verdict.fill_(True) if guard is not None and guard_bad is not None: @@ -1107,6 +1393,16 @@ def poll_finite_validation(self, *, force: bool = False) -> bool: return True raise ValueError("TorchEnv deferred finite validation detected NaN or Inf") + def _named_verdict_sets(self) -> list[DeviceVerdicts]: + """Env-owned and backend-owned named verdicts, for the single poll read.""" + sets = [self._verdicts] + backend_verdicts = getattr(self._backend, "tensor_validation_verdicts", None) + if callable(backend_verdicts): + backend_set = backend_verdicts() + if isinstance(backend_set, DeviceVerdicts): + sets.append(backend_set) + return [named for named in sets if len(named)] + def set_nan_guard(self, guard: NanGuard) -> None: """Attach the owner-provided diagnostic guard. diff --git a/src/unilab/envs/manager_based_rl_env.py b/src/unilab/envs/manager_based_rl_env.py index a2ecb8036..e84556b9b 100644 --- a/src/unilab/envs/manager_based_rl_env.py +++ b/src/unilab/envs/manager_based_rl_env.py @@ -25,6 +25,7 @@ TensorLifecycleCapabilities, TensorProcessTopology, ) +from unisim.backend.graph_capture import GraphCaptureError from unilab.base.backend_factory import ( SELECTIVE_TRACKED_BODY_BACKENDS, @@ -45,6 +46,7 @@ SceneTensorReadPlan, SceneTensorReadSpec, ) +from unilab.base.reset_selection import ResetSelection from unilab.base.reset_selectors import ResetMaskBuffer, resolve_reset_env_ids from unilab.base.reset_state import ResetStateTransaction from unilab.base.scene import SceneCfg, resolve_scene_default_qpos @@ -131,6 +133,28 @@ class ManagerBasedRlEnvCfg(EnvCfg): critic_observation_group: str | None = None defer_finite_validation: bool = True finite_validation_interval: int = 1 + # Fixed-shape, host-sync-free autoreset (opt-in). Every reset participant + # must accept a ``ResetSelection`` (see unilab.base.reset_selection); the + # env reports unsupported participants by name at init_state. The reset + # lifecycle runs on every step (an all-false mask is a no-op), so + # full-width reset samplers also draw RNG on steps without a done row and + # same-seed trajectories differ from the selected-row lifecycle; + # ``final_observation`` is published every step and is meaningful on done + # rows only. + masked_autoreset: bool = False + # Persistent step carriers updated in place (opt-in). Consumers receive a + # ring of ``published_buffer_depth`` persistent publication slots; a + # published state stays valid for ``published_buffer_depth - 1`` further + # steps. ``0`` publishes the working carriers themselves (no copy; the + # next step overwrites them in place). + persistent_step_buffers: bool = False + published_buffer_depth: int = 2 + # Capture the masked autoreset segment into one CUDA graph (opt-in; + # requires masked_autoreset, persistent_step_buffers, a CUDA collector + # device and a backend declaring ``external_graph_capture``). + autoreset_cuda_graph: bool = False + autoreset_graph_warmup: int = 2 + autoreset_graph_capture_mode: str = "thread_local" def validate(self) -> None: if self.collector_tensor_device is not None: @@ -197,9 +221,32 @@ def validate(self) -> None: "auto_reset", "scale_rewards_by_dt", "defer_finite_validation", + "masked_autoreset", + "persistent_step_buffers", + "autoreset_cuda_graph", ): if not isinstance(getattr(self, name), bool): raise TypeError(f"ManagerBasedRlEnvCfg {name} must be bool") + for name, minimum in (("published_buffer_depth", 0), ("autoreset_graph_warmup", 1)): + value = getattr(self, name) + if isinstance(value, bool) or not isinstance(value, int) or value < minimum: + raise ValueError(f"ManagerBasedRlEnvCfg {name} must be an integer >= {minimum}") + if self.autoreset_graph_capture_mode not in ("global", "thread_local", "relaxed"): + raise ValueError( + "ManagerBasedRlEnvCfg autoreset_graph_capture_mode must be 'global', " + "'thread_local' or 'relaxed'" + ) + if self.masked_autoreset and not self.defer_finite_validation: + raise ValueError( + "ManagerBasedRlEnvCfg masked_autoreset requires defer_finite_validation" + ) + if self.autoreset_cuda_graph and not ( + self.masked_autoreset and self.persistent_step_buffers + ): + raise ValueError( + "ManagerBasedRlEnvCfg autoreset_cuda_graph requires masked_autoreset and " + "persistent_step_buffers" + ) if ( isinstance(self.finite_validation_interval, bool) or not isinstance(self.finite_validation_interval, int) @@ -405,6 +452,8 @@ def __init__(self, cfg: ManagerBasedRlEnvCfg, backend: SimBackend, num_envs: int self._reset_state.set_finite_verdict_sink( self._finite_verdict if self.defer_finite_validation else None ) + if self.masked_autoreset: + self._reset_state.set_deferred_mask_checks(self.require) self.scene = EntityScene.from_scene_cfg( cfg.scene, backend, @@ -1300,6 +1349,33 @@ def _initial_episode_steps(self) -> torch.Tensor: device=self.device, dtype=torch.int64 ) + def reset_selection_blockers(self) -> list[str]: + """Collect every reset participant that cannot take a ``ResetSelection``.""" + blockers: list[str] = [] + for name in ( + "curriculum_manager", + "event_manager", + "command_manager", + "observation_manager", + "action_manager", + "reward_manager", + "metrics_manager", + "termination_manager", + "recorder_manager", + "reset_owner_manager", + ): + manager = getattr(self, name) + report = getattr(manager, "reset_selection_blockers", None) + if not callable(report): + blockers.append(f"{type(manager).__qualname__} has no reset_selection_blockers") + continue + blockers.extend(report()) + return blockers + + def _graph_generators(self) -> list[torch.Generator]: + generator = self.rng.generator + return [generator] if generator.device.type == "cuda" else [] + def init_state(self) -> TorchEnvState: state = super().init_state() # TorchEnv's initial state allocates counters, marks every row @@ -1314,6 +1390,11 @@ def init_state(self) -> TorchEnvState: state.info["steps"].copy_(initial_steps) self.episode_length_buf.copy_(state.info["steps"]) self.extras = state.info + if self.persistent_step_buffers: + # The working observation carriers are the env's persistent + # ``obs_buf``: update_state and the masked reset write them in place. + assert self._state is not None + self.obs_buf = self._state.obs return state def step(self, actions: torch.Tensor) -> TorchEnvState: @@ -1483,24 +1564,31 @@ def _update_state_in_read_phase(self, state: TorchEnvState) -> TorchEnvState: phase_started = time.perf_counter() mapped_obs = self._map_observations(manager_obs) - self.obs_buf = mapped_obs + if not self.persistent_step_buffers: + self.obs_buf = mapped_obs self._has_transition = True timing["update_state_map_ms"] = (time.perf_counter() - phase_started) * 1000.0 phase_started = time.perf_counter() publish_started = time.perf_counter() - public_obs, public_reward, public_terminated, public_truncated = self._publish_step_outputs( - mapped_obs - ) + if self.persistent_step_buffers: + self._copy_step_outputs(state, mapped_obs) + else: + public_obs, public_reward, public_terminated, public_truncated = ( + self._publish_step_outputs(mapped_obs) + ) timing["update_state_publication_ms"] = (time.perf_counter() - publish_started) * 1000.0 replace_started = time.perf_counter() - replacement = state.replace( - obs=public_obs, - reward=public_reward, - terminated=public_terminated, - truncated=public_truncated, - ) + if self.persistent_step_buffers: + replacement = state + else: + replacement = state.replace( + obs=public_obs, + reward=public_reward, + terminated=public_terminated, + truncated=public_truncated, + ) timing["update_state_state_replace_ms"] = (time.perf_counter() - replace_started) * 1000.0 timing["update_state_publish_ms"] = (time.perf_counter() - phase_started) * 1000.0 update_children = ( @@ -1652,6 +1740,16 @@ def _publish_step_outputs( eager_truncated = eager_truncated.clone() return eager_obs, eager_reward, eager_terminated, eager_truncated + def _copy_step_outputs( + self, state: TorchEnvState, mapped_obs: Mapping[str, torch.Tensor] + ) -> None: + """Write the step outputs into the persistent working carriers.""" + for name, values in mapped_obs.items(): + state.obs[name].copy_(values) + state.reward.copy_(self.reward_buf) + state.terminated.copy_(self.reset_terminated) + state.truncated.copy_(self.reset_time_outs) + def _compute_truncated(self, state: TorchEnvState) -> torch.Tensor: return state.truncated @@ -1677,17 +1775,29 @@ def _collect_step_core_drain_ms(self, state: TorchEnvState) -> float: def reset( self, - env_ids: torch.Tensor | None = None, + env_ids: torch.Tensor | ResetSelection | None = None, *, seed: int | None = None, ) -> tuple[dict[str, torch.Tensor], dict[str, Any]]: - """Execute the selected-row (``env_ids``) Manager reset lifecycle.""" + """Execute the selected-row (``env_ids``) Manager reset lifecycle. + + ``env_ids`` may be a ``ResetSelection`` (fixed-shape rows under a + persistent device mask) when every participant supports it; the + masked autoreset uses this form. + """ if seed is not None: self.seed(seed) if self._state is None: state = self.init_state() return state.obs, {"log": state.info.get("log", {})} - if env_ids is None: + selection = env_ids if isinstance(env_ids, ResetSelection) else None + if selection is not None: + if selection is not self._reset_selection: + raise ValueError( + "ManagerBasedRlEnv.reset accepts only the env-owned ResetSelection" + ) + reset_mask = self._reset_mask_buffer.scatter(selection) + elif env_ids is None: env_ids = torch.arange(self.num_envs, device=self.device) reset_mask = self._all_env_reset_mask else: @@ -1701,7 +1811,9 @@ def reset( reset_timing: dict[str, float] = {} read_plan = self.scene.read_plan # An empty explicit reset does not publish a pre-reset recorder boundary. - if self._has_transition and env_ids.numel() > 0: + if self._has_transition and ( + selection is not None or (isinstance(env_ids, torch.Tensor) and env_ids.numel() > 0) + ): self.recorder_manager.record_pre_reset(env_ids) # Match the selected-row lifecycle: curriculum consumes the terminal # episode state before reset clears the public counters. @@ -1815,6 +1927,11 @@ def reset( torch.where(reset_mask[:, None], reset_obs[name], self._state.obs[name]) ) self._autoreset_scattered_reset_obs = True + if selection is not None and log and self._active_graph_capture is not None: + raise GraphCaptureError( + "a captured reset cannot replay host log entries published by reset " + f"participants: {sorted(log)}" + ) step_log = self._state.info.get("log") if step_log: log = {**step_log, **log} @@ -1825,15 +1942,22 @@ def reset( self.reset_buf.masked_fill_(reset_mask, False) self.reset_terminated.masked_fill_(reset_mask, False) self.reset_time_outs.masked_fill_(reset_mask, False) - for name, values in mapped_obs.items(): - if set(self.obs_buf) != set(mapped_obs): - self.obs_buf = mapped_obs - else: - self.obs_buf[name].copy_( - torch.where(reset_mask[:, None], values, self.obs_buf[name]) - ) + if self.obs_buf is not self._state.obs: + for name, values in mapped_obs.items(): + if set(self.obs_buf) != set(mapped_obs): + self.obs_buf = mapped_obs + else: + self.obs_buf[name].copy_( + torch.where(reset_mask[:, None], values, self.obs_buf[name]) + ) self.extras = self._state.info self.recorder_manager.record_post_reset(env_ids) + if self.persistent_step_buffers and not self._autoreset_reset_active: + # A caller-driven reset is its own publication boundary. + self._publish_step_state() + published = self.state + assert published is not None + return published.obs, {"log": log} return self._state.obs, {"log": log} def _map_observations( diff --git a/src/unilab/envs/mdp/actions/actions.py b/src/unilab/envs/mdp/actions/actions.py index 6af44b1c7..014263dca 100644 --- a/src/unilab/envs/mdp/actions/actions.py +++ b/src/unilab/envs/mdp/actions/actions.py @@ -14,6 +14,7 @@ import torch +from unilab.base.reset_selection import ResetSelection, fill_rows_ from unilab.base.torch_env import TorchEnv from unilab.managers._buffers.delay_buffer import DelayBuffer from unilab.managers.action_manager import ActionTerm, ActionTermCfg @@ -81,6 +82,11 @@ class BaseAction(ActionTerm): cfg: BaseActionCfg _entity: Entity + # Raw-action clearing and the optional delay buffer are fixed-shape masked + # writes under a ResetSelection; subclasses that add reset state must + # re-declare support. + supports_reset_selection: ClassVar[bool] = True + def __init__(self, cfg: BaseActionCfg, env: ManagerBasedRlEnv): super().__init__(cfg=cfg, env=env) target_ids, target_names = self._entity.find_joints_by_actuator_names(cfg.actuator_names) @@ -220,12 +226,9 @@ def _apply_affine(self, source: torch.Tensor, destination: torch.Tensor) -> None ) self._processed_actions_finite_valid.copy_(torch.isfinite(destination).all()) - def reset(self, env_ids: torch.Tensor | None = None) -> None: + def reset(self, env_ids: torch.Tensor | ResetSelection | None = None) -> None: """Zero the raw actions of the selected envs.""" - if env_ids is None: - self._raw_actions.zero_() - return - self._raw_actions[env_ids] = 0.0 + fill_rows_(self._raw_actions, env_ids, 0.0) def _entity_values(self, values: torch.Tensor) -> torch.Tensor: """Publish processed controls on the public Torch control boundary.""" @@ -262,6 +265,8 @@ class JointPositionAction(BaseAction): """Convert policy actions into entity joint-position targets.""" cfg: JointPositionActionCfg + # reset() forwards the selection to BaseAction and the (masked) DelayBuffer. + supports_reset_selection: ClassVar[bool] = True def __init__(self, cfg: JointPositionActionCfg, env: ManagerBasedRlEnv): if not isinstance(cfg.use_default_offset, bool): @@ -296,7 +301,7 @@ def requires_substep_state_feedback(self) -> bool: def requires_substep_state_feedback(self, value: bool) -> None: self._requires_substep_state_feedback = value - def reset(self, env_ids: torch.Tensor | None = None) -> None: + def reset(self, env_ids: torch.Tensor | ResetSelection | None = None) -> None: super().reset(env_ids) if self._delay_buffer is not None: self._delay_buffer.reset(env_ids) diff --git a/src/unilab/envs/mdp/observations.py b/src/unilab/envs/mdp/observations.py index de259beca..14b38045e 100644 --- a/src/unilab/envs/mdp/observations.py +++ b/src/unilab/envs/mdp/observations.py @@ -82,6 +82,8 @@ class _NamedSensorObservation(ManagerTermBase): _term_name = "named_sensor" _tensor_sensor_names: tuple[str, ...] = () + # Stateless: reset is a no-op for any selector. + supports_reset_selection = True def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv): super().__init__(env) diff --git a/src/unilab/managers/_buffers/circular_buffer.py b/src/unilab/managers/_buffers/circular_buffer.py index 2e019c429..3ed6b8f30 100644 --- a/src/unilab/managers/_buffers/circular_buffer.py +++ b/src/unilab/managers/_buffers/circular_buffer.py @@ -81,6 +81,8 @@ import torch +from unilab.base.reset_selection import ResetSelection + def _selector_row_count(ids: Sequence[int] | torch.Tensor | slice, batch_size: int) -> int: """Return a selector's row count from host metadata only (no device read).""" @@ -94,6 +96,12 @@ def _selector_row_count(ids: Sequence[int] | torch.Tensor | slice, batch_size: i class CircularBuffer: """Fixed-length circular buffer for batched tensor history. + ``reset`` and ``backfill`` also accept a ``ResetSelection`` (fixed-shape + masked rows, no host read). A selection reset must be followed by a + selection ``backfill`` in the same reset lifecycle, as the observation + manager does, so the host-side first-push counter stays exact without + knowing how many rows were selected. + Stores history with shape (max_len, batch_size, ...) internally. The `buffer` property returns chronologically ordered data (oldest to newest) with shape (batch_size, max_len, ...). @@ -132,6 +140,8 @@ class CircularBuffer: batch_size: Size of the batch dimension. """ + supports_reset_selection = True + def __init__(self, max_len: int, batch_size: int) -> None: if max_len < 1: raise ValueError(f"Buffer size must be >= 1, got {max_len}") @@ -187,7 +197,10 @@ def buffer(self) -> torch.Tensor: idx = (torch.arange(self._max_len, device=self._buffer.device) + start) % self._max_len return self._buffer[idx].transpose(0, 1) - def reset(self, batch_ids: Sequence[int] | torch.Tensor | slice | None = None) -> None: + def reset( + self, + batch_ids: Sequence[int] | torch.Tensor | slice | ResetSelection | None = None, + ) -> None: """Zero out values and counters for specified batch rows. Args: @@ -196,6 +209,13 @@ def reset(self, batch_ids: Sequence[int] | torch.Tensor | slice | None = None) - # Reset may run before the first append (env bootstrap); bind the # counters to the selecting rows' device so later appends, which bind # to the observation device, stay consistent. + if isinstance(batch_ids, ResetSelection): + self._initialize_counters(batch_ids.device) + assert self._num_pushes is not None + self._num_pushes.masked_fill_(batch_ids.mask, 0) + if self._buffer is not None: + self._buffer.masked_fill_(batch_ids.row_mask(self._buffer.ndim, row_dim=1), 0.0) + return device = batch_ids.device if isinstance(batch_ids, torch.Tensor) else None self._initialize_counters(device) assert self._num_pushes is not None @@ -208,7 +228,7 @@ def reset(self, batch_ids: Sequence[int] | torch.Tensor | slice | None = None) - self._pending_first_push_rows + _selector_row_count(ids, self._batch_size), ) - def backfill(self, data: torch.Tensor, batch_ids: torch.Tensor) -> None: + def backfill(self, data: torch.Tensor, batch_ids: torch.Tensor | ResetSelection) -> None: """Fill the given rows' entire history with one frame, without advancing time. Unlike append, the global pointer does not move and other rows are @@ -229,6 +249,11 @@ def backfill(self, data: torch.Tensor, batch_ids: torch.Tensor) -> None: raise TypeError(f"buffer data must be torch.Tensor, got {type(data).__name__}") self._initialize_counters(data.device) assert self._num_pushes is not None + if isinstance(batch_ids, ResetSelection): + rows = batch_ids.row_mask(self._buffer.ndim, row_dim=1) + self._buffer.copy_(torch.where(rows, data.unsqueeze(0), self._buffer)) + self._num_pushes.masked_fill_(batch_ids.mask, 1) + return self._buffer[:, batch_ids] = data[batch_ids].unsqueeze(0) self._num_pushes[batch_ids] = 1 self._pending_first_push_rows = max( diff --git a/src/unilab/managers/_buffers/delay_buffer.py b/src/unilab/managers/_buffers/delay_buffer.py index b73c68b05..a338e282c 100644 --- a/src/unilab/managers/_buffers/delay_buffer.py +++ b/src/unilab/managers/_buffers/delay_buffer.py @@ -9,6 +9,7 @@ import torch +from unilab.base.reset_selection import ResetSelection from unilab.managers._buffers import CircularBuffer @@ -120,6 +121,8 @@ class DelayBuffer: >>> # But each refresh has 20% chance to keep previous lag """ + supports_reset_selection = True + def __init__( self, min_lag: int = 0, @@ -186,12 +189,26 @@ def set_lags( values = torch.as_tensor(lags, dtype=torch.int64, device=self._device) self._current_lags[idx] = values.clamp(self.min_lag, self.max_lag) - def reset(self, batch_ids: Sequence[int] | torch.Tensor | slice | None = None) -> None: + def reset( + self, + batch_ids: Sequence[int] | torch.Tensor | slice | ResetSelection | None = None, + ) -> None: """Reset specified environments to initial state. Args: - batch_ids: Batch indices to reset, or None to reset all. + batch_ids: Batch indices to reset, or None to reset all. A + ``ResetSelection`` resets the masked rows with fixed shapes. """ + if isinstance(batch_ids, ResetSelection): + self._buffer.reset(batch_ids=batch_ids) + self._current_lags.masked_fill_(batch_ids.mask, 0) + self._step_count.masked_fill_(batch_ids.mask, 0) + if self.update_period > 0 and self.per_env_phase: + new_phases = self._sample_integers(0, self.update_period, shape=(self.batch_size,)) + self._phase_offsets.copy_( + torch.where(batch_ids.mask, new_phases, self._phase_offsets) + ) + return if isinstance(batch_ids, slice): indices = range(*batch_ids.indices(self.batch_size)) batch_ids = list(indices) @@ -212,7 +229,7 @@ def append(self, data: torch.Tensor) -> None: """ self._buffer.append(data) - def backfill(self, data: torch.Tensor, batch_ids: torch.Tensor) -> None: + def backfill(self, data: torch.Tensor, batch_ids: torch.Tensor | ResetSelection) -> None: """Backfill the given rows with one frame, without advancing time. Used after a partial reset: the reset rows (whose lags and step counters diff --git a/src/unilab/managers/action_manager.py b/src/unilab/managers/action_manager.py index 27c5c869a..0695676f3 100644 --- a/src/unilab/managers/action_manager.py +++ b/src/unilab/managers/action_manager.py @@ -12,6 +12,7 @@ import torch from prettytable import PrettyTable +from unilab.base.reset_selection import fill_rows_, selection_blockers from unilab.base.reset_selectors import resolve_reset_env_ids from unilab.managers.manager_base import ManagerBase, ManagerTermBase @@ -59,6 +60,9 @@ class ActionTerm(ManagerTermBase): uses_tensor_state_packet: ClassVar[bool] = False """Whether ``apply_actions`` consumes the current packed tensor state packet.""" + supports_reset_selection: ClassVar[bool] = False + """Whether ``reset`` accepts a fixed-shape ``ResetSelection`` (opt-in).""" + # Contract invariant: raw actions are Torch tensors. ``process_actions`` and # the public ``raw_action`` carrier do not provide a NumPy fallback. @@ -234,10 +238,15 @@ def uses_tensor_state_packet(self) -> bool: def get_term(self, name: str) -> ActionTerm: return self._terms[name] + supports_reset_selection = True + def reset(self, env_ids: torch.Tensor | slice | None = None) -> dict[str, float]: self.clear_action_state(env_ids) return {} + def reset_selection_blockers(self) -> list[str]: + return selection_blockers("ActionManager", self._terms.items()) + def clear_action_state(self, env_ids: torch.Tensor | slice | None = None) -> None: """Clear manager and term action state for the selected envs without logging. @@ -251,9 +260,9 @@ def clear_action_state(self, env_ids: torch.Tensor | slice | None = None) -> Non device=self._device, owner="ActionManager.clear_action_state", ) - self._prev_action[env_ids] = 0.0 - self._prev_prev_action[env_ids] = 0.0 - self._action[env_ids] = 0.0 + fill_rows_(self._prev_action, env_ids, 0.0) + fill_rows_(self._prev_prev_action, env_ids, 0.0) + fill_rows_(self._action, env_ids, 0.0) for term in self._terms.values(): term.reset(env_ids) diff --git a/src/unilab/managers/command_manager.py b/src/unilab/managers/command_manager.py index 5c5d77f6b..0a50e47af 100644 --- a/src/unilab/managers/command_manager.py +++ b/src/unilab/managers/command_manager.py @@ -9,11 +9,13 @@ import math import time from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Sequence, cast +from typing import TYPE_CHECKING, Any, ClassVar, Sequence, cast import torch from prettytable import PrettyTable +from unilab.base.host_transitions import record_host_transition +from unilab.base.reset_selection import ResetSelection, selection_blockers from unilab.base.reset_selectors import ResetMaskBuffer, resolve_reset_env_ids from unilab.managers.manager_base import ManagerBase, ManagerTermBase from unilab.managers.rng import ManagerRng @@ -58,6 +60,13 @@ def build(self, env: ManagerBasedRlEnv) -> CommandTerm: class CommandTerm(ManagerTermBase): """Base class for command terms.""" + supports_reset_selection: ClassVar[bool] = False + """Whether the whole reset lifecycle (``reset``, ``_resample_command``, + ``refresh_reset``, ``post_compute_reset``) accepts a ``ResetSelection``: + full-width masked updates with no host read of the reset mask. The base + ``reset``/``_resample`` already are; a concrete term opts in once its own + hooks are.""" + def __init__(self, cfg: CommandTermCfg, env: ManagerBasedRlEnv): self.cfg = cfg super().__init__(env) @@ -91,10 +100,14 @@ def command(self): def reset( self, - env_ids: torch.Tensor | slice | None = None, + env_ids: torch.Tensor | slice | ResetSelection | None = None, *, publish_metrics: bool = True, ) -> dict[str, float]: + if isinstance(env_ids, ResetSelection) and not self._deferred_validation_enabled(): + raise RuntimeError( + f"CommandTerm '{self.name}' ResetSelection requires defer_finite_validation" + ) reset_mask = self._reset_mask_buffer.resolve( env_ids, owner=f"CommandTerm '{self.name}' reset", @@ -136,7 +149,11 @@ def reset( pending = torch.zeros((), dtype=torch.float32, device=self._device) self._pending_metric_sums[metric_name] = pending pending.add_(selected.sum()) - self._pending_metric_dirty = True + record_host_transition( + self._env, + f"command.{self.name}.pending_metrics_dirty", + self._mark_metrics_dirty, + ) else: reduced = torch.stack( ( @@ -179,6 +196,9 @@ def reset( def _deferred_validation_enabled(self) -> bool: return bool(getattr(self._env, "defer_finite_validation", False)) + def _mark_metrics_dirty(self) -> None: + self._pending_metric_dirty = True + def publish_pending_reset_metrics(self) -> dict[str, float]: """Drain completed episodes only at the environment's host boundary.""" if not self._pending_metric_dirty: @@ -456,7 +476,20 @@ def get_active_iterable_terms(self, env_idx: int) -> Sequence[tuple[str, Sequenc terms.append((name, command[env_idx].detach().cpu().tolist())) return terms - def reset(self, env_ids: torch.Tensor | slice | None = None) -> dict[str, float]: + supports_reset_selection = True + + def reset_selection_blockers(self) -> list[str]: + blockers = selection_blockers("CommandManager", self._terms.items()) + if self._terms and not bool(getattr(self._env, "defer_finite_validation", False)): + blockers.append( + "CommandManager: reset metrics are read on the host unless " + "defer_finite_validation is enabled" + ) + return blockers + + def reset( + self, env_ids: torch.Tensor | slice | ResetSelection | None = None + ) -> dict[str, float]: """Reset command terms for the selected environment rows. Every term consumes the public ``env_ids`` row selector; the packed @@ -617,6 +650,11 @@ def _validate_command_contract(self, name: str, command: torch.Tensor) -> None: class NullCommandManager: """Placeholder for absent command manager that safely no-ops all operations.""" + supports_reset_selection = True + + def reset_selection_blockers(self) -> list[str]: + return [] + def __init__(self): self.active_terms: list[str] = [] self._terms: dict[str, Any] = {} diff --git a/src/unilab/managers/curriculum_manager.py b/src/unilab/managers/curriculum_manager.py index 3f6c4763d..3056c8fbb 100644 --- a/src/unilab/managers/curriculum_manager.py +++ b/src/unilab/managers/curriculum_manager.py @@ -95,6 +95,16 @@ def get_active_iterable_terms(self, env_idx: int) -> Sequence[tuple[str, Sequenc terms.append((term_name, data)) return terms + supports_reset_selection = True + + def reset_selection_blockers(self) -> list[str]: + # Curriculum state is published to the reset log as host scalars, so + # configured terms cannot run inside a host-sync-free reset. + return [ + f"CurriculumManager: term '{name}' publishes host curriculum state on reset" + for name in self._term_names + ] + def reset(self, env_ids: torch.Tensor | slice | None = None) -> dict[str, float]: """Publish curriculum state and reset term state for the selected envs.""" env_ids = resolve_reset_env_ids( @@ -162,6 +172,11 @@ def _prepare_terms(self) -> None: class NullCurriculumManager: """Placeholder for absent curriculum manager that safely no-ops all operations.""" + supports_reset_selection = True + + def reset_selection_blockers(self) -> list[str]: + return [] + def __init__(self): self.active_terms: list[str] = [] self._curriculum_state: dict[str, Any] = {} diff --git a/src/unilab/managers/event_manager.py b/src/unilab/managers/event_manager.py index 4d737ed8d..8f21b332b 100644 --- a/src/unilab/managers/event_manager.py +++ b/src/unilab/managers/event_manager.py @@ -8,11 +8,19 @@ import math from copy import deepcopy from dataclasses import dataclass +from functools import partial from typing import TYPE_CHECKING, Literal, cast import torch from prettytable import PrettyTable +from unilab.base.host_transitions import record_host_transition +from unilab.base.reset_selection import ( + ResetSelection, + assign_rows_, + fill_rows_, + selection_blockers, +) from unilab.base.reset_selectors import resolve_reset_env_ids from unilab.managers.manager_base import ManagerBase, ManagerTermBaseCfg from unilab.managers.rng import ManagerRng @@ -150,7 +158,26 @@ def get_term_cfg(self, term_name: str) -> EventTermCfg: return self._mode_term_cfgs[mode][index] raise ValueError(f"Event term '{term_name}' not found in active terms.") - def reset(self, env_ids: torch.Tensor | slice | None = None) -> dict[str, float]: + supports_reset_selection = True + + def reset_selection_blockers(self) -> list[str]: + """Reset-mode terms and stateful step/interval terms must opt in.""" + participants = [ + (name, cfg.func) + for name, cfg in zip( + self._mode_term_names.get("reset", ()), + self._mode_term_cfgs.get("reset", ()), + strict=True, + ) + ] + for mode, mode_cfg in self._mode_class_term_cfgs.items(): + if mode != "reset": + participants.extend((f"{mode}:{cfg.func!r}", cfg.func) for cfg in mode_cfg) + return selection_blockers("EventManager", participants) + + def reset( + self, env_ids: torch.Tensor | slice | ResetSelection | None = None + ) -> dict[str, float]: """Clear stateful event terms for the selected environment rows.""" ids = resolve_reset_env_ids( env_ids, num_envs=self.num_envs, device=self._device, owner="EventManager.reset" @@ -166,15 +193,26 @@ def reset(self, env_ids: torch.Tensor | slice | None = None) -> dict[str, float] if not term_cfg.is_global_time: assert term_cfg.interval_range_s is not None lower, upper = term_cfg.interval_range_s - sampled = self._sample_interval(lower, upper, ids.numel()) - self._interval_term_time_left[index][ids] = sampled + assign_rows_( + self._interval_term_time_left[index], + ids, + partial(self._sample_interval, lower, upper), + ) # Resampled rows restart at >= lower; untouched rows keep the - # previous host lower bound. - self._interval_host_time_left_lb[index] = min( - self._interval_host_time_left_lb[index], lower + # previous host lower bound. Interval apply lowers the bound + # between resets, so a captured reset must replay this update. + record_host_transition( + self._env, + f"event.interval_lower_bound.{index}", + partial(self._lower_interval_bound, index, lower), ) return {} + def _lower_interval_bound(self, index: int, lower: float) -> None: + self._interval_host_time_left_lb[index] = min( + self._interval_host_time_left_lb[index], lower + ) + def apply( self, mode: EventMode, @@ -207,6 +245,9 @@ def apply( # instead of leaking a slice into event functions. if env_ids is None or isinstance(env_ids, slice): raise ValueError("Event mode 'reset' requires concrete environment indices.") + if isinstance(env_ids, ResetSelection): + self._apply_reset_selection(env_ids) + return self._apply_reset(env_ids, global_env_step_count=cast(int, global_env_step_count)) return @@ -286,7 +327,46 @@ def _apply_reset(self, env_ids: torch.Tensor, *, global_env_step_count: int) -> ) term_cfg.func(self._env, valid_env_ids, **term_cfg.params) + def _apply_reset_selection(self, selection: ResetSelection) -> None: + """Fixed-shape reset dispatch: every term sees a ``ResetSelection``. + + The trigger step comes from the env's device step counter, not a host + integer, so a captured dispatch stays correct on replay. Terms with a + ``min_step_count_between_reset`` receive their own selection whose mask + is narrowed to the rows allowed to trigger (possibly empty). + """ + step_counter = getattr(self._env, "step_counter_tensor", None) + if not isinstance(step_counter, torch.Tensor): + raise RuntimeError( + "EventManager ResetSelection dispatch requires env.step_counter_tensor" + ) + step_rows = step_counter.expand(self.num_envs) + for index, term_cfg in enumerate(self._mode_term_cfgs["reset"]): + last_step = self._reset_term_last_triggered_step_id[index] + once = self._reset_term_last_triggered_once[index] + min_step_count = term_cfg.min_step_count_between_reset + if min_step_count == 0: + term_selection = selection + else: + valid_trigger = (step_rows - last_step) >= min_step_count + valid_trigger |= (last_step == 0) & ~once + term_selection = self._term_reset_selection(index, selection) + torch.logical_and(selection.mask, valid_trigger, out=term_selection.mask) + assign_rows_(last_step, term_selection, step_rows) + fill_rows_(once, term_selection, True) + term_cfg.func(self._env, term_selection, **term_cfg.params) + + def _term_reset_selection(self, index: int, selection: ResetSelection) -> ResetSelection: + scratch = self._reset_term_selections.get(index) + if scratch is None: + scratch = ResetSelection( + mask=torch.zeros_like(selection.mask), row_ids=selection.row_ids + ) + self._reset_term_selections[index] = scratch + return scratch + def _prepare_terms(self) -> None: + self._reset_term_selections: dict[int, ResetSelection] = {} self._interval_term_time_left: list[torch.Tensor] = list() self._interval_host_time_left_lb: list[float] = list() self._reset_term_last_triggered_step_id: list[torch.Tensor] = list() diff --git a/src/unilab/managers/manager_base.py b/src/unilab/managers/manager_base.py index c4d0f5d65..4334c5928 100644 --- a/src/unilab/managers/manager_base.py +++ b/src/unilab/managers/manager_base.py @@ -119,11 +119,26 @@ def reset(self, env_ids: torch.Tensor | slice | None = None) -> dict[str, Any]: The public selector follows the IsaacLab/mjlab ``env_ids`` row-index form (``None`` selects every row); the v1.4 fixed-width bool mask form - is removed. + is removed. Managers that report no ``reset_selection_blockers`` also + accept a ``unilab.base.reset_selection.ResetSelection``. """ del env_ids return {} + def reset_selection_blockers(self) -> list[str]: + """Name what prevents this manager from taking a ``ResetSelection``. + + The base ``reset`` is a no-op and accepts any selector. A subclass that + overrides ``reset`` is a blocker until it sets + ``supports_reset_selection = True`` (and then reports its own + unsupported terms by overriding this method). + """ + if type(self).reset is ManagerBase.reset or getattr( + type(self), "supports_reset_selection", False + ): + return [] + return [f"{type(self).__qualname__}.reset does not accept a ResetSelection"] + def get_active_iterable_terms(self, env_idx: int) -> Sequence[tuple[str, Sequence[float]]]: raise NotImplementedError diff --git a/src/unilab/managers/metrics_manager.py b/src/unilab/managers/metrics_manager.py index 2eb04dc47..407b9a872 100644 --- a/src/unilab/managers/metrics_manager.py +++ b/src/unilab/managers/metrics_manager.py @@ -12,6 +12,8 @@ import torch from prettytable import PrettyTable +from unilab.base.host_transitions import record_host_transition +from unilab.base.reset_selection import ResetSelection, fill_rows_, selection_blockers from unilab.base.reset_selectors import resolve_reset_env_ids from unilab.managers.manager_base import ManagerBase, ManagerTermBaseCfg @@ -169,18 +171,40 @@ def clear_episode_state(self, env_ids: torch.Tensor | slice | None = None) -> No ) self._clear_episode_state(ids) - def reset(self, env_ids: torch.Tensor | slice | None = None) -> dict[str, float]: + supports_reset_selection = True + + def reset_selection_blockers(self) -> list[str]: + blockers = selection_blockers( + "MetricsManager", ((repr(cfg.func), cfg.func) for cfg in self._class_term_cfgs) + ) + if self._term_names and not self._defer_finite_validation: + blockers.append( + "MetricsManager: episode summaries are read on the host unless " + "defer_finite_validation is enabled" + ) + return blockers + + def reset( + self, env_ids: torch.Tensor | slice | ResetSelection | None = None + ) -> dict[str, float]: """Publish episode summaries and clear state for the selected rows.""" ids = resolve_reset_env_ids( env_ids, num_envs=self.num_envs, device=self._device, owner="MetricsManager.reset" ) + if isinstance(ids, ResetSelection) and not self._defer_finite_validation: + raise RuntimeError("MetricsManager ResetSelection requires defer_finite_validation") if self._defer_finite_validation: - reset_mask = torch.zeros(self.num_envs, dtype=torch.bool, device=self._device) - reset_mask[ids] = True + if isinstance(ids, ResetSelection): + reset_mask = ids.mask + selected_rows = True + else: + reset_mask = torch.zeros(self.num_envs, dtype=torch.bool, device=self._device) + reset_mask[ids] = True + selected_rows = ids.numel() > 0 runtime_reset = int(getattr(self._env, "step_counter", 0)) > 0 or bool( getattr(self._env, "_has_transition", False) ) - if runtime_reset and ids.numel() > 0: + if runtime_reset and selected_rows: completed = reset_mask & (self._step_count > 0) self._pending_episode_count += completed.sum() safe_counts = self._step_count.to(torch.float32).clamp_min(1.0) @@ -195,9 +219,13 @@ def reset(self, env_ids: torch.Tensor | slice | None = None) -> dict[str, float] else: values = self._episode_sums[key] / safe_counts self._pending_episode_sums[key].add_(torch.where(completed, values, 0.0).sum()) - self._pending_episode_dirty = True + record_host_transition( + self._env, "metrics.pending_episode_dirty", self._mark_episode_dirty + ) self._clear_episode_state(ids) return {} + # A ResetSelection without deferred validation was rejected above. + assert not isinstance(ids, ResetSelection) extras: dict[str, float] = {} counts = self._step_count[ids].to(torch.float32) safe_counts = torch.clamp(counts, min=1.0) @@ -236,15 +264,18 @@ def publish_pending_reset_metrics(self) -> dict[str, float]: for name, value in zip(self._term_names, means, strict=True) } - def _clear_episode_state(self, env_ids: torch.Tensor) -> None: + def _mark_episode_dirty(self) -> None: + self._pending_episode_dirty = True + + def _clear_episode_state(self, env_ids: torch.Tensor | ResetSelection) -> None: """Clear episode accumulators for the selected rows.""" for key in self._episode_sums: - self._episode_sums[key][env_ids] = 0.0 + fill_rows_(self._episode_sums[key], env_ids, 0.0) if key in self._episode_max: - self._episode_max[key][env_ids] = float("-inf") - self._step_count[env_ids] = 0 + fill_rows_(self._episode_max[key], env_ids, float("-inf")) + fill_rows_(self._step_count, env_ids, 0) for buf in self._substep_accum: - buf[env_ids] = 0.0 + fill_rows_(buf, env_ids, 0.0) for term_cfg in self._class_term_cfgs: if not callable(getattr(term_cfg.func, "reset", None)): raise NotImplementedError( @@ -347,6 +378,11 @@ def _log_mean(values: torch.Tensor) -> float: class NullMetricsManager: """Placeholder for absent metrics manager that safely no-ops all operations.""" + supports_reset_selection = True + + def reset_selection_blockers(self) -> list[str]: + return [] + def __init__(self): self.active_terms: list[str] = [] self.cfg = None diff --git a/src/unilab/managers/observation_manager.py b/src/unilab/managers/observation_manager.py index 5198a1dde..08047a7e8 100644 --- a/src/unilab/managers/observation_manager.py +++ b/src/unilab/managers/observation_manager.py @@ -9,7 +9,7 @@ from collections.abc import Callable from copy import deepcopy from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Literal, Sequence, cast +from typing import TYPE_CHECKING, Any, Literal, Sequence, cast import torch from prettytable import PrettyTable @@ -18,6 +18,8 @@ CONFIG_MAPPING_POLICY_KEY, MANAGER_TERM_MAPPING_POLICY, ) +from unilab.base.host_transitions import record_host_transition +from unilab.base.reset_selection import ResetSelection, selection_blockers from unilab.base.reset_selectors import ResetMaskBuffer, resolve_reset_env_ids from unilab.managers._buffers import CircularBuffer, DelayBuffer from unilab.managers._noise import noise_cfg, noise_model @@ -410,7 +412,31 @@ def get_term_cfg(self, group_name: str, term_name: str) -> ObservationTermCfg: index = self._group_obs_term_names[group_name].index(term_name) return self._group_obs_term_cfgs[group_name][index] - def reset(self, env_ids: torch.Tensor | slice | None = None) -> dict[str, float]: + supports_reset_selection = True + + def reset_selection_blockers(self) -> list[str]: + """Class terms must opt in; buffers and noise models are fixed-shape.""" + participants = [ + (f"{group}:{cfg.func!r}", cfg.func) + for group, group_cfg in self._group_obs_class_term_cfgs.items() + for cfg in group_cfg + ] + blockers = selection_blockers("ObservationManager", participants) + if not self._defer_finite_validation and any( + cfg.nan_policy != "disabled" for cfg in self._group_cfgs_for_selection() + ): + blockers.append( + "ObservationManager: NaN policies read the host unless " + "defer_finite_validation is enabled" + ) + return blockers + + def _group_cfgs_for_selection(self) -> list[Any]: + return [group for group in self.cfg.values() if group is not None] + + def reset( + self, env_ids: torch.Tensor | slice | ResetSelection | None = None + ) -> dict[str, float]: """Reset observation state for the selected environment rows. Class-term state, delay/history buffers, and noise models are cleared @@ -425,7 +451,9 @@ def reset(self, env_ids: torch.Tensor | slice | None = None) -> dict[str, float] owner="ObservationManager.reset", ) reset_mask = self._reset_mask_buffer.scatter(ids) - self._obs_buffer = None + # The per-step compute re-fills the cache between resets, so a captured + # reset must replay the invalidation. + record_host_transition(self._env, "observation.invalidate_cache", self._invalidate_cache) for group_cfg in self._group_obs_class_term_cfgs.values(): for term_cfg in group_cfg: if not callable(getattr(term_cfg.func, "reset", None)): @@ -452,6 +480,9 @@ def reset(self, env_ids: torch.Tensor | slice | None = None) -> dict[str, float] mod.reset(reset_mask) return {} + def _invalidate_cache(self) -> None: + self._obs_buffer = None + def bind_read_phase(self) -> None: """Bind class terms after the scene's packed tensor read plan exists.""" for group_cfg in self._group_obs_class_term_cfgs.values(): @@ -544,7 +575,7 @@ def _row_env_ids(mask: torch.Tensor) -> list[int]: def compute( self, update_history: bool = False, - env_ids: torch.Tensor | None = None, + env_ids: torch.Tensor | ResetSelection | None = None, ) -> dict[str, torch.Tensor | dict[str, torch.Tensor]]: """Compute observations for all groups. @@ -610,7 +641,7 @@ def compute_group( self, group_name: str, update_history: bool = False, - env_ids: torch.Tensor | None = None, + env_ids: torch.Tensor | ResetSelection | None = None, *, share_cache: dict[tuple, torch.Tensor] | None = None, ) -> torch.Tensor | dict[str, torch.Tensor]: diff --git a/src/unilab/managers/recorder_manager.py b/src/unilab/managers/recorder_manager.py index 76d3133dd..149833839 100644 --- a/src/unilab/managers/recorder_manager.py +++ b/src/unilab/managers/recorder_manager.py @@ -12,6 +12,7 @@ import torch from prettytable import PrettyTable +from unilab.base.reset_selection import selection_blockers from unilab.managers.manager_base import ManagerBase, ManagerTermBase, ManagerTermBaseCfg if TYPE_CHECKING: @@ -191,6 +192,11 @@ def get_term(self, name: str) -> RecorderTerm: msg = f"No recorder term named '{name}'. Active terms: {self.active_terms}" raise KeyError(msg) from None + supports_reset_selection = True + + def reset_selection_blockers(self) -> list[str]: + return selection_blockers("RecorderManager", self._terms.items()) + def record_pre_reset(self, env_ids: torch.Tensor) -> None: """Forward to each term's pre-reset boundary.""" for term in self._terms.values(): @@ -233,6 +239,11 @@ class NullRecorderManager: so it carries zero overhead. """ + supports_reset_selection = True + + def reset_selection_blockers(self) -> list[str]: + return [] + def __init__(self): self.active_terms: list[str] = [] self.cfg = None diff --git a/src/unilab/managers/reset_owner.py b/src/unilab/managers/reset_owner.py index 9ecc6ddf4..ac202bfc2 100644 --- a/src/unilab/managers/reset_owner.py +++ b/src/unilab/managers/reset_owner.py @@ -7,6 +7,7 @@ import torch +from unilab.base.reset_selection import selection_blockers from unilab.managers.manager_base import ManagerTermBase, ManagerTermBaseCfg if TYPE_CHECKING: @@ -84,6 +85,11 @@ def __init__(self, cfg: dict[str, ResetOwnerCfg | None], env: ManagerBasedRlEnv) def active_terms(self) -> list[str]: return list(self._terms) + supports_reset_selection = True + + def reset_selection_blockers(self) -> list[str]: + return selection_blockers("ResetOwnerManager", self._terms.items()) + @property def owner(self) -> ResetOwner | None: if not self._terms: @@ -118,6 +124,11 @@ def __str__(self) -> str: class NullResetOwnerManager: """No-op peer for owners without a reset-owner declaration.""" + supports_reset_selection = True + + def reset_selection_blockers(self) -> list[str]: + return [] + def __init__(self) -> None: self.active_terms: list[str] = [] self.cfg = None diff --git a/src/unilab/managers/reward_manager.py b/src/unilab/managers/reward_manager.py index 83460cd1b..9223330a3 100644 --- a/src/unilab/managers/reward_manager.py +++ b/src/unilab/managers/reward_manager.py @@ -14,6 +14,7 @@ import torch from prettytable import PrettyTable +from unilab.base.reset_selection import selection_blockers from unilab.base.reset_selectors import resolve_reset_env_ids from unilab.managers.manager_base import ManagerBase, ManagerTermBaseCfg @@ -154,6 +155,18 @@ def poll_finite_validation(self) -> bool: # Methods. + supports_reset_selection = True + + def reset_selection_blockers(self) -> list[str]: + return selection_blockers( + "RewardManager", + ( + (name, cfg.func) + for name, cfg in zip(self._term_names, self._term_cfgs, strict=True) + if any(cfg is class_cfg for class_cfg in self._class_term_cfgs) + ), + ) + def reset(self, env_ids: torch.Tensor | slice | None = None) -> dict[str, float]: """Clear stateful reward terms for the selected environment rows.""" ids = resolve_reset_env_ids( diff --git a/src/unilab/managers/termination_manager.py b/src/unilab/managers/termination_manager.py index 3a89f6b74..dfa9f89de 100644 --- a/src/unilab/managers/termination_manager.py +++ b/src/unilab/managers/termination_manager.py @@ -13,6 +13,8 @@ import torch from prettytable import PrettyTable +from unilab.base.host_transitions import record_host_transition +from unilab.base.reset_selection import ResetSelection, selection_blockers from unilab.base.reset_selectors import resolve_reset_env_ids from unilab.managers.manager_base import ManagerBase, ManagerTermBaseCfg @@ -99,25 +101,54 @@ def terminated(self) -> torch.Tensor: # Methods. - def reset(self, env_ids: torch.Tensor | slice | None = None) -> dict[str, int]: + supports_reset_selection = True + + def reset_selection_blockers(self) -> list[str]: + blockers = selection_blockers( + "TerminationManager", + ( + (name, cfg.func) + for name, cfg in zip(self._term_names, self._term_cfgs, strict=True) + if any(cfg is class_cfg for class_cfg in self._class_term_cfgs) + ), + ) + if not self._deferred_validation_enabled(): + blockers.append( + "TerminationManager: episode counts are read on the host unless " + "defer_finite_validation is enabled" + ) + return blockers + + def reset(self, env_ids: torch.Tensor | slice | ResetSelection | None = None) -> dict[str, int]: """Publish episodic counts and clear terms for the selected rows.""" ids = resolve_reset_env_ids( env_ids, num_envs=self.num_envs, device=self._device, owner="TerminationManager.reset" ) extras: dict[str, int] = {} keys = list(self._term_dones) - counts = [self._term_dones[key][ids].sum() for key in keys] + if isinstance(ids, ResetSelection): + if not self._deferred_validation_enabled(): + raise RuntimeError( + "TerminationManager ResetSelection requires defer_finite_validation" + ) + counts = [(self._term_dones[key] & ids.mask).sum() for key in keys] + selected_rows = True + else: + counts = [self._term_dones[key][ids].sum() for key in keys] + selected_rows = ids.numel() > 0 if self._deferred_validation_enabled(): # Accumulate counts on the device; the env's poll boundary drains # them through publish_pending_reset_metrics (the command/metrics # pending-buffer pattern), so resets never read the device. - if counts and ids.numel() > 0: + if counts and selected_rows: pending = self._pending_term_counts if pending is None: pending = torch.zeros(len(keys), dtype=torch.int64, device=self._device) self._pending_term_counts = pending pending.add_(torch.stack(counts)) - self._pending_counts_dirty = True + record_host_transition( + self._env, "termination.pending_counts_dirty", self._mark_counts_dirty + ) else: host_counts = torch.stack(counts).detach().cpu().tolist() if counts else [] for key, count in zip(keys, host_counts, strict=True): @@ -130,6 +161,9 @@ def reset(self, env_ids: torch.Tensor | slice | None = None) -> dict[str, int]: term_cfg.func.reset(ids) return extras + def _mark_counts_dirty(self) -> None: + self._pending_counts_dirty = True + def publish_pending_reset_metrics(self) -> dict[str, int]: """Drain accumulated episode counts at the environment's host boundary.""" if not self._pending_counts_dirty or self._pending_term_counts is None: diff --git a/src/unilab/tasks/motion_tracking/common/manager_terms.py b/src/unilab/tasks/motion_tracking/common/manager_terms.py index 807cd9c06..02f5c8202 100644 --- a/src/unilab/tasks/motion_tracking/common/manager_terms.py +++ b/src/unilab/tasks/motion_tracking/common/manager_terms.py @@ -6,10 +6,12 @@ import time from collections.abc import Sequence from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Callable, Literal, cast +from typing import TYPE_CHECKING, Callable, ClassVar, Literal, cast import torch +from unilab.base.host_transitions import record_host_transition +from unilab.base.reset_selection import ResetSelection, fill_rows_, is_reset_selection from unilab.envs.mdp.actions import JointPositionAction, JointPositionActionCfg from unilab.managers import ( CommandTerm, @@ -499,6 +501,11 @@ class MotionCommand(CommandTerm): """Device-resident motion reference command for Manager owners.""" cfg: MotionCommandCfg + # Every reset write is a full-width masked dispatch, so a fixed-shape + # ``ResetSelection`` is accepted. Under it the reset always samples (an + # all-false mask leaves state unchanged but consumes RNG) instead of + # probing the mask on the host. + supports_reset_selection: ClassVar[bool] = True def __init__(self, cfg: MotionCommandCfg, env: ManagerBasedRlEnv): _validate_motion_command_cfg(cfg) @@ -881,11 +888,15 @@ def anchor_ang_vel_w(self) -> torch.Tensor: return cast("torch.Tensor", self._motion_data.body_ang_vel_w)[:, self.anchor_body_idx] def reset( - self, env_ids: torch.Tensor | slice | None = None, *, publish_metrics: bool = True + self, + env_ids: torch.Tensor | slice | ResetSelection | None = None, + *, + publish_metrics: bool = True, ) -> dict[str, float]: """Selected-row motion command reset.""" del publish_metrics reset_mask = self._reset_mask_buffer.resolve(env_ids, owner="MotionCommand.reset") + fixed_shape = is_reset_selection(env_ids) if self._env.rng is None: raise NotImplementedError( "MotionCommand reset requires the Manager-owned Torch generator" @@ -910,7 +921,7 @@ def reset( resample_lower, resample_upper, (self.num_envs,), dtype=torch.float32 ) self.time_left.copy_(torch.where(reset_mask, sampled, self.time_left)) - self._resample_command(reset_mask) + self._resample_command(reset_mask, fixed_shape=fixed_shape) if ( env_ids is None or ( @@ -997,11 +1008,17 @@ def _resample(self, reset_mask: torch.Tensor) -> None: finally: self._skip_timer_resample = False - def _resample_command(self, reset_mask: torch.Tensor) -> None: - """Full-width motion sampling and masked state publication.""" + def _resample_command(self, reset_mask: torch.Tensor, *, fixed_shape: bool = False) -> None: + """Full-width motion sampling and masked state publication. + + ``fixed_shape`` (a ``ResetSelection`` reset, or any dispatch of an env + running ``masked_autoreset``) always samples: the mask is never read on + the host, so the dispatch is sync-free and CUDA-graph capturable. + """ if getattr(self, "_skip_timer_resample", False): return - if not bool(reset_mask.any()): + fixed_shape = fixed_shape or bool(getattr(self._env, "masked_autoreset", False)) + if not fixed_shape and not bool(reset_mask.any()): # An empty selected-reset mask is a no-op semantic boundary. It must # not consume sampler or task RNG merely to preserve a fixed-width # device dispatch (v1.3.5 contract). @@ -1152,7 +1169,11 @@ def post_compute(self) -> None: joint_view = read_plan.joint_tensor_view(self.robot) self._robot_joint_pos = joint_view.joint_pos self._robot_joint_vel = joint_view.joint_vel - self._robot_cache_step = self._env.common_step_counter + # The robot cache is fresh for the current step; a graph replay of the + # reset segment must re-stamp it with the replaying step. + record_host_transition( + self._env, "motion_command.robot_cache_step", self._stamp_robot_cache_step + ) rebind_ms = (time.perf_counter() - rebind_started) * 1000.0 self.last_post_compute_timing_ms = { "reset_done_motion_robot_refresh_ms": kernel_ms, @@ -1162,6 +1183,9 @@ def post_compute(self) -> None: } self.last_post_compute_timing_ms["reset_done_motion_relative_refresh_ms"] = 0.0 + def _stamp_robot_cache_step(self) -> None: + self._robot_cache_step = self._env.common_step_counter + def _sync_tensor_sampler_state(self) -> None: cast(torch.Tensor, self.time_steps).copy_(self.tensor_sampler.current_frames) cast(torch.Tensor, self.current_clip_end_frames).copy_( @@ -1413,12 +1437,14 @@ def validate_actions(self) -> None: """Publish the current latency-adjusted raw-input finite verdict.""" self.publish_validation_verdict() - def reset(self, env_ids: torch.Tensor | None = None) -> None: + supports_reset_selection: ClassVar[bool] = True + + def reset(self, env_ids: torch.Tensor | ResetSelection | None = None) -> None: super().reset(env_ids) if env_ids is None: self._previous_raw_actions.zero_() return - self._previous_raw_actions[env_ids] = 0.0 + fill_rows_(self._previous_raw_actions, env_ids, 0.0) def apply_actions(self) -> None: if isinstance(self._entity.data.control_buffer, torch.Tensor): @@ -1553,6 +1579,9 @@ def motion_global_anchor_orientation_error_exp( class _BodyTerm(ManagerTermBase): + # Stateless reads of the command carrier: the reset boundary is a no-op. + supports_reset_selection: ClassVar[bool] = True + def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv): super().__init__(env) command_name = cfg.params.get("command_name") @@ -1584,7 +1613,7 @@ def _validate(self, command_name: str, std: float) -> tuple[MotionCommand, float ) return _command(self._env, command_name), _positive_std(std, term_name=type(self).__name__) - def reset(self, env_ids: torch.Tensor) -> None: + def reset(self, env_ids: torch.Tensor | ResetSelection) -> None: """Row boundary for stateless motion reward reads.""" del env_ids @@ -1718,6 +1747,8 @@ def __call__( class BadAnchorPosZOnly(ManagerTermBase): """Anchor-height termination on the command's Torch carrier.""" + supports_reset_selection: ClassVar[bool] = True + def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv): super().__init__(env) command_name = cfg.params.get("command_name") @@ -1745,7 +1776,7 @@ def __call__( - robot_anchor_pos[:, command.anchor_body_idx, 2] ).abs() > threshold - def reset(self, env_ids: torch.Tensor) -> None: + def reset(self, env_ids: torch.Tensor | ResetSelection) -> None: """Row boundary for the stateless termination read.""" del env_ids diff --git a/src/unilab/tasks/motion_tracking/common/tensor_sampler.py b/src/unilab/tasks/motion_tracking/common/tensor_sampler.py index 1a2adfa83..93a38ccfb 100644 --- a/src/unilab/tasks/motion_tracking/common/tensor_sampler.py +++ b/src/unilab/tasks/motion_tracking/common/tensor_sampler.py @@ -401,7 +401,9 @@ def _publish_sampling_metrics(self, probabilities: torch.Tensor) -> None: self.sampling_entropy.fill_(1.0) top1_index = self._sampling_top1_index torch.argmax(probabilities, out=top1_index) - self.sampling_top1_prob.copy_(probabilities[top1_index]) + # The top-1 mass is the maximum; indexing with the 0-d argmax tensor + # would read it on the host. + self.sampling_top1_prob.copy_(torch.amax(probabilities)) self.sampling_top1_bin.copy_(top1_index) self.sampling_top1_bin.div_(self.bin_count) diff --git a/tests/envs/test_fixed_shape_autoreset.py b/tests/envs/test_fixed_shape_autoreset.py new file mode 100644 index 000000000..064964ac8 --- /dev/null +++ b/tests/envs/test_fixed_shape_autoreset.py @@ -0,0 +1,678 @@ +"""Fixed-shape (``ResetSelection``) autoreset, persistent step buffers, deferred +fail-closed checks and CUDA graph capture of the autoreset segment. + +The env under test runs on a device-resident fake backend whose physics is pure +Torch, so the same scenario runs on CPU and, when available, on CUDA (where the +autoreset segment is captured into a CUDA graph). +""" + +from __future__ import annotations + +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import dataclass +from typing import Any, cast + +import pytest +import torch +from unisim.backend.base import ( + SelectedResetPublication, + SimBackend, + TensorDataPlane, + TensorExecution, + TensorLifecycleCapabilities, + TensorProcessTopology, +) +from unisim.backend.graph_capture import ExternalGraphCapture, GraphCaptureError, closing_capture + +from tests.envs.test_manager_based_rl_env import _ResetBackend +from unilab.base.entity import EntityCfg +from unilab.base.reset_selection import ( + ResetSelection, + assign_rows_, + fill_rows_, + reset_selection_capable, + supports_reset_selection, +) +from unilab.base.reset_selectors import ResetMaskBuffer, resolve_reset_env_ids +from unilab.base.scene import SceneCfg +from unilab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg +from unilab.managers import ( + ActionTerm, + ActionTermCfg, + CircularBuffer, + CommandTerm, + CommandTermCfg, + CurriculumTermCfg, + EventTermCfg, + MetricsTermCfg, + ObservationGroupCfg, + ObservationTermCfg, + RewardTermCfg, + TerminationTermCfg, +) + +NUM_ENVS = 8 +CUDA = torch.cuda.is_available() +DEVICES = ["cpu", pytest.param("cuda", marks=pytest.mark.skipif(not CUDA, reason="CUDA"))] + + +# --------------------------------------------------------------------------- +# Device-resident fake backend +# --------------------------------------------------------------------------- + + +class _TensorBackend(_ResetBackend): + """Pure-Torch device-resident backend: no host reads on the hot path.""" + + backend_type = "fake_tensor" + + def __init__(self, num_envs: int, device: str) -> None: + super().__init__(num_envs) + self.device = torch.device(device) + self.qpos = torch.zeros((num_envs, 4), dtype=torch.float32, device=self.device) + self.qvel = torch.zeros((num_envs, 3), dtype=torch.float32, device=self.device) + self.capture_scopes = 0 + + def get_tensor_execution(self) -> TensorExecution: + return TensorExecution.DEVICE_RESIDENT + + def get_tensor_capabilities(self) -> TensorLifecycleCapabilities: + return TensorLifecycleCapabilities( + execution=TensorExecution.DEVICE_RESIDENT, + state_views=True, + state_fields=frozenset({"qpos", "qvel"}), + stepping=True, + selected_reset=True, + mask_selected_reset=True, + reset_randomization=True, + packed_reset=True, + device_reset_randomization=True, + packed_reset_fields=frozenset(("body_mass", "body_ipos", "geom_friction", "kp", "kd")), + selected_reset_publication=SelectedResetPublication.AUTHORITATIVE_VIEWS, + process_topology=TensorProcessTopology.IN_PROCESS, + data_plane=TensorDataPlane.DIRECT, + stream_event_ownership="test", + torch_devices=(self.device.type,), + external_graph_capture=self.device.type == "cuda", + ) + + def get_state_views(self, fields, device=None) -> dict[str, torch.Tensor]: + views: dict[str, torch.Tensor] = {} + if "qpos" in fields: + views["qpos"] = self.qpos + if "qvel" in fields: + views["qvel"] = self.qvel + return views + + def set_state_tensor(self, env_indices, qpos, qvel, randomization=None) -> None: + del randomization + self.qpos[env_indices] = qpos + self.qvel[env_indices] = qvel + + def set_state_tensor_mask(self, reset_mask, qpos, qvel, randomization=None, **_kwargs): + del randomization + self.qpos.copy_(torch.where(reset_mask[:, None], qpos, self.qpos)) + self.qvel.copy_(torch.where(reset_mask[:, None], qvel, self.qvel)) + return {"timing": {}} + + def step_tensor(self, ctrl: torch.Tensor, nsteps: int = 1) -> dict[str, Any]: + self.qpos[:, 3:4].add_(ctrl * (0.01 * nsteps)) + return {"timing": {}} + + @contextmanager + def torch_graph_capture_scope( + self, torch_stream: Any | None = None, *, capture_error_mode: str = "global" + ) -> Iterator[ExternalGraphCapture]: + del capture_error_mode + if not torch.cuda.is_current_stream_capturing(): + raise GraphCaptureError("scope requires an active Torch capture") + self.capture_scopes += 1 + stream = torch_stream or torch.cuda.current_stream(self.device) + with closing_capture(ExternalGraphCapture(torch_stream=stream)) as capture: + yield capture + + +# --------------------------------------------------------------------------- +# Selection-capable manager terms +# --------------------------------------------------------------------------- + + +@dataclass(kw_only=True) +class _DriveCfg(ActionTermCfg): + def build(self, env) -> ActionTerm: + return _DriveAction(self, env) + + +class _DriveAction(ActionTerm): + supports_reset_selection = True + + def __init__(self, cfg: _DriveCfg, env) -> None: + super().__init__(cfg, env) + self._processed = torch.zeros((self.num_envs, 1), dtype=torch.float32, device=env.device) + + @property + def action_dim(self) -> int: + return 1 + + @property + def raw_action(self) -> torch.Tensor: + return self._processed + + def process_actions(self, actions: torch.Tensor) -> None: + self._processed.copy_(actions) + + def reset(self, env_ids: torch.Tensor | ResetSelection | None = None) -> None: + fill_rows_(self._processed, env_ids, 0.0) + + def apply_actions(self) -> None: + self._entity.data.write_ctrl(self._processed) + + +@dataclass(kw_only=True) +class _TargetCfg(CommandTermCfg): + def build(self, env) -> CommandTerm: + return _Target(self, env) + + +class _Target(CommandTerm): + supports_reset_selection = True + + def __init__(self, cfg: _TargetCfg, env) -> None: + super().__init__(cfg, env) + self._command = torch.zeros((self.num_envs, 1), dtype=torch.float32, device=env.device) + self.metrics["error"] = torch.zeros(self.num_envs, dtype=torch.float32, device=env.device) + + @property + def command(self) -> torch.Tensor: + return self._command + + def _update_metrics(self) -> None: + self.metrics["error"].add_(self._command[:, 0].abs()) + + def _resample_command(self, reset_mask: torch.Tensor) -> None: + samples = self._env.rng.uniform(-1.0, 1.0, self._command.shape, dtype=torch.float32) + self._command.copy_(torch.where(reset_mask[:, None], samples, self._command)) + + def refresh_reset(self, reset_mask: torch.Tensor) -> None: + del reset_mask + + def post_compute_reset(self) -> None: + return None + + def _update_command(self) -> None: + return None + + +@reset_selection_capable +def _reset_joint(env: ManagerBasedRlEnv, env_ids: torch.Tensor | ResetSelection) -> None: + """Reset event: joint position from the selection-capable command sample.""" + target = env.command_manager.get_command("target") + env.scene["robot"].write_joint_state_to_sim_mask( + target.clone(), + torch.zeros((env.num_envs, 1), dtype=torch.float32, device=env.device), + reset_mask=env.event_term_mask_buffer.scatter(env_ids), + ) + + +def _joint_reset_ids_only(env: ManagerBasedRlEnv, env_ids: torch.Tensor) -> None: + """Undeclared reset event (ids only).""" + del env, env_ids + + +def _policy_obs(env: ManagerBasedRlEnv) -> torch.Tensor: + return torch.column_stack( + ( + env.episode_length_buf.to(torch.float32), + env.action_manager.action[:, 0], + env.command_manager.get_command("target")[:, 0], + # The fake backend's joint qpos column (device view, no host read). + cast(_TensorBackend, env.backend).qpos[:, 3], + ) + ) + + +def _critic_obs(env: ManagerBasedRlEnv) -> torch.Tensor: + return env.episode_length_buf[:, None].to(torch.float32) + + +def _reward(env: ManagerBasedRlEnv) -> torch.Tensor: + return env.action_manager.action[:, 0] - env.command_manager.get_command("target")[:, 0] + + +def _failure(env: ManagerBasedRlEnv) -> torch.Tensor: + return env.action_manager.action[:, 0] > 0.8 + + +def _time_out(env: ManagerBasedRlEnv) -> torch.Tensor: + return env.episode_length_buf >= 4 + + +def _metric(env: ManagerBasedRlEnv) -> torch.Tensor: + return env.episode_length_buf.to(torch.float32) + + +def _curriculum(env: ManagerBasedRlEnv, env_ids: Any) -> float: + del env, env_ids + return 0.0 + + +def _make_cfg(device: str, **overrides: Any) -> ManagerBasedRlEnvCfg: + cfg = ManagerBasedRlEnvCfg( + scene=SceneCfg( + model_file="fake.xml", + entities={"robot": EntityCfg(joint_names=("joint",), actuator_names=("motor",))}, + ), + sim_dt=0.01, + ctrl_dt=0.02, + max_episode_seconds=1.0, + seed=11, + collector_tensor_device=device, + observations={ + "actor": ObservationGroupCfg( + terms={"policy": ObservationTermCfg(func=_policy_obs, history_length=3)} + ), + "value": ObservationGroupCfg(terms={"critic": ObservationTermCfg(func=_critic_obs)}), + }, + actions={"drive": _DriveCfg(entity_name="robot")}, + rewards={"track": RewardTermCfg(func=_reward, weight=1.0)}, + terminations={ + "failure": TerminationTermCfg(func=_failure), + "time_out": TerminationTermCfg(func=_time_out, time_out=True), + }, + events={"joint": EventTermCfg(func=_reset_joint, mode="reset")}, + commands={"target": _TargetCfg(resampling_time_range=(10.0, 10.0))}, + metrics={"length": MetricsTermCfg(func=_metric, reduce="mean")}, + policy_observation_group="actor", + critic_observation_group="value", + defer_finite_validation=True, + ) + for name, value in overrides.items(): + setattr(cfg, name, value) + return cfg + + +def _make_env(device: str, **overrides: Any) -> tuple[ManagerBasedRlEnv, _TensorBackend]: + backend = _TensorBackend(NUM_ENVS, device) + env = ManagerBasedRlEnv(_make_cfg(device, **overrides), cast(SimBackend, backend), NUM_ENVS) + return env, backend + + +def _actions(device: str, steps: int) -> list[torch.Tensor]: + """Random actions; row 0 fails on every step. + + The masked autoreset dispatches its reset participants on every step, so + full-width reset samplers draw RNG even when no row is done, whereas the + selected-row path skips steps without a done row. Keeping one row failing + every step makes both lifecycles consume the RNG stream identically, so + same-seed trajectories can be compared exactly. + """ + generator = torch.Generator().manual_seed(3) + actions = [] + for _ in range(steps): + action = torch.rand((NUM_ENVS, 1), generator=generator) * 2.0 - 1.0 + action[0] = 0.9 + actions.append(action.to(device)) + return actions + + +def _snapshot(env: ManagerBasedRlEnv, backend: _TensorBackend) -> dict[str, torch.Tensor]: + state = env.state + assert state is not None + out = {f"obs.{k}": v.detach().clone() for k, v in state.obs.items()} + out["reward"] = state.reward.detach().clone() + out["terminated"] = state.terminated.detach().clone() + out["truncated"] = state.truncated.detach().clone() + out["steps"] = state.info["steps"].detach().clone() + out["qpos"] = backend.qpos.detach().clone() + out["command"] = env.command_manager.get_command("target").detach().clone() + # The masked autoreset publishes final_observation every step (meaningful + # on done rows only); the selected-row path publishes it only when some + # row reset. Compare the done rows. + done = state.terminated | state.truncated + final = state.final_observation or {k: torch.zeros_like(v) for k, v in state.obs.items()} + for name, value in final.items(): + out[f"final.{name}"] = torch.where(done[:, None], value, 0.0).detach().clone() + return out + + +def _run(env: ManagerBasedRlEnv, backend: _TensorBackend, actions: list[torch.Tensor]): + env.reset() + trajectory = [_snapshot(env, backend)] + logs = [] + for action in actions: + env.step(action) + trajectory.append(_snapshot(env, backend)) + logs.append(dict(env.state.info.get("log", {}))) + return trajectory, logs + + +def _assert_same(left, right) -> None: + assert len(left) == len(right) + for step, (a, b) in enumerate(zip(left, right, strict=True)): + assert a.keys() == b.keys(), step + for key in a: + torch.testing.assert_close(a[key].cpu(), b[key].cpu(), msg=f"step {step} {key}") + + +# --------------------------------------------------------------------------- +# ResetSelection primitives +# --------------------------------------------------------------------------- + + +def test_reset_selection_helpers_write_only_selected_rows() -> None: + selection = ResetSelection.allocate(4, "cpu") + selection.mask.copy_(torch.tensor([True, False, True, False])) + values = torch.arange(8, dtype=torch.float32).reshape(4, 2) + + fill_rows_(values, selection, -1.0) + assert values[:, 0].tolist() == [-1.0, 2.0, -1.0, 6.0] + + history = torch.zeros((3, 4, 2)) + fill_rows_(history, selection, 5.0, row_dim=1) + assert history[:, 1].abs().sum() == 0 and bool((history[:, 0] == 5.0).all()) + + timers = torch.zeros(4) + assign_rows_(timers, selection, lambda count: torch.full((count,), 7.0)) + assert timers.tolist() == [7.0, 0.0, 7.0, 0.0] + assign_rows_(timers, torch.tensor([1]), lambda count: torch.full((count,), 3.0)) + assert timers.tolist() == [7.0, 3.0, 7.0, 0.0] + + +def test_reset_override_must_redeclare_selection_support() -> None: + class Parent: + supports_reset_selection = True + + def reset(self, env_ids=None) -> None: + del env_ids + + class InheritsReset(Parent): + pass + + class OverridesReset(Parent): + def reset(self, env_ids=None) -> None: + del env_ids + + class Redeclares(OverridesReset): + supports_reset_selection = True + + class Withdraws(Parent): + supports_reset_selection = False + + assert supports_reset_selection(Parent()) + assert supports_reset_selection(InheritsReset()) + assert not supports_reset_selection(OverridesReset()) + assert not supports_reset_selection(OverridesReset) + assert supports_reset_selection(Redeclares()) + assert not supports_reset_selection(Withdraws()) + + def term(env, env_ids) -> None: + del env, env_ids + + assert not supports_reset_selection(term) + assert supports_reset_selection(reset_selection_capable(term)) + + +def test_selection_is_rejected_as_an_index_and_validated_by_resolvers() -> None: + selection = ResetSelection.allocate(4, "cpu") + with pytest.raises((TypeError, IndexError)): + torch.zeros(4)[selection] # type: ignore[index] + assert resolve_reset_env_ids(selection, num_envs=4, device="cpu", owner="t") is selection + with pytest.raises(ValueError, match="covers 4 rows"): + resolve_reset_env_ids(selection, num_envs=5, device="cpu", owner="t") + buffer = ResetMaskBuffer(4, torch.device("cpu")) + selection.mask[2] = True + assert buffer.scatter(selection).tolist() == [False, False, True, False] + with pytest.raises(TypeError): + ResetSelection(mask=torch.zeros(4), row_ids=torch.arange(4)) + + +def test_circular_buffer_selection_reset_and_backfill_match_ids() -> None: + data = torch.arange(12, dtype=torch.float32).reshape(4, 3) + by_ids, by_mask = CircularBuffer(3, 4), CircularBuffer(3, 4) + for buffer in (by_ids, by_mask): + buffer.append(data) + buffer.append(data + 100.0) + ids = torch.tensor([1, 3]) + selection = ResetSelection.allocate(4, "cpu") + selection.mask[ids] = True + by_ids.reset(ids) + by_ids.backfill(data * 2.0, ids) + by_mask.reset(selection) + by_mask.backfill(data * 2.0, selection) + torch.testing.assert_close(by_ids.buffer, by_mask.buffer) + by_ids.append(data + 7.0) + by_mask.append(data + 7.0) + torch.testing.assert_close(by_ids.buffer, by_mask.buffer) + + +# --------------------------------------------------------------------------- +# Masked autoreset: negotiation, equivalence, buffers, deferred checks +# --------------------------------------------------------------------------- + + +def test_masked_autoreset_requires_deferred_validation() -> None: + with pytest.raises(ValueError, match="defer_finite_validation"): + _make_cfg("cpu", masked_autoreset=True, defer_finite_validation=False).validate() + with pytest.raises(ValueError, match="requires masked_autoreset"): + _make_cfg("cpu", autoreset_cuda_graph=True).validate() + + +def test_masked_autoreset_reports_every_unsupported_participant() -> None: + env, _ = _make_env( + "cpu", + masked_autoreset=True, + events={"ids_only": EventTermCfg(func=_joint_reset_ids_only, mode="reset")}, + curriculum={"level": CurriculumTermCfg(func=_curriculum)}, + ) + with pytest.raises(NotImplementedError) as raised: + env.reset() + message = str(raised.value) + assert "_joint_reset_ids_only" in message + assert "CurriculumManager: term 'level'" in message + assert supports_reset_selection(_reset_joint) + assert not supports_reset_selection(_joint_reset_ids_only) + + +@pytest.mark.parametrize("device", DEVICES) +@pytest.mark.parametrize("depth", [0, 2]) +def test_masked_autoreset_matches_selected_row_autoreset(device: str, depth: int) -> None: + actions = _actions(device, 24) + eager_env, eager_backend = _make_env(device) + masked_env, masked_backend = _make_env( + device, + masked_autoreset=True, + persistent_step_buffers=True, + published_buffer_depth=depth, + ) + eager, eager_logs = _run(eager_env, eager_backend, actions) + masked, masked_logs = _run(masked_env, masked_backend, actions) + # Both lifecycles draw reset randomness with full-width masked draws, so + # same-seed trajectories are identical, including reset rows. + _assert_same(eager, masked) + assert eager_logs == masked_logs + assert any(bool((snap["terminated"] | snap["truncated"]).any()) for snap in masked) + + +def test_persistent_buffers_keep_addresses_and_publish_a_stable_ring() -> None: + env, backend = _make_env( + "cpu", masked_autoreset=True, persistent_step_buffers=True, published_buffer_depth=2 + ) + env.reset() + working = env._state + assert working is not None + working_ptrs = {name: value.data_ptr() for name, value in working.obs.items()} + working_ptrs["reward"] = working.reward.data_ptr() + assert env.obs_buf is working.obs + + published = [] + for action in _actions("cpu", 6): + state = env.step(action) + published.append((state, {k: v.clone() for k, v in state.obs.items()})) + assert env._state is working + assert {n: v.data_ptr() for n, v in working.obs.items()} == { + n: p for n, p in working_ptrs.items() if n != "reward" + } + assert working.reward.data_ptr() == working_ptrs["reward"] + # Ring of two persistent slots: alternate objects, and a published state is + # unchanged by the following step (the off-policy collector's obs_t). + assert published[0][0] is published[2][0] and published[1][0] is published[3][0] + assert published[0][0] is not published[1][0] + last_state, last_values = published[-1] + previous_state, previous_values = published[-2] + for name in previous_values: + torch.testing.assert_close(previous_state.obs[name], previous_values[name]) + torch.testing.assert_close(last_state.obs[name], last_values[name]) + del backend + + +def test_masked_step_has_no_reset_gate_and_reports_reset_call_timing() -> None: + env, _ = _make_env("cpu", masked_autoreset=True) + env.reset() + env.set_episode_length_buf(torch.zeros(NUM_ENVS, dtype=torch.int64)) + calls = [] + original = env.reset + + def counting_reset(env_ids=None, **kwargs): + calls.append(type(env_ids).__name__) + return original(env_ids, **kwargs) + + env.reset = counting_reset # type: ignore[method-assign] + state = env.step(torch.zeros((NUM_ENVS, 1))) + assert calls == ["ResetSelection"] + assert not bool((state.terminated | state.truncated).any()) + assert state.info["timing"]["reset_done_count"] == 0.0 + assert state.info["timing"]["reset_done_reset_call_ms"] >= 0.0 + + +def test_deferred_named_checks_fail_closed_at_the_poll_boundary() -> None: + env, _ = _make_env("cpu", masked_autoreset=True) + env.reset() + env.require(torch.tensor(True), "test.ok", "never raised") + env.require(torch.tensor(False), "test.bad", "custom deferred check failed") + with pytest.raises(ValueError, match="custom deferred check failed"): + env.step(torch.zeros((NUM_ENVS, 1))) + # The failed flag was cleared by the raise; the next boundary is clean. + env.step(torch.zeros((NUM_ENVS, 1))) + + +def test_transaction_mask_mismatch_is_deferred_not_read() -> None: + env, _ = _make_env("cpu", masked_autoreset=True) + env.reset() + transaction = env._reset_state + active = torch.tensor([True] + [False] * (NUM_ENVS - 1)) + other = torch.tensor([False, True] + [False] * (NUM_ENVS - 2)) + assert transaction._same_mask(active, other, "mask mismatch") + assert not transaction._mask_is_empty(torch.zeros(NUM_ENVS, dtype=torch.bool)) + with pytest.raises(ValueError, match="mask mismatch"): + env.poll_finite_validation(force=True) + + +def test_host_transitions_run_eagerly_outside_capture() -> None: + env, _ = _make_env("cpu", masked_autoreset=True) + env.reset() + seen = [] + env.record_host_transition("probe", lambda: seen.append(1)) + assert seen == [1] + + +# --------------------------------------------------------------------------- +# CUDA graph capture of the autoreset segment +# --------------------------------------------------------------------------- + +cuda_only = pytest.mark.skipif(not CUDA, reason="requires CUDA") + + +@cuda_only +def test_autoreset_graph_replay_matches_eager_masked_autoreset() -> None: + actions = _actions("cuda", 30) + eager_env, eager_backend = _make_env( + "cuda", masked_autoreset=True, persistent_step_buffers=True + ) + graph_env, graph_backend = _make_env( + "cuda", + masked_autoreset=True, + persistent_step_buffers=True, + autoreset_cuda_graph=True, + autoreset_graph_warmup=2, + ) + eager, eager_logs = _run(eager_env, eager_backend, actions) + graphed, graph_logs = _run(graph_env, graph_backend, actions) + graph = graph_env.autoreset_graph + assert graph is not None and graph.captured + assert graph.replays == len(actions) - 2 + assert graph_backend.capture_scopes == 1 + _assert_same(eager, graphed) + assert eager_logs == graph_logs + names = graph.host_transition_names + assert "env.final_observation" in names + assert "env.observation.invalidate_cache" in names + assert "env.metrics.pending_episode_dirty" in names + assert any(bool((snap["terminated"] | snap["truncated"]).any()) for snap in graphed[3:]) + + +@cuda_only +def test_graphed_masked_step_never_synchronizes_inside_the_autoreset() -> None: + env, _ = _make_env( + "cuda", + masked_autoreset=True, + persistent_step_buffers=True, + autoreset_cuda_graph=True, + autoreset_graph_warmup=1, + ) + env.reset() + for action in _actions("cuda", 3): + env.step(action) + graph = env.autoreset_graph + assert graph is not None and graph.captured + torch.cuda.set_sync_debug_mode("error") + try: + env._run_masked_autoreset() + finally: + torch.cuda.set_sync_debug_mode("default") + + +@cuda_only +def test_eager_masked_autoreset_body_never_synchronizes() -> None: + env, _ = _make_env("cuda", masked_autoreset=True, persistent_step_buffers=True) + env.reset() + for action in _actions("cuda", 3): + env.step(action) + torch.cuda.synchronize() + torch.cuda.set_sync_debug_mode("error") + try: + env._masked_autoreset_body() + finally: + torch.cuda.set_sync_debug_mode("default") + + +@cuda_only +def test_captured_reset_rejects_host_published_log_entries() -> None: + env, _ = _make_env( + "cuda", + masked_autoreset=True, + persistent_step_buffers=True, + autoreset_cuda_graph=True, + autoreset_graph_warmup=1, + ) + env.reset() + original = env.metrics_manager.reset + + def logging_reset(env_ids=None): + original(env_ids) + return {"Episode_Metrics/host": 1.0} + + env.metrics_manager.reset = logging_reset # type: ignore[method-assign] + env.step(_actions("cuda", 1)[0]) + with pytest.raises(GraphCaptureError, match="host log entries"): + env.step(_actions("cuda", 1)[0]) + + +def test_autoreset_graph_requires_cuda_and_backend_capability() -> None: + with pytest.raises(ValueError, match="CUDA env device"): + _make_env( + "cpu", + masked_autoreset=True, + persistent_step_buffers=True, + autoreset_cuda_graph=True, + ) diff --git a/tests/tasks/test_g1_motion_fixed_shape.py b/tests/tasks/test_g1_motion_fixed_shape.py new file mode 100644 index 000000000..c8099d3f8 --- /dev/null +++ b/tests/tasks/test_g1_motion_fixed_shape.py @@ -0,0 +1,108 @@ +"""G1 motion tracking on the fixed-shape (masked / CUDA-graph) autoreset path.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +import torch + +from unilab.base.reset_selection import supports_reset_selection +from unilab.envs.mdp.actions import JointPositionAction +from unilab.tasks.motion_tracking.common import manager_terms as motion_terms + +ROOT_DIR = Path(__file__).parents[2] + + +@pytest.mark.parametrize( + "term_cls", + [ + motion_terms.MotionCommand, + motion_terms.MotionJointPositionAction, + motion_terms.MotionRelativeBodyPositionErrorExp, + motion_terms.MotionRelativeBodyOrientationErrorExp, + motion_terms.MotionGlobalBodyLinearVelocityErrorExp, + motion_terms.MotionGlobalBodyAngularVelocityErrorExp, + motion_terms.UndesiredBodyContacts, + motion_terms.BadAnchorPosZOnly, + motion_terms.BadMotionBodyPosZOnly, + JointPositionAction, + ], +) +def test_g1_motion_terms_declare_reset_selection(term_cls: type) -> None: + assert supports_reset_selection(term_cls) + + +def _make_env(num_envs: int, **flags): + from hydra import compose, initialize_config_dir + from hydra.core.global_hydra import GlobalHydra + + from unilab.base.config_adapter import EnvConfigAdapter + from unilab.base.registry import apply_cfg_overrides + from unilab.envs import ManagerBasedRlEnvCfg, make_manager_based_rl_env + + GlobalHydra.instance().clear() + with initialize_config_dir( + config_dir=str(ROOT_DIR / "src" / "unilab" / "conf" / "flashsac"), version_base="1.3" + ): + owner = compose( + "config", + overrides=[ + "task=g1_motion_tracking/mjwarp", + "hydra.run.dir=.", + "hydra.output_subdir=null", + ], + ) + cfg = ManagerBasedRlEnvCfg() + apply_cfg_overrides( + cfg, + EnvConfigAdapter( + owner, root_dir=ROOT_DIR, algo_name="flashsac" + ).build_task_env_cfg_override(), + ) + for name, value in flags.items(): + setattr(cfg, name, value) + return make_manager_based_rl_env(cfg, num_envs=num_envs, backend_type="mjwarp") + + +@pytest.mark.slow +def test_g1_motion_tracking_captures_and_replays_autoreset_without_host_sync() -> None: + pytest.importorskip("mujoco_warp") + if not torch.cuda.is_available(): + pytest.skip("CUDA graph capture requires a CUDA device") + env = _make_env( + 64, + mjwarp_defer_tensor_validation=True, + masked_autoreset=True, + persistent_step_buffers=True, + autoreset_cuda_graph=True, + # Poll explicitly below so every step is checked for host syncs. + finite_validation_interval=10**6, + ) + assert env.reset_selection_blockers() == [] + env.seed(1) + env.init_state() + generator = torch.Generator(device=env.device).manual_seed(0) + dim = env.action_manager.total_action_dim + + def action() -> torch.Tensor: + return torch.randn((env.num_envs, dim), generator=generator, device=env.device) * 0.3 + + for _ in range(4): # warmup + capture + env.step(action()) + graph = env.autoreset_graph + assert graph is not None and graph.captured + done_rows = torch.zeros((), dtype=torch.int64, device=env.device) + torch.cuda.synchronize() + torch.cuda.set_sync_debug_mode("error") + try: + for _ in range(40): + state = env.step(action()) + done_rows += (state.terminated | state.truncated).sum() + finally: + torch.cuda.set_sync_debug_mode("default") + assert graph.replays >= 40 + assert int(done_rows) > 0, "the captured reset was not exercised" + assert env.poll_finite_validation(force=True) + for value in state.obs.values(): + assert bool(torch.isfinite(value).all()) diff --git a/tests/tasks/test_tensor_motion_sampler.py b/tests/tasks/test_tensor_motion_sampler.py index 2dad82431..c8f67d199 100644 --- a/tests/tasks/test_tensor_motion_sampler.py +++ b/tests/tasks/test_tensor_motion_sampler.py @@ -131,8 +131,8 @@ def test_adaptive_probabilities_and_metrics_are_numerically_correct() -> None: torch.testing.assert_close(repeat, expected, rtol=0, atol=1.0e-7) torch.testing.assert_close(repeat.sum(), torch.ones(()), rtol=0, atol=1.0e-7) assert bool(torch.isfinite(sampler.sampling_entropy)) - assert bool(torch.isfinite(sampler.sampling_top1_prob)) - assert bool(torch.isfinite(sampler.sampling_top1_bin)) + assert sampler.sampling_top1_prob.item() == pytest.approx(0.5625, abs=1.0e-7) + assert sampler.sampling_top1_bin.item() == 0.0 def test_mixed_metrics_publish_expected_values_across_samples() -> None: