Skip to content

Add generic Newton coupled-solver manager with Proxy + ADMM variants#5834

Draft
mmichelis wants to merge 23 commits into
isaac-sim:developfrom
mmichelis:mym/coupledsolver
Draft

Add generic Newton coupled-solver manager with Proxy + ADMM variants#5834
mmichelis wants to merge 23 commits into
isaac-sim:developfrom
mmichelis:mym/coupledsolver

Conversation

@mmichelis
Copy link
Copy Markdown
Collaborator

Description

Summary

  • New isaaclab_contrib.coupling submodule. One manager, NewtonCoupledSolverManager, that dispatches between two coupling algorithms by config subclass:
    • CoupledProxySolverCfgSolverCoupledProxy: lagged-impulse virtual-proxy coupling. Best when only a small set of rigid bodies (e.g. a gripper) needs to touch the deformable; the bulk of the articulation is solved purely by the rigid solver.
    • CoupledAdmmSolverCfgSolverCoupledAdmm: linearized ADMM coupling with tunable rho / gamma / baumgarte penalty knobs and joint attachment stiffness/damping.
  • Sub-solvers are generic. The pair is resolved from the cfg type, so MJWarp + VBD, Featherstone + VBD, XPBD + Kamino, etc. all work without manager changes.
  • Body selectors accept raw prim-path regex strings (e.g. "/World/envs/env_.*/Robot") alongside SceneEntityCfg — useful for assets that aren't registered as named scene entities.
  • Isaac-Lift-Soft-Franka-v0 now defaults to proxy-coupled MJWarp + VBD. Includes a breaking rename in the env's scene/MDP: deformableobject (entity, command, and all MDP functions). Update env configs, checkpoints, and RL configs accordingly.
  • Refactor: Solver-kwargs filtering deduplicated into NewtonManager._filter_solver_kwargs so subclasses (including the new coupled manager) can reuse it.

Builds on top of Newton Coupling in gdaviet/coupled-solver-framework.

How to use

Pick the coupling algorithm by choosing the config subclass; the manager handles dispatch:

from isaaclab_contrib.coupling import CoupledProxySolverCfg, CoupledAdmmSolverCfg

# Proxy coupling
solver_cfg = CoupledProxySolverCfg(
    src_solver_cfg=MJWarpSolverCfg(...),
    dst_solver_cfg=VBDSolverCfg(iterations=10),
    src_bodies=["/World/envs/env_.*/Robot"],
    dst_bodies=["/World/envs/env_.*/Object"],
    proxy_bodies=["/World/envs/env_.*/Robot/panda_hand", ...],
    proxy_collide_interval=5,
)

# ADMM coupling
solver_cfg = CoupledAdmmSolverCfg(
    src_solver_cfg=MJWarpSolverCfg(...),
    dst_solver_cfg=VBDSolverCfg(iterations=10),
    src_bodies=["/World/envs/env_.*/Robot"],
    dst_bodies=["/World/envs/env_.*/Object"],
    iterations=10, rho=1.0, gamma=0.0, baumgarte=0.0, joint_stiffness=1e4, joint_angular_stiffness=1e4,
)

Type of change

  • New feature (non-breaking change which adds functionality)
  • Documentation update

Test plan

  • ./isaaclab.sh -f (pre-commit) green
  • ./isaaclab.sh -p -m pytest source/isaaclab_contrib/test/coupling/test_coupled_manager.py — 16/16 pass
  • Isaac-Lift-Soft-Franka-v0 with the migrated newton_mjwarp_vbd_proxy preset runs stably (matches feat/proxycoupledsolver behavior)

Checklist

  • I have read and understood the contribution guidelines
  • I have run the pre-commit checks with ./isaaclab.sh --format
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • I have updated the changelog and the corresponding version in the extension's config/extension.toml file
  • I have added my name to the CONTRIBUTORS.md or my name already exists there

