Reviewing files that changed from the base of the PR and between 08ff1bf and 6c7e891.

• newton/_src/solvers/mujoco/solver_mujoco.py (3 hunks)
• newton/tests/test_rigid_contact.py (1 hunks)

• GitHub Check: Run GPU Benchmarks (Pull Request)
• GitHub Check: run-newton-tests / newton-unittests (windows-latest)
• GitHub Check: run-newton-tests / newton-unittests (ubuntu-latest)

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

2424-2431: Use MjModel.nconmax and warnings.warn (replace Data.ncon / print)

Verified: MjModel.nconmax exists and is the runtime maximum number of contacts. Using self.mj_data.ncon can be zero at init — use int(self.mj_model.nconmax) * nworld instead; replace prints with warnings.warn and remove the now-unused rigid_contact_max assignment.

File: newton/_src/solvers/mujoco/solver_mujoco.py Lines: 2424-2431

-            # TODO find better heuristics to determine nconmax and njmax
-            if disable_contacts:
-                nconmax = 0
-            else:
-                if ncon_per_env is not None:
-                    rigid_contact_max = nworld * ncon_per_env
-                    if rigid_contact_max < self.mj_data.ncon * nworld:
-                        print(f"[WARNING] Value for ncon_per_env is changed from {ncon_per_env} to {self.mj_data.ncon}")
-                        nconmax = self.mj_data.ncon * nworld
-                    else:
-                        nconmax = rigid_contact_max
-                else:
-                    nconmax = self.mj_data.ncon * nworld
-                    rigid_contact_max = model.rigid_contact_max
+            # Default nconmax to MuJoCo model size; allow per‑env override with clamping
+            if disable_contacts:
+                nconmax = 0
+            else:
+                if ncon_per_env is not None:
+                    rigid_contact_max = nworld * ncon_per_env
+                    mj_nconmax_total = int(self.mj_model.nconmax) * nworld
+                    if rigid_contact_max < mj_nconmax_total:
+                        warnings.warn(
+                            f"ncon_per_env={ncon_per_env} is below MuJoCo nconmax per env ({int(self.mj_model.nconmax)}); clamping to {int(self.mj_model.nconmax)}.",
+                            stacklevel=2,
+                        )
+                        nconmax = mj_nconmax_total
+                    else:
+                        nconmax = rigid_contact_max
+                else:
+                    nconmax = int(self.mj_model.nconmax) * nworld

Reviewing files that changed from the base of the PR and between 6c7e891 and cb7143d.

• newton/examples/robot/example_robot_anymal_c_walk.py (1 hunks)
• newton/examples/robot/example_robot_anymal_d.py (1 hunks)
• newton/examples/robot/example_robot_h1.py (1 hunks)
• newton/examples/robot/example_robot_humanoid.py (1 hunks)
• newton/examples/robot/example_robot_policy.py (1 hunks)
• newton/examples/selection/example_selection_articulations.py (1 hunks)
• newton/examples/selection/example_selection_materials.py (1 hunks)
• newton/examples/sensors/example_sensor_contact.py (1 hunks)

• GitHub Check: Run GPU Unit Tests on AWS EC2 (Pull Request)
• GitHub Check: Run GPU Benchmarks (Pull Request)
• GitHub Check: run-newton-tests / newton-unittests (windows-latest)
• GitHub Check: run-newton-tests / newton-unittests (ubuntu-latest)

newton/examples/robot/example_robot_policy.py (1)

268-274: LGTM! The explicit njmax parameter improves clarity.

The addition of njmax=100 to the SolverMuJoCo constructor is appropriate and aligns with the PR's API changes to make constraint limits explicit rather than relying on implicit defaults.

newton/examples/robot/example_robot_anymal_c_walk.py (1)

145-145: LGTM! Appropriate constraint limit for this simulation.

Setting njmax=50 is reasonable for a single Anymal robot simulation and aligns with the PR's API changes.

newton/examples/robot/example_robot_humanoid.py (1)

70-70: LGTM! Consistent with other multi-environment examples.

The njmax=100 value is appropriate for multiple humanoid environments and consistent with other robot examples in the PR.

