asv/benchmarks/setup/bench_solver_mujoco.py (4)

45-49: Comment does not match code behavior.
The comment says “Load a small model,” but you instantiate the solver with self._model (which may be very large). Either update the comment or actually use a small model for warm-up.

Apply one of:

• Update comment (preferred, since kernel specialization often depends on shape):

-        # Load a small model to cache the kernels
+        # Warm up: instantiate once to JIT-compile kernels for this model shape

• Or construct a tiny warm-up model (if acceptable).

---

16-16: Remove unused import.
gc is not used in this module.

-import gc

---

27-35: Optional: sanity-check large param sizes.
num_envs of 4096/8192 can OOM typical 16–24GB GPUs during solver construction. If CI runs these benches, consider adding an environment toggle or reducing defaults.

---

52-52: Optional: align method names with PEP 8 / other benches.
Consider time_initialize_solver_mujoco for consistency with time_initialize_model.

Also applies to: 79-79

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

asv.conf.json (1)

22-22: SHA verified; packaging metadata present in pyproject.toml
Optional:

• Add --no-deps to your pip install line to lock transitive dependencies:

-    "python -m pip install git+https://github.com/google-deepmind/mujoco_warp@64bb938f87b6aaf840753ef9ddc198db655020c0",
+    "python -m pip install --no-deps git+https://github.com/google-deepmind/mujoco_warp@64bb938f87b6aaf840753ef9ddc198db655020c0",

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

asv/benchmarks/setup/bench_solver_mujoco.py (5)

52-56: Fix Ruff F841 in timed section and ensure deterministic teardown.

 @skip_benchmark_if(wp.get_cuda_device_count() == 0)
 def time_initialize_solverMuJoCo(self, robot, num_envs):
-    _solver = SolverMuJoCo(self._model)
-    wp.synchronize_device()
+    _ = SolverMuJoCo(self._model)
+    wp.synchronize_device()
+    del _

---

73-78: Warm-up: same deterministic lifetime pattern as above.

-# Load a small model to cache the kernels
-solver = SolverMuJoCo(self._model)
-del solver
-
-wp.synchronize_device()
+# Warm up: instantiate once to JIT-compile kernels for this model shape
+_ = SolverMuJoCo(self._model)
+wp.synchronize_device()
+del _

---

79-83: Fix Ruff F841 in fast timed section and delete after sync.

 @skip_benchmark_if(wp.get_cuda_device_count() == 0)
 def time_initialize_solverMuJoCo(self, robot, num_envs):
-    _solver = SolverMuJoCo(self._model)
-    wp.synchronize_device()
+    _ = SolverMuJoCo(self._model)
+    wp.synchronize_device()
+    del _

---

38-45: Guard setup on CPU-only and finalize model on CUDA explicitly.

Without a guard, setup still creates GPU-dependent resources on CPU-only runners; also make the model device explicit to avoid accidental CPU placement skewing results.

 def setup(self, robot, num_envs):
-    wp.init()
-
-    builder = Example.create_model_builder(robot, num_envs, randomize=True, seed=123)
+    wp.init()
+    if wp.get_cuda_device_count() == 0:
+        # CPU-only: time method is skipped; avoid creating GPU-dependent resources.
+        return
+    builder = Example.create_model_builder(robot, num_envs, randomize=True, seed=123)
 
     # finalize model
-    self._model = builder.finalize()
+    self._model = builder.finalize(device="cuda:0")

---

67-72: Mirror the CPU guard and explicit CUDA device in the fast benchmark setup.

 def setup(self, robot, num_envs):
-    wp.init()
-    builder = Example.create_model_builder(robot, num_envs, randomize=True, seed=123)
+    wp.init()
+    if wp.get_cuda_device_count() == 0:
+        return
+    builder = Example.create_model_builder(robot, num_envs, randomize=True, seed=123)
     # finalize model
-    self._model = builder.finalize()
+    self._model = builder.finalize(device="cuda:0")

asv/benchmarks/setup/bench_solver_mujoco.py (3)

46-51: Warm-up: use throwaway name, sync, then delete for deterministic lifetime.

Also fix the stale comment (“small model”) to reflect actual behavior.

-# Load a small model to cache the kernels
-solver = SolverMuJoCo(self._model)
-del solver
-
-wp.synchronize_device()
+# Warm up: instantiate once to JIT-compile kernels for this model shape
+_ = SolverMuJoCo(self._model)
+wp.synchronize_device()
+del _

