newton/examples/diffsim/example_diffsim_ball.py (2)

43-47: Consider extracting the gradient descent update logic to a generic utility

The gradient descent step pattern is duplicated across multiple DiffSim examples (ball, cloth, etc.). Consider extracting this to a shared utility function or base class to promote code reuse and maintainability.

---

196-241: Move non-runtime dependencies inside the function

The numpy import should be moved to the top of the file with other imports for consistency.

+import numpy as np
 import warp as wp

 import newton
 import newton.examples

...

 def check_grad(self):
-    import numpy as np  # noqa: PLC0415
-
     param = self.states[0].particle_qd

newton/examples/diffsim/example_diffsim_cloth.py (2)

36-40: Ensure thread-safe atomic operations in the COM kernel

The com_kernel uses atomic operations to accumulate the center of mass across particles. While this is correct, the division by n happens within each thread, which could lead to numerical precision issues with large particle counts. Consider accumulating the sum first and then dividing once.

 @wp.kernel
 def com_kernel(positions: wp.array(dtype=wp.vec3), n: int, com: wp.array(dtype=wp.vec3)):
     tid = wp.tid()
 
     # compute center of mass
-    wp.atomic_add(com, 0, positions[tid] / float(n))
+    wp.atomic_add(com, 0, positions[tid])

Then divide after the kernel launch:

# After kernel launch in forward()
self.com[0] = self.com[0] / float(self.model.particle_count)

---

121-124: Consider making camera positioning configurable

The hard-coded camera position should be configurable through command-line arguments or configuration to improve usability across different viewing scenarios.

 if isinstance(self.viewer, newton.viewer.ViewerGL):
-    pos = type(self.viewer.camera.pos)(12.5, 0.0, 2.0)
+    # Allow camera position override via args
+    default_pos = (12.5, 0.0, 2.0)
+    pos = type(self.viewer.camera.pos)(*default_pos)
     self.viewer.camera.pos = pos

newton/examples/diffsim/example_diffsim_spring_cage.py (1)

204-204: Address the TODO comment for drawing springs

The TODO comment indicates that springs should be drawn inside log_state(). This functionality appears to be missing from the framework.

Would you like me to help implement the spring visualization functionality within the log_state() method or create a new issue to track this enhancement?

newton/examples/diffsim/example_diffsim_bear.py (3)

33-36: Organize imports following Python conventions

The imports should be organized in the standard order: standard library, third-party packages, then local imports. Also, warp.optim should be grouped with the main warp import.

 import numpy as np
-import warp as wp
-import warp.optim
 from pxr import Usd, UsdGeom
+
+import warp as wp
+import warp.optim
 
 import newton
 import newton.examples

---

75-77: Missing docstring for the tanh function

The tanh function applies activation scaling but lacks documentation explaining the purpose of ACTIVATION_STRENGTH.

 @wp.func
 def tanh(x: float):
+    """Apply tanh activation with scaling factor for controlling activation magnitude."""
     return wp.tanh(x) * ACTIVATION_STRENGTH

---

51-58: Document the loss function's objective

The loss kernel implements a specific objective function that rewards forward velocity while penalizing lateral/vertical movement. This should be documented for clarity.

 @wp.kernel
 def loss_kernel(com: wp.array(dtype=wp.vec3), loss: wp.array(dtype=float)):
+    """Compute loss that encourages forward (x) velocity while penalizing lateral (y) and vertical (z) movement."""
     tid = wp.tid()
     vx = com[tid][0]
     vy = com[tid][1]
     vz = com[tid][2]
+    # Maximize forward velocity (vx) while minimizing lateral and vertical velocities
     delta = wp.sqrt(vy * vy) + wp.sqrt(vz * vz) - vx
 
     wp.atomic_add(loss, 0, delta)

newton/tests/test_examples.py (1)

356-409: Consider adding a comment explaining the frame calculation.

The num_frames calculations (e.g., 250 * 36, 64 * 120) represent train_iters * sim_steps, which may not be immediately obvious to future maintainers.

 add_example_test(
     TestDiffSimExamples,
     name="diffsim.example_diffsim_ball",
     devices=test_devices,
-    test_options={"num_frames": 250 * 36},  # train_iters * sim_steps
+    test_options={"num_frames": 250 * 36},  # train_iters * sim_steps = 250 * 36
     test_options_cpu={"num_frames": 2 * 36},
     use_viewer=True,
 )

newton/examples/diffsim/example_diffsim_soft_body.py (2)

