-
Notifications
You must be signed in to change notification settings - Fork 206
Extended state attributes and IMU Sensor #1161
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Conversation
📝 WalkthroughWalkthroughAdds 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
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
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes
Suggested reviewers
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
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. Comment |
There was a problem hiding this 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 whenStatewas 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 helpStateobjects that were created before the solver was constructed. With the new allocation scheme, such states can havestate_in.particle_q is Noneeven whenmodel.particle_count > 0, which matches the pipeline failure (AttributeError: 'NoneType' object has no attribute 'device'when cloningstate_in.particle_q).Two concrete follow-ups are advisable:
- 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 bothstate_inandstate_out.
- Clarify the lifecycle (docs or migration notes) that XPBD expects its required attributes to be requested before any
Stateused with it is created—either via the builder or by constructing the solver before callingmodel.state().Without at least the runtime validation, XPBD will continue to fail with cryptic
NoneType/deviceerrors when used with pre-existing states under the new system.newton/_src/solvers/mujoco/solver_mujoco.py (1)
361-401:enable_rne_postconstraintgating is effectively unconditionalThe
hasattr()check will always returnTruebecauseStatedefinesbody_qddandbody_parent_fas fields initialized toNone. 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 = TrueThis 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 ofStateobjectsRequesting
"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 whatstep()uses.Be aware this only applies to
Stateinstances created afterSolverFeatherstoneis constructed;States created earlier will still haveNonefor 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 notNonewhenmodel.body_count/model.joint_countare nonzero.newton/_src/solvers/semi_implicit/solver_semi_implicit.py (1)
67-97: Appropriate attribute requests; consider guarding against pre-existingStatesGating the requests on
model.body_count/model.particle_countand asking for"body_q","body_qd","body_f"plus"particle_q","particle_qd","particle_f"matches whatSolverSemiImplicit.step()uses and keeps allocations minimal.As with XPBD/Featherstone, these requests only affect
Stateinstances created after they run. If any callers constructStates before the solver, those states can still haveNonefor 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 fieldsbody_qdd/body_parent_flook consistent with existing spatial conventionsAdding
body_qddandbody_parent_fas optionalwp.spatial_vectorarrays on theStatesurface 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_qddocstring wording (linear vs. angular ordering) forbody_qddand clarify the frame in whichbody_parent_fis 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; noteStatecreation order caveatRequesting
"body_q","body_qd","body_f","joint_q","joint_qd"from theModelis the right set for this solver and aligns with the new attribute-request API. The updated docstring forto_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 constructStates first and then attachSolverMuJoCo, 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()orupdate_newton_state()that required arrays aren’tNonewhenmodel.body_count/model.joint_countare nonzero.newton/_src/sim/builder.py (1)
497-498: State attribute request plumbing in ModelBuilder looks good; consider optional validationThe
_requested_state_attributesset, its aggregation inadd_builder(), and propagation intoModelinfinalize()form a coherent path for per-builder state-attribute requests. Therequest_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 outputsIn
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 inrender()) holds the post-step state, butIMUSensor.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_0This 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 quantityThe
ViewerPlotwindow 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 callingrequest_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-createdStateinstances are not retrofitted.- Optionally, when
assignment == STATEin_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_fcomputation is not yet implemented.The
body_parent_fattribute is currently not populated (no-op withpass). 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
📒 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.pynewton/_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.pynewton/_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.pynewton/_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.pynewton/_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.pynewton/_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: ExposeIMUSensorvianewton.sensorslooks correctImporting
IMUSensorfrom 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 soundThe 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_linandvel_anginto the sensor frame usingq = body_quat * site_quatandquat_rotate_invmatches 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 missingmujoco_warpdependency.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,
BODYrefers to the inertial/center-of-mass frame, whileXBODYrefers 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_framelinacccomputation and clarify which frame type aligns with the intended behavior, then update or resolve this TODO accordingly.
1250-1266: Verify the_framelinaccfunction signature and parameter correctness.The call to
_framelinaccpassesmjw_body_rootidfour times andmjw_xposfive 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.
There was a problem hiding this 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 statesRequesting
"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 fromModel.state()allocating no arrays when_requested_state_attributesis still empty. That needs to be fixed inModel.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 explicitis not Nonechecks when deciding whether to launchconvert_rigid_forces_from_mj_kernelThe 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.arrayobjects, which is non-obvious and inconsistent with the rest of this file (where you usually testis not Nonefor 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 requestedRight now
state()usesrequested = self._requested_state_attributesdirectly, so when the set is empty (the default afterModelBuilder.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 createsStateobjects before any solver or sensor callsrequest_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: Guardbody_qdd/body_parent_fallocation on the presence ofbody_qdAllocating
body_qddandbody_parent_fviawp.zeros_like(self.body_qd, ...)assumesself.body_qdhas 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 duplicatingModel.request_state_attributesdefinitionsThe
request_state_attributes()method here looks good, but the external context shows it is defined twice innewton/_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 presentThe new guard:
if assignment == ModelAttributeAssignment.STATE and full_name not in self._requested_state_attributes: continuemeans 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 + ): + continueThat 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_kernelspatial mapping looks consistent;body_parent_fremains TODOThe new kernel correctly follows the MuJoCo Warp → Newton spatial convention:
mjw_caccis 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 wherespatial_topis 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_fis intentionally left as a stub (pass), which is acceptable for the prototype but meansstate.body_parent_fwill stay unused even if allocated. When you move beyond prototype stage, you’ll likely want to follow up by wiring the appropriate RNE wrench intobody_parent_f.newton/_src/solvers/mujoco/solver_mujoco.py (1)
373-401:enable_rne_postconstraintshould check for allocated fields, not just attribute existence
Statealways hasbody_qddandbody_parent_fattributes (initialized toNone), so:if any(hasattr(state_out, field) for field in rne_postconstraint_fields):is always
True, andsensor_rne_postconstraintwill 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-
Noneor 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
📒 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.pynewton/_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.pynewton/_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.pynewton/_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.pynewton/_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.pynewton/_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.pynewton/_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.pynewton/_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_attributesis fine and matches the new APIStoring requested state attributes in a dedicated
set[str]is an appropriate internal representation and fits the newrequest_state_attributes()API.newton/_src/solvers/mujoco/solver_mujoco.py (2)
45-54: Kernel import ofconvert_rigid_forces_from_mj_kernelis consistent with new RNE extraction pathAdding
convert_rigid_forces_from_mj_kernelto the imported kernel list is appropriate and matches the newupdate_newton_state()logic that populatesstate.body_qdd/state.body_parent_ffrom MuJoCo data.
285-286: Requesting core body/joint state attributes inSolverMuJoCo.__init__is correct but depends on Model.state fallbackHaving the MuJoCo solver request
"body_q","body_qd","body_f","joint_q", and"joint_qd"up front is the right behavior for anyStateobjects 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 innewton/_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: Makemujoco_warpimports optional to avoid hard dependency at module import timeImporting
mujoco_warp._src.sensorandmujoco_warp._src.typesat the top level causesModuleNotFoundErrorwhennewton._src.solvers.mujoco.kernelsis imported in environments wheremujoco_warpis not installed. This breaks users that never touch the MuJoCo solver.Recommend wrapping these imports in a
try/exceptand providing a minimal stub, relying onSolverMuJoCo.import_mujoco()to enforce the hard dependency when the solver is actually used. This pattern is already used innewton/_src/utils/import_usd.pyfor 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 = 0This unblocks tests and non-MuJoCo usage while preserving the requirement that
mujoco_warpmust be installed to actually run the MuJoCo backend.
adenzler-nvidia
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks, some comments.
225ea4f to
2c7f58a
Compare
There was a problem hiding this 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 inIMUSensor.update
IMUSensor.update()assumesstate.body_q,state.body_qd, and especiallystate.body_qddare allocated:
model.request_state_attributes("body_qdd")only affectsStateobjects created after that call.- Users can easily pass an older State or one produced by a solver that doesn’t populate
body_qddyet, 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 truthinessRight now:
enable_rne_postconstraintchecksany(hasattr(state_out, field) for field in {"body_qdd", "body_parent_f"}), which is alwaysTruebecauseStatedefines those attributes (initialized toNone). This meanssensor_rne_postconstraintis effectively enabled unconditionally on the mujoco_warp path after the first step, even when no extended fields were requested.update_newton_stategates the force/acceleration kernel withif any((state.body_qdd, state.body_parent_f)), which relies on truthiness ofwp.arrayvsNoneand is harder to read than an explicitis not Nonecheck.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 || trueAlso 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 withassignment == ModelAttributeAssignment.STATE, so even if callers register custom STATE attributes viaadd_attribute(...)and request them, they will never be copied ontoStateinstances.If the intent is to gate STATE attributes on the new
_requested_state_attributesmechanism 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: ExposeIMUSensorvianewton.sensors– looks goodImporting
IMUSensorand 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 sensorscomment 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 inadd_builderis well‑aligned with reuse semanticsMerging
builder._requested_state_attributesintoself._requested_state_attributesensures that sub‑builders bring along their extended‑state requirements when composed into a larger scene, and thesetkeeps this idempotent. Consider briefly noting this behavior in theadd_builderdocstring so users understand that state‑attribute requirements from sub‑builders are aggregated automatically, but functionally this is sound.
5169-5177:ModelBuilder.request_state_attributesAPI is a clean mirror ofModel.request_state_attributesThe method is minimal and correctly accumulates attributes into the internal set, matching the behavior of
Model.request_state_attributesshown innewton/_src/sim/model.py. You might optionally add a tiny note in the docstring that these requests are forwarded to theModelduringfinalize, to make the lifecycle explicit, but behavior-wise this is fine.newton/examples/example_solver_data.py (4)
33-33: Inconsistent text label.The
ViewerPlotis initialized with title"Parent force", but the label insiderender()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.titleproperty 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
📒 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.pynewton/_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.pynewton/_src/solvers/mujoco/solver_mujoco.pynewton/_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.pynewton/_src/solvers/mujoco/solver_mujoco.pynewton/_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.pynewton/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.pynewton/_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.pynewton/_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.pynewton/_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.pynewton/_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.pynewton/_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: Newbody_qdd/body_parent_ffields are consistent with existing State designAdding
body_qddandbody_parent_fas optional spatial-vector arrays onStatefits the existing pattern for body fields and will be handled automatically byassign()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 orderThe 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 correctStoring 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 placeCalling
m.request_state_attributes(*self._requested_state_attributes)immediately after constructing theModelensures the model’s internal requested‑attributes set is populated before anyStateobjects are created. This cleanly connects the builder‑side API to the model‑side implementation without affecting existing users who never callrequest_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 wherespatial_top()extracts angular andspatial_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. Thebody_parent_fTODO is acknowledged for this prototype.newton/examples/example_solver_data.py (2)
110-117: LGTM!The
render()method follows the standard example pattern. The emptytest()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.
There was a problem hiding this 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 functionalityThe 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; fixListtype and guardbody_qdd/body_parent_fallocations
The combination of
_requested_state_attributes,get_requested_state_attributes(), and the gated allocations instate()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 likebody_qdd/body_parent_fare opt-in viarequest_state_attributes(). This addresses the earlier regression where defaultStateobjects lacked core arrays. ✔In
get_requested_state_attributes()at line 849, the return annotation usesList[str]butListis never imported, which Ruff correctly flags (F821). Given you already rely on Python 3.9+ generics, switch to the built-inlist[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 plottingimu_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 usingself.titleconsistently: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 intotest_examples.py, you’ll likely want it to run a short headless simulation (e.g., a fewstep()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
📒 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 consistentThe kernel arguments, indexing (
sensor_idxbounds 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.
- 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
09217b3 to
c8bfcf0
Compare
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
There was a problem hiding this 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 modeThe new
enable_rne_postconstraint()+convert_rigid_forces_from_mj_kernelpath works correctly for the mjwarp backend:enable_rne_postconstraint()is called before stepping (line 479), enabling RNE computation whenstate_out.body_qddorstate_out.body_parent_fis allocated, and then the extended update block populates these fields fromself.mjw_model/self.mjw_data.However, in
update_newton_state()the extended update (lines 799–820) runs wheneverstate.body_qddorstate.body_parent_fis notNone, 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_mjwarpis false and the step usesself.mj_data, but this block still reads fromself.mjw_model/self.mjw_data, which are never stepped in that mode. If a caller requestsbody_qddwhile 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 commentThe new
SensorIMUimport and inclusion in__all__are consistent with the implementation module and make the IMU sensor available vianewton.sensorsas 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 enhancementThe 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(...)) beforeimu.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 behaviorThis page cleanly explains:
- which attributes are “core” vs “extended,”
- how to request them on both
ModelBuilderandModel,- their idempotent behavior,
- the current set (
body_qdd,body_parent_f) and solver support,- and how components like
SensorIMUtransparently request what they need.Everything is consistent with the new
ModelBuilder.request_state_attributesandModel.request_state_attributesAPIs.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 finalizeThe new
_requested_state_attributesplumbing looks coherent:
__init__initializes a per‑builder set, andrequest_state_attributes()simplyupdates it, so repeated calls are idempotent.add_builder()unionsbuilder._requested_state_attributesinto the parent, which means imports,add_world(), andreplicate()all correctly aggregate extended attribute requests from sub‑builders.finalize()forwards the accumulated set intoModel.request_state_attributes(...)before anyStateis created, so builder‑level requests take effect as intended.This matches the behavior documented in
extended_state_attributes.rstand lines up with the existingModel.request_state_attributesAPI.If you’d like, you could add a short note to the
ModelBuilderclass docstring pointing users atrequest_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 guardThe
_requested_state_attributesset,request_state_attributes(), andget_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 ofbody_qdd/body_parent_finstate()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 withbody_count == 0), you could mirror the core-body guards and additionally checkself.body_countbefore 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 countThe 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 intoSensorIMU.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
📒 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.rstdocs/concepts/sensors.rstnewton/_src/solvers/mujoco/kernels.pynewton/_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.pynewton/_src/solvers/mujoco/kernels.pynewton/_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.pynewton/examples/sensors/example_sensor_imu.pynewton/_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.pynewton/examples/sensors/example_sensor_imu.pynewton/_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.pynewton/_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.pynewton/_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.pynewton/_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.pynewton/_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: ExposeSensorIMUin public sensor API autosummaryAdding
SensorIMUhere matches the new public export innewton.sensorsand 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()andconvert_rigid_forces_from_mj_kernel()look consistent with the documented conventions:
- For MuJoCo Warp, you treat
spatial_topas angular andspatial_bottomas linear, then reconstruct COM linear acceleration with the expected offset and Coriolis terms.- For Newton, you build
body_qddas(linear, angular)viawp.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_fvia simple truthiness checks is also aligned with existing kernels in this module.Also applies to: 1334-1370
There was a problem hiding this 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_modelandself.mjw_datawithout explicit backend checks. Whileupdate_newton_stateis called from both CPU and mjwarp paths, the CPU path would haveself.mjw_model = Nonewhenmjw_model/mjw_dataaren'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
📒 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.pynewton/_src/sim/state.pynewton/_src/solvers/mujoco/solver_mujoco.pynewton/_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.pynewton/_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.pynewton/_src/solvers/mujoco/solver_mujoco.pynewton/_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.rstnewton/_src/solvers/mujoco/solver_mujoco.pynewton/_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.pynewton/_src/solvers/mujoco/solver_mujoco.pynewton/_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.pynewton/_src/solvers/mujoco/solver_mujoco.pynewton/_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.pynewton/_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.pynewton/_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_qddattribute for rigid body accelerations is correctly declared with appropriate type annotation and docstring.
65-66: LGTM!The
body_parent_fattribute 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_qddand 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_postconstraintmethod 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_attributesset 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_qddandbody_parent_fare correctly allocated only when explicitly requested, matching the prototype design.
703-710: LGTM!The
request_state_attributesmethod provides a simple API for sensors and users to declare their state requirements.
838-860: LGTM!The
get_requested_state_attributesmethod 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_colorkernel 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 ofwp.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_fimplementation 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_qddmust be requested beforestatecreation.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
rwith 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_qddattribute during initialization (line 130)- Validates site indices and flags (lines 144-152)
- Checks for
body_qddavailability 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.
There was a problem hiding this 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
📒 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 thesensor.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_qddis 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.
There was a problem hiding this 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: Newrequest_state_attributesAPI is sound; consider optional introspection helperThe API shape and behavior (simple set update) are consistent with
Model.request_state_attributesinnewton/_src/sim/model.pyand 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 toModel.get_requested_state_attributes()), you might optionally add a smallget_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_fcomputation 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_fcomputation, or is this intentionally deferred to a future PR?
📜 Review details
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
newton/_src/sim/builder.pynewton/_src/sim/model.pynewton/_src/solvers/mujoco/kernels.pynewton/_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.pynewton/_src/solvers/mujoco/solver_mujoco.pynewton/_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.pynewton/_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.pynewton/_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.pynewton/_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.pynewton/_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.pynewton/_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.pynewton/_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_attributesas 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:
- Starting with default core attributes based on entity counts
- 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_kernelis appropriate for populating extended state attributes.
479-479: LGTM with prototype limitation notedThe
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 correctInitializing
_requested_state_attributesas aset[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‑buildersTaking the union of
self._requested_state_attributeswithbuilder._requested_state_attributesensures 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 ModelCalling
m.request_state_attributes(*self._requested_state_attributes)immediately after constructing theModelwires 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_qddin Newton's(linear, angular)convention:
lin + mjw_gravity[0]: linear acceleration including gravitywp.spatial_top(cacc): extracts angular acceleration from MuJoCo Warp's(angular, linear)formatThe 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.
There was a problem hiding this 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
⛔ Files ignored due to path filters (1)
docs/images/examples/example_sensor_imu.jpgis excluded by!**/*.jpg
📒 Files selected for processing (3)
README.mdnewton/examples/sensors/example_sensor_imu.pynewton/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.
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.
docs/migration.rstis up-to dateBefore your PR is "Ready for review"
newton/tests/test_examples.py)pre-commit run -aSummary by CodeRabbit
New Features
Documentation
Tests
✏️ Tip: You can customize this high-level summary in your review settings.
Summary by CodeRabbit
New Features
Documentation
Examples
Tests
Chore
✏️ Tip: You can customize this high-level summary in your review settings.