[Newton] Backend-agnostic task-space accessors for IK/OSC#5400
[Newton] Backend-agnostic task-space accessors for IK/OSC#5400hujc7 wants to merge 5 commits intoisaac-sim:developfrom
Conversation
Task-space controllers (IK, OSC, RMPFlow) previously called PhysX-only
methods directly on `asset.root_view` — `get_jacobians`,
`get_generalized_mass_matrices`, `get_gravity_compensation_forces`. On
Newton these crashed with AttributeError because the corresponding
view (`newton.selection.ArticulationView`) exposes a different API
(`eval_jacobian`, `eval_mass_matrix`, no gravity-compensation
primitive).
Add three abstract methods on `BaseArticulation`:
* `get_jacobians()` — per-env spatial Jacobians.
* `get_mass_matrix()` — per-env joint-space mass matrix.
* `get_gravity_compensation_forces()` — per-env joint torques to hold
against gravity.
PhysX implements all three as thin passthroughs. Newton wraps
`eval_jacobian` and `eval_mass_matrix`, with all source / output /
scratch buffers pre-allocated in `_create_buffers` (CUDA-graph-capture
safe). Two Warp gather kernels copy the rows owned by each
ArticulationView out of Newton's model-sized output, applying a
`link_offset` of 1 for fixed-base articulations to drop Newton's
fixed-root row and match the PhysX shape contract. Newton's gravity
stub raises NotImplementedError; ovphysx mirrors with stubs since it
has no view-level access at all.
Migrate every direct caller to the new accessors:
* Action terms: DifferentialInverseKinematicsAction,
OperationalSpaceControllerAction, RMPFlowTaskSpaceAction,
PinkTaskSpaceAction.
* Direct-workflow envs: factory, automate (assembly + disassembly),
deploy events.
* Tests: test_differential_ik, test_operational_space.
Gate OSC's per-step fetches in `_compute_dynamic_quantities`: mass
matrix only when `inertial_dynamics_decoupling` or
`nullspace_control != "none"`; gravity comp only when
`gravity_compensation`. Previously both were fetched unconditionally,
which (a) crashed on Newton even when the values were unused and
(b) was wasted work on PhysX.
Remove the `self.sim.physics = PhysxCfg(...)` workaround from
`Isaac-Reach-Franka-{IK-Abs,IK-Rel,OSC}-v0`. The previous comment
*"{IK,OSC} control is not supported with Newton physics; use PhysX
only"* no longer holds. All three tasks now inherit the
`ReachPhysicsCfg` preset, so `presets=newton` selects `NewtonCfg` and
the env runs under Newton without further overrides. The default OSC
config keeps `gravity_compensation=False` so the Newton path doesn't
hit the upstream-blocked gravity primitive.
Pin the upstream Newton gap with a regression test
(`test_get_gravity_compensation_forces_not_implemented_on_newton`)
that asserts `NotImplementedError` is raised. When upstream Newton
ships an `eval_gravity_compensation` API and the wrapper switches to
a real implementation, the test fails with `DID NOT RAISE` — directing
the maintainer to remove the stub and the OSC config-time guidance.
Tighten the floating-base OSC indexing regression test by enabling
both `inertial_dynamics_decoupling=True` and
`gravity_compensation=True` in its OSCCfg, so the new flag-gating
doesn't short-circuit the indexing checks the test was written to
verify.
Validation:
* PhysX: test_differential_ik 2/2, test_operational_space 18/18
(including the floating-base regression).
* Newton: Isaac-Reach-Franka-{IK-Abs,IK-Rel,OSC}-v0 random_agent
smokes ran zero-error through thousands of physics substeps.
* Codex review pass (4 P1/P2 findings — all addressed).
![isaaclab-review-bot[bot]](https://avatars.githubusercontent.com/u/157846462?s=60&v=4){kind=small}
There was a problem hiding this comment.
🤖 Isaac Lab Review Bot
Summary
This PR introduces backend-agnostic task-space accessors (get_jacobians, get_mass_matrix, get_gravity_compensation_forces) on BaseArticulation to enable IK/OSC/RMPFlow controllers to work under both PhysX and Newton backends. The implementation is well-designed with proper CUDA graph capture safety, pre-allocated buffers, and correct gating of dynamic quantity fetches. The core approach is sound, though there are a few issues worth addressing.
Architecture Impact
High impact, well-contained. The changes span the articulation abstraction layer (BaseArticulation), all three backend implementations (PhysX, Newton, ovphysx), the task-space action terms, and downstream task configs. The Newton implementation introduces new gather kernels that extract view-specific rows from model-sized buffers. The OSC action term now conditionally fetches mass matrix and gravity compensation based on controller config flags, which is the correct pattern for backend compatibility.
Implementation Verdict
Minor fixes needed — One potential correctness issue with the mass matrix gating logic and a few smaller concerns.
Test Coverage
- ✅
test_get_gravity_compensation_forces_not_implemented_on_newtonpins the known Newton gap - ✅
test_operational_space.pyupdated to enablegravity_compensation=Truefor floating-base test ⚠️ Missing: No direct unit test forget_jacobians()orget_mass_matrix()on Newton verifying shape/value correctness against PhysX⚠️ Missing: No test exercising OSC withinertial_dynamics_decoupling=Falseandnullspace_control="position"to verify the new gating logic
CI Status
No CI checks available yet — recommend waiting for CI before merge.
Findings
🔴 Critical: source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py:659-663 — Mass matrix gating may skip required data for nullspace control
The condition needs_mass_matrix = self.cfg.controller_cfg.inertial_dynamics_decoupling or (self.cfg.controller_cfg.nullspace_control != "none") is correct, but the _mass_matrix tensor is initialized to zeros in __init__ (line 409) and never cleared between steps. If a user toggles config at runtime (unlikely but possible), or if the controller internally assumes the mass matrix is always fresh, stale zeros could cause issues. More importantly, if needs_mass_matrix is False, _mass_matrix remains zeros and gets passed to self._osc.compute() at line 538. The OSC controller's compute() method receives mass_matrix unconditionally — verify the controller handles zero mass matrices gracefully when neither flag is set.
# Line 538 always passes self._mass_matrix regardless of gating
self._joint_efforts[:] = self._osc.compute(
...
mass_matrix=self._mass_matrix, # Could be zeros if not fetched
...
)🟡 Warning: source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py:226-240 — Redundant jacobian evaluation in get_mass_matrix
When get_jacobians() and get_mass_matrix() are called in the same step (common for OSC), eval_jacobian is called twice on the same state. The comment acknowledges this ("any earlier get_jacobians this step is overwritten with identical data"), but this is still a performance overhead. Consider caching the jacobian evaluation per-state or adding a combined accessor.
🟡 Warning: source/isaaclab_newton/isaaclab_newton/assets/articulation/kernels.py:581 — gather_jacobian_rows uses int instead of wp.int32
def gather_jacobian_rows(
...
link_offset: int, # Should be wp.int32 for Warp kernel type consistency
...
):While Warp may implicitly handle Python int, using wp.int32 would be more consistent with other kernel signatures in the codebase.
🟡 Warning: source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py:3320-3325 — Jacobian view buffer shape uses max_links from model
The Newton implementation uses model.max_joints_per_articulation for max_links, but the comment at line 3320 says "Source buffer covers every articulation in the model." For scenes with multiple articulation types (e.g., Franka + mobile base), max_links and max_dofs are model-wide maximums. This is correct for the source buffer, but the view buffer at line 3340 should ideally use this view's actual link/dof count rather than model maximums to avoid over-allocation. This isn't a bug (the gather kernel respects bounds), but wastes memory.
🔵 Improvement: source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py:654-668 — Consider caching needs_mass_matrix flag
The needs_mass_matrix computation happens every step in _compute_dynamic_quantities(). Since self.cfg.controller_cfg is immutable after init, this boolean could be computed once in __init__ and stored as self._needs_mass_matrix and self._needs_gravity.
🔵 Improvement: source/isaaclab_newton/test/assets/test_articulation.py:2604-2640 — Test should also verify jacobian/mass_matrix shapes
The new test only checks get_gravity_compensation_forces raises NotImplementedError. Consider adding companion tests that verify get_jacobians() and get_mass_matrix() return tensors with the expected shapes (num_instances, num_jacobi_bodies, 6, num_jacobi_joints) and (num_instances, num_jacobi_joints, num_jacobi_joints) respectively, matching the PhysX contract.
🔵 Improvement: source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py:182-190 — PhysX method naming inconsistency
The base class defines get_mass_matrix() but PhysX's underlying view method is get_generalized_mass_matrices() (plural). The passthrough is correct, but a brief comment noting this naming difference would help maintainability:
def get_mass_matrix(self) -> wp.array:
# PhysX view uses plural "matrices" naming convention
return self._root_view.get_generalized_mass_matrices()
|
@greptileai review |
Greptile SummaryThis PR introduces three abstract methods on Confidence Score: 5/5Safe to merge; only P2 documentation and naming nits found, no correctness or runtime issues. All changes are well-tested (shape contracts for both fixed/floating-base and heterogeneous scenes, xfail test pinning the Newton gap, smoke runs on Newton). The fetch-gating logic is correct, the kernel gather patterns are sound, and the PhysX passthroughs are trivial. Remaining findings are a docstring shape mismatch and a confusing variable name — neither affects runtime behavior. source/isaaclab_newton/isaaclab_newton/assets/articulation/kernels.py (docstring), source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py (max_links naming), source/isaaclab/isaaclab/envs/mdp/actions/pink_task_space_actions.py (PhysX +6 offset comment) Important Files Changed
Class Diagram%%{init: {'theme': 'neutral'}}%%
classDiagram
class BaseArticulation {
<<abstract>>
+get_jacobians() wp.array
+get_mass_matrix() wp.array
+get_gravity_compensation_forces() wp.array
}
class PhysXArticulation {
+get_jacobians() wp.array
+get_mass_matrix() wp.array
+get_gravity_compensation_forces() wp.array
}
class NewtonArticulation {
-_jacobian_buf_flat wp.array
-_jacobian_buf wp.array4d
-_mass_matrix_buf wp.array
-_jacobian_view_buf wp.array4d
-_mass_matrix_view_buf wp.array
-_jacobian_view_art_ids wp.array
+get_jacobians() wp.array
+get_mass_matrix() wp.array
+get_gravity_compensation_forces() NotImplementedError
}
class OvPhysXArticulation {
+get_jacobians() NotImplementedError
+get_mass_matrix() NotImplementedError
+get_gravity_compensation_forces() NotImplementedError
}
class OSCAction {
-_needs_mass_matrix bool
-_needs_gravity bool
+_compute_dynamic_quantities()
}
BaseArticulation <|-- PhysXArticulation
BaseArticulation <|-- NewtonArticulation
BaseArticulation <|-- OvPhysXArticulation
OSCAction --> BaseArticulation : get_jacobians / get_mass_matrix / get_gravity_compensation_forces
Reviews (4): Last reviewed commit: "Add shape-contract regression tests for ..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| if self._asset.is_fixed_base: | ||
| gravity = torch.zeros_like( | ||
| wp.to_torch(self._asset.root_view.get_gravity_compensation_forces())[ | ||
| :, self._controlled_joint_ids_tensor | ||
| ] | ||
| wp.to_torch(self._asset.get_gravity_compensation_forces())[:, self._controlled_joint_ids_tensor] | ||
| ) |
There was a problem hiding this comment.
get_gravity_compensation_forces() call on fixed-base
_apply_gravity_compensation calls get_gravity_compensation_forces() on a fixed-base robot solely to get the tensor shape for torch.zeros_like. On Newton, this immediately raises NotImplementedError, so any fixed-base robot using PinkInverseKinematicsAction with gravity not disabled (rigid_props.disable_gravity=False) will crash at runtime — even though the intended result is all zeros and the underlying value is never used.
| if self._asset.is_fixed_base: | |
| gravity = torch.zeros_like( | |
| wp.to_torch(self._asset.root_view.get_gravity_compensation_forces())[ | |
| :, self._controlled_joint_ids_tensor | |
| ] | |
| wp.to_torch(self._asset.get_gravity_compensation_forces())[:, self._controlled_joint_ids_tensor] | |
| ) | |
| if self._asset.is_fixed_base: | |
| gravity = torch.zeros( | |
| (self._asset.num_instances, len(self._controlled_joint_ids_tensor)), | |
| device=self.device, | |
| ) |
| gravity = wp.to_torch(self._asset.get_gravity_compensation_forces())[ | ||
| :, self._controlled_joint_ids_tensor + self._physx_floating_joint_indices_offset | ||
| ] |
There was a problem hiding this comment.
_physx_floating_joint_indices_offset embeds a PhysX convention (6 virtual free-floating DOFs prepended to gravity-comp output). When Newton eventually implements get_gravity_compensation_forces(), its output likely won't include those 6 virtual DOFs — so applying this offset here would silently read the wrong joints. Consider adding a comment or an assertion that ties this indexing to the PhysX backend, so a future implementor knows to revisit it.
* pink_task_space_actions.py: fixed-base gravity-comp path now constructs the zeros tensor directly. Previously called get_gravity_compensation_forces() solely to feed torch.zeros_like for shape, which crashes on Newton. * task_space_actions.py: cache _needs_mass_matrix / _needs_gravity at init (controller cfg is immutable). Pass None into OSC.compute() when those flags are False instead of forwarding the stale-zero buffer; the controller's existing None-checks then raise immediately on any misconfiguration rather than silently operating on zeros. * Newton gather_jacobian_rows: type link_offset as wp.int32 to match the kernel-signature convention used elsewhere in this module. * test_get_gravity_compensation_forces_not_implemented_on_newton: switch from pytest.raises to pytest.mark.xfail(strict=True, raises=...). Same contract (test fails when Newton ships the primitive), cleaner CI semantics — pytest reports XFAIL/XPASS instead of looking like a real failure when the gap is closed. * PhysX passthrough: comment noting that get_generalized_mass_matrices (plural) is PhysX's view-side name; the singular IsaacLab API (get_mass_matrix) hides the difference.
|
@greptileai review |
* Newton get_jacobians / get_mass_matrix output buffers were sized using ``model.max_joints_per_articulation`` and ``model.max_dofs_per_articulation`` — the model-wide maxima. In a homogeneous scene this matched per-asset counts, but in a heterogeneous scene where another articulation type has more joints/DoFs, this view's returned tensor would carry zero-padded rows/cols that don't belong to it. The padded mass-matrix was singular under torch.inverse. Output buffers now use ``self.num_bodies`` / ``self.num_joints`` so the returned shape always matches THIS view's asset, restoring the cross-backend contract. Source and scratch buffers stay model-sized — the kernel needs that to match Newton's writes. * Bump isaaclab_ovphysx 0.1.1 -> 0.1.2 with a CHANGELOG entry covering the new BaseArticulation.get_* stubs added in the previous commit. Per the per-extension changelog/version rule.
|
@greptileai review |
Pin the public shape contract for ``get_jacobians`` and
``get_mass_matrix`` so future regressions fail at test time instead of
in production.
* Newton: 5 new tests in test_articulation.py
- test_get_jacobians_shape_fixed_base — Franka fixed-base, locks
(N, num_bodies-1, 6, num_joints) and the link_offset=1 fix.
- test_get_mass_matrix_shape_and_nonsingular_fixed_base — Franka
fixed-base, locks (N, num_joints, num_joints) and that the matrix
has positive diagonals (a padded zero row would surface here).
- test_get_jacobians_shape_floating_base — Anymal-C floating-base,
locks (N, num_bodies, 6, num_joints) and link_offset=0 path.
- test_get_mass_matrix_shape_floating_base — Anymal-C, mass-matrix
shape on floating-base.
- test_heterogeneous_scene_per_view_shapes — co-resident Franka +
Anymal-C in one Newton model. Direct regression test for the
Codex round-2 finding: each view's output uses its own
num_bodies/num_joints, NOT model.max_*. ``num_per_type=1`` keeps
the test off Newton's pre-existing actuator-default replication
quirk in multi-instance multi-type scenes.
* PhysX: 3 new tests in test_articulation.py
- Mirror the Franka fixed-base + Anymal floating-base shape tests
so the cross-backend contract is locked at both ends. The
heterogeneous case is Newton-specific (PhysX views are
view-scoped at the engine level), so no PhysX-side equivalent.
Mass-matrix non-singularity check uses positive-diagonal rather than
``torch.linalg.det()``: a real Franka mass matrix has det~1e-13 by
numerical conditioning even when SPD, so the determinant threshold
would false-fail. Padded zero rows (the bug we want to catch) show up
clearly as zero diagonals.
|
@greptileai review |
* BaseArticulation.get_jacobians / get_mass_matrix / get_gravity_compensation_forces are now concrete with ``raise NotImplementedError`` bodies instead of ``@abstractmethod``. Out-of-tree subclasses no longer break at instantiation when these accessors are added — they fail only when the methods are actually invoked, matching AGENTS.md's deprecation policy. The in-tree PhysX / Newton / ovphysx subclasses still override explicitly. * BaseArticulation.get_jacobians docstring documents the floating-base joint-dim convention difference between PhysX (prepends 6 virtual DoFs to the joint dim, ``num_joints`` counts only actuated) and Newton (``ArticulationView.joint_dof_count`` already includes the base 6, so ``num_joints`` is the total). The total joint-dim is identical on both backends; only how ``num_joints`` is reported differs. Floating-base task-space callers should verify their joint indexing against the active backend's DoF ordering. No floating-base IK/OSC env exists today so no caller is affected.
1 participant
1. Summary
Make IK / OSC / RMPFlow task-space controllers backend-agnostic so Franka manipulation envs run under Newton.
The action terms previously called PhysX-only methods on
asset.root_viewdirectly, which crashed under Newton withAttributeError. The Franka reach envs worked around this by hardcodingself.sim.physics = PhysxCfg(...)with the comment "{IK,OSC} control is not supported with Newton physics; use PhysX only". This PR removes that workaround.End-to-end result:
Isaac-Reach-Franka-{IK-Abs,IK-Rel,OSC}-v0run on either backend, picked viapresets=newton.2. Design
Three new abstract methods on
BaseArticulation:get_jacobians()—(N, num_jacobi_bodies, 6, num_jacobi_joints)get_mass_matrix()—(N, num_jacobi_joints, num_jacobi_joints)get_gravity_compensation_forces()—(N, num_jacobi_joints)Per-backend implementations:
isaaclab_physxisaaclab_newtoneval_jacobian+ gather kerneleval_jacobian+eval_mass_matrix+ gather kernelNotImplementedError(no upstream primitive)isaaclab_ovphysxNotImplementedError(stubs to keep class concrete)NotImplementedErrorNotImplementedError3. Newton-side details
3.1 Pre-allocation for capture safety
Every buffer used at step time is pre-allocated in
_create_buffers: model-sized source buffers (_jacobian_buf_flat,_mass_matrix_buf), view-sized output buffers (_jacobian_view_buf,_mass_matrix_view_buf), kernel scratch (_jacobian_joint_S_s_buf,_mass_matrix_body_I_s_buf), and the flattened view-to-model articulation index map (_jacobian_view_art_ids). This keepsget_jacobians/get_mass_matrixallocation-free and safe under CUDA graph capture.3.2 View-level row gather
Newton's
eval_jacobian/eval_mass_matrixwrite every articulation in the model into a single buffer (shape includesmodel.articulation_count), regardless of whichArticulationViewinvoked them. PhysX returns view-scoped data already. Two Warp kernels (gather_jacobian_rows4-D,gather_mass_matrix_rows3-D) gather just this view's rows into a contiguous view-sized destination so the caller-facing shape contract matches PhysX. The view-to-model index map is reused across both gathers.3.3 Fixed-base row offset
Newton's
eval_jacobianincludes the fixed-root joint as row 0 (with zero motion subspace). PhysX's contract drops that row and reindexes bodies by −1 (_jacobi_body_idx = body_idx − 1). The Newton gather kernel takes alink_offset(1 for fixed-base, 0 for floating-base) so the output skips the fixed-root row and matches PhysX's body indexing.3.4
eval_mass_matrix's "J as input" gotchaNewton's
eval_mass_matrix(state, H, J=None, body_I_s=None, joint_S_s=None)treatsJas input when provided — it skips the internaleval_jacobianand uses the buffer as-is. Passing an empty pre-allocatedJproducesH = J^T·M·J = 0→ singular →LinAlgErrorin OSC'storch.inverse(mass_matrix). The wrapper explicitly populatesJby callingeval_jacobianfirst; we reuse_jacobian_buf_flat(same shape) so no separate scratch is needed.4. Action-term gating
OperationalSpaceControllerAction._compute_dynamic_quantitiespreviously fetched mass matrix and gravity-compensation forces unconditionally on every step. Under the new abstraction this would still call Newton's gravity-comp stub even when the user disabled gravity compensation in the controller config.The fetches are now gated to match what the controller actually consumes:
inertial_dynamics_decoupling=Trueornullspace_control != "none"(the null-space torque term inOperationalSpaceController.compute()consumes mass matrix independently of inertial decoupling).gravity_compensation=True.Side benefit: skips a per-step engine call on PhysX when neither flag is set.
5. Known limitation: gravity compensation on Newton
Newton's
ArticulationViewexposeseval_fk/eval_jacobian/eval_mass_matrixonly — no gravity-compensation primitive.get_gravity_compensation_forces()raisesNotImplementedErroron Newton with a message pointing at the upstream gap. OSC users on Newton must setgravity_compensation=False(the default in the currentIsaac-Reach-Franka-OSC-v0config); other flag combinations work normally.The gap is pinned by
test_get_gravity_compensation_forces_not_implemented_on_newton, which is deliberately written to fail when upstream Newton closes the gap (pytest.raisesreportsDID NOT RAISE). The failure points the maintainer at the wrapper, the test, and the OSC guidance to remove.6. Reach-env cfg cleanup
Removed the
self.sim.physics = PhysxCfg(bounce_threshold_velocity=0.2)override fromIsaac-Reach-Franka-{IK-Abs,IK-Rel,OSC}-v0. Tasks now inherit the parentReachPhysicsCfgpreset, sopresets=newtonselectsNewtonCfgandpresets=physx(the default) keeps the previous behavior — samebounce_threshold_velocity=0.2lives onReachPhysicsCfg.default = PhysxCfg(bounce_threshold_velocity=0.2). No information lost.7. Test plan
test_differential_ik.pymigrated to callrobot.get_jacobians(). Passes on PhysX (2/2).test_operational_space.pymigrated to callrobot.get_jacobians(),get_mass_matrix(),get_gravity_compensation_forces(). Passes on PhysX (18/18).test_floating_base_osc_action_term_indexingcfg updated to enable bothinertial_dynamics_decoupling=Trueandgravity_compensation=True, so the new flag gating doesn't short-circuit the indexing the test was written to verify.test_get_gravity_compensation_forces_not_implemented_on_newtonadded — pins the Newton gap.random_agentsmoke onIsaac-Reach-Franka-IK-Abs-v0under Newton: ran 5,205 physics substeps zero-error.random_agentsmoke onIsaac-Reach-Franka-IK-Rel-v0under Newton: ran 306 substeps zero-error.random_agentsmoke onIsaac-Reach-Franka-OSC-v0under Newton: ran 290 substeps zero-error.8. Files
Touched four extensions, version-bumped each:
isaaclab4.6.12 → 4.6.13: abstract methods, action-term gating, doc updates.isaaclab_physx0.5.22 → 0.5.23: three passthroughs.isaaclab_newton0.5.21 → 0.5.22: real wrappers + gather kernels + gravity stub + gap-pin test.isaaclab_tasks1.5.24 → 1.5.25: hardcode removals + direct-workflow caller migrations.26 files changed, +470 / −47.
9. Status
Draft / WIP. Open questions for reviewers:
_jacobian_view_art_idsbe onBaseArticulation(since both backends would benefit from a view→model index map for any future gathers) or stay private to Newton?link_offsetthe right contract on the gather kernel, or should the offset be baked into the destination shape and the kernel always usedst[i, l, ...] = src[id, l, ...]?