Skip to content

solverMuJoCo: change nefc_per_env to njmax#663

Merged
eric-heiden merged 2 commits into
newton-physics:mainfrom
adenzler-nvidia:dev/adenzler/nefc-njmax
Aug 28, 2025
Merged

solverMuJoCo: change nefc_per_env to njmax#663
eric-heiden merged 2 commits into
newton-physics:mainfrom
adenzler-nvidia:dev/adenzler/nefc-njmax

Conversation

@adenzler-nvidia

@adenzler-nvidia adenzler-nvidia commented Aug 28, 2025

Copy link
Copy Markdown
Member

This makes it more aligned with the MuJoCo language and makes the error messages useful

Before your PR is "Ready for review"

  • Necessary tests have been added and new examples are tested (see newton/tests/test_examples.py)
  • Documentation is up-to-date
  • Code passes formatting and linting checks with pre-commit run -a

Summary by CodeRabbit

  • Refactor

    • Renamed public solver parameter nefc_per_env to njmax for consistent per-environment constraint budgeting.
    • Constructor and conversion calls updated to use njmax.
  • Bug Fixes

    • Constraint budget logic now respects both configured njmax and observed constraints to avoid under-allocation.
  • Documentation

    • Examples updated to use the njmax parameter.
  • Tests

    • Tests updated to use njmax and verify behavior remains unchanged.

Signed-off-by: Alain Denzler <adenzler@nvidia.com>
@coderabbitai

coderabbitai Bot commented Aug 28, 2025

Copy link
Copy Markdown
Contributor
📝 Walkthrough

Walkthrough

Renamed the SolverMuJoCo parameter nefc_per_env to njmax across the solver constructor and convert_to_mjc; propagated njmax into MuJoCo Warp data allocation and adjusted final budget calculation to njmax = max(njmax, self.mj_data.nefc). Updated examples and tests to use njmax.

Changes

Cohort / File(s) Summary
Core MuJoCo solver API and propagation
newton/_src/solvers/mujoco/solver_mujoco.py
Rename parameter nefc_per_envnjmax in __init__ and convert_to_mjc; pass njmax into convert_to_mjc and mujoco_warp.put_data; compute final njmax = max(njmax, self.mj_data.nefc).
Examples update to new API
newton/examples/example_mujoco.py, newton/examples/robot/example_robot_g1.py
Replace SolverMuJoCo(..., nefc_per_env=...) with njmax=...; no other logic changes.
Tests update to new API
newton/tests/test_anymal_reset.py, newton/tests/test_equality_constraints.py, newton/tests/test_mujoco_solver.py
Update SolverMuJoCo constructor calls from nefc_per_env to njmax, preserving values and test logic.

Sequence Diagram(s)

sequenceDiagram
  autonumber
  actor User
  participant Solver as SolverMuJoCo
  participant Converter as convert_to_mjc
  participant MJW as mujoco_warp.put_data

  Note over User,Solver: Instantiate solver with njmax
  User->>Solver: __init__(..., njmax)
  Solver->>Converter: convert_to_mjc(..., njmax)
  Note over Converter: final_njmax = max(njmax, mj_data.nefc)
  Converter->>MJW: put_data(..., njmax=final_njmax)
  MJW-->>Converter: ack
  Converter-->>Solver: (MjWarpModel, MjWarpData, MjModel, MjData)
  Solver-->>User: solver instance ready
Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Possibly related PRs

Suggested reviewers

  • AntoineRichard

Tip

🔌 Remote MCP (Model Context Protocol) integration is now available!

Pro plan users can now connect to remote MCP servers from the Integrations page. Connect with popular remote MCPs such as Notion and Linear to add more context to your reviews and chats.


📜 Recent review details

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.

📥 Commits

Reviewing files that changed from the base of the PR and between f745c11 and ed68a79.

📒 Files selected for processing (1)
  • newton/tests/test_anymal_reset.py (1 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
  • newton/tests/test_anymal_reset.py
⏰ Context from checks skipped due to timeout of 900000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)
  • GitHub Check: Run GPU Benchmarks (Pull Request)
  • GitHub Check: Run GPU Unit Tests on AWS EC2 (Pull Request)
  • GitHub Check: run-newton-tests / newton-unittests (windows-latest)
  • GitHub Check: run-newton-tests / newton-unittests (ubuntu-latest)
✨ Finishing Touches
  • 📝 Generate Docstrings
🧪 Generate unit tests
  • Create PR with unit tests
  • Post copyable unit tests in a comment

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share
🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

  • Review comments: Directly reply to a review comment made by CodeRabbit. Example:
    • I pushed a fix in commit <commit_id>, please review it.
    • Open a follow-up GitHub issue for this discussion.
  • Files and specific lines of code (under the "Files changed" tab): Tag @coderabbit in a new review comment at the desired location with your query.
  • PR comments: Tag @coderabbit in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
    • @coderabbit gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
    • @coderabbit read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.

Support

Need help? Create a ticket on our support page for assistance with any issues or questions.

CodeRabbit Commands (Invoked using PR/Issue comments)

Type @coderabbit help to get the list of available commands.

Other keywords and placeholders

  • Add @coderabbit ignore or @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
  • Add @coderabbit summary or @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
  • Add @coderabbit or @coderabbitai anywhere in the PR title to generate the title automatically.

Status, Documentation and Community

  • Visit our Status Page to check the current availability of CodeRabbit.
  • Visit our Documentation for detailed information on how to use CodeRabbit.
  • Join our Discord Community to get help, request features, and share feedback.
  • Follow us on X/Twitter for updates and announcements.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🧹 Nitpick comments (4)
newton/_src/solvers/mujoco/solver_mujoco.py (3)

1162-1162: Consider keeping a backwards-compatible alias for nefc_per_env.

To avoid breaking downstream users immediately, accept nefc_per_env as a deprecated alias and map it to njmax, emitting a DeprecationWarning.

Apply within the signature:

-        njmax: int = 100,
+        njmax: int = 100,
+        nefc_per_env: int | None = None,

Then just after argument parsing:

# Back-compat shim for deprecated argument
if nefc_per_env is not None:
    import warnings
    warnings.warn("`nefc_per_env` is deprecated; use `njmax`.", DeprecationWarning, stacklevel=2)
    njmax = nefc_per_env if njmax == 100 else njmax  # preserve explicit njmax if set

1188-1189: Doc clarity: “per environment (world)” can be ambiguous when separate_envs_to_worlds=False.

Suggest: “Maximum number of constraints per world (when environments are not separated, world == all envs).”


2356-2359: Heuristic njmax = max(njmax, mj_data.nefc) is okay, but won’t catch later spikes.

If constraints can exceed the initial budget at runtime (e.g., more contacts later), consider detecting allocation overflow and reinitializing mjw_data with a larger njmax, or expose guidance in docs to set njmax conservatively.

Would you like a follow-up patch that grows mjw_data on demand when MuJoCo reports allocation errors?

newton/tests/test_anymal_reset.py (1)

100-104: Avoid magic number for njmax (optional).

Define a constant to improve readability and reuse across tests.

 class TestAnymalReset(unittest.TestCase):
+    NJMAX = 200
@@
-        self.solver = newton.solvers.SolverMuJoCo(
-            self.model, solver=2, cone=cone_type, impratio=impratio, iterations=100, ls_iterations=50, njmax=200
-        )
+        self.solver = newton.solvers.SolverMuJoCo(
+            self.model,
+            solver=2,
+            cone=cone_type,
+            impratio=impratio,
+            iterations=100,
+            ls_iterations=50,
+            njmax=self.NJMAX,
+        )
📜 Review details

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.

📥 Commits

Reviewing files that changed from the base of the PR and between ace430d and f745c11.

📒 Files selected for processing (6)
  • newton/_src/solvers/mujoco/solver_mujoco.py (5 hunks)
  • newton/examples/example_mujoco.py (1 hunks)
  • newton/examples/robot/example_robot_g1.py (1 hunks)
  • newton/tests/test_anymal_reset.py (1 hunks)
  • newton/tests/test_equality_constraints.py (1 hunks)
  • newton/tests/test_mujoco_solver.py (1 hunks)
