newton/_src/utils/import_mjcf.py (3)

198-203: *Bug: xform is already a Transform; re-wrapping with wp.transform(xform) will fail

Earlier (Lines 89-93), xform is normalized to a wp.transform. Here you re-wrap it using wp.transform(*xform), but a Transform is not an iterable tuple. This will likely raise at runtime. Compose with axis_xform directly.

Apply this diff:

-    if xform is None:
-        xform = axis_xform
-    else:
-        xform = wp.transform(*xform) * axis_xform
+    xform = xform * axis_xform

Optional cleanup: You can also remove the earlier normalization block (Lines 89-93) and handle normalization once here, but the above minimal fix is sufficient.

---

290-312: Double application of incoming_xform for world geoms (link == -1)

incoming_xform is applied to tf twice when link == -1: once at Lines 291-295 and again at 310-312. This doubles the transform for worldbody geoms. Plane geoms depend on geom_pos/geom_rot (used to compute plane equation), so we should apply incoming_xform once to tf and refresh geom_pos/geom_rot only once for planes.

Apply this diff:

-            tf = wp.transform(geom_pos, geom_rot)
-            if link == -1 and incoming_xform is not None:
-                tf = incoming_xform * tf
-                geom_pos = tf.p
-                geom_rot = tf.q
+            tf = wp.transform(geom_pos, geom_rot)
@@
-            if incoming_xform is not None:
-                tf = incoming_xform * tf
+            if incoming_xform is not None:
+                tf = incoming_xform * tf
+                if link == -1:
+                    # keep geom_pos/geom_rot in sync for types (e.g., plane) that use them directly
+                    geom_pos = tf.p
+                    geom_rot = tf.q

---

577-584: Bug: base_joint dict case uses child=root (XML root), not link

When base_joint is a dict, you set base_joint["child"] = root, which is the XML root element, not the body index. This will break joint creation; it must be the current link.

Apply this diff:

-                    base_joint["parent"] = -1
-                    base_joint["child"] = root
+                    base_joint["parent"] = -1
+                    base_joint["child"] = link

newton/_src/utils/import_mjcf.py (2)

434-440: Use consistent type annotations: prefer Transform over wp.transform

The rest of the file (e.g., parse_mjcf’s xform param) uses the alias Transform. Using wp.transform in annotations is inconsistent and not a valid typing construct; prefer Transform for readability and type checking.

Apply this diff:

-        parent_world_xform: wp.transform | None = None,
+        parent_world_xform: Transform | None = None,

---

254-432: General: world_xform should not be passed to per-body geoms (correct), but do validate worldbody geoms with tests

Your per-body shapes are correctly authored in the body-local frame (leveraging the body’s world transform from add_body). Worldbody geoms (link == -1) get incoming_xform applied, which is correct. Please add tests that validate:

• Nested bodies with pos/quat compose to the expected world_xform.
• Free joints’ joint_q reflect the composed world_xform.
• Worldbody plane, box, and mesh geoms are transformed by xform exactly once.

I can open a test PR or generate pytest cases for a small 2-body chain (root + child), including a free-root variant. Do you want me to draft them?

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.

newton/_src/utils/import_mjcf.py (3)

198-203: Bug: incorrect re-derivation of xform with axis_xform; will TypeError when xform is already a Transform

At this point, xform has already been normalized to a wp.transform at Lines 89-93. Re-applying wp.transform(*xform) assumes xform is iterable (tuple of (p, q)), which will raise when xform is a Transform. Also, the None-guard is now unreachable since xform cannot be None anymore, making the branch misleading.

Replace with a simple post-multiplication by axis_xform.

Apply this diff:

-    axis_xform = wp.transform(wp.vec3(0.0), quat_between_axes(up_axis, builder.up_axis))
-    if xform is None:
-        xform = axis_xform
-    else:
-        xform = wp.transform(*xform) * axis_xform
+    axis_xform = wp.transform(wp.vec3(0.0), quat_between_axes(up_axis, builder.up_axis))
+    xform = xform * axis_xform

---

578-584: Bug: dict-based base_joint uses child=root instead of child=link

When base_joint is a dict, you set base_joint["child"] = root, but root here is the parsed XML root element, not the body index. This will cause incorrect behavior or errors. It should be the newly created body index (link).

Apply this diff:

