newton/tests/test_anymal_reset.py (6)

15-15: Remove unused attribute

The torch_device attribute is set but never used anywhere in the test.

     def setUp(self):
         self.device = wp.get_device()
-        self.torch_device = "cuda" if self.device.is_cuda else "cpu"
         self.num_envs = 1000

---

239-239: Remove stray print statement

There's an empty print() statement that appears to be leftover debugging code.

                 if initial_value != current_value:
                     differences.append(f"{attr_name}: {initial_value} -> {current_value}")
-                    print()
                 else:

---

300-305: Extract magic number to class constant

The iteration difference threshold of 50 is used multiple times throughout the code. Consider extracting it to a class-level constant for better maintainability.

 class TestAnymalReset(unittest.TestCase):
+    ITERATION_DIFF_THRESHOLD = 50
     
     def setUp(self):
     
     # ... later in assertions ...
     
                 self.assertGreaterEqual(
                     iteration_diff, 
-                    50, 
+                    self.ITERATION_DIFF_THRESHOLD, 
-                    f"Iteration difference ({iteration_diff}) is below 50, "
+                    f"Iteration difference ({iteration_diff}) is below {self.ITERATION_DIFF_THRESHOLD}, "
                     "indicating solver is approaching maximum iteration limit"
                 )

---

311-311: Remove unused loop variable

The loop variable i is defined but never used within the loop body.

-        for i in range(50):
+        for _ in range(50):
             self.simulate_step()

---

1-343: Add comprehensive test documentation

The test file implements complex multi-environment Anymal reset functionality but lacks docstrings explaining the test objectives, setup, and assertions. Consider adding class and method docstrings to improve maintainability.

Add docstrings like:

class TestAnymalReset(unittest.TestCase):
    """Test suite for validating Anymal robot reset functionality in multi-environment simulations.
    
    Tests ensure that:
    - Robot states can be properly reset to initial configurations
    - MJW solver data remains consistent after resets
    - Solver iterations remain within acceptable bounds
    """
    
    def test_reset_functionality_elliptic(self):
        """Test reset functionality with ELLIPTIC cone type.
        
        Validates that robot states can be reset and solver maintains
        stability with elliptic contact cone configuration.
        """

---

1-343: Clean up whitespace and formatting issues

The file has numerous whitespace issues including trailing whitespace, blank lines with whitespace, and missing newline at end of file.

Run a formatter or manually clean up:

• Remove trailing whitespace on lines 64, 79-84, 94, 136, 223, 229-230, 280, 301-302, 307, 318-319, 343
• Remove whitespace from blank lines
• Add newline at end of file (line 343)

Configuration used: .coderabbit.yml
Review profile: CHILL
Plan: Pro

Learnt from: shi-eric
PR: newton-physics/newton#521
File: newton/examples/example_cloth_hanging.py:36-36
Timestamp: 2025-08-12T05:17:34.404Z
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/tests/test_anymal_reset.py (2)

139-153: Boolean return from step() is inverted/misleading and unused

Returning True when iteration_difference < 50 suggests an error state as “success.” Also, step() isn’t used elsewhere. Either invert the boolean to reflect “safe to continue,” or remove the return (or the method) and rely on get_iteration_difference() as you already do.

Apply this diff if you want to keep it:

-        return iteration_difference < 50
+        # True means we're safely away from the iteration cap
+        return iteration_difference >= 50

---

195-200: Avoid bare except; catch specific exceptions when deepcopy fails

Bare except masks bugs and is flagged by Ruff.

Apply this diff:

-                except:
-                    pass
+                except (TypeError, AttributeError, ValueError):
+                    # Some attributes are not deep-copyable; skip them.
+                    pass

Optional: saved_count is never used; consider removing it to reduce noise.

newton/tests/test_anymal_reset.py (8)

37-39: Right-size the test (1000 envs is CI-hostile); make it configurable

1000 environments will be very slow and memory-heavy in CI. Make this adjustable via env vars and default to a smaller number. Also expose headless via env var.

Apply this diff:

@@
-import copy
+import os
+import copy
 import unittest
@@
-        self.num_envs = 1000
-        self.headless = True
+        self.num_envs = int(os.getenv("NEWTON_TEST_NUM_ENVS", "256"))
+        self.headless = os.getenv("NEWTON_HEADLESS", "1") != "0"

Also applies to: 16-18

---

