[Cudrix-dev]

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 flags default to off, so existing behavior is unchanged.

Requires the unisim backend contract torch_graph_capture_scope (unisim PR: ).

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.

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.
• Full suite: no new failures relative to the unmodified base.