mmichelis added 23 commits May 28, 2026 14:47
Drop the local DeformableNewtonCfg in franka_soft_env_cfg.py and use the
shared CoupledNewtonCfg from isaaclab_contrib.deformable.newton_manager_cfg
instead, so the existing soft/cloth envs match the proxy-coupled manager.
Rename the scene entity, command, event, reward, termination, and mdp
function identifiers from 'deformable' to 'object' in the lift_franka_soft
environments and the lift_franka_soft state machine demo, so the proxy
coupled solver's SceneEntityCfg("object") selectors and the rest of the
lift task family share one asset key. Asset types (DeformableObjectCfg)
and USD prim paths are unchanged.
@github-actions github-actions Bot added documentation Improvements or additions to documentation isaac-lab Related to Isaac Lab team labels May 28, 2026
@mmichelis mmichelis marked this pull request as ready for review May 28, 2026 12:54
@mmichelis mmichelis marked this pull request as draft May 28, 2026 12:55
Copy link
Copy Markdown

@isaaclab-review-bot isaaclab-review-bot Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review: Add generic Newton coupled-solver manager with Proxy + ADMM variants

Thank you for this comprehensive addition to the Newton physics coupling framework! The design is well-structured and the implementation is clean. Below are my observations and suggestions.


✅ Strengths

  1. Excellent extensibility - The _SOLVER_CLASS_BY_CFG_TYPE registry pattern in NewtonCoupledSolverManager makes it trivial to add new solver combinations without modifying the manager.

  2. Clean dispatch pattern - Using config subclass type (CoupledProxySolverCfg vs CoupledAdmmSolverCfg) to select the coupling algorithm is elegant and type-safe.

  3. Good separation of concerns - Body selector resolution, model partitioning, and proxy selection are cleanly separated into well-named static methods.

  4. Comprehensive unit tests - The 358-line test file with fake objects is excellent for testing partitioning logic without launching Isaac Sim.

  5. Helpful error messages - The partition validation errors include specific body IDs and labels, making debugging much easier.

  6. Refactoring win - Extracting _filter_solver_kwargs into NewtonManager removes duplication across all solver managers.


🔴 Issues to Address

1. Breaking Rename Needs Migration Guide

The rename from deformable to object in lift_franka_soft is a breaking change that affects:

  • Scene entity: scene.deformablescene.object
  • Command: deformable_poseobject_pose
  • All MDP functions: deformable_*object_*

The changelog mentions this, but consider adding a migration note in the documentation or a deprecation warning for downstream users who have trained models referencing the old keys.

2. proxy_mode Should Be Validated

File: coupled_manager.py line ~161

proxies.append(
    SolverCoupledProxy.Proxy(
        ...
        mode=solver_cfg.proxy_mode,  # No validation
        ...
    )
)

A typo like "laggeed" would fail deep in Newton code. Add validation:

if solver_cfg.proxy_mode not in {"lagged", "staggered"}:
    raise ValueError(f"Invalid proxy_mode={solver_cfg.proxy_mode!r}; expected 'lagged' or 'staggered'.")

🟡 Suggestions

1. Add ADMM-Specific Tests

The unit tests thoroughly cover proxy coupling, but there are no tests for:

  • _build_admm_coupled_solver parameter forwarding
  • ContactPair construction when enable_contacts=True/False
  • The contact_distance/detection_margin None handling

2. Defensive Error for Missing scene_cfg

File: coupled_manager.py line ~87

outer_cfg = PhysicsManager._cfg
scene_cfg = outer_cfg.scene_cfg if isinstance(outer_cfg, CoupledNewtonCfg) else None

If someone uses SceneEntityCfg selectors without CoupledNewtonCfg, they get a confusing error later. Consider failing fast:

if any(isinstance(s, SceneEntityCfg) for s in solver_cfg.src_bodies + solver_cfg.dst_bodies):
    if scene_cfg is None:
        raise ValueError(
            "SceneEntityCfg selectors require physics.scene_cfg to be set "
            "(use CoupledNewtonCfg and set scene_cfg in __post_init__)."
        )

3. Consider Bounds for Numerical Parameters

Parameters like rho, gamma, proxy_mass_scale are passed through without validation. While Newton may handle invalid values, early validation with helpful error messages improves UX:

if solver_cfg.proxy_mass_scale <= 0:
    raise ValueError(f"proxy_mass_scale must be positive, got {solver_cfg.proxy_mass_scale}")

📝 Minor Observations

  1. Documentation is solid - The using-vbd-solver.rst additions explain proxy coupling well with good examples.

  2. Coupling mode validation - Good that coupled_mjwarp_vbd_manager.py and coupled_featherstone_vbd_manager.py now raise for unknown coupling modes.

  3. Static shape routing - The decision to route body == -1 shapes to destination for ground collision is correct. A brief inline comment explaining this would help future maintainers.