112-112: Consider making soft_contact_margin configurable.

The hard-coded soft_contact_margin=0.001 appears both here and in line 248. Consider making this a class attribute for easier tuning.

 class Example:
     def __init__(self, viewer, material_behavior="anisotropic", verbose=False):
         # setup simulation parameters first
         self.fps = 60
         self.frame = 0
         self.frame_dt = 1.0 / self.fps
         self.sim_steps = 60  # 1.0 seconds
         self.sim_substeps = 16
         self.sim_substeps_count = self.sim_steps * self.sim_substeps
         self.sim_dt = self.frame_dt / self.sim_substeps
+        self.soft_contact_margin = 0.001
 
         # ... rest of __init__ ...
 
-        self.contacts = self.model.collide(self.states[0], soft_contact_margin=0.001)
+        self.contacts = self.model.collide(self.states[0], soft_contact_margin=self.soft_contact_margin)

And update line 248 similarly:

-            self.contacts = self.model.collide(self.states[i], soft_contact_margin=0.001)
+            self.contacts = self.model.collide(self.states[i], soft_contact_margin=self.soft_contact_margin)

---

323-324: Remove empty test method or add TODO comment.

The empty test method doesn't serve any purpose. Either remove it or add a TODO comment explaining what testing functionality is planned.

-    def test(self):
-        pass

newton/examples/diffsim/example_diffsim_drone.py (3)

207-207: Address the TODO comment about mesh scaling.

The comment indicates that mesh scaling needs proper handling. This could lead to incorrect collision detection.

The TODO comment suggests mesh scaling isn't handled properly, which could affect collision accuracy. Would you like me to help implement proper mesh scaling or create an issue to track this?

---

549-564: Check for USD viewer type before USD-specific operations.

The USD-specific code assumes pxr is available whenever the viewer is a ViewerUSD. Consider adding error handling for the import.

         if isinstance(self.viewer, newton.viewer.ViewerUSD):
-            from pxr import UsdGeom  # noqa: PLC0415
+            try:
+                from pxr import UsdGeom  # noqa: PLC0415
+            except ImportError:
+                print("Warning: USD Python bindings not available, skipping drone geometry customization")
+                return
 
             # Remove the default drone geometries.

---

776-776: Use more robust state swapping.

Direct tuple assignment for swapping could be more clearly expressed and is more maintainable using standard Python idioms.

-            (self.drone.states[0], self.drone.states[1]) = (self.drone.states[1], self.drone.states[0])
+            self.drone.states[0], self.drone.states[1] = self.drone.states[1], self.drone.states[0]

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: shi-eric
PR: newton-physics/newton#521
File: newton/examples/example_cloth_hanging.py:36-36
Timestamp: 2025-08-12T05:17:34.423Z
Learning: The Newton migration guide (docs/migration.rst) is specifically for documenting how to migrate existing warp.sim functionality to Newton equivalents. New Newton-only features that didn't exist in warp.sim do not need migration documentation.

newton/examples/diffsim/example_diffsim_bear.py (1)

112-116: Ensure the bear.usd asset is present and add graceful error handling

I searched the repository (using fd -t f bear.usd) and didn’t locate bear.usd under the expected paths, so it’s unclear whether this asset is bundled with the repo or supplied at runtime. Please verify where bear.usd lives and update the code to handle its absence gracefully.

• Confirm that newton.examples.get_asset("bear.usd") returns a valid filesystem path—if the asset is hosted outside the repo, document its location or include it.
• Before opening the stage, check for file existence and raise a clear exception if missing. For example:

import os
from pxr import Usd, UsdGeom

asset_path = newton.examples.get_asset("bear.usd")
if not os.path.exists(asset_path):
    raise FileNotFoundError(f"Required USD asset not found: {asset_path}")
usd_stage = Usd.Stage.Open(asset_path)
usd_geom = UsdGeom.Mesh(usd_stage.GetPrimAtPath("/root/bear"))

tet_vertices = np.array(usd_geom.GetPointsAttr().Get())
tet_indices = np.array(usd_geom.GetPrim().GetAttribute("tetraIndices").Get())

• Optionally, wrap Usd.Stage.Open in a try/except to catch USD errors and provide additional context.

newton/examples/__init__.py (1)

186-191: LGTM! Clean integration of diffsim examples.

The new diffsim example mappings follow the established pattern and are properly organized under the newton.examples.diffsim module namespace.

newton/tests/test_examples.py (1)

352-354: LGTM!

The test class follows the same pattern as other test classes in the file.