-                    base_joint["parent"] = -1
-                    base_joint["child"] = root
+                    base_joint["parent"] = -1
+                    base_joint["child"] = link
                     base_joint["parent_xform"] = base_parent_xform
                     base_joint["child_xform"] = base_child_xform
                     base_joint["key"] = "base_joint"
                     builder.add_joint(**base_joint)

---

291-313: Double application of incoming_xform for world geoms (link == -1)

tf is multiplied by incoming_xform at Line 291 for world geoms, then again at Line 310 unconditionally when incoming_xform is set. This doubles the transform for world-level shapes and leads to incorrect placement, especially noticeable for planes where geom_pos/rot are also altered.

To avoid double transforms while retaining the plane normal computation, condition the second multiplication on link != -1:

-            if incoming_xform is not None:
-                tf = incoming_xform * tf
+            if incoming_xform is not None and link != -1:
+                tf = incoming_xform * tf

Alternatively, remove the early multiplication and compute plane normal and p using tf after a single application at this site.

newton/_src/utils/import_mjcf.py (1)

549-606: Edge case: root body with and floating/base_joint flags

The control flow here is nuanced. For root with :

• base_joint string: OK (D6 setup with position-only parent frame; rotation applied in child frame)
• floating True: OK (free joint added and joint_q set from world_xform)
• else: fixed base with _xform

Looks good overall; just consider explicitly documenting precedence when both base_joint and floating are provided (base_joint currently takes precedence).

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.

newton/tests/test_import_mjcf.py (11)

227-227: Make importer up-axis explicit to avoid hidden axis-correction

The importer may apply an axis correction if its default up_axis differs from the builder’s up_axis. To make the test stable across defaults, pass up_axis="Z" explicitly.

-        newton.utils.parse_mjcf(mjcf_content, builder)
+        newton.utils.parse_mjcf(mjcf_content, builder, up_axis="Z")

If you prefer, apply the same tweak to all parse_mjcf calls in this file to future-proof these tests.

---

236-239: Compare quaternions modulo sign to reduce brittleness

Numerically equivalent rotations can flip sign (q ≡ -q). Use a helper assertion that accepts both, to reduce intermittent failures across platforms and refactors.

Add this helper to TestImportMjcf (outside changed lines):

def assert_quat_allclose(self, q_actual, q_expected, atol=1e-6, rtol=1e-6):
    qa = np.asarray(q_actual)
    qe = np.asarray(q_expected)
    if np.allclose(qa, qe, rtol=rtol, atol=atol):
        return
    if np.allclose(qa, -qe, rtol=rtol, atol=atol):
        return
    self.fail(f"Quaternions not equal up to sign: {qa} vs {qe}")

Then replace usages like:

np.testing.assert_allclose(body_quat, expected_quat, atol=1e-6)

with:

self.assert_quat_allclose(body_quat, expected_quat, atol=1e-6)

This applies to all quaternion equality checks below.

---

241-280: Root body + custom xform composition is correct; minor robustness tweaks

The transform composition expected_xform = custom_xform * body_xform aligns with the importer’s semantics. Consider two small refinements:

-        newton.utils.parse_mjcf(mjcf_content, builder, xform=custom_xform)
+        newton.utils.parse_mjcf(mjcf_content, builder, xform=custom_xform, up_axis="Z")

And ensure dtype compatibility in asserts (Warp types generally work, but np.asarray is safer):

-        np.testing.assert_allclose(body_pos, expected_pos, atol=1e-6)
-        np.testing.assert_allclose(body_quat, expected_quat, atol=1e-6)
+        np.testing.assert_allclose(body_pos, np.asarray(expected_pos), atol=1e-6)
+        self.assert_quat_allclose(body_quat, np.asarray(expected_quat), atol=1e-6)

---

301-301: Stabilize hierarchy test against up-axis defaults

To avoid implicit axis-correction affecting expectations, make up-axis explicit.

-        newton.utils.parse_mjcf(mjcf_content, builder)
+        newton.utils.parse_mjcf(mjcf_content, builder, up_axis="Z")

---

320-324: Remove unused placeholder initialization

The temporary quat_child placeholder is assigned then immediately overwritten. Drop it to reduce noise.

