Add generic Newton coupled-solver manager with Proxy + ADMM variants#5834
Add generic Newton coupled-solver manager with Proxy + ADMM variants#5834mmichelis wants to merge 23 commits into
Conversation
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.
There was a problem hiding this comment.
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
-
Excellent extensibility - The
_SOLVER_CLASS_BY_CFG_TYPEregistry pattern inNewtonCoupledSolverManagermakes it trivial to add new solver combinations without modifying the manager. -
Clean dispatch pattern - Using config subclass type (
CoupledProxySolverCfgvsCoupledAdmmSolverCfg) to select the coupling algorithm is elegant and type-safe. -
Good separation of concerns - Body selector resolution, model partitioning, and proxy selection are cleanly separated into well-named static methods.
-
Comprehensive unit tests - The 358-line test file with fake objects is excellent for testing partitioning logic without launching Isaac Sim.
-
Helpful error messages - The partition validation errors include specific body IDs and labels, making debugging much easier.
-
Refactoring win - Extracting
_filter_solver_kwargsintoNewtonManagerremoves 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.deformable→scene.object - Command:
deformable_pose→object_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_solverparameter forwardingContactPairconstruction whenenable_contacts=True/False- The
contact_distance/detection_marginNone 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 NoneIf 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
-
Documentation is solid - The
using-vbd-solver.rstadditions explain proxy coupling well with good examples. -
Coupling mode validation - Good that
coupled_mjwarp_vbd_manager.pyandcoupled_featherstone_vbd_manager.pynow raise for unknown coupling modes. -
Static shape routing - The decision to route
body == -1shapes 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 SummaryThis PR introduces a new
Confidence Score: 4/5The 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
Sequence DiagramsequenceDiagram
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"
|
|
|
||
| import torch | ||
| import warp as wp | ||
|
|
There was a problem hiding this comment.
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!
Description
Summary
isaaclab_contrib.couplingsubmodule. One manager,NewtonCoupledSolverManager, that dispatches between two coupling algorithms by config subclass:CoupledProxySolverCfg→SolverCoupledProxy: 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.CoupledAdmmSolverCfg→SolverCoupledAdmm: linearized ADMM coupling with tunablerho/gamma/baumgartepenalty knobs and joint attachment stiffness/damping."/World/envs/env_.*/Robot") alongsideSceneEntityCfg— useful for assets that aren't registered as named scene entities.Isaac-Lift-Soft-Franka-v0now defaults to proxy-coupled MJWarp + VBD. Includes a breaking rename in the env's scene/MDP:deformable→object(entity, command, and all MDP functions). Update env configs, checkpoints, and RL configs accordingly.NewtonManager._filter_solver_kwargsso 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:
Type of change
Test plan
Checklist
pre-commitchecks with./isaaclab.sh --formatconfig/extension.tomlfileCONTRIBUTORS.mdor my name already exists there