From 3c21abc0e544bfafb140f8ef5506f4fa8f6aa616 Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Mon, 28 Oct 2024 04:24:11 -0700 Subject: [PATCH] Add ray intersection with ellipsoid to MJX. PiperOrigin-RevId: 690557087 Change-Id: Icd2f0be1243574720baeac82b8bdd92e57c6a61c --- doc/changelog.rst | 1 + mjx/mujoco/mjx/_src/ray.py | 24 ++++++++++++++++++++++ mjx/mujoco/mjx/test_data/sensor/sensor.xml | 6 ++++++ 3 files changed, 31 insertions(+) diff --git a/doc/changelog.rst b/doc/changelog.rst index 6aa3b745..436ffce3 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -20,6 +20,7 @@ MJX - Added ``apply_ft``, ``jac``, and ``xfrc_accumulate`` as public functions. - Added ``TOUCH`` sensor. - Added support for ``eq_active``. Fixes :github:issue:`2173`. +- Added ray intersection with ellipsoid. Bug fixes ^^^^^^^^^ diff --git a/mjx/mujoco/mjx/_src/ray.py b/mjx/mujoco/mjx/_src/ray.py index 2591385a..9335f4e9 100644 --- a/mjx/mujoco/mjx/_src/ray.py +++ b/mjx/mujoco/mjx/_src/ray.py @@ -108,6 +108,29 @@ def _ray_capsule( return x +def _ray_ellipsoid( + size: jax.Array, + pnt: jax.Array, + vec: jax.Array, +) -> jax.Array: + """Returns the distance at which a ray intersects with an ellipsoid.""" + + # invert size^2 + s = 1 / jp.square(size) + + # (x*lvec+lpnt)' * diag(1/size^2) * (x*lvec+lpnt) = 1 + svec = s * vec + a = svec @ vec + b = svec @ pnt + c = (s * pnt) @ pnt - 1 + + # solve a*x^2 + 2*b*x + c = 0 + x0, x1 = _ray_quad(a, b, c) + x = jp.where(jp.isinf(x0), x1, x0) + + return x + + def _ray_box( size: jax.Array, pnt: jax.Array, @@ -201,6 +224,7 @@ _RAY_FUNC = { GeomType.PLANE: _ray_plane, GeomType.SPHERE: _ray_sphere, GeomType.CAPSULE: _ray_capsule, + GeomType.ELLIPSOID: _ray_ellipsoid, GeomType.BOX: _ray_box, GeomType.MESH: _ray_mesh, } diff --git a/mjx/mujoco/mjx/test_data/sensor/sensor.xml b/mjx/mujoco/mjx/test_data/sensor/sensor.xml index 71532743..b6d3fa2a 100644 --- a/mjx/mujoco/mjx/test_data/sensor/sensor.xml +++ b/mjx/mujoco/mjx/test_data/sensor/sensor.xml @@ -119,6 +119,11 @@ + + + + + @@ -178,6 +183,7 @@ +