Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

Appearance settings

Conversation

@camevor
Copy link
Member

@camevor camevor commented Dec 1, 2025

Description

This PR introduces Extended State Attributes. These are state attributes containing time-varying data in a canonical form that are computed only upon request (e.g., body accelerations, parent interaction forces, etc.).

This change harmonizes the handling of state attributes, providing a unified mechanism for importers, sensors, and users to define what State attributes they require and or support.

This PR also contains a simple IMU sensor implementation to showcase the extended attributes.
Resolves #969 , resolves #1148

Newton Migration Guide

Please ensure the migration guide for warp.sim users is up-to-date with the changes made in this PR.

  • The migration guide in docs/migration.rst is up-to date

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

  • New Features

    • Added IMU sensor for per-site accelerometer and gyroscope outputs.
    • Added optional extended state attributes (e.g., body accelerations and parent forces) requestable via builders or models.
    • MuJoCo solver can populate extended rigid-body acceleration/force outputs when requested.
  • Documentation

    • Added Extended State Attributes guide and IMU sensor docs and examples.
    • Updated sensor catalog and API references to include SensorIMU.
  • Tests

    • Added comprehensive unit tests for SensorIMU behavior and validation.

✏️ Tip: You can customize this high-level summary in your review settings.

Summary by CodeRabbit

  • New Features

    • IMU sensor producing per-site accelerometer and gyroscope outputs.
    • Optional extended state attributes (body accelerations and parent forces) requestable from builders or models.
    • Solver support to populate extended rigid-body acceleration/force outputs when requested.
  • Documentation

    • Added Extended State Attributes guide, IMU sensor docs, examples, and updated sensor catalog/API references.
  • Examples

    • New IMU sensor example demonstrating visualization and usage.
  • Tests

    • Comprehensive unit tests for SensorIMU creation, validation, and runtime behavior.
  • Chore

    • Exposed IMU sensor in public module and updated README sensor tiles.

✏️ Tip: You can customize this high-level summary in your review settings.

@coderabbitai
Copy link
Contributor

coderabbitai bot commented Dec 1, 2025

Caution

Review failed

The head commit changed during the review from 1107d3a to 04b2124.

📝 Walkthrough

Walkthrough

Adds an opt-in extended state attribute mechanism (builder→model→state), two optional State fields (body_qdd, body_parent_f), MuJoCo kernels and solver integration to populate them, a SensorIMU that requests/uses body_qdd, tests, an example, and documentation for the feature.

Changes

Cohort / File(s) Summary
Core state & builder
newton/_src/sim/builder.py, newton/_src/sim/model.py, newton/_src/sim/state.py
ModelBuilder gains _requested_state_attributes and request_state_attributes(); add_builder() accumulates requests; finalize() forwards requests to Model. Model adds _requested_state_attributes, request_state_attributes(), get_requested_state_attributes() and uses them to allocate extended state fields. State adds optional body_qdd and body_parent_f.
MuJoCo solver (kernels + integration)
newton/_src/solvers/mujoco/kernels.py, newton/_src/solvers/mujoco/solver_mujoco.py
New Warp functions mj_body_acceleration and convert_rigid_forces_from_mj_kernel (note: duplicated insertion observed). Solver imports the kernel, adds enable_rne_postconstraint(), and conditionally launches the conversion kernel in update_newton_state() to populate body_qdd and body_parent_f.
IMU sensor implementation
newton/_src/sensors/sensor_imu.py, newton/sensors.py
New SensorIMU and Warp kernel compute_sensor_imu_kernel. SensorIMU validates sites, requests body_qdd by default, exposes accelerometer and gyroscope arrays, and computes readings from state. Exported via newton/sensors.py.
Tests & example
newton/tests/test_sensor_imu.py, newton/examples/sensors/example_sensor_imu.py, newton/tests/test_examples.py
New unit tests covering SensorIMU creation, validation, update behavior and rotated/world-frame cases. New example demonstrating IMU integration, visualization, and an acc_to_color kernel. Test harness registers the example.
Docs & API / README
docs/concepts/extended_state_attributes.rst, docs/concepts/sensors.rst, docs/api/newton_sensors.rst, docs/index.rst, README.md
Added docs for extended state attributes and SensorIMU, updated sensors doc and API autosummary, added navigation entry, and updated README sensor example tiles.

Sequence Diagram(s)

sequenceDiagram
    autonumber
    actor User
    participant Builder as ModelBuilder
    participant Model
    participant State
    participant Solver as SolverMuJoCo
    participant Sensor as SensorIMU

    User->>Builder: request_state_attributes("body_qdd","body_parent_f")
    Builder->>Builder: store in _requested_state_attributes
    User->>Builder: add_builder(sub)
    Builder->>Builder: accumulate sub._requested_state_attributes

    User->>Builder: finalize()
    Builder->>Model: create Model
    Builder->>Model: request_state_attributes(*accumulated)
    Model->>Model: record requests

    User->>Model: state()
    Model->>Model: get_requested_state_attributes()
    Model->>State: allocate body_qdd, body_parent_f (if requested)

    User->>Sensor: SensorIMU(model, sites)
    Sensor->>Model: request_state_attributes("body_qdd")

    loop simulation step
        User->>Solver: step(state)
        Solver->>Solver: enable_rne_postconstraint(state_out)
        Solver->>Solver: convert_rigid_forces_from_mj_kernel(mjw_data)
        Solver->>State: write body_qdd, body_parent_f
        User->>Sensor: sensor.update(state)
        Sensor->>Sensor: compute_sensor_imu_kernel(...)
        Sensor-->>User: accelerometer, gyroscope
    end
Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

  • Areas needing extra attention:
    • Warp kernels: numerical correctness, indexing, and duplicated definitions in kernels.py.
    • State request propagation timing (builder → model → state allocation) and API semantics.
    • Solver integration: correct enabling of RNE postconstraint and guarded kernel launches.
    • SensorIMU: site validation, device/array lifetimes, and error handling.
    • Tests and example: realism of assertions and resource/setup correctness.

Suggested reviewers

  • mmacklin
  • eric-heiden
  • nvlukasz

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)
Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 60.32% which is insufficient. The required threshold is 80.00%. You can run @coderabbitai generate docstrings to improve docstring coverage.
✅ Passed checks (4 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title 'Extended state attributes and IMU Sensor' directly reflects the main objective of the PR, which adds extended state attributes framework and an IMU sensor implementation.
Linked Issues check ✅ Passed The PR comprehensively addresses both linked issues: #969 requires IMU sensor implementation with accelerometer/gyroscope outputs (delivered via SensorIMU class), and #1148 requires State attribute request mechanism (delivered via request_state_attributes API and extended attributes framework).
Out of Scope Changes check ✅ Passed All changes are directly scoped to the linked issues: extended state attribute framework, IMU sensor implementation, MuJoCo solver integration, documentation, and tests. SensorIMU export in newton/sensors.py and README updates are supporting changes for the new sensor.

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

Comment @coderabbitai help to get the list of available commands and usage tips.

@camevor camevor changed the title Prototype for extended state attributes Prototype for extended state attributes and IMU Sensor Dec 1, 2025
Copy link
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

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

Actionable comments posted: 5

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (2)
newton/_src/solvers/xpbd/solver_xpbd.py (1)

64-108: XPBD can crash when State was created before attributes were requested

SolverXPBD.__init__ now does:

if model.body_count > 0:
    model.request_state_attributes("body_q", "body_qd", "body_f")

if model.particle_count > 0:
    model.request_state_attributes("particle_q", "particle_qd", "particle_f")

This is the right set of attributes for what XPBD’s step() uses, but it doesn’t help State objects that were created before the solver was constructed. With the new allocation scheme, such states can have state_in.particle_q is None even when model.particle_count > 0, which matches the pipeline failure (AttributeError: 'NoneType' object has no attribute 'device' when cloning state_in.particle_q).

Two concrete follow-ups are advisable:

  1. Defensive check in step() (or a helper) before using body/particle fields, so failures are explicit and point to the fix:
def _check_required_state_fields(self, model: Model, state: State) -> None:
    if model.particle_count > 0 and state.particle_q is None:
        raise RuntimeError(
            "SolverXPBD requires particle state attributes but this State was created "
            "without them. Ensure `model.request_state_attributes('particle_q', "
            "'particle_qd', 'particle_f')` is called before creating the State, or "
            "construct the solver before calling `model.state()`."
        )
    if model.body_count > 0 and (state.body_q is None or state.body_qd is None or state.body_f is None):
        raise RuntimeError(
            "SolverXPBD requires body state attributes but this State was created "
            "without them. Ensure the model requested 'body_q', 'body_qd', and 'body_f' "
            "before creating the State."
        )

and call this early in step() for both state_in and state_out.

  1. Clarify the lifecycle (docs or migration notes) that XPBD expects its required attributes to be requested before any State used with it is created—either via the builder or by constructing the solver before calling model.state().

Without at least the runtime validation, XPBD will continue to fail with cryptic NoneType / device errors when used with pre-existing states under the new system.

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

361-401: enable_rne_postconstraint gating is effectively unconditional

The hasattr() check will always return True because State defines body_qdd and body_parent_f as fields initialized to None. This means the sensor is enabled on the first call regardless of whether these fields are actually allocated, incurring unnecessary computational cost on every subsequent step.

To gate on actual field allocation rather than existence:

if (state_out.body_qdd is not None) or (state_out.body_parent_f is not None):
    if not m.sensor_rne_postconstraint:
        if wp.config.verbose:
            print("Setting model.sensor_rne_postconstraint True")
        m.sensor_rne_postconstraint = True

This preserves behavior for users who request these fields while avoiding the sensor path when they are not present.

🧹 Nitpick comments (9)
newton/_src/solvers/featherstone/solver_featherstone.py (1)

95-119: Good declaration of required state attributes, but mind creation order of State objects

Requesting "body_q", "body_qd", "body_f", "joint_q", "joint_qd" (and particle attributes when present) aligns the solver with the new attribute-request mechanism and matches what step() uses.

Be aware this only applies to State instances created after SolverFeatherstone is constructed; States created earlier will still have None for these attributes under the new allocation scheme. If that pattern exists in examples/tests, consider either:

  • Documenting that solvers should be constructed before calling model.state(), or
  • Adding explicit validation in step() that required attributes are not None when model.body_count/model.joint_count are nonzero.
newton/_src/solvers/semi_implicit/solver_semi_implicit.py (1)

67-97: Appropriate attribute requests; consider guarding against pre-existing States

Gating the requests on model.body_count / model.particle_count and asking for "body_q", "body_qd", "body_f" plus "particle_q", "particle_qd", "particle_f" matches what SolverSemiImplicit.step() uses and keeps allocations minimal.

As with XPBD/Featherstone, these requests only affect State instances created after they run. If any callers construct States before the solver, those states can still have None for these fields under the new allocation scheme. It’s worth either:

  • Documenting the expected creation order (solver or builder requests before model.state()), and/or
  • Adding a light assertion/validation at the top of step() to fail fast with a clear message if required arrays are missing.
newton/_src/sim/state.py (1)

49-51: Extended state fields body_qdd / body_parent_f look consistent with existing spatial conventions

Adding body_qdd and body_parent_f as optional wp.spatial_vector arrays on the State surface matches how the rest of the code (MuJoCo conversion and IMU kernel) expects to consume them, and keeps them clearly documented.

If you want to be extra explicit, you could mirror the body_qd docstring wording (linear vs. angular ordering) for body_qdd and clarify the frame in which body_parent_f is expressed, but functionally this is fine as-is.

Also applies to: 61-63

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

229-301: Correct attribute requests and mapping docstring; note State creation order caveat

Requesting "body_q", "body_qd", "body_f", "joint_q", "joint_qd" from the Model is the right set for this solver and aligns with the new attribute-request API. The updated docstring for to_mjc_body_index (“Newton body index to MuJoCo body index”) also now matches how the mapping is constructed in _convert_to_mjc.

As with the other solvers, this only influences States created after the constructor runs. If callers construct States first and then attach SolverMuJoCo, those states can still lack the requested fields under the new scheme. Consider:

  • Documenting that the solver (or builder) should be set up before model.state(), and/or
  • Adding a short validation at the top of step() or update_newton_state() that required arrays aren’t None when model.body_count / model.joint_count are nonzero.
newton/_src/sim/builder.py (1)

497-498: State attribute request plumbing in ModelBuilder looks good; consider optional validation

The _requested_state_attributes set, its aggregation in add_builder(), and propagation into Model in finalize() form a coherent path for per-builder state-attribute requests. The request_state_attributes() API is minimal and appropriate for a prototype.

One improvement to consider (not blocking): optionally validate requested names early (e.g., at finalize() time) against a known whitelist (Model/State attribute names or a registry) to catch typos before they turn into missing fields at runtime.

Also applies to: 1389-1390, 4864-4872, 4966-4966

newton/examples/example_solver_data.py (2)

77-91: IMU sensor currently reads from pre-step state; likely off-by-one vs solver outputs

In simulate(), the call order is:

  • solver.step(self.state_0, self.state_1, ...)
  • self.imu.update(self.state_0)
  • then swap self.state_0, self.state_1.

After the swap, self.state_0 (logged in render()) holds the post-step state, but IMUSensor.update() saw the pre-step state. For solvers that compute and write accelerations/forces (e.g., via RNE) into the output state, this will give stale or uninitialized IMU readings.

Consider updating the IMU from the same state that the solver advances, e.g.:

-            self.solver.step(self.state_0, self.state_1, self.control, self.contacts, self.sim_dt)
-
-            self.imu.update(self.state_0)
-
-            # swap states
-            self.state_0, self.state_1 = self.state_1, self.state_0
+            self.solver.step(self.state_0, self.state_1, self.control, self.contacts, self.sim_dt)
+
+            # update IMU from the advanced state
+            self.imu.update(self.state_1)
+
+            # swap states
+            self.state_0, self.state_1 = self.state_1, self.state_0

This keeps the IMU data aligned with what the viewer later logs.

Also applies to: 94-108


119-159: Plot window text is inconsistent with plotted quantity

The ViewerPlot window title is configured as "Parent force", while the body text uses "Flap contact force" and a generic "Force" series label, yet the example is plotting IMU accelerations (imu_acc[2]).

For clarity, consider updating the labels to reflect the actual signal (e.g., "IMU acceleration (z)") or making them configurable via constructor kwargs.

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

699-707: Clarify and enforce request semantics for STATE-assigned attributes

request_state_attributes() simply adds raw strings to _requested_state_attributes, and _add_custom_attributes() then skips STATE-assigned attributes whose full name (including namespace, if any) is not in this set.

This is fine for a prototype, but it’s easy for users to silently misconfigure things (e.g., forgetting the namespace: prefix or never calling request_state_attributes() for a STATE custom attribute).

Two incremental improvements to consider (optional):

  • Document explicitly (in the request_state_attributes() docstring) that:
    • Names must match the full attribute key used in attribute_frequency / add_attribute ("namespace:attr" for namespaced attributes).
    • The method only affects subsequent state() calls; already-created State instances are not retrofitted.
  • Optionally, when assignment == STATE in _add_custom_attributes, emit a clear error or warning if an attribute is registered but not present in _requested_state_attributes, instead of silently skipping it.

This would make configuration errors around extended/STATE attributes easier to diagnose.

Also applies to: 724-731

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

1270-1271: Note: body_parent_f computation is not yet implemented.

The body_parent_f attribute is currently not populated (no-op with pass). This is consistent with the PR's acknowledged prototype limitations, but ensure this is tracked for completion.

📜 Review details

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 14001a7 and 616a160.

📒 Files selected for processing (11)
  • newton/_src/sensors/imu_sensor.py (1 hunks)
  • newton/_src/sim/builder.py (4 hunks)
  • newton/_src/sim/model.py (4 hunks)
  • newton/_src/sim/state.py (2 hunks)
  • newton/_src/solvers/featherstone/solver_featherstone.py (1 hunks)
  • newton/_src/solvers/mujoco/kernels.py (2 hunks)
  • newton/_src/solvers/mujoco/solver_mujoco.py (6 hunks)
  • newton/_src/solvers/semi_implicit/solver_semi_implicit.py (1 hunks)
  • newton/_src/solvers/xpbd/solver_xpbd.py (1 hunks)
  • newton/examples/example_solver_data.py (1 hunks)
  • newton/sensors.py (2 hunks)
🧰 Additional context used
🧠 Learnings (14)
📚 Learning: 2025-08-12T17:58:00.815Z
Learnt from: nvlukasz
Repo: newton-physics/newton PR: 519
File: newton/_src/solvers/featherstone/kernels.py:18-20
Timestamp: 2025-08-12T17:58:00.815Z
Learning: In newton/_src/core/__init__.py, the function transform_twist is properly re-exported from the spatial module via "from .spatial import (..., transform_twist, ...)" and included in __all__, making it available for import as "from ...core import transform_twist" from other modules.

Applied to files:

  • newton/sensors.py
📚 Learning: 2025-08-12T17:51:37.474Z
Learnt from: nvlukasz
Repo: newton-physics/newton PR: 519
File: newton/geometry.py:16-22
Timestamp: 2025-08-12T17:51:37.474Z
Learning: In the Newton public API refactor, common geometry symbols like ParticleFlags and ShapeFlags are now exposed at the top level in newton/__init__.py rather than through newton.geometry. The newton/geometry.py module is intentionally focused only on broad-phase collision detection classes (BroadPhaseAllPairs, BroadPhaseExplicit, BroadPhaseSAP).

Applied to files:

  • newton/sensors.py
📚 Learning: 2025-08-14T17:38:36.106Z
Learnt from: nvlukasz
Repo: newton-physics/newton PR: 519
File: newton/__init__.py:25-29
Timestamp: 2025-08-14T17:38:36.106Z
Learning: The Newton project prefers incremental __all__ building using __all__ += [...] pattern to group exports with their related imports, rather than a single consolidated __all__ at the end of the file.

Applied to files:

  • newton/sensors.py
📚 Learning: 2025-08-12T05:17:34.423Z
Learnt from: shi-eric
Repo: newton-physics/newton PR: 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.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
  • newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-09-22T21:08:31.901Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/examples/ik/example_ik_franka.py:121-123
Timestamp: 2025-09-22T21:08:31.901Z
Learning: In the newton physics framework, when creating warp arrays for IK solver joint variables using wp.array(self.model.joint_q, shape=(1, coord_count)), the resulting array acts as a reference/pointer to the original model's joint coordinates, so updates from the IK solver automatically reflect in the model's joint_q buffer used for rendering.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
  • newton/_src/sim/state.py
📚 Learning: 2025-09-22T21:08:31.901Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/examples/ik/example_ik_franka.py:121-123
Timestamp: 2025-09-22T21:08:31.901Z
Learning: In the newton physics framework, when creating warp arrays for IK solver joint variables using wp.array(self.model.joint_q, shape=(1, coord_count)), the resulting array acts as a reference/pointer to the original model's joint coordinates, so updates from the IK solver automatically reflect in the model's joint_q buffer used for rendering, despite the general warp documentation suggesting copies are made by default.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
  • newton/_src/sim/state.py
📚 Learning: 2025-08-18T15:56:26.587Z
Learnt from: adenzler-nvidia
Repo: newton-physics/newton PR: 552
File: newton/_src/solvers/mujoco/solver_mujoco.py:0-0
Timestamp: 2025-08-18T15:56:26.587Z
Learning: In Newton's MuJoCo solver, when transforming joint axes from Newton's internal frame to MuJoCo's expected frame, use wp.quat_rotate(joint_rot, axis) not wp.quat_rotate_inv(joint_rot, axis). The joint_rot represents rotation from joint-local to body frame, so forward rotation is correct.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
  • newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-09-25T16:14:22.033Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_objectives.py:0-0
Timestamp: 2025-09-25T16:14:22.033Z
Learning: In NVIDIA Warp's Newton physics library, wp.transform supports direct numerical indexing (e.g., body_tf[0], body_tf[1], body_tf[2] for position and body_tf[3], body_tf[4], body_tf[5], body_tf[6] for quaternion components) to access its elements. This is the standard API used throughout the codebase.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-11-28T11:12:40.784Z
Learnt from: camevor
Repo: newton-physics/newton PR: 1146
File: newton/examples/basic/example_basic_joints.py:206-213
Timestamp: 2025-11-28T11:12:40.784Z
Learning: In Newton's standard solvers (XPBD, SemiImplicit, SolverMuJoCo), spatial vectors in State.body_qd use the ordering (linear, angular): wp.spatial_top(qd) returns linear velocity and wp.spatial_bottom(qd) returns angular velocity. This is the opposite of typical Modern Robotics spatial twist notation where (angular, linear) is used.

Applied to files:

  • newton/_src/sim/state.py
📚 Learning: 2025-08-19T12:27:30.629Z
Learnt from: preist-nvidia
Repo: newton-physics/newton PR: 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/examples/example_solver_data.py
  • newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-08-25T21:41:45.795Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 634
File: newton/tests/test_examples.py:329-333
Timestamp: 2025-08-25T21:41:45.795Z
Learning: Newton examples use a centralized CLI argument parsing system in newton/examples/__init__.py. The create_parser() function defines common arguments like --device, --viewer, --output-path, and --num-frames, while init(parser) creates viewers based on parsed arguments. Individual example scripts don't need to define these flags themselves - they inherit them from the centralized system via the example_map and runpy execution.

Applied to files:

  • newton/examples/example_solver_data.py
📚 Learning: 2025-08-19T12:27:30.629Z
Learnt from: preist-nvidia
Repo: newton-physics/newton PR: 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/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-11-24T08:05:21.390Z
Learnt from: adenzler-nvidia
Repo: newton-physics/newton PR: 1107
File: newton/_src/solvers/mujoco/kernels.py:973-974
Timestamp: 2025-11-24T08:05:21.390Z
Learning: In Newton's MuJoCo solver integration (newton/_src/solvers/mujoco/), the mapping between Newton joints and MuJoCo joints is not 1-to-1. Instead, each Newton DOF maps to a distinct MuJoCo joint. This means that for multi-DOF joints (like D6 with 6 DOFs), there will be 6 corresponding MuJoCo joints, each with its own properties (margin, solimp, solref, etc.). The mapping is done via dof_to_mjc_joint array, ensuring each DOF's properties are written to its own MuJoCo joint without overwriting.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-07-14T03:57:29.670Z
Learnt from: Kenny-Vilella
Repo: newton-physics/newton PR: 398
File: newton/examples/example_mujoco.py:352-352
Timestamp: 2025-07-14T03:57:29.670Z
Learning: The use_mujoco option in newton/examples/example_mujoco.py is currently unsupported and causes crashes. The code automatically disables this option with a warning message when users attempt to enable it. This is intentionally kept as a placeholder for future implementation.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
🧬 Code graph analysis (5)
newton/_src/solvers/featherstone/solver_featherstone.py (2)
newton/_src/sim/builder.py (2)
  • request_state_attributes (4864-4871)
  • particle_count (937-941)
newton/_src/sim/model.py (1)
  • request_state_attributes (699-706)
newton/sensors.py (1)
newton/_src/sensors/imu_sensor.py (1)
  • IMUSensor (59-93)
newton/_src/solvers/xpbd/solver_xpbd.py (2)
newton/_src/sim/builder.py (3)
  • body_count (923-927)
  • request_state_attributes (4864-4871)
  • particle_count (937-941)
newton/_src/sim/model.py (1)
  • request_state_attributes (699-706)
newton/_src/solvers/semi_implicit/solver_semi_implicit.py (2)
newton/_src/sim/builder.py (3)
  • body_count (923-927)
  • request_state_attributes (4864-4871)
  • particle_count (937-941)
newton/_src/sim/model.py (1)
  • request_state_attributes (699-706)
newton/_src/solvers/mujoco/solver_mujoco.py (6)
newton/_src/solvers/mujoco/kernels.py (1)
  • convert_rigid_forces_from_mj_kernel (1230-1271)
newton/_src/sim/builder.py (2)
  • request_state_attributes (4864-4871)
  • body_count (923-927)
newton/_src/sim/model.py (2)
  • request_state_attributes (699-706)
  • state (517-571)
newton/_src/sim/state.py (2)
  • State (21-118)
  • body_count (93-97)
newton/_src/geometry/broad_phase_nxn.py (2)
  • launch (295-361)
  • launch (378-435)
newton/_src/geometry/broad_phase_sap.py (1)
  • launch (506-650)
🪛 GitHub Actions: Pull Request
newton/_src/solvers/xpbd/solver_xpbd.py

[error] 242-242: AttributeError: 'NoneType' object has no attribute 'device' during SolverXPBD.step (Warp context). Likely state_in.particle_q is None when cloning state.

newton/_src/sim/builder.py

[error] 1-1: Doctest run failed in ModelBuilder: AttributeError: 'NoneType' object has no attribute 'device' during SolverXPBD.step (XPBD solver).

newton/_src/solvers/mujoco/kernels.py

[error] 23-23: ModuleNotFoundError: No module named 'mujoco_warp' while importing MuJoCo kernels. Ensure the mujoco_warp package is installed and available in the environment.

⏰ 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). (1)
  • GitHub Check: Run GPU Benchmarks / Run GPU Benchmarks on AWS EC2