newton/examples/selection/example_selection_materials.py (1)

97-97: Consider if njmax=50 is sufficient for 16 ant environments.

The example defaults to 16 environments with ant robots. Given that other multi-environment examples use njmax=100, you may want to verify that njmax=50 provides adequate constraint capacity for all environments, especially if users increase --num-envs.

newton/examples/robot/example_robot_h1.py (1)

79-79: LGTM! Comprehensive solver configuration.

The addition of njmax=100 alongside the existing iterations and ls_iterations parameters provides complete control over the solver's constraint handling.

Reviewing files that changed from the base of the PR and between cb7143d and a3269d8.

• newton/_src/solvers/mujoco/solver_mujoco.py (4 hunks)

• GitHub Check: Run GPU Unit Tests on AWS EC2 (Pull Request)
• GitHub Check: Run GPU Benchmarks (Pull Request)

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

1213-1214: LGTM on the API change.

The docstring accurately describes the parameter as optional and mentions that the maximum value is estimated from the model when None is provided.

---

1239-1240: LGTM on updated docstring and parameter descriptions.

The updated docstrings clearly explain that larger of user-provided or default values will be used, which accurately reflects the clamping behavior implemented.

---

2479-2488: Use model.njmax instead of mj_data.ncon for constraint estimation.

The current logic uses mj_data.ncon which represents the current number of contacts, but for constraint estimation we should use the model's maximum contact capacity. Also, the scaling by nworld should happen consistently for both user-provided and default values.

Based on the past review comments, this should use self.mj_data.ncon as the per-environment baseline and scale by nworld for the total constraint capacity.

             if disable_contacts:
                 nconmax = 0
             else:
                 if ncon_per_env is not None:
-                    rigid_contact_max = nworld * ncon_per_env
-                    if rigid_contact_max < self.mj_data.ncon * nworld:
+                    if ncon_per_env < self.mj_data.ncon:
                         warnings.warn(
-                            f"[WARNING] Value for ncon_per_env is changed from {ncon_per_env} to {self.mj_data.ncon} following a MjWarp requirement.",
+                            f"Value for ncon_per_env is changed from {ncon_per_env} to {self.mj_data.ncon} following a MjWarp requirement.",
                             stacklevel=2,
                         )
-                        nconmax = self.mj_data.ncon * nworld
-                    else:
-                        nconmax = rigid_contact_max
+                        ncon_per_env = self.mj_data.ncon
+                    nconmax = ncon_per_env * nworld
                 else:
                     nconmax = self.mj_data.ncon * nworld

---

2490-2499: Use model.njmax and scale properly for total constraint allocation.

The current implementation uses mj_data.nefc (current constraint count) instead of the model's maximum constraint capacity. This can under-allocate constraints for vectorized environments.

Based on past review feedback, this should be corrected to use self.mj_model.njmax and scale by nworld:

             if njmax is not None:
-                if njmax < self.mj_data.nefc:
+                mj_njmax = int(self.mj_model.njmax)
+                if njmax < mj_njmax:
                     warnings.warn(
-                        f"[WARNING] Value for njmax is changed from {njmax} to {self.mj_data.nefc} following a MjWarp requirement.",
+                        f"njmax={njmax} is below MuJoCo njmax per env ({mj_njmax}); clamping to {mj_njmax}.",
                         stacklevel=2,
                     )
-                    njmax = self.mj_data.nefc
+                    njmax = mj_njmax
+                njmax_total = njmax * nworld
             else:
-                njmax = self.mj_data.nefc
+                njmax_total = int(self.mj_model.njmax) * nworld
```Based on the search results, I can see that MuJoCo's Python bindings expose model fields, but let me search for more specific information about njmax field:Let me search more specifically for the njmax field in mjModel:Based on my search results and the existing review comments, I can confirm that the approach should indeed use `mj_model.njmax` instead of `mj_data.nefc`. The search results confirm that MuJoCo preallocates all the memory needed at runtime in mjData, and the model structure contains allocation parameters. However, I need to also consider that the past review comment uses `self.mj_data.nefc` as the baseline. The final approach should pass the correctly scaled total to `put_data`.

The key issue is that `njmax` is expected to be a total value for `put_data`, not per-environment, so the final line should pass `njmax_total` instead of `njmax`.

```diff
+            njmax_total = njmax * nworld
             self.mjw_data = mujoco_warp.put_data(
                 self.mj_model,
                 self.mj_data,
                 nworld=nworld,
                 nconmax=nconmax,
-                njmax=njmax,
+                njmax=njmax_total,
             )

