Implementation:Google deepmind Mujoco mjx put model

Knowledge Sources	MuJoCo MJX Documentation
Domains	GPU_Computing, Physics_Simulation, JAX
Last Updated	2026-02-15 06:00 GMT

Overview

Concrete tool for transferring a MuJoCo model to a JAX device provided by the MJX API.

Description

The mjx.put_model function converts a CPU mujoco.MjModel into an mjx.Model placed on the specified device. It supports multiple backend implementations: JAX (default), Warp, C, and CPP. The backend is auto-detected based on available hardware or can be specified explicitly.

Usage

Call after loading a model with mujoco.MjModel.from_xml_path. The default device is the first available GPU.

Code Reference

Source Location

Repository: mujoco
File: mjx/mujoco/mjx/_src/io.py
Lines: 530-566

Signature

def put_model(
    m: mujoco.MjModel,
    device: Optional[jax.Device] = None,
    impl: Optional[Union[str, types.Impl]] = None,
    graph_mode: Optional[mjxw.types.GraphMode] = None,
) -> types.Model:

Import

import mujoco.mjx as mjx

I/O Contract

Inputs

Name	Type	Required	Description
m	mujoco.MjModel	Yes	CPU-resident model
device	jax.Device	No	Target device (default: first GPU)
impl	str or types.Impl	No	Backend: 'jax', 'warp', 'c', 'cpp'
graph_mode	GraphMode	No	CUDA graph mode (Warp only)

Outputs

Name	Type	Description
return	types.Model	MJX model on device as JAX pytree

Usage Examples

import mujoco
import mujoco.mjx as mjx

# Load on CPU
m = mujoco.MjModel.from_xml_path("humanoid.xml")

# Transfer to default GPU
mx = mjx.put_model(m)

# Transfer to specific device
mx = mjx.put_model(m, device=jax.devices("gpu")[0])

# Use Warp backend
mx = mjx.put_model(m, impl="warp")

Related Pages

Implements Principle

Principle:Google_deepmind_Mujoco_MJX_Model_Transfer

Requires Environment

Uses Heuristic

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment