Add more public support functions for MJX.

PiperOrigin-RevId: 687469475
Change-Id: I4218cc2f6d51d3f2f63b1a9a96953d550a5503a5

doc/changelog.rst

@@ -2,6 +2,14 @@
 Changelog
 =========
 
+Upcoming version (not yet released)
+-----------------------------------
+
+MJX
+^^^
+
+1. Added ``apply_ft``, ``jac``, and  ``xfrc_accumulate`` as public functions.
+
 Version 3.2.4 (Oct 15, 2024)
 ----------------------------
 

mjx/mujoco/mjx/__init__.py

@@ -48,10 +48,13 @@ from mujoco.mjx._src.smooth import subtree_vel
 from mujoco.mjx._src.smooth import tendon
 from mujoco.mjx._src.smooth import transmission
 from mujoco.mjx._src.solver import solve
+from mujoco.mjx._src.support import apply_ft
 from mujoco.mjx._src.support import full_m
 from mujoco.mjx._src.support import id2name
 from mujoco.mjx._src.support import is_sparse
+from mujoco.mjx._src.support import jac
 from mujoco.mjx._src.support import mul_m
 from mujoco.mjx._src.support import name2id
+from mujoco.mjx._src.support import xfrc_accumulate
 from mujoco.mjx._src.test_util import benchmark
 from mujoco.mjx._src.types import *

mjx/mujoco/mjx/_src/support.py

@@ -348,7 +348,7 @@ def contact_force_dim(
     raise ValueError(f'Unknown cone type: {m.opt.cone}.')
 
 
-def length_circle(
+def _length_circle(
     p0: jax.Array, p1: jax.Array, ind: jax.Array, rad: jax.Array
 ) -> jax.Array:
   """Compute length of circle."""
@@ -366,7 +366,7 @@ def length_circle(
   return rad * angle
 
 
-def is_intersect(
+def _is_intersect(
     p1: jax.Array, p2: jax.Array, p3: jax.Array, p4: jax.Array
 ) -> jax.Array:
   """Check for intersection between two lines defined by their endpoints."""
@@ -429,7 +429,7 @@ def wrap_circle(
     good = jp.where(sidesite, good0, good1)
 
     # penalize for intersection
-    intersect = is_intersect(d[:2], sol[0], d[2:], sol[1])
+    intersect = _is_intersect(d[:2], sol[0], d[2:], sol[1])
     good = jp.where(intersect, -10000, good)
 
     return sol, good
@@ -442,10 +442,10 @@ def wrap_circle(
   pnt = sol.reshape(-1)
 
   # check for intersection
-  intersect = is_intersect(d[:2], pnt[:2], d[2:], pnt[2:])
+  intersect = _is_intersect(d[:2], pnt[:2], d[2:], pnt[2:])
 
   # compute curve length
-  wlen = length_circle(sol[0], sol[1], i, rad)
+  wlen = _length_circle(sol[0], sol[1], i, rad)
 
   # check cases
   invalid = (