MJX ray for planes, spheres, capsules, and boxes.
PiperOrigin-RevId: 596004453 Change-Id: I6c15268f7ec6afc1aadd370a78ef225ab9c5fd04
This commit is contained in:
Erik Frey
committed by
Copybara-Service
Copybara-Service
parent
c59a33dd81
commit
feb92bf535
+3
-1
@@ -9,10 +9,12 @@ MJX
|
||||
^^^
|
||||
1. Added :ref:`dyntype<actuator-general-dyntype>` ``filterexact``.
|
||||
2. Added :at:`site` transmission.
|
||||
3. Updated MJX colab tutorial with more stable quadruped environment.
|
||||
4. Added ``mjx.ray`` which mirrors :ref:`mj_ray` for planes, spheres, capsules, and boxes.
|
||||
|
||||
Bug fixes
|
||||
^^^^^^^^^
|
||||
3. Fixed a bug that prevented the use of pins with plugins if flexes are not in the worldbody. Fixes :github:issue:`1270`.
|
||||
5. Fixed a bug that prevented the use of pins with plugins if flexes are not in the worldbody. Fixes :github:issue:`1270`.
|
||||
|
||||
|
||||
Version 3.1.1 (December 18, 2023)
|
||||
|
||||
@@ -33,6 +33,7 @@ from mujoco.mjx._src.io import make_data
|
||||
from mujoco.mjx._src.io import put_data
|
||||
from mujoco.mjx._src.io import put_model
|
||||
from mujoco.mjx._src.passive import passive
|
||||
from mujoco.mjx._src.ray import ray
|
||||
from mujoco.mjx._src.smooth import com_pos
|
||||
from mujoco.mjx._src.smooth import com_vel
|
||||
from mujoco.mjx._src.smooth import crb
|
||||
|
||||
@@ -333,7 +333,7 @@ def put_data(m: mujoco.MjModel, d: mujoco.MjData, device=None) -> types.Data:
|
||||
efc_j[i, d.efc_J_colind[rowadr + j]] = fields['efc_J'][rowadr + j]
|
||||
fields['efc_J'] = efc_j
|
||||
else:
|
||||
fields['efc_J'] = fields['efc_J'].reshape((-1, m.nv))
|
||||
fields['efc_J'] = fields['efc_J'].reshape((-1 if m.nv else 0, m.nv))
|
||||
|
||||
for fname in ('efc_J', 'efc_frictionloss', 'efc_D', 'efc_aref', 'efc_force'):
|
||||
value = np.zeros((nefc, m.nv)) if fname == 'efc_J' else np.zeros(nefc)
|
||||
|
||||
@@ -0,0 +1,180 @@
|
||||
# Copyright 2023 DeepMind Technologies Limited
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# ==============================================================================
|
||||
"""Functions for ray interesection testing."""
|
||||
|
||||
from typing import Tuple
|
||||
|
||||
import jax
|
||||
from jax import numpy as jp
|
||||
import mujoco
|
||||
# pylint: disable=g-importing-member
|
||||
from mujoco.mjx._src.types import Data
|
||||
from mujoco.mjx._src.types import GeomType
|
||||
from mujoco.mjx._src.types import Model
|
||||
# pylint: enable=g-importing-member
|
||||
import numpy as np
|
||||
|
||||
|
||||
def _ray_quad(
|
||||
a: jax.Array, b: jax.Array, c: jax.Array
|
||||
) -> Tuple[jax.Array, jax.Array]:
|
||||
"""Returns two solutions for quadratic: a*x^2 + 2*b*x + c = 0."""
|
||||
det = b * b - a * c
|
||||
det_2 = jp.sqrt(det)
|
||||
|
||||
x0, x1 = (-b - det_2) / a, (-b + det_2) / a
|
||||
x0 = jp.where((det < mujoco.mjMINVAL) | (x0 < 0), jp.inf, x0)
|
||||
x1 = jp.where((det < mujoco.mjMINVAL) | (x1 < 0), jp.inf, x1)
|
||||
|
||||
return x0, x1
|
||||
|
||||
|
||||
def _ray_plane(
|
||||
size: jax.Array,
|
||||
pnt: jax.Array,
|
||||
vec: jax.Array,
|
||||
) -> jax.Array:
|
||||
"""Returns the distance at which a ray intersects with a plane."""
|
||||
x = -pnt[2] / vec[2]
|
||||
|
||||
valid = vec[2] <= -mujoco.mjMINVAL # z-vec pointing towards front face
|
||||
valid &= x >= 0
|
||||
# only within rendered rectangle
|
||||
p = pnt[0:2] + x * vec[0:2]
|
||||
valid &= jp.all((size[0:2] <= 0) | (jp.abs(p) <= size[0:2]))
|
||||
|
||||
return jp.where(valid, x, jp.inf)
|
||||
|
||||
|
||||
def _ray_sphere(
|
||||
size: jax.Array,
|
||||
pnt: jax.Array,
|
||||
vec: jax.Array,
|
||||
) -> jax.Array:
|
||||
"""Returns the distance at which a ray intersects with a sphere."""
|
||||
x0, x1 = _ray_quad(vec @ vec, vec @ pnt, pnt @ pnt - size[0] * size[0])
|
||||
x = jp.where(jp.isinf(x0), x1, x0)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
def _ray_capsule(
|
||||
size: jax.Array,
|
||||
pnt: jax.Array,
|
||||
vec: jax.Array,
|
||||
) -> jax.Array:
|
||||
"""Returns the distance at which a ray intersects with a capsule."""
|
||||
|
||||
# cylinder round side: (x*lvec+lpnt)'*(x*lvec+lpnt) = size[0]*size[0]
|
||||
a = vec[0:2] @ vec[0:2]
|
||||
b = vec[0:2] @ pnt[0:2]
|
||||
c = pnt[0:2] @ pnt[0:2] - size[0] * size[0]
|
||||
|
||||
# solve a*x^2 + 2*b*x + c = 0
|
||||
x0, x1 = _ray_quad(a, b, c)
|
||||
x = jp.where(jp.isinf(x0), x1, x0)
|
||||
|
||||
# make sure round solution is between flat sides
|
||||
x = jp.where(jp.abs(pnt[2] + x * vec[2]) <= size[1], x, jp.inf)
|
||||
|
||||
# top cap
|
||||
dif = pnt - jp.array([0, 0, size[1]])
|
||||
x0, x1 = _ray_quad(vec @ vec, vec @ dif, dif @ dif - size[0] * size[0])
|
||||
# accept only top half of sphere
|
||||
x = jp.where((pnt[2] + x0 * vec[2] >= size[1]) & (x0 < x), x0, x)
|
||||
x = jp.where((pnt[2] + x1 * vec[2] >= size[1]) & (x1 < x), x1, x)
|
||||
|
||||
# bottom cap
|
||||
dif = pnt + jp.array([0, 0, size[1]])
|
||||
x0, x1 = _ray_quad(vec @ vec, vec @ dif, dif @ dif - size[0] * size[0])
|
||||
|
||||
# accept only bottom half of sphere
|
||||
x = jp.where((pnt[2] + x0 * vec[2] <= -size[1]) & (x0 < x), x0, x)
|
||||
x = jp.where((pnt[2] + x1 * vec[2] <= -size[1]) & (x1 < x), x1, x)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
def _ray_box(
|
||||
size: jax.Array,
|
||||
pnt: jax.Array,
|
||||
vec: jax.Array,
|
||||
) -> jax.Array:
|
||||
"""Returns the distance at which a ray intersects with a box."""
|
||||
|
||||
iface = jp.array([(1, 2), (0, 2), (0, 1), (1, 2), (0, 2), (0, 1)])
|
||||
|
||||
# side +1, -1
|
||||
# solution of pnt[i] + x * vec[i] = side * size[i]
|
||||
x = jp.concatenate([(size - pnt) / vec, (-size - pnt) / vec])
|
||||
|
||||
# intersection with face
|
||||
p0 = pnt[iface[:, 0]] + x * vec[iface[:, 0]]
|
||||
p1 = pnt[iface[:, 1]] + x * vec[iface[:, 1]]
|
||||
valid = jp.abs(p0) <= size[iface[:, 0]]
|
||||
valid &= jp.abs(p1) <= size[iface[:, 1]]
|
||||
|
||||
return jp.min(jp.where(valid, x, jp.inf))
|
||||
|
||||
|
||||
def _ray_mesh(
|
||||
size: jax.Array,
|
||||
pnt: jax.Array,
|
||||
vec: jax.Array,
|
||||
) -> jax.Array:
|
||||
"""Returns the distance at which a ray intersects with a mesh."""
|
||||
del size, pnt, vec
|
||||
raise NotImplementedError("ray <> mesh not implemented yet")
|
||||
|
||||
|
||||
_RAY_FUNC = {
|
||||
GeomType.PLANE: _ray_plane,
|
||||
GeomType.SPHERE: _ray_sphere,
|
||||
GeomType.CAPSULE: _ray_capsule,
|
||||
GeomType.BOX: _ray_box,
|
||||
# GeomType.MESH: _ray_mesh,
|
||||
}
|
||||
|
||||
|
||||
def ray(
|
||||
m: Model, d: Data, pnt: jax.Array, vec: jax.Array
|
||||
) -> Tuple[jax.Array, jax.Array]:
|
||||
"""Returns the geom id and distance at which a ray intersects with a geom."""
|
||||
|
||||
ids = []
|
||||
dists = []
|
||||
|
||||
# map ray to local geom frames
|
||||
geom_pnts = jax.vmap(lambda x, y: x.T @ (pnt - y))(d.geom_xmat, d.geom_xpos)
|
||||
geom_vecs = jax.vmap(lambda x: x.T @ vec)(d.geom_xmat)
|
||||
|
||||
for geom_type, fn in _RAY_FUNC.items():
|
||||
if not np.any(m.geom_type == geom_type):
|
||||
continue
|
||||
|
||||
geom_ids = jp.array(np.nonzero(m.geom_type == geom_type)[0])
|
||||
geom_dists = jax.vmap(fn)(
|
||||
m.geom_size[geom_ids], geom_pnts[geom_ids], geom_vecs[geom_ids]
|
||||
)
|
||||
ids.append(geom_ids)
|
||||
dists.append(geom_dists)
|
||||
|
||||
ids = jp.concatenate(ids)
|
||||
dists = jp.concatenate(dists)
|
||||
min_id = jp.argmin(dists)
|
||||
id_ = jp.where(jp.isinf(dists[min_id]), -1, ids[min_id])
|
||||
dist = jp.where(jp.isinf(dists[min_id]), -1, dists[min_id])
|
||||
|
||||
return id_, dist
|
||||
@@ -0,0 +1,149 @@
|
||||
# Copyright 2023 DeepMind Technologies Limited
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# ==============================================================================
|
||||
"""Tests for ray functions."""
|
||||
|
||||
from absl.testing import absltest
|
||||
import jax
|
||||
from jax import numpy as jp
|
||||
import mujoco
|
||||
from mujoco import mjx
|
||||
from mujoco.mjx._src import test_util
|
||||
import numpy as np
|
||||
|
||||
# tolerance for difference between MuJoCo and MJX ray calculations - mostly
|
||||
# due to float precision
|
||||
_TOLERANCE = 5e-5
|
||||
|
||||
|
||||
def _assert_eq(a, b, name):
|
||||
tol = _TOLERANCE * 10 # avoid test noise
|
||||
err_msg = f'mismatch: {name}'
|
||||
np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
|
||||
|
||||
|
||||
class RayTest(absltest.TestCase):
|
||||
|
||||
def test_ray_nothing(self):
|
||||
"""Tests that MJX ray returns -1 when nothing is hit."""
|
||||
m = test_util.load_test_file('ray.xml')
|
||||
d = mujoco.MjData(m)
|
||||
mujoco.mj_forward(m, d)
|
||||
mx, dx = mjx.put_model(m), mjx.put_data(m, d)
|
||||
|
||||
pnt, vec = jp.array([12.146, 1.865, 3.895]), jp.array([0, 0, -1.0])
|
||||
geomid, dist = jax.jit(mjx.ray)(mx, dx, pnt, vec)
|
||||
_assert_eq(geomid, jp.array([-1]), 'geom_id')
|
||||
_assert_eq(dist, jp.array([-1]), 'dist')
|
||||
|
||||
def test_ray_plane(self):
|
||||
"""Tests MJX ray<>plane matches MuJoCo."""
|
||||
m = test_util.load_test_file('ray.xml')
|
||||
d = mujoco.MjData(m)
|
||||
mujoco.mj_forward(m, d)
|
||||
mx, dx = mjx.put_model(m), mjx.put_data(m, d)
|
||||
|
||||
# looking down at a slight angle
|
||||
pnt, vec = jp.array([2, 1, 3.0]), jp.array([0.1, 0.2, -1.0])
|
||||
vec /= jp.linalg.norm(vec)
|
||||
geomid, dist = jax.jit(mjx.ray)(mx, dx, pnt, vec)
|
||||
_assert_eq(geomid, jp.array([0]), 'geom_id')
|
||||
pnt, vec, unused = np.array(pnt), np.array(vec), np.zeros(1, dtype=np.int32)
|
||||
mj_dist = mujoco.mj_ray(m, d, pnt, vec, None, 1, -1, unused)
|
||||
_assert_eq(dist, mj_dist, 'dist')
|
||||
|
||||
# looking on wrong side of plane
|
||||
pnt = jp.array([0, 0, -0.5])
|
||||
geomid, dist = jax.jit(mjx.ray)(mx, dx, pnt, vec)
|
||||
_assert_eq(geomid, jp.array([-1]), 'geom_id')
|
||||
_assert_eq(dist, jp.array([-1]), 'dist')
|
||||
|
||||
def test_ray_sphere(self):
|
||||
"""Tests MJX ray<>sphere matches MuJoCo."""
|
||||
m = test_util.load_test_file('ray.xml')
|
||||
d = mujoco.MjData(m)
|
||||
mujoco.mj_forward(m, d)
|
||||
mx, dx = mjx.put_model(m), mjx.put_data(m, d)
|
||||
|
||||
# looking down at sphere at a slight angle
|
||||
pnt, vec = jp.array([0, 0, 1.6]), jp.array([0.1, 0.2, -1.0])
|
||||
vec /= jp.linalg.norm(vec)
|
||||
geomid, dist = jax.jit(mjx.ray)(mx, dx, pnt, vec)
|
||||
_assert_eq(geomid, jp.array([1]), 'geom_id')
|
||||
pnt, vec, unused = np.array(pnt), np.array(vec), np.zeros(1, dtype=np.int32)
|
||||
mj_dist = mujoco.mj_ray(m, d, pnt, vec, None, 1, -1, unused)
|
||||
_assert_eq(dist, mj_dist, 'dist')
|
||||
|
||||
def test_ray_capsule(self):
|
||||
"""Tests MJX ray<>capsule matches MuJoCo."""
|
||||
m = test_util.load_test_file('ray.xml')
|
||||
d = mujoco.MjData(m)
|
||||
mujoco.mj_forward(m, d)
|
||||
mx, dx = mjx.put_model(m), mjx.put_data(m, d)
|
||||
|
||||
# looking down at capsule at a slight angle
|
||||
pnt, vec = jp.array([0.5, 1, 1.6]), jp.array([0, 0.05, -1.0])
|
||||
vec /= jp.linalg.norm(vec)
|
||||
geomid, dist = jax.jit(mjx.ray)(mx, dx, pnt, vec)
|
||||
_assert_eq(geomid, jp.array([2]), 'geom_id')
|
||||
pnt, vec, unused = np.array(pnt), np.array(vec), np.zeros(1, dtype=np.int32)
|
||||
mj_dist = mujoco.mj_ray(m, d, pnt, vec, None, 1, -1, unused)
|
||||
_assert_eq(dist, mj_dist, 'dist')
|
||||
|
||||
# looking up at capsule from below
|
||||
pnt, vec = jp.array([-0.5, 1, 0.05]), jp.array([0, 0.05, 1.0])
|
||||
vec /= jp.linalg.norm(vec)
|
||||
geomid, dist = jax.jit(mjx.ray)(mx, dx, pnt, vec)
|
||||
_assert_eq(geomid, jp.array([2]), 'geom_id')
|
||||
pnt, vec, unused = np.array(pnt), np.array(vec), np.zeros(1, dtype=np.int32)
|
||||
mj_dist = mujoco.mj_ray(m, d, pnt, vec, None, 1, -1, unused)
|
||||
_assert_eq(dist, mj_dist, 'dist')
|
||||
|
||||
# looking at cylinder of capsule from the side
|
||||
pnt, vec = jp.array([0, 1, 0.75]), jp.array([1, 0, 0])
|
||||
vec /= jp.linalg.norm(vec)
|
||||
geomid, dist = jax.jit(mjx.ray)(mx, dx, pnt, vec)
|
||||
_assert_eq(geomid, jp.array([2]), 'geom_id')
|
||||
pnt, vec, unused = np.array(pnt), np.array(vec), np.zeros(1, dtype=np.int32)
|
||||
mj_dist = mujoco.mj_ray(m, d, pnt, vec, None, 1, -1, unused)
|
||||
_assert_eq(dist, mj_dist, 'dist')
|
||||
|
||||
def test_ray_box(self):
|
||||
"""Tests MJX ray<>box matches MuJoCo."""
|
||||
m = test_util.load_test_file('ray.xml')
|
||||
d = mujoco.MjData(m)
|
||||
mujoco.mj_forward(m, d)
|
||||
mx, dx = mjx.put_model(m), mjx.put_data(m, d)
|
||||
|
||||
# looking down at box at a slight angle
|
||||
pnt, vec = jp.array([1, 0, 1.6]), jp.array([0, 0.05, -1.0])
|
||||
vec /= jp.linalg.norm(vec)
|
||||
geomid, dist = jax.jit(mjx.ray)(mx, dx, pnt, vec)
|
||||
_assert_eq(geomid, jp.array([3]), 'geom_id')
|
||||
pnt, vec, unused = np.array(pnt), np.array(vec), np.zeros(1, dtype=np.int32)
|
||||
mj_dist = mujoco.mj_ray(m, d, pnt, vec, None, 1, -1, unused)
|
||||
_assert_eq(dist, mj_dist, 'dist')
|
||||
|
||||
# looking up at box from below
|
||||
pnt, vec = jp.array([1, 0, 0.05]), jp.array([0, 0.05, 1.0])
|
||||
vec /= jp.linalg.norm(vec)
|
||||
geomid, dist = jax.jit(mjx.ray)(mx, dx, pnt, vec)
|
||||
_assert_eq(geomid, jp.array([3]), 'geom_id')
|
||||
pnt, vec, unused = np.array(pnt), np.array(vec), np.zeros(1, dtype=np.int32)
|
||||
mj_dist = mujoco.mj_ray(m, d, pnt, vec, None, 1, -1, unused)
|
||||
_assert_eq(dist, mj_dist, 'dist')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
absltest.main()
|
||||
@@ -26,6 +26,7 @@ TEST_FILES: List[str] = [
|
||||
'constraints.xml',
|
||||
'convex.xml',
|
||||
'pendula.xml',
|
||||
'ray.xml',
|
||||
]
|
||||
|
||||
_ACTUATOR_TYPES = ['motor', 'velocity', 'position', 'general', 'intvelocity']
|
||||
|
||||
@@ -0,0 +1,16 @@
|
||||
<mujoco model="ray">
|
||||
<asset>
|
||||
<mesh name="tetrahedron" file="meshes/tetrahedron.stl" scale="0.4 0.4 0.4" />
|
||||
<texture builtin="checker" height="100" name="texplane" rgb1="0 0 0" rgb2="0.8 0.8 0.8" type="2d" width="100"/>
|
||||
<material name="MatPlane" reflectance="0.5" shininess="1" specular="1" texrepeat="60 60" texture="texplane"/>
|
||||
</asset>
|
||||
<worldbody>
|
||||
<light cutoff="100" diffuse="1 1 1" dir="-0 0 -1.3" directional="true" exponent="1" pos="0 0 1.3" specular=".1 .1 .1"/>
|
||||
|
||||
<geom name="plane" pos="0 0 0" quat="1 0 0 0" size="4 4 4" type="plane" rgba="0.1 0.1 0.1 1"/>
|
||||
<geom name="sphere" pos="0 0 1" quat="1 0 0 0" size="0.5" type="sphere" rgba="1 0 0 1"/>
|
||||
<geom name="capsule" pos="0 1 1" quat="0 0.3826834 0 0.9238795 " size="0.25 0.5" type="capsule" rgba="0 1 0 1"/>
|
||||
<geom name="box" pos="1 0 1" quat="0 0.3826834 0 0.9238795" size="0.5 0.25 0.3" type="box" rgba="0 0 1 1"/>
|
||||
<geom name="mesh" pos="1 1 1" quat="0 0 0.3826834 0.9238795" type="mesh" mesh="tetrahedron" rgba="1 1 0 1"/>
|
||||
</worldbody>
|
||||
</mujoco>
|
||||
Reference in New Issue
Block a user