Add ray intersection with ellipsoid to MJX.

PiperOrigin-RevId: 690557087
Change-Id: Icd2f0be1243574720baeac82b8bdd92e57c6a61c

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
 ^^^^^^^^^

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,
 }

mjx/mujoco/mjx/test_data/sensor/sensor.xml

@@ -119,6 +119,11 @@
       <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>
+    <body pos="-20 -20.75 -20">
+      <freejoint/>
+      <geom type="capsule" fromto="-.5 0 0 .5 0 0" size="0.0125"/>
+      <site name="touch_ellipsoid" type="ellipsoid" pos="0.5 0 0.0" size="0.05 0.01 0.02"/>
+    </body>
   </worldbody>
 
   <actuator>
@@ -178,6 +183,7 @@
     <framepos name="framepos0" objtype="site" objname="site0"/>
     <velocimeter name="velocimeter1" site="site1"/>
     <gyro name="gyro1" site="site1"/>
+    <touch name="touch_ellipsoid" site="touch_ellipsoid"/>
     <frameangvel name="frameangvel3" objtype="site" objname="site3"/>
     <frameangvel name="frameangvel3ref" objtype="site" objname="site3" reftype="body" refname="body2"/>
     <frameangvel name="frameangvel3cutoff" objtype="site" objname="site3" cutoff="1e-3"/>