---

62-66: Consider setting a timeout for symmetry with KPI class.

Optional, but keeps benchmark configs consistent.

 rounds = 1
 number = 1
 repeat = 3
 min_run_count = 1
+timeout = 3600

---

90-93: Rename CLI keys for clarity (they benchmark SolverMuJoCo, not “Model”).

Purely cosmetic but reduces confusion when invoking via --bench.

 benchmark_list = {
-    "KpiInitializeModel": KpiInitializeSolverMuJoCo,
-    "FastInitializeModel": FastInitializeSolverMuJoCo,
+    "KpiInitializeSolverMuJoCo": KpiInitializeSolverMuJoCo,
+    "FastInitializeSolverMuJoCo": FastInitializeSolverMuJoCo,
 }

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

asv/benchmarks/setup/bench_solver_mujoco.py (5)

61-65: Free the solver deterministically after timing.

Helps keep GPU memory steady across repeats.

     def time_initialize_solverMuJoCo(self, robot, num_envs):
-        _ = SolverMuJoCo(self._model, ncon_per_env=nconmax[robot])
-        wp.synchronize_device()
+        _ = SolverMuJoCo(self._model, ncon_per_env=nconmax[robot])
+        wp.synchronize_device()
+        del _

---

88-91: Deterministic cleanup in fast timing too.

     def time_initialize_solverMuJoCo(self, robot, num_envs):
-        _ = SolverMuJoCo(self._model, ncon_per_env=nconmax[robot])
-        wp.synchronize_device()
+        _ = SolverMuJoCo(self._model, ncon_per_env=nconmax[robot])
+        wp.synchronize_device()
+        del _

---

68-69: “Fast” still feels large; confirm intent or downscale.

128–256 envs for humanoid/ant/quadruped may still be heavy. Either rename to “Medium” or drop to, e.g., [64, 128].

---

47-60: Guard setup on CPU-only runners and finalize on CUDA explicitly.

Without a guard, setup will still instantiate the solver on CPU-only CI and likely fail. Also, finalize the model on a CUDA device to avoid skewing timings with CPU models. Replace the warm-up variable and order to avoid Ruff F841 and ensure deterministic lifetime after a device sync.

     def setup(self, robot, num_envs):
         wp.init()
-
-        builder = Example.create_model_builder(robot, num_envs, randomize=True, seed=123)
+        if wp.get_cuda_device_count() == 0:
+            # CPU-only: timing is skipped; avoid creating GPU-dependent resources.
+            return
+
+        builder = Example.create_model_builder(robot, num_envs, randomize=True, seed=123)
 
         # finalize model
-        self._model = builder.finalize()
+        self._model = builder.finalize(device="cuda:0")
 
-        # Load the model to cache the kernels
-        solver = SolverMuJoCo(self._model, ncon_per_env=nconmax[robot])
-        del solver
-
-        wp.synchronize_device()
+        # Warm up: instantiate once to JIT-compile kernels for this model shape
+        _ = SolverMuJoCo(self._model, ncon_per_env=nconmax[robot])
+        wp.synchronize_device()
+        del _

---

76-87: Mirror the CPU guard and explicit CUDA device in the fast benchmark setup.

Same failure mode as above; also fix the warm-up variable and ordering.

     def setup(self, robot, num_envs):
         wp.init()
-        builder = Example.create_model_builder(robot, num_envs, randomize=True, seed=123)
+        if wp.get_cuda_device_count() == 0:
+            return
+        builder = Example.create_model_builder(robot, num_envs, randomize=True, seed=123)
         # finalize model
-        self._model = builder.finalize()
+        self._model = builder.finalize(device="cuda:0")
 
-        # Load the model to cache the kernels
-        solver = SolverMuJoCo(self._model, ncon_per_env=nconmax[robot])
-        del solver
-
-        wp.synchronize_device()
+        # Warm up: instantiate once to JIT-compile kernels for this model shape
+        _ = SolverMuJoCo(self._model, ncon_per_env=nconmax[robot])
+        wp.synchronize_device()
+        del _

asv/benchmarks/setup/bench_solver_mujoco.py (4)

27-34: Revisit ncon_per_env for g1/h1 to avoid OOM at large env counts.

At 4096–8192 envs, 150 contacts/env yields extremely large allocations (nworld * ncon_per_env) and may OOM on 16–24GB GPUs. Consider scaling by env count, capping, or skipping large envs for these robots.

