[dundysm]

related to #1804

mujoco 4.0 (google-deepmind/mujoco@fdf1412) now reports gravity in accelerometer and framelinacc on static and mocap bodies, which is what warp already does on main, so the original sensor.py zeroing is no longer needed. per review, this pr now only adds a parity test.

tests: sensor_test.test_acc_sensors_dofless_body checks accelerometer, framelinacc and frameangacc sites on the world, a mocap body, a static body and a hinged body against mujoco 4.0.2, pre-filling sensordata with wp.inf and calling mjw.sensor_acc directly. it passes on main and fails with the old zeroing change applied. pytest mujoco_warp/_src/sensor_test.py gives 70 passed on cpu; ruff format clean.

[thowell]

Thanks @dundysm! Good catch fixing acceleration sensor parity for static and mocap bodies; a few minor recommendations below before merging.

1. mujoco_warp/_src/sensor.py (lines 1812--1841, 1933--1942)

• Issue: In smooth.rne_postconstraint (_rne_cacc_world and _cacc_branch), d.cacc[worldid, 0] is initialized to wp.spatial_vector(wp.vec3(0.0), -gravity), and any DOF-less body inherits d.cacc[worldid, 0] unchanged because body_dofnum == 0 along its ancestor chain. Since gravity is purely linear (wp.spatial_bottom), the angular component wp.spatial_top(cacc_in[worldid, bodyid]) is already wp.vec3(0.0) for all world, static, and mocap bodies. Adding the check to _frameangacc introduces unnecessary global memory reads on every frameangacc evaluation.
• Recommendation: Revert the changes to _frameangacc and its call site in _sensor_acc.

@wp.func
def _frameangacc(
  # Model:
  geom_bodyid: wp.array[int],
  site_bodyid: wp.array[int],
  cam_bodyid: wp.array[int],
  # Data in:
  cacc_in: wp.array2d[wp.spatial_vector],
  # In:
  worldid: int,
  objid: int,
  objtype: int,
) -> wp.vec3:
  if objtype == ObjType.BODY or objtype == ObjType.XBODY:
    bodyid = objid
  elif objtype == ObjType.GEOM:
    bodyid = geom_bodyid[objid]
  elif objtype == ObjType.SITE:
    bodyid = site_bodyid[objid]
  elif objtype == ObjType.CAMERA:
    bodyid = cam_bodyid[objid]
  else:  # UNKNOWN
    bodyid = 0

  return wp.spatial_top(cacc_in[worldid, bodyid])

2. mujoco_warp/_src/sensor.py (lines 1586--1602, 1756--1800, 1848--1932, 2830--2832)

• Issue: Checking body_dofnum[body_weldid[bodyid]] == 0 performs two dependent global memory reads across two arrays (body_weldid then body_dofnum) and threads both arrays through _sensor_acc, _accelerometer, and _framelinacc. In MuJoCo, m.body_treeid[i] is precomputed as -1 if and only if m.body_dofnum[m.body_weldid[i]] == 0 (and >= 0 otherwise).
• Recommendation: Pass m.body_treeid instead of m.body_weldid and m.body_dofnum, and check if body_treeid[bodyid] < 0: in _accelerometer and _framelinacc.

@wp.func
def _accelerometer(
  # Model:
  body_rootid: wp.array[int],
  body_treeid: wp.array[int],
  site_bodyid: wp.array[int],
  # Data in:
  site_xpos_in: wp.array2d[wp.vec3],
  site_xmat_in: wp.array2d[wp.mat33],
  subtree_com_in: wp.array2d[wp.vec3],
  cvel_in: wp.array2d[wp.spatial_vector],
  cacc_in: wp.array2d[wp.spatial_vector],
  # In:
  worldid: int,
  objid: int,
) -> wp.vec3:
  bodyid = site_bodyid[objid]
  if body_treeid[bodyid] < 0:
    return wp.vec3(0.0)
  rot = site_xmat_in[worldid, objid]

3. mujoco_warp/_src/sensor_test.py (lines 386--425)

