MjLab Framework¶

This section contains documentation for the MjLab framework adapter, which provides tools for exporting policies from MjLab to ONNX format using exploy.

Environment¶

The exportable environment adapter for MjLab.

class exploy.exporter.frameworks.mjlab.env.MjlabExportableEnvironment(env, policy_obs_group_name='actor')[source]¶

Bases: ExportableEnvironment

Wraps a MjLab ManagerBasedRlEnv for ONNX export.

Replaces each entity’s _data with an ONNX-traceable EntityDataSource proxy so that all observation computations use managed tensors that become ONNX graph inputs. Only entity data is replaced — scene sensors and other state are left intact.

After export, cleanup() restores the original entity data objects.

__init__(env, policy_obs_group_name='actor')[source]¶

property env: mjlab.envs.ManagerBasedRlEnv¶

property decimation: int¶: Return metadata about the environment required for export.

prepare_export()[source]¶

Prepare the environment for export. Called before each export.

Return type:: None

cleanup()[source]¶

Restore original entity data and raycaster sensor data.

Return type:: None

compute_observations()[source]¶

Compute and return the observations of the environment.

Return type:: Tensor

process_actions(actions)[source]¶

Process actions.

Return type:: None

apply_actions()[source]¶

Apply processed actions (e.g., joint targets) to the environment

Return type:: None

empty_actor_observations()[source]¶

Return type:: Tensor

empty_actions()[source]¶

Return type:: Tensor

validate()[source]¶

Return type:: bool

metadata()[source]¶

Return metadata about the environment required for export.

Return type:: dict[str, str]

get_observation_names()[source]¶

Get the names of the observations in the environment.

Return type:: list[str]

observations_reset()[source]¶

Get the observations after an environment reset.

Return type:: Tensor

register_evaluation_hooks(update, evaluate_substep)[source]¶

MjLab iterates scene._sensors at each substep and calls sensor.update(), and calls command_manager.compute() after each step. We inject dummy objects into those dictionaries so our callbacks fire at the right times.

Return type:: None

step(actions)[source]¶

Step the environment forward, delegating to mjlab’s own step loop.

Return type:: tuple[Tensor, bool]

Actor¶

Actor wrappers for RSL-RL v5 policy networks.

class exploy.exporter.frameworks.mjlab.actor.RslRlV5Actor(model)[source]¶

Bases: ExportableActor

Wraps an RSL-RL v5 MLPModel (non-recurrent) for ONNX export.

Converts the flat observation tensor into the TensorDict format that MLPModel.forward expects, using the model’s own obs_groups list to know which keys are needed.

__init__(model)[source]¶: Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(obs)[source]¶

Given a batch of observations, compute the corresponding actions.

Parameters:: obs (Tensor) – A tensor of shape (batch_size, obs_dim) containing the observations.
Return type:: Tensor

class exploy.exporter.frameworks.mjlab.actor.RslRlV5RecurrentActor(model)[source]¶

Bases: ExportableActor

Wraps an RSL-RL v5 RNNModel (recurrent) for ONNX export.

Same TensorDict wrapping as RslRlV5Actor but also exposes the RNN’s hidden state via get_state() so that add_actor_memory can register it as an ONNX input/output.

__init__(model)[source]¶: Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(obs)[source]¶

Given a batch of observations, compute the corresponding actions.

Parameters:: obs (Tensor) – A tensor of shape (batch_size, obs_dim) containing the observations.
Return type:: Tensor

reset(dones)[source]¶

Reset the actor’s internal state (e.g., RNN hidden states) based on the done flags.

Parameters:: dones (Tensor) – A tensor of shape (batch_size,) containing boolean flags indicating which environments have been reset.
Return type:: None

get_state()[source]¶

Get the actor’s internal state as a tuple of tensors, or None if there is no state.

Return type:: tuple[Tensor, ...] | None

exploy.exporter.frameworks.mjlab.actor.make_exportable_actor(env, actor_model, device)[source]¶

Create an exportable actor from an RSL-RL v5 MLPModel or RNNModel.

Parameters:

env (ExportableEnvironment) – The exportable environment that the actor will be used in.
actor_model (MLPModel | RNNModel) – The RSL-RL v5 actor model (PPO.actor).
device (str) – The device to place the actor on.

Return type:

ExportableActor

Inputs¶

Utilities for registering sensor and command inputs from MjLab managers.

exploy.exporter.frameworks.mjlab.inputs.add_base_com_vel(entities, context_manager)[source]¶

Add root COM linear and angular velocity inputs for all entities.

Reads from entity.data.root_com_lin_vel_b and root_com_ang_vel_b, which go through the EntityDataSource proxy and are therefore traced as ONNX graph inputs.

Parameters:

entities (dict) – Dict mapping entity name to Entity objects.
context_manager (ContextManager) – The context manager to add inputs to.

Return type:

None

exploy.exporter.frameworks.mjlab.inputs.add_body_pos_and_quat(entities, context_manager)[source]¶

Add body position and orientation inputs for all entities.

Reads from entity.data.body_link_pos_w and body_link_quat_w sliced per body, going through the proxy so the slice is an ONNX graph input.

Parameters:

entities (dict) – Dict mapping entity name to Entity objects.
context_manager (ContextManager) – The context manager to add inputs to.

Return type:

None

exploy.exporter.frameworks.mjlab.inputs.add_joint_pos_and_vel(entities, context_manager)[source]¶

Add joint position and velocity inputs for all entities.

Parameters:

entities (dict) – Dict mapping entity name to Entity objects.
context_manager (ContextManager) – The context manager to add inputs to.

Return type:

None

exploy.exporter.frameworks.mjlab.inputs.add_commands(command_manager, context_manager)[source]¶

Add command inputs from the command manager to the context manager.

Iterates all active command terms and adds them as inputs. Currently supports UniformVelocityCommand.

Parameters:

command_manager – The MjLab command manager containing active command terms.
context_manager (ContextManager) – The context manager to add command inputs to.

Return type:

None

exploy.exporter.frameworks.mjlab.inputs.add_sensor_inputs(sensors, context_manager)[source]¶

Add sensor inputs to the context manager.

Iterates all sensors and adds their data as inputs. Currently supports RayCastSensor.

Parameters:

sensors (dict) – Dictionary mapping sensor names to sensor objects.
context_manager (ContextManager) – The context manager to add sensor inputs to.

Return type:

None

Outputs¶

Utilities for registering action outputs from MjLab managers.

exploy.exporter.frameworks.mjlab.outputs.add_outputs(action_manager, context_manager)[source]¶

Add joint target outputs for all action terms in the action manager.

Iterates all active action terms and adds their joint target outputs. Currently supports BaseAction (and all its subclasses such as JointPositionAction, JointVelocityAction, JointEffortAction).

Parameters:

action_manager (ActionManager) – The MjLab action manager containing active action terms.
context_manager (ContextManager) – The context manager to add components to.

Return type:

None

Memory¶

Utilities for registering memory components from MjLab managers.

exploy.exporter.frameworks.mjlab.memory.add_memory(env, context_manager)[source]¶

Add memory components for the action manager.

Registers the latest actions and processed actions (one per active action term) as Memory components so they are available as ONNX graph inputs.

Parameters:

env – The MjLab ManagerBasedRlEnv.
context_manager (ContextManager) – The context manager to add components to.

Return type:

None

exploy.exporter.frameworks.mjlab.memory.add_term_memory(action_manager, context_manager, action_term_name, attr_name)[source]¶

Parameters:

action_manager – The MjLab action manager.
context_manager (ContextManager) – The context manager to add components to.
action_term_name (str) – Name of the action term in the action manager.
attr_name (str) – Attribute name on the action term to expose as memory.

Return type:

None

Data Sources¶

Adaptor classes that mirror MjLab data interfaces while managing their own tensor data for ONNX export.

class exploy.exporter.frameworks.mjlab.entity_data.EntityDataSource(entity)[source]¶

Bases: object

Mimic the interface of an EntityData, but manage its own ONNX-traceable tensor data.

