[software] 添加16DOF早期训练仿真与Sim2Real闭环
|
After Width: | Height: | Size: 27 KiB |
|
After Width: | Height: | Size: 1.5 MiB |
|
After Width: | Height: | Size: 253 KiB |
|
After Width: | Height: | Size: 104 KiB |
|
After Width: | Height: | Size: 285 KiB |
|
After Width: | Height: | Size: 26 KiB |
|
After Width: | Height: | Size: 18 KiB |
|
After Width: | Height: | Size: 751 KiB |
|
After Width: | Height: | Size: 1.9 MiB |
|
After Width: | Height: | Size: 78 KiB |
|
After Width: | Height: | Size: 20 KiB |
@@ -0,0 +1,172 @@
|
||||
/*
|
||||
* PyData Sphinx Theme — Option A (Indigo/Teal)
|
||||
* Aesthetic: modern lab — indigo primary, teal accent, neutral grays
|
||||
*/
|
||||
|
||||
/* LIGHT THEME */
|
||||
html[data-theme="light"] {
|
||||
/* Brand */
|
||||
--pst-color-primary: #4F46E5;
|
||||
/* Indigo-600 */
|
||||
--pst-color-secondary: #14B8A6;
|
||||
/* Teal-500 */
|
||||
--pst-color-secondary-highlight: #2DD4BF;
|
||||
/* Teal-400 */
|
||||
|
||||
/* Links / code links */
|
||||
--pst-color-inline-code-links: #0D9488;
|
||||
/* Teal-600 */
|
||||
--pst-color-link: var(--pst-color-primary);
|
||||
--pst-color-link-hover: #4338CA;
|
||||
/* Indigo-700 */
|
||||
|
||||
/* Semantic */
|
||||
--pst-color-info: var(--pst-color-secondary);
|
||||
--pst-color-info-highlight: var(--pst-color-secondary);
|
||||
--pst-color-info-bg: #D1FAE5;
|
||||
/* Teal-50 */
|
||||
--pst-color-attention: #F59E0B;
|
||||
/* Amber-500 */
|
||||
--pst-color-target: #EEF2FF;
|
||||
/* Indigo-50 */
|
||||
|
||||
/* Text */
|
||||
--pst-color-text-base: #1F2937;
|
||||
/* Slate-800 */
|
||||
--pst-color-text-muted: #6B7280;
|
||||
/* Slate-500 */
|
||||
|
||||
/* Surfaces */
|
||||
--pst-color-background: #FFFFFF;
|
||||
--pst-color-on-background: #FFFFFF;
|
||||
--pst-color-surface: #F3F4F6;
|
||||
/* Gray-100 */
|
||||
--pst-color-on-surface: #E5E7EB;
|
||||
/* Gray-200 */
|
||||
--pst-color-shadow: #D1D5DB;
|
||||
--pst-color-border: #E5E7EB;
|
||||
|
||||
/* Inline code */
|
||||
--pst-color-inline-code: #0D9488;
|
||||
/* Teal-600 */
|
||||
|
||||
/* Tables / hovers */
|
||||
--pst-color-table-row-hover-bg: #EEF2FF;
|
||||
/* Indigo-50 */
|
||||
|
||||
/* Accent (sparingly) */
|
||||
--pst-color-accent: #10B981;
|
||||
/* Emerald-500 */
|
||||
}
|
||||
|
||||
/* DARK THEME */
|
||||
html[data-theme="dark"] {
|
||||
/* Brand */
|
||||
--pst-color-primary: #A5B4FC;
|
||||
/* Indigo-300/200 mix for readability */
|
||||
--pst-color-secondary: #5EEAD4;
|
||||
/* Teal-300 */
|
||||
--pst-color-secondary-highlight: #2DD4BF;
|
||||
|
||||
/* Links / code links */
|
||||
--pst-color-inline-code-links: #93C5FD;
|
||||
/* Indigo-300 */
|
||||
--pst-color-link: var(--pst-color-primary);
|
||||
--pst-color-link-hover: #818CF8;
|
||||
/* Indigo-400 */
|
||||
|
||||
/* Semantic */
|
||||
--pst-color-info: var(--pst-color-secondary);
|
||||
--pst-color-info-highlight: var(--pst-color-secondary);
|
||||
--pst-color-info-bg: #042F2E;
|
||||
/* Deep teal */
|
||||
--pst-color-attention: #F59E0B;
|
||||
--pst-color-target: #1B1C2A;
|
||||
/* Indigo-tinted surface */
|
||||
|
||||
/* Text */
|
||||
--pst-color-text-base: #E5E7EB;
|
||||
/* Gray-200 */
|
||||
--pst-color-text-muted: #9CA3AF;
|
||||
/* Gray-400 */
|
||||
|
||||
/* Surfaces */
|
||||
--pst-color-background: #0B0C10;
|
||||
/* Deep graphite */
|
||||
--pst-color-on-background: #12131A;
|
||||
--pst-color-surface: #111827;
|
||||
/* Slate-900 */
|
||||
--pst-color-on-surface: #1F2937;
|
||||
/* Slate-800 */
|
||||
--pst-color-shadow: #0F172A;
|
||||
--pst-color-border: #2A2D3A;
|
||||
|
||||
/* Inline code */
|
||||
--pst-color-inline-code: #5EEAD4;
|
||||
/* Teal-300 */
|
||||
|
||||
/* Tables / hovers */
|
||||
--pst-color-table-row-hover-bg: #1B1C2A;
|
||||
|
||||
/* Accent */
|
||||
--pst-color-accent: #34D399;
|
||||
/* Emerald-400 */
|
||||
}
|
||||
|
||||
/* General tweaks */
|
||||
a {
|
||||
text-decoration: none !important;
|
||||
}
|
||||
|
||||
.bd-header-announcement a,
|
||||
.bd-header-version-warning a {
|
||||
color: #5EEAD4;
|
||||
}
|
||||
|
||||
.form-control {
|
||||
border-radius: 0 !important;
|
||||
border: none !important;
|
||||
outline: none !important;
|
||||
}
|
||||
|
||||
.navbar-brand,
|
||||
.navbar-icon-links {
|
||||
padding-top: 0rem !important;
|
||||
padding-bottom: 0rem !important;
|
||||
}
|
||||
|
||||
/* Version switcher */
|
||||
.sidebar-version-switcher {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
gap: 0.5rem;
|
||||
padding: 0.4rem 1rem;
|
||||
margin-bottom: 0.5rem;
|
||||
}
|
||||
|
||||
.sidebar-version-label {
|
||||
font-size: 0.8rem;
|
||||
font-weight: 600;
|
||||
color: var(--pst-color-text-muted);
|
||||
white-space: nowrap;
|
||||
}
|
||||
|
||||
.sidebar-version-select {
|
||||
flex: 1;
|
||||
font-size: 0.8rem;
|
||||
padding: 0.25rem 0.5rem;
|
||||
border: 1px solid var(--pst-color-border);
|
||||
border-radius: 4px;
|
||||
background: var(--pst-color-background);
|
||||
color: var(--pst-color-text-base);
|
||||
cursor: pointer;
|
||||
}
|
||||
|
||||
.sidebar-version-select:hover {
|
||||
border-color: var(--pst-color-primary);
|
||||
}
|
||||
|
||||
/* Sidebar section spacing */
|
||||
.bd-sidebar .navbar-icon-links {
|
||||
padding: 0 1rem 0.25rem !important;
|
||||
}
|
||||
|
After Width: | Height: | Size: 442 KiB |
|
After Width: | Height: | Size: 838 KiB |
|
After Width: | Height: | Size: 76 KiB |
|
After Width: | Height: | Size: 1.3 MiB |
|
After Width: | Height: | Size: 751 KiB |
|
After Width: | Height: | Size: 1.9 MiB |
|
After Width: | Height: | Size: 335 KiB |
|
After Width: | Height: | Size: 34 KiB |
|
After Width: | Height: | Size: 67 KiB |
|
After Width: | Height: | Size: 72 KiB |
|
After Width: | Height: | Size: 93 KiB |
|
After Width: | Height: | Size: 94 KiB |
|
After Width: | Height: | Size: 90 KiB |
|
After Width: | Height: | Size: 150 KiB |
|
After Width: | Height: | Size: 109 KiB |
|
After Width: | Height: | Size: 85 KiB |
|
After Width: | Height: | Size: 114 KiB |
|
After Width: | Height: | Size: 125 KiB |
|
After Width: | Height: | Size: 923 KiB |
|
After Width: | Height: | Size: 178 KiB |
|
After Width: | Height: | Size: 707 KiB |
|
After Width: | Height: | Size: 338 KiB |
|
After Width: | Height: | Size: 1.1 MiB |
|
After Width: | Height: | Size: 733 KiB |
|
After Width: | Height: | Size: 46 KiB |
|
After Width: | Height: | Size: 1.1 MiB |
|
After Width: | Height: | Size: 936 KiB |
@@ -0,0 +1,173 @@
|
||||
.. _actions:
|
||||
|
||||
Actions
|
||||
=======
|
||||
|
||||
Actions define how the policy controls the simulation. The action
|
||||
manager receives the policy's output tensor each step, splits it across
|
||||
registered action terms, and routes each slice to the appropriate
|
||||
entity's actuators. Each term maps a contiguous segment of the policy
|
||||
output to a control mode (position, velocity, effort) on a set of
|
||||
joints, tendons, or sites.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.envs.mdp.actions import JointPositionActionCfg
|
||||
|
||||
actions = {
|
||||
"joint_pos": JointPositionActionCfg(
|
||||
entity_name="robot",
|
||||
actuator_names=(".*",), # regex matching actuator names
|
||||
scale=0.5,
|
||||
use_default_offset=True, # action 0 = default pose
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
Common parameters
|
||||
-----------------
|
||||
|
||||
All action types share a base set of parameters inherited from
|
||||
``BaseActionCfg``.
|
||||
|
||||
``entity_name`` identifies the scene entity to control. ``actuator_names``
|
||||
is a tuple of regex patterns matched against actuator (or tendon/site)
|
||||
names to select the controlled targets.
|
||||
|
||||
``scale`` multiplies the raw policy output before any offset is applied.
|
||||
It accepts a scalar or a dict mapping actuator name patterns to
|
||||
per-target values. This keeps policy outputs in a normalized range while
|
||||
mapping to physically meaningful units. ``offset`` is added after
|
||||
scaling; joint action types also provide ``use_default_offset``, which
|
||||
automatically loads the entity's default joint positions or velocities
|
||||
as the offset so that a raw output of zero produces the default pose.
|
||||
|
||||
``clip`` optionally clamps the processed action (after scale and offset)
|
||||
before it reaches the actuator. It accepts a dict mapping actuator name
|
||||
patterns to ``(min, max)`` tuples, resolved the same way as ``scale``
|
||||
and ``offset``.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
JointPositionActionCfg(
|
||||
entity_name="robot",
|
||||
actuator_names=(".*",),
|
||||
scale=0.5,
|
||||
clip={".*_hip_.*": (-1.0, 1.0), ".*_knee_.*": (-0.5, 2.0)},
|
||||
)
|
||||
|
||||
Actions are written to actuator targets on every decimation substep
|
||||
(physics step), not just once per policy step. This is in contrast to
|
||||
observation delay, which operates in units of policy steps.
|
||||
|
||||
|
||||
Action types
|
||||
------------
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 28 72
|
||||
|
||||
* - Type
|
||||
- Description
|
||||
* - ``JointPositionAction``
|
||||
- Sets joint position targets. With ``use_default_offset=True``
|
||||
(the default), a policy output of zero commands the default pose.
|
||||
Encoder bias from ``dr.encoder_bias`` is subtracted automatically
|
||||
so that randomized offsets propagate correctly to the control
|
||||
command.
|
||||
* - ``RelativeJointPositionAction``
|
||||
- Sets joint position targets relative to the current joint positions.
|
||||
The target is ``current_pos + action * scale``, so a policy output of
|
||||
zero holds the robot in place regardless of its current configuration.
|
||||
* - ``JointVelocityAction``
|
||||
- Sets joint velocity targets. ``use_default_offset=True`` uses the
|
||||
default joint velocities (typically zero).
|
||||
* - ``JointEffortAction``
|
||||
- Sets joint effort (torque) targets directly. No default offset.
|
||||
* - ``TendonLengthAction``
|
||||
- Sets tendon length targets. Targets are resolved by matching
|
||||
``actuator_names`` against tendon names.
|
||||
* - ``TendonVelocityAction``
|
||||
- Sets tendon velocity targets.
|
||||
* - ``TendonEffortAction``
|
||||
- Sets tendon effort targets.
|
||||
* - ``SiteEffortAction``
|
||||
- Applies forces and torques at named sites. Useful for
|
||||
quadrotors and drones where thrust is applied at rotor sites
|
||||
rather than through joint actuators.
|
||||
|
||||
|
||||
Task-space actions
|
||||
------------------
|
||||
|
||||
``DifferentialIKAction`` converts Cartesian position and orientation
|
||||
commands into joint-space position targets via damped least-squares
|
||||
inverse kinematics. One IK step is executed per decimation substep, so
|
||||
the end-effector tracks the target continuously across substeps rather
|
||||
than only at policy frequency.
|
||||
|
||||
The action dimension is selected automatically based on configuration:
|
||||
|
||||
- ``orientation_weight == 0``: **3D** (position only)
|
||||
- ``orientation_weight > 0, use_relative_mode=True``: **6D** (delta
|
||||
position + delta axis-angle)
|
||||
- ``orientation_weight > 0, use_relative_mode=False``: **7D** (absolute
|
||||
position + quaternion)
|
||||
|
||||
All objectives (position, orientation, joint limits, posture) are
|
||||
stacked into a single DLS system. Setting a weight to zero disables
|
||||
that objective with no overhead in the solve.
|
||||
|
||||
The ``compute_dq()`` method returns joint displacements without writing
|
||||
to actuator targets, enabling multi-iteration IK in standalone scripts
|
||||
outside of RL training.
|
||||
|
||||
|
||||
Action dimensions and history
|
||||
------------------------------
|
||||
|
||||
The total action dimension presented to the policy is the sum of each
|
||||
registered term's ``action_dim``. For joint, tendon, and site actions
|
||||
this equals the number of matched targets. For ``DifferentialIKAction``
|
||||
it is 3, 6, or 7 depending on the active objectives.
|
||||
|
||||
The action manager tracks the three most recent action vectors:
|
||||
``action``, ``prev_action``, and ``prev_prev_action``. Observation terms
|
||||
such as ``last_action`` and reward terms such as ``action_rate_l2`` and
|
||||
``action_acc_l2`` read from these buffers. Action history is zeroed on
|
||||
environment reset so that episode boundaries do not leak information.
|
||||
|
||||
|
||||
Multiple action terms
|
||||
---------------------
|
||||
|
||||
An environment can register any number of terms. The action manager
|
||||
concatenates their dimensions in registration order, splits the
|
||||
policy's output tensor at the corresponding boundaries, and routes
|
||||
each slice independently.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.envs.mdp.actions import (
|
||||
JointPositionActionCfg,
|
||||
JointVelocityActionCfg,
|
||||
)
|
||||
|
||||
actions = {
|
||||
"arm_joints": JointPositionActionCfg(
|
||||
entity_name="robot",
|
||||
actuator_names=(".*_arm_.*",),
|
||||
scale=0.5,
|
||||
),
|
||||
"wheel_joints": JointVelocityActionCfg(
|
||||
entity_name="robot",
|
||||
actuator_names=(".*_wheel_.*",),
|
||||
scale=10.0,
|
||||
),
|
||||
}
|
||||
|
||||
The policy outputs a tensor whose width equals the total number of
|
||||
matched targets across all terms. Terms can also target different
|
||||
entities, for example one term for a robot and another for an object
|
||||
being manipulated.
|
||||
@@ -0,0 +1,450 @@
|
||||
.. _actuators:
|
||||
|
||||
Actuators
|
||||
=========
|
||||
|
||||
Actuators convert high-level commands (position, velocity, effort) into
|
||||
low-level efforts that drive joints. They are configured through the
|
||||
``articulation`` field of :ref:`EntityCfg <entity>`. mjlab provides
|
||||
**built-in** actuators that leverage the physics engine's implicit
|
||||
integration for best stability, and **explicit** actuators for custom
|
||||
control laws and actuator dynamics.
|
||||
|
||||
|
||||
Quick start
|
||||
-----------
|
||||
|
||||
Basic PD control with ``BuiltinPositionActuator``, the most common
|
||||
starting point.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.actuator import BuiltinPositionActuatorCfg
|
||||
from mjlab.entity import EntityCfg, EntityArticulationInfoCfg
|
||||
|
||||
robot_cfg = EntityCfg(
|
||||
spec_fn=lambda: load_robot_spec(),
|
||||
articulation=EntityArticulationInfoCfg(
|
||||
actuators=(
|
||||
BuiltinPositionActuatorCfg(
|
||||
target_names_expr=(".*_hip_.*", ".*_knee_.*"),
|
||||
stiffness=80.0,
|
||||
damping=10.0,
|
||||
effort_limit=100.0,
|
||||
),
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
Add delay fields directly on any actuator config to model communication
|
||||
latency.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.actuator import BuiltinPositionActuatorCfg
|
||||
|
||||
BuiltinPositionActuatorCfg(
|
||||
target_names_expr=(".*",),
|
||||
stiffness=80.0,
|
||||
damping=10.0,
|
||||
delay_min_lag=2, # Minimum 2 physics steps
|
||||
delay_max_lag=5, # Maximum 5 physics steps
|
||||
)
|
||||
|
||||
|
||||
Built-in vs explicit actuators
|
||||
------------------------------
|
||||
|
||||
The key design decision when configuring actuators is whether to use
|
||||
**built-in** or **explicit** types. The difference comes down to how
|
||||
MuJoCo's integrator handles velocity-dependent forces.
|
||||
|
||||
**Built-in actuators** (``BuiltinPositionActuator``,
|
||||
``BuiltinVelocityActuator``, ``BuiltinMotorActuator``,
|
||||
``BuiltinMuscleActuator``) create native MuJoCo actuator elements in the
|
||||
MjSpec. The physics engine computes the control law and integrates
|
||||
velocity-dependent damping forces implicitly. This provides the best
|
||||
numerical stability, particularly with high gains or large timesteps.
|
||||
|
||||
**Explicit actuators** (``IdealPdActuator``, ``DcMotorActuator``,
|
||||
``LearnedMlpActuator``) compute torques in user code and forward them
|
||||
through a ``<motor>`` actuator acting as a passthrough. Because the
|
||||
integrator cannot account for the velocity derivatives of these
|
||||
externally computed forces, they are less numerically robust than built-in
|
||||
types. Use explicit actuators when you need custom control laws or actuator
|
||||
dynamics that cannot be expressed with built-in types (e.g.,
|
||||
velocity-dependent torque limits, learned actuator networks).
|
||||
|
||||
The two approaches match closely in the linear, unconstrained regime at
|
||||
small timesteps. At larger timesteps or higher gains, built-in actuators
|
||||
are more forgiving.
|
||||
|
||||
**Integrator choice.** mjlab places damping inside the actuator rather than
|
||||
in joints. The ``euler`` integrator treats joint damping implicitly but
|
||||
actuator damping explicitly, limiting stability. The ``implicitfast``
|
||||
integrator treats all known velocity-dependent forces implicitly, handling
|
||||
both proportional and damping terms of the actuator without additional cost.
|
||||
|
||||
.. note::
|
||||
|
||||
mjlab defaults to ``implicitfast``, as it is MuJoCo's recommended
|
||||
integrator and provides superior stability for actuator-side damping.
|
||||
|
||||
|
||||
Actuator types
|
||||
--------------
|
||||
|
||||
All actuator configs share a few common fields inherited from
|
||||
``ActuatorCfg``:
|
||||
|
||||
- ``target_names_expr``: Tuple of regex patterns matched against joint
|
||||
names (or tendon/site names when using a different
|
||||
``transmission_type``).
|
||||
- ``armature``: Reflected rotor inertia added to the target joint.
|
||||
- ``frictionloss``: Static friction (stiction) modeled as a constraint
|
||||
on the target joint. See MuJoCo's
|
||||
`frictionloss <https://mujoco.readthedocs.io/en/stable/XMLreference.html#body-joint-frictionloss>`_.
|
||||
|
||||
Built-in actuators
|
||||
^^^^^^^^^^^^^^^^^^
|
||||
|
||||
Built-in actuators use MuJoCo's native actuator types via the MjSpec API.
|
||||
|
||||
**BuiltinPositionActuator**: Creates ``<position>`` actuators for PD
|
||||
control.
|
||||
|
||||
**BuiltinVelocityActuator**: Creates ``<velocity>`` actuators for velocity
|
||||
control.
|
||||
|
||||
**BuiltinMotorActuator**: Creates ``<motor>`` actuators for direct torque
|
||||
control.
|
||||
|
||||
**BuiltinMuscleActuator**: Creates ``<muscle>`` actuators for
|
||||
biologically-inspired muscle dynamics with force-length-velocity
|
||||
characteristics.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.actuator import BuiltinPositionActuatorCfg, BuiltinVelocityActuatorCfg
|
||||
|
||||
# Mobile manipulator: PD for arm joints, velocity control for wheels.
|
||||
actuators = (
|
||||
BuiltinPositionActuatorCfg(
|
||||
target_names_expr=(".*_shoulder_.*", ".*_elbow_.*", ".*_wrist_.*"),
|
||||
stiffness=100.0,
|
||||
damping=10.0,
|
||||
effort_limit=150.0,
|
||||
),
|
||||
BuiltinVelocityActuatorCfg(
|
||||
target_names_expr=(".*_wheel_.*",),
|
||||
damping=20.0,
|
||||
effort_limit=50.0,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
Explicit actuators
|
||||
^^^^^^^^^^^^^^^^^^
|
||||
|
||||
Explicit actuators compute efforts and forward them to an underlying
|
||||
``<motor>`` actuator acting as a passthrough. See
|
||||
`Built-in vs explicit actuators`_ above for stability implications.
|
||||
|
||||
**IdealPdActuator**: Implements an ideal PD controller. Computes torques
|
||||
as ``tau = Kp * pos_error + Kd * vel_error``.
|
||||
|
||||
**DcMotorActuator**: Extends ``IdealPdActuator`` with velocity-dependent
|
||||
torque saturation to model DC motor torque-speed curves (back-EMF
|
||||
effects). Implements a linear torque-speed curve: maximum torque at zero
|
||||
velocity, zero torque at maximum velocity.
|
||||
|
||||
**LearnedMlpActuator**: Neural network-based actuator that uses a
|
||||
trained MLP to predict torque outputs from joint state history. Useful
|
||||
when analytical models cannot capture complex actuator dynamics like
|
||||
delays, nonlinearities, and friction effects. Inherits DC motor
|
||||
velocity-based torque limits.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.actuator import IdealPdActuatorCfg, DcMotorActuatorCfg
|
||||
|
||||
# Ideal PD for hips, DC motor model with torque-speed curve for knees.
|
||||
actuators = (
|
||||
IdealPdActuatorCfg(
|
||||
target_names_expr=(".*_hip_.*",),
|
||||
stiffness=80.0,
|
||||
damping=10.0,
|
||||
effort_limit=100.0,
|
||||
),
|
||||
DcMotorActuatorCfg(
|
||||
target_names_expr=(".*_knee_.*",),
|
||||
stiffness=80.0,
|
||||
damping=10.0,
|
||||
effort_limit=25.0, # Continuous torque limit
|
||||
saturation_effort=50.0, # Peak torque at stall
|
||||
velocity_limit=30.0, # No-load speed (rad/s)
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
XML actuators
|
||||
^^^^^^^^^^^^^
|
||||
|
||||
XML actuators wrap actuators already defined in your robot's XML file. The
|
||||
config finds existing actuators by matching their ``target`` joint name
|
||||
against the ``target_names_expr`` patterns. Each joint must have exactly one
|
||||
matching actuator.
|
||||
|
||||
**XmlActuator**: Wraps any actuator already defined in the XML. The
|
||||
actuator type (position, velocity, motor, muscle) is auto detected from
|
||||
the XML element, or you can set ``command_field`` explicitly.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.actuator import XmlActuatorCfg
|
||||
|
||||
# Robot XML already has:
|
||||
# <actuator>
|
||||
# <position name="hip_joint" joint="hip_joint" kp="100"/>
|
||||
# </actuator>
|
||||
|
||||
# Wrap existing XML actuators.
|
||||
actuators = (
|
||||
XmlActuatorCfg(target_names_expr=("hip_joint",)),
|
||||
)
|
||||
|
||||
Actuator delays
|
||||
^^^^^^^^^^^^^^^
|
||||
|
||||
Any actuator config supports inline delay fields for modeling command
|
||||
latency. On a real robot, the onboard PD loop runs at KHz with direct
|
||||
encoder access, but the position target from the policy arrives late due
|
||||
to inference time and communication bus cycles. Actuator
|
||||
delay models this: the command target is delayed, but the control law
|
||||
still sees fresh joint state.
|
||||
|
||||
This is distinct from observation delay, which models sensor pipeline
|
||||
latency (stale state going into the policy). Together they cover both
|
||||
legs of the round trip: sensor to policy to motor.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.actuator import IdealPdActuatorCfg
|
||||
|
||||
# Add 2-5 step delay to position commands.
|
||||
actuators = (
|
||||
IdealPdActuatorCfg(
|
||||
target_names_expr=(".*",),
|
||||
stiffness=80.0,
|
||||
damping=10.0,
|
||||
delay_min_lag=2,
|
||||
delay_max_lag=5,
|
||||
delay_hold_prob=0.3, # 30% chance to keep current lag
|
||||
delay_update_period=10, # Resample lag every 10 steps
|
||||
),
|
||||
)
|
||||
|
||||
Each step, a lag is sampled uniformly from ``[delay_min_lag,
|
||||
delay_max_lag]``. Delays are quantized to physics timesteps. For
|
||||
example, with 500Hz physics (2ms/step), ``delay_min_lag=2`` represents
|
||||
a 4ms minimum delay.
|
||||
|
||||
|
||||
Authoring actuator configs
|
||||
--------------------------
|
||||
|
||||
Since actuator parameters are uniform within each config, use separate
|
||||
actuator configs for joints that need different parameters:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.actuator import BuiltinPositionActuatorCfg
|
||||
|
||||
# G1 humanoid with different gains per joint group.
|
||||
G1_ACTUATORS = (
|
||||
BuiltinPositionActuatorCfg(
|
||||
target_names_expr=(".*_hip_.*", "waist_yaw_joint"),
|
||||
stiffness=180.0,
|
||||
damping=18.0,
|
||||
effort_limit=88.0,
|
||||
armature=0.0015,
|
||||
),
|
||||
BuiltinPositionActuatorCfg(
|
||||
target_names_expr=("left_hip_pitch_joint", "right_hip_pitch_joint"),
|
||||
stiffness=200.0,
|
||||
damping=20.0,
|
||||
effort_limit=88.0,
|
||||
armature=0.0015,
|
||||
),
|
||||
BuiltinPositionActuatorCfg(
|
||||
target_names_expr=(".*_knee_joint",),
|
||||
stiffness=150.0,
|
||||
damping=15.0,
|
||||
effort_limit=139.0,
|
||||
armature=0.0025,
|
||||
),
|
||||
BuiltinPositionActuatorCfg(
|
||||
target_names_expr=(".*_ankle_.*",),
|
||||
stiffness=40.0,
|
||||
damping=5.0,
|
||||
effort_limit=25.0,
|
||||
armature=0.0008,
|
||||
),
|
||||
)
|
||||
|
||||
This design choice reflects a deliberate simplification in mjlab: each
|
||||
``ActuatorCfg`` represents a single actuator type (e.g., a specific
|
||||
motor/gearbox model) applied uniformly across all joints it drives.
|
||||
Hardware parameters such as ``armature`` (reflected rotor inertia) and
|
||||
``gear`` describe properties of the actuator hardware, even though they
|
||||
are implemented in MuJoCo as joint or actuator fields. In other frameworks
|
||||
(like Isaac Lab), these fields may accept ``float | dict[str, float]`` to
|
||||
support per-joint variation. mjlab instead encourages one config per
|
||||
actuator type or per joint group, keeping the hardware model physically
|
||||
consistent and explicit. The main trade-off is verbosity in special cases,
|
||||
such as parallel linkages, where per-joint overrides could have been
|
||||
convenient, but the benefit is clearer semantics and simpler maintenance.
|
||||
|
||||
See :ref:`actions` for how action terms route policy outputs to actuators
|
||||
(including DifferentialIK for task-space control), and
|
||||
:ref:`domain_randomization` for randomizing gains and effort limits.
|
||||
|
||||
|
||||
Computing hardware parameters
|
||||
------------------------------
|
||||
|
||||
This section is relevant when configuring actuators from real motor
|
||||
datasheets. If you are using manually tuned gains, you can skip ahead.
|
||||
|
||||
mjlab provides utilities in ``mjlab.utils.actuator`` to compute actuator
|
||||
parameters from physical motor specifications. This is particularly
|
||||
useful for computing reflected inertia (``armature``) and deriving
|
||||
appropriate control gains from hardware datasheets.
|
||||
|
||||
**Example: Unitree G1 motor configuration**
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from math import pi
|
||||
|
||||
from mjlab.utils.actuator import (
|
||||
reflected_inertia_from_two_stage_planetary,
|
||||
ElectricActuator
|
||||
)
|
||||
|
||||
# Motor specs from manufacturer datasheet.
|
||||
ROTOR_INERTIAS_7520_14 = (
|
||||
0.489e-4, # Motor rotor inertia (kg*m**2)
|
||||
0.098e-4, # Planet carrier inertia
|
||||
0.533e-4, # Output stage inertia
|
||||
)
|
||||
GEARS_7520_14 = (
|
||||
1, # First stage (motor to planet)
|
||||
4.5, # Second stage (planet to carrier)
|
||||
1 + (48/22), # Third stage (carrier to output)
|
||||
)
|
||||
|
||||
# Compute reflected inertia at joint output.
|
||||
# J_reflected = J_motor*(N1*N2)**2 + J_carrier*N2**2 + J_output.
|
||||
ARMATURE_7520_14 = reflected_inertia_from_two_stage_planetary(
|
||||
ROTOR_INERTIAS_7520_14, GEARS_7520_14
|
||||
)
|
||||
|
||||
# Create motor spec container.
|
||||
ACTUATOR_7520_14 = ElectricActuator(
|
||||
reflected_inertia=ARMATURE_7520_14,
|
||||
velocity_limit=32.0, # rad/s at joint
|
||||
effort_limit=88.0, # N*m continuous torque
|
||||
)
|
||||
|
||||
# Derive PD gains from natural frequency and damping ratio.
|
||||
NATURAL_FREQ = 10 * 2*pi # 10 Hz bandwidth.
|
||||
DAMPING_RATIO = 2.0 # Overdamped, see note below.
|
||||
STIFFNESS = ARMATURE_7520_14 * NATURAL_FREQ**2
|
||||
DAMPING = 2 * DAMPING_RATIO * ARMATURE_7520_14 * NATURAL_FREQ
|
||||
|
||||
# Use in actuator config.
|
||||
from mjlab.actuator import BuiltinPositionActuatorCfg
|
||||
|
||||
actuator = BuiltinPositionActuatorCfg(
|
||||
target_names_expr=(".*_hip_pitch_joint",),
|
||||
stiffness=STIFFNESS,
|
||||
damping=DAMPING,
|
||||
effort_limit=ACTUATOR_7520_14.effort_limit,
|
||||
armature=ACTUATOR_7520_14.reflected_inertia,
|
||||
)
|
||||
|
||||
.. note::
|
||||
|
||||
The example uses ``DAMPING_RATIO = 2.0``
|
||||
(overdamped) rather than the critically damped value of 1.0. This is
|
||||
because the reflected inertia calculation only accounts for the motor's
|
||||
rotor inertia, not the apparent inertia of the links being moved. In
|
||||
practice, the total effective inertia at the joint is higher than just
|
||||
the reflected motor inertia, so using an overdamped ratio provides
|
||||
better stability margins when the true system inertia is
|
||||
underestimated.
|
||||
|
||||
**Parallel linkage approximation:**
|
||||
|
||||
For joints driven by parallel linkages (like the G1's ankles with dual
|
||||
motors), the effective armature in the nominal configuration can be
|
||||
approximated as the sum of the individual motor armatures:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# Two 5020 motors driving ankle through parallel linkage.
|
||||
G1_ACTUATOR_ANKLE = BuiltinPositionActuatorCfg(
|
||||
target_names_expr=(".*_ankle_pitch_joint", ".*_ankle_roll_joint"),
|
||||
stiffness=STIFFNESS_5020 * 2,
|
||||
damping=DAMPING_5020 * 2,
|
||||
effort_limit=ACTUATOR_5020.effort_limit * 2,
|
||||
armature=ACTUATOR_5020.reflected_inertia * 2,
|
||||
)
|
||||
|
||||
|
||||
Extending: custom actuators
|
||||
----------------------------
|
||||
|
||||
All actuators implement a unified ``compute()`` interface that receives an
|
||||
``ActuatorCmd`` (containing position, velocity, and effort targets) and
|
||||
returns control signals for the low-level MuJoCo actuators driving each
|
||||
joint.
|
||||
|
||||
**Core interface:**
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
def compute(self, cmd: ActuatorCmd) -> torch.Tensor:
|
||||
"""Convert high-level commands to control signals.
|
||||
|
||||
Args:
|
||||
cmd: Command containing position_target, velocity_target,
|
||||
effort_target (each is a [num_envs, num_targets] tensor
|
||||
or None)
|
||||
|
||||
Returns:
|
||||
Control signals for this actuator
|
||||
([num_envs, num_targets] tensor)
|
||||
"""
|
||||
|
||||
**Lifecycle hooks:**
|
||||
|
||||
- ``edit_spec``: Modify MjSpec before compilation (add actuators, set
|
||||
gains)
|
||||
- ``initialize``: Post-compilation setup (resolve indices, allocate
|
||||
buffers)
|
||||
- ``reset``: Per-environment reset logic
|
||||
- ``update``: Pre-step updates
|
||||
- ``compute``: Convert commands to control signals
|
||||
|
||||
**Properties:**
|
||||
|
||||
- ``target_ids``: Tensor of local target indices controlled by this
|
||||
actuator
|
||||
- ``target_names``: List of target names controlled by this actuator
|
||||
- ``ctrl_ids``: Tensor of global control input indices for this actuator
|
||||
|
||||
``IdealPdActuator`` is the recommended base class for custom explicit
|
||||
actuators. ``DcMotorActuator`` and ``LearnedMlpActuator`` are both
|
||||
built on top of it and serve as examples of the extension pattern.
|
||||
@@ -0,0 +1,141 @@
|
||||
mjlab.actuator
|
||||
==============
|
||||
|
||||
.. automodule:: mjlab.actuator
|
||||
|
||||
.. rubric:: Classes
|
||||
|
||||
.. hlist::
|
||||
:columns: 3
|
||||
|
||||
- :class:`Actuator`
|
||||
- :class:`ActuatorCfg`
|
||||
- :class:`ActuatorCmd`
|
||||
- :class:`BuiltinActuatorGroup`
|
||||
- :class:`BuiltinMotorActuator`
|
||||
- :class:`BuiltinMotorActuatorCfg`
|
||||
- :class:`BuiltinPositionActuator`
|
||||
- :class:`BuiltinPositionActuatorCfg`
|
||||
- :class:`BuiltinVelocityActuator`
|
||||
- :class:`BuiltinVelocityActuatorCfg`
|
||||
- :class:`BuiltinMuscleActuator`
|
||||
- :class:`BuiltinMuscleActuatorCfg`
|
||||
- :class:`XmlActuator`
|
||||
- :class:`XmlActuatorCfg`
|
||||
- :class:`IdealPdActuator`
|
||||
- :class:`IdealPdActuatorCfg`
|
||||
- :class:`DcMotorActuator`
|
||||
- :class:`DcMotorActuatorCfg`
|
||||
- :class:`LearnedMlpActuator`
|
||||
- :class:`LearnedMlpActuatorCfg`
|
||||
|
||||
Base
|
||||
----
|
||||
|
||||
.. autoclass:: Actuator
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: ActuatorCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
|
||||
.. autoclass:: ActuatorCmd
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
Builtin Actuators
|
||||
-----------------
|
||||
|
||||
.. autoclass:: BuiltinActuatorGroup
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: BuiltinMotorActuator
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: BuiltinMotorActuatorCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
|
||||
.. autoclass:: BuiltinPositionActuator
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: BuiltinPositionActuatorCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
|
||||
.. autoclass:: BuiltinVelocityActuator
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: BuiltinVelocityActuatorCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
|
||||
.. autoclass:: BuiltinMuscleActuator
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: BuiltinMuscleActuatorCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
XML Actuators
|
||||
-------------
|
||||
|
||||
.. autoclass:: XmlActuator
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: XmlActuatorCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
Ideal PD Actuator
|
||||
-----------------
|
||||
|
||||
.. autoclass:: IdealPdActuator
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: IdealPdActuatorCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
DC Motor Actuator
|
||||
-----------------
|
||||
|
||||
.. autoclass:: DcMotorActuator
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: DcMotorActuatorCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
Learned MLP Actuator
|
||||
--------------------
|
||||
|
||||
.. autoclass:: LearnedMlpActuator
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: LearnedMlpActuatorCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
@@ -0,0 +1,45 @@
|
||||
mjlab.entity
|
||||
============
|
||||
|
||||
.. automodule:: mjlab.entity
|
||||
|
||||
.. rubric:: Classes
|
||||
|
||||
.. hlist::
|
||||
:columns: 3
|
||||
|
||||
- :class:`Entity`
|
||||
- :class:`EntityCfg`
|
||||
- :class:`EntityArticulationInfoCfg`
|
||||
- :class:`EntityIndexing`
|
||||
- :class:`EntityData`
|
||||
|
||||
Entity
|
||||
------
|
||||
|
||||
.. autoclass:: Entity
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: EntityCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
|
||||
.. autoclass:: EntityArticulationInfoCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
EntityIndexing
|
||||
--------------
|
||||
|
||||
.. autoclass:: EntityIndexing
|
||||
:members:
|
||||
|
||||
EntityData
|
||||
----------
|
||||
|
||||
.. autoclass:: EntityData
|
||||
:members:
|
||||
@@ -0,0 +1,36 @@
|
||||
mjlab.envs
|
||||
==========
|
||||
|
||||
.. automodule:: mjlab.envs
|
||||
|
||||
.. rubric:: Classes
|
||||
|
||||
.. hlist::
|
||||
:columns: 3
|
||||
|
||||
- :class:`ManagerBasedRlEnv`
|
||||
- :class:`ManagerBasedRlEnvCfg`
|
||||
- :data:`VecEnvObs`
|
||||
- :data:`VecEnvStepReturn`
|
||||
|
||||
ManagerBasedRlEnv
|
||||
-----------------
|
||||
|
||||
.. autoclass:: ManagerBasedRlEnv
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: ManagerBasedRlEnvCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
VecEnvObs
|
||||
---------
|
||||
|
||||
.. autodata:: VecEnvObs
|
||||
|
||||
VecEnvStepReturn
|
||||
----------------
|
||||
|
||||
.. autodata:: VecEnvStepReturn
|
||||
@@ -0,0 +1,19 @@
|
||||
API Reference
|
||||
=============
|
||||
|
||||
This section provides detailed API documentation for all public modules in mjlab.
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
|
||||
envs
|
||||
scene
|
||||
sim
|
||||
entity
|
||||
actuator
|
||||
sensor
|
||||
managers
|
||||
terrains
|
||||
rl
|
||||
viewer
|
||||
tasks
|
||||
@@ -0,0 +1,208 @@
|
||||
mjlab.managers
|
||||
==============
|
||||
|
||||
.. automodule:: mjlab.managers
|
||||
|
||||
.. rubric:: Classes
|
||||
|
||||
.. hlist::
|
||||
:columns: 3
|
||||
|
||||
- :class:`ManagerBase`
|
||||
- :class:`ManagerTermBase`
|
||||
- :class:`ManagerTermBaseCfg`
|
||||
- :class:`SceneEntityCfg`
|
||||
- :class:`ActionManager`
|
||||
- :class:`ActionTerm`
|
||||
- :class:`ActionTermCfg`
|
||||
- :class:`ObservationManager`
|
||||
- :class:`ObservationGroupCfg`
|
||||
- :class:`ObservationTermCfg`
|
||||
- :class:`RewardManager`
|
||||
- :class:`RewardTermCfg`
|
||||
- :class:`TerminationManager`
|
||||
- :class:`TerminationTermCfg`
|
||||
- :class:`CommandManager`
|
||||
- :class:`NullCommandManager`
|
||||
- :class:`CommandTerm`
|
||||
- :class:`CommandTermCfg`
|
||||
- :class:`CurriculumManager`
|
||||
- :class:`NullCurriculumManager`
|
||||
- :class:`CurriculumTermCfg`
|
||||
- :class:`EventManager`
|
||||
- :class:`EventMode`
|
||||
- :class:`EventTermCfg`
|
||||
- :class:`MetricsManager`
|
||||
- :class:`NullMetricsManager`
|
||||
- :class:`MetricsTermCfg`
|
||||
- :class:`RecorderManager`
|
||||
- :class:`NullRecorderManager`
|
||||
- :class:`RecorderTerm`
|
||||
- :class:`RecorderTermCfg`
|
||||
|
||||
Base
|
||||
----
|
||||
|
||||
.. autoclass:: ManagerBase
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: ManagerTermBase
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: ManagerTermBaseCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
|
||||
.. autoclass:: SceneEntityCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
Action Manager
|
||||
--------------
|
||||
|
||||
.. autoclass:: ActionManager
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: ActionTerm
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: ActionTermCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
Observation Manager
|
||||
-------------------
|
||||
|
||||
.. autoclass:: ObservationManager
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: ObservationGroupCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
|
||||
.. autoclass:: ObservationTermCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
Reward Manager
|
||||
--------------
|
||||
|
||||
.. autoclass:: RewardManager
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: RewardTermCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
Termination Manager
|
||||
-------------------
|
||||
|
||||
.. autoclass:: TerminationManager
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: TerminationTermCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
Command Manager
|
||||
---------------
|
||||
|
||||
.. autoclass:: CommandManager
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: NullCommandManager
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: CommandTerm
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: CommandTermCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
Curriculum Manager
|
||||
------------------
|
||||
|
||||
.. autoclass:: CurriculumManager
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: NullCurriculumManager
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: CurriculumTermCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
Event Manager
|
||||
-------------
|
||||
|
||||
.. autoclass:: EventManager
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: EventMode
|
||||
:members:
|
||||
:undoc-members:
|
||||
|
||||
.. autoclass:: EventTermCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
Metrics Manager
|
||||
---------------
|
||||
|
||||
.. autoclass:: MetricsManager
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: NullMetricsManager
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: MetricsTermCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
|
||||
Recorder Manager
|
||||
----------------
|
||||
|
||||
.. autoclass:: RecorderManager
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: NullRecorderManager
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: RecorderTerm
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: RecorderTermCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
@@ -0,0 +1,52 @@
|
||||
mjlab.rl
|
||||
========
|
||||
|
||||
.. automodule:: mjlab.rl
|
||||
|
||||
.. rubric:: Classes
|
||||
|
||||
.. hlist::
|
||||
:columns: 3
|
||||
|
||||
- :class:`MjlabOnPolicyRunner`
|
||||
- :class:`RslRlVecEnvWrapper`
|
||||
- :class:`RslRlOnPolicyRunnerCfg`
|
||||
- :class:`RslRlPpoAlgorithmCfg`
|
||||
- :class:`RslRlModelCfg`
|
||||
- :class:`RslRlBaseRunnerCfg`
|
||||
|
||||
Runner
|
||||
------
|
||||
|
||||
.. autoclass:: MjlabOnPolicyRunner
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: RslRlVecEnvWrapper
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
Configuration
|
||||
-------------
|
||||
|
||||
.. autoclass:: RslRlOnPolicyRunnerCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
|
||||
.. autoclass:: RslRlPpoAlgorithmCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
|
||||
.. autoclass:: RslRlModelCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
|
||||
.. autoclass:: RslRlBaseRunnerCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
@@ -0,0 +1,23 @@
|
||||
mjlab.scene
|
||||
===========
|
||||
|
||||
.. automodule:: mjlab.scene
|
||||
|
||||
.. rubric:: Classes
|
||||
|
||||
.. hlist::
|
||||
:columns: 3
|
||||
|
||||
- :class:`Scene`
|
||||
- :class:`SceneCfg`
|
||||
|
||||
Scene
|
||||
-----
|
||||
|
||||
.. autoclass:: Scene
|
||||
:members:
|
||||
|
||||
.. autoclass:: SceneCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
@@ -0,0 +1,126 @@
|
||||
mjlab.sensor
|
||||
============
|
||||
|
||||
.. automodule:: mjlab.sensor
|
||||
|
||||
.. rubric:: Classes
|
||||
|
||||
.. hlist::
|
||||
:columns: 3
|
||||
|
||||
- :class:`Sensor`
|
||||
- :class:`SensorCfg`
|
||||
- :class:`SensorContext`
|
||||
- :class:`BuiltinSensor`
|
||||
- :class:`BuiltinSensorCfg`
|
||||
- :class:`ObjRef`
|
||||
- :class:`ContactSensor`
|
||||
- :class:`ContactSensorCfg`
|
||||
- :class:`ContactData`
|
||||
- :class:`ContactMatch`
|
||||
- :class:`RayCastSensor`
|
||||
- :class:`RayCastSensorCfg`
|
||||
- :class:`RayCastData`
|
||||
- :class:`GridPatternCfg`
|
||||
- :class:`PinholeCameraPatternCfg`
|
||||
- :class:`CameraSensor`
|
||||
- :class:`CameraSensorCfg`
|
||||
- :class:`CameraSensorData`
|
||||
|
||||
Base
|
||||
----
|
||||
|
||||
.. autoclass:: Sensor
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: SensorCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
|
||||
.. autoclass:: SensorContext
|
||||
:members:
|
||||
|
||||
Builtin Sensor
|
||||
--------------
|
||||
|
||||
.. autoclass:: BuiltinSensor
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: BuiltinSensorCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
|
||||
.. autoclass:: ObjRef
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
Contact Sensor
|
||||
--------------
|
||||
|
||||
.. autoclass:: ContactSensor
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: ContactSensorCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
|
||||
.. autoclass:: ContactData
|
||||
:members:
|
||||
|
||||
.. autoclass:: ContactMatch
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
Ray Cast Sensor
|
||||
---------------
|
||||
|
||||
.. autoclass:: RayCastSensor
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: RayCastSensorCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
|
||||
.. autoclass:: RayCastData
|
||||
:members:
|
||||
|
||||
.. autoclass:: GridPatternCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
|
||||
.. autoclass:: PinholeCameraPatternCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
Camera Sensor
|
||||
-------------
|
||||
|
||||
.. autoclass:: CameraSensor
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: CameraSensorCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
|
||||
.. autoclass:: CameraSensorData
|
||||
:members:
|
||||
@@ -0,0 +1,44 @@
|
||||
mjlab.sim
|
||||
=========
|
||||
|
||||
.. automodule:: mjlab.sim
|
||||
|
||||
.. rubric:: Classes
|
||||
|
||||
.. hlist::
|
||||
:columns: 3
|
||||
|
||||
- :class:`Simulation`
|
||||
- :class:`SimulationCfg`
|
||||
- :class:`MujocoCfg`
|
||||
- :class:`TorchArray`
|
||||
- :class:`WarpBridge`
|
||||
|
||||
Simulation
|
||||
----------
|
||||
|
||||
.. autoclass:: Simulation
|
||||
:members:
|
||||
|
||||
.. autoclass:: SimulationCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
|
||||
.. autoclass:: MujocoCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
TorchArray
|
||||
----------
|
||||
|
||||
.. autoclass:: TorchArray
|
||||
:members:
|
||||
|
||||
WarpBridge
|
||||
----------
|
||||
|
||||
.. autoclass:: WarpBridge
|
||||
:members:
|
||||
@@ -0,0 +1,25 @@
|
||||
mjlab.tasks
|
||||
===========
|
||||
|
||||
.. automodule:: mjlab.tasks.registry
|
||||
|
||||
.. rubric:: Functions
|
||||
|
||||
.. hlist::
|
||||
:columns: 3
|
||||
|
||||
- :func:`register_mjlab_task`
|
||||
- :func:`list_tasks`
|
||||
- :func:`load_env_cfg`
|
||||
- :func:`load_rl_cfg`
|
||||
- :func:`load_runner_cls`
|
||||
|
||||
.. autofunction:: register_mjlab_task
|
||||
|
||||
.. autofunction:: list_tasks
|
||||
|
||||
.. autofunction:: load_env_cfg
|
||||
|
||||
.. autofunction:: load_rl_cfg
|
||||
|
||||
.. autofunction:: load_runner_cls
|
||||
@@ -0,0 +1,167 @@
|
||||
mjlab.terrains
|
||||
==============
|
||||
|
||||
.. automodule:: mjlab.terrains
|
||||
|
||||
.. rubric:: Classes
|
||||
|
||||
.. hlist::
|
||||
:columns: 3
|
||||
|
||||
- :class:`TerrainEntity`
|
||||
- :class:`TerrainEntityCfg`
|
||||
- :class:`TerrainGenerator`
|
||||
- :class:`TerrainGeneratorCfg`
|
||||
- :class:`SubTerrainCfg`
|
||||
- :class:`FlatPatchSamplingCfg`
|
||||
- :class:`HfDiscreteObstaclesTerrainCfg`
|
||||
- :class:`HfPerlinNoiseTerrainCfg`
|
||||
- :class:`HfPyramidSlopedTerrainCfg`
|
||||
- :class:`HfRandomUniformTerrainCfg`
|
||||
- :class:`HfWaveTerrainCfg`
|
||||
- :class:`BoxFlatTerrainCfg`
|
||||
- :class:`BoxInvertedPyramidStairsTerrainCfg`
|
||||
- :class:`BoxNarrowBeamsTerrainCfg`
|
||||
- :class:`BoxNestedRingsTerrainCfg`
|
||||
- :class:`BoxOpenStairsTerrainCfg`
|
||||
- :class:`BoxPyramidStairsTerrainCfg`
|
||||
- :class:`BoxRandomGridTerrainCfg`
|
||||
- :class:`BoxRandomSpreadTerrainCfg`
|
||||
- :class:`BoxRandomStairsTerrainCfg`
|
||||
- :class:`BoxSteppingStonesTerrainCfg`
|
||||
- :class:`BoxTiltedGridTerrainCfg`
|
||||
|
||||
Core
|
||||
----
|
||||
|
||||
.. autoclass:: TerrainEntity
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: TerrainEntityCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
|
||||
.. autoclass:: TerrainGenerator
|
||||
:members:
|
||||
|
||||
.. autoclass:: TerrainGeneratorCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
|
||||
.. autoclass:: SubTerrainCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
|
||||
.. autoclass:: FlatPatchSamplingCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
Heightfield Terrains
|
||||
--------------------
|
||||
|
||||
.. autoclass:: HfDiscreteObstaclesTerrainCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: HfPerlinNoiseTerrainCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: HfRandomUniformTerrainCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: HfPyramidSlopedTerrainCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: HfWaveTerrainCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
:show-inheritance:
|
||||
|
||||
Primitive (Box) Terrains
|
||||
------------------------
|
||||
|
||||
.. autoclass:: BoxFlatTerrainCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: BoxInvertedPyramidStairsTerrainCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: BoxNarrowBeamsTerrainCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: BoxNestedRingsTerrainCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: BoxOpenStairsTerrainCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: BoxPyramidStairsTerrainCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: BoxRandomGridTerrainCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: BoxRandomSpreadTerrainCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: BoxRandomStairsTerrainCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: BoxSteppingStonesTerrainCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: BoxTiltedGridTerrainCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
:show-inheritance:
|
||||
@@ -0,0 +1,72 @@
|
||||
mjlab.viewer
|
||||
============
|
||||
|
||||
.. automodule:: mjlab.viewer
|
||||
|
||||
.. rubric:: Classes
|
||||
|
||||
.. hlist::
|
||||
:columns: 3
|
||||
|
||||
- :class:`ViewerConfig`
|
||||
- :class:`BaseViewer`
|
||||
- :class:`NativeMujocoViewer`
|
||||
- :class:`ViserPlayViewer`
|
||||
- :class:`OffscreenRenderer`
|
||||
|
||||
.. rubric:: Protocols
|
||||
|
||||
.. hlist::
|
||||
:columns: 3
|
||||
|
||||
- :class:`EnvProtocol`
|
||||
- :class:`PolicyProtocol`
|
||||
- :class:`VerbosityLevel`
|
||||
|
||||
ViewerConfig
|
||||
------------
|
||||
|
||||
.. autoclass:: ViewerConfig
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
BaseViewer
|
||||
----------
|
||||
|
||||
.. autoclass:: BaseViewer
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
NativeMujocoViewer
|
||||
------------------
|
||||
|
||||
.. autoclass:: NativeMujocoViewer
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
ViserPlayViewer
|
||||
---------------
|
||||
|
||||
.. autoclass:: ViserPlayViewer
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
OffscreenRenderer
|
||||
-----------------
|
||||
|
||||
.. autoclass:: OffscreenRenderer
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
Protocols
|
||||
---------
|
||||
|
||||
.. autoclass:: EnvProtocol
|
||||
:members:
|
||||
|
||||
.. autoclass:: PolicyProtocol
|
||||
:members:
|
||||
|
||||
.. autoclass:: VerbosityLevel
|
||||
:members:
|
||||
@@ -0,0 +1,191 @@
|
||||
.. _architecture_overview:
|
||||
|
||||
Architecture Overview
|
||||
=====================
|
||||
|
||||
mjlab is organized into two layers: a **simulation layer** that models
|
||||
the robot and world, and a **manager layer** that defines the
|
||||
reinforcement learning problem on top of it. Understanding this separation
|
||||
is the fastest way to build a mental map of the system.
|
||||
|
||||
.. figure:: _static/architecture_diagram.png
|
||||
:width: 60%
|
||||
:align: center
|
||||
:alt: mjlab architecture diagram
|
||||
|
||||
Entities are composed into an MjSpec, compiled, and transferred to
|
||||
MuJoCo Warp for GPU simulation. The ManagerBasedRlEnv orchestrates the
|
||||
MDP; RSL-RL handles training.
|
||||
|
||||
|
||||
The simulation layer
|
||||
--------------------
|
||||
|
||||
**Scene pipeline.**
|
||||
mjlab constructs scenes by composing entity descriptions into a single
|
||||
`MjSpec <https://mujoco.readthedocs.io/en/stable/programming/modeledit.html>`_.
|
||||
Each entity starts from an
|
||||
`MJCF <https://mujoco.readthedocs.io/en/latest/modeling.html>`_ file
|
||||
loaded via ``MjSpec.from_file()``. Users who define everything in XML can
|
||||
use this directly. For more control, Python dataclasses can extend or
|
||||
override properties on the loaded spec: actuators, collision rules,
|
||||
materials, sensors, and initial state. This hybrid approach lets users
|
||||
start from existing MuJoCo models and layer on task-specific configuration
|
||||
without modifying the original XML. The composed specification is compiled
|
||||
into an ``MjModel`` on the CPU, then transferred to the GPU via
|
||||
`MuJoCo Warp <https://mujoco.readthedocs.io/en/stable/mjwarp/index.html>`_,
|
||||
which is built on `NVIDIA Warp <https://nvidia.github.io/warp/>`_.
|
||||
|
||||
**MuJoCo Warp.**
|
||||
MuJoCo Warp is a GPU-accelerated backend for MuJoCo. It preserves
|
||||
MuJoCo's ``MjModel``/``MjData`` paradigm but adds a leading *world*
|
||||
dimension: a single ``MjData`` object holds the state of N independent
|
||||
simulation instances in parallel, enabling thousands of environments to
|
||||
be stepped simultaneously. Model parameters are shared across all worlds
|
||||
by default, and individual fields can be expanded to vary per-world when
|
||||
domain randomization requires it. mjlab captures the simulation step as a
|
||||
`CUDA graph <https://developer.nvidia.com/blog/cuda-graphs>`_: the kernel
|
||||
execution sequence is recorded once and replayed on subsequent calls,
|
||||
eliminating CPU-side dispatch overhead.
|
||||
|
||||
.. note::
|
||||
|
||||
CUDA graph capture is a one-time cost at environment startup. Per-episode
|
||||
resets and domain randomization events run as regular Python between graph
|
||||
replays and do not break the capture.
|
||||
|
||||
**Components.**
|
||||
The simulation layer provides four core components, each with its own
|
||||
documentation page:
|
||||
|
||||
- :ref:`entity`: a robot, a manipulated object, or a static object such
|
||||
as :ref:`terrain <terrain>`, defined by an MJCF description plus
|
||||
optional Python configuration for actuators, collision rules, and
|
||||
initial state.
|
||||
- :ref:`actuators`: how entities are controlled. Users can wrap actuators
|
||||
already defined in MJCF or create new ones from Python configuration.
|
||||
- :ref:`sensors`: how the world is observed. Includes MuJoCo-native
|
||||
sensors as well as custom sensors like RGB-D cameras and raycasters.
|
||||
- :ref:`scene`: scene composition and environment placement.
|
||||
|
||||
|
||||
The manager layer
|
||||
-----------------
|
||||
|
||||
On top of the simulation layer, mjlab adopts the manager-based environment
|
||||
design introduced by Isaac Lab. Users define their environment by composing
|
||||
small, self-contained *terms* (reward functions, observation computations,
|
||||
domain randomization events) and register them with the appropriate manager.
|
||||
Each manager handles the lifecycle of its terms: calling them at the right
|
||||
point in the simulation loop, aggregating their outputs, and exposing
|
||||
diagnostics.
|
||||
|
||||
Terms can be plain functions for stateless computations, or classes that
|
||||
inherit from ``ManagerTermBase`` when they need to cache expensive setup
|
||||
(such as resolving regex patterns to joint indices at initialization) or
|
||||
maintain per-episode state through a ``reset()`` hook.
|
||||
|
||||
Environments are configured through ``ManagerBasedRlEnvCfg``, a plain
|
||||
dataclass that holds term configuration dictionaries for each manager.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.envs import ManagerBasedRlEnvCfg
|
||||
|
||||
cfg = ManagerBasedRlEnvCfg(
|
||||
decimation=4, # 4 physics steps per policy step
|
||||
episode_length_s=20.0,
|
||||
scene=..., # SceneCfg: terrain, entities, sensors
|
||||
sim=..., # SimulationCfg: timestep, solver, integrator
|
||||
observations={...}, # ObservationManager terms
|
||||
actions={...}, # ActionManager terms
|
||||
rewards={...}, # RewardManager terms
|
||||
terminations={...}, # TerminationManager terms
|
||||
events={...}, # EventManager terms (resets, DR)
|
||||
commands={...}, # CommandManager terms (velocity targets, etc.)
|
||||
curriculum={...}, # CurriculumManager terms
|
||||
metrics={...}, # MetricsManager terms
|
||||
)
|
||||
|
||||
.. rubric:: The eight managers
|
||||
|
||||
- **ObservationManager**: assembles observation groups with configurable
|
||||
processing (clipping, noise, delay, history). Supports asymmetric
|
||||
actor-critic. See :ref:`observations`.
|
||||
- **ActionManager**: routes the policy's output tensor to entity actuators,
|
||||
handling scaling and offset. See :ref:`actions`.
|
||||
- **RewardManager**: computes a weighted sum of reward terms, scaled by step
|
||||
duration for frequency invariance. See :ref:`rewards`.
|
||||
- **TerminationManager**: evaluates stop conditions, distinguishing terminal
|
||||
resets from timeouts. See :ref:`terminations`.
|
||||
- **EventManager**: fires terms at lifecycle points (startup, reset,
|
||||
interval). Domain randomization is implemented through event terms.
|
||||
See :ref:`events` and :ref:`domain_randomization`.
|
||||
- **CommandManager**: generates and resamples goal signals (velocity
|
||||
targets, pose targets). See :ref:`commands`.
|
||||
- **CurriculumManager**: adjusts training conditions based on policy
|
||||
performance. See :ref:`curriculum`.
|
||||
- **MetricsManager**: logs custom per-step values as episode averages.
|
||||
See :ref:`metrics`.
|
||||
|
||||
For the full configuration reference covering all managers, see
|
||||
:ref:`environment_config`.
|
||||
|
||||
|
||||
The environment lifecycle
|
||||
-------------------------
|
||||
|
||||
Each environment instance passes through four phases.
|
||||
|
||||
1. **Build.** ``Scene`` composes entity MJCF files via ``MjSpec`` and
|
||||
compiles ``MjModel`` on the CPU. ``Simulation`` uploads the model to the
|
||||
GPU via MuJoCo Warp, allocating a single ``MjData`` with N parallel
|
||||
worlds. CUDA graphs for ``step``, ``forward``, ``reset``, and ``sense``
|
||||
are captured.
|
||||
|
||||
2. **Initialize.** Managers are constructed from the term configuration
|
||||
dictionaries. Regex patterns are matched to joint, body, and geom
|
||||
indices. Observation history and delay buffers are allocated. Model
|
||||
fields required by domain randomization terms are expanded from shared
|
||||
to per-world storage, and CUDA graphs are rebuilt to reflect the new
|
||||
layout. Startup events are fired once.
|
||||
|
||||
3. **Reset.** Called at the start of training and whenever an environment
|
||||
terminates or times out. The ``EventManager`` fires ``reset`` terms,
|
||||
which return the scene to an initial state with optional randomization.
|
||||
Command targets are resampled. Observation history buffers are cleared.
|
||||
|
||||
4. **Step.** The policy action is processed by the ``ActionManager``. The
|
||||
physics simulation advances ``decimation`` times, with actuator commands
|
||||
applied and entity state updated each sub-step. After the decimation
|
||||
loop, the ``TerminationManager`` checks stop conditions, the
|
||||
``RewardManager`` computes the reward signal, and any terminated
|
||||
environments are reset. A single ``forward()`` call refreshes derived
|
||||
quantities for all environments. The ``CommandManager`` advances or
|
||||
resamples goals. Interval events fire if scheduled. Sensors update. The
|
||||
``ObservationManager`` assembles the observation for the next policy
|
||||
query.
|
||||
|
||||
The step sequence in order:
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
action_manager.process_action(action)
|
||||
for _ in range(decimation):
|
||||
action_manager.apply_action()
|
||||
sim.step()
|
||||
scene.update()
|
||||
termination_manager.compute()
|
||||
reward_manager.compute()
|
||||
metrics_manager.compute()
|
||||
[reset terminated envs]
|
||||
sim.forward()
|
||||
command_manager.compute()
|
||||
event_manager.apply(mode="interval")
|
||||
sim.sense()
|
||||
observation_manager.compute()
|
||||
|
||||
With this mental model in place, the Concepts pages cover each simulation
|
||||
layer component in detail, and The Manager Layer pages walk through each
|
||||
manager's configuration and built-in terms. If you are coming from Isaac
|
||||
Lab, :ref:`migration_isaac_lab` describes the key API differences.
|
||||
@@ -0,0 +1,603 @@
|
||||
=========
|
||||
Changelog
|
||||
=========
|
||||
|
||||
Upcoming version (not yet released)
|
||||
-----------------------------------
|
||||
|
||||
Added
|
||||
^^^^^
|
||||
|
||||
- Added ``--log-root`` CLI option to ``train``, ``play``, and ``evaluate``
|
||||
scripts for choosing where training logs are stored. Defaults to
|
||||
``logs/rsl_rl`` (unchanged behavior). Useful for directing outputs to a
|
||||
scratch disk or shared mount.
|
||||
- ``RewardManager``, ``TerminationManager``, and ``MetricsManager`` now
|
||||
validate that every term function returns a tensor of shape
|
||||
``(num_envs,)`` when evaluated, raising a clear ``ValueError``
|
||||
naming the offending term instead of silently broadcasting or crashing
|
||||
with an opaque error later during training.
|
||||
- Added ``ContactSensor.primary_names`` property to expose the resolved
|
||||
primary names in the order they appear along the per-contact axis of the
|
||||
output tensors. This makes it possible to map a contact-data column back
|
||||
to the primary it belongs to (:issue:`914`).
|
||||
- Added per-world mesh variant support via ``VariantEntityCfg`` and
|
||||
``VariantCfg``. Each world in a batched simulation can now use a
|
||||
different mesh asset for the same logical entity (e.g. world 0 holds a
|
||||
cube, world 1 a sphere), with weights controlling the proportion of
|
||||
worlds assigned to each variant. Mesh-derived constants (collision
|
||||
bounds, body inertials, subtree mass, inverse weights) are compiled
|
||||
per-variant and stored as per-world arrays in the Warp model, so domain
|
||||
randomization, the native viewer, the offscreen renderer, and the Viser
|
||||
viewer all pick up the variant assignment automatically. Variants must
|
||||
share the same kinematic structure (same bodies, joints, joint types);
|
||||
only mesh geoms may differ. Assignment is fixed at simulation init.
|
||||
See :ref:`per_world_mesh` for usage. With help from @XiangruiJiang.
|
||||
|
||||
Changed
|
||||
^^^^^^^
|
||||
|
||||
- Bumped ``mujoco`` to 3.8 and ``mujoco-warp`` to 3.8.0. The ``multiccd``
|
||||
enable flag was removed in mujoco 3.8 (it became default-on), so configs
|
||||
that listed ``"multiccd"`` in ``MujocoCfg.enableflags`` need to drop it.
|
||||
- Camera segmentation now matches ``mujoco_warp``'s typed segmentation
|
||||
output. ``CameraSensorData.segmentation`` stores ``(object_id,
|
||||
object_type)`` pairs in shape ``[B, H, W, 2]`` instead of the previous
|
||||
legacy geom-id-only layout. Contribution by @tkelestemur.
|
||||
- Sped up ``RayCaster`` post-processing by removing boolean-mask indexing
|
||||
operations and replacing them with ``masked_fill_`` plus a clamped-distance
|
||||
formulation of ``hit_pos_w`` that places misses at the world origin. This
|
||||
removes all CUDA syncs from the ray post-process, letting the CPU thread
|
||||
proceed while GPU-based sensing runs. Contribution by @bd-pdomanico.
|
||||
- Bumped ``rsl-rl-lib`` from 5.0.1 to 5.2.0. This brings ``torch.compile`` support for
|
||||
PPO and Distillation, and optional std clamping and constant std in
|
||||
``GaussianDistribution``. No code changes required on the mjlab side.
|
||||
- ``TerrainEntityCfg`` debug visualization sites (environment origins,
|
||||
terrain origins, flat patches) are now off by default. Set
|
||||
``debug_vis=True`` to re-enable them. The sites inflated ``nsite`` and
|
||||
caused a measurable slowdown in the per-step ``site_local_to_global``
|
||||
kernel (:issue:`942`).
|
||||
- Task package load failures during ``mjlab`` import now print the full
|
||||
traceback (and the entry point's module path) to ``stderr`` instead of
|
||||
just the exception message, making it easier to pinpoint the source of
|
||||
import errors when running commands like ``list-envs`` (:issue:`910`).
|
||||
Contribution by @saikishor.
|
||||
- Clarified ``ContactSensor`` shape conventions: per-contact fields
|
||||
(``found``, ``force``, ``torque``, ``dist``, ``pos``, ``normal``,
|
||||
``tangent``) have shape ``[B, P * num_slots, ...]`` while per-primary
|
||||
air-time fields (``current_air_time``, ``last_air_time``,
|
||||
``current_contact_time``, ``last_contact_time``) have shape ``[B, P]``,
|
||||
where ``P`` is the number of resolved primaries (:issue:`914`).
|
||||
|
||||
Fixed
|
||||
^^^^^
|
||||
|
||||
- Fixed ``ManagerBasedRlEnv`` initializing Warp on all visible CUDA devices
|
||||
even when constructed with ``device="cpu"``. ``seed_rng`` now accepts a
|
||||
``device`` argument and skips ``wp.rand_init`` on CPU devices, so a
|
||||
CPU-only env no longer claims a CUDA context on machines with a visible
|
||||
GPU (:issue:`949`).
|
||||
- Fixed ``ContactSensor.compute_first_contact`` and ``compute_first_air``
|
||||
occasionally missing events when a contact began or ended right at the
|
||||
last physics substep of a control step. ``current_contact_time`` /
|
||||
``current_air_time`` accumulate in float32 and can drift a few ULPs past
|
||||
``dt``, but the default ``abs_tol`` of ``1e-8`` sat at the noise floor
|
||||
and rejected the comparison. Raised the default to ``1e-6``, which stays
|
||||
well below typical control ``dt`` while comfortably covering float32
|
||||
accumulation noise (:issue:`933`). Contribution by @paLeziart.
|
||||
- Fixed ``out_of_terrain_bounds`` using stale terrain dimensions. It read
|
||||
``TerrainGeneratorCfg.num_cols`` directly, which is ignored in curriculum
|
||||
mode (the generator uses ``len(sub_terrains)`` columns instead), and it
|
||||
did not account for ``border_width``. The termination now reads the
|
||||
effective grid shape from ``terrain.terrain_origins`` and includes the
|
||||
border in the footprint, so robots no longer reset while still on valid
|
||||
terrain (or fail to reset after running off it) (:issue:`923`).
|
||||
- ``ObservationManager`` now skips observation groups that end up with
|
||||
zero active terms (e.g. all terms set to ``None``) with a log message,
|
||||
instead of crashing later in ``torch.stack``/``torch.cat``. This lets
|
||||
a shared runner config define groups that become empty under certain
|
||||
runtime flags (e.g. model-specific terms all disabled for one variant).
|
||||
The whole group can still be set to ``None`` to disable it explicitly.
|
||||
- Fixed a runtime broadcast error in ``ContactSensor`` when combining
|
||||
``num_slots > 1`` with ``track_air_time=True`` and more than one primary.
|
||||
Air-time tracking now reduces ``found`` across slots so that a primary is
|
||||
considered in contact when any of its slots reports a match (:issue:`914`).
|
||||
- Updated the ``create_new_task.ipynb`` Colab tutorial to import
|
||||
``XmlActuatorCfg`` instead of the removed ``XmlVelocityActuatorCfg``.
|
||||
Added a regression test (``tests/test_notebooks.py``) that parses each
|
||||
notebook cell and verifies that every ``from mjlab... import X``
|
||||
reference resolves, so future renames in the mjlab public API can't
|
||||
silently rot the tutorials (:issue:`913`).
|
||||
- Fixed ``ObservationManager`` silently sharing a single ``NoiseModelCfg``
|
||||
instance across observation groups that declared terms with the same
|
||||
name. ``_group_obs_class_instances`` was keyed by term name alone, so
|
||||
the last group processed in ``_prepare_terms`` overwrote earlier
|
||||
groups' instances. Symptoms included the wrong noise config being
|
||||
applied, shared per-episode state for ``NoiseModelWithAdditiveBias``
|
||||
(e.g. bias drawn from the wrong ``bias_noise_cfg``), and missed
|
||||
``reset()`` calls for overwritten instances. Instances are now keyed
|
||||
by ``(group_name, term_name)`` so each group owns its own noise model.
|
||||
- Fixed ``CurriculumManager.get_active_iterable_terms`` raising
|
||||
``TypeError`` when a term's state was a dict. The dict branch indexed
|
||||
the output list by term name instead of appending to the local ``data``
|
||||
list. No in-tree caller currently invokes this method, so the bug was
|
||||
latent.
|
||||
|
||||
Version 1.3.0 (April 14, 2026)
|
||||
------------------------------
|
||||
|
||||
Added
|
||||
^^^^^
|
||||
|
||||
- Added ``ManagerBasedRlEnvCfg.auto_reset`` flag. When ``True`` (default),
|
||||
``step()`` continues to reset done environments in place and returns the
|
||||
post-reset observation. When ``False``, ``step()`` skips the reset block
|
||||
and returns the terminal observation directly; the caller must call
|
||||
``reset(env_ids=...)`` for done environments before the next ``step()``
|
||||
or a ``RuntimeError`` is raised. Enables access to the true terminal
|
||||
state for algorithms that need it. Note that mjlab's bundled ``train.py``
|
||||
uses rsl_rl's ``OnPolicyRunner``, which does not drive manual resets, so
|
||||
``auto_reset=False`` is intended for custom training loops (:issue:`900`).
|
||||
- Added ``ActuatorCfg.viscous_damping`` for passive velocity proportional
|
||||
damping (``f = -b·v``), distinct from the PD derivative gain ``damping``
|
||||
used by position and velocity actuators. Maps to ``<joint damping>`` for
|
||||
JOINT transmission and ``<tendon damping>`` for TENDON transmission.
|
||||
Defaults to ``None`` (preserves the XML value).
|
||||
- Added :class:`~mjlab.managers.RecorderManager` for logging observations,
|
||||
actions, or arbitrary environment data during rollouts. Implement a
|
||||
:class:`~mjlab.managers.RecorderTerm` subclass and register it in the
|
||||
``recorders`` dict on ``ManagerBasedRlEnvCfg``. The manager provides
|
||||
``record_pre_reset``, ``record_post_reset``, and ``record_post_step``
|
||||
lifecycle hooks with no opinion on how data is stored.
|
||||
- Added :func:`~mjlab.envs.mdp.curriculums.termination_curriculum` for
|
||||
scheduling changes to termination term parameters during training,
|
||||
matching the existing ``reward_curriculum`` pattern. Both now share a
|
||||
single internal engine with init-time validation of stage ordering,
|
||||
field existence, and param keys.
|
||||
- Added ``reduce`` field to ``MetricsTermCfg``. Setting ``reduce="last"``
|
||||
reports the value from the final step of the episode rather than the
|
||||
episode mean, which is useful for binary success metrics.
|
||||
- Added :class:`~mjlab.envs.mdp.actions.RelativeJointPositionAction` for
|
||||
joint position control relative to the current configuration. The target is
|
||||
``current_pos + action * scale``, so a zero action holds the current
|
||||
configuration rather than commanding the default pose.
|
||||
- Added :func:`~mjlab.envs.mdp.dr.pair_friction` for randomizing geom-pair
|
||||
friction overrides (``pair_friction`` in ``mjModel``), with an
|
||||
``isotropic=True`` option that mirrors the symmetric tangent and roll
|
||||
axes so single-axis randomization does not leave the paired axis stale.
|
||||
- Added ``STAIRS_TERRAINS_CFG`` terrain preset for progressive stair
|
||||
curriculum training and ``@terrain_preset`` decorator for composing
|
||||
terrain configurations from reusable presets.
|
||||
- Added cartpole balance and swingup tasks (``Mjlab-Cartpole-Balance`` and
|
||||
``Mjlab-Cartpole-Swingup``) with a :ref:`tutorial <tutorial-cartpole>`
|
||||
that walks through building an environment from scratch.
|
||||
- Added :ref:`motion imitation <motion-imitation>` documentation with
|
||||
preprocessing instructions. The README now links here instead of the
|
||||
BeyondMimic repository, which produced incompatible NPZ files when used
|
||||
with mjlab (:issue:`777`).
|
||||
- Added ``margin``, ``gap``, and ``solmix`` fields to ``CollisionCfg``
|
||||
for per geom contact parameter configuration (:issue:`766`).
|
||||
- NaN guard now captures mocap body poses (``mocap_pos``, ``mocap_quat``)
|
||||
when the model has mocap bodies, enabling full state reconstruction in
|
||||
the dump viewer for fixed-base entities.
|
||||
- Implemented ``ActionTermCfg.clip`` for clamping processed actions after
|
||||
scale and offset (:issue:`771`).
|
||||
- Added ``qfrc_actuator`` and ``qfrc_external`` generalized force accessors
|
||||
to ``EntityData``. ``qfrc_actuator`` gives actuator forces in joint space
|
||||
(projected through the transmission). ``qfrc_external`` recovers the
|
||||
generalized force from body external wrenches (``xfrc_applied``)
|
||||
(:issue:`776`).
|
||||
- Added ``RewardBarPanel`` to the Viser viewer, showing horizontal bars for
|
||||
each reward term with a running mean over ~1 second (:issue:`800`).
|
||||
- Added ``per_substep`` flag to ``MetricsTermCfg`` for evaluating metrics
|
||||
once per physics substep inside the decimation loop. The per substep
|
||||
values are averaged within each environment step, so episode averages
|
||||
remain comparable to regular per step metrics.
|
||||
- Added ``project-instinct/InstinctMJ`` to the research page's list of
|
||||
projects built on mjlab.
|
||||
- Added a Checkpoints tab to the Viser play viewer for hot-swapping
|
||||
checkpoints without restarting. Works with local directories and W&B
|
||||
runs (:issue:`751`). Contribution by @omarrayyann.
|
||||
- Added ``"segmentation"`` camera data type for per-pixel geom ID output
|
||||
alongside RGB and depth, and a multi-cube goal-conditioned lifting task
|
||||
(``Mjlab-Multi-Cube-Seg-Yam``) that uses it (:issue:`862`).
|
||||
Contribution by @pthangeda.
|
||||
|
||||
Changed
|
||||
^^^^^^^
|
||||
|
||||
- Renamed the ``list_envs`` console script to ``list-envs`` for consistency
|
||||
with the other hyphenated entry points (``viz-nan``, ``export-scene``).
|
||||
Invoke via ``uv run list-envs``.
|
||||
- ``ActuatorCfg.armature`` and ``ActuatorCfg.frictionloss`` now default to
|
||||
``None`` instead of ``0.0``. ``None`` preserves the value defined in the
|
||||
XML. Previously, builtin actuators would silently overwrite XML joint and
|
||||
tendon properties with zero when these fields were not explicitly set.
|
||||
To restore the old behavior, pass ``armature=0.0`` or ``frictionloss=0.0``
|
||||
explicitly.
|
||||
- Actuator delay is now configured inline on any ``ActuatorCfg`` subclass
|
||||
(e.g. ``BuiltinPositionActuatorCfg(..., delay_min_lag=2, delay_max_lag=5)``)
|
||||
instead of wrapping with ``DelayedActuatorCfg``. ``DelayedActuator``,
|
||||
``DelayedActuatorCfg``, and ``DelayedBuiltinActuatorGroup`` are removed.
|
||||
- Removed ``delay_target`` from ``ActuatorCfg``. Delay now always applies to
|
||||
the actuator's ``command_field`` automatically. Multi-target delay
|
||||
(``delay_target=("position", "velocity")``) is no longer supported.
|
||||
- ``XmlPositionActuatorCfg``, ``XmlVelocityActuatorCfg``, ``XmlMotorActuatorCfg``,
|
||||
and ``XmlMuscleActuatorCfg`` are replaced by a single ``XmlActuatorCfg`` that auto
|
||||
detects the actuator type from XML. Pass ``command_field=...`` to override detection.
|
||||
- Replaced the viser viewer internals with the ``mjviser`` package. Scene
|
||||
creation, mesh conversion, and overlay rendering (contacts, forces,
|
||||
inertia, tendons, joints, frames) are now provided by mjviser. The viewer
|
||||
exposes a new Visualization tab for overlay controls and a Groups tab for
|
||||
geom/site visibility. Debug visualization and warp tensor conversion remain
|
||||
in mjlab's ``MjlabViserScene`` subclass (:issue:`839`).
|
||||
- In curriculum terrain mode, each terrain type now gets exactly one column
|
||||
(``num_cols`` is set to ``len(sub_terrains)``). The ``proportion`` field
|
||||
now controls robot spawning distribution across columns rather than column
|
||||
count. Random mode is unchanged (:issue:`811`).
|
||||
- ``BoxSteppingStonesTerrainCfg`` stone size now decreases with difficulty,
|
||||
interpolating from the large end of ``stone_size_range`` at difficulty 0
|
||||
to the small end at difficulty 1 (:issue:`785`).
|
||||
- Removed deprecated ``TerrainImporter`` and ``TerrainImporterCfg`` aliases.
|
||||
Use ``TerrainEntity`` and ``TerrainEntityCfg`` instead (:issue:`667`).
|
||||
- ``Entity.clear_state()`` is deprecated. Use ``Entity.reset()`` instead.
|
||||
``clear_state`` only zeroed actuator targets without resetting actuator
|
||||
internal state (e.g. delay buffers), which could cause stale commands
|
||||
after teleporting the robot to a new pose.
|
||||
- Removed ``EntityData.generalized_force``. The property was bugged (indexed
|
||||
free joint DOFs instead of articulated DOFs) and the name was ambiguous.
|
||||
Use ``qfrc_actuator`` or ``qfrc_external`` instead (:issue:`776`).
|
||||
- ``get_wandb_checkpoint_path`` now filters checkpoints server-side via the
|
||||
``pattern`` parameter, avoiding unnecessary pagination and tolerance to
|
||||
corrupted metadata (:issue:`898`).
|
||||
|
||||
Fixed
|
||||
^^^^^
|
||||
|
||||
- ``train`` and ``play`` now print a top-level usage message when invoked
|
||||
with ``-h`` / ``--help`` and no task argument, pointing users at
|
||||
``list-envs`` and ``<TASK> --help`` (:issue:`905`).
|
||||
- Fixed ghost geom filtering in the Viser viewer. Ghost geoms were selected
|
||||
by collision flags, so collision-disabled robot geoms appeared as ghosts.
|
||||
The viewer now uses visual alpha to determine which geoms to render.
|
||||
- Scene now warns when an attached entity or terrain spec has non-default
|
||||
``<option>`` fields (e.g. ``<flag contact="disable"/>``), which are
|
||||
silently dropped by ``MjSpec.attach()``. Use ``MujocoCfg`` to set
|
||||
simulation options instead (:issue:`885`).
|
||||
- Fixed ``SceneEntityCfg`` names and IDs ordering mismatch when
|
||||
``preserve_order=False`` (:issue:`876`). Contribution by @jsw7460.
|
||||
- Fixed ONNX export path resolution in the velocity, manipulation, and
|
||||
tracking runners when a parent directory name contains the word
|
||||
``"model"`` (:issue:`867`). Contribution by @gokulp01.
|
||||
- ``export-scene`` now writes only referenced assets and places them
|
||||
correctly under the output directory. Previously, asset keys containing
|
||||
path traversal could write files outside the output directory, and all
|
||||
spec assets were included regardless of whether the scene XML referenced
|
||||
them (:issue:`858`).
|
||||
- ``electrical_power_cost`` now uses ``qfrc_actuator`` (joint space) instead
|
||||
of ``actuator_force`` (actuation space) for mechanical power computation.
|
||||
Previously the reward was incorrect for actuators with gear ratios other
|
||||
than 1 (:issue:`776`).
|
||||
- ``create_velocity_actuator`` no longer sets ``ctrllimited=True`` with
|
||||
``inheritrange=1.0``. This caused a ``ValueError`` for continuous joints
|
||||
(e.g. wheels) that have no position range defined (:issue:`787`).
|
||||
- ``write_root_com_velocity_to_sim`` no longer fails with tensor ``env_ids``
|
||||
on floating base entities (:issue:`793`).
|
||||
- Joint limits for unlimited joints are now set to [-inf, inf] instead of
|
||||
[0, 0]. Previously the zero range caused incorrect clamping for entities
|
||||
with unlimited hinge or slide joints.
|
||||
- Contact force visualization now copies ``ctrl`` into the CPU ``MjData``
|
||||
before calling ``mj_forward``. Actuators that compute torques in Python
|
||||
(``DcMotorActuator``, ``IdealPdActuator``) previously showed incorrect
|
||||
contact forces because the viewer ran with ``ctrl=0``
|
||||
(:issue:`786`).
|
||||
- ``BoxSteppingStonesTerrainCfg`` no longer creates a large gap around the
|
||||
platform. Stones are now only skipped when their center falls inside the
|
||||
platform; edges that extend under the platform are allowed since the
|
||||
platform covers them (:issue:`785`).
|
||||
- ``dr.pseudo_inertia`` no longer loads cuSOLVER, eliminating ~4 GB of
|
||||
persistent GPU memory overhead. Cholesky and eigendecomposition are now
|
||||
computed analytically for the small matrices involved (4x4 and 3x3)
|
||||
(:issue:`753`).
|
||||
- Set terrain geom mass to zero so that the static terrain body does not
|
||||
inflate ``stat.meanmass``, which made force arrow visualization invisible
|
||||
on rough terrain (:issue:`734`, :issue:`537`).
|
||||
- Native viewer now syncs ``qpos0`` when domain randomized, fixing incorrect
|
||||
body positions after ``dr.joint_default_pos`` randomization
|
||||
(:issue:`760`).
|
||||
- ``command_manager.compute()`` is now called during ``reset()`` so that
|
||||
derived command state (e.g. relative body positions in tracking
|
||||
environments) is populated before the first observation is returned
|
||||
(:issue:`761`).
|
||||
- ``RayCastSensor`` with ``ray_alignment="yaw"`` or ``"world"`` now correctly
|
||||
aligns the frame offset when attached to a site or geom with a local offset
|
||||
from its parent body. Previously only ray directions and pattern offsets were
|
||||
aligned, causing the frame position to swing with body pitch/roll
|
||||
(:issue:`775`).
|
||||
|
||||
Version 1.2.0 (March 6, 2026)
|
||||
-----------------------------
|
||||
|
||||
.. admonition:: Breaking API changes
|
||||
:class: attention
|
||||
|
||||
- ``randomize_field`` no longer exists. Replace calls with typed functions
|
||||
from the new ``dr`` module (e.g. ``dr.geom_friction``, ``dr.body_mass``).
|
||||
- ``EventTermCfg`` no longer accepts ``domain_randomization``. The
|
||||
``@requires_model_fields`` decorator on each ``dr`` function takes care
|
||||
of field expansion automatically.
|
||||
- ``Scene.to_zip()`` is deprecated. Use ``Scene.write(path, zip=True)``.
|
||||
- ``RslRlModelCfg`` no longer accepts ``stochastic``, ``init_noise_std``,
|
||||
or ``noise_std_type``. Use ``distribution_cfg`` instead
|
||||
(e.g. ``{"class_name": "GaussianDistribution", "init_std": 1.0,
|
||||
"std_type": "scalar"}``). Existing checkpoints are automatically
|
||||
migrated on load.
|
||||
|
||||
Added
|
||||
^^^^^
|
||||
|
||||
- Added ``"step"`` event mode that fires every environment step.
|
||||
- Added ``apply_body_impulse`` event for applying transient external wrenches
|
||||
to bodies with configurable duration and optional application point offset.
|
||||
- ONNX auto-export and metadata attachment for manipulation tasks (lift cube)
|
||||
on every checkpoint save, matching the velocity and tracking task behavior.
|
||||
- Multi-frame ``RayCastSensor``: pass a tuple of ``ObjRef`` to ``frame`` for
|
||||
per-site raycasting with independent body exclusion. New properties:
|
||||
``num_frames``, ``num_rays_per_frame``. New ``RayCastData`` fields:
|
||||
``frame_pos_w`` and ``frame_quat_w``.
|
||||
- ``RingPatternCfg`` ray pattern for concentric ring sampling around each
|
||||
frame.
|
||||
- ``TerrainHeightSensor``, a ``RayCastSensor`` subclass that computes
|
||||
per-frame vertical clearance above terrain (``sensor.data.heights``).
|
||||
Velocity task configs now use it for ``feet_clearance``,
|
||||
``feet_swing_height``, and ``foot_height``, replacing the previous
|
||||
world-Z proxy that was incorrect on rough terrain.
|
||||
- Cloud training support via `SkyPilot <https://skypilot.readthedocs.io/>`_
|
||||
and Lambda Cloud, with documentation covering setup, monitoring, and
|
||||
cost management.
|
||||
- W&B hyperparameter sweep scripts that distribute one agent per GPU
|
||||
across a multi-GPU instance.
|
||||
- Contributing guide with documentation for shared Claude Code commands
|
||||
(``/update-mjwarp``, ``/commit-push-pr``).
|
||||
- Added optional ``ViewerConfig.fovy`` and apply it in native viewer camera
|
||||
setup when provided.
|
||||
- Native viewer now tracks the first non-fixed body by default (matching
|
||||
the Viser viewer behavior introduced in
|
||||
``716aaaa58ad7bfaf34d2f771549d461204d1b4ba``).
|
||||
- New ``dr`` module (``mjlab.envs.mdp.dr``) replacing ``randomize_field``
|
||||
with typed per-field domain randomization functions. Each function
|
||||
automatically recomputes derived fields via ``set_const``. Highlights:
|
||||
|
||||
- Camera and light randomization: ``dr.cam_fovy``, ``dr.cam_pos``,
|
||||
``dr.cam_quat``, ``dr.cam_intrinsic``, ``dr.light_pos``,
|
||||
``dr.light_dir``. Camera and light names are now supported in
|
||||
``SceneEntityCfg`` (``camera_names`` / ``light_names``).
|
||||
- ``dr.pseudo_inertia`` for physics-consistent randomization of
|
||||
``body_mass``, ``body_ipos``, ``body_inertia``, and ``body_iquat``
|
||||
via the pseudo-inertia matrix parameterization (Rucker & Wensing
|
||||
2022). Replaces the removed ``dr.body_inertia`` /
|
||||
``dr.body_iquat``.
|
||||
- ``dr.geom_size`` with automatic recomputation of ``geom_rbound``
|
||||
and ``geom_aabb`` for broadphase consistency.
|
||||
- ``dr.tendon_armature`` and ``dr.tendon_frictionloss``.
|
||||
- ``dr.body_quat``, ``dr.geom_quat``, and ``dr.site_quat`` with RPY
|
||||
perturbation composed onto the default quaternion.
|
||||
- Extensible ``Operation`` and ``Distribution`` types. Users can define
|
||||
custom operations and distributions as class instances and pass them
|
||||
anywhere a string is accepted. Built-in instances (``dr.abs``,
|
||||
``dr.scale``, ``dr.add``, ``dr.uniform``, ``dr.log_uniform``,
|
||||
``dr.gaussian``) are exported from the ``dr`` module.
|
||||
- ``dr.mat_rgba`` for per-world material color randomization. Tints
|
||||
the texture color, useful for randomizing appearance of textured
|
||||
surfaces. Material names are now supported in ``SceneEntityCfg``
|
||||
(``material_names``).
|
||||
- Fixed ``dr.effort_limits`` drifting on repeated randomization.
|
||||
- Fixed ``dr.body_com_offset`` not triggering ``set_const``.
|
||||
|
||||
- ``export-scene`` CLI script to export any task scene or asset_zoo entity
|
||||
(``g1``, ``go1``, ``yam``) to a directory or zip archive for inspection
|
||||
and debugging.
|
||||
|
||||
- ``yam_lift_cube_vision_env_cfg`` now randomizes cube color (``dr.geom_rgba``)
|
||||
on every reset when ``cam_type="rgb"``.
|
||||
|
||||
- The native viewer now reflects per-world DR changes to visual model fields
|
||||
on each reset. Geom appearance, body and site poses, camera parameters,
|
||||
and light positions are all synced from the GPU model before rendering.
|
||||
Inertia boxes (press ``I``) and camera frustums (press ``Q``) update
|
||||
correctly when the corresponding fields are randomized. See
|
||||
:doc:`randomization` for viewer-specific caveats.
|
||||
|
||||
- ``MaterialCfg.geom_names_expr`` for assigning materials to geoms by
|
||||
name pattern during ``edit_spec``.
|
||||
|
||||
- ``TerrainEntityCfg`` now exposes ``textures``, ``materials``, and
|
||||
``lights`` as configurable fields (previously hardcoded). Set
|
||||
``textures=()``, ``materials=()`` to use flat ``dr.geom_rgba``
|
||||
instead of the default checker texture.
|
||||
|
||||
- ``DebugVisualizer`` now supports ellipsoid visualization via
|
||||
``add_ellipsoid``.
|
||||
|
||||
- Interactive velocity joystick sliders in the Viser viewer. Enable the
|
||||
joystick under Commands/Twist to override velocity commands with manual
|
||||
sliders for ``lin_vel_x``, ``lin_vel_y``, and ``ang_vel_z``
|
||||
(`#666 <https://github.com/mujocolab/mjlab/issues/666>`_).
|
||||
- Per-term debug visualization toggles in the Viser viewer. Individual
|
||||
command term visualizers (e.g. velocity arrows) can now be toggled
|
||||
independently under Scene/Debug Viz.
|
||||
- Viewer single-step mode: press RIGHT arrow (native) or click "Step"
|
||||
(Viser) to advance exactly one physics step while paused.
|
||||
- Viewer error recovery: exceptions during stepping now pause the viewer
|
||||
and log the traceback instead of crashing the process.
|
||||
- Native viewer runs forward kinematics while paused, keeping
|
||||
perturbation visuals accurate.
|
||||
- Viewer speed multipliers use clean power-of-2 fractions (1/32x to 1x).
|
||||
|
||||
- Visualizers display the realtime factor alongside FPS.
|
||||
|
||||
- ``joint_torques_l2`` now respects ``SceneEntityCfg.actuator_ids``,
|
||||
allowing penalization of a subset of actuators instead of all of them
|
||||
(`#703 <https://github.com/mujocolab/mjlab/pull/703>`_). Contribution by
|
||||
`@saikishor <https://github.com/saikishor>`_.
|
||||
|
||||
- Terrain is now a proper ``Entity`` subclass (``TerrainEntity``). This
|
||||
allows domain randomization functions to target terrain parameters
|
||||
(friction, cameras, lights) via ``SceneEntityCfg("terrain", ...)``.
|
||||
``TerrainImporter`` / ``TerrainImporterCfg`` remain as aliases but will be
|
||||
deprecated in a future version.
|
||||
- Added ``upload_model`` option to ``RslRlBaseRunnerCfg`` to control W&B model
|
||||
file uploads (``.pt`` and ``.onnx``) while keeping metric logging enabled
|
||||
(`#654 <https://github.com/mujocolab/mjlab/pull/654>`_).
|
||||
- ``Scene.write(output_dir, zip=False)`` exports the scene XML and mesh
|
||||
assets to a directory (or zip archive). Replaces ``Scene.to_zip()``.
|
||||
- ``Entity.write_xml()`` and ``Scene.write()`` now apply XML fixups
|
||||
(empty defaults, duplicate nested defaults) and strip buffer textures
|
||||
that ``MjSpec.to_xml()`` cannot serialize.
|
||||
- ``fix_spec_xml`` and ``strip_buffer_textures`` utilities in
|
||||
``mjlab.utils.xml``.
|
||||
|
||||
Changed
|
||||
^^^^^^^
|
||||
|
||||
- Native viewer now syncs ``xfrc_applied`` to the render buffer and draws
|
||||
arrows for any nonzero applied forces. Mouse perturbation forces are
|
||||
converted to ``qfrc_applied`` (generalized joint space) so they coexist
|
||||
with programmatic forces on ``xfrc_applied`` without conflict.
|
||||
- ``ViewerConfig.OriginType.WORLD`` now configures a free camera at the
|
||||
specified lookat point instead of auto tracking a body. A new ``AUTO``
|
||||
origin type (now the default) preserves the previous auto tracking
|
||||
behavior.
|
||||
- Upgraded ``rsl-rl-lib`` from 4.0.1 to 5.0.1. ``RslRlModelCfg`` now
|
||||
uses ``distribution_cfg`` dict instead of ``stochastic`` /
|
||||
``init_noise_std`` / ``noise_std_type``. Existing checkpoints are
|
||||
automatically migrated on load.
|
||||
- Reorganized the Viser Controls tab into a cleaner folder hierarchy:
|
||||
Info, Simulation, Commands, Scene (with Environment, Camera, Debug Viz,
|
||||
Contacts sub-folders), and Camera Feeds. The Environment folder is
|
||||
hidden for single-env tasks and the Commands folder is hidden when no
|
||||
command terms are active.
|
||||
- Viser camera tracking is now enabled by default so the agent stays in
|
||||
frame on launch.
|
||||
- Self collision and illegal contact sensors now use ``history_length`` to
|
||||
catch contacts across decimation substeps. Reward and termination functions
|
||||
read ``force_history`` with a configurable ``force_threshold``.
|
||||
- Replaced the single ``scale`` parameter in ``DifferentialIKActionCfg`` with
|
||||
separate ``delta_pos_scale`` and ``delta_ori_scale`` for independent scaling
|
||||
of position and orientation components.
|
||||
- Improved offscreen multi environment framing by selecting neighboring
|
||||
environments around the focused env instead of first N envs.
|
||||
- Tuned tracking task viewer defaults for tighter camera framing.
|
||||
- Disabled shadow casting on the G1 tracking light to avoid duplicate
|
||||
stacked shadows when robots are close.
|
||||
|
||||
Fixed
|
||||
^^^^^
|
||||
|
||||
- Fixed actuator target resolution for entities whose ``spec_fn`` uses
|
||||
internal ``MjSpec.attach(prefix=...)``
|
||||
(`#709 <https://github.com/mujocolab/mjlab/issues/709>`_).
|
||||
- Fixed viewer physics loop starving the renderer by replacing the single
|
||||
sim-time budget with a two-clock design (tracked vs actual sim time).
|
||||
Physics now self-corrects after overshooting, keeping FPS smooth at all
|
||||
speed multipliers.
|
||||
- Bundled ``ffmpeg`` for ``mediapy`` via ``imageio-ffmpeg``, removing the
|
||||
requirement for a system ``ffmpeg`` install. Thanks to
|
||||
`@rdeits-bd <https://github.com/rdeits-bd>`_ for the suggestion.
|
||||
- Fixed ``height_scan`` returning ~0 for missed rays; now defaults to
|
||||
``max_distance``. Replaced ``clip=(-1, 1)`` with ``scale`` normalization
|
||||
in the velocity task config. Thanks to `@eufrizz <https://github.com/eufrizz>`_
|
||||
for reporting and the initial fix (`#642 <https://github.com/mujocolab/mjlab/pull/642>`_).
|
||||
- Fixed ghost mesh visualization for fixed-base entities by extending
|
||||
``DebugVisualizer.add_ghost_mesh`` to optionally accept ``mocap_pos`` and
|
||||
``mocap_quat`` (`#645 <https://github.com/mujocolab/mjlab/pull/645>`_).
|
||||
- Fixed viser viewer crashing on scenes with no mocap bodies by adding
|
||||
an ``nmocap`` guard, matching the native viewer behavior.
|
||||
- Fixed offscreen rendering artifacts in large vectorized scenes by applying
|
||||
a render local extent override in ``OffscreenRenderer`` and restoring the
|
||||
original extent on close.
|
||||
- Fixed ``RslRlVecEnvWrapper.unwrapped`` to return the base environment,
|
||||
ensuring checkpoint state restore and logging work correctly when wrappers
|
||||
such as ``VideoRecorder`` are enabled.
|
||||
|
||||
Version 1.1.1 (February 14, 2026)
|
||||
---------------------------------
|
||||
|
||||
Added
|
||||
^^^^^
|
||||
|
||||
- Added reward term visualization to the native viewer (toggle with ``P``) (`#629 <https://github.com/mujocolab/mjlab/pull/629>`_).
|
||||
- Added ``DifferentialIKAction`` for task-space control via damped
|
||||
least-squares IK. Supports weighted position/orientation tracking,
|
||||
soft joint-limit avoidance, and null-space posture regularization.
|
||||
Includes an interactive viser demo (``scripts/demos/differential_ik.py``) (`#632 <https://github.com/mujocolab/mjlab/pull/632>`_).
|
||||
|
||||
Fixed
|
||||
^^^^^
|
||||
|
||||
- Fixed ``play.py`` defaulting to the base rsl-rl ``OnPolicyRunner`` instead
|
||||
of ``MjlabOnPolicyRunner``, which caused a ``TypeError`` from an unexpected
|
||||
``cnn_cfg`` keyword argument (`#626 <https://github.com/mujocolab/mjlab/pull/626>`_). Contribution by
|
||||
`@griffinaddison <https://github.com/griffinaddison>`_.
|
||||
|
||||
Changed
|
||||
^^^^^^^
|
||||
|
||||
- Removed ``body_mass``, ``body_inertia``, ``body_pos``, and ``body_quat``
|
||||
from ``FIELD_SPECS`` in domain randomization. These fields have derived
|
||||
quantities that require ``set_const`` to recompute; without that call,
|
||||
randomizing them silently breaks physics (`#631 <https://github.com/mujocolab/mjlab/pull/631>`_).
|
||||
- Replaced ``moviepy`` with ``mediapy`` for video recording. ``mediapy``
|
||||
handles cloud storage paths (GCS, S3) natively (`#637 <https://github.com/mujocolab/mjlab/pull/637>`_).
|
||||
|
||||
.. figure:: _static/changelog/native_reward.png
|
||||
:width: 80%
|
||||
|
||||
Version 1.1.0 (February 12, 2026)
|
||||
---------------------------------
|
||||
|
||||
Added
|
||||
^^^^^
|
||||
|
||||
- Added RGB and depth camera sensors and BVH-accelerated raycasting (`#597 <https://github.com/mujocolab/mjlab/pull/597>`_).
|
||||
- Added ``MetricsManager`` for logging custom metrics during training (`#596 <https://github.com/mujocolab/mjlab/pull/596>`_).
|
||||
- Added terrain visualizer (`#609 <https://github.com/mujocolab/mjlab/pull/609>`_). Contribution by
|
||||
`@mktk1117 <https://github.com/mktk1117>`_.
|
||||
|
||||
.. figure:: _static/changelog/terrain_visualizer.jpg
|
||||
:width: 80%
|
||||
|
||||
- Added many new terrains including ``HfDiscreteObstaclesTerrainCfg``,
|
||||
``HfPerlinNoiseTerrainCfg``, ``BoxSteppingStonesTerrainCfg``,
|
||||
``BoxNarrowBeamsTerrainCfg``, ``BoxRandomStairsTerrainCfg``, and
|
||||
more. Added flat patch sampling for heightfield terrains (`#542 <https://github.com/mujocolab/mjlab/pull/542>`_, `#581 <https://github.com/mujocolab/mjlab/pull/581>`_).
|
||||
- Added site group visualization to the Viser viewer (Geoms and Sites
|
||||
tabs unified into a single Groups tab) (`#551 <https://github.com/mujocolab/mjlab/pull/551>`_).
|
||||
- Added ``env_ids`` parameter to ``Entity.write_ctrl_to_sim`` (`#567 <https://github.com/mujocolab/mjlab/pull/567>`_).
|
||||
|
||||
Changed
|
||||
^^^^^^^
|
||||
|
||||
- Upgraded ``rsl-rl-lib`` to 4.0.0 and replaced the custom ONNX
|
||||
exporter with rsl-rl's built-in ``as_onnx()`` (`#589 <https://github.com/mujocolab/mjlab/pull/589>`_, `#595 <https://github.com/mujocolab/mjlab/pull/595>`_).
|
||||
- ``sim.forward()`` is now called unconditionally after the decimation
|
||||
loop. See :ref:`faq-sim-forward` for details (`#591 <https://github.com/mujocolab/mjlab/pull/591>`_).
|
||||
- Unnamed freejoints are now automatically named to prevent
|
||||
``KeyError`` during entity init (`#545 <https://github.com/mujocolab/mjlab/pull/545>`_).
|
||||
|
||||
Fixed
|
||||
^^^^^
|
||||
|
||||
- Fixed ``randomize_pd_gains`` crash with ``num_envs > 1`` (`#564 <https://github.com/mujocolab/mjlab/pull/564>`_).
|
||||
- Fixed ``ctrl_ids`` index error with multiple actuated entities (`#573 <https://github.com/mujocolab/mjlab/pull/573>`_).
|
||||
Reported by `@bwrooney82 <https://github.com/bwrooney82>`_.
|
||||
- Fixed Viser viewer rendering textured robots as gray (`#544 <https://github.com/mujocolab/mjlab/pull/544>`_).
|
||||
- Fixed Viser plane rendering ignoring MuJoCo size parameter (`#540 <https://github.com/mujocolab/mjlab/pull/540>`_).
|
||||
- Fixed ``HfDiscreteObstaclesTerrainCfg`` spawn height (`#552 <https://github.com/mujocolab/mjlab/pull/552>`_).
|
||||
- Fixed ``RaycastSensor`` visualization ignoring the all-envs toggle (`#607 <https://github.com/mujocolab/mjlab/pull/607>`_).
|
||||
Contribution by `@oxkitsune <https://github.com/oxkitsune>`_.
|
||||
|
||||
Version 1.0.0 (January 28, 2026)
|
||||
--------------------------------
|
||||
|
||||
Initial release of mjlab.
|
||||
@@ -0,0 +1,109 @@
|
||||
.. _commands:
|
||||
|
||||
Commands
|
||||
========
|
||||
|
||||
Commands specify what the policy should achieve at each moment: a target
|
||||
velocity, a reference trajectory, a goal position. The command manager
|
||||
generates these signals, resamples them at configurable intervals, and
|
||||
passes them to the policy through the observation system.
|
||||
|
||||
|
||||
Registration
|
||||
------------
|
||||
|
||||
Commands are registered in ``ManagerBasedRlEnvCfg`` as a dictionary
|
||||
mapping string names to ``CommandTermCfg`` instances. Unlike the
|
||||
function-based terms used by other managers, every command term is a
|
||||
class that inherits from ``CommandTerm``.
|
||||
|
||||
The ``resampling_time_range`` field controls how often the command
|
||||
changes. After each resample the term draws a new timer value uniformly
|
||||
from the given ``(min, max)`` range in seconds. Commands are also
|
||||
resampled unconditionally on every episode reset.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
commands = {
|
||||
"twist": UniformVelocityCommandCfg(
|
||||
entity_name="robot",
|
||||
resampling_time_range=(3.0, 8.0),
|
||||
ranges=UniformVelocityCommandCfg.Ranges(
|
||||
lin_vel_x=(-1.0, 1.0),
|
||||
lin_vel_y=(-1.0, 1.0),
|
||||
ang_vel_z=(-0.5, 0.5),
|
||||
),
|
||||
),
|
||||
}
|
||||
|
||||
The ``generated_commands`` observation function reads the current
|
||||
command tensor by name and passes it to the policy:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
ObservationTermCfg(
|
||||
func=mdp.generated_commands,
|
||||
params={"command_name": "twist"},
|
||||
)
|
||||
|
||||
If the environment has no commands, the manager no-ops all operations
|
||||
and returns empty tensors. There is no special handling required.
|
||||
|
||||
|
||||
Included command terms
|
||||
----------------------
|
||||
|
||||
Each task ships with its own command terms tailored to its objective.
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 28 72
|
||||
|
||||
* - Term
|
||||
- Description
|
||||
* - ``UniformVelocityCommand``
|
||||
- Generates planar velocity commands ``[v_x, v_y, omega_z]``
|
||||
sampled uniformly from configurable ranges. Supports a standing
|
||||
mode (fraction of environments receive zero velocity) and a
|
||||
heading mode (yaw rate replaced by a proportional controller
|
||||
tracking a sampled heading angle). Used by the velocity task.
|
||||
* - ``LiftingCommand``
|
||||
- Generates a 3D target position for a manipulated object.
|
||||
Supports fixed and dynamic difficulty modes. Tracks metrics
|
||||
including position error and episode success rate. Used by the
|
||||
manipulation task.
|
||||
* - ``MotionCommand``
|
||||
- Streams reference joint positions, velocities, and body poses
|
||||
from a pre-recorded ``.npz`` motion clip. Supports three
|
||||
start-frame sampling modes: ``"start"`` (always frame 0),
|
||||
``"uniform"`` (random), and ``"adaptive"`` (biased toward
|
||||
difficult regions). At reset the robot is initialized from the
|
||||
sampled frame with optional perturbations. Used by the tracking
|
||||
task.
|
||||
|
||||
Each term can render debug visualizations in the interactive viewer
|
||||
when ``debug_vis=True`` is set in the configuration. The image below
|
||||
shows the ghost visualization from ``MotionCommand``, which renders a
|
||||
translucent copy of the robot at the reference pose alongside the
|
||||
actual robot.
|
||||
|
||||
.. figure:: _static/ghost_visualization.png
|
||||
:align: center
|
||||
:width: 100%
|
||||
|
||||
Viser visualization of the commanded reference motion for the G1 tracking task.
|
||||
|
||||
|
||||
Writing custom command terms
|
||||
-----------------------------
|
||||
|
||||
A custom command term is a class inheriting from ``CommandTerm`` paired
|
||||
with a configuration dataclass inheriting from ``CommandTermCfg``. The
|
||||
term must implement four methods: ``_resample_command(env_ids)`` to
|
||||
sample new goals, ``_update_command()`` for per-step updates,
|
||||
``_update_metrics()`` for logging, and a ``command`` property returning
|
||||
the current goal tensor. The base class manages the resampling timer
|
||||
and reset logic automatically.
|
||||
|
||||
The configuration must implement a ``build(env)`` method that
|
||||
constructs the paired term instance.
|
||||
@@ -0,0 +1,108 @@
|
||||
Contributing
|
||||
============
|
||||
|
||||
Bug fixes and documentation improvements are always welcome.
|
||||
|
||||
.. important::
|
||||
|
||||
For new features, please
|
||||
`open an issue <https://github.com/mujocolab/mjlab/issues>`_ first so
|
||||
we can discuss whether it fits the project scope.
|
||||
|
||||
|
||||
Development setup
|
||||
-----------------
|
||||
|
||||
Clone the repository and sync dependencies:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
git clone https://github.com/mujocolab/mjlab.git && cd mjlab
|
||||
uv sync
|
||||
|
||||
Install pre-commit hooks to catch formatting and lint issues before each
|
||||
commit:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
uvx pre-commit install
|
||||
|
||||
|
||||
Common commands
|
||||
---------------
|
||||
|
||||
The ``Makefile`` provides shortcuts for the most common development tasks:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
make format # Format code and fix lint errors (ruff)
|
||||
make type # Type check (ty + pyright)
|
||||
make check # Format + type check
|
||||
make test-fast # Run tests, excluding slow ones
|
||||
make test # Run the full test suite
|
||||
make test-all # Format + type check + full test suite
|
||||
|
||||
You can also run individual tests for faster iteration:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
uv run pytest tests/test_rewards.py
|
||||
|
||||
Type checking (``make type``) is required. PRs that do not pass will be
|
||||
blocked.
|
||||
|
||||
|
||||
Building the docs
|
||||
-----------------
|
||||
|
||||
Build the documentation locally:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
make docs
|
||||
|
||||
The HTML output is written to ``docs/_build/``. For live reload during
|
||||
editing:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
make docs-watch
|
||||
|
||||
|
||||
Submitting a pull request
|
||||
-------------------------
|
||||
|
||||
1. Fork the repository and create a feature branch.
|
||||
2. Make your changes.
|
||||
3. Run ``make test-all`` to verify formatting, type checking, and tests
|
||||
pass.
|
||||
4. Add an entry to the "Upcoming version" section in
|
||||
``docs/source/changelog.rst`` under the appropriate category
|
||||
(Added / Changed / Fixed), following
|
||||
`Keep a Changelog <https://keepachangelog.com/>`_ conventions.
|
||||
5. Submit a pull request.
|
||||
|
||||
|
||||
Development with Claude Code
|
||||
----------------------------
|
||||
|
||||
The repository includes a ``CLAUDE.md`` file at the project root. This file
|
||||
defines development conventions, style guidelines, and common commands for
|
||||
`Claude Code <https://claude.com/claude-code>`_. It is also a useful
|
||||
reference for human contributors since it captures the same rules enforced
|
||||
in CI.
|
||||
|
||||
The project also includes shared commands in ``.claude/commands/``.
|
||||
Any contributor with Claude Code installed can invoke them as slash commands.
|
||||
|
||||
``/update-mjwarp <commit-hash>``
|
||||
Update the ``mujoco-warp`` dependency to a specific commit. This edits
|
||||
``pyproject.toml``, runs ``uv lock``, and opens a PR in one step.
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
/update-mjwarp e28c6038cdf8a353b4146974e4cf37e74dda809a
|
||||
|
||||
``/commit-push-pr``
|
||||
Stage current changes, commit, push, and open a PR.
|
||||
|
||||
@@ -0,0 +1,184 @@
|
||||
.. _curriculum:
|
||||
|
||||
Curriculum
|
||||
==========
|
||||
|
||||
The curriculum manager adjusts training conditions based on policy
|
||||
performance. Training begins with an easier problem and difficulty
|
||||
increases as the policy demonstrates it can handle the current
|
||||
conditions. Common uses include advancing robots to harder terrain,
|
||||
widening command velocity ranges, and ramping reward penalty weights
|
||||
over the course of training.
|
||||
|
||||
Curriculum terms are called at each environment reset. Each term
|
||||
receives the environment and the set of resetting environment IDs,
|
||||
examines some performance signal, and applies changes to environment
|
||||
parameters directly.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.managers.curriculum_manager import CurriculumTermCfg
|
||||
|
||||
curriculum = {
|
||||
"terrain_levels": CurriculumTermCfg(
|
||||
func=mdp.terrain_levels_vel,
|
||||
params={"command_name": "twist"},
|
||||
),
|
||||
}
|
||||
|
||||
The return value of a curriculum function is logged under
|
||||
``Curriculum/<term_name>`` in the training metrics.
|
||||
|
||||
|
||||
Built-in curriculum functions
|
||||
------------------------------
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 24 76
|
||||
|
||||
* - Function
|
||||
- Description
|
||||
* - ``terrain_levels_vel``
|
||||
- Measures how far each robot traveled during the episode. Robots
|
||||
that covered enough distance move up one difficulty row in the
|
||||
terrain grid; those that fell short move down. See the terrain
|
||||
curriculum section below.
|
||||
* - ``commands_vel``
|
||||
- Widens velocity command ranges based on training step count.
|
||||
Each stage specifies a step threshold and the new ranges to
|
||||
apply once that threshold is exceeded.
|
||||
* - ``reward_curriculum``
|
||||
- Adjusts a reward term's weight and/or params according to
|
||||
training step thresholds. Replaces the older ``reward_weight``
|
||||
function and also supports modifying reward function parameters.
|
||||
* - ``termination_curriculum``
|
||||
- Adjusts a termination term's params according to training step
|
||||
thresholds. Useful for gradually tightening termination
|
||||
conditions (e.g. energy limits) as training progresses.
|
||||
|
||||
|
||||
Reward curriculum
|
||||
-----------------
|
||||
|
||||
``reward_curriculum`` schedules changes to a reward term's weight or
|
||||
keyword arguments as training progresses. Each stage specifies a
|
||||
``step`` threshold and an optional ``weight`` or ``params`` update.
|
||||
Stages are evaluated in order, and each one whose threshold has been
|
||||
reached is applied.
|
||||
|
||||
**Ramping a penalty weight**
|
||||
|
||||
A common pattern is to introduce a penalty term at low weight early in
|
||||
training and increase it once the policy has learned the basics:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.managers.curriculum_manager import CurriculumTermCfg
|
||||
|
||||
curriculum = {
|
||||
"joint_vel_hinge_weight": CurriculumTermCfg(
|
||||
func=mdp.reward_curriculum,
|
||||
params={
|
||||
"reward_name": "joint_vel_hinge",
|
||||
"stages": [
|
||||
{"step": 0, "weight": -0.01},
|
||||
{"step": 12000, "weight": -0.1},
|
||||
{"step": 24000, "weight": -1.0},
|
||||
],
|
||||
},
|
||||
),
|
||||
}
|
||||
|
||||
**Adjusting reward parameters**
|
||||
|
||||
You can also change the parameters passed to the reward function. For
|
||||
example, tightening a tracking tolerance as training progresses:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
curriculum = {
|
||||
"track_lin_vel_tighten": CurriculumTermCfg(
|
||||
func=mdp.reward_curriculum,
|
||||
params={
|
||||
"reward_name": "track_linear_velocity",
|
||||
"stages": [
|
||||
{"step": 0, "params": {"std": 0.5}},
|
||||
{"step": 20000, "params": {"std": 0.3}},
|
||||
{"step": 50000, "params": {"std": 0.1}},
|
||||
],
|
||||
},
|
||||
),
|
||||
}
|
||||
|
||||
**Combining weight and params**
|
||||
|
||||
A single stage can update both weight and params at once:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
{"step": 24000, "weight": -1.0, "params": {"max_vel": 1.0}}
|
||||
|
||||
|
||||
Termination curriculum
|
||||
----------------------
|
||||
|
||||
``termination_curriculum`` schedules changes to a termination term's
|
||||
parameters as training progresses. This is useful for gradually
|
||||
tightening termination conditions once the policy has learned basic
|
||||
behaviors.
|
||||
|
||||
**Tightening an energy limit**
|
||||
|
||||
Start with a permissive energy threshold and reduce it over training:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.managers.curriculum_manager import CurriculumTermCfg
|
||||
|
||||
curriculum = {
|
||||
"energy_threshold": CurriculumTermCfg(
|
||||
func=mdp.termination_curriculum,
|
||||
params={
|
||||
"termination_name": "energy",
|
||||
"stages": [
|
||||
{"step": 12000, "params": {"threshold": 1000.0}},
|
||||
{"step": 24000, "params": {"threshold": 700.0}},
|
||||
{"step": 36000, "params": {"threshold": 400.0}},
|
||||
],
|
||||
},
|
||||
),
|
||||
}
|
||||
|
||||
The ``time_out`` field on ``TerminationTermCfg`` can also be toggled
|
||||
via stages if needed, though this is uncommon in practice.
|
||||
|
||||
|
||||
Terrain curriculum
|
||||
------------------
|
||||
|
||||
The terrain grid used with procedural terrain is a
|
||||
``num_rows x num_cols`` matrix of patches. Columns represent terrain
|
||||
type variants; rows represent difficulty levels, with row 0 being the
|
||||
easiest and row ``num_rows - 1`` the hardest. When
|
||||
``TerrainGeneratorCfg.curriculum=True``, each column is assigned exactly
|
||||
one terrain type so that difficulty increases monotonically along rows.
|
||||
|
||||
At environment construction each environment is assigned a random
|
||||
starting row within ``[0, max_init_terrain_level]``. The
|
||||
``terrain_levels_vel`` curriculum term promotes or demotes environments
|
||||
on each reset based on distance traveled during the episode.
|
||||
Environments that reach the maximum level are randomly reassigned to
|
||||
any row, maintaining coverage across all difficulty levels. See
|
||||
:ref:`terrain` for details on configuring the terrain grid itself.
|
||||
|
||||
|
||||
Writing custom curriculum functions
|
||||
------------------------------------
|
||||
|
||||
A curriculum function accepts ``env`` and ``env_ids``, applies
|
||||
parameter changes, and returns a value to log (a scalar tensor, a dict
|
||||
of tensors, or ``None``). A typical implementation reads a performance
|
||||
metric, decides whether to increase or decrease difficulty, mutates the
|
||||
relevant configuration in place, and returns the current difficulty
|
||||
level. See :ref:`env-config-term-pattern` for the general pattern.
|
||||
@@ -0,0 +1,50 @@
|
||||
.. _export-scene:
|
||||
|
||||
Export Scene
|
||||
============
|
||||
|
||||
The ``export-scene`` script writes a complete scene (XML and mesh assets) to a
|
||||
directory for inspection, sharing, or loading in standalone MuJoCo.
|
||||
|
||||
Quick start
|
||||
-----------
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
# Export a built-in entity by alias.
|
||||
uv run export-scene g1 --output-dir /tmp/g1
|
||||
|
||||
# Export a registered task scene.
|
||||
uv run export-scene Mjlab-Velocity-Flat-Unitree-Go1 --output-dir /tmp/task
|
||||
|
||||
# Export as a zip archive.
|
||||
uv run export-scene yam --output-dir /tmp/yam --zip True
|
||||
|
||||
# Export a custom entity via import path.
|
||||
uv run export-scene my_pkg.robots:get_my_robot_cfg --output-dir /tmp/custom
|
||||
|
||||
The output directory contains a ``scene.xml`` and an ``assets/`` subdirectory
|
||||
with all referenced mesh files. The XML can be loaded directly with
|
||||
``mujoco.MjModel.from_xml_path()`` or dropped into the
|
||||
`simulate viewer <https://mujoco.readthedocs.io/en/stable/programming/samples.html#sasimulate>`_.
|
||||
|
||||
Target resolution
|
||||
-----------------
|
||||
|
||||
The positional ``target`` argument is resolved in order:
|
||||
|
||||
1. **Task ID**: checked against the task registry (``import mjlab.tasks``).
|
||||
2. **Entity alias**: one of the built-in shorthands (``g1``, ``go1``, ``yam``).
|
||||
3. **Import path**: a ``module:attribute`` string pointing to any callable
|
||||
that returns an ``EntityCfg``.
|
||||
|
||||
If none match, the script prints available task IDs and aliases.
|
||||
|
||||
Options
|
||||
-------
|
||||
|
||||
``--output-dir DIR`` *(default: "export")*
|
||||
Destination directory. Cleaned before each export to prevent stale assets.
|
||||
|
||||
``--zip True`` *(default: False)*
|
||||
Compress the output into a ``.zip`` archive and remove the directory.
|
||||
@@ -0,0 +1,145 @@
|
||||
.. _nan-guard:
|
||||
|
||||
NaN Guard
|
||||
=========
|
||||
|
||||
The NaN guard captures simulation states when NaN/Inf is detected, helping
|
||||
debug numerical instability issues.
|
||||
|
||||
|
||||
Quick start
|
||||
-----------
|
||||
|
||||
Enable the NaN guard with a single CLI flag:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
uv run train <task-name> --enable-nan-guard True
|
||||
|
||||
This automatically captures and saves simulation states when NaN/Inf is
|
||||
detected. You can also enable it programmatically:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.sim.sim import SimulationCfg
|
||||
from mjlab.utils.nan_guard import NanGuardCfg
|
||||
|
||||
cfg = SimulationCfg(
|
||||
nan_guard=NanGuardCfg(
|
||||
enabled=True,
|
||||
buffer_size=100,
|
||||
output_dir="/tmp/mjlab/nan_dumps",
|
||||
max_envs_to_dump=5,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
Configuration
|
||||
-------------
|
||||
|
||||
``enabled`` *(default: False)*
|
||||
Enable/disable NaN detection and dumping.
|
||||
|
||||
``buffer_size`` *(default: 100)*
|
||||
Number of recent simulation states to keep in the rolling buffer.
|
||||
|
||||
``output_dir`` *(default: "/tmp/mjlab/nan_dumps")*
|
||||
Directory where NaN dump files are saved.
|
||||
|
||||
``max_envs_to_dump`` *(default: 5)*
|
||||
Maximum number of NaN environments to dump to disk. All environments are
|
||||
tracked in the buffer, but only the first N are saved to reduce dump
|
||||
size.
|
||||
|
||||
|
||||
Behavior
|
||||
--------
|
||||
|
||||
- **Captures** simulation state before each step (``qpos``, ``qvel``,
|
||||
``act`` if the model has actuator activations, and ``mocap_pos``/``mocap_quat``
|
||||
if the model has mocap bodies)
|
||||
- **Detects** NaN/Inf in ``qpos``, ``qvel``, ``qacc``,
|
||||
``qacc_warmstart``, and ``sensordata`` after each step
|
||||
- **Dumps** the rolling buffer and model to disk on first detection
|
||||
- **Stops** after the first dump to avoid spam
|
||||
|
||||
When disabled, all operations are no-ops with negligible overhead.
|
||||
|
||||
|
||||
Output format
|
||||
-------------
|
||||
|
||||
Each NaN detection creates timestamped files plus latest symlinks:
|
||||
|
||||
- ``nan_dump_TIMESTAMP.npz``: compressed state buffer
|
||||
|
||||
- ``states_step_NNNNNN``: captured states per step
|
||||
(shape: ``[num_envs_dumped, state_size]``)
|
||||
- ``_metadata``: dict with ``num_envs_total``, ``nan_env_ids``,
|
||||
``dumped_env_ids``, etc.
|
||||
|
||||
- ``model_TIMESTAMP.mjb``: MuJoCo model in binary format
|
||||
- ``nan_dump_latest.npz``: symlink to most recent dump
|
||||
- ``model_latest.mjb``: symlink to most recent model
|
||||
|
||||
|
||||
Visualizing dumps
|
||||
-----------------
|
||||
|
||||
Use the interactive viewer to scrub through captured states:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
# View latest dump.
|
||||
uv run viz-nan /tmp/mjlab/nan_dumps/nan_dump_latest.npz
|
||||
|
||||
# View a specific dump.
|
||||
uv run viz-nan /tmp/mjlab/nan_dumps/nan_dump_20251014_123456.npz
|
||||
|
||||
|
||||
.. figure:: ../_static/content/nan_debug.gif
|
||||
:alt: NaN Debug Viewer
|
||||
|
||||
NaN debug viewer.
|
||||
|
||||
The viewer provides:
|
||||
|
||||
- Step slider to scrub through the buffer
|
||||
- Environment slider to compare different environments
|
||||
- Info panel showing which environments have NaN/Inf
|
||||
- 3D visualization of the robot and terrain at each state
|
||||
|
||||
|
||||
NaN detection termination
|
||||
-------------------------
|
||||
|
||||
While the NaN guard helps **debug** NaN issues by capturing states, you can
|
||||
also **prevent** training crashes using the ``nan_detection`` termination
|
||||
term. This marks NaN environments as terminated, allowing them to reset
|
||||
while training continues:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.envs.mdp.terminations import nan_detection
|
||||
from mjlab.managers.termination_manager import TerminationTermCfg
|
||||
|
||||
nan_term: TerminationTermCfg = field(
|
||||
default_factory=lambda: TerminationTermCfg(
|
||||
func=nan_detection,
|
||||
time_out=False,
|
||||
)
|
||||
)
|
||||
|
||||
Terminations are logged as ``Episode_Termination/nan_term`` in your metrics.
|
||||
|
||||
.. important::
|
||||
|
||||
``nan_detection`` is a band-aid, not a cure. If NaNs occur during your
|
||||
task objective (e.g., NaNs happen when grasping), the policy will never
|
||||
learn to complete the task since it resets before receiving rewards.
|
||||
Monitor your ``Episode_Termination/nan_term`` metrics carefully.
|
||||
|
||||
**When to use which:**
|
||||
|
||||
- ``nan_guard``: debug and understand why NaNs occur (always do this first)
|
||||
- ``nan_detection``: keep training stable while working on a permanent fix
|
||||
@@ -0,0 +1,507 @@
|
||||
.. _entity_data:
|
||||
|
||||
Entity Data
|
||||
===========
|
||||
|
||||
This page is the property reference for ``EntityData``. For an overview
|
||||
of how ``entity.data`` fits into the broader data access story, see
|
||||
:ref:`entity`.
|
||||
|
||||
All properties are PyTorch tensors backed by MuJoCo Warp's GPU buffers
|
||||
with no copy overhead. The first dimension is always ``num_envs``, the
|
||||
number of parallel simulation worlds.
|
||||
|
||||
.. warning::
|
||||
|
||||
Read properties reflect the state after ``sim.forward()`` is called.
|
||||
If you write simulation state and then read a derived property in the
|
||||
same event term, call ``sim.forward()`` between the write and the
|
||||
read. The environment step sequence already does this; the warning
|
||||
applies only when writing custom event terms that mix reads and
|
||||
writes. See the :ref:`FAQ <faq-sim-forward>` for a detailed
|
||||
explanation.
|
||||
|
||||
|
||||
Reference: root state
|
||||
---------------------
|
||||
|
||||
Root properties describe the position, orientation, and velocity of the
|
||||
entity's root body. Properties ending in ``_w`` are expressed in the world
|
||||
frame. Properties ending in ``_b`` are expressed in the entity's base frame.
|
||||
See :ref:`frame-conventions` for details.
|
||||
|
||||
Each entity has two root reference points: the **link origin** (the body
|
||||
frame origin defined in the MJCF) and the **center of mass (COM)**.
|
||||
Which one is relevant depends on the task.
|
||||
|
||||
.. admonition:: MuJoCo's mixed-frame ``qvel``
|
||||
|
||||
For floating-base entities, the free joint stores 6 DOFs in
|
||||
``qvel``. MuJoCo expresses the **linear** components in the
|
||||
**world frame** but the **angular** components in the **local body
|
||||
frame**. EntityData avoids this pitfall: all ``_w`` velocity
|
||||
properties are computed from ``cvel`` (see
|
||||
:ref:`cvel-section` below) and are fully world-frame. If you
|
||||
read ``env.sim.data.qvel`` directly, be aware of the mixed
|
||||
convention.
|
||||
|
||||
.. rubric:: Root link properties
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 35 20 15 30
|
||||
|
||||
* - Property
|
||||
- Shape
|
||||
- Frame
|
||||
- Description
|
||||
* - ``root_link_pose_w``
|
||||
- ``[num_envs, 7]``
|
||||
- world
|
||||
- Root link position (3) and quaternion (4) concatenated
|
||||
* - ``root_link_pos_w``
|
||||
- ``[num_envs, 3]``
|
||||
- world
|
||||
- Root link position
|
||||
* - ``root_link_quat_w``
|
||||
- ``[num_envs, 4]``
|
||||
- world
|
||||
- Root link orientation as quaternion (w, x, y, z)
|
||||
* - ``root_link_vel_w``
|
||||
- ``[num_envs, 6]``
|
||||
- world
|
||||
- Root link linear (3) and angular (3) velocity concatenated
|
||||
* - ``root_link_lin_vel_w``
|
||||
- ``[num_envs, 3]``
|
||||
- world
|
||||
- Root link linear velocity
|
||||
* - ``root_link_ang_vel_w``
|
||||
- ``[num_envs, 3]``
|
||||
- world
|
||||
- Root link angular velocity
|
||||
* - ``root_link_lin_vel_b``
|
||||
- ``[num_envs, 3]``
|
||||
- body
|
||||
- Root link linear velocity in base frame
|
||||
* - ``root_link_ang_vel_b``
|
||||
- ``[num_envs, 3]``
|
||||
- body
|
||||
- Root link angular velocity in base frame
|
||||
|
||||
.. rubric:: Root COM properties
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 35 20 15 30
|
||||
|
||||
* - Property
|
||||
- Shape
|
||||
- Frame
|
||||
- Description
|
||||
* - ``root_com_pose_w``
|
||||
- ``[num_envs, 7]``
|
||||
- world
|
||||
- Root COM position (3) and quaternion (4) concatenated
|
||||
* - ``root_com_pos_w``
|
||||
- ``[num_envs, 3]``
|
||||
- world
|
||||
- Root COM position
|
||||
* - ``root_com_quat_w``
|
||||
- ``[num_envs, 4]``
|
||||
- world
|
||||
- Root COM orientation as quaternion (w, x, y, z)
|
||||
* - ``root_com_vel_w``
|
||||
- ``[num_envs, 6]``
|
||||
- world
|
||||
- Root COM linear (3) and angular (3) velocity concatenated
|
||||
* - ``root_com_lin_vel_w``
|
||||
- ``[num_envs, 3]``
|
||||
- world
|
||||
- Root COM linear velocity
|
||||
* - ``root_com_ang_vel_w``
|
||||
- ``[num_envs, 3]``
|
||||
- world
|
||||
- Root COM angular velocity
|
||||
* - ``root_com_lin_vel_b``
|
||||
- ``[num_envs, 3]``
|
||||
- body
|
||||
- Root COM linear velocity in base frame
|
||||
* - ``root_com_ang_vel_b``
|
||||
- ``[num_envs, 3]``
|
||||
- body
|
||||
- Root COM angular velocity in base frame
|
||||
|
||||
.. rubric:: Derived root properties
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 35 20 15 30
|
||||
|
||||
* - Property
|
||||
- Shape
|
||||
- Frame
|
||||
- Description
|
||||
* - ``projected_gravity_b``
|
||||
- ``[num_envs, 3]``
|
||||
- body
|
||||
- Gravity vector (0, 0, -1) rotated into the base frame. Used to measure
|
||||
tilt: a perfectly upright robot reads ``[0, 0, -1]``.
|
||||
* - ``heading_w``
|
||||
- ``[num_envs]``
|
||||
- world
|
||||
- Heading angle (radians) of the root body's forward axis projected onto
|
||||
the XY plane.
|
||||
|
||||
|
||||
Reference: body state
|
||||
---------------------
|
||||
|
||||
Body properties give per-body kinematic state for all bodies belonging to the
|
||||
entity. The second dimension is ``num_bodies``, which counts all non-world
|
||||
bodies in the entity's kinematic tree.
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 35 25 15 25
|
||||
|
||||
* - Property
|
||||
- Shape
|
||||
- Frame
|
||||
- Description
|
||||
* - ``body_link_pose_w``
|
||||
- ``[num_envs, num_bodies, 7]``
|
||||
- world
|
||||
- Per-body link position (3) and quaternion (4)
|
||||
* - ``body_link_pos_w``
|
||||
- ``[num_envs, num_bodies, 3]``
|
||||
- world
|
||||
- Per-body link positions
|
||||
* - ``body_link_quat_w``
|
||||
- ``[num_envs, num_bodies, 4]``
|
||||
- world
|
||||
- Per-body link orientations
|
||||
* - ``body_link_vel_w``
|
||||
- ``[num_envs, num_bodies, 6]``
|
||||
- world
|
||||
- Per-body link linear (3) and angular (3) velocity
|
||||
* - ``body_link_lin_vel_w``
|
||||
- ``[num_envs, num_bodies, 3]``
|
||||
- world
|
||||
- Per-body link linear velocities
|
||||
* - ``body_link_ang_vel_w``
|
||||
- ``[num_envs, num_bodies, 3]``
|
||||
- world
|
||||
- Per-body link angular velocities
|
||||
* - ``body_com_pose_w``
|
||||
- ``[num_envs, num_bodies, 7]``
|
||||
- world
|
||||
- Per-body COM position (3) and quaternion (4)
|
||||
* - ``body_com_pos_w``
|
||||
- ``[num_envs, num_bodies, 3]``
|
||||
- world
|
||||
- Per-body COM positions
|
||||
* - ``body_com_quat_w``
|
||||
- ``[num_envs, num_bodies, 4]``
|
||||
- world
|
||||
- Per-body COM orientations
|
||||
* - ``body_com_vel_w``
|
||||
- ``[num_envs, num_bodies, 6]``
|
||||
- world
|
||||
- Per-body COM linear (3) and angular (3) velocity
|
||||
* - ``body_com_lin_vel_w``
|
||||
- ``[num_envs, num_bodies, 3]``
|
||||
- world
|
||||
- Per-body COM linear velocities
|
||||
* - ``body_com_ang_vel_w``
|
||||
- ``[num_envs, num_bodies, 3]``
|
||||
- world
|
||||
- Per-body COM angular velocities
|
||||
* - ``body_external_wrench``
|
||||
- ``[num_envs, num_bodies, 6]``
|
||||
- world
|
||||
- External force (3) and torque (3) applied to each body
|
||||
* - ``body_external_force``
|
||||
- ``[num_envs, num_bodies, 3]``
|
||||
- world
|
||||
- External forces applied to each body
|
||||
* - ``body_external_torque``
|
||||
- ``[num_envs, num_bodies, 3]``
|
||||
- world
|
||||
- External torques applied to each body
|
||||
|
||||
|
||||
Reference: joint state
|
||||
----------------------
|
||||
|
||||
Joint properties cover 1-DOF revolute and prismatic joints. The free joint
|
||||
(root floating-base DOF) is excluded; use root state properties for that.
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 35 25 40
|
||||
|
||||
* - Property
|
||||
- Shape
|
||||
- Description
|
||||
* - ``joint_pos``
|
||||
- ``[num_envs, num_joints]``
|
||||
- Joint positions in radians (revolute) or metres (prismatic)
|
||||
* - ``joint_pos_biased``
|
||||
- ``[num_envs, num_joints]``
|
||||
- Joint positions with encoder bias added. Used when simulating
|
||||
encoder calibration errors via domain randomization.
|
||||
* - ``joint_vel``
|
||||
- ``[num_envs, num_joints]``
|
||||
- Joint velocities in rad/s or m/s
|
||||
* - ``joint_acc``
|
||||
- ``[num_envs, num_joints]``
|
||||
- Joint accelerations in rad/s² or m/s²
|
||||
* - ``actuator_force``
|
||||
- ``[num_envs, num_actuators]``
|
||||
- Scalar actuator output in actuation space (per actuator). This is
|
||||
the force before projection through the transmission Jacobian. For
|
||||
actuator forces in joint space, use ``qfrc_actuator`` instead.
|
||||
|
||||
|
||||
.. _generalized-forces:
|
||||
|
||||
Reference: generalized forces
|
||||
-----------------------------
|
||||
|
||||
These properties expose selected components of MuJoCo's generalized
|
||||
force decomposition, sliced to this entity's articulated joint DOFs.
|
||||
Free joint DOFs are excluded. All shapes are ``[num_envs, nv]`` where
|
||||
``nv`` is the number of articulated DOFs belonging to this entity.
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 30 70
|
||||
|
||||
* - Property
|
||||
- Description
|
||||
* - ``qfrc_actuator``
|
||||
- Forces produced by all actuators, mapped into joint space. For
|
||||
motors this is the commanded torque times the gear ratio. For
|
||||
position and velocity actuators this is the force computed by
|
||||
the internal PD law. When ``actuatorgravcomp`` is enabled on a
|
||||
joint, the gravity compensation force is included here.
|
||||
* - ``qfrc_external``
|
||||
- Forces on joints due to Cartesian wrenches applied to bodies
|
||||
via ``xfrc_applied``. This is the :math:`J^\top F` mapping.
|
||||
MuJoCo does not store this term separately; the property
|
||||
recovers it from other force components after ``forward()``.
|
||||
|
||||
Reference: geom and site state
|
||||
-------------------------------
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 35 25 40
|
||||
|
||||
* - Property
|
||||
- Shape
|
||||
- Description
|
||||
* - ``geom_pose_w``
|
||||
- ``[num_envs, num_geoms, 7]``
|
||||
- Per-geom position (3) and quaternion (4) in world frame
|
||||
* - ``geom_pos_w``
|
||||
- ``[num_envs, num_geoms, 3]``
|
||||
- Per-geom positions in world frame
|
||||
* - ``geom_quat_w``
|
||||
- ``[num_envs, num_geoms, 4]``
|
||||
- Per-geom orientations in world frame
|
||||
* - ``geom_vel_w``
|
||||
- ``[num_envs, num_geoms, 6]``
|
||||
- Per-geom linear (3) and angular (3) velocity in world frame
|
||||
* - ``geom_lin_vel_w``
|
||||
- ``[num_envs, num_geoms, 3]``
|
||||
- Per-geom linear velocities in world frame
|
||||
* - ``geom_ang_vel_w``
|
||||
- ``[num_envs, num_geoms, 3]``
|
||||
- Per-geom angular velocities in world frame
|
||||
* - ``site_pose_w``
|
||||
- ``[num_envs, num_sites, 7]``
|
||||
- Per-site position (3) and quaternion (4) in world frame
|
||||
* - ``site_pos_w``
|
||||
- ``[num_envs, num_sites, 3]``
|
||||
- Per-site positions in world frame
|
||||
* - ``site_quat_w``
|
||||
- ``[num_envs, num_sites, 4]``
|
||||
- Per-site orientations in world frame
|
||||
* - ``site_vel_w``
|
||||
- ``[num_envs, num_sites, 6]``
|
||||
- Per-site linear (3) and angular (3) velocity in world frame
|
||||
* - ``site_lin_vel_w``
|
||||
- ``[num_envs, num_sites, 3]``
|
||||
- Per-site linear velocities in world frame
|
||||
* - ``site_ang_vel_w``
|
||||
- ``[num_envs, num_sites, 3]``
|
||||
- Per-site angular velocities in world frame
|
||||
|
||||
|
||||
Reference: tendon state
|
||||
-----------------------
|
||||
|
||||
Tendon properties are only populated for entities that have tendon-driven
|
||||
actuators.
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 35 25 40
|
||||
|
||||
* - Property
|
||||
- Shape
|
||||
- Description
|
||||
* - ``tendon_len``
|
||||
- ``[num_envs, num_tendons]``
|
||||
- Tendon lengths
|
||||
* - ``tendon_vel``
|
||||
- ``[num_envs, num_tendons]``
|
||||
- Tendon velocities
|
||||
|
||||
|
||||
.. _frame-conventions:
|
||||
|
||||
Frame conventions
|
||||
-----------------
|
||||
|
||||
Property names encode their reference frame with a suffix.
|
||||
|
||||
``_w`` (world frame)
|
||||
A fixed global frame. The origin is typically at the scene origin and
|
||||
its axes are constant throughout the episode. World-frame quantities are
|
||||
useful when you need absolute position, such as checking whether the
|
||||
robot has fallen below a height threshold.
|
||||
|
||||
``_b`` (body frame / base frame)
|
||||
The entity's root body frame. It translates and rotates with the robot.
|
||||
Most observation terms use body-frame quantities because they are
|
||||
invariant to the robot's heading direction. A velocity expressed in the
|
||||
body frame reads the same whether the robot faces north or south, which
|
||||
makes it easier for the policy to generalize.
|
||||
|
||||
``projected_gravity_b`` is a good example of why the frame suffix
|
||||
matters. It takes the world-frame gravity vector ``[0, 0, -1]`` and
|
||||
rotates it into the base frame. When the robot is upright the result is
|
||||
``[0, 0, -1]``; as the robot tilts, the x and y components grow,
|
||||
giving the policy a direct signal for orientation correction.
|
||||
|
||||
Quaternion convention
|
||||
^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
All quaternions use the ``(w, x, y, z)`` convention, matching MuJoCo.
|
||||
|
||||
Reduced state vs. derived quantities
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
EntityData properties fall into two categories that behave differently
|
||||
with respect to ``sim.forward()``:
|
||||
|
||||
**Reduced state.** ``joint_pos`` and ``joint_vel`` read directly from
|
||||
MuJoCo's ``qpos`` and ``qvel`` arrays. Write methods such as
|
||||
``write_joint_state_to_sim()`` modify these arrays directly, so reads
|
||||
are always current.
|
||||
|
||||
**Derived quantities.** All pose and velocity properties (``*_pose_w``,
|
||||
``*_vel_w``, ``*_vel_b``) are computed from MuJoCo's internal arrays
|
||||
(``xpos``, ``xquat``, ``cvel``, ``subtree_com``, etc.) which are only
|
||||
updated when ``sim.forward()`` runs. If you write to ``qpos``/``qvel``
|
||||
and then read a derived property without an intervening ``forward()``,
|
||||
the read will return stale values.
|
||||
|
||||
The environment step sequence calls ``forward()`` at the right time, so
|
||||
this only matters if you write custom event terms that both write and
|
||||
read in the same function. See the :ref:`FAQ <faq-sim-forward>` for
|
||||
details.
|
||||
|
||||
.. _cvel-section:
|
||||
|
||||
How velocity properties are computed from ``cvel``
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
MuJoCo does not store world-frame linear velocities directly. Instead,
|
||||
it stores a 6D spatial velocity per body called ``cvel`` (com-based
|
||||
velocity), laid out as ``(angular[3], linear[3])``. This vector is
|
||||
expressed in the **c-frame**: a frame centered at ``subtree_com`` (the
|
||||
center of mass of the body's kinematic subtree) and oriented like the
|
||||
world frame. MuJoCo uses this representation to improve numerical
|
||||
precision for mechanisms far from the world origin. See
|
||||
`c-frame variables <https://mujoco.readthedocs.io/en/stable/APIreference/APItypes.html#c-frame-variables>`_
|
||||
and Featherstone's
|
||||
`Spatial Algebra <http://royfeatherstone.org/spatial/>`_ for background.
|
||||
|
||||
To recover the world-frame linear velocity at an arbitrary point
|
||||
:math:`\mathbf{p}` on a rigid body, we apply the standard rigid-body
|
||||
velocity transfer formula. Let :math:`\boldsymbol{\omega}` and
|
||||
:math:`\mathbf{v}_c` denote the angular and linear components of
|
||||
``cvel``, and let :math:`\mathbf{c}` denote ``subtree_com``. Because
|
||||
the c-frame is world-aligned, :math:`\boldsymbol{\omega}` is already in
|
||||
the world frame. The linear velocity at :math:`\mathbf{p}` is:
|
||||
|
||||
.. math::
|
||||
|
||||
\mathbf{v}_p
|
||||
= \mathbf{v}_c
|
||||
- \boldsymbol{\omega} \times (\mathbf{c} - \mathbf{p})
|
||||
|
||||
EntityData applies this formula in ``compute_velocity_from_cvel()``:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
def compute_velocity_from_cvel(pos, subtree_com, cvel):
|
||||
lin_vel_c = cvel[..., 3:6]
|
||||
ang_vel_c = cvel[..., 0:3]
|
||||
offset = subtree_com - pos
|
||||
lin_vel_w = lin_vel_c - torch.cross(ang_vel_c, offset, dim=-1)
|
||||
ang_vel_w = ang_vel_c
|
||||
return torch.cat([lin_vel_w, ang_vel_w], dim=-1)
|
||||
|
||||
Every velocity property in EntityData (``root_link_vel_w``,
|
||||
``body_link_vel_w``, ``geom_vel_w``, ``site_vel_w``, and their COM
|
||||
variants) uses this function, substituting the appropriate point:
|
||||
|
||||
- **Link velocities** use ``xpos`` (body frame origin).
|
||||
- **COM velocities** use ``xipos`` (body center of mass).
|
||||
- **Geom/site velocities** use ``geom_xpos``/``site_xpos``, with
|
||||
``cvel`` looked up from the parent body.
|
||||
|
||||
|
||||
Default pose and relative quantities
|
||||
--------------------------------------
|
||||
|
||||
``entity.data.default_joint_pos`` holds the joint positions from the entity's
|
||||
initial-state configuration (the ``init_state.joint_pos`` field of
|
||||
``EntityCfg``). It has shape ``[num_envs, num_joints]`` and is replicated
|
||||
across all environments at initialization time.
|
||||
|
||||
The relative joint position is the deviation of the current joint position
|
||||
from this default:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
joint_pos_rel = joint_pos - default_joint_pos
|
||||
|
||||
This is what the ``joint_pos_rel`` observation function computes:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
def joint_pos_rel(env, asset_cfg):
|
||||
asset = env.scene[asset_cfg.name]
|
||||
jnt_ids = asset_cfg.joint_ids
|
||||
return (
|
||||
asset.data.joint_pos[:, jnt_ids]
|
||||
- asset.data.default_joint_pos[:, jnt_ids]
|
||||
)
|
||||
|
||||
Relative joint positions give the policy a compact representation of posture
|
||||
deviation. When the robot is at its default pose, every element is zero.
|
||||
|
||||
Similarly, ``default_joint_vel`` is used by the ``joint_vel_rel`` observation
|
||||
function. For most configurations the default velocity is zero, so
|
||||
``joint_vel_rel`` is identical to ``joint_vel``. The indirection exists to
|
||||
allow non-zero reference velocities in tasks such as motion imitation.
|
||||
|
||||
The ``use_default_offset=True`` option in joint position action configs uses
|
||||
``default_joint_pos`` as the zero point for the action space, so a network
|
||||
output of zero commands the robot to its default pose. This is the standard
|
||||
configuration for locomotion tasks.
|
||||
@@ -0,0 +1,340 @@
|
||||
.. _entity:
|
||||
|
||||
Entity
|
||||
======
|
||||
|
||||
An ``Entity`` represents a physical object in the simulation: a robot, a
|
||||
manipulated object, or a fixed fixture like a table. It is the central
|
||||
abstraction in mjlab's physics layer.
|
||||
|
||||
A single ``Entity`` class covers all variants (contrast Isaac Lab, which
|
||||
splits this across ``Articulation``, ``RigidObject``, and several other
|
||||
subclasses of ``AssetBase``). Two orthogonal boolean properties classify
|
||||
each instance:
|
||||
|
||||
**Base type.**
|
||||
A *fixed-base* entity is welded to the world and has no free joint. A
|
||||
*floating-base* entity has a free joint giving it 6-DOF movement.
|
||||
|
||||
**Articulation.**
|
||||
An *articulated* entity has internal joints (revolute, prismatic, etc.).
|
||||
A *non-articulated* entity has none beyond a possible free joint.
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 30 25 15 15 15
|
||||
|
||||
* - Type
|
||||
- Example
|
||||
- ``is_fixed_base``
|
||||
- ``is_articulated``
|
||||
- ``is_actuated``
|
||||
* - Fixed non-articulated
|
||||
- Table, wall
|
||||
- True
|
||||
- False
|
||||
- False
|
||||
* - Fixed articulated
|
||||
- Robot arm, door
|
||||
- True
|
||||
- True
|
||||
- True/False
|
||||
* - Floating non-articulated
|
||||
- Box, ball, mug
|
||||
- False
|
||||
- False
|
||||
- False
|
||||
* - Floating articulated
|
||||
- Humanoid, quadruped
|
||||
- False
|
||||
- True
|
||||
- True/False
|
||||
|
||||
.. note::
|
||||
|
||||
mjlab automatically wraps every fixed-base entity in a
|
||||
`mocap body <https://mujoco.readthedocs.io/en/stable/modeling.html#mocap-bodies>`_
|
||||
so that each parallel environment can place the entity at a different
|
||||
position. Without this wrapping, all fixed-base entities would be
|
||||
welded to the world origin. The wrapping is transparent, but
|
||||
**positioning only happens when a reset event runs**. You must
|
||||
include a reset event such as ``reset_root_state_uniform`` in your
|
||||
event config; without one, every fixed-base entity will remain at
|
||||
the origin. See the :ref:`FAQ <faq>` for a full example. Mocap
|
||||
entities can also be repositioned at runtime via
|
||||
``entity.write_mocap_pose_to_sim()``.
|
||||
|
||||
|
||||
Configuring an entity
|
||||
---------------------
|
||||
|
||||
Every entity is described by an ``EntityCfg``. Only ``spec_fn`` is
|
||||
required in practice; all other fields have sensible defaults. A passive
|
||||
floating object needs nothing more than:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.entity import EntityCfg
|
||||
|
||||
cube_cfg = EntityCfg(spec_fn=get_cube_spec)
|
||||
|
||||
An actuated robot uses more of the interface:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.entity import EntityCfg, EntityArticulationInfoCfg
|
||||
from mjlab.actuator import IdealPDActuatorCfg
|
||||
|
||||
robot_cfg = EntityCfg(
|
||||
spec_fn=get_spec,
|
||||
init_state=EntityCfg.InitialStateCfg(
|
||||
pos=(0.0, 0.0, 0.8),
|
||||
joint_pos={".*_hip_.*": 0.5, ".*": 0.0},
|
||||
),
|
||||
articulation=EntityArticulationInfoCfg(
|
||||
actuators=(
|
||||
IdealPDActuatorCfg(
|
||||
target_names_expr=(".*",),
|
||||
stiffness={".*": 50.0},
|
||||
damping={".*": 5.0},
|
||||
),
|
||||
),
|
||||
),
|
||||
collisions=(my_collision_cfg,),
|
||||
)
|
||||
|
||||
The following sections describe each field.
|
||||
|
||||
``spec_fn``
|
||||
^^^^^^^^^^^
|
||||
|
||||
A callable that returns an ``mujoco.MjSpec``. The scene calls it during
|
||||
composition, attaches the returned spec with a name prefix, and compiles
|
||||
everything into a shared ``MjModel``.
|
||||
|
||||
For simple cases a lambda suffices:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
spec_fn = lambda: mujoco.MjSpec.from_file("robot.xml")
|
||||
|
||||
For anything more involved, use a regular function. MuJoCo resolves mesh
|
||||
assets from disk automatically, so ``get_spec`` only needs to load the
|
||||
XML:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
def get_spec() -> mujoco.MjSpec:
|
||||
return mujoco.MjSpec.from_file(str(ROBOT_XML))
|
||||
|
||||
Because ``spec_fn`` is an arbitrary callable, you can perform any
|
||||
`MjSpec edits <https://mujoco.readthedocs.io/en/stable/python.html#spec>`_
|
||||
before returning: add bodies, change joint limits, swap materials,
|
||||
or build the entire model programmatically without an XML file at all.
|
||||
|
||||
``init_state``
|
||||
^^^^^^^^^^^^^^
|
||||
|
||||
Default root pose, root velocity, and joint positions/velocities. These
|
||||
values are stored as a MuJoCo keyframe and used by reset events to
|
||||
return the entity to its initial configuration.
|
||||
|
||||
``joint_pos`` and ``joint_vel`` are dicts mapping regex patterns to
|
||||
values. Patterns are matched against joint names in order, so later
|
||||
entries override earlier ones for any joint that matches both:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
init_state = EntityCfg.InitialStateCfg(
|
||||
pos=(0.0, 0.0, 0.8), # root position
|
||||
rot=(1.0, 0.0, 0.0, 0.0), # root quaternion (w, x, y, z)
|
||||
joint_pos={
|
||||
".*": 0.0, # all joints to zero
|
||||
".*_hip_.*": 0.5, # then override hips to 0.5
|
||||
},
|
||||
)
|
||||
|
||||
Set ``joint_pos=None`` to use an existing keyframe from the MJCF model
|
||||
instead of defining values here.
|
||||
|
||||
``articulation``
|
||||
^^^^^^^^^^^^^^^^
|
||||
|
||||
Actuator configuration. Only needed for entities that have actuated
|
||||
joints. Passive objects (boxes, tables, walls) can omit this field
|
||||
entirely. See :ref:`actuators` for details on actuator types.
|
||||
|
||||
``soft_joint_pos_limit_factor`` (default 1.0) shrinks the joint range
|
||||
used by soft-limit penalty rewards, so the policy is penalized before
|
||||
reaching the physical hard stop. This does not modify the actual joint
|
||||
limits in the MuJoCo model.
|
||||
|
||||
Spec editors
|
||||
^^^^^^^^^^^^
|
||||
|
||||
The remaining fields are optional tuples of spec editor configs that
|
||||
modify the ``MjSpec`` before compilation:
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 20 80
|
||||
|
||||
* - Field
|
||||
- Purpose
|
||||
* - ``collisions``
|
||||
- Set contact parameters (contype, conaffinity, friction) per geom.
|
||||
* - ``lights``
|
||||
- Add lights to specific bodies.
|
||||
* - ``cameras``
|
||||
- Add cameras to specific bodies.
|
||||
* - ``textures``
|
||||
- Add procedural textures (checker, gradient, etc.).
|
||||
* - ``materials``
|
||||
- Add materials and optionally assign them to geoms by regex.
|
||||
|
||||
Each editor accepts regex patterns to target specific elements. For
|
||||
example, a ``CollisionCfg`` with ``geom_names_expr=(".*_foot.*",)``
|
||||
sets contact parameters only on foot geoms. See the asset zoo
|
||||
(``mjlab.asset_zoo.robots``) for complete examples.
|
||||
|
||||
Per-world mesh variants
|
||||
^^^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
For scenes that need different mesh assets in different parallel worlds
|
||||
(for example, training a manipulation policy that generalizes across
|
||||
object shapes), use ``VariantEntityCfg`` instead of ``EntityCfg``. Each
|
||||
world is assigned a variant proportional to a configurable weight, and
|
||||
mesh-dependent compiled constants (collision bounds, body inertials,
|
||||
subtree mass) are stored as per-world arrays so domain randomization and
|
||||
viewers stay consistent. See :ref:`per_world_mesh`.
|
||||
|
||||
Subclassing Entity
|
||||
^^^^^^^^^^^^^^^^^^
|
||||
|
||||
``Entity`` and ``EntityCfg`` can be subclassed for specialized behavior.
|
||||
mjlab itself does this for terrain: ``TerrainEntity`` extends ``Entity``
|
||||
with procedural terrain generation and per-environment origin
|
||||
computation, and ``TerrainEntityCfg`` adds fields like
|
||||
``terrain_type``, ``env_spacing``, and ``terrain_generator``. The same
|
||||
pattern works for any domain-specific entity that needs logic beyond
|
||||
what ``EntityCfg`` and spec editors provide.
|
||||
|
||||
Finding elements
|
||||
^^^^^^^^^^^^^^^^
|
||||
|
||||
Entity provides ``find_*`` methods that accept regex patterns and return
|
||||
matched element indices and names:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
ids, names = entity.find_joints((".*_hip_.*", ".*_knee_.*"))
|
||||
ids, names = entity.find_geoms((".*foot.*",))
|
||||
ids, names = entity.find_bodies((".*",))
|
||||
|
||||
Available methods: ``find_bodies()``, ``find_joints()``,
|
||||
``find_geoms()``, ``find_sites()``, ``find_tendons()``.
|
||||
These are used internally during scene construction and manager
|
||||
initialization. In reward and observation terms, prefer
|
||||
``SceneEntityCfg`` with name patterns as described below.
|
||||
|
||||
|
||||
Reading runtime state
|
||||
---------------------
|
||||
|
||||
Once entities are added to a ``SceneCfg`` and the environment is
|
||||
constructed, their state is accessible through three interfaces at
|
||||
decreasing levels of abstraction.
|
||||
|
||||
EntityData
|
||||
^^^^^^^^^^
|
||||
|
||||
``entity.data`` is the primary interface for reward, observation, and
|
||||
termination functions. It exposes kinematic state (poses, velocities, accelerations), actuator forces,
|
||||
generalized forces, and derived body-frame quantities such as projected
|
||||
gravity, all as PyTorch tensors with
|
||||
shape ``(num_envs, ...)``. See :ref:`entity_data` for the full property
|
||||
reference.
|
||||
|
||||
``SceneEntityCfg`` selects which entity and which elements within it a
|
||||
term operates on. Regex patterns in ``joint_names``, ``body_names``,
|
||||
``site_names``, etc. are resolved to integer indices once at manager
|
||||
initialization, so there is no regex overhead at runtime:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.managers.scene_entity_config import SceneEntityCfg
|
||||
|
||||
def flat_orientation_l2(
|
||||
env,
|
||||
asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
|
||||
) -> torch.Tensor:
|
||||
"""Penalize non-flat base orientation using projected gravity."""
|
||||
asset = env.scene[asset_cfg.name]
|
||||
return torch.sum(
|
||||
torch.square(asset.data.projected_gravity_b[:, :2]), dim=1
|
||||
)
|
||||
|
||||
``SceneEntityCfg`` also supports regex element selection through
|
||||
``joint_names``, ``body_names``, ``site_names``, etc. The resolved
|
||||
integer indices (e.g., ``asset_cfg.joint_ids``) make the runtime read a
|
||||
single tensor slice with no regex overhead.
|
||||
|
||||
Sensors
|
||||
^^^^^^^
|
||||
|
||||
Sensors are configured on the **scene**, not on individual entities.
|
||||
A sensor can reference an entity element (e.g., a contact sensor on the
|
||||
robot's feet, an accelerometer attached to a body site), but it can also
|
||||
be independent of any entity. This is why sensors live in ``SceneCfg``
|
||||
rather than ``EntityCfg``.
|
||||
|
||||
At runtime, sensors are accessed by name through ``env.scene``, the same
|
||||
way entities are:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
def angular_momentum_penalty(env, sensor_name: str) -> torch.Tensor:
|
||||
sensor = env.scene[sensor_name]
|
||||
return torch.sum(torch.square(sensor.data), dim=-1)
|
||||
|
||||
Builtin sensors wrap MuJoCo sensor types (accelerometer, gyro, framepos,
|
||||
subtreeangmom, etc.). ``ContactSensor``, ``RayCastSensor``, and
|
||||
``CameraSensor`` provide higher-level abstractions for contact detection,
|
||||
terrain scanning, and RGB-D rendering. See :ref:`sensors` for details.
|
||||
|
||||
Raw simulation data
|
||||
^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
For anything not covered by ``EntityData`` or sensors, the underlying
|
||||
MuJoCo Warp arrays are accessible through ``env.sim.data`` and
|
||||
``env.sim.model``. These expose the full ``mjData`` and ``mjModel``
|
||||
fields as PyTorch tensors (zero-copy), indexed by global MuJoCo IDs
|
||||
rather than per-entity IDs:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# Global joint positions across all entities.
|
||||
qpos = env.sim.data.qpos # (num_envs, nq)
|
||||
|
||||
# All body positions.
|
||||
xpos = env.sim.data.xpos # (num_envs, nbody, 3)
|
||||
|
||||
# Model-level constants.
|
||||
body_mass = env.sim.model.body_mass # (nbody,)
|
||||
|
||||
This is useful for low-level operations or when you need quantities
|
||||
that span multiple entities.
|
||||
|
||||
.. note::
|
||||
|
||||
The main limitation of raw sim data is that you must manage global
|
||||
MuJoCo indices yourself. In the future, we plan to support MuJoCo's
|
||||
`bind <https://mujoco.readthedocs.io/en/latest/python.html#relationship-to-pymjcf-and-bind>`_
|
||||
functionality, which will allow binding spec elements directly to
|
||||
their corresponding data views without manual index bookkeeping.
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
|
||||
entity_data
|
||||
per_world_mesh
|
||||
@@ -0,0 +1,200 @@
|
||||
.. _per_world_mesh:
|
||||
|
||||
Mesh Variants
|
||||
=============
|
||||
|
||||
Mesh variants let a single batched simulation run with different mesh
|
||||
assets in different parallel worlds. World 0 may simulate a cube, world
|
||||
1 a sphere, and world 2 a bowl, all sharing the same compiled scene
|
||||
and the same kinematic structure. The result is a heterogeneous batch
|
||||
in which the mesh and its derived constants vary across worlds while
|
||||
everything else (the body tree, the joint structure, the contact and
|
||||
solver setup) is fixed.
|
||||
|
||||
Mesh variants are configured at the entity level through
|
||||
``VariantEntityCfg`` and ``VariantCfg``. Once configured,
|
||||
domain randomization, the native viewer, the offscreen renderer, and
|
||||
the Viser viewer all pick up the variant assignment automatically.
|
||||
|
||||
|
||||
How it works
|
||||
------------
|
||||
|
||||
A standard ``EntityCfg`` provides a single ``spec_fn`` that returns one
|
||||
``MjSpec``. A ``VariantEntityCfg`` provides a dictionary of named
|
||||
variants, each with its own ``spec_fn`` and a weight controlling the
|
||||
proportion of worlds that use it.
|
||||
|
||||
**All variants must declare the same kinematic structure.** The batched
|
||||
simulator assumes a single topology across worlds; per-world variation
|
||||
is confined to mesh assets and the constants derived from them. mjlab
|
||||
uses the first variant's body tree as the template and copies mesh
|
||||
assets and explicit body inertials from the others. Geom-level
|
||||
properties on later variants such as ``rgba``, friction, and material
|
||||
assignments are not propagated; control per-world appearance through
|
||||
domain randomization on ``geom_rgba`` or ``mat_rgba``. The structural
|
||||
check is enforced at construction time and raises a ``ValueError``
|
||||
describing the first mismatch. Variants must also be floating-base
|
||||
(declare a free joint on the root body); fixed-base variants are
|
||||
rejected.
|
||||
|
||||
A minimal two-variant config:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import mujoco
|
||||
|
||||
from mjlab.entity import EntityCfg, VariantCfg, VariantEntityCfg
|
||||
|
||||
def make_sphere_spec() -> mujoco.MjSpec:
|
||||
spec = mujoco.MjSpec()
|
||||
mesh = spec.add_mesh(name="visual")
|
||||
mesh.make_sphere(subdivision=3)
|
||||
mesh.scale[:] = (0.05,) * 3
|
||||
body = spec.worldbody.add_body(name="prop")
|
||||
body.add_freejoint()
|
||||
body.add_geom(type=mujoco.mjtGeom.mjGEOM_MESH, meshname="visual")
|
||||
return spec
|
||||
|
||||
# ``make_cone_spec`` follows the same shape with
|
||||
# ``mesh.make_cone(nedge=16, radius=0.04)`` in place of the sphere call.
|
||||
|
||||
object_cfg = VariantEntityCfg(
|
||||
variants={
|
||||
"sphere": VariantCfg(spec_fn=make_sphere_spec, weight=1.0),
|
||||
"cone": VariantCfg(spec_fn=make_cone_spec, weight=2.0),
|
||||
},
|
||||
init_state=EntityCfg.InitialStateCfg(pos=(0.0, 0.0, 0.2)),
|
||||
)
|
||||
|
||||
During scene construction mjlab merges the per-variant specs into a
|
||||
single ``MjSpec`` whose mesh slots are padded to the maximum count any
|
||||
variant uses, then writes a per-world ``geom_dataid`` table that
|
||||
selects the right mesh for each world. In the merged scene
|
||||
``geom_dataid`` is no longer a flat ``(ngeom,)`` vector but a
|
||||
``(num_envs, ngeom)`` table whose rows differ by variant. A value of
|
||||
``-1`` marks a disabled mesh slot, used for variants with fewer mesh
|
||||
geoms than the maximum.
|
||||
|
||||
Mesh choice is entangled with several other compiled-model constants:
|
||||
geom collision bounds, geom local frames, body inertials, subtree mass,
|
||||
and inverse weights. mjlab compiles each unique row of the
|
||||
``geom_dataid`` table on the host and copies the relevant compiled
|
||||
fields into per-world arrays on the GPU, so each world's compiled
|
||||
constants stay consistent with that world's mesh selection. The full
|
||||
list of fields handled this way is in
|
||||
``mjlab.sim.mesh_variants.VARIANT_DEPENDENT_FIELDS``.
|
||||
|
||||
|
||||
World assignment
|
||||
----------------
|
||||
|
||||
mjlab assigns variants to worlds proportionally by weight using the
|
||||
`largest remainder method
|
||||
<https://en.wikipedia.org/wiki/Largest_remainder_method>`_. Each
|
||||
variant's quota is ``q_i = (w_i / sum(w)) * num_envs``; each variant
|
||||
first receives ``floor(q_i)`` worlds, and the remaining
|
||||
``num_envs - sum(floors)`` worlds go to the variants with the largest
|
||||
fractional remainders, with ties broken by declaration order. For
|
||||
``num_envs = 10`` and weights ``(1.0, 2.0, 1.0)`` this gives
|
||||
``(3, 5, 2)`` worlds per variant. Weights are normalized internally,
|
||||
so ``(1, 2, 1)`` and ``(0.25, 0.5, 0.25)`` produce identical
|
||||
assignments. A weight of zero is allowed and produces zero worlds for
|
||||
that variant; at least one variant must have a positive weight.
|
||||
|
||||
Variant assignment is fixed at simulation initialization and does not
|
||||
resample on episode reset. The intended use is heterogeneous training
|
||||
across the batch, not per-episode mesh randomization. To inspect the
|
||||
assignment from user code, read ``env.sim.world_to_variant``:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
>>> env.sim.world_to_variant["object"]
|
||||
tensor([0, 0, 0, 1, 1, 1, 1, 1, 1, 1])
|
||||
|
||||
The mapping is keyed by entity name (without trailing slash) and
|
||||
returns a ``(num_envs,)`` tensor of variant indices in the order
|
||||
variants were declared in ``VariantEntityCfg.variants``. The dict is
|
||||
empty for non-variant scenes.
|
||||
|
||||
|
||||
Domain randomization
|
||||
--------------------
|
||||
|
||||
Domain randomization on variant scenes preserves per-variant baselines
|
||||
automatically. When the simulation initializes, mjlab snapshots the
|
||||
variant-dependent fields (``body_mass``, ``body_inertia``,
|
||||
``geom_size``, and others listed in ``VARIANT_DEPENDENT_FIELDS``) as
|
||||
``(num_envs, ...)`` tensors and registers them in
|
||||
``sim.per_world_default_fields``. Domain randomization operations that
|
||||
read defaults (scale, additive offsets) detect this registration and
|
||||
index the per-world default array by environment, so a 10% mass scale
|
||||
applied across a batch containing a 100 g sphere variant and a 1 kg
|
||||
cube variant produces 10% perturbations around each variant's own
|
||||
mass, not 10% of a shared template mass. Fields that are not
|
||||
variant-dependent (``geom_friction``, ``dof_armature``,
|
||||
``dof_damping``, and so on) behave identically on variant and
|
||||
non-variant scenes.
|
||||
|
||||
For inertial randomization the recommended path is
|
||||
``dr.pseudo_inertia``, which jointly randomizes mass, COM offset,
|
||||
principal moments of inertia, and principal frame orientation through
|
||||
the pseudo-inertia matrix factorization of `Rucker and Wensing (2022)
|
||||
<https://par.nsf.gov/servlets/purl/10347458>`_. It is exact for any
|
||||
perturbation magnitude and remains physically consistent across
|
||||
variants of different scale. ``dr.body_mass`` modifies ``body_mass``
|
||||
without touching the inertia tensor and emits a ``UserWarning`` when
|
||||
called; it is appropriate only for modeling a point mass added at the
|
||||
COM, not for density-like randomization. The distinction matters more
|
||||
on variant scenes than on single-asset scenes because variants often
|
||||
differ in mass by an order of magnitude.
|
||||
|
||||
|
||||
Viewers
|
||||
-------
|
||||
|
||||
The native viewer, offscreen renderer, and Viser viewer all sync the
|
||||
selected environment's per-world fields into the host ``MjModel``
|
||||
before rendering, so the rendered geometry matches the variant
|
||||
assigned to the viewed environment. Switching environments in the
|
||||
native viewer (the ``,`` and ``.`` keys) updates the displayed mesh
|
||||
accordingly.
|
||||
|
||||
Viser bakes mesh data into batched handles and cannot rely on a live
|
||||
view of ``geom_dataid``. It groups worlds by visual fingerprint (mesh
|
||||
selection, local geom frames, baked appearance) and builds one batched
|
||||
handle per group, with each environment assigned to its handle. A
|
||||
scene with N variants typically produces up to N handles per body.
|
||||
Convex hull visualization is computed per variant from the variant's
|
||||
mesh vertices.
|
||||
|
||||
|
||||
Performance considerations
|
||||
--------------------------
|
||||
|
||||
Mesh variants do not add per-step overhead in the GPU kernels.
|
||||
Variant-dependent fields are stored as per-world arrays accessed by
|
||||
world index in the existing kernels, with no branching or dispatch
|
||||
on variant.
|
||||
|
||||
Initialization is the main consideration. mjlab compiles each unique
|
||||
row of the ``geom_dataid`` table by taking a fresh ``MjSpec.copy()``,
|
||||
editing the mesh selection and (if applicable) the explicit body
|
||||
inertials, and calling ``spec.compile()``. This work scales with the
|
||||
number of unique variant combinations rather than with ``num_envs``.
|
||||
For a scene with one variant entity declaring k variants, this is k
|
||||
host compiles regardless of how many worlds use each variant. With
|
||||
multiple variant entities the unique-row count is bounded by the
|
||||
product of their variant counts in the worst case, so a scene with
|
||||
two variant entities of 5 variants each could trigger up to 25 host
|
||||
compiles at init.
|
||||
|
||||
``MjSpec.copy()`` and ``spec.compile()`` are non-trivial operations,
|
||||
and their cost grows with scene size. For a scene with many variant
|
||||
entities or many variants per entity, the cumulative initialization
|
||||
cost can be measured in seconds. This cost is paid once at startup
|
||||
and does not affect training throughput.
|
||||
|
||||
The merged spec contains every variant's mesh assets simultaneously.
|
||||
Memory footprint at scene-build time scales with the total number of
|
||||
mesh vertices and faces across all declared variants.
|
||||
@@ -0,0 +1,484 @@
|
||||
.. _environment_config:
|
||||
|
||||
Environment Configuration
|
||||
=========================
|
||||
|
||||
A single ``ManagerBasedRlEnvCfg`` dataclass fully specifies an mjlab
|
||||
environment: the physical world, the agent's interface to it, and the
|
||||
MDP defined on top. Because everything lives in one flat
|
||||
object, an environment can be inspected, copied, and modified without
|
||||
navigating a class hierarchy.
|
||||
|
||||
For a broad orientation to mjlab before reading this page, start with
|
||||
:ref:`architecture_overview`.
|
||||
|
||||
|
||||
.. _env-config-skeleton:
|
||||
|
||||
Annotated skeleton
|
||||
------------------
|
||||
|
||||
The complete set of fields on ``ManagerBasedRlEnvCfg`` is shown below with
|
||||
inline comments. The fields marked with ``...`` must be provided; all others
|
||||
have defaults.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from mjlab.envs import ManagerBasedRlEnvCfg
|
||||
from mjlab.managers.action_manager import ActionTermCfg
|
||||
from mjlab.managers.command_manager import CommandTermCfg
|
||||
from mjlab.managers.curriculum_manager import CurriculumTermCfg
|
||||
from mjlab.managers.event_manager import EventTermCfg
|
||||
from mjlab.managers.metrics_manager import MetricsTermCfg
|
||||
from mjlab.managers.observation_manager import ObservationGroupCfg
|
||||
from mjlab.managers.reward_manager import RewardTermCfg
|
||||
from mjlab.managers.termination_manager import TerminationTermCfg
|
||||
from mjlab.scene.scene import SceneCfg
|
||||
from mjlab.sim.sim import SimulationCfg
|
||||
from mjlab.viewer.viewer_config import ViewerConfig
|
||||
|
||||
|
||||
@dataclass
|
||||
class MyEnvCfg(ManagerBasedRlEnvCfg):
|
||||
|
||||
# --- Physics ---
|
||||
|
||||
decimation: int = 4
|
||||
# Number of physics steps per policy step.
|
||||
# Environment step duration = sim.mujoco.timestep * decimation.
|
||||
|
||||
sim: SimulationCfg = field(default_factory=SimulationCfg)
|
||||
# Physics parameters: timestep, integrator, solver, contact settings.
|
||||
# Default timestep is 0.002 s (500 Hz). Override with MujocoCfg.
|
||||
|
||||
scene: SceneCfg = ...
|
||||
# Terrain, entities, and sensors. Also sets num_envs.
|
||||
# Required; there is no default.
|
||||
|
||||
# --- Episode ---
|
||||
|
||||
episode_length_s: float = 20.0
|
||||
# Episode duration in seconds.
|
||||
# Steps = ceil(episode_length_s / (sim.mujoco.timestep * decimation)).
|
||||
|
||||
is_finite_horizon: bool = False
|
||||
# False (default): time limit is an artificial cutoff. The agent
|
||||
# receives a truncated signal and bootstraps value beyond the limit.
|
||||
# True: time limit defines the task boundary. The agent receives a
|
||||
# terminal done signal with no future value beyond it.
|
||||
|
||||
scale_rewards_by_dt: bool = True
|
||||
# When True (default), each reward term is multiplied by step_dt so
|
||||
# that cumulative episodic sums are invariant to simulation frequency.
|
||||
# Set to False for algorithms that expect unscaled reward signals.
|
||||
|
||||
# --- Managers ---
|
||||
|
||||
observations: dict[str, ObservationGroupCfg] = field(default_factory=dict)
|
||||
# Observation groups. Each key is a group name (e.g. "actor", "critic").
|
||||
# Groups can differ in noise, history, delay, and concatenation.
|
||||
|
||||
actions: dict[str, ActionTermCfg] = field(default_factory=dict)
|
||||
# Action terms. Each term controls one slice of the policy output
|
||||
# and routes it to a specific entity's actuators.
|
||||
|
||||
rewards: dict[str, RewardTermCfg] = field(default_factory=dict)
|
||||
# Reward terms. The manager computes a weighted sum each step.
|
||||
|
||||
terminations: dict[str, TerminationTermCfg] = field(default_factory=dict)
|
||||
# Termination conditions. If empty, episodes never terminate early.
|
||||
# Add a time_out term to enforce the episode length limit.
|
||||
|
||||
events: dict[str, EventTermCfg] = field(
|
||||
default_factory=lambda: {
|
||||
"reset_scene_to_default": EventTermCfg(
|
||||
func=reset_scene_to_default,
|
||||
mode="reset",
|
||||
)
|
||||
}
|
||||
)
|
||||
# Event terms for domain randomization and state resets.
|
||||
# The default includes reset_scene_to_default, which resets all
|
||||
# entities to their initial pose each episode. Override this dict
|
||||
# to replace or extend the default reset behavior.
|
||||
|
||||
commands: dict[str, CommandTermCfg] = field(default_factory=dict)
|
||||
# Command generators (e.g. velocity targets for locomotion).
|
||||
# Commands are resampled at configurable intervals and on reset.
|
||||
|
||||
curriculum: dict[str, CurriculumTermCfg] = field(default_factory=dict)
|
||||
# Curriculum terms that adjust training conditions based on performance.
|
||||
|
||||
metrics: dict[str, MetricsTermCfg] = field(default_factory=dict)
|
||||
# Custom metrics logged as episode averages alongside reward terms.
|
||||
|
||||
# --- Misc ---
|
||||
|
||||
seed: int | None = None
|
||||
# Random seed for reproducibility. If None, a random seed is chosen
|
||||
# and stored back into this field after initialization.
|
||||
|
||||
viewer: ViewerConfig = field(default_factory=ViewerConfig)
|
||||
# Camera position, resolution, and tracking target for rendering.
|
||||
|
||||
|
||||
.. _env-config-term-pattern:
|
||||
|
||||
Term configuration pattern
|
||||
--------------------------
|
||||
|
||||
All manager dictionaries follow the same pattern. Each entry maps a string
|
||||
name to a term configuration object. The configuration always carries at
|
||||
minimum a ``func`` field pointing to the callable that implements the term,
|
||||
and a ``params`` dict of extra keyword arguments forwarded to that callable.
|
||||
|
||||
The manager calls ``func(env, **params)`` each step (or ``term(env, **params)``
|
||||
when ``func`` is a class that has been instantiated). Term names are arbitrary;
|
||||
they appear in training logs and are used only for identification.
|
||||
|
||||
.. rubric:: Reward terms
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.envs import mdp
|
||||
from mjlab.managers.reward_manager import RewardTermCfg
|
||||
from mjlab.managers.scene_entity_config import SceneEntityCfg
|
||||
|
||||
rewards = {
|
||||
"alive": RewardTermCfg(
|
||||
func=mdp.is_alive,
|
||||
weight=1.0,
|
||||
),
|
||||
"joint_torques": RewardTermCfg(
|
||||
func=mdp.joint_torques_l2,
|
||||
weight=-1e-4,
|
||||
params={"asset_cfg": SceneEntityCfg("robot")},
|
||||
),
|
||||
"action_rate": RewardTermCfg(
|
||||
func=mdp.action_rate_l2,
|
||||
weight=-0.1,
|
||||
),
|
||||
}
|
||||
|
||||
``weight`` scales the function's output before it is summed into the total
|
||||
reward. Negative weights produce penalties.
|
||||
|
||||
``params`` maps to keyword arguments of the function. For example,
|
||||
``mdp.joint_torques_l2(env, asset_cfg=...)`` receives ``asset_cfg`` from the
|
||||
``params`` dict. Any argument not listed in ``params`` must have a default
|
||||
value in the function signature.
|
||||
|
||||
.. rubric:: Termination terms
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.envs import mdp
|
||||
from mjlab.managers.termination_manager import TerminationTermCfg
|
||||
|
||||
terminations = {
|
||||
"time_out": TerminationTermCfg(
|
||||
func=mdp.time_out,
|
||||
time_out=True, # marks this as a truncation, not a failure
|
||||
),
|
||||
"fell_over": TerminationTermCfg(
|
||||
func=mdp.bad_orientation,
|
||||
params={"limit_angle": 1.22}, # ~70 degrees in radians
|
||||
),
|
||||
}
|
||||
|
||||
The ``time_out`` flag on ``TerminationTermCfg`` tells the manager to treat
|
||||
this condition as a truncation rather than a terminal failure. Truncations
|
||||
map to the ``truncated`` signal in the Gym interface; failures map to
|
||||
``terminated``. This distinction matters for value bootstrapping in RL
|
||||
algorithms.
|
||||
|
||||
.. rubric:: Event terms
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.managers.event_manager import EventTermCfg
|
||||
|
||||
events = {
|
||||
"reset_base": EventTermCfg(
|
||||
func=mdp.reset_root_state_uniform,
|
||||
mode="reset",
|
||||
params={
|
||||
"pose_range": {"yaw": (-3.14, 3.14)},
|
||||
"velocity_range": {},
|
||||
},
|
||||
),
|
||||
}
|
||||
|
||||
The ``mode`` field on ``EventTermCfg`` controls when the term fires:
|
||||
at startup, on episode reset, or at regular intervals. See :ref:`events`
|
||||
for the full treatment of lifecycle modes, built-in event functions, and
|
||||
the relationship between events and domain randomization.
|
||||
|
||||
.. rubric:: Function-based vs. class-based terms
|
||||
|
||||
Terms can be plain functions or classes. Functions are suitable for stateless
|
||||
computations; classes are useful when a term needs to cache expensive setup or
|
||||
maintain state across steps.
|
||||
|
||||
A function-based term has the signature ``func(env, **params) -> Tensor``. A
|
||||
class-based term is instantiated once with ``(cfg, env)`` and then called with
|
||||
the same signature. Classes can optionally implement a ``reset(env_ids)`` hook
|
||||
for per-episode state clearing.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# Function-based (stateless)
|
||||
RewardTermCfg(func=mdp.joint_torques_l2, weight=-0.01)
|
||||
|
||||
# Class-based (caches joint indices at init)
|
||||
class MyReward:
|
||||
def __init__(self, cfg, env):
|
||||
self.joint_ids = resolve_joint_ids(cfg.params, env)
|
||||
|
||||
def __call__(self, env) -> torch.Tensor:
|
||||
return compute_reward(env, self.joint_ids)
|
||||
|
||||
RewardTermCfg(func=MyReward, weight=1.0)
|
||||
|
||||
|
||||
.. _env-config-timing:
|
||||
|
||||
Timing: decimation, timestep, and episode length
|
||||
-------------------------------------------------
|
||||
|
||||
Three parameters jointly determine the temporal structure of the environment.
|
||||
|
||||
``sim.mujoco.timestep``
|
||||
The physics integration step in seconds. The default is 0.002 s (500 Hz).
|
||||
This is one of the most important parameters in any environment: smaller
|
||||
values produce more stable physics but slow down simulation. See the MuJoCo
|
||||
`performance tuning <https://mujoco.readthedocs.io/en/stable/modeling.html#performance-tuning>`_
|
||||
guide for practical advice on choosing timesteps and solver settings.
|
||||
|
||||
``decimation``
|
||||
The number of physics steps executed per policy step. The policy runs at
|
||||
``1 / (timestep * decimation)`` Hz.
|
||||
|
||||
``episode_length_s``
|
||||
The episode duration in seconds. The maximum number of policy steps per
|
||||
episode is ``ceil(episode_length_s / (timestep * decimation))``.
|
||||
|
||||
**Concrete example.** The velocity task uses ``timestep=0.005`` (200 Hz
|
||||
physics) and ``decimation=4``, giving a policy frequency of 50 Hz. With
|
||||
``episode_length_s=20.0``, each episode runs for exactly 1000 policy steps.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
physics_dt = 0.005 # seconds per physics step (200 Hz)
|
||||
decimation = 4 # physics steps per policy step
|
||||
step_dt = 0.005 * 4 # = 0.02 s per policy step (50 Hz)
|
||||
episode_len = 20.0 / 0.02 # = 1000 policy steps per episode
|
||||
|
||||
To read these values at runtime, use the environment properties:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
env.physics_dt # = cfg.sim.mujoco.timestep
|
||||
env.step_dt # = cfg.sim.mujoco.timestep * cfg.decimation
|
||||
env.max_episode_length # steps (int)
|
||||
env.max_episode_length_s # seconds (float)
|
||||
|
||||
When ``scale_rewards_by_dt=True`` (the default), each reward term is
|
||||
multiplied by ``step_dt`` before being returned. A reward function that
|
||||
returns a constant value of 1.0 contributes ``step_dt`` per step and
|
||||
approximately ``episode_length_s`` over a full episode, regardless of how
|
||||
``decimation`` and ``timestep`` are set. Changing the simulation frequency
|
||||
without disabling this scaling leaves reward magnitudes unchanged.
|
||||
|
||||
|
||||
.. _env-config-subclassing:
|
||||
|
||||
Subclassing pattern
|
||||
-------------------
|
||||
|
||||
mjlab uses plain dataclass inheritance rather than deeply nested class
|
||||
hierarchies. To build a task-specific configuration, subclass
|
||||
``ManagerBasedRlEnvCfg`` and override fields.
|
||||
|
||||
The recommended approach is to define the full configuration in a factory
|
||||
function, then call it from robot-specific configs that override only the
|
||||
fields that differ. The velocity task uses this pattern: ``make_velocity_env_cfg``
|
||||
returns a fully assembled ``ManagerBasedRlEnvCfg``, and each robot
|
||||
configuration calls the factory and patches in robot-specific values such
|
||||
as the scene, joint name patterns, and action scale.
|
||||
|
||||
A condensed version of the factory illustrates the full assembly pattern:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import math
|
||||
from dataclasses import replace
|
||||
|
||||
from mjlab.envs import ManagerBasedRlEnvCfg
|
||||
from mjlab.envs.mdp import dr
|
||||
from mjlab.envs.mdp.actions import JointPositionActionCfg
|
||||
from mjlab.managers.event_manager import EventTermCfg
|
||||
from mjlab.managers.observation_manager import ObservationGroupCfg, ObservationTermCfg
|
||||
from mjlab.managers.reward_manager import RewardTermCfg
|
||||
from mjlab.managers.scene_entity_config import SceneEntityCfg
|
||||
from mjlab.managers.termination_manager import TerminationTermCfg
|
||||
from mjlab.scene import SceneCfg
|
||||
from mjlab.sim import MujocoCfg, SimulationCfg
|
||||
from mjlab.tasks.velocity import mdp
|
||||
from mjlab.tasks.velocity.mdp import UniformVelocityCommandCfg
|
||||
from mjlab.terrains import TerrainEntityCfg
|
||||
from mjlab.terrains.config import ROUGH_TERRAINS_CFG
|
||||
from mjlab.viewer import ViewerConfig
|
||||
|
||||
|
||||
def make_velocity_env_cfg() -> ManagerBasedRlEnvCfg:
|
||||
|
||||
observations = {
|
||||
"actor": ObservationGroupCfg(
|
||||
terms={
|
||||
"base_lin_vel": ObservationTermCfg(
|
||||
func=mdp.builtin_sensor,
|
||||
params={"sensor_name": "robot/imu_lin_vel"},
|
||||
),
|
||||
"joint_pos": ObservationTermCfg(func=mdp.joint_pos_rel),
|
||||
"command": ObservationTermCfg(
|
||||
func=mdp.generated_commands,
|
||||
params={"command_name": "twist"},
|
||||
),
|
||||
# additional terms omitted for brevity
|
||||
},
|
||||
concatenate_terms=True,
|
||||
enable_corruption=True,
|
||||
),
|
||||
"critic": ObservationGroupCfg(
|
||||
terms={...},
|
||||
concatenate_terms=True,
|
||||
enable_corruption=False,
|
||||
),
|
||||
}
|
||||
|
||||
actions = {
|
||||
"joint_pos": JointPositionActionCfg(
|
||||
entity_name="robot",
|
||||
actuator_names=(".*",),
|
||||
scale=0.5,
|
||||
use_default_offset=True,
|
||||
)
|
||||
}
|
||||
|
||||
commands = {
|
||||
"twist": UniformVelocityCommandCfg(
|
||||
entity_name="robot",
|
||||
resampling_time_range=(3.0, 8.0),
|
||||
ranges=UniformVelocityCommandCfg.Ranges(
|
||||
lin_vel_x=(-1.0, 1.0),
|
||||
lin_vel_y=(-1.0, 1.0),
|
||||
ang_vel_z=(-0.5, 0.5),
|
||||
heading=(-math.pi, math.pi),
|
||||
),
|
||||
)
|
||||
}
|
||||
|
||||
events = {
|
||||
"reset_base": EventTermCfg(
|
||||
func=mdp.reset_root_state_uniform,
|
||||
mode="reset",
|
||||
params={
|
||||
"pose_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5), "yaw": (-3.14, 3.14)},
|
||||
"velocity_range": {},
|
||||
},
|
||||
),
|
||||
"foot_friction": EventTermCfg(
|
||||
mode="startup",
|
||||
func=dr.geom_friction,
|
||||
params={
|
||||
"asset_cfg": SceneEntityCfg("robot", geom_names=[]),
|
||||
"operation": "abs",
|
||||
"ranges": (0.3, 1.2),
|
||||
},
|
||||
),
|
||||
"push_robot": EventTermCfg(
|
||||
func=mdp.push_by_setting_velocity,
|
||||
mode="interval",
|
||||
interval_range_s=(1.0, 3.0),
|
||||
params={"velocity_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5)}},
|
||||
),
|
||||
}
|
||||
|
||||
rewards = {
|
||||
"track_linear_velocity": RewardTermCfg(
|
||||
func=mdp.track_linear_velocity,
|
||||
weight=2.0,
|
||||
params={"command_name": "twist", "std": math.sqrt(0.25)},
|
||||
),
|
||||
"dof_pos_limits": RewardTermCfg(func=mdp.joint_pos_limits, weight=-1.0),
|
||||
"action_rate_l2": RewardTermCfg(func=mdp.action_rate_l2, weight=-0.1),
|
||||
}
|
||||
|
||||
terminations = {
|
||||
"time_out": TerminationTermCfg(func=mdp.time_out, time_out=True),
|
||||
"fell_over": TerminationTermCfg(
|
||||
func=mdp.bad_orientation,
|
||||
params={"limit_angle": math.radians(70.0)},
|
||||
),
|
||||
}
|
||||
|
||||
return ManagerBasedRlEnvCfg(
|
||||
decimation=4,
|
||||
episode_length_s=20.0,
|
||||
sim=SimulationCfg(
|
||||
nconmax=35,
|
||||
njmax=1500,
|
||||
mujoco=MujocoCfg(timestep=0.005, iterations=10, ls_iterations=20),
|
||||
),
|
||||
scene=SceneCfg(
|
||||
terrain=TerrainEntityCfg(
|
||||
terrain_type="generator",
|
||||
terrain_generator=replace(ROUGH_TERRAINS_CFG),
|
||||
max_init_terrain_level=5,
|
||||
),
|
||||
num_envs=1,
|
||||
),
|
||||
observations=observations,
|
||||
actions=actions,
|
||||
commands=commands,
|
||||
events=events,
|
||||
rewards=rewards,
|
||||
terminations=terminations,
|
||||
)
|
||||
|
||||
Robot-specific configs call this factory and patch fields using
|
||||
``dataclasses.replace`` or direct assignment. Common per-robot overrides
|
||||
include ``scene`` (to add the robot entity and sensors), joint name patterns
|
||||
inside ``SceneEntityCfg``, action ``scale``, and body names for reward terms.
|
||||
|
||||
.. note::
|
||||
|
||||
Isaac Lab uses deeply nested ``__post_init__`` overrides for configuration
|
||||
inheritance. mjlab avoids that pattern: each ``ManagerBasedRlEnvCfg`` is a
|
||||
flat, inspectable dataclass. A misspelled field name raises a ``TypeError``
|
||||
at construction rather than silently creating a new attribute. See
|
||||
:ref:`migration_isaac_lab` for a full comparison.
|
||||
|
||||
|
||||
Where to go next
|
||||
----------------
|
||||
|
||||
The remaining pages in the Manager Layer section cover each manager in
|
||||
detail:
|
||||
|
||||
- :ref:`observations`: observation groups, the processing pipeline
|
||||
(clip, scale, noise, delay, history), and built-in observation functions.
|
||||
- :ref:`actions`: action types and how the action manager routes policy
|
||||
output to actuators.
|
||||
- :ref:`rewards`: reward terms and scaling by dt.
|
||||
- :ref:`terminations`: episode end conditions and the truncation/failure
|
||||
distinction.
|
||||
- :ref:`commands`: command generators and goal-conditioned task setup.
|
||||
- :ref:`events`: the event manager lifecycle (startup, reset, interval).
|
||||
- :ref:`domain_randomization`: the full ``dr`` module for domain
|
||||
randomization.
|
||||
- :ref:`curriculum`: difficulty progression based on policy performance.
|
||||
- :ref:`metrics`: custom per-step metrics logged as episode averages.
|
||||
@@ -0,0 +1,212 @@
|
||||
.. _events:
|
||||
|
||||
Events
|
||||
======
|
||||
|
||||
The event manager executes hooks at specific points in the environment
|
||||
lifecycle. Any logic that should run at startup, on episode reset, or at
|
||||
regular intervals during training is registered as an event term. Common
|
||||
examples include resetting entities to an initial state, applying domain
|
||||
randomization to model parameters, pushing the robot with random velocity
|
||||
perturbations, and initializing robot state from a reference motion clip.
|
||||
All of these are configured through the same ``EventTermCfg`` interface,
|
||||
differing only in the ``mode`` field that controls when each term fires.
|
||||
|
||||
Domain randomization, one of the most common uses of events, has its own
|
||||
dedicated reference page. See :ref:`domain_randomization` for the full
|
||||
``dr`` module, available functions, and internals.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.envs.mdp import events as event_fns, dr
|
||||
from mjlab.managers.event_manager import EventTermCfg
|
||||
from mjlab.managers.scene_entity_config import SceneEntityCfg
|
||||
|
||||
events = {
|
||||
# Reset all entities to their default state each episode.
|
||||
"reset_scene": EventTermCfg(
|
||||
func=event_fns.reset_scene_to_default,
|
||||
mode="reset",
|
||||
),
|
||||
# Randomize foot friction once at startup.
|
||||
"foot_friction": EventTermCfg(
|
||||
func=dr.geom_friction,
|
||||
mode="startup",
|
||||
params={
|
||||
"asset_cfg": SceneEntityCfg("robot", geom_names=[".*foot.*"]),
|
||||
"ranges": (0.3, 1.2),
|
||||
"operation": "abs",
|
||||
},
|
||||
),
|
||||
# Push the robot at random intervals during the episode.
|
||||
"push_robot": EventTermCfg(
|
||||
func=event_fns.push_by_setting_velocity,
|
||||
mode="interval",
|
||||
interval_range_s=(1.0, 3.0),
|
||||
params={
|
||||
"velocity_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5)},
|
||||
},
|
||||
),
|
||||
# Transient random impulses with duration and cooldown.
|
||||
"impulse": EventTermCfg(
|
||||
func=event_fns.apply_body_impulse,
|
||||
mode="step",
|
||||
params={
|
||||
"force_range": (-50.0, 50.0),
|
||||
"torque_range": (0.0, 0.0),
|
||||
"duration_s": (0.1, 0.2),
|
||||
"cooldown_s": (1.0, 3.0),
|
||||
"asset_cfg": SceneEntityCfg("robot", body_names=("base",)),
|
||||
},
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
Lifecycle modes
|
||||
---------------
|
||||
|
||||
The ``mode`` field on ``EventTermCfg`` determines when the term fires. The
|
||||
four modes correspond to the timescales of an RL training run: once at
|
||||
process startup, once per episode, periodically within an episode, and on
|
||||
every environment step.
|
||||
|
||||
``"startup"``
|
||||
Fires once during environment initialization, after all managers are
|
||||
constructed. Every environment receives the event simultaneously. This
|
||||
mode is intended for parameters that should differ across environments
|
||||
but remain fixed for the entire training run, such as link masses or
|
||||
joint armatures randomized via the ``dr`` module.
|
||||
|
||||
``"reset"``
|
||||
Fires on every episode reset, for each environment being reset. This is
|
||||
the most common mode. State initialization (writing the robot back to
|
||||
its default pose) and episode-level domain randomization both belong
|
||||
here.
|
||||
|
||||
The optional ``min_step_count_between_reset`` field prevents the term
|
||||
from firing too frequently when episodes are very short. The term is
|
||||
skipped for any environment that has not taken at least that many steps
|
||||
since its last trigger. The first invocation always fires regardless.
|
||||
|
||||
``"interval"``
|
||||
Fires at regular time intervals during training, independent of episode
|
||||
boundaries. The trigger frequency is controlled by ``interval_range_s``,
|
||||
a ``(min, max)`` range in seconds. After each trigger the manager
|
||||
samples a new wait time uniformly from that range. Each environment has
|
||||
its own independent timer by default; setting ``is_global_time=True``
|
||||
synchronizes all environments to a single shared timer. Interval events
|
||||
are the natural home for mid-episode perturbations such as external
|
||||
pushes or drifting model parameters.
|
||||
|
||||
``"step"``
|
||||
Fires on every environment step, for all environments. This mode is
|
||||
intended for continuous effects that must be evaluated each step, such
|
||||
as ``apply_body_impulse`` which manages its own internal duration and
|
||||
cooldown timers. Because step events run every step, they should be
|
||||
lightweight or manage their own activation logic internally to avoid
|
||||
unnecessary computation.
|
||||
|
||||
As with all manager terms, ``func`` points to the callable and ``params``
|
||||
holds keyword arguments forwarded to it alongside ``env`` and ``env_ids``.
|
||||
Any ``SceneEntityCfg`` values inside ``params`` are resolved once at
|
||||
manager construction (regex patterns are matched to model indices at that
|
||||
point, not on every call). Terms can be plain functions or classes; see
|
||||
:ref:`env-config-term-pattern` for the general pattern.
|
||||
|
||||
|
||||
Built-in event functions
|
||||
------------------------
|
||||
|
||||
The functions below are available in ``mjlab.envs.mdp.events``.
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 30 70
|
||||
|
||||
* - Function
|
||||
- Description
|
||||
* - ``reset_scene_to_default``
|
||||
- Resets all entities to their default states: root pose and velocity
|
||||
for floating-base entities, mocap pose for fixed-base entities, and
|
||||
joint positions and velocities for articulated entities. Environment
|
||||
origins are applied automatically. This is the default event on
|
||||
``ManagerBasedRlEnvCfg``; most environments keep it and add
|
||||
additional terms alongside it.
|
||||
* - ``reset_root_state_uniform``
|
||||
- Resets a single entity's root pose and velocity with uniform random
|
||||
offsets from the default. Accepts ``pose_range`` and
|
||||
``velocity_range`` dictionaries with keys ``"x"``, ``"y"``,
|
||||
``"z"``, ``"roll"``, ``"pitch"``, ``"yaw"``. Orientation
|
||||
perturbations compose with the default quaternion. For fixed-base
|
||||
robots, this is the only way to position them at their environment
|
||||
origins; without it they stack at the world origin.
|
||||
* - ``reset_root_state_from_flat_patches``
|
||||
- Places an entity on a randomly chosen flat terrain patch based on
|
||||
the environment's assigned terrain level and type. Falls back to
|
||||
``reset_root_state_uniform`` when no flat patches are available.
|
||||
Useful for locomotion tasks where robots should spawn on level
|
||||
ground within their assigned sub-terrain.
|
||||
* - ``reset_joints_by_offset``
|
||||
- Resets joint positions and velocities by adding a uniform random
|
||||
offset to the entity's defaults, clamped to soft joint limits.
|
||||
* - ``push_by_setting_velocity``
|
||||
- Adds a random velocity increment to the entity's current root
|
||||
velocity, simulating an external push. Typically used with
|
||||
``mode="interval"`` to test disturbance rejection.
|
||||
* - ``apply_external_force_torque``
|
||||
- Applies random forces and torques to one or more bodies via the
|
||||
MuJoCo external wrench mechanism.
|
||||
* - ``apply_body_impulse``
|
||||
- Applies transient external wrenches to bodies with configurable
|
||||
duration and cooldown. Each environment independently samples a
|
||||
random force direction and holds it for a sampled duration, then
|
||||
waits through a cooldown before firing again. Supports an optional
|
||||
``body_point_offset`` to shift the application point away from the
|
||||
center of mass. Includes built in debug visualization that draws
|
||||
force arrows in the viewer. Use with ``mode="step"``.
|
||||
* - ``randomize_terrain``
|
||||
- Assigns each environment to a random sub-terrain row and column,
|
||||
ignoring the curriculum. Useful for evaluation or play mode.
|
||||
|
||||
|
||||
Writing custom event terms
|
||||
--------------------------
|
||||
|
||||
An event function takes ``env`` and ``env_ids`` as its first two arguments
|
||||
and any additional parameters from ``EventTermCfg.params``. It modifies
|
||||
simulation state in place and returns nothing. For terms that need
|
||||
expensive one-time setup (such as loading data from disk), use a class
|
||||
so that the setup runs once at construction rather than on every call.
|
||||
For example, the following custom event term resets the robot to a
|
||||
random pose sampled from a pre-recorded dataset:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import torch
|
||||
from mjlab.managers.manager_base import ManagerTermBase
|
||||
from mjlab.managers.scene_entity_config import SceneEntityCfg
|
||||
|
||||
class ResetFromDataset(ManagerTermBase):
|
||||
"""Reset the robot to a random pose from a dataset."""
|
||||
|
||||
def __init__(self, cfg, env):
|
||||
super().__init__(env)
|
||||
self._robot = env.scene["robot"]
|
||||
self._poses = torch.load(
|
||||
cfg.params["dataset_path"],
|
||||
map_location=env.device,
|
||||
)
|
||||
|
||||
def __call__(self, env, env_ids, **kwargs):
|
||||
# Sample with replacement: each env gets an independent pose.
|
||||
indices = torch.randint(
|
||||
len(self._poses), (len(env_ids),), device=env.device,
|
||||
)
|
||||
self._robot.write_joint_position_to_sim(
|
||||
self._poses[indices], env_ids=env_ids,
|
||||
)
|
||||
|
||||
When a term needs to maintain state or perform expensive setup, implement
|
||||
it as a class. See :ref:`env-config-term-pattern` for the general
|
||||
pattern. For custom DR terms that write to model fields, see
|
||||
:ref:`domain_randomization`.
|
||||
@@ -0,0 +1,523 @@
|
||||
.. _faq:
|
||||
|
||||
FAQ & Troubleshooting
|
||||
=====================
|
||||
|
||||
This page collects common questions about **platform support**, **performance**,
|
||||
**training stability**, and **visualization**, along with practical debugging
|
||||
tips and links to further resources.
|
||||
|
||||
Platform Support
|
||||
----------------
|
||||
|
||||
Does it work on macOS?
|
||||
~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Yes, but only with limited performance. mjlab runs on macOS
|
||||
using **CPU-only** execution through MuJoCo Warp.
|
||||
|
||||
- **Training is not recommended on macOS**, as it lacks GPU acceleration.
|
||||
- **Evaluation works**, but is significantly slower than on Linux with CUDA.
|
||||
|
||||
For serious training workloads, we strongly recommend **Linux with an NVIDIA GPU**.
|
||||
|
||||
Does it work on Windows?
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
We have performed preliminary testing on **Windows** and **WSL**, but some
|
||||
workflows are not guaranteed to be stable.
|
||||
|
||||
- Windows support may **lag behind** Linux.
|
||||
- Windows will be **tested less frequently**, since Linux is the primary
|
||||
development and deployment platform.
|
||||
- Community contributions that improve Windows support are very welcome.
|
||||
|
||||
CUDA Compatibility
|
||||
~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Not all CUDA versions are supported by MuJoCo Warp.
|
||||
|
||||
- See `mujoco_warp#101 <https://github.com/google-deepmind/mujoco_warp/issues/101>`_
|
||||
for details on CUDA compatibility.
|
||||
- **Recommended**: CUDA **12.4+** (for conditional execution support in CUDA
|
||||
graphs).
|
||||
|
||||
Performance
|
||||
-----------
|
||||
|
||||
Is it faster than Isaac Lab?
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Based on our experience over the last few months, mjlab is **on par or
|
||||
faster** than Isaac Lab.
|
||||
|
||||
What GPU do you recommend?
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
- **RTX 40-series GPUs** (or newer)
|
||||
- **L40s, H100**
|
||||
|
||||
Does mjlab support multi-GPU training?
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Yes, mjlab supports **multi-GPU distributed training** using
|
||||
`torchrunx <https://github.com/apoorvkh/torchrunx>`_.
|
||||
|
||||
- Use ``--gpu-ids "[0, 1]"`` (or ``--gpu-ids all``) when running the ``train``
|
||||
command.
|
||||
- See the :doc:`training/distributed_training` for configuration details and examples.
|
||||
|
||||
Training & Debugging
|
||||
--------------------
|
||||
|
||||
My training crashes with NaN errors
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
A typical error when using ``rsl_rl`` looks like:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
RuntimeError: normal expects all elements of std >= 0.0
|
||||
|
||||
This occurs when NaN/Inf values in the **physics state** propagate to the
|
||||
policy network, causing its output standard deviation to become negative or NaN.
|
||||
|
||||
There are many possible causes, including potential bugs in **MuJoCo Warp**
|
||||
(which is still in beta). mjlab offers two complementary mechanisms to help
|
||||
you handle this:
|
||||
|
||||
1. **For training stability** - NaN termination
|
||||
|
||||
Add a ``nan_detection`` termination to reset environments that hit NaN:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.envs.mdp import terminations as mdp_term
|
||||
from mjlab.managers.termination_manager import TerminationTermCfg
|
||||
|
||||
# In your ManagerBasedRlEnvCfg subclass:
|
||||
terminations = {
|
||||
# Your other terminations...
|
||||
"nan_term": TerminationTermCfg(func=mdp_term.nan_detection),
|
||||
}
|
||||
|
||||
This marks NaN environments as terminated so they can reset while training
|
||||
continues. Terminations are logged as
|
||||
``Episode_Termination/nan_term`` in your metrics.
|
||||
|
||||
.. warning::
|
||||
|
||||
This is a **band-aid solution**. If NaNs correlate with your task objective
|
||||
(for example, NaNs occur exactly when the agent tries to grasp an object),
|
||||
the policy will never learn to complete that part of the task. Always
|
||||
investigate the **root cause** using ``nan_guard`` in addition to this
|
||||
termination.
|
||||
|
||||
2. **For debugging** - NaN guard
|
||||
|
||||
Enable ``nan_guard`` to capture the simulation state when NaNs occur:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
uv run train.py --enable-nan-guard True
|
||||
|
||||
See the :doc:`NaN Guard documentation <debugging/nan_guard>` for details.
|
||||
|
||||
The ``nan_guard`` tool makes it easier to:
|
||||
|
||||
- Inspect the simulation state at the moment NaNs appear.
|
||||
- Build a minimal reproducible example (MRE).
|
||||
- Report potential framework bugs to the
|
||||
`MuJoCo Warp team <https://github.com/google-deepmind/mujoco_warp/issues>`_.
|
||||
|
||||
Reporting well-isolated issues helps improve the framework for everyone.
|
||||
|
||||
How can I inspect the generated scene XML?
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Use the ``export-scene`` script to write the full scene (XML and mesh assets)
|
||||
to a directory:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
uv run export-scene g1 --output-dir /tmp/g1
|
||||
|
||||
The exported ``scene.xml`` can be loaded directly in MuJoCo for visual
|
||||
inspection or diffing. This is useful for verifying that task configuration
|
||||
and physics are set up correctly, and for creating minimal reproducible
|
||||
examples to share with mjlab or MuJoCo Warp developers. The script accepts task IDs,
|
||||
entity aliases (``g1``, ``go1``, ``yam``), or arbitrary import paths. See
|
||||
:doc:`debugging/export_scene` for full details.
|
||||
|
||||
My contact sensor misses collisions when using decimation
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
With ``decimation > 1`` the physics runs multiple substeps per policy
|
||||
step. A brief contact (e.g. a self collision or an illegal ground touch)
|
||||
can appear and disappear within the substep loop, so by the time the
|
||||
sensor is read, ``found`` is zero and the event is invisible to
|
||||
rewards and terminations.
|
||||
|
||||
Set ``history_length`` on the ``ContactSensorCfg`` equal to your
|
||||
decimation value. The sensor then stores force, torque, and distance
|
||||
for the last *N* substeps. Your reward or termination function can
|
||||
inspect the history to detect contacts that would otherwise be missed:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
ContactSensorCfg(
|
||||
name="self_collision",
|
||||
...,
|
||||
fields=("found", "force"),
|
||||
history_length=4, # matches decimation=4
|
||||
)
|
||||
|
||||
# In the reward/termination function:
|
||||
force_mag = torch.norm(sensor.data.force_history, dim=-1) # [B, N, H]
|
||||
had_contact = (force_mag > 10.0).any(dim=1).any(dim=-1) # [B]
|
||||
|
||||
See :ref:`contact-sensor-history` for full details.
|
||||
|
||||
.. note::
|
||||
|
||||
Feet ground sensors with ``track_air_time=True`` already accumulate
|
||||
contact state across substeps, so they do not need history.
|
||||
|
||||
.. _faq-sim-forward:
|
||||
|
||||
When do I need to call ``sim.forward()``?
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Short answer: you almost certainly don't.
|
||||
|
||||
``sim.forward()`` wraps MuJoCo's ``mj_forward``, which runs the full forward
|
||||
dynamics pipeline (kinematics, contacts, forces, constraint solving, sensors)
|
||||
but skips integration, leaving ``qpos``/``qvel`` unchanged. It brings all
|
||||
derived quantities in ``mjData`` (``xpos``, ``xquat``, ``site_xpos``,
|
||||
``cvel``, ``sensordata``, etc.) into a consistent state with the current
|
||||
``qpos``/``qvel``.
|
||||
The environment's ``step()`` method calls it once per step, right before
|
||||
observation computation, so observations, commands, and interval events
|
||||
always see fresh derived quantities. Termination and reward managers run
|
||||
*before* this call and therefore see derived quantities that are stale by
|
||||
one physics substep, a deliberate tradeoff that avoids a second
|
||||
``forward()`` call while keeping the MDP well-defined (the staleness is
|
||||
consistent across all envs and all steps).
|
||||
|
||||
The one case where this matters is if you write an event or command that
|
||||
both writes state and reads derived quantities in the same function. For
|
||||
example, if Event A calls ``entity.write_root_velocity_to_sim()`` (which
|
||||
modifies ``qvel``) and then immediately reads ``entity.data.root_link_vel_w``
|
||||
(which comes from ``cvel``), the read will see stale values from before the
|
||||
write.
|
||||
|
||||
.. warning::
|
||||
|
||||
Write methods (``write_root_state_to_sim``, ``write_joint_state_to_sim``,
|
||||
etc.) modify ``qpos``/``qvel`` directly. Read properties
|
||||
(``root_link_pose_w``, ``body_link_vel_w``, etc.) return derived
|
||||
quantities that are only current as of the last ``sim.forward()`` call.
|
||||
If you need to write then read in the same function, call
|
||||
``env.sim.forward()`` between them.
|
||||
|
||||
For a deeper explanation, see `Discussion #289
|
||||
<https://github.com/mujocolab/mjlab/discussions/289>`_.
|
||||
|
||||
Why aren't my training runs reproducible even with a fixed seed?
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
MuJoCo Warp does not yet guarantee determinism, so running the same
|
||||
simulation with identical inputs may produce slightly different outputs.
|
||||
This is a known limitation being tracked in
|
||||
`mujoco_warp#562 <https://github.com/google-deepmind/mujoco_warp/issues/562>`_.
|
||||
|
||||
Until determinism is implemented upstream, mjlab training runs will not be
|
||||
perfectly reproducible even when setting a seed.
|
||||
|
||||
My XML ``<option>`` flags are not taking effect
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
If you set simulation options like ``<flag contact="disable"/>`` in your
|
||||
entity XML, they will be silently ignored. This is because mjlab composes
|
||||
scenes by attaching entity specs into a parent scene spec using
|
||||
``MjSpec.attach()``, which does not propagate ``<option>`` settings from
|
||||
the child to the parent. This is a MuJoCo design decision: there is no
|
||||
sensible way to merge engine options (timestep, gravity, solver settings,
|
||||
etc.) across multiple attached models.
|
||||
|
||||
To configure simulation options, use :class:`~mjlab.sim.sim.MujocoCfg` in
|
||||
your task's Python config:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.sim.sim import MujocoCfg, SimulationCfg
|
||||
|
||||
sim=SimulationCfg(
|
||||
mujoco=MujocoCfg(
|
||||
disableflags=("contact",),
|
||||
# timestep=0.01, gravity=(0, 0, -9.81), etc.
|
||||
),
|
||||
)
|
||||
|
||||
``MujocoCfg`` applies options directly to the compiled model, so they
|
||||
always take effect. mjlab will emit a warning if it detects non-default
|
||||
``<option>`` fields on an attached entity spec.
|
||||
|
||||
Rendering & Visualization
|
||||
-------------------------
|
||||
|
||||
What visualization options are available?
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
mjlab currently supports two visualizers for policy evaluation and
|
||||
debugging:
|
||||
|
||||
- **Native MuJoCo visualizer** - the built-in visualizer that ships with MuJoCo.
|
||||
- **Viser** - `Viser <https://github.com/nerfstudio-project/viser>`_,
|
||||
a web-based 3D visualization tool.
|
||||
|
||||
We are exploring **training-time visualization** (e.g., live rollout viewers),
|
||||
but this is not yet available.
|
||||
|
||||
As an alternative, mjlab supports **video logging to Weights & Biases
|
||||
(W&B)**, so you can monitor rollout videos directly in the experiment dashboard.
|
||||
|
||||
How many environments can I visualize at once?
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Viewers render a small number of environments for performance reasons.
|
||||
|
||||
- **Offscreen renderer** (for video recording): Renders the tracked
|
||||
environment plus its nearest neighbors. The count is controlled by
|
||||
``ViewerConfig.max_extra_envs`` (default 2).
|
||||
- **Native/Viser viewers**: Limited by MuJoCo's geometry buffer
|
||||
(default 10,000 geoms). The viewer shows whichever environments fit
|
||||
within the geometry budget.
|
||||
|
||||
Why are my fixed-base robots all stacked at the origin instead of in a grid?
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Fixed-base robots require an **explicit reset event** to position them at
|
||||
their ``env_origins``. If your robots appear stacked at (0, 0, 0):
|
||||
|
||||
**Common causes:**
|
||||
|
||||
1. **Missing reset event** - Most common issue.
|
||||
2. **env_spacing is 0 or very small** - Check your ``SceneCfg(env_spacing=...)``.
|
||||
Even with proper reset events, if ``env_spacing=0.0``, all robots will
|
||||
be at the same position. If ``env_spacing`` is very small (e.g., 0.01),
|
||||
they'll be clustered in a tiny area that looks like a line from a distance.
|
||||
|
||||
**Solution**: Add a reset event that calls ``reset_root_state_uniform``:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# In your ManagerBasedRlEnvCfg
|
||||
events = {
|
||||
# For positioning the base of the robot at env_origins.
|
||||
"reset_base": EventTermCfg(
|
||||
func=mdp.reset_root_state_uniform,
|
||||
mode="reset",
|
||||
params={
|
||||
"pose_range": {}, # Empty = use default pose + env_origins
|
||||
"velocity_range": {},
|
||||
},
|
||||
),
|
||||
# ... other events
|
||||
}
|
||||
|
||||
This pattern is used in the example manipulation task (see ``lift_cube_env_cfg.py:85-94``).
|
||||
|
||||
**Why this is needed**: Fixed-base robots are automatically wrapped in mocap
|
||||
bodies by ``auto_wrap_fixed_base_mocap()``, but mocap positioning only happens
|
||||
when you explicitly call a reset event. The ``env_origins`` offset is applied
|
||||
inside ``reset_root_state_uniform()`` at line 131 of ``envs/mdp/events.py``.
|
||||
|
||||
See `issue #560 <https://github.com/mujocolab/mjlab/issues/560>`_ for examples.
|
||||
|
||||
How does env_origins determine robot layout?
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Robot spacing depends on your terrain configuration:
|
||||
|
||||
**Plane terrain** (``terrain_type="plane"``):
|
||||
- Creates an approximately square grid automatically
|
||||
- Grid size: ``ceil(sqrt(num_envs))`` rows x cols
|
||||
- Spacing controlled by ``env_spacing`` parameter (default: 2.0m)
|
||||
- Examples with ``env_spacing=2.0``:
|
||||
- 32 envs → 7x5 grid spanning 12m x 8m
|
||||
- 4096 envs → 64x64 grid spanning 126m x 126m
|
||||
- **Important**: If ``env_spacing=0``, all robots will be at (0, 0, 0)
|
||||
- Implementation: ``terrain_importer.py:_compute_env_origins_grid()``
|
||||
|
||||
**Procedural terrain** (``terrain_type="generator"``):
|
||||
- Origins loaded from pre-generated terrain sub-patches
|
||||
- Grid size: ``TerrainGeneratorCfg.num_rows x num_cols``
|
||||
- Row index = difficulty level (curriculum mode)
|
||||
- Column index = terrain type variant
|
||||
- **Important allocation behavior**: Columns (terrain types) are evenly distributed
|
||||
across environments, but rows (difficulty levels) are randomly sampled. This means
|
||||
multiple environments can spawn on the same (row, col) patch, leaving others unoccupied,
|
||||
even when ``num_envs > num_patches``.
|
||||
- Example: 5x5 grid (25 patches), 100 envs → each column gets exactly 20 envs,
|
||||
but those 20 are randomly distributed across 5 rows, so some patches remain empty.
|
||||
- Supports ``randomize_env_origins()`` to shuffle positions during training
|
||||
|
||||
How do I ensure each terrain type gets its own column?
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Set ``curriculum=True`` in your ``TerrainGeneratorCfg``. This makes column
|
||||
allocation deterministic, with each column getting one terrain type based on
|
||||
normalized proportions.
|
||||
|
||||
Example with 2 terrain types:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
TerrainGeneratorCfg(
|
||||
num_rows=3,
|
||||
num_cols=2,
|
||||
curriculum=True, # Required for deterministic column allocation!
|
||||
sub_terrains={
|
||||
"flat": BoxFlatTerrainCfg(proportion=0.5), # Gets column 0
|
||||
"pillars": HfDiscreteObstaclesTerrainCfg(
|
||||
proportion=0.5, # Gets column 1
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
Without ``curriculum=True``, every patch is randomly sampled and you'll get
|
||||
a random mix of both terrain types scattered across all patches.
|
||||
|
||||
**Note**: When ``num_cols`` equals the number of terrain types, each terrain
|
||||
gets exactly one column regardless of proportion values (they're normalized).
|
||||
When ``num_cols > num_terrain_types``, proportions determine how many columns
|
||||
each terrain type occupies.
|
||||
|
||||
What is flat patch sampling and how does it affect robot spawning?
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Flat patch sampling detects flat regions on heightfield terrains where robots
|
||||
can safely spawn. It uses morphological filtering on the heightfield to find
|
||||
circular areas where height variation is within a tolerance.
|
||||
|
||||
Configure it on any sub-terrain via ``flat_patch_sampling``:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.terrains.terrain_generator import FlatPatchSamplingCfg
|
||||
|
||||
"obstacles": HfDiscreteObstaclesTerrainCfg(
|
||||
...,
|
||||
flat_patch_sampling={
|
||||
"spawn": FlatPatchSamplingCfg(
|
||||
num_patches=10, # patches to sample per sub-terrain
|
||||
patch_radius=0.5, # flatness check radius (meters)
|
||||
max_height_diff=0.05, # max height variation within radius
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
Then use ``reset_root_state_from_flat_patches`` as your reset event to spawn
|
||||
robots on detected patches instead of at the sub-terrain center.
|
||||
|
||||
**Key details:**
|
||||
|
||||
- Only heightfield (``Hf*``) terrains support actual flat patch detection.
|
||||
Box terrains (``Box*``) don't have heightfield data to analyze.
|
||||
- If any sub-terrain in the grid configures ``flat_patch_sampling``, the
|
||||
flat patches array is allocated for **all** cells. Sub-terrains that don't
|
||||
produce patches have their slots filled with the sub-terrain's spawn origin,
|
||||
so ``reset_root_state_from_flat_patches`` always gets valid positions.
|
||||
- Without ``flat_patch_sampling``, use ``reset_root_state_uniform`` which
|
||||
spawns at the sub-terrain origin (``env_origins``) plus an optional random
|
||||
offset.
|
||||
|
||||
Development & Extensions
|
||||
------------------------
|
||||
|
||||
Can I develop custom tasks in my own repository?
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Yes, mjlab has a **plugin system** that lets you develop tasks in separate
|
||||
repositories while still integrating seamlessly with the core:
|
||||
|
||||
- Your tasks appear as regular entries for the ``train`` and ``play`` commands.
|
||||
- You can version and maintain your task repositories independently.
|
||||
|
||||
A complete guide will be available in a future release.
|
||||
|
||||
Assets & Compatibility
|
||||
----------------------
|
||||
|
||||
What robots are included?
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
mjlab includes two **reference robots**:
|
||||
|
||||
- **Unitree Go1** (quadruped).
|
||||
- **Unitree G1** (humanoid).
|
||||
|
||||
These robots serve as:
|
||||
|
||||
- Minimal examples for **robot integration**.
|
||||
- Stable, well-tested baselines for **benchmark tasks**.
|
||||
|
||||
To keep the core library lean, we do **not** plan to aggressively expand the
|
||||
built-in robot library. Additional robots may be provided in separate
|
||||
repositories or community-maintained packages.
|
||||
|
||||
Can I use USD or URDF models?
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
No, mjlab expects **MJCF (MuJoCo XML)** models.
|
||||
|
||||
- You will need to **convert** USD or URDF assets to MJCF.
|
||||
- For many common robots, you can directly use
|
||||
`MuJoCo Menagerie <https://github.com/google-deepmind/mujoco_menagerie>`_,
|
||||
which ships high-quality MJCF models and assets.
|
||||
|
||||
Getting Help
|
||||
------------
|
||||
|
||||
GitHub Issues
|
||||
~~~~~~~~~~~~~
|
||||
|
||||
Use GitHub issues for:
|
||||
|
||||
- **Bug reports**
|
||||
- **Performance regressions**
|
||||
- **Documentation gaps**
|
||||
|
||||
When filing a bug, please include:
|
||||
|
||||
- CUDA driver and runtime versions
|
||||
- GPU model
|
||||
- A minimal reproduction script
|
||||
- Complete error logs and stack traces
|
||||
- Appropriate labels (for example: ``bug``, ``performance``, ``docs``)
|
||||
|
||||
`Open an issue <https://github.com/mujocolab/mjlab/issues>`_
|
||||
|
||||
Discussions
|
||||
~~~~~~~~~~~
|
||||
|
||||
Use GitHub Discussions for:
|
||||
|
||||
- Usage questions (config, debugging, best practices)
|
||||
- Performance tuning tips
|
||||
- Asset conversion and modeling questions
|
||||
- Design discussions and roadmap ideas
|
||||
|
||||
`Start a discussion <https://github.com/mujocolab/mjlab/discussions>`_
|
||||
|
||||
Known Limitations
|
||||
-----------------
|
||||
|
||||
We're tracking missing features for the stable release in
|
||||
https://github.com/mujocolab/mjlab/issues/100. Check our
|
||||
`open issues <https://github.com/mujocolab/mjlab/issues>`_ to see what's actively
|
||||
being worked on.
|
||||
|
||||
If something isn't working or if we've missed something, please
|
||||
`file a bug report <https://github.com/mujocolab/mjlab/issues/new>`_.
|
||||
@@ -0,0 +1,233 @@
|
||||
.. _installation:
|
||||
|
||||
Installation Guide
|
||||
==================
|
||||
|
||||
This guide presents different installation paths so you can
|
||||
choose the one that best fits your use case.
|
||||
|
||||
.. contents::
|
||||
:local:
|
||||
:depth: 1
|
||||
|
||||
.. note::
|
||||
|
||||
**System Requirements**
|
||||
|
||||
- **Training**: Linux + NVIDIA GPU (CUDA 12.4+ recommended)
|
||||
- **Evaluation**: Linux, macOS, or Windows (WSL)
|
||||
- **Python**: 3.10 or higher
|
||||
|
||||
See :ref:`faq` for more details on what is exactly supported.
|
||||
|
||||
|
||||
How to choose an installation method?
|
||||
-------------------------------------
|
||||
|
||||
Select the card that best matches how you plan to use ``mjlab``.
|
||||
|
||||
.. grid:: 2
|
||||
:gutter: 2
|
||||
|
||||
.. grid-item-card:: Method 1 - Use mjlab as a dependency (uv)
|
||||
:link: install-uv-dependency
|
||||
:link-type: ref
|
||||
|
||||
You are **using mjlab as a dependency** in your own project managed by ``uv``. **(Recommended for most users)**
|
||||
|
||||
.. grid-item-card:: Method 2 - Develop / contribute (uv)
|
||||
:link: install-uv-develop
|
||||
:link-type: ref
|
||||
|
||||
You are **trying mjlab** or **contributing to mjlab itself** directly from inside the mjlab repository, with ``uv`` managing the environment.
|
||||
|
||||
.. grid-item-card:: Method 3 - Classic pip / venv / conda
|
||||
:link: install-pip
|
||||
:link-type: ref
|
||||
|
||||
You are using **classic tools** (``pip`` / ``venv`` / ``conda``) and **do not use uv**.
|
||||
|
||||
.. grid-item-card:: Method 4 - Docker / clusters
|
||||
:link: install-docker
|
||||
:link-type: ref
|
||||
|
||||
You are **running in containers or on clusters** and prefer a **Docker-based** setup.
|
||||
|
||||
|
||||
.. _install-uv-dependency:
|
||||
|
||||
Method 1 - Use mjlab as a dependency (uv)
|
||||
-----------------------------------------
|
||||
|
||||
This is our recommended way to use ``mjlab``. You have
|
||||
your own project and want to use ``mjlab`` as a dependency
|
||||
using ``uv``.
|
||||
|
||||
1. Install uv
|
||||
^^^^^^^^^^^^^
|
||||
|
||||
If you do not have ``uv`` installed, run:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
curl -LsSf https://astral.sh/uv/install.sh | sh
|
||||
|
||||
2. Initialize your project
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
Initialize a managed Python project:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
# Create a new package-based project
|
||||
uv init --package my_mjlab_project
|
||||
cd my_mjlab_project
|
||||
|
||||
3. Add mjlab dependencies
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
There are different options to add ``mjlab`` as a dependency.
|
||||
We recommend using the latest stable version from PyPI. If you need
|
||||
the latest features, use the direct GitHub installation. Finally, if you
|
||||
need to use a feature you have developed locally, use the local editable
|
||||
install. These options are interchangeable: you can switch at any time.
|
||||
|
||||
.. tab-set::
|
||||
|
||||
.. tab-item:: PyPI
|
||||
|
||||
Once in your project, install the latest snapshot from PyPI:
|
||||
|
||||
.. code:: bash
|
||||
|
||||
uv add mjlab
|
||||
|
||||
.. tab-item:: Source
|
||||
|
||||
Once in your project, install directly from GitHub without cloning:
|
||||
|
||||
.. code:: bash
|
||||
|
||||
uv add "mjlab @ git+https://github.com/mujocolab/mjlab"
|
||||
|
||||
.. tab-item:: Local
|
||||
|
||||
Clone the repository:
|
||||
|
||||
.. code:: bash
|
||||
|
||||
git clone https://github.com/mujocolab/mjlab.git
|
||||
|
||||
Once in your project, add it as an editable dependency:
|
||||
|
||||
.. code:: bash
|
||||
|
||||
uv add --editable /path/to/cloned/mjlab
|
||||
|
||||
.. tip::
|
||||
|
||||
For a complete example of how to structure a project that integrates a custom robot
|
||||
with an existing ``mjlab`` task, check out the
|
||||
`ANYmal C Velocity Tracking <https://github.com/mujocolab/anymal_c_velocity>`_ repository.
|
||||
|
||||
Verification
|
||||
^^^^^^^^^^^^
|
||||
|
||||
After installation, verify that ``mjlab`` is working by running the demo:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
uv run demo
|
||||
|
||||
|
||||
.. _install-uv-develop:
|
||||
|
||||
Method 2 - Develop / contribute (uv)
|
||||
------------------------------------
|
||||
|
||||
This method is for developing ``mjlab`` itself or contributing to the project.
|
||||
|
||||
.. code:: bash
|
||||
|
||||
git clone https://github.com/mujocolab/mjlab.git && cd mjlab
|
||||
uv sync
|
||||
|
||||
Verification
|
||||
^^^^^^^^^^^^
|
||||
|
||||
After installation, verify that ``mjlab`` is working by running the demo:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
uv run demo
|
||||
|
||||
|
||||
.. _install-pip:
|
||||
|
||||
Method 3 - Classic pip / venv / conda
|
||||
-------------------------------------
|
||||
|
||||
Activate your virtual environment (``venv``, ``conda``, etc.), then install:
|
||||
|
||||
.. code:: bash
|
||||
|
||||
pip install mjlab
|
||||
|
||||
|
||||
Verification
|
||||
^^^^^^^^^^^^
|
||||
|
||||
After installation, verify that ``mjlab`` is working by running the demo:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
demo
|
||||
|
||||
|
||||
.. _install-docker:
|
||||
|
||||
Method 4 - Docker / clusters
|
||||
----------------------------
|
||||
|
||||
Prerequisites:
|
||||
|
||||
- Install Docker: `Docker installation guide <https://docs.docker.com/engine/install/>`_.
|
||||
- Install an appropriate NVIDIA driver for your system and the
|
||||
`NVIDIA Container Toolkit <https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html>`_.
|
||||
|
||||
- Be sure to register the container runtime with Docker and restart,
|
||||
as described in the Docker configuration section of the NVIDIA
|
||||
install guide.
|
||||
|
||||
.. tab-set::
|
||||
|
||||
.. tab-item:: Pre-built image (recommended)
|
||||
|
||||
Pull and run the latest image from the GitHub Container Registry:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
docker run --rm --runtime=nvidia --gpus all \
|
||||
ghcr.io/mujocolab/mjlab uv run demo
|
||||
|
||||
The image is rebuilt on every push to ``main``.
|
||||
|
||||
.. tab-item:: Local build
|
||||
|
||||
Build from source and run:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
./scripts/run_docker.sh uv run demo
|
||||
|
||||
|
||||
Having some troubles?
|
||||
---------------------
|
||||
|
||||
1. **Check the FAQ**
|
||||
|
||||
Consult the mjlab :ref:`faq` for answers to common installation and runtime issues
|
||||
|
||||
2. **Still stuck?**
|
||||
|
||||
Open an issue on GitHub: https://github.com/mujocolab/mjlab/issues
|
||||
@@ -0,0 +1,103 @@
|
||||
.. _metrics:
|
||||
|
||||
Metrics
|
||||
=======
|
||||
|
||||
The metrics manager logs per-step scalar values as episode averages. Unlike
|
||||
rewards, metrics carry no weight and are not scaled by the step duration.
|
||||
They exist purely for diagnostics: tracking quantities such as tracking
|
||||
error, contact forces, or energy consumption alongside reward curves
|
||||
without influencing the optimization.
|
||||
|
||||
Metrics are computed every environment step, accumulated per environment,
|
||||
and averaged over the episode length when the environment resets. The
|
||||
resulting averages are written to the training logger (TensorBoard or
|
||||
Weights & Biases) under the ``Episode_Metrics/`` prefix.
|
||||
|
||||
If the ``metrics`` dictionary on ``ManagerBasedRlEnvCfg`` is empty, the
|
||||
environment substitutes a lightweight no-op manager with zero overhead.
|
||||
|
||||
|
||||
Registration
|
||||
------------
|
||||
|
||||
Each metric term is registered by name in the ``metrics`` dictionary of
|
||||
``ManagerBasedRlEnvCfg``. The configuration is minimal: a callable and an
|
||||
optional ``params`` dictionary.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.managers.metrics_manager import MetricsTermCfg
|
||||
|
||||
metrics = {
|
||||
"base_height": MetricsTermCfg(
|
||||
func=base_height,
|
||||
params={"asset_cfg": SceneEntityCfg("robot")},
|
||||
),
|
||||
}
|
||||
|
||||
The callable receives ``env`` as its first argument and any entries in
|
||||
``params`` as keyword arguments. It must return a tensor of shape
|
||||
``[num_envs]``, one scalar per environment per step.
|
||||
|
||||
|
||||
How metrics are computed
|
||||
-------------------------
|
||||
|
||||
The manager maintains a running sum and a step counter for each
|
||||
environment. On every call to ``compute()``:
|
||||
|
||||
1. The step counter increments for all environments.
|
||||
2. Each term function is called with the current environment state.
|
||||
3. The returned per-environment values are added to the running sums.
|
||||
|
||||
When an environment resets, the manager reduces each term's accumulated
|
||||
values to a scalar, averages the result across all resetting environments,
|
||||
and returns it under the key ``Episode_Metrics/<term_name>``. The sums and
|
||||
counters are then zeroed for the reset environments.
|
||||
|
||||
The reduction is controlled by the ``reduce`` field on ``MetricsTermCfg``:
|
||||
|
||||
- ``"mean"`` (default): divides the accumulated sum by the step count for
|
||||
each environment. Division is per-environment, so environments that
|
||||
terminated early are not diluted by longer-running ones.
|
||||
- ``"last"``: reports the value from the final step of the episode. This is
|
||||
useful for binary success metrics (such as whether the robot is standing)
|
||||
that should not be averaged over time.
|
||||
|
||||
These scalars flow through ``env.extras["log"]`` into the training runner,
|
||||
which writes them to the configured logger. In a typical training run they
|
||||
appear as:
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
Episode_Metrics/base_height
|
||||
Episode_Metrics/contact_force
|
||||
|
||||
alongside the ``Episode_Reward/`` entries produced by the reward manager.
|
||||
|
||||
|
||||
Writing custom metric functions
|
||||
--------------------------------
|
||||
|
||||
A metric function follows the same pattern as reward and observation
|
||||
functions. It takes the environment as its first argument, reads whatever
|
||||
state it needs, and returns a ``[num_envs]`` tensor.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import torch
|
||||
from mjlab.envs import ManagerBasedRlEnv
|
||||
from mjlab.managers.scene_entity_config import SceneEntityCfg
|
||||
|
||||
def base_height(
|
||||
env: ManagerBasedRlEnv,
|
||||
asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
|
||||
) -> torch.Tensor:
|
||||
robot = env.scene[asset_cfg.name]
|
||||
return robot.data.root_link_pos_w[:, 2]
|
||||
|
||||
For metrics that require cached setup or per-episode state, implement the
|
||||
term as a class with ``__init__(self, cfg, env)`` and a ``__call__``
|
||||
method. If the class defines a ``reset(env_ids)`` method, the manager
|
||||
calls it automatically on episode resets.
|
||||
@@ -0,0 +1,285 @@
|
||||
.. _migration_isaac_lab:
|
||||
|
||||
Migrating from Isaac Lab
|
||||
========================
|
||||
|
||||
.. warning::
|
||||
|
||||
This guide is a work in progress. As more users migrate, we will update this
|
||||
page with additional patterns and edge cases. If something is not covered,
|
||||
please open an issue on GitHub or start a discussion:
|
||||
|
||||
- Issues: https://github.com/mujocolab/mjlab/issues
|
||||
- Discussions: https://github.com/mujocolab/mjlab/discussions
|
||||
|
||||
TL;DR
|
||||
-----
|
||||
|
||||
Most Isaac Lab *manager-based* task configs can be ported to ``mjlab`` with
|
||||
only small changes:
|
||||
|
||||
- The overall **MDP structure is the same** (managers for rewards, observations,
|
||||
actions, commands, terminations, events, curriculum).
|
||||
- The **environment base classes are similar**, but naming is slightly
|
||||
different.
|
||||
- The biggest change is **configuration style**: Isaac Lab uses nested
|
||||
``@configclass`` definitions; ``mjlab`` uses dictionaries of config objects.
|
||||
|
||||
If you are familiar with Isaac Lab's manager-based API, migration is mostly
|
||||
mechanical.
|
||||
|
||||
Key Differences
|
||||
---------------
|
||||
|
||||
1. Import Paths
|
||||
~~~~~~~~~~~~~~~
|
||||
|
||||
Isaac Lab:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from isaaclab.envs import ManagerBasedRLEnv
|
||||
|
||||
mjlab:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.envs import ManagerBasedRlEnvCfg
|
||||
|
||||
.. note::
|
||||
|
||||
``mjlab`` uses a consistent ``CamelCase`` naming convention (for example,
|
||||
``RlEnv`` instead of ``RLEnv``).
|
||||
|
||||
2. Configuration Structure
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Isaac Lab uses nested ``@configclass`` blocks for manager terms. ``mjlab``
|
||||
instead uses **plain dictionaries** mapping names to config objects, which makes
|
||||
it easy to construct variants, merge configs, or generate them programmatically.
|
||||
For the full context behind this design decision, see
|
||||
`PR #292 <https://github.com/mujocolab/mjlab/pull/292>`_.
|
||||
|
||||
**Isaac Lab:**
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@configclass
|
||||
class RewardsCfg:
|
||||
"""Reward terms for the MDP."""
|
||||
|
||||
motion_global_anchor_pos = RewTerm(
|
||||
func=mdp.motion_global_anchor_position_error_exp,
|
||||
weight=0.5,
|
||||
params={"command_name": "motion", "std": 0.3},
|
||||
)
|
||||
motion_global_anchor_ori = RewTerm(
|
||||
func=mdp.motion_global_anchor_orientation_error_exp,
|
||||
weight=0.5,
|
||||
params={"command_name": "motion", "std": 0.4},
|
||||
)
|
||||
|
||||
**mjlab:**
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
rewards = {
|
||||
"motion_global_anchor_pos": RewardTermCfg(
|
||||
func=mdp.motion_global_anchor_position_error_exp,
|
||||
weight=0.5,
|
||||
params={"command_name": "motion", "std": 0.3},
|
||||
),
|
||||
"motion_global_anchor_ori": RewardTermCfg(
|
||||
func=mdp.motion_global_anchor_orientation_error_exp,
|
||||
weight=0.5,
|
||||
params={"command_name": "motion", "std": 0.4},
|
||||
),
|
||||
}
|
||||
|
||||
cfg = ManagerBasedRlEnvCfg(
|
||||
scene=scene,
|
||||
rewards=rewards,
|
||||
# ... other manager dictionaries:
|
||||
# observations=..., actions=..., commands=..., terminations=...,
|
||||
# events=..., curriculum=...
|
||||
)
|
||||
|
||||
This pattern applies to all managers:
|
||||
|
||||
- ``rewards``
|
||||
- ``observations``
|
||||
- ``actions``
|
||||
- ``commands``
|
||||
- ``terminations``
|
||||
- ``events``
|
||||
- ``curriculum``
|
||||
|
||||
3. Scene Configuration
|
||||
~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Scene setup is **simpler** in ``mjlab``:
|
||||
|
||||
- No Omniverse / USD scene graph, no ``prim_path`` management.
|
||||
- Assets are pure MuJoCo (MJCF) with modifier dataclasses applied to
|
||||
``mujoco.MjSpec``.
|
||||
- Lights, materials, textures, and sensors are configured as part of
|
||||
``SceneCfg`` and robot configs.
|
||||
|
||||
**Isaac Lab:**
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from whole_body_tracking.robots.g1 import G1_ACTION_SCALE, G1_CYLINDER_CFG
|
||||
from isaaclab.scene import InteractiveSceneCfg
|
||||
from isaaclab.sensors import ContactSensorCfg
|
||||
from isaaclab.terrains import TerrainImporterCfg
|
||||
import isaaclab.sim as sim_utils
|
||||
from isaaclab.assets import ArticulationCfg, AssetBaseCfg
|
||||
|
||||
@configclass
|
||||
class MySceneCfg(InteractiveSceneCfg):
|
||||
"""Configuration for the terrain scene with a legged robot."""
|
||||
|
||||
# ground terrain
|
||||
terrain = TerrainEntityCfg(
|
||||
prim_path="/World/ground",
|
||||
terrain_type="plane",
|
||||
collision_group=-1,
|
||||
physics_material=sim_utils.RigidBodyMaterialCfg(
|
||||
friction_combine_mode="multiply",
|
||||
restitution_combine_mode="multiply",
|
||||
static_friction=1.0,
|
||||
dynamic_friction=1.0,
|
||||
),
|
||||
visual_material=sim_utils.MdlFileCfg(
|
||||
mdl_path="{NVIDIA_NUCLEUS_DIR}/Materials/Base/Architecture/Shingles_01.mdl",
|
||||
project_uvw=True,
|
||||
),
|
||||
)
|
||||
# lights
|
||||
light = AssetBaseCfg(
|
||||
prim_path="/World/light",
|
||||
spawn=sim_utils.DistantLightCfg(
|
||||
color=(0.75, 0.75, 0.75), intensity=3000.0
|
||||
),
|
||||
)
|
||||
sky_light = AssetBaseCfg(
|
||||
prim_path="/World/skyLight",
|
||||
spawn=sim_utils.DomeLightCfg(
|
||||
color=(0.13, 0.13, 0.13), intensity=1000.0
|
||||
),
|
||||
)
|
||||
robot = G1_CYLINDER_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot")
|
||||
|
||||
**mjlab:**
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from dataclasses import replace
|
||||
|
||||
from mjlab.scene import SceneCfg
|
||||
from mjlab.asset_zoo.robots.unitree_g1.g1_constants import get_g1_robot_cfg
|
||||
from mjlab.utils.spec_config import ContactSensorCfg
|
||||
from mjlab.terrains import TerrainEntityCfg
|
||||
|
||||
# Configure contact sensor
|
||||
self_collision_sensor = ContactSensorCfg(
|
||||
name="self_collision",
|
||||
subtree1="pelvis",
|
||||
subtree2="pelvis",
|
||||
data=("found",),
|
||||
reduce="netforce",
|
||||
num=10, # report up to 10 contacts
|
||||
)
|
||||
|
||||
# Add sensor to robot config
|
||||
g1_cfg = replace(get_g1_robot_cfg(), sensors=(self_collision_sensor,))
|
||||
|
||||
# Create scene
|
||||
SCENE_CFG = SceneCfg(
|
||||
terrain=TerrainEntityCfg(terrain_type="plane"),
|
||||
entities={"robot": g1_cfg},
|
||||
)
|
||||
|
||||
Key changes:
|
||||
|
||||
- No USD ``prim_path`` or cloning; the scene is described directly in MuJoCo.
|
||||
- Materials, lights, and visual properties are applied via
|
||||
``MjSpec``-modifier dataclasses.
|
||||
- See ``mjlab.utils.spec_config`` in the repository for helpers that apply
|
||||
these changes for you.
|
||||
- ``asset_name`` has been unified to ``entity_name`` across all configurations.
|
||||
|
||||
Complete Example Comparison
|
||||
---------------------------
|
||||
|
||||
A good way to learn the pattern is to compare concrete tasks that have already
|
||||
been ported:
|
||||
|
||||
- Isaac Lab implementation (Beyond Mimic):
|
||||
|
||||
- https://github.com/HybridRobotics/whole_body_tracking/blob/main/source/whole_body_tracking/whole_body_tracking/tasks/tracking/tracking_env_cfg.py
|
||||
|
||||
- mjlab implementation:
|
||||
|
||||
- https://github.com/mujocolab/mjlab/blob/main/src/mjlab/tasks/tracking/tracking_env_cfg.py
|
||||
|
||||
You will see that:
|
||||
|
||||
- Manager dictionaries in ``mjlab`` mirror Isaac Lab's config classes,
|
||||
- Reward, observation, command, and termination logic is almost identical,
|
||||
- Scene and asset setup are simplified to pure MuJoCo.
|
||||
|
||||
Migration Checklist
|
||||
-------------------
|
||||
|
||||
Use this as a quick checklist when porting a task:
|
||||
|
||||
1. **Base class and imports**
|
||||
|
||||
- Replace Isaac Lab imports (for example,
|
||||
``from isaaclab.envs import ManagerBasedRLEnv``) with the corresponding
|
||||
``mjlab`` imports (for example,
|
||||
``from mjlab.envs import ManagerBasedRlEnvCfg``).
|
||||
|
||||
2. **Manager configuration**
|
||||
|
||||
- Convert each Isaac Lab ``@configclass`` manager (``RewardsCfg``,
|
||||
``ObservationsCfg``, etc.) into a dictionary of config objects.
|
||||
- Pass these dictionaries into ``ManagerBasedRlEnvCfg``.
|
||||
|
||||
3. **Scene and assets**
|
||||
|
||||
- Replace ``InteractiveSceneCfg`` with a ``SceneCfg`` instance.
|
||||
- Replace USD / ``prim_path`` logic with MuJoCo asset configs and scene
|
||||
entities (for example, a robot from ``asset_zoo``).
|
||||
|
||||
4. **Sensors and contact handling**
|
||||
|
||||
- Convert Isaac Lab ``ContactSensorCfg`` to
|
||||
``mjlab.utils.spec_config.ContactSensorCfg`` and attach it to the robot
|
||||
config.
|
||||
|
||||
5. **RL entry points**
|
||||
|
||||
- Make sure your training script or entry point uses the correct task id and
|
||||
environment config (for example, via Gymnasium registration or direct
|
||||
construction, depending on how your project is structured).
|
||||
|
||||
Tips and Support
|
||||
----------------
|
||||
|
||||
1. Check the examples in the repository under:
|
||||
|
||||
- ``src/mjlab/tasks/``
|
||||
|
||||
2. If you get stuck:
|
||||
|
||||
- Open an issue: https://github.com/mujocolab/mjlab/issues
|
||||
- Start a discussion: https://github.com/mujocolab/mjlab/discussions
|
||||
|
||||
3. Keep in mind MuJoCo vs Isaac Sim differences:
|
||||
|
||||
- Some Omniverse / USD rendering features do not have direct equivalents.
|
||||
- Focus first on matching the **physics and observations**, then polish
|
||||
visuals if needed.
|
||||
@@ -0,0 +1,88 @@
|
||||
.. _motivation:
|
||||
|
||||
Why mjlab?
|
||||
==========
|
||||
|
||||
Reinforcement learning has become a powerful tool for training robot
|
||||
controllers in simulation and transferring them to real hardware. The
|
||||
fidelity of this pipeline hinges on getting simulation details right.
|
||||
|
||||
Several frameworks address this.
|
||||
`Isaac Lab <https://github.com/isaac-sim/IsaacLab>`_
|
||||
provides a comprehensive manager-based API for composing RL environments,
|
||||
but requires the Omniverse runtime, which adds installation complexity and
|
||||
startup latency.
|
||||
`MuJoCo Playground <https://playground.mujoco.org/>`_ takes the opposite
|
||||
approach: minimal
|
||||
abstractions and monolithic environment definitions that are easy to hack
|
||||
and quick to prototype, but code duplication across robots and tasks makes
|
||||
multi-task codebases difficult to maintain. There remains a gap for a
|
||||
framework that is both lightweight and built on a proven orchestration API
|
||||
with access to best-in-class physics.
|
||||
|
||||
mjlab fills this gap. It adopts Isaac Lab's manager-based design, where
|
||||
users compose self-contained building blocks for observations, rewards,
|
||||
events, and commands, and pairs it with MuJoCo Warp for GPU-accelerated
|
||||
physics simulation. The result is a framework with minimal dependencies,
|
||||
fast startup, direct access to native MuJoCo model and data structures,
|
||||
and a PyTorch-native interface for policy training.
|
||||
|
||||
|
||||
Design philosophy
|
||||
-----------------
|
||||
|
||||
mjlab is designed around three core engineering commitments:
|
||||
|
||||
1. **Minimal installation friction.** A single
|
||||
``uvx --from mjlab --refresh demo`` command is enough to run the
|
||||
framework. No heavyweight runtimes, no multi-gigabyte downloads. The
|
||||
dependency footprint is kept intentionally small.
|
||||
|
||||
2. **Transparent and inspectable physics.** mjlab targets a single physics
|
||||
stack, MuJoCo Warp, to prioritize simulation transparency and
|
||||
debuggability. The framework exposes MuJoCo-native ``MjModel`` and
|
||||
``MjData`` structures for direct inspection and state access.
|
||||
Cross-simulator portability is a non-goal; mjlab favors precise control
|
||||
and interpretability over backend generality.
|
||||
|
||||
3. **Tight MuJoCo ecosystem integration.** Users work directly with MuJoCo
|
||||
models and conventions. MJCF files, MuJoCo Menagerie assets, and
|
||||
standard MuJoCo tooling all work without translation layers.
|
||||
|
||||
|
||||
Scope
|
||||
-----
|
||||
|
||||
mjlab provides infrastructure for rigid-body robot learning. It includes
|
||||
depth and raycast sensors for geometric perception. High-fidelity RGB
|
||||
rendering is out of scope. This does not preclude vision-based policies:
|
||||
a common approach is to train privileged policies using full state, then
|
||||
distill into vision-based controllers using external rendering.
|
||||
|
||||
The framework is intended to be extended to custom robots, tasks, sensors,
|
||||
and actuators. It ships with reference implementations of velocity tracking,
|
||||
motion imitation, and manipulation tasks.
|
||||
|
||||
|
||||
Comparison
|
||||
----------
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 25 25 50
|
||||
|
||||
* - Framework
|
||||
- Strengths
|
||||
- Best for
|
||||
* - **mjlab**
|
||||
- Lightweight, fast iteration, native MuJoCo, PyTorch
|
||||
- MuJoCo users who want structured RL environments with GPU acceleration
|
||||
* - **Isaac Lab**
|
||||
- Photorealistic rendering, USD pipeline, Omniverse ecosystem
|
||||
- Projects that need Isaac Sim capabilities
|
||||
* - **MuJoCo Playground**
|
||||
- Minimal abstractions, easy to hack, quick prototyping
|
||||
- One-off experiments and rapid iteration on single tasks
|
||||
* - **Newton**
|
||||
- Multi-physics solvers (deformables, VBD), differentiable simulation
|
||||
- Projects that need solver flexibility beyond rigid-body MuJoCo
|
||||
@@ -0,0 +1,287 @@
|
||||
.. _observations:
|
||||
|
||||
Observations
|
||||
============
|
||||
|
||||
Observations define what the agent perceives at each step. The
|
||||
observation manager assembles individual observation terms into the
|
||||
tensor the policy receives as input. Each term passes through a
|
||||
configurable processing pipeline: noise injection, clipping, scaling,
|
||||
sensor delay, and history stacking.
|
||||
|
||||
|
||||
Observation groups
|
||||
------------------
|
||||
|
||||
Each group is an ``ObservationGroupCfg`` that holds a ``terms`` dict
|
||||
mapping string names to ``ObservationTermCfg`` entries. The manager
|
||||
concatenates term outputs in registration order along the last dimension.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.managers.observation_manager import (
|
||||
ObservationGroupCfg,
|
||||
ObservationTermCfg,
|
||||
)
|
||||
from mjlab.envs.mdp import observations as obs_fns
|
||||
|
||||
observations = {
|
||||
"policy": ObservationGroupCfg(
|
||||
terms={
|
||||
"base_lin_vel": ObservationTermCfg(func=obs_fns.base_lin_vel),
|
||||
"base_ang_vel": ObservationTermCfg(func=obs_fns.base_ang_vel),
|
||||
"projected_gravity": ObservationTermCfg(
|
||||
func=obs_fns.projected_gravity
|
||||
),
|
||||
"joint_pos": ObservationTermCfg(func=obs_fns.joint_pos_rel),
|
||||
"joint_vel": ObservationTermCfg(func=obs_fns.joint_vel_rel),
|
||||
"last_action": ObservationTermCfg(func=obs_fns.last_action),
|
||||
},
|
||||
enable_corruption=True,
|
||||
),
|
||||
}
|
||||
|
||||
This dictionary is passed to ``ManagerBasedRlEnvCfg(observations=...)``.
|
||||
The observation manager resolves term functions at initialization and
|
||||
allocates any required history or delay buffers at that point.
|
||||
|
||||
By default, term outputs within a group are concatenated along the last
|
||||
dimension into a single ``[num_envs, D]`` tensor. Set
|
||||
``concatenate_terms=False`` to receive a dict mapping term names to
|
||||
individual tensors instead.
|
||||
|
||||
The ``enable_corruption`` flag gates noise application for the entire
|
||||
group: when ``False``, noise configs on individual terms are ignored.
|
||||
This makes it straightforward to share term definitions between a noisy
|
||||
actor group and a noise-free critic group, as shown in the
|
||||
:ref:`asymmetric actor-critic <obs-asymmetric>` section below.
|
||||
|
||||
History and delay can also be set at the group level to apply uniformly
|
||||
across all terms; see :ref:`obs-history-delay`.
|
||||
|
||||
|
||||
Processing pipeline
|
||||
-------------------
|
||||
|
||||
Each step, every term in every group passes through the following
|
||||
pipeline in order:
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
compute → noise → clip → scale → delay → history
|
||||
|
||||
1. **compute**: the term function is called. It must return a
|
||||
``[num_envs, D]`` tensor.
|
||||
|
||||
2. **noise**: if ``enable_corruption=True`` on the group and the term
|
||||
has a ``noise`` config, noise is applied. Stateless noise
|
||||
(``NoiseCfg``) is applied directly; stateful noise (``NoiseModelCfg``)
|
||||
is maintained by the manager across steps.
|
||||
|
||||
3. **clip**: if ``clip=(lo, hi)`` is set on the term, values are clamped
|
||||
to that range.
|
||||
|
||||
4. **scale**: if ``scale`` is set, the output is multiplied
|
||||
element-wise. Accepts a scalar, a tuple, or a tensor.
|
||||
|
||||
5. **delay**: if ``delay_max_lag > 0``, the term's output is stored in a
|
||||
ring buffer and a value from an earlier step is returned. See
|
||||
:ref:`obs-history-delay`.
|
||||
|
||||
6. **history**: if ``history_length > 0``, past outputs are stacked.
|
||||
See :ref:`obs-history-delay`.
|
||||
|
||||
.. note::
|
||||
|
||||
Delay is applied before history. This models real systems where old
|
||||
sensor readings are buffered: the history stacks delayed observations,
|
||||
not future ones.
|
||||
|
||||
|
||||
.. _obs-history-delay:
|
||||
|
||||
Observation history and delay
|
||||
------------------------------
|
||||
|
||||
Observations support two temporal features: history and delay. History
|
||||
stacks past frames to give the policy temporal context; delay models
|
||||
sensor latency by returning observations from earlier timesteps.
|
||||
|
||||
Both are configured per term via fields on ``ObservationTermCfg``.
|
||||
They can also be set at the group level on ``ObservationGroupCfg``,
|
||||
which applies uniformly to all terms in the group. Term-level settings
|
||||
override group-level settings.
|
||||
|
||||
History
|
||||
^^^^^^^
|
||||
|
||||
Setting ``history_length=N`` stacks the N most recent outputs of a term.
|
||||
When ``flatten_history_dim=True`` (the default), the history dimension
|
||||
is folded into the feature dimension, producing a ``[num_envs, N * D]``
|
||||
tensor suitable for MLPs. When ``flatten_history_dim=False``, the output
|
||||
retains the time dimension as ``[num_envs, N, D]``, suitable for RNNs.
|
||||
|
||||
History buffers are cleared on environment reset. The first observation
|
||||
after reset is backfilled across all history slots, so the policy
|
||||
receives valid data from step zero.
|
||||
|
||||
When ``flatten_history_dim=True`` and ``concatenate_terms=True``, mjlab
|
||||
uses **term-major** ordering: each term's full history is flattened
|
||||
before concatenating across terms.
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
Term A (D=4, history=3), Term B (D=2, history=3):
|
||||
[A_t0, A_t1, A_t2, B_t0, B_t1, B_t2]
|
||||
└─ A history ──┘ └─ B history ─┘
|
||||
|
||||
Some frameworks use **time-major** ordering instead, where full frames
|
||||
are built at each timestep before concatenating across time. Transferring
|
||||
policies between frameworks with different orderings requires reindexing
|
||||
the observation vector.
|
||||
|
||||
Delay
|
||||
^^^^^
|
||||
|
||||
Setting ``delay_max_lag > 0`` enables a ring buffer that stores past
|
||||
outputs and returns one from an earlier step. The lag is sampled
|
||||
uniformly from ``[delay_min_lag, delay_max_lag]`` in integer steps.
|
||||
A lag of zero returns the current observation; a lag of two returns the
|
||||
observation from two steps ago.
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
50Hz control (20ms/step), lag=2:
|
||||
|
||||
Sensor captures: A B C D E F G H
|
||||
Control steps: 0 1 2 3 4 5 6 7
|
||||
|
||||
Policy sees: A A A B C D E F
|
||||
└clamp┘ └ 40ms delay from here on
|
||||
|
||||
Steps 0-1: lag clamped because the buffer is not yet full.
|
||||
Step 2 onward: each step returns the observation from 2 steps ago.
|
||||
|
||||
To convert real-world latency to lag steps:
|
||||
``lag = latency_seconds / step_dt``. At 50 Hz control (20 ms per step),
|
||||
a 40 ms sensor latency corresponds to a lag of 2. Delays are quantized
|
||||
to integer steps; to approximate a latency that falls between steps, set
|
||||
``delay_min_lag`` and ``delay_max_lag`` to the two nearest integers.
|
||||
|
||||
By default each environment samples its own lag independently
|
||||
(``delay_per_env=True``). Additional parameters control resampling
|
||||
frequency (``delay_update_period``), hold probability
|
||||
(``delay_hold_prob``), and phase staggering
|
||||
(``delay_per_env_phase``).
|
||||
|
||||
Both history and delay buffers are allocated only when enabled; terms
|
||||
with default settings incur no overhead.
|
||||
|
||||
|
||||
Built-in observation functions
|
||||
--------------------------------
|
||||
|
||||
The functions below live in ``mjlab.envs.mdp.observations`` (also
|
||||
re-exported as ``mjlab.envs.mdp``). All return ``[num_envs, D]``
|
||||
tensors.
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 26 74
|
||||
|
||||
* - Function
|
||||
- Description
|
||||
* - ``base_lin_vel``
|
||||
- Linear velocity of the robot base in the base frame.
|
||||
* - ``base_ang_vel``
|
||||
- Angular velocity of the robot base in the base frame.
|
||||
* - ``projected_gravity``
|
||||
- Gravity vector projected into the base frame. Provides roll and
|
||||
pitch information without an explicit orientation representation.
|
||||
* - ``joint_pos_rel``
|
||||
- Joint positions relative to the default pose. Pass
|
||||
``biased=True`` for encoder-biased positions (for sim2real with
|
||||
``dr.encoder_bias``).
|
||||
* - ``joint_vel_rel``
|
||||
- Joint velocities relative to the default velocities.
|
||||
* - ``last_action``
|
||||
- The most recent action tensor. Optionally pass ``action_name``
|
||||
to select a single action term.
|
||||
* - ``generated_commands``
|
||||
- The current command tensor from a named command term. Requires
|
||||
``params={"command_name": "<name>"}``.
|
||||
* - ``builtin_sensor``
|
||||
- Raw data from a named ``BuiltinSensor`` (MuJoCo ``sensordata``
|
||||
slice). Requires ``params={"sensor_name": "<entity>/<sensor>"}``.
|
||||
* - ``height_scan``
|
||||
- Height above each raycast hit point from a ``RayCastSensor``.
|
||||
Requires ``params={"sensor_name": "<name>"}``.
|
||||
|
||||
For ``builtin_sensor`` and ``height_scan``, the ``sensor_name`` parameter
|
||||
must match a sensor registered in the scene. See :ref:`sensors` for how
|
||||
to configure sensors.
|
||||
|
||||
|
||||
.. _obs-asymmetric:
|
||||
|
||||
Asymmetric actor-critic
|
||||
-----------------------
|
||||
|
||||
Multiple observation groups enable asymmetric actor-critic
|
||||
architectures. The actor group contains only the observations that
|
||||
would be available on real hardware; the critic group can include
|
||||
privileged simulation state that is only accessible during training.
|
||||
|
||||
The velocity locomotion task uses this pattern. The actor group
|
||||
receives noisy IMU readings and joint state; the critic group adds
|
||||
noise-free height scan data and foot contact information. The
|
||||
``enable_corruption`` flag makes this separation clean: actor terms
|
||||
carry noise configs but the critic group disables them entirely.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
observations = {
|
||||
"actor": ObservationGroupCfg(
|
||||
terms=actor_terms,
|
||||
concatenate_terms=True,
|
||||
enable_corruption=True, # Noise active during training.
|
||||
),
|
||||
"critic": ObservationGroupCfg(
|
||||
terms={**actor_terms, **privileged_terms},
|
||||
concatenate_terms=True,
|
||||
enable_corruption=False, # No noise on critic.
|
||||
),
|
||||
}
|
||||
|
||||
The training framework receives both groups. The policy network reads
|
||||
``obs["actor"]`` at inference time; the value network reads
|
||||
``obs["critic"]`` during training only.
|
||||
|
||||
|
||||
Writing custom observation functions
|
||||
--------------------------------------
|
||||
|
||||
An observation function accepts ``env`` as its first argument and
|
||||
returns a ``[num_envs, D]`` tensor. Additional parameters are declared
|
||||
as function arguments and supplied via
|
||||
``ObservationTermCfg(params={...})``.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import torch
|
||||
from mjlab.envs import ManagerBasedRlEnv
|
||||
from mjlab.managers.scene_entity_config import SceneEntityCfg
|
||||
|
||||
|
||||
def my_observation(
|
||||
env: ManagerBasedRlEnv,
|
||||
asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
|
||||
) -> torch.Tensor:
|
||||
robot = env.scene[asset_cfg.name]
|
||||
return robot.data.root_lin_vel_b
|
||||
|
||||
When a term needs to cache setup work or maintain per-episode state,
|
||||
implement it as a class with ``__init__(self, cfg, env)`` and
|
||||
``__call__(self, env, ...)``. If the class has a ``reset(env_ids)``
|
||||
method, the manager calls it automatically on episode resets. See
|
||||
:ref:`env-config-term-pattern` for the general pattern.
|
||||
@@ -0,0 +1,121 @@
|
||||
.. _recorders:
|
||||
|
||||
Recorders
|
||||
=========
|
||||
|
||||
The recorder manager provides lifecycle hooks for logging data during
|
||||
rollouts. Unlike rewards, recorders have no effect on the optimization
|
||||
loop. They exist purely to let you capture observations, actions, or any
|
||||
other environment state without modifying mjlab internals.
|
||||
|
||||
Each recorder term is a class that you implement. mjlab calls its methods
|
||||
at the right moments and leaves all I/O decisions to you. If the ``recorders``
|
||||
dictionary on ``ManagerBasedRlEnvCfg`` is empty, the environment substitutes a
|
||||
lightweight no-op manager with zero overhead.
|
||||
|
||||
|
||||
Lifecycle hooks
|
||||
---------------
|
||||
|
||||
The manager exposes three hooks per term:
|
||||
|
||||
``record_pre_reset(env_ids)``
|
||||
Called inside ``env.step()`` before terminated environments are reset.
|
||||
``obs_buf`` holds the observation from the end of the *previous* step
|
||||
(the input the agent used to choose the terminal action, not the
|
||||
post-action terminal state). ``action_manager.action`` holds the
|
||||
terminal action and is still valid here; it will be zeroed for these
|
||||
environments by ``_reset_idx`` immediately after. ``reward_buf`` holds
|
||||
the terminal reward. This is the right place to record the terminal
|
||||
transition ``(obs_t, action_t, reward_t, done=True)``.
|
||||
|
||||
``record_post_reset(env_ids)``
|
||||
Called after a reset completes and fresh observations are available.
|
||||
Fires at the end of ``env.reset()`` (all environments) and within
|
||||
``env.step()`` after each batch of terminated environments is reset.
|
||||
``obs_buf[env_ids]`` holds the initial observation of the new episode;
|
||||
``action_manager.action[env_ids]`` is zero. Use this to initialize
|
||||
per-episode state or record the first observation.
|
||||
|
||||
``record_post_step()``
|
||||
Called at the end of every ``env.step()`` with fresh observations.
|
||||
For environments that reset during this step, ``action_manager.action``
|
||||
has been zeroed and ``obs_buf`` holds the initial state of the new
|
||||
episode rather than the post-action terminal observation. Use
|
||||
``record_pre_reset`` for those environments' terminal transitions and
|
||||
``self._env.reset_buf`` to identify which environments reset.
|
||||
|
||||
``close()``
|
||||
Called when the environment closes. Release file handles and flush
|
||||
buffers here.
|
||||
|
||||
|
||||
Writing a recorder term
|
||||
------------------------
|
||||
|
||||
Subclass :class:`~mjlab.managers.RecorderTerm` and override whichever
|
||||
hooks you need. The environment is available as ``self._env``, giving
|
||||
access to ``self._env.obs_buf``, ``self._env.action_manager.action``, and
|
||||
all other managers.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import csv
|
||||
from mjlab.managers import RecorderTerm, RecorderTermCfg
|
||||
|
||||
class CsvRecorder(RecorderTerm):
|
||||
def __init__(self, cfg, env):
|
||||
super().__init__(cfg, env)
|
||||
self._file = open(cfg.params["path"], "w", newline="")
|
||||
self._writer = csv.writer(self._file)
|
||||
|
||||
def record_pre_reset(self, env_ids):
|
||||
# Terminal transition: action is still intact here.
|
||||
# It will be zeroed by _reset_idx immediately after this returns.
|
||||
obs = self._env.obs_buf["actor"][env_ids].cpu().numpy()
|
||||
act = self._env.action_manager.action[env_ids].cpu().numpy()
|
||||
for o, a in zip(obs, act):
|
||||
self._writer.writerow(o.tolist() + a.tolist())
|
||||
|
||||
def record_post_step(self):
|
||||
# Skip envs that just reset: their terminal pair was written
|
||||
# in record_pre_reset and their action is now zeroed.
|
||||
mask = ~self._env.reset_buf
|
||||
obs = self._env.obs_buf["actor"][mask].cpu().numpy()
|
||||
act = self._env.action_manager.action[mask].cpu().numpy()
|
||||
for o, a in zip(obs, act):
|
||||
self._writer.writerow(o.tolist() + a.tolist())
|
||||
|
||||
def close(self):
|
||||
self._file.close()
|
||||
|
||||
The term receives the full ``cfg`` object so it can read any values you
|
||||
put in ``cfg.params``.
|
||||
|
||||
|
||||
Registration
|
||||
------------
|
||||
|
||||
Add the term to the ``recorders`` dictionary on your environment config:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from mjlab.managers import RecorderTermCfg
|
||||
|
||||
@dataclass
|
||||
class MyEnvCfg(SomeTaskEnvCfg):
|
||||
recorders: dict = field(default_factory=lambda: {
|
||||
"csv": RecorderTermCfg(
|
||||
func=CsvRecorder,
|
||||
params={"path": "rollout.csv"},
|
||||
)
|
||||
})
|
||||
|
||||
Multiple terms can be registered under different keys and run together.
|
||||
|
||||
.. note::
|
||||
|
||||
``func`` must be a :class:`~mjlab.managers.RecorderTerm` subclass.
|
||||
Function-based terms are not supported because recorder terms are
|
||||
stateful.
|
||||
@@ -0,0 +1,75 @@
|
||||
.. _research:
|
||||
|
||||
Research
|
||||
========
|
||||
|
||||
Citing mjlab
|
||||
------------
|
||||
|
||||
If you use mjlab in your research, please cite:
|
||||
|
||||
.. code-block:: bibtex
|
||||
|
||||
@article{Zakka_mjlab_A_Lightweight_2026,
|
||||
author = {Zakka, Kevin and Liao, Qiayuan and Yi, Brent and Le Lay, Louis and Sreenath, Koushil and Abbeel, Pieter},
|
||||
title = {{mjlab: A Lightweight Framework for GPU-Accelerated Robot Learning}},
|
||||
url = {https://arxiv.org/abs/2601.22074},
|
||||
year = {2026}
|
||||
}
|
||||
|
||||
Publications
|
||||
------------
|
||||
|
||||
Papers that use mjlab. To add your work, open a pull request or post in
|
||||
`Show and Tell <https://github.com/mujocolab/mjlab/discussions/categories/show-and-tell>`_.
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 60 30 10
|
||||
|
||||
* - Title
|
||||
- Authors
|
||||
- Year
|
||||
* - `HUSKY: Humanoid Skateboarding System via Physics-Aware Whole-Body Control
|
||||
<https://arxiv.org/abs/2602.03205>`_
|
||||
- Han, Wang, Zhang, Liu, Luo, Bai, Li
|
||||
- 2026
|
||||
* - `DynaRetarget: Dynamically-Feasible Retargeting using Sampling-Based
|
||||
Trajectory Optimization <https://arxiv.org/abs/2602.06827>`_
|
||||
- Dhedin, Taouil, Omar, Yu, Tao, Dai, Khadiv
|
||||
- 2026
|
||||
* - `CLOT: Closed-Loop Global Motion Tracking for Whole-Body Humanoid
|
||||
Teleoperation <https://arxiv.org/abs/2602.15060>`_
|
||||
- Zhu, Cai, Yang, Ren, Xie, Wang, Wu, et al.
|
||||
- 2026
|
||||
|
||||
Projects
|
||||
--------
|
||||
|
||||
Projects built on mjlab. To add yours, open a pull request or post in
|
||||
`Show and Tell <https://github.com/mujocolab/mjlab/discussions/categories/show-and-tell>`_.
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 35 65
|
||||
|
||||
* - Project
|
||||
- Description
|
||||
* - `menloresearch/asimov-mjlab <https://github.com/menloresearch/asimov-mjlab>`_
|
||||
- Locomotion fork for the Asimov bipedal robot.
|
||||
* - `Nagi-ovo/mjlab-homierl <https://github.com/Nagi-ovo/mjlab-homierl>`_
|
||||
- H1 locomotion across multiple tasks with robustness to upper body disturbances.
|
||||
* - `MyoHub/mjlab_myosuite <https://github.com/MyoHub/mjlab_myosuite>`_
|
||||
- Musculoskeletal simulation integration with MyoSuite.
|
||||
* - `MarcDcls/mjlab_upkie <https://github.com/MarcDcls/mjlab_upkie>`_
|
||||
- Velocity control for the Upkie wheeled biped.
|
||||
* - `unitreerobotics/unitree_rl_mjlab <https://github.com/unitreerobotics/unitree_rl_mjlab>`_
|
||||
- Official Unitree RL environments for Go2, G1, and H1\_2.
|
||||
* - `pal-robotics/pal_mjlab <https://github.com/pal-robotics/pal_mjlab>`_
|
||||
- PAL Robotics robots and tasks.
|
||||
* - `Msornerrrr/in-hand-rotation-mjlab <https://github.com/Msornerrrr/in-hand-rotation-mjlab>`_
|
||||
- Sim to real RL for in hand cube rotation with the LEAP Hand.
|
||||
* - `project-instinct/InstinctMJ <https://github.com/project-instinct/InstinctMJ>`_
|
||||
- mjlab version of Project-Instinct, a whole-body control toolchain to study Instinct-Level intelligence.
|
||||
* - `lzyang2000/twist2_mjlab <https://github.com/lzyang2000/twist2_mjlab>`_
|
||||
- mjlab port of `TWIST2 <https://arxiv.org/abs/2511.02832>`_.
|
||||
@@ -0,0 +1,102 @@
|
||||
.. _rewards:
|
||||
|
||||
Rewards
|
||||
=======
|
||||
|
||||
Rewards are the training signal that shapes policy behavior. Each reward
|
||||
term is a function that returns a per-environment scalar every step. The
|
||||
reward manager computes a weighted sum of all terms and returns it to
|
||||
the training framework.
|
||||
|
||||
Each term is registered by name with a ``RewardTermCfg`` that carries
|
||||
the callable and a ``weight``. Negative weights produce penalties.
|
||||
Additional keyword arguments are supplied through ``params``.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.envs.mdp import rewards
|
||||
from mjlab.managers.reward_manager import RewardTermCfg
|
||||
from mjlab.managers.scene_entity_config import SceneEntityCfg
|
||||
|
||||
rewards_cfg = {
|
||||
"alive": RewardTermCfg(func=rewards.is_alive, weight=1.0),
|
||||
"joint_torques": RewardTermCfg(
|
||||
func=rewards.joint_torques_l2,
|
||||
weight=-1e-4,
|
||||
params={"asset_cfg": SceneEntityCfg("robot")},
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
Built-in reward functions
|
||||
-------------------------
|
||||
|
||||
The functions below are available in ``mjlab.envs.mdp.rewards`` and are
|
||||
shared across tasks. Individual tasks also define their own reward
|
||||
functions tailored to the task objective (e.g. velocity tracking for
|
||||
locomotion). All reward functions return a tensor of shape
|
||||
``[num_envs]``.
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 28 72
|
||||
|
||||
* - Function
|
||||
- Description
|
||||
* - ``is_alive``
|
||||
- Returns ``1.0`` for environments that have not terminated this
|
||||
step. Use with a positive weight as a survival bonus.
|
||||
* - ``is_terminated``
|
||||
- Returns ``1.0`` for environments that terminated due to a
|
||||
non-timeout condition. Use with a negative weight to penalize
|
||||
failure.
|
||||
* - ``joint_torques_l2``
|
||||
- Sum of squared actuator forces. Penalizes energy-intensive
|
||||
actions.
|
||||
* - ``joint_vel_l2``
|
||||
- Sum of squared joint velocities.
|
||||
* - ``joint_acc_l2``
|
||||
- Sum of squared joint accelerations.
|
||||
* - ``action_rate_l2``
|
||||
- Sum of squared differences between the current and previous
|
||||
action. Penalizes rapid changes in the policy output.
|
||||
* - ``action_acc_l2``
|
||||
- Sum of squared second-order action differences. Penalizes
|
||||
high-frequency jitter in the action signal.
|
||||
* - ``joint_pos_limits``
|
||||
- Penalty for joint positions exceeding the soft limits. Zero
|
||||
when all joints are within limits.
|
||||
* - ``posture`` *(class)*
|
||||
- Exponential kernel measuring deviation from the default joint
|
||||
positions: ``exp(-mean(error^2 / std^2))``.
|
||||
* - ``electrical_power_cost`` *(class)*
|
||||
- Sum of positive mechanical power consumed by actuators.
|
||||
Regenerative power is not penalized.
|
||||
* - ``flat_orientation_l2``
|
||||
- Sum of squares of the x and y components of the projected
|
||||
gravity vector in the base frame. Zero when perfectly upright.
|
||||
|
||||
|
||||
Reward scaling by dt
|
||||
--------------------
|
||||
|
||||
``ManagerBasedRlEnvCfg.scale_rewards_by_dt`` is ``True`` by default.
|
||||
When enabled, the reward manager multiplies each term by the environment
|
||||
step duration before accumulating it. This makes episodic reward totals
|
||||
invariant to simulation frequency: a task running at 50 Hz produces the
|
||||
same expected episode return as the same task at 200 Hz, because each
|
||||
step contributes proportionally less to the total.
|
||||
|
||||
Per-term episodic sums are logged as ``Episode_Reward/<term_name>`` and
|
||||
are always divided by the episode duration, giving a reward rate that is
|
||||
comparable across runs with different episode lengths.
|
||||
|
||||
|
||||
Writing custom reward functions
|
||||
-------------------------------
|
||||
|
||||
A reward function accepts ``env`` as its first argument and returns a
|
||||
``[num_envs]`` tensor. Additional parameters are declared as function
|
||||
arguments and supplied via ``RewardTermCfg(params={...})``. When a term
|
||||
needs to cache setup work or maintain per-episode state, implement it as
|
||||
a class. See :ref:`env-config-term-pattern` for the general pattern.
|
||||
@@ -0,0 +1,175 @@
|
||||
.. _scene:
|
||||
|
||||
Scene
|
||||
=====
|
||||
|
||||
The scene merges entities, terrain, and sensors into a single
|
||||
simulation. ``SceneCfg`` describes the contents of the world, and the
|
||||
``Scene`` class handles MJCF composition, compilation, and runtime
|
||||
state management.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.scene import SceneCfg
|
||||
from mjlab.terrains import TerrainEntityCfg
|
||||
|
||||
# A robot on a flat ground plane with 4096 parallel environments.
|
||||
scene_cfg = SceneCfg(
|
||||
num_envs=4096,
|
||||
env_spacing=2.5,
|
||||
terrain=TerrainEntityCfg(terrain_type="plane"),
|
||||
entities={"robot": robot_cfg},
|
||||
)
|
||||
|
||||
A scene with procedural terrain, sensors, and multiple entities:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.scene import SceneCfg
|
||||
from mjlab.terrains import TerrainEntityCfg
|
||||
from mjlab.terrains.config import ROUGH_TERRAINS_CFG
|
||||
from mjlab.sensor import RayCastSensorCfg, ContactSensorCfg
|
||||
|
||||
scene_cfg = SceneCfg(
|
||||
num_envs=4096,
|
||||
terrain=TerrainEntityCfg(
|
||||
terrain_type="generator",
|
||||
terrain_generator=ROUGH_TERRAINS_CFG,
|
||||
max_init_terrain_level=5,
|
||||
),
|
||||
entities={
|
||||
"robot": robot_cfg,
|
||||
"cube": cube_cfg,
|
||||
},
|
||||
sensors=(
|
||||
RayCastSensorCfg(name="terrain_scan", ...),
|
||||
ContactSensorCfg(name="feet_contact", ...),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
Composition
|
||||
-----------
|
||||
|
||||
The scene starts from a root ``MjSpec`` and
|
||||
`attaches <https://mujoco.readthedocs.io/en/stable/python.html#attachment>`_
|
||||
each entity's spec into it with a unique name prefix. A robot entity named ``"robot"`` has all its
|
||||
internal MuJoCo elements (bodies, joints, geoms, actuators, sensors)
|
||||
prefixed with ``robot/``, so ``base_link`` becomes ``robot/base_link``,
|
||||
``joint0`` becomes ``robot/joint0``, and so on. Prefixing prevents name
|
||||
collisions when multiple entities share element names and provides a
|
||||
consistent namespace for observation and reward terms.
|
||||
|
||||
Terrain, when present, is attached without a prefix (its elements live
|
||||
in the global namespace). Sensors are added after entities and can
|
||||
reference entity elements by their prefixed names.
|
||||
|
||||
``scene.compile()`` converts the composed ``MjSpec`` into a single
|
||||
``MjModel``. The ``Simulation`` class then uploads this model to the
|
||||
GPU via MuJoCo Warp. After the simulation is created,
|
||||
``scene.initialize()`` resolves each entity's element indices into the
|
||||
compiled model, allocates state buffers, and sets up GPU rendering
|
||||
resources for any camera or raycast sensors.
|
||||
|
||||
``scene.to_zip(path)`` exports the compiled model as a ``.zip`` file
|
||||
for offline inspection in the standalone MuJoCo viewer. Each entity's
|
||||
initial state keyframe is merged into the export, so the model opens
|
||||
in its default pose.
|
||||
|
||||
At runtime, entities and sensors are accessible by name:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
robot = env.scene["robot"] # Entity
|
||||
scan = env.scene["terrain_scan"] # Sensor
|
||||
contact = env.scene["feet_contact"] # Sensor
|
||||
|
||||
robot.data.joint_pos # [B, num_joints]
|
||||
scan.data.distances # [B, N]
|
||||
contact.data.force # [B, N, 3]
|
||||
|
||||
Builtin sensors defined in an entity's XML are auto-discovered during
|
||||
composition and accessible with the entity name prefix:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
imu = env.scene["robot/trunk_imu"] # Auto-discovered sensor
|
||||
|
||||
|
||||
Environment origins
|
||||
-------------------
|
||||
|
||||
Each environment in MuJoCo Warp is an independent world with its own
|
||||
state. Environments do not share physical space and cannot interact with
|
||||
each other. Environment origins exist for two purposes: spreading
|
||||
entities across the world for visualization (so the viewer shows robots
|
||||
side by side rather than stacked at the origin), and for locomotion
|
||||
tasks with procedural terrain, placing each environment at a specific
|
||||
sub-terrain patch.
|
||||
|
||||
**Flat terrain.** Origins form a regular grid centered at the world
|
||||
origin with ``env_spacing`` meters between neighbors.
|
||||
|
||||
**Procedural terrain.** The terrain generator produces a
|
||||
``num_rows x num_cols`` grid of sub-terrain patches, each with its own
|
||||
center point. Each environment is assigned to one patch, and the
|
||||
terrain curriculum system moves environments to harder patches as
|
||||
performance improves. See :ref:`terrain` for details.
|
||||
|
||||
.. note::
|
||||
|
||||
All environments currently share the same ``MjModel`` (identical
|
||||
meshes, geometries, and kinematic trees). Heterogeneous simulation,
|
||||
where different worlds can have different meshes or geometries, is
|
||||
`in progress in MuJoCo Warp <https://github.com/google-deepmind/mujoco_warp/pull/1009>`_.
|
||||
mjlab will support this once it lands upstream.
|
||||
|
||||
Reset event terms read ``scene.env_origins`` to position entities:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# Inside a reset event term.
|
||||
robot.write_root_pose_to_sim(
|
||||
default_root_pose + env_origins[env_ids]
|
||||
)
|
||||
|
||||
Each origin is marked with an invisible sphere site (geom group 4) that
|
||||
appears in the MuJoCo viewer when group 4 is enabled, useful for
|
||||
verifying placement during development.
|
||||
|
||||
|
||||
Custom spec editing
|
||||
-------------------
|
||||
|
||||
Most scenes are fully described by their entities, terrain, and sensors.
|
||||
Occasionally a modification spans multiple entities. A tendon connecting
|
||||
a ceiling gantry to the robot, for example, cannot be defined inside
|
||||
either entity's MJCF because it references sites from both.
|
||||
|
||||
The ``spec_fn`` callback on ``SceneCfg`` handles this case. It receives
|
||||
the fully composed ``MjSpec`` after all entities and sensors have been
|
||||
attached with their prefixed names, but before compilation:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import mujoco
|
||||
|
||||
def add_gantry(spec: mujoco.MjSpec):
|
||||
spec.worldbody.add_site(name="gantry", pos=(0, 0, 2))
|
||||
for side in ["left", "right"]:
|
||||
tendon = spec.add_tendon(
|
||||
name=f"{side}_rope",
|
||||
limited=True,
|
||||
range=(0, 1),
|
||||
)
|
||||
tendon.wrap_site("gantry")
|
||||
tendon.wrap_site(f"robot/{side}_hook")
|
||||
|
||||
scene_cfg = SceneCfg(
|
||||
entities={"robot": robot_cfg},
|
||||
spec_fn=add_gantry,
|
||||
)
|
||||
|
||||
Other common uses include global equality constraints, custom
|
||||
visualization geometry, and any modification that requires access to the
|
||||
fully composed scene.
|
||||
@@ -0,0 +1,455 @@
|
||||
.. _sensors:
|
||||
|
||||
Sensors
|
||||
=======
|
||||
|
||||
As described in :ref:`entity`, sensors sit between ``EntityData`` and
|
||||
raw simulation arrays in mjlab's data access hierarchy. At their
|
||||
simplest, they wrap MuJoCo sensor primitives with a clean interface
|
||||
that maps to real robot hardware. Beyond wrapping, they are a general
|
||||
abstraction for transforming simulation data into structured outputs:
|
||||
``ContactSensor`` aggregates contact pairs with reduction and air time
|
||||
tracking, ``RayCastSensor`` performs GPU-accelerated terrain scanning,
|
||||
``CameraSensor`` renders RGB and depth images on the GPU, and the base
|
||||
``Sensor`` class can be subclassed for custom measurement logic.
|
||||
|
||||
Sensors are configured at the **scene level**, not on individual entities. A
|
||||
sensor can reference an entity element (a contact sensor on the robot's
|
||||
feet, an accelerometer attached to a body site), but it can also be
|
||||
independent of any entity entirely. This is why sensors live in
|
||||
``SceneCfg`` rather than ``EntityCfg``.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.sensor import (
|
||||
BuiltinSensorCfg, ContactSensorCfg, ContactMatch, ObjRef,
|
||||
)
|
||||
|
||||
# A robot with an IMU accelerometer and foot contact detection.
|
||||
scene_cfg = SceneCfg(
|
||||
entities={"robot": robot_cfg},
|
||||
sensors=(
|
||||
BuiltinSensorCfg(
|
||||
name="imu_acc",
|
||||
sensor_type="accelerometer",
|
||||
obj=ObjRef(type="site", name="imu_site", entity="robot"),
|
||||
),
|
||||
ContactSensorCfg(
|
||||
name="feet_contact",
|
||||
primary=ContactMatch(
|
||||
mode="geom", pattern=r".*_foot$", entity="robot",
|
||||
),
|
||||
secondary=ContactMatch(mode="body", pattern="terrain"),
|
||||
fields=("found", "force"),
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
# Access at runtime.
|
||||
imu = env.scene["robot/imu_acc"].data # [B, 3] acceleration
|
||||
feet = env.scene["feet_contact"].data # ContactData
|
||||
feet.found # [B, P] contact count per foot
|
||||
feet.force # [B, P, 3] contact force per foot
|
||||
|
||||
mjlab provides four sensor types: ``BuiltinSensor`` for native MuJoCo
|
||||
measurements, ``ContactSensor`` for structured contact detection,
|
||||
``RayCastSensor`` for GPU-accelerated raycasting, and ``CameraSensor``
|
||||
for RGB-D rendering. The base ``Sensor`` class can be subclassed for
|
||||
custom measurement logic; see `Extending: custom sensors`_ below.
|
||||
|
||||
|
||||
BuiltinSensor
|
||||
-------------
|
||||
|
||||
``BuiltinSensor`` wraps MuJoCo's native sensor types. Each sensor is
|
||||
attached to a MuJoCo element (site, joint, body, etc.) via ``ObjRef``
|
||||
and returns a ``torch.Tensor`` with shape ``[num_envs, dim]`` where
|
||||
``dim`` depends on the sensor type (3 for vectors, 4 for quaternions,
|
||||
1 for scalars).
|
||||
|
||||
+-----------+----------------------------------------------------------------------------------------------------------------------------------------------------+
|
||||
| Category | Available Sensors |
|
||||
+===========+====================================================================================================================================================+
|
||||
| **Site** | ``accelerometer``, ``velocimeter``, ``gyro``, ``force``, ``torque``, ``magnetometer``, ``rangefinder`` |
|
||||
+-----------+----------------------------------------------------------------------------------------------------------------------------------------------------+
|
||||
| **Joint** | ``jointpos``, ``jointvel``, ``jointlimitpos``, ``jointlimitvel``, ``jointlimitfrc``, ``jointactuatorfrc`` |
|
||||
+-----------+----------------------------------------------------------------------------------------------------------------------------------------------------+
|
||||
| **Frame** | ``framepos``, ``framequat``, ``framexaxis``, ``frameyaxis``, ``framezaxis``, ``framelinvel``, ``frameangvel``, ``framelinacc``, ``frameangacc`` |
|
||||
+-----------+----------------------------------------------------------------------------------------------------------------------------------------------------+
|
||||
| **Other** | ``actuatorpos``, ``actuatorvel``, ``actuatorfrc``, ``subtreecom``, ``subtreelinvel``, ``subtreeangmom``, ``clock``, ``e_potential``, ``e_kinetic`` |
|
||||
+-----------+----------------------------------------------------------------------------------------------------------------------------------------------------+
|
||||
|
||||
``ObjRef`` identifies which MuJoCo element the sensor attaches to. The
|
||||
``entity`` field scopes the lookup to a specific entity's namespace, and
|
||||
the sensor name is auto-prefixed accordingly (e.g., ``"imu_acc"`` on
|
||||
entity ``"robot"`` becomes ``"robot/imu_acc"``).
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.sensor import BuiltinSensorCfg, ObjRef
|
||||
|
||||
# Accelerometer attached to a site.
|
||||
BuiltinSensorCfg(
|
||||
name="imu_acc",
|
||||
sensor_type="accelerometer",
|
||||
obj=ObjRef(type="site", name="imu_site", entity="robot"),
|
||||
)
|
||||
|
||||
# Joint limit sensor with output clamping.
|
||||
BuiltinSensorCfg(
|
||||
name="knee_limit",
|
||||
sensor_type="jointlimitpos",
|
||||
obj=ObjRef(type="joint", name="knee_joint", entity="robot"),
|
||||
cutoff=0.1,
|
||||
)
|
||||
|
||||
# Relative frame position (end-effector w.r.t. base).
|
||||
BuiltinSensorCfg(
|
||||
name="ee_pos",
|
||||
sensor_type="framepos",
|
||||
obj=ObjRef(type="body", name="end_effector", entity="robot"),
|
||||
ref=ObjRef(type="body", name="base", entity="robot"),
|
||||
)
|
||||
|
||||
|
||||
Auto-discovery
|
||||
^^^^^^^^^^^^^^
|
||||
|
||||
Sensors already defined in an entity's XML are automatically discovered
|
||||
during scene composition and prefixed with the entity name. There is no
|
||||
need to create a ``BuiltinSensorCfg`` for these.
|
||||
|
||||
.. code-block:: xml
|
||||
|
||||
<!-- In robot.xml -->
|
||||
<sensor>
|
||||
<accelerometer name="trunk_imu" site="imu_site"/>
|
||||
<jointpos name="hip_sensor" joint="hip_joint"/>
|
||||
</sensor>
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# Access by prefixed name.
|
||||
imu = env.scene["robot/trunk_imu"]
|
||||
hip = env.scene["robot/hip_sensor"]
|
||||
|
||||
|
||||
ContactSensor
|
||||
-------------
|
||||
|
||||
Each physics step, MuJoCo produces a flat, unstructured list of contact
|
||||
pairs across the entire scene. A single foot geom might generate several
|
||||
simultaneous contacts with the ground, interleaved with contacts from
|
||||
other entities. ``ContactSensor`` filters this raw list to the pairs you
|
||||
care about, reduces multiple contacts per element down to a fixed count,
|
||||
and packages the result into clean, batched tensors your policy can
|
||||
consume directly. It builds on MuJoCo's native
|
||||
`contact sensor <https://mujoco.readthedocs.io/en/stable/XMLreference.html#sensor-contact>`_.
|
||||
|
||||
Primary and secondary
|
||||
^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
Contacts are pairwise: you typically want to know "did the robot's feet
|
||||
touch the terrain?", not just "did something touch something."
|
||||
``primary`` defines the elements you are measuring (the feet).
|
||||
``secondary`` optionally restricts what they are contacting (the
|
||||
terrain). When ``secondary`` is ``None``, any contact with a primary
|
||||
element counts.
|
||||
|
||||
Each side is specified with a ``ContactMatch``. The ``mode`` selects the
|
||||
MuJoCo element type (``"geom"``, ``"body"``, or ``"subtree"``) and the
|
||||
``pattern`` accepts a regex or tuple of regexes matched against element
|
||||
names within the entity.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.sensor import ContactSensorCfg, ContactMatch
|
||||
|
||||
# Foot geoms contacting the terrain body.
|
||||
ContactSensorCfg(
|
||||
name="feet_ground",
|
||||
primary=ContactMatch(
|
||||
mode="geom", pattern=r".*_foot$", entity="robot",
|
||||
),
|
||||
secondary=ContactMatch(mode="body", pattern="terrain"),
|
||||
fields=("found", "force"),
|
||||
)
|
||||
|
||||
# Self-collision: pelvis subtree against itself.
|
||||
ContactSensorCfg(
|
||||
name="self_collision",
|
||||
primary=ContactMatch(
|
||||
mode="subtree", pattern="pelvis", entity="robot",
|
||||
),
|
||||
secondary=ContactMatch(
|
||||
mode="subtree", pattern="pelvis", entity="robot",
|
||||
),
|
||||
fields=("found",),
|
||||
)
|
||||
|
||||
Output shape
|
||||
^^^^^^^^^^^^
|
||||
|
||||
A pattern like ``r".*_foot$"`` resolves to ``P`` primary elements (e.g.
|
||||
four feet on a quadruped). Each primary becomes one column on the
|
||||
per-contact axis of the output tensors:
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 35 25 40
|
||||
|
||||
* - Field group
|
||||
- Shape
|
||||
- Notes
|
||||
* - Per-contact
|
||||
(``found``, ``force``, ``torque``, ``dist``, ``pos``, ``normal``,
|
||||
``tangent``)
|
||||
- ``[B, P * num_slots, ...]``
|
||||
- Primary-major: indices
|
||||
``[i * num_slots : (i + 1) * num_slots]`` belong to primary ``i``.
|
||||
* - Per-primary
|
||||
(``current_air_time``, ``last_air_time``,
|
||||
``current_contact_time``, ``last_contact_time``)
|
||||
- ``[B, P]``
|
||||
- Air-time fields are accumulated per primary, reducing across
|
||||
slots (any slot in contact counts as the primary in contact).
|
||||
|
||||
With the default ``num_slots=1`` the two shape families coincide
|
||||
(``N == P``), which is why most code can treat both as ``[B, P, ...]``.
|
||||
|
||||
Use :attr:`mjlab.sensor.contact_sensor.ContactSensor.primary_names` to
|
||||
recover the index-to-name mapping after pattern expansion:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
sensor = env.scene["feet_contact"]
|
||||
sensor.primary_names # ["FR_foot", "FL_foot", "RR_foot", "RL_foot"]
|
||||
sensor.data.current_air_time[:, 0] # air time for FR_foot
|
||||
|
||||
Reduction
|
||||
^^^^^^^^^
|
||||
|
||||
A single primary element can have many simultaneous contacts with the
|
||||
secondary (e.g. a flat foot resting on rough terrain has multiple
|
||||
contact points). The ``reduce`` mode collapses those raw contacts down
|
||||
to ``num_slots`` representative contacts:
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 20 80
|
||||
|
||||
* - Mode
|
||||
- Behavior
|
||||
* - ``"none"``
|
||||
- Fast, non-deterministic selection of up to ``num_slots`` contacts.
|
||||
* - ``"mindist"``
|
||||
- Keep the deepest ``num_slots`` contacts.
|
||||
* - ``"maxforce"``
|
||||
- Keep the strongest ``num_slots`` contacts by force magnitude.
|
||||
* - ``"netforce"``
|
||||
- Sum all contacts into a single net wrench at the force-weighted
|
||||
centroid. Always emits one slot per primary regardless of
|
||||
``num_slots``.
|
||||
|
||||
When to set ``num_slots > 1``
|
||||
"""""""""""""""""""""""""""""
|
||||
|
||||
Almost all configurations leave ``num_slots`` at its default of ``1``,
|
||||
because pattern expansion already produces one column per element of
|
||||
interest (one per foot, one per finger, one per body link). Increase
|
||||
``num_slots`` only when a single primary may have several physically
|
||||
distinct contact points that you want to inspect separately, for
|
||||
example:
|
||||
|
||||
- Computing a center of pressure on a flat foot from its corner
|
||||
contacts.
|
||||
- Reasoning about grasp quality from multiple fingertip-to-object
|
||||
contact points.
|
||||
- Detecting tipping by watching whether one corner of a contact patch
|
||||
loses contact.
|
||||
|
||||
In those cases pair ``num_slots`` with ``reduce`` in
|
||||
``{"mindist", "maxforce", "none"}``. With ``reduce="netforce"`` it has
|
||||
no effect.
|
||||
|
||||
.. note::
|
||||
|
||||
``num_slots`` is a ceiling on what the sensor will store, not a
|
||||
guarantee that MuJoCo will produce that many contacts. The collision
|
||||
detector caps the number of contact points generated for each
|
||||
geom-pair according to the geom types involved. For example a
|
||||
sphere-vs-plane pair produces at most one contact, and a box-vs-box
|
||||
pair produces at most four. Setting ``num_slots=8`` on a sphere
|
||||
primary against a plane secondary therefore leaves seven slots
|
||||
permanently zero. Check
|
||||
`MuJoCo's collision documentation <https://mujoco.readthedocs.io/en/stable/computation/index.html#collision-detection>`_
|
||||
for the per-pair limits.
|
||||
|
||||
Fields
|
||||
^^^^^^
|
||||
|
||||
The ``fields`` tuple selects which contact quantities to extract. Only
|
||||
requested fields are allocated; the rest are ``None`` on the output
|
||||
dataclass. Available fields are ``"found"``, ``"force"``,
|
||||
``"torque"``, ``"dist"``, ``"pos"``, ``"normal"``, and ``"tangent"``.
|
||||
|
||||
.. note::
|
||||
|
||||
``torque`` and the friction-tangent component of ``force`` are zero
|
||||
unless the contact pair has friction enabled, which requires
|
||||
``condim >= 3`` on at least one geom in the pair. With ``condim=1``
|
||||
(frictionless), the contact only produces a normal force. This is a
|
||||
physics property of the contact, not a sensor limitation.
|
||||
|
||||
Air time tracking
|
||||
^^^^^^^^^^^^^^^^^
|
||||
|
||||
Locomotion tasks often need to know when feet land and take off for
|
||||
gait rewards. Setting ``track_air_time=True`` enables per-primary
|
||||
timing. The sensor maintains four additional tensors on
|
||||
``ContactData``, each shaped ``[B, P]``: ``current_air_time``,
|
||||
``last_air_time``, ``current_contact_time``, and
|
||||
``last_contact_time``. Two helper methods provide edge detection for
|
||||
transition events:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
sensor = env.scene["feet_air"]
|
||||
first_contact = sensor.compute_first_contact(dt) # [B, P], True for primaries that just landed
|
||||
first_air = sensor.compute_first_air(dt) # [B, P], True for primaries that just took off
|
||||
|
||||
Air time is per primary even when ``num_slots > 1``: the sensor reduces
|
||||
``found`` across slots so that any slot in contact counts as the
|
||||
primary being in contact.
|
||||
|
||||
.. _contact-sensor-history:
|
||||
|
||||
History (decimation safe contacts)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
When using decimation (multiple physics substeps per policy step), a
|
||||
brief collision can occur and resolve entirely within the substep loop.
|
||||
By the time the policy reads the sensor, the contact is gone and
|
||||
``found`` reports zero. Setting ``history_length`` on the sensor config
|
||||
tells the sensor to keep a rolling buffer of the last *N* substeps for
|
||||
force, torque, and distance fields. The policy can then inspect the
|
||||
full history and decide whether a real contact occurred.
|
||||
|
||||
Set ``history_length`` equal to your decimation value so the buffer
|
||||
covers exactly one policy step:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
ContactSensorCfg(
|
||||
name="self_collision",
|
||||
primary=ContactMatch(mode="subtree", pattern="pelvis", entity="robot"),
|
||||
secondary=ContactMatch(mode="subtree", pattern="pelvis", entity="robot"),
|
||||
fields=("found", "force"),
|
||||
history_length=4, # matches decimation=4
|
||||
)
|
||||
|
||||
The history tensors live on ``ContactData`` alongside the regular
|
||||
fields:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
data = sensor.data
|
||||
data.force_history # [B, N, H, 3] (H = history_length)
|
||||
data.torque_history # [B, N, H, 3]
|
||||
data.dist_history # [B, N, H]
|
||||
|
||||
Index 0 is the most recent substep. To check whether any substep had
|
||||
a contact force above a threshold:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
force_mag = torch.norm(data.force_history, dim=-1) # [B, N, H]
|
||||
had_contact = (force_mag > 10.0).any(dim=1).any(dim=-1) # [B]
|
||||
|
||||
.. note::
|
||||
|
||||
``track_air_time=True`` already accumulates contact state across
|
||||
substeps for gait rewards, so feet ground sensors typically do not
|
||||
need ``history_length``. Use history for sensors where you need to
|
||||
detect brief collisions that would otherwise be missed (self
|
||||
collisions, illegal contact terminations).
|
||||
|
||||
|
||||
Output
|
||||
^^^^^^
|
||||
|
||||
``ContactData`` is a dataclass whose fields correspond to the
|
||||
``fields`` tuple on the config. Unrequested fields are ``None``.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@dataclass
|
||||
class ContactData:
|
||||
found: Tensor | None # [B, N] contact count
|
||||
force: Tensor | None # [B, N, 3]
|
||||
torque: Tensor | None # [B, N, 3]
|
||||
dist: Tensor | None # [B, N] penetration depth
|
||||
pos: Tensor | None # [B, N, 3] contact position
|
||||
normal: Tensor | None # [B, N, 3] surface normal
|
||||
tangent: Tensor | None # [B, N, 3]
|
||||
|
||||
# With track_air_time=True.
|
||||
current_air_time: Tensor | None
|
||||
last_air_time: Tensor | None
|
||||
current_contact_time: Tensor | None
|
||||
last_contact_time: Tensor | None
|
||||
|
||||
|
||||
RayCastSensor
|
||||
-------------
|
||||
|
||||
``RayCastSensor`` provides GPU-accelerated raycasting for terrain
|
||||
scanning and depth sensing. It supports grid and pinhole camera ray
|
||||
patterns with configurable alignment modes. See :ref:`raycast_sensor`
|
||||
for full documentation.
|
||||
|
||||
|
||||
RGB-D Camera
|
||||
------------
|
||||
|
||||
``CameraSensor`` renders RGB and depth images from MuJoCo cameras. See
|
||||
:ref:`rgbd_camera` for full documentation.
|
||||
|
||||
|
||||
Extending: custom sensors
|
||||
-------------------------
|
||||
|
||||
All sensors inherit from ``Sensor[T]``, a generic base class where
|
||||
``T`` is the data type returned by the ``data`` property (e.g.,
|
||||
``torch.Tensor`` for ``BuiltinSensor``, ``ContactData`` for
|
||||
``ContactSensor``).
|
||||
|
||||
The base class provides automatic per-step caching. The ``data``
|
||||
property calls ``_compute_data()`` on first access each step and
|
||||
caches the result. The cache is invalidated automatically when
|
||||
``update()`` or ``reset()`` is called, so multiple reads within the
|
||||
same step (from different observation or reward terms) pay the
|
||||
computation cost only once.
|
||||
|
||||
**Lifecycle methods:**
|
||||
|
||||
- ``edit_spec``: Add sensor elements to the MjSpec during scene
|
||||
construction.
|
||||
- ``initialize``: Post-compilation setup. Cache sensor indices,
|
||||
allocate buffers, resolve references.
|
||||
- ``update``: Called each physics step. Invalidates the data cache.
|
||||
Override to maintain per-step state (e.g., air time counters).
|
||||
- ``reset``: Called on environment reset. Invalidates the data cache.
|
||||
Override to clear per-environment state.
|
||||
- ``_compute_data``: Compute and return the sensor output. Called
|
||||
lazily by the ``data`` property when the cache is stale.
|
||||
|
||||
``ContactSensor`` and ``RayCastSensor`` are the most complete
|
||||
reference implementations for custom sensor development.
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
:hidden:
|
||||
|
||||
raycast_sensor
|
||||
rgbd_camera
|
||||
@@ -0,0 +1,299 @@
|
||||
.. _raycast_sensor:
|
||||
|
||||
RayCast Sensor
|
||||
==============
|
||||
|
||||
``RayCastSensor`` provides GPU-accelerated raycasting for terrain
|
||||
scanning, obstacle detection, and depth sensing. Rays are emitted from
|
||||
a frame attached to a body, site, or geom in the scene, and the sensor
|
||||
reports hit distances, world-space hit positions, and surface normals.
|
||||
|
||||
.. raw:: html
|
||||
|
||||
<video controls style="display: block; margin: 0 auto; max-width: 100%; height: auto;">
|
||||
<source src="../../_static/raycast_demo.mp4" type="video/mp4">
|
||||
</video>
|
||||
|
||||
|
||||
Quick start
|
||||
-----------
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.sensor import RayCastSensorCfg, GridPatternCfg, ObjRef
|
||||
|
||||
# Downward-facing grid for terrain height scanning.
|
||||
raycast_cfg = RayCastSensorCfg(
|
||||
name="terrain_scan",
|
||||
frame=ObjRef(type="body", name="base", entity="robot"),
|
||||
pattern=GridPatternCfg(size=(1.0, 1.0), resolution=0.1),
|
||||
max_distance=5.0,
|
||||
)
|
||||
|
||||
scene_cfg = SceneCfg(
|
||||
entities={"robot": robot_cfg},
|
||||
sensors=(raycast_cfg,),
|
||||
)
|
||||
|
||||
# Access at runtime.
|
||||
data = env.scene["terrain_scan"].data
|
||||
data.distances # [B, N] distance to hit, -1 if miss
|
||||
data.hit_pos_w # [B, N, 3] world-space hit positions
|
||||
data.normals_w # [B, N, 3] surface normals
|
||||
|
||||
|
||||
Ray patterns
|
||||
------------
|
||||
|
||||
Ray patterns define the spatial distribution and direction of rays
|
||||
emitted from the sensor frame.
|
||||
|
||||
.. grid:: 2
|
||||
|
||||
.. grid-item-card:: Grid pattern
|
||||
|
||||
Parallel rays in a 2D grid with fixed spatial resolution. The
|
||||
ground footprint does not change with sensor height because ray
|
||||
spacing is defined in world units (meters). The natural choice for
|
||||
height maps and terrain scanning.
|
||||
|
||||
.. raw:: html
|
||||
|
||||
<video autoplay loop muted playsinline style="width: 100%; height: auto;">
|
||||
<source src="../../_static/pattern_grid.mp4" type="video/mp4">
|
||||
</video>
|
||||
|
||||
.. grid-item-card:: Pinhole camera pattern
|
||||
|
||||
Diverging rays emitted from a single origin, analogous to a depth
|
||||
camera. The ground coverage increases with sensor
|
||||
height because the field of view is fixed in angular units.
|
||||
|
||||
.. raw:: html
|
||||
|
||||
<video autoplay loop muted playsinline style="width: 100%; height: auto;">
|
||||
<source src="../../_static/pattern_pinhole.mp4" type="video/mp4">
|
||||
</video>
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.sensor import GridPatternCfg, PinholeCameraPatternCfg
|
||||
|
||||
# Parallel grid: fixed footprint, height-invariant.
|
||||
grid = GridPatternCfg(
|
||||
size=(1.0, 1.0), # Grid dimensions in meters
|
||||
resolution=0.1, # Spacing between rays
|
||||
direction=(0.0, 0.0, -1.0), # Ray direction (down)
|
||||
)
|
||||
|
||||
# Pinhole: perspective projection, diverging rays.
|
||||
pinhole = PinholeCameraPatternCfg(
|
||||
width=16,
|
||||
height=12,
|
||||
fovy=45.0, # Vertical FOV in degrees
|
||||
)
|
||||
|
||||
# Pinhole from a MuJoCo camera definition.
|
||||
pinhole = PinholeCameraPatternCfg.from_mujoco_camera("robot/depth_cam")
|
||||
|
||||
# Pinhole from an intrinsic matrix.
|
||||
pinhole = PinholeCameraPatternCfg.from_intrinsic_matrix(
|
||||
intrinsic_matrix=[500, 0, 320, 0, 500, 240, 0, 0, 1],
|
||||
width=640,
|
||||
height=480,
|
||||
)
|
||||
|
||||
|
||||
Pattern comparison
|
||||
^^^^^^^^^^^^^^^^^^
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 20 40 40
|
||||
|
||||
* - Aspect
|
||||
- Grid
|
||||
- Pinhole
|
||||
* - Ray direction
|
||||
- Parallel
|
||||
- Diverging
|
||||
* - Spacing unit
|
||||
- Meters
|
||||
- Degrees (FOV)
|
||||
* - Height affects coverage
|
||||
- No
|
||||
- Yes
|
||||
* - Projection model
|
||||
- Orthographic
|
||||
- Perspective
|
||||
|
||||
|
||||
Frame attachment
|
||||
----------------
|
||||
|
||||
Rays are emitted from a frame in the scene specified via ``ObjRef``.
|
||||
The frame can be a body, site, or geom on any entity.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
frame = ObjRef(type="body", name="base", entity="robot")
|
||||
frame = ObjRef(type="site", name="scan_site", entity="robot")
|
||||
frame = ObjRef(type="geom", name="sensor_mount", entity="robot")
|
||||
|
||||
``exclude_parent_body`` (default ``True``) prevents rays from hitting
|
||||
the body to which the sensor is attached.
|
||||
|
||||
|
||||
Ray alignment
|
||||
-------------
|
||||
|
||||
The ``ray_alignment`` setting controls how rays orient relative to the
|
||||
attached frame when the body rotates.
|
||||
|
||||
.. raw:: html
|
||||
|
||||
<video autoplay loop muted playsinline
|
||||
style="display: block; margin: 0 auto; max-width: 100%; height: auto;">
|
||||
<source src="../../_static/ray_alignment_comparison.mp4" type="video/mp4">
|
||||
</video>
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 15 45 40
|
||||
|
||||
* - Mode
|
||||
- Description
|
||||
- Use case
|
||||
* - ``"base"``
|
||||
- Full position and rotation tracking
|
||||
- Body-mounted sensors
|
||||
* - ``"yaw"``
|
||||
- Follows yaw, ignores pitch and roll
|
||||
- Terrain height maps
|
||||
* - ``"world"``
|
||||
- Fixed world-frame direction
|
||||
- Gravity-aligned sensing
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
RayCastSensorCfg(
|
||||
name="height_scan",
|
||||
frame=ObjRef(type="body", name="base", entity="robot"),
|
||||
pattern=GridPatternCfg(size=(1.0, 1.0), resolution=0.1),
|
||||
ray_alignment="yaw",
|
||||
)
|
||||
|
||||
|
||||
Geom group filtering
|
||||
--------------------
|
||||
|
||||
MuJoCo assigns geoms to groups 0 through 5. Use
|
||||
``include_geom_groups`` to restrict which geoms rays can hit. This is
|
||||
useful for ignoring visual-only geoms or isolating terrain geometry.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
RayCastSensorCfg(
|
||||
name="terrain_only",
|
||||
frame=ObjRef(type="body", name="base", entity="robot"),
|
||||
pattern=GridPatternCfg(),
|
||||
include_geom_groups=(0, 1),
|
||||
)
|
||||
|
||||
|
||||
Output
|
||||
------
|
||||
|
||||
``RayCastData`` is a dataclass with shape annotations relative to
|
||||
``B`` (number of environments) and ``N`` (number of rays).
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@dataclass
|
||||
class RayCastData:
|
||||
distances: Tensor # [B, N] distance to hit, -1 if miss
|
||||
hit_pos_w: Tensor # [B, N, 3] world-space hit positions
|
||||
normals_w: Tensor # [B, N, 3] surface normals
|
||||
pos_w: Tensor # [B, 3] sensor frame position
|
||||
quat_w: Tensor # [B, 4] sensor frame orientation (w, x, y, z)
|
||||
|
||||
.. note::
|
||||
|
||||
Set ``debug_vis=True`` on the config to visualize ray hits at
|
||||
runtime.
|
||||
|
||||
|
||||
Examples
|
||||
--------
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.sensor import (
|
||||
RayCastSensorCfg, GridPatternCfg, PinholeCameraPatternCfg, ObjRef,
|
||||
)
|
||||
|
||||
# Dense height map for terrain-aware locomotion.
|
||||
height_scan = RayCastSensorCfg(
|
||||
name="height_scan",
|
||||
frame=ObjRef(type="body", name="base", entity="robot"),
|
||||
pattern=GridPatternCfg(
|
||||
size=(1.6, 1.0),
|
||||
resolution=0.1,
|
||||
direction=(0.0, 0.0, -1.0),
|
||||
),
|
||||
ray_alignment="yaw",
|
||||
max_distance=2.0,
|
||||
)
|
||||
|
||||
# Simulated depth camera using pinhole projection.
|
||||
depth_cam = RayCastSensorCfg(
|
||||
name="depth",
|
||||
frame=ObjRef(type="site", name="camera_site", entity="robot"),
|
||||
pattern=PinholeCameraPatternCfg.from_mujoco_camera("robot/depth_cam"),
|
||||
max_distance=10.0,
|
||||
)
|
||||
|
||||
# Forward-facing obstacle scan.
|
||||
obstacle_scan = RayCastSensorCfg(
|
||||
name="obstacle",
|
||||
frame=ObjRef(type="body", name="head", entity="robot"),
|
||||
pattern=GridPatternCfg(
|
||||
size=(0.5, 0.3),
|
||||
resolution=0.1,
|
||||
direction=(-1.0, 0.0, 0.0),
|
||||
),
|
||||
max_distance=3.0,
|
||||
include_geom_groups=(0,),
|
||||
)
|
||||
|
||||
|
||||
TerrainHeightSensor
|
||||
-------------------
|
||||
|
||||
``TerrainHeightSensor`` is a thin ``RayCastSensor`` subclass that adds
|
||||
per-frame vertical clearance to the sensor data. It computes
|
||||
``frame_z - hit_z`` for each ray, replaces misses with ``max_distance``,
|
||||
and reduces across rays per frame.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.sensor import TerrainHeightSensorCfg, RingPatternCfg, ObjRef
|
||||
|
||||
cfg = TerrainHeightSensorCfg(
|
||||
name="foot_height",
|
||||
frame=(
|
||||
ObjRef(type="site", name="left_foot", entity="robot"),
|
||||
ObjRef(type="site", name="right_foot", entity="robot"),
|
||||
),
|
||||
pattern=RingPatternCfg.single_ring(radius=0.04, num_samples=4),
|
||||
max_distance=1.0,
|
||||
include_geom_groups=(0,),
|
||||
)
|
||||
|
||||
# At runtime:
|
||||
sensor = env.scene["foot_height"]
|
||||
sensor.data.heights # [B, F] vertical clearance per foot
|
||||
sensor.data.distances # [B, N] raw ray distances (inherited)
|
||||
|
||||
The ``reduction`` config field controls how rays are aggregated within
|
||||
each frame: ``"min"`` (default), ``"max"``, or ``"mean"``.
|
||||
@@ -0,0 +1,219 @@
|
||||
.. _rgbd_camera:
|
||||
|
||||
RGB-D Camera
|
||||
============
|
||||
|
||||
``CameraSensor`` renders RGB and depth images on the GPU using MuJoCo
|
||||
Warp's ray-traced rendering pipeline. It can either wrap an existing
|
||||
MuJoCo camera defined in your XML or create a new one programmatically.
|
||||
|
||||
|
||||
Quick start
|
||||
-----------
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.sensor import CameraSensorCfg
|
||||
|
||||
# Wrap an existing MuJoCo camera from the robot's XML.
|
||||
cam = CameraSensorCfg(
|
||||
name="wrist_cam",
|
||||
camera_name="robot/wrist_camera",
|
||||
data_types=("rgb", "depth"),
|
||||
width=160,
|
||||
height=120,
|
||||
)
|
||||
|
||||
scene_cfg = SceneCfg(
|
||||
entities={"robot": robot_cfg},
|
||||
sensors=(cam,),
|
||||
)
|
||||
|
||||
# Access at runtime.
|
||||
data = env.scene["wrist_cam"].data
|
||||
data.rgb # [B, 120, 160, 3] uint8
|
||||
data.depth # [B, 120, 160, 1] float32
|
||||
|
||||
|
||||
Creating vs wrapping cameras
|
||||
-----------------------------
|
||||
|
||||
There are two ways to set up a camera sensor.
|
||||
|
||||
**Wrap an existing camera.** If your MJCF model already defines a
|
||||
camera, pass its name via ``camera_name``. The sensor uses the
|
||||
camera's position, orientation, and field of view from the model.
|
||||
You can optionally override ``fovy`` or switch to orthographic
|
||||
projection.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# Wrap the camera named "front_cam" in the robot's XML.
|
||||
CameraSensorCfg(
|
||||
name="front",
|
||||
camera_name="robot/front_cam",
|
||||
data_types=("rgb",),
|
||||
)
|
||||
|
||||
**Create a new camera.** When ``camera_name`` is ``None`` (the
|
||||
default), the sensor adds a new camera to the MjSpec during scene
|
||||
construction. Specify ``pos``, ``quat``, and optionally ``fovy`` to
|
||||
place it.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# Fixed overhead camera on the worldbody.
|
||||
CameraSensorCfg(
|
||||
name="overhead",
|
||||
pos=(0.0, 0.0, 2.0),
|
||||
quat=(0.0, 0.707, 0.707, 0.0),
|
||||
fovy=60.0,
|
||||
width=320,
|
||||
height=240,
|
||||
data_types=("rgb", "depth"),
|
||||
)
|
||||
|
||||
|
||||
Camera parameterization
|
||||
-----------------------
|
||||
|
||||
MuJoCo supports two ways to define a camera's projection, and both
|
||||
work with ``CameraSensor``. See the `MuJoCo camera documentation
|
||||
<https://mujoco.readthedocs.io/en/stable/XMLreference.html#body-camera>`_
|
||||
for full details.
|
||||
|
||||
**FOV-based.** The simpler approach. A single ``fovy`` (vertical field
|
||||
of view in degrees) combined with the image resolution defines the
|
||||
projection. This is the default when creating cameras
|
||||
programmatically via ``CameraSensorCfg``.
|
||||
|
||||
**Intrinsic-based.** For matching real camera hardware, MuJoCo cameras
|
||||
can be parameterized with ``sensorsize``, ``focal`` (or
|
||||
``focalpixel``), and ``principal`` (or ``principalpixel``). These
|
||||
fields are set in the MJCF XML and provide direct control over the
|
||||
intrinsic matrix. When intrinsic parameters are present, ``fovy`` is
|
||||
ignored by MuJoCo.
|
||||
|
||||
When wrapping an existing camera, the sensor inherits whichever
|
||||
parameterization the XML defines. When creating a new camera, the
|
||||
sensor uses ``fovy``. To use intrinsic parameters for a new camera,
|
||||
define it in your XML and wrap it with ``camera_name``.
|
||||
|
||||
.. note::
|
||||
|
||||
If you plan to randomize the field of view with domain
|
||||
randomization, use ``dr.cam_fovy`` for FOV-based cameras or
|
||||
``dr.cam_intrinsic`` for intrinsic-based cameras. Randomizing
|
||||
``cam_fovy`` has no effect on cameras that use intrinsic parameters.
|
||||
|
||||
|
||||
Body-mounted cameras
|
||||
--------------------
|
||||
|
||||
Set ``parent_body`` to attach a new camera to a specific body rather
|
||||
than the worldbody. The ``pos`` and ``quat`` are then relative to the
|
||||
parent body frame, so the camera moves with the body.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# Camera mounted on the robot's end-effector.
|
||||
CameraSensorCfg(
|
||||
name="ee_cam",
|
||||
parent_body="robot/link_6",
|
||||
pos=(0.0, 0.0, 0.05),
|
||||
quat=(1.0, 0.0, 0.0, 0.0),
|
||||
fovy=45.0,
|
||||
width=160,
|
||||
height=120,
|
||||
data_types=("rgb", "depth"),
|
||||
)
|
||||
|
||||
|
||||
Data types
|
||||
----------
|
||||
|
||||
The ``data_types`` tuple selects which image modalities to render.
|
||||
Only requested types are allocated; the other field on
|
||||
``CameraSensorData`` is ``None``.
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 15 20 65
|
||||
|
||||
* - Type
|
||||
- Shape
|
||||
- Description
|
||||
* - ``"rgb"``
|
||||
- ``[B, H, W, 3]`` uint8
|
||||
- Color image. Rendered as packed ABGR uint32 by MuJoCo Warp,
|
||||
then unpacked to RGB channels on the GPU.
|
||||
* - ``"depth"``
|
||||
- ``[B, H, W, 1]`` float32
|
||||
- Depth image. Values are distances from the camera plane.
|
||||
* - ``"segmentation"``
|
||||
- ``[B, H, W, 2]`` int32
|
||||
- Typed segmentation. Channel 0 stores object IDs and channel 1 stores
|
||||
MuJoCo object types. Background pixels are ``(-1, -1)``.
|
||||
|
||||
|
||||
Render settings
|
||||
---------------
|
||||
|
||||
All camera sensors in a scene must share identical values for
|
||||
``use_textures``, ``use_shadows``, and ``enabled_geom_groups``. This
|
||||
is a constraint of the underlying MuJoCo Warp rendering system, which
|
||||
uses a single ``RenderContext`` for all cameras. Mismatched settings
|
||||
raise a ``ValueError`` at scene construction.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# These two cameras must agree on render settings.
|
||||
cam_a = CameraSensorCfg(
|
||||
name="cam_a",
|
||||
camera_name="robot/front_cam",
|
||||
use_textures=True,
|
||||
use_shadows=False,
|
||||
enabled_geom_groups=(0, 1, 2),
|
||||
data_types=("rgb",),
|
||||
)
|
||||
cam_b = CameraSensorCfg(
|
||||
name="cam_b",
|
||||
camera_name="robot/wrist_cam",
|
||||
use_textures=True, # Must match cam_a
|
||||
use_shadows=False, # Must match cam_a
|
||||
enabled_geom_groups=(0, 1, 2), # Must match cam_a
|
||||
data_types=("depth",),
|
||||
)
|
||||
|
||||
|
||||
Output
|
||||
------
|
||||
|
||||
``CameraSensorData`` is a dataclass with one field per data type.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@dataclass
|
||||
class CameraSensorData:
|
||||
rgb: Tensor | None # [B, H, W, 3] uint8
|
||||
depth: Tensor | None # [B, H, W, 1] float32
|
||||
segmentation: Tensor | None # [B, H, W, 2] int32
|
||||
|
||||
By default, the returned tensors are zero-copy views into the render
|
||||
buffer. Set ``clone_data=True`` on the config if you modify them in
|
||||
place, to avoid corrupting the shared buffer.
|
||||
|
||||
|
||||
Visualization in Viser
|
||||
----------------------
|
||||
|
||||
The Viser viewer automatically discovers all ``CameraSensor`` instances
|
||||
in the scene and displays their RGB and depth outputs as live image
|
||||
panels in the GUI sidebar. A camera frustum is rendered in the 3D
|
||||
viewport showing the camera's position, orientation, and field of view.
|
||||
Depth images include an interactive scale slider for adjusting the
|
||||
visualization range.
|
||||
|
||||
.. image:: ../_static/viser_camera_pane.png
|
||||
:align: center
|
||||
:alt: Viser viewer showing camera image panels and frustum visualization
|
||||
@@ -0,0 +1,70 @@
|
||||
.. _terminations:
|
||||
|
||||
Terminations
|
||||
============
|
||||
|
||||
Termination terms define when an episode ends. Each term is a function
|
||||
that returns a boolean per-environment tensor every step. The
|
||||
termination manager aggregates all terms and reports the result to the
|
||||
training framework as either a terminal failure or a truncation.
|
||||
|
||||
Each term is registered by name with a ``TerminationTermCfg``. Setting
|
||||
``time_out=True`` marks the condition as a truncation rather than a
|
||||
terminal failure. Truncations map to the ``truncated`` signal in the
|
||||
Gym interface; failures map to ``terminated``. This distinction matters
|
||||
for value bootstrapping: the agent should estimate future value beyond
|
||||
a truncation but not beyond a failure.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.envs.mdp import terminations
|
||||
from mjlab.managers.termination_manager import TerminationTermCfg
|
||||
|
||||
terminations_cfg = {
|
||||
"time_out": TerminationTermCfg(
|
||||
func=terminations.time_out, time_out=True,
|
||||
),
|
||||
"fallen": TerminationTermCfg(
|
||||
func=terminations.bad_orientation,
|
||||
params={"limit_angle": 1.0},
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
Built-in termination functions
|
||||
-------------------------------
|
||||
|
||||
The functions below are available in ``mjlab.envs.mdp.terminations`` and
|
||||
are shared across tasks. Individual tasks may define additional
|
||||
termination functions specific to their objective. All termination
|
||||
functions return a boolean tensor of shape ``[num_envs]``.
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 28 72
|
||||
|
||||
* - Function
|
||||
- Description
|
||||
* - ``time_out``
|
||||
- Returns ``True`` when the episode length reaches
|
||||
``env.max_episode_length``. Register with ``time_out=True`` so
|
||||
the manager treats it as a truncation.
|
||||
* - ``bad_orientation``
|
||||
- Returns ``True`` when the angle between the asset's up axis and
|
||||
world up exceeds ``limit_angle`` (radians).
|
||||
* - ``root_height_below_minimum``
|
||||
- Returns ``True`` when the asset's root link height is below
|
||||
``minimum_height`` (meters).
|
||||
* - ``nan_detection``
|
||||
- Returns ``True`` when NaN or Inf values appear anywhere in the
|
||||
physics state. A safety net to terminate diverged simulations
|
||||
cleanly.
|
||||
|
||||
|
||||
Writing custom termination functions
|
||||
-------------------------------------
|
||||
|
||||
Custom termination functions follow the same patterns as reward
|
||||
functions. A plain function accepts ``env`` and returns a boolean
|
||||
``[num_envs]`` tensor. See :ref:`env-config-term-pattern` for the
|
||||
general pattern.
|
||||
@@ -0,0 +1,357 @@
|
||||
.. _terrain:
|
||||
|
||||
Terrain
|
||||
=======
|
||||
|
||||
The terrain is the shared ground surface for all environments in a scene.
|
||||
mjlab supports two modes: a flat ground plane for tasks that do not need
|
||||
varying terrain, and a procedural terrain generator that assembles a grid
|
||||
of sub-terrain patches with configurable difficulty. Procedural terrain
|
||||
is particularly useful for training locomotion policies, where a
|
||||
curriculum of increasing ground difficulty drives robust walking and
|
||||
climbing behaviors.
|
||||
|
||||
Terrain is configured through ``TerrainEntityCfg`` and passed to the
|
||||
scene via the ``terrain`` field of ``SceneCfg``. See :ref:`scene` for
|
||||
how the terrain integrates with the rest of the scene.
|
||||
|
||||
|
||||
Flat terrain
|
||||
------------
|
||||
|
||||
The default mode. A single ground plane modeled as a MuJoCo plane geom
|
||||
with no procedural geometry. Environments are arranged in a regular grid
|
||||
with spacing controlled by ``env_spacing`` on ``SceneCfg``.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.terrains import TerrainEntityCfg
|
||||
|
||||
terrain = TerrainEntityCfg(terrain_type="plane")
|
||||
|
||||
|
||||
Procedural terrain
|
||||
------------------
|
||||
|
||||
For tasks that benefit from terrain variety (locomotion, navigation),
|
||||
``TerrainGeneratorCfg`` assembles a rectangular grid of sub-terrain
|
||||
patches. Each patch is generated from a ``SubTerrainCfg`` that defines
|
||||
the geometry and how it scales with difficulty.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.terrains import TerrainEntityCfg
|
||||
from mjlab.terrains.terrain_generator import TerrainGeneratorCfg
|
||||
import mjlab.terrains as terrain_gen
|
||||
|
||||
terrain = TerrainEntityCfg(
|
||||
terrain_type="generator",
|
||||
terrain_generator=TerrainGeneratorCfg(
|
||||
size=(8.0, 8.0),
|
||||
num_rows=10,
|
||||
num_cols=20,
|
||||
border_width=20.0,
|
||||
curriculum=True,
|
||||
sub_terrains={
|
||||
"flat": terrain_gen.BoxFlatTerrainCfg(proportion=0.2),
|
||||
"stairs": terrain_gen.BoxPyramidStairsTerrainCfg(
|
||||
proportion=0.4,
|
||||
step_height_range=(0.0, 0.15),
|
||||
step_width=0.3,
|
||||
platform_width=2.0,
|
||||
),
|
||||
"rough": terrain_gen.HfRandomUniformTerrainCfg(
|
||||
proportion=0.4,
|
||||
noise_range=(0.02, 0.10),
|
||||
noise_step=0.02,
|
||||
),
|
||||
},
|
||||
),
|
||||
max_init_terrain_level=5,
|
||||
)
|
||||
|
||||
The generator creates a ``num_rows x num_cols`` grid of patches. The
|
||||
``sub_terrains`` dictionary maps names to ``SubTerrainCfg`` instances,
|
||||
and each sub-terrain's ``proportion`` weight controls how many columns
|
||||
(curriculum mode) or sampling probability (random mode) it receives.
|
||||
|
||||
|
||||
Grid layout
|
||||
^^^^^^^^^^^
|
||||
|
||||
Two generation modes control how terrain types are distributed across
|
||||
the grid:
|
||||
|
||||
**Curriculum mode** (``curriculum=True``). Columns are deterministically
|
||||
assigned to terrain types based on their ``proportion`` weights. A type
|
||||
with proportion 0.4 in a 20-column grid gets 8 columns. All patches in
|
||||
a column share the same terrain type, and difficulty increases from row 0
|
||||
(easiest) to row ``num_rows - 1`` (hardest). This structured layout is
|
||||
what enables the curriculum system to advance environments to harder rows
|
||||
as performance improves.
|
||||
|
||||
**Random mode** (``curriculum=False``). Every patch independently samples
|
||||
a terrain type weighted by ``proportion`` and a difficulty from
|
||||
``difficulty_range``. This provides maximum variety but no structured
|
||||
difficulty progression.
|
||||
|
||||
|
||||
The difficulty parameter
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
Each sub-terrain's generation function receives a ``difficulty`` value
|
||||
in ``[0, 1]``. This value linearly interpolates the terrain's
|
||||
configurable ranges. For example, a ``BoxPyramidStairsTerrainCfg`` with
|
||||
``step_height_range=(0.0, 0.2)`` produces flat ground at difficulty 0
|
||||
and 20 cm steps at difficulty 1. In curriculum mode, difficulty is
|
||||
determined by the row: row 0 gets the minimum, row ``num_rows - 1`` gets
|
||||
the maximum.
|
||||
|
||||
|
||||
Sub-terrain types
|
||||
-----------------
|
||||
|
||||
mjlab provides two families of sub-terrain types: **primitive terrains**
|
||||
built from box geoms, and **heightfield terrains** built from continuous
|
||||
elevation grids. All types inherit from ``SubTerrainCfg`` and accept a
|
||||
``proportion`` weight and optional ``flat_patch_sampling`` configuration.
|
||||
|
||||
|
||||
Primitive terrains
|
||||
^^^^^^^^^^^^^^^^^^
|
||||
|
||||
Procedural patches built entirely from box geoms. The discrete geometry
|
||||
makes them well suited for staircases, stepping stones, and other
|
||||
structured obstacles. Most primitive types share common parameters:
|
||||
``platform_width`` (central flat area), ``border_width`` (flat margin),
|
||||
and one or more difficulty-scaled ranges.
|
||||
|
||||
.. grid:: 3
|
||||
|
||||
.. grid-item-card:: Flat
|
||||
|
||||
.. image:: _static/terrains/box_flat.png
|
||||
|
||||
Flat box patch. Useful as an easy baseline in a curriculum grid.
|
||||
|
||||
.. grid-item-card:: Pyramid Stairs
|
||||
|
||||
.. image:: _static/terrains/box_pyramid_stairs.png
|
||||
|
||||
Pyramid staircase with steps descending inward toward a central
|
||||
platform.
|
||||
|
||||
.. grid-item-card:: Inverted Pyramid Stairs
|
||||
|
||||
.. image:: _static/terrains/box_inverted_pyramid_stairs.png
|
||||
|
||||
Inverted pyramid with steps ascending from the outside inward.
|
||||
|
||||
.. grid-item-card:: Random Stairs
|
||||
|
||||
.. image:: _static/terrains/box_random_stairs.png
|
||||
|
||||
Pyramid staircase with random per-step heights.
|
||||
|
||||
.. grid-item-card:: Open Stairs
|
||||
|
||||
.. image:: _static/terrains/box_open_stairs.png
|
||||
|
||||
Concentric step rings. Can be a bowl or pyramid depending on the
|
||||
``inverted`` flag.
|
||||
|
||||
.. grid-item-card:: Random Grid
|
||||
|
||||
.. image:: _static/terrains/box_random_grid.png
|
||||
|
||||
Grid of boxes at randomly sampled heights.
|
||||
|
||||
.. grid-item-card:: Random Spread
|
||||
|
||||
.. image:: _static/terrains/box_random_spread.png
|
||||
|
||||
Randomly positioned and rotated boxes of varying sizes scattered
|
||||
across the patch.
|
||||
|
||||
.. grid-item-card:: Stepping Stones
|
||||
|
||||
.. image:: _static/terrains/box_stepping_stones.png
|
||||
|
||||
Stepping-stone columns rising from a deep pit.
|
||||
|
||||
.. grid-item-card:: Narrow Beams
|
||||
|
||||
.. image:: _static/terrains/box_narrow_beams.png
|
||||
|
||||
Radial beams extending outward from a central platform above a
|
||||
pit.
|
||||
|
||||
.. grid-item-card:: Tilted Grid
|
||||
|
||||
.. image:: _static/terrains/box_tilted_grid.png
|
||||
|
||||
Grid of independently tilted mesh tiles.
|
||||
|
||||
.. grid-item-card:: Nested Rings
|
||||
|
||||
.. image:: _static/terrains/box_nested_rings.png
|
||||
|
||||
Concentric ring structures at random heights.
|
||||
|
||||
|
||||
Heightfield terrains
|
||||
^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
Continuous terrain profiles built from MuJoCo heightfield geoms. The
|
||||
surface is a dense grid of elevation samples, producing smooth slopes
|
||||
and undulating ground that box geoms cannot represent.
|
||||
|
||||
.. grid:: 3
|
||||
|
||||
.. grid-item-card:: Pyramid Slope
|
||||
|
||||
.. image:: _static/terrains/hf_pyramid_slope.png
|
||||
|
||||
Smooth pyramid slope with a flat platform at the peak.
|
||||
``inverted=True`` places the platform at the bottom.
|
||||
|
||||
.. grid-item-card:: Random Uniform
|
||||
|
||||
.. image:: _static/terrains/hf_random_uniform.png
|
||||
|
||||
Random uniform noise, optionally downsampled and interpolated to
|
||||
control feature size.
|
||||
|
||||
.. grid-item-card:: Wave
|
||||
|
||||
.. image:: _static/terrains/hf_wave.png
|
||||
|
||||
Sinusoidal wave profile.
|
||||
|
||||
.. grid-item-card:: Discrete Obstacles
|
||||
|
||||
.. image:: _static/terrains/hf_discrete_obstacles.png
|
||||
|
||||
Rectangular bumps and pits scattered across a flat base.
|
||||
|
||||
.. grid-item-card:: Perlin Noise
|
||||
|
||||
.. image:: _static/terrains/hf_perlin_noise.png
|
||||
|
||||
Fractal Perlin noise producing natural terrain undulation.
|
||||
|
||||
|
||||
Preset configurations
|
||||
---------------------
|
||||
|
||||
mjlab ships two ready-made ``TerrainGeneratorCfg`` presets in
|
||||
``mjlab.terrains.config``:
|
||||
|
||||
``ROUGH_TERRAINS_CFG``
|
||||
A 10x20 grid with seven terrain types (flat, stairs, inverted
|
||||
stairs, slopes, inverted slopes, random rough, waves). Designed for
|
||||
locomotion training with a moderate difficulty range.
|
||||
|
||||
``ALL_TERRAINS_CFG``
|
||||
A 10x16 grid with all sixteen terrain types at equal proportion.
|
||||
Useful for training on maximum terrain variety.
|
||||
|
||||
Both can be used directly or customized with ``dataclasses.replace()``:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from dataclasses import replace
|
||||
from mjlab.terrains.config import ROUGH_TERRAINS_CFG
|
||||
|
||||
my_terrains = replace(ROUGH_TERRAINS_CFG, num_rows=5)
|
||||
|
||||
|
||||
Terrain curriculum
|
||||
------------------
|
||||
|
||||
In curriculum mode the terrain grid provides a natural axis for
|
||||
progressive training: rows represent difficulty levels, and the
|
||||
curriculum system moves environments up or down the grid based on
|
||||
performance. See :ref:`curriculum` for full details on configuring
|
||||
curriculum terms.
|
||||
|
||||
The key concepts:
|
||||
|
||||
- Each environment tracks a ``terrain_level`` (row index) and
|
||||
``terrain_type`` (column index).
|
||||
- ``TerrainEntityCfg.max_init_terrain_level`` controls how high
|
||||
environments can start at their first reset. Setting it to 5 means
|
||||
environments begin on rows 0 through 5.
|
||||
- The built-in ``terrain_levels_vel`` curriculum term promotes
|
||||
environments that track commanded velocity well and demotes
|
||||
environments that fall or fail to make progress.
|
||||
- When an environment reaches the maximum row, it is randomly reassigned
|
||||
to a lower row to prevent the policy from collapsing to a single
|
||||
difficulty level.
|
||||
|
||||
|
||||
Flat patch detection
|
||||
--------------------
|
||||
|
||||
Heightfield terrains can pre-compute flat regions on their surface during
|
||||
generation. These flat patches are useful as safe spawn points for tasks
|
||||
that require the robot to start on level ground, even on otherwise rough
|
||||
terrain.
|
||||
|
||||
Flat patch detection is configured per sub-terrain via the
|
||||
``flat_patch_sampling`` field on ``SubTerrainCfg``:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.terrains.terrain_generator import FlatPatchSamplingCfg
|
||||
|
||||
rough = terrain_gen.HfRandomUniformTerrainCfg(
|
||||
proportion=0.5,
|
||||
noise_range=(0.02, 0.10),
|
||||
flat_patch_sampling={
|
||||
"spawn": FlatPatchSamplingCfg(
|
||||
num_patches=10,
|
||||
patch_radius=0.5,
|
||||
max_height_diff=0.05,
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
The detection algorithm uses morphological filtering to find circular
|
||||
regions where height variation stays within ``max_height_diff``. Detected
|
||||
patches are accessible at runtime through
|
||||
``scene.terrain.flat_patches["spawn"]``.
|
||||
|
||||
To spawn robots on detected patches instead of at the sub-terrain center,
|
||||
use ``reset_root_state_from_flat_patches`` as the reset event term. See
|
||||
:ref:`events` for details.
|
||||
|
||||
.. note::
|
||||
|
||||
Only heightfield (``Hf*``) terrains support flat patch detection.
|
||||
Primitive (``Box*``) terrains do not have heightfield data to analyze.
|
||||
If any sub-terrain in the grid configures ``flat_patch_sampling``,
|
||||
the flat patches array is allocated for all cells; sub-terrains
|
||||
without patches have their slots filled with the sub-terrain's spawn
|
||||
origin so that the reset event always receives valid positions.
|
||||
|
||||
|
||||
Debug visualization
|
||||
-------------------
|
||||
|
||||
The terrain entity adds debug sites to three geom groups that can be
|
||||
toggled in the MuJoCo native viewer or Viser viewer:
|
||||
|
||||
- **Group 3**: flat patch sites (yellow boxes marking safe spawn regions)
|
||||
- **Group 4**: environment origin sites (green spheres at each
|
||||
environment's position)
|
||||
- **Group 5**: terrain origin sites (blue spheres at each sub-terrain
|
||||
patch center)
|
||||
|
||||
.. figure:: _static/terrains/flat_patch_group.png
|
||||
:width: 100%
|
||||
:align: center
|
||||
:alt: Flat patch visualization
|
||||
|
||||
Flat patches (group 3) overlaid on a procedural terrain grid in the
|
||||
Viser viewer.
|
||||
@@ -0,0 +1,290 @@
|
||||
.. _cloud-training:
|
||||
|
||||
Cloud Training
|
||||
==============
|
||||
|
||||
This guide walks through launching training jobs on
|
||||
`Lambda Cloud <https://lambdalabs.com/>`_ using
|
||||
`SkyPilot <https://skypilot.readthedocs.io/>`_. SkyPilot provisions a GPU
|
||||
instance, syncs your code, runs the job, and tears down the machine when
|
||||
it finishes.
|
||||
|
||||
Two SkyPilot task files live in ``scripts/cloud/``:
|
||||
|
||||
.. list-table::
|
||||
:widths: 30 70
|
||||
:header-rows: 1
|
||||
|
||||
* - File
|
||||
- Description
|
||||
* - ``train.yaml``
|
||||
- Installs mjlab directly with uv.
|
||||
* - ``train-docker.yaml``
|
||||
- Pulls the pre-built Docker image from GHCR for a reproducible
|
||||
environment.
|
||||
|
||||
|
||||
Prerequisites
|
||||
-------------
|
||||
|
||||
**1. Install SkyPilot**
|
||||
|
||||
SkyPilot is a local CLI tool, not a project dependency. Install it with:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
uv tool install "skypilot[lambda]"
|
||||
|
||||
**2. Lambda Cloud API key**
|
||||
|
||||
Generate a key at `Lambda Cloud API keys
|
||||
<https://cloud.lambda.ai/api-keys/cloud-api>`_. Name it after your
|
||||
machine (e.g. ``kevins-macbook``) so you can tell keys apart later.
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
mkdir -p ~/.lambda_cloud && chmod 700 ~/.lambda_cloud
|
||||
echo "api_key = <your-api-key>" > ~/.lambda_cloud/lambda_keys
|
||||
chmod 600 ~/.lambda_cloud/lambda_keys
|
||||
|
||||
**3. Verify setup**
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
sky check lambda
|
||||
|
||||
You should see Lambda listed as an enabled cloud.
|
||||
|
||||
**4. W&B credentials** *(optional)*
|
||||
|
||||
If you log to Weights & Biases, install the ``wandb`` CLI and log in:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
uv tool install wandb
|
||||
wandb login
|
||||
|
||||
This stores your credentials in ``~/.netrc``. The SkyPilot task files
|
||||
mount this file onto the remote instance via ``file_mounts`` so that
|
||||
``wandb`` authenticates automatically, no environment variable needed.
|
||||
|
||||
|
||||
Quick start
|
||||
-----------
|
||||
|
||||
From the repo root:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
sky launch scripts/cloud/train.yaml \
|
||||
--env TASK=Mjlab-Velocity-Flat-Unitree-G1
|
||||
|
||||
# Or with Docker:
|
||||
sky launch scripts/cloud/train-docker.yaml \
|
||||
--env TASK=Mjlab-Velocity-Flat-Unitree-G1
|
||||
What happens behind the scenes:
|
||||
|
||||
1. SkyPilot finds an available Lambda instance with the requested GPU.
|
||||
2. It provisions the instance and uploads your local code via rsync.
|
||||
3. The ``setup`` step runs (uv install or Docker pull).
|
||||
4. The ``run`` step runs (training).
|
||||
5. After 5 minutes of idle time the instance is terminated automatically.
|
||||
|
||||
.. warning::
|
||||
|
||||
Lambda instances can only be **launched** or **terminated**. There is
|
||||
no pause or suspend. Do not run ``sudo shutdown`` from inside the
|
||||
instance; it will put the machine in an alert state and billing will
|
||||
continue. Always use ``sky down`` to terminate.
|
||||
|
||||
|
||||
Common operations
|
||||
-----------------
|
||||
|
||||
**List available GPUs**
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
sky show-gpus --infra lambda
|
||||
|
||||
**Choose a different GPU**
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
sky launch scripts/cloud/train.yaml --gpus H100:1 # 1x H100
|
||||
sky launch scripts/cloud/train.yaml --gpus A100:8 # 8x A100
|
||||
sky launch scripts/cloud/train.yaml --gpus A10:1 # 1x A10 (cheaper)
|
||||
|
||||
.. note::
|
||||
|
||||
Both task files pass ``--gpu-ids all``, so multi-GPU instances
|
||||
automatically use :ref:`distributed training <distributed-training>`.
|
||||
When requesting more than one GPU, consider scaling
|
||||
``MAX_ITERATIONS`` down proportionally. See
|
||||
:ref:`distributed-training` for details on scaling behavior.
|
||||
|
||||
**Override training parameters**
|
||||
|
||||
Every variable in the YAML ``envs`` block can be overridden from the
|
||||
command line with ``--env``:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
sky launch scripts/cloud/train.yaml \
|
||||
--env TASK=Mjlab-Velocity-Flat-Unitree-Go1 \
|
||||
--env NUM_ENVS=8192 \
|
||||
--env MAX_ITERATIONS=10000
|
||||
**Run your own task**
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
sky launch scripts/cloud/train.yaml \
|
||||
--env TASK=Mjlab-Velocity-Flat-Unitree-Go1
|
||||
|
||||
To see all registered tasks:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
uv run list-envs
|
||||
uv run list-envs --keyword Velocity # filter by keyword
|
||||
|
||||
|
||||
Hyperparameter sweeps
|
||||
---------------------
|
||||
|
||||
Use `W&B Sweeps <https://docs.wandb.ai/models/sweeps/>`_ with SkyPilot
|
||||
to search hyperparameters across a multi-GPU instance. The sweep
|
||||
controller lives on the W&B servers; each GPU on the instance runs an
|
||||
independent sweep agent that pulls a hyperparameter configuration,
|
||||
trains, and reports metrics.
|
||||
|
||||
The example uses ``method: random``, where each agent samples
|
||||
independently. Bayesian search also works well with parallel agents.
|
||||
Agents report results back as they finish and the controller updates its
|
||||
model between rounds. If using Bayesian, set ``run_cap`` high enough for
|
||||
the optimizer to go through several rounds.
|
||||
|
||||
Four files are involved:
|
||||
|
||||
.. list-table::
|
||||
:widths: 35 65
|
||||
:header-rows: 1
|
||||
|
||||
* - File
|
||||
- Description
|
||||
* - ``sweep.yaml``
|
||||
- W&B sweep configuration (parameters, search method, metric).
|
||||
* - ``sweep-cluster.yaml``
|
||||
- SkyPilot cluster definition (resources, setup, no run section).
|
||||
* - ``sweep-agent.yaml``
|
||||
- SkyPilot job definition that runs ``wandb agent`` on one GPU.
|
||||
* - ``sweep-launch.sh``
|
||||
- Convenience script that creates the sweep, provisions the
|
||||
cluster, and submits one agent per GPU.
|
||||
|
||||
**Quick start**
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
./scripts/cloud/sweep-launch.sh A100:8 # 8 agents on an 8xA100
|
||||
|
||||
This creates a W&B sweep, provisions a cluster, and submits one agent
|
||||
per GPU. Each agent runs training with a different set of
|
||||
hyperparameters sampled by the sweep controller.
|
||||
|
||||
**Manual steps** (if you prefer more control):
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
# 1. Create the sweep (returns a SWEEP_ID).
|
||||
wandb sweep scripts/cloud/sweep.yaml
|
||||
|
||||
# 2. Provision the cluster (runs setup, no agents yet).
|
||||
sky launch scripts/cloud/sweep-cluster.yaml \
|
||||
-c mjlab-sweep --gpus A100:8
|
||||
|
||||
# 3. Submit one agent per GPU.
|
||||
sky exec mjlab-sweep scripts/cloud/sweep-agent.yaml \
|
||||
--gpus A100:1 --env SWEEP_ID=<entity/project/sweep_id> -d
|
||||
|
||||
Monitor progress on the W&B dashboard or with ``sky queue mjlab-sweep``.
|
||||
When done, tear down the cluster with ``sky down mjlab-sweep``.
|
||||
|
||||
|
||||
Monitoring
|
||||
----------
|
||||
|
||||
Provisioning can take five minutes or more while Lambda allocates the
|
||||
instance. Open a second terminal to keep an eye on things:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
sky status # cluster state (INIT, UP, ...)
|
||||
sky logs sky-<cluster-name> # stream logs in real time
|
||||
sky logs sky-<cluster-name> --no-follow # print current logs and exit
|
||||
sky queue sky-<cluster-name> # job queue for the cluster
|
||||
|
||||
.. tip::
|
||||
|
||||
If the cluster stays in ``INIT`` for a long time, the GPU type is
|
||||
likely sold out. Cancel with ``sky down`` and try a different GPU, or
|
||||
add ``--retry-until-up`` to let SkyPilot keep polling until capacity
|
||||
opens up.
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
sky down sky-<cluster-name>
|
||||
sky launch scripts/cloud/train.yaml --retry-until-up
|
||||
|
||||
|
||||
Iterating on a failed job
|
||||
-------------------------
|
||||
|
||||
When a job fails the cluster keeps running (and billing). You can fix
|
||||
the problem locally and resubmit without waiting for a new instance:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
sky exec sky-<cluster-name> scripts/cloud/train.yaml
|
||||
.. important::
|
||||
|
||||
``sky exec`` rsyncs your latest code and reruns the ``run`` step
|
||||
only. It does **not** rerun ``setup``. If your fix involves
|
||||
dependency changes, use ``sky launch`` again or SSH in and run the
|
||||
setup commands manually.
|
||||
|
||||
Other useful commands:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
sky down sky-<cluster-name> # terminate the instance immediately
|
||||
ssh sky-<cluster-name> # SSH in (SkyPilot configures this for you)
|
||||
|
||||
|
||||
Cost management
|
||||
---------------
|
||||
|
||||
.. warning::
|
||||
|
||||
Always run ``sky status`` after each session to confirm nothing is
|
||||
still running. Forgotten instances are the most common source of
|
||||
unexpected charges. To terminate everything at once: ``sky down -a``.
|
||||
|
||||
- Instances auto-terminate after 5 minutes of idle time by default.
|
||||
You can change this in the YAML (``idle_minutes``) or at launch time
|
||||
with ``--idle-minutes-to-autostop``.
|
||||
- The ``down: true`` setting in the YAML means the instance is fully
|
||||
terminated when it stops, not just paused. Billing stops completely.
|
||||
|
||||
|
||||
Troubleshooting
|
||||
---------------
|
||||
|
||||
**No instances available**
|
||||
|
||||
Lambda GPUs sell out frequently. A few things to try:
|
||||
|
||||
- Use ``--retry-until-up`` to poll automatically.
|
||||
- Try a different GPU type: ``--gpus A100:1``, ``--gpus A10:1``, etc.
|
||||
- If you have credentials for other clouds (GCP, AWS), SkyPilot can fall
|
||||
back to them automatically.
|
||||
@@ -0,0 +1,109 @@
|
||||
.. _distributed-training:
|
||||
|
||||
Distributed Training
|
||||
====================
|
||||
|
||||
mjlab supports multi-GPU distributed training using
|
||||
`torchrunx <https://github.com/apoorvkh/torchrunx>`_. Each GPU runs
|
||||
independent rollouts with its own environments, and gradients are
|
||||
synchronized during policy updates. Throughput scales nearly linearly with
|
||||
GPU count.
|
||||
|
||||
|
||||
Usage
|
||||
-----
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
# Single GPU (default).
|
||||
uv run train <task-name> --gpu-ids "[0]"
|
||||
|
||||
# Two GPUs.
|
||||
uv run train <task-name> --gpu-ids "[0, 1]"
|
||||
|
||||
# All available GPUs.
|
||||
uv run train <task-name> --gpu-ids all
|
||||
|
||||
# CPU mode.
|
||||
uv run train <task-name> --gpu-ids None
|
||||
|
||||
Key points:
|
||||
|
||||
- GPU indices are relative to ``CUDA_VISIBLE_DEVICES`` if set. For example,
|
||||
``CUDA_VISIBLE_DEVICES=2,3 uv run train ... --gpu-ids "[0, 1]"`` uses physical
|
||||
GPUs 2 and 3.
|
||||
- Single-GPU and CPU modes run directly without torchrunx.
|
||||
|
||||
|
||||
Scaling behavior
|
||||
----------------
|
||||
|
||||
Multi-GPU training is **data-parallel, not work-splitting**. Each GPU runs
|
||||
the full ``num-envs`` count independently, so the total experience collected
|
||||
per iteration is:
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
experience per iteration = num_envs x num_steps_per_env x num_gpus
|
||||
|
||||
Iteration speed stays roughly the same because each GPU does the same amount
|
||||
of work. The benefit is that each policy update sees more diverse experience,
|
||||
so the policy converges faster in wall-clock time.
|
||||
|
||||
.. important::
|
||||
|
||||
Because ``max-iterations`` is not automatically adjusted, training with
|
||||
more GPUs runs for proportionally longer. If you want the same total
|
||||
training time, scale ``max-iterations`` down by the number of GPUs
|
||||
(e.g., halve it when doubling from 1 to 2 GPUs).
|
||||
|
||||
|
||||
How it works
|
||||
------------
|
||||
|
||||
mjlab's role is to **isolate MuJoCo Warp simulations on each GPU** using
|
||||
``wp.ScopedDevice``. torchrunx handles the rest.
|
||||
|
||||
**Process spawning.** ``torchrunx.Launcher`` spawns one process per GPU and
|
||||
sets ``RANK``, ``LOCAL_RANK``, and ``WORLD_SIZE`` to coordinate them. Each
|
||||
process executes the training function with its assigned GPU.
|
||||
|
||||
**Independent rollouts.** Each process maintains its own:
|
||||
|
||||
- Environment instances (with ``num-envs`` parallel environments), isolated
|
||||
on its assigned GPU via ``wp.ScopedDevice``
|
||||
- Policy network copy
|
||||
- Experience buffer (sized ``num_steps_per_env * num_envs``)
|
||||
|
||||
Each process uses ``seed = cfg.seed + local_rank`` to ensure different
|
||||
random experiences across GPUs, increasing sample diversity.
|
||||
|
||||
**Gradient synchronization.** During the update phase, RSL-RL synchronizes
|
||||
gradients after each mini-batch through its ``reduce_parameters()`` method:
|
||||
|
||||
1. Each process computes gradients independently on its local mini-batch
|
||||
2. All policy gradients are flattened into a single tensor
|
||||
3. ``torch.distributed.all_reduce`` averages gradients across all GPUs
|
||||
4. Averaged gradients are copied back to each parameter, keeping policies
|
||||
synchronized
|
||||
|
||||
**Single-writer I/O.** Only rank 0 writes config files, videos, and W&B
|
||||
logs to avoid race conditions.
|
||||
|
||||
|
||||
Logging
|
||||
-------
|
||||
|
||||
By default, torchrunx process logs are saved to ``{log_dir}/torchrunx/``.
|
||||
This can be customized:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
# Disable torchrunx file logging.
|
||||
uv run train <task-name> --gpu-ids "[0, 1]" --torchrunx-log-dir ""
|
||||
|
||||
# Custom log directory.
|
||||
uv run train <task-name> --gpu-ids "[0, 1]" --torchrunx-log-dir /path/to/logs
|
||||
|
||||
# Environment variable (takes precedence over the flag).
|
||||
TORCHRUNX_LOG_DIR=/tmp/logs uv run train <task-name> --gpu-ids "[0, 1]"
|
||||
@@ -0,0 +1,65 @@
|
||||
.. _motion-imitation:
|
||||
|
||||
Motion Imitation
|
||||
================
|
||||
|
||||
mjlab can train humanoid policies to imitate reference motions. This page
|
||||
covers motion data preprocessing and training.
|
||||
|
||||
WandB registry setup
|
||||
--------------------
|
||||
|
||||
mjlab uses `Weights & Biases <https://wandb.ai/>`_ to store and load
|
||||
reference motions. Before preprocessing any motions, create a WandB registry
|
||||
by following the
|
||||
`BeyondMimic instructions <https://github.com/HybridRobotics/whole_body_tracking/blob/main/README.md#motion-preprocessing--registry-setup>`_
|
||||
(only the registry creation step; skip the ``csv_to_npz.py`` command shown
|
||||
there).
|
||||
|
||||
Motion preprocessing
|
||||
--------------------
|
||||
|
||||
Reference motions are retargeted CSV files in Unitree's generalized
|
||||
coordinate convention (base position, base quaternion in xyzw, then joint
|
||||
angles).
|
||||
|
||||
Convert a CSV to the NPZ format mjlab expects:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
MUJOCO_GL=egl uv run -m mjlab.scripts.csv_to_npz \
|
||||
--input-file <PATH_TO_CSV> \
|
||||
--output-name <MOTION_NAME> \
|
||||
--input-fps 30 \
|
||||
--output-fps 50 \
|
||||
--render True
|
||||
|
||||
The script plays the motion through MuJoCo Warp, computes forward kinematics
|
||||
for every body, and uploads the resulting NPZ to your WandB registry.
|
||||
|
||||
.. warning::
|
||||
|
||||
You **must** use mjlab's converter (``mjlab.scripts.csv_to_npz``).
|
||||
Converters from other frameworks such as IsaacLab produce NPZ files with
|
||||
incompatible body orderings. The NPZ stores precomputed body positions and
|
||||
quaternions indexed by body number, and different physics engines assign
|
||||
body indices differently (MuJoCo uses depth first traversal, PhysX uses
|
||||
breadth first). A mismatched NPZ will map tracking targets to the wrong
|
||||
bodies and training will not converge.
|
||||
|
||||
Training
|
||||
--------
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
uv run train Mjlab-Tracking-Flat-Unitree-G1 \
|
||||
--registry-name your-org/motions/motion-name \
|
||||
--env.scene.num-envs 4096
|
||||
|
||||
Evaluation
|
||||
----------
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
uv run play Mjlab-Tracking-Flat-Unitree-G1 \
|
||||
--wandb-run-path your-org/mjlab/run-id
|
||||
@@ -0,0 +1,277 @@
|
||||
.. _rsl_rl:
|
||||
|
||||
Training with RSL-RL
|
||||
====================
|
||||
|
||||
mjlab uses `RSL-RL <https://github.com/leggedrobotics/rsl_rl>`_ for on-policy
|
||||
reinforcement learning. The integration has three parts: a **task registry**
|
||||
that bundles environment and training configs under a single name, a
|
||||
**VecEnv wrapper** that adapts mjlab environments to the interface RSL-RL
|
||||
expects, and a set of **configuration dataclasses** that control the training
|
||||
run.
|
||||
|
||||
|
||||
Task registry
|
||||
-------------
|
||||
|
||||
Every task in mjlab is a pair: an environment configuration
|
||||
(``ManagerBasedRlEnvCfg``) and a training configuration
|
||||
(``RslRlOnPolicyRunnerCfg``). The task registry maps a string name to this
|
||||
pair so that training can be launched by name from the CLI.
|
||||
|
||||
Tasks are registered by calling ``register_mjlab_task`` in the task's
|
||||
``__init__.py``:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.tasks.registry import register_mjlab_task
|
||||
from mjlab.tasks.velocity.rl import VelocityOnPolicyRunner
|
||||
|
||||
from .env_cfgs import unitree_g1_rough_env_cfg, unitree_g1_flat_env_cfg
|
||||
from .rl_cfg import unitree_g1_ppo_runner_cfg
|
||||
|
||||
register_mjlab_task(
|
||||
task_id="Mjlab-Velocity-Rough-Unitree-G1",
|
||||
env_cfg=unitree_g1_rough_env_cfg(),
|
||||
play_env_cfg=unitree_g1_rough_env_cfg(play=True),
|
||||
rl_cfg=unitree_g1_ppo_runner_cfg(),
|
||||
runner_cls=VelocityOnPolicyRunner,
|
||||
)
|
||||
|
||||
Each registration takes:
|
||||
|
||||
- ``task_id``: a unique name following the convention
|
||||
``Mjlab-{Category}-{Terrain}-{Robot}``
|
||||
- ``env_cfg``: the ``ManagerBasedRlEnvCfg`` used for training
|
||||
- ``play_env_cfg``: a variant with randomization disabled and episode length
|
||||
set to infinity, used for evaluation
|
||||
- ``rl_cfg``: the ``RslRlOnPolicyRunnerCfg`` with PPO hyperparameters and
|
||||
network architecture
|
||||
- ``runner_cls``: an optional custom runner class (defaults to
|
||||
``MjlabOnPolicyRunner``)
|
||||
|
||||
All task packages under ``src/mjlab/tasks/`` are auto-discovered at import
|
||||
time, so adding a new task only requires creating the config package and
|
||||
calling ``register_mjlab_task``.
|
||||
|
||||
|
||||
Training and playback
|
||||
---------------------
|
||||
|
||||
**Launching a training run:**
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
uv run train Mjlab-Velocity-Flat-Unitree-G1 --num-envs 4096
|
||||
|
||||
The task name is the first positional argument. The entire configuration
|
||||
hierarchy (environment, scene, rewards, PPO hyperparameters, etc.) is
|
||||
exposed as CLI flags through `tyro <https://brentyi.github.io/tyro/>`_.
|
||||
Every field in ``ManagerBasedRlEnvCfg`` and ``RslRlOnPolicyRunnerCfg`` can
|
||||
be overridden from the command line using dot-separated paths:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
uv run train Mjlab-Velocity-Flat-Unitree-G1 \
|
||||
--num-envs 4096 \
|
||||
--agent.max-iterations 10000 \
|
||||
--agent.algorithm.learning-rate 3e-4 \
|
||||
--env.decimation 2
|
||||
|
||||
.. important::
|
||||
|
||||
- **Hyphens, not underscores**: Python field names use underscores
|
||||
(``num_envs``), but CLI flags use POSIX-style hyphens (``--num-envs``).
|
||||
- **Explicit booleans**: boolean flags require an explicit ``True`` or
|
||||
``False`` value (e.g., ``--agent.resume True``, not ``--agent.resume``).
|
||||
This is intentional for compatibility with W&B sweep configs.
|
||||
|
||||
To discover available flags, use ``--help`` and pipe through ``grep``:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
# See all flags.
|
||||
uv run train Mjlab-Velocity-Flat-Unitree-G1 --help
|
||||
|
||||
# Search for a specific field.
|
||||
uv run train Mjlab-Velocity-Flat-Unitree-G1 --help | grep learning-rate
|
||||
|
||||
Some commonly used top-level flags:
|
||||
|
||||
``--num-envs``
|
||||
Number of parallel simulation environments.
|
||||
|
||||
``--gpu-ids``
|
||||
GPU indices to use. Pass multiple indices for multi-GPU training (see
|
||||
:ref:`distributed-training`), or ``None`` for CPU mode.
|
||||
|
||||
``--video``
|
||||
Record training rollout videos to ``{log_dir}/videos/train/``.
|
||||
|
||||
``--enable-nan-guard``
|
||||
Enable NaN detection and state capture (see :ref:`nan-guard`).
|
||||
|
||||
|
||||
**Playing back a trained policy:**
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
# From W&B.
|
||||
uv run play Mjlab-Velocity-Flat-Unitree-G1 \
|
||||
--wandb-run-path your-entity/mjlab/run-id
|
||||
|
||||
# From a local checkpoint.
|
||||
uv run play Mjlab-Velocity-Flat-Unitree-G1 \
|
||||
--checkpoint-file logs/rsl_rl/g1_velocity/2025-01-27_14-30-00/model_1000.pt
|
||||
|
||||
Key ``play`` arguments:
|
||||
|
||||
``--agent``
|
||||
Policy mode: ``"trained"`` (default), ``"zero"`` (zero actions), or
|
||||
``"random"`` (uniform random).
|
||||
|
||||
``--viewer``
|
||||
Viewer backend: ``"native"`` (MuJoCo viewer) or ``"viser"``
|
||||
(browser-based).
|
||||
|
||||
``--no-terminations``
|
||||
Disable termination conditions so the policy runs indefinitely.
|
||||
|
||||
|
||||
VecEnv wrapper
|
||||
--------------
|
||||
|
||||
``RslRlVecEnvWrapper`` adapts a ``ManagerBasedRlEnv`` to RSL-RL's ``VecEnv``
|
||||
interface. It handles three things:
|
||||
|
||||
1. **Observation format**: translates observation dictionaries into the
|
||||
``TensorDict`` format RSL-RL expects.
|
||||
2. **Done signal**: merges ``terminated`` and ``truncated`` into a single
|
||||
``dones`` tensor and passes ``time_outs`` through ``extras`` so RSL-RL can
|
||||
bootstrap correctly on truncated episodes.
|
||||
3. **Action clipping**: applies optional action clipping when ``clip_actions``
|
||||
is set in the runner config.
|
||||
|
||||
The wrapper also calls ``env.reset()`` during construction because RSL-RL does
|
||||
not call reset before beginning rollout collection.
|
||||
|
||||
In normal usage you do not interact with the wrapper directly. The training
|
||||
script handles wrapping automatically.
|
||||
|
||||
|
||||
Configuration
|
||||
-------------
|
||||
|
||||
``RslRlOnPolicyRunnerCfg`` is the top-level training configuration. It groups
|
||||
runner settings, network architecture (``RslRlModelCfg``), and PPO
|
||||
hyperparameters (``RslRlPpoAlgorithmCfg``). The following example from the
|
||||
Unitree G1 velocity task shows a typical configuration:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.rl import (
|
||||
RslRlModelCfg,
|
||||
RslRlOnPolicyRunnerCfg,
|
||||
RslRlPpoAlgorithmCfg,
|
||||
)
|
||||
|
||||
def unitree_g1_ppo_runner_cfg() -> RslRlOnPolicyRunnerCfg:
|
||||
return RslRlOnPolicyRunnerCfg(
|
||||
actor=RslRlModelCfg(
|
||||
hidden_dims=(512, 256, 128),
|
||||
activation="elu",
|
||||
obs_normalization=True,
|
||||
),
|
||||
critic=RslRlModelCfg(
|
||||
hidden_dims=(512, 256, 128),
|
||||
activation="elu",
|
||||
obs_normalization=True,
|
||||
),
|
||||
algorithm=RslRlPpoAlgorithmCfg(
|
||||
value_loss_coef=1.0,
|
||||
use_clipped_value_loss=True,
|
||||
clip_param=0.2,
|
||||
entropy_coef=0.01,
|
||||
num_learning_epochs=5,
|
||||
num_mini_batches=4,
|
||||
learning_rate=1.0e-3,
|
||||
schedule="adaptive",
|
||||
gamma=0.99,
|
||||
lam=0.95,
|
||||
desired_kl=0.01,
|
||||
max_grad_norm=1.0,
|
||||
),
|
||||
experiment_name="g1_velocity",
|
||||
save_interval=50,
|
||||
num_steps_per_env=24,
|
||||
max_iterations=30_000,
|
||||
)
|
||||
|
||||
All fields have sensible defaults and can be overridden from the command line
|
||||
(e.g., ``--agent.algorithm.learning-rate 3e-4``). Use ``--help`` to see the
|
||||
full list of available fields and their defaults.
|
||||
|
||||
|
||||
Checkpoints and logging
|
||||
-----------------------
|
||||
|
||||
Training artifacts are written to:
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
logs/rsl_rl/{experiment_name}/{timestamp}/
|
||||
model_{iteration}.pt # policy checkpoints
|
||||
params/
|
||||
env.yaml # full environment config
|
||||
agent.yaml # full runner config
|
||||
|
||||
Checkpoints are saved every ``save_interval`` iterations and uploaded to W&B
|
||||
as model artifacts by default. Set ``upload_model=False`` in the runner
|
||||
config to disable uploads while keeping metric logging.
|
||||
|
||||
.. rubric:: Resuming from a checkpoint
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
uv run train Mjlab-Velocity-Flat-Unitree-G1 \
|
||||
--num-envs 4096 \
|
||||
--agent.resume True
|
||||
|
||||
The runner searches for the most recent run directory under
|
||||
``logs/rsl_rl/{experiment_name}/`` and loads the highest-numbered checkpoint.
|
||||
Narrow the search with ``--agent.load-run`` (regex on directory names) and
|
||||
``--agent.load-checkpoint`` (regex on checkpoint filenames).
|
||||
|
||||
``--agent.max-iterations`` controls how many *additional* iterations to run
|
||||
from the checkpoint. If you are resuming from iteration 11500 with
|
||||
``--agent.max-iterations 300`` (the default), training will run iterations
|
||||
11500 through 11800. Set this to the number of new iterations you want.
|
||||
|
||||
To resume from a W&B run:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
uv run train Mjlab-Velocity-Flat-Unitree-G1 \
|
||||
--num-envs 4096 \
|
||||
--agent.resume True \
|
||||
--wandb-run-path your-entity/mjlab/run-id
|
||||
|
||||
|
||||
Citation
|
||||
--------
|
||||
|
||||
If you use RSL-RL in your research, consider citing:
|
||||
|
||||
.. code-block:: bibtex
|
||||
|
||||
@article{schwarke2025rslrl,
|
||||
title={RSL-RL: A Learning Library for Robotics Research},
|
||||
author={Schwarke, Clemens and Mittal, Mayank and Rudin, Nikita and Hoeller, David and Hutter, Marco},
|
||||
journal={arXiv preprint arXiv:2509.10771},
|
||||
year={2025}
|
||||
}
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
|
||||
motion_imitation
|
||||
@@ -0,0 +1,11 @@
|
||||
.. _tutorials:
|
||||
|
||||
Tutorials
|
||||
=========
|
||||
|
||||
Tutorials to help you get started with mjlab.
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
|
||||
tutorials/cartpole
|
||||
@@ -0,0 +1,439 @@
|
||||
.. _tutorial-cartpole:
|
||||
|
||||
Cartpole: Building Your First Environment
|
||||
=========================================
|
||||
|
||||
This tutorial walks through building a cartpole swingup task from scratch.
|
||||
A cart slides along a rail with a pole attached by a hinge. The agent
|
||||
applies force to the cart to swing the pole up and balance it.
|
||||
|
||||
.. raw:: html
|
||||
|
||||
<video style="width:80%; display:block; margin:0 auto;" autoplay loop muted playsinline>
|
||||
<source src="../../_static/tutorials/cartpole_swingup.mp4" type="video/mp4">
|
||||
</video>
|
||||
<p style="text-align:center; color:#666; font-size:0.9em; margin-top:0.5em;">
|
||||
A trained agent performing the swingup task.
|
||||
</p>
|
||||
|
||||
The entire task lives in two files: an XML model and a Python module. We
|
||||
will build both piece by piece, then snap them together at the end.
|
||||
|
||||
|
||||
The XML model
|
||||
-------------
|
||||
|
||||
Every environment starts with a MuJoCo XML that defines the physical
|
||||
system. For cartpole that means two bodies, two joints, and one motor:
|
||||
|
||||
.. code-block:: xml
|
||||
|
||||
<!-- A cart on a rail with a pole attached by a hinge. -->
|
||||
<body name="cart" pos="0 0 1">
|
||||
<joint name="slider" type="slide" axis="1 0 0"
|
||||
limited="true" range="-1.8 1.8" damping="5e-4"/>
|
||||
<geom name="cart" type="box" size="0.2 0.15 0.1" mass="1"/>
|
||||
<body name="pole_1" childclass="pole">
|
||||
<joint name="hinge_1"/>
|
||||
<geom name="pole_1"/>
|
||||
</body>
|
||||
</body>
|
||||
|
||||
<!-- A motor that pushes the cart along the rail. -->
|
||||
<actuator>
|
||||
<motor name="slide" joint="slider" gear="10"
|
||||
ctrllimited="true" ctrlrange="-1 1"/>
|
||||
</actuator>
|
||||
|
||||
The motor has gear ratio 10 and control range [-1, 1], so the maximum
|
||||
force is 10 N. ``ctrllimited`` tells MuJoCo to clamp the control signal
|
||||
internally, so policy outputs outside this range are safe.
|
||||
|
||||
The full XML is at ``src/mjlab/tasks/cartpole/cartpole.xml``.
|
||||
|
||||
|
||||
Building the environment
|
||||
------------------------
|
||||
|
||||
Everything else lives in a single file, ``cartpole_env_cfg.py``. An
|
||||
mjlab environment is made of small, composable pieces. We will define
|
||||
each piece, then assemble them into a complete config at the end.
|
||||
|
||||
Entity: wrapping the XML
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
An entity is a simulated object in the scene. It can be anything from a
|
||||
static table to an articulated robot. The ``EntityCfg`` wraps a MuJoCo
|
||||
XML and, optionally, actuator and initial state configurations. At
|
||||
runtime, the entity exposes simulation data (joint positions, velocities,
|
||||
etc.) as batched PyTorch tensors.
|
||||
|
||||
The cartpole is an articulated entity with one actuator, so we need a
|
||||
function that loads the XML, an actuator configuration, and an initial
|
||||
state.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# Load the XML.
|
||||
_CARTPOLE_XML = Path(__file__).parent / "cartpole.xml"
|
||||
|
||||
def _get_spec() -> mujoco.MjSpec:
|
||||
return mujoco.MjSpec.from_file(str(_CARTPOLE_XML))
|
||||
|
||||
# Tell mjlab to use the motor defined in the XML as is.
|
||||
_CARTPOLE_ARTICULATION = EntityArticulationInfoCfg(
|
||||
actuators=(XmlActuatorCfg(target_names_expr=("slider",)),),
|
||||
)
|
||||
|
||||
The initial joint state depends on the task variant:
|
||||
|
||||
.. tab-set::
|
||||
|
||||
.. tab-item:: Swingup
|
||||
|
||||
The pole starts pointing down (``hinge = pi``). The agent must swing
|
||||
it up and balance it.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
_SWINGUP_INIT = EntityCfg.InitialStateCfg(
|
||||
joint_pos={"slider": 0.0, "hinge_1": math.pi},
|
||||
joint_vel={".*": 0.0},
|
||||
)
|
||||
|
||||
.. tab-item:: Balance
|
||||
|
||||
The pole starts upright (``hinge = 0``). The agent only needs to
|
||||
keep it balanced.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
_BALANCE_INIT = EntityCfg.InitialStateCfg(
|
||||
joint_pos={"slider": 0.0, "hinge_1": 0.0},
|
||||
joint_vel={".*": 0.0},
|
||||
)
|
||||
|
||||
Now we can snap these together into an ``EntityCfg``:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# Bundle the spec loader, actuator, and initial state into one config.
|
||||
def _get_cartpole_cfg(swing_up: bool = False) -> EntityCfg:
|
||||
return EntityCfg(
|
||||
spec_fn=_get_spec,
|
||||
articulation=_CARTPOLE_ARTICULATION,
|
||||
init_state=_SWINGUP_INIT if swing_up else _BALANCE_INIT,
|
||||
)
|
||||
|
||||
That is the entity done. Later, we will pass it to the scene so the
|
||||
environment knows what to simulate.
|
||||
|
||||
Observations: what the agent sees
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
Each observation term is a function that reads from the simulation and
|
||||
returns a tensor. The observation manager concatenates them into a single
|
||||
vector for the policy. mjlab provides common terms in ``mjlab.envs.mdp``
|
||||
(joint positions, velocities, etc.), but you can always define your own.
|
||||
|
||||
The cartpole has two moving parts, so its physical state is fully
|
||||
described by two positions and two velocities:
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 22 12 66
|
||||
|
||||
* - Term
|
||||
- Dim
|
||||
- Description
|
||||
* - ``cart_pos``
|
||||
- 1
|
||||
- Where is the cart on the rail?
|
||||
* - ``pole_angle``
|
||||
- 2
|
||||
- Which way is the pole pointing? (cosine and sine)
|
||||
* - ``cart_vel``
|
||||
- 1
|
||||
- How fast is the cart moving?
|
||||
* - ``pole_vel``
|
||||
- 1
|
||||
- How fast is the pole rotating?
|
||||
|
||||
.. tip::
|
||||
|
||||
The pole angle is encoded as cosine and sine rather than a raw angle.
|
||||
MuJoCo's unlimited hinge does not wrap the angle, so as the pole
|
||||
spins the raw value keeps growing. Cosine and sine give the same
|
||||
output for the same physical angle regardless of how many rotations
|
||||
have occurred.
|
||||
|
||||
This is the one custom observation function:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
def pole_angle_cos_sin(env, asset_cfg) -> torch.Tensor:
|
||||
asset: Entity = env.scene[asset_cfg.name]
|
||||
angle = asset.data.joint_pos[:, asset_cfg.joint_ids]
|
||||
return torch.cat([torch.cos(angle), torch.sin(angle)], dim=-1)
|
||||
|
||||
.. note::
|
||||
|
||||
All data in mjlab is batched: tensors have shape
|
||||
``[num_envs, ...]`` because many environments run in parallel.
|
||||
Every function you write should accept and return tensors with this
|
||||
leading batch dimension.
|
||||
|
||||
To wire these up, we create ``ObservationTermCfg`` entries and group
|
||||
them. ``SceneEntityCfg`` scopes each function to specific joints on
|
||||
the entity:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
cart_cfg = SceneEntityCfg("cartpole", joint_names=("slider",))
|
||||
hinge_cfg = SceneEntityCfg("cartpole", joint_names=("hinge_1",))
|
||||
|
||||
cart_pos = ObservationTermCfg(
|
||||
func=joint_pos_rel, params={"asset_cfg": cart_cfg},
|
||||
)
|
||||
pole_angle = ObservationTermCfg(
|
||||
func=pole_angle_cos_sin, params={"asset_cfg": hinge_cfg},
|
||||
)
|
||||
cart_vel = ObservationTermCfg(
|
||||
func=joint_vel_rel, params={"asset_cfg": cart_cfg},
|
||||
)
|
||||
pole_vel = ObservationTermCfg(
|
||||
func=joint_vel_rel, params={"asset_cfg": hinge_cfg},
|
||||
)
|
||||
|
||||
Each term pairs a function with the parameters to call it with. Now
|
||||
we group them. The RL algorithm expects an ``"actor"`` and ``"critic"``
|
||||
group; they share the same terms here, but when you add noise later
|
||||
you can give the critic clean observations
|
||||
(asymmetric actor-critic [#aac]_).
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
actor_terms = {
|
||||
"cart_pos": cart_pos,
|
||||
"pole_angle": pole_angle,
|
||||
"cart_vel": cart_vel,
|
||||
"pole_vel": pole_vel,
|
||||
}
|
||||
|
||||
observations = {
|
||||
"actor": ObservationGroupCfg(actor_terms),
|
||||
"critic": ObservationGroupCfg({**actor_terms}),
|
||||
}
|
||||
|
||||
Actions: what the agent does
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
The agent outputs a single scalar: the force on the cart.
|
||||
``JointEffortActionCfg`` writes the policy output to the actuator's
|
||||
effort target. The ``XmlActuator`` passes it to MuJoCo's ``ctrl``
|
||||
buffer, which clamps it to [-1, 1] and multiplies by the gear ratio:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
actions = {
|
||||
"effort": JointEffortActionCfg(
|
||||
entity_name="cartpole",
|
||||
actuator_names=("slider",),
|
||||
scale=1.0,
|
||||
),
|
||||
}
|
||||
|
||||
Rewards: the training signal
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
Each reward term is a function that returns a scalar per environment.
|
||||
The reward manager computes a weighted sum of all terms each step.
|
||||
|
||||
The cartpole reward reproduces dm_control's smooth reward as a single
|
||||
multiplicative term:
|
||||
|
||||
.. math::
|
||||
|
||||
r = \underbrace{\frac{\cos\theta + 1}{2}}_{\text{upright}}
|
||||
\times \underbrace{\frac{1 + g(x)}{2}}_{\text{centered}}
|
||||
\times \underbrace{\frac{4 + q(u)}{5}}_{\text{small control}}
|
||||
\times \underbrace{\frac{1 + g(\dot\theta)}{2}}_{\text{small velocity}}
|
||||
|
||||
Each factor is between 0 and 1. The product is high only when all four
|
||||
conditions hold simultaneously, preventing the agent from trading off
|
||||
one factor against another.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
rewards = {
|
||||
"smooth_reward": RewardTermCfg(
|
||||
func=cartpole_smooth_reward,
|
||||
weight=1.0,
|
||||
params={"cart_cfg": cart_cfg, "hinge_cfg": hinge_cfg},
|
||||
),
|
||||
}
|
||||
|
||||
Terminations: when to stop
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
The cartpole has no failure states, so the only termination is a time
|
||||
limit. Setting ``time_out=True`` tells the RL algorithm this is a
|
||||
truncation, not a true terminal state, so it bootstraps the value
|
||||
function past the episode boundary:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
terminations = {
|
||||
"time_out": TerminationTermCfg(func=time_out, time_out=True),
|
||||
}
|
||||
|
||||
Events: resetting the state
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
At the start of each episode, reset events randomize joint positions and
|
||||
velocities around the initial state we defined in the entity:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
events = {
|
||||
"reset_slider": EventTermCfg(
|
||||
func=reset_joints_by_offset,
|
||||
mode="reset",
|
||||
params={
|
||||
"position_range": (-0.1, 0.1),
|
||||
"velocity_range": (-0.01, 0.01),
|
||||
"asset_cfg": SceneEntityCfg("cartpole", joint_names=("slider",)),
|
||||
},
|
||||
),
|
||||
"reset_hinge": EventTermCfg(
|
||||
func=reset_joints_by_offset,
|
||||
mode="reset",
|
||||
params={
|
||||
"position_range": (-0.034, 0.034),
|
||||
"velocity_range": (-0.01, 0.01),
|
||||
"asset_cfg": SceneEntityCfg("cartpole", joint_names=("hinge_1",)),
|
||||
},
|
||||
),
|
||||
}
|
||||
|
||||
The offsets are relative to the entity's initial state. For swingup the
|
||||
hinge starts at pi, so the noise keeps it near pointing down.
|
||||
|
||||
Snapping everything together
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
``ManagerBasedRlEnvCfg`` is where all the pieces come together. The
|
||||
scene holds the entity, and the config holds everything else:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
return ManagerBasedRlEnvCfg(
|
||||
scene=SceneCfg(
|
||||
terrain=TerrainEntityCfg(terrain_type="plane"),
|
||||
entities={"cartpole": _get_cartpole_cfg(swing_up=swing_up)},
|
||||
num_envs=1,
|
||||
env_spacing=4.0,
|
||||
),
|
||||
observations=observations,
|
||||
actions=actions,
|
||||
events=events,
|
||||
rewards=rewards,
|
||||
terminations=terminations,
|
||||
sim=SimulationCfg(
|
||||
mujoco=MujocoCfg(timestep=0.01, disableflags=("contact",)),
|
||||
),
|
||||
decimation=5,
|
||||
episode_length_s=50.0,
|
||||
)
|
||||
|
||||
``decimation=5`` means the physics runs five substeps per policy step,
|
||||
giving a 20 Hz control frequency. ``disableflags=("contact",)`` skips
|
||||
contact computation since cartpole has no collisions. ``num_envs=1`` is
|
||||
the default; override it from the CLI with ``--num-envs``.
|
||||
|
||||
|
||||
Registration and training
|
||||
-------------------------
|
||||
|
||||
The last step is to register the task so it can be launched by name.
|
||||
Each registration pairs an environment config with an RL config that
|
||||
specifies the network architecture and PPO hyperparameters. For cartpole
|
||||
a small network of two 64 unit hidden layers is plenty. The full RL
|
||||
config is in ``cartpole_env_cfg.py`` alongside the environment config.
|
||||
|
||||
This goes in ``__init__.py``:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
register_mjlab_task(
|
||||
task_id="Mjlab-Cartpole-Swingup",
|
||||
env_cfg=cartpole_swingup_env_cfg(),
|
||||
play_env_cfg=cartpole_swingup_env_cfg(play=True),
|
||||
rl_cfg=cartpole_ppo_runner_cfg(),
|
||||
)
|
||||
|
||||
Train:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
uv run train Mjlab-Cartpole-Swingup --env.scene.num-envs 4096
|
||||
|
||||
Play back a trained checkpoint, either from a local file or a W&B run:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
uv run play Mjlab-Cartpole-Swingup --checkpoint-file logs/rsl_rl/cartpole/model_500.pt
|
||||
uv run play Mjlab-Cartpole-Swingup --wandb-run-path <user/project/run_id>
|
||||
|
||||
.. figure:: ../_static/tutorials/cartpole_training_curve.png
|
||||
:width: 70%
|
||||
:align: center
|
||||
:alt: Cartpole swingup training curve
|
||||
|
||||
Mean reward over 5 seeds (shaded: one standard deviation).
|
||||
|
||||
Config fields can be overridden from the CLI:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
uv run train Mjlab-Cartpole-Swingup \
|
||||
--num-envs 8192 \
|
||||
--agent.algorithm.learning-rate 3e-4 \
|
||||
--agent.algorithm.entropy-coef 0.005
|
||||
|
||||
|
||||
Next steps
|
||||
----------
|
||||
|
||||
**Add observation noise.** The current config has no noise, so the
|
||||
policy is brittle. Add noise to any observation term to train a more
|
||||
robust policy:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.utils.noise import UniformNoiseCfg
|
||||
|
||||
ObservationTermCfg(
|
||||
func=joint_pos_rel,
|
||||
params={"asset_cfg": cart_cfg},
|
||||
noise=UniformNoiseCfg(n_min=-0.05, n_max=0.05),
|
||||
)
|
||||
|
||||
**Randomize the physics.** Use the :ref:`domain_randomization` system
|
||||
to vary pole mass or joint damping across environments, training a
|
||||
policy that transfers across physical variations.
|
||||
|
||||
**Explore other tasks.** The library ships with locomotion,
|
||||
manipulation, and motion tracking tasks you can run out of the box:
|
||||
``Mjlab-Velocity-Flat-Unitree-Go1``, ``Mjlab-Lift-Cube-Yam``, and
|
||||
``Mjlab-Tracking-Flat-Unitree-G1``, among others. Reading their source
|
||||
shows how more complex observation and reward structures are composed.
|
||||
|
||||
**Build something new.** The cartpole is intentionally minimal. Once
|
||||
you are comfortable with the pieces, try designing your own robot model
|
||||
and task from scratch. The same pattern applies regardless of how
|
||||
complex the system becomes.
|
||||
|
||||
.. rubric:: References
|
||||
|
||||
.. [#aac] Pinto, L., Andrychowicz, M., Welinder, P., Zaremba, W., & Abbeel, P. (2018). `Asymmetric Actor Critic for Image-Based Robot Learning <https://www.roboticsproceedings.org/rss14/p08.pdf>`_. *Robotics: Science and Systems XIV*.
|
||||
@@ -0,0 +1,315 @@
|
||||
.. _viewers:
|
||||
|
||||
Viewers
|
||||
=======
|
||||
|
||||
mjlab ships two interactive viewers for evaluating trained policies and
|
||||
debugging environment behavior: a **native viewer** built on MuJoCo's
|
||||
`passive viewer <https://mujoco.readthedocs.io/en/stable/python.html#passive-viewer>`_
|
||||
that opens a desktop window, and a `Viser <https://viser.studio/main/>`_ **viewer** that runs in the
|
||||
browser. Both share a common ``ViewerConfig`` and execute the same
|
||||
simulation loop; they differ in interface, feature set, and where they
|
||||
shine.
|
||||
|
||||
Launching a viewer
|
||||
------------------
|
||||
|
||||
The ``play`` script accepts a ``--viewer`` flag:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
# Desktop window (MuJoCo native viewer).
|
||||
uv run play Mjlab-Velocity-Flat-Unitree-G1 --viewer native \
|
||||
--wandb-run-path your-entity/your-project/run_id
|
||||
|
||||
# Browser-based viewer (opens localhost:8080).
|
||||
uv run play Mjlab-Velocity-Flat-Unitree-G1 --viewer viser \
|
||||
--wandb-run-path your-entity/your-project/run_id
|
||||
|
||||
The default is ``auto``, which selects native when a display server is
|
||||
available (``DISPLAY`` or ``WAYLAND_DISPLAY``) and falls back to Viser
|
||||
on headless machines.
|
||||
|
||||
For quick exploration without a trained checkpoint, pass ``--agent zero``
|
||||
or ``--agent random`` to use a dummy policy:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
uv run play Mjlab-Velocity-Flat-Unitree-G1 --agent zero --viewer viser
|
||||
|
||||
|
||||
Viewer configuration
|
||||
--------------------
|
||||
|
||||
Camera position, tracking target, and rendering options live in
|
||||
``ViewerConfig``, set through the ``viewer`` field of
|
||||
``ManagerBasedRlEnvCfg``:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.viewer import ViewerConfig
|
||||
|
||||
viewer = ViewerConfig(
|
||||
lookat=(0.0, 0.0, 0.5),
|
||||
distance=3.0,
|
||||
elevation=-20.0,
|
||||
azimuth=135.0,
|
||||
)
|
||||
|
||||
The ``origin_type`` field controls the camera reference frame:
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 22 78
|
||||
|
||||
* - Origin type
|
||||
- Behavior
|
||||
* - ``WORLD``
|
||||
- Free camera anchored at world origin (default).
|
||||
* - ``ASSET_ROOT``
|
||||
- Camera tracks the root body of the entity named by
|
||||
``entity_name``. Good for locomotion tasks where the robot moves
|
||||
through the world.
|
||||
* - ``ASSET_BODY``
|
||||
- Camera tracks a specific body (``body_name``) within the entity
|
||||
named by ``entity_name``. Useful for close-up views of an
|
||||
end-effector or head.
|
||||
|
||||
Example with asset tracking:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
viewer = ViewerConfig(
|
||||
origin_type=ViewerConfig.OriginType.ASSET_ROOT,
|
||||
entity_name="robot",
|
||||
distance=2.5,
|
||||
elevation=-15.0,
|
||||
)
|
||||
|
||||
Additional fields:
|
||||
|
||||
- ``enable_shadows`` and ``enable_reflections`` toggle rendering
|
||||
quality.
|
||||
- ``height`` and ``width`` set the offscreen render resolution (used
|
||||
by ``OffscreenRenderer`` and video recording).
|
||||
- ``env_idx`` selects which environment to display at startup.
|
||||
|
||||
|
||||
Native MuJoCo viewer
|
||||
---------------------
|
||||
|
||||
.. figure:: _static/native_viewer.png
|
||||
:width: 100%
|
||||
:align: center
|
||||
:alt: Native MuJoCo viewer with reward plots
|
||||
|
||||
The native viewer opens MuJoCo's
|
||||
`passive viewer <https://mujoco.readthedocs.io/en/stable/python.html#passive-viewer>`_
|
||||
in a desktop window. It provides the fastest, most faithful rendering
|
||||
with full MuJoCo visual fidelity. Choose this viewer for local
|
||||
iteration and interactive perturbation testing. The
|
||||
MuJoCo team has a
|
||||
`video tutorial <https://www.youtube.com/watch?v=P83tKA1iz2Y>`_
|
||||
covering the viewer's built-in controls and navigation.
|
||||
|
||||
**Keyboard controls.**
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 18 82
|
||||
|
||||
* - Key
|
||||
- Action
|
||||
* - ``Space``
|
||||
- Pause or resume simulation.
|
||||
* - ``Enter``
|
||||
- Reset the environment.
|
||||
* - ``+`` / ``-``
|
||||
- Increase or decrease playback speed.
|
||||
* - ``<`` / ``>``
|
||||
- Cycle through environments (when ``num_envs > 1``).
|
||||
* - ``A``
|
||||
- Toggle rendering all environments simultaneously. Debug
|
||||
visualization draws for all environments when this is active.
|
||||
* - ``P``
|
||||
- Toggle reward plots.
|
||||
* - ``R``
|
||||
- Toggle debug visualization.
|
||||
|
||||
**Reward plots.**
|
||||
Press ``P`` to display per-term reward curves in a strip along the right
|
||||
edge of the window. Each term gets its own plot with an autoscaling
|
||||
y-axis. The plots update live and clear on environment reset. This is
|
||||
the fastest way to diagnose which reward terms dominate or misbehave
|
||||
during a rollout.
|
||||
|
||||
**Interactive perturbations.**
|
||||
Click and drag any body in the scene to apply external forces during
|
||||
playback. The force transfers into the simulation on the next step,
|
||||
making it easy to test balance recovery, grasp robustness, or
|
||||
disturbance rejection without writing any code.
|
||||
|
||||
Mouse perturbation forces are kept separate from programmatic forces
|
||||
(e.g. ``apply_body_impulse``) by routing them through different MuJoCo
|
||||
channels: programmatic forces use ``xfrc_applied`` (Cartesian body
|
||||
forces), while mouse forces are converted to ``qfrc_applied``
|
||||
(generalized joint forces) via ``mj_applyFT``. Both channels are summed
|
||||
during forward dynamics, so they coexist without conflict.
|
||||
|
||||
**Domain randomization visualization.**
|
||||
The native viewer syncs all visual DR fields from GPU to CPU each
|
||||
frame. Randomized geom colors, sizes, positions, material colors,
|
||||
body poses, camera parameters, light positions, and inertia ellipsoids
|
||||
all render faithfully. If a DR event changes a visual property, the
|
||||
native viewer shows it.
|
||||
|
||||
|
||||
Viser (browser-based)
|
||||
---------------------
|
||||
|
||||
.. figure:: _static/viser_viewer.png
|
||||
:width: 100%
|
||||
:align: center
|
||||
:alt: Viser browser-based viewer
|
||||
|
||||
The `Viser <https://viser.studio/main/>`_ viewer opens an interactive 3D scene
|
||||
in the browser at ``localhost:8080``. It works on remote machines over
|
||||
SSH tunnels, making it the natural choice for headless GPU servers and
|
||||
shared debugging sessions. Its web-based architecture makes it far more customizable than the
|
||||
native viewer. It also provides dedicated panels for camera sensor
|
||||
output that the native viewer does not.
|
||||
|
||||
**Tab-based interface.**
|
||||
The sidebar organizes controls into tabs:
|
||||
|
||||
- **Controls**: play/pause, reset, speed adjustment, environment
|
||||
selection, and display settings (FOV, contacts, geom groups, camera
|
||||
tracking).
|
||||
- **Rewards**: live per-term reward charts, toggled by a checkbox.
|
||||
- **Metrics**: live per-term metric charts when a ``MetricsManager``
|
||||
is present.
|
||||
- **Camera Feeds**: live RGB and depth image panels for every
|
||||
``CameraSensor`` in the scene. A depth scale slider adjusts the
|
||||
visualization range, and a frustum toggle draws the camera's field of
|
||||
view in the 3D scene.
|
||||
- **Groups**: show or hide MuJoCo geom and site groups.
|
||||
|
||||
**Camera sensor integration.**
|
||||
Viser auto-discovers all ``CameraSensor`` instances in the scene and
|
||||
displays their output as live image panels. Each camera also gets a
|
||||
frustum visualization in the 3D viewport, so you can see exactly what
|
||||
the sensor covers. This makes Viser the best tool for debugging camera
|
||||
placement, field of view, and depth sensing.
|
||||
|
||||
**Contact visualization.**
|
||||
The Controls tab exposes contact rendering options. When enabled, contact
|
||||
points appear as colored markers and contact forces as red arrows,
|
||||
giving immediate visual feedback on collision behavior.
|
||||
|
||||
.. note::
|
||||
|
||||
The Viser viewer does not support interactive perturbations (applying
|
||||
wrenches to bodies). Use the native viewer for that, or set up
|
||||
perturbations via :ref:`events <events>`.
|
||||
|
||||
.. note::
|
||||
|
||||
Viser reads world-space body positions directly from GPU each frame,
|
||||
so body poses update correctly. However, ``geom_rgba`` and
|
||||
``geom_size`` are baked into GLB meshes at scene construction time
|
||||
and will not reflect per-world DR changes. This will be addressed in
|
||||
a future release. For now, use the native viewer when you need to
|
||||
verify visual DR.
|
||||
|
||||
|
||||
Debug visualization
|
||||
-------------------
|
||||
|
||||
Both viewers support a shared ``DebugVisualizer`` interface that manager
|
||||
terms can draw into. Available primitives:
|
||||
|
||||
- **Arrows**: velocity commands, force vectors, heading indicators.
|
||||
- **Spheres**: target positions, contact points.
|
||||
- **Cylinders**: limb targets, distance markers.
|
||||
- **Ellipsoids**: inertia visualization.
|
||||
- **Coordinate frames**: body frame orientation, end-effector targets.
|
||||
- **Ghost meshes**: transparent renderings of a robot at a target pose,
|
||||
useful for motion tracking or goal visualization.
|
||||
|
||||
In the native viewer, toggle debug visualization with ``R`` and press
|
||||
``A`` to show debug draws for all environments at once. In Viser, the
|
||||
Controls tab has toggles for both. The ``DebugVisualizer`` abstraction
|
||||
means that reward and command terms draw once, and both viewers display
|
||||
the result without any viewer-specific code.
|
||||
|
||||
|
||||
Offscreen renderer
|
||||
------------------
|
||||
|
||||
For recording videos without a display, ``OffscreenRenderer`` renders
|
||||
frames using MuJoCo's offscreen rendering pipeline. It supports the
|
||||
same ``ViewerConfig`` camera configuration and accepts a debug
|
||||
visualization callback. The renderer is hard-capped at 32 environments
|
||||
to keep memory and rendering time manageable.
|
||||
|
||||
The ``play`` script uses ``OffscreenRenderer`` when the ``--video`` flag
|
||||
is set:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
uv run play Mjlab-Velocity-Flat-Unitree-G1 --video --video-length 300 \
|
||||
--wandb-run-path your-entity/your-project/run_id
|
||||
|
||||
|
||||
Quick comparison
|
||||
----------------
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 24 38 38
|
||||
|
||||
* -
|
||||
- Native
|
||||
- Viser
|
||||
* - Interface
|
||||
- Desktop window
|
||||
- Browser (``localhost:8080``)
|
||||
* - Best for
|
||||
- Local iteration, perturbations
|
||||
- Customization, remote dev, cameras
|
||||
* - Reward plots
|
||||
- ``P`` key, right-side strip
|
||||
- Rewards tab, uPlot charts
|
||||
* - Metrics plots
|
||||
-
|
||||
- Metrics tab
|
||||
* - Camera feeds
|
||||
-
|
||||
- Auto-discovered, with frustum
|
||||
* - Perturbations
|
||||
- Click and drag
|
||||
-
|
||||
* - DR visualization
|
||||
- Full (all visual fields synced)
|
||||
- Partial (body poses only)
|
||||
* - Contact rendering
|
||||
-
|
||||
- Contact points and forces
|
||||
* - Multi-environment
|
||||
- ``<`` ``>`` to cycle, ``A`` for all
|
||||
- Dropdown selector
|
||||
|
||||
|
||||
Citation
|
||||
--------
|
||||
|
||||
If you use the Viser viewer in your research, consider citing:
|
||||
|
||||
.. code-block:: bibtex
|
||||
|
||||
@article{yi2025viser,
|
||||
title={Viser: Imperative, web-based 3d visualization in python},
|
||||
author={Yi, Brent and Kim, Chung Min and Kerr, Justin and Wu, Gina and Feng, Rebecca and Zhang, Anthony and Kulhanek, Jonas and Choi, Hongsuk and Ma, Yi and Tancik, Matthew and Kanazawa, Angjoo},
|
||||
journal={arXiv preprint arXiv:2507.22885},
|
||||
year={2025}
|
||||
}
|
||||