Would you like me to propose a helper (e.g., compute_ncon_per_env(robot, num_envs)) and wire it through setup/time for safer defaults?

---

19-21: Avoid global config side-effects at import time.

Setting wp.config globally on import can leak into other ASV modules in the same process. Move these into setup or guard with “if not already set”.

-wp.config.enable_backward = False
-wp.config.quiet = True
+wp.config.enable_backward = False  # consider moving into setup() to limit scope
+wp.config.quiet = True             # or gate with an env flag if needed

---

71-75: Add a timeout for the fast benchmark too.

Keeps runaway runs contained on constrained GPUs.

     rounds = 1
     number = 1
     repeat = 3
     min_run_count = 1
+    timeout = 1200

---

99-102: Rename CLI keys for clarity.

These benchmarks initialize the MuJoCo solver, not the model. Rename keys to match class names.

-        "KpiInitializeModel": KpiInitializeSolverMuJoCo,
-        "FastInitializeModel": FastInitializeSolverMuJoCo,
+        "KpiInitializeSolverMuJoCo": KpiInitializeSolverMuJoCo,
+        "FastInitializeSolverMuJoCo": FastInitializeSolverMuJoCo,

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

asv/benchmarks/setup/bench_model.py (2)

127-137: Same warm-up release issue as KPI.

Delete _model/builder and GC after warm-up to avoid cumulative memory use.

Apply:

             solver = SolverMuJoCo(_model, ncon_per_env=nconmax[robot])
             del solver
+            del _model
+            del builder
+            gc.collect()
 
             wp.synchronize_device()

---

141-159: Same per-repeat cleanup improvements as KPI.

Delete builder post-finalize and GC after releasing solver/model.

Apply:

                 model = builder.finalize()
                 wp.synchronize_device()
                 modelBuilder_end = time.perf_counter()
                 timings_modelBuilder.append(modelBuilder_end - modelBuilder_beg)
+                del builder
                 ...
                 del solver
                 del model
+                gc.collect()

asv/benchmarks/setup/bench_model.py (4)

27-36: Validate and document ncon_per_env overrides.

Hard-coding 150 for g1/h1 may be too low/high depending on contact-rich scenes and large num_envs. Please confirm against worst-case contacts per env, and add a short comment with rationale or derive this once from a 1-env probe.

---

73-92: Tighten per-repeat cleanup to avoid transient pressure.

Delete builder after finalize and GC after releasing solver/model.

Apply:

                 model = builder.finalize()
                 wp.synchronize_device()
                 modelBuilder_end = time.perf_counter()
                 timings_modelBuilder.append(modelBuilder_end - modelBuilder_beg)
+                del builder
 
                 solver_beg = time.perf_counter()
                 solver = SolverMuJoCo(model, ncon_per_env=nconmax[robot])
                 wp.synchronize_device()
                 solver_end = time.perf_counter()
                 timings_solver.append(solver_end - solver_beg)
 
                 del solver
                 del model
+                gc.collect()

---

49-95: Deduplicate the timing logic.

Both classes duplicate setup_cache/loop code. Factor into a small helper to reduce drift and ease maintenance.

Example (sketch):

def _compute_timings(robots, envs, repeats):
    ...  # shared body here
    return timings

Then call it from each setup_cache with the class-specific params.

Also applies to: 117-164

---

39-48: Guard against OOM for largest configs.

4096/8192 envs for humanoid/g1/h1 can exceed common GPU RAM. Consider a try/except around finalize/solver to skip combos gracefully (record None and short-circuit track_*), or downscale per-robot env counts.

Also applies to: 108-116

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

newton/_src/utils/benchmark.py (4)

20-30: Add type hints and improve docstring for better maintainability.

The new helper function would benefit from type hints and a more comprehensive docstring to clarify its purpose and expected input/output types.

-def _convert_params_to_dict(params, param_names):
-    """Convert params to keyword arguments using param_names."""
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+def _convert_params_to_dict(
+    params: Union[Any, List, Tuple],
+    param_names: Optional[List[str]]
+) -> Optional[Dict[str, Any]]:
+    """Convert params to keyword arguments using param_names.
+    
+    Args:
+        params: Parameter values, can be a single value or tuple/list of values
+        param_names: List of parameter names to map to the values
+        
+    Returns:
+        Dictionary mapping parameter names to values, or None if no param_names
+    """

