Run all Newton actuators through one flat-layout adapter path#6616
Draft
ooctipus wants to merge 8 commits into
Draft
Run all Newton actuators through one flat-layout adapter path#6616ooctipus wants to merge 8 commits into
ooctipus wants to merge 8 commits into
Conversation
Collaborator
Author
|
Rework in progress before review: the legacy Lab actuator compute path on the Newton backend will be removed in this PR (the adapter becomes the only actuation path on Newton; PhysX keeps its existing default path unchanged), and the test suite will be rewritten to assert current behavior directly (tracking, clamping, decimation, reset isolation, golden captures) instead of A/B-comparing against the removed path. Keeping as draft until then. |
NewtonManager ran two parallel step programs (_simulate_full, _simulate_physics_only) and a capture decision forked three ways (init-time wp.ScopedCapture, RTX deferral, Newton-actuator skip). The skip arm relied on set_decimation to capture later, so callers that never set a decimation - plain SimulationContext loops - ran eager forever (measured 5.06 ms vs 0.91 ms per tick for the same workload). Collapse both forks: - _simulate() is the only step program, composed of _iteration_stages (collide -> actuators -> post-actuator hooks -> substeps) and _tail_stages (post-step hooks -> sensors). It runs the decimation loop when the manager owns it and exactly one sub-step otherwise; mixed scenes step the adapter eagerly and replay a per-iteration graph (_iteration_graph) inside the loop. - handles_decimation() now follows the Newton actuator path (loop ownership), no longer aliased to graph-safety, so mixed scenes fold the decimation loop into one step() call as well. - Decision sites (initialize_solver, set_decimation, hard reset) only call _invalidate_graph(): drop the graphs, owe a re-capture. The sole capture site is step()'s pending-consume block, and the one capture routine _capture_relaxed_graph(device, warmup, capture) works both headless and with RTX active. Its warmup runs _simulate eagerly and counts as the capturing tick's step: the execute branch is skipped and sim_time advances at a single site, so the first step never double-advances. The Kamino replay-once-to-pin lives inside the capture routine instead of being duplicated at two sites. The never-captured-scenes bug is fixed as a corollary: decision sites only invalidate and the first step captures, on every path. Exactly one capture per invalidation also makes the back-to-back capture CUDA faults of the contact + actuator pipeline structurally impossible. The regression test (test_deferred_graph_capture.py) was verified to fail against the original skip-instead-of-defer behavior.
Replace the adapter's env-0-stride assumption (managed DOFs derived from the first env slice with a constant per-env stride) with a flat DOF layout: the adapter takes the model's total DOF count and a DOF-to-env table, and per-articulation binding supplies a flat dof_index_map built from the articulation's own DOF layout. Homogeneous and heterogeneous scenes construct the adapter identically. Gains, managed masks, and computed-effort telemetry move to flat per-DOF arrays with scatter/gather kernels; reset resolves env ids through the DOF-to-env table.
Freeze the adapter's engine-channel outputs (applied and computed effort) on a fixed heterogeneous ANYmal+Cartpole scenario so adapter refactors that keep the same Newton kernels are refereed bitwise, independently of the Lab actuator models. Each recorded tick starts from an explicitly written joint state: the adapter writes joint_f before the solver integrates, so single-tick captures are bitwise-stable even though multi-step trajectories are not (MuJoCo-Warp's FP-atomic solve is not run-to-run deterministic). The second tick follows a partial reset to pin the reset-mask path. Tolerance-based Lab-vs-Newton equivalence continues to live in test_newton_actuators_newton.py; together they form the parity suite that later adapter changes cite.
Delete the stored torch mirrors of the flat gain and managed-DOF snapshots (bind-window views are created where needed: a held torch twin of a warp buffer is a stale-alias hazard), preallocate the reset scratch masks so resets allocate nothing, and unify the three fill-shaped scatter kernels (zero_at_indices, set_mask, mark_flat) into one generic fill_at_indices_kernel with typed overloads. Verified bitwise against the adapter golden captures: same kernels, identical outputs; the actuator parity suite passes unchanged.
A scene reset chain called the model-global adapter reset once per articulation with the same env set — 22 identical whole-model resets per event. Actuator-state resets are idempotent, so an identical repeat within one reset event (same sim time, same env set) is now a no-op. Every actuated DOF must have exactly one writer for the effort scatter to be order-free: overlapping actuator index sets now fail at build with a clear error instead of silently corrupting efforts (previously overlapping claims were undetected, and unclaimed joints could go dead without a diagnostic).
The Newton backend now always steps actuators through the physics engine: NewtonActuator USD authoring runs whenever the active backend is Newton, and the articulation unconditionally takes the Newton actuator route. The per-step Lab actuator compute path (_apply_actuator_model) on Newton is deleted, together with the SimulationCfg.use_newton_actuators flag and every consultation site. On PhysX the flag-gated, opt-in Newton-actuator option is deleted so PhysX always uses its existing Isaac Lab actuator models — its default behavior is unchanged. The PhysX-only adapter plumbing (_apply_actuator_model_newton, PhysxActuatorWrapper allocation, the write_actuator_*_to_sim DR writes) goes with it, as does the PhysX A/B test suite that exercised the removed option. NewtonManager drops the _use_newton_actuators_active plumbing: whether the decimation loop can be folded into one CUDA graph (handles_decimation/_is_all_graphable) is derived from actuator-adapter presence, and the start-time graph capture is skipped exactly when set_decimation would re-capture.
With the legacy Lab actuator path removed there is no second implementation to A/B against, so every Lab-vs-Newton equivalence class becomes a direct assertion of current behavior: - PD tracking: with gravity disabled the commanded pose is reached within tolerance and held; computed_torque follows the exact PD law reconstructed from the recorded pre-step state. - Clamp semantics: applied_torque is checked against the exact max-effort box, DC-motor speed-torque curve (saturation provably engages), and position-based lookup clamp formulas. - Delay semantics: a mid-run target switch is invisible for exactly the authored delay_steps and takes effect then; the active target is recovered by inverting the PD law on telemetry. - Implicit drives: all-implicit scenes build no adapter and still track; telemetry matches the drive-law shadow; a feedforward shifts the zero-gravity equilibrium by exactly FF/kp. - Decimation: one folded manager step covering d sub-steps must match d explicit single sub-steps of the same path (self-consistency). - Heterogeneous scenes: each robot's torques must follow its own PD law, which pins the adapter's per-articulation DOF mapping. Ordering, reset isolation, DR gain writes, gain snapshots, and USD authoring tests are kept, now flag-free. The Lab-path reset-isolation test is deleted with the path it inspected.
Spawns many ANYmal-D copies per environment, cycling the analytic actuator models (DC motor, ideal PD, solver-implicit PD), and measures steady-state per-step time on Newton MJWarp. Actuator dispatch cost scales with the number of articulation groups rather than environments, which single-robot benchmarks cannot expose. Every environment has an identical composition, so the scene needs no multi-world grouping and the same command reproduces the baseline on any branch.
ooctipus
force-pushed
the
zhengyuz/actuator-adapter-flat-layout
branch
from
July 20, 2026 08:05
2f1df28 to
a009c7e
Compare
ooctipus
changed the title
Collaborator
Author
|
Rebased onto #6614 — the dependency is functional: without its invalidate-then-capture policy, removing the actuator flag would leave plain 🤖 Generated with Claude Code |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Newton actuation previously ran as two parallel code paths selected by
SimulationCfg.use_newton_actuators: the default path computed Lab actuator modelseagerly on the host once per articulation per step, and the opt-in fast path ran Newton
actuators inside the physics step but assumed every environment repeats env-0's DOF
rectangle, so it broke on any scene whose articulations do not tile uniformly.
This PR makes the fast path the only path and removes the layout assumption:
joint_world/joint_qd_starttables instead of anenv_stride × num_envsrectangle. Per-articulationstate is assembled by pure gathers through each articulation's flat DOF index map
(including the view layout's global
offset— previously only articulations placedfirst in their world resolved correctly), so homogeneous and heterogeneous scenes
construct identically.
SimulationCfg.use_newton_actuatorsis deleted. Newton articulationsalways bind their explicit actuators (ideal PD, DC motor, delayed PD, remotized PD,
actuator nets) to Newton actuators that run inside the physics step — inside the CUDA
graph when every actuator is graph-safe, stepped eagerly per iteration otherwise.
_is_all_graphableis derived from the adapter alone, and the manager always owns thedecimation loop; the activation flag state is gone.
_apply_actuator_modeloverride), the flag plumbing inschemas_actuators, and thePhysX-vs-Newton A/B flag test are removed.
BaseArticulation._apply_actuator_modelis aconcrete no-op hook: backends that compute actuator models on the host (PhysX) override
it; backends whose actuators run inside the physics step (Newton) do not.
Dependency
This branch is rebased over #6614 and must merge after it. The dependency is
functional, not just textual: without #6614's invalidate-then-capture policy, removing the
flag makes every explicit-actuator scene take the old "skip initial capture, recapture on
set_decimation" arm — and plainSimulationContextloops (demos, standalone scripts,benchmarks) never call
set_decimation, so they run eager forever (measured 7.0 ms/stepvs 2.9 ms baseline on the benchmark below). On #6614's manager, decision sites only
invalidate and the first
step()captures, on every path.Performance
Actuator dispatch cost scales with the number of articulation groups, not environments —
single-robot benchmarks cannot show it. This PR adds
scripts/benchmarks/benchmark_actuator_dispatch.py: N ANYmal-D per environment cycling DCmotor, ideal PD, and solver-implicit PD actuator configs, stepped steady-state on Newton
MJWarp. Every environment is identical, so the same command runs on any branch without
multi-world grouping.
./isaaclab.sh -p scripts/benchmarks/benchmark_actuator_dispatch.py --headless \ --num_envs 1024 --num_articulations 8Baseline =
develop+ the benchmark commit cherry-picked; identical commands, singleRTX 5090, per-step time over 500 steps after 100 warmup steps, median of 3 interleaved
runs per configuration:
The #6614-only column sits at develop parity, so the gain is this PR's in-graph actuators
rather than the capture-policy change. The win tracks the ratio of per-articulation host
dispatch to GPU solver time: largest at low and mid environment counts, smaller once
per-world solver cost dominates (in this benchmark, adding robots per environment also
grows per-world solver work, which is why 16 robots/world gains less than 8).
Tests
test_newton_actuators_newton.pyis rewritten as behavior checks: 73 passed, 3 xfailed.DC-motor speed-torque clamp, delay semantics, implicit drive telemetry — plus
joint-ordering permutation cases and a heterogeneous ANYmal + cartpole scene whose
per-robot laws catch cross-articulation DOF-map clobbering (this suite caught the view
offsetbug fixed in this PR).test_actuator_adapter_golden.py): single-tick bitwise captures fromexplicitly written states; machine-generated on first run, not committed.
expectedFailure, documenting an upstream MJWarpdefect this PR does not change: implicit joint drives are baked into the MuJoCo actuator
set at model finalize from USD-authored gains, so assets authored with zero drive gains
(e.g. ANYmal) get no MuJoCo actuator, and gains written at initialization land in
model.joint_target_kewhere the solver never reads them back. Reproduced identicallyon
develop.Breaking change & migration
SimulationCfg.use_newton_actuatorsis removed. Newton scenes always use the Newtonactuator path; delete the field from configs. PhysX behavior is unchanged.