mjlab.sim

Classes

Simulation SimulationCfg

MujocoCfg TorchArray

WarpBridge

Simulation

class mjlab.sim.Simulation

GPU-accelerated MuJoCo simulation powered by MJWarp.

CUDA Graph Capture

On CUDA devices with memory pools enabled, the simulation captures CUDA graphs for step(), forward(), and reset() operations. Graph capture records a sequence of GPU kernels and their memory addresses, then replays the entire sequence with a single kernel launch, eliminating CPU overhead from repeated kernel dispatches.

Important: A captured graph holds pointers to the GPU arrays that existed at capture time. If those arrays are later replaced (e.g., via expand_model_fields()), the graph will still read from the old arrays, silently ignoring any new values. The expand_model_fields() method handles this automatically by calling create_graph() after replacing arrays.

If you write code that replaces model or data arrays after simulation initialization, you must call create_graph() afterward to re-capture the graphs with the new memory addresses.

__init__(num_envs: int, cfg: SimulationCfg, model: MjModel, device: str)

create_graph() → None

Capture CUDA graphs for step, forward, and reset operations.

This method must be called whenever GPU arrays in the model or data are replaced after initialization. The captured graphs hold pointers to the arrays that existed at capture time. If those arrays are replaced, the graphs will silently read from the old arrays, ignoring any new values.

Called automatically by: - __init__() during simulation initialization - expand_model_fields() after replacing model arrays

On CPU devices or when memory pools are disabled, this is a no-op.

property default_model_fields: dict[str, Tensor]

Default values for expanded model fields, used in domain randomization.

property expanded_fields: set[str]

Names of model fields that have been expanded for per-env DR.

expand_model_fields(fields: tuple[str, ...]) → None

Expand model fields to support per-environment parameters.

get_default_field(field: str) → Tensor

Get the default value for a model field, caching for reuse.

Returns the original values from the C MuJoCo model (mj_model), obtained from the final compiled scene spec before any randomization is applied. Not to be confused with the GPU Warp model (wp_model) which may have randomized values.

recompute_constants(level: RecomputeLevel) → None

Recompute derived model constants after domain randomization.

Parameters:

level – Which constants to recompute. set_const is the most expensive (covers body_mass changes), set_const_0 covers qpos0/body_inertia/dof_armature changes, and set_const_fixed is the cheapest (covers body_gravcomp changes).

set_sensor_context(ctx: SensorContext) → None

Wire a SensorContext for camera/raycast sensing.

Automatically re-captures CUDA graphs so the sense_graph includes the new sensor kernels.

sense() → None

Execute the sense pipeline: prepare -> graph -> finalize.

Runs BVH refit, camera rendering, and raycasting in a single CUDA graph launch. Should be called once per env step, right before observation computation.

class mjlab.sim.SimulationCfg

nconmax: int | None = None

Number of contacts to allocate per world.

Contacts exist in large heterogenous arrays: one world may have more than nconmax contacts. If None, a heuristic value is used.

njmax: int | None = None

Number of constraints to allocate per world.

Constraint arrays are batched by world: no world may have more than njmax constraints. If None, a heuristic value is used.

ls_parallel: bool = True

contact_sensor_maxmatch: int = 64

mujoco: MujocoCfg

nan_guard: NanGuardCfg

class mjlab.sim.MujocoCfg

Configuration for MuJoCo simulation parameters.

timestep: float = 0.002

integrator: Literal['euler', 'implicitfast'] = 'implicitfast'

impratio: float = 1.0

cone: Literal['pyramidal', 'elliptic'] = 'pyramidal'

jacobian: Literal['auto', 'dense', 'sparse'] = 'auto'

solver: Literal['newton', 'cg', 'pgs'] = 'newton'

iterations: int = 100

tolerance: float = 1e-08

ls_iterations: int = 50

ls_tolerance: float = 0.01

ccd_iterations: int = 50

gravity: tuple[float, float, float] = (0.0, 0.0, -9.81)

multiccd: bool = False

apply(model: MjModel) → None

Apply configuration settings to a compiled MjModel.

TorchArray

class mjlab.sim.TorchArray

Warp array that behaves like a torch.Tensor with shared memory.

__init__(wp_array: warp.array, nworld: int | None = None) → None

Initialize the tensor proxy with a Warp array.

WarpBridge

class mjlab.sim.WarpBridge

Wraps mjwarp objects to expose Warp arrays as PyTorch tensors.

Automatically converts Warp array attributes to TorchArray objects on access, enabling direct PyTorch operations on simulation data. Recursively wraps nested structures that contain Warp arrays.

IMPORTANT: This wrapper is read-only. To modify array data, use in-place operations like obj.field[:] = value. Direct assignment like obj.field = new_array will raise an AttributeError to prevent accidental memory address changes that break CUDA graphs.

__init__(struct: T, nworld: int | None = None) → None

property struct: T

Access the underlying wrapped struct.

clear_cache() → None

Clear the wrapped cache to force re-wrapping of arrays.

This should be called after operations that modify the underlying warp arrays, such as expand_model_fields(), to ensure the cache reflects the updated arrays.