rigidax 1.0.0__py3-none-any.whl

rigidax/__init__.py

@@ -0,0 +1,64 @@
+"""rigidax: differentiable rigid-body kinematics and dynamics in JAX"""
+
+from rigidax.dynamics import (
+    bias_forces,
+    body_spatial_inertias,
+    coriolis_forces,
+    difference,
+    dof_inertias,
+    dynamics_terms,
+    forward_dynamics,
+    gravity_forces,
+    integrate,
+    mass_matrix,
+    mass_matrix_inverse,
+    rnea,
+)
+from rigidax.kinematics import (
+    Kinematics,
+    body_coms,
+    body_jacobian,
+    body_pose,
+    center_of_mass,
+    center_of_mass_jacobian,
+    forward_kinematics,
+    point_jacobian,
+    points_jacobians,
+    site_jacobian,
+    site_pose,
+)
+from rigidax.model import Model, Params, Topology, from_mujoco, from_xml_string, load
+
+__version__ = "1.0.0"
+
+__all__ = [
+    "Kinematics",
+    "Model",
+    "Params",
+    "Topology",
+    "bias_forces",
+    "body_coms",
+    "body_jacobian",
+    "body_pose",
+    "body_spatial_inertias",
+    "center_of_mass",
+    "center_of_mass_jacobian",
+    "coriolis_forces",
+    "difference",
+    "dof_inertias",
+    "dynamics_terms",
+    "forward_dynamics",
+    "forward_kinematics",
+    "from_mujoco",
+    "from_xml_string",
+    "gravity_forces",
+    "integrate",
+    "load",
+    "mass_matrix",
+    "mass_matrix_inverse",
+    "point_jacobian",
+    "points_jacobians",
+    "rnea",
+    "site_jacobian",
+    "site_pose",
+]

rigidax/dynamics.py

@@ -0,0 +1,212 @@
+"""Inverse/forward dynamics and configuration-space integration.
+
+All algorithms are written as dense, vectorized passes over the DOF ancestor mask,
+so there are no sequential loops over joints.
+"""
+
+import jax
+import jax.numpy as jnp
+import jax.scipy.linalg as jsl
+import numpy as np
+from jax import Array
+
+from rigidax.kinematics import Kinematics, _kin, body_coms
+from rigidax.model import Model, derived
+from rigidax.spatial import (
+    cross_force,
+    cross_motion,
+    einsum,
+    mm,
+    quat_conj,
+    quat_integrate,
+    quat_log,
+    quat_mul,
+    quat_normalize,
+    quat_to_mat,
+    skew,
+)
+
+
+def body_spatial_inertias(model: Model, kin: Kinematics, bodies=None) -> Array:
+    """(nbody, 6, 6) spatial inertias in the world frame, about the origin"""
+    P = model.params
+    b = slice(None) if bodies is None else bodies
+    with jax.ensure_compile_time_eval():  # parameter-only: constant if not traced
+        Ri = quat_to_mat(jnp.asarray(P.body_iquat)[b])
+        I_local = einsum("bij,bj,bkj->bik", Ri, jnp.asarray(P.body_inertia)[b], Ri)
+        ipos = jnp.asarray(P.body_ipos)[b]
+        m = jnp.asarray(P.body_mass)[b][:, None, None]
+        m_eye = m * jnp.eye(3)
+    R, p = kin.body_R[b], kin.body_p[b]
+    Ic = mm(mm(R, I_local), jnp.swapaxes(R, -1, -2))
+    C = skew(p + einsum("bij,bj->bi", R, ipos))
+    mC = m * C
+    top = jnp.concatenate([Ic - mm(mC, C), mC], -1)
+    bot = jnp.concatenate([-mC, jnp.broadcast_to(m_eye, mC.shape)], -1)
+    return jnp.concatenate([top, bot], -2)
+
+
+def _moving_inertias(model, kin):
+    """Spatial inertias of the bodies that some DOF moves (static subset)"""
+    return body_spatial_inertias(model, kin, derived(model.topo).owned)
+
+
+def dof_inertias(model: Model, kin: Kinematics) -> Array:
+    """(nv, 6, 6) inertia carried by each DOF (welded bodies fused, virtual DOFs zero)"""
+    D = derived(model.topo)
+    own = np.zeros((model.nv, len(D.owned)))
+    own[D.owned_dof, np.arange(len(D.owned))] = 1.0
+    return einsum("ki,iab->kab", own, _moving_inertias(model, kin))
+
+
+def _ext_wrenches(model: Model, kin: Kinematics, xfrc: Array) -> Array:
+    """Per-body [force, torque] at the COM (world axes, MuJoCo's xfrc_applied layout)
+    -> spatial forces about the origin, for each moving body"""
+    D = derived(model.topo)
+    c = body_coms(model, None, kin)[D.owned]
+    f, tq = xfrc[D.owned, :3], xfrc[D.owned, 3:]
+    return jnp.concatenate([tq + jnp.cross(c, f), f], -1)
+
+
+def _rnea(model, kin, Ib, v, a, gravity, F_ext):
+    D = derived(model.topo)
+    A, W = jnp.asarray(D.A), jnp.asarray(D.W)
+    S = kin.S
+    acc = None
+    if v is not None:
+        Sv = S * v[:, None]
+        V = mm(A, Sv)
+        if np.array_equal(D.cross, np.arange(model.nv)):
+            Vc = V  # hinge/slide only: each axis moves with its own frame
+        else:
+            Vc = jnp.concatenate([V, jnp.zeros((1, 6), V.dtype)])[D.cross]
+        acc = cross_motion(Vc, Sv)
+    if a is not None:
+        Sa = S * a[:, None]
+        acc = Sa if acc is None else acc + Sa
+    Acc = jnp.zeros((len(D.owned), 6), S.dtype) if acc is None else mm(A, acc)[D.owned_dof]
+    if gravity:
+        g = jnp.asarray(model.params.gravity, S.dtype)
+        Acc = Acc + jnp.concatenate([jnp.zeros(3, S.dtype), -g])
+    F = einsum("iab,ib->ia", Ib, Acc)
+    if v is not None:
+        Vb = V[D.owned_dof]
+        F = F + cross_force(Vb, einsum("iab,ib->ia", Ib, Vb))
+    if F_ext is not None:
+        F = F - F_ext
+    tau = jnp.sum(S * mm(W, F), -1)
+    if a is not None:
+        tau = tau + jnp.asarray(model.params.dof_armature) * a
+    return tau
+
+
+def rnea(model: Model, q: Array, v: Array | None, a: Array | None, *,
+         xfrc: Array | None = None, gravity: bool = True,
+         kin: Kinematics | None = None) -> Array:
+    """Inverse dynamics: tau = M(q) a + C(q, v) v + g(q) - J^T f_ext.
+
+    v / a may be None (treated as zero). `xfrc` is (nbody, 6) [force, torque] applied
+    at each body's COM in world axes, like MuJoCo's xfrc_applied.
+    Armature is included in the M(q) a term.
+    """
+    kin = _kin(model, q, kin)
+    F_ext = None if xfrc is None else _ext_wrenches(model, kin, xfrc)
+    return _rnea(model, kin, _moving_inertias(model, kin), v, a, gravity, F_ext)
+
+
+def _crba(model, kin, Ib):
+    D = derived(model.topo)
+    A = jnp.asarray(D.A)
+    Ic = einsum("ki,iab->kab", D.W, Ib)  # composite inertia of each DOF's subtree
+    F = einsum("kab,kb->ka", Ic, kin.S)
+    M_all = mm(F, kin.S.T)  # M_all[k, j] = s_j . Ic_k s_k
+    L = A * M_all  # keep (k, ancestor j) pairs
+    return L + jnp.tril(L, -1).T + jnp.diag(jnp.asarray(model.params.dof_armature))
+
+
+def mass_matrix(model: Model, q: Array, kin: Kinematics | None = None) -> Array:
+    """Joint-space inertia matrix (CRBA), including armature, (nv, nv)"""
+    kin = _kin(model, q, kin)
+    return _crba(model, kin, _moving_inertias(model, kin))
+
+
+def bias_forces(model: Model, q: Array, v: Array, kin: Kinematics | None = None) -> Array:
+    """C(q, v) v + g(q), (nv,)"""
+    return rnea(model, q, v, None, kin=kin)
+
+
+def gravity_forces(model: Model, q: Array, kin: Kinematics | None = None) -> Array:
+    return rnea(model, q, None, None, kin=kin)
+
+
+def coriolis_forces(model: Model, q: Array, v: Array, kin: Kinematics | None = None) -> Array:
+    return rnea(model, q, v, None, gravity=False, kin=kin)
+
+
+def dynamics_terms(model: Model, q: Array, v: Array, kin: Kinematics | None = None):
+    """(M, bias) from a single kinematics and inertia pass"""
+    kin = _kin(model, q, kin)
+    Ib = _moving_inertias(model, kin)
+    return _crba(model, kin, Ib), _rnea(model, kin, Ib, v, None, True, None)
+
+
+def forward_dynamics(model: Model, q: Array, v: Array, tau: Array, *,
+                     xfrc: Array | None = None, kin: Kinematics | None = None) -> Array:
+    """qacc = M^-1 (tau + J^T f_ext - bias), via one Cholesky factorization"""
+    kin = _kin(model, q, kin)
+    Ib = _moving_inertias(model, kin)
+    M = _crba(model, kin, Ib)
+    F_ext = None if xfrc is None else _ext_wrenches(model, kin, xfrc)
+    bias = _rnea(model, kin, Ib, v, None, True, F_ext)
+    return jsl.cho_solve(jsl.cho_factor(M, lower=True), tau - bias)
+
+
+def mass_matrix_inverse(M: Array) -> Array:
+    c = jsl.cho_factor(M, lower=True)
+    return jsl.cho_solve(c, jnp.eye(M.shape[0], dtype=M.dtype))
+
+
+# ----- configuration space -----
+
+
+def integrate(model: Model, q: Array, v: Array, dt: float) -> Array:
+    """q (+) v*dt, respecting quaternions (MuJoCo's mj_integratePos)"""
+    D = derived(model.topo)
+    q = jnp.asarray(q)
+    if len(D.scalar_j):
+        j = D.scalar_j
+        q = q.at[D.qadr[j]].add(v[D.dadr[j]] * dt)
+    if len(D.free_j):
+        qa, da = D.qadr[D.free_j][:, None], D.dadr[D.free_j][:, None]
+        q = q.at[qa + jnp.arange(3)].add(v[da + jnp.arange(3)] * dt)
+        quat = quat_normalize(q[qa + 3 + jnp.arange(4)])
+        q = q.at[qa + 3 + jnp.arange(4)].set(quat_integrate(quat, v[da + 3 + jnp.arange(3)], dt))
+    if len(D.ball_j):
+        qa, da = D.qadr[D.ball_j][:, None], D.dadr[D.ball_j][:, None]
+        quat = quat_normalize(q[qa + jnp.arange(4)])
+        q = q.at[qa + jnp.arange(4)].set(quat_integrate(quat, v[da + jnp.arange(3)], dt))
+    return q
+
+
+def difference(model: Model, q0: Array, q1: Array) -> Array:
+    """v such that integrate(q0, v, 1) == q1 (MuJoCo's mj_differentiatePos, dt=1)"""
+    D = derived(model.topo)
+    q0, q1 = jnp.asarray(q0), jnp.asarray(q1)
+    v = jnp.zeros(model.nv, jnp.result_type(q0, q1))
+    if len(D.scalar_j):
+        j = D.scalar_j
+        v = v.at[D.dadr[j]].set(q1[D.qadr[j]] - q0[D.qadr[j]])
+
+    def rot(qa):
+        a = quat_normalize(q0[qa + jnp.arange(4)])
+        b = quat_normalize(q1[qa + jnp.arange(4)])
+        return quat_log(quat_mul(quat_conj(a), b))
+
+    if len(D.free_j):
+        qa, da = D.qadr[D.free_j][:, None], D.dadr[D.free_j][:, None]
+        v = v.at[da + jnp.arange(3)].set(q1[qa + jnp.arange(3)] - q0[qa + jnp.arange(3)])
+        v = v.at[da + 3 + jnp.arange(3)].set(rot(qa + 3))
+    if len(D.ball_j):
+        qa, da = D.qadr[D.ball_j][:, None], D.dadr[D.ball_j][:, None]
+        v = v.at[da + jnp.arange(3)].set(rot(qa))
+    return v

