Fix xform handling for floating=False in MJCF/URDF importers (#1168)

eric-heiden · web-flow · commit a4eaa354f5cf · 2025-12-02T11:25:47.000-08:00
Signed-off-by: Eric Heiden &lt;eric-heiden@outlook.com&gt;
diff --git a/newton/_src/utils/import_mjcf.py b/newton/_src/utils/import_mjcf.py
@@ -726,7 +726,7 @@ def parse_body(
             elif floating is not None and floating:
                 joint_indices.append(builder.add_joint_free(link, key="floating_base"))
             else:
-                joint_indices.append(builder.add_joint_fixed(-1, link, parent_xform=_xform, key="fixed_base"))
+                joint_indices.append(builder.add_joint_fixed(-1, link, parent_xform=world_xform, key="fixed_base"))
 
         else:
             joint_pos = joint_pos[0] if len(joint_pos) > 0 else wp.vec3(0.0, 0.0, 0.0)
diff --git a/newton/tests/test_import_mjcf.py b/newton/tests/test_import_mjcf.py
@@ -1037,6 +1037,91 @@ def test_joint_stiffness_damping(self):
             self.assertAlmostEqual(joint_target_ke[dof_idx], expected["target_ke"], places=1)
             self.assertAlmostEqual(joint_target_kd[dof_idx], expected["target_kd"], places=1)
 
+    def test_xform_with_floating_false(self):
+        """Test that xform parameter is respected when floating=False"""
+        local_pos = wp.vec3(1.0, 2.0, 3.0)
+        local_quat = wp.quat_rpy(0.5, -0.8, 0.7)
+        local_xform = wp.transform(local_pos, local_quat)
+
+        # Create a simple MJCF with a body that has a freejoint
+        mjcf_content = f"""<?xml version="1.0" encoding="utf-8"?>
+<mujoco model="test_xform">
+    <worldbody>
+        <body name="test_body" pos="{local_pos.x} {local_pos.y} {local_pos.z}" quat="{local_quat.w} {local_quat.x} {local_quat.y} {local_quat.z}">
+            <freejoint/>
+            <geom type="sphere" size="0.1"/>
+        </body>
+    </worldbody>
+</mujoco>
+"""
+        # Create a non-identity transform to apply
+        xform_pos = wp.vec3(5.0, 10.0, 15.0)
+        xform_quat = wp.quat_from_axis_angle(wp.vec3(0.0, 0.0, 1.0), wp.pi / 4.0)  # 45 degree rotation around Z
+        xform = wp.transform(xform_pos, xform_quat)
+
+        # Parse with floating=False and the xform
+        builder = newton.ModelBuilder()
+        builder.add_mjcf(mjcf_content, floating=False, xform=xform)
+        model = builder.finalize()
+
+        # Verify the model has a fixed joint
+        self.assertEqual(model.joint_count, 1)
+        joint_type = model.joint_type.numpy()[0]
+        self.assertEqual(joint_type, newton.JointType.FIXED)
+
+        # Verify the fixed joint has the correct parent_xform
+        # The joint_X_p should match the world_xform (xform * local_xform)
+        joint_X_p = model.joint_X_p.numpy()[0]
+
+        expected_xform = xform * local_xform
+
+        # Check position
+        np.testing.assert_allclose(
+            joint_X_p[:3],
+            [expected_xform.p[0], expected_xform.p[1], expected_xform.p[2]],
+            rtol=1e-5,
+            atol=1e-5,
+            err_msg="Fixed joint parent_xform position does not match expected xform",
+        )
+
+        # Check quaternion (note: quaternions can be negated and still represent the same rotation)
+        expected_quat = np.array([expected_xform.q[0], expected_xform.q[1], expected_xform.q[2], expected_xform.q[3]])
+        actual_quat = joint_X_p[3:7]
+
+        # Check if quaternions match (accounting for q and -q representing the same rotation)
+        quat_match = np.allclose(actual_quat, expected_quat, rtol=1e-5, atol=1e-5) or np.allclose(
+            actual_quat, -expected_quat, rtol=1e-5, atol=1e-5
+        )
+        self.assertTrue(
+            quat_match,
+            f"Fixed joint parent_xform quaternion does not match expected xform.\n"
+            f"Expected: {expected_quat}\nActual: {actual_quat}",
+        )
+
+        # Verify body_q after eval_fk also matches the expected transform
+        state = model.state()
+        newton.eval_fk(model, model.joint_q, model.joint_qd, state)
+
+        body_q = state.body_q.numpy()[0]
+        np.testing.assert_allclose(
+            body_q[:3],
+            [expected_xform.p[0], expected_xform.p[1], expected_xform.p[2]],
+            rtol=1e-5,
+            atol=1e-5,
+            err_msg="Body position after eval_fk does not match expected xform",
+        )
+
+        # Check body quaternion
+        body_quat = body_q[3:7]
+        quat_match = np.allclose(body_quat, expected_quat, rtol=1e-5, atol=1e-5) or np.allclose(
+            body_quat, -expected_quat, rtol=1e-5, atol=1e-5
+        )
+        self.assertTrue(
+            quat_match,
+            f"Body quaternion after eval_fk does not match expected xform.\n"
+            f"Expected: {expected_quat}\nActual: {body_quat}",
+        )
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)
diff --git a/newton/tests/test_import_urdf.py b/newton/tests/test_import_urdf.py
@@ -19,6 +19,7 @@
 from unittest.mock import patch
 
 import numpy as np
+import warp as wp
 
 import newton
 import newton.examples
@@ -204,8 +205,12 @@ def parse_urdf(urdf: str, builder: newton.ModelBuilder, res_dir: dict[str, str]
         urdf: URDF file to parse
         res_dir: dict[str, str]: (filename, content): extra resources files to include in the directory"""
 
+        # Default to up_axis="Y" if not specified in kwargs
+        if "up_axis" not in kwargs:
+            kwargs["up_axis"] = "Y"
+
         if not res_dir:
-            builder.add_urdf(urdf, up_axis="Y", **kwargs)
+            builder.add_urdf(urdf, **kwargs)
             return
 
         urdf_filename = "robot.urdf"
@@ -220,7 +225,7 @@ def parse_urdf(urdf: str, builder: newton.ModelBuilder, res_dir: dict[str, str]
 
             # Parse the URDF file
             urdf_path = Path(temp_dir) / urdf_filename
-            builder.add_urdf(str(urdf_path), up_axis="Y", **kwargs)
+            builder.add_urdf(str(urdf_path), **kwargs)
 
     def test_sphere_urdf(self):
         # load a urdf containing a sphere with r=0.5 and pos=(1.0,2.0,3.0)
@@ -462,6 +467,95 @@ def test_joint_body_ordering_bfs(self):
         ]
         assert builder.joint_key == ["fixed_base", "joint_2", "joint_1", "joint_2b", "joint_2a", "joint_1a", "joint_1b"]
 
+    def test_xform_with_floating_false(self):
+        """Test that xform parameter is respected when floating=False"""
+
+        # Create a simple URDF with a link (no position/orientation in URDF for root link)
+        urdf_content = """<?xml version="1.0" encoding="utf-8"?>
+<robot name="test_xform">
+    <link name="base_link">
+        <inertial>
+            <mass value="1.0"/>
+            <inertia ixx="0.1" ixy="0.0" ixz="0.0" iyy="0.1" iyz="0.0" izz="0.1"/>
+        </inertial>
+        <visual>
+            <geometry>
+                <sphere radius="0.1"/>
+            </geometry>
+        </visual>
+    </link>
+</robot>
+"""
+        # Create a non-identity transform to apply
+        xform_pos = wp.vec3(5.0, 10.0, 15.0)
+        xform_quat = wp.quat_from_axis_angle(wp.vec3(0.0, 0.0, 1.0), wp.pi / 4.0)  # 45 degree rotation around Z
+        xform = wp.transform(xform_pos, xform_quat)
+
+        # Parse with floating=False and the xform
+        # Use up_axis="Z" to match builder default and avoid axis transformation
+        builder = newton.ModelBuilder()
+        self.parse_urdf(urdf_content, builder, floating=False, xform=xform, up_axis="Z")
+        model = builder.finalize()
+
+        # Verify the model has a fixed joint
+        self.assertEqual(model.joint_count, 1)
+        joint_type = model.joint_type.numpy()[0]
+        self.assertEqual(joint_type, newton.JointType.FIXED)
+
+        # Verify the fixed joint has the correct parent_xform
+        # In URDF, the xform is applied directly to the root body (no local transform)
+        joint_X_p = model.joint_X_p.numpy()[0]
+
+        # Expected transform is just xform (URDF root links don't have position/orientation)
+        expected_xform = xform
+
+        # Check position
+        np.testing.assert_allclose(
+            joint_X_p[:3],
+            [expected_xform.p[0], expected_xform.p[1], expected_xform.p[2]],
+            rtol=1e-5,
+            atol=1e-5,
+            err_msg="Fixed joint parent_xform position does not match expected xform",
+        )
+
+        # Check quaternion (note: quaternions can be negated and still represent the same rotation)
+        expected_quat = np.array([expected_xform.q[0], expected_xform.q[1], expected_xform.q[2], expected_xform.q[3]])
+        actual_quat = joint_X_p[3:7]
+
+        # Check if quaternions match (accounting for q and -q representing the same rotation)
+        quat_match = np.allclose(actual_quat, expected_quat, rtol=1e-5, atol=1e-5) or np.allclose(
+            actual_quat, -expected_quat, rtol=1e-5, atol=1e-5
+        )
+        self.assertTrue(
+            quat_match,
+            f"Fixed joint parent_xform quaternion does not match expected xform.\n"
+            f"Expected: {expected_quat}\nActual: {actual_quat}",
+        )
+
+        # Verify body_q after eval_fk also matches the expected transform
+        state = model.state()
+        newton.eval_fk(model, model.joint_q, model.joint_qd, state)
+
+        body_q = state.body_q.numpy()[0]
+        np.testing.assert_allclose(
+            body_q[:3],
+            [expected_xform.p[0], expected_xform.p[1], expected_xform.p[2]],
+            rtol=1e-5,
+            atol=1e-5,
+            err_msg="Body position after eval_fk does not match expected xform",
+        )
+
+        # Check body quaternion
+        body_quat = body_q[3:7]
+        quat_match = np.allclose(body_quat, expected_quat, rtol=1e-5, atol=1e-5) or np.allclose(
+            body_quat, -expected_quat, rtol=1e-5, atol=1e-5
+        )
+        self.assertTrue(
+            quat_match,
+            f"Body quaternion after eval_fk does not match expected xform.\n"
+            f"Expected: {expected_quat}\nActual: {body_quat}",
+        )
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)