🔇 Additional comments (5)
newton/sensors.py (1)

28-31: Expose IMUSensor via newton.sensors looks correct

Importing IMUSensor from the internal module and adding it to __all__ matches how other sensors are surfaced, so this extension of the public API is consistent.

Also applies to: 38-45

newton/_src/sensors/imu_sensor.py (1)

24-57: IMU kernel math and spatial layout look sound

The kernel’s computation

  • r = R_body * (site_pos_body - com_body)
  • vel_ang = spatial_bottom(body_qd)
  • acc_lin = a_com + alpha × r + w × (w × r)
  • rotation of both acc_lin and vel_ang into the sensor frame using q = body_quat * site_quat and quat_rotate_inv

matches the standard rigid-body kinematics relation for point acceleration and respects the existing (linear, angular) spatial-vector ordering.

No functional changes needed here; any further refinement (e.g., specific-force vs. absolute acceleration convention) can be handled at the API/docs level.

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

23-24: Resolve the missing mujoco_warp dependency.

The pipeline is failing with ModuleNotFoundError: No module named 'mujoco_warp'. Ensure the dependency is properly installed in the CI environment.

Additionally, importing from private modules (_src.sensor, _src.types) may be fragile if the mujoco_warp package changes its internal structure.


1267-1267: Resolve the TODO: Clarify the correct frame type for linear acceleration computation.

The TODO comment reflects a genuine ambiguity. In MuJoCo, BODY refers to the inertial/center-of-mass frame, while XBODY refers to the kinematic/local body frame. The choice depends on whether the linear acceleration should be computed in the body's inertial frame or its local kinematic frame. Review the semantic requirements of the _framelinacc computation and clarify which frame type aligns with the intended behavior, then update or resolve this TODO accordingly.


1250-1266: Verify the _framelinacc function signature and parameter correctness.

The call to _framelinacc passes mjw_body_rootid four times and mjw_xpos five times. While this repetition could be correct if the function signature semantically requires distinct parameters that happen to use the same values, it warrants verification against the actual function definition. Confirm the parameter count and order match the expected signature.

newton/_src/sensors/imu_sensor.py Outdated Show resolved Hide resolved
newton/_src/sensors/imu_sensor.py Outdated Show resolved Hide resolved
newton/_src/sim/model.py Show resolved Hide resolved
newton/_src/solvers/mujoco/kernels.py Outdated Show resolved Hide resolved
newton/_src/solvers/mujoco/solver_mujoco.py Show resolved Hide resolved
Copy link
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

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

Actionable comments posted: 0

♻️ Duplicate comments (3)
newton/_src/solvers/xpbd/solver_xpbd.py (1)

103-107: SolverXPBD requests are correct but rely on Model.state’s legacy fallback to avoid None states

Requesting "body_q", "body_qd", "body_f" and "particle_q", "particle_qd", "particle_f" here is the right thing for new States created after the solver is constructed. However, the current pipeline failure (state_in.particle_q is None) comes from Model.state() allocating no arrays when _requested_state_attributes is still empty. That needs to be fixed in Model.state (see comment there); once the legacy fallback is in place, this constructor code will work as intended without changes.

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

663-680: Prefer explicit is not None checks when deciding whether to launch convert_rigid_forces_from_mj_kernel

The new block:

# Update rigid force fields on state.
if any((state.body_qdd, state.body_parent_f)):
    ...
    wp.launch(
        convert_rigid_forces_from_mj_kernel,
        (nworld, bodies_per_world),
        ...
        outputs=[state.body_qdd, state.body_parent_f],
    )

relies on Python truthiness of wp.array objects, which is non-obvious and inconsistent with the rest of this file (where you usually test is not None for optional arrays, e.g. if state.body_f is not None:).

For clarity and to avoid any surprises if wp.array.__bool__ ever changes, I’d strongly suggest:

-        if any((state.body_qdd, state.body_parent_f)):
+        if (state.body_qdd is not None) or (state.body_parent_f is not None):

This also mirrors the way the kernel itself treats optional outputs via its if body_qdd: / if body_parent_f: guards.

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

533-563: Restore backward-compatible state allocation when no attributes are explicitly requested

Right now state() uses requested = self._requested_state_attributes directly, so when the set is empty (the default after ModelBuilder.finalize()), no core state arrays are allocated at all. This is the root cause of the XPBD failure (state_in.particle_q is None), and it breaks existing user code that creates State objects before any solver or sensor calls request_state_attributes() (as documented in the earlier review).

To keep the new gating while preserving old behavior, treat an empty request set as “legacy mode”: allocate the standard core fields by default, and only apply fine-grained gating once something has been explicitly requested. For example:

-        requested = self._requested_state_attributes
+        requested = self._requested_state_attributes
+
+        # Backwards-compatibility: if no attributes have been requested explicitly,
+        # allocate the standard core solver-facing fields as before. Fine-grained
+        # gating only applies once callers start requesting attributes.
+        if not requested:
+            requested = {
+                "particle_q",
+                "particle_qd",
+                "particle_f",
+                "body_q",
+                "body_qd",
+                "body_f",
+                "joint_q",
+                "joint_qd",
+            }

This fixes the regression for existing solvers and doctests while keeping extended attributes (body_qdd, body_parent_f, custom STATE attributes) opt-in.

🧹 Nitpick comments (5)
newton/_src/sim/model.py (3)

564-568: Guard body_qdd / body_parent_f allocation on the presence of body_qd

Allocating body_qdd and body_parent_f via wp.zeros_like(self.body_qd, ...) assumes self.body_qd has been allocated—i.e., that "body_qd" is also in the requested set. If someone (or a sensor) requests "body_qdd" without "body_qd", this will raise at state creation time.

Consider either documenting that "body_qd" must be requested together with "body_qdd"/"body_parent_f" or adding a defensive check/assert, e.g.:

-        if "body_qdd" in requested:
-            s.body_qdd = wp.zeros_like(self.body_qd, requires_grad=requires_grad)
+        if "body_qdd" in requested:
+            if self.body_qd is None:
+                raise RuntimeError("Requesting 'body_qdd' requires 'body_qd' to be available on the model.")
+            s.body_qdd = wp.zeros_like(self.body_qd, requires_grad=requires_grad)

(and similarly for "body_parent_f").


701-708: Avoid duplicating Model.request_state_attributes definitions

The request_state_attributes() method here looks good, but the external context shows it is defined twice in newton/_src/sim/model.py (around lines 700–707 and again around 5024–5031). Duplicating this public API in the same class is confusing and risks divergence if one copy is edited later.

Recommend keeping a single definition (preferably near other public API methods) and removing the duplicate.


730-732: Only gate STATE custom attributes when a non-empty request set is present

The new guard:

if assignment == ModelAttributeAssignment.STATE and full_name not in self._requested_state_attributes:
    continue

means no STATE-assigned custom attributes are ever attached unless they are explicitly listed in _requested_state_attributes. For pre-existing custom STATE attributes, this is a silent behavior change similar to the core-field gating issue.

To keep backward compatibility while still allowing opt-in gating, consider only enforcing this filter once some attributes have been requested:

-            if assignment == ModelAttributeAssignment.STATE and full_name not in self._requested_state_attributes:
-                continue
+            if (
+                assignment == ModelAttributeAssignment.STATE
+                and self._requested_state_attributes
+                and full_name not in self._requested_state_attributes
+            ):
+                continue

That way, existing code that never calls request_state_attributes() continues to see its STATE custom attributes, and selective materialization kicks in only after callers start using the new API.

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

1231-1273: convert_rigid_forces_from_mj_kernel spatial mapping looks consistent; body_parent_f remains TODO

The new kernel correctly follows the MuJoCo Warp → Newton spatial convention:

  • mjw_cacc is in MuJoCo Warp’s (angular, linear) ordering.
  • _framelinacc(...) returns the COM linear acceleration in Newton/world coordinates.
  • wp.spatial_vector(lin, wp.spatial_top(cacc)) thus builds a Newton spatial vector (linear_accel, angular_accel) (per the recorded convention where spatial_top is angular for MuJoCo Warp data and linear for Newton).

This matches the documented convention difference between MuJoCo Warp and Newton’s spatial vectors (see retrieved learnings for this PR).

body_parent_f is intentionally left as a stub (pass), which is acceptable for the prototype but means state.body_parent_f will stay unused even if allocated. When you move beyond prototype stage, you’ll likely want to follow up by wiring the appropriate RNE wrench into body_parent_f.

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

373-401: enable_rne_postconstraint should check for allocated fields, not just attribute existence

State always has body_qdd and body_parent_f attributes (initialized to None), so:

if any(hasattr(state_out, field) for field in rne_postconstraint_fields):

is always True, and sensor_rne_postconstraint will be enabled unconditionally once per solver instance, regardless of whether those fields were actually requested/allocated.

If the intent is to only pay the cost of RNE postconstraint when the extended fields are in use, consider checking for non-None or tying this directly to the model’s requested attributes, e.g.:

-        if any(hasattr(state_out, field) for field in rne_postconstraint_fields):
+        if any(getattr(state_out, field, None) is not None for field in rne_postconstraint_fields):

or even:

if any(name in self.model._requested_state_attributes for name in rne_postconstraint_fields):
    ...

This keeps the behavior aligned with the new request-based allocation mechanism.

📜 Review details

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 616a160 and 225ea4f.