106-114: Use consistent source for initial articulation state (state_0 vs model)

You simulate once before snapshotting. For consistency, capture both root transforms and root velocities from the same state you just simulated (state_0). Passing the model to getters may fetch different “initial” values than the post-sim state.

Apply this diff:

@@
-        self.anymal = ArticulationView(
-            self.model, "/World/envs/env_0/Robot/base", verbose=False, exclude_joint_types=[newton.JOINT_FREE]
-        )
-        self.default_root_transforms = wp.to_torch(self.anymal.get_root_transforms(self.model)).clone()
-        self.default_root_velocities = wp.to_torch(self.anymal.get_root_velocities(self.model)).clone()
+        self.anymal = ArticulationView(
+            self.model, "/World/envs/env_0/Robot/base", verbose=False, exclude_joint_types=[newton.JOINT_FREE]
+        )
+        self.default_root_transforms = wp.to_torch(self.anymal.get_root_transforms(self.state_0)).clone()
+        self.default_root_velocities = wp.to_torch(self.anymal.get_root_velocities(self.state_0)).clone()

If you intentionally wanted the “model defaults” here, consider adding a short comment to clarify why transforms/velocities come from the model while DOF state comes from state_0.

---

116-123: Minor: avoid redundant get_device() call

Use the already-fetched device for mempool check.

Apply this diff:

-        self.use_cuda_graph = self.device.is_cuda and wp.is_mempool_enabled(wp.get_device())
+        self.use_cuda_graph = self.device.is_cuda and wp.is_mempool_enabled(self.device)

---

154-160: Deduplicate render() and render_frame()

Both methods implement the same logic. Keep one. If you want both names, have render() delegate to render_frame().

Apply this diff:

     def render(self):
-        if self.renderer is None:
-            return
-        self.renderer.begin_frame(self.sim_time)
-        self.renderer.render(self.state_0)
-        self.renderer.end_frame()
+        self.render_frame()

Also applies to: 260-266

---

247-251: Remove stray debug print

There’s an empty print() which will spam test output.

Apply this diff:

-                    print()

---

276-278: Misleading variable name

zero_dof_velocities actually reuses the initial velocities. Rename to avoid confusion, or set zeros if that was intended.

Apply this diff to clarify the intent:

-        zero_dof_velocities = wp.from_torch(self.initial_dof_velocities, dtype=wp.float32)
-        self.anymal.set_dof_velocities(self.state_0, zero_dof_velocities)
+        initial_dof_velocities = wp.from_torch(self.initial_dof_velocities, dtype=wp.float32)
+        self.anymal.set_dof_velocities(self.state_0, initial_dof_velocities)

If you intended zero velocities: use wp.zeros_like on the target buffer instead.

---

322-327: Use underscore for unused loop variable

Silences the Ruff warning and signals intent.

Apply this diff:

-        for i in range(50):
+        for _ in range(50):

---

19-21: Optional: make mujoco import non-fatal when unavailable

If CI doesn’t have mujoco installed, the whole test module will fail import. Consider a guarded import and skipping tests when unavailable.

Example:

try:
    import mujoco
except Exception as e:
    mujoco = None

@unittest.skipIf(mujoco is None, "mujoco not available")
class TestAnymalReset(unittest.TestCase):
    ...

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.

newton/tests/test_anymal_reset.py (1)

1-15: License header added correctly (addresses prior request)

SPDX headers and license block look good and satisfy the “copyright notice required” feedback.

newton/tests/test_anymal_reset.py (9)

31-33: 1000 envs is very heavy for CI — make it configurable and default lighter

A 1000-env test with MJW can be slow and memory-intensive. Consider reading num envs from an env var and defaulting to a smaller value for CI runs.

Apply this diff (includes import):

+import os
 import copy
 import unittest
@@
     def setUp(self):
         self.device = wp.get_device()
-        self.num_envs = 1000
+        # Allow overriding for CI/local runs; default lighter to avoid OOM/timeouts
+        self.num_envs = int(os.getenv("NEWTON_TEST_NUM_ENVS", "64"))
         self.headless = True

---

88-91: Potential OOM/time cost with nefc_per_env=200 at large env counts

At 1000 envs this allocates very large contact-related buffers. Unless strictly required, reduce nefc_per_env or make it configurable similar to num_envs.

Example:

-        self.solver = newton.solvers.MuJoCoSolver(
-            self.model, solver=2, cone=cone_type, impratio=100.0, iterations=100, ls_iterations=50, nefc_per_env=200
-        )
+        nefc_per_env = int(os.getenv("NEWTON_TEST_NEFC_PER_ENV", "64"))
+        self.solver = newton.solvers.MuJoCoSolver(
+            self.model, solver=2, cone=cone_type, impratio=100.0, iterations=100, ls_iterations=50, nefc_per_env=nefc_per_env
+        )

---

110-116: Use existing device instance when checking mempool

Minor nit: avoid calling wp.get_device() again; use self.device.

-        self.use_cuda_graph = self.device.is_cuda and wp.is_mempool_enabled(wp.get_device())
+        self.use_cuda_graph = self.device.is_cuda and wp.is_mempool_enabled(self.device)

---

140-146: Drop unused render() method (duplicate of render_frame)

render() duplicates render_frame and is not used elsewhere.

-    def render(self):
-        if self.renderer is None:
-            return
-        self.renderer.begin_frame(self.sim_time)
-        self.renderer.render(self.state_0)
-        self.renderer.end_frame()

---

208-211: Use rtol with isclose to be robust across scales

Pure atol can over-flag diffs for larger magnitude values. Add a modest rtol.

-                        tolerance = 1e-4
-                        diff_mask = ~np.isclose(initial_value, current_value, atol=tolerance)
+                        atol = 1e-4
+                        rtol = 1e-5
+                        diff_mask = ~np.isclose(initial_value, current_value, rtol=rtol, atol=atol)

---

223-227: Remove stray print()

Empty print risks noisy test logs.

-                    differences.append(f"{attr_name}: {initial_value} -> {current_value}")
-                    print()
+                    differences.append(f"{attr_name}: {initial_value} -> {current_value}")

---

252-253: Misleading variable name; not zeros

You’re restoring initial DOF velocities, not zero velocities. Rename for clarity.

-        zero_dof_velocities = wp.from_torch(self.initial_dof_velocities, dtype=wp.float32)
-        self.anymal.set_dof_velocities(self.state_0, zero_dof_velocities)
+        initial_dof_velocities = wp.from_torch(self.initial_dof_velocities, dtype=wp.float32)
+        self.anymal.set_dof_velocities(self.state_0, initial_dof_velocities)

---

274-279: Prefer using step() to avoid duplicating capture/sim logic

This loop reimplements the logic from step(). Consolidate to reduce drift.

-            if self.use_cuda_graph and self.graph:
-                wp.capture_launch(self.graph)
-            else:
-                self.simulate()
-            self.sim_time += self.frame_dt
+            self.step()

---

294-301: Prefer using step() to avoid duplicating capture/sim logic (second loop)

Same duplication as above.

-            if self.use_cuda_graph and self.graph:
-                wp.capture_launch(self.graph)
-            else:
-                self.simulate()
-            self.sim_time += self.frame_dt
+            self.step()

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.

newton/tests/test_anymal_reset.py (6)

99-106: Assert articulation selection covers all envs

Guard against pattern mismatches by asserting that the view matches exactly one articulation per env. This avoids false negatives when comparing states and ensures reset affects all envs.

         self.anymal = ArticulationView(
             self.model, "/World/envs/*/Robot/base", verbose=False, exclude_joint_types=[newton.JOINT_FREE]
         )
         self.default_root_transforms = wp.to_torch(self.anymal.get_root_transforms(self.model)).clone()
         self.default_root_velocities = wp.to_torch(self.anymal.get_root_velocities(self.model)).clone()
 
+        # Sanity check: selection must match one articulation per environment
+        self.assertEqual(
+            int(self.default_root_transforms.shape[0]),
+            int(self.num_envs),
+            "ArticulationView pattern did not select all envs",
+        )
+
         self.initial_dof_positions = wp.to_torch(self.anymal.get_dof_positions(self.state_0)).clone()
         self.initial_dof_velocities = wp.to_torch(self.anymal.get_dof_velocities(self.state_0)).clone()

---

110-116: Make CUDA-graph toggle robust across Warp versions

Depending on Warp version, wp.get_device() may return a string or an object. Use a resilient CUDA check and reuse the same device handle.

