Implementation:Google deepmind Mujoco mjx put data make data

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

Overview

Concrete tools for creating or transferring MuJoCo simulation state to a JAX device provided by the MJX API.

Description

mjx.put_data transfers an existing CPU mujoco.MjData to the device, preserving the current simulation state. mjx.make_data creates a fresh, zero-initialized mjx.Data on the device without requiring a CPU MjData. Both return an mjx.Data pytree suitable for JAX transformations.

Usage

Use mjx.make_data for fresh simulations (standard for RL). Use mjx.put_data when transferring specific state from CPU.

Code Reference

Source Location

• Repository: mujoco
• File: mjx/mujoco/mjx/_src/io.py
• Lines (put_data): 1368-1420
• Lines (make_data): 958-1019

Signature

def put_data(
    m: mujoco.MjModel,
    d: mujoco.MjData,
    device: Optional[jax.Device] = None,
    impl: Optional[Union[str, types.Impl]] = None,
    nconmax: Optional[int] = None,
    naconmax: Optional[int] = None,
    njmax: Optional[int] = None,
) -> types.Data:

def make_data(
    m: Union[types.Model, mujoco.MjModel],
    device: Optional[jax.Device] = None,
    impl: Optional[Union[str, types.Impl]] = None,
    nconmax: Optional[int] = None,
    naconmax: Optional[int] = None,
    njmax: Optional[int] = None,
) -> types.Data:

Import

import mujoco.mjx as mjx

I/O Contract

Inputs (put_data)

Name	Type	Required	Description
m	mujoco.MjModel	Yes	CPU model
d	mujoco.MjData	Yes	CPU data to transfer
device	jax.Device	No	Target device
impl	str	No	Backend: 'jax', 'warp', 'c', 'cpp'

Inputs (make_data)

Name	Type	Required	Description
m	Model or MjModel	Yes	Model (MJX or CPU)
device	jax.Device	No	Target device
impl	str	No	Backend: 'jax', 'warp'
naconmax	int	No	Max contacts across all worlds (Warp)
njmax	int	No	Max constraints across all worlds (Warp)

Outputs

Name	Type	Description
return	types.Data	MJX data on device as JAX pytree

Usage Examples

import mujoco
import mujoco.mjx as mjx

m = mujoco.MjModel.from_xml_path("humanoid.xml")
mx = mjx.put_model(m)

# Option 1: Make fresh data on GPU
dx = mjx.make_data(mx)

# Option 2: Transfer existing CPU state
d = mujoco.MjData(m)
d.qpos[2] = 1.0  # Set initial height
dx = mjx.put_data(m, d)

Related Pages

Implements Principle

• Principle:Google_deepmind_Mujoco_MJX_Data_Transfer

Requires Environment

• Environment:Google_deepmind_Mujoco_MJX_JAX_Environment