📒 Files selected for processing (5)
  • newton/_src/sim/builder.py (4 hunks)
  • newton/_src/sim/model.py (4 hunks)
  • newton/_src/solvers/mujoco/kernels.py (2 hunks)
  • newton/_src/solvers/mujoco/solver_mujoco.py (6 hunks)
  • newton/_src/solvers/xpbd/solver_xpbd.py (1 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
  • newton/_src/sim/builder.py
🧰 Additional context used
🧠 Learnings (21)
📚 Learning: 2025-08-12T05:17:34.423Z
Learnt from: shi-eric
Repo: newton-physics/newton PR: 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.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
  • newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-11-28T11:12:40.784Z
Learnt from: camevor
Repo: newton-physics/newton PR: 1146
File: newton/examples/basic/example_basic_joints.py:206-213
Timestamp: 2025-11-28T11:12:40.784Z
Learning: In Newton's standard solvers (XPBD, SemiImplicit, SolverMuJoCo), spatial vectors in State.body_qd use the ordering (linear, angular): wp.spatial_top(qd) returns linear velocity and wp.spatial_bottom(qd) returns angular velocity. This is the opposite of typical Modern Robotics spatial twist notation where (angular, linear) is used.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-12-01T17:00:28.841Z
Learnt from: camevor
Repo: newton-physics/newton PR: 1161
File: newton/_src/solvers/mujoco/kernels.py:1268-1268
Timestamp: 2025-12-01T17:00:28.841Z
Learning: MuJoCo Warp spatial vectors (e.g., cacc, cvel) use the convention (angular, linear), which is opposite to Newton's (linear, angular) convention. When converting from MuJoCo Warp to Newton format, spatial_top() on MuJoCo Warp data extracts the angular component, and spatial_bottom() extracts the linear component.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-08-18T15:56:26.587Z
Learnt from: adenzler-nvidia
Repo: newton-physics/newton PR: 552
File: newton/_src/solvers/mujoco/solver_mujoco.py:0-0
Timestamp: 2025-08-18T15:56:26.587Z
Learning: In Newton's MuJoCo solver, when transforming joint axes from Newton's internal frame to MuJoCo's expected frame, use wp.quat_rotate(joint_rot, axis) not wp.quat_rotate_inv(joint_rot, axis). The joint_rot represents rotation from joint-local to body frame, so forward rotation is correct.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
  • newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-11-28T11:12:40.006Z
Learnt from: camevor
Repo: newton-physics/newton PR: 1146
File: newton/examples/basic/example_basic_joints.py:215-222
Timestamp: 2025-11-28T11:12:40.006Z
Learning: In Warp's spatial vector convention used by Newton: spatial_bottom() extracts the angular velocity component (omega) and spatial_top() extracts the linear velocity component (v) from a spatial twist (qd). Similarly for spatial wrenches: spatial_bottom() is torque and spatial_top() is force.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-12-01T16:21:36.558Z
Learnt from: jvonmuralt
Repo: newton-physics/newton PR: 1160
File: newton/_src/sim/articulation.py:0-0
Timestamp: 2025-12-01T16:21:36.558Z
Learning: In Newton's Featherstone solver, during forward kinematics for FREE/DISTANCE joints, the computed spatial velocity v_wc follows Featherstone's spatial twist convention where the linear component (spatial_top) represents the velocity at the world origin (0,0,0), not at the joint anchor or COM. To convert to State.body_qd (which stores COM velocity), use v_com = v_origin + ω × x_com, where x_com is the absolute world position of the COM.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-09-22T21:03:39.624Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_lbfgs_optimizer.py:739-752
Timestamp: 2025-09-22T21:03:39.624Z
Learning: The L-BFGS optimizer in newton/_src/sim/ik/ik_lbfgs_optimizer.py currently intentionally only supports additive updates (assuming n_coords == n_dofs). Velocity space integration for joints with mismatched coordinate/DOF dimensions (like free/ball joints) is planned for future work and should not be flagged as an issue in current reviews.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
  • newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-09-25T16:14:22.033Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_objectives.py:0-0
Timestamp: 2025-09-25T16:14:22.033Z
Learning: wp.transform in NVIDIA Warp supports indexing using numerical indices (e.g., body_tf[0], body_tf[3], etc.) to access translation and rotation components. The indexing approach used in the Newton codebase is valid and compiles correctly.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-09-25T16:14:22.033Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_objectives.py:0-0
Timestamp: 2025-09-25T16:14:22.033Z
Learning: In NVIDIA Warp's Newton physics library, wp.transform supports direct numerical indexing (e.g., body_tf[0], body_tf[1], body_tf[2] for position and body_tf[3], body_tf[4], body_tf[5], body_tf[6] for quaternion components) to access its elements. This is the standard API used throughout the codebase.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-09-22T21:08:31.901Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/examples/ik/example_ik_franka.py:121-123
Timestamp: 2025-09-22T21:08:31.901Z
Learning: In the newton physics framework, when creating warp arrays for IK solver joint variables using wp.array(self.model.joint_q, shape=(1, coord_count)), the resulting array acts as a reference/pointer to the original model's joint coordinates, so updates from the IK solver automatically reflect in the model's joint_q buffer used for rendering, despite the general warp documentation suggesting copies are made by default.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
  • newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-09-22T21:08:31.901Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/examples/ik/example_ik_franka.py:121-123
Timestamp: 2025-09-22T21:08:31.901Z
Learning: In the newton physics framework, when creating warp arrays for IK solver joint variables using wp.array(self.model.joint_q, shape=(1, coord_count)), the resulting array acts as a reference/pointer to the original model's joint coordinates, so updates from the IK solver automatically reflect in the model's joint_q buffer used for rendering.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
  • newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-11-24T08:05:21.390Z
Learnt from: adenzler-nvidia
Repo: newton-physics/newton PR: 1107
File: newton/_src/solvers/mujoco/kernels.py:973-974
Timestamp: 2025-11-24T08:05:21.390Z
Learning: In Newton's MuJoCo solver integration (newton/_src/solvers/mujoco/), the mapping between Newton joints and MuJoCo joints is not 1-to-1. Instead, each Newton DOF maps to a distinct MuJoCo joint. This means that for multi-DOF joints (like D6 with 6 DOFs), there will be 6 corresponding MuJoCo joints, each with its own properties (margin, solimp, solref, etc.). The mapping is done via dof_to_mjc_joint array, ensuring each DOF's properties are written to its own MuJoCo joint without overwriting.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
  • newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-08-12T18:04:06.577Z
Learnt from: nvlukasz
Repo: newton-physics/newton PR: 519
File: newton/_src/solvers/featherstone/kernels.py:75-75
Timestamp: 2025-08-12T18:04:06.577Z
Learning: The Newton physics framework requires nightly Warp builds, which means compatibility concerns with older stable Warp versions (like missing functions such as wp.spatial_adjoint) are not relevant for this project.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
  • newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-08-19T12:27:30.629Z
Learnt from: preist-nvidia
Repo: newton-physics/newton PR: 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/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-08-19T12:27:30.629Z
Learnt from: preist-nvidia
Repo: newton-physics/newton PR: 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/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-07-14T03:57:29.670Z
Learnt from: Kenny-Vilella
Repo: newton-physics/newton PR: 398
File: newton/examples/example_mujoco.py:352-352
Timestamp: 2025-07-14T03:57:29.670Z
Learning: The use_mujoco option in newton/examples/example_mujoco.py is currently unsupported and causes crashes. The code automatically disables this option with a warning message when users attempt to enable it. This is intentionally kept as a placeholder for future implementation.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-09-09T12:50:42.929Z
Learnt from: gdaviet
Repo: newton-physics/newton PR: 750
File: newton/_src/solvers/implicit_mpm/render_grains.py:142-160
Timestamp: 2025-09-09T12:50:42.929Z
Learning: The `update_render_grains` function in newton/_src/solvers/implicit_mpm/render_grains.py should return early when `state.velocity_field is None` because its purpose is specifically to advect grains in the velocity field. Without a velocity field, the function should be a no-op.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-09-26T05:58:21.013Z
Learnt from: WenchaoHuang
Repo: newton-physics/newton PR: 835
File: newton/_src/solvers/style3d/collision/collision.py:49-53
Timestamp: 2025-09-26T05:58:21.013Z
Learning: In newton/_src/solvers/style3d/collision/collision.py, the broad_phase buffers (broad_phase_ee, broad_phase_ef, broad_phase_vf) created with wp.array() don't need zero-initialization because they follow a write-first-read-later pattern: BVH query kernels in frame_begin() populate these buffers completely (write-only operations), and collision kernels only read from them afterward in accumulate_contact_force(), ensuring no undefined reads occur.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-09-18T07:05:56.836Z
Learnt from: gdaviet
Repo: newton-physics/newton PR: 750
File: newton/_src/solvers/implicit_mpm/solve_rheology.py:969-986
Timestamp: 2025-09-18T07:05:56.836Z
Learning: In newton/_src/solvers/implicit_mpm/solve_rheology.py, transposed_strain_mat parameter cannot be None - the type signature was corrected to reflect this guarantee, eliminating the need for None checks when accessing transposed_strain_mat.offsets.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-09-26T05:58:21.013Z
Learnt from: WenchaoHuang
Repo: newton-physics/newton PR: 835
File: newton/_src/solvers/style3d/collision/collision.py:49-53
Timestamp: 2025-09-26T05:58:21.013Z
Learning: In newton/_src/solvers/style3d/collision/collision.py, the broad_phase buffers (broad_phase_ee, broad_phase_ef, broad_phase_vf) created with wp.array() don't need zero-initialization because the BVH query kernels that use them always write the count to query_results[0, tid] = 0 first before populating any data, ensuring no undefined reads occur.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-09-22T21:03:18.367Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_lbfgs_optimizer.py:638-648
Timestamp: 2025-09-22T21:03:18.367Z
Learning: The L-BFGS optimizer in newton/_src/sim/ik/ik_lbfgs_optimizer.py currently only supports additive updates, meaning it assumes n_coords == n_dofs. Velocity space integration for joints with mismatched coordinate/DOF dimensions (like free/ball joints) is planned for future work.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
🧬 Code graph analysis (2)
newton/_src/solvers/xpbd/solver_xpbd.py (2)
newton/_src/sim/builder.py (3)
  • body_count (925-929)
  • request_state_attributes (5025-5032)
  • particle_count (939-943)
newton/_src/sim/model.py (1)
  • request_state_attributes (701-708)
newton/_src/solvers/mujoco/solver_mujoco.py (4)
newton/_src/solvers/mujoco/kernels.py (1)
  • convert_rigid_forces_from_mj_kernel (1232-1273)
newton/_src/sim/builder.py (2)
  • request_state_attributes (5025-5032)
  • body_count (925-929)
newton/_src/sim/model.py (2)
  • request_state_attributes (701-708)
  • state (519-573)
newton/_src/sim/state.py (2)
  • State (21-118)
  • body_count (93-97)
🪛 GitHub Actions: Pull Request
newton/_src/solvers/xpbd/solver_xpbd.py

[error] 242-242: XPBD solver.step() failed: AttributeError: 'NoneType' object has no attribute 'device' (state_in.particle_q is None).

newton/_src/solvers/mujoco/kernels.py

[error] 23-23: ModuleNotFoundError: No module named 'mujoco_warp'. Ensure 'mujoco_warp' is installed.

⏰ 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). (1)
  • GitHub Check: Run GPU Benchmarks / Run GPU Benchmarks on AWS EC2
🔇 Additional comments (4)
newton/_src/sim/model.py (1)

456-456: Initialization of _requested_state_attributes is fine and matches the new API

Storing requested state attributes in a dedicated set[str] is an appropriate internal representation and fits the new request_state_attributes() API.

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

45-54: Kernel import of convert_rigid_forces_from_mj_kernel is consistent with new RNE extraction path

Adding convert_rigid_forces_from_mj_kernel to the imported kernel list is appropriate and matches the new update_newton_state() logic that populates state.body_qdd/state.body_parent_f from MuJoCo data.


285-286: Requesting core body/joint state attributes in SolverMuJoCo.__init__ is correct but depends on Model.state fallback

Having the MuJoCo solver request "body_q", "body_qd", "body_f", "joint_q", and "joint_qd" up front is the right behavior for any State objects created after the solver is constructed.

However, just like in XPBD, this does not fix existing code that creates states before constructing the solver; that regression must be addressed in Model.state() via the empty-request legacy fallback (see comment in newton/_src/sim/model.py). Once that is in place, this request call ensures all later states remain consistent with the solver’s expectations.

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

22-25: Make mujoco_warp imports optional to avoid hard dependency at module import time

Importing mujoco_warp._src.sensor and mujoco_warp._src.types at the top level causes ModuleNotFoundError when newton._src.solvers.mujoco.kernels is imported in environments where mujoco_warp is not installed. This breaks users that never touch the MuJoCo solver.

Recommend wrapping these imports in a try/except and providing a minimal stub, relying on SolverMuJoCo.import_mujoco() to enforce the hard dependency when the solver is actually used. This pattern is already used in newton/_src/utils/import_usd.py for handling optional USD dependencies. For example:

-import warp as wp
-from mujoco_warp._src.sensor import _framelinacc
-from mujoco_warp._src.types import ObjType
+import warp as wp
+try:
+    from mujoco_warp._src.sensor import _framelinacc  # type: ignore[import]
+    from mujoco_warp._src.types import ObjType  # type: ignore[import]
+except ImportError:
+    # Keep the module importable without mujoco_warp; SolverMuJoCo.import_mujoco()
+    # will still raise if the MuJoCo backend is actually used.
+    _framelinacc = None  # type: ignore[assignment]
+
+    class ObjType:  # minimal stub for type/attribute access in kernels
+        BODY = 0

This unblocks tests and non-MuJoCo usage while preserving the requirement that mujoco_warp must be installed to actually run the MuJoCo backend.

Copy link
Member

@adenzler-nvidia adenzler-nvidia left a comment

Choose a reason for hiding this comment

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

Thanks, some comments.

newton/_src/solvers/featherstone/solver_featherstone.py Outdated Show resolved Hide resolved
newton/_src/solvers/mujoco/kernels.py Show resolved Hide resolved
newton/_src/solvers/mujoco/solver_mujoco.py Outdated Show resolved Hide resolved
newton/_src/solvers/mujoco/solver_mujoco.py Show resolved Hide resolved
newton/_src/sensors/imu_sensor.py Outdated Show resolved Hide resolved
newton/_src/sim/model.py Show resolved Hide resolved
@camevor camevor force-pushed the extended-state-attributes branch from 225ea4f to 2c7f58a Compare December 11, 2025 17:37
Copy link
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

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

Actionable comments posted: 1

♻️ Duplicate comments (3)
newton/_src/sensors/imu_sensor.py (2)

16-17: Update module docstring to describe IMU behavior (not frame transforms)

The top-level docstring still describes a “Frame Transform Sensor”, which no longer matches this module. Consider updating it to reflect IMU functionality, e.g.:

-"""Frame Transform Sensor - measures transforms relative to sites."""
+"""IMU sensor – computes linear acceleration and angular velocity at reference sites."""

59-93: Defensively validate required State fields in IMUSensor.update

IMUSensor.update() assumes state.body_q, state.body_qd, and especially state.body_qdd are allocated:

  • model.request_state_attributes("body_qdd") only affects State objects created after that call.
  • Users can easily pass an older State or one produced by a solver that doesn’t populate body_qdd yet, leading to hard-to-debug Warp launch errors.