• Issue: d.sensordata.zero_() pre-fills d.sensordata with 0.0, which is the exact expected output for the DOF-less acceleration sensors; tested float fields should be pre-filled with wp.inf. In addition, calling mjw.sensor_acc(m, d) directly tests the acceleration sensor stage without re-running mjw.forward(m, d), default type="hinge" should be omitted, and the <mujoco> / </mujoco> tags and closing """ delimiter should align with xml=""".
• Recommendation: Use d.sensordata.fill_(wp.inf), call mjw.sensor_acc(m, d), omit default type="hinge", and align the inline XML indentation.

  def test_acc_sensors_dofless_body(self):
    """Test acceleration sensors on static and mocap bodies are zero."""
    _, mjd, m, d = test_data.fixture(
      xml="""
      <mujoco>
        <worldbody>
          <site name="world" pos=".1 0 0"/>
          <body name="mocap" mocap="true" pos="0 0 1">
            <site name="mocap"/>
          </body>
          <body name="static" pos="0 1 1">
            <geom size=".1"/>
            <site name="static"/>
          </body>
          <body pos="1 0 1">
            <joint axis="0 1 0"/>
            <geom size=".1" pos=".2 0 0"/>
            <site name="hinge" pos=".2 0 0"/>
          </body>
        </worldbody>
        <sensor>
          <accelerometer site="world"/>
          <accelerometer site="mocap"/>
          <accelerometer site="static"/>
          <accelerometer site="hinge"/>
          <framelinacc objtype="site" objname="world"/>
          <framelinacc objtype="body" objname="mocap"/>
          <frameangacc objtype="xbody" objname="static"/>
          <framelinacc objtype="site" objname="hinge"/>
        </sensor>
      </mujoco>
      """
    )

    d.sensordata.fill_(wp.inf)

    mjw.sensor_acc(m, d)

    _assert_eq(d.sensordata.numpy()[0], mjd.sensordata, "sensordata")
    np.testing.assert_array_equal(d.sensordata.numpy()[0, :9], 0.0)

[dundysm]

thanks @thowell, all three are in 3741c69:

• reverted the _frameangacc change, you're right that the angular part of cacc is already zero for dof-less bodies
• _accelerometer and _framelinacc now check body_treeid[bodyid] < 0 instead of reading body_weldid and body_dofnum
• the test pre-fills sensordata with wp.inf, calls mjw.sensor_acc directly, drops the default type="hinge" and lines the xml up with xml="""

[thowell]

lgtm, thanks @dundysm!

[thowell]

@dundysm lets update the branch in order to resolve the ci failures. thanks!

[dundysm]

thanks @thowell; merged main into the branch in 2b826f0 (no rebase), and the checks are green now.

[dundysm]

@adenzler-nvidia this is approved by thowell and green after the main merge; it zeros acceleration sensors on static and mocap bodies. could you merge it when you get a chance?

[dundysm]

@thowell @adenzler-nvidia correction to my last comment: the stable-dependencies (3.14) failure does come from this pr. mujoco 4.0 (google-deepmind/mujoco@fdf1412) now reports gravity in accelerometer and framelinacc on static and mocap bodies, which is what warp already does on main, so the zeroing here now breaks parity (test_acc_sensors_dofless_body gets 0 where mujoco gives 9.81). the rest of the failures come from the branch being behind main.

so i think this pr is superseded. i'm fine closing it; or i can drop the sensor.py change and keep only the test as a parity check against the new behavior, if that's useful.

[thowell]

@dundysm lets refactor the pr to just include the test. thanks!

[dundysm]

thanks @thowell; done in 794da06. it merges main and reverts the sensor.py change, so the pr is now just test_acc_sensors_dofless_body as a parity check against mujoco 4.0 (gravity on the static and mocap accelerometer and framelinacc sites). it passes on main and fails with the old zeroing applied; sensor_test.py gives 70 passed on cpu. title and description are updated to match.

[thowell]

thanks @dundysm!