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

1-274: Add a "Differentiable (diffsim)" migration section to docs/migration.rst — required

The PR adds a diffsim example (newton/examples/diffsim/example_diffsim_ball.py) but docs/migration.rst contains no guidance for migrating warp.sim differentiable workflows; please add a migration subsection and link the example.

Files/locations to update:

• docs/migration.rst — add a "Differentiable simulation (diffsim)" subsection that documents how to migrate warp.sim diffsim code to Newton and links to the new example.
• Reference example: newton/examples/diffsim/example_diffsim_ball.py

Suggested minimal items to include in the new subsection (can be expanded into prose/examples):

• Enable gradients: builder.finalize(requires_grad=True)
• State & control: use model.state() for per-timestep states and model.control() for controls
• Contacts: use model.collide(state, soft_contact_margin=...) and pass contacts into solver.step
• Stepping: solver.step(state_i, state_{i+1}, control, contacts, dt) (see SolverSemiImplicit usage)
• Gradients: capture forward pass with wp.Tape() and call tape.backward(loss); optionally use wp.ScopedCapture + wp.capture_launch(graph) for CUDA-graph acceleration
• Point readers to newton/examples/diffsim/example_diffsim_ball.py for a complete end-to-end diffsim training loop and numeric gradient check

I can draft the exact rst text and a patch for docs/migration.rst if you want.

---

1-274: Add a unit test for the diffsim example and fix CLI flag mismatch

The tests file does not register newton.examples.diffsim.example_diffsim_ball — add a test entry. Also the example uses underscore-style argparse flags (e.g. --stage_path, --train_iters) while the test harness passes hyphen-style flags (e.g. --stage-path). Fix one of these (prefer adding hyphen aliases in the example) so the test runner's arguments are recognized.

Files / changes:

• Add a test registration to newton/tests/test_examples.py (place under TestOtherExamples):

add_example_test(
    TestOtherExamples,
    name="diffsim.example_diffsim_ball",
    devices=test_devices,
    test_options={"stage_path": "None", "train_iters": 0},
)

• Update flags in newton/examples/diffsim/example_diffsim_ball.py so the parser accepts hyphen-style options (add aliases), e.g.:

parser.add_argument(
    "--stage-path", "--stage_path",
    type=lambda x: None if x == "None" else str(x),
    default="example_bounce.usd",
    help="Path to the output USD file.",
)
parser.add_argument(
    "--train-iters", "--train_iters",
    type=int,
    default=250,
    help="Total number of training iterations.",
)

These changes will (1) register the example in the test suite and (2) ensure the test harness' hyphenated options are accepted.

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

17-17: Fix the outdated comment header.

The comment header says "Example Sim Grad Bounce" but the filename is example_diffsim_ball.py. Update the comment to match the filename for consistency.

-# Example Sim Grad Bounce
+# Example DiffSim Ball

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.

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_soft_body.py (3)

17-26: Update the comment header to reflect the actual optimization target

The comment states this example optimizes material parameters for a bounce trajectory to hit a target, but the code actually optimizes to drive the center of mass toward a target position without any bouncing behavior requirement.

-# Example FEM Bounce
+# Example Soft Body Material Optimization
 #
-# Shows how to use Warp to optimize for the material parameters of a soft body,
-# such that it bounces off the wall and floor in order to hit a target.
+# Shows how to use Warp to optimize for the material parameters of a soft body
+# to drive its center of mass toward a target position.

---

36-41: Add bounds checking for array indexing

The kernel accesses params array using modulo arithmetic but doesn't validate the resulting indices are within bounds. While the modulo operation should keep indices valid, adding explicit bounds checking would make the code more robust.

 @wp.kernel
 def assign_param(params: wp.array(dtype=wp.float32), tet_materials: wp.array2d(dtype=wp.float32)):
     tid = wp.tid()
     params_idx = 2 * wp.tid() % params.shape[0]
+    # Ensure we have at least 2 elements available from params_idx
+    if params_idx + 1 >= params.shape[0]:
+        return
     tet_materials[tid, 0] = params[params_idx]
     tet_materials[tid, 1] = params[params_idx + 1]

---

300-315: Add boundary checks for numpy array operations

The code reshapes gradient arrays without verifying the shape is compatible with the reshape operation, which could raise an exception if the array size isn't divisible by 2.

 def log_step(self):
+    if self.material_params.shape[0] % 2 != 0:
+        print(f"Warning: material_params size {self.material_params.shape[0]} is not divisible by 2")
+        return
+        
     x = self.material_params.numpy().reshape(-1, 2)
     x_grad = self.material_params.grad.numpy().reshape(-1, 2)

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.

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.

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 (2)

61-64: Replace sqrt(v^2) with abs(v) in loss computation

sqrt(v*v) is equivalent to abs(v) but more expensive and less readable. Behavior is the same (subgrad at 0).

-    delta = wp.sqrt(vy * vy) + wp.sqrt(vz * vz) - vx
+    delta = wp.abs(vy) + wp.abs(vz) - vx

---

212-228: Optional: initialize loss/com buffers within forward to reduce reliance on external zeroing

You zero self.loss/self.coms in step(), which works. For clarity and local correctness, zeroing com[frame] and loss inside forward before accumulation makes forward self-contained.

Example tweak (inside forward, before com/loss launches):

self.coms[frame].zero_()
# self.loss.zero_()  # if you want per-frame loss; otherwise keep as accumulator

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

38-39: Prefer wp.float32 cast inside kernels for type consistency

Using float(n) can promote to host float64 on some backends. Casting with wp.float32 avoids implicit widening.

-    wp.atomic_add(com, 0, positions[tid] / float(n))
+    wp.atomic_add(com, 0, positions[tid] / wp.float32(n))

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

34-36: Import SDF kernels from a public module
I didn’t find any public re-exports of box_sdf, sphere_sdf, etc. (they only live in newton._src.geometry.kernels). To avoid coupling your example to private internals, please:

• File: newton/examples/diffsim/example_diffsim_drone.py (lines 34–36)
• Add public re-exports of the SDF kernels in something like
newton/geometry/kernels.py or newton/geometry/__init__.py, for example:
python # newton/geometry/kernels.py from newton._src.geometry.kernels import ( box_sdf, box_sdf_grad, capsule_sdf, capsule_sdf_grad, cone_sdf, cone_sdf_grad, cylinder_sdf, cylinder_sdf_grad, mesh_sdf, plane_sdf, sphere_sdf, sphere_sdf_grad )
• Update your example to import from the public API:
diff - from newton._src.geometry.kernels import box_sdf, capsule_sdf, … + from newton.geometry.kernels import box_sdf, capsule_sdf, …

This will ensure the example remains stable as the internal layout evolves.

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

117-128: Add minimal tests and a docs note per PR checklist

Examples are great; the PR notes tests/docs are pending. Consider adding a smoke test per example (CPU device, tiny frame/iter counts) and updating docs/migration.rst to clarify these are Newton-native examples (not direct warp.sim API).

I can add a newton/tests/test_examples.py with parametrized runs like:

• spring_cage: 1 frame, 1 iter, device="cpu", no stage_path
• bear: 1 frame, verbose=False, device="cpu", no stage_path
• cloth: 1-2 iters, device="cpu", no stage_path
• drone: 2 frames, 2 rollouts, headless, device="cpu", no stage_path

Would you like me to draft that?

Also applies to: 186-207, 237-244

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.

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)