Add an explicit check up front and raise a clear error when prerequisites are not met, e.g.:

     def update(self, state: State):
         """Update the IMU sensor.
@@
-        wp.launch(
+        if state.body_q is None or state.body_qd is None or state.body_qdd is None:
+            raise RuntimeError(
+                "IMUSensor requires 'body_q', 'body_qd', and 'body_qdd' on the provided State. "
+                "Ensure the model requested 'body_qdd' before creating the State and that the "
+                "active solver populates 'body_qdd' (currently SolverMuJoCo with RNE enabled)."
+            )
+
+        wp.launch(
             compute_imu_sensor_kernel,
             dim=self.n_sensors,
@@
             device=self.model.device,
         )

This will make prototype usage much easier to diagnose and won’t affect the happy path.

Check whether Warp currently provides a clearer exception when a kernel input argument is None or uninitialized, or if a manual Python-side check (as above) is still the recommended way to surface such configuration errors.
newton/_src/solvers/mujoco/solver_mujoco.py (1)

466-467: Make RNE/extended-force path depend on allocated fields, not attribute presence / generic truthiness

Right now:

  • enable_rne_postconstraint checks any(hasattr(state_out, field) for field in {"body_qdd", "body_parent_f"}), which is always True because State defines those attributes (initialized to None). This means sensor_rne_postconstraint is effectively enabled unconditionally on the mujoco_warp path after the first step, even when no extended fields were requested.
  • update_newton_state gates the force/acceleration kernel with if any((state.body_qdd, state.body_parent_f)), which relies on truthiness of wp.array vs None and is harder to read than an explicit is not None check.

To better match the “compute on demand” intent and avoid paying for RNE when nothing consumes the results, consider:

     def enable_rne_postconstraint(self, state_out: State):
         """Request computation of RNE forces if required for state fields."""
         rne_postconstraint_fields = {"body_qdd", "body_parent_f"}
         # TODO: handle use_mujoco_cpu
         m = self.mjw_model
         if m.sensor_rne_postconstraint:
             return
-        if any(hasattr(state_out, field) for field in rne_postconstraint_fields):
+        # Only enable if at least one extended field was actually allocated on the State.
+        if (state_out.body_qdd is not None) or (state_out.body_parent_f is not None):
             if wp.config.verbose:
                 print("Setting model.sensor_rne_postconstraint True")
             m.sensor_rne_postconstraint = True
             # required for cfrc_ext, cfrc_int, cacc
@@
-        # Update rigid force fields on state.
-        if any((state.body_qdd, state.body_parent_f)):
+        # Update rigid force fields on state only if the corresponding outputs were allocated.
+        if (state.body_qdd is not None) or (state.body_parent_f is not None):
             # Launch over MuJoCo bodies
             nbody = self.mjc_body_to_newton.shape[1]
             wp.launch(
                 convert_rigid_forces_from_mj_kernel,
                 (nworld, nbody),
                 inputs=[
                     self.mjc_body_to_newton,
                     self.mjw_model.body_rootid,
                     self.mjw_data.xipos,
                     self.mjw_data.subtree_com,
                     self.mjw_data.cacc,
                     self.mjw_data.cvel,
                     # self.mjw_data.cfrc_int,
                 ],
                 outputs=[state.body_qdd, state.body_parent_f],
             )

This keeps the prototype simple, avoids accidental always‑on RNE for mujoco_warp, and makes the extended-force path clearly conditional on the new State attributes actually being allocated.

#!/bin/bash
# Verify where body_qdd/body_parent_f are allocated and used to ensure the
# above gating matches actual usage across the repo.

set -euo pipefail

echo "== Find allocations/usages of body_qdd and body_parent_f =="
rg -n "body_qdd|body_parent_f" newton -C3 || true

echo
echo "== Show current enable_rne_postconstraint and extended-force launch =="
rg -n "enable_rne_postconstraint|convert_rigid_forces_from_mj_kernel" newton/_src/solvers/mujoco/solver_mujoco.py -C5 || true

Also applies to: 482-494, 784-801

🧹 Nitpick comments (8)
newton/_src/sim/model.py (1)

740-742: STATE-assigned custom attributes are currently never attached to State

_add_custom_attributes() now unconditionally skips attributes with assignment == ModelAttributeAssignment.STATE, so even if callers register custom STATE attributes via add_attribute(...) and request them, they will never be copied onto State instances.

If the intent is to gate STATE attributes on the new _requested_state_attributes mechanism rather than disable them entirely, consider changing this to check against the effective requested set instead of early‑continuing on all STATE attributes. For example, something along the lines of:

-        for full_name, _freq in self.attribute_frequency.items():
-            if self.attribute_assignment.get(full_name, ModelAttributeAssignment.MODEL) != assignment:
-                continue
-
-            if assignment == ModelAttributeAssignment.STATE:
-                continue
+        requested = set()
+        if assignment == ModelAttributeAssignment.STATE:
+            requested = set(self.get_requested_state_attributes())
+
+        for full_name, _freq in self.attribute_frequency.items():
+            if self.attribute_assignment.get(full_name, ModelAttributeAssignment.MODEL) != assignment:
+                continue
+
+            if assignment == ModelAttributeAssignment.STATE and full_name not in requested:
+                continue

(adapting the exact membership test if you want to map from namespaced keys to base attribute names). This keeps STATE custom attributes usable while still honoring the new request‑based gating.

Also applies to: 581-581

newton/sensors.py (1)

28-31: Expose IMUSensor via newton.sensors – looks good

Importing IMUSensor and adding it to __all__ is consistent with how other sensors are surfaced from this module; no functional issues here. One tiny nit: you now have two consecutive # Raycast sensors comment lines (Lines 33–34); you can drop one for clarity, but it’s purely cosmetic.

Also applies to: 42-50

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

1571-1572: Propagation of requested attributes in add_builder is well‑aligned with reuse semantics

Merging builder._requested_state_attributes into self._requested_state_attributes ensures that sub‑builders bring along their extended‑state requirements when composed into a larger scene, and the set keeps this idempotent. Consider briefly noting this behavior in the add_builder docstring so users understand that state‑attribute requirements from sub‑builders are aggregated automatically, but functionally this is sound.


5169-5177: ModelBuilder.request_state_attributes API is a clean mirror of Model.request_state_attributes

The method is minimal and correctly accumulates attributes into the internal set, matching the behavior of Model.request_state_attributes shown in newton/_src/sim/model.py. You might optionally add a tiny note in the docstring that these requests are forwarded to the Model during finalize, to make the lifecycle explicit, but behavior-wise this is fine.

newton/examples/example_solver_data.py (4)

33-33: Inconsistent text label.

The ViewerPlot is initialized with title "Parent force", but the label inside render() at line 158 says "Flap contact force". Consider making these consistent.

-        self.plot_window = ViewerPlot(viewer, "Parent force", scale_min=0, graph_size=(400, 200))
+        self.plot_window = ViewerPlot(viewer, "IMU Acceleration", scale_min=0, graph_size=(400, 200))

92-93: Remove commented-out dead code.

These commented lines appear to be remnants from development. Consider removing them before the PR is marked ready for review.

             # swap states
             self.state_0, self.state_1 = self.state_1, self.state_0
-        # for update in step graph
-        # self.solver.update_data()

101-107: Remove dead code and consider removing or guarding the print statement.

Lines 101-104 contain commented-out code. Line 107 prints to console every frame (200 Hz), which can spam output and impact performance. Consider removing the print or guarding it behind a debug flag.

     def step(self):
         if self.graph:
             wp.capture_launch(self.graph)
         else:
             self.simulate()

-        # self.solver.update_data()  # implicit
-        # pendulum_acc = self.state_0.body_qdd.numpy()[0]
-        # self.plot_window.add_point(pendulum_acc[2])
-        # print(f"Pendulum acceleration: {pendulum_acc}")
         imu_acc = self.imu.sensor_qdd.numpy()[0]
         self.plot_window.add_point(imu_acc[2])
-        print(f"IMU acceleration: {imu_acc}")
         self.sim_time += self.frame_dt

156-160: Misleading plot label.

The hardcoded text "Flap contact force" doesn't match the actual data being plotted (IMU acceleration). Consider using the self.title property or updating the text.

         if imgui.begin(self.title, flags=flags):
-            imgui.text("Flap contact force")
+            imgui.text(self.title)
             imgui.plot_lines("Force", self.data, **self.plot_kwargs)
         imgui.end()
📜 Review details

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 225ea4f and 2c7f58a.

📒 Files selected for processing (8)
  • newton/_src/sensors/imu_sensor.py (1 hunks)
  • newton/_src/sim/builder.py (4 hunks)
  • newton/_src/sim/model.py (5 hunks)
  • newton/_src/sim/state.py (2 hunks)
  • newton/_src/solvers/mujoco/kernels.py (1 hunks)
  • newton/_src/solvers/mujoco/solver_mujoco.py (5 hunks)
  • newton/examples/example_solver_data.py (1 hunks)
  • newton/sensors.py (2 hunks)
🧰 Additional context used
🧠 Learnings (25)
📚 Learning: 2025-11-28T11:12:40.805Z
Learnt from: camevor
Repo: newton-physics/newton PR: 1146
File: newton/examples/basic/example_basic_joints.py:206-213
Timestamp: 2025-11-28T11:12:40.805Z
Learning: In Newton's standard solvers (XPBD, SemiImplicit, SolverMuJoCo), spatial vectors in State.body_qd use the ordering (linear, angular): wp.spatial_top(qd) returns linear velocity and wp.spatial_bottom(qd) returns angular velocity. This is the opposite of typical Modern Robotics spatial twist notation where (angular, linear) is used.

Applied to files:

  • newton/_src/sim/state.py
  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-09-22T21:08:31.901Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/examples/ik/example_ik_franka.py:121-123
Timestamp: 2025-09-22T21:08:31.901Z
Learning: In the newton physics framework, when creating warp arrays for IK solver joint variables using wp.array(self.model.joint_q, shape=(1, coord_count)), the resulting array acts as a reference/pointer to the original model's joint coordinates, so updates from the IK solver automatically reflect in the model's joint_q buffer used for rendering.

Applied to files:

  • newton/_src/sim/state.py
  • newton/_src/solvers/mujoco/solver_mujoco.py
  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-09-22T21:08:31.901Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/examples/ik/example_ik_franka.py:121-123
Timestamp: 2025-09-22T21:08:31.901Z
Learning: In the newton physics framework, when creating warp arrays for IK solver joint variables using wp.array(self.model.joint_q, shape=(1, coord_count)), the resulting array acts as a reference/pointer to the original model's joint coordinates, so updates from the IK solver automatically reflect in the model's joint_q buffer used for rendering, despite the general warp documentation suggesting copies are made by default.

Applied to files:

  • newton/_src/sim/state.py
  • newton/_src/solvers/mujoco/solver_mujoco.py
  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-08-19T12:27:30.629Z
Learnt from: preist-nvidia
Repo: newton-physics/newton PR: 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/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-08-19T12:27:30.629Z
Learnt from: preist-nvidia
Repo: newton-physics/newton PR: 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/_src/solvers/mujoco/solver_mujoco.py
  • newton/examples/example_solver_data.py
📚 Learning: 2025-11-24T08:05:21.390Z
Learnt from: adenzler-nvidia
Repo: newton-physics/newton PR: 1107
File: newton/_src/solvers/mujoco/kernels.py:973-974
Timestamp: 2025-11-24T08:05:21.390Z
Learning: In Newton's MuJoCo solver integration (newton/_src/solvers/mujoco/), the mapping between Newton joints and MuJoCo joints is not 1-to-1. Instead, each Newton DOF maps to a distinct MuJoCo joint. This means that for multi-DOF joints (like D6 with 6 DOFs), there will be 6 corresponding MuJoCo joints, each with its own properties (margin, solimp, solref, etc.). The mapping is done via dof_to_mjc_joint array, ensuring each DOF's properties are written to its own MuJoCo joint without overwriting.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-09-09T12:50:42.929Z
Learnt from: gdaviet
Repo: newton-physics/newton PR: 750
File: newton/_src/solvers/implicit_mpm/render_grains.py:142-160
Timestamp: 2025-09-09T12:50:42.929Z
Learning: The `update_render_grains` function in newton/_src/solvers/implicit_mpm/render_grains.py should return early when `state.velocity_field is None` because its purpose is specifically to advect grains in the velocity field. Without a velocity field, the function should be a no-op.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-09-22T21:03:39.624Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_lbfgs_optimizer.py:739-752
Timestamp: 2025-09-22T21:03:39.624Z
Learning: The L-BFGS optimizer in newton/_src/sim/ik/ik_lbfgs_optimizer.py currently intentionally only supports additive updates (assuming n_coords == n_dofs). Velocity space integration for joints with mismatched coordinate/DOF dimensions (like free/ball joints) is planned for future work and should not be flagged as an issue in current reviews.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-09-26T05:58:21.013Z
Learnt from: WenchaoHuang
Repo: newton-physics/newton PR: 835
File: newton/_src/solvers/style3d/collision/collision.py:49-53
Timestamp: 2025-09-26T05:58:21.013Z
Learning: In newton/_src/solvers/style3d/collision/collision.py, the broad_phase buffers (broad_phase_ee, broad_phase_ef, broad_phase_vf) created with wp.array() don't need zero-initialization because they follow a write-first-read-later pattern: BVH query kernels in frame_begin() populate these buffers completely (write-only operations), and collision kernels only read from them afterward in accumulate_contact_force(), ensuring no undefined reads occur.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-09-18T07:05:56.836Z
Learnt from: gdaviet
Repo: newton-physics/newton PR: 750
File: newton/_src/solvers/implicit_mpm/solve_rheology.py:969-986
Timestamp: 2025-09-18T07:05:56.836Z
Learning: In newton/_src/solvers/implicit_mpm/solve_rheology.py, transposed_strain_mat parameter cannot be None - the type signature was corrected to reflect this guarantee, eliminating the need for None checks when accessing transposed_strain_mat.offsets.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-08-18T15:56:26.587Z
Learnt from: adenzler-nvidia
Repo: newton-physics/newton PR: 552
File: newton/_src/solvers/mujoco/solver_mujoco.py:0-0
Timestamp: 2025-08-18T15:56:26.587Z
Learning: In Newton's MuJoCo solver, when transforming joint axes from Newton's internal frame to MuJoCo's expected frame, use wp.quat_rotate(joint_rot, axis) not wp.quat_rotate_inv(joint_rot, axis). The joint_rot represents rotation from joint-local to body frame, so forward rotation is correct.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-09-26T05:58:21.013Z
Learnt from: WenchaoHuang
Repo: newton-physics/newton PR: 835
File: newton/_src/solvers/style3d/collision/collision.py:49-53
Timestamp: 2025-09-26T05:58:21.013Z
Learning: In newton/_src/solvers/style3d/collision/collision.py, the broad_phase buffers (broad_phase_ee, broad_phase_ef, broad_phase_vf) created with wp.array() don't need zero-initialization because the BVH query kernels that use them always write the count to query_results[0, tid] = 0 first before populating any data, ensuring no undefined reads occur.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-09-22T21:03:18.367Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_lbfgs_optimizer.py:638-648
Timestamp: 2025-09-22T21:03:18.367Z
Learning: The L-BFGS optimizer in newton/_src/sim/ik/ik_lbfgs_optimizer.py currently only supports additive updates, meaning it assumes n_coords == n_dofs. Velocity space integration for joints with mismatched coordinate/DOF dimensions (like free/ball joints) is planned for future work.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-08-12T05:17:34.423Z
Learnt from: shi-eric
Repo: newton-physics/newton PR: 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.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-08-12T18:04:06.577Z
Learnt from: nvlukasz
Repo: newton-physics/newton PR: 519
File: newton/_src/solvers/featherstone/kernels.py:75-75
Timestamp: 2025-08-12T18:04:06.577Z
Learning: The Newton physics framework requires nightly Warp builds, which means compatibility concerns with older stable Warp versions (like missing functions such as wp.spatial_adjoint) are not relevant for this project.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-08-25T21:41:45.795Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 634
File: newton/tests/test_examples.py:329-333
Timestamp: 2025-08-25T21:41:45.795Z
Learning: Newton examples use a centralized CLI argument parsing system in newton/examples/__init__.py. The create_parser() function defines common arguments like --device, --viewer, --output-path, and --num-frames, while init(parser) creates viewers based on parsed arguments. Individual example scripts don't need to define these flags themselves - they inherit them from the centralized system via the example_map and runpy execution.

Applied to files:

  • newton/examples/example_solver_data.py
📚 Learning: 2025-08-27T23:32:48.670Z
Learnt from: shi-eric
Repo: newton-physics/newton PR: 660
File: pyproject.toml:54-59
Timestamp: 2025-08-27T23:32:48.670Z
Learning: In the newton project's pyproject.toml, self-referential extras (e.g., `examples = ["newton[sim]", "newton[importers]", ...]`) are normal and encouraged for dependency management, rather than inlining the dependencies.

Applied to files:

  • newton/examples/example_solver_data.py
📚 Learning: 2025-08-14T17:38:36.106Z
Learnt from: nvlukasz
Repo: newton-physics/newton PR: 519
File: newton/__init__.py:25-29
Timestamp: 2025-08-14T17:38:36.106Z
Learning: The Newton project prefers incremental __all__ building using __all__ += [...] pattern to group exports with their related imports, rather than a single consolidated __all__ at the end of the file.

Applied to files:

  • newton/sensors.py
📚 Learning: 2025-08-12T17:58:00.815Z
Learnt from: nvlukasz
Repo: newton-physics/newton PR: 519
File: newton/_src/solvers/featherstone/kernels.py:18-20
Timestamp: 2025-08-12T17:58:00.815Z
Learning: In newton/_src/core/__init__.py, the function transform_twist is properly re-exported from the spatial module via "from .spatial import (..., transform_twist, ...)" and included in __all__, making it available for import as "from ...core import transform_twist" from other modules.

Applied to files:

  • newton/sensors.py
📚 Learning: 2025-08-12T17:51:37.474Z
Learnt from: nvlukasz
Repo: newton-physics/newton PR: 519
File: newton/geometry.py:16-22
Timestamp: 2025-08-12T17:51:37.474Z
Learning: In the Newton public API refactor, common geometry symbols like ParticleFlags and ShapeFlags are now exposed at the top level in newton/__init__.py rather than through newton.geometry. The newton/geometry.py module is intentionally focused only on broad-phase collision detection classes (BroadPhaseAllPairs, BroadPhaseExplicit, BroadPhaseSAP).

Applied to files:

  • newton/sensors.py
📚 Learning: 2025-12-01T17:00:28.889Z
Learnt from: camevor
Repo: newton-physics/newton PR: 1161
File: newton/_src/solvers/mujoco/kernels.py:1268-1268
Timestamp: 2025-12-01T17:00:28.889Z
Learning: MuJoCo Warp spatial vectors (e.g., cacc, cvel) use the convention (angular, linear), which is opposite to Newton's (linear, angular) convention. When converting from MuJoCo Warp to Newton format, spatial_top() on MuJoCo Warp data extracts the angular component, and spatial_bottom() extracts the linear component.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-11-28T11:12:40.030Z
Learnt from: camevor
Repo: newton-physics/newton PR: 1146
File: newton/examples/basic/example_basic_joints.py:215-222
Timestamp: 2025-11-28T11:12:40.030Z
Learning: In Warp's spatial vector convention used by Newton: spatial_bottom() extracts the angular velocity component (omega) and spatial_top() extracts the linear velocity component (v) from a spatial twist (qd). Similarly for spatial wrenches: spatial_bottom() is torque and spatial_top() is force.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-12-01T16:21:36.581Z
Learnt from: jvonmuralt
Repo: newton-physics/newton PR: 1160
File: newton/_src/sim/articulation.py:0-0
Timestamp: 2025-12-01T16:21:36.581Z
Learning: In Newton's Featherstone solver, during forward kinematics for FREE/DISTANCE joints, the computed spatial velocity v_wc follows Featherstone's spatial twist convention where the linear component (spatial_top) represents the velocity at the world origin (0,0,0), not at the joint anchor or COM. To convert to State.body_qd (which stores COM velocity), use v_com = v_origin + ω × x_com, where x_com is the absolute world position of the COM.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-09-25T16:14:22.033Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_objectives.py:0-0
Timestamp: 2025-09-25T16:14:22.033Z
Learning: wp.transform in NVIDIA Warp supports indexing using numerical indices (e.g., body_tf[0], body_tf[3], etc.) to access translation and rotation components. The indexing approach used in the Newton codebase is valid and compiles correctly.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-09-25T16:14:22.033Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_objectives.py:0-0
Timestamp: 2025-09-25T16:14:22.033Z
Learning: In NVIDIA Warp's Newton physics library, wp.transform supports direct numerical indexing (e.g., body_tf[0], body_tf[1], body_tf[2] for position and body_tf[3], body_tf[4], body_tf[5], body_tf[6] for quaternion components) to access its elements. This is the standard API used throughout the codebase.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
🧬 Code graph analysis (5)
newton/_src/solvers/mujoco/solver_mujoco.py (3)
newton/_src/solvers/mujoco/kernels.py (1)
  • convert_rigid_forces_from_mj_kernel (1314-1348)
newton/_src/sim/state.py (1)
  • State (21-151)
newton/_src/sim/model.py (1)
  • state (530-583)
newton/_src/sim/model.py (1)
newton/_src/sim/state.py (4)
  • State (21-151)
  • requires_grad (117-123)
  • body_count (126-130)
  • particle_count (133-137)
newton/sensors.py (1)
newton/_src/sensors/imu_sensor.py (1)
  • IMUSensor (59-93)
newton/_src/sim/builder.py (1)
newton/_src/sim/model.py (1)
  • request_state_attributes (711-718)
newton/_src/sensors/imu_sensor.py (2)
newton/_src/sim/model.py (3)
  • Model (96-872)
  • state (530-583)
  • request_state_attributes (711-718)
newton/_src/sim/state.py (1)
  • State (21-151)
🪛 Ruff (0.14.8)
newton/_src/sim/model.py

849-849: Undefined name List

(F821)

🔇 Additional comments (8)
newton/_src/sim/state.py (1)

49-51: New body_qdd / body_parent_f fields are consistent with existing State design

Adding body_qdd and body_parent_f as optional spatial-vector arrays on State fits the existing pattern for body fields and will be handled automatically by assign() and the derived-count properties. No changes needed here.

Also applies to: 61-63

newton/_src/sensors/imu_sensor.py (1)

24-57: Fix spatial_vector construction argument order

The spatial_vector at line 56 has arguments in the wrong order. Warp's spatial_vector(angular, linear) expects angular components first, but the code passes (acc_lin, vel_ang) which is (linear, angular).

Change line 56 from:

imu_sensor[sensor_idx] = wp.spatial_vector(wp.quat_rotate_inv(q, acc_lin), wp.quat_rotate_inv(q, vel_ang))

To:

imu_sensor[sensor_idx] = wp.spatial_vector(wp.quat_rotate_inv(q, vel_ang), wp.quat_rotate_inv(q, acc_lin))

This ensures the output follows Warp's standard spatial vector layout where indices [0-2] contain angular components and [3-5] contain linear components. The example code currently accesses index [2] by coincidence rather than design; with this fix, linear acceleration would be at indices [3-5] as expected.

⛔ Skipped due to learnings
Learnt from: camevor
Repo: newton-physics/newton PR: 1161
File: newton/_src/solvers/mujoco/kernels.py:1268-1268
Timestamp: 2025-12-01T17:00:28.889Z
Learning: MuJoCo Warp spatial vectors (e.g., cacc, cvel) use the convention (angular, linear), which is opposite to Newton's (linear, angular) convention. When converting from MuJoCo Warp to Newton format, spatial_top() on MuJoCo Warp data extracts the angular component, and spatial_bottom() extracts the linear component.
Learnt from: camevor
Repo: newton-physics/newton PR: 1146
File: newton/examples/basic/example_basic_joints.py:206-213
Timestamp: 2025-11-28T11:12:40.805Z
Learning: In Newton's standard solvers (XPBD, SemiImplicit, SolverMuJoCo), spatial vectors in State.body_qd use the ordering (linear, angular): wp.spatial_top(qd) returns linear velocity and wp.spatial_bottom(qd) returns angular velocity. This is the opposite of typical Modern Robotics spatial twist notation where (angular, linear) is used.
Learnt from: adenzler-nvidia
Repo: newton-physics/newton PR: 552
File: newton/_src/solvers/mujoco/solver_mujoco.py:0-0
Timestamp: 2025-08-18T15:56:26.587Z
Learning: In Newton's MuJoCo solver, when transforming joint axes from Newton's internal frame to MuJoCo's expected frame, use wp.quat_rotate(joint_rot, axis) not wp.quat_rotate_inv(joint_rot, axis). The joint_rot represents rotation from joint-local to body frame, so forward rotation is correct.
Learnt from: daniela-hase
Repo: newton-physics/newton PR: 1044
File: newton/_src/sensors/tiled_camera_sensor.py:451-466
Timestamp: 2025-11-12T23:53:29.478Z
Learning: In newton/_src/sensors/tiled_camera_sensor.py, the convert_camera_to_warp_arrays method extracts [:3, :3] from camera.get_view_matrix().reshape(4, 4) without transposing. This is correct because Pyglet's Mat4.look_at returns a column-major view matrix, and reshaping it as row-major in NumPy creates an implicit transpose, converting the world→camera rotation into the camera→world rotation that the renderer expects.
Learnt from: camevor
Repo: newton-physics/newton PR: 1146
File: newton/examples/basic/example_basic_joints.py:215-222
Timestamp: 2025-11-28T11:12:40.030Z
Learning: In Warp's spatial vector convention used by Newton: spatial_bottom() extracts the angular velocity component (omega) and spatial_top() extracts the linear velocity component (v) from a spatial twist (qd). Similarly for spatial wrenches: spatial_bottom() is torque and spatial_top() is force.
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_objectives.py:0-0
Timestamp: 2025-09-25T16:14:22.033Z
Learning: wp.transform in NVIDIA Warp supports indexing using numerical indices (e.g., body_tf[0], body_tf[3], etc.) to access translation and rotation components. The indexing approach used in the Newton codebase is valid and compiles correctly.
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_objectives.py:0-0
Timestamp: 2025-09-25T16:14:22.033Z
Learning: In NVIDIA Warp's Newton physics library, wp.transform supports direct numerical indexing (e.g., body_tf[0], body_tf[1], body_tf[2] for position and body_tf[3], body_tf[4], body_tf[5], body_tf[6] for quaternion components) to access its elements. This is the standard API used throughout the codebase.
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/examples/ik/example_ik_franka.py:121-123
Timestamp: 2025-09-22T21:08:31.901Z
Learning: In the newton physics framework, when creating warp arrays for IK solver joint variables using wp.array(self.model.joint_q, shape=(1, coord_count)), the resulting array acts as a reference/pointer to the original model's joint coordinates, so updates from the IK solver automatically reflect in the model's joint_q buffer used for rendering, despite the general warp documentation suggesting copies are made by default.
newton/_src/sim/builder.py (2)

519-520: Per‑builder tracking of requested state attributes looks correct

Storing requested state attributes as a set[str] on the builder is appropriate and keeps the data structure simple and deduplicated. No issues here.


5388-5389: Forwarding requested attributes from builder to model at finalize is in the right place

Calling m.request_state_attributes(*self._requested_state_attributes) immediately after constructing the Model ensures the model’s internal requested‑attributes set is populated before any State objects are created. This cleanly connects the builder‑side API to the model‑side implementation without affecting existing users who never call request_state_attributes. Looks good.

newton/_src/solvers/mujoco/kernels.py (2)

1291-1310: LGTM! Spatial vector convention handling is correct.

Based on learnings, MuJoCo Warp uses (angular, linear) convention where spatial_top() extracts angular and spatial_bottom() extracts linear. This function correctly applies that convention to compute linear body acceleration.


1313-1348: LGTM! Convention conversion is correct.

The kernel correctly converts from MuJoCo Warp's (angular, linear) format to Newton's (linear, angular) format at line 1344. The body_parent_f TODO is acknowledged for this prototype.

newton/examples/example_solver_data.py (2)

110-117: LGTM!

The render() method follows the standard example pattern. The empty test() is acknowledged in the PR description as pending work before marking ready.


163-170: LGTM!

Standard Newton example entry point following the centralized CLI argument parsing system.

newton/_src/sim/model.py Outdated Show resolved Hide resolved
Copy link
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

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

Actionable comments posted: 1

♻️ Duplicate comments (2)
newton/_src/sensors/imu_sensor.py (1)

16-17: Update module docstring to describe IMU sensor functionality

The module docstring still says “Frame Transform Sensor – measures transforms relative to sites.” but this file now implements an IMU sensor. Please update it to reflect the IMU behavior, e.g. “IMU sensor – computes linear acceleration and angular velocity at reference sites.”

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

467-467: State-attribute gating design looks good; fix List type and guard body_qdd/body_parent_f allocations

  • The combination of _requested_state_attributes, get_requested_state_attributes(), and the gated allocations in state() nicely preserves backward compatibility: when the model has particles/bodies/joints, you still allocate the core solver-facing fields (particle_q/qd/f, body_q/qd/f, joint_q/qd) even if nothing was explicitly requested, and extended fields like body_qdd/body_parent_f are opt-in via request_state_attributes(). This addresses the earlier regression where default State objects lacked core arrays. ✔

  • In get_requested_state_attributes() at line 849, the return annotation uses List[str] but List is never imported, which Ruff correctly flags (F821). Given you already rely on Python 3.9+ generics, switch to the built-in list[str]:

  • def get_requested_state_attributes(self) -> List[str]:

  • def get_requested_state_attributes(self) -> list[str]:

- `_add_custom_attributes()` now skips all attributes with `assignment == ModelAttributeAssignment.STATE` (line 740). This means any custom state attributes registered via `add_attribute(..., assignment=ModelAttributeAssignment.STATE, ...)` will no longer be attached to `State` objects. If this is intentional (all state attributes must now be explicitly materialized via `request_state_attributes()`), document this behavior change. If not, consider a more selective skip that only affects attributes not in `get_requested_state_attributes()`.

- **Critical robustness issue at lines 574–578**: `body_qdd` and `body_parent_f` allocations call `wp.zeros_like(self.body_qd)` without any guard. Unlike `body_q`/`body_qd`/`body_f` (lines 559–565), they are outside the `if self.body_count:` block. If someone calls `request_state_attributes("body_qdd")` on a model with `body_count == 0`, `self.body_qd` is `None` and the allocation will fail. Move these two blocks inside the `if self.body_count:` guard.

</blockquote></details>

</blockquote></details>

<details>
<summary>🧹 Nitpick comments (4)</summary><blockquote>

<details>
<summary>newton/sensors.py (1)</summary><blockquote>

`28-31`: **Expose `SensorIMU` via `newton.sensors` – LGTM; clean up duplicate comment**

Importing `SensorIMU` and adding it to `__all__` correctly surfaces the new IMU sensor at the top level. One small nit: Line 33 introduces a second `# Raycast sensors` comment directly above the existing one on Line 34; you can drop one of them to avoid duplication.





Also applies to: 33-50

</blockquote></details>
<details>
<summary>newton/_src/sensors/imu_sensor.py (2)</summary><blockquote>

`59-71`: **Document `SensorIMU`’s state-attribute requirement in the class docstring**

`SensorIMU.__init__` correctly requests `"body_qdd"` on the model and sets up buffers. To make the usage contract clearer, consider expanding the class docstring to mention that callers must construct the sensor before creating `State` instances so that `body_qdd` is allocated, and that the active solver must populate `body_qdd` for meaningful measurements.

---

`72-79`: **Long exception message in `update()` – optional style tweak**

The explicit `body_qdd` `None` check is good and gives a clear error. Ruff’s TRY003 hint here is purely stylistic: if you want to appease it, you could shorten the inline message or assign it to a constant/variable before raising, but there’s no functional issue.

Example:

```diff
-        if state.body_qdd is None:
-            raise ValueError("SensorIMU requires a State with body_qdd allocated. Create SensorIMU before State.")
+        if state.body_qdd is None:
+            msg = (
+                "SensorIMU requires a State with 'body_qdd' allocated. "
+                "Create SensorIMU before calling Model.state()."
+            )
+            raise ValueError(msg)

Also applies to: 81-95

newton/examples/example_solver_data.py (1)

95-109: Align plot/UI labels with IMU acceleration and wire up example test later

  • In step(), you’re plotting imu_acc[2] and printing "IMU acceleration", but the plot window is constructed as "Parent force" and the UI text says "Flap contact force" / "Force". This looks like copy-paste from a force example and will be confusing. Consider renaming the window and labels to something like "IMU acceleration (z)" and/or using self.title consistently:

  •    self.plot_window = ViewerPlot(viewer, "Parent force", scale_min=0, graph_size=(400, 200))
    
  •    self.plot_window = ViewerPlot(viewer, "IMU acceleration (z)", scale_min=0, graph_size=(400, 200))
    
    ...
  •        imgui.text("Flap contact force")
    
  •        imgui.plot_lines("Force", self.data, **self.plot_kwargs)
    
  •        imgui.text("IMU acceleration (z)")
    
  •        imgui.plot_lines("Acceleration", self.data, **self.plot_kwargs)
    
    
    
  • Example.test() is currently a no-op. For when this moves out of draft and into test_examples.py, you’ll likely want it to run a short headless simulation (e.g., a few step() calls) to catch basic wiring errors (solver, IMU, viewer logging) in CI.

Also applies to: 120-160

📜 Review details

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 2c7f58a and 09217b3.

📒 Files selected for processing (4)
  • newton/_src/sensors/imu_sensor.py (1 hunks)
  • newton/_src/sim/model.py (5 hunks)
  • newton/examples/example_solver_data.py (1 hunks)
  • newton/sensors.py (2 hunks)
🧰 Additional context used
🧠 Learnings (7)
📓 Common learnings
Learnt from: shi-eric
Repo: newton-physics/newton PR: 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.
📚 Learning: 2025-08-14T17:38:36.106Z
Learnt from: nvlukasz
Repo: newton-physics/newton PR: 519
File: newton/__init__.py:25-29
Timestamp: 2025-08-14T17:38:36.106Z
Learning: The Newton project prefers incremental __all__ building using __all__ += [...] pattern to group exports with their related imports, rather than a single consolidated __all__ at the end of the file.

Applied to files:

  • newton/sensors.py
📚 Learning: 2025-08-12T17:58:00.815Z
Learnt from: nvlukasz
Repo: newton-physics/newton PR: 519
File: newton/_src/solvers/featherstone/kernels.py:18-20
Timestamp: 2025-08-12T17:58:00.815Z
Learning: In newton/_src/core/__init__.py, the function transform_twist is properly re-exported from the spatial module via "from .spatial import (..., transform_twist, ...)" and included in __all__, making it available for import as "from ...core import transform_twist" from other modules.

Applied to files:

  • newton/sensors.py
📚 Learning: 2025-08-12T17:51:37.474Z
Learnt from: nvlukasz
Repo: newton-physics/newton PR: 519
File: newton/geometry.py:16-22
Timestamp: 2025-08-12T17:51:37.474Z
Learning: In the Newton public API refactor, common geometry symbols like ParticleFlags and ShapeFlags are now exposed at the top level in newton/__init__.py rather than through newton.geometry. The newton/geometry.py module is intentionally focused only on broad-phase collision detection classes (BroadPhaseAllPairs, BroadPhaseExplicit, BroadPhaseSAP).

Applied to files:

  • newton/sensors.py
📚 Learning: 2025-08-19T12:27:30.629Z
Learnt from: preist-nvidia
Repo: newton-physics/newton PR: 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/examples/example_solver_data.py
📚 Learning: 2025-08-25T21:41:45.795Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 634
File: newton/tests/test_examples.py:329-333
Timestamp: 2025-08-25T21:41:45.795Z
Learning: Newton examples use a centralized CLI argument parsing system in newton/examples/__init__.py. The create_parser() function defines common arguments like --device, --viewer, --output-path, and --num-frames, while init(parser) creates viewers based on parsed arguments. Individual example scripts don't need to define these flags themselves - they inherit them from the centralized system via the example_map and runpy execution.

Applied to files:

  • newton/examples/example_solver_data.py
📚 Learning: 2025-08-19T12:27:30.629Z
Learnt from: preist-nvidia
Repo: newton-physics/newton PR: 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/examples/example_solver_data.py
🧬 Code graph analysis (3)
newton/sensors.py (1)
newton/_src/sensors/imu_sensor.py (1)
  • SensorIMU (59-95)
newton/_src/sensors/imu_sensor.py (2)
newton/_src/sim/model.py (3)
  • Model (96-872)
  • state (530-583)
  • request_state_attributes (711-718)
newton/_src/sim/state.py (1)
  • State (21-151)
newton/examples/example_solver_data.py (4)
newton/_src/viewer/viewer_gl.py (1)
  • register_ui_callback (143-161)
newton/_src/sensors/imu_sensor.py (2)
  • SensorIMU (59-95)
  • update (72-95)
newton/_src/viewer/viewer.py (1)
  • log_contacts (272-335)
newton/examples/__init__.py (1)
  • init (344-387)
🪛 Ruff (0.14.8)
newton/_src/sensors/imu_sensor.py

79-79: Avoid specifying long messages outside the exception class

(TRY003)

newton/_src/sim/model.py

849-849: Undefined name List

(F821)

🔇 Additional comments (1)
newton/_src/sensors/imu_sensor.py (1)

24-57: IMU Warp kernel shape and math look consistent

The kernel arguments, indexing (sensor_idx bounds check), and use of spatial top/bottom components to compute point acceleration and rotate into the sensor frame are consistent with the Model/State layout and expected IMU output. No issues from a correctness or safety standpoint in this kernel.

newton/examples/sensors/example_sensor_imu.py Show resolved Hide resolved
- Adapt terms: require->request; fields->attributes
- Allocate custom attributes only when requested
- Remove requested fields from SolverMuJoCo
- Remove custom attributes from state field filtering
- Rename to SensorIMU
- Use accelerometer and gyroscope attributes
- Fix duplicate requested attributes
- Compute kinematic body_qdd
- Add sensor validation, documentation
- Add example
@camevor camevor force-pushed the extended-state-attributes branch from 09217b3 to c8bfcf0 Compare December 16, 2025 17:20
@codecov
Copy link

codecov bot commented Dec 16, 2025

Codecov Report

❌ Patch coverage is 93.42105% with 5 lines in your changes missing coverage. Please review.
✅ All tests successful. No failed tests found.

Files with missing lines Patch % Lines
newton/_src/sensors/sensor_imu.py 94.28% 2 Missing ⚠️
newton/_src/sim/builder.py 80.00% 1 Missing ⚠️
newton/_src/sim/model.py 94.44% 1 Missing ⚠️
newton/_src/solvers/mujoco/solver_mujoco.py 92.30% 1 Missing ⚠️

📢 Thoughts on this report? Let us know!

@camevor camevor marked this pull request as ready for review December 16, 2025 17:51
@camevor camevor changed the title Prototype for extended state attributes and IMU Sensor Extended state attributes and IMU Sensor Dec 16, 2025
Copy link
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

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

Actionable comments posted: 3

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)
newton/_src/solvers/mujoco/solver_mujoco.py (1)

45-55: Gate rigid force/acceleration field updates on the mjwarp backend to prevent stale data in CPU mode

The new enable_rne_postconstraint() + convert_rigid_forces_from_mj_kernel path works correctly for the mjwarp backend: enable_rne_postconstraint() is called before stepping (line 479), enabling RNE computation when state_out.body_qdd or state_out.body_parent_f is allocated, and then the extended update block populates these fields from self.mjw_model/self.mjw_data.

However, in update_newton_state() the extended update (lines 799–820) runs whenever state.body_qdd or state.body_parent_f is not None, regardless of backend:

# Update rigid force fields on state.
if state.body_qdd is not None or state.body_parent_f is not None:
    nbody = self.mjc_body_to_newton.shape[1]
    wp.launch(
        convert_rigid_forces_from_mj_kernel,
        (nworld, nbody),
        inputs=[
            self.mjc_body_to_newton,
            self.mjw_model.body_rootid,
            self.mjw_model.opt.gravity,
            self.mjw_data.xipos,
            self.mjw_data.subtree_com,
            self.mjw_data.cacc,
            self.mjw_data.cvel,
        ],
        outputs=[state.body_qdd, state.body_parent_f],
    )

On the CPU path (use_mujoco_cpu=True), is_mjwarp is false and the step uses self.mj_data, but this block still reads from self.mjw_model/self.mjw_data, which are never stepped in that mode. If a caller requests body_qdd while using the CPU backend, they'll silently get accelerations derived from stale mjwarp buffers.

Gate this block on is_mjwarp:

-        # Update rigid force fields on state.
-        if state.body_qdd is not None or state.body_parent_f is not None:
+        # Update rigid acceleration/force fields only for the mjwarp backend.
+        if is_mjwarp and (state.body_qdd is not None or state.body_parent_f is not None):
             nbody = self.mjc_body_to_newton.shape[1]
             wp.launch(
                 convert_rigid_forces_from_mj_kernel,
                 (nworld, nbody),
                 inputs=[
                     self.mjc_body_to_newton,
                     self.mjw_model.body_rootid,
                     self.mjw_model.opt.gravity,
                     self.mjw_data.xipos,
                     self.mjw_data.subtree_com,
                     self.mjw_data.cacc,
                     self.mjw_data.cvel,
                 ],
                 outputs=[state.body_qdd, state.body_parent_f],
             )

This keeps semantics clear: extended IMU-style outputs are only supported on the mjwarp path, and CPU mode won't read inconsistent mjwarp state.

🧹 Nitpick comments (6)
newton/sensors.py (1)

28-37: SensorIMU export wiring is correct; drop duplicated Raycast comment

The new SensorIMU import and inclusion in __all__ are consistent with the implementation module and make the IMU sensor available via newton.sensors as intended.

Minor nit: Lines 28–29 both read # Raycast sensors; you can safely remove one of them to avoid redundant commentary.

Also applies to: 42-48

docs/concepts/sensors.rst (1)

164-201: SensorIMU documentation is aligned with implementation; tiny optional enhancement

The SensorIMU section correctly documents:

  • per-site IMU frames,
  • accelerometer/gyroscope outputs,
  • dependency on the extended state attribute body_qdd,
  • the requirement to construct the sensor before calling model.state() when relying on its auto-request behavior.

Optionally, you could add a short note to the basic usage snippet that in a real simulation loop you need to run a solver that populates state.body_qdd (e.g., SolverMuJoCo.step(...)) before imu.update(state) to get non‑gravity accelerations, but this is a documentation nicety only.

Also applies to: 202-209

docs/concepts/extended_state_attributes.rst (1)

1-46: Extended State Attributes docs match the new API and sensor behavior

This page cleanly explains:

  • which attributes are “core” vs “extended,”
  • how to request them on both ModelBuilder and Model,
  • their idempotent behavior,
  • the current set (body_qdd, body_parent_f) and solver support,
  • and how components like SensorIMU transparently request what they need.

Everything is consistent with the new ModelBuilder.request_state_attributes and Model.request_state_attributes APIs.

Optionally, you could use :class:~newton.sensors.SensorIMU`` for shorter link text, but that’s purely cosmetic.

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

528-529: Builder-level state attribute requests are wired correctly through composition and finalize

The new _requested_state_attributes plumbing looks coherent:

  • __init__ initializes a per‑builder set, and request_state_attributes() simply updates it, so repeated calls are idempotent.
  • add_builder() unions builder._requested_state_attributes into the parent, which means imports, add_world(), and replicate() all correctly aggregate extended attribute requests from sub‑builders.
  • finalize() forwards the accumulated set into Model.request_state_attributes(...) before any State is created, so builder‑level requests take effect as intended.

This matches the behavior documented in extended_state_attributes.rst and lines up with the existing Model.request_state_attributes API.

If you’d like, you could add a short note to the ModelBuilder class docstring pointing users at request_state_attributes() and the new concepts page, but the current implementation itself looks good.

Also applies to: 583-584, 5286-5294, 5505-5506

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

466-467: Extended state-attribute request plumbing looks coherent; consider a small defensive guard

The _requested_state_attributes set, request_state_attributes(), and get_requested_state_attributes() together give a clear opt‑in path for extended fields while keeping core particle/body/joint arrays always allocated, so existing solver patterns remain valid. The conditional allocation of body_qdd/body_parent_f in state() based on the requested set is also consistent with the new IMU/sensor use cases.

If you want to harden this further against odd usages (e.g. someone calling request_state_attributes("body_qdd") on a model with body_count == 0), you could mirror the core-body guards and additionally check self.body_count before allocating the extended arrays; otherwise this looks fine as‑is for the intended workflows.

Also applies to: 529-571, 703-711, 838-861

newton/examples/sensors/example_sensor_imu.py (1)

24-46: Tie color buffers and kernel launch dimension to the IMU sensor count

The example currently hard‑codes a single sensor:

  • self.buffer = wp.zeros(1, ...) / self.colors = wp.zeros(1, ...)
  • wp.launch(acc_to_color, dim=1, ...)

This is fine for the current [imu_site] usage, but it will silently misbehave if someone later passes multiple sites into SensorIMU.

You can make this robust at essentially no cost by sizing the buffers and launch dimension from self.imu.n_sensors, e.g.:

-        self.buffer = wp.zeros(1, dtype=wp.vec3)
-        self.colors = wp.zeros(1, dtype=wp.vec3)
+        self.buffer = wp.zeros(self.imu.n_sensors, dtype=wp.vec3)
+        self.colors = wp.zeros(self.imu.n_sensors, dtype=wp.vec3)
...
-            wp.launch(acc_to_color, dim=1, inputs=[0.01, self.imu.accelerometer, self.buffer, self.colors])
+            wp.launch(
+                acc_to_color,
+                dim=self.imu.n_sensors,
+                inputs=[0.01, self.imu.accelerometer, self.buffer, self.colors],
+            )

The rest of the wiring (swap before IMU update, use of self.colors.numpy()[0] for the cube) already matches the intended single‑sensor visualization.

Also applies to: 80-105, 124-141

📜 Review details

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 09217b3 and 141a447.

📒 Files selected for processing (13)
  • docs/api/newton_sensors.rst (1 hunks)
  • docs/concepts/extended_state_attributes.rst (1 hunks)
  • docs/concepts/sensors.rst (2 hunks)
  • docs/index.rst (1 hunks)
  • newton/_src/sensors/sensor_imu.py (1 hunks)
  • newton/_src/sim/builder.py (4 hunks)
  • newton/_src/sim/model.py (5 hunks)
  • newton/_src/sim/state.py (2 hunks)
  • newton/_src/solvers/mujoco/kernels.py (2 hunks)
  • newton/_src/solvers/mujoco/solver_mujoco.py (5 hunks)
  • newton/examples/sensors/example_sensor_imu.py (1 hunks)
  • newton/sensors.py (2 hunks)
  • newton/tests/test_sensor_imu.py (1 hunks)
✅ Files skipped from review due to trivial changes (1)
  • docs/index.rst
🚧 Files skipped from review as they are similar to previous changes (1)
  • newton/_src/sim/state.py
🧰 Additional context used
🧠 Learnings (23)
📓 Common learnings
Learnt from: shi-eric
Repo: newton-physics/newton PR: 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.
📚 Learning: 2025-08-12T05:17:34.423Z
Learnt from: shi-eric
Repo: newton-physics/newton PR: 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.

Applied to files:

  • docs/concepts/extended_state_attributes.rst
  • docs/concepts/sensors.rst
  • newton/_src/solvers/mujoco/kernels.py
  • newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-12-12T08:45:43.428Z
Learnt from: nvtw
Repo: newton-physics/newton PR: 1221
File: newton/examples/example_sdf.py:277-287
Timestamp: 2025-12-12T08:45:43.428Z
Learning: In Newton examples (e.g., example_sdf.py), computing contacts once per frame and reusing them across multiple substeps is an intentional design choice, not a bug. The contacts are computed before the substep loop and intentionally kept constant throughout all substeps within that frame.

Applied to files:

  • docs/concepts/sensors.rst
📚 Learning: 2025-09-22T21:08:31.901Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/examples/ik/example_ik_franka.py:121-123
Timestamp: 2025-09-22T21:08:31.901Z
Learning: In the newton physics framework, when creating warp arrays for IK solver joint variables using wp.array(self.model.joint_q, shape=(1, coord_count)), the resulting array acts as a reference/pointer to the original model's joint coordinates, so updates from the IK solver automatically reflect in the model's joint_q buffer used for rendering, despite the general warp documentation suggesting copies are made by default.

Applied to files:

  • newton/_src/sim/model.py
  • newton/_src/solvers/mujoco/kernels.py
  • newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-08-12T17:58:00.815Z
Learnt from: nvlukasz
Repo: newton-physics/newton PR: 519
File: newton/_src/solvers/featherstone/kernels.py:18-20
Timestamp: 2025-08-12T17:58:00.815Z
Learning: In newton/_src/core/__init__.py, the function transform_twist is properly re-exported from the spatial module via "from .spatial import (..., transform_twist, ...)" and included in __all__, making it available for import as "from ...core import transform_twist" from other modules.

Applied to files:

  • newton/sensors.py
📚 Learning: 2025-11-28T11:12:40.805Z
Learnt from: camevor
Repo: newton-physics/newton PR: 1146
File: newton/examples/basic/example_basic_joints.py:206-213
Timestamp: 2025-11-28T11:12:40.805Z
Learning: In Newton's standard solvers (XPBD, SemiImplicit, SolverMuJoCo), spatial vectors in State.body_qd use the ordering (linear, angular): wp.spatial_top(qd) returns linear velocity and wp.spatial_bottom(qd) returns angular velocity. This is the opposite of typical Modern Robotics spatial twist notation where (angular, linear) is used.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-12-01T17:00:28.889Z
Learnt from: camevor
Repo: newton-physics/newton PR: 1161
File: newton/_src/solvers/mujoco/kernels.py:1268-1268
Timestamp: 2025-12-01T17:00:28.889Z
Learning: MuJoCo Warp spatial vectors (e.g., cacc, cvel) use the convention (angular, linear), which is opposite to Newton's (linear, angular) convention. When converting from MuJoCo Warp to Newton format, spatial_top() on MuJoCo Warp data extracts the angular component, and spatial_bottom() extracts the linear component.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-08-18T15:56:26.587Z
Learnt from: adenzler-nvidia
Repo: newton-physics/newton PR: 552
File: newton/_src/solvers/mujoco/solver_mujoco.py:0-0
Timestamp: 2025-08-18T15:56:26.587Z
Learning: In Newton's MuJoCo solver, when transforming joint axes from Newton's internal frame to MuJoCo's expected frame, use wp.quat_rotate(joint_rot, axis) not wp.quat_rotate_inv(joint_rot, axis). The joint_rot represents rotation from joint-local to body frame, so forward rotation is correct.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
  • newton/examples/sensors/example_sensor_imu.py
  • newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-11-28T11:12:40.030Z
Learnt from: camevor
Repo: newton-physics/newton PR: 1146
File: newton/examples/basic/example_basic_joints.py:215-222
Timestamp: 2025-11-28T11:12:40.030Z
Learning: In Warp's spatial vector convention used by Newton: spatial_bottom() extracts the angular velocity component (omega) and spatial_top() extracts the linear velocity component (v) from a spatial twist (qd). Similarly for spatial wrenches: spatial_bottom() is torque and spatial_top() is force.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-12-01T16:21:36.581Z
Learnt from: jvonmuralt
Repo: newton-physics/newton PR: 1160
File: newton/_src/sim/articulation.py:0-0
Timestamp: 2025-12-01T16:21:36.581Z
Learning: In Newton's Featherstone solver, during forward kinematics for FREE/DISTANCE joints, the computed spatial velocity v_wc follows Featherstone's spatial twist convention where the linear component (spatial_top) represents the velocity at the world origin (0,0,0), not at the joint anchor or COM. To convert to State.body_qd (which stores COM velocity), use v_com = v_origin + ω × x_com, where x_com is the absolute world position of the COM.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-09-22T21:03:39.624Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_lbfgs_optimizer.py:739-752
Timestamp: 2025-09-22T21:03:39.624Z
Learning: The L-BFGS optimizer in newton/_src/sim/ik/ik_lbfgs_optimizer.py currently intentionally only supports additive updates (assuming n_coords == n_dofs). Velocity space integration for joints with mismatched coordinate/DOF dimensions (like free/ball joints) is planned for future work and should not be flagged as an issue in current reviews.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
  • newton/examples/sensors/example_sensor_imu.py
  • newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-12-12T16:25:51.464Z
Learnt from: vastsoun
Repo: newton-physics/newton PR: 1019
File: newton/_src/solvers/kamino/core/math.py:248-255
Timestamp: 2025-12-12T16:25:51.464Z
Learning: In newton/_src/solvers/kamino/core/math.py, the quat_conj function's current implementation (negating w instead of x,y,z) is a known issue that will be revised in a future PR per vastsoun's decision, not to be flagged in current Kamino reviews.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-09-25T16:14:22.033Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_objectives.py:0-0
Timestamp: 2025-09-25T16:14:22.033Z
Learning: wp.transform in NVIDIA Warp supports indexing using numerical indices (e.g., body_tf[0], body_tf[3], etc.) to access translation and rotation components. The indexing approach used in the Newton codebase is valid and compiles correctly.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-09-25T16:14:22.033Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_objectives.py:0-0
Timestamp: 2025-09-25T16:14:22.033Z
Learning: In NVIDIA Warp's Newton physics library, wp.transform supports direct numerical indexing (e.g., body_tf[0], body_tf[1], body_tf[2] for position and body_tf[3], body_tf[4], body_tf[5], body_tf[6] for quaternion components) to access its elements. This is the standard API used throughout the codebase.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-09-22T21:08:31.901Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/examples/ik/example_ik_franka.py:121-123
Timestamp: 2025-09-22T21:08:31.901Z
Learning: In the newton physics framework, when creating warp arrays for IK solver joint variables using wp.array(self.model.joint_q, shape=(1, coord_count)), the resulting array acts as a reference/pointer to the original model's joint coordinates, so updates from the IK solver automatically reflect in the model's joint_q buffer used for rendering.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
  • newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-11-24T08:05:21.390Z
Learnt from: adenzler-nvidia
Repo: newton-physics/newton PR: 1107
File: newton/_src/solvers/mujoco/kernels.py:973-974
Timestamp: 2025-11-24T08:05:21.390Z
Learning: In Newton's MuJoCo solver integration (newton/_src/solvers/mujoco/), the mapping between Newton joints and MuJoCo joints is not 1-to-1. Instead, each Newton DOF maps to a distinct MuJoCo joint. This means that for multi-DOF joints (like D6 with 6 DOFs), there will be 6 corresponding MuJoCo joints, each with its own properties (margin, solimp, solref, etc.). The mapping is done via dof_to_mjc_joint array, ensuring each DOF's properties are written to its own MuJoCo joint without overwriting.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
  • newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-08-12T18:04:06.577Z
Learnt from: nvlukasz
Repo: newton-physics/newton PR: 519
File: newton/_src/solvers/featherstone/kernels.py:75-75
Timestamp: 2025-08-12T18:04:06.577Z
Learning: The Newton physics framework requires nightly Warp builds, which means compatibility concerns with older stable Warp versions (like missing functions such as wp.spatial_adjoint) are not relevant for this project.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
  • newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-08-19T12:27:30.629Z
Learnt from: preist-nvidia
Repo: newton-physics/newton PR: 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/examples/sensors/example_sensor_imu.py
  • newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-12-12T08:45:43.428Z
Learnt from: nvtw
Repo: newton-physics/newton PR: 1221
File: newton/examples/example_sdf.py:277-287
Timestamp: 2025-12-12T08:45:43.428Z
Learning: In Newtown (Newton) example code, specifically files under newton/examples, computing contacts once per frame and reusing them across all substeps is an intentional design choice, not a bug. Reviewers should verify that contacts are computed before the substep loop and reused for every substep within the same frame. This pattern reduces redundant work and preserves frame-consistency; do not flag as a regression unless the behavior is changed for correctness or performance reasons.

Applied to files:

  • newton/examples/sensors/example_sensor_imu.py
📚 Learning: 2025-08-19T12:27:30.629Z
Learnt from: preist-nvidia
Repo: newton-physics/newton PR: 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/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-09-09T12:50:42.929Z
Learnt from: gdaviet
Repo: newton-physics/newton PR: 750
File: newton/_src/solvers/implicit_mpm/render_grains.py:142-160
Timestamp: 2025-09-09T12:50:42.929Z
Learning: The `update_render_grains` function in newton/_src/solvers/implicit_mpm/render_grains.py should return early when `state.velocity_field is None` because its purpose is specifically to advect grains in the velocity field. Without a velocity field, the function should be a no-op.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-09-18T07:05:56.836Z
Learnt from: gdaviet
Repo: newton-physics/newton PR: 750
File: newton/_src/solvers/implicit_mpm/solve_rheology.py:969-986
Timestamp: 2025-09-18T07:05:56.836Z
Learning: In newton/_src/solvers/implicit_mpm/solve_rheology.py, transposed_strain_mat parameter cannot be None - the type signature was corrected to reflect this guarantee, eliminating the need for None checks when accessing transposed_strain_mat.offsets.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-09-26T05:58:21.013Z
Learnt from: WenchaoHuang
Repo: newton-physics/newton PR: 835
File: newton/_src/solvers/style3d/collision/collision.py:49-53
Timestamp: 2025-09-26T05:58:21.013Z
Learning: In newton/_src/solvers/style3d/collision/collision.py, the broad_phase buffers (broad_phase_ee, broad_phase_ef, broad_phase_vf) created with wp.array() don't need zero-initialization because they follow a write-first-read-later pattern: BVH query kernels in frame_begin() populate these buffers completely (write-only operations), and collision kernels only read from them afterward in accumulate_contact_force(), ensuring no undefined reads occur.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
🧬 Code graph analysis (5)
newton/_src/sim/builder.py (1)
newton/_src/sim/model.py (1)
  • request_state_attributes (703-710)
newton/_src/sim/model.py (3)
newton/_src/sim/state.py (3)
  • requires_grad (121-127)
  • particle_count (137-141)
  • body_count (130-134)
newton/_src/sim/builder.py (3)
  • request_state_attributes (5286-5293)
  • particle_count (964-968)
  • body_count (950-954)
newton/_src/sensors/sensor_imu.py (1)
  • update (153-177)
newton/sensors.py (1)
newton/_src/sensors/sensor_imu.py (1)
  • SensorIMU (74-177)
newton/_src/sensors/sensor_imu.py (4)
newton/_src/geometry/flags.py (1)
  • ShapeFlags (30-45)
newton/_src/sim/model.py (3)
  • Model (98-861)
  • state (529-575)
  • request_state_attributes (703-710)
newton/_src/sim/state.py (1)
  • State (21-155)
newton/_src/sim/builder.py (1)
  • shape_count (943-947)
newton/_src/solvers/mujoco/solver_mujoco.py (3)
newton/_src/solvers/mujoco/kernels.py (1)
  • convert_rigid_forces_from_mj_kernel (1335-1370)
newton/_src/sim/state.py (1)
  • State (21-155)
newton/_src/sim/model.py (1)
  • state (529-575)
🪛 Ruff (0.14.8)
newton/_src/sensors/sensor_imu.py

126-126: Avoid specifying long messages outside the exception class

(TRY003)


149-149: Avoid specifying long messages outside the exception class

(TRY003)


151-151: Avoid specifying long messages outside the exception class

(TRY003)


160-160: Avoid specifying long messages outside the exception class

(TRY003)

⏰ 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). (2)
  • GitHub Check: Run GPU Benchmarks / Run GPU Benchmarks on AWS EC2
  • GitHub Check: Run GPU Tests / Run GPU Unit Tests on AWS EC2
🔇 Additional comments (2)
docs/api/newton_sensors.rst (1)

8-16: Expose SensorIMU in public sensor API autosummary

Adding SensorIMU here matches the new public export in newton.sensors and keeps the API docs in sync with the code. Looks good.

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

1291-1311: RNE conversion helper and kernel correctly bridge MuJoCo Warp → Newton spatial conventions

mj_body_acceleration() and convert_rigid_forces_from_mj_kernel() look consistent with the documented conventions:

  • For MuJoCo Warp, you treat spatial_top as angular and spatial_bottom as linear, then reconstruct COM linear acceleration with the expected offset and Coriolis terms.
  • For Newton, you build body_qdd as (linear, angular) via wp.spatial_vector(lin + mjw_gravity[0], wp.spatial_top(cacc)), which matches the State.body_qdd semantics.

The optional handling of body_qdd / body_parent_f via simple truthiness checks is also aligned with existing kernels in this module.

Also applies to: 1334-1370

docs/concepts/sensors.rst Show resolved Hide resolved
newton/_src/sensors/sensor_imu.py Show resolved Hide resolved
newton/tests/test_sensor_imu.py Show resolved Hide resolved
Copy link
Contributor

@coderabbitai coderabbitai bot left a comment

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 (1)
newton/_src/solvers/mujoco/solver_mujoco.py (1)

799-817: Verify mjwarp-only access is guarded upstream.

This block accesses self.mjw_model and self.mjw_data without explicit backend checks. While update_newton_state is called from both CPU and mjwarp paths, the CPU path would have self.mjw_model = None when mjw_model/mjw_data aren't provided.

The condition gate prevents the kernel launch when state attributes are None, but if someone later adds CPU-mode support for extended attributes, this could fail. Consider adding a guard or documenting that extended attributes are mjwarp-only.

         # Update rigid force fields on state.
-        if state.body_qdd is not None or state.body_parent_f is not None:
+        if (state.body_qdd is not None or state.body_parent_f is not None) and not self.use_mujoco_cpu:
             # Launch over MuJoCo bodies
             nbody = self.mjc_body_to_newton.shape[1]
📜 Review details

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 09217b3 and 141a447.

📒 Files selected for processing (13)
  • docs/api/newton_sensors.rst (1 hunks)
  • docs/concepts/extended_state_attributes.rst (1 hunks)
  • docs/concepts/sensors.rst (2 hunks)
  • docs/index.rst (1 hunks)
  • newton/_src/sensors/sensor_imu.py (1 hunks)
  • newton/_src/sim/builder.py (4 hunks)
  • newton/_src/sim/model.py (5 hunks)
  • newton/_src/sim/state.py (2 hunks)
  • newton/_src/solvers/mujoco/kernels.py (2 hunks)
  • newton/_src/solvers/mujoco/solver_mujoco.py (5 hunks)
  • newton/examples/sensors/example_sensor_imu.py (1 hunks)
  • newton/sensors.py (2 hunks)
  • newton/tests/test_sensor_imu.py (1 hunks)
✅ Files skipped from review due to trivial changes (1)
  • docs/concepts/extended_state_attributes.rst
🚧 Files skipped from review as they are similar to previous changes (1)
  • newton/_src/sim/builder.py
🧰 Additional context used
🧠 Learnings (25)
📓 Common learnings
Learnt from: shi-eric
Repo: newton-physics/newton PR: 521
File: newton/examples/example_cloth_hanging.py:36-36
Timestamp: 2025-08-12T05:17:34.423Z
Learning: The Newton migration guide (docs/migration.rst) is specifically for documenting how to migrate existing warp.sim functionality to Newton equivalents. New Newton-only features that didn't exist in warp.sim do not need migration documentation.
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_lbfgs_optimizer.py:739-752
Timestamp: 2025-09-22T21:03:39.624Z
Learning: The L-BFGS optimizer in newton/_src/sim/ik/ik_lbfgs_optimizer.py currently intentionally only supports additive updates (assuming n_coords == n_dofs). Velocity space integration for joints with mismatched coordinate/DOF dimensions (like free/ball joints) is planned for future work and should not be flagged as an issue in current reviews.
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_lbfgs_optimizer.py:638-648
Timestamp: 2025-09-22T21:03:18.367Z
Learning: The L-BFGS optimizer in newton/_src/sim/ik/ik_lbfgs_optimizer.py currently only supports additive updates, meaning it assumes n_coords == n_dofs. Velocity space integration for joints with mismatched coordinate/DOF dimensions (like free/ball joints) is planned for future work.
📚 Learning: 2025-08-12T17:58:00.815Z
Learnt from: nvlukasz
Repo: newton-physics/newton PR: 519
File: newton/_src/solvers/featherstone/kernels.py:18-20
Timestamp: 2025-08-12T17:58:00.815Z
Learning: In newton/_src/core/__init__.py, the function transform_twist is properly re-exported from the spatial module via "from .spatial import (..., transform_twist, ...)" and included in __all__, making it available for import as "from ...core import transform_twist" from other modules.

Applied to files:

  • newton/sensors.py
📚 Learning: 2025-08-14T17:38:36.106Z
Learnt from: nvlukasz
Repo: newton-physics/newton PR: 519
File: newton/__init__.py:25-29
Timestamp: 2025-08-14T17:38:36.106Z
Learning: The Newton project prefers incremental __all__ building using __all__ += [...] pattern to group exports with their related imports, rather than a single consolidated __all__ at the end of the file.

Applied to files:

  • newton/sensors.py
📚 Learning: 2025-08-12T17:51:37.474Z
Learnt from: nvlukasz
Repo: newton-physics/newton PR: 519
File: newton/geometry.py:16-22
Timestamp: 2025-08-12T17:51:37.474Z
Learning: In the Newton public API refactor, common geometry symbols like ParticleFlags and ShapeFlags are now exposed at the top level in newton/__init__.py rather than through newton.geometry. The newton/geometry.py module is intentionally focused only on broad-phase collision detection classes (BroadPhaseAllPairs, BroadPhaseExplicit, BroadPhaseSAP).

Applied to files:

  • newton/sensors.py
📚 Learning: 2025-09-22T21:08:31.901Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/examples/ik/example_ik_franka.py:121-123
Timestamp: 2025-09-22T21:08:31.901Z
Learning: In the newton physics framework, when creating warp arrays for IK solver joint variables using wp.array(self.model.joint_q, shape=(1, coord_count)), the resulting array acts as a reference/pointer to the original model's joint coordinates, so updates from the IK solver automatically reflect in the model's joint_q buffer used for rendering, despite the general warp documentation suggesting copies are made by default.

Applied to files:

  • newton/_src/sim/model.py
  • newton/_src/sim/state.py
  • newton/_src/solvers/mujoco/solver_mujoco.py
  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-11-28T11:12:40.805Z
Learnt from: camevor
Repo: newton-physics/newton PR: 1146
File: newton/examples/basic/example_basic_joints.py:206-213
Timestamp: 2025-11-28T11:12:40.805Z
Learning: In Newton's standard solvers (XPBD, SemiImplicit, SolverMuJoCo), spatial vectors in State.body_qd use the ordering (linear, angular): wp.spatial_top(qd) returns linear velocity and wp.spatial_bottom(qd) returns angular velocity. This is the opposite of typical Modern Robotics spatial twist notation where (angular, linear) is used.

Applied to files:

  • newton/_src/sim/state.py
  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-09-22T21:08:31.901Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/examples/ik/example_ik_franka.py:121-123
Timestamp: 2025-09-22T21:08:31.901Z
Learning: In the newton physics framework, when creating warp arrays for IK solver joint variables using wp.array(self.model.joint_q, shape=(1, coord_count)), the resulting array acts as a reference/pointer to the original model's joint coordinates, so updates from the IK solver automatically reflect in the model's joint_q buffer used for rendering.

Applied to files:

  • newton/_src/sim/state.py
  • newton/_src/solvers/mujoco/solver_mujoco.py
  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-12-12T08:45:43.428Z
Learnt from: nvtw
Repo: newton-physics/newton PR: 1221
File: newton/examples/example_sdf.py:277-287
Timestamp: 2025-12-12T08:45:43.428Z
Learning: In Newton examples (e.g., example_sdf.py), computing contacts once per frame and reusing them across multiple substeps is an intentional design choice, not a bug. The contacts are computed before the substep loop and intentionally kept constant throughout all substeps within that frame.

Applied to files:

  • docs/concepts/sensors.rst
📚 Learning: 2025-08-12T05:17:34.423Z
Learnt from: shi-eric
Repo: newton-physics/newton PR: 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.

Applied to files:

  • docs/concepts/sensors.rst
  • newton/_src/solvers/mujoco/solver_mujoco.py
  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-09-22T21:03:39.624Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_lbfgs_optimizer.py:739-752
Timestamp: 2025-09-22T21:03:39.624Z
Learning: The L-BFGS optimizer in newton/_src/sim/ik/ik_lbfgs_optimizer.py currently intentionally only supports additive updates (assuming n_coords == n_dofs). Velocity space integration for joints with mismatched coordinate/DOF dimensions (like free/ball joints) is planned for future work and should not be flagged as an issue in current reviews.

Applied to files:

  • newton/examples/sensors/example_sensor_imu.py
  • newton/_src/solvers/mujoco/solver_mujoco.py
  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-08-18T15:56:26.587Z
Learnt from: adenzler-nvidia
Repo: newton-physics/newton PR: 552
File: newton/_src/solvers/mujoco/solver_mujoco.py:0-0
Timestamp: 2025-08-18T15:56:26.587Z
Learning: In Newton's MuJoCo solver, when transforming joint axes from Newton's internal frame to MuJoCo's expected frame, use wp.quat_rotate(joint_rot, axis) not wp.quat_rotate_inv(joint_rot, axis). The joint_rot represents rotation from joint-local to body frame, so forward rotation is correct.

Applied to files:

  • newton/examples/sensors/example_sensor_imu.py
  • newton/_src/solvers/mujoco/solver_mujoco.py
  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-12-12T08:45:43.428Z
Learnt from: nvtw
Repo: newton-physics/newton PR: 1221
File: newton/examples/example_sdf.py:277-287
Timestamp: 2025-12-12T08:45:43.428Z
Learning: In Newtown (Newton) example code, specifically files under newton/examples, computing contacts once per frame and reusing them across all substeps is an intentional design choice, not a bug. Reviewers should verify that contacts are computed before the substep loop and reused for every substep within the same frame. This pattern reduces redundant work and preserves frame-consistency; do not flag as a regression unless the behavior is changed for correctness or performance reasons.

Applied to files:

  • newton/examples/sensors/example_sensor_imu.py
📚 Learning: 2025-08-19T12:27:30.629Z
Learnt from: preist-nvidia
Repo: newton-physics/newton PR: 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/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-11-24T08:05:21.390Z
Learnt from: adenzler-nvidia
Repo: newton-physics/newton PR: 1107
File: newton/_src/solvers/mujoco/kernels.py:973-974
Timestamp: 2025-11-24T08:05:21.390Z
Learning: In Newton's MuJoCo solver integration (newton/_src/solvers/mujoco/), the mapping between Newton joints and MuJoCo joints is not 1-to-1. Instead, each Newton DOF maps to a distinct MuJoCo joint. This means that for multi-DOF joints (like D6 with 6 DOFs), there will be 6 corresponding MuJoCo joints, each with its own properties (margin, solimp, solref, etc.). The mapping is done via dof_to_mjc_joint array, ensuring each DOF's properties are written to its own MuJoCo joint without overwriting.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-08-19T12:27:30.629Z
Learnt from: preist-nvidia
Repo: newton-physics/newton PR: 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/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-09-09T12:50:42.929Z
Learnt from: gdaviet
Repo: newton-physics/newton PR: 750
File: newton/_src/solvers/implicit_mpm/render_grains.py:142-160
Timestamp: 2025-09-09T12:50:42.929Z
Learning: The `update_render_grains` function in newton/_src/solvers/implicit_mpm/render_grains.py should return early when `state.velocity_field is None` because its purpose is specifically to advect grains in the velocity field. Without a velocity field, the function should be a no-op.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-09-18T07:05:56.836Z
Learnt from: gdaviet
Repo: newton-physics/newton PR: 750
File: newton/_src/solvers/implicit_mpm/solve_rheology.py:969-986
Timestamp: 2025-09-18T07:05:56.836Z
Learning: In newton/_src/solvers/implicit_mpm/solve_rheology.py, transposed_strain_mat parameter cannot be None - the type signature was corrected to reflect this guarantee, eliminating the need for None checks when accessing transposed_strain_mat.offsets.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-09-26T05:58:21.013Z
Learnt from: WenchaoHuang
Repo: newton-physics/newton PR: 835
File: newton/_src/solvers/style3d/collision/collision.py:49-53
Timestamp: 2025-09-26T05:58:21.013Z
Learning: In newton/_src/solvers/style3d/collision/collision.py, the broad_phase buffers (broad_phase_ee, broad_phase_ef, broad_phase_vf) created with wp.array() don't need zero-initialization because they follow a write-first-read-later pattern: BVH query kernels in frame_begin() populate these buffers completely (write-only operations), and collision kernels only read from them afterward in accumulate_contact_force(), ensuring no undefined reads occur.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-08-12T18:04:06.577Z
Learnt from: nvlukasz
Repo: newton-physics/newton PR: 519
File: newton/_src/solvers/featherstone/kernels.py:75-75
Timestamp: 2025-08-12T18:04:06.577Z
Learning: The Newton physics framework requires nightly Warp builds, which means compatibility concerns with older stable Warp versions (like missing functions such as wp.spatial_adjoint) are not relevant for this project.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-12-01T17:00:28.889Z
Learnt from: camevor
Repo: newton-physics/newton PR: 1161
File: newton/_src/solvers/mujoco/kernels.py:1268-1268
Timestamp: 2025-12-01T17:00:28.889Z
Learning: MuJoCo Warp spatial vectors (e.g., cacc, cvel) use the convention (angular, linear), which is opposite to Newton's (linear, angular) convention. When converting from MuJoCo Warp to Newton format, spatial_top() on MuJoCo Warp data extracts the angular component, and spatial_bottom() extracts the linear component.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-11-28T11:12:40.030Z
Learnt from: camevor
Repo: newton-physics/newton PR: 1146
File: newton/examples/basic/example_basic_joints.py:215-222
Timestamp: 2025-11-28T11:12:40.030Z
Learning: In Warp's spatial vector convention used by Newton: spatial_bottom() extracts the angular velocity component (omega) and spatial_top() extracts the linear velocity component (v) from a spatial twist (qd). Similarly for spatial wrenches: spatial_bottom() is torque and spatial_top() is force.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-12-01T16:21:36.581Z
Learnt from: jvonmuralt
Repo: newton-physics/newton PR: 1160
File: newton/_src/sim/articulation.py:0-0
Timestamp: 2025-12-01T16:21:36.581Z
Learning: In Newton's Featherstone solver, during forward kinematics for FREE/DISTANCE joints, the computed spatial velocity v_wc follows Featherstone's spatial twist convention where the linear component (spatial_top) represents the velocity at the world origin (0,0,0), not at the joint anchor or COM. To convert to State.body_qd (which stores COM velocity), use v_com = v_origin + ω × x_com, where x_com is the absolute world position of the COM.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-12-12T16:25:51.464Z
Learnt from: vastsoun
Repo: newton-physics/newton PR: 1019
File: newton/_src/solvers/kamino/core/math.py:248-255
Timestamp: 2025-12-12T16:25:51.464Z
Learning: In newton/_src/solvers/kamino/core/math.py, the quat_conj function's current implementation (negating w instead of x,y,z) is a known issue that will be revised in a future PR per vastsoun's decision, not to be flagged in current Kamino reviews.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-09-25T16:14:22.033Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_objectives.py:0-0
Timestamp: 2025-09-25T16:14:22.033Z
Learning: wp.transform in NVIDIA Warp supports indexing using numerical indices (e.g., body_tf[0], body_tf[3], etc.) to access translation and rotation components. The indexing approach used in the Newton codebase is valid and compiles correctly.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-09-25T16:14:22.033Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_objectives.py:0-0
Timestamp: 2025-09-25T16:14:22.033Z
Learning: In NVIDIA Warp's Newton physics library, wp.transform supports direct numerical indexing (e.g., body_tf[0], body_tf[1], body_tf[2] for position and body_tf[3], body_tf[4], body_tf[5], body_tf[6] for quaternion components) to access its elements. This is the standard API used throughout the codebase.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
🧬 Code graph analysis (5)
newton/sensors.py (1)
newton/_src/sensors/sensor_imu.py (1)
  • SensorIMU (74-177)
newton/_src/sim/model.py (2)
newton/_src/sim/state.py (3)
  • requires_grad (121-127)
  • particle_count (137-141)
  • body_count (130-134)
newton/_src/sim/builder.py (3)
  • request_state_attributes (5286-5293)
  • particle_count (964-968)
  • body_count (950-954)
newton/tests/test_sensor_imu.py (1)
newton/_src/sensors/sensor_imu.py (2)
  • SensorIMU (74-177)
  • update (153-177)
newton/_src/sensors/sensor_imu.py (3)
newton/_src/geometry/flags.py (1)
  • ShapeFlags (30-45)
newton/_src/sim/model.py (3)
  • Model (98-861)
  • state (529-575)
  • request_state_attributes (703-710)
newton/_src/sim/state.py (1)
  • State (21-155)
newton/examples/sensors/example_sensor_imu.py (2)
newton/_src/utils/import_utils.py (1)
  • transform (139-140)
newton/_src/sensors/sensor_imu.py (2)
  • SensorIMU (74-177)
  • update (153-177)
🪛 Ruff (0.14.8)
newton/_src/sensors/sensor_imu.py

126-126: Avoid specifying long messages outside the exception class

(TRY003)


149-149: Avoid specifying long messages outside the exception class

(TRY003)


151-151: Avoid specifying long messages outside the exception class

(TRY003)


160-160: Avoid specifying long messages outside the exception class

(TRY003)

⏰ 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). (1)
  • GitHub Check: Run GPU Benchmarks / Run GPU Benchmarks on AWS EC2
🔇 Additional comments (25)
docs/index.rst (1)

32-32: LGTM!

The new documentation entry is correctly placed in the Concepts section toctree alongside related topics.

docs/api/newton_sensors.rst (1)

16-16: LGTM!

SensorIMU is correctly added to the API documentation class listing.

newton/sensors.py (2)

34-37: LGTM!

SensorIMU import follows the established pattern for sensor imports.


47-47: LGTM!

SensorIMU correctly added to __all__ exports, maintaining alphabetical order.

newton/_src/sim/state.py (2)

52-53: LGTM!

The body_qdd attribute for rigid body accelerations is correctly declared with appropriate type annotation and docstring.


65-66: LGTM!

The body_parent_f attribute for parent interaction forces follows the same pattern as existing optional state arrays.

docs/concepts/sensors.rst (2)

164-201: LGTM!

The SensorIMU documentation is comprehensive, covering outputs, usage patterns, and the important ordering requirement (create sensor before state). The example code correctly demonstrates the intended usage.


202-217: LGTM!

The State/Solver requirements section clearly documents the dependency on body_qdd and the need for a compatible solver like SolverMuJoCo. The See Also references are appropriately updated.

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

54-54: LGTM!

The new kernel import follows the existing pattern for MuJoCo solver kernels.


479-479: LGTM!

Enabling RNE postconstraint before stepping ensures body accelerations are computed when extended state attributes are requested.


495-506: LGTM!

The enable_rne_postconstraint method correctly:

  • Returns early if already enabled (avoiding redundant work)
  • Checks for presence of extended state attributes before enabling
  • Includes helpful verbose logging

The TODO for CPU mode is acceptable for a prototype.

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

466-467: LGTM!

The _requested_state_attributes set provides a clean mechanism for tracking extended attribute requests without impacting existing behavior.


542-544: LGTM!

Using get_requested_state_attributes() ensures backward compatibility by always including core state fields while allowing opt-in extended attributes.


566-570: LGTM!

Extended state attributes body_qdd and body_parent_f are correctly allocated only when explicitly requested, matching the prototype design.


703-710: LGTM!

The request_state_attributes method provides a simple API for sensors and users to declare their state requirements.


838-860: LGTM!

The get_requested_state_attributes method correctly:

  • Builds the default attribute list based on model configuration (particle/body/joint counts)
  • Appends explicitly requested attributes without duplicates using set difference
  • Returns a predictable list for state allocation
newton/examples/sensors/example_sensor_imu.py (4)

24-45: LGTM!

The acc_to_color kernel correctly implements exponential smoothing with clamping. The color mapping provides visual feedback for acceleration magnitude.


48-114: LGTM!

The Example initialization correctly:

  • Creates the sensor before calling model.state() (ensuring body_qdd allocation)
  • Sets up the model with appropriate shapes and an IMU site
  • Initializes the SolverMuJoCo which supports body_qdd computation

116-140: LGTM!

The simulation logic correctly:

  • Uses CUDA graph capture for performance when available
  • Places IMU update after the state swap (reading from latest state)
  • Computes contacts per-substep which is the standard Newton examples pattern

142-168: LGTM!

The step/render/main structure follows the standard Newton examples pattern. The visual color update provides real-time feedback of IMU readings.

newton/_src/solvers/mujoco/kernels.py (2)

1291-1310: LGTM - Body acceleration computation is correct.

The function correctly handles MuJoCo Warp spatial vector convention (angular, linear) when extracting velocity and acceleration components, and properly computes the frame-dependent linear acceleration with centripetal correction.


1334-1370: LGTM - Correct spatial vector conversion from MuJoCo Warp to Newton.

The kernel properly converts MuJoCo Warp's (angular, linear) convention to Newton's (linear, angular) convention. The use of wp.spatial_top(cacc) correctly extracts the angular component from MuJoCo Warp data, and the gravity correction is appropriately applied to the linear component.

Note: body_parent_f implementation is deferred (line 1368-1370), which is documented as a prototype limitation.

newton/tests/test_sensor_imu.py (1)

28-184: LGTM - Comprehensive test coverage for SensorIMU.

The test suite thoroughly covers:

  • Creation and multi-site scenarios
  • Validation of site indices and types
  • Error handling for missing body_qdd (correct order enforcement)
  • Functional behavior with gravity and rotations
  • Edge cases (world-frame sites, rotated site transforms)

The tests correctly verify the sensor's requirement that body_qdd must be requested before state creation.

newton/_src/sensors/sensor_imu.py (2)

25-72: LGTM - Correct IMU kinematics with proper spatial vector handling.

The kernel correctly:

  • Handles world-frame sites (lines 50-53) with gravity-only acceleration
  • Extracts velocity/acceleration components using Newton's (linear, angular) convention
  • Computes proper rigid-body kinematics at offset r with centripetal terms (line 66)
  • Transforms results into sensor frame (lines 69-71)

74-177: LGTM - Well-structured sensor class with proper validation.

The implementation correctly:

  • Requests body_qdd attribute during initialization (line 130)
  • Validates site indices and flags (lines 144-152)
  • Checks for body_qdd availability before use (line 159-160)
  • Launches computation kernel with proper inputs (lines 162-177)

Static analysis hints about long error messages (TRY003) can be safely ignored—the messages provide clear, user-facing context appropriate for validation errors in a prototype.

newton/sensors.py Outdated Show resolved Hide resolved
Copy link
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

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

Actionable comments posted: 0

🧹 Nitpick comments (1)
newton/_src/sensors/sensor_imu.py (1)

139-143: Offer to implement the TODO: body per site logging.

The verbose initialization logging is helpful. The TODO at line 142 indicates that per-site body information could be added to the output.

Would you like me to generate code that logs which body each site is attached to? This could be implemented by iterating over the sites and reading model.shape_body[site_idx].

📜 Review details

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 141a447 and cd42647.

📒 Files selected for processing (5)
  • docs/concepts/extended_state_attributes.rst (1 hunks)
  • docs/concepts/sensors.rst (2 hunks)
  • newton/_src/sensors/sensor_imu.py (1 hunks)
  • newton/_src/sim/state.py (1 hunks)
  • newton/sensors.py (2 hunks)
🚧 Files skipped from review as they are similar to previous changes (4)
  • newton/sensors.py
  • docs/concepts/extended_state_attributes.rst
  • docs/concepts/sensors.rst
  • newton/_src/sim/state.py
🧰 Additional context used
🧬 Code graph analysis (1)
newton/_src/sensors/sensor_imu.py (3)
newton/_src/geometry/flags.py (1)
  • ShapeFlags (30-45)
newton/_src/sim/model.py (3)
  • Model (98-861)
  • state (529-575)
  • request_state_attributes (703-710)
newton/_src/sim/state.py (1)
  • State (21-154)
🪛 Ruff (0.14.8)
newton/_src/sensors/sensor_imu.py

126-126: Avoid specifying long messages outside the exception class

(TRY003)


149-149: Avoid specifying long messages outside the exception class

(TRY003)


151-151: Avoid specifying long messages outside the exception class

(TRY003)


160-160: Avoid specifying long messages outside the exception class

(TRY003)

⏰ 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). (2)
  • GitHub Check: Run GPU Benchmarks / Run GPU Benchmarks on AWS EC2
  • GitHub Check: Run GPU Tests / Run GPU Unit Tests on AWS EC2
🔇 Additional comments (4)
newton/_src/sensors/sensor_imu.py (4)

25-72: Solid rigid-body kinematics implementation.

The kernel correctly handles both world-fixed sites (using gravity transformation alone) and body-attached sites (full rigid-body acceleration with centrifugal and Euler terms). The physics formula a_com - g + α×r + ω×(ω×r) and sensor-frame transformation are accurate.


74-102: Well-documented sensor interface.

The class docstring clearly explains the extended state attribute requirement (body_qdd), and the example correctly shows the sensor.update(state) signature. The attribute documentation is concise and helpful.


144-152: Thorough site validation.

The validation correctly checks both bounds and the SITE flag, preventing invalid sensor configurations with clear error messages.


153-177: Proper runtime validation and kernel invocation.

The update method correctly validates that body_qdd is allocated before launching the kernel, providing a helpful error message that guides users to instantiate the sensor before creating the state. All required inputs are properly passed to the kernel.

Copy link
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

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

Actionable comments posted: 0

🧹 Nitpick comments (2)
newton/_src/sim/builder.py (1)

5596-5604: New request_state_attributes API is sound; consider optional introspection helper

The API shape and behavior (simple set update) are consistent with Model.request_state_attributes in newton/_src/sim/model.py and provide a clear entry point for importers/sensors to declare needs. If users are expected to inspect what’s been requested on the builder (similar to Model.get_requested_state_attributes()), you might optionally add a small get_requested_state_attributes() accessor here for symmetry.

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

1402-1404: TODO for body_parent_f aligns with documented prototype limitations.

The placeholder for implementing body_parent_f computation is consistent with the PR objectives, which document that this is a prototype with incomplete features.

Would you like me to help implement the body_parent_f computation, or is this intentionally deferred to a future PR?

📜 Review details

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between cd42647 and af1770e.

📒 Files selected for processing (4)
  • newton/_src/sim/builder.py
  • newton/_src/sim/model.py
  • newton/_src/solvers/mujoco/kernels.py
  • newton/_src/solvers/mujoco/solver_mujoco.py
🧰 Additional context used
🧠 Learnings (19)
📚 Learning: 2025-09-22T21:08:31.901Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/examples/ik/example_ik_franka.py:121-123
Timestamp: 2025-09-22T21:08:31.901Z
Learning: In the newton physics framework, when creating warp arrays for IK solver joint variables using wp.array(self.model.joint_q, shape=(1, coord_count)), the resulting array acts as a reference/pointer to the original model's joint coordinates, so updates from the IK solver automatically reflect in the model's joint_q buffer used for rendering, despite the general warp documentation suggesting copies are made by default.

Applied to files:

  • newton/_src/sim/model.py
  • newton/_src/solvers/mujoco/solver_mujoco.py
  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-08-19T12:27:30.629Z
Learnt from: preist-nvidia
Repo: newton-physics/newton PR: 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/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-11-24T08:05:21.390Z
Learnt from: adenzler-nvidia
Repo: newton-physics/newton PR: 1107
File: newton/_src/solvers/mujoco/kernels.py:973-974
Timestamp: 2025-11-24T08:05:21.390Z
Learning: In Newton's MuJoCo solver integration (newton/_src/solvers/mujoco/), the mapping between Newton joints and MuJoCo joints is not 1-to-1. Instead, each Newton DOF maps to a distinct MuJoCo joint. This means that for multi-DOF joints (like D6 with 6 DOFs), there will be 6 corresponding MuJoCo joints, each with its own properties (margin, solimp, solref, etc.). The mapping is done via dof_to_mjc_joint array, ensuring each DOF's properties are written to its own MuJoCo joint without overwriting.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-08-19T12:27:30.629Z
Learnt from: preist-nvidia
Repo: newton-physics/newton PR: 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/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-08-18T15:56:26.587Z
Learnt from: adenzler-nvidia
Repo: newton-physics/newton PR: 552
File: newton/_src/solvers/mujoco/solver_mujoco.py:0-0
Timestamp: 2025-08-18T15:56:26.587Z
Learning: In Newton's MuJoCo solver, when transforming joint axes from Newton's internal frame to MuJoCo's expected frame, use wp.quat_rotate(joint_rot, axis) not wp.quat_rotate_inv(joint_rot, axis). The joint_rot represents rotation from joint-local to body frame, so forward rotation is correct.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-09-09T12:50:42.929Z
Learnt from: gdaviet
Repo: newton-physics/newton PR: 750
File: newton/_src/solvers/implicit_mpm/render_grains.py:142-160
Timestamp: 2025-09-09T12:50:42.929Z
Learning: The `update_render_grains` function in newton/_src/solvers/implicit_mpm/render_grains.py should return early when `state.velocity_field is None` because its purpose is specifically to advect grains in the velocity field. Without a velocity field, the function should be a no-op.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-09-22T21:03:39.624Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_lbfgs_optimizer.py:739-752
Timestamp: 2025-09-22T21:03:39.624Z
Learning: The L-BFGS optimizer in newton/_src/sim/ik/ik_lbfgs_optimizer.py currently intentionally only supports additive updates (assuming n_coords == n_dofs). Velocity space integration for joints with mismatched coordinate/DOF dimensions (like free/ball joints) is planned for future work and should not be flagged as an issue in current reviews.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-09-18T07:05:56.836Z
Learnt from: gdaviet
Repo: newton-physics/newton PR: 750
File: newton/_src/solvers/implicit_mpm/solve_rheology.py:969-986
Timestamp: 2025-09-18T07:05:56.836Z
Learning: In newton/_src/solvers/implicit_mpm/solve_rheology.py, transposed_strain_mat parameter cannot be None - the type signature was corrected to reflect this guarantee, eliminating the need for None checks when accessing transposed_strain_mat.offsets.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-09-26T05:58:21.013Z
Learnt from: WenchaoHuang
Repo: newton-physics/newton PR: 835
File: newton/_src/solvers/style3d/collision/collision.py:49-53
Timestamp: 2025-09-26T05:58:21.013Z
Learning: In newton/_src/solvers/style3d/collision/collision.py, the broad_phase buffers (broad_phase_ee, broad_phase_ef, broad_phase_vf) created with wp.array() don't need zero-initialization because they follow a write-first-read-later pattern: BVH query kernels in frame_begin() populate these buffers completely (write-only operations), and collision kernels only read from them afterward in accumulate_contact_force(), ensuring no undefined reads occur.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-08-12T05:17:34.423Z
Learnt from: shi-eric
Repo: newton-physics/newton PR: 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.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-09-22T21:08:31.901Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/examples/ik/example_ik_franka.py:121-123
Timestamp: 2025-09-22T21:08:31.901Z
Learning: In the newton physics framework, when creating warp arrays for IK solver joint variables using wp.array(self.model.joint_q, shape=(1, coord_count)), the resulting array acts as a reference/pointer to the original model's joint coordinates, so updates from the IK solver automatically reflect in the model's joint_q buffer used for rendering.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-08-12T18:04:06.577Z
Learnt from: nvlukasz
Repo: newton-physics/newton PR: 519
File: newton/_src/solvers/featherstone/kernels.py:75-75
Timestamp: 2025-08-12T18:04:06.577Z
Learning: The Newton physics framework requires nightly Warp builds, which means compatibility concerns with older stable Warp versions (like missing functions such as wp.spatial_adjoint) are not relevant for this project.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-11-28T11:12:40.805Z
Learnt from: camevor
Repo: newton-physics/newton PR: 1146
File: newton/examples/basic/example_basic_joints.py:206-213
Timestamp: 2025-11-28T11:12:40.805Z
Learning: In Newton's standard solvers (XPBD, SemiImplicit, SolverMuJoCo), spatial vectors in State.body_qd use the ordering (linear, angular): wp.spatial_top(qd) returns linear velocity and wp.spatial_bottom(qd) returns angular velocity. This is the opposite of typical Modern Robotics spatial twist notation where (angular, linear) is used.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-12-01T17:00:28.889Z
Learnt from: camevor
Repo: newton-physics/newton PR: 1161
File: newton/_src/solvers/mujoco/kernels.py:1268-1268
Timestamp: 2025-12-01T17:00:28.889Z
Learning: MuJoCo Warp spatial vectors (e.g., cacc, cvel) use the convention (angular, linear), which is opposite to Newton's (linear, angular) convention. When converting from MuJoCo Warp to Newton format, spatial_top() on MuJoCo Warp data extracts the angular component, and spatial_bottom() extracts the linear component.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-11-28T11:12:40.030Z
Learnt from: camevor
Repo: newton-physics/newton PR: 1146
File: newton/examples/basic/example_basic_joints.py:215-222
Timestamp: 2025-11-28T11:12:40.030Z
Learning: In Warp's spatial vector convention used by Newton: spatial_bottom() extracts the angular velocity component (omega) and spatial_top() extracts the linear velocity component (v) from a spatial twist (qd). Similarly for spatial wrenches: spatial_bottom() is torque and spatial_top() is force.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-12-01T16:21:36.581Z
Learnt from: jvonmuralt
Repo: newton-physics/newton PR: 1160
File: newton/_src/sim/articulation.py:0-0
Timestamp: 2025-12-01T16:21:36.581Z
Learning: In Newton's Featherstone solver, during forward kinematics for FREE/DISTANCE joints, the computed spatial velocity v_wc follows Featherstone's spatial twist convention where the linear component (spatial_top) represents the velocity at the world origin (0,0,0), not at the joint anchor or COM. To convert to State.body_qd (which stores COM velocity), use v_com = v_origin + ω × x_com, where x_com is the absolute world position of the COM.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-12-12T16:25:54.685Z
Learnt from: vastsoun
Repo: newton-physics/newton PR: 1019
File: newton/_src/solvers/kamino/core/math.py:248-255
Timestamp: 2025-12-12T16:25:54.685Z
Learning: In newton/_src/solvers/kamino/core/math.py, the quat_conj function's current implementation (negating w instead of x,y,z) is a known issue that will be revised in a future PR per vastsoun's decision, not to be flagged in current Kamino reviews.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-09-25T16:14:22.033Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_objectives.py:0-0
Timestamp: 2025-09-25T16:14:22.033Z
Learning: wp.transform in NVIDIA Warp supports indexing using numerical indices (e.g., body_tf[0], body_tf[3], etc.) to access translation and rotation components. The indexing approach used in the Newton codebase is valid and compiles correctly.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-09-25T16:14:22.033Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_objectives.py:0-0
Timestamp: 2025-09-25T16:14:22.033Z
Learning: In NVIDIA Warp's Newton physics library, wp.transform supports direct numerical indexing (e.g., body_tf[0], body_tf[1], body_tf[2] for position and body_tf[3], body_tf[4], body_tf[5], body_tf[6] for quaternion components) to access its elements. This is the standard API used throughout the codebase.

Applied to files:

  • newton/_src/solvers/mujoco/kernels.py
🧬 Code graph analysis (1)
newton/_src/sim/builder.py (1)
newton/_src/sim/model.py (1)
  • request_state_attributes (712-719)
⏰ 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 Tests / Run GPU Unit Tests on AWS EC2
  • GitHub Check: Run GPU Benchmarks / Run GPU Benchmarks on AWS EC2
  • GitHub Check: run-newton-tests / newton-unittests (ubuntu-latest)
  • GitHub Check: run-newton-tests / newton-unittests (windows-latest)
🔇 Additional comments (11)
newton/_src/sim/model.py (4)

473-474: LGTM!

The initialization of _requested_state_attributes as an empty set is appropriate for tracking which extended state attributes are requested.


551-552: LGTM!

The selective allocation of extended state attributes is implemented correctly. Core attributes (particle_q, body_q, joint_q, etc.) are allocated unconditionally when entity counts are non-zero (maintaining backward compatibility), while extended attributes (body_qdd, body_parent_f) are only allocated when explicitly requested.

Also applies to: 575-580


712-720: LGTM!

The request_state_attributes() method provides a clean API for requesting extended state attributes. Note that attribute names are not validated here; invalid names will simply be ignored during state allocation.


847-870: LGTM!

The get_requested_state_attributes() method correctly builds the set of attributes to allocate by:

  1. Starting with default core attributes based on entity counts
  2. Adding explicitly requested attributes that aren't already included

The use of difference() on line 869 ensures no duplicates are added.

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

54-54: LGTM!

Import of convert_rigid_forces_from_mj_kernel is appropriate for populating extended state attributes.


479-479: LGTM with prototype limitation noted

The enable_rne_postconstraint() integration correctly enables RNE sensor data when extended state attributes are present. The TODO comment (line 498) explicitly documents that CPU mode is not yet supported, which is acceptable for this prototype.

Also applies to: 495-507

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

555-556: Storing requested state attributes on the builder looks correct

Initializing _requested_state_attributes as a set[str] is appropriate here and keeps the builder’s configuration self-contained without affecting existing behavior.


1611-1612: Correctly merging requested state attributes when adding sub‑builders

Taking the union of self._requested_state_attributes with builder._requested_state_attributes ensures extended state requirements from sub‑builders are preserved when composing scenes, which matches how other builder data is merged.


5838-5838: Finalize correctly propagates builder requests into the Model

Calling m.request_state_attributes(*self._requested_state_attributes) immediately after constructing the Model wires the builder’s configuration into the model cleanly and aligns with the extended state attribute flow (builder → model → state) without impacting existing initialization logic.

newton/_src/solvers/mujoco/kernels.py (2)

1325-1344: LGTM! Correct handling of MuJoCo Warp spatial vector convention.

The function correctly extracts angular (spatial_top) and linear (spatial_bottom) components from MuJoCo Warp's (angular, linear) convention and applies appropriate frame transformation corrections for offset-based acceleration computation.

Based on learnings, MuJoCo Warp spatial vectors use (angular, linear) ordering, opposite to Newton's convention.


1368-1400: LGTM! Correct convention conversion from MuJoCo Warp to Newton format.

Line 1400 correctly constructs body_qdd in Newton's (linear, angular) convention:

  • lin + mjw_gravity[0]: linear acceleration including gravity
  • wp.spatial_top(cacc): extracts angular acceleration from MuJoCo Warp's (angular, linear) format

The convention handling matches the resolved discussion in past review comments.

Based on learnings, this correctly converts between MuJoCo Warp's (angular, linear) and Newton's (linear, angular) spatial vector conventions.

Copy link
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

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

Actionable comments posted: 1

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)
README.md (1)

373-375: Add the command for the IMU sensor example.

The third column in the command row remains commented as "Future Sensor example" but should now display the command to run the newly added IMU sensor example. While these lines aren't part of the annotated changes, completing this row would make the table consistent with the new IMU sensor tile above.

🔎 Suggested addition
     <td align="center">
-      <!-- Future Sensor example -->
+      <code>uv run -m newton.examples sensor_imu</code>
     </td>
📜 Review details

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between af1770e and 8f25d32.

⛔ Files ignored due to path filters (1)
  • docs/images/examples/example_sensor_imu.jpg is excluded by !**/*.jpg
📒 Files selected for processing (3)
  • README.md
  • newton/examples/sensors/example_sensor_imu.py
  • newton/tests/test_examples.py
🚧 Files skipped from review as they are similar to previous changes (1)
  • newton/examples/sensors/example_sensor_imu.py
🧰 Additional context used
🧠 Learnings (1)
📚 Learning: 2025-08-27T19:05:44.697Z
Learnt from: Milad-Rakhsha-NV
Repo: newton-physics/newton PR: 535
File: newton/tests/test_examples.py:320-414
Timestamp: 2025-08-27T19:05:44.697Z
Learning: In newton/examples/__init__.py, the robot policy example is registered with the key "robot_policy" (not "robot.example_robot_policy"), so tests should reference it as name="robot_policy".

Applied to files:

  • newton/tests/test_examples.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 / Run GPU Benchmarks on AWS EC2
  • GitHub Check: Run GPU Tests / Run GPU Unit Tests on AWS EC2
  • GitHub Check: run-newton-tests / newton-unittests (windows-latest)
  • GitHub Check: run-newton-tests / newton-unittests (ubuntu-latest)
🔇 Additional comments (1)
newton/tests/test_examples.py (1)

628-634: LGTM!

The test registration follows the established pattern for sensor examples and includes appropriate test parameters. The 200-frame duration with the clarifying comment is helpful.

README.md Outdated Show resolved Hide resolved
@camevor camevor enabled auto-merge December 22, 2025 17:30
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.

Implement prototype for requesting State attributes Add IMU Sensor

3 participants

Morty Proxy This is a proxified and sanitized view of the page, visit original site.