Add touch sensor to MJX.
PiperOrigin-RevId: 688056696 Change-Id: I50b0ebf3141aab0a5b07e5ae770101ff8b7737d4
This commit is contained in:
Taylor Howell
committed by
Copybara-Service
Copybara-Service
parent
680fb3e5dd
commit
b00a7c6705
+2
-1
@@ -8,7 +8,8 @@ Upcoming version (not yet released)
|
||||
MJX
|
||||
^^^
|
||||
|
||||
1. Added ``apply_ft``, ``jac``, and ``xfrc_accumulate`` as public functions.
|
||||
- Added ``apply_ft``, ``jac``, and ``xfrc_accumulate`` as public functions.
|
||||
- Added ``TOUCH`` sensor.
|
||||
|
||||
Version 3.2.4 (Oct 15, 2024)
|
||||
----------------------------
|
||||
|
||||
+2
-2
@@ -220,8 +220,8 @@ The following features are **fully supported** in MJX:
|
||||
- ``MAGNETOMETER``, ``CAMPROJECTION``, ``RANGEFINDER``, ``JOINTPOS``, ``TENDONPOS``, ``ACTUATORPOS``, ``BALLQUAT``,
|
||||
``FRAMEPOS``, ``FRAMEXAXIS``, ``FRAMEYAXIS``, ``FRAMEZAXIS``, ``FRAMEQUAT``, ``SUBTREECOM``, ``CLOCK``,
|
||||
``VELOCIMETER``, ``GYRO``, ``JOINTVEL``, ``TENDONVEL``, ``ACTUATORVEL``, ``BALLANGVEL``, ``FRAMELINVEL``,
|
||||
``FRAMEANGVEL``, ``SUBTREELINVEL``, ``SUBTREEANGMOM``, ``ACCELEROMETER``, ``FORCE``, ``TORQUE``, ``ACTUATORFRC``,
|
||||
``JOINTACTFRC``, ``FRAMELINACC``, ``FRAMEANGACC``.
|
||||
``FRAMEANGVEL``, ``SUBTREELINVEL``, ``SUBTREEANGMOM``, ``TOUCH``, ``ACCELEROMETER``, ``FORCE``, ``TORQUE``,
|
||||
``ACTUATORFRC``, ``JOINTACTFRC``, ``FRAMELINACC``, ``FRAMEANGACC``.
|
||||
|
||||
The following features are **in development** and coming soon:
|
||||
|
||||
|
||||
@@ -126,6 +126,7 @@ def _ray_box(
|
||||
p1 = pnt[iface[:, 1]] + x * vec[iface[:, 1]]
|
||||
valid = jp.abs(p0) <= size[iface[:, 0]]
|
||||
valid &= jp.abs(p1) <= size[iface[:, 1]]
|
||||
valid &= x >= 0
|
||||
|
||||
return jp.min(jp.where(valid, x, jp.inf))
|
||||
|
||||
@@ -268,3 +269,20 @@ def ray(
|
||||
id_ = jp.where(jp.isinf(dists[min_id]), -1, ids[min_id])
|
||||
|
||||
return dist, id_
|
||||
|
||||
|
||||
def ray_geom(
|
||||
size: jax.Array, pnt: jax.Array, vec: jax.Array, geomtype: GeomType
|
||||
) -> jax.Array:
|
||||
"""Returns the distance at which a ray intersects with a primitive geom.
|
||||
|
||||
Args:
|
||||
size: geom size (1,), (2,), or (3,)
|
||||
pnt: ray origin point (3,)
|
||||
vec: ray direction (3,)
|
||||
geomtype: type of geom
|
||||
|
||||
Returns:
|
||||
dist: distance from ray origin to geom surface
|
||||
"""
|
||||
return _RAY_FUNC[geomtype](size, pnt, vec)
|
||||
|
||||
@@ -21,6 +21,7 @@ import mujoco
|
||||
from mujoco.mjx._src import math
|
||||
from mujoco.mjx._src import ray
|
||||
from mujoco.mjx._src import smooth
|
||||
from mujoco.mjx._src import support
|
||||
from mujoco.mjx._src.types import Data
|
||||
from mujoco.mjx._src.types import DisableBit
|
||||
from mujoco.mjx._src.types import Model
|
||||
@@ -452,7 +453,65 @@ def sensor_acc(m: Model, d: Data) -> Data:
|
||||
cutoff = m.sensor_cutoff[idx]
|
||||
data_type = m.sensor_datatype[idx]
|
||||
|
||||
if sensor_type == SensorType.ACCELEROMETER:
|
||||
if sensor_type == SensorType.TOUCH:
|
||||
# compute contact forces
|
||||
forces = []
|
||||
condim_ids = []
|
||||
for dim in set(d.contact.dim):
|
||||
force, condim_id = support.contact_force_dim(m, d, dim)
|
||||
forces.append(force)
|
||||
condim_ids.append(condim_id)
|
||||
forces = jp.concatenate(forces)[jp.concatenate(condim_ids)]
|
||||
|
||||
# get bodies of contact geoms
|
||||
conbody = jp.array(m.geom_bodyid)[d.contact.geom]
|
||||
|
||||
# get site information
|
||||
site_bodyid = m.site_bodyid[objid]
|
||||
site_size = m.site_size[objid]
|
||||
site_xpos = d.site_xpos[objid]
|
||||
site_xmat = d.site_xmat[objid]
|
||||
site_type = m.site_type[objid]
|
||||
conbody0 = site_bodyid[:, None] == conbody[:, 0]
|
||||
conbody1 = site_bodyid[:, None] == conbody[:, 1]
|
||||
contacts = (d.contact.efc_address >= 0)[None] & (conbody0 | conbody1)
|
||||
|
||||
# compute conray, flip if second body
|
||||
conray = jax.vmap(
|
||||
lambda frame, force: math.normalize(frame[0] * force[0])
|
||||
)(d.contact.frame, forces)
|
||||
conray = jp.where(conbody1[..., None], -conray, conray)
|
||||
|
||||
# compute distance, mapping over sites and contacts
|
||||
def _distance(
|
||||
site_size, site_xpos, site_xmat, site_type, contact_pos, conray
|
||||
):
|
||||
return jax.vmap(
|
||||
lambda site_size, site_xpos, site_xmat, conray: jax.vmap(
|
||||
lambda pnt, vec: ray.ray_geom(site_size, pnt, vec, site_type)
|
||||
)((contact_pos - site_xpos) @ site_xmat, conray @ site_xmat)
|
||||
)(site_size, site_xpos, site_xmat, conray)
|
||||
|
||||
dist = []
|
||||
dist_id = []
|
||||
for st in set(site_type):
|
||||
(dist_id_site,) = np.nonzero(st == site_type)
|
||||
dist_site = _distance(
|
||||
site_size[dist_id_site],
|
||||
site_xpos[dist_id_site],
|
||||
site_xmat[dist_id_site],
|
||||
st,
|
||||
d.contact.pos,
|
||||
conray[dist_id_site],
|
||||
)
|
||||
dist.append(jp.where(jp.isinf(dist_site), 0, dist_site))
|
||||
dist_id.append(dist_id_site)
|
||||
|
||||
dist = jp.vstack(dist)[np.concatenate(dist_id)]
|
||||
|
||||
# accumulate normal forces for each site
|
||||
sensor = jp.dot((dist > 0) & contacts, forces[:, 0])
|
||||
elif sensor_type == SensorType.ACCELEROMETER:
|
||||
|
||||
@jax.vmap
|
||||
def _accelerometer(cvel, cacc, diff, rot):
|
||||
|
||||
@@ -331,6 +331,7 @@ class SensorType(enum.IntEnum):
|
||||
FRAMEANGVEL: 3D angular velocity
|
||||
SUBTREELINVEL: subtree linear velocity
|
||||
SUBTREEANGMOM: subtree angular momentum
|
||||
TOUCH: scalar contact normal forces summed over the sensor zone
|
||||
ACCELEROMETER: accelerometer
|
||||
FORCE: force
|
||||
TORQUE: torque
|
||||
@@ -363,6 +364,7 @@ class SensorType(enum.IntEnum):
|
||||
FRAMEANGVEL = mujoco.mjtSensor.mjSENS_FRAMEANGVEL
|
||||
SUBTREELINVEL = mujoco.mjtSensor.mjSENS_SUBTREELINVEL
|
||||
SUBTREEANGMOM = mujoco.mjtSensor.mjSENS_SUBTREEANGMOM
|
||||
TOUCH = mujoco.mjtSensor.mjSENS_TOUCH
|
||||
ACCELEROMETER = mujoco.mjtSensor.mjSENS_ACCELEROMETER
|
||||
FORCE = mujoco.mjtSensor.mjSENS_FORCE
|
||||
TORQUE = mujoco.mjtSensor.mjSENS_TORQUE
|
||||
|
||||
@@ -27,6 +27,7 @@
|
||||
-subtreelinvel
|
||||
-subtreeangmom
|
||||
* acceleration/force-dependent sensors:
|
||||
-touch
|
||||
-accelerometer
|
||||
-force
|
||||
-torque
|
||||
@@ -100,6 +101,17 @@
|
||||
<site name="site_torque"/>
|
||||
<geom size="0.1"/>
|
||||
</body>
|
||||
|
||||
<!-- plane and body for touch sensors -->
|
||||
<geom type="plane" size="1 1 .1" pos="-20 -20 -20"/>
|
||||
<body pos="-20 -20 -19.985">
|
||||
<joint type="slide" axis="0 0 1"/>
|
||||
<joint type="hinge" axis="0 1 0"/>
|
||||
<geom type="capsule" fromto="-.5 0 0 .5 0 0" size="0.0125"/>
|
||||
<site name="touch_sphere" type="sphere" pos="-0.5 0 0" size="0.025"/>
|
||||
<site name="touch_capsule" type="capsule" fromto="-0.1 0 0.01 0.1 0 0.01" size="0.025"/>
|
||||
<site name="touch_box" pos="0.5 0 0" type="box" size="0.025 0.025 0.025"/>
|
||||
</body>
|
||||
</worldbody>
|
||||
|
||||
<actuator>
|
||||
@@ -121,6 +133,7 @@
|
||||
</tendon>
|
||||
|
||||
<sensor>
|
||||
<touch name="touch_sphere" site="touch_sphere"/>
|
||||
<torque site="site_torque"/>
|
||||
<tendonpos tendon="fixed"/>
|
||||
<tendonvel tendon="fixed"/>
|
||||
@@ -128,6 +141,7 @@
|
||||
<velocimeter name="velocimeter0" site="site0"/>
|
||||
<velocimeter name="velocimeter0cutoff" site="site0" cutoff="3e-4"/>
|
||||
<gyro name="gyro0" site="site0"/>
|
||||
<touch name="touch_box" site="touch_box"/>
|
||||
<gyro name="gyro0cutoff" site="site0" cutoff="2e-3"/>
|
||||
<rangefinder name="rangefinder0" site="site_rangefinder0"/>
|
||||
<accelerometer name="accelerometer1" site="site1"/>
|
||||
@@ -138,6 +152,7 @@
|
||||
<jointpos name="jointpos0cutoff" joint="hinge0" cutoff="1e-4"/>
|
||||
<jointvel name="jointvel0" joint="hinge0"/>
|
||||
<jointvel name="jointvel0cutoff" joint="hinge0" cutoff="1e-3"/>
|
||||
<touch name="touch_capsule" site="touch_capsule"/>
|
||||
<actuatorfrc name="actuatorfrc0" actuator="motor0"/>
|
||||
<framelinvel name="framelinvel3" objtype="site" objname="site3"/>
|
||||
<framelinvel name="framelinvel3ref" objtype="site" objname="site3" reftype="body" refname="body2"/>
|
||||
|
||||
Reference in New Issue
Block a user