chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions 4 Packages Projects Releases Wiki Activity

Fix MJX touch sensor.

Browse Source 
PiperOrigin-RevId: 690544118
Change-Id: Ic909c6f6ce0e31237cbe317fbbfbb1fd267bd462
This commit is contained in:
Taylor Howell
2024-10-28 03:38:52 -07:00
committed by Copybara-Service
parent 7dc8aef807
commit 1c424644dd
3 changed files with 28 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/sensor.py
+10 -11
View File
@@ -461,7 +461,7 @@ def sensor_acc(m: Model, d: Data) -> Data:
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)]
forces = jp.concatenate(forces)[np.argsort(np.concatenate(condim_ids))]
# get bodies of contact geoms
conbody = jp.array(m.geom_bodyid)[d.contact.geom]
@@ -483,14 +483,14 @@ def sensor_acc(m: Model, d: Data) -> Data:
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)
def _distance(site_size, site_xpos, site_xmat, site_type, pos, conray):
def dist(size, xpos, xmat, conray):
pnt = (pos - xpos) @ xmat
vec = conray @ xmat
ray_geom_ = lambda pnt, vec: ray.ray_geom(size, pnt, vec, site_type)
return jax.vmap(ray_geom_)(pnt, vec)
return jax.vmap(dist)(site_size, site_xpos, site_xmat, conray)
dist = []
dist_id = []
@@ -506,8 +506,7 @@ def sensor_acc(m: Model, d: Data) -> Data:
)
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)]
dist = jp.vstack(dist)[np.argsort(np.concatenate(dist_id))]
# accumulate normal forces for each site
sensor = jp.dot((dist > 0) & contacts, forces[:, 0])
mjx/mujoco/mjx/_src/sensor_test.py
-3
View File
@@ -101,17 +101,14 @@ class SensorTest(parameterized.TestCase):
<body>
<joint type="hinge"/>
<geom name="geom0" size="0.1"/>
<site name="site0"/>
<body>
<joint type="hinge"/>
<geom name="geom1" size="0.25"/>
<site name="site1"/>
</body>
</body>
</worldbody>
<sensor>
<distance name="distance" geom1="geom0" geom2="geom1"/>
<touch name="touch" site="site0"/>
</sensor>
</mujoco>
""")
mjx/mujoco/mjx/test_data/sensor/sensor.xml
+18 -9
View File
@@ -102,15 +102,22 @@
<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"/>
<!-- plane and bodies for touch sensors -->
<geom type="plane" size="2 2 .1" pos="-20 -20 -20"/>
<body pos="-20 -20 -20">
<freejoint/>
<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>
<body pos="-20 -20.25 -20">
<freejoint/>
<geom type="capsule" fromto="-.5 0 0 .5 0 0" size="0.0125"/>
<site name="touch_capsule" type="capsule" fromto="0.4 0 0.0 0.5 0 0.0" size="0.025"/>
</body>
<body pos="-20 -20.5 -20">
<freejoint/>
<geom type="capsule" fromto="-.5 0 0 .5 0 0" size="0.0125"/>
<site name="touch_box" pos="-0.5 0 0" type="box" size="0.025 0.025 0.025"/>
</body>
</worldbody>
@@ -141,7 +148,8 @@
<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"/>
<touch name="touch_box1" site="touch_box"/>
<touch name="touch_box2" site="touch_box"/>
<gyro name="gyro0cutoff" site="site0" cutoff="2e-3"/>
<rangefinder name="rangefinder0" site="site_rangefinder0"/>
<accelerometer name="accelerometer1" site="site1"/>
@@ -152,7 +160,8 @@
<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"/>
<touch name="touch_capsule1" site="touch_capsule"/>
<touch name="touch_capsule2" 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"/>