mjinx: GPU-accelerated numerical IK

[domrachev03]

Dear MuJoCo community,

I am delighted to announce a release of mjinx -- tool for solving numerical IK using jax and MuJoCo MJX to allow full auto-differentiation and GPU-accelerated batching.

The inverse kinematics problems are formulated in terms of tasks and barriers (inequality constraints) and solved either via differentiable QP solver to find locally optimal velocities, or via gradient descent to achieve global solution in joints.

Right now the following examples are available:

1. Batched IK for manipulator with constraints (local IK, global IK)
2. Batched IK for Go2 quadruped torso movement (local IK)
3. Batched IK for G1 biped hand movements with collision avoidance (local IK)

We also have plans to add examples with sampling-based Cassie walking via MPPI and learning policy together with numerical IK in MuJoCo Playground, so follow us for more news!

The whole idea was inspired by Kevin Zakka and his repository mink and also by similar pinocchio-based repository pink by Stéphane Caron. I am especially thankful to @kevinzakka, his work was the main foundation for the majority of functionality in mjinx. Finally, the project heavily relies on jaxlie for spacial transformations and optax for OSQP and nonlinear solvers implementation.

Happy hacking,

[kevinzakka]

An awesome job and visualization, congrats @domrachev03 and team! 👏