rigidax/ik/__init__.py

@@ -0,0 +1,48 @@
+"""rigidax.ik: differential IK, batched retargeting and self-collision avoidance"""
+
+from rigidax.ik.collision import SphereModel
+from rigidax.ik.qp import INF, QPResult, QPState, solve_qp, spd_inverse
+from rigidax.ik.solver import FloorPoints, IKSolver
+from rigidax.ik.tasks import (
+    DirectionTask,
+    HeightTask,
+    JointTask,
+    OrientationTask,
+    PositionTask,
+    PostureTask,
+    TwistTask,
+    direction,
+    height,
+    joints,
+    orientation,
+    position,
+    rotation_error,
+    twist,
+)
+from rigidax.ik.trajectory import smooth_trajectory
+
+__all__ = [
+    "DirectionTask",
+    "FloorPoints",
+    "HeightTask",
+    "IKSolver",
+    "INF",
+    "JointTask",
+    "OrientationTask",
+    "PositionTask",
+    "PostureTask",
+    "QPResult",
+    "QPState",
+    "SphereModel",
+    "TwistTask",
+    "direction",
+    "height",
+    "joints",
+    "orientation",
+    "position",
+    "rotation_error",
+    "smooth_trajectory",
+    "solve_qp",
+    "spd_inverse",
+    "twist",
+]

