newton/tests/test_mujoco_solver.py (1)

472-485: Updated inertias are diagonal-only; summary claims SPD with off-diagonals.

Tests still pass, but this weakens coverage for orientation updates after notify_model_changed, and contradicts the PR summary.

Consider mirroring the initial SPD construction (with off-diagonals and PD fix) for updated_inertias:

-        updated_inertias = np.zeros((self.model.body_count, 3, 3), dtype=np.float32)
-        for i in range(self.model.body_count):
-            env_idx = i // bodies_per_env
-            if env_idx == 0:
-                a = np.float32(2.5 + self.rng.uniform(0.0, 0.5))
-                b = np.float32(3.5 + self.rng.uniform(0.0, 0.5))
-                c = np.float32(min(a + b - 0.1, 4.5))
-                updated_inertias[i] = np.diag([a, b, c]).astype(np.float32)
-            else:
-                a = np.float32(3.5 + self.rng.uniform(0.0, 0.5))
-                b = np.float32(4.5 + self.rng.uniform(0.0, 0.5))
-                c = np.float32(min(a + b - 0.1, 5.5))
-                updated_inertias[i] = np.diag([a, b, c]).astype(np.float32)
+        updated_inertias = np.zeros((self.model.body_count, 3, 3), dtype=np.float32)
+        for i in range(self.model.body_count):
+            env_idx = i // bodies_per_env
+            if env_idx == 0:
+                a_base, b_base, c_max = 2.5, 3.5, 4.5
+            else:
+                a_base, b_base, c_max = 3.5, 4.5, 5.5
+            a = np.float32(a_base + self.rng.uniform(0.0, 0.5))
+            b = np.float32(b_base + self.rng.uniform(0.0, 0.5))
+            c = np.float32(min(a + b - 0.1, c_max))
+            ab = np.float32(self.rng.uniform(-0.2, 0.2))
+            ac = np.float32(self.rng.uniform(-0.2, 0.2))
+            bc = np.float32(self.rng.uniform(-0.2, 0.2))
+            inertia = np.array([[a, ab, ac], [ab, b, bc], [ac, bc, c]], dtype=np.float32)
+            eigvals = np.linalg.eigvalsh(inertia)
+            if np.any(eigvals <= 0):
+                inertia += np.eye(3, dtype=np.float32) * (np.abs(np.min(eigvals)) + 0.1)
+            updated_inertias[i] = inertia

Configuration used: Path: .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.

Learnt from: adenzler-nvidia
PR: newton-physics/newton#552
File: newton/_src/solvers/mujoco/solver_mujoco.py:0-0
Timestamp: 2025-08-18T15:56:26.587Z
Learning: In Newton's MuJoCo solver, when transforming joint axes from Newton's internal frame to MuJoCo's expected frame, use wp.quat_rotate(joint_rot, axis) not wp.quat_rotate_inv(joint_rot, axis). The joint_rot represents rotation from joint-local to body frame, so forward rotation is correct.

Learnt from: adenzler-nvidia
PR: newton-physics/newton#552
File: newton/_src/solvers/mujoco/solver_mujoco.py:0-0
Timestamp: 2025-08-18T15:56:26.587Z
Learning: In Newton's MuJoCo solver, when transforming joint axes from Newton's internal frame to MuJoCo's expected frame, use wp.quat_rotate(joint_rot, axis) not wp.quat_rotate_inv(joint_rot, axis). The joint_rot represents rotation from joint-local to body frame, so forward rotation is correct.

newton/tests/test_mujoco_solver.py (1)

420-447: Fix eigenvector reordering and make quaternion comparison sign-invariant.

• Reversing rows (newton_eigvecs[::-1]) is incorrect; reorder columns with the eigenvalues.
• Use NumPy’s eigh for symmetric matrices on host.
• Enforce right-handed basis (det > 0).
• Account for q ≡ −q and skip orientation check when eigenvalues are near-degenerate.

