newton/examples/mpm/example_anymal_c_walk_on_sand.py (3)

41-57: Docstring formatting + optional quaternion normalization for numerical robustness.

• The docstring is squashed; reformat for readability.
• If inputs may drift from unit quaternions, normalize q to avoid subtle frame drift. TorchScript-friendly normalization shown below.

 @torch.jit.script
 def quat_rotate_inverse(q: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
-    """Rotate a vector by the inverse of a quaternion along the last dimension of q and v.    Args:
-    q: The quaternion in (x, y, z, w). Shape is (..., 4).
-    v: The vector in (x, y, z). Shape is (..., 3).    Returns:
-    The rotated vector in (x, y, z). Shape is (..., 3).
-    """
-    q_w = q[..., 3]  # w component is at index 3 for XYZW format
+    """
+    Rotate vectors by the inverse of a quaternion (last dim of q/v).
+    Args:
+        q: Quaternion in XYZW order, shape (..., 4).
+        v: Vectors to rotate, shape (..., 3).
+    Returns:
+        Rotated vectors in body frame, shape (..., 3).
+    """
+    # Optional: normalize to guard against drift
+    denom = torch.sqrt(torch.clamp(torch.sum(q * q, dim=-1, keepdim=True), min=1.0e-12))
+    q = q / denom
+
+    q_w = q[..., 3]  # w component is at index 3 for XYZW format
     q_vec = q[..., :3]  # xyz components are at indices 0, 1, 2
     a = v * (2.0 * q_w**2 - 1.0).unsqueeze(-1)
     b = torch.cross(q_vec, v, dim=-1) * q_w.unsqueeze(-1) * 2.0
     # for two-dimensional tensors, bmm is faster than einsum
     if q_vec.dim() == 2:
         c = q_vec * torch.bmm(q_vec.view(q.shape[0], 1, 3), v.view(q.shape[0], 3, 1)).squeeze(-1) * 2.0
     else:
         c = q_vec * torch.einsum("...i,...i->...", q_vec, v).unsqueeze(-1) * 2.0
     return a - b + c

---

275-283: Use index_select instead of gather for action reordering.

index_select is simpler and avoids constructing a broadcasted index tensor.

-            self.rearranged_act = torch.gather(self.act, 1, self.mujoco_to_lab_indices.unsqueeze(0))
+            self.rearranged_act = torch.index_select(self.act, 1, self.mujoco_to_lab_indices)

---

350-361: Fix redundant prod and clarify particle mass computation.

cell_volume is already a scalar; np.prod(cell_volume) is redundant. Also, name the variable to reflect it’s a mass proxy.

-    cell_volume = np.prod(cell_size)
+    cell_volume = np.prod(cell_size)
@@
-    volume = np.prod(cell_volume) * packing_fraction
+    particle_volume = cell_volume * packing_fraction
@@
-    builder.particle_mass = np.full(points.shape[0], volume)
+    builder.particle_mass = np.full(points.shape[0], particle_volume)

Note: this sets each particle’s “mass” proportional to cell volume; if you intend physical mass, multiply by a material density.

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: preist-nvidia
PR: newton-physics/newton#579
File: newton/examples/example_mujoco.py:350-354
Timestamp: 2025-08-19T12:27:30.629Z
Learning: In Newton examples, there's a distinction between solver parameters and Example class attributes. The Example class can have its own use_mujoco attribute for controlling example-level behavior (like CUDA graphs, rendering logic), while the solver uses use_mujoco_cpu for backend selection. These serve different purposes and should not be conflated during API renames.

newton/examples/mpm/example_anymal_c_walk_on_sand.py (1)

59-76: Great fix: observations now use zero-copy Warp→Torch views.

Switching to wp.to_torch(state.joint_q/joint_qd) eliminates host transfers and ensures tensors stay on the active device. This addresses the earlier feedback about unsafe torch.tensor(...) conversions.

newton/examples/mpm/example_anymal_c_walk_on_sand.py (3)

40-56: Tidy up docstring; optionally clarify unit-quaternion assumption.

• The docstring lines include inline “Args/Returns” on the same line; minor formatting polish helps readability.
• The math assumes q is unit-length. If that’s guaranteed upstream, mention it in the docstring; otherwise consider normalizing or asserting near-unit length to avoid drift.

Example minimal change:

-    """Rotate a vector by the inverse of a quaternion along the last dimension of q and v.    Args:
-    q: The quaternion in (x, y, z, w). Shape is (..., 4).
-    v: The vector in (x, y, z). Shape is (..., 3).    Returns:
-    The rotated vector in (x, y, z). Shape is (..., 3).
-    """
+    """Rotate vector(s) by the inverse of a quaternion.
+
+    Args:
+        q: Quaternion(s) in (x, y, z, w) format, shape (..., 4). Assumes unit length.
+        v: Vector(s), shape (..., 3).
+    Returns:
+        Rotated vector(s), shape (..., 3).
+    """

---

233-235: Avoid redundant asset download; reuse the earlier asset_path.

You already have asset_path = newton.utils.download_asset("anybotics_anymal_c") above. Reuse it to construct policy_path to save an extra lookup.

-policy_asset_path = newton.utils.download_asset("anybotics_anymal_c")
-policy_path = str(policy_asset_path / "rl_policies" / "anymal_walking_policy_physx.pt")
+policy_path = str(asset_path / "rl_policies" / "anymal_walking_policy_physx.pt")

---

260-266: Confirm CUDA-graph compatibility of viewer-driven forces.

viewer.apply_forces(self.state_0) is inside simulate_robot() which is captured. If the viewer updates forces interactively each frame on the CPU side, ensure the captured kernels read the updated device buffers on each wp.capture_launch (and that no CPU-side logic in apply_forces is required per frame).

If forces don’t update under capture, we can gate the call and inject a pre-step hook:

-            self.viewer.apply_forces(self.state_0)
+            # During graph replay, ensure forces are already staged in device buffers.
+            if self.graph is None:
+                self.viewer.apply_forces(self.state_0)

And before launching the graph each frame (after we’ve updated control but before physics):

 def step(self):
     # compute control before graph/step
     self.apply_control()
+    # Stage any interactive forces to device before graph replay
+    if self.graph:
+        self.viewer.stage_forces(self.state_0)  # implement as a no-op if not needed

Only adopt if you observe stale forces during CUDA graph replay.

Also applies to: 285-293, 299-306

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: preist-nvidia
PR: newton-physics/newton#579
File: newton/examples/example_mujoco.py:350-354
Timestamp: 2025-08-19T12:27:30.629Z
Learning: In Newton examples, there's a distinction between solver parameters and Example class attributes. The Example class can have its own use_mujoco attribute for controlling example-level behavior (like CUDA graphs, rendering logic), while the solver uses use_mujoco_cpu for backend selection. These serve different purposes and should not be conflated during API renames.

Learnt from: gdaviet
PR: newton-physics/newton#631
File: newton/examples/mpm/example_anymal_c_walk_on_sand.py:164-169
Timestamp: 2025-08-25T11:07:47.779Z
Learning: In Newton's MuJoCo solver implementation, setting TARGET_POSITION mode on floating base DOFs (first 6 DOFs) does not overconstrain the base - the solver handles this appropriately and maintains expected behavior for floating robots.