[dundysm]

contact sensors that request linvel or angvel raised NotImplementedError (the TODO(team) in io.py). this computes them in the velocity stage, matching mujoco's copySensorData for geom contacts: the velocity of geom[1] relative to geom[0] at the contact point, from cvel shifted from each body's subtree com, rotated into the contact frame, with z flipped when the match direction is reversed.

• adds sensor_contact_vel_adr and runs _eval_sensor_contact in sensor_vel. contact sensors are excluded from sensor_vel_adr
• the tangent slot now advances the write offset so later fields land correctly
• models with flex still raise for these fields, since c uses the contact jacobian there (left as a todo)

tests: test_contact_sensor_velocity checks against mujoco c across data combos, obj/ref specs including flipped direction, mindist/maxforce/default reduce, and 1 and 2 worlds. full suite passes on cpu (2190 passed).

@thowell would appreciate a look when you have time.

[dundysm]

@thowell when you have a moment, could you take a look? this adds linvel and angvel to contact sensors, computed in the velocity stage to match mujoco c copySensorData, and removes the todo NotImplementedError. i also applied the patterns from #1805 and #1809 here: fewer global reads, sensordata pre-filled with wp.inf in the test, and test_data.fixture.

[thowell]

Thanks @dundysm! Nice work adding linvel and angvel contact sensor support and fixing the tangent slot offset -- a few recommendations below before merging.

1. mujoco_warp/_src/sensor.py (lines 1575--1578, 2434--2443, 2769--2799)

• Issue: In _sensor_contact, body_rootid and geom_bodyid are placed after sensor_adr under # Model:, which fails kernel_analyzer because body_rootid and geom_bodyid are declared before sensor_intprm, sensor_dim, and sensor_adr on types.Model. In sensor_vel, _eval_sensor_contact was also inserted below the # apply sensor delay/interval for velocity sensors comment instead of above it.
• Recommendation: Place body_rootid and geom_bodyid before sensor_intprm in _sensor_contact and at the wp.launch(_sensor_contact, ...) call site, and move _eval_sensor_contact(m, d, m.sensor_contact_vel_adr) above the delay comment in sensor_vel.

  _eval_sensor_contact(m, d, m.sensor_contact_vel_adr)

  # apply sensor delay/interval for velocity sensors
  history.apply_sensor_delay(m, d, m.sensor_stage_vel_adr)