234-240: Out-of-bounds risk in tile GEMM when tet_count is not a multiple of TILE_TETS

dim=ceil(tet_count/8) launches a partially filled tile for the last block. tile_load/store expect full tiles; without a mask you risk invalid reads/writes.

Apply this diff to launch only full tiles and handle the tail separately:

-            wp.launch_tiled(
-                kernel=network,
-                dim=math.ceil(self.model.tet_count / TILE_TETS),
-                inputs=[self.phases[frame].reshape((self.phase_count, 1)), self.weights],
-                outputs=[self.tet_activations[frame].reshape((self.model.tet_count, 1))],
-                block_dim=TILE_THREADS,
-            )
+            # Launch full tiles
+            tile_size = 8  # TILE_TETS
+            full_tiles = self.model.tet_count // tile_size
+            if full_tiles > 0:
+                wp.launch_tiled(
+                    kernel=network,
+                    dim=full_tiles,
+                    inputs=[self.phases[frame].reshape((self.phase_count, 1)), self.weights],
+                    outputs=[self.tet_activations[frame].reshape((self.model.tet_count, 1))],
+                    block_dim=TILE_THREADS,
+                )
+            # Handle remaining rows that don't fill a tile
+            tail = self.model.tet_count - full_tiles * tile_size
+            if tail > 0:
+                wp.launch(
+                    kernel=network_tail,
+                    dim=tail,
+                    inputs=[
+                        self.phases[frame],
+                        self.weights,  # [tet_count, phase_count]
+                        self.tet_activations[frame],
+                        full_tiles * tile_size,  # row offset
+                        self.phase_count,
+                    ],
+                    outputs=[],
+                )

Add this helper kernel near the existing network kernel:

@wp.kernel
def network_tail(
    phases: wp.array(dtype=float),
    weights: wp.array2d(dtype=float),
    tet_activations: wp.array(dtype=float),
    row_offset: int,
    phase_count: int,
):
    r = row_offset + wp.tid()  # global row index
    acc = float(0.0)
    for c in range(phase_count):
        acc += weights[r, c] * phases[c]
    tet_activations[r] = wp.tanh(acc) * ACTIVATION_STRENGTH

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

39-40: Guard pxr/USD imports to avoid import-time crashes for users without USD

This example hard-imports pxr at module import time. If USD isn't installed, simply importing this module will crash. Make the dependency explicit and fail with a clear message.

Apply this diff:

-from pxr import Usd, UsdGeom
+try:
+    from pxr import Usd, UsdGeom
+except ImportError as e:
+    raise RuntimeError(
+        "OpenUSD (pxr) is required to run this example. Install USD to proceed."
+    ) from e

---

55-64: Use abs() instead of sqrt(x*x) for clarity and speed in the loss

sqrt(vyvy) and sqrt(vzvz) compute |vy| and |vz| with extra work and poor readability.

Apply this diff:

-    vx = com[tid][0]
-    vy = com[tid][1]
-    vz = com[tid][2]
-    delta = wp.sqrt(vy * vy) + wp.sqrt(vz * vz) - vx
+    v = com[tid]
+    # L1 penalty on lateral/vertical speed; reward forward speed
+    delta = wp.abs(v[1]) + wp.abs(v[2]) - v[0]

---

291-299: Resetting state0 is good; consider also zeroing tet_activations per frame (optional)

You already zero coms and loss. tet activations are overwritten each forward, so not strictly necessary, but zeroing can help avoid debugging confusion during partial runs.

No code change required if you prefer current behavior.

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

116-136: Validate material parameter lengths to match expected shapes

In the anisotropic branch you rely on tet_materials[:, :2].flatten() to produce exactly 2 * tet_count parameters; in unusual cases this assumption can break.

Consider asserting the expected length at construction time:

params = self.model.tet_materials.numpy()[:, :2].reshape(-1, 2)
assert params.shape[0] == self.model.tet_count, "Unexpected tet_materials shape"
self.material_params = wp.array(params.flatten(), dtype=wp.float32, requires_grad=True)

For isotropic mode the current approach (2 params) works with assign_param as written.

---

347-352: Guard inline pxr imports in render()

This will crash if USD isn't installed. Make rendering resilient when pxr is unavailable.

Apply this diff:

-                from pxr import Gf, UsdGeom  # noqa: PLC0415
-
-                particles_prim = self.renderer.stage.GetPrimAtPath("/root/particles")
-                particles = UsdGeom.Points.Get(self.renderer.stage, particles_prim.GetPath())
-                particles.CreateDisplayColorAttr().Set([Gf.Vec3f(1.0, 1.0, 1.0)], time=self.renderer.time)
+                try:
+                    from pxr import Gf, UsdGeom  # noqa: PLC0415
+                    particles_prim = self.renderer.stage.GetPrimAtPath("/root/particles")
+                    if particles_prim:
+                        particles = UsdGeom.Points.Get(self.renderer.stage, particles_prim.GetPath())
+                        particles.CreateDisplayColorAttr().Set([Gf.Vec3f(1.0, 1.0, 1.0)], time=self.renderer.time)
+                except ImportError:
+                    pass  # USD not available; skip coloring

---

268-276: Tape usage is fine; optional: prefer recreating over zeroing

You already recreate the tape per iteration when not using graphs. The later call to tape.zero() (Line 295) is redundant; consider dropping it to simplify.

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.

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.

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.

Learnt from: nvlukasz
PR: newton-physics/newton#519
File: newton/utils.py:30-33
Timestamp: 2025-08-12T18:07:17.267Z
Learning: In newton/_src/utils/import_usd.py, the PXR/USD modules are lazily loaded using try-except blocks within the functions (lines 109-111 and 1289-1291) rather than at module import time, making it safe to import parse_usd at the module level without causing import-time failures for users without USD dependencies.

Learnt from: nvlukasz
PR: newton-physics/newton#519
File: newton/utils.py:30-33
Timestamp: 2025-08-12T18:07:17.267Z
Learning: In newton/_src/utils/import_usd.py, the PXR/USD modules are lazily loaded using try-except blocks within the functions rather than at module import time, making it safe to import parse_usd at the module level without causing import-time failures for users without USD dependencies.

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

40-49: Clamp spring rest lengths to non-negative after gradient step

Unconstrained updates can drive rest lengths negative, destabilizing the solver.

Apply clamp in the update kernel:

 @wp.kernel(enable_backward=False)
 def apply_gradient_kernel(
     spring_rest_lengths_grad: wp.array(dtype=float),
     train_rate: float,
     spring_rest_lengths: wp.array(dtype=float),
 ):
     tid = wp.tid()
-
-    spring_rest_lengths[tid] -= spring_rest_lengths_grad[tid] * train_rate
+    # gradient descent with non-negativity clamp
+    spring_rest_lengths[tid] = wp.max(
+        0.0, spring_rest_lengths[tid] - spring_rest_lengths_grad[tid] * train_rate
+    )

---

109-111: Enable gradients on states to ensure backprop through trajectory

States are created without requires_grad=True, which can drop gradients on some backends.

-        self.states = tuple(self.model.state() for _ in range(self.num_frames * self.sim_substep_count + 1))
+        self.states = tuple(
+            self.model.state(requires_grad=True) for _ in range(self.num_frames * self.sim_substep_count + 1)
+        )

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

147-151: States should be created with requires_grad=True

