From 9805df616bd3f8190d80f32e8789ebe2a6215bf5 Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Sat, 31 Aug 2024 06:11:57 -0700 Subject: [PATCH] Add framelinvel and frameangvel sensors to MJX. PiperOrigin-RevId: 669647542 Change-Id: I612230e81b456c26f29173d1cb5a60975791875a --- doc/changelog.rst | 3 +- mjx/mujoco/mjx/_src/sensor.py | 74 ++++++++++++++++++++++++++++- mjx/mujoco/mjx/_src/types.py | 4 ++ mjx/mujoco/mjx/test_data/sensor.xml | 14 +++++- 4 files changed, 91 insertions(+), 4 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index a1529bb9..3b050c75 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -32,7 +32,8 @@ MJX - Added position-dependent sensors: ``MAGNETOMETER``, ``CAMPROJECTION``, ``RANGEFINDER``, ``JOINTPOS``, ``ACTUATORPOS``, ``BALLQUAT``, ``FRAMEPOS``, ``FRAMEXAXIS``, ``FRAMEYAXIS``, ``FRAMEZAXIS``, ``FRAMEQUAT``, ``SUBTREECOM``, ``CLOCK``. -- Added velocity-dependent sensors: ``VELOCIMETER``, ``GYRO``, ``JOINTVEL``, ``ACTUATORVEL``, ``BALLANGVEL``. +- Added velocity-dependent sensors: ``VELOCIMETER``, ``GYRO``, ``JOINTVEL``, ``ACTUATORVEL``, ``BALLANGVEL``, + ``FRAMELINVEL``, ``FRAMEANGVEL``. - Added acceleration/force-dependent sensors: ``ACTUATORFRC``, ``JOINTACTFRC``. - Changed default policy to avoid placing unused (MuJoCo-only) arrays on device. - Added ``device`` parameter to ``mjx.make_data`` to bring it to parity with ``mjx.put_model`` and ``mjx.put_data``. diff --git a/mjx/mujoco/mjx/_src/sensor.py b/mjx/mujoco/mjx/_src/sensor.py index e5cbf604..05bd721d 100644 --- a/mjx/mujoco/mjx/_src/sensor.py +++ b/mjx/mujoco/mjx/_src/sensor.py @@ -55,8 +55,8 @@ def sensor_pos(m: Model, d: Data) -> Data: # position and orientation by object type objtype_data = { ObjType.UNKNOWN: ( - np.expand_dims(np.eye(3), axis=0), np.zeros((1, 3)), + np.expand_dims(np.eye(3), axis=0), ), # world ObjType.BODY: (d.xipos, d.ximat), ObjType.XBODY: (d.xpos, d.xmat), @@ -281,6 +281,20 @@ def sensor_vel(m: Model, d: Data) -> Data: if m.opt.disableflags & DisableBit.SENSOR: return d + # position and orientation by object type + objtype_data = { + ObjType.UNKNOWN: ( + np.zeros((1, 3)), + np.expand_dims(np.eye(3), axis=0), + np.arange(1), + ), # world + ObjType.BODY: (d.xipos, d.ximat, np.arange(m.nbody)), + ObjType.XBODY: (d.xpos, d.xmat, np.arange(m.nbody)), + ObjType.GEOM: (d.geom_xpos, d.geom_xmat, m.geom_bodyid), + ObjType.SITE: (d.site_xpos, d.site_xmat, m.site_bodyid), + ObjType.CAMERA: (d.cam_xpos, d.cam_xmat, m.cam_bodyid), + } + stage_vel = m.sensor_needstage == mujoco.mjtStage.mjSTAGE_VEL sensors, adrs = [], [] @@ -315,9 +329,65 @@ def sensor_vel(m: Model, d: Data) -> Data: jnt_dotadr = m.jnt_dofadr[objid, None] + np.arange(3)[None] sensor = d.qvel[jnt_dotadr] adr = (adr[:, None] + np.arange(3)[None]).reshape(-1) + elif sensor_type in (SensorType.FRAMELINVEL, SensorType.FRAMEANGVEL): + objtype = m.sensor_objtype[idx] + reftype = m.sensor_reftype[idx] + refid = m.sensor_refid[idx] + + # evaluate for valid object and reference object type pairs + for ot, rt in set(zip(objtype, reftype)): + idxt = (objtype == ot) & (reftype == rt) + objidt = objid[idxt] + refidt = refid[idxt] + cutofft = cutoff[idxt] + + xpos, _, _ = objtype_data[ot] + xposref, xmatref, _ = objtype_data[rt] + xpos = xpos[objidt] + xposref = xposref[refidt] + xmatref = xmatref[refidt] + + def _cvel_offset(otype, oid): + pos, _, bodyid = objtype_data[otype] + pos = pos[oid] + bodyid = bodyid[oid] + return d.cvel[bodyid], pos - d.subtree_com[m.body_rootid[bodyid]] + + cvel, offset = _cvel_offset(ot, objidt) + cvelref, offsetref = _cvel_offset(rt, refidt) + cangvel = cvel[:, :3] + cangvelref = cvelref[:, :3] + + if sensor_type == SensorType.FRAMELINVEL: + clinvel = cvel[:, 3:] + clinvelref = cvelref[:, 3:] + xlinvel = clinvel - jp.cross(offset, cangvel) + xlinvelref = clinvelref - jp.cross(offsetref, cangvelref) + rvec = xpos - xposref + rel_vel = xlinvel - xlinvelref + jp.cross(rvec, cangvelref) + sensor = jp.where( + (refidt > -1)[:, None], + jax.vmap(lambda mat, vec: mat.T @ vec)(xmatref, rel_vel), + xlinvel, + ) + elif sensor_type == SensorType.FRAMEANGVEL: + rel_vel = cangvel - cangvelref + sensor = jp.where( + (refidt > -1)[:, None], + jax.vmap(lambda mat, vec: mat.T @ vec)(xmatref, rel_vel), + cangvel, + ) + else: + raise ValueError(f'Unknown sensor type: {sensor_type}') + + adrt = adr[idxt, None] + np.arange(3)[None] + + sensors.append(apply_cutoff(sensor, cutofft, data_type[0]).reshape(-1)) + adrs.append(adrt.reshape(-1)) + continue # avoid adding to sensors/adrs list a second time else: # TODO(taylorhowell): raise error after adding sensor check to io.py - continue # unsupported sensor typ + continue # unsupported sensor type sensors.append(apply_cutoff(sensor, cutoff, data_type[0]).reshape(-1)) adrs.append(adr) diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index 8f47d048..7dbe68dc 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -312,6 +312,8 @@ class SensorType(enum.IntEnum): JOINTVEL: joint velocity ACTUATORVEL: actuator velocity BALLANGVEL: ball joint angular velocity + FRAMELINVEL: 3D linear velocity + FRAMEANGVEL: 3D angular velocity ACTUATORFRC: scalar actuator force JOINTACTFRC: scalar actuator force, measured at the joint """ @@ -333,6 +335,8 @@ class SensorType(enum.IntEnum): JOINTVEL = mujoco.mjtSensor.mjSENS_JOINTVEL ACTUATORVEL = mujoco.mjtSensor.mjSENS_ACTUATORVEL BALLANGVEL = mujoco.mjtSensor.mjSENS_BALLANGVEL + FRAMELINVEL = mujoco.mjtSensor.mjSENS_FRAMELINVEL + FRAMEANGVEL = mujoco.mjtSensor.mjSENS_FRAMEANGVEL ACTUATORFRC = mujoco.mjtSensor.mjSENS_ACTUATORFRC JOINTACTFRC = mujoco.mjtSensor.mjSENS_JOINTACTFRC diff --git a/mjx/mujoco/mjx/test_data/sensor.xml b/mjx/mujoco/mjx/test_data/sensor.xml index d1417b31..350c7bd3 100644 --- a/mjx/mujoco/mjx/test_data/sensor.xml +++ b/mjx/mujoco/mjx/test_data/sensor.xml @@ -20,6 +20,8 @@ -jointvel -actuatorvel -ballangvel +-framelinvel +-frameangvel * acceleration/force-dependent sensors: -actuatorfrc -jointactfrc @@ -46,7 +48,7 @@ - + @@ -94,6 +96,8 @@ + + @@ -106,6 +110,9 @@ + + + @@ -121,10 +128,15 @@ + + + + +