diff --git a/doc/changelog.rst b/doc/changelog.rst index 0b89de4b..9665e7a5 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -9,10 +9,12 @@ MJX ^^^ 1. Added :ref:`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) diff --git a/mjx/mujoco/mjx/__init__.py b/mjx/mujoco/mjx/__init__.py index d8e3f437..f4953c60 100644 --- a/mjx/mujoco/mjx/__init__.py +++ b/mjx/mujoco/mjx/__init__.py @@ -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 diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index 788c3899..d7ae4a1d 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -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) diff --git a/mjx/mujoco/mjx/_src/ray.py b/mjx/mujoco/mjx/_src/ray.py new file mode 100644 index 00000000..582d6e4f --- /dev/null +++ b/mjx/mujoco/mjx/_src/ray.py @@ -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 diff --git a/mjx/mujoco/mjx/_src/ray_test.py b/mjx/mujoco/mjx/_src/ray_test.py new file mode 100644 index 00000000..712aca63 --- /dev/null +++ b/mjx/mujoco/mjx/_src/ray_test.py @@ -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() diff --git a/mjx/mujoco/mjx/_src/test_util.py b/mjx/mujoco/mjx/_src/test_util.py index ab27c2b1..d350c11a 100644 --- a/mjx/mujoco/mjx/_src/test_util.py +++ b/mjx/mujoco/mjx/_src/test_util.py @@ -26,6 +26,7 @@ TEST_FILES: List[str] = [ 'constraints.xml', 'convex.xml', 'pendula.xml', + 'ray.xml', ] _ACTUATOR_TYPES = ['motor', 'velocity', 'position', 'general', 'intvelocity'] diff --git a/mjx/mujoco/mjx/test_data/ray.xml b/mjx/mujoco/mjx/test_data/ray.xml new file mode 100644 index 00000000..a6424ec4 --- /dev/null +++ b/mjx/mujoco/mjx/test_data/ray.xml @@ -0,0 +1,16 @@ + + + + + + + + + + + + + + + +