frax

After defining your collision and self-collision models, use the provided visualization script to ensure they are loaded correctly in frax.

Run the following script to visualize the model:

python scripts/visualize_collision_model.py

Visual Cues in the Viewer:

• Yellow spheres: Standard collision spheres representing the robot's geometry.
• Red spheres: Spheres specifically included in the self-collision model.
• Lines: Indicate the specific collision pairs being monitored.

Once you have a collision model (a set of spheres), you must define specific self-collision pairs to monitor. Instead of checking every possible pair of spheres, define a subset of pairs that are most critical for practical use.

In your robot definition (e.g., frax/robots/your_robot.py), you specify pairs by providing:

• The link name for the first sphere.
• The sphere index for the first sphere.
• The link name for the second sphere.
• The sphere index for the second sphere.
• An optional tolerance/inflation factor to expand the collision volume of a sphere.

Best Practice: You can simplify the self-collision model by using a single 'inflated' sphere to represent a complex part (like an end-effector). This reduces the number of pairs to check while maintaining conservative collision behavior.

Follow these guidelines to maximize performance in frax:

• JIT Compilation: frax does not automatically wrap every method in @jax.jit. Always wrap your top-most function calls in a jitted region.
• Precision: Use double precision (jax.config.update("jax_enable_x64", True)) for high accuracy, especially for QP-based controllers. Note that on GPUs, double precision can cause a 2-6x slowdown.
• CPU vs GPU: If you are only simulating a single robot instance, using a CPU is often significantly faster than a GPU. Force CPU usage with jax.config.update("jax_platforms", "cpu").
• Threading: For typical robot controller design on CPU, restricting JAX/XLA to a single thread can improve performance and resource sharing.
• Library usage: Inside a jitted region, use jax.numpy. Outside of a jitted region, use standard numpy.

You can install frax via PyPI or from source. For GPU/TPU support, you can specify JAX installation tags like [cuda12], [cuda13], or [tpu]. If you want to run the included examples, install from source with the [examples] tag.

# From PyPI
pip install frax

# From source (recommended for examples)
git clone https://github.com/danielpmorton/frax
cd frax
pip install -e "[examples]"

If you are running FRAX on a single CPU backend, you can improve precision and speed by configuring specific environment variables. If these are not set, FRAX will issue a warning at runtime.

To optimize CPU performance, set the following environment variables:

Note: For best CPU performance, it is also recommended to use a JAX version earlier than 0.4.32 if possible.

export JAX_ENABLE_X64=1
export XLA_FLAGS='--xla_cpu_multi_thread_eigen=false'
export OPENBLAS_NUM_THREADS=1

To compute a robot's mass matrix (joint-space inertia matrix), load a robot using a URDF file and call the mass_matrix method. It is highly recommended to enable 64-bit precision for high accuracy and to wrap your calls in a jax.jit decorated function for performance.

import frax
import jax
import numpy as np

# Recommended for high accuracy
jax.config.update("jax_enable_x64", True)

robot = frax.Robot("path/to/your/robot.urdf")
q = np.zeros(robot.num_joints)

@jax.jit
def jit_mass_matrix(q_):
    return robot.mass_matrix(q_)

M = jit_mass_matrix(q)
print(M)

The following core classes are exported by the frax package for defining and interacting with robot models:

• Robot: The base class for robot models.
• Manipulator: A specialized class for manipulator-style robots.
• Humanoid: A specialized class for humanoid robots.

from frax import Robot, Manipulator, Humanoid

Frax provides helper functions to quickly load common robot models. You can import these directly from the frax package.

Available loaders:

• load_panda(): Loads a Franka Panda robot.
• load_g1(): Loads a Unitree G1 humanoid robot.

from frax import load_panda, load_g1

panda = load_panda()
g1 = load_g1()