-        self.use_cuda_graph = self.device.is_cuda and wp.is_mempool_enabled(wp.get_device())
+        dev = self.device
+        # Robust check: support both device objects and string device names
+        is_cuda = bool(getattr(dev, "is_cuda", False))
+        if not is_cuda:
+            is_cuda = str(dev).startswith("cuda")
+        self.use_cuda_graph = is_cuda and wp.is_mempool_enabled(dev)
         if self.use_cuda_graph:
             with wp.ScopedCapture() as capture:
                 self.simulate()
             self.graph = capture.graph
         else:
             self.graph = None

---

238-244: Remove duplicate rendering method and use a single render()

render_frame() duplicates render(). Consolidate to avoid drift.

-    def render_frame(self):
-        if self.renderer is None:
-            return
-        self.renderer.begin_frame(self.sim_time)
-        self.renderer.render(self.state_0)
-        self.renderer.end_frame()
+    # render() already exists; render_frame() is redundant. Use render() instead.

And switch call sites:

-            if not self.headless:
-                self.render_frame()
+            if not self.headless:
+                self.render()

-            if not self.headless:
-                self.render_frame()
+            if not self.headless:
+                self.render()

Also applies to: 281-283, 302-303

---

225-229: Remove stray print() in non-equality branch

This prints a blank line on every scalar mismatch, cluttering CI logs.

                 if initial_value != current_value:
                     differences.append(f"{attr_name}: {initial_value} -> {current_value}")
-                    print()
                 else:
                     identical_count += 1

---

31-31: Consider reducing env count or making it configurable for CI stability

1000 envs with MuJoCo + mjwarp can be heavy in CI. Allow overriding via env var to keep runtime predictable.

-        self.num_envs = 1000
+        import os
+        # Default high local coverage; CI can reduce with NEWTON_RESET_TEST_NUM_ENVS
+        self.num_envs = int(os.getenv("NEWTON_RESET_TEST_NUM_ENVS", "1000"))

---

245-259: Reset semantics are consistent

Root transforms/velocities and DOF states are restored to captured initial values, and sim time is reset. Naming nit: zero_dof_velocities holds initial velocities (likely zeros), which is fine but slightly misleading.

-        zero_dof_velocities = wp.from_torch(self.initial_dof_velocities, dtype=wp.float32)
-        self.anymal.set_dof_velocities(self.state_0, zero_dof_velocities)
+        reset_dof_velocities = wp.from_torch(self.initial_dof_velocities, dtype=wp.float32)
+        self.anymal.set_dof_velocities(self.state_0, reset_dof_velocities)

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/tests/test_anymal_reset.py (9)

105-111: Good: ArticulationView selects all envs (not just env_0); consider asserting selection size

Selecting all envs via "/World/envs/*/Robot/base" fixes the earlier single-env reset issue. Add a quick assertion to ensure the selection matches num_envs to harden the test against asset/path regressions.

         self.anymal = ArticulationView(
             self.model, "/World/envs/*/Robot/base", verbose=False, exclude_joint_types=[newton.JOINT_FREE]
         )
         self.default_root_transforms = wp.to_torch(self.anymal.get_root_transforms(self.model)).clone()
         self.default_root_velocities = wp.to_torch(self.anymal.get_root_velocities(self.model)).clone()
+        # Sanity check: one articulation per env
+        self.assertEqual(
+            self.default_root_transforms.shape[0],
+            self.num_envs,
+            "ArticulationView pattern did not select all envs",
+        )

---

57-78: Avoid slicing joint_q to place the robot; set root z via the instance xform

Directly mutating joint_q with [:3] assumes a specific DOF layout for the free root (order/size can vary). It’s safer to set the initial height in the xform when instantiating the articulation(s).

-        articulation_builder.joint_q[:3] = [0.0, 0.0, 0.8]
+        z0 = 0.8
@@
-        builder.add_builder(articulation_builder, xform=wp.transform(wp.vec3(0.0, 0.0, 0.0), wp.quat_identity()))
+        builder.add_builder(articulation_builder, xform=wp.transform(wp.vec3(0.0, 0.0, z0), wp.quat_identity()))
@@
-            builder.add_builder(articulation_builder, xform=wp.transform(wp.vec3(x, y, 0.0), wp.quat_identity()))
+            builder.add_builder(articulation_builder, xform=wp.transform(wp.vec3(x, y, z0), wp.quat_identity()))

If you intended to modify a specific joint (not the free root), consider using a named lookup rather than relying on positional slicing.

---