---

32-44: Add type hints and improve documentation for the parameter calling helper.

The function handles complex parameter passing logic that would benefit from explicit type annotations.

-def _call_with_params(method, param_dict, params, cached_data):
-    """Call a method with appropriate parameters (keyword or positional)."""
+def _call_with_params(
+    method: callable,
+    param_dict: Optional[Dict[str, Any]],
+    params: Tuple,
+    cached_data: Optional[Any]
+) -> Any:
+    """Call a method with appropriate parameters (keyword or positional).
+    
+    Args:
+        method: The method to call
+        param_dict: Keyword arguments if param_names was provided, None otherwise
+        params: Positional arguments to pass
+        cached_data: Cached data from setup_cache, passed as first positional arg
+        
+    Returns:
+        The result of the method call
+    """

---

74-75: Consider handling the enumerate index more explicitly.

While the variable i is used correctly, consider making its purpose clearer with a more descriptive name.

-    for i, params in enumerate(combinations):
+    for combination_index, params in enumerate(combinations):

Then update line 88:

-        if i == 0 and hasattr(benchmark_cls, "setup_cache"):
+        if combination_index == 0 and hasattr(benchmark_cls, "setup_cache"):

---

99-99: Consider improving parameter display for better readability.

When params is a tuple with complex objects, the print output might be hard to read. Consider formatting it more nicely.

-                print(f"\n[Benchmark] Running {benchmark_cls.__name__}.{attr} with parameters {params}")
+                param_str = ", ".join(str(p) for p in params) if isinstance(params, (list, tuple)) else str(params)
+                print(f"\n[Benchmark] Running {benchmark_cls.__name__}.{attr} with parameters ({param_str})")

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

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

82-85: Fix: params_tuple may be undefined when params is already a list/tuple

Define params_tuple in both branches to avoid UnboundLocalError on subsequent calls.

-        # Ensure params is always a tuple for consistent handling
-        if not isinstance(params, (list | tuple)):
-            params_tuple = (params,)
+        # Ensure params is always a tuple for consistent handling
+        params_tuple = tuple(params) if isinstance(params, (list | tuple)) else (params,)

newton/_src/utils/benchmark.py (3)

20-30: Validate param_names/params length to prevent silent mismatches

Guard against dropped/extra parameters; avoid relying on zip(strict=False).

 def _convert_params_to_dict(params, param_names):
     """Convert params to keyword arguments using param_names."""
     if not param_names:
         return None
 
     # Handle single value case by wrapping in tuple
     if not isinstance(params, (list | tuple)):
         params = (params,)
 
-    return dict(zip(param_names, params, strict=False))
+    if len(param_names) != len(params):
+        raise ValueError(
+            f"param_names length ({len(param_names)}) != params length ({len(params)})"
+        )
+    return dict(zip(param_names, params))

---

87-90: Reduce noise: print setup_cache message only when printing results

Keeps output quiet for programmatic runs.

-            print(f"\n[Benchmark] Running {benchmark_cls.__name__}.setup_cache")
+            if print_results:
+                print(f"\n[Benchmark] Running {benchmark_cls.__name__}.setup_cache")

---

100-115: Gracefully skip NotImplemented benchmarks; avoid ZeroDivisionError

Catching NotImplementedError lets the runner continue; also skip averaging when there are no samples.

                 if attr.startswith("time_"):
                     # Run timing benchmarks multiple times and measure elapsed time.
                     for _ in range(number):
                         start = time.perf_counter()
-                        _call_with_params(method, param_dict, params_tuple, cached_data)
-                        t = time.perf_counter() - start
-                        samples.append(t)
+                        try:
+                            _call_with_params(method, param_dict, params_tuple, cached_data)
+                            t = time.perf_counter() - start
+                            samples.append(t)
+                        except NotImplementedError:
+                            samples.clear()
+                            break
                 elif attr.startswith("track_"):
                     # Run tracking benchmarks multiple times and record returned values.
                     for _ in range(number):
-                        val = _call_with_params(method, param_dict, params_tuple, cached_data)
-                        samples.append(val)
+                        try:
+                            val = _call_with_params(method, param_dict, params_tuple, cached_data)
+                            samples.append(val)
+                        except NotImplementedError:
+                            samples.clear()
+                            break
                 # Compute the average result.