Replaces entity._data during ONNX export so that observation functions access managed tensors instead of live simulation buffers. These managed tensors become ONNX graph inputs because they are the same Python objects returned by entity.data.* during tracing.

After export, the original entity._data is restored by MjlabExportableEnvironment.cleanup().

__init__(entity)[source]¶

update(*args, **kwargs)[source]¶

No-op update to mimic simulator data-source interfaces.

Return type:: None

reset(*args, **kwargs)[source]¶

No-op reset to mimic simulator data-source interfaces.

Return type:: None

property root_link_pose_w: Tensor¶: Root link pose [pos, quat] in simulation world frame. Shape is (num_instances, 7).

property root_com_pose_w: Tensor¶: Root COM pose [pos, quat] in simulation world frame. Shape is (num_instances, 7).

property projected_gravity_b: Tensor¶: Projection of gravity direction on base frame. Shape is (num_instances, 3).

property heading_w: Tensor¶: Yaw heading of the base frame (in radians). Shape is (num_instances,).

property joint_pos: Tensor¶: Joint positions. Shape is (num_instances, num_joints).

property joint_vel: Tensor¶: Joint velocities. Shape is (num_instances, num_joints).

property joint_acc: Tensor¶: Joint accelerations. Shape is (num_instances, num_joints).

property root_link_pos_w: Tensor¶

Root link position in world frame. Shape is (num_instances, 3).

This is the base stored data - position of the root body’s link frame.

property root_link_quat_w: Tensor¶

Root link orientation (w, x, y, z) in world frame. Shape is (num_instances, 4).

This is the base stored data - orientation of the root body’s link frame.

property root_link_vel_w: Tensor¶: Root link velocity in world frame. Shape is (num_instances, 6).

property root_link_lin_vel_w: Tensor¶

Root link linear velocity in world frame. Shape is (num_instances, 3).

Computed from COM velocity by subtracting the velocity offset due to COM offset.

property root_link_ang_vel_w: Tensor¶

Root link angular velocity in world frame. Shape is (num_instances, 3).

Angular velocity is the same for link and COM frames.

property root_link_lin_vel_b: Tensor¶: Root link linear velocity in base frame. Shape is (num_instances, 3).

property root_link_ang_vel_b: Tensor¶: Root link angular velocity in base frame. Shape is (num_instances, 3).

property root_com_pos_w: Tensor¶: Root COM position in world frame. Shape is (num_instances, 3).

property root_com_quat_w: Tensor¶: Root COM orientation (w, x, y, z) in world frame. Shape is (num_instances, 4).

property root_com_vel_w: Tensor¶: Root COM velocity in world frame. Shape is (num_instances, 6).

property root_com_lin_vel_w: Tensor¶

Root COM linear velocity in world frame. Shape is (num_instances, 3).

Computed by transforming from body frame to world frame.

property root_com_ang_vel_w: Tensor¶

Root COM angular velocity in world frame. Shape is (num_instances, 3).

Computed by transforming from body frame to world frame.

property root_com_lin_vel_b: Tensor¶

Root COM linear velocity in base frame. Shape is (num_instances, 3).

This is the base stored data - linear velocity of the root body’s COM in body frame.

property root_com_ang_vel_b: Tensor¶

Root COM angular velocity in base frame. Shape is (num_instances, 3).

This is the base stored data - angular velocity of the root body’s COM in body frame.

property body_link_pose_w: Tensor¶: Poses [pos, quat] of all body links in world frame. Shape is (num_instances, num_bodies, 7).

property body_link_pos_w: Tensor¶

Positions of all body links in world frame.

This is the base stored data - positions of all bodies’ link frames.

property body_link_quat_w: Tensor¶

Orientations (w, x, y, z) of all body links in world frame.

This is the base stored data - orientations of all bodies’ link frames.

property body_link_vel_w: Tensor¶: Velocities of all body links in world frame.

property body_link_lin_vel_w: Tensor¶

Linear velocities of all body links in world frame.

Computed from COM velocity by subtracting the velocity offset due to COM offset.

