Robocon/RC_WheelLeg
4
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[train] 更新新MJCF与第一版完整训练框架

Browse Source
This commit is contained in:
eastzio
2026-07-27 12:31:21 +08:00
parent dfd646cf47
commit b08956aec7
116 changed files with 7649 additions and 6760 deletions
Download Patch File Download Diff File Expand all files Collapse all files
05_software/train/rc_mjlab/mjlab/docs/source/actuators.rst
+26
View File
@@ -61,6 +61,7 @@ MuJoCo's integrator handles velocity-dependent forces.
**Built-in actuators** (``BuiltinPositionActuator``,
``BuiltinVelocityActuator``, ``BuiltinMotorActuator``,
``BuiltinPdActuator``, ``BuiltinDcMotorActuator``,
``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
@@ -119,6 +120,31 @@ control.
**BuiltinMotorActuator**: Creates ``<motor>`` actuators for direct torque
control.
**BuiltinPdActuator**: Native PD that closes on both a position and a
velocity target, implemented as paired ``<position>`` + ``<velocity>``
actuators summing to ``kp * (p_target - q) + kd * (v_target - qdot)``.
``BuiltinPositionActuator`` puts kd on the ``<position>`` element and
implicitly assumes a zero velocity reference; use this when the policy
emits a non-zero velocity target. Native delivery lets
``implicit`` / ``implicitfast`` see the kd term in their velocity update,
unlike ``IdealPdActuator`` which forwards Python-computed torque through
an opaque ``<motor>``.
**BuiltinDcMotorActuator**: Wraps MuJoCo's native
`<dcmotor> <https://mujoco.readthedocs.io/en/stable/XMLreference.html#actuator-dcmotor>`_
element. Torque is ``tau = K * (V - K * omega) / R``; the back-EMF runs
through the native bias path, so ``implicit`` / ``implicitfast`` pick up
its velocity derivative as effective damping. Three input modes pick what
``ctrl`` carries: VOLTAGE drives the motor directly; POSITION / VELOCITY
close an internal PID (with anti-windup and slew limiting) against a
single setpoint, whose Vmax-clamped output becomes torque. POSITION mode
pins v_target = 0 (the kd term acts on raw velocity). Optional physics:
inductance,
thermal model with I^2R heating, cogging ripple, LuGre friction.
``DcMotorActuator`` (the explicit version) is a software PD with a
velocity-dependent torque clamp on top of a ``<motor>``; this is the real
electrical model.
**BuiltinMuscleActuator**: Creates ``<muscle>`` actuators for
biologically-inspired muscle dynamics with force-length-velocity
characteristics.
05_software/train/rc_mjlab/mjlab/docs/source/api/actuator.rst
+41
View File
@@ -18,6 +18,13 @@ mjlab.actuator
- :class:`BuiltinPositionActuatorCfg`
- :class:`BuiltinVelocityActuator`
- :class:`BuiltinVelocityActuatorCfg`
- :class:`BuiltinPdActuator`
- :class:`BuiltinPdActuatorCfg`
- :class:`BuiltinDcMotorActuator`
- :class:`BuiltinDcMotorActuatorCfg`
- :class:`DcMotorInputMode`
- :class:`DcMotorDatasheetParams`
- :class:`DcMotorPhysicalParams`
- :class:`BuiltinMuscleActuator`
- :class:`BuiltinMuscleActuatorCfg`
- :class:`XmlActuator`
@@ -84,6 +91,40 @@ Builtin Actuators
:undoc-members:
.. autoclass:: BuiltinPdActuator
:members:
:show-inheritance:
.. autoclass:: BuiltinPdActuatorCfg
:members:
:exclude-members: __init__
:undoc-members:
.. autoclass:: BuiltinDcMotorActuator
:members:
:show-inheritance:
.. autoclass:: BuiltinDcMotorActuatorCfg
:members:
:exclude-members: __init__
:undoc-members:
.. autoclass:: DcMotorInputMode
:members:
:show-inheritance:
.. autoclass:: DcMotorDatasheetParams
:members:
:exclude-members: __init__
:undoc-members:
.. autoclass:: DcMotorPhysicalParams
:members:
:exclude-members: __init__
:undoc-members:
.. autoclass:: BuiltinMuscleActuator
:members:
:show-inheritance:
05_software/train/rc_mjlab/mjlab/docs/source/changelog.rst
+167 -17
View File
@@ -8,6 +8,89 @@ Upcoming version (not yet released)
Added
^^^^^
- Added ``BuiltinDcMotorActuator``, a native MuJoCo ``<dcmotor>`` wrapper.
Supports voltage / position / velocity input modes with back-EMF,
configurable motor constants, and optional integral, slew, inductance,
thermal, LuGre, and cogging extensions.
- Added ``scale_with_difficulty`` to ``HfRandomUniformTerrainCfg``. When
enabled, the noise amplitude scales with difficulty (flat at 0, full
``noise_range`` at 1) so the terrain progresses in a curriculum. Defaults to
``False``, preserving the previous difficulty-independent behavior.
Changed
^^^^^^^
- Bumped ``rsl-rl-lib`` from 5.2.0 to 5.4.0.
- Curriculum-mode terrain difficulty is now deterministic across rows
and reaches the configured ``difficulty_range`` endpoints
(:issue:`1027`).
- Heightfield terrains now color by absolute height with a diverging palette
(cool below the ground plane, green at ground level, warm above) on a fixed
scale, replacing the per-patch normalization. Color is now consistent across
terrains, and low-amplitude terrain such as ``random_rough`` reads as gently
tinted ground instead of high-contrast noise.
- ``BoxNestedRingsTerrainCfg`` now builds uniform-height concentric ridges
whose separating gaps widen with difficulty, replacing the random per-ring
heights. Rings are colored by height (like the other terrains) and the outer
border matches the ring height.
- Terrain generation no longer prints timing information to stdout.
Fixed
^^^^^
- Fixed ``select_gpus`` crashing when ``CUDA_VISIBLE_DEVICES`` contains MIG UUIDs instead of numeric indices.
- Fixed pyramid-stairs terrains (``BoxPyramidStairsTerrainCfg``,
``BoxInvertedPyramidStairsTerrainCfg``, and ``BoxOpenStairsTerrainCfg``)
leaving an empty, geometry-free border at difficulty 0, where the step
height collapses to zero. The flat border frame is now always generated as
solid geometry flush with the ground (:issue:`1033`).
- Fixed ``HfPerlinNoiseTerrainCfg`` failing to compile at difficulty 0, where
the target height collapses to zero and MuJoCo rejects the non-positive
heightfield size.
- Fixed ``BoxRandomGridTerrainCfg`` producing NaN colors (and failing to build)
at difficulty 0, where the grid height is zero and the color normalization
divided by zero.
- Fixed the center platform z-fighting with surrounding geometry in
``BoxRandomGridTerrainCfg`` (grid cells were left underneath the platform) and
``BoxRandomSpreadTerrainCfg`` (the platform duplicated the floor surface).
- Fixed ``BoxNarrowBeamsTerrainCfg`` square platform corners protruding between
the beams at high difficulty; the platform now shrinks to stay within the
beams' angular coverage.
- Fixed ``BoxSteppingStonesTerrainCfg`` reconfiguring abruptly at a difficulty
threshold, where the stone grid re-tiled as its spacing crossed an integer
boundary, and leaving an oversized gap around the center platform. The grid is
now difficulty-independent and the platform snaps to it as a clean island.
- Fixed ``train --video``, ``play``, and ``demo`` crashing with ``OpenGL
platform library not loaded`` on headless Linux hosts that don't pre-set
``MUJOCO_GL``. The default is now applied in ``mjlab/__init__.py`` (Linux
only) so it takes effect before mujoco's GL backend selection runs.
Version 1.4.0 (May 26, 2026)
----------------------------
Added
^^^^^
- Added ``BuiltinPdActuator``, the implicit-integration version of
``IdealPdActuator``. Same interface (position + velocity targets,
kp/kd gains), but expresses the PD as native MuJoCo ``<position>``
and ``<velocity>`` elements so the ``implicit`` / ``implicitfast``
integrators include the kp/kd derivatives in their velocity update.
The actuator stays stable at gain/timestep combinations where
explicit Python PD would diverge, which matters when you want to
run a real motor's stiff on-board PD gains in sim. ``effort_limit``
is enforced as a sum-clamp on the two PD terms via
``jnt_actfrcrange`` (or ``tendon_actfrcrange``). Supported by
``dr.pd_gains`` and ``dr.effort_limits``.
- Added ``mdp.projected_gravity_from_sensor``, an observation that derives
projected gravity from a ``framezaxis`` up-vector sensor (negated) rather
than from the root body orientation. Unlike ``mdp.projected_gravity``, it
reflects the sensor's site frame, so it can observe IMU mounting domain
randomization (e.g. via ``dr.site_quat``). Go1 and G1 ship an
``imu_upvector`` sensor for this.
- Added ``DebugVisualizer.add_box`` for drawing an axis-oriented box
primitive, mirroring ``add_ellipsoid``. Supported by both the native
and Viser viewers. ``size`` is the box half-extents (:issue:`992`).
- 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
@@ -21,22 +104,41 @@ Added
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.
- Added per-world mesh variant support via ``VariantEntityCfg``. 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). Variants are passed as a ``dict[str, Callable]`` of named
spec callables; the optional ``assignment`` field controls how worlds
map to variants and accepts ``None`` (uniform), a ``dict[str, float]``
of per-variant weights, or a custom ``Callable[[int], Sequence[int]]``.
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:`heterogeneous_worlds` for usage. With help from @XiangruiJiang.
- Per-world mesh variants now support per-variant materials and textures.
Each variant can reference its own named material, which is automatically
prefixed and scattered via ``geom_matid`` alongside the existing
``geom_dataid`` table. Variants without a material get ``matid = -1``.
Contribution by @omarrayyann.
Changed
^^^^^^^
- ``Entity`` now raises a clear error at construction when its spec contains
more than one freejoint. An entity models a single system rooted at one
body, so it has at most one freejoint; a second one was previously accepted
silently and only surfaced later as a cryptic shape mismatch when writing
root state. Model each detached floating body as its own entry in
``SceneCfg.entities`` instead.
- Changed ``compute_root_relative_mpkpe`` to re-anchor the reference to the
robot's root each step, removing yaw drift as well as translation so it
measures intrinsic body pose error.
- Changed ``compute_joint_velocity_error`` from an L2 norm to a per-joint
RMS, so it no longer scales with the number of joints.
- 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.
@@ -68,15 +170,63 @@ Changed
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`).
- Event functions now share a single ``resolve_env_ids`` helper to expand
``env_ids=None`` to all environments, replacing five copies of the same
guard. ``push_by_setting_velocity`` and ``apply_external_force_torque``
accept ``env_ids=None`` too, so they work as global-time interval terms.
Documented when to use ``apply_external_force_torque`` (a constant,
self-managed wrench) versus ``apply_body_impulse`` (transient, automatic
impulses) versus ``push_by_setting_velocity`` (an instantaneous velocity
kick).
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`).
- Removed use of deprecated ``warp-lang`` symbols (``wp.context.runtime``
and ``wp.context.Device``) that were dropped in newer ``warp-lang``
releases, causing ``AttributeError: module 'warp' has no attribute
'context'`` at import/runtime. mjlab now uses
``wp.get_cuda_driver_version()`` and ``wp.Device`` instead
(:issue:`967`). Contribution by @rdeits.
- Fixed the tracking ``evaluate`` script scoring each metric against the
next motion frame; the reference is now snapshotted before each step to
match the reward.
- Fixed the tracking end-effector metrics silently scoring zero for an
unknown body name; they now raise ``ValueError``.
- Fixed ``compute_mpkpe`` measuring root-relative instead of global error;
it now uses the global reference ``body_pos_w`` (:issue:`1006`).
- Fixed heavy flicker in offscreen training videos on rough-terrain tasks.
The renderer recomputed its context "neighbor" robots every frame from
``env_origins``, which the terrain curriculum mutates on reset, so the
neighbor set kept changing and robots popped in and out. The neighbor
set is now computed once and cached (:issue:`979`).
- Fixed command delay only applying to an actuator's position target.
``IdealPdActuator`` and ``DcMotorActuator`` also use velocity and effort, which
arrived undelayed and out of sync; all command targets now share one delay.
Zero-reference setups are unaffected.
- Fixed duplicate random seeds across nodes in multi-node training. The
per-process seed offset in ``scripts/train.py`` now uses the global
``RANK`` instead of ``LOCAL_RANK``. Contribution by @bd-pdomanico.
- Fixed ``apply_body_impulse`` firing an impulse on the very first step (and
the first step after every reset) instead of starting with a cooldown as
documented. The cooldown is now sampled lazily on the first call so impulse
timing is decorrelated from episode resets (:issue:`973`).
- Fixed ``dr.pd_gains`` and ``dr.effort_limits`` silently no-oping when
passed an ``Operation`` object (e.g. ``dr.scale``) instead of a string.
Both functions now accept ``Operation | str`` like every other DR event
and raise ``ValueError`` for unsupported operations (:issue:`971`).
- Fixed ``ContactSensor`` with ``global_frame=True`` and
``reduce`` ∈ {``"none"``, ``"mindist"``, ``"maxforce"``} producing forces
rotated onto the wrong axis. The contact-frame→world rotation matrix had
its columns ordered ``[tangent, tangent2, normal]`` instead of
``[normal, tangent, tangent2]``, projecting the normal-force component
onto a tangent direction. Contribution by @bd-pdomanico.
- Fixed ``extras["log"]`` entries written by reward terms (e.g. ``Metrics/*``
values in velocity tasks) being silently discarded on any step where at
least one environment resets. ``_reset_idx`` was clearing the dict after
``reward_manager.compute()`` had already populated it. The clear now
happens at the top of ``step()`` and ``reset()`` so that all entries
survive (:issue:`957`).
- 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`` /
05_software/train/rc_mjlab/mjlab/docs/source/entity/index.rst
+3 -3
View File
@@ -197,8 +197,8 @@ 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
^^^^^^^^^^^^^^^^^^^^^^^
Heterogeneous worlds
^^^^^^^^^^^^^^^^^^^^
For scenes that need different mesh assets in different parallel worlds
(for example, training a manipulation policy that generalizes across
@@ -206,7 +206,7 @@ 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`.
viewers stay consistent. See :ref:`heterogeneous_worlds`.
Subclassing Entity
^^^^^^^^^^^^^^^^^^
05_software/train/rc_mjlab/mjlab/docs/source/entity/per_world_mesh.rst
+391 -101
View File
@@ -1,50 +1,35 @@
.. _per_world_mesh:
.. _heterogeneous_worlds:
Mesh Variants
=============
Heterogeneous Worlds
====================
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.
mjlab can run a single batched simulation in which different parallel
worlds use different mesh assets for the same logical entity. World 0
may simulate a cube, world 1 a sphere, world 2 a bowl. All worlds
share the same compiled scene and the same body and joint structure;
only the meshes and the per-geom attributes that travel with them
(friction, contact bits, mass, density, and a few more) differ across
worlds. Articulated props work too (you can have a hinge or slide
below the variant's root), as long as the joint topology matches
across variants. The feature is exposed through ``VariantEntityCfg``.
The full breakdown of what can and cannot vary across variants is in
the next section.
How it works
------------
Quickstart
----------
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:
Say you want some parallel worlds to hold a sphere and others to hold
a cone, with a single shared scene running both at once. Define each
variant as a function that returns an ``MjSpec``, then group them
under one ``VariantEntityCfg``:
.. code-block:: python
import mujoco
from mjlab.entity import EntityCfg, VariantCfg, VariantEntityCfg
from mjlab.entity import EntityCfg, VariantEntityCfg
def make_sphere_spec() -> mujoco.MjSpec:
spec = mujoco.MjSpec()
@@ -56,40 +41,267 @@ A minimal two-variant config:
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.
def make_cone_spec() -> mujoco.MjSpec:
spec = mujoco.MjSpec()
mesh = spec.add_mesh(name="visual")
mesh.make_cone(nedge=16, radius=0.04)
body = spec.worldbody.add_body(name="prop")
body.add_freejoint()
body.add_geom(type=mujoco.mjtGeom.mjGEOM_MESH, meshname="visual")
return spec
object_cfg = VariantEntityCfg(
variants={
"sphere": VariantCfg(spec_fn=make_sphere_spec, weight=1.0),
"cone": VariantCfg(spec_fn=make_cone_spec, weight=2.0),
"sphere": make_sphere_spec,
"cone": make_cone_spec,
},
assignment={"cone": 2.0}, # twice as many cones as spheres
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.
Plug the variant entity into a :ref:`scene` exactly like a regular
``EntityCfg``:
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``.
.. code-block:: python
from mjlab.scene import SceneCfg
scene_cfg = SceneCfg(
num_envs=4096,
entities={"object": object_cfg},
)
Twice as many worlds will hold a cone as a sphere. Variants not listed
in the ``assignment`` dict default to weight 1.0; omit ``assignment``
entirely for uniform allocation across all variants.
What variants can differ in
---------------------------
**Free to vary across variants:** the mesh asset assigned to each
slot, the number of mesh geoms per ``(body, role)`` bucket on the
variant body (one variant can have more collision meshes than
another), the per-mesh-geom attributes that travel with the mesh
(friction, contact bits, mass, density, ``condim``, and a handful of
others), and explicit body inertial values within whichever single
inertial mode the variants agree on per body.
**Must match across variants:** the body tree, joint topology,
primitive (non-mesh) geoms, and any actuators / sensors / tendons /
equalities. Variants must also agree on the inertial representation
per body (mesh-derived, diagonal, or fullinertia), and may not use the
reserved ``mjlab/pad/`` name prefix on any element. Variant entities
must also be floating-base: the root body declares a freejoint.
The validator runs at entity build time and raises ``ValueError``
naming the offending variant and the exact mismatch.
How variants are assembled
--------------------------
mjlab merges every variant's mesh assets into a single ``MjSpec`` and
gives the variant body enough mesh-geom *slots* to cover the maximum
mesh count any variant uses for each ``(body, role)`` bucket. A slot
is identified by ``(body_path, role, ordinal)``. ``role`` is "visual"
or "collision", derived from ``contype``/``conaffinity``;
mujoco_warp's ``geom_contype``/``geom_conaffinity`` are 1D shared
(not per-world), so a slot's role is fixed across worlds by
construction.
A worked example
~~~~~~~~~~~~~~~~
Say variant ``sphere`` has 1 visual mesh geom and 2 collision mesh
geoms on the prop body, and variant ``cone`` has 1 visual mesh geom
and 4 collision mesh geoms on the same body.
.. code-block:: text
sphere variant body cone variant body
------------------- -------------------
prop body prop body
[visual] sphere_vis [visual] cone_vis
[coll] sphere_col_0 [coll] cone_col_0
[coll] sphere_col_1 [coll] cone_col_1
[coll] cone_col_2
[coll] cone_col_3
mjlab walks each variant's body tree, buckets mesh geoms by
``(body_path, role)``, and lays the union out as slots:
.. list-table::
:header-rows: 1
:widths: 8 18 8 12 27 27
* - Slot
- body_path
- role
- ordinal
- sphere fills with
- cone fills with
* - 0
- /prop
- visual
- 0
- sphere_vis
- cone_vis
* - 1
- /prop
- collision
- 0
- sphere_col_0
- cone_col_0
* - 2
- /prop
- collision
- 1
- sphere_col_1
- cone_col_1
* - 3
- /prop
- collision
- 2
- *(unfilled)*
- cone_col_2
* - 4
- /prop
- collision
- 3
- *(unfilled)*
- cone_col_3
Five slots total. The merged scene's prop body has five mesh geoms:
slot 0 plus four collision slots (the union of sphere's two and
cone's four). At merge time, every variant's mesh asset is added to
the merged spec under a unique name (e.g.
``sphere/sphere_vis``, ``cone/cone_col_2``).
The merged scene compiles once into a single canonical ``MjModel``
that every world in the batch agrees on layout-wise: same nbody,
ngeom, same body and geom IDs. mjlab's per-world overrides on top of
that one model are what make worlds heterogeneous.
What each world sees at runtime
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Worlds where ``sphere`` is active see only its three meshes; the two
extra collision slots are disabled via per-world ``geom_dataid = -1``,
and mujoco_warp skips them. Worlds where ``cone`` is active see all
five meshes wired up.
.. list-table::
:header-rows: 1
:widths: 14 14 12 12 12 12 12
* - World
- variant
- slot 0
- slot 1
- slot 2
- slot 3
- slot 4
* - 0
- sphere
- sphere_vis
- sphere_col_0
- sphere_col_1
- **off (-1)**
- **off (-1)**
* - 1
- cone
- cone_vis
- cone_col_0
- cone_col_1
- cone_col_2
- cone_col_3
Three categories of per-world override carry the variation:
* **geom_dataid** is a ``(num_envs, ngeom)`` table. Its row for
world W picks which compiled mesh each slot points at. ``-1`` is
the "skip me" sentinel mujoco_warp already understands.
* **Mesh-derived fields** (``geom_size``, ``geom_rbound``,
``geom_aabb``, ``geom_pos``, ``geom_quat``, ``body_mass``,
``body_subtreemass``, ``body_inertia``, ``body_invweight0``,
``body_ipos``, ``body_iquat``) are stored as ``(num_envs, ...)``
arrays. The values for sphere worlds reflect a sphere-shaped
inertia tensor and sphere-sized AABBs; the values for cone worlds
reflect the cone. The full list is in
``mjlab.entity.variants.VARIANT_DEPENDENT_FIELDS``.
* **Per-mesh-geom attributes** (contact bits, friction, mass,
density, condim, group, priority, rgba, solref, solimp, margin,
gap) are captured per variant in ``VariantGeomSpec`` at merge time
and restored verbatim on the slot geom during the per-variant
reference compile. So if sphere's collision geoms have
``friction=0.5`` and cone's have ``friction=1.2``, world W's
per-step friction reflects the assigned variant's source value.
The one exception is ``material``, which is not propagated across
variants; if you need per-world appearance variation use DR on
``geom_rgba`` / ``mat_rgba``.
If ``sphere`` adds a body that ``cone`` lacks (or vice versa), the
validator rejects the configuration before any of the merge logic
runs. The slot mechanism only flexes mesh geom counts within
matching bodies; everything structural above the geom level must
agree.
.. note::
**Doesn't compiling the merged scene ruin the prop body's
inertia?**
No, but it's worth understanding why, because the naive intuition
says it should. If you stuck every variant's mesh geoms on the
prop body and called ``spec.compile()``, MuJoCo would sum each
geom's inertial contribution, and you would get a body whose mass
and inertia tensor are a meaningless mix of every variant's shape.
mjlab avoids this in two layers:
* **The merged scene does not stick every variant's geoms on the
body.** The prop body in the merged spec carries variant 0's
mesh geoms (with their original mass and density) plus, for any
slot variant 0 doesn't fill, a synthesized padding geom that has
``mass = 0`` and ``density = 0``. Padding contributes nothing to
body inertia. Other variants' meshes are present in the merged
spec only as **mesh assets** (in the assets section, not as geoms
on any body). They get wired in at runtime via per-world
``geom_dataid`` and never affect the host compile's inertial
sums.
* **Per-world overrides come from per-variant source compiles.**
Even with the above, the merged-scene compile's prop body inertia
is only correct for variant 0. For every other variant, mjlab
compiles that variant's original source spec in isolation (one
body, one variant's worth of meshes), reads the resulting
``body_mass``, ``body_inertia``, ``body_ipos``, ``body_iquat``,
``body_invweight0``, and ``body_subtreemass``, and writes them
into the per-world arrays at the prop body's index.
Net result: world W's prop body inertia is byte-equal to what you
would get by compiling variant W's source spec on its own. There
is a regression test
(``test_visual_collision_split_inertia_matches_independent_compile``
in ``tests/test_variants.py``) that asserts exactly this against
independent per-variant compiles.
World assignment
----------------
mjlab assigns variants to worlds proportionally by weight using the
How worlds get mapped to variants is controlled by the ``assignment``
field on ``VariantEntityCfg``. It accepts three shapes:
* ``None`` (default): uniform allocation across variants.
* ``dict[str, float]``: per-variant weights. Variants not listed
default to weight 1.0.
* ``Callable[[int], Sequence[int]]``: an explicit assignment function
called with ``num_envs`` at simulation init.
Both the ``None`` and dict cases use 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
@@ -98,14 +310,28 @@ first receives ``floor(q_i)`` worlds, and the remaining
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.
so ``{"a": 1, "b": 2, "c": 1}`` and ``{"a": 0.25, "b": 0.5, "c": 0.25}``
produce identical assignments. A weight of zero is allowed and
produces zero worlds for that variant; at least one variant must end
up with 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``:
The default and dict paths are purely deterministic given
``(assignment, num_envs)``. With ``assignment={"a": 1, "b": 1}`` and
``num_envs = 8`` you always get ``[0, 0, 0, 0, 1, 1, 1, 1]``. There is
no seed involved; rerunning the same config produces the same
partition every time. Note that the partition's *boundaries* depend
on ``num_envs``, so world W's variant is not necessarily stable when
you change ``num_envs``. If you need explicit per-world stability
across batch sizes (e.g. "world 0 is always variant 0, world 1 is
always variant 1, regardless of how many envs I launch"), use a
callable assignment as below.
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.
Read the resolved assignment from user code via
``env.sim.world_to_variant``:
.. code-block:: python
@@ -118,23 +344,53 @@ variants were declared in ``VariantEntityCfg.variants``. The dict is
empty for non-variant scenes.
Custom assignment with a callable
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
When the weighted default is not what you want, pass a callable to
``assignment``. The callable receives ``num_envs`` and must return a
length-``num_envs`` sequence of variant indices in
``[0, len(variants))``. The returned sequence's length and bounds are
validated at sim init; mismatches raise a ``ValueError`` naming the
offending entity.
A few patterns:
**Round-robin** - cycle through variants by world index.
.. code-block:: python
cfg = VariantEntityCfg(
variants={"a": make_a, "b": make_b, "c": make_c},
assignment=lambda n: [w % 3 for w in range(n)],
)
**Stratified halves** - first half is variant 0, second half is
variant 1.
.. code-block:: python
cfg = VariantEntityCfg(
variants={"easy": make_easy, "hard": make_hard},
assignment=lambda n: [0] * (n // 2) + [1] * (n - n // 2),
)
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
variant-dependent fields as ``(num_envs, ...)`` tensors and registers
them in ``sim.per_world_default_fields``. DR 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.
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,
@@ -169,32 +425,66 @@ Convex hull visualization is computed per variant from the variant's
mesh vertices.
Performance considerations
--------------------------
Performance
-----------
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.
**Per-step cost is unaffected by variant count.** 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.
**Construction cost is linear in the total variant count.** mjlab
compiles the merged scene once to produce the canonical ``MjModel``,
then compiles each variant's original (un-merged) source spec in
isolation to recover that variant's per-body and per-geom mesh-derived
fields. Each per-variant compile sees only that variant's single body
and mesh, so its cost is independent of the total number of variants
in the scene.
``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.
For a scene with one variant entity declaring k variants, construction
runs ``1 + k`` compiles. With multiple variant entities, compiles
decouple across entities: two variant entities of 5 variants each cost
``1 + 5 + 5 = 11`` compiles, not ``1 + 5 * 5 = 26``. As an order of
magnitude on CPU with typical procedural meshes, each per-variant
compile takes around 1-2 ms, so a scene with 100 variants pays a few
hundred milliseconds at startup and a scene with 1000 variants pays
roughly two seconds.
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.
The merged spec contains every variant's mesh assets simultaneously,
so memory at scene-build time scales with the total mesh vertex /
face count across all variants. This is paid once at startup and does
not affect training throughput.
Limitations
-----------
**Floating-base only.** Each variant's root body must declare a free
joint. Fixed-base variants are rejected; mocap auto-wrapping that
applies to non-variant entities is not applied here.
**Material assets are not propagated.** Each variant's ``contype``,
``conaffinity``, ``condim``, ``friction``, ``mass``, ``density``,
``group``, ``priority``, ``rgba``, ``solref``, ``solimp``, ``margin``,
and ``gap`` are restored per-world during compile, but the
``material`` reference on slot geoms inherits whichever material the
template variant set. Use DR on ``geom_rgba`` / ``mat_rgba`` for
per-world appearance variation.
**Assignment is fixed at sim init.** There is no API to swap a world
to a different variant on episode reset. World W's mesh asset is
whatever it was assigned at init for the lifetime of the simulation.
Per-episode mesh randomization is not supported today; DR can vary
scalar properties (mass, friction, color, scale) on a fixed variant
but cannot swap one mesh for another.
**No support for per-world differing kinematic topology.** Variants
must share the same body tree, joints, and actuator/sensor counts,
so you cannot configure things like:
* a different number of objects per world (world 0 has two props on
the table, world 1 has three);
* different articulation per world (world 0's prop is an articulated
drawer with a slider joint, world 1's prop is a rigid block).
True heterogeneous topology requires upstream support in mujoco_warp
that does not currently exist.
05_software/train/rc_mjlab/mjlab/docs/source/faq.rst
+21
View File
@@ -42,6 +42,27 @@ Not all CUDA versions are supported by MuJoCo Warp.
- **Recommended**: CUDA **12.4+** (for conditional execution support in CUDA
graphs).
How do I run on CPU without touching the GPU?
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Passing ``device="cpu"`` puts all mjlab computation on the CPU, but it does
**not** stop Warp from initializing the GPU. The first time Warp's runtime
comes up, it eagerly enumerates and creates a CUDA context on **every**
visible device, regardless of which device you requested. So on a machine
with a visible GPU, a ``device="cpu"`` run still claims VRAM.
This happens inside Warp and cannot be prevented from Python once the
package is imported. To keep the process entirely off the GPU, hide the
devices from CUDA before launching:
.. code-block:: bash
CUDA_VISIBLE_DEVICES="" uv run train.py ...
With no visible CUDA devices, Warp initializes CPU-only and never allocates
on the GPU. See `issue #949
<https://github.com/mujocolab/mjlab/issues/949>`_ for background.
Performance
-----------
05_software/train/rc_mjlab/mjlab/docs/source/terrain.rst
+53 -28
View File
@@ -49,7 +49,6 @@ the geometry and how it scales with difficulty.
terrain_generator=TerrainGeneratorCfg(
size=(8.0, 8.0),
num_rows=10,
num_cols=20,
border_width=20.0,
curriculum=True,
sub_terrains={
@@ -70,10 +69,12 @@ the geometry and how it scales with difficulty.
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.
The generator creates a grid of patches sized ``num_rows`` by either
``num_cols`` (random mode) or ``len(sub_terrains)`` (curriculum mode,
where ``num_cols`` is ignored). The ``sub_terrains`` dictionary maps
names to ``SubTerrainCfg`` instances; each sub-terrain's ``proportion``
controls robot spawning distribution across columns in curriculum mode,
or per-patch sampling probability in random mode.
Grid layout
@@ -82,30 +83,48 @@ 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.
**Curriculum mode** (``curriculum=True``). Each terrain type gets exactly
one column; the generator uses ``len(sub_terrains)`` columns regardless of
``num_cols``. All patches in a column share the same terrain type, and
difficulty increases from row 0 (easiest) to row ``num_rows - 1``
(hardest). The ``proportion`` field controls how robots are distributed
across columns at spawn time, not column count. 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.
``difficulty_range``. ``num_cols`` is honored. 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
that 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.
and 20 cm steps at difficulty 1.
In curriculum mode, difficulty is determined by the row:
``difficulty = lower + (upper - lower) * row / max(num_rows - 1, 1)``,
where ``(lower, upper) = difficulty_range``. Row 0 is exactly
``lower``, row ``num_rows - 1`` is exactly ``upper``, and intermediate
rows are evenly spaced between them. All columns in a given row share
the same difficulty scalar; the visible variation across columns comes
from each sub-terrain type generating different geometry at the same
difficulty.
.. note::
With ``num_rows=1`` and ``curriculum=True``, every patch is generated
at ``difficulty = lower`` (the easiest configured difficulty). Use
``curriculum=False`` if you want a single grid of randomly sampled
difficulties instead.
In random mode, difficulty is sampled uniformly from
``difficulty_range`` independently for every patch.
Sub-terrain types
@@ -244,17 +263,23 @@ and undulating ground that box geoms cannot represent.
Preset configurations
---------------------
mjlab ships two ready-made ``TerrainGeneratorCfg`` presets in
mjlab ships three 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.
A 10x20 random-mode grid with seven terrain types (flat, stairs,
inverted stairs, slopes, inverted slopes, random rough, waves).
Designed for locomotion training with a moderate difficulty range.
Set ``curriculum=True`` via ``dataclasses.replace`` to use it as a
curriculum grid (one column per terrain type).
``STAIRS_TERRAINS_CFG``
A 10-row curriculum grid focused on stair traversal: flat plus
three pyramid-stair variants of increasing difficulty.
``ALL_TERRAINS_CFG``
A 10x16 grid with all sixteen terrain types at equal proportion.
Useful for training on maximum terrain variety.
A 10-row random-mode grid covering all available terrain types at
equal proportion. Useful for training on maximum terrain variety.
Both can be used directly or customized with ``dataclasses.replace()``:
@@ -285,9 +310,9 @@ The key concepts:
- 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.
- When an environment is promoted past the hardest row, it is randomly
reassigned to any row in ``[0, num_rows)`` to prevent the policy from
collapsing to a single difficulty level.
Flat patch detection