Reviewing files that changed from the base of the PR and between a3269d8 and ab73ce2.

• newton/_src/solvers/mujoco/solver_mujoco.py (6 hunks)

• GitHub Check: Run GPU Unit Tests on AWS EC2 (Pull Request)
• GitHub Check: Run GPU Benchmarks (Pull Request)

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

1498-1502: LGTM on clarifying comments.

The added comments make the CPU vs Warp branches clearer.

Also applies to: 1503-1507

---

1543-1547: LGTM on clarifying comments.

Clearer distinction between MjWarp Data and MjData paths.

Also applies to: 1551-1558

---

2478-2489: nconmax derived from runtime ncon — can under‑allocate and trigger put_data errors; clamp to model.nconmax and scale.

Using self.mj_data.ncon (current contact count) as the default/lower bound can be 0 at init and below MuJoCo’s compiled limit; mujoco_warp.put_data will reject too‑small allocations. Use self.mj_model.nconmax per‑env, scale by nworld, and warn without the “[WARNING]” prefix.

Apply:

-                if ncon_per_env is not None:
-                    rigid_contact_max = nworld * ncon_per_env
-                    if rigid_contact_max < self.mj_data.ncon * nworld:
-                        warnings.warn(
-                            f"[WARNING] Value for ncon_per_env is changed from {ncon_per_env} to {self.mj_data.ncon} following an MjWarp requirement.",
-                            stacklevel=2,
-                        )
-                        nconmax = self.mj_data.ncon * nworld
-                    else:
-                        nconmax = rigid_contact_max
-                else:
-                    nconmax = self.mj_data.ncon * nworld
+                if ncon_per_env is not None:
+                    mj_nconmax = int(self.mj_model.nconmax)
+                    if ncon_per_env < mj_nconmax:
+                        warnings.warn(
+                            f"ncon_per_env={ncon_per_env} is below MuJoCo per-env nconmax ({mj_nconmax}); clamping to {mj_nconmax}.",
+                            stacklevel=2,
+                        )
+                        ncon_per_env = mj_nconmax
+                    nconmax = ncon_per_env * nworld
+                else:
+                    nconmax = int(self.mj_model.nconmax) * nworld

#!/bin/bash
rg -nC2 -e '\bncon\b' --type=py

---

2490-2499: njmax clamped to Data.nefc and not scaled — risks under‑allocation; use model.njmax and scale by nworld.

Baseline must be self.mj_model.njmax (per‑env). Also compute a total (* nworld) for put_data. Current behavior can be below model size and fail; also repeats a previously flagged issue.

Apply:

-            if njmax is not None:
-                if njmax < self.mj_data.nefc:
-                    warnings.warn(
-                        f"[WARNING] Value for njmax is changed from {njmax} to {self.mj_data.nefc} following an MjWarp requirement.",
-                        stacklevel=2,
-                    )
-                    njmax = self.mj_data.nefc
-            else:
-                njmax = self.mj_data.nefc
+            if njmax is not None:
+                mj_njmax = int(self.mj_model.njmax)
+                if njmax < mj_njmax:
+                    warnings.warn(
+                        f"njmax={njmax} is below MuJoCo per-env njmax ({mj_njmax}); clamping to {mj_njmax}.",
+                        stacklevel=2,
+                    )
+                    njmax = mj_njmax
+                njmax_total = njmax * nworld
+            else:
+                njmax_total = int(self.mj_model.njmax) * nworld

And pass the total to MjWarp:

-                njmax=njmax,
+                njmax=njmax_total,

#!/bin/bash
rg -n 'nefc|njmax_total|put_data\(.+njmax=' --type=py

---

1239-1241: Clarify MuJoCo/MjWarp defaults, per‑env scaling, and upward clamping in docstrings