property body_link_ang_vel_w: Tensor¶

Angular velocities of all body links in world frame.

Angular velocity is the same for link and COM frames.

property body_com_pose_w: Tensor¶: Poses [pos, quat] of all body COMs in world frame. Shape is (num_instances, num_bodies, 7).

property body_com_pos_w: Tensor¶: Positions of all body COMs in world frame.

property body_com_quat_w: Tensor¶: Orientations (w, x, y, z) of all body COMs in world frame.

property body_com_vel_w: Tensor¶: Velocities of all body COMs in world frame.

property body_com_lin_vel_w: Tensor¶

Linear velocities of all body COMs in world frame.

Computed by transforming from body frame to world frame.

property body_com_ang_vel_w: Tensor¶

Angular velocities of all body COMs in world frame.

Computed by transforming from body frame to world frame.

property actuator_force: Tensor¶

property site_pose_w: Tensor¶

property site_pos_w: Tensor¶

property site_quat_w: Tensor¶

property site_vel_w: Tensor¶

property site_lin_vel_w: Tensor¶

property site_ang_vel_w: Tensor¶

property geom_pose_w: Tensor¶

property geom_pos_w: Tensor¶

property geom_quat_w: Tensor¶

property geom_vel_w: Tensor¶

property geom_lin_vel_w: Tensor¶

property geom_ang_vel_w: Tensor¶

property tendon_len: Tensor¶

property tendon_vel: Tensor¶

property joint_torques: Tensor¶

property joint_pos_biased: Tensor¶

property generalized_force: Tensor¶

property body_external_force: Tensor¶

property body_external_torque: Tensor¶

property body_external_wrench: Tensor¶

class exploy.exporter.frameworks.mjlab.raycaster_data.RayCasterDataSource(sensor, entities)[source]¶

Bases: object

Proxy for raycaster sensor data that holds ONNX-traceable tensors.

Replaces sensor._data during ONNX export so that raycaster height observations use managed tensors that become ONNX graph inputs.

__init__(sensor, entities)[source]¶

property hit_pos_w: Tensor¶

update(*args, **kwargs)[source]¶

Return type:: None

class exploy.exporter.frameworks.mjlab.sensor_proxy.MjLabBuiltinSensorProxy(original_sensor)[source]¶

Bases: object

Duck-type proxy for BuiltinSensor that returns a managed ONNX-traceable tensor.

Replaces a BuiltinSensor in scene._sensors so that calls to scene["robot/imu_lin_vel"].data (and similar) return a plain PyTorch tensor that can be traced for ONNX export.

After each simulation sense step, sync_from(original_sensor) must be called to refresh the managed tensor with the current sensor reading.

__init__(original_sensor)[source]¶

property data: Tensor¶

sync_from(original_sensor)[source]¶

Refresh the managed tensor from the live warp sensor buffer.

Reads directly from original_sensor._data_view (the raw warp tensor view) so that the most recent sim.sense() result is captured without going through the sensor’s internal cache.

Return type:: None

Utilities¶

Helper functions for MjLab integration.

exploy.exporter.frameworks.mjlab.utils.get_observation_names(observation_manager, group_name='actor')[source]¶

Compute a list of observation names for the flattened observation buffer.

Parameters:

observation_manager (ObservationManager) – The MjLab observation manager.
group_name (str) – The observation group (default: "actor").

Return type:

list[str]

Returns:

A list of names for each scalar in the flattened observation buffer.

exploy.exporter.frameworks.mjlab.utils.get_entity_actuator_gains(entity)[source]¶

Extract stiffness and damping gains from a MjLab entity’s actuator config.

Iterates the entity’s actuators and returns a dict mapping joint name to {"stiffness": float, "damping": float}. Only actuators that expose cfg.stiffness / cfg.damping (e.g. IdealPdActuator) contribute gains; other actuator types map to zeros.

Parameters:: entity (Entity) – A MjLab entity (robot) with actuators.
Return type:: dict[str, dict[str, float]]
Returns:: Dict mapping joint/target name to gain values.