Skip to content

mjx: zero acceleration sensors on static and mocap bodies - #3717

Open
dundysm wants to merge 1 commit into
google-deepmind:mainfrom
dundysm:fix/mjx-static-body-acc-sensors
Open

dundysm wants to merge 1 commit into
google-deepmind:mainfrom
dundysm:fix/mjx-static-body-acc-sensors

Conversation

@dundysm

@dundysm dundysm commented Oct 20, 2026

Copy link
Copy Markdown

fixes #3716

native mj_objectAcceleration returns zero when the body is welded to a dof-less body. mjx sensor_acc read cacc directly, so accelerometer, framelinacc and frameangacc on a world, static or mocap site reported -gravity (e.g. menagerie barkour_v0 origin accelerometer). this masks those rows with body_dofnum[body_weldid[bodyid]] == 0, matching c.

tests: the new sensor_test.test_acc_sensors_dofless_body compares against mujoco; it fails on main and passes here. pytest mjx/_src/sensor_test.py gives 13 passed.

the same fix for mujoco_warp is google-deepmind/mujoco_warp#1805 (issue google-deepmind/mujoco_warp#1804).

@thowell could you take a look? cc @erikfrey

mj_objectAcceleration returns zero when the body is welded to a dof-less
body. mjx read cacc directly, so accelerometer, framelinacc and frameangacc
on a world, static or mocap site reported -gravity.

Signed-off-by: Dundy Pasupuleti <dundysm@gmail.com>
@manfromnowhere143

Copy link
Copy Markdown

I checked e1f96e8545260920a3c17313daa1d627720aeb62 against native MuJoCo, including a jointless child attached to a moving joint. The weld-root predicate matters: your sensor_acc implementation matches all 408 named-sensor comparisons in both float32 and float64 CPU runs; an intentionally incorrect local-body-only mask loses all 132 moving-link readings. These are 12 constructed states across 34 sensors, not independent statistical trials. Your unchanged new regression also fails the parent and passes the candidate.

There is an adjacent, inherited full-forward gap worth covering separately. In forward.py, the efc_J.size == 0 branch returns before sensor_acc. This same code is present at current base efdd8cbc2745bb69443167e2ea9b29bec480f907; it is not introduced by your mask.

With collisions disabled and no other constraints, acceleration sensors retain their input buffer. Importing already-computed native sensor values can conceal this path, so the regression below initializes them to a sentinel. It fails on the current candidate and passes with the one-line change below.

 if d._impl.efc_J.size == 0:
   d = d.replace(qacc=d.qacc_smooth)
-  return d
+  return sensor.sensor_acc(m, d)

The constrained path still computes sensors after the solver, preserving the ordering established in #2767. With this local change, all 408 full-forward comparisons pass in each CPU precision, the existing 13 sensor tests pass, and ForwardTest.test_forward passes.

Standalone no-constraint regression
import jax
import jax.numpy as jp
import mujoco
from mujoco import mjx
import numpy as np

m = mujoco.MjModel.from_xml_string("""
<mujoco>
  <default><geom contype="0" conaffinity="0"/></default>
  <worldbody>
    <site name="world"/>
    <body>
      <joint type="hinge" axis="0 1 0"/>
      <geom size=".1" mass="1" pos=".2 0 0"/>
      <body pos=".3 0 0">
        <geom size=".05" mass=".2"/>
        <site name="moving"/>
      </body>
    </body>
  </worldbody>
  <sensor>
    <accelerometer site="world"/>
    <accelerometer site="moving"/>
    <framelinacc objtype="site" objname="moving"/>
    <frameangacc objtype="site" objname="moving"/>
  </sensor>
</mujoco>
""")
d = mujoco.MjData(m)
d.qpos[0], d.qvel[0], d.qfrc_applied[0] = 0.4, 1.2, 0.7
mujoco.mj_forward(m, d)
mx, dx = mjx.put_model(m), mjx.put_data(m, d)
assert dx._impl.efc_J.size == 0
dx = dx.replace(sensordata=jp.full_like(dx.sensordata, 1234.5))
actual = jax.jit(mjx.forward)(mx, dx).sensordata
np.testing.assert_allclose(actual, d.sensordata, rtol=5e-5, atol=5e-5)

Checked on macOS arm64, Python 3.12.13, native MuJoCo 3.17.0, JAX 0.12.0.1 and NumPy 2.5.4, with the complete MJX source pinned to your revision. The broader sensor-stage comparison clears the derived acceleration/force buffers as well as sensor output. This is CPU software evidence; GPU, Warp and physical sensors were not exercised. I retained the original failures, the moving-link negative control and the corrected runs as a scoped Nisayon qualification case.

@dundysm

dundysm commented Oct 28, 2026

Copy link
Copy Markdown
Author

thanks @manfromnowhere143 for checking this against native mujoco; good to see all 408 sensor comparisons match in float32 and float64, and that the regression fails on the parent. the efc_J.size == 0 early return in forward.py is a real gap, but it predates this pr and sits outside the mask fix, so i would rather keep it out of scope here to keep this one easy to review. happy to open a follow-up issue and pr with your one-line change and sentinel regression once this lands, if the maintainers want it.

@manfromnowhere143

Copy link
Copy Markdown

Agreed. The no-constraint forward-pass gap should stay separate from this mask fix. Your current head is still e1f96e8, so my earlier mask comparisons remain tied to that revision.

For the follow-up, the sentinel is the useful guard: mjx.put_data can import already-computed native sensor values, allowing a comparison to pass even when the MJX sensor stage never runs. The regression above forces the unconstrained forward path to produce fresh output. The broader retained probe also clears the derived acceleration/force buffers and checks the moving weld-root case and constrained path.

The reproducer and checked one-line correction are available for that separate follow-up when the maintainers are ready. Thank you for keeping the scope clear.

@dundysm

dundysm commented Nov 11, 2026

Copy link
Copy Markdown
Author

@btaba could you take a look when you get a chance? mjx now zeros acceleration sensors on static and mocap bodies; manfromnowhere143 checked all 408 sensor comparisons against native mujoco.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[MJX] accelerometer, framelinacc and frameangacc report -gravity on static and mocap bodies

2 participants