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

Fix bug in `mjx.ray` and speed up ray-mesh by ~30% for large meshes.

Browse Source 
PiperOrigin-RevId: 629209303
Change-Id: I39466e9a2dc55ec48dcca1780e8abef903d1c754
This commit is contained in:
Erik Frey
2024-04-29 15:46:44 -07:00
committed by Copybara-Service
parent 618d879904
commit 719476c274
3 changed files with 32 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/ray.py
+19 -28
View File
@@ -134,32 +134,27 @@ def _ray_triangle(
vert: jax.Array,
pnt: jax.Array,
vec: jax.Array,
b0: jax.Array,
b1: jax.Array,
basis: jax.Array,
) -> jax.Array:
"""Returns the distance at which a ray intersects with a triangle."""
# project difference vectors in ray normal plane
planar = jp.dot(jp.array([b0, b1]), (vert - pnt).T)
planar = jp.dot(vert - pnt, basis)
# determine if origin is inside planar projection of triangle
# A = (p0-p2, p1-p2), b = -p2, solve A*t = b
A = jp.array( # pylint: disable=invalid-name
[planar[:, 0] - planar[:, 2], planar[:, 1] - planar[:, 2]]
).T.flatten()
b = -planar[:, 2]
det = A[0] * A[3] - A[1] * A[2]
valid = jp.abs(det) >= mujoco.mjMINVAL
A = planar[0:2] - planar[2] # pylint: disable=invalid-name
b = -planar[2]
det = A[0, 0] * A[1, 1] - A[1, 0] * A[0, 1]
t0 = (A[3] * b[0] - A[1] * b[1]) / det
t1 = (-A[2] * b[0] + A[0] * b[1]) / det
valid &= (t0 >= 0) & (t1 >= 0) & (t0 + t1 <= 1)
t0 = (A[1, 1] * b[0] - A[1, 0] * b[1]) / det
t1 = (-A[0, 1] * b[0] + A[0, 0] * b[1]) / det
valid = (t0 >= 0) & (t1 >= 0) & (t0 + t1 <= 1)
# intersect ray with plane of triangle
nrm = jp.cross(vert[0] - vert[2], vert[1] - vert[2])
denom = jp.dot(vec, nrm)
valid &= jp.abs(denom) >= mujoco.mjMINVAL
dist = jp.where(valid, -jp.dot(pnt - vert[2], nrm) / denom, jp.inf)
dist = jp.dot(vert[2] - pnt, nrm) / jp.dot(vec, nrm)
valid &= dist >= 0
dist = jp.where(valid, dist, jp.inf)
return dist
@@ -174,33 +169,29 @@ def _ray_mesh(
"""Returns the best distance and geom_id for ray mesh intersections."""
data_id = m.geom_dataid[geom_id]
ray_basis = lambda x: math.orthogonals(math.normalize(x))
b0, b1 = jax.vmap(ray_basis)(vec)
ray_basis = lambda x: jp.array(math.orthogonals(math.normalize(x))).T
basis = jax.vmap(ray_basis)(vec)
faceadr = np.append(m.mesh_faceadr, m.nmeshface)
vertadr = np.append(m.mesh_vertadr, m.nmeshvert)
dists = []
dists, geom_ids = [], []
for i, id_ in enumerate(data_id):
face = m.mesh_face[faceadr[id_] : faceadr[id_ + 1]]
vert = m.mesh_vert[vertadr[id_] : vertadr[id_ + 1]]
dist = jax.vmap(_ray_triangle, in_axes=(0, None, None, None, None))(
vert[face], pnt[i], vec[i], b0[i], b1[i]
vert = jp.array(vert[face])
dist = jax.vmap(_ray_triangle, in_axes=(0, None, None, None))(
vert, pnt[i], vec[i], basis[i]
)
dists.append(dist)
# map the triangle id to data id
tri_id = np.append(0, (faceadr[data_id + 1] - faceadr[data_id]).cumsum())
tri_data_id = np.zeros(tri_id[-1], dtype=np.int32)
tri_data_id[tri_id[:-1]] = 1
tri_data_id = tri_data_id.cumsum() - 1
geom_ids.append(np.repeat(geom_id[i], dist.size))
dists = jp.concatenate(dists)
min_id = jp.argmin(dists)
# Grab the best distance amongst all meshes, bypassing the argmin in `ray`.
# This avoids having to compute the best distance per mesh.
dist = dists[min_id, None]
id_ = jp.array(geom_id)[jp.array(tri_data_id)[min_id], None]
id_ = jp.array(np.concatenate(geom_ids))[min_id, None]
return dist, id_
mjx/mujoco/mjx/_src/ray_test.py
+8 -10
View File
@@ -152,11 +152,7 @@ class RayTest(absltest.TestCase):
mx, dx = mjx.put_model(m), mjx.put_data(m, d)
# look at the tetrahedron
pnt, vec = jp.array([2.0, 2.0, 2.0]), -jp.array([
1.0,
1.0,
1.0,
])
pnt, vec = jp.array([2.0, 2.0, 2.0]), -jp.array([1.0, 1.0, 1.0])
vec /= jp.linalg.norm(vec)
dist, geomid = jax.jit(mjx.ray)(mx, dx, pnt, vec)
_assert_eq(geomid, 4, 'geom_id')
@@ -166,12 +162,14 @@ class RayTest(absltest.TestCase):
_assert_eq(geomid, 4, 'geom_id')
_assert_eq(dist, mj_dist, 'dist-tetrahedron')
# look away from the dodecahedron
pnt, vec = jp.array([4.0, 2.0, 2.0]), jp.array([2.0, 1.0, 1.0])
vec /= jp.linalg.norm(vec)
_, geomid = jax.jit(mjx.ray)(mx, dx, pnt, vec)
_assert_eq(geomid, -1, 'geom_id')
# look at the dodecahedron
pnt, vec = jp.array([4.0, 2.0, 2.0]), -jp.array([
2.0,
1.0,
1.0,
])
pnt, vec = jp.array([4.0, 2.0, 2.0]), -jp.array([2.0, 1.0, 1.0])
vec /= jp.linalg.norm(vec)
dist, geomid = jax.jit(mjx.ray)(mx, dx, pnt, vec)
_assert_eq(geomid, 5, 'geom_id')