116-123: Minor: reuse the resolved device when checking mempool

Avoid calling wp.get_device() twice and pass the already-resolved device instance.

-        self.use_cuda_graph = self.device.is_cuda and wp.is_mempool_enabled(wp.get_device())
+        self.use_cuda_graph = self.device.is_cuda and wp.is_mempool_enabled(self.device)

---

153-171: Harden mjw_data capture: filter out volatile attributes by prefix

The blacklist covers common counters, but MuJoCo’s constraint/contact fields (efc_, con_, etc.) and some timing/diagnostic fields are inherently volatile and can differ due to internal ordering even when state is equivalent. Filtering by prefixes makes the test more robust.

         all_attributes = [attr for attr in dir(mjw_data) if not attr.startswith("_")]
 
-        skip_attributes = {
+        skip_attributes = {
             "time",
             "solver_niter",
             "ncollision",
             "nsolving",
             "collision_pair",
             "collision_pairid",
             "solver_nisland",
             "nefc",
             "ncon",
             "cfrc_int",
             "collision_worldid",
         }
 
+        # Volatile families: contact/constraint internals, timers, warnings, etc.
+        skip_prefixes = (
+            "efc_",      # constraint Jacobian/solver internals
+            "con_",      # contact arrays (indices/order can vary)
+            "timer_",    # timing diagnostics
+            "warning",   # warning counters/flags
+        )
+
         for attr_name in all_attributes:
-            if attr_name in skip_attributes:
+            if attr_name in skip_attributes or any(attr_name.startswith(p) for p in skip_prefixes):
                 continue
             attr_value = getattr(mjw_data, attr_name)

Also applies to: 173-185

---

200-218: Make numeric comparison tolerant to NaNs in identical positions

When comparing arrays that may contain NaNs (rare but possible in diagnostic buffers), use equal_nan=True.

-                        tolerance = 1e-3
-                        diff_mask = ~np.isclose(initial_value, current_value, atol=tolerance)
+                        tolerance = 1e-3
+                        diff_mask = ~np.isclose(initial_value, current_value, atol=tolerance, equal_nan=True)

---

243-257: Type-safety for transforms; fix misleading variable name

set_root_transforms will be more robust with a wp.transform-typed array. Also, the local name zero_dof_velocities is misleading since you apply initial (likely zero) velocities.

     def reset_robot_state(self):
-        self.anymal.set_root_transforms(self.state_0, self.default_root_transforms)
+        initial_root_transforms = wp.from_torch(self.default_root_transforms, dtype=wp.transform)
+        self.anymal.set_root_transforms(self.state_0, initial_root_transforms)
@@
-        zero_dof_velocities = wp.from_torch(self.initial_dof_velocities, dtype=wp.float32)
-        self.anymal.set_dof_velocities(self.state_0, zero_dof_velocities)
+        initial_dof_velocities = wp.from_torch(self.initial_dof_velocities, dtype=wp.float32)
+        self.anymal.set_dof_velocities(self.state_0, initial_dof_velocities)

If you actually wanted zero velocities irrespective of initial values, explicitly create a zeros buffer with the same shape as initial_dof_velocities.

---

81-87: Remove unused sim_step field

self.sim_step is assigned but never read.

-        self.sim_step = 0

---

271-285: Iteration threshold policy: consider ratio-based check to reduce flakiness

A fixed threshold of 50 assumes opt.iterations=100. If this changes, the check may become too weak/strict. Consider asserting a margin like max_iterations <= 0.5 * opt_iterations instead.

Example:

iteration_diff = self.get_iteration_difference()
self.assertGreaterEqual(
    iteration_diff,
    int(0.5 * int(self.solver.mjw_model.opt.iterations)),
    "Solver is approaching the max iteration limit",
)

Also applies to: 290-302

---

319-321: Optional: drop the main block in test modules

Tests are run by the test runner; keeping an entry point is harmless but unnecessary and can complicate discovery in some setups.

-if __name__ == "__main__":
-    wp.init()
-    unittest.main()

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/tests/test_anymal_reset.py (6)

16-16: Fix docstring grammar for clarity

Small grammar nit to improve readability.

-"""Tests that reset results in the same data and converge of solver is preserved."""
+"""Tests that reset results in the same data and that convergence of the solver is preserved."""

---

105-111: Assert that all envs were selected by ArticulationView

Prevents silent mismatches when the selection pattern changes or asset paths differ.

         self.anymal = ArticulationView(
             self.model, "/World/envs/*/Robot/base", verbose=False, exclude_joint_types=[newton.JOINT_FREE]
         )
         self.default_root_transforms = wp.to_torch(self.anymal.get_root_transforms(self.model)).clone()
         self.default_root_velocities = wp.to_torch(self.anymal.get_root_velocities(self.model)).clone()
+        # Ensure the pattern matched one articulation per env
+        self.assertEqual(
+            self.default_root_transforms.shape[0],
+            self.num_envs,
+            "ArticulationView pattern did not select all envs",
+        )

---

113-115: Document the warm-up step before capturing the baseline snapshot

Consider adding a brief comment so future readers understand why the first simulate occurs before saving mjw_data.

-        self.simulate()
+        # Warm up one frame to initialize internal solver/model buffers before snapshotting mjw_data
+        self.simulate()
         self.save_initial_mjw_data()

---

182-184: Handle NumPy scalar attributes in mjw_data snapshot

Some mjw_data fields may be NumPy scalars (np.generic). Capture them consistently to avoid silent omissions.

-            elif isinstance(attr_value, (int, float, bool)):
-                self.initial_mjw_data[attr_name] = copy.deepcopy(attr_value)
+            elif isinstance(attr_value, (int, float, bool, np.generic)):
+                val = attr_value.item() if isinstance(attr_value, np.generic) else attr_value
+                self.initial_mjw_data[attr_name] = copy.deepcopy(val)

Also applies to: 157-171

---

200-229: Make numeric diffing NaN-safe and a bit more robust

• Use nan-aware reductions to avoid NaN-propagation in max/mean.
• Count differences via boolean sum for correctness across shapes.
• Optionally add a small rtol to tolerate tiny relative deviations.

-                else:
-                    if not np.array_equal(initial_value, current_value):
-                        max_diff = np.max(np.abs(initial_value - current_value))
-                        mean_diff = np.mean(np.abs(initial_value - current_value))
-                        tolerance = 1e-3
-                        diff_mask = ~np.isclose(initial_value, current_value, atol=tolerance, equal_nan=True)
-
-                        diff_indices = np.where(diff_mask)
-                        num_different = len(diff_indices[0])
+                else:
+                    if not np.array_equal(initial_value, current_value):
+                        tolerance = 1e-3
+                        rtolerance = 1e-5
+                        diff_mask = ~np.isclose(
+                            initial_value, current_value, atol=tolerance, rtol=rtolerance, equal_nan=True
+                        )
+                        abs_diff = np.abs(initial_value - current_value)
+                        max_diff = np.nanmax(abs_diff)
+                        mean_diff = np.nanmean(abs_diff)
+                        num_different = int(diff_mask.sum())
                         percent_different = (num_different / initial_value.size) * 100
                         if num_different > 0:
                             differences.append(
                                 f"{attr_name}: max_diff={max_diff:.10f}, mean_diff={mean_diff:.10f}, shape={initial_value.shape}, {num_different}/{initial_value.size} different values ({percent_different:.2f}%)"
                             )
                         else:
                             identical_count += 1

---

31-35: Make the test configurable to avoid CI timeouts on heavy runs

64 envs and 100 total frames (2×50) at 2 substeps can be heavy in CI. Allow overriding env count and steps via env vars. This keeps defaults intact locally while enabling faster CI.

-import copy
+import copy
+import os
 import unittest

     def setUp(self):
         self.device = wp.get_device()
-        self.num_envs = 64
+        self.num_envs = int(os.getenv("NEWTON_NUM_ENVS", "64"))
         self.headless = True
+        self.num_steps = int(os.getenv("NEWTON_RESET_TEST_STEPS", "50"))

-        for i in range(50):
+        for i in range(self.num_steps):
             self.step()
             if not self.headless:
                 self.render()

-        for i in range(50):
+        for i in range(self.num_steps):
             self.step()
             if not self.headless:
                 self.render()

Also applies to: 273-287, 291-303, 18-20

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/tests/test_anymal_reset.py (2)

147-153: Fix CUDA graph replay: missing state ping-pong breaks simulation progress

simulate() swaps state_0/state_1 each frame; CUDA graph capture replays only GPU launches and omits this Python-level swap. Without mirroring the swap on replay, steps will use stale pointers and the sim may stall or diverge.

Apply this diff:

 def step(self):
     if self.use_cuda_graph:
         wp.capture_launch(self.graph)
+        # Mirror simulate() ping-pong of states when replaying a captured graph
+        self.state_0, self.state_1 = self.state_1, self.state_0
     else:
         self.simulate()
     self.sim_time += self.frame_dt

---

270-276: Mirror state ping-pong after CUDA graph replay in propagate_reset_state()

Same issue as in step(): without swapping state_0/state_1 after wp.capture_launch(self.graph), state buffers don’t advance.

Apply this diff:

 def propagate_reset_state(self):
     if self.use_cuda_graph and self.graph:
         wp.capture_launch(self.graph)
+        # Keep state progression consistent with simulate()
+        self.state_0, self.state_1 = self.state_1, self.state_0
     else:
         self.simulate()
     self.sim_time += self.frame_dt

newton/tests/test_anymal_reset.py (6)

124-130: Graph capture advances physics by one frame; confirm baseline snapshot timing

The capture block runs simulate() once, advancing the state by an extra frame before the main loop. Since initial_mjw_data is captured before this (Line 121-123), the baseline corresponds to “post-1-frame,” while subsequent runs include an additional captured frame. This is fine if intentional; otherwise consider resetting to the baseline after capture to avoid an off-by-one drift in total steps.

Possible adjustment:

• After building the graph, call reset_robot_state() and propagate_reset_state() once to re-align to the intended baseline before entering the main loop.

---

16-16: Docstring grammar: clarify intent

Nit: “converge of solver” → “convergence of the solver”.

-"""Tests that reset results in the same data and converge of solver is preserved."""
+"""Tests that reset reproduces the same mjw_data and preserves solver convergence."""

---

33-34: Remove unused num_envs or wire it into the setup

self.num_envs is defined but unused. Either remove it or make the setup honor the count (e.g., by building multiple envs).

-        self.num_envs = 1
+        # Single-environment test; keep only if you plan to expand to multi-env
+        # self.num_envs = 1

---

89-95: Targeting all DOFs may inadvertently include free joint DOFs

You set every DOF to TARGET_POSITION with zero gains. If builder.joint_dof_mode includes the free root joint’s DOFs, this may be nonsensical. Prefer restricting to actuated/non-free joints.

-        for i in range(len(builder.joint_dof_mode)):
-            builder.joint_dof_mode[i] = newton.JointMode.TARGET_POSITION
+        # Only apply to non-free joints (adjust if API exposes joint types per dof)
+        for i in range(len(builder.joint_dof_mode)):
+            # If available, gate by joint type
+            # if builder.joint_type[i] != newton.JointType.FREE:
+            builder.joint_dof_mode[i] = newton.JointMode.TARGET_POSITION

If joint types aren’t aligned per-DOF at this point, consider setting modes via ArticulationView or a builder helper that skips free joints.

---

103-105: Prefer symbolic enums/strings for readability over magic numbers

solver=2 is unclear to readers. Use the corresponding string or enum value (e.g., "cg") for clarity.

-        self.solver = newton.solvers.SolverMuJoCo(
-            self.model, solver=2, cone=cone_type, impratio=impratio, iterations=100, ls_iterations=50, nefc_per_env=200
-        )
+        self.solver = newton.solvers.SolverMuJoCo(
+            self.model,
+            solver="cg",
+            cone=cone_type,
+            impratio=impratio,
+            iterations=100,
+            ls_iterations=50,
+            nefc_per_env=200,
+        )

---

319-322: Surface detailed diffs in assertion failure

compare_mjw_data_with_initial() prints diffs to stdout. For CI diagnostics, it’s better to include a summary in the assertion message.

One option: have compare_mjw_data_with_initial return (ok, summary) and use it in the assert. Example:

-        mjw_data_matches = self.compare_mjw_data_with_initial()
+        mjw_data_matches, summary = self.compare_mjw_data_with_initial()

-        self.assertTrue(
-            mjw_data_matches,
-            f"mjw_data after reset does not match initial state with {self._cone_type_name(cone_type)} cone",
-        )
+        self.assertTrue(
+            mjw_data_matches,
+            f"mjw_data mismatch after reset ({self._cone_type_name(cone_type)} cone):\n{summary}",
+        )

And adapt compare_mjw_data_with_initial() to build and return a concise summary string when differences exist.

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.