rigidax/ik/bvh.py

@@ -0,0 +1,99 @@
+"""Minimal BVH motion-capture reader (Y-up files converted to Z-up)"""
+
+from pathlib import Path
+
+import numpy as np
+
+
+def load_bvh(path):
+    """-> joint names, parents, offsets (J, 3), channels per joint, motion (T, C)"""
+    names, parents, offsets, channels, stack = [], [], [], [], []
+    lines = iter(Path(path).read_text().split("\n"))
+    for line in lines:
+        tok = line.split()
+        if not tok:
+            continue
+        if tok[0] in ("ROOT", "JOINT"):
+            names.append(tok[1])
+            parents.append(stack[-1] if stack else -1)
+            offsets.append(None)
+            channels.append([])
+        elif tok[0] == "End":
+            names.append(None)  # end site: parsed for nesting, dropped later
+            parents.append(stack[-1])
+            offsets.append(None)
+            channels.append([])
+        elif tok[0] == "{":
+            stack.append(len(names) - 1)
+        elif tok[0] == "}":
+            stack.pop()
+        elif tok[0] == "OFFSET":
+            offsets[stack[-1]] = [float(x) for x in tok[1:4]]
+        elif tok[0] == "CHANNELS":
+            channels[stack[-1]] = tok[2:]
+        elif tok[0] == "MOTION":
+            break
+    next(lines)  # Frames:
+    next(lines)  # Frame Time:
+    motion = np.array([[float(x) for x in ln.split()] for ln in lines if ln.strip()])
+    keep = [i for i, n in enumerate(names) if n is not None]
+    remap = {old: new for new, old in enumerate(keep)}
+    return ([names[i] for i in keep], [remap.get(parents[i], -1) for i in keep],
+            np.array([offsets[i] for i in keep]), [channels[i] for i in keep], motion)
+
+
+def _axis_rot(axis, deg):
+    a = np.radians(deg)
+    c, s = np.cos(a), np.sin(a)
+    o, z = np.ones_like(a), np.zeros_like(a)
+    if axis == "X":
+        m = [[o, z, z], [z, c, -s], [z, s, c]]
+    elif axis == "Y":
+        m = [[c, z, s], [z, o, z], [-s, z, c]]
+    else:
+        m = [[c, -s, z], [s, c, z], [z, z, o]]
+    return np.moveaxis(np.array(m), [0, 1], [-2, -1])
+
+
+def bvh_rotations(path):
+    """Global joint orientations {name: (T, 3, 3)}, Z-up (same frame as bvh_positions)"""
+    names, parents, offsets, channels, motion = load_bvh(path)
+    T, col = len(motion), 0
+    R = [None] * len(names)
+    for j, ch in enumerate(channels):
+        Rl = np.broadcast_to(np.eye(3), (T, 3, 3))
+        for c in ch:
+            if not c.endswith("position"):
+                Rl = Rl @ _axis_rot(c[0], motion[:, col])
+            col += 1
+        R[j] = Rl if parents[j] < 0 else R[parents[j]] @ Rl
+    C = np.array([[1.0, 0, 0], [0, 0, -1], [0, 1, 0]])  # Y-up -> Z-up
+    return {n: C @ R[i] for i, n in enumerate(names)}
+
+
+def bvh_positions(path, unit: float = 0.01):
+    """Global joint positions {name: (T, 3)}, Z-up, in meters (unit: BVH units -> m)"""
+    names, parents, offsets, channels, motion = load_bvh(path)
+    T, col = len(motion), 0
+    R = [None] * len(names)
+    P = [None] * len(names)
+    for j, ch in enumerate(channels):
+        Rl = np.broadcast_to(np.eye(3), (T, 3, 3))
+        pos = np.broadcast_to(offsets[j], (T, 3))
+        for c in ch:
+            if c.endswith("position"):
+                pos = pos.copy() if pos.base is not None else pos
+                pos = np.array(pos)
+                pos[:, "XYZ".index(c[0])] = motion[:, col]
+            else:
+                Rl = Rl @ _axis_rot(c[0], motion[:, col])
+            col += 1
+        p = parents[j]
+        if p < 0:
+            R[j], P[j] = Rl, pos
+        else:
+            R[j] = R[p] @ Rl
+            P[j] = P[p] + np.einsum("tij,j->ti", R[p], offsets[j])
+    pos = np.stack(P, 1) * unit
+    C = np.array([[1.0, 0, 0], [0, 0, -1], [0, 1, 0]])  # Y-up -> Z-up
+    return {n: pos[:, i] @ C.T for i, n in enumerate(names)}

rigidax/ik/collision.py

@@ -0,0 +1,235 @@
+"""Automatic sphere collision models for any MuJoCo model.
+
+Each body's collision geometry (meshes, boxes, capsules, cylinders, spheres, ellipsoids)
+is sampled into points. The points are sliced along their principal axis, and each
+slice gets a sphere at its centroid covering most of its cross-section. Pairs are
+generated automatically: different bodies, not parent/child, optionally filtered.
+Pairs that already overlap in the reference pose get a relaxed minimum distance, so
+the robot never fights its own default posture.
+"""
+
+from collections.abc import Callable
+from dataclasses import dataclass
+
+import jax.numpy as jnp
+import numpy as np
+from jax import Array
+
+from rigidax.kinematics import Kinematics, forward_kinematics
+from rigidax.model import Model, derived
+
+
+@dataclass(frozen=True)
+class SphereModel:
+    body: np.ndarray  # (S,) body index of each sphere
+    center: np.ndarray  # (S, 3) center in the body frame
+    radius: np.ndarray  # (S,)
+    pairs: np.ndarray  # (P, 2) sphere indices
+    min_dist: np.ndarray  # (P,) required clearance between sphere surfaces
+    # Row structure for the solver: each pair is owned by one "row" sphere, and each
+    # row constrains its closest partner. Fixed shapes, no sorting, stable meanings.
+    rows: np.ndarray = None  # (R,) sphere indices that own pairs
+    partner: np.ndarray = None  # (R, S) bool: partner spheres of each row
+    row_min_dist: np.ndarray = None  # (R, S) clearance per (row, partner)
+
+    @property
+    def num_spheres(self) -> int:
+        return len(self.radius)
+
+    def centers(self, kin: Kinematics) -> Array:
+        b = self.body
+        return kin.body_p[b] + jnp.einsum("sij,sj->si", kin.body_R[b], self.center)
+
+    def gaps(self, kin: Kinematics) -> Array:
+        """Surface distance minus required clearance per pair (< 0: violated)"""
+        C = self.centers(kin)
+        i, j = self.pairs[:, 0], self.pairs[:, 1]
+        return (jnp.linalg.norm(C[i] - C[j], axis=1) - self.radius[i] - self.radius[j]
+                - self.min_dist)
+
+    @staticmethod
+    def fit(m, model: Model, bodies: Callable[[str], bool] | None = None,
+            pair_filter: Callable[[str, str], bool] | None = None,
+            spacing: float = 0.05, cover: float = 95.0, margin: float = 0.01,
+            geom_group: int | None = 3, q_ref=None, extra=None) -> "SphereModel":
+        """Fit spheres to a compiled mujoco.MjModel `m` (the same model as `model`).
+
+        bodies: which bodies to include (by name); default: all with geometry.
+        pair_filter: which body pairs to check (by names); default: all eligible.
+        geom_group: geom group holding collision geometry (None: all geoms).
+        q_ref: reference pose for relaxing pre-existing overlaps (default qpos0).
+        extra: optional [(body name, centers (n, 3) in that body's frame, radii (n,))] of
+            spheres to attach rigidly, e.g. the envelope of parts not in the model
+            (fingers whose joints are not solved).
+        """
+        import mujoco
+
+        names = [mujoco.mj_id2name(m, mujoco.mjtObj.mjOBJ_BODY, b) or "" for b in range(m.nbody)]
+        sb, sc, sr = [], [], []
+        for b in range(1, m.nbody):
+            if bodies is not None and not bodies(names[b]):
+                continue
+            V = _body_points(m, b, geom_group)
+            if len(V) < 4:
+                continue
+            c, r = _fit_spheres(V, spacing, cover)
+            sb += [b] * len(r)
+            sc.append(c)
+            sr.append(r)
+        for name, c, r in (extra or ()):
+            b = m.body(name).id
+            if bodies is not None and not bodies(names[b]):
+                continue
+            sb += [b] * len(r)
+            sc.append(np.asarray(c, float).reshape(-1, 3))
+            sr.append(np.asarray(r, float).reshape(-1))
+        sb, sc, sr = np.array(sb, int), np.concatenate(sc), np.concatenate(sr)
+        parent = m.body_parentid
+        pairs = [(i, j) for i in range(len(sb)) for j in range(i + 1, len(sb))
+                 if sb[i] != sb[j] and parent[sb[i]] != sb[j] and parent[sb[j]] != sb[i]
+                 and (pair_filter is None or pair_filter(names[sb[i]], names[sb[j]]))]
+        pairs = np.array(pairs, int).reshape(-1, 2)
+        model_ = SphereModel(sb, sc, sr, pairs, np.zeros(len(pairs)))
+        kin = forward_kinematics(model, m.qpos0 if q_ref is None else q_ref)
+        rest = np.asarray(model_.gaps(kin))
+        min_dist = np.minimum(margin, rest - 0.005)
+        return SphereModel.with_rows(sb, sc, sr, pairs, min_dist)
+
+    @staticmethod
+    def with_rows(body, center, radius, pairs, min_dist) -> "SphereModel":
+        """Assign each pair to the endpoint with more partners (fewest rows overall)"""
+        S = len(radius)
+        degree = np.bincount(pairs.ravel(), minlength=S)
+        owner = np.where(degree[pairs[:, 0]] >= degree[pairs[:, 1]], pairs[:, 0], pairs[:, 1])
+        other = np.where(owner == pairs[:, 0], pairs[:, 1], pairs[:, 0])
+        rows = np.unique(owner)
+        r_of = {r: k for k, r in enumerate(rows)}
+        partner = np.zeros((len(rows), S), bool)
+        row_min = np.zeros((len(rows), S))
+        for o, t, md in zip(owner, other, min_dist):
+            partner[r_of[o], t] = True
+            row_min[r_of[o], t] = md
+        return SphereModel(body, center, radius, pairs, min_dist, rows, partner, row_min)
+
+    def num_rows(self, k: int = 1) -> int:
+        return len(self.rows) * min(k, self.partner.shape[1])
+
+    def constraint_rows(self, model: Model, kin: Kinematics, k: int = 1):
+        """Linearized constraints for each row sphere against its k closest partners:
+        rows a (R*k, nv) and bounds b (R*k,) such that  a @ dq >= b  (b = -gap)."""
+        import jax
+
+        C = self.centers(kin)
+        Cr = C[self.rows]
+        D = (jnp.linalg.norm(Cr[:, None] - C[None], axis=-1)
+             - self.radius[self.rows][:, None] - self.radius[None] - self.row_min_dist)
+        D = jnp.where(self.partner, D, jnp.inf)
+        k = min(k, D.shape[1])
+        if k == 1:
+            j = jnp.argmin(D, axis=1)[:, None]
+        else:
+            j = jax.lax.top_k(-D, k)[1]  # (R, k) closest partners
+        gap = jnp.take_along_axis(D, j, axis=1)  # (R, k)
+        gap = jnp.where(jnp.isfinite(gap), gap, 1.0)  # rows with < k partners: inactive
+        Cj = C[j]  # (R, k, 3)
+        delta = Cr[:, None] - Cj
+        n = delta / jnp.maximum(jnp.linalg.norm(delta, axis=-1, keepdims=True), 1e-9)
+        # a = n^T (J(c_i) - J(c_j)), with J(c) = (u + w x c) masked by the body's DOFs.
+        # Per DOF: n.(u + w x c) = u.n + w.(c x n): two small mat-vecs, no (S, 3, nv)
+        # Jacobian tensor.
+        w, u = kin.S[:, :3], kin.S[:, 3:]
+        mask = jnp.asarray(derived(model.topo).body_mask)  # (nbody, nv)
+        un = jnp.einsum("vk,rqk->rqv", u, n)
+        ti = jnp.einsum("vk,rqk->rqv", w, jnp.cross(Cr[:, None], n))
+        tj = jnp.einsum("vk,rqk->rqv", w, jnp.cross(Cj, n))
+        mi = mask[self.body[self.rows]][:, None]  # (R, 1, nv)
+        mj = mask[jnp.asarray(self.body)[j]]  # (R, k, nv)
+        a = mi * (un + ti) - mj * (un + tj)
+        return a.reshape(-1, a.shape[-1]), -gap.reshape(-1)
+
+
+def _quat_mat(q):
+    import mujoco
+
+    R = np.zeros(9)
+    mujoco.mju_quat2Mat(R, q)
+    return R.reshape(3, 3)
+
+
+def _primitive_points(gtype, size, n=400):
+    import mujoco
+
+    rng = np.random.default_rng(0)
+    G = mujoco.mjtGeom
+    u = rng.normal(size=(n, 3))
+    u /= np.linalg.norm(u, axis=1, keepdims=True)
+    if gtype == G.mjGEOM_SPHERE:
+        return u * size[0]
+    if gtype == G.mjGEOM_ELLIPSOID:
+        return u * size[:3]
+    if gtype == G.mjGEOM_BOX:
+        return rng.uniform(-1, 1, (n, 3)) * size[:3]
+    if gtype in (G.mjGEOM_CAPSULE, G.mjGEOM_CYLINDER):
+        z = rng.uniform(-size[1], size[1], n)
+        ring = u[:, :2] / np.linalg.norm(u[:, :2], axis=1, keepdims=True) * size[0]
+        pts = np.c_[ring, z]
+        if gtype == G.mjGEOM_CAPSULE:
+            pts = np.r_[pts, u * size[0] + [0, 0, size[1]], u * size[0] - [0, 0, size[1]]]
+        return pts
+    return np.zeros((0, 3))
+
+
+def _body_points(m, b, group):
+    import mujoco
+
+    out = []
+    for g in np.where(m.geom_bodyid == b)[0]:
+        if group is not None and m.geom_group[g] != group:
+            continue
+        if m.geom_type[g] == mujoco.mjtGeom.mjGEOM_MESH:
+            mid = m.geom_dataid[g]
+            v = m.mesh_vert[m.mesh_vertadr[mid]: m.mesh_vertadr[mid] + m.mesh_vertnum[mid]]
+        else:
+            v = _primitive_points(m.geom_type[g], m.geom_size[g])
+        if len(v):
+            out.append(v @ _quat_mat(m.geom_quat[g]).T + m.geom_pos[g])
+    return np.concatenate(out) if out else np.zeros((0, 3))
+
+
+def _fit_spheres(V, spacing, cover):
+    """Slice along the principal axis; split wide (anisotropic) slices along their
+    second axis so flat bodies like a torso get several tight spheres, not one fat one"""
+    mu = V.mean(0)
+    _, _, Vt = np.linalg.svd(V - mu, full_matrices=False)
+    axis = Vt[0]
+    t = (V - mu) @ axis
+    k = max(1, int(np.ceil((t.max() - t.min()) / spacing)))
+    edges = np.linspace(t.min(), t.max(), k + 1)
+    centers, radii = [], []
+    for a, b in zip(edges[:-1], edges[1:]):
+        sel = (t >= a) & (t <= b)
+        if sel.sum() < 4:
+            continue
+        P = V[sel]
+        c0 = P.mean(0)
+        flat = P - c0
+        flat -= np.outer(flat @ axis, axis)  # cross-section points
+        _, sv, Wt = np.linalg.svd(flat, full_matrices=False)
+        w_axis = Wt[0]
+        s_w = flat @ w_axis
+        ext_w = s_w.max() - s_w.min()
+        ext_d = np.ptp(flat @ Wt[1]) if len(Wt) > 1 else ext_w
+        # number of sub-spheres across the wide direction
+        n_sub = int(np.clip(np.round(ext_w / max(ext_d, 1e-6)), 1, 4))
+        sub_edges = np.linspace(s_w.min(), s_w.max(), n_sub + 1)
+        for sa, sb in zip(sub_edges[:-1], sub_edges[1:]):
+            ss = (s_w >= sa) & (s_w <= sb)
+            if ss.sum() < 4:
+                continue
+            c = P[ss].mean(0)
+            radial = P[ss] - c
+            radial -= np.outer(radial @ axis, axis)
+            r = max(np.percentile(np.linalg.norm(radial, axis=1), cover), 0.5 * (b - a))
+            centers.append(c)
+            radii.append(r)
+    return np.array(centers), np.array(radii)