-                avg = sum(samples) / len(samples)
-                results[(attr, params)] = avg
+                if not samples:
+                    if print_results:
+                        print(f"[Benchmark] Skipping {benchmark_cls.__name__}.{attr} (NotImplemented).")
+                    continue
+                avg = sum(samples) / len(samples)
+                results[(attr, params)] = avg

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

asv/benchmarks/setup/bench_model.py (2)

64-71: Warm-up cleanup for KPI looks correct and resolves prior OOM risk.

del solver/_model/builder + gc.collect() + wp.synchronize_device() is in place. Good fix.

---

132-142: Free warm-up model/builder in FastInitializeModel to prevent memory growth across (robot, num_envs).

_model and builder persist after warm-up; this can accumulate. Mirror KPI cleanup and sync after deletes.

-            solver = SolverMuJoCo(_model, ncon_per_env=nconmax[robot])
-            del solver
-
-            wp.synchronize_device()
+            solver = SolverMuJoCo(_model, ncon_per_env=nconmax[robot])
+            del solver
+            del _model
+            del builder
+            gc.collect()
+            wp.synchronize_device()

asv/benchmarks/setup/bench_model.py (5)

23-24: Avoid global Warp config side effects; set per-benchmark.

Move wp.config.enable_backward = False and wp.config.quiet = True into each setup_cache to prevent global impact when this module is imported.

Apply:

- wp.config.enable_backward = False
- wp.config.quiet = True

And add at the beginning of both setup methods:

@@
-    def setup_cache(self):
+    def setup_cache(self):
+        wp.config.enable_backward = False
+        wp.config.quiet = True
         wp.init()

---

146-165: Delete builder inside the repeat loop (Fast) to keep host memory flat.

Currently only solver and model are deleted.

                 del solver
                 del model
+                del builder
+                gc.collect()

---

29-36: Document/verify nconmax heuristics.

If 150 is empirical for g1/h1, add a brief comment; otherwise consider deriving from model or leaving None for consistency.

---

117-121: Align benchmark metadata across classes.

Optional: add timeout to FastInitializeModel for consistency with KPI.

     rounds = 1
     number = 1
     repeat = 3
     min_run_count = 1
+    timeout = 1200

---

180-191: Release builder in peakmem bench to avoid skewing measurements.

Delete builder and force GC after finalizing/del model.

         with wp.ScopedDevice("cpu"):
             builder = Example.create_model_builder(robot, num_envs, randomize=True, seed=123)

             # finalize model
             model = builder.finalize()

         del model
+        del builder
+        gc.collect()

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

asv/benchmarks/setup/bench_model.py (2)

132-141: Missing resource cleanup compared to KpiInitializeModel.

The FastInitializeModel class is missing the del _model and gc.collect() cleanup that exists in KpiInitializeModel after the kernel preloading step. This could lead to memory accumulation across iterations.

Apply this diff to add consistent resource cleanup:

             # Load the model to cache the kernels
             solver = SolverMuJoCo(_model, ncon_per_env=nconmax[robot])
             del solver
+            del _model
+            del builder
+            gc.collect()

             wp.synchronize_device()

---

76-100: Fix variable name inconsistencies.

Several variable naming issues need to be addressed:

1. modelBuilder_beg should use snake_case as model_builder_begin
2. The timing collection uses inconsistent variable names

Apply this diff to fix the variable naming:

             for _ in range(self.repeat):
-                modelBuilder_beg = time.perf_counter()
+                model_builder_begin = time.perf_counter()
                 builder = Example.create_model_builder(robot, num_envs, randomize=True, seed=123)

                 # finalize model
                 model = builder.finalize()
                 wp.synchronize_device()
-                modelBuilder_end = time.perf_counter()
-                timings_modelBuilder.append(modelBuilder_end - modelBuilder_beg)
+                model_builder_end = time.perf_counter()
+                timings_modelBuilder.append(model_builder_end - model_builder_begin)

-                solver_beg = time.perf_counter()
+                solver_begin = time.perf_counter()
                 solver = SolverMuJoCo(model, ncon_per_env=nconmax[robot])
                 wp.synchronize_device()
                 solver_end = time.perf_counter()
-                timings_solver.append(solver_end - solver_beg)
+                timings_solver.append(solver_end - solver_begin)

asv/benchmarks/setup/bench_model.py (2)

29-36: Consider making the constraint configuration more explicit.

The nconmax dictionary mixes None values with specific integers, making the intent unclear. Consider documenting why some robots use None while others have specific values, or use a more explicit configuration structure.

