[Cudrix-dev]

Important

Hard dependency: unisim#533 and a unisim release that includes it. unilab.base.torch_env, unilab.base.autoreset_graph and unilab.envs.manager_based_rl_env import unisim.backend.graph_capture at module level, with no fallback. CI installs the released unisim-core==2.7.1, which does not have this module, so every CI job fails at import unilab until unisim#533 is merged and released and the unisim pin is bumped here. Please re-run CI after that release; until then the failures are expected.

Summary

Adds an opt-in, fixed-shape step path for TorchEnv / ManagerBasedRlEnv. It removes the per-step host synchronizations of the done-row reset and can capture the reset segment as one CUDA graph. All new flags default to False, so tasks that do not opt in are unaffected.

Fixed-shape reset selection

ResetSelection (unilab/base/reset_selection.py) is a persistent full-width bool mask passed through the existing env_ids reset APIs.

• Participants opt in with supports_reset_selection = True or @reset_selection_capable.
• Helpers fill_rows_ and assign_rows_ cover the common masked writes.
• The env negotiates support at init_state and names every unsupported participant (reset_selection_blockers()).
• A selection is not an index, so unconverted code fails closed.
• A declaration counts only at or below the class that defines the effective reset(). A subclass that overrides reset without converting it is reported as a blocker; it does not silently inherit the parent's opt-in.

Env flags

• masked_autoreset:
  • the done-row reset runs every step, with no .any() / .nonzero() gate;
  • final_observation is published every step and is valid on done rows;
  • requires defer_finite_validation.
• Deferred fail-closed checks:
  • TorchEnv.require(ok, name, message) latches a named device verdict;
  • reset-state mask checks and backend tensor validation join the single poll_finite_validation host read.
• persistent_step_buffers:
  • the working obs/reward/done carriers are updated in place;
  • a publication ring (published_buffer_depth, default 2) keeps a consumer's held observation valid across the next step.
• autoreset_cuda_graph, with autoreset_graph_warmup and autoreset_graph_capture_mode:
  • AutoresetGraph captures the whole masked reset body (managers plus backend reset) into one torch.cuda.graph, inside the backend's external capture scope;
  • it registers the manager RNG generator;
  • recorded host transitions (unilab/base/host_transitions.py) replay after each launch;
  • there is no eager epilogue.

Managers and buffers

• Command, event, curriculum, metrics, observation, recorder, reward, termination and reset-owner managers accept a selection.
• CircularBuffer and DelayBuffer reset through the mask.
• Class-level defaults for the fixed-shape state let constructor-free (__new__) test shells keep the ordinary path.

G1 motion tracking opts in

• MotionCommand:
  • under a selection, or with masked_autoreset, always dispatches its full-width sampling instead of probing the mask;
  • the robot-cache step stamp is a declared host transition.
• MotionJointPositionAction, JointPositionAction, the motion body reward terms and the anchor/body height terminations declare support.
• MotionSampler's top-1 probability uses amax; indexing with the device argmax caused a host read.

Docs

New user guide page docs/sphinx/source/en/2-user_guide/1-training/8-sync_free_autoreset.md (linked from the training index):

• the flags and their requirements, with a ++env.* CLI example;
• behavior notes: RNG draw order, deferred validation polling, and reset diagnostics that read 0 or update only at the poll;
• how to opt custom terms into ResetSelection and read reset_selection_blockers.

Results (G1MotionTracking, MJWarp, 4096 envs, RTX 6000D)

mode	ms / env.step	host syncs in the reset segment
selected-row (current)	12.9	yes
masked_autoreset	10.5	0
masked_autoreset + persistent_step_buffers + autoreset_cuda_graph	10.4	0 (graph replayed, 5 host transitions)

• Reward and done statistics match the selected-row path within seed noise. MJWarp is not bitwise deterministic, so equivalence was checked statistically, not value by value.
• A downstream task with heavier reset work (SONIC, 512 envs): 12.8 → 7.6 ms/step, host syncs per step 58 → 0.1. The remaining 0.1 is the every-32-steps validation poll.

Tests

• tests/envs/test_fixed_shape_autoreset.py: negotiation, the MRO opt-in rule, masked reset equivalence, publication ring lifetime, deferred verdicts, and graph capture/replay with host transitions.
• tests/tasks/test_g1_motion_fixed_shape.py, including a slow real-MJWarp G1 graph test.
• CI-equivalent run with unisim#533 installed (Linux, Python 3.11, CPU torch 2.10.0, locked deps):
  • ruff check, ruff format, mypy and pyright are clean;
  • test-rl, test-environments and test-applications pass, as do the benchmark smoke and the docs checks (tests/scripts/test_check_docs.py, sphinx-build -n);
  • test-core has one failure, test_unisim_dependency_uses_an_approved_source. It occurs only because the unreleased unisim was installed from a local, non-editable copy, and it goes away once unisim is installed from the registry.