[DRAFT] Add `Model.joint_act` that reads from `MjData.qfrc_actuator` #871

etaoxing · Oct 2, 2025

Description

Adds attribute State.joint_act to read actuation from MjData.qfrc_actuator. This PR is a draft to illustrate the current challenge of modifying (ModelBuilder, Model, State, SolverMujoco) in order to read attributes from MjData that are not currently supported in Newton.

Newton Migration Guide

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

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

Before your PR is "Ready for review"

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

Summary by CodeRabbit

New Features
- Added per-DOF joint activation (joint_act) across the builder, model, and state, enabling storage and retrieval of actuation values.
- Included joint_act in model export and state cloning, keeping it aligned with joint positions and velocities.
- Preserved joint_act through joint collapse/restore workflows.
- Integrated actuator state with MuJoCo and MuJoCo Warp backends, synchronizing joint_act during coordinate conversions and state updates for consistent data exchange.

coderabbitai · Oct 2, 2025

📝 Walkthrough

Walkthrough

Adds per-DOF joint activation tracking across builder, model, and state as joint_act. Wires joint_act through builder creation, collapse/restore, add_builder transfer, and finalize export. Extends Model and State to store joint_act. Updates MuJoCo solver to convert and synchronize joint_act with MuJoCo qact in both MJ→Newton and Newton→MJ flows and kernels.

Changes

Cohort / File(s)	Summary of changes
Simulation builder: joint_act plumbing `newton/_src/sim/builder.py`	Introduces builder-level `joint_act` list; appends per-DOF entries on joint creation; includes in collapse/restore axis data; adds to builder-to-model copy lists; exports to `Model` during finalize.
Model and State: joint_act storage `newton/_src/sim/model.py`, `newton/_src/sim/state.py`	Adds `joint_act` attribute to `Model` and `State`; clones `joint_act` in `Model.state()`; registers `joint_act` in `attribute_frequency` as "joint_dof".
MuJoCo solver: actuator/qact integration `newton/_src/solvers/mujoco/solver_mujoco.py`	Extends conversion kernels’ signatures and flows to map `joint_act`↔`qact`; threads qact through MJ data paths (classic and Warp); updates state synchronization to read/write `joint_act`/`qact` alongside qpos/qvel.

Sequence Diagram(s)

sequenceDiagram
  autonumber
  participant MJ as MuJoCo mj_data
  participant ConvF as convert_mj_coords_to_warp_kernel
  participant NW as Newton State/Model

  Note over MJ,NW: MJ → Newton update (read)
  MJ->>ConvF: qpos, qvel, qact
  ConvF->>NW: joint_q, joint_qd, joint_act
  Note right of NW: State now holds joint_act aligned with DOFs

sequenceDiagram
  autonumber
  participant NW as Newton State/Model
  participant ConvB as convert_warp_coords_to_mj_kernel
  participant MJ as MuJoCo mj_data

  Note over NW,MJ: Newton → MJ update (write)
  NW->>ConvB: joint_q, joint_qd, joint_act
  ConvB->>MJ: qpos, qvel, qact
  Note right of MJ: qact (e.g., qfrc_actuator) updated from joint_act

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 61.54% which is insufficient. The required threshold is 80.00%.	You can run `@coderabbitai generate docstrings` to improve docstring coverage.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title Check	✅ Passed	The title succinctly summarizes the primary change by stating the addition of Model.joint_act and its linkage to MjData.qfrc_actuator, making it clear and focused on the key update.

✨ Finishing touches

📝 Generate Docstrings

🧪 Generate unit tests

Create PR with unit tests
Post copyable unit tests in a comment

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

❤️ Share

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

coderabbitai

Actionable comments posted: 1

📜 Review details

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 709422d and 65ba706.

📒 Files selected for processing (4)

newton/_src/sim/builder.py (6 hunks)
newton/_src/sim/model.py (3 hunks)
newton/_src/sim/state.py (1 hunks)
newton/_src/solvers/mujoco/solver_mujoco.py (13 hunks)

🧰 Additional context used

🧬 Code graph analysis (2)

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

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

clear (67-77)

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

requires_grad (81-87)

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

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

joint_coord_count (104-108)

joint_dof_count (111-115)

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

state (443-478)

🔇 Additional comments (10)

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

64-65: LGTM!

The new joint_act field follows the established pattern for joint state attributes (matching joint_q and joint_qd). The documentation clearly specifies the shape and dtype, maintaining consistency with other joint-related fields.

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

229-230: LGTM!

The joint_act field is properly added to the Model class with clear documentation following the same pattern as joint_q and joint_qd.

410-410: LGTM!

Correctly registers joint_act in the attribute frequency mapping as a "joint_dof" attribute, consistent with joint_qd and other per-DOF attributes.

476-476: LGTM!

The state() method properly clones joint_act when creating a new State object, maintaining consistency with the cloning of joint_q and joint_qd.

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

420-420: LGTM!

The joint_act list is properly initialized in ModelBuilder.__init__, following the same pattern as joint_qd and joint_f.

1105-1105: LGTM!

Including joint_act in more_builder_attrs ensures it is correctly copied when using add_builder() to compose models, maintaining consistency with joint_qd and joint_f.

1341-1341: LGTM!

The per-DOF initialization of joint_act with 0.0 in add_joint is consistent with the initialization of joint_qd and joint_f on adjacent lines (1340, 1342).