Verdict

This is a well-designed and well-implemented feature. The architecture is sound and the code quality is high. Addressing the proxy_mode validation and adding a migration note for the breaking rename would make this PR ready to merge.

Overall: 👍 Approve with minor suggestions

@greptile-apps
Copy link
Copy Markdown
Contributor

greptile-apps Bot commented May 28, 2026

Greptile Summary

This PR introduces a new isaaclab_contrib.coupling submodule providing NewtonCoupledSolverManager, which dispatches between two experimental Newton coupled solvers — SolverCoupledProxy (lagged-impulse virtual-proxy) and SolverCoupledAdmm (linearized ADMM) — selected by config subclass. It also refactors the shared solver-kwargs filtering into NewtonManager._filter_solver_kwargs, migrates Isaac-Lift-Soft-Franka-v0 to default to the new proxy-coupled preset, and performs a breaking rename of deformableobject throughout the lift-franka-soft MDP.

  • New coupling submodule: NewtonCoupledSolverManager partitions Newton model bodies/joints/shapes between source and destination entries, builds sub-solvers generically via _SOLVER_CLASS_BY_CFG_TYPE, and supports both raw prim-path regex strings and SceneEntityCfg body selectors; 16 pure-Python unit tests cover the partitioning and selection logic.
  • NewtonManager._filter_solver_kwargs refactor: Deduplicates the inspect-based kwargs pattern that was copy-pasted across Featherstone, MJWarp, XPBD, and VBD managers.
  • Breaking rename and default-preset change: deformableobject across scene entities, commands, and all MDP reward/observation/termination terms in both the soft-body and cloth env variants; FrankaSoftEnvCfg now defaults to the proxy-coupled MJWarp + VBD preset.

Confidence Score: 4/5

The core solver manager and partitioning logic is well-tested and the refactoring is clean; the main concern is a misleading documentation claim about automatic scene-cfg wiring that doesn't match the code.

The new coupling module is thoroughly unit-tested and the body-partitioning logic is sound. The one meaningful issue is that the VBD-solver guide states the Franka soft env automatically injects scene_cfg via PresetCfg plumbing, but FrankaSoftEnvCfg.post_init never does this and no PresetCfg mechanism does either. Users who write SceneEntityCfg-based body selectors while relying on the documented automatic wiring will get scene_cfg=None at solver-build time and a ValueError for every asset lookup. The current env config avoids this only because all selectors are raw strings. The two style findings are non-blocking.

docs/source/overview/core-concepts/physical-backends/newton/using-vbd-solver.rst — the scene_cfg injection claim needs correction or implementation

Important Files Changed

Filename Overview
source/isaaclab_contrib/isaaclab_contrib/coupling/coupled_manager.py New manager dispatching to SolverCoupledProxy and SolverCoupledAdmm; partitioning and proxy-body resolution logic is well-structured and thoroughly tested.
source/isaaclab_contrib/isaaclab_contrib/coupling/coupled_manager_cfg.py Config dataclasses for CoupledSolverCfg, CoupledProxySolverCfg, and CoupledAdmmSolverCfg; clean with sensible defaults and good docstrings.
docs/source/overview/core-concepts/physical-backends/newton/using-vbd-solver.rst Adds proxy-coupling section and usage guide; contains a factually incorrect claim about automatic scene_cfg injection via PresetCfg plumbing that doesn't match the implementation.
source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py Adds _filter_solver_kwargs static helper that deduplicates the inspect-based kwargs pattern across all manager subclasses; clean refactor.
source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift_franka_soft/franka_soft_env_cfg.py Migrates to CoupledNewtonCfg, adds the proxy-coupled MJWarp+VBD preset as default, renames deformable->object throughout; breaking rename is clearly documented in changelog.
source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift_franka_soft/mdp/rewards.py Renames deformable_* reward/termination functions to object_*; extracts _com_w/_points_w private helpers that are then imported by observations.py (cross-module private import).
source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift_franka_soft/mdp/observations.py Renames DeformableSampledPointsInRobotRootFrame to ObjectSampledPointsInRobotRootFrame and generalizes its implementation; imports private helpers _com_w/_points_w from rewards.py.
source/isaaclab_contrib/test/coupling/test_coupled_manager.py 16 pure-Python unit tests covering body partitioning, proxy-body selection, overlap/unclaimed error cases, and mixed string/SceneEntityCfg selectors; comprehensive and well-structured.

