mjinx 0.1.0

Numerical Inverse Kinematics based on JAX + MJX

MJINX

Mjinx is a library for auto-differentiable numerical inverse kinematics, powered by JAX and Mujoco MJX. The library was heavily inspired by the similar Pinocchio-based tool pink and Mujoco-based analogue mink.

Key features

1. Flexibility. Each control problem is assembled via Components, which enforce desired behaviour or keeps system in a safety set.
2. Different solution approaches. JAX (i.e. it's efficient sampling and autodifferentiation) allows to implement variety of solvers, which might be more beneficial in different scenarios.
3. Fully Jax-compatible. Both optimal control problem and its solver are jax-compatible: jit-compilation and automatic vectorization are available for the whole problem.
4. Convinience. The functionality is nicely wrapped to make the interaction with it easier.

Installation

The package is available in PyPI registry, and could be installed via pip:

pip install mjinx

Different installation versions:

1. Visualization tool mjinx.visualization.BatchVisualizer is available in mjinx[visual]
2. To run examples, install mjinx[examples]
3. To install development version, install mjinx[dev] (preferably in editable mode)
4. To build docs, install mjinx[docs]
5. To install the repository with all dependencies, install mjikx[all]

Usage

Here is the example of mjinx usage:

from mujoco import mjx mjx
from mjinx.problem import Problem

# Initialize the robot model using MuJoCo
MJCF_PATH = "path_to_mjcf.xml"
mj_model = mj.MjModel.from_xml_path(MJCF_PATH)
mjx_model = mjx.put_model(mj_model)

# Create instance of the problem
problem = Problem(mjx_model)

# Add tasks to track desired behavior
frame_task = FrameTask("ee_task", cost=1, gain=20, body_name="link7")
problem.add_component(frame_task)

# Add barriers to keep robot in a safety set
joints_barrier = JointBarrier("jnt_range", gain=10)
problem.add_component(joints_barrier)

# Initialize the solver
solver = LocalIKSolver(mjx_model)

# Initializing initial condition
q0 = np.zeros(7)

# Initialize solver data
solver_data = solver.init()

# jit-compiling solve and integrate 
solve_jit = jax.jit(solver.solve)
integrate_jit = jax.jit(integrate, static_argnames=["dt"])

# === Control loop ===
for t in np.arange(0, 5, 1e-2):
    # Changing problem and compiling it
    frame_task.target_frame = np.array([0.1 * np.sin(t), 0.1 * np.cos(t), 0.1, 1, 0, 0,])
    problem_data = problem.compile()

    # Solving the instance of the problem
    opt_solution, solver_data = solve_jit(q, solver_data, problem_data)

    # Integrating
    q = integrate_jit(
        mjx_model,
        q,
        opt_solution.v_opt,
        dt,
    )

Examples

The list of examples includes:

1. Kuka iiwa local inverse kinematics (single item, vmap over desired trajectory)
2. Kuka iiwa global inverse kinematics (single item, vmap over desired trajectory)
3. Go2 batched squats example

Contributing

We are always open for the suggestions and contributions. For contribution guidelines, see the CONTRIBUTING.md file.

Acknowledgement

The repository was highly inspired by pink and mink. Both authors, Stéphane Caron and Kevin Zakka, deeply inspire me to study robotics and contribute to the open source. Without them, this repository would not exist.

Some utility functions in this code are taken from source code of MuJoCo MJX. Apart from being a wonderful tool for batched computations and ML, the source code is compact yet readable and informative, and I encourage everyone to take a look at it to learn a little bit more about physical simulations, jax, and MuJoCo in general.