Replace the two parameter descriptions so they state defaults come from the MuJoCo/MjWarp model/data limits, are per‑environment and scaled by nworld, and that user values below the model limit are clamped upward with a warning.

-            njmax (int): Maximum number of constraints per environment (world). If None, a default value is estimated from the initial state. Note that the larger of the user-provided value or the default value is used.
-            ncon_per_env (int | None): Number of contact points per environment (world). If None, a default value is estimated from the initial state. Note that the larger of the user-provided value or the default value is used.
+            njmax (int | None): Maximum number of scalar constraint rows per environment. If None, defaults to the MuJoCo/MjWarp data/model limit (``mj_data.nefc``). The per-env value is scaled by the number of worlds (``nworld``) when allocating in MuJoCo Warp; if a user value is lower than the model limit it is clamped upward and a warning is emitted.
+            ncon_per_env (int | None): Maximum number of contacts per environment. If None, defaults to the MuJoCo/MjWarp data/model limit (``mj_data.nconmax``). The per-env value is scaled by ``nworld`` for allocation; user values below the model limit are clamped upward with a warning.

Also update the Migration Guide (docs/) to mention this behaviour. Quick check command:
rg -n 'njmax|nconmax|ncon_per_env' docs/ -S || true

Reviewing files that changed from the base of the PR and between a3269d8 and ab73ce2.

• newton/_src/solvers/mujoco/solver_mujoco.py (6 hunks)

• GitHub Check: Run GPU Unit Tests on AWS EC2 (Pull Request)
• GitHub Check: Run GPU Benchmarks (Pull Request)

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

1213-1215: Making njmax optional is good.

Signature change aligns with PR goals; matches convert_to_mjc parameter. LGTM.

---

1764-1766: Units mismatch: comment says “per world”. It’s per‑env; put_data needs totals.

Align comment with per‑env semantics and ensure code multiplies by nworld before passing to put_data.

-        njmax: int | None = None,  # number of constraints per world
+        njmax: int | None = None,  # max scalar constraints per environment (per‑env; scaled by nworld)

---

2490-2499: njmax clamp must use MjModel.njmax (model limit) and scale per‑env to total; current use of Data.nefc is unsafe.

nefc is the current number of efc rows, not the capacity. This can under‑allocate and crash under different contact/limit configurations.

-            if njmax is not None:
-                if njmax < self.mj_data.nefc:
-                    warnings.warn(
-                        f"[WARNING] Value for njmax is changed from {njmax} to {self.mj_data.nefc} following an MjWarp requirement.",
-                        stacklevel=2,
-                    )
-                    njmax = self.mj_data.nefc
-            else:
-                njmax = self.mj_data.nefc
+            mj_njmax = int(self.mj_model.njmax)
+            if njmax is not None:
+                if njmax < mj_njmax:
+                    warnings.warn(
+                        f"njmax={njmax} is below MuJoCo njmax per env ({mj_njmax}); clamping to {mj_njmax}.",
+                        stacklevel=2,
+                    )
+                    njmax = mj_njmax
+                njmax_total = njmax * nworld
+            else:
+                njmax_total = mj_njmax * nworld

---

1239-1241: Docs/tests/migration: verify all call sites and docs align with per‑env semantics and clamping.

• Confirmed in code: ncon_per_env is multiplied by number of environments (rigid_contact_max = nworld * ncon_per_env — newton/_src/solvers/mujoco/solver_mujoco.py:2478) and njmax is clamped to MuJoCo limits with warning/assignment logic (newton/_src/solvers/mujoco/solver_mujoco.py:2490–2496). Constructor docstring already refers to "per environment (world)" (newton/_src/solvers/mujoco/solver_mujoco.py:1239–1240).
• Action required: ensure examples/tests and the migration guide explicitly state that njmax and ncon_per_env are per‑env, that ncon_per_env is scaled by num_envs, and that both are clamped to MuJoCo model limits. Check/update call sites such as newton/tests/test_equality_constraints.py:48–49, newton/tests/test_mujoco_solver.py:309, newton/tests/test_rigid_contact.py:185 and newton/examples/sensors/example_sensor_contact.py:142, and add/update docs/migration.rst accordingly.