rigidax 1.0.0__tar.gz

rigidax-1.0.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Nishanth Sundaran
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

rigidax-1.0.0/PKG-INFO

@@ -0,0 +1,257 @@
+Metadata-Version: 2.4
+Name: rigidax
+Version: 1.0.0
+Summary: Differentiable rigid-body kinematics, dynamics and batched IK / motion retargeting in JAX (MJCF + URDF)
+Author-email: Nishanth Sundaran <sundharnishanth@gmail.com>
+License: MIT
+Project-URL: Homepage, https://github.com/NishanthSundaran/Rigidax
+Project-URL: Repository, https://github.com/NishanthSundaran/Rigidax
+Project-URL: Issues, https://github.com/NishanthSundaran/Rigidax/issues
+Project-URL: Changelog, https://github.com/NishanthSundaran/Rigidax/blob/master/CHANGELOG.md
+Keywords: robotics,kinematics,dynamics,inverse-kinematics,retargeting,jax,mujoco
+Classifier: Development Status :: 5 - Production/Stable
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Physics
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy
+Requires-Dist: jax>=0.4.30
+Requires-Dist: mujoco>=3.1
+Provides-Extra: test
+Requires-Dist: pytest; extra == "test"
+Requires-Dist: pin; extra == "test"
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: pin; extra == "dev"
+Requires-Dist: ruff; extra == "dev"
+Dynamic: license-file
+
+# rigidax
+
+Differentiable rigid-body kinematics and dynamics in JAX. Its design starts from
+[frax](https://github.com/StanfordASL/frax): the dense, loop-free formulation of the
+dynamics over a DOF ancestor mask. It fills in frax's gaps:
+
+| | frax 0.2 | rigidax |
+|---|---|---|
+| Model formats | URDF | **MJCF and URDF** (compiled by MuJoCo) |
+| Joint types | revolute, prismatic, free | hinge, slide, **ball**, free, **several joints per body** |
+| Gravity | hardcoded 9.81 in `forward_dynamics` | **model parameter**, any vector |
+| Physical parameters | frozen into the compiled function | **pytree leaves**: `jax.grad` / `vmap` over masses, inertias, offsets, gravity; swap them without recompiling |
+| Armature | no | yes (in M and inverse dynamics) |
+| External wrenches | per-link | per-body, MuJoCo `xfrc_applied` layout |
+| Ground truth tests | Pinocchio / MuJoCo | every quantity vs MuJoCo on random trees + G1/A2/Panda, plus Pinocchio |
+
+![Operational space control demo](https://raw.githubusercontent.com/NishanthSundaran/Rigidax/master/examples/osc_demo.gif)
+
+The Panda tracks a figure-eight with operational space control (`examples/osc_demo.py`).
+rigidax computes the kinematics, Jacobian, mass matrix and bias forces in about 54 µs
+per step, Python dispatch included. Mean tracking error is 0.9 mm. MuJoCo only
+simulates and renders.
+
+```
+env -u PYTHONPATH MUJOCO_GL=egl python examples/osc_demo.py --out osc_demo.mp4
+```
+
+![G1 + Dex3 hands, retargeted LAFAN1 motions](https://raw.githubusercontent.com/NishanthSundaran/Rigidax/master/examples/g1_lafan1.gif)
+
+Two output variants are produced from the same solve settings. Speed on the same
+laptop with self-collision: 22.8k frames/s for the plain G1 and 18.8k for G1 + Dex3. The
+Dex3 version is slower because the finger envelope adds collision spheres. frax, with
+self-collision, runs at 18.3k.
+
+| Variant | Command | Output | GIF |
+|---|---|---|---|
+| G1 + Dex3 hands (default) | `retarget_lafan1.py <clips>` | `data/g1_lafan1_hands/` (nq 50) | `examples/g1_lafan1.gif` |
+| Plain G1, rigid hands | `retarget_lafan1.py <clips> --no-hands` | `data/g1_lafan1_nohands/` (nq 36) | `examples/g1_lafan1_nohands.gif` |
+
+Both variants track the forearm twist. Palm floor contact and fingers are Dex3-only.
+Render a variant with `render_motions.py --data data/g1_lafan1_nohands ...`.
+
+## `rigidax.ik`: batched IK and retargeting
+
+A dedicated differential IK solver for retargeting whole datasets on CPU:
+
+- **Each Gauss-Newton step is a small QP**, solved by a purpose-built fixed-iteration
+  ADMM (`rigidax.ik.solve_qp`). It has no branches, so it vectorizes over thousands of
+  frames. It uses:
+  - row equilibration,
+  - one explicit inverse per solve, then matrix-vector products (batched triangular
+    solves are slow on XLA:CPU),
+  - warm starts from the previous frame,
+  - optional polishing.
+- **Constraints:**
+  - hard joint limits,
+  - a per-iteration trust region, so a joint can never flip across its range in one
+    frame,
+  - a hard floor (`FloorPoints`),
+  - elastic self-collision, an exact L1 penalty that keeps the problem always feasible,
+    on an automatic sphere model of any MuJoCo model (`SphereModel.fit`).
+- **Tasks:** position, orientation, limb direction (retargets between skeletons with
+  different proportions) and posture.
+- **Batching:** `solve_trajectory` splits clips into segments that run in parallel (vmap
+  within a device, pmap across devices). Each segment is seeded by a multi-start solve,
+  and consecutive segments are crossfaded so seams are smooth.
+
+```python
+from rigidax import ik
+tasks = [ik.position(model, bodies, w), ik.direction(model, links, bases, w),
+         ik.orientation(model, frames, w), ik.PostureTask(posture_w)]
+solver = ik.IKSolver(model, tasks, floor=ik.FloorPoints.on(model, feet, sole_points),
+                     collision=ik.SphereModel.fit(mj_model, model, pair_filter=...))
+qs, costs = solver.solve_trajectory(seed_fn, [clip1_targets, clip2_targets, ...])
+```
+
+### vs frax on LAFAN1 → Unitree G1
+
+Both libraries ran frax's own LAFAN1 settings on the same i9-11900H laptop (8 cores, one
+XLA device per core), on 8 clips (53,621 frames, 30 min of motion). Both outputs are scored
+with the same metrics:
+- **Motion kept:** robot vs human limb angular speed (1.0 = all of it).
+- **Lag:** in frames.
+- **Joint acceleration:** worst 5% of frames.
+- **Skating:** the robot's toe contact sliding while the human toe is planted.
+- **Penetration:** MuJoCo's own mesh collision check.
+
+`bench/bench_frax_ik.py` runs frax.
+
+| | frames/s | arm / leg dir. error p95 | hand / forearm rot. p95 | arms / legs kept | lag (frames) | >20° jumps | skate median / p95 | touchdown p95 | penetration | through floor |
+|---|---|---|---|---|---|---|---|---|---|---|
+| frax, self-collision | 18,327 | 17.4° / 11.6° | 13° / 20° | 0.88 / 0.95 | 0.16 | 1.9% | 4.6 / 32 cm/s | 40 | 0.0% | 0.6 cm max |
+| **rigidax, self-collision** (plain G1) | **22,836** | **14.0°** / 11.7° | **12°** / 24° | **0.92** / 0.94 | **-0.01** | **0.1%** | **0.9 / 21** | **35** | 0.0% | 3.8 cm max (lying) |
+| frax, no self-collision | 37,490 | 17.1° / 11.3° | 11° / 19° | 0.88 / 0.95 | 0.15 | 1.8% | 4.4 / 31 | 38 | 3.3% | |
+| **rigidax, no self-collision** | **46,808** | **3.9°** / 11.7° | 10° / 22° | **0.96** / 0.94 | -0.10 | **0.1%** | **0.9 / 21** | **35** | 2.2% | |
+
+The rigidax task set covers everything frax tracks: pelvis, the thigh / shin / foot links,
+torso, and the upper arm / forearm / hand orientations. It adds limb directions and
+toe contact points on top.
+
+Where frax is still slightly ahead:
+- forearm twist (24° vs 20°),
+- leg motion kept (0.94 vs 0.95).
+
+frax's lower joint-acceleration number (8.4 vs 9.3°/f²) is mostly smoothed-away motion.
+On the knees, rigidax matches the human's own knee acceleration and frax sits below it.
+
+Hands: the output uses the G1 with Dex3-1 hands (`g1_with_hands.xml`, 7 finger joints
+per hand). The body is solved on the 35-DOF model and the finger angles are written
+afterwards, since nothing depends on them, which keeps the solve as fast as before.
+- **Synthesized fingers:** LAFAN1 records no fingers, so the grip is inferred from context
+  and is not tracked motion:
+  - fist in a fighting guard (hand in front of the chest, in the chest frame) or during
+    a strike, in `fight` clips only;
+  - flat open hand while the human hand rests on the floor;
+  - relaxed otherwise.
+- **Palm contact:** a height task puts the palm on the floor while the human hand rests
+  there. The previous hover was 6.2 cm; now it is about 0.
+- **Forearm twist:** a twist-only task about the forearm axis.
+
+Diagnostics:
+- `examples/render_overlay.py`: draws the solver's targets over the robot (visible red
+  means a real tracking error).
+- `examples/floor_report.py`: whole-body floor penetration.
+- `examples/collision_report.py`: self-collision.
+
+```
+# CPU, one XLA device per core
+XLA_FLAGS="--xla_force_host_platform_device_count=8 --xla_cpu_multi_thread_eigen=false" \
+  python examples/retarget_lafan1.py dance1_subject1 walk1_subject1 ...   [--no-collision]
+python examples/collision_report.py dance1_subject1 ...   # MuJoCo mesh ground truth
+```
+
+LAFAN1 is CC BY-NC-ND 4.0 (non-commercial). Get it with:
+
+```
+curl -L -o data/lafan1.zip https://github.com/ubisoft/ubisoft-laforge-animation-dataset/raw/master/lafan1/lafan1.zip
+unzip -j data/lafan1.zip '*.bvh' -d data/lafan1
+MUJOCO_GL=egl python examples/render_motions.py --out g1_lafan1.gif dance1_subject1:DANCE \
+  dance2_subject3:"DANCE 2" fightAndSports1_subject1@13.8:"FIGHT + KICK" jumps1_subject1:JUMPS \
+  run1_subject2:RUN sprint1_subject2@177:SPRINT walk1_subject1:WALK fallAndGetUp1_subject1@63:"FALL + GET UP"
+# clip@seconds picks the window; without it the most active 10 s are used
+```
+
+![G1 motions side by side](https://raw.githubusercontent.com/NishanthSundaran/Rigidax/master/examples/g1_all.gif)
+
+![G1 whole-body control demo](https://raw.githubusercontent.com/NishanthSundaran/Rigidax/master/examples/g1_demo.gif)
+
+A free-floating Unitree G1 (35 DOF) squats, twists and waves while balancing on its
+feet (`examples/g1_demo.py`). Per step, rigidax computes:
+- M and h, the foot contact Jacobians, and Jc_dot·v (via autodiff)
+- reference accelerations that keep the feet planted (KKT solve)
+- the contact wrenches that balance the floating base
+- the feedforward joint torques
+
+Joint PD supplies the feedback. A step takes about 190 µs. MuJoCo's position servos
+are replaced by pure torque motors.
+
+```
+env -u PYTHONPATH MUJOCO_GL=egl python examples/g1_demo.py --out g1_demo.mp4
+# one motion set, frax-style grey studio: --motion squat|twist|wave|squat,wave|all --style light
+```
+
+## Use
+
+```python
+import jax, rigidax as rx
+jax.config.update("jax_enable_x64", True)
+
+model = rx.load("robot.xml")              # or .urdf; strip_meshes=True if meshes are missing
+q, v = model.q0, jax.numpy.zeros(model.nv)
+
+M, bias = jax.jit(rx.dynamics_terms)(model, q, v)            # one kinematics pass
+qacc = jax.jit(rx.forward_dynamics)(model, q, v, tau)
+moon = model.replace(gravity=jax.numpy.array([0, 0, -1.62]))  # no recompile
+
+# gradients w.r.t. physical parameters (e.g. system identification)
+dL_dmass = jax.grad(lambda m: loss(model.replace(body_mass=m)))(model.params.body_mass)
+```
+
+Available functions: `forward_kinematics`, `site_pose`, `body_pose`, `point_jacobian`,
+`site_jacobian`, `body_jacobian`, `center_of_mass`, `center_of_mass_jacobian`, `rnea`,
+`mass_matrix`, `bias_forces`, `gravity_forces`, `coriolis_forces`, `dynamics_terms`,
+`forward_dynamics`, `mass_matrix_inverse`, `integrate`, `difference`.
+Every function takes an optional `kin=` argument to reuse a forward-kinematics pass.
+
+Conventions follow MuJoCo: `q` uses `[pos, wxyz quat]` for free joints and `wxyz` for
+ball joints. `v` uses world-frame linear and body-frame angular velocity.
+
+## Tests
+
+```
+env -u PYTHONPATH python -m pytest      # unset ROS's PYTHONPATH if it is sourced
+```
+
+## Speed vs frax (CPU, float64, single thread, jax 0.4.30)
+
+Median µs per call, from `bench/bench_vs_frax.py`, two runs. A ratio above 1 means rigidax is faster.
+
+| | Panda single | G1 (35 dof) single | Panda batch 1024 | G1 batch 1024 |
+|---|---|---|---|---|
+| mass matrix | 0.96–0.98x | 0.97x | 1.09x | 1.02–1.15x |
+| bias forces | 0.85–0.91x | 0.78–0.80x | 0.95–1.02x | 0.91–0.92x |
+| forward dynamics | 0.82–0.87x | 0.95–0.96x | 0.97–1.11x | 0.94–0.97x |
+
+The two are about equal overall. The single-call bias computation is still 10–20% slower.
+GPU has not been benchmarked.
+
+## Known limitations (1.0)
+
+- **Synthesized fingers:** Dex3 finger motion is inferred from context (LAFAN1 has no
+  finger data). Finger collisions are handled with rigid envelopes, so fists held
+  against the torso in the fight guard still touch it in about 1% of frames.
+- **GPU:** the GPU path works but is not tuned (~6k frames/s on an RTX 3060 laptop).
+- **G1 forearm pronation:** on the G1 this lives on `wrist_roll_link`. The twist task
+  currently targets `elbow_link`.
+
+## Roadmap
+
+- Foot skating p95 (contact-aware per-frame weights), GPU throughput
+- Capsule/box collision primitives
+- Close the remaining bias-forces gap; GPU benchmarks
+- Analytic derivatives of RNEA and ABA

rigidax-1.0.0/README.md

@@ -0,0 +1,223 @@
+# rigidax
+
+Differentiable rigid-body kinematics and dynamics in JAX. Its design starts from
+[frax](https://github.com/StanfordASL/frax): the dense, loop-free formulation of the
+dynamics over a DOF ancestor mask. It fills in frax's gaps:
+
+| | frax 0.2 | rigidax |
+|---|---|---|
+| Model formats | URDF | **MJCF and URDF** (compiled by MuJoCo) |
+| Joint types | revolute, prismatic, free | hinge, slide, **ball**, free, **several joints per body** |
+| Gravity | hardcoded 9.81 in `forward_dynamics` | **model parameter**, any vector |
+| Physical parameters | frozen into the compiled function | **pytree leaves**: `jax.grad` / `vmap` over masses, inertias, offsets, gravity; swap them without recompiling |
+| Armature | no | yes (in M and inverse dynamics) |
+| External wrenches | per-link | per-body, MuJoCo `xfrc_applied` layout |
+| Ground truth tests | Pinocchio / MuJoCo | every quantity vs MuJoCo on random trees + G1/A2/Panda, plus Pinocchio |
+
+![Operational space control demo](https://raw.githubusercontent.com/NishanthSundaran/Rigidax/master/examples/osc_demo.gif)
+
+The Panda tracks a figure-eight with operational space control (`examples/osc_demo.py`).
+rigidax computes the kinematics, Jacobian, mass matrix and bias forces in about 54 µs
+per step, Python dispatch included. Mean tracking error is 0.9 mm. MuJoCo only
+simulates and renders.
+
+```
+env -u PYTHONPATH MUJOCO_GL=egl python examples/osc_demo.py --out osc_demo.mp4
+```
+
+![G1 + Dex3 hands, retargeted LAFAN1 motions](https://raw.githubusercontent.com/NishanthSundaran/Rigidax/master/examples/g1_lafan1.gif)
+
+Two output variants are produced from the same solve settings. Speed on the same
+laptop with self-collision: 22.8k frames/s for the plain G1 and 18.8k for G1 + Dex3. The
+Dex3 version is slower because the finger envelope adds collision spheres. frax, with
+self-collision, runs at 18.3k.
+
+| Variant | Command | Output | GIF |
+|---|---|---|---|
+| G1 + Dex3 hands (default) | `retarget_lafan1.py <clips>` | `data/g1_lafan1_hands/` (nq 50) | `examples/g1_lafan1.gif` |
+| Plain G1, rigid hands | `retarget_lafan1.py <clips> --no-hands` | `data/g1_lafan1_nohands/` (nq 36) | `examples/g1_lafan1_nohands.gif` |
+
+Both variants track the forearm twist. Palm floor contact and fingers are Dex3-only.
+Render a variant with `render_motions.py --data data/g1_lafan1_nohands ...`.
+
+## `rigidax.ik`: batched IK and retargeting
+
+A dedicated differential IK solver for retargeting whole datasets on CPU:
+
+- **Each Gauss-Newton step is a small QP**, solved by a purpose-built fixed-iteration
+  ADMM (`rigidax.ik.solve_qp`). It has no branches, so it vectorizes over thousands of
+  frames. It uses:
+  - row equilibration,
+  - one explicit inverse per solve, then matrix-vector products (batched triangular
+    solves are slow on XLA:CPU),
+  - warm starts from the previous frame,
+  - optional polishing.
+- **Constraints:**
+  - hard joint limits,
+  - a per-iteration trust region, so a joint can never flip across its range in one
+    frame,
+  - a hard floor (`FloorPoints`),
+  - elastic self-collision, an exact L1 penalty that keeps the problem always feasible,
+    on an automatic sphere model of any MuJoCo model (`SphereModel.fit`).
+- **Tasks:** position, orientation, limb direction (retargets between skeletons with
+  different proportions) and posture.
+- **Batching:** `solve_trajectory` splits clips into segments that run in parallel (vmap
+  within a device, pmap across devices). Each segment is seeded by a multi-start solve,
+  and consecutive segments are crossfaded so seams are smooth.
+
+```python
+from rigidax import ik
+tasks = [ik.position(model, bodies, w), ik.direction(model, links, bases, w),
+         ik.orientation(model, frames, w), ik.PostureTask(posture_w)]
+solver = ik.IKSolver(model, tasks, floor=ik.FloorPoints.on(model, feet, sole_points),
+                     collision=ik.SphereModel.fit(mj_model, model, pair_filter=...))
+qs, costs = solver.solve_trajectory(seed_fn, [clip1_targets, clip2_targets, ...])
+```
+
+### vs frax on LAFAN1 → Unitree G1
+
+Both libraries ran frax's own LAFAN1 settings on the same i9-11900H laptop (8 cores, one
+XLA device per core), on 8 clips (53,621 frames, 30 min of motion). Both outputs are scored
+with the same metrics:
+- **Motion kept:** robot vs human limb angular speed (1.0 = all of it).
+- **Lag:** in frames.
+- **Joint acceleration:** worst 5% of frames.
+- **Skating:** the robot's toe contact sliding while the human toe is planted.
+- **Penetration:** MuJoCo's own mesh collision check.
+
+`bench/bench_frax_ik.py` runs frax.
+
+| | frames/s | arm / leg dir. error p95 | hand / forearm rot. p95 | arms / legs kept | lag (frames) | >20° jumps | skate median / p95 | touchdown p95 | penetration | through floor |
+|---|---|---|---|---|---|---|---|---|---|---|
+| frax, self-collision | 18,327 | 17.4° / 11.6° | 13° / 20° | 0.88 / 0.95 | 0.16 | 1.9% | 4.6 / 32 cm/s | 40 | 0.0% | 0.6 cm max |
+| **rigidax, self-collision** (plain G1) | **22,836** | **14.0°** / 11.7° | **12°** / 24° | **0.92** / 0.94 | **-0.01** | **0.1%** | **0.9 / 21** | **35** | 0.0% | 3.8 cm max (lying) |
+| frax, no self-collision | 37,490 | 17.1° / 11.3° | 11° / 19° | 0.88 / 0.95 | 0.15 | 1.8% | 4.4 / 31 | 38 | 3.3% | |
+| **rigidax, no self-collision** | **46,808** | **3.9°** / 11.7° | 10° / 22° | **0.96** / 0.94 | -0.10 | **0.1%** | **0.9 / 21** | **35** | 2.2% | |
+
+The rigidax task set covers everything frax tracks: pelvis, the thigh / shin / foot links,
+torso, and the upper arm / forearm / hand orientations. It adds limb directions and
+toe contact points on top.
+
+Where frax is still slightly ahead:
+- forearm twist (24° vs 20°),
+- leg motion kept (0.94 vs 0.95).
+
+frax's lower joint-acceleration number (8.4 vs 9.3°/f²) is mostly smoothed-away motion.
+On the knees, rigidax matches the human's own knee acceleration and frax sits below it.
+
+Hands: the output uses the G1 with Dex3-1 hands (`g1_with_hands.xml`, 7 finger joints
+per hand). The body is solved on the 35-DOF model and the finger angles are written
+afterwards, since nothing depends on them, which keeps the solve as fast as before.
+- **Synthesized fingers:** LAFAN1 records no fingers, so the grip is inferred from context
+  and is not tracked motion:
+  - fist in a fighting guard (hand in front of the chest, in the chest frame) or during
+    a strike, in `fight` clips only;
+  - flat open hand while the human hand rests on the floor;
+  - relaxed otherwise.
+- **Palm contact:** a height task puts the palm on the floor while the human hand rests
+  there. The previous hover was 6.2 cm; now it is about 0.
+- **Forearm twist:** a twist-only task about the forearm axis.
+
+Diagnostics:
+- `examples/render_overlay.py`: draws the solver's targets over the robot (visible red
+  means a real tracking error).
+- `examples/floor_report.py`: whole-body floor penetration.
+- `examples/collision_report.py`: self-collision.
+
+```
+# CPU, one XLA device per core
+XLA_FLAGS="--xla_force_host_platform_device_count=8 --xla_cpu_multi_thread_eigen=false" \
+  python examples/retarget_lafan1.py dance1_subject1 walk1_subject1 ...   [--no-collision]
+python examples/collision_report.py dance1_subject1 ...   # MuJoCo mesh ground truth
+```
+
+LAFAN1 is CC BY-NC-ND 4.0 (non-commercial). Get it with:
+
+```
+curl -L -o data/lafan1.zip https://github.com/ubisoft/ubisoft-laforge-animation-dataset/raw/master/lafan1/lafan1.zip
+unzip -j data/lafan1.zip '*.bvh' -d data/lafan1
+MUJOCO_GL=egl python examples/render_motions.py --out g1_lafan1.gif dance1_subject1:DANCE \
+  dance2_subject3:"DANCE 2" fightAndSports1_subject1@13.8:"FIGHT + KICK" jumps1_subject1:JUMPS \
+  run1_subject2:RUN sprint1_subject2@177:SPRINT walk1_subject1:WALK fallAndGetUp1_subject1@63:"FALL + GET UP"
+# clip@seconds picks the window; without it the most active 10 s are used
+```
+
+![G1 motions side by side](https://raw.githubusercontent.com/NishanthSundaran/Rigidax/master/examples/g1_all.gif)
+
+![G1 whole-body control demo](https://raw.githubusercontent.com/NishanthSundaran/Rigidax/master/examples/g1_demo.gif)
+
+A free-floating Unitree G1 (35 DOF) squats, twists and waves while balancing on its
+feet (`examples/g1_demo.py`). Per step, rigidax computes:
+- M and h, the foot contact Jacobians, and Jc_dot·v (via autodiff)
+- reference accelerations that keep the feet planted (KKT solve)
+- the contact wrenches that balance the floating base
+- the feedforward joint torques
+
+Joint PD supplies the feedback. A step takes about 190 µs. MuJoCo's position servos
+are replaced by pure torque motors.
+
+```
+env -u PYTHONPATH MUJOCO_GL=egl python examples/g1_demo.py --out g1_demo.mp4
+# one motion set, frax-style grey studio: --motion squat|twist|wave|squat,wave|all --style light
+```
+
+## Use
+
+```python
+import jax, rigidax as rx
+jax.config.update("jax_enable_x64", True)
+
+model = rx.load("robot.xml")              # or .urdf; strip_meshes=True if meshes are missing
+q, v = model.q0, jax.numpy.zeros(model.nv)
+
+M, bias = jax.jit(rx.dynamics_terms)(model, q, v)            # one kinematics pass
+qacc = jax.jit(rx.forward_dynamics)(model, q, v, tau)
+moon = model.replace(gravity=jax.numpy.array([0, 0, -1.62]))  # no recompile
+
+# gradients w.r.t. physical parameters (e.g. system identification)
+dL_dmass = jax.grad(lambda m: loss(model.replace(body_mass=m)))(model.params.body_mass)
+```
+
+Available functions: `forward_kinematics`, `site_pose`, `body_pose`, `point_jacobian`,
+`site_jacobian`, `body_jacobian`, `center_of_mass`, `center_of_mass_jacobian`, `rnea`,
+`mass_matrix`, `bias_forces`, `gravity_forces`, `coriolis_forces`, `dynamics_terms`,
+`forward_dynamics`, `mass_matrix_inverse`, `integrate`, `difference`.
+Every function takes an optional `kin=` argument to reuse a forward-kinematics pass.
+
+Conventions follow MuJoCo: `q` uses `[pos, wxyz quat]` for free joints and `wxyz` for
+ball joints. `v` uses world-frame linear and body-frame angular velocity.
+
+## Tests
+
+```
+env -u PYTHONPATH python -m pytest      # unset ROS's PYTHONPATH if it is sourced
+```
+
+## Speed vs frax (CPU, float64, single thread, jax 0.4.30)
+
+Median µs per call, from `bench/bench_vs_frax.py`, two runs. A ratio above 1 means rigidax is faster.
+
+| | Panda single | G1 (35 dof) single | Panda batch 1024 | G1 batch 1024 |
+|---|---|---|---|---|
+| mass matrix | 0.96–0.98x | 0.97x | 1.09x | 1.02–1.15x |
+| bias forces | 0.85–0.91x | 0.78–0.80x | 0.95–1.02x | 0.91–0.92x |
+| forward dynamics | 0.82–0.87x | 0.95–0.96x | 0.97–1.11x | 0.94–0.97x |
+
+The two are about equal overall. The single-call bias computation is still 10–20% slower.
+GPU has not been benchmarked.
+
+## Known limitations (1.0)
+
+- **Synthesized fingers:** Dex3 finger motion is inferred from context (LAFAN1 has no
+  finger data). Finger collisions are handled with rigid envelopes, so fists held
+  against the torso in the fight guard still touch it in about 1% of frames.
+- **GPU:** the GPU path works but is not tuned (~6k frames/s on an RTX 3060 laptop).
+- **G1 forearm pronation:** on the G1 this lives on `wrist_roll_link`. The twist task
+  currently targets `elbow_link`.
+
+## Roadmap
+
+- Foot skating p95 (contact-aware per-frame weights), GPU throughput
+- Capsule/box collision primitives
+- Close the remaining bias-forces gap; GPU benchmarks
+- Analytic derivatives of RNEA and ABA

rigidax-1.0.0/pyproject.toml

@@ -0,0 +1,48 @@
+[build-system]
+requires = ["setuptools>=61"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "rigidax"
+version = "1.0.0"
+description = "Differentiable rigid-body kinematics, dynamics and batched IK / motion retargeting in JAX (MJCF + URDF)"
+readme = "README.md"
+license = { text = "MIT" }
+authors = [{ name = "Nishanth Sundaran", email = "sundharnishanth@gmail.com" }]
+requires-python = ">=3.10"
+keywords = ["robotics", "kinematics", "dynamics", "inverse-kinematics", "retargeting", "jax", "mujoco"]
+classifiers = [
+    "Development Status :: 5 - Production/Stable",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Topic :: Scientific/Engineering :: Physics",
+]
+dependencies = ["numpy", "jax>=0.4.30", "mujoco>=3.1"]
+
+[project.urls]
+Homepage = "https://github.com/NishanthSundaran/Rigidax"
+Repository = "https://github.com/NishanthSundaran/Rigidax"
+Issues = "https://github.com/NishanthSundaran/Rigidax/issues"
+Changelog = "https://github.com/NishanthSundaran/Rigidax/blob/master/CHANGELOG.md"
+
+[project.optional-dependencies]
+test = ["pytest", "pin"]
+dev = ["pytest", "pin", "ruff"]
+
+[tool.setuptools.packages.find]
+include = ["rigidax*"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+
+[tool.ruff]
+line-length = 100
+target-version = "py310"
+
+[tool.ruff.lint]
+select = ["E", "F", "W", "I"]
+ignore = ["E731", "E741"]

rigidax-1.0.0/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",
+]