Implementation:Google deepmind Mujoco MJX Warp Smooth

Knowledge Sources	Google_deepmind_Mujoco
Domains	Physics Simulation, Kinematics, GPU Computing, JAX-Warp Bridge
Last Updated	2026-02-15 04:00 GMT

Overview

Auto-generated bridge module that dispatches MuJoCo smooth dynamics (kinematics) computations from JAX to NVIDIA Warp GPU kernels.

Description

smooth.py is an auto-generated shim file that bridges MJX's JAX pipeline to Warp for smooth dynamics computations, primarily forward kinematics. It defines _kinematics_shim decorated with @ffi.format_args_for_warp that accepts body kinematic tree parameters (positions, orientations, parent IDs, weld IDs, joint definitions) and computes body, geom, and site transforms. The function maps the hierarchical body tree traversal using branch start/end arrays, processing joint types (free, ball, slide, hinge) to produce world-frame positions and rotation matrices.

Usage

This module is called during the forward simulation pipeline to compute kinematic quantities such as body positions (xpos), orientations (xquat, xmat), geom transforms, and site transforms. It is invoked before collision detection and force computation.

Code Reference

Source Location

• Repository: Google_deepmind_Mujoco
• File: mjx/mujoco/mjx/warp/smooth.py
• Lines: 1-471

Key Functions

@ffi.format_args_for_warp
def _kinematics_shim(
    # Model
    nworld: int,
    body_branch_start: wp.array(dtype=int),
    body_branches: wp.array(dtype=int),
    body_ipos: wp.array2d(dtype=wp.vec3),
    body_iquat: wp.array2d(dtype=wp.quat),
    body_jntadr: wp.array(dtype=int),
    body_jntnum: wp.array(dtype=int),
    body_mocapid: wp.array(dtype=int),
    body_parentid: wp.array(dtype=int),
    body_pos: wp.array2d(dtype=wp.vec3),
    body_quat: wp.array2d(dtype=wp.quat),
    ...
    # Data
    geom_xmat: wp.array2d(dtype=wp.mat33),
    ...
)

Import

from mujoco.mjx.warp import smooth

I/O Contract

Inputs

Name	Type	Required	Description
nworld	int	Yes	Number of parallel simulation worlds
body_branch_start	wp.array(int)	Yes	Start indices for kinematic tree branches
body_branches	wp.array(int)	Yes	Branch body indices for parallel tree traversal
body_ipos	wp.array2d(vec3)	Yes	Body inertial frame positions
body_iquat	wp.array2d(quat)	Yes	Body inertial frame orientations
body_parentid	wp.array(int)	Yes	Parent body IDs in the kinematic tree
jnt_axis	wp.array2d(vec3)	Yes	Joint rotation/translation axes
jnt_type	wp.array(int)	Yes	Joint type identifiers (free, ball, slide, hinge)
qpos0	wp.array2d(float)	Yes	Reference configuration joint positions

Outputs

Name	Type	Description
geom_xmat	wp.array2d(mat33)	World-frame geom rotation matrices
geom_xpos	wp.array2d(vec3)	World-frame geom positions
xpos	wp.array2d(vec3)	World-frame body positions
xquat	wp.array2d(quat)	World-frame body orientations
site_xmat	wp.array2d(mat33)	World-frame site rotation matrices
site_xpos	wp.array2d(vec3)	World-frame site positions

Related Pages

• Google_deepmind_Mujoco_MJX_Warp_Forward - Forward dynamics pipeline that invokes kinematics
• Google_deepmind_Mujoco_MJX_Warp_FFI - FFI decorator and argument flattening used by this module
• Google_deepmind_Mujoco_MJX_Warp_Types - Type definitions for Model, Data, and related dataclasses