Implementation:Google deepmind Mujoco MJWarp IO

Knowledge Sources	Google_deepmind_Mujoco MJWarp Documentation
Domains	Physics_Simulation, GPU_Computing, NVIDIA_Warp
Last Updated	2026-02-15 04:00 GMT

Overview

MJWarp IO handles model and data transfer between MuJoCo CPU structures and Warp GPU arrays, plus constant initialization and render context creation.

Description

This module provides the I/O bridge between MuJoCo's C-based MjModel/MjData and Warp's GPU-resident Model/Data structures. It includes functions for uploading models to the GPU (put_model), creating and populating simulation data (make_data, put_data), retrieving results back to CPU (get_data_into), and resetting simulation state (reset_data). The module also computes model constants (set_const) including subtree masses, inertia-weighted quantities, and Jacobian-derived metrics, and can create render contexts with ray arrays for GPU rendering.

Usage

Used at initialization to upload MuJoCo models and data to the GPU, and during simulation to transfer data back for visualization or analysis. Also called when reinitializing simulation state.

Code Reference

Source Location

• Repository: Google_deepmind_Mujoco
• File: mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py
• Lines: 1-2437

Key Functions

# Model transfer
def put_model(mjm: mujoco.MjModel) -> Model
def is_sparse(mjm: mujoco.MjModel) -> bool

# Data transfer
def make_data(mjm: mujoco.MjModel, ...) -> Data
def put_data(mjm: mujoco.MjModel, mjd: mujoco.MjData, ...) -> Data
def get_data_into(mjd: mujoco.MjData, mjm: mujoco.MjModel, d: Data)
def reset_data(m: Model, d: Data, reset: Optional[wp.array] = None)

# Constants initialization
def set_const(m: Model, d: Data)
def set_const_0(m: Model, d: Data)
def set_const_fixed(m: Model, d: Data)

# Rendering
def create_render_context(...)

# Model overrides
def override_model(model, overrides)
def find_keys(model, keyname_prefix) -> list[int]
def make_trajectory(model, keys) -> np.ndarray

Import

from mujoco.mjx.third_party.mujoco_warp._src.io import put_model
from mujoco.mjx.third_party.mujoco_warp._src.io import put_data
from mujoco.mjx.third_party.mujoco_warp._src.io import make_data
from mujoco.mjx.third_party.mujoco_warp._src.io import get_data_into

I/O Contract

Inputs

Name	Type	Required	Description
mjm	mujoco.MjModel	Yes	MuJoCo C model to transfer to GPU
mjd	mujoco.MjData	No	MuJoCo C data to transfer to GPU
nconmax	int	No	Maximum number of contacts (auto-computed if not provided)
njmax	int	No	Maximum number of constraints per world
nworld	int	No	Number of parallel simulation worlds (default 1)

Outputs

Name	Type	Description
Model	types.Model	GPU-resident model with all parameters
Data	types.Data	GPU-resident simulation data
RenderContext	types.RenderContext	Render context with BVH and ray data

Related Pages

• Environment:Google_deepmind_Mujoco_MJX_Warp_CUDA_Environment