Apply this diff:

-                        # Get eigenvalues of both tensors
-                        newton_eigvecs, newton_eigvals = wp.eig3(newton_inertia)
-                        newton_eigvecs = np.array(newton_eigvecs)
-                        newton_eigvecs = newton_eigvecs.reshape((3, 3))
-
-                        newton_eigvals = np.array(newton_eigvals)
+                        # Symmetric eigendecomposition on host (vals asc., vecs in columns)
+                        vals, vecs = np.linalg.eigh(newton_inertia)

                         mjc_eigvals = mjc_inertia  # Already in diagonal form
                         mjc_iquat = np.roll(solver.mjw_model.body_iquat.numpy()[env_idx, mjc_idx].astype(np.float32), 1)

-                        # Sort eigenvalues in ascending order and shuffle eigenvectors accordingly
-                        sort_indices = np.argsort(newton_eigvals)
-                        newton_eigvals = newton_eigvals[sort_indices]
-                        newton_eigvecs = newton_eigvecs[:, sort_indices]
-
-                        # reverse because we want descending order of eigenvalues
-                        newton_eigvals = newton_eigvals[::-1]
-                        newton_eigvecs = newton_eigvecs[::-1]
-
-                        newton_quat = wp.quat_from_matrix(
-                            wp.matrix_from_cols(
-                                wp.vec3(newton_eigvecs[0]), wp.vec3(newton_eigvecs[1]), wp.vec3(newton_eigvecs[2])
-                            )
-                        )
-                        newton_quat = wp.normalize(newton_quat)
+                        # Sort to descending order and reorder eigenvectors by columns
+                        order = vals.argsort()[::-1]
+                        newton_eigvals = vals[order].astype(np.float32)
+                        vecs = vecs[:, order].astype(np.float32)
+                        # enforce right-handed basis
+                        if np.linalg.det(vecs) < 0.0:
+                            vecs[:, 2] *= -1.0
+                        # Build rotation from column eigenvectors and convert to a normalized quaternion (xyzw)
+                        R = wp.matrix_from_cols(
+                            wp.vec3(float(vecs[0, 0]), float(vecs[1, 0]), float(vecs[2, 0])),
+                            wp.vec3(float(vecs[0, 1]), float(vecs[1, 1]), float(vecs[2, 1])),
+                            wp.vec3(float(vecs[0, 2]), float(vecs[1, 2]), float(vecs[2, 2])),
+                        )
+                        newton_quat = wp.normalize(wp.quat_from_matrix(R))

@@
-                        for dim in range(4):
-                            self.assertAlmostEqual(
-                                float(mjc_iquat[dim]),
-                                float(newton_quat[dim]),
-                                places=5,
-                                msg=f"{msg_prefix}Inertia quaternion mismatch for body {body_idx} in environment {env_idx}",
-                            )
+                        # Quaternion compare: sign-invariant; skip when eigenvalues are nearly degenerate
+                        rel_gap12 = abs(float(newton_eigvals[0] - newton_eigvals[1])) / max(1.0, float(abs(newton_eigvals[0])))
+                        rel_gap23 = abs(float(newton_eigvals[1] - newton_eigvals[2])) / max(1.0, float(abs(newton_eigvals[1])))
+                        if min(rel_gap12, rel_gap23) > 1e-4:
+                            newton_quat_np = np.array(
+                                [float(newton_quat[0]), float(newton_quat[1]), float(newton_quat[2]), float(newton_quat[3])],
+                                dtype=np.float32,
+                            )
+                            if float(np.dot(mjc_iquat, newton_quat_np)) < 0.0:
+                                mjc_iquat = -mjc_iquat
+                            for dim in range(4):
+                                self.assertAlmostEqual(
+                                    float(mjc_iquat[dim]),
+                                    float(newton_quat_np[dim]),
+                                    places=5,
+                                    msg=f"{msg_prefix}Inertia quaternion mismatch for body {body_idx} in environment {env_idx}",
+                                )

Also applies to: 456-462

newton/tests/test_mujoco_solver.py (3)

386-406: Prefer sampling SPD inertias via principal moments + random rotation (fewer edge cases).

Your SPD “shift” fix works, but it can yield near-degenerate eigenvalues and unstable eigenvectors, making quaternion checks flaky. Consider generating inertias as I = R diag(λ1≥λ2≥λ3) Rᵀ with a random right-handed R.

Example (drop-in inside the loop; no extra deps):

# Sample principal moments (triangle inequality on principal values)
l1 = np.float32(a_base + self.rng.uniform(0.0, 0.5))
l2 = np.float32(b_base + self.rng.uniform(0.0, 0.5))
l3 = np.float32(min(l1 + l2 - 0.1, c_max))
lam = np.array(sorted([l1, l2, l3], reverse=True), dtype=np.float32)

# Random right-handed rotation
Q, _ = np.linalg.qr(self.rng.normal(size=(3, 3)).astype(np.float32))
if np.linalg.det(Q) < 0.0:
    Q[:, 2] *= -1.0

inertia = (Q @ np.diag(lam) @ Q.T).astype(np.float32)
new_inertias[i] = inertia

---

451-455: Tolerance may be too strict for float32 paths; consider places=5.

If CI flakes due to ordering ties or fp32 accumulation, relax to places=5.

---

472-489: DRY the inertia sampling logic (factor into a helper).

The updated-inertia block duplicates the initial generation. Extract to a local helper to keep tests concise.

Add near the top of the test method:

def _make_spd_inertia(self, a_base, b_base, c_max):
    a = np.float32(a_base + self.rng.uniform(0.0, 0.5))
    b = np.float32(b_base + self.rng.uniform(0.0, 0.5))
    c = np.float32(min(a + b - 0.1, c_max))
    ab = np.float32(self.rng.uniform(-0.2, 0.2))
    ac = np.float32(self.rng.uniform(-0.2, 0.2))
    bc = np.float32(self.rng.uniform(-0.2, 0.2))
    inertia = np.array([[a, ab, ac], [ab, b, bc], [ac, bc, c]], dtype=np.float32)
    eigvals = np.linalg.eigvalsh(inertia)
    if np.any(eigvals <= 0):
        inertia += np.eye(3, dtype=np.float32) * (np.abs(np.min(eigvals)) + 0.1)
    return inertia

Then replace the block with:

-            a = np.float32(a_base + self.rng.uniform(0.0, 0.5))
-            b = np.float32(b_base + self.rng.uniform(0.0, 0.5))
-            c = np.float32(min(a + b - 0.1, c_max))
-            ab = np.float32(self.rng.uniform(-0.2, 0.2))
-            ac = np.float32(self.rng.uniform(-0.2, 0.2))
-            bc = np.float32(self.rng.uniform(-0.2, 0.2))
-            inertia = np.array([[a, ab, ac], [ab, b, bc], [ac, bc, c]], dtype=np.float32)
-            eigvals = np.linalg.eigvalsh(inertia)
-            if np.any(eigvals <= 0):
-                inertia += np.eye(3, dtype=np.float32) * (np.abs(np.min(eigvals)) + 0.1)
-            updated_inertias[i] = inertia
+            updated_inertias[i] = self._make_spd_inertia(a_base, b_base, c_max)

Configuration used: Path: .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.

Learnt from: adenzler-nvidia
PR: newton-physics/newton#552
File: newton/_src/solvers/mujoco/solver_mujoco.py:0-0
Timestamp: 2025-08-18T15:56:26.587Z
Learning: In Newton's MuJoCo solver, when transforming joint axes from Newton's internal frame to MuJoCo's expected frame, use wp.quat_rotate(joint_rot, axis) not wp.quat_rotate_inv(joint_rot, axis). The joint_rot represents rotation from joint-local to body frame, so forward rotation is correct.