Principle:Google deepmind Mujoco MJX Data Transfer
| Knowledge Sources | |
|---|---|
| Domains | GPU_Computing, Physics_Simulation, JAX |
| Last Updated | 2026-02-15 06:00 GMT |
Overview
Mechanism for creating or transferring simulation state data to a JAX device for GPU-accelerated physics computation.
Description
MJX Data Transfer provides two paths: mjx.put_data transfers an existing CPU MjData to the device (preserving its current state), while mjx.make_data allocates a fresh, zero-initialized mjx.Data on the device directly. The resulting mjx.Data is a JAX pytree that can be transformed with jit, vmap, and other JAX primitives.
Usage
Use mjx.make_data when starting a fresh simulation (most common for RL). Use mjx.put_data when transferring a specific simulation state from CPU to GPU.
Theoretical Basis
Data transfer involves:
- Field extraction: All state arrays (qpos, qvel, ctrl, etc.) are extracted from MjData
- Type conversion: C arrays are converted to JAX arrays with appropriate dtypes
- Device placement: Arrays are placed on the target device
- Pytree wrapping: Arrays are wrapped in an mjx.Data dataclass that JAX can trace
Related Pages
Implemented By
Page Connections
Double-click a node to navigate. Hold to expand connections.
Principle
Implementation
Heuristic
Environment