Without requires_grad=True, gradients through the trajectory and contact generation can be dropped.

-        self.states = []
-        for _i in range(self.sim_steps + 1):
-            self.states.append(self.model.state())
+        self.states = []
+        for _i in range(self.sim_steps + 1):
+            self.states.append(self.model.state(requires_grad=True))

---

306-314: Gradient/tape lifecycle and state reset between iterations

• Tape: Reusing the same Tape and calling tape.zero() can be OK, but the common pattern is to recreate per iteration (or use CUDA graphs). Ensure you also clear parameter grads explicitly.
• Reset: The simulation is not reset to initial conditions between iterations, which can bias optimization. If you’re using CUDA graphs, reset in-place to keep capture valid.

Example:

-            # clear grads for next iteration
-            self.tape.zero()
-            self.loss.zero_()
-            self.com.zero_()
-            self.pos_error.zero_()
+            # clear grads for next iteration
+            self.tape.zero()
+            self.material_params.grad.zero_()
+            self.loss.zero_()
+            self.com.zero_()
+            self.pos_error.zero_()
+
+            # reset trajectory initial conditions in-place (CUDA-graph friendly)
+            init = self.model.state(requires_grad=True)
+            wp.copy(init.particle_q, self.states[0].particle_q)
+            wp.copy(init.particle_qd, self.states[0].particle_qd)

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

251-257: Tiled GEMM tail can access out-of-bounds when tet_count % TILE_TETS != 0

dim=ceil(tet_count/8) with tile load/store assumes full tiles. The final partial tile will access/write past the end.

Split the work: full tiles via launch_tiled and tail rows via a scalar kernel.

Add the tail kernel:

@wp.kernel
def network_tail(
    phases: wp.array(dtype=float),
    weights: wp.array2d(dtype=float),
    tet_activations: wp.array(dtype=float),
    row_offset: int,
    phase_count: int,
):
    r = row_offset + wp.tid()
    acc = float(0.0)
    for c in range(phase_count):
        acc += weights[r, c] * phases[c]
    tet_activations[r] = wp.tanh(acc) * ACTIVATION_STRENGTH

And modify the launch:

-            wp.launch_tiled(
-                kernel=network,
-                dim=math.ceil(self.model.tet_count / TILE_TETS),
-                inputs=[self.phases[frame].reshape((self.phase_count, 1)), self.weights],
-                outputs=[self.tet_activations[frame].reshape((self.model.tet_count, 1))],
-                block_dim=TILE_THREADS,
-            )
+            full_tiles = int(self.model.tet_count // int(TILE_TETS))
+            if full_tiles > 0:
+                wp.launch_tiled(
+                    kernel=network,
+                    dim=full_tiles,
+                    inputs=[self.phases[frame].reshape((self.phase_count, 1)), self.weights],
+                    outputs=[self.tet_activations[frame].reshape((self.model.tet_count, 1))],
+                    block_dim=TILE_THREADS,
+                )
+            tail = int(self.model.tet_count - full_tiles * int(TILE_TETS))
+            if tail > 0:
+                wp.launch(
+                    kernel=network_tail,
+                    dim=tail,
+                    inputs=[
+                        self.phases[frame],
+                        self.weights,
+                        self.tet_activations[frame],
+                        full_tiles * int(TILE_TETS),
+                        self.phase_count,
+                    ],
+                    outputs=[],
+                )

---

260-270: Recompute contacts per substep; one-time contacts become stale as the mesh moves

Contacts are computed once in init and reused. For deformables, recompute each substep.

-            for i in range(self.sim_substeps):
-                self.states[frame * self.sim_substeps + i].clear_forces()
-                self.solver.step(
-                    self.states[frame * self.sim_substeps + i],
-                    self.states[frame * self.sim_substeps + i + 1],
-                    self.control,
-                    self.contacts,
-                    self.sim_dt,
-                )
+            for i in range(self.sim_substeps):
+                idx = frame * self.sim_substeps + i
+                self.states[idx].clear_forces()
+                self.contacts = self.model.collide(self.states[idx], soft_contact_margin=40.0)
+                self.solver.step(
+                    self.states[idx],
+                    self.states[idx + 1],
+                    self.control,
+                    self.contacts,
+                    self.sim_dt,
+                )

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

125-135: Non-CUDA path now initializes self.graph; good catch addressed.

The else branch setting self.graph = None prevents accidental AttributeError later. This resolves the earlier concern.

---

153-172: Accessing self.tape outside capture context is brittle; rely on x.grad instead.

When using CUDA graph replay, the Tape instance created during capture isn’t re-created per step. Accessing self.tape.gradients after capture_launch may be invalid or stale. Use x.grad consistently (Warp arrays carry .grad) and guard logging when gradients are available.

Apply this diff:

             # gradient descent step
             x = self.states[0].particle_qd
             wp.launch(step_kernel, dim=len(x), inputs=[x, x.grad, self.train_rate])

-            x_grad = self.tape.gradients[self.states[0].particle_qd]
+            # Prefer array-backed gradients; Tape dictionary may be invalid after CUDA graph replay
+            x_grad = x.grad if hasattr(x, "grad") else None

-            if self.verbose:
-                print(f"Train iter: {self.train_iter} Loss: {self.loss}")
-                print(f"    x: {x} g: {x_grad}")
+            if self.verbose:
+                print(f"Train iter: {self.train_iter} Loss: {self.loss}")
+                if x_grad is not None:
+                    print(f"    x: {x} g: {x_grad}")
+                else:
+                    print(f"    x: {x} (no grad available)")
 
             # clear grads for next iteration
             self.tape.zero()

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

136-137: Fix invalid cast: wp.float is not a valid Warp type (use wp.float32 or Python float)

As written, this will fail JIT compilation. Use wp.float32(...) or cast with float(...).

-    discount = 0.8 ** wp.float(horizon_length - step - 1) / wp.float(horizon_length) ** 2.0
+    discount = wp.pow(wp.float32(0.8), wp.float32(horizon_length - step - 1)) / (wp.float32(horizon_length) ** 2.0)

newton/_src/viewer/gl/icon.py (8)

1-1: Filename/header mismatch – align module name with usage.

Header says "generate_icons.py" while the file is icon.py. Either rename the file to generate_icons.py or update the header comment to avoid confusion and broken references.

---

6-9: Clarify expectations in the docstring (size validation and platform font caveats).

Briefly document that size must be a positive integer and that color rendering depends on availability of a platform emoji font.

Apply this diff to add a simple input check:

 def create_and_save_emoji_png(character: str, size: int, filename: str):
     """
     Renders a Unicode character onto a transparent PNG and saves it.
     """
+    if size <= 0:
+        raise ValueError(f"size must be a positive integer, got {size}")

---

18-27: Avoid hardcoding a Windows-only font; add cross-platform font discovery with graceful fallback.

As written, most environments won’t find seguiemj.ttf, causing monochrome fallback. Provide OS-specific candidates and warn if none found.

Apply this diff:

-    font_path = "seguiemj.ttf"
-    font_size = int(size * 0.8)  # Adjust font size relative to image size
-
-    try:
-        font = ImageFont.truetype(font_path, size=font_size)
-    except OSError:
-        print(f"Warning: Font '{font_path}' not found. Using default font.")
-        print("The icon may not render in color.")
-        font = ImageFont.load_default()
+    font_size = int(size * 0.8)  # Adjust font size relative to image size
+
+    # Try common color-emoji fonts by platform
+    import sys
+    font = None
+    candidates = []
+    if sys.platform.startswith("win"):
+        candidates = [
+            r"C:\Windows\Fonts\seguiemj.ttf",
+            "seguiemj.ttf",
+        ]
+    elif sys.platform == "darwin":
+        candidates = ["/System/Library/Fonts/Apple Color Emoji.ttc"]
+    else:
+        candidates = [
+            "/usr/share/fonts/truetype/noto/NotoColorEmoji.ttf",
+            "/usr/share/fonts/NotoColorEmoji.ttf",
+            "/usr/local/share/fonts/NotoColorEmoji.ttf",
+        ]
+
+    for path in candidates:
+        try:
+            font = ImageFont.truetype(path, size=font_size)
+            break
+        except OSError:
+            continue
+
+    if font is None:
+        import warnings
+        warnings.warn(
+            "No color-emoji font found; falling back to default font. "
+            "Icons may render in monochrome."
+        )
+        font = ImageFont.load_default()

---

28-32: Remove no-op textbbox and “+2” fudge; center reliably.

The textbbox result is unused, and the +2 offset is a magic number. Rely on anchor="mm" for centering.

Apply this diff:

-    # Calculate position to center the character
-    draw.textbbox((0, 0), character, font=font, spacing=0)
-    x = size // 2
-    y = size // 2 + 2  # +2 fudge factor
+    # Center the character using midpoint anchoring
+    x = size // 2
+    y = size // 2

---

33-35: Be robust to older Pillow versions lacking embedded_color.

Some environments may not support the embedded_color parameter; fall back gracefully.

Apply this diff:

-    # Draw the character onto the image
-    draw.text((x, y), anchor="mm", text=character, font=font, embedded_color=True)
+    # Draw the character onto the image
+    try:
+        draw.text((x, y), text=character, font=font, anchor="mm", embedded_color=True)
+    except TypeError:
+        # Pillow without embedded_color support
+        draw.text((x, y), text=character, font=font, anchor="mm")

---

37-38: Avoid prints in library code for non-errors.

Either return the path and let the caller log, or use logging at INFO level. If you keep prints, consider making them optional.

Example minimal change:

-    image.save(filename, "PNG")
-    print(f"Successfully created {filename}")
+    image.save(filename, "PNG")
+    # Consider: use logging.info(...) or return filename instead of printing.

---

3-3: Optional: Lazy-import Pillow to avoid hard dependency on import.

If this module is imported as part of the viewer, top-level PIL imports can fail in environments without Pillow. Moving them inside the function localizes the dependency.

Replace the top import with local imports inside the function:

-from PIL import Image, ImageDraw, ImageFont

And at the start of the function body:

from PIL import Image, ImageDraw, ImageFont

---

41-47: Consider moving the CLI example to a scripts/ tool or guard it behind a small CLI.

Keeping main here is fine for a dev utility; if this becomes part of the installed package, a dedicated entry point (e.g., newton-generate-icon) improves UX and avoids stray prints.

Happy to scaffold a tiny argparse-based CLI if you want.

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

344-345: Message/logic inconsistency around mesh size handling.

Comment says “ignore geom size attribute” but asserts “need to specify size for mesh geom”. Either remove the assert or clarify the comment (MJCF mesh size often unused; scale is taken from asset/geom.mesh scale).

newton/tests/test_import_mjcf.py (1)

73-119: Make the “default maxhullvert” assertion robust (avoid magic number 64).

Instead of asserting against a hardcoded 64, import the constant used by Mesh (MESH_MAXHULLVERT) to keep tests aligned with defaults.

Apply this diff:

-                self.assertEqual(meshes[2].maxhullvert, 64)  # Default value
+                from newton._src.geometry import MESH_MAXHULLVERT
+                self.assertEqual(meshes[2].maxhullvert, MESH_MAXHULLVERT)  # Default value

newton/examples/example_cartpole.py (1)

42-47: Scope USD import to the cartpole prim for robustness and performance

Specifying the root path limits parsing to only the /cartPole subtree (as confirmed in examples/assets/cartpole.usda), avoiding unrelated prims and speeding up load time.

Apply this optional refactor in newton/examples/example_cartpole.py (around lines 42–47):

 newton.utils.parse_usd(
-    newton.examples.get_asset("cartpole.usda"),
+    newton.examples.get_asset("cartpole.usda"),
     articulation_builder,
     enable_self_collisions=False,
     collapse_fixed_joints=True,
+    root_path="/cartPole",
 )

newton/_src/geometry/inertia.py (1)

42-43: Docstrings should reflect COM-centric returns for consistency across shapes

compute_cone_inertia correctly documents (mass, center of mass, inertia). The sphere, capsule, cylinder, and box docstrings still say “(mass, inertia) … around the origin” despite returning (m, wp.vec3(), I). Align these to reduce confusion.

Apply these docstring tweaks:

@@
-    Returns:
-
-        A tuple of (mass, inertia) with inertia specified around the origin
+    Returns:
+        A tuple of (mass, center of mass, inertia) with inertia specified around the center of mass.
+        For a sphere, the center of mass is at the origin.
@@
-    Returns:
-
-        A tuple of (mass, inertia) with inertia specified around the origin
+    Returns:
+        A tuple of (mass, center of mass, inertia) with inertia specified around the center of mass.
+        For this capsule (aligned with z), the center of mass is at the origin.
@@
-    Returns:
-
-        A tuple of (mass, inertia) with inertia specified around the origin
+    Returns:
+        A tuple of (mass, center of mass, inertia) with inertia specified around the center of mass.
+        For this cylinder (aligned with z), the center of mass is at the origin.
@@
-    Returns:
-
-        A tuple of (mass, inertia) with inertia specified around the origin
+    Returns:
+        A tuple of (mass, center of mass, inertia) with inertia specified around the center of mass.
+        For an axis-aligned box centered at the origin, the center of mass is at the origin.

Also applies to: 65-66, 94-95, 171-172

newton/_src/viewer/__init__.py (1)

41-49: Make viewer backend imports lazy to avoid hard import-time dependency on optional packages

Directly importing ViewerGL / ViewerRerun / ViewerUSD at module import time can fail if optional deps (e.g., pyglet, rerun, USD) aren’t installed. Prefer lazy attribute resolution so imports happen only on first use.

Apply this diff:

-from .viewer_gl import ViewerGL
-from .viewer_rerun import ViewerRerun
-from .viewer_usd import ViewerUSD
+def __getattr__(name):
+    if name == "ViewerGL":
+        from .viewer_gl import ViewerGL as _ViewerGL
+        return _ViewerGL
+    if name == "ViewerRerun":
+        from .viewer_rerun import ViewerRerun as _ViewerRerun
+        return _ViewerRerun
+    if name == "ViewerUSD":
+        from .viewer_usd import ViewerUSD as _ViewerUSD
+        return _ViewerUSD
+    raise AttributeError(name)

newton/viewer.py (1)

20-31: Top-level viewer re-exports are fine; consider guarding optional deps

Re-exporting the new viewer classes is good. If any backend has heavy/optional dependencies, consider mirroring the lazy import pattern here (via a module-level getattr) to avoid import-time failures for users who only need a subset.

Would you like a patch to add a getattr that lazily imports ViewerGL/ViewerRerun/ViewerUSD and emits helpful ImportError messages?

newton/examples/example_selection_cartpole.py (1)

56-62: Limit USD parsing to the /cartPole prim for faster imports

Verified that the cartpole.usda asset defines its bodies under the /cartPole prim. Passing root_path="/cartPole" to parse_usd will restrict parsing to the CartPole articulation, improving import robustness and performance on larger stages.

• File: newton/examples/example_selection_cartpole.py, lines 56–62
Add the root_path parameter to the parse_usd call:

         newton.utils.parse_usd(
             newton.examples.get_asset("cartpole.usda"),
             articulation_builder,
             xform=wp.transform((0.0, 0.0, 2.0), wp.quat_identity()),
             collapse_fixed_joints=COLLAPSE_FIXED_JOINTS,
             enable_self_collisions=False,
+            root_path="/cartPole",
         )

newton/tests/test_inertia.py (1)

253-323: Cone inertia mesh-vs-analytic test looks correct; consider minor robustness tweaks.

The mesh construction (apex at +half_height, base at -half_height) and comparisons against compute_cone_inertia are consistent. Two small suggestions:

• Use a tolerance that scales with the matrix norm to treat off-diagonals (near zero) proportionally, e.g., atol = 1e-3 * np.linalg.norm(I_cone).
• 500 segments is fine but may be heavy in CI; consider parametrizing resolution if runtime becomes an issue.

Example for a norm-based tolerance:

tol_I = 1e-3 * np.linalg.norm(np.array(I_cone))
assert_np_equal(np.array(I_mesh), np.array(I_cone), tol=tol_I)

newton/tests/test_pendulum_revolute_vs_d6.py (1)

55-61: Minor: qd0_model is computed but unused.

Can be removed to keep the test tidy.

newton/tests/test_import_usd.py (1)

431-527: Force limits propagation test is good; small robustness nit.

Asserting joint_effort_limit via joint_qd_start is fine here since each joint exposes one relevant DOF. If future scenes include multi-DOF joints, consider locating the exact DOF by axis/type to avoid assuming the first DOF belongs to the drive being tested.

newton/_src/sim/model.py (1)

291-293: articulation_group attribute added consistently; consider adding attribute_frequency entries.

To fully integrate with attribute introspection (Model.get_attribute_frequency), add these group attributes to attribute_frequency. Without this, get_attribute_frequency("joint_group") would raise.

Suggested additions (place alongside existing mappings):

# attributes per particle
self.attribute_frequency["particle_group"] = "particle"

# attributes per shape
self.attribute_frequency["shape_group"] = "shape"

# attributes per body
self.attribute_frequency["body_group"] = "body"

# attributes per joint
self.attribute_frequency["joint_group"] = "joint"

# attributes per articulation
self.attribute_frequency["articulation_group"] = "articulation"

newton/tests/test_cone_orientation.py (1)

45-79: SDF tests are consistent with kernels.cone_sdf orientation and formula

The selected points (apex, base center, base edge, origin, mid-height edge) are well-chosen and match the current SDF implementation.

Minor nit: consider being explicit with Warp dtype (e.g., wp.float32) to avoid device-dependent defaults creeping in CI.

-        wp_sdf_values = wp.zeros(len(points), dtype=float)
+        wp_sdf_values = wp.zeros(len(points), dtype=wp.float32)

newton/_src/solvers/mujoco/solver_mujoco.py (3)

2262-2271: Potential type mismatch: assigning a wp.transform into a NumPy float array

shape_incoming_xform is an np.ndarray of shape (shape_count, 7). Assigning a wp.transform directly to a row relies on implicit conversion that may not be guaranteed. Be explicit to avoid dtype issues.

-                        shape_incoming_xform[global_shape_idx] = mjc_tf * wp.transform_inverse(original_tf)
+                        incoming = mjc_tf * wp.transform_inverse(original_tf)
+                        incoming_np = np.array(incoming, dtype=np.float32)
+                        shape_incoming_xform[global_shape_idx, :] = incoming_np

Apply the same change to the single-env path below.

---

2272-2280: Single-env path: also use explicit assignment for shape_incoming_xform

Mirror the explicit conversion to avoid ambiguous assignment behavior.

-                    shape_incoming_xform[shape_idx] = mjc_tf * wp.transform_inverse(original_tf)
+                    incoming = mjc_tf * wp.transform_inverse(original_tf)
+                    shape_incoming_xform[shape_idx, :] = np.array(incoming, dtype=np.float32)

---

2313-2332: Avoid assigning a dict then an array to model.to_mjc_geom_index

You set model.to_mjc_geom_index twice with incompatible types: first to a dict (Line 2313), then to a wp.array. The first assignment is redundant and may confuse downstream tooling.

-            model.to_mjc_geom_index = shape_mapping  # pyright: ignore[reportAttributeAccessIssue]
+            # model.to_mjc_geom_index will be set to a dense array below.

newton/examples/example_viewer.py (2)

123-156: Mismatch between comment and positions: you lay out along Y but comment says X

The code varies the second coordinate (Y) while the comment claims “line along X-axis”. Either fix the comment or move the positions to X for clarity.

-        # Clean layout: arrange objects in a line along X-axis
+        # Clean layout: arrange objects in a line along X-axis
...
-        base_left = -4.0
+        base_x = -4.0
...
-        sphere_pos = wp.vec3(0.0, base_left, base_height + 0.3 * abs(math.sin(1.2 * t)))
+        sphere_pos = wp.vec3(base_x, 0.0, base_height + 0.3 * abs(math.sin(1.2 * t)))
...
-        base_left += spacing
+        base_x += spacing
...
-        x_box_anim = wp.array([wp.transform([0.0, base_left, base_height], qx_slow)], dtype=wp.transform)
+        x_box_anim = wp.array([wp.transform([base_x, 0.0, base_height], qx_slow)], dtype=wp.transform)
-        base_left += spacing
+        base_x += spacing
...
-        x_cone_anim = wp.array([wp.transform([0.0, base_left, base_height], qy_slow)], dtype=wp.transform)
+        x_cone_anim = wp.array([wp.transform([base_x, 0.0, base_height], qy_slow)], dtype=wp.transform)
-        base_left += spacing
+        base_x += spacing
...
-        x_cyl_anim = wp.array([wp.transform([0.0, base_left, base_height], qz_slow)], dtype=wp.transform)
+        x_cyl_anim = wp.array([wp.transform([base_x, 0.0, base_height], qz_slow)], dtype=wp.transform)
-        base_left += spacing
+        base_x += spacing
...
-        capsule_pos = wp.vec3(0.3 * math.sin(0.8 * t), base_left, base_height)
+        capsule_pos = wp.vec3(base_x, 0.3 * math.sin(0.8 * t), base_height)
-        base_left += spacing
+        base_x += spacing

---

199-199: Typo in lines path: use a sensible entity name

"coordinate_self.axes" looks accidental. Suggest renaming to “/axes” or a similar path.

-        self.viewer.log_lines("coordinate_self.axes", self.axes_begins, self.axes_ends, self.axes_colors)
+        self.viewer.log_lines("/axes", self.axes_begins, self.axes_ends, self.axes_colors)

newton/_src/viewer/viewer_rerun.py (1)

208-214: is_running() relies on a viewer process that is never spawned

self._viewer_process is never set; as written, is_running() returns True indefinitely unless close() is called. If that’s intentional (e.g., driven by an outer loop), consider documenting the behavior. Otherwise, either spawn a process for the viewer or base liveness on rerun’s connection status.

Would you like me to wire a simple health check (e.g., maintain a “last frame time” and stop after N frames when in non-interactive rerun mode), or to spawn the web-viewer in a subprocess to derive a real process handle?

newton/_src/viewer/picking.py (3)

40-54: Device/host allocation for pick_body and pick_state: verify CUDA write semantics

pick_body is allocated as pinned host memory when on CUDA, while pick_state lives on model.device. Kernels write to both. Confirm Warp kernels can safely write to the pinned host array (pick_body) from CUDA in your target versions; if not, allocate both on model.device and mirror to host only when needed.

If needed, align device placement:

-        if model and model.device.is_cuda:
-            self.pick_body = wp.array([-1], dtype=int, pinned=True)
-        else:
-            self.pick_body = wp.array([-1], dtype=int, device="cpu")
+        # Keep pick_body on the simulation device for consistent kernel writes
+        self.pick_body = wp.array([-1], dtype=int, device=(model.device if model else "cpu"))

---

77-82: Avoid host round-trip just to check active pick

is_picking() calls .numpy() which forces a host read/sync. If pick_body lives on device (see earlier comment), this incurs a transfer every call.

Consider tracking a small host-side flag updated when pick/release happens, or keep pick_body as a device array and use it only inside kernels; if you need a quick host predicate, store a Python-side shadow boolean you update in pick()/release() (and optionally after compute_pick_state_kernel).

---

83-97: Type assumptions for ray inputs

update() assumes ray_start/ray_dir are wp.vec3. If backend passes tuples/np arrays, the kernel call may fail or incur implicit conversions. Minimal validation or conversion at the boundary would harden this.

Example:

-    def update(self, ray_start, ray_dir):
+    def update(self, ray_start, ray_dir):
+        # Coerce to wp.vec3 for robustness
+        ray_start = ray_start if isinstance(ray_start, wp.vec3) else wp.vec3(*ray_start)
+        ray_dir = ray_dir if isinstance(ray_dir, wp.vec3) else wp.vec3(*ray_dir)

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

178-181: Print numeric loss value instead of array object

print(f"... Loss: {self.loss}") prints the array wrapper, not its value.

-                print(f"Train iter: {self.train_iter} Loss: {self.loss}")
+                print(f"Train iter: {self.train_iter} Loss: {self.loss.numpy()[0]:.6f}")

---

221-235: --device argument is unused

The CLI parses --device but never applies it. Align with other examples by scoping the device for the run.

-    example = Example(viewer_type=args.viewer, stage_path=args.stage_path, verbose=args.verbose)
+    with wp.ScopedDevice(args.device):
+        example = Example(viewer_type=args.viewer, stage_path=args.stage_path, verbose=args.verbose)

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

38-39: Unconditional pxr/USD import prevents running with non-USD setups

This example imports Usd/UsdGeom at module import and always loads a USD asset, even for GL/Rerun viewers. Enforce lazy/optional USD usage to avoid import-time failures for users without USD.

Leverage the project’s pattern of guarded USD imports:

-import warp.optim
-from pxr import Usd, UsdGeom
+import warp.optim
+try:
+    from pxr import Usd, UsdGeom
+except ImportError:
+    Usd = UsdGeom = None

Then guard usage in init:

-        asset_stage = Usd.Stage.Open(os.path.join(newton.examples.get_asset_directory(), "bear.usd"))
+        if Usd is None:
+            raise ImportError("usd-core package is required for this example. Install with: pip install usd-core")
+        asset_stage = Usd.Stage.Open(os.path.join(newton.examples.get_asset_directory(), "bear.usd"))

Alternatively, move the import and open into the block where you need them to avoid import-time failures entirely.

newton/_src/viewer/kernels.py (1)

191-265: Joint basis lines: supported joint types gated; colors and transforms look good.

The axis extraction and color coding are correct. Consider scaling by local geometry or collision radii later for better readability.

newton/_src/viewer/viewer_usd.py (2)

122-151: Prototype resolution is path-based; ensure mesh argument matches stored prototype path.

log_instances expects mesh to be a key in self._meshes. ViewerBase.log_shapes caches geometry and returns a path; as long as that same path is used here, this will work. If you pass a different alias/name, it raises at runtime. Add a short note to caller contract or normalize to the stored path.

---

248-272: end_frame is a no-op; consider flushing/saving periodically for long runs.

Not blocking, but giving users an option to SaveCheckpoint per N frames can be helpful for large recordings.

newton/_src/sim/builder.py (1)

3695-3746: Model finalization emits particle_group; max envs defaulted to at least 1.

Minor note: when only global (-1) entities exist, num_envs becomes 1. If true “no env” semantics are desired, consider leaving it at 0. Otherwise this is fine and simplifies downstream code.

newton/examples/assets/cartpole.usda (2)

47-66: Mass/inertia attributes on 'cart' without centerOfMass: ensure consistency

You authored diagonalInertia/mass/principalAxes on the cart but omitted physics:centerOfMass. If the importer assumes COM at the local origin, this is fine; otherwise inertia about a non-origin COM could lead to incorrect dynamics.

• If COM is assumed at origin: no change needed.
• If COM is estimated from geom: consider removing manual inertia and letting the importer compute.
• If authoring custom inertia: consider also setting physics:centerOfMass explicitly to document intent.

---

139-148: New pole-to-pole joint: check limits/damping if instability observed

You added a second revolute joint between pole1 and pole2 with axis X and identity rotations. If the double-pendulum dynamics are too fast/unstable in your solver, consider adding drive or damping attributes, or tightening limits.

newton/_src/viewer/viewer.py (3)

98-106: Respect "show_contacts" option and clear lines when disabled

log_contacts disregards the viewer option. Add an early-out so UI toggles are effective and previously drawn lines get cleared.

 def log_contacts(self, contacts, state):
     """
     Creates line segments along contact normals for rendering.
@@
-        # Get contact count (handle case where it might be zero)
+        # Honor option toggle
+        if not self.options.get("show_contacts", False):
+            # Clear any previously drawn contact lines
+            self.log_lines("/contacts", None, None, None)
+            return
+
+        # Get contact count (handle case where it might be zero)
         num_contacts = contacts.rigid_contact_count.numpy()[0]

---

150-153: Color comment/code mismatch for contacts

Comment says orange-red, code uses green. Align for clarity.

-        # Use orange-red color for contact normals
-        line_colors = (0.0, 1.0, 0.0)
+        # Use orange-red color for contact normals
+        line_colors = (1.0, 0.4, 0.0)

---

83-91: Always pass instance attribute arrays to log_instances
Passing None for scales, colors, or materials can lead to runtime errors if a backend implementation assumes all inputs are array-like. Instead, always forward the arrays and let each backend decide whether to skip re-uploading unchanged data.

• File: newton/_src/viewer/viewer.py (lines 83–91)

-            self.log_instances(
-                shapes.name,
-                shapes.mesh,
-                shapes.world_xforms,
-                shapes.scales if self.model_changed else None,
-                shapes.colors if self.model_changed else None,
-                shapes.materials if self.model_changed else None,
-            )
+            self.log_instances(
+                shapes.name,
+                shapes.mesh,
+                shapes.world_xforms,
+                shapes.scales,
+                shapes.colors,
+                shapes.materials,
+            )

This change guarantees backends always receive valid arrays and can internally handle unchanged fields.

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

631-638: Solver.step signature: verify None placeholders match API

You pass None for intermediate arguments. Ensure SolverSemiImplicit.step matches (state_in, state_out, control, contacts, dt) or equivalent in current API. If not, wire control state explicitly.

newton/_src/viewer/viewer_gl.py (3)

198-206: Cache of control data is never populated; joint_f values will be unavailable

log_state caches self._last_state but not control, yet _render_attribute_values tries to read joint_f from _last_control. Consider adding a log_control(control) or passing control into log_state to cache it, or gate UI accordingly.

---

379-384: Remove debug print of ray on right-click

Printing each pick ray will spam stdout. Drop it or guard behind a debug flag.

-            ray_start, ray_dir = self.camera.get_world_ray(x, y)
-            print(f"ray: start {ray_start}, dir {ray_dir}")
-            self.picking.pick(self._last_state, ray_start, ray_dir)
+            ray_start, ray_dir = self.camera.get_world_ray(x, y)
+            self.picking.pick(self._last_state, ray_start, ray_dir)

---

313-316: Fix truncated docstring text ("ent")

The docstring line breaks mid-word; clean it up for readability.

-                 String examples: 'w', 'a', 's', 'd', 'space', 'escape', 'ent
+                 String examples: 'w', 'a', 's', 'd', 'space', 'escape', 'enter'

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

251-257: Out-of-bounds risk in tiled GEMM when tet_count % 8 != 0

dim=ceil(tet_count / TILE_TETS) schedules a full tile for the tail without masking; tile_load/store will read/write past the end. Split full tiles and handle the tail separately.

Apply this diff:

-            wp.launch_tiled(
-                kernel=network,
-                dim=math.ceil(self.model.tet_count / TILE_TETS),
-                inputs=[self.phases[frame].reshape((self.phase_count, 1)), self.weights],
-                outputs=[self.tet_activations[frame].reshape((self.model.tet_count, 1))],
-                block_dim=TILE_THREADS,
-            )
+            full_tiles = self.model.tet_count // 8  # TILE_TETS
+            if full_tiles > 0:
+                wp.launch_tiled(
+                    kernel=network,
+                    dim=full_tiles,
+                    inputs=[self.phases[frame].reshape((self.phase_count, 1)), self.weights],
+                    outputs=[self.tet_activations[frame].reshape((self.model.tet_count, 1))],
+                    block_dim=TILE_THREADS,
+                )
+            tail = self.model.tet_count - full_tiles * 8  # TILE_TETS
+            if tail > 0:
+                wp.launch(
+                    kernel=network_tail,
+                    dim=tail,
+                    inputs=[
+                        self.phases[frame],
+                        self.weights,  # [tet_count, phase_count]
+                        self.tet_activations[frame],
+                        full_tiles * 8,  # row offset
+                        self.phase_count,
+                    ],
+                    outputs=[],
+                )

Add this helper kernel somewhere above Example.forward:

@wp.kernel
def network_tail(
    phases: wp.array(dtype=float),
    weights: wp.array2d(dtype=float),         # [tet_count, phase_count]
    tet_activations: wp.array(dtype=float),    # [tet_count]
    row_offset: int,
    phase_count: int,
):
    r = row_offset + wp.tid()
    acc = float(0.0)
    for c in range(phase_count):
        acc += weights[r, c] * phases[c]
    tet_activations[r] = wp.tanh(acc) * ACTIVATION_STRENGTH

---

262-270: Contacts must be recomputed each substep; current code reuses stale manifolds

Contacts are computed once in __init__ (Line 192) and reused, which produces incorrect dynamics as the mesh moves.

Apply this diff:

-            for i in range(self.sim_substeps):
-                self.states[frame * self.sim_substeps + i].clear_forces()
-                self.solver.step(
-                    self.states[frame * self.sim_substeps + i],
-                    self.states[frame * self.sim_substeps + i + 1],
-                    self.control,
-                    self.contacts,
-                    self.sim_dt,
-                )
+            for i in range(self.sim_substeps):
+                idx = frame * self.sim_substeps + i
+                self.states[idx].clear_forces()
+                # Recompute contacts at current state
+                self.contacts = self.model.collide(self.states[idx], soft_contact_margin=0.05)
+                self.solver.step(
+                    self.states[idx],
+                    self.states[idx + 1],
+                    self.control,
+                    self.contacts,
+                    self.sim_dt,
+                )

---

303-306: CUDA graph correctness: don’t reassign state objects after capture

Reassigning self.states[0] breaks captured pointers; the graph will still write to old buffers. Reset in-place.

Apply this diff:

-        # reset sim
-        self.sim_time = 0.0
-        self.states[0] = self.model.state(requires_grad=True)
+        # reset sim (in-place to preserve captured pointers)
+        self.sim_time = 0.0
+        init = self.model.state(requires_grad=True)
+        wp.copy(init.particle_q,  self.states[0].particle_q)
+        wp.copy(init.particle_qd, self.states[0].particle_qd)
+        wp.copy(init.body_q,      self.states[0].body_q)
+        wp.copy(init.body_qd,     self.states[0].body_qd)

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

136-137: Replace invalid wp.float casts; use wp.float32 (or Python float) to avoid JIT failure

As-is, this breaks Warp JIT. Use wp.float32(...) or float(...).

Apply:

-    discount = 0.8 ** wp.float(horizon_length - step - 1) / wp.float(horizon_length) ** 2.0
+    discount = 0.8 ** wp.float32(horizon_length - step - 1) / (wp.float32(horizon_length) ** 2.0)

---

814-826: Fix rollout rendering: degenerate lines (start == end) and incorrect body index per env

Nothing is drawn because begins == ends; also indexing uses body index i instead of env i’s first body (i*bodies_per_env).

Apply:

-            positions = np.array([x.body_q.numpy()[:, :3] for x in self.rollouts.states])
+            # World translations for the first body of each rollout env across time
+            bodies_per_env = self.rollouts.body_count
+            positions = np.stack(
+                [s.body_q.numpy()[0 : self.rollout_count * bodies_per_env : bodies_per_env, :3]
+                 for s in self.rollouts.states],
+                axis=0,
+            )
@@
-            for i in range(self.rollout_count):
-                # Flip colors, so red means best trajectory, blue worst.
-                color = wp.render.bourke_color_map(-max_cost, -min_cost, -costs[i])
-                self.viewer.log_lines(
-                    f"/rollout_{i}",
-                    wp.array(positions[:, i], dtype=wp.vec3),
-                    wp.array(positions[:, i], dtype=wp.vec3),
-                    wp.array([color], dtype=wp.vec3),
-                )
+            for i in range(self.rollout_count):
+                # Flip colors, so red means best trajectory, blue worst.
+                color = wp.render.bourke_color_map(-max_cost, -min_cost, -costs[i])
+                path = positions[:, i]  # [T, 3]
+                if path.shape[0] >= 2:
+                    begins = wp.array(path[:-1], dtype=wp.vec3)
+                    ends = wp.array(path[1:], dtype=wp.vec3)
+                    self.viewer.log_lines(
+                        f"/rollout_{i}",
+                        begins,
+                        ends,
+                        wp.array([color], dtype=wp.vec3),
+                    )

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

318-323: Guard viewer logging; USD/Rerun viewers may not expose log_state

ViewerUSD/ViewerRerun might not implement log_state; calling it unconditionally can crash. Guard or provide viewer-specific logging.

Apply this diff:

-            self.viewer.begin_frame(self.render_time)
-            self.viewer.log_state(self.states[i * self.sim_substeps])
-            self.viewer.end_frame()
+            self.viewer.begin_frame(self.render_time)
+            if hasattr(self.viewer, "log_state"):
+                self.viewer.log_state(self.states[i * self.sim_substeps])
+            else:
+                # TODO: implement per-viewer logging (e.g., log_mesh/log_instances)
+                pass
+            self.viewer.end_frame()

Please confirm that ViewerUSD and ViewerRerun expose log_state; if they do, we can drop the guard.

---

52-61: Smoother loss: prefer L2 lateral/vertical speed penalty

sqrt(vy*vy) + sqrt(vz*vz) is L1; switch to L2 (sqrt(vy*vy + vz*vz)) for smoother gradients at zero.

Apply this diff:

-    delta = wp.sqrt(vy * vy) + wp.sqrt(vz * vz) - vx
+    delta = wp.sqrt(vy * vy + vz * vz) - vx

---

48-50: Host/device constant split to avoid type gotchas

TILE_TETS as a wp.constant is used in host math (ceil/div). Use a Python int for host code; keep kernel shapes literal-8 or introduce a separate TILE_TETS_WP if desired.

Apply this diff:

-TILE_TETS = wp.constant(8)
-TILE_THREADS = 64
+TILE_TETS = 8           # host-side int; use literal 8 in kernel tile shapes
+TILE_THREADS = 64

---

183-189: Memory footprint: states allocation scales as O(num_frames * substeps)

Allocating ~24k states for defaults (300*80) is heavy. If training allows, consider a rolling buffer (sim_substeps+1 states) per frame or chunked unroll to reduce memory.

---

126-143: Asset robustness: handle missing bear.usd gracefully

Opening the USD without existence check will throw; consider checking path and raising a clear error or falling back to a procedural mesh for tests.

---

327-359: Tests and CLI ergonomics

Per PR objectives, example tests are still missing. Add a minimal smoke test that runs 1 training iter on CPU with tiny num_frames (e.g., 2) and sim_substeps (e.g., 2) to exercise forward/backward without viewers.

I can draft newton/tests/test_examples.py that:

• Instantiates Example(viewer_type="gl", num_frames=2) with wp.ScopedDevice("cpu")
• Calls example.step() once and asserts finite loss.
• Skips if assets are unavailable.

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

249-257: Guard against out-of-bounds when interpolating final control point

If t reaches the last knot, t_id + 1 can index past the last element. It’s unlikely due to how sim_tick is advanced, but add a clamp for robustness.

Apply:

-    t_id = int(t)
-    frac = t - wp.floor(t)
-    control_left = control_points[env_id, t_id, control_id]
-    control_right = control_points[env_id, t_id + 1, control_id]
+    t_id = int(t)
+    frac = t - wp.floor(t)
+    t_right = wp.min(t_id + 1, control_points.shape[1] - 1)  # or pass max index in as a kernel arg
+    control_left = control_points[env_id, t_id, control_id]
+    control_right = control_points[env_id, t_right, control_id]

If arr.shape[...] isn’t supported in kernels for your Warp version, add control_points_count as an input and clamp against control_points_count - 1.

---

551-563: Headless flag is unused and cannot be set False via CLI

headless is accepted in Example.init but not used, and argparse sets it to True by default with store_true (so it can’t be disabled). Wire it to viewers that can run headless, or drop it.

Apply:

-    parser.add_argument(
-        "--headless",
-        default=True,
-        action="store_true",
-        help="Run in headless mode, suppressing the opening of any graphical windows.",
-    )
+    parser.add_argument(
+        "--headless",
+        action="store_true",
+        default=False,
+        help="Run in headless mode, suppressing the opening of any graphical windows.",
+    )

And in init, respect headless (e.g., pass through to ViewerGL/ViewerRerun if supported).

Also applies to: 456-464

---

739-757: Seed handling and optimizer capture: ensure determinism across frames when desired

Incrementing a single-seed array once per frame is fine, but CUDA graph capture will replay the same random sequence unless the seed is mutated between capture launches. Also, reusing a single Tape across many optim steps without re-capture is correct here, but be cautious if you later add data-dependent control flow.

If you need per-frame stochasticity with graph capture:

• Include seed value as a graph input (or re-capture per target switch), or
• Move increment_seed before capture and pass the updated seed into sample_gaussian.

Also applies to: 536-547

---

160-218: Collision SDF scaling for meshes is a placeholder; clarify or fix scaling behavior

The mesh branch scales by min_scale after sampling in scaled coordinates with a TODO note. This can mis-estimate distances for non-uniform scales.

Consider:

• Uniform-only assumption with an assert, or
• Properly apply scale via transforming the query and normalizing by the correct Jacobian term, or
• Pre-bake scale into geometry and keep geo_scale = (1,1,1) for meshes.

---

798-809: Use current target value when rendering to avoid transient -1 index pitfalls

Render is called after step() in main, so this is safe, but defensively consider clamping or defaulting when self.target_idx == -1 (e.g., first ever render call outside main).

Example:

-            (0.05,),
-            wp.array([wp.transform(self.targets[self.target_idx], wp.quat_identity())], dtype=wp.transform),
+            (0.05,),
+            wp.array([wp.transform(self.targets[max(self.target_idx, 0)], wp.quat_identity())], dtype=wp.transform),

---

128-135: Control effort term: consider normalizing by 4 and/or using per-motor limits

Minor: with 4 props, wp.dot(control, control) grows with dimension. Either keep as-is and tune weights, or divide by control_dim so weights are scale-invariant across different rotor counts.

---

471-479: Consider using separate sim_dt for optimizer vs execution

Currently, optimizer and execution share sim_dt; sometimes a finer dt for optimization improves stability. Not required, but worth noting.

---

16-24: Docs/migration.rst: add a short note that this example is Newton-only, not a warp.sim parity item

Per our previous learning, migration docs should cover warp.sim-to-Newton parity; this example can be referenced as a “new example,” not a migrated API surface.

I can draft a brief migration note and a short “How to run” section for docs/examples if helpful.

---

831-887: CLI plumbing and main loop are serviceable; consider gating per-frame loss prints behind --verbose

Currently prints every frame regardless of --verbose.

Apply:

-            loss = np.min(example.rollout_costs.numpy())
-            print(f"[{example.frame:3d}/{example.num_frames}] loss={loss:.8f}")
+            if args.verbose:
+                loss = np.min(example.rollout_costs.numpy())
+                print(f"[{example.frame:3d}/{example.num_frames}] loss={loss:.8f}")

---

171-176: Indexing into body_q by env_id assumes exactly one body per env

This holds today (bodies_per_env == 1), but if the drone gains more links/bodies, these kernels will under-address. Document or generalize to env_id * bodies_per_env + body_idx.

---

62-84: Gaussian sampling with rejection is fine; you can avoid the loop by sampling truncated normal

Optional: replace rejection with a closed-form truncated normal sampler if bounds get tight.

---

771-785: Copying best rollout back to reference: verify intended index [-1] is reserved for nominal

You sample only rollouts.shape[0] - 1 envs and then assign [-1] from the reference. That matches the “reserve last rollout for nominal” pattern; just confirm tests cover this.

I can add a minimal test to assert rollouts.trajectories[-1] == drone.trajectories[0] post-step.