2090-2090: LGTM!

The joint_act slice is properly captured in the collapsed joint data structure within collapse_fixed_joints, matching the handling of qd on line 2089.

2284-2284: LGTM!

Clearing joint_act during the collapse_fixed_joints restructuring is necessary and consistent with the clearing of other joint attributes like joint_qd (line 2283).

4339-4339: LGTM!

The finalization of joint_act into a wp.array with dtype=wp.float32 and requires_grad parameter is correct and follows the same pattern as joint_qd (line 4321) and joint_f (line 4340).

coderabbitai · Oct 2, 2025

newton/_src/solvers/mujoco/solver_mujoco.py

+        joint_act[wqd_i + 0] = qact[worldid, q_i + 0]
+        joint_act[wqd_i + 1] = qact[worldid, q_i + 1]
+        joint_act[wqd_i + 2] = qact[worldid, q_i + 2]
+        joint_act[wqd_i + 3] = qact[worldid, q_i + 3]
+        joint_act[wqd_i + 4] = qact[worldid, q_i + 4]
+        joint_act[wqd_i + 5] = qact[worldid, q_i + 5]
+


⚠️ Potential issue | 🔴 Critical

Fix qact indexing for free joints.

qact/qfrc_actuator follow the velocity layout (stride = nv), but here we index them with q_i (qpos stride). On the second and later free joints this walks past the end of the actuator array, corrupting data. Use qd_i just like we already do for velocities.

Apply this diff:

- joint_act[wqd_i + 0] = qact[worldid, q_i + 0] - joint_act[wqd_i + 1] = qact[worldid, q_i + 1] - joint_act[wqd_i + 2] = qact[worldid, q_i + 2] - joint_act[wqd_i + 3] = qact[worldid, q_i + 3] - joint_act[wqd_i + 4] = qact[worldid, q_i + 4] - joint_act[wqd_i + 5] = qact[worldid, q_i + 5] + joint_act[wqd_i + 0] = qact[worldid, qd_i + 0] + joint_act[wqd_i + 1] = qact[worldid, qd_i + 1] + joint_act[wqd_i + 2] = qact[worldid, qd_i + 2] + joint_act[wqd_i + 3] = qact[worldid, qd_i + 3] + joint_act[wqd_i + 4] = qact[worldid, qd_i + 4] + joint_act[wqd_i + 5] = qact[worldid, qd_i + 5]

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

joint_act[wqd_i + 0] = qact[worldid, q_i + 0]

joint_act[wqd_i + 1] = qact[worldid, q_i + 1]

joint_act[wqd_i + 2] = qact[worldid, q_i + 2]

joint_act[wqd_i + 3] = qact[worldid, q_i + 3]

joint_act[wqd_i + 4] = qact[worldid, q_i + 4]

joint_act[wqd_i + 5] = qact[worldid, q_i + 5]

joint_act[wqd_i + 0] = qact[worldid, qd_i + 0]

joint_act[wqd_i + 1] = qact[worldid, qd_i + 1]

joint_act[wqd_i + 2] = qact[worldid, qd_i + 2]

joint_act[wqd_i + 3] = qact[worldid, qd_i + 3]

joint_act[wqd_i + 4] = qact[worldid, qd_i + 4]

joint_act[wqd_i + 5] = qact[worldid, qd_i + 5]

Add Model.joint_act that reads from MjData.qfrc_actuator

65ba706

etaoxing had a problem deploying to external-pr-approval October 2, 2025 17:34 — with GitHub Actions Failure

coderabbitai bot reviewed Oct 2, 2025

View reviewed changes

etaoxing mentioned this pull request Oct 2, 2025

Monolithic Model/State vs Solver-specific Model/State #872

Open

shi-eric marked this pull request as draft October 2, 2025 23:31

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[DRAFT] Add `Model.joint_act` that reads from `MjData.qfrc_actuator` #871

[DRAFT] Add `Model.joint_act` that reads from `MjData.qfrc_actuator` #871

etaoxing commented Oct 2, 2025 •

edited by coderabbitai bot

Loading

Uh oh!

coderabbitai bot commented Oct 2, 2025 •

edited

Loading

Walkthrough

Changes

Sequence Diagram(s)

Estimated code review effort

Uh oh!

coderabbitai bot left a comment

Uh oh!

coderabbitai bot Oct 2, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

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

[DRAFT] Add Model.joint_act that reads from MjData.qfrc_actuator #871

Are you sure you want to change the base?

[DRAFT] Add Model.joint_act that reads from MjData.qfrc_actuator #871

Conversation

etaoxing commented Oct 2, 2025 • edited by coderabbitai bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Newton Migration Guide

Before your PR is "Ready for review"

Summary by CodeRabbit

Uh oh!

coderabbitai bot commented Oct 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Sequence Diagram(s)

Estimated code review effort

Pre-merge checks and finishing touches

Uh oh!

coderabbitai bot left a comment

Choose a reason for hiding this comment

Uh oh!

coderabbitai bot Oct 2, 2025

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

[DRAFT] Add `Model.joint_act` that reads from `MjData.qfrc_actuator` #871

[DRAFT] Add `Model.joint_act` that reads from `MjData.qfrc_actuator` #871

etaoxing commented Oct 2, 2025 •

edited by coderabbitai bot

Loading

coderabbitai bot commented Oct 2, 2025 •

edited

Loading