-
Notifications
You must be signed in to change notification settings - Fork 87
Add support for Mujoco Condim parameter in the model / builder #538 #553
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Add support for Mujoco Condim parameter in the model / builder #538 #553
Conversation
…-physics#538 Signed-off-by: Viktor Reutskyy <[email protected]>
📝 WalkthroughWalkthroughAdds per-shape contact dimensionality (condim): ShapeConfig.condim and Model.shape_material_condim, parsing from MJCF/USD, propagation into MuJoCo geom creation and the update kernel, and unit tests validating parsing and propagation (with some duplicated test methods). Changes
Sequence Diagram(s)sequenceDiagram
participant Importer as MJCF / USD Importer
participant Builder as ModelBuilder
participant Model as Model
participant Solver as SolverMuJoCo
participant Kernel as update_geom_properties_kernel
participant MJW as MuJoCo Warp Model
Importer->>Builder: parse shape -> read condim (attr / warp:contact_condim)
Builder->>Builder: add_shape(cfg.condim) -> append shape_material_condim[]
Builder->>Model: finalize() -> Model.shape_material_condim
Solver->>Solver: add_geoms() -> set geom_params["condim"] from Model.shape_material_condim
Solver->>Kernel: update_geom_properties(..., shape_condim)
Kernel-->>MJW: write geom_condim (guard: worldid==0)
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Assessment against linked issues
Assessment against linked issues: Out-of-scope changes
Possibly related PRs
Suggested reviewers
Tip 🔌 Remote MCP (Model Context Protocol) integration is now available!Pro plan users can now connect to remote MCP servers from the Integrations page. Connect with popular remote MCPs such as Notion and Linear to add more context to your reviews and chats. ✨ Finishing Touches
🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
Status, Documentation and Community
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Actionable comments posted: 2
🔭 Outside diff range comments (2)
newton/sim/builder.py (2)
651-723
: Critical: condim is not cloned in add_builder (env cloning, templating) — values silently lostWhen composing scenes via add_builder, most shape properties are extended from the source builder through more_builder_attrs, but shape_material_condim is missing. This leads to:
- Length mismatch between shape arrays and shape_material_condim.
- Default/undefined condim in composed models.
- Incorrect propagation to MuJoCo geoms in multi-env or templated builds.
Apply this diff to include condim in the cloning list:
more_builder_attrs = [ "articulation_key", "body_inertia", "body_mass", "body_inv_inertia", "body_inv_mass", "body_com", "body_qd", "body_key", "joint_type", "joint_enabled", "joint_X_c", "joint_armature", "joint_axis", "joint_dof_dim", "joint_dof_mode", "joint_key", "joint_qd", "joint_f", "joint_target", "joint_limit_lower", "joint_limit_upper", "joint_limit_ke", "joint_limit_kd", "joint_target_ke", "joint_target_kd", "joint_effort_limit", "joint_velocity_limit", "joint_friction", "shape_key", "shape_flags", "shape_type", "shape_scale", "shape_source", "shape_is_solid", "shape_thickness", "shape_material_ke", "shape_material_kd", "shape_material_kf", "shape_material_ka", "shape_material_mu", "shape_material_restitution", + "shape_material_condim", "shape_collision_radius", "particle_qd", "particle_mass", "particle_radius", "particle_flags", "edge_rest_angle", "edge_rest_length", "edge_bending_properties", "spring_rest_length", "spring_stiffness", "spring_damping", "spring_control", "tri_poses", "tri_activations", "tri_materials", "tri_areas", "tet_poses", "tet_activations", "tet_materials", "equality_constraint_type", "equality_constraint_body1", "equality_constraint_body2", "equality_constraint_anchor", "equality_constraint_torquescale", "equality_constraint_relpose", "equality_constraint_joint1", "equality_constraint_joint2", "equality_constraint_polycoef", "equality_constraint_key", "equality_constraint_enabled", ]
2439-2451
: Preserve condim when approximate_meshes spawns additional convex partsWhen convex decomposition yields multiple parts, new shapes are created with a freshly constructed ShapeConfig. condim is currently omitted, defaulting to 3 and breaking parity with the source shape.
Apply this diff to copy condim from the original shape:
cfg = ModelBuilder.ShapeConfig( density=0.0, # do not add extra mass / inertia ke=self.shape_material_ke[shape], kd=self.shape_material_kd[shape], kf=self.shape_material_kf[shape], ka=self.shape_material_ka[shape], mu=self.shape_material_mu[shape], restitution=self.shape_material_restitution[shape], + condim=self.shape_material_condim[shape], thickness=self.shape_thickness[shape], is_solid=self.shape_is_solid[shape], collision_group=self.shape_collision_group[shape], collision_filter_parent=self.default_shape_cfg.collision_filter_parent, )
🧹 Nitpick comments (5)
newton/sim/model.py (1)
108-109
: Add attribute_frequency entry for condim to keep metadata consistentYou added the public field and docstring, good. To maintain consistency with other per-shape material fields (ke/kd/kf/ka/mu/restitution), also register
shape_material_condim
inattribute_frequency
. This keeps tooling relying on attribute frequency metadata working as expected.Apply this patch in the attributes-per-shape section:
self.attribute_frequency["shape_material_mu"] = "shape" self.attribute_frequency["shape_material_restitution"] = "shape" + self.attribute_frequency["shape_material_condim"] = "shape" self.attribute_frequency["shape_type"] = "shape"
newton/solvers/mujoco/solver_mujoco.py (2)
984-1006
: Plumbing shape_condim → geom_condim is correct; keep dtype usage consistentThe kernel signature changes look correct: taking
shape_condim: wp.array(dtype=wp.int32)
and addinggeom_condim: wp.array(dtype=wp.int32)
as output matches the intended 1D per-geom storage. Consider usingdtype=int
consistently across kernels (as used elsewhere for condim arrays) to avoid confusion, since Warp aliasesint
to 32-bit anyway.
1975-1980
: Per-shape condim propagation on creation looks good; optionally validate rangeAssigning
geom_params["condim"]
from the model is correct. Optionally validate that condim is in MuJoCo’s supported range [1..6] to catch malformed inputs early.For example:
- geom_params["condim"] = int(shape_condim[shape]) + cdim = int(shape_condim[shape]) + if not (1 <= cdim <= 6): + warnings.warn(f"Invalid condim={cdim} for shape {shape}, clamping to [1,6].", stacklevel=2) + cdim = 1 if cdim < 1 else 6 + geom_params["condim"] = cdimnewton/utils/import_mjcf.py (1)
301-310
: Condim parsing and propagation are correct; consider light validationReading
condim
with a default frombuilder.default_shape_cfg.condim
and storing it intoshape_cfg.condim
is correct. Optionally validate that the value is within MuJoCo’s expected range [1..6] to avoid surprises from malformed files.- geom_condim = parse_int(geom_attrib, "condim", builder.default_shape_cfg.condim) + geom_condim = parse_int(geom_attrib, "condim", builder.default_shape_cfg.condim) + if geom_condim < 1 or geom_condim > 6: + if verbose: + print(f"Warning: clamping invalid condim={geom_condim} on geom '{geom_name}' to [1,6]") + geom_condim = min(6, max(1, geom_condim))newton/utils/import_usd.py (1)
1007-1011
: USD condim support wired correctly; optional range checkUsing
parse_float_with_fallback(..., builder.default_shape_cfg.condim)
and casting toint
is reasonable given attribute typing on USD. Consider validating to [1..6] as a guardrail (and optionally surfacing a warning when clamping).- condim=int( - parse_float_with_fallback( - prim_and_scene, "warp:contact_condim", builder.default_shape_cfg.condim - ) - ), + condim=max(1, min(6, int( + parse_float_with_fallback( + prim_and_scene, "warp:contact_condim", builder.default_shape_cfg.condim + ) + ))),
📜 Review details
Configuration used: .coderabbit.yml
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these settings in your CodeRabbit configuration.
📒 Files selected for processing (8)
newton/sim/builder.py
(4 hunks)newton/sim/model.py
(1 hunks)newton/solvers/mujoco/solver_mujoco.py
(6 hunks)newton/tests/test_import_mjcf.py
(1 hunks)newton/tests/test_import_usd.py
(1 hunks)newton/tests/test_mujoco_solver.py
(1 hunks)newton/utils/import_mjcf.py
(2 hunks)newton/utils/import_usd.py
(1 hunks)
🧰 Additional context used
🧠 Learnings (2)
📚 Learning: 2025-08-12T18:07:17.267Z
Learnt from: nvlukasz
PR: newton-physics/newton#519
File: newton/utils.py:30-33
Timestamp: 2025-08-12T18:07:17.267Z
Learning: In newton/_src/utils/import_usd.py, the PXR/USD modules are lazily loaded using try-except blocks within the functions (lines 109-111 and 1289-1291) rather than at module import time, making it safe to import parse_usd at the module level without causing import-time failures for users without USD dependencies.
Applied to files:
newton/tests/test_import_usd.py
📚 Learning: 2025-08-12T18:07:17.267Z
Learnt from: nvlukasz
PR: newton-physics/newton#519
File: newton/utils.py:30-33
Timestamp: 2025-08-12T18:07:17.267Z
Learning: In newton/_src/utils/import_usd.py, the PXR/USD modules are lazily loaded using try-except blocks within the functions rather than at module import time, making it safe to import parse_usd at the module level without causing import-time failures for users without USD dependencies.
Applied to files:
newton/tests/test_import_usd.py
🧬 Code Graph Analysis (3)
newton/utils/import_mjcf.py (1)
newton/sim/builder.py (1)
copy
(180-181)
newton/tests/test_import_mjcf.py (3)
newton/tests/test_import_usd.py (1)
test_condim_parsing
(333-399)newton/sim/builder.py (2)
ModelBuilder
(80-3928)finalize
(3564-3896)newton/utils/import_mjcf.py (1)
parse_mjcf
(33-928)
newton/tests/test_import_usd.py (2)
newton/tests/test_import_mjcf.py (1)
test_condim_parsing
(120-142)newton/sim/builder.py (2)
ModelBuilder
(80-3928)finalize
(3564-3896)
⏰ Context from checks skipped due to timeout of 900000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)
- GitHub Check: Run GPU Benchmarks (Pull Request)
- GitHub Check: Run GPU Unit Tests on AWS EC2 (Pull Request)
- GitHub Check: run-newton-tests / newton-unittests (windows-latest)
- GitHub Check: run-newton-tests / newton-unittests (ubuntu-latest)
🔇 Additional comments (9)
newton/solvers/mujoco/solver_mujoco.py (1)
2545-2572
: Update path wiring is correct; ensure condim array is always availableWiring
self.model.shape_material_condim
into the kernel and writing toself.mjw_model.geom_condim
is correct. Verify that all code paths building a Model (especially outside the builder) populateshape_material_condim
to avoid passingNone
into the kernel.Would you like a small fallback shim in
update_geom_properties()
to allocate a defaultwp.array
filled with 3 whenshape_material_condim
isNone
?newton/utils/import_mjcf.py (1)
209-214
: Helper parse_int is fineSimple and clear utility; matches the style of
parse_float
. LGTM.newton/tests/test_import_mjcf.py (1)
120-143
: Good coverage for condim parsingThe test asserts presence of all targeted condim values in the model after parsing. Clear and sufficient.
newton/tests/test_mujoco_solver.py (1)
1089-1128
: LGTM: Solid end-to-end test for condim storage in the ModelGood coverage: you create shapes with distinct condims, finalize, and assert the per-shape values are stored in model.shape_material_condim as expected.
newton/tests/test_import_usd.py (1)
332-401
: LGTM: USD condim parsing test is comprehensive and correctly guardedThe test constructs in-memory USD prims with warp:contact_condim, parses them, and validates propagation into model.shape_material_condim. Properly guarded by @skipUnless(USD_AVAILABLE) and consistent with the lazy USD imports note in import_usd.
newton/sim/builder.py (4)
146-147
: Add ShapeConfig.condim (default 3): Good API addition and docClear semantics and sensible default align with MuJoCo expectations.
334-334
: Track shape_material_condim in builder stateThis complements other per-shape material arrays and is required for finalization.
1955-1955
: Persist condim at shape creationAppending cfg.condim here ensures per-shape values survive to Model finalization.
3661-3661
: Expose Model.shape_material_condimCorrect dtype (int32) for MuJoCo geom condim propagation. Good.
Signed-off-by: Viktor Reutskyy <[email protected]>
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Actionable comments posted: 0
♻️ Duplicate comments (1)
newton/solvers/mujoco/solver_mujoco.py (1)
1038-1041
: Data race on 1D condim avoided by gating writes to world 0 — good fixGuarding the write with
if worldid == 0:
removes the cross-world write race when launching with(nworld, ngeom)
. This addresses the earlier review concern.
🧹 Nitpick comments (3)
newton/solvers/mujoco/solver_mujoco.py (3)
989-989
: Kernel now accepts per-shape condim: ensure dtype consistency and non-None input
shape_condim
is declared aswp.int32
. Elsewhere (e.g., contact kernels) condim/dim arrays sometimes usedtype=int
. For Warp, mixingint
andwp.int32
can be fragile across backends. I recommend standardizing onwp.int32
for all condim/dim arrays.Also, make sure
self.model.shape_material_condim
is always defined (non-None) to match this signature; otherwise the kernel launch will fail.Would you like me to sweep the codebase and normalize all condim/dim arrays to
wp.int32
? I can generate a script to find and propose changes.
1976-1981
: Validate per-shape condim values; clamp or warn on invalid inputYou’re correctly propagating per-shape condim into MJCF geoms. Add a small validation to guard against invalid values and improve debuggability.
Apply this diff:
- # use per-shape condim if available - if model.shape_material_condim is not None: - shape_condim = model.shape_material_condim.numpy() - if shape < len(shape_condim): - geom_params["condim"] = int(shape_condim[shape]) + # use per-shape condim if available + if model.shape_material_condim is not None: + shape_condim = model.shape_material_condim.numpy() + if shape < len(shape_condim): + condim_val = int(shape_condim[shape]) + if condim_val not in (1, 3, 4, 6): + warnings.warn( + f"Invalid condim {condim_val} for shape {shape}; falling back to 3", + stacklevel=2, + ) + condim_val = 3 + geom_params["condim"] = condim_val
2553-2553
: Kernel launch assumes shape_material_condim is present; add a safe fallbackIf
model.shape_material_condim
is everNone
(e.g., older models or custom builders), this launch will fail. Pass a default (all 3’s) when it’s missing.Apply this diff within the inputs list:
- self.model.shape_material_condim, + shape_condim_in,And define the helper just before the
wp.launch
:# fallback to default condim=3 if array is missing shape_condim_in = ( self.model.shape_material_condim if self.model.shape_material_condim is not None else wp.zeros(self.model.shape_count, dtype=wp.int32, device=self.model.device) + 3 )If
shape_material_condim
is guaranteed by the builder/model invariants, feel free to ignore this. Otherwise, this guard prevents runtime errors.
📜 Review details
Configuration used: .coderabbit.yml
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these settings in your CodeRabbit configuration.
📒 Files selected for processing (1)
newton/solvers/mujoco/solver_mujoco.py
(6 hunks)
⏰ Context from checks skipped due to timeout of 900000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)
- GitHub Check: Run GPU Unit Tests on AWS EC2 (Pull Request)
- GitHub Check: Run GPU Benchmarks (Pull Request)
- GitHub Check: run-newton-tests / newton-unittests (ubuntu-latest)
- GitHub Check: run-newton-tests / newton-unittests (windows-latest)
🔇 Additional comments (2)
newton/solvers/mujoco/solver_mujoco.py (2)
1002-1002
: geom_condim as 1D output: OK, but verify MjWarp model field availabilityDeclaring
geom_condim
as 1D output matches the “write from world 0 only” pattern below and avoids cross-world write races. Please confirm thatmujoco_warp.put_model(self.mj_model)
exposesgeom_condim
and that its dtype iswp.int32
. If using an older mujoco_warp, this attribute may not exist.I can provide a quick repo check script to assert
geom_condim
exists onMjWarpModel
instances if helpful.
2567-2567
: Ensure MjWarp’s geom_condim is initialized and not expanded across worldsWriting to
self.mjw_model.geom_condim
here is correct (1D, per-geom). Given it’s intentionally not inexpand_model_fields
, it stays single and avoids races. Just verify thatgeom_condim
is indeed a field on the current mujoco_warp version used in CI.I can add a quick CI-time assertion to catch missing fields if wanted.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
approving - some comments and the prefix change to newton but this looks good! Thank you!
…in-the-model-builder
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I would hold off on merging this for now until we have decided whether we need this MuJoCo property now and if it makes sense to be part of the core Newton Model.
Signed-off-by: Viktor Reutskyy <[email protected]>
…in-the-model-builder
wasn't it decided that we need this for the IsaacLab integration? |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Actionable comments posted: 0
🔭 Outside diff range comments (1)
newton/_src/sim/builder.py (1)
2569-2591
: Preserve condim when creating new shapes during convex decomposition.When approximate_meshes creates additional convex parts, the new shapes inherit most material params but drop condim, implicitly reverting to default (3). This can change contact behavior for decomposed shapes.
Add condim to the derived cfg so all parts retain the original shape’s condim:
cfg = ModelBuilder.ShapeConfig( density=0.0, # do not add extra mass / inertia ke=self.shape_material_ke[shape], kd=self.shape_material_kd[shape], kf=self.shape_material_kf[shape], ka=self.shape_material_ka[shape], mu=self.shape_material_mu[shape], restitution=self.shape_material_restitution[shape], + condim=self.shape_material_condim[shape], thickness=self.shape_thickness[shape], is_solid=self.shape_is_solid[shape], collision_group=self.shape_collision_group[shape], collision_filter_parent=self.default_shape_cfg.collision_filter_parent, )
🧹 Nitpick comments (4)
newton/_src/sim/model.py (1)
121-122
: Add attribute_frequency entry and tighten the docstring dtype.Nice addition. Two follow-ups:
- Register this attribute in attribute_frequency so downstream code that relies on the frequency map can classify it correctly.
- Docstring currently says int; everywhere else this is int32. Consider being explicit.
Apply this doc tweak:
- """Shape contact dimensionality for MuJoCo solver, shape [shape_count], int.""" + """Shape contact dimensionality for MuJoCo solver, shape [shape_count], int32."""And add the frequency mapping near the other per-shape attributes (around Line 426):
# attributes per shape self.attribute_frequency["shape_material_ke"] = "shape" ... self.attribute_frequency["shape_material_restitution"] = "shape" self.attribute_frequency["shape_material_condim"] = "shape" # add this self.attribute_frequency["shape_type"] = "shape"newton/_src/utils/import_mjcf.py (1)
302-311
: Coerce and validate geom condim to supported values (1, 3, 4, 6).Parsing works, but MuJoCo only recognizes a small set of condim values. Add a small guard to avoid surprising inputs in source MJCF.
- geom_condim = parse_int(geom_attrib, "condim", builder.default_shape_cfg.condim) + geom_condim = parse_int(geom_attrib, "condim", builder.default_shape_cfg.condim) + # enforce valid MuJoCo condim values + if geom_condim not in (1, 3, 4, 6): + if verbose: + print(f"Warning: Unsupported condim={geom_condim} for geom '{geom_name}', " + f"falling back to default {builder.default_shape_cfg.condim}.") + geom_condim = builder.default_shape_cfg.condimnewton/_src/utils/import_usd.py (1)
1014-1019
: Use integer parsing and validate condim for USD import.Casting float->int may truncate silently. Prefer explicit integer parse with validation to ensure only supported MuJoCo condim values pass through.
- condim=int( - parse_float_with_fallback( - prim_and_scene, "warp:contact_condim", builder.default_shape_cfg.condim - ) - ), + # read and validate MuJoCo condim + condim=(lambda _v: ( + _v if _v in (1, 3, 4, 6) + else (builder.default_shape_cfg.condim) + ))( + int(parse_float_with_fallback( + prim_and_scene, "warp:contact_condim", builder.default_shape_cfg.condim + )) + ),Optionally, emit a warning when coercing an invalid value if verbose is enabled at the call-site.
newton/_src/sim/builder.py (1)
2084-2084
: LGTM, with optional guard: Persist condim per shape at creation.Appending cfg.condim is correct. Consider clamping to the supported MuJoCo set here as a generalized safety net (besides importer paths), e.g.:
- self.shape_material_condim.append(cfg.condim) + self.shape_material_condim.append(cfg.condim if cfg.condim in (1, 3, 4, 6) else 3)Or, better, a small helper that centralizes condim validation.
📜 Review details
Configuration used: .coderabbit.yml
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (7)
newton/_src/sim/builder.py
(5 hunks)newton/_src/sim/model.py
(1 hunks)newton/_src/solvers/mujoco/solver_mujoco.py
(6 hunks)newton/_src/utils/import_mjcf.py
(2 hunks)newton/_src/utils/import_usd.py
(1 hunks)newton/tests/test_import_usd.py
(1 hunks)newton/tests/test_mujoco_solver.py
(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (2)
- newton/tests/test_import_usd.py
- newton/tests/test_mujoco_solver.py
⏰ Context from checks skipped due to timeout of 900000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)
- GitHub Check: Run GPU Unit Tests on AWS EC2 (Pull Request)
- GitHub Check: Run GPU Benchmarks (Pull Request)
- GitHub Check: run-newton-tests / newton-unittests (ubuntu-latest)
- GitHub Check: run-newton-tests / newton-unittests (windows-latest)
🔇 Additional comments (9)
newton/_src/sim/builder.py (4)
162-164
: LGTM: Per-shape condim in ShapeConfig is well-scoped and documented.Defaulting to 3 matches MuJoCo’s typical friction model; docstring is clear.
351-351
: LGTM: New builder storage for shape_material_condim.This aligns with the other material arrays and sets you up for clean propagation into Model.
789-789
: LGTM: Propagating shape_material_condim in more_builder_attrs.Ensures add_builder and cloning paths carry condim correctly. Good catch including it here.
3795-3796
: LGTM: Finalize creates Model.shape_material_condim with int32 dtype.Matches the Model attribute and will interop cleanly with the MuJoCo solver path.
newton/_src/solvers/mujoco/solver_mujoco.py (5)
999-999
: LGTM! Kernel signature update is correct.The addition of
shape_condim
parameter to the kernel signature is properly positioned and correctly typed aswp.int32
.
1012-1012
: LGTM! Kernel output parameter added correctly.The
geom_condim
output parameter is properly added to the kernel signature to receive the condim values.
1985-1989
: LGTM! Per-shape condim support integrated correctly.The implementation correctly checks for the availability of
shape_material_condim
, retrieves the value, and sets it in the geom parameters. The conversion toint()
ensures type compatibility with MuJoCo.
2561-2561
: LGTM! Kernel invocation updated with condim parameters.The
update_geom_properties_kernel
invocation correctly passesmodel.shape_material_condim
as input andself.mjw_model.geom_condim
as output, maintaining consistency with the kernel signature changes.Also applies to: 2575-2575
1048-1051
: Verify condim’s intended scope
The guard onworldid == 0
serializes writes into the single‐dimensionalgeom_condim
array—correct if condim is global, but potentially incorrect if each world requires its own condim value. Please confirm the design intent:
- If condim is truly a global property (identical across worlds), no changes are needed.
- If condim must vary per world, refactor
geom_condim
into a 2D array (e.g.[n_worlds][n_geoms]
) and index it byworldid
.Location to review:
- newton/_src/solvers/mujoco/solver_mujoco.py:1048–1051
@eric-heiden - it probably does not make sense as a core newton property, but it would be a MJC specific attribute. Should we not have a base model derived MJC model? Or how do we handle MJC specific parameters? |
…in-the-model-builder
I think we need a mechanism generally to handle these solver-specific attributes (and parse them from USD). I've added some notes to the architecture backlog. |
…in-the-model-builder Signed-off-by: Viktor Reutskyy <[email protected]>
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Actionable comments posted: 0
♻️ Duplicate comments (1)
newton/tests/test_mujoco_solver.py (1)
1087-1146
: Make geom_condim indexing robust across single-/multi-env layoutsgeom_condim may be 1D (per-geom) or 2D ([env, geom]) depending on backend/version. Current test assumes 1D. Index by ndim to keep the test stable.
Apply this diff:
- expected_condim = condim_values[shape_idx] - actual_condim = geom_condim[geom_idx] + expected_condim = condim_values[shape_idx] + # Support both [ngeom] and [num_envs, ngeom] layouts + actual_condim = ( + geom_condim[world_idx, geom_idx] + if getattr(geom_condim, "ndim", None) == 2 + else geom_condim[geom_idx] + )
🧹 Nitpick comments (3)
newton/_src/utils/import_mjcf.py (2)
210-215
: Harden integer parsing to avoid ValueError and noisy crashesparse_int() directly casts attribute strings to int. MJCF values can occasionally be malformed or provided as floats (e.g., "3.0"). Gracefully fallback to default and optionally warn when verbose.
Apply this diff:
- def parse_int(attrib, key, default) -> int: - if key in attrib: - return int(attrib[key]) - else: - return default + def parse_int(attrib, key, default) -> int: + if key in attrib: + try: + return int(attrib[key]) + except (ValueError, TypeError): + if verbose: + print(f"MJCF parse warning: invalid int for '{key}'='{attrib[key]}', using default {default}") + return default + else: + return default
303-311
: Validate geom condim and fallback to default for unsupported valuesYou correctly parse condim and propagate it to shape_cfg. To prevent invalid values, validate against supported MuJoCo condim semantics {1, 3, 4, 6} and fallback (with a verbose warning) to the builder default when out-of-range.
Apply this diff:
- geom_condim = parse_int(geom_attrib, "condim", builder.default_shape_cfg.condim) + condim_default = builder.default_shape_cfg.condim + geom_condim = parse_int(geom_attrib, "condim", condim_default) + if geom_condim not in (1, 3, 4, 6): + if verbose: + print(f"MJCF parse warning: geom '{geom_name}' condim={geom_condim} not in {{1,3,4,6}}. Using default {condim_default}.") + geom_condim = condim_defaultDo you expect any additional legal values for condim beyond {1,3,4,6} in your target MuJoCo version? If yes, adjust the validation set accordingly.
newton/_src/solvers/mujoco/solver_mujoco.py (1)
2546-2571
: Ensure mujoco_warp exposes geom_condim; consider notifier coverage for condim-only updates
- You pass model.shape_material_condim into the kernel and write to mjw_model.geom_condim. Confirm your mujoco_warp version exposes geom_condim; otherwise this will raise at runtime.
- notify_model_changed currently triggers update_geom_properties when shape_material_mu is present. If a client updates condim only, consider calling update_geom_properties as well (either always or based on a dedicated flag).
Suggested minimal guard for notifier (optional):
- flags = SolverNotifyFlags.BODY_INERTIAL_PROPERTIES | SolverNotifyFlags.JOINT_DOF_PROPERTIES - - if model.shape_material_mu is not None: - flags |= SolverNotifyFlags.SHAPE_PROPERTIES + flags = SolverNotifyFlags.BODY_INERTIAL_PROPERTIES | SolverNotifyFlags.JOINT_DOF_PROPERTIES + # Always refresh geom properties; they are cheap and cover mu/kd/condim/size/pose in one pass. + flags |= SolverNotifyFlags.SHAPE_PROPERTIESVerification request:
- Please confirm that solver.mjw_model has a geom_condim field at runtime (mujoco_warp.put_model). If it doesn’t, we’ll need to add it upstream or gate the writes.
📜 Review details
Configuration used: .coderabbit.yml
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (7)
newton/_src/sim/builder.py
(5 hunks)newton/_src/solvers/mujoco/solver_mujoco.py
(6 hunks)newton/_src/utils/import_mjcf.py
(2 hunks)newton/_src/utils/import_usd.py
(1 hunks)newton/tests/test_import_mjcf.py
(1 hunks)newton/tests/test_import_usd.py
(1 hunks)newton/tests/test_mujoco_solver.py
(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (4)
- newton/tests/test_import_mjcf.py
- newton/_src/sim/builder.py
- newton/_src/utils/import_usd.py
- newton/tests/test_import_usd.py
⏰ Context from checks skipped due to timeout of 900000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)
- GitHub Check: Run GPU Benchmarks (Pull Request)
- GitHub Check: Run GPU Unit Tests on AWS EC2 (Pull Request)
🔇 Additional comments (3)
newton/_src/solvers/mujoco/solver_mujoco.py (3)
993-1016
: Kernel signature extension looks correct and type-safeAdding shape_condim (int32) as input and geom_condim (int32) as output is consistent with other buffers and aligns with per-shape material semantics.
1048-1051
: Avoid data races via single-writer for geom_condimWriting geom_condim only from world 0 avoids multi-world write contention. Good choice given geom_condim is per-geom, not per-world.
1977-1982
: Propagate per-shape condim at geom creation — LGTMAssigning geom_params["condim"] from model.shape_material_condim ensures template-time correctness and matches MuJoCo’s per-geom condim semantics.
Description
This PR adds support for the MuJoCo
condim
parameter in the Newton model/builder system. Thecondim
parameter controls contact dimensionality, allowing fine-grained control over collision behavior on a per-shape basis.What's new:
condim
property toModelBuilder.ShapeConfig
with a default value of 3 (normal + 2D friction)shape_material_condim
array to theModel
class to store condim values for each shapecondim
attribute from MJCF fileswarp:contact_condim
attribute in USD filesCondim values:
Example usage:
Closes #538
Newton Migration Guide
Please ensure the migration guide for warp.sim users is up-to-date with the changes made in this MR.
docs/migration.rst
is up-to dateBefore your PR is "Ready for review"
newton/tests/test_examples.py
)pre-commit run -a
Summary by CodeRabbit
New Features
Tests