Implementation:Google deepmind Mujoco MJX Types
| Knowledge Sources | |
|---|---|
| Domains | Physics_Simulation, JAX, Type_System |
| Last Updated | 2026-02-15 04:00 GMT |
Overview
Base type definitions for MJX including all enums, PyTreeNode dataclasses for Model, Data, Option, Contact, and the Impl backend selector.
Description
This module defines the complete type hierarchy used throughout MJX. It provides enum types mirroring MuJoCo's C enums (DisableBit, JointType, GeomType, SolverType, ConstraintType, SensorType, etc.) and PyTreeNode dataclasses for the core simulation state: Model (static model parameters), Data (mutable simulation state), Option (solver/integrator options), and Contact (collision results). The Impl enum selects between JAX, Warp, C, and CPP backend implementations. Separate ModelJAX/DataJAX and ModelC/DataC inner classes hold backend-specific fields, wrapped by the unified Model and Data outer classes.
Usage
Imported by every other MJX module. These types define the data structures that flow through the entire mjx.step() pipeline and are the primary interface for users interacting with MJX.
Code Reference
Source Location
- Repository: Google_deepmind_Mujoco
- File: mjx/mujoco/mjx/_src/types.py
- Lines: 1-1363
Key Classes
class Impl(enum.Enum): # Backend selector: JAX, WARP, C, CPP
class DisableBit(enum.IntFlag)
class EnableBit(enum.IntFlag)
class JointType(enum.IntEnum)
class IntegratorType(enum.IntEnum)
class GeomType(enum.IntEnum)
class ConvexMesh(PyTreeNode)
class ConeType(enum.IntEnum)
class JacobianType(enum.IntEnum)
class SolverType(enum.IntEnum)
class EqType(enum.IntEnum)
class WrapType(enum.IntEnum)
class TrnType(enum.IntEnum)
class DynType(enum.IntEnum)
class GainType(enum.IntEnum)
class BiasType(enum.IntEnum)
class ConstraintType(enum.IntEnum)
class CamLightType(enum.IntEnum)
class SensorType(enum.IntEnum)
class ObjType(PyTreeNode)
class Statistic(PyTreeNode)
class OptionJAX(PyTreeNode)
class Option(PyTreeNode)
class ModelJAX(PyTreeNode)
class Model(PyTreeNode)
class Contact(PyTreeNode)
class DataJAX(PyTreeNode)
class Data(PyTreeNode)
Import
from mujoco.mjx._src.types import Model
from mujoco.mjx._src.types import Data
from mujoco.mjx._src.types import GeomType
from mujoco.mjx._src.types import JointType
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| (N/A - type definitions) | -- | -- | This module defines types rather than processing inputs |
Outputs
| Name | Type | Description |
|---|---|---|
| Model | PyTreeNode | JAX-compatible model holding all static simulation parameters |
| Data | PyTreeNode | JAX-compatible mutable simulation state (positions, velocities, forces, contacts) |
| Contact | PyTreeNode | Contact data with dist, pos, frame, geom pairs, and friction |
| Option | PyTreeNode | Solver and integrator configuration options |