@wp.kernel
def _sensor_contact(
  # Model:
  opt_cone: int,
  opt_contact_sensor_maxmatch: int,
  body_rootid: wp.array[int],
  geom_bodyid: wp.array[int],
  sensor_intprm: wp.array2d[int],
  sensor_dim: wp.array[int],
  sensor_adr: wp.array[int],
  # Data in:

2. mujoco_warp/_src/types.py, mujoco_warp/__init__.py, mujoco_warp/_src/io.py (lines 364--366), and mujoco_warp/_src/sensor.py (lines 2473--2515)

• Issue: io.py ((1 << 7) | (1 << 8)) and _sensor_contact (for i in range(9): ...) use raw integer bit literals instead of wrapping mujoco.mjtConDataField.
• Recommendation: Wrap mujoco.mjtConDataField as ConDataField(enum.IntEnum) in mujoco_warp/_src/types.py (and re-export in mujoco_warp/__init__.py, matching RayDataField on main), use types.ConDataField in io.py, and unpack dataspec directly with ConDataField in _sensor_contact. (Rebasing onto main is also recommended.)

# mujoco_warp/_src/types.py
class ConDataField(enum.IntEnum):
  """Contact data fields for contact sensors."""

  FOUND = mujoco.mjtConDataField.mjCONDATA_FOUND
  FORCE = mujoco.mjtConDataField.mjCONDATA_FORCE
  TORQUE = mujoco.mjtConDataField.mjCONDATA_TORQUE
  DIST = mujoco.mjtConDataField.mjCONDATA_DIST
  POS = mujoco.mjtConDataField.mjCONDATA_POS
  NORMAL = mujoco.mjtConDataField.mjCONDATA_NORMAL
  TANGENT = mujoco.mjtConDataField.mjCONDATA_TANGENT
  LINVEL = mujoco.mjtConDataField.mjCONDATA_LINVEL
  ANGVEL = mujoco.mjtConDataField.mjCONDATA_ANGVEL

# mujoco_warp/_src/io.py
  if mjm.nflex and (contact_dataspec & ((1 << types.ConDataField.LINVEL) | (1 << types.ConDataField.ANGVEL))).any():
    raise NotImplementedError("Contact sensor linvel and angvel not supported for models with flexes.")

# mujoco_warp/_src/sensor.py
  found = bool(dataspec & (1 << ConDataField.FOUND))
  force = bool(dataspec & (1 << ConDataField.FORCE))
  torque = bool(dataspec & (1 << ConDataField.TORQUE))
  dist = bool(dataspec & (1 << ConDataField.DIST))
  pos = bool(dataspec & (1 << ConDataField.POS))
  normal = bool(dataspec & (1 << ConDataField.NORMAL))
  tangent = bool(dataspec & (1 << ConDataField.TANGENT))
  linvel = bool(dataspec & (1 << ConDataField.LINVEL))
  angvel = bool(dataspec & (1 << ConDataField.ANGVEL))

  size = (
    (1 if found else 0)
    + (3 if force else 0)
    + (3 if torque else 0)
    + (1 if dist else 0)
    + (3 if pos else 0)
    + (3 if normal else 0)
    + (3 if tangent else 0)
    + (3 if linvel else 0)
    + (3 if angvel else 0)
  )

3. mujoco_warp/_src/io_test.py (lines 1532--1549)

• Issue: test_contact_sensor_vel_flex_unsupported manually constructs MjModel with data="found" and patches mjm.sensor_intprm[0, 0] |= 1 << bit, which is no longer needed now that linvel and angvel are supported in MuJoCo XML.
• Recommendation: Parameterize over ("linvel", "angvel") and use test_data.fixture(xml=...) with <contact data="{data}"/>.

  @parameterized.parameters("linvel", "angvel")
  def test_contact_sensor_vel_flex_unsupported(self, data):
    """Contact sensor linvel or angvel with flex raises NotImplementedError."""
    with self.assertRaises(NotImplementedError):
      test_data.fixture(
        xml=f"""
        <mujoco>
          <worldbody>
            <geom type="plane" size="1 1 1"/>
            <flexcomp name="rope" type="grid" count="3 1 1" spacing=".1 .1 .1" pos="0 0 .5" dim="1" radius=".01"/>
          </worldbody>
          <sensor>
            <contact data="{data}"/>
          </sensor>
        </mujoco>
        """
      )

4. mujoco_warp/_src/sensor_test.py (lines 732--797)

• Issue: test_contact_sensor_velocity compiles and runs mjw.forward 96 separate times with a single <contact/> sensor at sensor_adr == 0, calls sensor_pos, sensor_vel, and sensor_acc together without verifying stage isolation, omits self.assertFalse(np.allclose(...)) between worlds when nworld == 2, and specifies default type="sphere".
• Recommendation: Combine the <contact/> sensors into a single multi-sensor fixture parameterized over @parameterized.parameters(1, 2), verify mjw.sensor_vel and mjw.sensor_acc stage isolation separately, assert distinct world outputs when nworld == 2, and drop default type="sphere".

  @parameterized.parameters(1, 2)
  def test_contact_sensor_velocity(self, nworld):
    """Test contact sensor linvel and angvel against MuJoCo."""
    mjm, mjd, m, d = test_data.fixture(
      xml="""
      <mujoco>
        <option cone="elliptic"/>
        <worldbody>
          <body name="plane">
            <geom name="plane" type="plane" size="10 10 .001"/>
          </body>
          <body name="ball" pos="0 0 .09">
            <freejoint/>
            <geom name="ball" size=".1"/>
          </body>
          <body name="arm" pos=".5 0 .3">
            <joint axis="0 1 0"/>
            <geom type="capsule" size=".02" fromto="0 0 0 0 0 -.1" contype="0" conaffinity="0"/>
            <body name="box" pos="0 0 -.21">
              <joint type="slide" axis="1 0 0"/>
              <geom name="box" type="box" size=".1 .08 .1" euler="3 4 5"/>
            </body>
          </body>
        </worldbody>
        <sensor>
          <contact geom1="plane" num="4" data="linvel"/>
          <contact geom1="ball" geom2="plane" num="4" data="angvel"/>
          <contact body1="box" num="4" data="pos normal tangent linvel angvel" reduce="mindist"/>
          <contact body2="ball" num="4" data="linvel angvel" reduce="mindist"/>
          <contact geom1="plane" num="4" data="found force dist linvel"/>
          <contact body1="box" num="4" data="pos normal tangent linvel angvel" reduce="maxforce"/>
          <contact body2="ball" num="4" data="found force dist linvel" reduce="maxforce"/>
        </sensor>
        <keyframe>
          <key qpos="0 0 .09 1 0 0 0 0 0" qvel=".3 -.2 -.4 1.1 -.7 .5 .8 -.6"/>
        </keyframe>
      </mujoco>
      """,
      keyframe=0,
      nworld=nworld,
    )

    mjds = [mjd]
    if nworld == 2:
      mjd1 = mujoco.MjData(mjm)
      mujoco.mj_resetDataKeyframe(mjm, mjd1, 0)
      mjd1.qvel[:] = -1.3 * mjd1.qvel + 0.2
      mujoco.mj_forward(mjm, mjd1)
      qvel = d.qvel.numpy()
      qvel[1] = mjd1.qvel
      d.qvel.assign(qvel)
      mjw.forward(m, d)
      mjds.append(mjd1)

    vel_dim = int(mjm.sensor_adr[4])

    # VEL stage populates only VEL-stage contact sensors
    d.sensordata.fill_(wp.inf)
    mjw.sensor_vel(m, d)
    sensordata_vel = d.sensordata.numpy()
    for w in range(nworld):
      _assert_eq(sensordata_vel[w, :vel_dim], mjds[w].sensordata[:vel_dim], f"vel_sensordata_world_{w}")
      self.assertTrue(np.abs(sensordata_vel[w, :vel_dim]).sum() > 0.0)
      self.assertTrue(np.isinf(sensordata_vel[w, vel_dim:]).all())

    if nworld == 2:
      self.assertFalse(np.allclose(sensordata_vel[0, :vel_dim], sensordata_vel[1, :vel_dim]))

    # ACC stage populates only ACC-stage contact sensors
    d.sensordata.fill_(wp.inf)
    mjw.sensor_acc(m, d)
    sensordata_acc = d.sensordata.numpy()
    for w in range(nworld):
      self.assertTrue(np.isinf(sensordata_acc[w, :vel_dim]).all())
      _assert_eq(sensordata_acc[w, vel_dim:], mjds[w].sensordata[vel_dim:], f"acc_sensordata_world_{w}")
      self.assertTrue(np.abs(sensordata_acc[w, vel_dim:]).sum() > 0.0)

    if nworld == 2:
      self.assertFalse(np.allclose(sensordata_acc[0, vel_dim:], sensordata_acc[1, vel_dim:]))

[dundysm]

thanks @thowell; merged main (no rebase, to avoid a force-push) and addressed everything in 4b4b0e6:

1. body_rootid and geom_bodyid now come before sensor_intprm in _sensor_contact and at the launch site; _eval_sensor_contact runs above the velocity delay comment. kernel_analyzer is clean.
2. added ConDataField(enum.IntEnum) wrapping mujoco.mjtConDataField in types.py, re-exported in __init__.py like RayDataField; io.py and _sensor_contact use it, and _sensor_contact unpacks dataspec directly instead of the range(9) loop.
3. test_contact_sensor_vel_flex_unsupported is parameterized over linvel/angvel with the xml fixture.
4. test_contact_sensor_velocity is now your consolidated multi-sensor model over nworld 1 and 2, with separate vel/acc stage isolation checks and distinct world asserts.

sensor, io and types tests pass on cpu (248 passed, 4 skipped); ruff is clean.

[thowell]

thanks @dundysm!