Sequence Diagram

sequenceDiagram
    participant EnvCfg as FrankaSoftEnvCfg
    participant PhysicsCfg as PhysicsCfg (PresetCfg)
    participant NewtonCoupled as NewtonCoupledSolverManager
    participant Partition as _partition_model_by_entities
    participant ProxySolver as SolverCoupledProxy
    participant AdmmSolver as SolverCoupledAdmm

    EnvCfg->>PhysicsCfg: resolve preset (newton_mjwarp_vbd_proxy)
    PhysicsCfg->>NewtonCoupled: _build_solver(model, CoupledProxySolverCfg)
    NewtonCoupled->>NewtonCoupled: _resolve_solver_class(src_solver_cfg)
    NewtonCoupled->>NewtonCoupled: _resolve_solver_class(dst_solver_cfg)
    NewtonCoupled->>Partition: _partition_model_by_entities(src_bodies, dst_bodies)
    Partition-->>NewtonCoupled: src/dst bodies, joints, shapes
    alt CoupledProxySolverCfg
        NewtonCoupled->>NewtonCoupled: _select_proxy_bodies(proxy_bodies)
        NewtonCoupled->>ProxySolver: SolverCoupledProxy(model, entries, coupling)
        ProxySolver-->>NewtonCoupled: solver instance
    else CoupledAdmmSolverCfg
        NewtonCoupled->>AdmmSolver: SolverCoupledAdmm(model, entries, coupling)
        AdmmSolver-->>NewtonCoupled: solver instance
    end
    NewtonCoupled->>NewtonCoupled: "NewtonManager._solver = solver instance"
Loading

Comments Outside Diff (2)

  1. docs/source/overview/core-concepts/physical-backends/newton/using-vbd-solver.rst, line 191-195 (link)

    P1 Documentation incorrectly claims automatic scene_cfg wiring

    The guide states that "The env's __post_init__ is responsible for setting self.sim.physics.scene_cfg = self.scene — the Franka soft env does this automatically via the PresetCfg plumbing." However, FrankaSoftEnvCfg.__post_init__ (and FrankaClothEnvCfg.__post_init__) never sets scene_cfg, and there is no code path in PresetCfg that does so either. A search for scene_cfg in the entire task directory returns zero matches.

    The current env config works only because all body-selector values happen to be raw prim-path strings, which bypass the scene_cfg lookup entirely. Any user who follows the docs and attempts to use SceneEntityCfg selectors while relying on the "automatic" wiring will silently get scene_cfg=None, causing every SceneEntityCfg lookup to raise ValueError at solver-build time. Either implement the automatic injection or correct the documentation to clearly state that users must set self.sim.physics.scene_cfg = self.scene themselves in __post_init__.

  2. source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift_franka_soft/mdp/rewards.py, line 107-114 (link)

    P2 Inconsistent data-accessor style for Articulation tensors

    observations.py was updated to access robot articulation data via the .torch property (robot.data.root_pos_w.torch), while object_com_goal_distance here still calls wp.to_torch(robot.data.root_pos_w). The two forms should produce the same tensor, but the inconsistency makes it harder to move away from direct warp imports in the future and reads as if one pattern is "old" and the other "new".

    Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

Reviews (1): Last reviewed commit: "Fix: No need to integrate with external ..." | Re-trigger Greptile


import torch
import warp as wp

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Cross-module import of private helpers

observations.py imports _com_w and _points_w from rewards.py. Both are prefixed with _, meaning they are private to that module. Importing private symbols across module boundaries makes refactoring fragile and breaks the encapsulation contract that underscore-prefixed names imply. These shared utilities should be extracted into a small internal helpers module (e.g. _deformable_utils.py) so both rewards.py and observations.py can import from a common, intentionally public location.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

@kellyguo11 kellyguo11 moved this to In progress in Isaac Lab May 30, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentation Improvements or additions to documentation isaac-lab Related to Isaac Lab team

Projects

Status: In progress

Development

Successfully merging this pull request may close these issues.

2 participants