-        quat_child = np.array([0, 0, 0, 1])  # placeholder, will be overwritten
-        quat_child_mjcf = np.array([0.7071068, 0, 0, 0.7071068])
-        # MJCF: [w, x, y, z] → warp: [x, y, z, w]
-        quat_child = np.array([quat_child_mjcf[1], quat_child_mjcf[2], quat_child_mjcf[3], quat_child_mjcf[0]])
+        quat_child_mjcf = np.array([0.7071068, 0, 0, 0.7071068])  # [w, x, y, z]
+        # MJCF [w, x, y, z] → Warp [x, y, z, w]
+        quat_child = np.array([quat_child_mjcf[1], quat_child_mjcf[2], quat_child_mjcf[3], quat_child_mjcf[0]])

---

351-351: Avoid accidental axis-correction in floating-base test

Make up-axis explicit for stability across environment defaults.

-        newton.utils.parse_mjcf(mjcf_content, builder, floating=True)
+        newton.utils.parse_mjcf(mjcf_content, builder, floating=True, up_axis="Z")

---

402-402: Make up-axis explicit in hierarchy + freejoint test

Same rationale as above; avoids hidden rotations from axis correction.

-        newton.utils.parse_mjcf(mjcf_content, builder)
+        newton.utils.parse_mjcf(mjcf_content, builder, up_axis="Z")

---

441-445: Remove debug prints from tests

These prints will spam test output and aren’t needed for assertions. Prefer logging behind a verbosity flag or drop entirely.

-        # Debug output to see what's happening
-        print(f"Root transform: pos={root_xform.p}, quat={root_xform.q}")
-        print(f"Child local transform: pos={child_local_xform.p}, quat={child_local_xform.q}")
-        print(f"Expected composed: pos={expected_pos}, quat={expected_quat}")
-        print(f"Actual joint_q: pos={joint_pos}, quat={joint_quat}")

---

469-469: Explicit up-axis makes chain test resilient

-        newton.utils.parse_mjcf(mjcf_content, builder)
+        newton.utils.parse_mjcf(mjcf_content, builder, up_axis="Z")

---

526-526: Explicit up-axis in scale test

-        newton.utils.parse_mjcf(mjcf_content, builder, scale=2.0)
+        newton.utils.parse_mjcf(mjcf_content, builder, scale=2.0, up_axis="Z")

---

570-570: Explicit up-axis in branching tree test

-        newton.utils.parse_mjcf(mjcf_content, builder)
+        newton.utils.parse_mjcf(mjcf_content, builder, up_axis="Z")

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.

newton/_src/utils/import_mjcf.py (1)

578-583: Bug: base_joint dict path sets child to the XML root instead of the link

Setting base_joint["child"] = root will pass the XML root element to add_joint, which should be an integer link index. This will mis-wire the joint or raise at runtime.

Apply:

-                    base_joint["child"] = root
+                    base_joint["child"] = link

newton/_src/utils/import_mjcf.py (1)

594-595: Use wp.vec3 default here too (remaining tuple default)

This site still defaults joint_pos to a Python tuple, which can trigger Warp deprecation warnings and potential type mismatch.

Apply:

-            joint_pos = joint_pos[0] if len(joint_pos) > 0 else (0.0, 0.0, 0.0)
+            joint_pos = joint_pos[0] if len(joint_pos) > 0 else wp.vec3(0.0, 0.0, 0.0)

newton/_src/utils/import_mjcf.py (2)

459-471: Simplify world transform composition to one branch

You can collapse the three-branch logic into a single expression to reduce cognitive load and avoid edge-case drift.

Apply:

-        # Compose transforms properly through the hierarchy
-        if parent == -1:
-            # Root body: apply the root transform
-            world_xform = xform * local_xform
-        elif parent_world_xform is not None:
-            # Child body: compose with parent's world transform
-            world_xform = parent_world_xform * local_xform
-        else:
-            # Fallback (shouldn't happen in normal cases)
-            world_xform = local_xform
+        # Compose with either the passed parent world transform or the import root xform
+        world_xform = (parent_world_xform or xform) * local_xform

And call parse_body for roots with xform as the initial parent transform:

-    for body in world.findall("body"):
-        parse_body(body, -1, world_defaults, parent_world_xform=None)
+    for body in world.findall("body"):
+        parse_body(body, -1, world_defaults, parent_world_xform=xform)

---

598-603: Verify: parent_xform for free joints is necessary?

ModelBuilder.add_joint_free typically initializes joint_q from the child body's xform; passing a non-identity parent_xform can rotate the joint frame relative to the parent unnecessarily and risks double-applying orientation.

Please confirm this is intended. If not needed, drop parent_xform for clarity and to align with builder semantics.

Proposed change:

-                builder.add_joint_free(
-                    link,
-                    key="_".join(joint_name),
-                    parent_xform=wp.transform(wp.vec3(0.0, 0.0, 0.0), body_ori_for_joints),
-                )
+                builder.add_joint_free(
+                    link,
+                    key="_".join(joint_name),
+                )

newton/tests/test_import_mjcf.py (2)

113-119: Make mesh maxhullvert test robust to shape ordering

Iterating range(3) assumes the first three shapes are meshes. Filter across all shapes for those with maxhullvert instead.

Apply:

-            # Check that meshes have correct maxhullvert values
-            # Note: This assumes meshes are added in order they appear in MJCF
-            meshes = [model.shape_source[i] for i in range(3) if hasattr(model.shape_source[i], "maxhullvert")]
-
-            if len(meshes) >= 3:
-                self.assertEqual(meshes[0].maxhullvert, 32)
-                self.assertEqual(meshes[1].maxhullvert, 128)
-                self.assertEqual(meshes[2].maxhullvert, 64)  # Default value
+            # Check that meshes have correct maxhullvert values
+            meshes = [src for src in model.shape_source if hasattr(src, "maxhullvert")]
+            self.assertEqual(len(meshes), 3)
+            self.assertEqual(meshes[0].maxhullvert, 32)
+            self.assertEqual(meshes[1].maxhullvert, 128)
+            self.assertEqual(meshes[2].maxhullvert, 64)  # Default value

---

441-445: Remove debug prints from tests

These prints add noise to CI logs; assertions already capture mismatches.

Apply:

-        # Debug output to see what's happening
-        print(f"Root transform: pos={root_xform.p}, quat={root_xform.q}")
-        print(f"Child local transform: pos={child_local_xform.p}, quat={child_local_xform.q}")
-        print(f"Expected composed: pos={expected_pos}, quat={expected_quat}")
-        print(f"Actual joint_q: pos={joint_pos}, quat={joint_quat}")

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.

newton/_src/utils/import_mjcf.py (2)

577-584: base_joint dict branch sets child to XML root instead of link

Passing child=root (the XML element) into builder.add_joint is a bug. It should be the body index link.

Apply this diff:

-                    base_joint["child"] = root
+                    base_joint["child"] = link

---

606-614: Rotate joint_pos into the parent frame when building parent_xform

parent_xform should be T_parent->joint = T_parent->child ∘ T_child->joint. Given joint_pos is specified in the child frame, you need to rotate it by body_ori_for_joints before adding to body_pos_for_joints. Without this, anchors are incorrect whenever the child has a rotation.

Apply this diff:

-                    parent_xform=wp.transform(body_pos_for_joints + joint_pos, body_ori_for_joints),
+                    parent_xform=wp.transform(
+                        body_pos_for_joints + wp.quat_rotate(body_ori_for_joints, joint_pos),
+                        body_ori_for_joints,
+                    ),

If helpful, I can add a minimal test where a child body has a nontrivial orientation and a joint_pos offset, verifying the computed anchor matches T_parent->child ∘ T_child->joint.

newton/_src/utils/import_mjcf.py (1)

593-595: Good: default joint_pos is now a wp.vec3

This resolves the Warp deprecation warning and prevents type-mismatch at runtime. Thanks for addressing the earlier feedback consistently in both branches.

newton/_src/utils/import_mjcf.py (2)

459-476: World transform composition is correct; minor naming/consistency nit

local_xform and world_xform are composed correctly, and body_pos_for_joints/body_ori_for_joints are cleanly separated for joint construction. No issues here.

If you want to avoid confusion later, consider renaming body_pos_for_joints to child_pos_in_parent and body_ori_for_joints to child_ori_in_parent to make the frame explicit.

---

89-93: Deduplicate and harden xform normalization with up-axis alignment

Minor cleanup: xform is normalized at Lines 89–93 and then re-normalized/checked again at Lines 198–203. Because xform is set to a non-None value at Line 90, the “if xform is None” branch at Line 199 is effectively dead. Also, re-applying wp.transform(*xform) when xform is already a wp.transform relies on unpacking semantics that may not be guaranteed.

Consider normalizing exactly once near the up-axis alignment and handling both tuple and wp.transform inputs explicitly.

Apply this diff to simplify:

@@
-    if xform is None:
-        xform = wp.transform()
-    else:
-        xform = wp.transform(*xform)
+    # Defer normalization until after up-axis alignment is known
@@
-    axis_xform = wp.transform(wp.vec3(0.0), quat_between_axes(up_axis, builder.up_axis))
-    if xform is None:
-        xform = axis_xform
-    else:
-        xform = wp.transform(*xform) * axis_xform
+    axis_xform = wp.transform(wp.vec3(0.0), quat_between_axes(up_axis, builder.up_axis))
+    # Normalize xform into a wp.transform and compose up-axis adjustment exactly once
+    if xform is None:
+        xform = axis_xform
+    elif isinstance(xform, wp.transform):
+        xform = xform * axis_xform
+    else:
+        xform = wp.transform(*xform) * axis_xform

Also applies to: 198-203

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.

newton/_src/utils/import_mjcf.py (2)

545-576: base_joint dict: child set to XML root instead of the body index

When base_joint is a dict, you assign child=root (the XML element), which will break joint creation. This should be the newly created body index (link).

Apply this diff:

-                elif isinstance(base_joint, dict):
-                    base_joint["parent"] = -1
-                    base_joint["child"] = root
-                    base_joint["parent_xform"] = base_parent_xform
-                    base_joint["child_xform"] = base_child_xform
-                    base_joint["key"] = "base_joint"
-                    builder.add_joint(**base_joint)
+                elif isinstance(base_joint, dict):
+                    base_joint["parent"] = -1
+                    base_joint["child"] = link
+                    base_joint["parent_xform"] = base_parent_xform
+                    base_joint["child_xform"] = base_child_xform
+                    base_joint["key"] = "base_joint"
+                    builder.add_joint(**base_joint)

---

596-606: Parent joint frame ignores child's rotation for joint_pos (misplaced anchors when child is rotated)

In the non-root joint path, parent_xform uses body_pos_for_joints + joint_pos without rotating joint_pos by the child body’s orientation. This yields incorrect parent-side joint anchor when the child has a local rotation.

Apply this diff:

-                    parent_xform=wp.transform(body_pos_for_joints + joint_pos, body_ori_for_joints),
+                    parent_xform=wp.transform(
+                        body_pos_for_joints + wp.quat_rotate(body_ori_for_joints, joint_pos),
+                        body_ori_for_joints
+                    ),

Note: The special handling in the root base_joint path to not rotate the base joint’s axis is orthogonal and doesn’t apply here; regular joints should rotate their local offsets into the parent frame.

newton/_src/utils/import_mjcf.py (1)

581-583: Avoid hardcoded joint name for floating-base free joint

Hardcoding "floating_base" is brittle. Prefer the MJCF-provided freejoint name if present (or default to "_freejoint") to keep naming consistent with the non-base path.

Apply this diff:

-            elif floating is not None and floating:
-                builder.add_joint_free(link, key="floating_base")
+            elif floating is not None and floating:
+                builder.add_joint_free(link, key=freejoint_tags[0].attrib.get("name", f"{body_name}_freejoint"))

newton/tests/test_import_mjcf.py (2)

332-379: Make joint lookup resilient to name changes

Relying on the exact name "floating_body_freejoint" can be brittle. Prefer a fallback: if that name is missing, find a free joint whose name ends with "_freejoint" and contains "floating_body" (or use a model API mapping if available).

Apply this diff:

-        joint_idx = model.joint_key.index("floating_body_freejoint")
+        joint_names = model.joint_key
+        if "floating_body_freejoint" in joint_names:
+            joint_idx = joint_names.index("floating_body_freejoint")
+        else:
+            joint_idx = next(
+                i for i, n in enumerate(joint_names)
+                if n.endswith("_freejoint") and "floating_body" in n
+            )

---

230-240: Optional: make quaternion asserts sign-agnostic

Quaternions q and -q represent the same rotation; numerical operations can occasionally flip signs. Current asserts may fail in such cases.

You can add a small helper and use it where you compare quats:

def assert_quat_allclose(testcase, actual, expected, **kwargs):
    actual = np.asarray(actual)
    expected = np.asarray(expected)
    if np.dot(actual, expected) < 0.0:
        expected = -expected
    np.testing.assert_allclose(actual, expected, **kwargs)

Then replace np.testing.assert_allclose(body_quat, expected, ...) with assert_quat_allclose(self, body_quat, expected, ...).

Also applies to: 278-280, 311-313, 329-331, 426-428, 437-439, 516-519, 531-533, 543-545, 555-557, 567-569

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.