Add a comment explaining the configuration:

+# Per-robot constraint limits for SolverMuJoCo initialization
+# None = use solver's automatic estimation, int = explicit limit
 nconmax = {
     "humanoid": None,
     "g1": 150,
     "h1": 150,
     "cartpole": None,
     "ant": None,
     "quadruped": None,
 }

---

180-180: Unused parameter should be named with underscore prefix.

The first parameter appears to be unused cached data. Consider making this more explicit by using a descriptive name or double underscore prefix.

Apply this diff for better clarity:

-    def peakmem_initialize_model_cpu(self, _, robot, num_envs):
+    def peakmem_initialize_model_cpu(self, _cached_data, robot, num_envs):

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro

asv/benchmarks/setup/bench_model.py (2)

59-72: Warm-up and cleanup looks good.

Preloading kernels and freeing _model/builder between (robot, num_envs) avoids memory growth and noisy timings. Nice.

---

159-164: Undefined name solver_beg → use solver_begin.

Typo breaks timing; use the variable you define.

                 solver = SolverMuJoCo(model, ncon_per_env=nconmax[robot])
                 wp.synchronize_device()
                 solver_end = time.perf_counter()
-                timings_solver.append(solver_end - solver_beg)
+                timings_solver.append(solver_end - solver_begin)

asv/benchmarks/setup/bench_model.py (5)

185-196: Also delete builder in peakmem benchmark to avoid lingering allocations.

Minor hygiene; helps stabilizing RSS measurements.

     with wp.ScopedDevice("cpu"):
         builder = Example.create_model_builder(robot, num_envs, randomize=True, seed=123)
 
         # finalize model
         model = builder.finalize()
 
     del model
+    del builder
+    gc.collect()

---

65-65: Defensive access for nconmax.

Future robot additions won’t crash the benchmark if the map isn’t updated.

-            solver = SolverMuJoCo(_model, ncon_per_env=nconmax[robot])
+            solver = SolverMuJoCo(_model, ncon_per_env=nconmax.get(robot))

-                solver = SolverMuJoCo(model, ncon_per_env=nconmax[robot])
+                solver = SolverMuJoCo(model, ncon_per_env=nconmax.get(robot))

Also applies to: 87-87, 138-138, 160-160

---

117-121: Consider adding timeout like KPI class for consistency.

Keeps runner behavior uniform across benchmarks.

     rounds = 1
     number = 1
     repeat = 3
     min_run_count = 1
+    timeout = 3600

---

23-25: Module-wide Warp config can leak to other benchmarks.

Optional: set/restore these in setup_cache to localize side effects.

---

59-101: DRY: factor shared benchmarking loop.

Both classes duplicate the same warm-up/measure/cleanup logic; a small helper would reduce maintenance.

If helpful, I can sketch a _collect_timings(robots, envs, repeat) helper and integrate it into both classes.

Also applies to: 132-173

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro

asv/benchmarks/setup/bench_model.py (2)

59-72: Warm-up + explicit cleanup looks good; prevents memory growth across sweeps.

Preloading kernels, deleting solver/model/builder, gc.collect(), and synchronizing are correctly in place.

---

159-164: Fix NameError in timing: use solver_begin, not solver_beg.

This currently crashes and also triggers Ruff F821/F841. Use the correct variable.

                 solver_begin = time.perf_counter()
                 solver = SolverMuJoCo(model, ncon_per_env=nconmax[robot])
                 wp.synchronize_device()
                 solver_end = time.perf_counter()
-                timings_solver.append(solver_end - solver_beg)
+                timings_solver.append(solver_end - solver_begin)

asv/benchmarks/setup/bench_model.py (2)

185-196: Release builder and force GC in peakmem path (post-measurement cleanup).

Minor cleanup parity with other paths; avoids lingering CPU allocations after peakmem sampling.

     with wp.ScopedDevice("cpu"):
         builder = Example.create_model_builder(robot, num_envs, randomize=True, seed=123)

         # finalize model
         model = builder.finalize()

     del model
+    del builder
+    gc.collect()

---

39-111: Optional: de-duplicate common benchmarking logic between classes.

Both classes share identical loops; extracting a small helper (e.g., _measure_timings(robots, envs, repeat)) would reduce maintenance and keep fixes (like the NameError above) in one place.

Also applies to: 113-184

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro