Implementation:Google deepmind Mujoco MJWarp Math
| Knowledge Sources | |
|---|---|
| Domains | Physics_Simulation, GPU_Computing, NVIDIA_Warp |
| Last Updated | 2026-02-15 04:00 GMT |
Overview
MJWarp Math provides core mathematical utility functions for quaternion arithmetic, rotation conversions, and spatial algebra used throughout MuJoCo Warp.
Description
This module contains foundational GPU-compatible mathematical functions implemented as Warp functions. It includes quaternion multiplication, quaternion-to-matrix conversion, axis-angle-to-quaternion conversion, rotation of vectors by quaternions, and spatial motion cross products. These functions form the mathematical backbone used by the kinematics, dynamics, and collision subsystems. All functions are decorated with @wp.func for GPU execution.
Usage
Imported extensively by nearly every other MJWarp module for quaternion operations, rotation matrix conversions, and spatial algebra computations in kinematics and dynamics.
Code Reference
Source Location
- Repository: Google_deepmind_Mujoco
- File: mjx/mujoco/mjx/third_party/mujoco_warp/_src/math.py
- Lines: 1-304
Key Functions
# Quaternion operations
def mul_quat(u: wp.quat, v: wp.quat) -> wp.quat
def quat_mul_axis(q: wp.quat, axis: wp.vec3f) -> wp.quat
def rot_vec_quat(vec: wp.vec3, quat: wp.quat) -> wp.vec3
def axis_angle_to_quat(axis: wp.vec3, angle: float) -> wp.quat
def quat_to_mat(quat: wp.quat) -> wp.mat33
# Spatial algebra
def motion_cross(...) # Motion cross product for spatial vectors
# Frame and indexing utilities
def make_frame(...)
def upper_trid_index(...)
def upper_tri_index(...)
def safe_div(...)
Import
from mujoco.mjx.third_party.mujoco_warp._src.math import mul_quat
from mujoco.mjx.third_party.mujoco_warp._src.math import quat_to_mat
from mujoco.mjx.third_party.mujoco_warp._src.math import rot_vec_quat
from mujoco.mjx.third_party.mujoco_warp._src.math import make_frame
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| quat | wp.quat | Yes | Quaternion (w, x, y, z) for rotation operations |
| vec | wp.vec3 | Yes | 3D vector for rotation or spatial operations |
| axis | wp.vec3 | No | Axis for axis-angle conversions |
| angle | float | No | Angle in radians for axis-angle conversions |
Outputs
| Name | Type | Description |
|---|---|---|
| result | wp.quat / wp.mat33 / wp.vec3 | Transformed quaternion, rotation matrix, or vector |