🧰 Additional context used
🧠 Learnings (3)
📓 Common learnings
Learnt from: preist-nvidia
PR: newton-physics/newton#579
File: newton/examples/example_mujoco.py:350-354
Timestamp: 2025-08-19T12:27:30.629Z
Learning: In Newton examples, Example.use_mujoco is a high-level attribute that controls whether to use MuJoCo solver vs other solvers (like XPBD), while SolverMuJoCo.use_mujoco_cpu controls the backend within MuJoCo (CPU vs Warp). These are separate concepts serving different purposes - the PR rename only applies to the solver parameter, not the Example class attributes.
📚 Learning: 2025-08-19T12:27:30.629Z
Learnt from: preist-nvidia
PR: newton-physics/newton#579
File: newton/examples/example_mujoco.py:350-354
Timestamp: 2025-08-19T12:27:30.629Z
Learning: In Newton examples, there's a distinction between solver parameters and Example class attributes. The Example class can have its own use_mujoco attribute for controlling example-level behavior (like CUDA graphs, rendering logic), while the solver uses use_mujoco_cpu for backend selection. These serve different purposes and should not be conflated during API renames.

Applied to files:

  • newton/tests/test_mujoco_solver.py
  • newton/examples/robot/example_robot_g1.py
  • newton/examples/example_mujoco.py
  • newton/_src/solvers/mujoco/solver_mujoco.py
  • newton/tests/test_equality_constraints.py
📚 Learning: 2025-08-19T12:27:30.629Z
Learnt from: preist-nvidia
PR: newton-physics/newton#579
File: newton/examples/example_mujoco.py:350-354
Timestamp: 2025-08-19T12:27:30.629Z
Learning: In Newton examples, Example.use_mujoco is a high-level attribute that controls whether to use MuJoCo solver vs other solvers (like XPBD), while SolverMuJoCo.use_mujoco_cpu controls the backend within MuJoCo (CPU vs Warp). These are separate concepts serving different purposes - the PR rename only applies to the solver parameter, not the Example class attributes.

Applied to files:

  • newton/tests/test_mujoco_solver.py
  • newton/examples/robot/example_robot_g1.py
  • newton/examples/example_mujoco.py
  • newton/_src/solvers/mujoco/solver_mujoco.py
⏰ Context from checks skipped due to timeout of 900000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)
  • 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 (ubuntu-latest)
  • GitHub Check: run-newton-tests / newton-unittests (windows-latest)
🔇 Additional comments (8)
newton/examples/example_mujoco.py (1)

246-246: Rename applied correctly (njmax passed into SolverMuJoCo).

Constructor call now matches the updated API.

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

80-80: Rename applied correctly (njmax=300).

Call site updated and consistent with solver API.

newton/tests/test_equality_constraints.py (1)

48-48: Test updated to new API (njmax).

Looks consistent with the rename; no logic changes.

newton/tests/test_mujoco_solver.py (1)

318-318: Rename applied correctly (njmax=1).

Note: runtime will upsize to max(njmax, nefc) at convert time, so this won’t under-allocate.

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

1253-1265: Propagation of njmax into convert_to_mjc is correct.

No issues spotted.


1647-1648: Signature rename to njmax is consistent.

Doc comment “number of constraints per world” matches MuJoCo terminology.


1155-1181: All nefc_per_env references have been removed. No further updates required.

newton/tests/test_anymal_reset.py (1)

102-104: Rename usage to njmax looks correct.

Constructor now passes njmax explicitly; aligns with the SolverMuJoCo API rename.

Comment thread newton/tests/test_anymal_reset.py
@codecov

codecov Bot commented Aug 28, 2025

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ All tests successful. No failed tests found.

📢 Thoughts on this report? Let us know!

@eric-heiden eric-heiden left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@eric-heiden eric-heiden enabled auto-merge (squash) August 28, 2025 21:50
@eric-heiden eric-heiden merged commit 8b9ee1a into newton-physics:main Aug 28, 2025
12 checks passed
maxkra15 pushed a commit to maxkra15/newton that referenced this pull request Sep 1, 2025
Signed-off-by: Alain Denzler <adenzler@nvidia.com>
Co-authored-by: Eric Heiden <eheiden@nvidia.com>
@coderabbitai coderabbitai Bot mentioned this pull request Oct 28, 2025
4 tasks
@coderabbitai coderabbitai Bot mentioned this pull request Jan 5, 2026
4 tasks
eric-heiden added a commit to eric-heiden/newton that referenced this pull request Jan 28, 2026
Signed-off-by: Alain Denzler <adenzler@nvidia.com>
Co-authored-by: Eric Heiden <eheiden@nvidia.com>
mmacklin pushed a commit to mmacklin/newton that referenced this pull request Apr 7, 2026
Signed-off-by: Alain Denzler <adenzler@nvidia.com>
Co-authored-by: Eric Heiden <eheiden@nvidia.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants