[train] 更新新MJCF与第一版完整训练框架
This commit is contained in:
@@ -28,9 +28,9 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v4
|
||||
uses: actions/checkout@v6
|
||||
- name: Setup uv
|
||||
uses: astral-sh/setup-uv@v6
|
||||
uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
version: "0.9.27"
|
||||
@@ -45,9 +45,9 @@ jobs:
|
||||
matrix:
|
||||
python-version: ["3.10", "3.11", "3.12", "3.13"]
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
- name: Setup uv
|
||||
uses: astral-sh/setup-uv@v6
|
||||
uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
python-version: ${{ matrix.python-version }}
|
||||
enable-cache: true
|
||||
@@ -68,9 +68,9 @@ jobs:
|
||||
matrix:
|
||||
python-version: ["3.10", "3.11", "3.12", "3.13"]
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
- name: Setup uv
|
||||
uses: astral-sh/setup-uv@v6
|
||||
uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
python-version: ${{ matrix.python-version }}
|
||||
enable-cache: true
|
||||
@@ -84,9 +84,9 @@ jobs:
|
||||
matrix:
|
||||
python-version: ["3.10", "3.11", "3.12", "3.13"]
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
- name: Setup uv
|
||||
uses: astral-sh/setup-uv@v6
|
||||
uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
python-version: ${{ matrix.python-version }}
|
||||
enable-cache: true
|
||||
|
||||
@@ -0,0 +1,44 @@
|
||||
name: Claude Code Review
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
types: [opened, synchronize, ready_for_review, reopened]
|
||||
# Optional: Only run on specific file changes
|
||||
# paths:
|
||||
# - "src/**/*.ts"
|
||||
# - "src/**/*.tsx"
|
||||
# - "src/**/*.js"
|
||||
# - "src/**/*.jsx"
|
||||
|
||||
jobs:
|
||||
claude-review:
|
||||
# Optional: Filter by PR author
|
||||
# if: |
|
||||
# github.event.pull_request.user.login == 'external-contributor' ||
|
||||
# github.event.pull_request.user.login == 'new-developer' ||
|
||||
# github.event.pull_request.author_association == 'FIRST_TIME_CONTRIBUTOR'
|
||||
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
contents: read
|
||||
pull-requests: read
|
||||
issues: read
|
||||
id-token: write
|
||||
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
fetch-depth: 1
|
||||
|
||||
- name: Run Claude Code Review
|
||||
id: claude-review
|
||||
uses: anthropics/claude-code-action@v1
|
||||
with:
|
||||
claude_code_oauth_token: ${{ secrets.CLAUDE_CODE_OAUTH_TOKEN }}
|
||||
plugin_marketplaces: 'https://github.com/anthropics/claude-code.git'
|
||||
plugins: 'code-review@claude-code-plugins'
|
||||
prompt: '/code-review:code-review ${{ github.repository }}/pull/${{ github.event.pull_request.number }}'
|
||||
# See https://github.com/anthropics/claude-code-action/blob/main/docs/usage.md
|
||||
# or https://code.claude.com/docs/en/cli-reference for available options
|
||||
|
||||
@@ -0,0 +1,50 @@
|
||||
name: Claude Code
|
||||
|
||||
on:
|
||||
issue_comment:
|
||||
types: [created]
|
||||
pull_request_review_comment:
|
||||
types: [created]
|
||||
issues:
|
||||
types: [opened, assigned]
|
||||
pull_request_review:
|
||||
types: [submitted]
|
||||
|
||||
jobs:
|
||||
claude:
|
||||
if: |
|
||||
(github.event_name == 'issue_comment' && contains(github.event.comment.body, '@claude')) ||
|
||||
(github.event_name == 'pull_request_review_comment' && contains(github.event.comment.body, '@claude')) ||
|
||||
(github.event_name == 'pull_request_review' && contains(github.event.review.body, '@claude')) ||
|
||||
(github.event_name == 'issues' && (contains(github.event.issue.body, '@claude') || contains(github.event.issue.title, '@claude')))
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
contents: read
|
||||
pull-requests: read
|
||||
issues: read
|
||||
id-token: write
|
||||
actions: read # Required for Claude to read CI results on PRs
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
fetch-depth: 1
|
||||
|
||||
- name: Run Claude Code
|
||||
id: claude
|
||||
uses: anthropics/claude-code-action@v1
|
||||
with:
|
||||
claude_code_oauth_token: ${{ secrets.CLAUDE_CODE_OAUTH_TOKEN }}
|
||||
|
||||
# This is an optional setting that allows Claude to read CI results on PRs
|
||||
additional_permissions: |
|
||||
actions: read
|
||||
|
||||
# Optional: Give a custom prompt to Claude. If this is not specified, Claude will perform the instructions specified in the comment that tagged it.
|
||||
# prompt: 'Update the pull request description to include a summary of changes.'
|
||||
|
||||
# Optional: Add claude_args to customize behavior and configuration
|
||||
# See https://github.com/anthropics/claude-code-action/blob/main/docs/usage.md
|
||||
# or https://code.claude.com/docs/en/cli-reference for available options
|
||||
# claude_args: '--allowed-tools Bash(gh pr *)'
|
||||
|
||||
@@ -17,7 +17,7 @@ jobs:
|
||||
build:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
|
||||
@@ -26,7 +26,7 @@ jobs:
|
||||
python-version: '3.13'
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v4
|
||||
uses: astral-sh/setup-uv@v7
|
||||
|
||||
- name: Build Sphinx Documentation
|
||||
run: uv run --group docs sphinx-multiversion docs docs/_build
|
||||
|
||||
@@ -15,9 +15,9 @@ jobs:
|
||||
contents: read
|
||||
steps:
|
||||
- name: Checkout
|
||||
uses: actions/checkout@v4
|
||||
uses: actions/checkout@v6
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v6
|
||||
uses: astral-sh/setup-uv@v7
|
||||
- name: Install Python 3.13
|
||||
run: uv python install 3.13
|
||||
- name: Build
|
||||
|
||||
@@ -33,9 +33,9 @@ keywords:
|
||||
- reinforcement-learning
|
||||
- robotics
|
||||
license: Apache-2.0
|
||||
commit: e2f33c6fb49caa26ec11f7b2de3c0c9aba71e9fd
|
||||
version: 1.3.0
|
||||
date-released: '2026-04-14'
|
||||
commit: 3cc461cd15e7155a8998b75ad767fae6dd448072
|
||||
version: 1.4.0
|
||||
date-released: '2026-05-26'
|
||||
preferred-citation:
|
||||
type: article
|
||||
title: >-
|
||||
|
||||
@@ -1,6 +1,10 @@
|
||||
.PHONY: sync
|
||||
sync:
|
||||
uv sync --all-extras --all-packages --group dev
|
||||
uv sync --all-packages --extra cu128 --group dev
|
||||
|
||||
.PHONY: sync-cpu
|
||||
sync-cpu:
|
||||
uv sync --all-packages --extra cpu --group dev
|
||||
|
||||
.PHONY: format
|
||||
format:
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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`` /
|
||||
|
||||
@@ -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
|
||||
^^^^^^^^^^^^^^^^^^
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
-----------
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "uv_build"
|
||||
|
||||
[project]
|
||||
name = "mjlab"
|
||||
version = "1.3.0"
|
||||
version = "1.4.0"
|
||||
license = "Apache-2.0"
|
||||
license-files = ["LICENSE"]
|
||||
readme = { file = "README.md", content-type = "text/markdown" }
|
||||
@@ -37,15 +37,16 @@ dependencies = [
|
||||
"torch>=2.7.0",
|
||||
"torchrunx>=0.3.4",
|
||||
"warp-lang>=1.12.0",
|
||||
"mujoco-warp>=3.8.0",
|
||||
"mujoco>=3.8.0",
|
||||
"mujoco-warp>=3.8.0.3,~=3.8.0",
|
||||
"mujoco~=3.8.0",
|
||||
"trimesh>=4.8.3",
|
||||
"viser>=1.0.26",
|
||||
"mjviser>=0.0.13",
|
||||
"scipy>=1.15",
|
||||
"viser>=1.0.27",
|
||||
"mjviser>=0.0.14",
|
||||
"mediapy>=1.2.6",
|
||||
"imageio-ffmpeg",
|
||||
"tensordict",
|
||||
"rsl-rl-lib==5.2.0",
|
||||
"rsl-rl-lib==5.4.0",
|
||||
"tensorboard>=2.20.0",
|
||||
"onnxscript>=0.5.4",
|
||||
"wandb>=0.22.3",
|
||||
@@ -98,11 +99,11 @@ conflicts = [
|
||||
[{extra = "cu128"}, {extra = "cpu"}],
|
||||
]
|
||||
# The nightly index (py.mujoco.org) only has dev builds, and PEP 440 ranks
|
||||
# 3.7.0.devN < 3.7.0, so the >=3.7.0 floor in [project.dependencies] would
|
||||
# 3.8.0.devN < 3.8.0, so the ~=3.8.0 floor in [project.dependencies] would
|
||||
# reject them. This override loosens the constraint for uv resolution only.
|
||||
override-dependencies = ["mujoco>=3.8.0.dev0"]
|
||||
constraint-dependencies = [
|
||||
"GitPython>=3.1.47",
|
||||
"GitPython>=3.1.49",
|
||||
"lxml>=6.1.0",
|
||||
]
|
||||
required-environments = [
|
||||
@@ -111,12 +112,6 @@ required-environments = [
|
||||
]
|
||||
|
||||
[[tool.uv.index]]
|
||||
name = "tsinghua"
|
||||
url = "https://pypi.tuna.tsinghua.edu.cn/simple"
|
||||
default = true
|
||||
|
||||
[[tool.uv.index]]
|
||||
name = "pypi"
|
||||
url = "https://pypi.org/simple"
|
||||
|
||||
[[tool.uv.index]]
|
||||
@@ -146,8 +141,7 @@ torch = [
|
||||
{ index = "pytorch-cpu", extra = "cpu", marker = "sys_platform != 'darwin'" },
|
||||
]
|
||||
mujoco = { index = "mujoco" }
|
||||
mujoco-warp = { git = "https://github.com/google-deepmind/mujoco_warp", rev = "6f235d4" }
|
||||
mjviser = { git = "https://github.com/mujocolab/mjviser", rev = "1bdfd6fe79066b847a5f430000fcfbb53ec31a6f" }
|
||||
mujoco-warp = { git = "https://github.com/google-deepmind/mujoco_warp", rev = "88b55fc2696960b927bc12584994bb8412b36558" }
|
||||
|
||||
[tool.ruff]
|
||||
src = ["src"] # Helpful for recognizing first-party imports.
|
||||
|
||||
@@ -215,6 +215,27 @@ def generate_dashboard_html(runs: list[dict], throughput_data: list[dict]) -> st
|
||||
border-color: var(--accent);
|
||||
color: white;
|
||||
}}
|
||||
.range-selector {{
|
||||
display: flex;
|
||||
gap: 0.4rem;
|
||||
margin-bottom: 1rem;
|
||||
}}
|
||||
.range-btn {{
|
||||
background: var(--bg-card);
|
||||
border: 1px solid var(--border);
|
||||
border-radius: 6px;
|
||||
padding: 0.3rem 0.75rem;
|
||||
cursor: pointer;
|
||||
color: var(--text);
|
||||
font-size: 0.8rem;
|
||||
font-weight: 500;
|
||||
}}
|
||||
.range-btn:hover {{ border-color: var(--accent); }}
|
||||
.range-btn.active {{
|
||||
background: var(--accent);
|
||||
border-color: var(--accent);
|
||||
color: white;
|
||||
}}
|
||||
.tab-content {{ display: none; }}
|
||||
.tab-content.active {{ display: block; }}
|
||||
.tab-description {{
|
||||
@@ -303,12 +324,24 @@ def generate_dashboard_html(runs: list[dict], throughput_data: list[dict]) -> st
|
||||
|
||||
<div id="tracking" class="tab-content active">
|
||||
<p class="tab-description">Nightly motion imitation training and evaluation on Unitree G1 (1024 trials per run).</p>
|
||||
<div class="range-selector" id="range-selector">
|
||||
<button class="range-btn" data-days="30">30d</button>
|
||||
<button class="range-btn active" data-days="90">90d</button>
|
||||
<button class="range-btn" data-days="180">180d</button>
|
||||
<button class="range-btn" data-days="0">All</button>
|
||||
</div>
|
||||
<div class="charts" id="charts"></div>
|
||||
</div>
|
||||
|
||||
<div id="throughput" class="tab-content">
|
||||
<p class="tab-description">Physics simulation throughput across tasks (4096 parallel envs, NVIDIA RTX 5090).</p>
|
||||
<div class="task-grid" id="task-grid"></div>
|
||||
<div class="range-selector" id="range-selector-tp">
|
||||
<button class="range-btn" data-days="30">30d</button>
|
||||
<button class="range-btn active" data-days="90">90d</button>
|
||||
<button class="range-btn" data-days="180">180d</button>
|
||||
<button class="range-btn" data-days="0">All</button>
|
||||
</div>
|
||||
<div id="task-chart-panels"></div>
|
||||
</div>
|
||||
|
||||
@@ -393,6 +426,8 @@ def generate_dashboard_html(runs: list[dict], throughput_data: list[dict]) -> st
|
||||
}};
|
||||
|
||||
let charts = [];
|
||||
let trackingCharts = [];
|
||||
let throughputCharts = [];
|
||||
|
||||
function updateChartColors() {{
|
||||
const style = getComputedStyle(root);
|
||||
@@ -449,7 +484,7 @@ def generate_dashboard_html(runs: list[dict], throughput_data: list[dict]) -> st
|
||||
`;
|
||||
chartsContainer.appendChild(card);
|
||||
|
||||
charts.push(new Chart(card.querySelector('canvas'), {{
|
||||
const chart = new Chart(card.querySelector('canvas'), {{
|
||||
type: 'line',
|
||||
data: {{
|
||||
datasets: [
|
||||
@@ -459,7 +494,8 @@ def generate_dashboard_html(runs: list[dict], throughput_data: list[dict]) -> st
|
||||
borderColor: color,
|
||||
backgroundColor: color + '20',
|
||||
borderWidth: 2,
|
||||
pointRadius: 4,
|
||||
pointRadius: 2,
|
||||
pointHoverRadius: 5,
|
||||
tension: 0.1,
|
||||
fill: true
|
||||
}},
|
||||
@@ -531,7 +567,9 @@ def generate_dashboard_html(runs: list[dict], throughput_data: list[dict]) -> st
|
||||
}}
|
||||
}}
|
||||
}}
|
||||
}}));
|
||||
}});
|
||||
charts.push(chart);
|
||||
trackingCharts.push(chart);
|
||||
}});
|
||||
|
||||
// Tab switching
|
||||
@@ -621,7 +659,8 @@ def generate_dashboard_html(runs: list[dict], throughput_data: list[dict]) -> st
|
||||
borderColor: '#58a6ff',
|
||||
backgroundColor: '#58a6ff20',
|
||||
borderWidth: 2,
|
||||
pointRadius: 4,
|
||||
pointRadius: 2,
|
||||
pointHoverRadius: 5,
|
||||
tension: 0.1,
|
||||
fill: true
|
||||
}},
|
||||
@@ -631,7 +670,8 @@ def generate_dashboard_html(runs: list[dict], throughput_data: list[dict]) -> st
|
||||
borderColor: '#3fb950',
|
||||
backgroundColor: '#3fb95020',
|
||||
borderWidth: 2,
|
||||
pointRadius: 4,
|
||||
pointRadius: 2,
|
||||
pointHoverRadius: 5,
|
||||
tension: 0.1,
|
||||
fill: true
|
||||
}}
|
||||
@@ -697,6 +737,7 @@ def generate_dashboard_html(runs: list[dict], throughput_data: list[dict]) -> st
|
||||
}}
|
||||
}});
|
||||
charts.push(chart);
|
||||
throughputCharts.push(chart);
|
||||
throughputChartInstances[task] = {{ chart, panelId: `task-panel-${{i}}` }};
|
||||
|
||||
// Card click handler
|
||||
@@ -712,6 +753,23 @@ def generate_dashboard_html(runs: list[dict], throughput_data: list[dict]) -> st
|
||||
}} else {{
|
||||
taskGrid.innerHTML = '<p style="color: var(--text-dim)">No throughput data available. Run measure_throughput.py to generate data.</p>';
|
||||
}}
|
||||
|
||||
// Date-range windowing across both tracking and throughput charts.
|
||||
// Setting min and clearing max also resets any zoom/pan.
|
||||
function setRange(days) {{
|
||||
const min = days > 0 ? Date.now() - days * 86400000 : undefined;
|
||||
[...trackingCharts, ...throughputCharts].forEach(c => {{
|
||||
c.options.scales.x.min = min;
|
||||
c.options.scales.x.max = undefined;
|
||||
c.update();
|
||||
}});
|
||||
document.querySelectorAll('.range-btn').forEach(b =>
|
||||
b.classList.toggle('active', parseInt(b.dataset.days) === days));
|
||||
}}
|
||||
document.querySelectorAll('.range-btn').forEach(btn => {{
|
||||
btn.addEventListener('click', () => setRange(parseInt(btn.dataset.days)));
|
||||
}});
|
||||
setRange(90);
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
@@ -764,7 +822,11 @@ def main(
|
||||
if run_id in eval_results_by_id:
|
||||
print(f"Using cached result for {run_id}")
|
||||
else:
|
||||
result = evaluate_run(run_path, num_envs)
|
||||
try:
|
||||
result = evaluate_run(run_path, num_envs)
|
||||
except RuntimeError as e:
|
||||
print(f"Skipping {run_path}: {e}")
|
||||
continue
|
||||
eval_results_by_id[run_id] = result
|
||||
new_evals += 1
|
||||
else:
|
||||
@@ -783,7 +845,11 @@ def main(
|
||||
print(f"Reached eval limit ({eval_limit}), skipping remaining new runs")
|
||||
break
|
||||
run_path = f"{entity}/{project}/{run.id}"
|
||||
result = evaluate_run(run_path, num_envs)
|
||||
try:
|
||||
result = evaluate_run(run_path, num_envs)
|
||||
except RuntimeError as e:
|
||||
print(f"Skipping {run.name} ({run.id}): {e}")
|
||||
continue
|
||||
eval_results_by_id[run.id] = result
|
||||
new_evals += 1
|
||||
|
||||
|
||||
@@ -0,0 +1,115 @@
|
||||
"""Interactive single-patch terrain explorer (Viser + MuJoCo MjSpec).
|
||||
|
||||
Run with:
|
||||
uv run python scripts/tools/terrain_explorer.py
|
||||
uv run python scripts/tools/terrain_explorer.py --port 8081
|
||||
|
||||
Then open the printed URL (default http://localhost:8080).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import time
|
||||
|
||||
import mujoco
|
||||
import numpy as np
|
||||
import viser
|
||||
from mjviser.conversions import merge_geoms
|
||||
|
||||
from mjlab.terrains.config import ALL_TERRAIN_PRESETS
|
||||
from mjlab.terrains.terrain_generator import TerrainGenerator, TerrainGeneratorCfg
|
||||
|
||||
PATCH_SIZE = (8.0, 8.0)
|
||||
|
||||
|
||||
# Per-preset overrides applied when building in the explorer (e.g. to surface
|
||||
# difficulty-driven behavior that is off by default).
|
||||
_PRESET_OVERRIDES: dict[str, dict] = {
|
||||
"random_rough": {"scale_with_difficulty": True},
|
||||
}
|
||||
|
||||
|
||||
def _build_terrain_mesh(preset_name: str, difficulty: float, seed: int):
|
||||
"""Generate a single terrain patch and return a merged trimesh (or raise)."""
|
||||
preset_fn = ALL_TERRAIN_PRESETS[preset_name]
|
||||
overrides = _PRESET_OVERRIDES.get(preset_name, {})
|
||||
generator_cfg = TerrainGeneratorCfg(
|
||||
seed=seed,
|
||||
size=PATCH_SIZE,
|
||||
num_rows=1,
|
||||
num_cols=1,
|
||||
border_width=0.0,
|
||||
curriculum=False,
|
||||
# A degenerate range pins the single patch to exactly this difficulty.
|
||||
difficulty_range=(difficulty, difficulty),
|
||||
color_scheme="height",
|
||||
sub_terrains={preset_name: preset_fn(proportion=1.0, **overrides)},
|
||||
)
|
||||
generator = TerrainGenerator(generator_cfg)
|
||||
spec = mujoco.MjSpec()
|
||||
generator.compile(spec)
|
||||
model = spec.compile()
|
||||
|
||||
terrain_body_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "terrain")
|
||||
geom_ids = [i for i in range(model.ngeom) if model.geom_bodyid[i] == terrain_body_id]
|
||||
return merge_geoms(model, geom_ids)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument(
|
||||
"--port", type=int, default=8080, help="Port for the viser server."
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
server = viser.ViserServer(port=args.port)
|
||||
preset_names = sorted(ALL_TERRAIN_PRESETS)
|
||||
|
||||
terrain_dropdown = server.gui.add_dropdown(
|
||||
"Terrain", options=preset_names, initial_value=preset_names[0]
|
||||
)
|
||||
difficulty_slider = server.gui.add_slider(
|
||||
"Difficulty", min=0.0, max=1.0, step=0.01, initial_value=0.0
|
||||
)
|
||||
seed_input = server.gui.add_number("Seed", initial_value=42, step=1)
|
||||
status = server.gui.add_markdown("**Status:** ready")
|
||||
|
||||
handle: viser.SceneNodeHandle | None = None
|
||||
|
||||
def update() -> None:
|
||||
nonlocal handle
|
||||
name = terrain_dropdown.value
|
||||
difficulty = float(difficulty_slider.value)
|
||||
seed = int(seed_input.value)
|
||||
status.content = f"**Status:** building `{name}` at difficulty {difficulty:.2f}..."
|
||||
try:
|
||||
mesh = _build_terrain_mesh(name, difficulty, seed)
|
||||
except Exception as e: # noqa: BLE001 - surface any generation failure in the UI.
|
||||
status.content = f"**Error:** {type(e).__name__}: {e}"
|
||||
print(f"Failed to build {name} at difficulty {difficulty}: {e}")
|
||||
return
|
||||
if handle is not None:
|
||||
handle.remove()
|
||||
handle = server.scene.add_mesh_trimesh("/terrain", mesh)
|
||||
status.content = (
|
||||
f"**Loaded** `{name}` at difficulty {difficulty:.2f} ({len(mesh.faces):,} faces)"
|
||||
)
|
||||
|
||||
terrain_dropdown.on_update(lambda _: update())
|
||||
difficulty_slider.on_update(lambda _: update())
|
||||
seed_input.on_update(lambda _: update())
|
||||
|
||||
# Top-down-ish initial camera.
|
||||
@server.on_client_connect
|
||||
def _(client: viser.ClientHandle) -> None:
|
||||
client.camera.position = np.array([10.0, 10.0, 8.0])
|
||||
client.camera.look_at = np.array([0.0, 0.0, 0.0])
|
||||
|
||||
update()
|
||||
while True:
|
||||
time.sleep(1.0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,5 +1,15 @@
|
||||
import os
|
||||
import sys
|
||||
|
||||
# Default to EGL for GPU-accelerated offscreen rendering on Linux. Must be set
|
||||
# before any mujoco import: mujoco's gl_context module captures MUJOCO_GL once
|
||||
# at load time. Override with e.g. MUJOCO_GL=osmesa on clusters without EGL.
|
||||
# Linux-only because mujoco's gl_context rejects "egl" on macOS/Windows and
|
||||
# raises at import. On those platforms we leave MUJOCO_GL alone so mujoco
|
||||
# defaults to GLFW.
|
||||
if sys.platform.startswith("linux"):
|
||||
os.environ.setdefault("MUJOCO_GL", "egl")
|
||||
|
||||
import traceback
|
||||
from importlib.metadata import entry_points
|
||||
from pathlib import Path
|
||||
|
||||
@@ -4,6 +4,12 @@ from mjlab.actuator.actuator import Actuator as Actuator
|
||||
from mjlab.actuator.actuator import ActuatorCfg as ActuatorCfg
|
||||
from mjlab.actuator.actuator import ActuatorCmd as ActuatorCmd
|
||||
from mjlab.actuator.actuator import CommandField as CommandField
|
||||
from mjlab.actuator.builtin_actuator import (
|
||||
BuiltinDcMotorActuator as BuiltinDcMotorActuator,
|
||||
)
|
||||
from mjlab.actuator.builtin_actuator import (
|
||||
BuiltinDcMotorActuatorCfg as BuiltinDcMotorActuatorCfg,
|
||||
)
|
||||
from mjlab.actuator.builtin_actuator import (
|
||||
BuiltinMotorActuator as BuiltinMotorActuator,
|
||||
)
|
||||
@@ -16,6 +22,12 @@ from mjlab.actuator.builtin_actuator import (
|
||||
from mjlab.actuator.builtin_actuator import (
|
||||
BuiltinMuscleActuatorCfg as BuiltinMuscleActuatorCfg,
|
||||
)
|
||||
from mjlab.actuator.builtin_actuator import (
|
||||
BuiltinPdActuator as BuiltinPdActuator,
|
||||
)
|
||||
from mjlab.actuator.builtin_actuator import (
|
||||
BuiltinPdActuatorCfg as BuiltinPdActuatorCfg,
|
||||
)
|
||||
from mjlab.actuator.builtin_actuator import (
|
||||
BuiltinPositionActuator as BuiltinPositionActuator,
|
||||
)
|
||||
@@ -28,6 +40,15 @@ from mjlab.actuator.builtin_actuator import (
|
||||
from mjlab.actuator.builtin_actuator import (
|
||||
BuiltinVelocityActuatorCfg as BuiltinVelocityActuatorCfg,
|
||||
)
|
||||
from mjlab.actuator.builtin_actuator import (
|
||||
DcMotorDatasheetParams as DcMotorDatasheetParams,
|
||||
)
|
||||
from mjlab.actuator.builtin_actuator import (
|
||||
DcMotorInputMode as DcMotorInputMode,
|
||||
)
|
||||
from mjlab.actuator.builtin_actuator import (
|
||||
DcMotorPhysicalParams as DcMotorPhysicalParams,
|
||||
)
|
||||
from mjlab.actuator.builtin_group import BuiltinActuatorGroup as BuiltinActuatorGroup
|
||||
from mjlab.actuator.dc_actuator import DcMotorActuator as DcMotorActuator
|
||||
from mjlab.actuator.dc_actuator import DcMotorActuatorCfg as DcMotorActuatorCfg
|
||||
|
||||
@@ -174,15 +174,6 @@ class Actuator(ABC, Generic[ActuatorCfgT]):
|
||||
"""Whether this actuator has delay configured."""
|
||||
return self.cfg.delay_max_lag > 0
|
||||
|
||||
@property
|
||||
def command_field(self) -> CommandField | None:
|
||||
"""The primary command field this actuator consumes.
|
||||
|
||||
Returns None by default. Subclasses should override to return the
|
||||
appropriate field.
|
||||
"""
|
||||
return None
|
||||
|
||||
@property
|
||||
def target_ids(self) -> torch.Tensor:
|
||||
"""Local indices of targets controlled by this actuator."""
|
||||
@@ -271,11 +262,6 @@ class Actuator(ABC, Generic[ActuatorCfgT]):
|
||||
"""Create delay buffer. Called during initialize()."""
|
||||
if not self.has_delay:
|
||||
return
|
||||
if self.command_field is None:
|
||||
raise ValueError(
|
||||
f"{self.__class__.__name__}: delay is configured (delay_max_lag="
|
||||
f"{self.cfg.delay_max_lag}) but command_field is not defined."
|
||||
)
|
||||
self._delay_buffer = DelayBuffer(
|
||||
min_lag=self.cfg.delay_min_lag,
|
||||
max_lag=self.cfg.delay_max_lag,
|
||||
@@ -287,19 +273,25 @@ class Actuator(ABC, Generic[ActuatorCfgT]):
|
||||
)
|
||||
|
||||
def apply_delay(self, cmd: ActuatorCmd) -> ActuatorCmd:
|
||||
"""Apply delay to the command_field target. No-op without delay."""
|
||||
"""Delay all command targets with one shared lag. No-op without delay.
|
||||
|
||||
Every target the policy issues (position, velocity, effort) travels the same
|
||||
command channel and experiences the same latency, so they are stacked and
|
||||
delayed together. Feedback fields (``pos``, ``vel``) are never delayed.
|
||||
"""
|
||||
if self._delay_buffer is None:
|
||||
return cmd
|
||||
cf = self.command_field
|
||||
if cf == "position":
|
||||
self._delay_buffer.append(cmd.position_target)
|
||||
return dataclasses.replace(cmd, position_target=self._delay_buffer.compute())
|
||||
elif cf == "velocity":
|
||||
self._delay_buffer.append(cmd.velocity_target)
|
||||
return dataclasses.replace(cmd, velocity_target=self._delay_buffer.compute())
|
||||
else:
|
||||
self._delay_buffer.append(cmd.effort_target)
|
||||
return dataclasses.replace(cmd, effort_target=self._delay_buffer.compute())
|
||||
targets = torch.stack(
|
||||
(cmd.position_target, cmd.velocity_target, cmd.effort_target), dim=-1
|
||||
)
|
||||
self._delay_buffer.append(targets)
|
||||
delayed = self._delay_buffer.compute()
|
||||
return dataclasses.replace(
|
||||
cmd,
|
||||
position_target=delayed[..., 0],
|
||||
velocity_target=delayed[..., 1],
|
||||
effort_target=delayed[..., 2],
|
||||
)
|
||||
|
||||
def set_lags(
|
||||
self,
|
||||
|
||||
@@ -7,19 +7,21 @@ created programmatically via the MjSpec API.
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from enum import IntEnum
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import mujoco
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from mjlab.actuator.actuator import (
|
||||
Actuator,
|
||||
ActuatorCfg,
|
||||
ActuatorCmd,
|
||||
CommandField,
|
||||
TransmissionType,
|
||||
)
|
||||
from mjlab.utils.spec import (
|
||||
apply_target_overrides,
|
||||
create_motor_actuator,
|
||||
create_muscle_actuator,
|
||||
create_position_actuator,
|
||||
@@ -63,10 +65,6 @@ class BuiltinPositionActuatorCfg(ActuatorCfg):
|
||||
class BuiltinPositionActuator(Actuator[BuiltinPositionActuatorCfg]):
|
||||
"""MuJoCo built-in position actuator."""
|
||||
|
||||
@property
|
||||
def command_field(self) -> CommandField:
|
||||
return "position"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
cfg: BuiltinPositionActuatorCfg,
|
||||
@@ -96,6 +94,102 @@ class BuiltinPositionActuator(Actuator[BuiltinPositionActuatorCfg]):
|
||||
return cmd.position_target
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class BuiltinPdActuatorCfg(ActuatorCfg):
|
||||
"""Implicit-integration version of IdealPdActuator.
|
||||
|
||||
Both consume a position target and a velocity target with kp/kd gains. The
|
||||
difference is in how the PD is delivered to MuJoCo: IdealPdActuator computes
|
||||
the PD force in Python and feeds it to a ``<motor>`` element, which MuJoCo
|
||||
sees as an opaque external force. This actuator expresses the PD as native
|
||||
MuJoCo elements (a ``<position>`` carrying kp, a ``<velocity>`` carrying kd),
|
||||
so the implicit and implicitfast integrators include the kp/kd derivatives
|
||||
in their velocity update. That makes the actuator numerically 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.
|
||||
"""
|
||||
|
||||
stiffness: float
|
||||
"""Proportional gain (kp)."""
|
||||
damping: float
|
||||
"""Derivative gain (kd)."""
|
||||
effort_limit: float | None = None
|
||||
"""Maximum total torque applied to the joint or tendon. Enforced as a
|
||||
sum-clamp on the two PD terms via jnt_actfrcrange (JOINT) or
|
||||
tendon_actfrcrange (TENDON). None leaves the limit unset."""
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
super().__post_init__()
|
||||
if self.transmission_type == TransmissionType.SITE:
|
||||
raise ValueError(
|
||||
"BuiltinPdActuatorCfg does not support SITE transmission. "
|
||||
"Use BuiltinMotorActuatorCfg for site transmission."
|
||||
)
|
||||
|
||||
def build(
|
||||
self, entity: Entity, target_ids: list[int], target_names: list[str]
|
||||
) -> BuiltinPdActuator:
|
||||
return BuiltinPdActuator(self, entity, target_ids, target_names)
|
||||
|
||||
|
||||
class BuiltinPdActuator(Actuator[BuiltinPdActuatorCfg]):
|
||||
"""MuJoCo native PD: paired <position> + <velocity> elements per target."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
cfg: BuiltinPdActuatorCfg,
|
||||
entity: Entity,
|
||||
target_ids: list[int],
|
||||
target_names: list[str],
|
||||
) -> None:
|
||||
super().__init__(cfg, entity, target_ids, target_names)
|
||||
|
||||
@property
|
||||
def num_targets(self) -> int:
|
||||
"""Number of targets. ``ctrl_ids`` is laid out as ``[pos..., vel...]``,
|
||||
each block of length ``num_targets``."""
|
||||
return len(self._target_ids_list)
|
||||
|
||||
def edit_spec(self, spec: mujoco.MjSpec, target_names: list[str]) -> None:
|
||||
# Position elements first, then velocity elements, so ctrl_ids is laid out
|
||||
# as [pos_0..pos_{N-1}, vel_0..vel_{N-1}].
|
||||
for target_name in target_names:
|
||||
pos_act = create_position_actuator(
|
||||
spec,
|
||||
target_name,
|
||||
actuator_name=f"{target_name}_pd_pos",
|
||||
stiffness=self.cfg.stiffness,
|
||||
damping=0.0, # damping lives on the <velocity> element.
|
||||
armature=self.cfg.armature,
|
||||
frictionloss=self.cfg.frictionloss,
|
||||
viscous_damping=self.cfg.viscous_damping,
|
||||
transmission_type=self.cfg.transmission_type,
|
||||
)
|
||||
self._mjs_actuators.append(pos_act)
|
||||
for target_name in target_names:
|
||||
vel_act = create_velocity_actuator(
|
||||
spec,
|
||||
target_name,
|
||||
actuator_name=f"{target_name}_pd_vel",
|
||||
damping=self.cfg.damping,
|
||||
transmission_type=self.cfg.transmission_type,
|
||||
)
|
||||
self._mjs_actuators.append(vel_act)
|
||||
# Effort limit: sum-clamp on the joint/tendon, not on each element.
|
||||
if self.cfg.effort_limit is not None:
|
||||
lim = self.cfg.effort_limit
|
||||
for target_name in target_names:
|
||||
if self.cfg.transmission_type == TransmissionType.JOINT:
|
||||
target = spec.joint(target_name)
|
||||
else:
|
||||
target = spec.tendon(target_name)
|
||||
target.actfrclimited = mujoco.mjtLimited.mjLIMITED_TRUE
|
||||
target.actfrcrange[:] = np.array([-lim, lim])
|
||||
|
||||
def compute(self, cmd: ActuatorCmd) -> torch.Tensor:
|
||||
return torch.cat((cmd.position_target, cmd.velocity_target), dim=1)
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class BuiltinMotorActuatorCfg(ActuatorCfg):
|
||||
"""Configuration for MuJoCo built-in motor actuator.
|
||||
@@ -119,10 +213,6 @@ class BuiltinMotorActuatorCfg(ActuatorCfg):
|
||||
class BuiltinMotorActuator(Actuator[BuiltinMotorActuatorCfg]):
|
||||
"""MuJoCo built-in motor actuator."""
|
||||
|
||||
@property
|
||||
def command_field(self) -> CommandField:
|
||||
return "effort"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
cfg: BuiltinMotorActuatorCfg,
|
||||
@@ -151,6 +241,248 @@ class BuiltinMotorActuator(Actuator[BuiltinMotorActuatorCfg]):
|
||||
return cmd.effort_target
|
||||
|
||||
|
||||
def _or_zeros(t: tuple[float, ...] | None, n: int) -> list[float]:
|
||||
return list(t) if t is not None else [0.0] * n
|
||||
|
||||
|
||||
class DcMotorInputMode(IntEnum):
|
||||
"""What the ``ctrl`` signal of a ``<dcmotor>`` represents.
|
||||
|
||||
Values match MuJoCo's enum, consumed by mjs_setToDCMotor and read as gainprm[8].
|
||||
"""
|
||||
|
||||
VOLTAGE = 0
|
||||
POSITION = 1
|
||||
VELOCITY = 2
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DcMotorDatasheetParams:
|
||||
"""Datasheet characterization of a DC motor."""
|
||||
|
||||
nominal_voltage: float
|
||||
"""Nominal (rated) voltage V_n [V]."""
|
||||
stall_torque: float
|
||||
"""Stall torque tau_stall at V_n [N*m]."""
|
||||
no_load_speed: float
|
||||
"""No-load angular velocity omega_no_load at V_n [rad/s]."""
|
||||
|
||||
def _pack(self) -> tuple[list[float], float, list[float]]:
|
||||
"""Returns (motorconst, resistance, nominal) for set_to_dcmotor."""
|
||||
return (
|
||||
[0.0, 0.0],
|
||||
0.0,
|
||||
[self.nominal_voltage, self.stall_torque, self.no_load_speed],
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DcMotorPhysicalParams:
|
||||
"""Physical characterization of a DC motor."""
|
||||
|
||||
kt: float
|
||||
"""Torque constant [N*m/A]."""
|
||||
ke: float
|
||||
"""Back-EMF constant [V*s/rad]."""
|
||||
resistance: float
|
||||
"""Terminal resistance R [Ohm]."""
|
||||
|
||||
def _pack(self) -> tuple[list[float], float, list[float]]:
|
||||
"""Returns (motorconst, resistance, nominal) for set_to_dcmotor."""
|
||||
return [self.kt, self.ke], self.resistance, [0.0, 0.0, 0.0]
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class BuiltinDcMotorActuatorCfg(ActuatorCfg):
|
||||
"""Native MuJoCo ``<dcmotor>`` wrapper.
|
||||
|
||||
Models a DC motor: torque is derived from voltage via the motor constant K and
|
||||
back-EMF, tau = K * (V - K * omega) / R. The back-EMF term lives in biasprm, so
|
||||
MuJoCo's implicit / implicitfast integrators pick up its velocity derivative as
|
||||
effective damping.
|
||||
|
||||
Three input modes select what ctrl carries:
|
||||
|
||||
* VOLTAGE: ctrl is the drive voltage. cmd.effort_target carries volts, not torque.
|
||||
* POSITION / VELOCITY: an internal PID closes on the setpoint and the motor produces
|
||||
torque from its (Vmax-clamped) voltage output.
|
||||
|
||||
Motor characterization: pass either DcMotorDatasheetParams or DcMotorPhysicalParams
|
||||
as motor_params. mjs_setToDCMotor derives K and R (including the viscous-damping
|
||||
correction) and packs the generic gainprm / biasprm / dynprm slots.
|
||||
|
||||
Optional extensions, off by default: integral_gain / integral_limit, slew_rate,
|
||||
inductance / electrical_time_constant, thermal, lugre, cogging.
|
||||
|
||||
dr.pd_gains randomizes only kp and kd; for DR over the extensions, write directly to
|
||||
actuator_gainprm or actuator_dynprm.
|
||||
"""
|
||||
|
||||
motor_params: DcMotorDatasheetParams | DcMotorPhysicalParams
|
||||
"""Motor characterization. Datasheet form: (V_n, tau_stall, omega_no_load).
|
||||
Physical form: (Kt, Ke, R)."""
|
||||
|
||||
mode: DcMotorInputMode = DcMotorInputMode.POSITION
|
||||
"""ctrl input semantics. See class docstring."""
|
||||
|
||||
stiffness: float = 0.0
|
||||
"""PID proportional gain kp. Required in POSITION / VELOCITY mode; must be
|
||||
0 in VOLTAGE mode."""
|
||||
|
||||
damping: float = 0.0
|
||||
"""PID derivative gain kd. Used in POSITION / VELOCITY mode; must be 0 in
|
||||
VOLTAGE mode."""
|
||||
|
||||
voltage_limit: float = 0.0
|
||||
"""Max drive voltage Vmax. Required in POSITION / VELOCITY mode (clamps the
|
||||
PID output). In VOLTAGE mode it is an optional clamp on ctrl; 0 disables."""
|
||||
|
||||
integral_gain: float = 0.0
|
||||
"""PID integral gain ki. In position mode the integrator tracks
|
||||
ki * integral(target - q); in velocity mode, ki * (integral(target) - q).
|
||||
Must be 0 in VOLTAGE mode."""
|
||||
|
||||
integral_limit: float = 0.0
|
||||
"""Anti-windup clamp Imax on the integrator state. 0 disables (the
|
||||
integrator can run away)."""
|
||||
|
||||
slew_rate: float = 0.0
|
||||
"""Max rate of change of ctrl per second. 0 disables."""
|
||||
|
||||
effort_limit: float | None = None
|
||||
"""Continuous torque cap [N*m]. Sets actuator_forcerange. None leaves the
|
||||
per-element forcerange unset."""
|
||||
|
||||
gear: float = 1.0
|
||||
"""Mechanical gear ratio."""
|
||||
|
||||
inductance: float = 0.0
|
||||
"""Winding inductance L [H]. Enables first-order electrical dynamics on the
|
||||
motor current. MuJoCo internally uses te = L / R; pass
|
||||
electrical_time_constant directly to skip the divide. 0 disables."""
|
||||
|
||||
electrical_time_constant: float = 0.0
|
||||
"""Alternative to inductance: specify te [s] directly. Ignored if
|
||||
inductance > 0. 0 disables."""
|
||||
|
||||
thermal: tuple[float, float, float, float, float, float] | None = None
|
||||
"""Thermal model (R_thermal, C_thermal, tau_thermal, alpha, T0, T_ambient).
|
||||
See MuJoCo's ``<dcmotor thermal=...>`` reference for units and which of the
|
||||
first three may be underspecified. Effective resistance becomes
|
||||
R * (1 + alpha * (T + T_ambient - T0)). None disables."""
|
||||
|
||||
cogging: tuple[float, float, float] | None = None
|
||||
"""Cogging torque (amplitude, periodicity, phase) in (N*m, cycles per unit
|
||||
length, rad). Models magnetic torque ripple from rotor-stator interaction;
|
||||
at joint angle q the contribution is amplitude * sin(periodicity * q + phase).
|
||||
|
||||
Added *after* effort_limit is enforced, matching MuJoCo's physical model:
|
||||
effort_limit bounds the electromagnetic torque (the current limit), not the
|
||||
mechanical torque. Total joint torque can exceed effort_limit by up to
|
||||
amplitude. None disables."""
|
||||
|
||||
lugre: tuple[float, float, float, float, float] | None = None
|
||||
"""LuGre friction (sigma0, sigma1, F_Coulomb, F_Stribeck, v_Stribeck).
|
||||
Stick-slip friction with bristle-deflection state. Subtracted from joint
|
||||
torque after the effort_limit clamp (mechanical, like cogging). None
|
||||
disables."""
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
super().__post_init__()
|
||||
if self.transmission_type == TransmissionType.SITE:
|
||||
raise ValueError(
|
||||
"BuiltinDcMotorActuatorCfg does not support SITE transmission. "
|
||||
"Use BuiltinMotorActuatorCfg for site transmission."
|
||||
)
|
||||
|
||||
if self.mode in (DcMotorInputMode.POSITION, DcMotorInputMode.VELOCITY):
|
||||
if self.stiffness <= 0.0:
|
||||
raise ValueError(f"{self.mode.name} mode requires stiffness > 0.")
|
||||
if self.voltage_limit <= 0.0:
|
||||
raise ValueError(f"{self.mode.name} mode requires voltage_limit > 0.")
|
||||
else:
|
||||
if self.stiffness != 0.0 or self.damping != 0.0 or self.integral_gain != 0.0:
|
||||
raise ValueError(
|
||||
"stiffness, damping, and integral_gain are unused in VOLTAGE mode."
|
||||
)
|
||||
|
||||
for name in (
|
||||
"integral_gain",
|
||||
"integral_limit",
|
||||
"slew_rate",
|
||||
"inductance",
|
||||
"electrical_time_constant",
|
||||
):
|
||||
if getattr(self, name) < 0.0:
|
||||
raise ValueError(f"{name} must be non-negative.")
|
||||
|
||||
def build(
|
||||
self, entity: Entity, target_ids: list[int], target_names: list[str]
|
||||
) -> BuiltinDcMotorActuator:
|
||||
return BuiltinDcMotorActuator(self, entity, target_ids, target_names)
|
||||
|
||||
|
||||
class BuiltinDcMotorActuator(Actuator[BuiltinDcMotorActuatorCfg]):
|
||||
"""MuJoCo native ``<dcmotor>``: one actuator per target."""
|
||||
|
||||
def edit_spec(self, spec: mujoco.MjSpec, target_names: list[str]) -> None:
|
||||
cfg = self.cfg
|
||||
motorconst, resistance, nominal = cfg.motor_params._pack()
|
||||
saturation = (
|
||||
[cfg.effort_limit, 0.0, 0.0] if cfg.effort_limit is not None else [0.0] * 3
|
||||
)
|
||||
controller = [
|
||||
cfg.stiffness, # kp
|
||||
cfg.integral_gain, # ki
|
||||
cfg.damping, # kd
|
||||
cfg.slew_rate, # slewmax
|
||||
cfg.integral_limit, # Imax (anti-windup)
|
||||
cfg.voltage_limit, # v_max
|
||||
]
|
||||
# SITE is rejected in __post_init__, so only JOINT and TENDON remain.
|
||||
trntype = (
|
||||
mujoco.mjtTrn.mjTRN_JOINT
|
||||
if cfg.transmission_type == TransmissionType.JOINT
|
||||
else mujoco.mjtTrn.mjTRN_TENDON
|
||||
)
|
||||
|
||||
for target_name in target_names:
|
||||
actuator = spec.add_actuator(name=target_name, target=target_name)
|
||||
actuator.trntype = trntype
|
||||
actuator.gear[0] = cfg.gear
|
||||
actuator.set_to_dcmotor(
|
||||
motorconst=motorconst,
|
||||
resistance=resistance,
|
||||
nominal=nominal,
|
||||
saturation=saturation,
|
||||
controller=controller,
|
||||
cogging=_or_zeros(cfg.cogging, 3),
|
||||
inductance=[cfg.inductance, cfg.electrical_time_constant],
|
||||
thermal=_or_zeros(cfg.thermal, 6),
|
||||
lugre=_or_zeros(cfg.lugre, 5),
|
||||
input_mode=cfg.mode,
|
||||
)
|
||||
|
||||
apply_target_overrides(
|
||||
spec,
|
||||
target_name,
|
||||
cfg.transmission_type,
|
||||
armature=cfg.armature,
|
||||
frictionloss=cfg.frictionloss,
|
||||
viscous_damping=cfg.viscous_damping,
|
||||
)
|
||||
|
||||
self._mjs_actuators.append(actuator)
|
||||
|
||||
def compute(self, cmd: ActuatorCmd) -> torch.Tensor:
|
||||
if self.cfg.mode == DcMotorInputMode.POSITION:
|
||||
return cmd.position_target
|
||||
if self.cfg.mode == DcMotorInputMode.VELOCITY:
|
||||
return cmd.velocity_target
|
||||
# voltage mode: ctrl is the drive voltage carried in effort_target.
|
||||
return cmd.effort_target
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class BuiltinVelocityActuatorCfg(ActuatorCfg):
|
||||
"""Configuration for MuJoCo built-in velocity actuator.
|
||||
@@ -182,10 +514,6 @@ class BuiltinVelocityActuatorCfg(ActuatorCfg):
|
||||
class BuiltinVelocityActuator(Actuator[BuiltinVelocityActuatorCfg]):
|
||||
"""MuJoCo built-in velocity actuator."""
|
||||
|
||||
@property
|
||||
def command_field(self) -> CommandField:
|
||||
return "velocity"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
cfg: BuiltinVelocityActuatorCfg,
|
||||
@@ -260,10 +588,6 @@ class BuiltinMuscleActuatorCfg(ActuatorCfg):
|
||||
class BuiltinMuscleActuator(Actuator[BuiltinMuscleActuatorCfg]):
|
||||
"""MuJoCo built-in muscle actuator."""
|
||||
|
||||
@property
|
||||
def command_field(self) -> CommandField:
|
||||
return "effort"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
cfg: BuiltinMuscleActuatorCfg,
|
||||
|
||||
@@ -33,6 +33,9 @@ class DcMotorActuatorCfg(IdealPdActuatorCfg):
|
||||
Note: effort_limit should be explicitly set to a realistic value for proper
|
||||
motor modeling. Using the default (inf) will trigger a warning. Use
|
||||
IdealPdActuator if unlimited torque is desired.
|
||||
|
||||
For a native MuJoCo ``<dcmotor>`` with back-EMF, voltage saturation, and
|
||||
configurable ``Kt`` / ``Ke`` / ``R``, see ``BuiltinDcMotorActuator``.
|
||||
"""
|
||||
|
||||
saturation_effort: float
|
||||
|
||||
@@ -9,7 +9,7 @@ import mujoco
|
||||
import mujoco_warp as mjwarp
|
||||
import torch
|
||||
|
||||
from mjlab.actuator.actuator import Actuator, ActuatorCfg, ActuatorCmd, CommandField
|
||||
from mjlab.actuator.actuator import Actuator, ActuatorCfg, ActuatorCmd
|
||||
from mjlab.utils.spec import create_motor_actuator
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -38,10 +38,6 @@ class IdealPdActuatorCfg(ActuatorCfg):
|
||||
class IdealPdActuator(Actuator, Generic[IdealPdCfgT]):
|
||||
"""Ideal PD control actuator."""
|
||||
|
||||
@property
|
||||
def command_field(self) -> CommandField:
|
||||
return "position"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
cfg: IdealPdCfgT,
|
||||
|
||||
@@ -102,11 +102,6 @@
|
||||
diaginertia="0.00167218 0.0016161 0.000217621"/>
|
||||
<joint name="left_ankle_roll_joint" axis="1 0 0" range="-0.2618 0.2618"/>
|
||||
<geom class="visual" material="black" mesh="left_ankle_roll_link"/>
|
||||
<!-- <geom name="left_foot1_collision" class="foot_capsule" fromto="0.1 -0.026 -0.025 0.05 -0.027
|
||||
-0.025"/>
|
||||
<geom name="left_foot2_collision" class="foot_capsule" fromto="-0.045 0 -0.015 0.12 0 -0.015"
|
||||
size="0.02"/>
|
||||
<geom name="left_foot3_collision" class="foot_capsule" fromto="0.1 0.026 -0.025 0.05 0.026 -0.025"/> -->
|
||||
<geom name="left_foot1_collision" class="foot_capsule" fromto="0.1 -0.026 -0.025 0.05 -0.027 -0.025"/>
|
||||
<geom name="left_foot2_collision" class="foot_capsule"
|
||||
fromto="-0.044 -0.018 -0.025 0.123 -0.018 -0.025"/>
|
||||
@@ -156,11 +151,6 @@
|
||||
diaginertia="0.00167218 0.0016161 0.000217621"/>
|
||||
<joint name="right_ankle_roll_joint" axis="1 0 0" range="-0.2618 0.2618"/>
|
||||
<geom class="visual" material="black" mesh="right_ankle_roll_link"/>
|
||||
<!-- <geom name="right_foot1_collision" class="foot_capsule" fromto="0.1 -0.026 -0.025 0.05 -0.026
|
||||
-0.025"/>
|
||||
<geom name="right_foot2_collision" class="foot_capsule" fromto="-0.045 0 -0.015 0.12 0 -0.015"
|
||||
size="0.02"/>
|
||||
<geom name="right_foot3_collision" class="foot_capsule" fromto="0.1 0.026 -0.025 0.05 0.026 -0.025"/> -->
|
||||
<geom name="right_foot1_collision" class="foot_capsule" fromto="0.1 -0.026 -0.025 0.05 -0.026 -0.025"/>
|
||||
<geom name="right_foot2_collision" class="foot_capsule"
|
||||
fromto="-0.044 -0.018 -0.025 0.123 -0.018 -0.025"/>
|
||||
@@ -312,6 +302,7 @@
|
||||
<gyro name="imu_ang_vel" site="imu_in_pelvis"/>
|
||||
<velocimeter name="imu_lin_vel" site="imu_in_pelvis"/>
|
||||
<accelerometer name="imu_lin_acc" site="imu_in_pelvis"/>
|
||||
<framezaxis name="imu_upvector" objtype="body" objname="world" reftype="site" refname="imu_in_pelvis"/>
|
||||
<subtreeangmom name="root_angmom" body="pelvis"/>
|
||||
</sensor>
|
||||
</mujoco>
|
||||
|
||||
@@ -165,6 +165,7 @@
|
||||
<gyro name="imu_ang_vel" site="imu"/>
|
||||
<velocimeter name="imu_lin_vel" site="imu"/>
|
||||
<accelerometer name="imu_lin_acc" site="imu"/>
|
||||
<framezaxis name="imu_upvector" objtype="body" objname="world" reftype="site" refname="imu"/>
|
||||
<subtreeangmom name="root_angmom" body="trunk"/>
|
||||
</sensor>
|
||||
</mujoco>
|
||||
|
||||
@@ -3,6 +3,5 @@ from mjlab.entity.entity import Entity as Entity
|
||||
from mjlab.entity.entity import EntityArticulationInfoCfg as EntityArticulationInfoCfg
|
||||
from mjlab.entity.entity import EntityCfg as EntityCfg
|
||||
from mjlab.entity.entity import EntityIndexing as EntityIndexing
|
||||
from mjlab.entity.entity import VariantCfg as VariantCfg
|
||||
from mjlab.entity.entity import VariantEntityCfg as VariantEntityCfg
|
||||
from mjlab.entity.entity import VariantMetadata as VariantMetadata
|
||||
from mjlab.entity.variants import VariantEntityCfg as VariantEntityCfg
|
||||
from mjlab.entity.variants import VariantMetadata as VariantMetadata
|
||||
|
||||
@@ -3,7 +3,7 @@ from __future__ import annotations
|
||||
import warnings
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Callable, Sequence
|
||||
from typing import TYPE_CHECKING, Callable, Sequence
|
||||
|
||||
import mujoco
|
||||
import mujoco_warp as mjwarp
|
||||
@@ -18,14 +18,13 @@ from mjlab.entity.data import EntityData
|
||||
from mjlab.utils import spec_config as spec_cfg
|
||||
from mjlab.utils.lab_api.string import resolve_matching_names
|
||||
from mjlab.utils.mujoco import dof_width, qpos_width
|
||||
from mjlab.utils.spec import (
|
||||
auto_wrap_fixed_base_mocap,
|
||||
copy_mesh_data,
|
||||
validate_variant_structure,
|
||||
)
|
||||
from mjlab.utils.spec import auto_wrap_fixed_base_mocap
|
||||
from mjlab.utils.string import resolve_expr
|
||||
from mjlab.utils.xml import fix_spec_xml, strip_buffer_textures
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from mjlab.entity.variants import VariantMetadata
|
||||
|
||||
|
||||
@dataclass(frozen=False)
|
||||
class EntityIndexing:
|
||||
@@ -67,79 +66,6 @@ class EntityIndexing:
|
||||
return self.bodies[0].id
|
||||
|
||||
|
||||
@dataclass
|
||||
class VariantCfg:
|
||||
"""One object variant for per-world mesh randomization.
|
||||
|
||||
Each variant provides a ``spec_fn`` that returns an MjSpec for one object.
|
||||
The ``weight`` controls what fraction of worlds use this variant.
|
||||
"""
|
||||
|
||||
spec_fn: Callable[[], mujoco.MjSpec]
|
||||
weight: float = 1.0
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BodyInertialMetadata:
|
||||
"""Explicit inertial properties for one body in a mesh variant."""
|
||||
|
||||
body_name: str
|
||||
mass: float
|
||||
ipos: tuple[float, float, float]
|
||||
inertia: tuple[float, float, float]
|
||||
iquat: tuple[float, float, float, float]
|
||||
|
||||
|
||||
@dataclass
|
||||
class VariantMetadata:
|
||||
"""Bookkeeping produced by Entity when merging variant specs."""
|
||||
|
||||
variant_names: tuple[str, ...]
|
||||
variant_weights: tuple[float, ...]
|
||||
# Per-variant ordered mesh names for each geom slot. Shorter variants
|
||||
# have None for padding slots that should be disabled (dataid = -1).
|
||||
variant_mesh_names: tuple[tuple[str | None, ...], ...]
|
||||
num_mesh_geoms: int # Max mesh geom count after padding.
|
||||
# Per-variant explicit body inertials. Names are local to the variant spec;
|
||||
# build_mesh_variant_model prefixes them with the scene entity name when
|
||||
# applying them.
|
||||
variant_body_inertials: tuple[tuple[BodyInertialMetadata, ...], ...] = ()
|
||||
|
||||
|
||||
def _iter_body_tree(body: mujoco.MjsBody):
|
||||
yield body
|
||||
for child in body.bodies:
|
||||
yield from _iter_body_tree(child)
|
||||
|
||||
|
||||
def _collect_explicit_body_inertials(
|
||||
root_body: mujoco.MjsBody,
|
||||
) -> tuple[BodyInertialMetadata, ...]:
|
||||
inertials: list[BodyInertialMetadata] = []
|
||||
for body in _iter_body_tree(root_body):
|
||||
if not body.name or not body.explicitinertial:
|
||||
continue
|
||||
inertials.append(
|
||||
BodyInertialMetadata(
|
||||
body_name=body.name,
|
||||
mass=float(body.mass),
|
||||
ipos=(float(body.ipos[0]), float(body.ipos[1]), float(body.ipos[2])),
|
||||
inertia=(
|
||||
float(body.inertia[0]),
|
||||
float(body.inertia[1]),
|
||||
float(body.inertia[2]),
|
||||
),
|
||||
iquat=(
|
||||
float(body.iquat[0]),
|
||||
float(body.iquat[1]),
|
||||
float(body.iquat[2]),
|
||||
float(body.iquat[3]),
|
||||
),
|
||||
)
|
||||
)
|
||||
return tuple(inertials)
|
||||
|
||||
|
||||
@dataclass
|
||||
class EntityCfg:
|
||||
@dataclass
|
||||
@@ -187,50 +113,6 @@ class EntityArticulationInfoCfg:
|
||||
soft_joint_pos_limit_factor: float = 1.0
|
||||
|
||||
|
||||
def _variant_spec_fn_unset() -> mujoco.MjSpec:
|
||||
"""Sentinel default for ``VariantEntityCfg.spec_fn``.
|
||||
|
||||
``VariantEntityCfg`` builds its spec from ``variants`` via
|
||||
``Entity._build_merged_spec``; the inherited ``spec_fn`` field is unused.
|
||||
Identity comparison against this sentinel detects accidental user overrides.
|
||||
"""
|
||||
raise AssertionError(
|
||||
"VariantEntityCfg.spec_fn should never be called; the merged spec is "
|
||||
"built from `variants`."
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class VariantEntityCfg(EntityCfg):
|
||||
"""Entity config for per-world mesh variants.
|
||||
|
||||
Instead of a single ``spec_fn``, provide a dict of named variants.
|
||||
Each world gets a variant assigned proportionally by weight. The
|
||||
merged spec (with all variant meshes and padded geoms) is built
|
||||
automatically.
|
||||
|
||||
All variants must share the same kinematic structure (same bodies,
|
||||
joints, joint types). Only mesh geoms can differ.
|
||||
|
||||
Variant assignment is fixed at ``Simulation`` initialization; it does
|
||||
not resample on episode reset. Pass the per-variant spec via
|
||||
:class:`VariantCfg` rather than setting ``spec_fn`` directly.
|
||||
"""
|
||||
|
||||
variants: dict[str, VariantCfg] = field(default_factory=dict)
|
||||
"""Named mesh variants with weights."""
|
||||
|
||||
spec_fn: Callable[[], mujoco.MjSpec] = field(default=_variant_spec_fn_unset)
|
||||
"""Unused on ``VariantEntityCfg``; the merged spec is built from ``variants``."""
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.spec_fn is not _variant_spec_fn_unset:
|
||||
raise ValueError(
|
||||
"VariantEntityCfg.spec_fn cannot be set; pass per-variant specs via "
|
||||
"VariantCfg(spec_fn=...) inside `variants` instead."
|
||||
)
|
||||
|
||||
|
||||
class Entity:
|
||||
"""An entity represents a physical object in the simulation.
|
||||
|
||||
@@ -272,130 +154,13 @@ class Entity:
|
||||
self._add_initial_state_keyframe()
|
||||
|
||||
def _build_spec(self) -> None:
|
||||
from mjlab.entity.variants import VariantEntityCfg, build_merged_variant_spec
|
||||
|
||||
if isinstance(self.cfg, VariantEntityCfg):
|
||||
self._build_merged_spec()
|
||||
self._spec, self._variant_metadata = build_merged_variant_spec(self.cfg)
|
||||
else:
|
||||
self._spec = auto_wrap_fixed_base_mocap(self.cfg.spec_fn)()
|
||||
|
||||
def _build_merged_spec(self) -> None:
|
||||
"""Build a merged spec from multiple variant specs.
|
||||
|
||||
Validates that all variants share the same kinematic structure,
|
||||
merges all mesh assets into a single spec, and pads the body to
|
||||
the max mesh geom count across variants.
|
||||
"""
|
||||
assert isinstance(self.cfg, VariantEntityCfg)
|
||||
variants = self.cfg.variants
|
||||
if not variants:
|
||||
raise ValueError("VariantEntityCfg.variants must contain at least one entry.")
|
||||
|
||||
variant_names: list[str] = []
|
||||
variant_weights: list[float] = []
|
||||
variant_specs: list[mujoco.MjSpec] = []
|
||||
for name, vcfg in variants.items():
|
||||
variant_names.append(name)
|
||||
variant_weights.append(vcfg.weight)
|
||||
variant_specs.append(vcfg.spec_fn())
|
||||
|
||||
# Find root body in each variant.
|
||||
variant_bodies: list[mujoco.MjsBody] = []
|
||||
for i, spec in enumerate(variant_specs):
|
||||
children = list(spec.worldbody.bodies)
|
||||
if len(children) != 1:
|
||||
raise ValueError(
|
||||
f"Variant '{variant_names[i]}' must have exactly one "
|
||||
f"root body under worldbody, got {len(children)}."
|
||||
)
|
||||
variant_bodies.append(children[0])
|
||||
|
||||
validate_variant_structure(variant_names, variant_bodies)
|
||||
|
||||
# Variant entities must be floating-base. Mocap auto-wrap is not applied
|
||||
# for variant entities, so fixed-base variants would silently stack at
|
||||
# the world origin. Variants share joint structure (validated above), so
|
||||
# checking the first is sufficient.
|
||||
ref_joints = list(variant_bodies[0].joints)
|
||||
if not ref_joints or ref_joints[0].type != mujoco.mjtJoint.mjJNT_FREE:
|
||||
raise ValueError(
|
||||
"VariantEntityCfg requires floating-base variants. Each variant's "
|
||||
"root body must declare a free joint via body.add_freejoint(); "
|
||||
"fixed-base variants are not supported."
|
||||
)
|
||||
|
||||
variant_body_inertials = tuple(
|
||||
_collect_explicit_body_inertials(body) for body in variant_bodies
|
||||
)
|
||||
|
||||
# Collect original mesh names per variant BEFORE any renaming.
|
||||
variant_orig_mesh_names: list[list[str]] = []
|
||||
variant_mesh_geom_counts: list[int] = []
|
||||
for body in variant_bodies:
|
||||
orig_names = [
|
||||
g.meshname for g in body.geoms if g.type == mujoco.mjtGeom.mjGEOM_MESH
|
||||
]
|
||||
variant_orig_mesh_names.append(orig_names)
|
||||
variant_mesh_geom_counts.append(len(orig_names))
|
||||
|
||||
max_mesh_geoms = max(variant_mesh_geom_counts)
|
||||
|
||||
# Use first variant as template. Prefix ALL mesh names with
|
||||
# variant name to avoid collisions across variants.
|
||||
template_spec = variant_specs[0]
|
||||
template_body = variant_bodies[0]
|
||||
|
||||
# Rename template meshes first.
|
||||
template_prefix = f"{variant_names[0]}/"
|
||||
old_to_new: dict[str, str] = {}
|
||||
for mesh in template_spec.meshes:
|
||||
new_name = f"{template_prefix}{mesh.name}"
|
||||
old_to_new[mesh.name] = new_name
|
||||
mesh.name = new_name
|
||||
for g in template_body.geoms:
|
||||
if g.meshname in old_to_new:
|
||||
g.meshname = old_to_new[g.meshname]
|
||||
|
||||
# Copy mesh assets from other variants.
|
||||
for i in range(1, len(variant_specs)):
|
||||
prefix = f"{variant_names[i]}/"
|
||||
for mesh in variant_specs[i].meshes:
|
||||
new_mesh = template_spec.add_mesh()
|
||||
new_mesh.name = f"{prefix}{mesh.name}"
|
||||
copy_mesh_data(mesh, new_mesh)
|
||||
|
||||
# Pad body to max mesh geom count.
|
||||
current_count = variant_mesh_geom_counts[0]
|
||||
if max_mesh_geoms > current_count:
|
||||
longest_idx = max(
|
||||
range(len(variant_mesh_geom_counts)),
|
||||
key=lambda j: variant_mesh_geom_counts[j],
|
||||
)
|
||||
longest_prefix = f"{variant_names[longest_idx]}/"
|
||||
longest_names = variant_orig_mesh_names[longest_idx]
|
||||
for k in range(current_count, max_mesh_geoms):
|
||||
geom = template_body.add_geom()
|
||||
geom.type = mujoco.mjtGeom.mjGEOM_MESH
|
||||
geom.meshname = f"{longest_prefix}{longest_names[k]}"
|
||||
geom.contype = 1
|
||||
geom.conaffinity = 1
|
||||
|
||||
# Build variant_mesh_names: use original names with variant prefix.
|
||||
variant_mesh_name_lists: list[tuple[str | None, ...]] = []
|
||||
for i, orig_names in enumerate(variant_orig_mesh_names):
|
||||
prefix = f"{variant_names[i]}/"
|
||||
names: list[str | None] = [f"{prefix}{n}" for n in orig_names]
|
||||
while len(names) < max_mesh_geoms:
|
||||
names.append(None)
|
||||
variant_mesh_name_lists.append(tuple(names))
|
||||
|
||||
self._variant_metadata = VariantMetadata(
|
||||
variant_names=tuple(variant_names),
|
||||
variant_weights=tuple(variant_weights),
|
||||
variant_mesh_names=tuple(variant_mesh_name_lists),
|
||||
num_mesh_geoms=max_mesh_geoms,
|
||||
variant_body_inertials=variant_body_inertials,
|
||||
)
|
||||
self._spec = template_spec
|
||||
|
||||
@property
|
||||
def variant_metadata(self) -> VariantMetadata | None:
|
||||
return self._variant_metadata
|
||||
@@ -404,6 +169,16 @@ class Entity:
|
||||
self._all_joints = self._spec.joints
|
||||
self._free_joint = None
|
||||
self._non_free_joints = tuple(self._all_joints)
|
||||
|
||||
free_joints = [j for j in self._all_joints if j.type == mujoco.mjtJoint.mjJNT_FREE]
|
||||
if len(free_joints) > 1:
|
||||
raise ValueError(
|
||||
f"Entity spec has {len(free_joints)} freejoints. An Entity models a "
|
||||
"single rigid- or articulated-body system with at most one freejoint, "
|
||||
"which serves as its root. Model each detached floating body as its own "
|
||||
"entry in SceneCfg.entities instead."
|
||||
)
|
||||
|
||||
if self._all_joints and self._all_joints[0].type == mujoco.mjtJoint.mjJNT_FREE:
|
||||
self._free_joint = self._all_joints[0]
|
||||
if not self._free_joint.name:
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -184,7 +184,7 @@ class ManagerBasedRlEnv:
|
||||
# Initialize base environment state.
|
||||
self.cfg = cfg
|
||||
if self.cfg.seed is not None:
|
||||
self.cfg.seed = self.seed(self.cfg.seed, device=device)
|
||||
self.cfg.seed = self.seed(self.cfg.seed)
|
||||
self._sim_step_counter = 0
|
||||
self.extras = {}
|
||||
self.obs_buf = {}
|
||||
@@ -194,21 +194,13 @@ class ManagerBasedRlEnv:
|
||||
|
||||
# Initialize scene and simulation.
|
||||
self.scene = Scene(self.cfg.scene, device=device)
|
||||
if self.scene.has_mesh_variants:
|
||||
self.sim = Simulation(
|
||||
num_envs=self.scene.num_envs,
|
||||
cfg=self.cfg.sim,
|
||||
spec=self.scene.spec,
|
||||
variant_info=self.scene.collect_variant_info(),
|
||||
device=device,
|
||||
)
|
||||
else:
|
||||
self.sim = Simulation(
|
||||
num_envs=self.scene.num_envs,
|
||||
cfg=self.cfg.sim,
|
||||
model=self.scene.compile(),
|
||||
device=device,
|
||||
)
|
||||
self.sim = Simulation(
|
||||
num_envs=self.scene.num_envs,
|
||||
cfg=self.cfg.sim,
|
||||
spec=self.scene.spec,
|
||||
variant_info=self.scene.collect_variant_info(),
|
||||
device=device,
|
||||
)
|
||||
|
||||
self.scene.initialize(
|
||||
mj_model=self.sim.mj_model,
|
||||
@@ -373,6 +365,7 @@ class ManagerBasedRlEnv:
|
||||
env_ids = torch.arange(self.num_envs, dtype=torch.int64, device=self.device)
|
||||
if seed is not None:
|
||||
self.seed(seed)
|
||||
self.extras["log"] = dict()
|
||||
self._reset_idx(env_ids)
|
||||
self.scene.write_data_to_sim()
|
||||
self.sim.forward()
|
||||
@@ -422,6 +415,7 @@ class ManagerBasedRlEnv:
|
||||
"reset(env_ids=...) before calling step() again when auto_reset=False."
|
||||
)
|
||||
|
||||
self.extras["log"] = dict()
|
||||
self.action_manager.process_action(action.to(self.device))
|
||||
|
||||
for _ in range(self.cfg.decimation):
|
||||
@@ -484,6 +478,9 @@ class ManagerBasedRlEnv:
|
||||
self.extras,
|
||||
)
|
||||
|
||||
def get_observations(self) -> dict:
|
||||
return self.observation_manager.compute()
|
||||
|
||||
def render(self) -> np.ndarray | None:
|
||||
if self.render_mode == "human" or self.render_mode is None:
|
||||
return None
|
||||
@@ -506,11 +503,12 @@ class ManagerBasedRlEnv:
|
||||
self._offline_renderer.close()
|
||||
self.recorder_manager.close()
|
||||
|
||||
def seed(self, seed: int = -1, device: str | torch.device | None = None) -> int:
|
||||
@staticmethod
|
||||
def seed(seed: int = -1) -> int:
|
||||
if seed == -1:
|
||||
seed = np.random.randint(0, 10_000)
|
||||
print_info(f"Setting seed: {seed}")
|
||||
random_utils.seed_rng(seed, device=device if device is not None else self.device)
|
||||
random_utils.seed_rng(seed)
|
||||
return seed
|
||||
|
||||
def update_visualizers(self, visualizer: DebugVisualizer) -> None:
|
||||
@@ -564,7 +562,6 @@ class ManagerBasedRlEnv:
|
||||
)
|
||||
|
||||
# NOTE: This is order sensitive.
|
||||
self.extras["log"] = dict()
|
||||
# observation manager.
|
||||
info = self.observation_manager.reset(env_ids)
|
||||
self.extras["log"].update(info)
|
||||
|
||||
@@ -7,23 +7,65 @@ from typing import TYPE_CHECKING, Literal
|
||||
import torch
|
||||
|
||||
from mjlab.actuator import (
|
||||
BuiltinPositionActuator,
|
||||
BuiltinVelocityActuator,
|
||||
BuiltinMotorActuator,
|
||||
IdealPdActuator,
|
||||
BuiltinDcMotorActuator,
|
||||
BuiltinPdActuator,
|
||||
BuiltinPositionActuator,
|
||||
IdealPdActuator,
|
||||
)
|
||||
from mjlab.actuator.actuator import TransmissionType
|
||||
from mjlab.actuator.builtin_actuator import DcMotorInputMode
|
||||
from mjlab.actuator.xml_actuator import XmlActuator
|
||||
from mjlab.entity import Entity
|
||||
from mjlab.managers.event_manager import requires_model_fields
|
||||
from mjlab.managers.scene_entity_config import SceneEntityCfg
|
||||
|
||||
from ._core import _DEFAULT_ASSET_CFG
|
||||
from ._types import resolve_distribution
|
||||
from ._types import Operation, resolve_distribution, resolve_operation
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from mjlab.envs import ManagerBasedRlEnv
|
||||
|
||||
|
||||
def _resolve_actuators(asset: Entity, asset_cfg: SceneEntityCfg) -> list:
|
||||
"""Resolve actuator objects from SceneEntityCfg.
|
||||
|
||||
SceneEntityCfg actuator_ids/name resolution is based on spec actuators, while
|
||||
runtime ``asset.actuators`` may contain grouped actuator objects. For grouped
|
||||
actuators, map matched actuator target names back to the owning actuator object.
|
||||
"""
|
||||
if asset_cfg.actuator_names is not None:
|
||||
matched_names = asset_cfg.actuator_names
|
||||
if isinstance(matched_names, str):
|
||||
matched_names = [matched_names]
|
||||
resolved = []
|
||||
for actuator in asset.actuators:
|
||||
if any(name in actuator.target_names for name in matched_names):
|
||||
resolved.append(actuator)
|
||||
return resolved
|
||||
|
||||
if isinstance(asset_cfg.actuator_ids, list):
|
||||
if all(0 <= i < len(asset.actuators) for i in asset_cfg.actuator_ids):
|
||||
return [asset.actuators[i] for i in asset_cfg.actuator_ids]
|
||||
|
||||
resolved = []
|
||||
seen = set()
|
||||
actuator_names = asset.actuator_names
|
||||
for i in asset_cfg.actuator_ids:
|
||||
if not (0 <= i < len(actuator_names)):
|
||||
continue
|
||||
target_name = actuator_names[i]
|
||||
for actuator_idx, actuator in enumerate(asset.actuators):
|
||||
if target_name in actuator.target_names and actuator_idx not in seen:
|
||||
resolved.append(actuator)
|
||||
seen.add(actuator_idx)
|
||||
break
|
||||
return resolved
|
||||
elif isinstance(asset_cfg.actuator_ids, slice):
|
||||
return asset.actuators[asset_cfg.actuator_ids]
|
||||
else:
|
||||
return [asset.actuators[asset_cfg.actuator_ids]]
|
||||
|
||||
|
||||
@requires_model_fields("actuator_gainprm", "actuator_biasprm")
|
||||
def pd_gains(
|
||||
env: ManagerBasedRlEnv,
|
||||
@@ -32,7 +74,7 @@ def pd_gains(
|
||||
kd_range: tuple[float, float],
|
||||
asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
|
||||
distribution: Literal["uniform", "log_uniform"] = "uniform",
|
||||
operation: Literal["scale", "abs"] = "scale",
|
||||
operation: Operation | str = "scale",
|
||||
) -> None:
|
||||
"""Randomize PD stiffness and damping gains.
|
||||
|
||||
@@ -46,6 +88,11 @@ def pd_gains(
|
||||
operation: "scale" multiplies default gains by sampled values, "abs" sets
|
||||
absolute values.
|
||||
"""
|
||||
op = resolve_operation(operation)
|
||||
if op.name not in ("scale", "abs"):
|
||||
raise ValueError(
|
||||
f"pd_gains only supports 'scale' and 'abs' operations, got {op.name!r}"
|
||||
)
|
||||
asset: Entity = env.scene[asset_cfg.name]
|
||||
|
||||
if env_ids is None:
|
||||
@@ -53,34 +100,36 @@ def pd_gains(
|
||||
else:
|
||||
env_ids = env_ids.to(env.device, dtype=torch.int)
|
||||
|
||||
if isinstance(asset_cfg.actuator_ids, list):
|
||||
actuators = [asset.actuators[i] for i in asset_cfg.actuator_ids]
|
||||
elif isinstance(asset_cfg.actuator_ids, slice):
|
||||
actuators = asset.actuators[asset_cfg.actuator_ids]
|
||||
else:
|
||||
actuators = [asset.actuators[asset_cfg.actuator_ids]]
|
||||
actuators = _resolve_actuators(asset, asset_cfg)
|
||||
|
||||
for actuator in actuators:
|
||||
ctrl_ids = actuator.global_ctrl_ids
|
||||
# Each target needs one kp draw and one kd draw. For single-element
|
||||
# actuators that's len(ctrl_ids) of each; for BuiltinPd the ctrl tensor
|
||||
# has 2*N entries but only N independent kp/kd values, so we sample
|
||||
# num_targets to avoid throwing the other half away.
|
||||
n_gains = (
|
||||
actuator.num_targets if isinstance(actuator, BuiltinPdActuator) else len(ctrl_ids)
|
||||
)
|
||||
|
||||
dist = resolve_distribution(distribution)
|
||||
kp_samples = dist.sample(
|
||||
torch.tensor(kp_range[0], device=env.device),
|
||||
torch.tensor(kp_range[1], device=env.device),
|
||||
(len(env_ids), len(ctrl_ids)),
|
||||
(len(env_ids), n_gains),
|
||||
env.device,
|
||||
)
|
||||
kd_samples = dist.sample(
|
||||
torch.tensor(kd_range[0], device=env.device),
|
||||
torch.tensor(kd_range[1], device=env.device),
|
||||
(len(env_ids), len(ctrl_ids)),
|
||||
(len(env_ids), n_gains),
|
||||
env.device,
|
||||
)
|
||||
|
||||
if isinstance(actuator, BuiltinPositionActuator) or (
|
||||
isinstance(actuator, XmlActuator) and actuator.command_field == "position"
|
||||
):
|
||||
if operation == "scale":
|
||||
if op.name == "scale":
|
||||
default_gainprm = env.sim.get_default_field("actuator_gainprm")
|
||||
default_biasprm = env.sim.get_default_field("actuator_biasprm")
|
||||
env.sim.model.actuator_gainprm[env_ids[:, None], ctrl_ids, 0] = (
|
||||
@@ -92,15 +141,69 @@ def pd_gains(
|
||||
env.sim.model.actuator_biasprm[env_ids[:, None], ctrl_ids, 2] = (
|
||||
default_biasprm[ctrl_ids, 2] * kd_samples
|
||||
)
|
||||
elif operation == "abs":
|
||||
else:
|
||||
assert op.name == "abs"
|
||||
env.sim.model.actuator_gainprm[env_ids[:, None], ctrl_ids, 0] = kp_samples
|
||||
env.sim.model.actuator_biasprm[env_ids[:, None], ctrl_ids, 1] = -kp_samples
|
||||
env.sim.model.actuator_biasprm[env_ids[:, None], ctrl_ids, 2] = -kd_samples
|
||||
|
||||
elif isinstance(actuator, BuiltinDcMotorActuator):
|
||||
if actuator.cfg.mode == DcMotorInputMode.VOLTAGE:
|
||||
raise ValueError(
|
||||
"dr.pd_gains does not apply to BuiltinDcMotorActuator in VOLTAGE "
|
||||
"mode (no internal PID gains to scale)."
|
||||
)
|
||||
# DC motor stores kp at gainprm[4] and kd at gainprm[6] (set via
|
||||
# set_to_dcmotor). The bias slots carry back-EMF / cogging, not the PD,
|
||||
# so we only touch gainprm.
|
||||
if op.name == "scale":
|
||||
default_gainprm = env.sim.get_default_field("actuator_gainprm")
|
||||
env.sim.model.actuator_gainprm[env_ids[:, None], ctrl_ids, 4] = (
|
||||
default_gainprm[ctrl_ids, 4] * kp_samples
|
||||
)
|
||||
env.sim.model.actuator_gainprm[env_ids[:, None], ctrl_ids, 6] = (
|
||||
default_gainprm[ctrl_ids, 6] * kd_samples
|
||||
)
|
||||
else:
|
||||
assert op.name == "abs"
|
||||
env.sim.model.actuator_gainprm[env_ids[:, None], ctrl_ids, 4] = kp_samples
|
||||
env.sim.model.actuator_gainprm[env_ids[:, None], ctrl_ids, 6] = kd_samples
|
||||
|
||||
elif isinstance(actuator, BuiltinPdActuator):
|
||||
# ctrl_ids is laid out as [pos_0..pos_{N-1}, vel_0..vel_{N-1}], so the
|
||||
# first N rows carry kp and the next N carry kd.
|
||||
n = actuator.num_targets
|
||||
pos_ids = ctrl_ids[:n]
|
||||
vel_ids = ctrl_ids[n:]
|
||||
if op.name == "scale":
|
||||
default_gainprm = env.sim.get_default_field("actuator_gainprm")
|
||||
default_biasprm = env.sim.get_default_field("actuator_biasprm")
|
||||
env.sim.model.actuator_gainprm[env_ids[:, None], pos_ids, 0] = (
|
||||
default_gainprm[pos_ids, 0] * kp_samples
|
||||
)
|
||||
env.sim.model.actuator_biasprm[env_ids[:, None], pos_ids, 1] = (
|
||||
default_biasprm[pos_ids, 1] * kp_samples
|
||||
)
|
||||
env.sim.model.actuator_gainprm[env_ids[:, None], vel_ids, 0] = (
|
||||
default_gainprm[vel_ids, 0] * kd_samples
|
||||
)
|
||||
env.sim.model.actuator_biasprm[env_ids[:, None], vel_ids, 2] = (
|
||||
default_biasprm[vel_ids, 2] * kd_samples
|
||||
)
|
||||
else:
|
||||
assert op.name == "abs"
|
||||
env.sim.model.actuator_gainprm[env_ids[:, None], pos_ids, 0] = kp_samples
|
||||
env.sim.model.actuator_biasprm[env_ids[:, None], pos_ids, 1] = -kp_samples
|
||||
env.sim.model.actuator_gainprm[env_ids[:, None], vel_ids, 0] = kd_samples
|
||||
env.sim.model.actuator_biasprm[env_ids[:, None], vel_ids, 2] = -kd_samples
|
||||
# biasprm[2] on the position half stays zero by construction. Writing
|
||||
# anything else here would inject damping into the position element on
|
||||
# top of the velocity element, silently double-counting kd.
|
||||
|
||||
elif isinstance(actuator, IdealPdActuator):
|
||||
assert actuator.stiffness is not None
|
||||
assert actuator.damping is not None
|
||||
if operation == "scale":
|
||||
if op.name == "scale":
|
||||
assert actuator.default_stiffness is not None
|
||||
assert actuator.default_damping is not None
|
||||
actuator.set_gains(
|
||||
@@ -108,25 +211,26 @@ def pd_gains(
|
||||
kp=actuator.default_stiffness[env_ids] * kp_samples,
|
||||
kd=actuator.default_damping[env_ids] * kd_samples,
|
||||
)
|
||||
elif operation == "abs":
|
||||
else:
|
||||
assert op.name == "abs"
|
||||
actuator.set_gains(env_ids, kp=kp_samples, kd=kd_samples)
|
||||
|
||||
else:
|
||||
raise TypeError(
|
||||
f"pd_gains only supports BuiltinPositionActuator, "
|
||||
f"XmlActuator (position), and IdealPdActuator, "
|
||||
f"got {type(actuator).__name__}"
|
||||
f"pd_gains only supports BuiltinPositionActuator, BuiltinPdActuator, "
|
||||
f"BuiltinDcMotorActuator (position/velocity mode), XmlActuator (position), "
|
||||
f"and IdealPdActuator, got {type(actuator).__name__}"
|
||||
)
|
||||
|
||||
|
||||
@requires_model_fields("actuator_forcerange")
|
||||
@requires_model_fields("actuator_forcerange", "jnt_actfrcrange", "tendon_actfrcrange")
|
||||
def effort_limits(
|
||||
env: ManagerBasedRlEnv,
|
||||
env_ids: torch.Tensor | None,
|
||||
effort_limit_range: tuple[float, float],
|
||||
asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
|
||||
distribution: Literal["uniform", "log_uniform"] = "uniform",
|
||||
operation: Literal["scale", "abs"] = "scale",
|
||||
operation: Operation | str = "scale",
|
||||
) -> None:
|
||||
"""Randomize actuator effort limits.
|
||||
|
||||
@@ -138,6 +242,11 @@ def effort_limits(
|
||||
distribution: Distribution type ("uniform" or "log_uniform").
|
||||
operation: "scale" multiplies default limits, "abs" sets absolute values.
|
||||
"""
|
||||
op = resolve_operation(operation)
|
||||
if op.name not in ("scale", "abs"):
|
||||
raise ValueError(
|
||||
f"effort_limits only supports 'scale' and 'abs' operations, got {op.name!r}"
|
||||
)
|
||||
asset: Entity = env.scene[asset_cfg.name]
|
||||
|
||||
if env_ids is None:
|
||||
@@ -145,36 +254,32 @@ def effort_limits(
|
||||
else:
|
||||
env_ids = env_ids.to(env.device, dtype=torch.int)
|
||||
|
||||
if isinstance(asset_cfg.actuator_ids, list):
|
||||
actuators = [asset.actuators[i] for i in asset_cfg.actuator_ids]
|
||||
else:
|
||||
actuators = asset.actuators[asset_cfg.actuator_ids]
|
||||
actuators = _resolve_actuators(asset, asset_cfg)
|
||||
|
||||
if not isinstance(actuators, list):
|
||||
actuators = [actuators]
|
||||
|
||||
for actuator in actuators:
|
||||
ctrl_ids = actuator.global_ctrl_ids
|
||||
num_actuators = len(ctrl_ids)
|
||||
# One effort sample per target. For single-element actuators this matches
|
||||
# ctrl_ids; for BuiltinPd the limit lives on the joint/tendon, so one
|
||||
# sample per target is sufficient regardless of the two-element ctrl.
|
||||
n_samples = (
|
||||
actuator.num_targets if isinstance(actuator, BuiltinPdActuator) else len(ctrl_ids)
|
||||
)
|
||||
|
||||
dist = resolve_distribution(distribution)
|
||||
effort_samples = dist.sample(
|
||||
torch.tensor(effort_limit_range[0], device=env.device),
|
||||
torch.tensor(effort_limit_range[1], device=env.device),
|
||||
(len(env_ids), num_actuators),
|
||||
(len(env_ids), n_samples),
|
||||
env.device,
|
||||
)
|
||||
|
||||
if isinstance(
|
||||
actuator,
|
||||
(
|
||||
BuiltinPositionActuator,
|
||||
BuiltinVelocityActuator,
|
||||
BuiltinMotorActuator,
|
||||
XmlActuator,
|
||||
),
|
||||
if isinstance(actuator, (BuiltinPositionActuator, BuiltinDcMotorActuator)) or (
|
||||
isinstance(actuator, XmlActuator) and actuator.command_field == "position"
|
||||
):
|
||||
if operation == "scale":
|
||||
if op.name == "scale":
|
||||
default_forcerange = env.sim.get_default_field("actuator_forcerange")
|
||||
env.sim.model.actuator_forcerange[env_ids[:, None], ctrl_ids, 0] = (
|
||||
default_forcerange[ctrl_ids, 0] * effort_samples
|
||||
@@ -182,7 +287,8 @@ def effort_limits(
|
||||
env.sim.model.actuator_forcerange[env_ids[:, None], ctrl_ids, 1] = (
|
||||
default_forcerange[ctrl_ids, 1] * effort_samples
|
||||
)
|
||||
elif operation == "abs":
|
||||
else:
|
||||
assert op.name == "abs"
|
||||
env.sim.model.actuator_forcerange[
|
||||
env_ids[:, None], ctrl_ids, 0
|
||||
] = -effort_samples
|
||||
@@ -192,18 +298,42 @@ def effort_limits(
|
||||
|
||||
elif isinstance(actuator, IdealPdActuator):
|
||||
assert actuator.force_limit is not None
|
||||
if operation == "scale":
|
||||
if op.name == "scale":
|
||||
assert actuator.default_force_limit is not None
|
||||
actuator.set_effort_limit(
|
||||
env_ids,
|
||||
effort_limit=actuator.default_force_limit[env_ids] * effort_samples,
|
||||
)
|
||||
elif operation == "abs":
|
||||
else:
|
||||
assert op.name == "abs"
|
||||
actuator.set_effort_limit(env_ids, effort_limit=effort_samples)
|
||||
|
||||
elif isinstance(actuator, BuiltinPdActuator):
|
||||
# BuiltinPd's effort_limit lives on the joint/tendon as a sum-clamp
|
||||
# (jnt_actfrcrange / tendon_actfrcrange), not on per-element forcerange.
|
||||
if actuator.transmission_type == TransmissionType.JOINT:
|
||||
field = "jnt_actfrcrange"
|
||||
target_global_ids = asset.indexing.joint_ids[actuator.target_ids]
|
||||
else:
|
||||
field = "tendon_actfrcrange"
|
||||
target_global_ids = asset.indexing.tendon_ids[actuator.target_ids]
|
||||
arr = getattr(env.sim.model, field)
|
||||
if op.name == "scale":
|
||||
default = env.sim.get_default_field(field)
|
||||
arr[env_ids[:, None], target_global_ids, 0] = (
|
||||
default[target_global_ids, 0] * effort_samples
|
||||
)
|
||||
arr[env_ids[:, None], target_global_ids, 1] = (
|
||||
default[target_global_ids, 1] * effort_samples
|
||||
)
|
||||
else:
|
||||
assert op.name == "abs"
|
||||
arr[env_ids[:, None], target_global_ids, 0] = -effort_samples
|
||||
arr[env_ids[:, None], target_global_ids, 1] = effort_samples
|
||||
|
||||
else:
|
||||
raise TypeError(
|
||||
f"effort_limits only supports BuiltinPositionActuator, BuiltinVelocityActuator, "
|
||||
f"BuiltinMotorActuator, XmlActuator, and IdealPdActuator, "
|
||||
f"effort_limits only supports BuiltinPositionActuator, BuiltinPdActuator, "
|
||||
f"BuiltinDcMotorActuator, XmlActuator (position), and IdealPdActuator, "
|
||||
f"got {type(actuator).__name__}"
|
||||
)
|
||||
|
||||
@@ -21,6 +21,38 @@ if TYPE_CHECKING:
|
||||
from mjlab.viewer.debug_visualizer import DebugVisualizer
|
||||
|
||||
_DEFAULT_ASSET_CFG = SceneEntityCfg("robot")
|
||||
_SE3_KEYS = ("x", "y", "z", "roll", "pitch", "yaw")
|
||||
|
||||
|
||||
def _sample_se3_range(
|
||||
range_dict: dict[str, tuple[float, float]] | None,
|
||||
shape: tuple[int, ...],
|
||||
device: str,
|
||||
) -> torch.Tensor:
|
||||
"""Sample uniform ``[x, y, z, roll, pitch, yaw]`` offsets.
|
||||
|
||||
``range_dict`` maps any subset of those keys to ``(min, max)`` ranges; missing
|
||||
keys default to ``(0.0, 0.0)`` (no offset). ``None`` is treated as empty. The
|
||||
returned tensor has the requested ``shape`` whose last dimension must be 6.
|
||||
"""
|
||||
range_dict = range_dict or {}
|
||||
range_list = [range_dict.get(key, (0.0, 0.0)) for key in _SE3_KEYS]
|
||||
ranges = torch.tensor(range_list, device=device)
|
||||
return sample_uniform(ranges[:, 0], ranges[:, 1], shape, device=device)
|
||||
|
||||
|
||||
def resolve_env_ids(
|
||||
env: ManagerBasedRlEnv, env_ids: torch.Tensor | None
|
||||
) -> torch.Tensor:
|
||||
"""Return ``env_ids`` unchanged, or all environment indices if ``None``.
|
||||
|
||||
Event functions receive ``env_ids=None`` to mean "all environments" (a full
|
||||
reset, or a global-time interval term). This normalizes that sentinel to a
|
||||
concrete index tensor so the function body can assume a real ``torch.Tensor``.
|
||||
"""
|
||||
if env_ids is None:
|
||||
return torch.arange(env.num_envs, device=env.device, dtype=torch.int)
|
||||
return env_ids
|
||||
|
||||
|
||||
def randomize_terrain(env: ManagerBasedRlEnv, env_ids: torch.Tensor | None) -> None:
|
||||
@@ -29,8 +61,7 @@ def randomize_terrain(env: ManagerBasedRlEnv, env_ids: torch.Tensor | None) -> N
|
||||
This picks a random terrain type (column) and difficulty level (row) for each
|
||||
environment. Useful for play/evaluation mode to test on varied terrains.
|
||||
"""
|
||||
if env_ids is None:
|
||||
env_ids = torch.arange(env.num_envs, device=env.device, dtype=torch.int)
|
||||
env_ids = resolve_env_ids(env, env_ids)
|
||||
|
||||
terrain = env.scene.terrain
|
||||
if terrain is not None:
|
||||
@@ -48,8 +79,7 @@ def reset_scene_to_default(
|
||||
|
||||
Automatically applies env_origins offset to position all entities correctly.
|
||||
"""
|
||||
if env_ids is None:
|
||||
env_ids = torch.arange(env.num_envs, device=env.device, dtype=torch.int)
|
||||
env_ids = resolve_env_ids(env, env_ids)
|
||||
|
||||
for entity in env.scene.entities.values():
|
||||
if not isinstance(entity, Entity):
|
||||
@@ -105,19 +135,12 @@ def reset_root_state_uniform(
|
||||
velocity_range: Velocity range (only used for floating-base entities).
|
||||
asset_cfg: Asset configuration.
|
||||
"""
|
||||
if env_ids is None:
|
||||
env_ids = torch.arange(env.num_envs, device=env.device, dtype=torch.int)
|
||||
env_ids = resolve_env_ids(env, env_ids)
|
||||
|
||||
asset: Entity = env.scene[asset_cfg.name]
|
||||
|
||||
# Pose.
|
||||
range_list = [
|
||||
pose_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]
|
||||
]
|
||||
ranges = torch.tensor(range_list, device=env.device)
|
||||
pose_samples = sample_uniform(
|
||||
ranges[:, 0], ranges[:, 1], (len(env_ids), 6), device=env.device
|
||||
)
|
||||
pose_samples = _sample_se3_range(pose_range, (len(env_ids), 6), env.device)
|
||||
|
||||
# Fixed-based entities with mocap=True.
|
||||
if asset.is_fixed_base:
|
||||
@@ -157,16 +180,7 @@ def reset_root_state_uniform(
|
||||
orientations = quat_mul(root_states[:, 3:7], orientations_delta)
|
||||
|
||||
# Velocities.
|
||||
if velocity_range is None:
|
||||
velocity_range = {}
|
||||
range_list = [
|
||||
velocity_range.get(key, (0.0, 0.0))
|
||||
for key in ["x", "y", "z", "roll", "pitch", "yaw"]
|
||||
]
|
||||
ranges = torch.tensor(range_list, device=env.device)
|
||||
vel_samples = sample_uniform(
|
||||
ranges[:, 0], ranges[:, 1], (len(env_ids), 6), device=env.device
|
||||
)
|
||||
vel_samples = _sample_se3_range(velocity_range, (len(env_ids), 6), env.device)
|
||||
velocities = root_states[:, 7:13] + vel_samples
|
||||
|
||||
asset.write_root_link_pose_to_sim(
|
||||
@@ -199,8 +213,7 @@ def reset_root_state_from_flat_patches(
|
||||
velocity_range: Optional velocity range (floating-base only).
|
||||
asset_cfg: Asset configuration.
|
||||
"""
|
||||
if env_ids is None:
|
||||
env_ids = torch.arange(env.num_envs, device=env.device, dtype=torch.int)
|
||||
env_ids = resolve_env_ids(env, env_ids)
|
||||
|
||||
terrain = env.scene.terrain
|
||||
if terrain is None or patch_name not in terrain.flat_patches:
|
||||
@@ -230,15 +243,7 @@ def reset_root_state_from_flat_patches(
|
||||
root_states = default_root_state[env_ids].clone()
|
||||
|
||||
# Apply optional pose range offset.
|
||||
if pose_range is None:
|
||||
pose_range = {}
|
||||
range_list = [
|
||||
pose_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]
|
||||
]
|
||||
ranges = torch.tensor(range_list, device=env.device)
|
||||
pose_samples = sample_uniform(
|
||||
ranges[:, 0], ranges[:, 1], (len(env_ids), 6), device=env.device
|
||||
)
|
||||
pose_samples = _sample_se3_range(pose_range, (len(env_ids), 6), env.device)
|
||||
|
||||
# Position: flat patch position + optional offset. Use patch z instead of default.
|
||||
final_positions = positions.clone()
|
||||
@@ -262,16 +267,7 @@ def reset_root_state_from_flat_patches(
|
||||
return
|
||||
|
||||
# Velocities.
|
||||
if velocity_range is None:
|
||||
velocity_range = {}
|
||||
vel_range_list = [
|
||||
velocity_range.get(key, (0.0, 0.0))
|
||||
for key in ["x", "y", "z", "roll", "pitch", "yaw"]
|
||||
]
|
||||
vel_ranges = torch.tensor(vel_range_list, device=env.device)
|
||||
vel_samples = sample_uniform(
|
||||
vel_ranges[:, 0], vel_ranges[:, 1], (len(env_ids), 6), device=env.device
|
||||
)
|
||||
vel_samples = _sample_se3_range(velocity_range, (len(env_ids), 6), env.device)
|
||||
velocities = root_states[:, 7:13] + vel_samples
|
||||
|
||||
asset.write_root_link_pose_to_sim(
|
||||
@@ -287,8 +283,7 @@ def reset_joints_by_offset(
|
||||
velocity_range: tuple[float, float],
|
||||
asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
|
||||
) -> None:
|
||||
if env_ids is None:
|
||||
env_ids = torch.arange(env.num_envs, device=env.device, dtype=torch.int)
|
||||
env_ids = resolve_env_ids(env, env_ids)
|
||||
|
||||
asset: Entity = env.scene[asset_cfg.name]
|
||||
default_joint_pos = asset.data.default_joint_pos
|
||||
@@ -320,28 +315,60 @@ def reset_joints_by_offset(
|
||||
|
||||
def push_by_setting_velocity(
|
||||
env: ManagerBasedRlEnv,
|
||||
env_ids: torch.Tensor,
|
||||
env_ids: torch.Tensor | None,
|
||||
velocity_range: dict[str, tuple[float, float]],
|
||||
asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
|
||||
) -> None:
|
||||
"""Push an entity by overwriting its root velocity with a sampled offset.
|
||||
|
||||
This is an *instantaneous, mass-independent* kick: it adds a uniformly sampled
|
||||
delta directly to the root velocity, ignoring inertia and contact dynamics. It
|
||||
is the cheapest disturbance and the standard locomotion "push the robot" term.
|
||||
Use with ``mode="interval"``.
|
||||
|
||||
For force-based disturbances that respect the entity's dynamics, see
|
||||
:func:`apply_external_force_torque` (a constant wrench you manage yourself) or
|
||||
:class:`apply_body_impulse` (transient, self-managing impulses).
|
||||
"""
|
||||
env_ids = resolve_env_ids(env, env_ids)
|
||||
asset: Entity = env.scene[asset_cfg.name]
|
||||
vel_w = asset.data.root_link_vel_w[env_ids]
|
||||
range_list = [
|
||||
velocity_range.get(key, (0.0, 0.0))
|
||||
for key in ["x", "y", "z", "roll", "pitch", "yaw"]
|
||||
]
|
||||
ranges = torch.tensor(range_list, device=env.device)
|
||||
vel_w += sample_uniform(ranges[:, 0], ranges[:, 1], vel_w.shape, device=env.device)
|
||||
vel_w += _sample_se3_range(velocity_range, vel_w.shape, env.device)
|
||||
asset.write_root_link_velocity_to_sim(vel_w, env_ids=env_ids)
|
||||
|
||||
|
||||
def apply_external_force_torque(
|
||||
env: ManagerBasedRlEnv,
|
||||
env_ids: torch.Tensor,
|
||||
env_ids: torch.Tensor | None,
|
||||
force_range: tuple[float, float],
|
||||
torque_range: tuple[float, float],
|
||||
asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
|
||||
) -> None:
|
||||
"""Apply a single *constant* external wrench to bodies.
|
||||
|
||||
Samples a force and torque once and writes them to ``xfrc_applied``. The wrench
|
||||
is **stateless and never expires**: MuJoCo holds it constant on every physics
|
||||
step until something overwrites or zeroes it. There is no duration, cooldown,
|
||||
or auto-clear.
|
||||
|
||||
**When to use this vs.** :class:`apply_body_impulse`:
|
||||
|
||||
- Use ``apply_external_force_torque`` for a *steady, episode-long* disturbance
|
||||
such as a fixed payload, a constant wind, or a sustained load. The intended
|
||||
pattern is ``mode="reset"``: re-randomize the wrench each episode so it holds
|
||||
for that episode's duration. Because it never turns itself off, **you are
|
||||
responsible for clearing or overwriting it** (e.g. via the next reset). It is
|
||||
*not* suited to transient bumps on its own.
|
||||
|
||||
- Use :class:`apply_body_impulse` for *transient, repeated, randomized*
|
||||
disturbances during an episode (bumps, gusts, collisions). It runs a full
|
||||
cooldown -> trigger -> sustain -> expire lifecycle per environment, zeroing
|
||||
the wrench automatically when each impulse ends, and ticks on ``mode="step"``.
|
||||
|
||||
For an instantaneous, mass-independent kick instead of a force, see
|
||||
:func:`push_by_setting_velocity`.
|
||||
"""
|
||||
env_ids = resolve_env_ids(env, env_ids)
|
||||
asset: Entity = env.scene[asset_cfg.name]
|
||||
num_bodies = (
|
||||
len(asset_cfg.body_ids)
|
||||
@@ -385,6 +412,10 @@ class apply_body_impulse:
|
||||
applied.
|
||||
|
||||
Use with ``mode="step"``.
|
||||
|
||||
For a *constant* episode-long wrench instead of transient impulses, see
|
||||
:func:`apply_external_force_torque`. For an instantaneous, mass-independent
|
||||
velocity kick, see :func:`push_by_setting_velocity`.
|
||||
"""
|
||||
|
||||
@dataclass
|
||||
@@ -422,9 +453,16 @@ class apply_body_impulse:
|
||||
else self._asset.num_bodies
|
||||
)
|
||||
|
||||
self._cooldown_s: tuple[float, float] = cfg.params["cooldown_s"]
|
||||
self._time_remaining = torch.zeros(self._num_envs, device=self._device)
|
||||
self._interval_time_left = torch.zeros(self._num_envs, device=self._device)
|
||||
self._active = torch.zeros(self._num_envs, device=self._device, dtype=torch.bool)
|
||||
# Pre-sample the initial cooldown so the first impulse is preceded by a cooldown
|
||||
# rather than firing immediately at t=0.
|
||||
self._interval_time_left = self._sample_cooldown(self._num_envs)
|
||||
|
||||
def _sample_cooldown(self, n: int) -> torch.Tensor:
|
||||
low, high = self._cooldown_s
|
||||
return sample_uniform(low, high, n, self._device)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
@@ -446,13 +484,14 @@ class apply_body_impulse:
|
||||
torque_range: ``(min, max)`` uniform range for each torque component (Nm).
|
||||
duration_s: ``(min, max)`` uniform range for impulse duration in seconds.
|
||||
cooldown_s: ``(min, max)`` uniform range for the cooldown between consecutive
|
||||
impulses in seconds.
|
||||
impulses in seconds. Captured at init so the first impulse can be
|
||||
preceded by a sampled cooldown; the kwarg passed here is unused.
|
||||
asset_cfg: Entity and body selection. ``body_ids`` on the config selects which
|
||||
bodies receive forces.
|
||||
body_point_offset: Optional ``(x, y, z)`` offset in the body frame where the
|
||||
force is applied. Generates additional torque via ``cross(offset, force)``.
|
||||
"""
|
||||
del env, env_ids, asset_cfg # Unused.
|
||||
del env, env_ids, asset_cfg, cooldown_s # Unused at call time.
|
||||
dt = self._step_dt
|
||||
|
||||
# Decrement timers for active envs.
|
||||
@@ -468,11 +507,7 @@ class apply_body_impulse:
|
||||
)
|
||||
self._active[expired_ids] = False
|
||||
self._time_remaining[expired_ids] = 0.0
|
||||
int_low, int_high = cooldown_s
|
||||
self._interval_time_left[expired_ids] = (
|
||||
torch.rand(len(expired_ids), device=self._device) * (int_high - int_low)
|
||||
+ int_low
|
||||
)
|
||||
self._interval_time_left[expired_ids] = self._sample_cooldown(len(expired_ids))
|
||||
|
||||
# Decrement interval timers.
|
||||
self._interval_time_left -= dt
|
||||
@@ -514,10 +549,7 @@ class apply_body_impulse:
|
||||
self._active[trigger_ids] = True
|
||||
|
||||
# Resample interval timers.
|
||||
int_low, int_high = cooldown_s
|
||||
self._interval_time_left[trigger_ids] = (
|
||||
torch.rand(n, device=self._device) * (int_high - int_low) + int_low
|
||||
)
|
||||
self._interval_time_left[trigger_ids] = self._sample_cooldown(n)
|
||||
|
||||
def debug_vis(self, visualizer: DebugVisualizer) -> None:
|
||||
"""Draw arrows for active impulse forces."""
|
||||
@@ -553,13 +585,7 @@ class apply_body_impulse:
|
||||
if env_ids is None:
|
||||
env_ids = slice(None)
|
||||
|
||||
# Clear forces for reset envs.
|
||||
if isinstance(env_ids, slice):
|
||||
reset_ids = env_ids
|
||||
else:
|
||||
reset_ids = env_ids
|
||||
|
||||
if self._active[reset_ids].any():
|
||||
if self._active[env_ids].any():
|
||||
if isinstance(env_ids, slice):
|
||||
active_ids = self._active.nonzero(as_tuple=False).squeeze(-1)
|
||||
else:
|
||||
@@ -573,6 +599,7 @@ class apply_body_impulse:
|
||||
zeros, zeros, env_ids=active_ids, body_ids=self._body_ids
|
||||
)
|
||||
|
||||
self._time_remaining[reset_ids] = 0.0
|
||||
self._interval_time_left[reset_ids] = 0.0
|
||||
self._active[reset_ids] = False
|
||||
n = self._num_envs if isinstance(env_ids, slice) else len(env_ids)
|
||||
self._time_remaining[env_ids] = 0.0
|
||||
self._interval_time_left[env_ids] = self._sample_cooldown(n)
|
||||
self._active[env_ids] = False
|
||||
|
||||
@@ -106,6 +106,23 @@ def builtin_sensor(env: ManagerBasedRlEnv, sensor_name: str) -> torch.Tensor:
|
||||
return sensor.data
|
||||
|
||||
|
||||
def projected_gravity_from_sensor(
|
||||
env: ManagerBasedRlEnv, sensor_name: str
|
||||
) -> torch.Tensor:
|
||||
"""Projected gravity from a ``framezaxis`` up-vector sensor.
|
||||
|
||||
The sensor is expected to output the world Z-axis expressed in the sensor's frame
|
||||
(e.g. ``framezaxis`` with ``objtype=body objname=world`` and ``reftype=site``). That
|
||||
is the body-frame "up" vector, so it is negated to point along gravity.
|
||||
|
||||
Unlike :func:`projected_gravity`, which uses the root body orientation, this reads
|
||||
the sensor's site frame and therefore reflects IMU site pose randomization.
|
||||
"""
|
||||
sensor = env.scene[sensor_name]
|
||||
assert isinstance(sensor, BuiltinSensor)
|
||||
return -sensor.data
|
||||
|
||||
|
||||
def height_scan(
|
||||
env: ManagerBasedRlEnv,
|
||||
sensor_name: str,
|
||||
|
||||
@@ -290,9 +290,12 @@ class EventManager(ManagerBase):
|
||||
fired = True
|
||||
elif mode == "reset":
|
||||
assert global_env_step_count is not None
|
||||
# Reset events require concrete indices: callers (e.g. ManagerBasedRlEnv)
|
||||
# resolve None to all environments upstream. Enforce that here so a future
|
||||
# caller passing None fails loudly instead of leaking a slice into event
|
||||
# functions, which only understand None or a tensor.
|
||||
assert env_ids is not None, "reset events require concrete env_ids, got None"
|
||||
min_step_count = term_cfg.min_step_count_between_reset
|
||||
if env_ids is None:
|
||||
env_ids = slice(None)
|
||||
if min_step_count == 0:
|
||||
self._reset_term_last_triggered_step_id[index][env_ids] = (
|
||||
global_env_step_count
|
||||
|
||||
@@ -11,7 +11,7 @@ import numpy as np
|
||||
import torch
|
||||
|
||||
from mjlab.entity import Entity, EntityCfg
|
||||
from mjlab.entity.entity import VariantMetadata
|
||||
from mjlab.entity.variants import VariantMetadata
|
||||
from mjlab.sensor import BuiltinSensor, RayCastSensor, Sensor, SensorCfg
|
||||
from mjlab.sensor.camera_sensor import CameraSensor
|
||||
from mjlab.sensor.sensor_context import SensorContext
|
||||
@@ -59,7 +59,7 @@ class Scene:
|
||||
self._default_env_origins: torch.Tensor | None = None
|
||||
self._sensor_context: SensorContext | None = None
|
||||
|
||||
self._spec = mujoco.MjSpec.from_string(_SCENE_XML.read_text())
|
||||
self._spec = mujoco.MjSpec.from_file(str(_SCENE_XML))
|
||||
if self._cfg.extent is not None:
|
||||
self._spec.stat.extent = self._cfg.extent
|
||||
self._add_terrain()
|
||||
@@ -132,11 +132,6 @@ class Scene:
|
||||
def device(self) -> str:
|
||||
return self._device
|
||||
|
||||
@property
|
||||
def has_mesh_variants(self) -> bool:
|
||||
"""True if any entity declares per-world mesh variants."""
|
||||
return any(ent.variant_metadata is not None for ent in self._entities.values())
|
||||
|
||||
def collect_variant_info(
|
||||
self,
|
||||
) -> list[tuple[str, VariantMetadata]]:
|
||||
|
||||
@@ -59,7 +59,7 @@ def run_train(task_id: str, cfg: TrainConfig, log_dir: Path) -> None:
|
||||
os.environ["MUJOCO_EGL_DEVICE_ID"] = str(local_rank)
|
||||
device = f"cuda:{local_rank}"
|
||||
# Set seed to have diversity in different processes.
|
||||
seed = cfg.agent.seed + local_rank
|
||||
seed = cfg.agent.seed + rank
|
||||
|
||||
configure_torch_backends()
|
||||
|
||||
@@ -197,7 +197,6 @@ def launch_training(task_id: str, args: TrainConfig | None = None):
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = ""
|
||||
else:
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(map(str, selected_gpus))
|
||||
os.environ["MUJOCO_GL"] = "egl"
|
||||
|
||||
if num_gpus <= 1:
|
||||
# CPU or single GPU: run directly without torchrunx.
|
||||
|
||||
@@ -428,7 +428,7 @@ class ContactSensor(Sensor[ContactData]):
|
||||
normal = data.normal
|
||||
tangent = data.tangent
|
||||
tangent2 = torch.cross(normal, tangent, dim=-1)
|
||||
R = torch.stack([tangent, tangent2, normal], dim=-1)
|
||||
R = torch.stack([normal, tangent, tangent2], dim=-1)
|
||||
|
||||
has_contact = torch.norm(normal, dim=-1, keepdim=True) > 1e-8
|
||||
|
||||
|
||||
@@ -435,7 +435,7 @@ class RayCastSensor(Sensor[RayCastData]):
|
||||
self._model: mjwarp.Model | None = None
|
||||
self._mj_model: mujoco.MjModel | None = None
|
||||
self._device: str | None = None
|
||||
self._wp_device: wp.context.Device | None = None
|
||||
self._wp_device: wp.Device | None = None
|
||||
|
||||
# Per-frame info: list of (frame_type, obj_id, body_id).
|
||||
self._frame_infos: list[tuple[Literal["body", "site", "geom"], int, int]] = []
|
||||
|
||||
@@ -1,335 +0,0 @@
|
||||
"""Per-world mesh variant support.
|
||||
|
||||
Sibling of :mod:`mjlab.sim.randomization`: that module expands singleton
|
||||
model fields into per-world arrays for DR; this one writes per-world
|
||||
arrays whose rows differ by mesh variant.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Callable
|
||||
|
||||
import mujoco
|
||||
import mujoco_warp as mjwarp
|
||||
import numpy as np
|
||||
import warp as wp
|
||||
|
||||
from mjlab.entity.entity import BodyInertialMetadata, VariantMetadata
|
||||
|
||||
# Fields that depend on mesh geometry and must be compiled per-variant.
|
||||
VARIANT_DEPENDENT_FIELDS = (
|
||||
"geom_size",
|
||||
"geom_rbound",
|
||||
"geom_aabb",
|
||||
"geom_pos",
|
||||
"geom_quat",
|
||||
"body_mass",
|
||||
"body_subtreemass",
|
||||
"body_inertia",
|
||||
"body_invweight0",
|
||||
"body_ipos",
|
||||
"body_iquat",
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class MeshVariantResult:
|
||||
"""Output of :func:`build_mesh_variant_model`."""
|
||||
|
||||
wp_model: mjwarp.Model
|
||||
mj_model: mujoco.MjModel
|
||||
# Maps entity prefix -> array of variant indices per world.
|
||||
world_to_variant: dict[str, np.ndarray]
|
||||
|
||||
|
||||
def _find_entity_mesh_geom_ids(
|
||||
model: mujoco.MjModel,
|
||||
entity_prefix: str,
|
||||
) -> list[int]:
|
||||
"""Find all mesh geom IDs belonging to an entity, including padding."""
|
||||
named_ids: list[int] = []
|
||||
for gid in range(model.ngeom):
|
||||
gname = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_GEOM, gid)
|
||||
if (
|
||||
gname
|
||||
and gname.startswith(entity_prefix)
|
||||
and model.geom_type[gid] == mujoco.mjtGeom.mjGEOM_MESH
|
||||
):
|
||||
named_ids.append(gid)
|
||||
if not named_ids:
|
||||
return []
|
||||
# Include unnamed padding geoms on the same body.
|
||||
body_id = model.geom_bodyid[named_ids[0]]
|
||||
all_ids = set(named_ids)
|
||||
for gid in range(model.ngeom):
|
||||
if (
|
||||
model.geom_bodyid[gid] == body_id
|
||||
and model.geom_type[gid] == mujoco.mjtGeom.mjGEOM_MESH
|
||||
):
|
||||
all_ids.add(gid)
|
||||
return sorted(all_ids)
|
||||
|
||||
|
||||
def allocate_worlds(
|
||||
weights: tuple[float, ...],
|
||||
nworld: int,
|
||||
) -> list[int]:
|
||||
"""Assign worlds proportionally by weight (largest-remainder method).
|
||||
|
||||
Returns a list of length *nworld* containing variant indices. Weights
|
||||
must be non-negative with at least one positive entry.
|
||||
"""
|
||||
if any(w < 0 for w in weights):
|
||||
raise ValueError(f"weights must be non-negative, got {weights}.")
|
||||
total = sum(weights)
|
||||
if total <= 0:
|
||||
raise ValueError(f"weights must have a positive sum, got {weights}.")
|
||||
quotas = [(w / total) * nworld for w in weights]
|
||||
floors = [int(q) for q in quotas]
|
||||
remainders = sorted(
|
||||
((quotas[i] - floors[i], i) for i in range(len(weights))),
|
||||
key=lambda x: -x[0],
|
||||
)
|
||||
allocated = sum(floors)
|
||||
for j in range(nworld - allocated):
|
||||
floors[remainders[j][1]] += 1
|
||||
assignment: list[int] = []
|
||||
for idx, count in enumerate(floors):
|
||||
assignment.extend([idx] * count)
|
||||
return assignment
|
||||
|
||||
|
||||
def build_mesh_variant_model(
|
||||
spec: mujoco.MjSpec,
|
||||
nworld: int,
|
||||
variant_info: list[tuple[str, VariantMetadata]],
|
||||
configure_model: Callable[[mujoco.MjModel], None] | None = None,
|
||||
) -> MeshVariantResult:
|
||||
"""Build a warp Model with per-world mesh assignments.
|
||||
|
||||
Args:
|
||||
spec: Scene spec (already merged with padded variant geoms).
|
||||
nworld: Number of simulation worlds.
|
||||
variant_info: List of ``(entity_prefix, metadata)`` pairs for
|
||||
entities that have mesh variants.
|
||||
configure_model: Optional callback to configure the compiled
|
||||
MjModel before ``put_model`` (e.g., setting solver options).
|
||||
|
||||
Returns:
|
||||
A :class:`MeshVariantResult` containing the warp model, host
|
||||
model, and per-entity world-to-variant mappings.
|
||||
"""
|
||||
spec = spec.copy()
|
||||
model = spec.compile()
|
||||
if configure_model is not None:
|
||||
configure_model(model)
|
||||
|
||||
# Start from base dataid tiled for all worlds.
|
||||
base_dataid = model.geom_dataid.copy()
|
||||
dataid_table = np.tile(base_dataid, (nworld, 1))
|
||||
|
||||
world_to_variant: dict[str, np.ndarray] = {}
|
||||
|
||||
for entity_prefix, metadata in variant_info:
|
||||
# Allocate worlds by weight.
|
||||
assignment = allocate_worlds(metadata.variant_weights, nworld)
|
||||
w2v = np.array(assignment, dtype=np.int32)
|
||||
world_to_variant[entity_prefix] = w2v
|
||||
|
||||
mesh_geom_ids = _find_entity_mesh_geom_ids(model, entity_prefix)
|
||||
nslots = len(mesh_geom_ids)
|
||||
|
||||
# Resolve every (variant, slot) -> mesh_id once. Mesh names in the merged
|
||||
# spec are variant-prefixed ("mug/visual_mesh"); after attaching to the
|
||||
# scene they also carry the entity prefix ("object/mug/visual_mesh").
|
||||
# Padding slots are -1.
|
||||
nvariants = len(metadata.variant_mesh_names)
|
||||
variant_slot_ids = np.full((nvariants, nslots), -1, dtype=np.int64)
|
||||
for v_idx, mesh_names in enumerate(metadata.variant_mesh_names):
|
||||
for slot in range(min(nslots, len(mesh_names))):
|
||||
name = mesh_names[slot]
|
||||
if name is None:
|
||||
continue
|
||||
full = f"{entity_prefix}{name}"
|
||||
mid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_MESH, full)
|
||||
if mid < 0:
|
||||
variant_label = metadata.variant_names[v_idx]
|
||||
raise ValueError(
|
||||
f"Mesh '{full}' (variant '{variant_label}', slot {slot}) "
|
||||
f"not found in compiled model."
|
||||
)
|
||||
variant_slot_ids[v_idx, slot] = mid
|
||||
|
||||
# Vectorized scatter: row-select by variant assignment, write into the
|
||||
# mesh-geom columns of the per-world dataid table.
|
||||
dataid_table[:, mesh_geom_ids] = variant_slot_ids[w2v]
|
||||
|
||||
# Build warp model.
|
||||
m = mjwarp.put_model(model)
|
||||
m.geom_dataid = wp.array(dataid_table, dtype=int)
|
||||
|
||||
# Populate dependent per-world fields.
|
||||
_populate_dependent_fields(
|
||||
m, spec, model, dataid_table, nworld, variant_info, world_to_variant
|
||||
)
|
||||
|
||||
return MeshVariantResult(
|
||||
wp_model=m,
|
||||
mj_model=model,
|
||||
world_to_variant=world_to_variant,
|
||||
)
|
||||
|
||||
|
||||
def _populate_dependent_fields(
|
||||
m: mjwarp.Model,
|
||||
spec: mujoco.MjSpec,
|
||||
padded_model: mujoco.MjModel,
|
||||
dataid_table: np.ndarray,
|
||||
nworld: int,
|
||||
variant_info: list[tuple[str, VariantMetadata]],
|
||||
world_to_variant: dict[str, np.ndarray],
|
||||
) -> None:
|
||||
"""Compile each unique variant and write per-world dependent fields.
|
||||
|
||||
Each unique variant is compiled from a fresh ``spec.copy()``; the input
|
||||
``spec`` is not mutated.
|
||||
"""
|
||||
# Find unique dataid rows.
|
||||
unique_rows: dict[tuple[int, ...], int] = {}
|
||||
for w in range(nworld):
|
||||
key = tuple(dataid_table[w])
|
||||
if key not in unique_rows:
|
||||
unique_rows[key] = w
|
||||
|
||||
# Map padded_model geom IDs to geom names (stable across spec copies).
|
||||
geom_id_to_name: dict[int, str] = {}
|
||||
for g in spec.geoms:
|
||||
if not g.name:
|
||||
continue
|
||||
gid = mujoco.mj_name2id(padded_model, mujoco.mjtObj.mjOBJ_GEOM, g.name)
|
||||
if gid >= 0:
|
||||
geom_id_to_name[gid] = g.name
|
||||
|
||||
# Collect all variant geom IDs in padded_model.
|
||||
all_variant_geom_ids: set[int] = set()
|
||||
for entity_prefix, _ in variant_info:
|
||||
all_variant_geom_ids.update(_find_entity_mesh_geom_ids(padded_model, entity_prefix))
|
||||
|
||||
# Bodies any variant marks as explicit-inertial: must be reset on the
|
||||
# fresh spec copy before applying this variant's inertials. Variants
|
||||
# without an explicit inertial fall back to MuJoCo's mesh-derived path
|
||||
# during compile, so we clear the diagonal inertial fields. Do NOT
|
||||
# assign ``body.fullinertia``: any assignment (even zeros) flags the
|
||||
# field as user-specified and ``spec.compile()`` then rejects it as
|
||||
# conflicting with ``body.inertia``.
|
||||
variant_inertial_body_names: set[str] = set()
|
||||
for entity_prefix, metadata in variant_info:
|
||||
for variant_inertials in metadata.variant_body_inertials:
|
||||
for inertial in variant_inertials:
|
||||
variant_inertial_body_names.add(f"{entity_prefix}{inertial.body_name}")
|
||||
|
||||
# Compile each unique variant from a fresh spec copy.
|
||||
compiled_variants: dict[tuple[int, ...], mujoco.MjModel] = {}
|
||||
for key, first_world in unique_rows.items():
|
||||
variant_spec = spec.copy()
|
||||
geoms_by_name = {g.name: g for g in variant_spec.geoms if g.name}
|
||||
bodies_by_name = {b.name: b for b in variant_spec.bodies if b.name}
|
||||
|
||||
# Apply this variant's mesh selection per geom slot.
|
||||
for gid in all_variant_geom_ids:
|
||||
name = geom_id_to_name.get(gid)
|
||||
if name is None:
|
||||
continue
|
||||
geom = geoms_by_name[name]
|
||||
mesh_id = int(dataid_table[first_world, gid])
|
||||
if mesh_id >= 0:
|
||||
mesh_name = mujoco.mj_id2name(padded_model, mujoco.mjtObj.mjOBJ_MESH, mesh_id)
|
||||
geom.meshname = mesh_name
|
||||
geom.contype = 1
|
||||
geom.conaffinity = 1
|
||||
else:
|
||||
geom.contype = 0
|
||||
geom.conaffinity = 0
|
||||
geom.mass = 0.0
|
||||
|
||||
for body_name in variant_inertial_body_names:
|
||||
body = bodies_by_name.get(body_name)
|
||||
if body is None:
|
||||
continue
|
||||
body.explicitinertial = 0
|
||||
body.mass = 0.0
|
||||
body.inertia = np.zeros(3, dtype=np.float64)
|
||||
body.ipos = np.zeros(3, dtype=np.float64)
|
||||
body.iquat = np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float64)
|
||||
|
||||
for entity_prefix, metadata in variant_info:
|
||||
variant_idx = int(world_to_variant[entity_prefix][first_world])
|
||||
if variant_idx >= len(metadata.variant_body_inertials):
|
||||
continue
|
||||
for inertial in metadata.variant_body_inertials[variant_idx]:
|
||||
_apply_body_inertial(
|
||||
bodies_by_name,
|
||||
f"{entity_prefix}{inertial.body_name}",
|
||||
inertial,
|
||||
)
|
||||
|
||||
compiled_variants[key] = variant_spec.compile()
|
||||
|
||||
# Build per-world numpy arrays.
|
||||
ngeom = padded_model.ngeom
|
||||
nbody = padded_model.nbody
|
||||
|
||||
geom_size = np.zeros((nworld, ngeom, 3), dtype=np.float32)
|
||||
geom_rbound = np.zeros((nworld, ngeom), dtype=np.float32)
|
||||
geom_aabb = np.zeros((nworld, ngeom, 2, 3), dtype=np.float32)
|
||||
geom_pos = np.zeros((nworld, ngeom, 3), dtype=np.float32)
|
||||
geom_quat = np.zeros((nworld, ngeom, 4), dtype=np.float32)
|
||||
body_mass = np.zeros((nworld, nbody), dtype=np.float32)
|
||||
body_subtreemass = np.zeros((nworld, nbody), dtype=np.float32)
|
||||
body_inertia = np.zeros((nworld, nbody, 3), dtype=np.float32)
|
||||
body_invweight0 = np.zeros((nworld, nbody, 2), dtype=np.float32)
|
||||
body_ipos = np.zeros((nworld, nbody, 3), dtype=np.float32)
|
||||
body_iquat = np.zeros((nworld, nbody, 4), dtype=np.float32)
|
||||
|
||||
for w in range(nworld):
|
||||
key = tuple(dataid_table[w])
|
||||
ref = compiled_variants[key]
|
||||
geom_size[w] = ref.geom_size
|
||||
geom_rbound[w] = ref.geom_rbound
|
||||
geom_aabb[w] = ref.geom_aabb.reshape(ngeom, 2, 3)
|
||||
geom_pos[w] = ref.geom_pos
|
||||
geom_quat[w] = ref.geom_quat
|
||||
body_mass[w] = ref.body_mass
|
||||
body_subtreemass[w] = ref.body_subtreemass
|
||||
body_inertia[w] = ref.body_inertia
|
||||
body_invweight0[w] = ref.body_invweight0
|
||||
body_ipos[w] = ref.body_ipos
|
||||
body_iquat[w] = ref.body_iquat
|
||||
|
||||
m.geom_size = wp.array(geom_size, dtype=wp.vec3)
|
||||
m.geom_rbound = wp.array(geom_rbound, dtype=float)
|
||||
m.geom_aabb = wp.array(geom_aabb, dtype=wp.vec3)
|
||||
m.geom_pos = wp.array(geom_pos, dtype=wp.vec3)
|
||||
m.geom_quat = wp.array(geom_quat, dtype=wp.quat)
|
||||
m.body_mass = wp.array(body_mass, dtype=float)
|
||||
m.body_subtreemass = wp.array(body_subtreemass, dtype=float)
|
||||
m.body_inertia = wp.array(body_inertia, dtype=wp.vec3)
|
||||
m.body_invweight0 = wp.array(body_invweight0, dtype=wp.vec2)
|
||||
m.body_ipos = wp.array(body_ipos, dtype=wp.vec3)
|
||||
m.body_iquat = wp.array(body_iquat, dtype=wp.quat)
|
||||
|
||||
|
||||
def _apply_body_inertial(
|
||||
bodies_by_name: dict[str, mujoco.MjsBody],
|
||||
body_name: str,
|
||||
inertial: BodyInertialMetadata,
|
||||
) -> None:
|
||||
body = bodies_by_name.get(body_name)
|
||||
if body is None:
|
||||
raise ValueError(f"Body '{body_name}' not found in compiled variant spec.")
|
||||
body.explicitinertial = 1
|
||||
body.mass = inertial.mass
|
||||
body.ipos = np.asarray(inertial.ipos, dtype=np.float64)
|
||||
body.inertia = np.asarray(inertial.inertia, dtype=np.float64)
|
||||
body.iquat = np.asarray(inertial.iquat, dtype=np.float64)
|
||||
@@ -10,14 +10,14 @@ import mujoco_warp as mjwarp
|
||||
import torch
|
||||
import warp as wp
|
||||
|
||||
from mjlab.entity.variants import VARIANT_DEPENDENT_FIELDS, build_variant_model
|
||||
from mjlab.managers.event_manager import RecomputeLevel
|
||||
from mjlab.sim.mesh_variants import VARIANT_DEPENDENT_FIELDS, build_mesh_variant_model
|
||||
from mjlab.sim.randomization import expand_model_fields
|
||||
from mjlab.sim.sim_data import TorchArray, WarpBridge
|
||||
from mjlab.utils.nan_guard import NanGuard, NanGuardCfg
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from mjlab.entity.entity import VariantMetadata
|
||||
from mjlab.entity.variants import VariantMetadata
|
||||
from mjlab.sensor.sensor_context import SensorContext
|
||||
|
||||
# Type aliases for better IDE support while maintaining runtime compatibility
|
||||
@@ -246,7 +246,7 @@ class Simulation:
|
||||
they are rendering or inspecting.
|
||||
"""
|
||||
with wp.ScopedDevice(self.wp_device):
|
||||
result = build_mesh_variant_model(
|
||||
result = build_variant_model(
|
||||
spec,
|
||||
self.num_envs,
|
||||
variant_info,
|
||||
@@ -275,9 +275,10 @@ class Simulation:
|
||||
# viewer syncs them per-world.
|
||||
self._expanded_fields.update(VARIANT_DEPENDENT_FIELDS)
|
||||
self._expanded_fields.add("geom_dataid")
|
||||
self._expanded_fields.add("geom_matid")
|
||||
|
||||
# Stash variant assignments as torch tensors keyed by bare entity name
|
||||
# (mesh_variants emits "<name>/" prefixes; strip the trailing slash for
|
||||
# (build_variant_model emits "<name>/" prefixes; strip the trailing slash for
|
||||
# the public API).
|
||||
for prefix, arr in result.world_to_variant.items():
|
||||
key = prefix.rstrip("/")
|
||||
@@ -525,7 +526,7 @@ class Simulation:
|
||||
if not self.wp_device.is_cuda:
|
||||
return False
|
||||
|
||||
driver_ver = wp.context.runtime.driver_version
|
||||
driver_ver = wp.get_cuda_driver_version()
|
||||
has_mempool = wp.is_mempool_enabled(self.wp_device)
|
||||
|
||||
if driver_ver is None:
|
||||
|
||||
@@ -14,9 +14,10 @@ def compute_mpkpe(command: MotionCommand) -> torch.Tensor:
|
||||
"""Compute Mean Per-Keybody Position Error (MPKPE).
|
||||
|
||||
MPKPE measures the average Euclidean distance between the reference and
|
||||
actual positions of all key bodies in world frame.
|
||||
actual key body positions in the global world frame. It captures all
|
||||
tracking error, including global translation and heading drift.
|
||||
"""
|
||||
pos_error = command.body_pos_relative_w - command.robot_body_pos_w
|
||||
pos_error = command.body_pos_w - command.robot_body_pos_w
|
||||
per_body_error = torch.norm(pos_error, dim=-1) # (num_envs, num_bodies)
|
||||
return per_body_error.mean(dim=-1) # (num_envs,)
|
||||
|
||||
@@ -24,29 +25,25 @@ def compute_mpkpe(command: MotionCommand) -> torch.Tensor:
|
||||
def compute_root_relative_mpkpe(command: MotionCommand) -> torch.Tensor:
|
||||
"""Compute Root-relative Mean Per-Keybody Position Error (R-MPKPE).
|
||||
|
||||
R-MPKPE measures pose error independent of global drift by computing
|
||||
positions relative to the root/anchor body.
|
||||
R-MPKPE measures intrinsic pose error independent of global drift. It
|
||||
uses ``body_pos_relative_w``, the reference re-anchored to the robot's
|
||||
current root position and heading each step (the same quantity the
|
||||
tracking reward optimizes), so both global translation and yaw drift are
|
||||
removed and only the local body pose error remains.
|
||||
"""
|
||||
# Compute reference positions relative to reference anchor.
|
||||
ref_anchor_pos = command.anchor_pos_w.unsqueeze(1) # (num_envs, 1, 3)
|
||||
ref_rel_pos = command.body_pos_w - ref_anchor_pos # (num_envs, num_bodies, 3)
|
||||
|
||||
# Compute robot positions relative to robot anchor.
|
||||
robot_anchor_pos = command.robot_anchor_pos_w.unsqueeze(1) # (num_envs, 1, 3)
|
||||
robot_rel_pos = (
|
||||
command.robot_body_pos_w - robot_anchor_pos
|
||||
) # (num_envs, num_bodies, 3)
|
||||
|
||||
# Compute error between relative positions.
|
||||
pos_error = ref_rel_pos - robot_rel_pos
|
||||
pos_error = command.body_pos_relative_w - command.robot_body_pos_w
|
||||
per_body_error = torch.norm(pos_error, dim=-1) # (num_envs, num_bodies)
|
||||
return per_body_error.mean(dim=-1) # (num_envs,)
|
||||
|
||||
|
||||
def compute_joint_velocity_error(command: MotionCommand) -> torch.Tensor:
|
||||
"""Compute average joint velocity error."""
|
||||
"""Compute root-mean-square joint velocity error.
|
||||
|
||||
Uses an RMS over joints (rather than a raw L2 norm) so the value is a
|
||||
per-joint quantity, comparable across robots with different DOF counts.
|
||||
"""
|
||||
vel_error = command.joint_vel - command.robot_joint_vel
|
||||
return torch.norm(vel_error, dim=-1) # (num_envs,)
|
||||
return torch.sqrt(torch.mean(vel_error**2, dim=-1)) # (num_envs,)
|
||||
|
||||
|
||||
def compute_ee_position_error(
|
||||
@@ -93,6 +90,18 @@ def _get_body_indices(
|
||||
body_names: Names of bodies to find.
|
||||
|
||||
Returns:
|
||||
List of indices into command.cfg.body_names.
|
||||
List of indices into command.cfg.body_names, in the order requested.
|
||||
|
||||
Raises:
|
||||
ValueError: If any requested body name is not tracked by the command.
|
||||
Silently dropping unknown names would otherwise report a spurious
|
||||
zero error for misconfigured end-effector lists.
|
||||
"""
|
||||
return [i for i, name in enumerate(command.cfg.body_names) if name in body_names]
|
||||
name_to_index = {name: i for i, name in enumerate(command.cfg.body_names)}
|
||||
missing = [name for name in body_names if name not in name_to_index]
|
||||
if missing:
|
||||
raise ValueError(
|
||||
f"Body names {missing} are not tracked by the command. "
|
||||
f"Available bodies: {tuple(command.cfg.body_names)}."
|
||||
)
|
||||
return [name_to_index[name] for name in body_names]
|
||||
|
||||
@@ -95,7 +95,9 @@ class MotionTrackingOnPolicyRunner(MjlabOnPolicyRunner):
|
||||
try:
|
||||
self.export_policy_to_onnx(str(policy_dir), filename)
|
||||
run_name: str = (
|
||||
wandb.run.name if self.logger.logger_type == "wandb" and wandb.run else "local"
|
||||
wandb.run.name
|
||||
if self.logger.logger_type in ("wandb", "WandbLogWriter") and wandb.run
|
||||
else "local"
|
||||
) # type: ignore[assignment]
|
||||
metadata = get_base_metadata(self.env.unwrapped, run_name)
|
||||
motion_term = cast(
|
||||
@@ -108,7 +110,10 @@ class MotionTrackingOnPolicyRunner(MjlabOnPolicyRunner):
|
||||
}
|
||||
)
|
||||
attach_metadata_to_onnx(str(onnx_path), metadata)
|
||||
if self.logger.logger_type in ["wandb"] and self.cfg["upload_model"]:
|
||||
if (
|
||||
self.logger.logger_type in ("wandb", "WandbLogWriter")
|
||||
and self.cfg["upload_model"]
|
||||
):
|
||||
wandb.save(str(onnx_path), base_path=str(policy_dir))
|
||||
if self.registry_name is not None:
|
||||
wandb.run.use_artifact(self.registry_name) # type: ignore
|
||||
|
||||
@@ -6,6 +6,7 @@ import json
|
||||
import sys
|
||||
from dataclasses import asdict, dataclass
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from typing import cast
|
||||
|
||||
import torch
|
||||
@@ -97,35 +98,61 @@ def run_evaluate(task_id: str, cfg: EvaluateConfig) -> dict[str, float]:
|
||||
all_joint_vel_error: list[torch.Tensor] = []
|
||||
all_ee_pos_error: list[torch.Tensor] = []
|
||||
all_ee_ori_error: list[torch.Tensor] = []
|
||||
all_active: list[torch.Tensor] = []
|
||||
|
||||
done_envs = torch.zeros(cfg.num_envs, dtype=torch.bool, device=device)
|
||||
success = torch.zeros(cfg.num_envs, dtype=torch.bool, device=device)
|
||||
|
||||
obs = env.get_observations()
|
||||
env.unwrapped.command_manager.compute(dt=env.unwrapped.step_dt)
|
||||
|
||||
print(f"[INFO] Running {cfg.num_envs} evaluation episodes...")
|
||||
|
||||
step = 0
|
||||
while not done_envs.all():
|
||||
# Snapshot the reference frame the upcoming step will be scored against.
|
||||
# env.step computes the reward (against the current reference) and only
|
||||
# afterwards advances the command's motion frame, so reading the
|
||||
# reference after stepping would pair the robot with the *next* frame.
|
||||
# We snapshot here and pair it with the post-step robot state below,
|
||||
# matching how the reward is computed.
|
||||
ref = SimpleNamespace(
|
||||
num_envs=command.num_envs,
|
||||
device=command.device,
|
||||
cfg=command.cfg,
|
||||
body_pos_w=command.body_pos_w.clone(),
|
||||
body_pos_relative_w=command.body_pos_relative_w.clone(),
|
||||
body_quat_relative_w=command.body_quat_relative_w.clone(),
|
||||
joint_vel=command.joint_vel.clone(),
|
||||
)
|
||||
|
||||
with torch.no_grad():
|
||||
actions = policy(obs)
|
||||
obs, _, dones, _ = env.step(actions)
|
||||
|
||||
# Compute metrics for active envs.
|
||||
# Pair the snapshotted reference with the post-step robot state.
|
||||
ref.robot_body_pos_w = command.robot_body_pos_w
|
||||
ref.robot_body_quat_w = command.robot_body_quat_w
|
||||
ref.robot_joint_vel = command.robot_joint_vel
|
||||
ref_command = cast(MotionCommand, ref)
|
||||
|
||||
# Accumulate metrics for envs still running this step. active.any() is
|
||||
# always true here: the loop runs only while some env is not done, and
|
||||
# done_envs is updated below after this point.
|
||||
active = ~done_envs
|
||||
if active.any():
|
||||
all_mpkpe.append(torch.where(active, compute_mpkpe(command), 0.0))
|
||||
all_r_mpkpe.append(torch.where(active, compute_root_relative_mpkpe(command), 0.0))
|
||||
all_joint_vel_error.append(
|
||||
torch.where(active, compute_joint_velocity_error(command), 0.0)
|
||||
)
|
||||
all_ee_pos_error.append(
|
||||
torch.where(active, compute_ee_position_error(command, ee_body_names), 0.0)
|
||||
)
|
||||
all_ee_ori_error.append(
|
||||
torch.where(active, compute_ee_orientation_error(command, ee_body_names), 0.0)
|
||||
)
|
||||
all_active.append(active.float())
|
||||
all_mpkpe.append(torch.where(active, compute_mpkpe(ref_command), 0.0))
|
||||
all_r_mpkpe.append(
|
||||
torch.where(active, compute_root_relative_mpkpe(ref_command), 0.0)
|
||||
)
|
||||
all_joint_vel_error.append(
|
||||
torch.where(active, compute_joint_velocity_error(ref_command), 0.0)
|
||||
)
|
||||
all_ee_pos_error.append(
|
||||
torch.where(active, compute_ee_position_error(ref_command, ee_body_names), 0.0)
|
||||
)
|
||||
all_ee_ori_error.append(
|
||||
torch.where(active, compute_ee_orientation_error(ref_command, ee_body_names), 0.0)
|
||||
)
|
||||
|
||||
# Track completions.
|
||||
terminated = env.unwrapped.termination_manager.terminated
|
||||
@@ -142,7 +169,7 @@ def run_evaluate(task_id: str, cfg: EvaluateConfig) -> dict[str, float]:
|
||||
)
|
||||
step += 1
|
||||
|
||||
# Compute mean metrics.
|
||||
# Compute mean metrics over the steps each env was active.
|
||||
stacks = [
|
||||
all_mpkpe,
|
||||
all_r_mpkpe,
|
||||
@@ -151,7 +178,7 @@ def run_evaluate(task_id: str, cfg: EvaluateConfig) -> dict[str, float]:
|
||||
all_ee_ori_error,
|
||||
]
|
||||
stacks = [torch.stack(s, dim=0) for s in stacks]
|
||||
active_steps = (stacks[0] != 0).sum(dim=0).float().clamp(min=1)
|
||||
active_steps = torch.stack(all_active, dim=0).sum(dim=0).clamp(min=1)
|
||||
means = [s.sum(dim=0) / active_steps for s in stacks]
|
||||
|
||||
metrics = {
|
||||
|
||||
@@ -24,14 +24,32 @@ from mjlab.terrains.terrain_generator import (
|
||||
)
|
||||
from mjlab.terrains.utils import find_flat_patches_from_heightfield
|
||||
|
||||
# Smallest positive hfield elevation/base size, in meters. MuJoCo rejects
|
||||
# non-positive hfield sizes, so flat heightfields (difficulty 0) are clamped to
|
||||
# this instead of zero.
|
||||
_MIN_HFIELD_HEIGHT = 1e-3
|
||||
|
||||
# Physical height (meters) that maps to full color saturation. Heights are
|
||||
# colored on this fixed absolute scale rather than normalized per patch, so a
|
||||
# given height reads the same color across every terrain and small-amplitude
|
||||
# terrain stays gently tinted instead of stretching into rainbow noise.
|
||||
_COLOR_SCALE = 0.75
|
||||
|
||||
|
||||
def color_by_height(
|
||||
spec: mujoco.MjSpec,
|
||||
noise: np.ndarray,
|
||||
unique_id: str,
|
||||
normalized_elevation: np.ndarray,
|
||||
physical_heights: np.ndarray,
|
||||
texture_size: int = 128,
|
||||
) -> str:
|
||||
"""Build a height-colored texture for a heightfield.
|
||||
|
||||
Diverging colormap anchored at the ground plane (z=0): cool blue below ground,
|
||||
green at z=0, warm red above. ``physical_heights`` is the surface height of
|
||||
each cell in meters relative to z=0; it is colored on the fixed ``_COLOR_SCALE``
|
||||
so color encodes absolute height consistently across all terrains.
|
||||
"""
|
||||
texture_name = f"hf_texture_{unique_id}"
|
||||
texture = spec.add_texture(
|
||||
name=texture_name,
|
||||
@@ -40,16 +58,20 @@ def color_by_height(
|
||||
height=texture_size,
|
||||
)
|
||||
|
||||
texture_elevation = ndimage.zoom(
|
||||
normalized_elevation,
|
||||
texture_height = ndimage.zoom(
|
||||
physical_heights,
|
||||
(texture_size / noise.shape[0], texture_size / noise.shape[1]),
|
||||
order=1,
|
||||
)
|
||||
texture_elevation = np.asarray(texture_elevation)
|
||||
texture_height = np.asarray(texture_height)
|
||||
|
||||
hue = 0.5 - texture_elevation * 0.45
|
||||
saturation = 0.6 - texture_elevation * 0.2
|
||||
value = 0.4 + texture_elevation * 0.3
|
||||
# Signed deviation from the ground plane in [-1, 1] on a fixed absolute scale.
|
||||
signed = np.clip(texture_height / _COLOR_SCALE, -1.0, 1.0)
|
||||
|
||||
# signed=+1 -> hue 0.0 (red, high), 0 -> 0.33 (green, ground), -1 -> 0.66 (blue, low).
|
||||
hue = 0.33 - 0.33 * signed
|
||||
saturation = 0.45 + 0.25 * np.abs(signed)
|
||||
value = 0.45 + 0.25 * np.abs(signed)
|
||||
|
||||
c = value * saturation
|
||||
x = c * (1 - np.abs((hue * 6) % 2 - 1))
|
||||
@@ -326,7 +348,8 @@ class HfPyramidSlopedTerrainCfg(SubTerrainCfg):
|
||||
else:
|
||||
hfield_z_offset = 0
|
||||
|
||||
material_name = color_by_height(spec, noise, unique_id, normalized_elevation)
|
||||
physical_heights = hfield_z_offset + normalized_elevation * max_physical_height
|
||||
material_name = color_by_height(spec, noise, unique_id, physical_heights)
|
||||
|
||||
hfield_geom = body.add_geom(
|
||||
type=mujoco.mjtGeom.mjGEOM_HFIELD,
|
||||
@@ -378,14 +401,24 @@ class HfRandomUniformTerrainCfg(SubTerrainCfg):
|
||||
border_width: float = 0.0
|
||||
"""Width of the flat border around the terrain edges, in meters. Must be >=
|
||||
horizontal_scale if non-zero."""
|
||||
scale_with_difficulty: bool = False
|
||||
"""If False (default), the roughness is fixed and ``difficulty`` is ignored,
|
||||
matching upstream behavior. If True, the noise amplitude scales linearly with
|
||||
difficulty (flat at 0, full ``noise_range`` at 1) so the terrain progresses in
|
||||
a curriculum."""
|
||||
|
||||
def function(
|
||||
self, difficulty: float, spec: mujoco.MjSpec, rng: np.random.Generator
|
||||
) -> TerrainOutput:
|
||||
del difficulty # Unused.
|
||||
|
||||
body = spec.body("terrain")
|
||||
|
||||
# When difficulty scaling is enabled, ramp the noise amplitude from flat (0)
|
||||
# to the full configured range (1). Otherwise use the full range regardless
|
||||
# of difficulty (difficulty is ignored).
|
||||
scale = difficulty if self.scale_with_difficulty else 1.0
|
||||
noise_lo = self.noise_range[0] * scale
|
||||
noise_hi = self.noise_range[1] * scale
|
||||
|
||||
if self.border_width > 0 and self.border_width < self.horizontal_scale:
|
||||
raise ValueError(
|
||||
f"Border width ({self.border_width}) must be >= horizontal scale "
|
||||
@@ -419,8 +452,8 @@ class HfRandomUniformTerrainCfg(SubTerrainCfg):
|
||||
width_downsampled = int(inner_size[0] / downsampled_scale)
|
||||
length_downsampled = int(inner_size[1] / downsampled_scale)
|
||||
|
||||
height_min = int(self.noise_range[0] / self.vertical_scale)
|
||||
height_max = int(self.noise_range[1] / self.vertical_scale)
|
||||
height_min = int(noise_lo / self.vertical_scale)
|
||||
height_max = int(noise_hi / self.vertical_scale)
|
||||
height_step = int(self.noise_step / self.vertical_scale)
|
||||
|
||||
height_range = np.arange(height_min, height_max + height_step, height_step)
|
||||
@@ -443,8 +476,8 @@ class HfRandomUniformTerrainCfg(SubTerrainCfg):
|
||||
else:
|
||||
width_downsampled = int(self.size[0] / downsampled_scale)
|
||||
length_downsampled = int(self.size[1] / downsampled_scale)
|
||||
height_min = int(self.noise_range[0] / self.vertical_scale)
|
||||
height_max = int(self.noise_range[1] / self.vertical_scale)
|
||||
height_min = int(noise_lo / self.vertical_scale)
|
||||
height_max = int(noise_hi / self.vertical_scale)
|
||||
height_step = int(self.noise_step / self.vertical_scale)
|
||||
|
||||
height_range = np.arange(height_min, height_max + height_step, height_step)
|
||||
@@ -489,7 +522,8 @@ class HfRandomUniformTerrainCfg(SubTerrainCfg):
|
||||
userdata=normalized_elevation.flatten().astype(np.float32).tolist(),
|
||||
)
|
||||
|
||||
material_name = color_by_height(spec, noise, unique_id, normalized_elevation)
|
||||
physical_heights = normalized_elevation * max_physical_height
|
||||
material_name = color_by_height(spec, noise, unique_id, physical_heights)
|
||||
|
||||
hfield_geom = body.add_geom(
|
||||
type=mujoco.mjtGeom.mjGEOM_HFIELD,
|
||||
@@ -498,7 +532,7 @@ class HfRandomUniformTerrainCfg(SubTerrainCfg):
|
||||
material=material_name,
|
||||
)
|
||||
|
||||
spawn_height = (self.noise_range[0] + self.noise_range[1]) / 2
|
||||
spawn_height = (noise_lo + noise_hi) / 2
|
||||
origin = np.array([self.size[0] / 2, self.size[1] / 2, spawn_height])
|
||||
|
||||
flat_patches = _compute_flat_patches(
|
||||
@@ -616,7 +650,11 @@ class HfWaveTerrainCfg(SubTerrainCfg):
|
||||
userdata=normalized_elevation.flatten().astype(np.float32).tolist(),
|
||||
)
|
||||
|
||||
material_name = color_by_height(spec, noise, unique_id, normalized_elevation)
|
||||
# The wave oscillates around z=0 (geom is offset down by half the range).
|
||||
physical_heights = (
|
||||
normalized_elevation * max_physical_height - max_physical_height / 2
|
||||
)
|
||||
material_name = color_by_height(spec, noise, unique_id, physical_heights)
|
||||
|
||||
hfield_geom = body.add_geom(
|
||||
type=mujoco.mjtGeom.mjGEOM_HFIELD,
|
||||
@@ -783,7 +821,9 @@ class HfDiscreteObstaclesTerrainCfg(SubTerrainCfg):
|
||||
else:
|
||||
hfield_z_offset = 0
|
||||
|
||||
material_name = color_by_height(spec, noise, unique_id, normalized_elevation)
|
||||
# Physical surface height per cell (pits negative, bumps positive about z=0).
|
||||
physical_heights = hfield_z_offset + normalized_elevation * max_physical_height
|
||||
material_name = color_by_height(spec, noise, unique_id, physical_heights)
|
||||
|
||||
hfield_geom = body.add_geom(
|
||||
type=mujoco.mjtGeom.mjGEOM_HFIELD,
|
||||
@@ -887,8 +927,13 @@ class HfPerlinNoiseTerrainCfg(SubTerrainCfg):
|
||||
noise_range = noise_max - noise_min if noise_max > noise_min else 1.0
|
||||
normalized_elevation = ((noise_raw - noise_min) / noise_range).astype(np.float32)
|
||||
|
||||
max_physical_height = target_height
|
||||
base_thickness = max_physical_height * self.base_thickness_ratio
|
||||
# MuJoCo requires positive hfield elevation and base sizes. At difficulty 0
|
||||
# (target_height == 0) the surface is flat; clamp to a small positive height
|
||||
# so compilation does not fail with "size parameter is not positive".
|
||||
max_physical_height = max(target_height, _MIN_HFIELD_HEIGHT)
|
||||
base_thickness = max(
|
||||
max_physical_height * self.base_thickness_ratio, _MIN_HFIELD_HEIGHT
|
||||
)
|
||||
|
||||
unique_id = uuid.uuid4().hex
|
||||
field = spec.add_hfield(
|
||||
@@ -904,8 +949,9 @@ class HfPerlinNoiseTerrainCfg(SubTerrainCfg):
|
||||
userdata=normalized_elevation.flatten().tolist(),
|
||||
)
|
||||
|
||||
physical_heights = normalized_elevation * max_physical_height
|
||||
material_name = color_by_height(
|
||||
spec, normalized_elevation, unique_id, normalized_elevation
|
||||
spec, normalized_elevation, unique_id, physical_heights
|
||||
)
|
||||
|
||||
hfield_geom = body.add_geom(
|
||||
|
||||
@@ -11,7 +11,6 @@ References:
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Tuple
|
||||
|
||||
import mujoco
|
||||
import numpy as np
|
||||
@@ -23,30 +22,19 @@ from mjlab.terrains.terrain_generator import (
|
||||
)
|
||||
from mjlab.terrains.utils import make_border, make_plane
|
||||
from mjlab.utils.color import (
|
||||
HSV,
|
||||
brand_ramp,
|
||||
clamp,
|
||||
darken_rgba,
|
||||
hsv_to_rgb,
|
||||
rgb_to_hsv,
|
||||
)
|
||||
|
||||
_MUJOCO_BLUE = (0.20, 0.45, 0.95)
|
||||
_MUJOCO_RED = (0.90, 0.30, 0.30)
|
||||
_MUJOCO_GREEN = (0.25, 0.80, 0.45)
|
||||
|
||||
|
||||
def _get_platform_color(
|
||||
base_rgb: Tuple[float, float, float],
|
||||
desaturation_factor: float = 0.4,
|
||||
lightening_factor: float = 0.25,
|
||||
) -> Tuple[float, float, float, float]:
|
||||
hsv = rgb_to_hsv(base_rgb)
|
||||
new_s = hsv.s * desaturation_factor
|
||||
new_v = clamp(hsv.v + lightening_factor)
|
||||
new_hsv = HSV(hsv.h, new_s, new_v)
|
||||
r, g, b = hsv_to_rgb(new_hsv)
|
||||
return (r, g, b, 1.0)
|
||||
# Minimum vertical extent of a flat border frame, in meters. The border top sits
|
||||
# flush at z=0 and extends downward, so this depth is not visible from above; it
|
||||
# only guarantees the frame is solid (never a degenerate zero-height geom) when
|
||||
# the step height collapses to zero at difficulty 0.
|
||||
_MIN_BORDER_HEIGHT = 0.05
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
@@ -107,13 +95,18 @@ class BoxPyramidStairsTerrainCfg(SubTerrainCfg):
|
||||
border_rgba = darken_rgba(first_step_rgba, 0.85)
|
||||
|
||||
if self.border_width > 0.0 and not self.holes:
|
||||
border_center = (0.5 * self.size[0], 0.5 * self.size[1], -step_height / 2)
|
||||
# Decouple the border's vertical extent from step_height so difficulty 0
|
||||
# (step_height == 0) still produces a solid, gap-free frame instead of
|
||||
# being skipped or generating degenerate zero-height geoms. The top stays
|
||||
# flush with the ground at z=0.
|
||||
border_height = max(step_height, _MIN_BORDER_HEIGHT)
|
||||
border_center = (0.5 * self.size[0], 0.5 * self.size[1], -border_height / 2)
|
||||
border_inner_size = (
|
||||
self.size[0] - 2 * self.border_width,
|
||||
self.size[1] - 2 * self.border_width,
|
||||
)
|
||||
border_boxes = make_border(
|
||||
body, self.size, border_inner_size, step_height, border_center
|
||||
body, self.size, border_inner_size, border_height, border_center
|
||||
)
|
||||
boxes.extend(border_boxes)
|
||||
for _ in range(len(border_boxes)):
|
||||
@@ -280,13 +273,16 @@ class BoxInvertedPyramidStairsTerrainCfg(BoxPyramidStairsTerrainCfg):
|
||||
border_rgba = darken_rgba(first_step_rgba, 0.85)
|
||||
|
||||
if self.border_width > 0.0 and not self.holes:
|
||||
border_center = (0.5 * self.size[0], 0.5 * self.size[1], -0.5 * step_height)
|
||||
# See BoxPyramidStairsTerrainCfg: keep the border solid and flush at z=0
|
||||
# even when step_height collapses to 0 at difficulty 0.
|
||||
border_height = max(step_height, _MIN_BORDER_HEIGHT)
|
||||
border_center = (0.5 * self.size[0], 0.5 * self.size[1], -0.5 * border_height)
|
||||
border_inner_size = (
|
||||
self.size[0] - 2 * self.border_width,
|
||||
self.size[1] - 2 * self.border_width,
|
||||
)
|
||||
border_boxes = make_border(
|
||||
body, self.size, border_inner_size, step_height, border_center
|
||||
body, self.size, border_inner_size, border_height, border_center
|
||||
)
|
||||
boxes.extend(border_boxes)
|
||||
for _ in range(len(border_boxes)):
|
||||
@@ -546,8 +542,7 @@ class BoxRandomGridTerrainCfg(SubTerrainCfg):
|
||||
pos=(self.size[0] / 2, self.size[1] / 2, platform_center_z),
|
||||
)
|
||||
boxes_list.append(box)
|
||||
platform_rgba = _get_platform_color(_MUJOCO_GREEN)
|
||||
box_colors.append(platform_rgba)
|
||||
box_colors.append(brand_ramp(_MUJOCO_GREEN, 0.5))
|
||||
|
||||
origin = np.array([self.size[0] / 2, self.size[1] / 2, grid_height])
|
||||
|
||||
@@ -575,6 +570,22 @@ class BoxRandomGridTerrainCfg(SubTerrainCfg):
|
||||
half_border_width = border_width / 2
|
||||
neg_half_terrain = -terrain_height / 2
|
||||
|
||||
# Mark cells under the center platform as visited so they are never emitted
|
||||
# or merged; the platform box covers that region and would otherwise z-fight
|
||||
# with the cells beneath it.
|
||||
platform_half = self.platform_width / 2
|
||||
terrain_center = self.size[0] / 2
|
||||
platform_min = terrain_center - platform_half
|
||||
platform_max = terrain_center + platform_half
|
||||
for i in range(num_boxes_x):
|
||||
cx = half_border_width + (i + 0.5) * self.grid_width
|
||||
if not (platform_min <= cx <= platform_max):
|
||||
continue
|
||||
for j in range(num_boxes_y):
|
||||
cy = half_border_width + (j + 0.5) * self.grid_width
|
||||
if platform_min <= cy <= platform_max:
|
||||
visited[i, j] = True
|
||||
|
||||
# Quantize heights to create more merging opportunities
|
||||
quantized_heights = (
|
||||
np.round(height_map / self.height_merge_threshold) * self.height_merge_threshold
|
||||
@@ -588,7 +599,12 @@ class BoxRandomGridTerrainCfg(SubTerrainCfg):
|
||||
# Find rectangular region with similar height
|
||||
height = quantized_heights[i, j]
|
||||
|
||||
normalized_height = (height + grid_height) / (2 * grid_height)
|
||||
# grid_height == 0 (difficulty 0) means a flat grid; use the midpoint
|
||||
# color and avoid dividing by zero.
|
||||
if grid_height > 0:
|
||||
normalized_height = (height + grid_height) / (2 * grid_height)
|
||||
else:
|
||||
normalized_height = 0.5
|
||||
t = float(np.clip(normalized_height, 0.0, 1.0))
|
||||
rgba = brand_ramp(_MUJOCO_GREEN, t)
|
||||
|
||||
@@ -653,14 +669,10 @@ class BoxRandomGridTerrainCfg(SubTerrainCfg):
|
||||
half_border_width = border_width / 2
|
||||
neg_half_terrain = -terrain_height / 2
|
||||
|
||||
if self.holes:
|
||||
platform_half = self.platform_width / 2
|
||||
terrain_center = self.size[0] / 2
|
||||
platform_min = terrain_center - platform_half
|
||||
platform_max = terrain_center + platform_half
|
||||
else:
|
||||
platform_min = None
|
||||
platform_max = None
|
||||
platform_half = self.platform_width / 2
|
||||
terrain_center = self.size[0] / 2
|
||||
platform_min = terrain_center - platform_half
|
||||
platform_max = terrain_center + platform_half
|
||||
|
||||
for i in range(num_boxes_x):
|
||||
box_center_x = half_border_width + (i + 0.5) * self.grid_width
|
||||
@@ -678,11 +690,24 @@ class BoxRandomGridTerrainCfg(SubTerrainCfg):
|
||||
if not (in_x_strip or in_y_strip):
|
||||
continue
|
||||
|
||||
# Skip cells under the center platform so the platform is the only
|
||||
# geometry there. Otherwise the platform box sits on top of these cells
|
||||
# and the coplanar faces z-fight.
|
||||
if (platform_min <= box_center_x <= platform_max) and (
|
||||
platform_min <= box_center_y <= platform_max
|
||||
):
|
||||
continue
|
||||
|
||||
height_noise = height_map[i, j]
|
||||
box_height = terrain_height + height_noise
|
||||
box_center_z = neg_half_terrain + height_noise / 2
|
||||
|
||||
normalized_height = (height_noise + grid_height) / (2 * grid_height)
|
||||
# grid_height == 0 (difficulty 0) means a flat grid; use the midpoint
|
||||
# color and avoid dividing by zero.
|
||||
if grid_height > 0:
|
||||
normalized_height = (height_noise + grid_height) / (2 * grid_height)
|
||||
else:
|
||||
normalized_height = 0.5
|
||||
t = float(np.clip(normalized_height, 0.0, 1.0))
|
||||
rgba = brand_ramp(_MUJOCO_GREEN, t)
|
||||
box_colors.append(rgba)
|
||||
@@ -744,13 +769,17 @@ class BoxRandomSpreadTerrainCfg(SubTerrainCfg):
|
||||
)
|
||||
geometries.append(TerrainGeometry(geom=floor_geom, color=(0.4, 0.4, 0.4, 1.0)))
|
||||
|
||||
# Platform
|
||||
platform_geom = body.add_geom(
|
||||
type=mujoco.mjtGeom.mjGEOM_BOX,
|
||||
size=(self.platform_width / 2, self.platform_width / 2, terrain_height / 2),
|
||||
pos=(self.size[0] / 2, self.size[1] / 2, -terrain_height / 2),
|
||||
)
|
||||
geometries.append(TerrainGeometry(geom=platform_geom, color=(0.4, 0.4, 0.4, 1.0)))
|
||||
# Center platform. When a floor is present it already provides flat ground at
|
||||
# z=0 across the (box-free) center, so an extra platform box would only
|
||||
# duplicate that surface and z-fight with the floor. Add the platform only
|
||||
# when there is no floor, where it is the sole ground at the spawn point.
|
||||
if not self.add_floor:
|
||||
platform_geom = body.add_geom(
|
||||
type=mujoco.mjtGeom.mjGEOM_BOX,
|
||||
size=(self.platform_width / 2, self.platform_width / 2, terrain_height / 2),
|
||||
pos=(self.size[0] / 2, self.size[1] / 2, -terrain_height / 2),
|
||||
)
|
||||
geometries.append(TerrainGeometry(geom=platform_geom, color=(0.4, 0.4, 0.4, 1.0)))
|
||||
|
||||
platform_half = self.platform_width / 2
|
||||
terrain_center = self.size[0] / 2
|
||||
@@ -840,13 +869,15 @@ class BoxOpenStairsTerrainCfg(SubTerrainCfg):
|
||||
border_rgba = darken_rgba(first_step_rgba, 0.85)
|
||||
|
||||
if self.border_width > 0.0:
|
||||
border_center = (0.5 * self.size[0], 0.5 * self.size[1], -step_height / 2)
|
||||
# Keep the border solid and flush at z=0 even if step_height is 0.
|
||||
border_height = max(step_height, _MIN_BORDER_HEIGHT)
|
||||
border_center = (0.5 * self.size[0], 0.5 * self.size[1], -border_height / 2)
|
||||
border_inner_size = (
|
||||
self.size[0] - 2 * self.border_width,
|
||||
self.size[1] - 2 * self.border_width,
|
||||
)
|
||||
border_boxes = make_border(
|
||||
body, self.size, border_inner_size, step_height, border_center
|
||||
body, self.size, border_inner_size, border_height, border_center
|
||||
)
|
||||
for box in border_boxes:
|
||||
geometries.append(TerrainGeometry(geom=box, color=border_rgba))
|
||||
@@ -1132,7 +1163,12 @@ class BoxRandomStairsTerrainCfg(SubTerrainCfg):
|
||||
@dataclass(kw_only=True)
|
||||
class BoxSteppingStonesTerrainCfg(SubTerrainCfg):
|
||||
stone_size_range: tuple[float, float] = (0.4, 0.8)
|
||||
"""Max and min stone side length, in meters. Stones shrink from the max toward
|
||||
the min as difficulty increases, which widens the gaps between them."""
|
||||
stone_distance_range: tuple[float, float] = (0.2, 0.5)
|
||||
"""Gap between stones, in meters. The lower bound seeds the (fixed) grid
|
||||
density; the gap itself grows with difficulty as the stones shrink, so the
|
||||
upper bound is no longer used directly."""
|
||||
stone_height: float = 0.2
|
||||
stone_height_variation: float = 0.1
|
||||
stone_size_variation: float = 0.1
|
||||
@@ -1152,23 +1188,56 @@ class BoxSteppingStonesTerrainCfg(SubTerrainCfg):
|
||||
displacement_range = self.displacement_range * difficulty
|
||||
stone_height_variation = self.stone_height_variation * difficulty
|
||||
|
||||
# Increase distance between stones with difficulty.
|
||||
d_low, d_high = self.stone_distance_range
|
||||
avg_distance = d_low + difficulty * (d_high - d_low)
|
||||
|
||||
# Decrease stone size with difficulty (larger stones are easier).
|
||||
# Decrease stone size with difficulty (larger stones are easier). With the
|
||||
# grid pitch held fixed (below), shrinking stones means the gaps between them
|
||||
# grow, which is the actual difficulty curriculum.
|
||||
s_min, s_max = self.stone_size_range
|
||||
avg_stone_size = s_max - difficulty * (s_max - s_min)
|
||||
spacing = avg_stone_size + avg_distance
|
||||
|
||||
# Aggressive grid density to reach borders.
|
||||
# Difficulty-INDEPENDENT grid. The count and pitch are fixed across difficulty
|
||||
# so the layout never re-tiles (previously, num = floor(inner / spacing) + 1
|
||||
# jumped by one as the difficulty-varying spacing crossed an integer boundary,
|
||||
# shifting every stone at once). The pitch exactly spans the inner region so
|
||||
# edge stones always reach the borders. Density is seeded by the tightest
|
||||
# nominal spacing (largest stones + smallest gap).
|
||||
inner_w = self.size[0] - 2 * self.border_width
|
||||
inner_h = self.size[1] - 2 * self.border_width
|
||||
num_x = int(np.floor(inner_w / spacing)) + 1
|
||||
num_y = int(np.floor(inner_h / spacing)) + 1
|
||||
nominal_spacing = s_max + self.stone_distance_range[0]
|
||||
num_x = max(2, int(np.floor(inner_w / nominal_spacing)) + 1)
|
||||
num_y = max(2, int(np.floor(inner_h / nominal_spacing)) + 1)
|
||||
pitch_x = inner_w / (num_x - 1)
|
||||
pitch_y = inner_h / (num_y - 1)
|
||||
|
||||
offset_x = self.border_width + (inner_w - (num_x - 1) * spacing) / 2
|
||||
offset_y = self.border_width + (inner_h - (num_y - 1) * spacing) / 2
|
||||
# Inter-stone gap (grows with difficulty as stones shrink).
|
||||
gap_x = max(0.0, pitch_x - avg_stone_size)
|
||||
gap_y = max(0.0, pitch_y - avg_stone_size)
|
||||
|
||||
# Snap the central platform out to the grid. It is at least the configured
|
||||
# width and reaches to exactly one gap before the nearest *full* stone, so the
|
||||
# ring of stones around it are whole (no clipped slivers that pop in and out
|
||||
# with difficulty) and sit one consistent gap away. The platform simply
|
||||
# absorbs the stones that would otherwise be partially under it.
|
||||
center_x, center_y = self.size[0] / 2, self.size[1] / 2
|
||||
half_stone = avg_stone_size / 2
|
||||
a0 = self.platform_width / 2
|
||||
|
||||
def _snapped_half(center: float, pitch: float, gap: float, num: int) -> float:
|
||||
# Nearest grid stone that can stay full while the platform is >= a0 wide.
|
||||
threshold = center + a0 + half_stone + gap
|
||||
i_keep = min(num - 1, int(np.ceil((threshold - self.border_width) / pitch)))
|
||||
c_keep = self.border_width + i_keep * pitch
|
||||
return max(a0, c_keep - half_stone - gap - center)
|
||||
|
||||
platform_half_x = _snapped_half(center_x, pitch_x, gap_x, num_x)
|
||||
platform_half_y = _snapped_half(center_y, pitch_y, gap_y, num_y)
|
||||
platform_min_x, platform_max_x = (
|
||||
center_x - platform_half_x,
|
||||
center_x + platform_half_x,
|
||||
)
|
||||
platform_min_y, platform_max_y = (
|
||||
center_y - platform_half_y,
|
||||
center_y + platform_half_y,
|
||||
)
|
||||
|
||||
border_rgba = darken_rgba(brand_ramp(_MUJOCO_GREEN, 0.0), 0.85)
|
||||
z_center = (self.stone_height - self.floor_depth) / 2
|
||||
@@ -1195,25 +1264,20 @@ class BoxSteppingStonesTerrainCfg(SubTerrainCfg):
|
||||
)
|
||||
geometries.append(TerrainGeometry(geom=floor_geom, color=(0.1, 0.1, 0.1, 1.0)))
|
||||
|
||||
# Platform Column.
|
||||
# Platform Column (grid-snapped, see above).
|
||||
platform_geom = body.add_geom(
|
||||
type=mujoco.mjtGeom.mjGEOM_BOX,
|
||||
size=(
|
||||
np.maximum(1e-6, self.platform_width / 2),
|
||||
np.maximum(1e-6, self.platform_width / 2),
|
||||
np.maximum(1e-6, platform_half_x),
|
||||
np.maximum(1e-6, platform_half_y),
|
||||
np.maximum(1e-6, half_height),
|
||||
),
|
||||
pos=(self.size[0] / 2, self.size[1] / 2, z_center),
|
||||
pos=(center_x, center_y, z_center),
|
||||
)
|
||||
geometries.append(
|
||||
TerrainGeometry(geom=platform_geom, color=brand_ramp(_MUJOCO_GREEN, 0.5))
|
||||
)
|
||||
|
||||
platform_half = self.platform_width / 2
|
||||
terrain_center = self.size[0] / 2
|
||||
platform_min = terrain_center - platform_half
|
||||
platform_max = terrain_center + platform_half
|
||||
|
||||
inner_min_x, inner_max_x = self.border_width, self.size[0] - self.border_width
|
||||
inner_min_y, inner_max_y = self.border_width, self.size[1] - self.border_width
|
||||
|
||||
@@ -1221,12 +1285,17 @@ class BoxSteppingStonesTerrainCfg(SubTerrainCfg):
|
||||
for j in range(num_y):
|
||||
base_size = avg_stone_size
|
||||
|
||||
# Proposed position with displacement.
|
||||
# Proposed position on the fixed grid with random displacement. Centers
|
||||
# span border to (size - border), so edge stones reach the borders.
|
||||
px = (
|
||||
offset_x + i * spacing + rng.uniform(-displacement_range, displacement_range)
|
||||
self.border_width
|
||||
+ i * pitch_x
|
||||
+ rng.uniform(-displacement_range, displacement_range)
|
||||
)
|
||||
py = (
|
||||
offset_y + j * spacing + rng.uniform(-displacement_range, displacement_range)
|
||||
self.border_width
|
||||
+ j * pitch_y
|
||||
+ rng.uniform(-displacement_range, displacement_range)
|
||||
)
|
||||
|
||||
# Randomized size.
|
||||
@@ -1237,10 +1306,11 @@ class BoxSteppingStonesTerrainCfg(SubTerrainCfg):
|
||||
x_min, x_max = px - size_x / 2, px + size_x / 2
|
||||
y_min, y_max = py - size_y / 2, py + size_y / 2
|
||||
|
||||
# Skip stones centered inside the platform. Stones whose edges
|
||||
# extend under the platform are kept; the platform covers the overlap.
|
||||
if (platform_min <= px <= platform_max) and (
|
||||
platform_min <= py <= platform_max
|
||||
# Drop stones whose center lies under the (grid-snapped) platform; the
|
||||
# platform absorbs them. Every remaining stone stays full size and sits
|
||||
# one gap from the platform, so there are no clipped slivers.
|
||||
if (platform_min_x <= px <= platform_max_x) and (
|
||||
platform_min_y <= py <= platform_max_y
|
||||
):
|
||||
continue
|
||||
|
||||
@@ -1296,6 +1366,7 @@ class BoxNarrowBeamsTerrainCfg(SubTerrainCfg):
|
||||
def function(
|
||||
self, difficulty: float, spec: mujoco.MjSpec, rng: np.random.Generator
|
||||
) -> TerrainOutput:
|
||||
del rng # Beam layout is deterministic.
|
||||
body = spec.body("terrain")
|
||||
geometries = []
|
||||
|
||||
@@ -1306,6 +1377,19 @@ class BoxNarrowBeamsTerrainCfg(SubTerrainCfg):
|
||||
w_min, w_max = self.beam_width_range
|
||||
beam_width = w_max - difficulty * (w_max - w_min)
|
||||
|
||||
# Shrink the square platform so its corners stay within the beams' angular
|
||||
# coverage rather than protruding into the pit between beams. A corner sits at
|
||||
# radius r*sqrt(2) and, in the worst case, pi/num_beams away from the nearest
|
||||
# beam, so it is covered while r*sqrt(2)*sin(pi/num_beams) <= beam_width/2.
|
||||
# Beams thin with difficulty, so the safe radius shrinks with it. The beams
|
||||
# attach at this same radius (below), so shrinking never opens a fall gap.
|
||||
spacing_sin = float(np.sin(np.pi / num_beams)) if num_beams > 1 else 0.0
|
||||
if spacing_sin > 1e-9:
|
||||
max_no_protrude = beam_width / (2.0 * np.sqrt(2.0) * spacing_sin)
|
||||
platform_radius = float(min(self.platform_width / 2.0, max_no_protrude))
|
||||
else:
|
||||
platform_radius = self.platform_width / 2.0
|
||||
|
||||
border_rgba = darken_rgba(brand_ramp(_MUJOCO_BLUE, 0.0), 0.85)
|
||||
z_center = (self.beam_height - self.floor_depth) / 2
|
||||
half_height = (self.beam_height + self.floor_depth) / 2
|
||||
@@ -1335,8 +1419,8 @@ class BoxNarrowBeamsTerrainCfg(SubTerrainCfg):
|
||||
platform_geom = body.add_geom(
|
||||
type=mujoco.mjtGeom.mjGEOM_BOX,
|
||||
size=(
|
||||
np.maximum(1e-6, self.platform_width / 2),
|
||||
np.maximum(1e-6, self.platform_width / 2),
|
||||
np.maximum(1e-6, platform_radius),
|
||||
np.maximum(1e-6, platform_radius),
|
||||
np.maximum(1e-6, half_height),
|
||||
),
|
||||
pos=(self.size[0] / 2, self.size[1] / 2, z_center),
|
||||
@@ -1347,7 +1431,6 @@ class BoxNarrowBeamsTerrainCfg(SubTerrainCfg):
|
||||
|
||||
inner_size = self.size[0] - 2 * self.border_width
|
||||
center_x, center_y = self.size[0] / 2, self.size[1] / 2
|
||||
platform_radius = self.platform_width / 2
|
||||
|
||||
# Radial beams as columns.
|
||||
angles = np.linspace(0, 2 * np.pi, num_beams, endpoint=False)
|
||||
@@ -1526,6 +1609,8 @@ class BoxNestedRingsTerrainCfg(SubTerrainCfg):
|
||||
ring_width_range: tuple[float, float] = (0.3, 0.6)
|
||||
gap_range: tuple[float, float] = (0.0, 0.2)
|
||||
height_range: tuple[float, float] = (0.1, 0.4)
|
||||
"""Min and max ring height, in meters. All rings share a single fixed height
|
||||
taken as the midpoint of this range; difficulty does not scale it."""
|
||||
platform_width: float = 1.0
|
||||
border_width: float = 0.25
|
||||
floor_depth: float = 2.0
|
||||
@@ -1533,20 +1618,25 @@ class BoxNestedRingsTerrainCfg(SubTerrainCfg):
|
||||
def function(
|
||||
self, difficulty: float, spec: mujoco.MjSpec, rng: np.random.Generator
|
||||
) -> TerrainOutput:
|
||||
del rng # Ring layout is deterministic.
|
||||
body = spec.body("terrain")
|
||||
geometries = []
|
||||
|
||||
# Difficulty scaling: wider width range and higher average height.
|
||||
h_scale = 1.0 + difficulty * 0.5
|
||||
# Concentric ridges of a single fixed height. Difficulty controls
|
||||
# gap-crossing only: gaps widen and rings narrow, so the terrain reads
|
||||
# consistently across difficulty instead of weakly scaling height.
|
||||
w_min, w_max = self.ring_width_range
|
||||
ring_width = w_max - difficulty * (w_max - w_min)
|
||||
|
||||
ring_height = 0.5 * (self.height_range[0] + self.height_range[1])
|
||||
ring_rgba = brand_ramp(_MUJOCO_BLUE, 0.6)
|
||||
|
||||
border_rgba = darken_rgba(brand_ramp(_MUJOCO_BLUE, 0.0), 0.85)
|
||||
# Use ground level z=0 as top of border/beams for consistency with NarrowBeams.
|
||||
# In beam terrain, border top was at beam_height.
|
||||
|
||||
if self.border_width > 0.0:
|
||||
border_h = 0.5
|
||||
# Outer border wall matches the ring height so there is no arbitrary
|
||||
# crossover between the two as difficulty changes.
|
||||
border_h = ring_height
|
||||
border_center = (
|
||||
0.5 * self.size[0],
|
||||
0.5 * self.size[1],
|
||||
@@ -1582,12 +1672,9 @@ class BoxNestedRingsTerrainCfg(SubTerrainCfg):
|
||||
gap_min, gap_max = self.gap_range
|
||||
gap = gap_min + difficulty * (gap_max - gap_min)
|
||||
|
||||
for k in range(self.num_rings):
|
||||
# Ring k: randomized height.
|
||||
h = rng.uniform(self.height_range[0], self.height_range[1]) * h_scale
|
||||
|
||||
t = k / max(self.num_rings - 1, 1)
|
||||
rgba = brand_ramp(_MUJOCO_BLUE, t)
|
||||
for _ in range(self.num_rings):
|
||||
h = ring_height
|
||||
rgba = ring_rgba
|
||||
|
||||
# Outer dimensions of this ring.
|
||||
ring_outer_size = (
|
||||
@@ -1655,7 +1742,8 @@ class BoxNestedRingsTerrainCfg(SubTerrainCfg):
|
||||
), # Fill the ring hole + gap area.
|
||||
np.maximum(1e-2, current_outer_size[1] + 2 * gap),
|
||||
)
|
||||
platform_h = 0.2
|
||||
# Center pad sits flush with the ring height.
|
||||
platform_h = ring_height
|
||||
|
||||
platform_half_h = (platform_h + self.floor_depth) / 2
|
||||
platform_z = (platform_h - self.floor_depth) / 2
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import abc
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Literal
|
||||
|
||||
@@ -194,16 +193,9 @@ class TerrainGenerator:
|
||||
body = spec.worldbody.add_body(name="terrain")
|
||||
|
||||
if self.cfg.curriculum:
|
||||
tic = time.perf_counter()
|
||||
self._generate_curriculum_terrains(spec)
|
||||
toc = time.perf_counter()
|
||||
print(f"Curriculum terrain generation took {toc - tic:.4f} seconds.")
|
||||
|
||||
else:
|
||||
tic = time.perf_counter()
|
||||
self._generate_random_terrains(spec)
|
||||
toc = time.perf_counter()
|
||||
print(f"Terrain generation took {toc - tic:.4f} seconds.")
|
||||
|
||||
self._add_terrain_border(spec)
|
||||
self._add_grid_lights(spec)
|
||||
@@ -257,11 +249,11 @@ class TerrainGenerator:
|
||||
# One column per terrain type — proportion is only for spawning.
|
||||
sub_terrains_cfgs = list(self.cfg.sub_terrains.values())
|
||||
|
||||
lower, upper = self.cfg.difficulty_range
|
||||
for sub_col in range(self._num_cols):
|
||||
for sub_row in range(self.cfg.num_rows):
|
||||
lower, upper = self.cfg.difficulty_range
|
||||
difficulty = (sub_row + self.np_rng.uniform()) / self.cfg.num_rows
|
||||
difficulty = lower + (upper - lower) * difficulty
|
||||
t = sub_row / max(self.cfg.num_rows - 1, 1)
|
||||
difficulty = lower + (upper - lower) * t
|
||||
world_position = self._get_sub_terrain_position(sub_row, sub_col)
|
||||
spawn_origin = self._create_terrain_geom(
|
||||
spec,
|
||||
|
||||
@@ -209,8 +209,9 @@ class CircularBuffer:
|
||||
|
||||
# Backfill entire history with first frame for newly initialized batches.
|
||||
is_first_push = self._num_pushes == 0
|
||||
if torch.any(is_first_push):
|
||||
self._buffer[:, is_first_push] = data[is_first_push]
|
||||
torch.where(
|
||||
is_first_push[None, :, None], data[None, :, :], self._buffer, out=self._buffer
|
||||
)
|
||||
|
||||
self._num_pushes += 1
|
||||
|
||||
@@ -236,8 +237,5 @@ class CircularBuffer:
|
||||
pushes = self._num_pushes.clamp_min(1)
|
||||
valid = torch.minimum(key, pushes - 1).clamp_min(0)
|
||||
|
||||
if torch.all(valid == 0):
|
||||
return self._buffer[self._pointer]
|
||||
|
||||
idx = torch.remainder(self._pointer - valid, self._max_len)
|
||||
return self._buffer[idx, self._all_indices]
|
||||
|
||||
@@ -3,10 +3,12 @@
|
||||
import os
|
||||
from typing import Literal
|
||||
|
||||
GpuId = int | str
|
||||
|
||||
|
||||
def select_gpus(
|
||||
gpu_ids: list[int] | Literal["all"] | None,
|
||||
) -> tuple[list[int] | None, int]:
|
||||
) -> tuple[list[GpuId] | None, int]:
|
||||
"""Select GPUs based on CUDA_VISIBLE_DEVICES and user specification.
|
||||
|
||||
This function treats the `gpu_ids` parameter as indices into the existing
|
||||
@@ -19,7 +21,8 @@ def select_gpus(
|
||||
|
||||
Returns:
|
||||
A tuple of (selected_gpu_ids, num_gpus) where:
|
||||
- selected_gpu_ids: List of physical GPU IDs to use, or None for CPU mode
|
||||
- selected_gpu_ids: List of physical GPU IDs (int for numeric, str for MIG
|
||||
UUIDs), or None for CPU mode
|
||||
- num_gpus: Number of GPUs selected (0 for CPU mode)
|
||||
|
||||
Examples:
|
||||
@@ -50,8 +53,11 @@ def select_gpus(
|
||||
|
||||
if existing_visible_devices is not None:
|
||||
# Parse existing CUDA_VISIBLE_DEVICES.
|
||||
available_gpus = [
|
||||
int(x.strip()) for x in existing_visible_devices.split(",") if x.strip()
|
||||
# Use int for numeric IDs, keep as string for MIG UUIDs.
|
||||
available_gpus: list[GpuId] = [
|
||||
int(x.strip()) if x.strip().isdigit() else x.strip()
|
||||
for x in existing_visible_devices.split(",")
|
||||
if x.strip()
|
||||
]
|
||||
# Empty CUDA_VISIBLE_DEVICES means CPU mode.
|
||||
if not available_gpus:
|
||||
@@ -60,15 +66,16 @@ def select_gpus(
|
||||
# If not set, default to all available GPUs.
|
||||
import torch.cuda
|
||||
|
||||
available_gpus = list(range(torch.cuda.device_count()))
|
||||
available_gpus: list[GpuId] = list(range(torch.cuda.device_count()))
|
||||
|
||||
# Map gpu_ids indices to actual GPU IDs.
|
||||
selected: list[GpuId]
|
||||
if gpu_ids == "all":
|
||||
selected_gpus = available_gpus
|
||||
selected = available_gpus
|
||||
else:
|
||||
# gpu_ids are indices into available_gpus.
|
||||
selected_gpus = [available_gpus[i] for i in gpu_ids]
|
||||
selected = [available_gpus[i] for i in gpu_ids]
|
||||
|
||||
num_gpus = len(selected_gpus)
|
||||
num_gpus = len(selected)
|
||||
|
||||
return selected_gpus, num_gpus
|
||||
return selected, num_gpus
|
||||
|
||||
@@ -6,18 +6,9 @@ import torch
|
||||
import warp as wp
|
||||
|
||||
|
||||
def seed_rng(
|
||||
seed: int,
|
||||
torch_deterministic: bool = False,
|
||||
device: str | torch.device | None = None,
|
||||
) -> None:
|
||||
def seed_rng(seed: int, torch_deterministic: bool = False) -> None:
|
||||
"""Seed all random number generators for reproducibility.
|
||||
|
||||
When ``device`` is a CPU device, ``wp.rand_init`` is skipped so that Warp's
|
||||
CUDA runtime is not initialized on machines where a GPU is visible but the
|
||||
caller has explicitly opted into CPU-only execution. When ``device`` is
|
||||
``None``, behavior is unchanged (Warp is seeded).
|
||||
|
||||
Note: MuJoCo Warp is not fully deterministic yet.
|
||||
See: https://github.com/google-deepmind/mujoco_warp/issues/562
|
||||
"""
|
||||
@@ -26,8 +17,7 @@ def seed_rng(
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
|
||||
if device is None or torch.device(device).type != "cpu":
|
||||
wp.rand_init(wp.int32(seed))
|
||||
wp.rand_init(wp.int32(seed))
|
||||
|
||||
# Ref: https://docs.pytorch.org/docs/stable/notes/randomness.html
|
||||
torch.manual_seed(seed) # Seed RNG for all devices.
|
||||
|
||||
@@ -120,6 +120,34 @@ _TRANSMISSION_TYPE_MAP = {
|
||||
}
|
||||
|
||||
|
||||
def apply_target_overrides(
|
||||
spec: mujoco.MjSpec,
|
||||
target_name: str,
|
||||
transmission_type: TransmissionType,
|
||||
*,
|
||||
armature: float | None,
|
||||
frictionloss: float | None,
|
||||
viscous_damping: float | None,
|
||||
) -> None:
|
||||
"""Apply joint- or tendon-level overrides. ``None`` preserves the XML value.
|
||||
|
||||
SITE transmission is a no-op (sites have no armature / frictionloss / damping);
|
||||
callers using SITE should not pass non-None overrides.
|
||||
"""
|
||||
if transmission_type == TransmissionType.JOINT:
|
||||
target = spec.joint(target_name)
|
||||
elif transmission_type == TransmissionType.TENDON:
|
||||
target = spec.tendon(target_name)
|
||||
else:
|
||||
return
|
||||
if armature is not None:
|
||||
target.armature = armature
|
||||
if frictionloss is not None:
|
||||
target.frictionloss = frictionloss
|
||||
if viscous_damping is not None:
|
||||
target.damping[0] = viscous_damping
|
||||
|
||||
|
||||
def auto_wrap_fixed_base_mocap(
|
||||
spec_fn: Callable[[], mujoco.MjSpec],
|
||||
) -> Callable[[], mujoco.MjSpec]:
|
||||
@@ -235,21 +263,14 @@ def create_motor_actuator(
|
||||
actuator.ctrllimited = True
|
||||
actuator.ctrlrange[:] = np.array([-effort_limit, effort_limit])
|
||||
|
||||
# Set armature, frictionloss, and viscous_damping (None = preserve XML value).
|
||||
if transmission_type == TransmissionType.JOINT:
|
||||
if armature is not None:
|
||||
spec.joint(joint_name).armature = armature
|
||||
if frictionloss is not None:
|
||||
spec.joint(joint_name).frictionloss = frictionloss
|
||||
if viscous_damping is not None:
|
||||
spec.joint(joint_name).damping[0] = viscous_damping
|
||||
elif transmission_type == TransmissionType.TENDON:
|
||||
if armature is not None:
|
||||
spec.tendon(joint_name).armature = armature
|
||||
if frictionloss is not None:
|
||||
spec.tendon(joint_name).frictionloss = frictionloss
|
||||
if viscous_damping is not None:
|
||||
spec.tendon(joint_name).damping[0] = viscous_damping
|
||||
apply_target_overrides(
|
||||
spec,
|
||||
joint_name,
|
||||
transmission_type,
|
||||
armature=armature,
|
||||
frictionloss=frictionloss,
|
||||
viscous_damping=viscous_damping,
|
||||
)
|
||||
|
||||
return actuator
|
||||
|
||||
@@ -265,14 +286,21 @@ def create_position_actuator(
|
||||
frictionloss: float | None = None,
|
||||
viscous_damping: float | None = None,
|
||||
transmission_type: TransmissionType = TransmissionType.JOINT,
|
||||
actuator_name: str | None = None,
|
||||
) -> mujoco.MjsActuator:
|
||||
"""Creates a <position> actuator.
|
||||
|
||||
An important note about this actuator is that we set `ctrllimited` to False. This is
|
||||
because we want to allow the policy to output setpoints that are outside the kinematic
|
||||
limits of the joint.
|
||||
|
||||
``actuator_name`` defaults to ``joint_name``; pass a distinct value when multiple
|
||||
actuators target the same joint (e.g. paired position+velocity elements).
|
||||
"""
|
||||
actuator = spec.add_actuator(name=joint_name, target=joint_name)
|
||||
actuator = spec.add_actuator(
|
||||
name=actuator_name if actuator_name is not None else joint_name,
|
||||
target=joint_name,
|
||||
)
|
||||
|
||||
actuator.trntype = _TRANSMISSION_TYPE_MAP[transmission_type]
|
||||
actuator.dyntype = mujoco.mjtDyn.mjDYN_NONE
|
||||
@@ -314,21 +342,14 @@ def create_position_actuator(
|
||||
actuator.forcelimited = False
|
||||
# No forcerange needed.
|
||||
|
||||
# Set armature, frictionloss, and viscous_damping (None = preserve XML value).
|
||||
if transmission_type == TransmissionType.JOINT:
|
||||
if armature is not None:
|
||||
spec.joint(joint_name).armature = armature
|
||||
if frictionloss is not None:
|
||||
spec.joint(joint_name).frictionloss = frictionloss
|
||||
if viscous_damping is not None:
|
||||
spec.joint(joint_name).damping[0] = viscous_damping
|
||||
elif transmission_type == TransmissionType.TENDON:
|
||||
if armature is not None:
|
||||
spec.tendon(joint_name).armature = armature
|
||||
if frictionloss is not None:
|
||||
spec.tendon(joint_name).frictionloss = frictionloss
|
||||
if viscous_damping is not None:
|
||||
spec.tendon(joint_name).damping[0] = viscous_damping
|
||||
apply_target_overrides(
|
||||
spec,
|
||||
joint_name,
|
||||
transmission_type,
|
||||
armature=armature,
|
||||
frictionloss=frictionloss,
|
||||
viscous_damping=viscous_damping,
|
||||
)
|
||||
|
||||
return actuator
|
||||
|
||||
@@ -343,14 +364,21 @@ def create_velocity_actuator(
|
||||
frictionloss: float | None = None,
|
||||
viscous_damping: float | None = None,
|
||||
transmission_type: TransmissionType = TransmissionType.JOINT,
|
||||
actuator_name: str | None = None,
|
||||
) -> mujoco.MjsActuator:
|
||||
"""Creates a <velocity> actuator.
|
||||
|
||||
Control inputs are not clamped so that velocity commands work for any joint,
|
||||
including continuous joints that have no range defined. Force output is still
|
||||
bounded when effort_limit is set.
|
||||
|
||||
``actuator_name`` defaults to ``joint_name``; pass a distinct value when multiple
|
||||
actuators target the same joint (e.g. paired position+velocity elements).
|
||||
"""
|
||||
actuator = spec.add_actuator(name=joint_name, target=joint_name)
|
||||
actuator = spec.add_actuator(
|
||||
name=actuator_name if actuator_name is not None else joint_name,
|
||||
target=joint_name,
|
||||
)
|
||||
|
||||
actuator.trntype = _TRANSMISSION_TYPE_MAP[transmission_type]
|
||||
actuator.dyntype = mujoco.mjtDyn.mjDYN_NONE
|
||||
@@ -369,21 +397,14 @@ def create_velocity_actuator(
|
||||
else:
|
||||
actuator.forcelimited = False
|
||||
|
||||
# Set armature, frictionloss, and viscous_damping (None = preserve XML value).
|
||||
if transmission_type == TransmissionType.JOINT:
|
||||
if armature is not None:
|
||||
spec.joint(joint_name).armature = armature
|
||||
if frictionloss is not None:
|
||||
spec.joint(joint_name).frictionloss = frictionloss
|
||||
if viscous_damping is not None:
|
||||
spec.joint(joint_name).damping[0] = viscous_damping
|
||||
elif transmission_type == TransmissionType.TENDON:
|
||||
if armature is not None:
|
||||
spec.tendon(joint_name).armature = armature
|
||||
if frictionloss is not None:
|
||||
spec.tendon(joint_name).frictionloss = frictionloss
|
||||
if viscous_damping is not None:
|
||||
spec.tendon(joint_name).damping[0] = viscous_damping
|
||||
apply_target_overrides(
|
||||
spec,
|
||||
joint_name,
|
||||
transmission_type,
|
||||
armature=armature,
|
||||
frictionloss=frictionloss,
|
||||
viscous_damping=viscous_damping,
|
||||
)
|
||||
|
||||
return actuator
|
||||
|
||||
@@ -467,54 +488,52 @@ def copy_mesh_data(src: mujoco.MjsMesh, dst: mujoco.MjsMesh) -> None:
|
||||
dst.smoothnormal = src.smoothnormal
|
||||
|
||||
|
||||
def validate_variant_structure(
|
||||
names: list[str],
|
||||
bodies: list[mujoco.MjsBody],
|
||||
) -> None:
|
||||
"""Validate that variant specs share the same kinematic structure.
|
||||
def copy_texture_data(src: mujoco.MjsTexture, dst: mujoco.MjsTexture) -> None:
|
||||
"""Copy texture data from *src* to *dst*.
|
||||
|
||||
Checks that all variants have the same number of child bodies, the same number of
|
||||
joints, the same joint types, and the same joint names. Raises ``ValueError`` with a
|
||||
descriptive message if any differ.
|
||||
Copies the file path or builtin/data fields, format, dimensions, and color
|
||||
settings. The ``name`` field is NOT copied; set it on *dst* before calling.
|
||||
"""
|
||||
ref_name = names[0]
|
||||
ref_body = bodies[0]
|
||||
ref_joints = list(ref_body.joints)
|
||||
ref_joint_types = [j.type for j in ref_joints]
|
||||
ref_joint_names = [j.name for j in ref_joints]
|
||||
ref_sub_bodies = list(ref_body.bodies)
|
||||
assert dst.name, "dst.name must be set before copy_texture_data."
|
||||
dst.type = src.type
|
||||
dst.colorspace = src.colorspace
|
||||
dst.builtin = src.builtin
|
||||
dst.mark = src.mark
|
||||
dst.rgb1[:] = src.rgb1
|
||||
dst.rgb2[:] = src.rgb2
|
||||
dst.markrgb[:] = src.markrgb
|
||||
dst.random = src.random
|
||||
dst.gridsize[:] = src.gridsize
|
||||
dst.gridlayout = src.gridlayout
|
||||
dst.width = src.width
|
||||
dst.height = src.height
|
||||
dst.nchannel = src.nchannel
|
||||
dst.hflip = src.hflip
|
||||
dst.vflip = src.vflip
|
||||
if src.file:
|
||||
dst.file = src.file
|
||||
if len(src.cubefiles) > 0:
|
||||
dst.cubefiles = src.cubefiles
|
||||
if len(src.data) > 0:
|
||||
dst.data = src.data
|
||||
if src.content_type:
|
||||
dst.content_type = src.content_type
|
||||
|
||||
for i in range(1, len(names)):
|
||||
other_name = names[i]
|
||||
other_body = bodies[i]
|
||||
|
||||
other_sub_bodies = list(other_body.bodies)
|
||||
if len(other_sub_bodies) != len(ref_sub_bodies):
|
||||
raise ValueError(
|
||||
f"Variant '{other_name}' has {len(other_sub_bodies)} "
|
||||
f"child bodies, but '{ref_name}' has "
|
||||
f"{len(ref_sub_bodies)}."
|
||||
)
|
||||
def copy_material_data(src: mujoco.MjsMaterial, dst: mujoco.MjsMaterial) -> None:
|
||||
"""Copy material data from *src* to *dst*.
|
||||
|
||||
other_joints = list(other_body.joints)
|
||||
if len(other_joints) != len(ref_joints):
|
||||
raise ValueError(
|
||||
f"Variant '{other_name}' has {len(other_joints)} "
|
||||
f"joints, but '{ref_name}' has {len(ref_joints)}."
|
||||
)
|
||||
|
||||
other_joint_types = [j.type for j in other_joints]
|
||||
if other_joint_types != ref_joint_types:
|
||||
raise ValueError(
|
||||
f"Variant '{other_name}' has joint types "
|
||||
f"{other_joint_types}, but '{ref_name}' has "
|
||||
f"{ref_joint_types}."
|
||||
)
|
||||
|
||||
other_joint_names = [j.name for j in other_joints]
|
||||
if other_joint_names != ref_joint_names:
|
||||
raise ValueError(
|
||||
f"Variant '{other_name}' has joint names "
|
||||
f"{other_joint_names}, but '{ref_name}' has "
|
||||
f"{ref_joint_names}."
|
||||
)
|
||||
Copies appearance settings (rgba, specular, shininess, ...) and texture
|
||||
bindings. The ``name`` field is NOT copied; set it on *dst* before calling.
|
||||
"""
|
||||
assert dst.name, "dst.name must be set before copy_material_data."
|
||||
dst.rgba[:] = src.rgba
|
||||
dst.emission = src.emission
|
||||
dst.specular = src.specular
|
||||
dst.shininess = src.shininess
|
||||
dst.reflectance = src.reflectance
|
||||
dst.roughness = src.roughness
|
||||
dst.metallic = src.metallic
|
||||
dst.texuniform = src.texuniform
|
||||
dst.texrepeat[:] = src.texrepeat
|
||||
dst.textures = list(src.textures)
|
||||
|
||||
@@ -179,6 +179,26 @@ class DebugVisualizer(ABC):
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def add_box(
|
||||
self,
|
||||
center: np.ndarray,
|
||||
size: np.ndarray,
|
||||
mat: np.ndarray,
|
||||
color: tuple[float, float, float, float],
|
||||
label: str | None = None,
|
||||
) -> None:
|
||||
"""Add an axis-oriented box visualization.
|
||||
|
||||
Args:
|
||||
center: Center position (3D vector).
|
||||
size: Half-extents along each local axis (3D vector: a, b, c).
|
||||
mat: 3x3 rotation matrix (or flattened 9-element array).
|
||||
color: RGBA color (values 0-1).
|
||||
label: Optional label for this box.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def clear(self) -> None:
|
||||
"""Clear all debug visualizations."""
|
||||
@@ -242,5 +262,8 @@ class NullDebugVisualizer:
|
||||
def add_ellipsoid(self, center, size, mat, color, label=None) -> None:
|
||||
pass
|
||||
|
||||
def add_box(self, center, size, mat, color, label=None) -> None:
|
||||
pass
|
||||
|
||||
def clear(self) -> None:
|
||||
pass
|
||||
|
||||
@@ -25,6 +25,7 @@ VIEWER_MODEL_FIELDS = frozenset(
|
||||
{
|
||||
"qpos0", # Needed for correct mj_forward kinematics (qpos - qpos0).
|
||||
"geom_dataid", # Per-world mesh variants.
|
||||
"geom_matid", # Per-world material variants.
|
||||
"geom_rgba",
|
||||
"geom_size",
|
||||
"geom_pos",
|
||||
|
||||
@@ -238,6 +238,31 @@ class MujocoNativeDebugVisualizer(DebugVisualizer):
|
||||
rgba=np.asarray(color, dtype=np.float32),
|
||||
)
|
||||
|
||||
@override
|
||||
def add_box(
|
||||
self,
|
||||
center: np.ndarray,
|
||||
size: np.ndarray,
|
||||
mat: np.ndarray,
|
||||
color: tuple[float, float, float, float],
|
||||
label: str | None = None,
|
||||
) -> None:
|
||||
"""Add a box visualization using MuJoCo's box geometry."""
|
||||
del label # Unused.
|
||||
|
||||
self.scn.ngeom += 1
|
||||
geom = self.scn.geoms[self.scn.ngeom - 1]
|
||||
geom.category = mujoco.mjtCatBit.mjCAT_DECOR
|
||||
|
||||
mujoco.mjv_initGeom(
|
||||
geom=geom,
|
||||
type=mujoco.mjtGeom.mjGEOM_BOX.value,
|
||||
size=np.asarray(size, dtype=np.float64),
|
||||
pos=np.asarray(center, dtype=np.float64),
|
||||
mat=np.asarray(mat, dtype=np.float64).flatten(),
|
||||
rgba=np.asarray(color, dtype=np.float32),
|
||||
)
|
||||
|
||||
@override
|
||||
def clear(self) -> None:
|
||||
"""Clear debug visualizations by resetting geom count."""
|
||||
|
||||
@@ -57,6 +57,7 @@ class OffscreenRenderer:
|
||||
self._opt = mujoco.MjvOption()
|
||||
self._pert = mujoco.MjvPerturb()
|
||||
self._catmask = mujoco.mjtCatBit.mjCAT_DYNAMIC
|
||||
self._extra_env_ids: list[int] | None = None
|
||||
|
||||
@property
|
||||
def renderer(self) -> mujoco.Renderer:
|
||||
@@ -134,9 +135,17 @@ class OffscreenRenderer:
|
||||
|
||||
We render a small local neighborhood around ``env_idx`` instead of the first
|
||||
N environments, so videos stay focused on the tracked robot and nearby peers.
|
||||
|
||||
The neighbor set is computed once and cached. ``env_origins`` can mutate during
|
||||
training (e.g. the terrain curriculum reassigns origins on reset), so recomputing
|
||||
every frame would make the context robots pop in and out, causing video flicker.
|
||||
"""
|
||||
if self._extra_env_ids is not None:
|
||||
return self._extra_env_ids
|
||||
|
||||
if self._cfg.max_extra_envs <= 0 or nworld <= 1:
|
||||
return []
|
||||
self._extra_env_ids = []
|
||||
return self._extra_env_ids
|
||||
|
||||
k = min(self._cfg.max_extra_envs, nworld - 1)
|
||||
origins = self._scene.env_origins[:nworld].cpu().numpy()
|
||||
@@ -146,7 +155,8 @@ class OffscreenRenderer:
|
||||
|
||||
nearest = np.argpartition(dist2, kth=k - 1)[:k]
|
||||
nearest = nearest[np.argsort(dist2[nearest])]
|
||||
return [int(i) for i in nearest]
|
||||
self._extra_env_ids = [int(i) for i in nearest]
|
||||
return self._extra_env_ids
|
||||
|
||||
def _sync_model_fields(self, env_idx: int) -> None:
|
||||
"""Sync visually relevant per-world model fields into the host MjModel."""
|
||||
|
||||
@@ -263,6 +263,7 @@ class MjlabViserScene(ViserMujocoScene, DebugVisualizer):
|
||||
self._queued_spheres: list = []
|
||||
self._queued_cylinders: list = []
|
||||
self._queued_ellipsoids: list = []
|
||||
self._queued_boxes: list = []
|
||||
|
||||
# Batched mesh handles for simple primitives.
|
||||
def _shaft_mesh() -> trimesh.Trimesh:
|
||||
@@ -287,12 +288,19 @@ class MjlabViserScene(ViserMujocoScene, DebugVisualizer):
|
||||
"ellipsoids",
|
||||
lambda: trimesh.creation.icosphere(subdivisions=2, radius=1.0),
|
||||
)
|
||||
# Unit half-extents so that scaling by the box size yields the requested
|
||||
# half-extents (extents=2 spans -1 to 1 along each axis).
|
||||
self._boxes = _BatchedPrimitive(
|
||||
"boxes",
|
||||
lambda: trimesh.creation.box(extents=(2.0, 2.0, 2.0)),
|
||||
)
|
||||
self._all_primitives = [
|
||||
self._arrow_shafts,
|
||||
self._arrow_heads,
|
||||
self._spheres,
|
||||
self._cylinders,
|
||||
self._ellipsoids,
|
||||
self._boxes,
|
||||
]
|
||||
|
||||
# Ghost mesh state.
|
||||
@@ -955,6 +963,27 @@ class MjlabViserScene(ViserMujocoScene, DebugVisualizer):
|
||||
)
|
||||
)
|
||||
|
||||
@override
|
||||
def add_box(
|
||||
self,
|
||||
center: np.ndarray | torch.Tensor,
|
||||
size: np.ndarray | torch.Tensor,
|
||||
mat: np.ndarray | torch.Tensor,
|
||||
color: tuple[float, float, float, float],
|
||||
label: str | None = None,
|
||||
) -> None:
|
||||
if not self.debug_visualization_enabled:
|
||||
return
|
||||
del label
|
||||
self._queued_boxes.append(
|
||||
(
|
||||
np.asarray(_to_numpy(center), dtype=np.float32).copy(),
|
||||
np.asarray(_to_numpy(size), dtype=np.float32).copy(),
|
||||
np.asarray(_to_numpy(mat), dtype=np.float32).reshape(3, 3).copy(),
|
||||
color,
|
||||
)
|
||||
)
|
||||
|
||||
@override
|
||||
def clear(self) -> None:
|
||||
"""Clear all debug visualization queues."""
|
||||
@@ -962,6 +991,7 @@ class MjlabViserScene(ViserMujocoScene, DebugVisualizer):
|
||||
self._queued_spheres.clear()
|
||||
self._queued_cylinders.clear()
|
||||
self._queued_ellipsoids.clear()
|
||||
self._queued_boxes.clear()
|
||||
self._queued_ghosts.clear()
|
||||
|
||||
def clear_debug_all(self) -> None:
|
||||
@@ -1039,6 +1069,7 @@ class MjlabViserScene(ViserMujocoScene, DebugVisualizer):
|
||||
self._sync_spheres()
|
||||
self._sync_cylinders()
|
||||
self._sync_ellipsoids()
|
||||
self._sync_boxes()
|
||||
|
||||
def _sync_spheres(self) -> None:
|
||||
if not self._queued_spheres:
|
||||
@@ -1125,6 +1156,32 @@ class MjlabViserScene(ViserMujocoScene, DebugVisualizer):
|
||||
opacity,
|
||||
)
|
||||
|
||||
def _sync_boxes(self) -> None:
|
||||
if not self._queued_boxes:
|
||||
self._boxes.remove()
|
||||
return
|
||||
n = len(self._queued_boxes)
|
||||
positions = np.zeros((n, 3), dtype=np.float32)
|
||||
wxyzs = np.zeros((n, 4), dtype=np.float32)
|
||||
scales = np.zeros((n, 3), dtype=np.float32)
|
||||
colors = np.zeros((n, 3), dtype=np.uint8)
|
||||
opacity = 1.0
|
||||
for i, (center, size, mat, color) in enumerate(self._queued_boxes):
|
||||
positions[i] = center + self._scene_offset
|
||||
wxyzs[i] = vtf.SO3.from_matrix(mat).wxyz
|
||||
scales[i] = size
|
||||
colors[i] = _color_uint8(color)
|
||||
opacity = color[3]
|
||||
self._boxes.sync(
|
||||
self.server,
|
||||
self.env_idx,
|
||||
positions,
|
||||
wxyzs,
|
||||
scales,
|
||||
colors,
|
||||
opacity,
|
||||
)
|
||||
|
||||
def _sync_ghosts(self) -> None:
|
||||
"""Render queued ghosts as one batched handle per (model, body)."""
|
||||
if not self._queued_ghosts:
|
||||
|
||||
@@ -106,17 +106,25 @@ def initialize_entity(entity: Entity, device: str, num_envs: int = 1):
|
||||
|
||||
def make_scene_and_sim(
|
||||
device: str,
|
||||
xml: str,
|
||||
xml: str | dict[str, str],
|
||||
sensors: tuple,
|
||||
num_envs: int = 1,
|
||||
sim_cfg: SimulationCfg | None = None,
|
||||
) -> tuple[Scene, Simulation]:
|
||||
"""Create a scene and simulation from inline XML with sensors wired up."""
|
||||
entity_cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(xml))
|
||||
"""Create a scene and simulation from inline XML with sensors wired up.
|
||||
|
||||
``xml`` may be a single XML string (registered as the ``robot`` entity) or a
|
||||
mapping of entity name to XML string for multi-entity scenes.
|
||||
"""
|
||||
xml_by_entity = {"robot": xml} if isinstance(xml, str) else xml
|
||||
entities = {
|
||||
name: EntityCfg(spec_fn=lambda s=s: mujoco.MjSpec.from_string(s))
|
||||
for name, s in xml_by_entity.items()
|
||||
}
|
||||
scene_cfg = SceneCfg(
|
||||
num_envs=num_envs,
|
||||
env_spacing=5.0,
|
||||
entities={"robot": entity_cfg},
|
||||
entities=entities,
|
||||
sensors=sensors,
|
||||
)
|
||||
scene = Scene(scene_cfg, device)
|
||||
|
||||
@@ -0,0 +1,653 @@
|
||||
"""Tests for BuiltinDcMotorActuator.
|
||||
|
||||
Covers wiring of MuJoCo's native ``<dcmotor>`` element through mjlab: the
|
||||
three input modes (voltage / position / velocity), torque saturation,
|
||||
config validation, and DR integration.
|
||||
"""
|
||||
|
||||
import math
|
||||
from unittest.mock import Mock
|
||||
|
||||
import mujoco
|
||||
import pytest
|
||||
import torch
|
||||
from conftest import (
|
||||
create_entity_with_actuator,
|
||||
get_test_device,
|
||||
initialize_entity,
|
||||
load_fixture_xml,
|
||||
)
|
||||
|
||||
from mjlab.actuator import (
|
||||
BuiltinDcMotorActuator,
|
||||
BuiltinDcMotorActuatorCfg,
|
||||
DcMotorDatasheetParams,
|
||||
DcMotorInputMode,
|
||||
DcMotorPhysicalParams,
|
||||
)
|
||||
from mjlab.actuator.actuator import TransmissionType
|
||||
from mjlab.entity import Entity, EntityArticulationInfoCfg, EntityCfg
|
||||
from mjlab.envs.mdp import dr
|
||||
from mjlab.managers.scene_entity_config import SceneEntityCfg
|
||||
from mjlab.scene import Scene, SceneCfg
|
||||
from mjlab.sim.sim import Simulation, SimulationCfg
|
||||
|
||||
ROBOT_XML = load_fixture_xml("floating_base_articulated")
|
||||
|
||||
# Motor characterization used throughout (resolves to K=0.24, R=2.88).
|
||||
V_NOM, TAU_STALL, OMEGA_NL = 24.0, 2.0, 100.0
|
||||
K = V_NOM / OMEGA_NL
|
||||
R = K * V_NOM / TAU_STALL
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def device():
|
||||
return get_test_device()
|
||||
|
||||
|
||||
DATASHEET = DcMotorDatasheetParams(
|
||||
nominal_voltage=V_NOM, stall_torque=TAU_STALL, no_load_speed=OMEGA_NL
|
||||
)
|
||||
|
||||
|
||||
def _make_cfg(
|
||||
*,
|
||||
mode: DcMotorInputMode = DcMotorInputMode.POSITION,
|
||||
stiffness=5.0,
|
||||
damping=0.5,
|
||||
voltage_limit=24.0,
|
||||
**extra,
|
||||
) -> BuiltinDcMotorActuatorCfg:
|
||||
"""Build a cfg with sensible PID defaults. ``extra`` forwards any other
|
||||
BuiltinDcMotorActuatorCfg kwarg (effort_limit, integral_gain, thermal,
|
||||
delay_*, etc.)."""
|
||||
return BuiltinDcMotorActuatorCfg(
|
||||
target_names_expr=("joint.*",),
|
||||
mode=mode,
|
||||
motor_params=DATASHEET,
|
||||
stiffness=stiffness,
|
||||
damping=damping,
|
||||
voltage_limit=voltage_limit,
|
||||
**extra,
|
||||
)
|
||||
|
||||
|
||||
def _make_entity(**kwargs) -> Entity:
|
||||
return create_entity_with_actuator(ROBOT_XML, _make_cfg(**kwargs))
|
||||
|
||||
|
||||
def _make_initialized(device, **kwargs):
|
||||
"""Build entity from cfg kwargs and initialize it through the sim."""
|
||||
return initialize_entity(_make_entity(**kwargs), device)
|
||||
|
||||
|
||||
def _drive(
|
||||
entity: Entity,
|
||||
sim,
|
||||
device: str,
|
||||
*,
|
||||
pos_target=None,
|
||||
vel_target=None,
|
||||
effort_target=None,
|
||||
q0=None,
|
||||
qd0=None,
|
||||
) -> None:
|
||||
zero = torch.zeros(1, 2, device=device)
|
||||
entity.write_joint_state_to_sim(
|
||||
position=q0 if q0 is not None else zero,
|
||||
velocity=qd0 if qd0 is not None else zero,
|
||||
)
|
||||
entity.set_joint_position_target(pos_target if pos_target is not None else zero)
|
||||
entity.set_joint_velocity_target(vel_target if vel_target is not None else zero)
|
||||
entity.set_joint_effort_target(effort_target if effort_target is not None else zero)
|
||||
entity.write_data_to_sim()
|
||||
sim.forward()
|
||||
|
||||
|
||||
# Wiring sanity.
|
||||
|
||||
|
||||
def test_kr_packed_into_gainprm(device):
|
||||
"""The XML compiler derives K and R from the nominal triplet."""
|
||||
_, sim = initialize_entity(_make_entity(effort_limit=1.5), device)
|
||||
m = sim.mj_model
|
||||
for i in range(2):
|
||||
assert m.actuator_gainprm[i, 0] == pytest.approx(R, abs=1e-6)
|
||||
assert m.actuator_gainprm[i, 1] == pytest.approx(K, abs=1e-6)
|
||||
assert m.actuator_gainprm[i, 4] == pytest.approx(5.0) # kp
|
||||
assert m.actuator_gainprm[i, 6] == pytest.approx(0.5) # kd
|
||||
assert m.actuator_gainprm[i, 7] == pytest.approx(24.0) # Vmax
|
||||
assert m.actuator_gainprm[i, 8] == pytest.approx(1.0) # input_mode=position
|
||||
assert m.actuator_gaintype[i] == mujoco.mjtGain.mjGAIN_DCMOTOR
|
||||
assert m.actuator_biastype[i] == mujoco.mjtBias.mjBIAS_DCMOTOR
|
||||
# No activation state: ki=0, no inductance, no thermal/lugre/slew.
|
||||
assert m.actuator_actnum[i] == 0
|
||||
|
||||
|
||||
def test_motor_const_path(device):
|
||||
"""Physical params pack K = sqrt(Kt*Ke) and R verbatim."""
|
||||
cfg = BuiltinDcMotorActuatorCfg(
|
||||
target_names_expr=("joint.*",),
|
||||
mode=DcMotorInputMode.VOLTAGE,
|
||||
motor_params=DcMotorPhysicalParams(kt=0.1, ke=0.05, resistance=2.0),
|
||||
)
|
||||
_, sim = initialize_entity(create_entity_with_actuator(ROBOT_XML, cfg), device)
|
||||
m = sim.mj_model
|
||||
for i in range(2):
|
||||
assert m.actuator_gainprm[i, 0] == pytest.approx(2.0, abs=1e-6)
|
||||
assert m.actuator_gainprm[i, 1] == pytest.approx((0.1 * 0.05) ** 0.5, abs=1e-6)
|
||||
|
||||
|
||||
# Stateless motor physics.
|
||||
|
||||
|
||||
def test_voltage_mode_steady_state(device):
|
||||
"""At rest, ctrl = V -> tau = K * V / R."""
|
||||
cfg = BuiltinDcMotorActuatorCfg(
|
||||
target_names_expr=("joint.*",),
|
||||
mode=DcMotorInputMode.VOLTAGE,
|
||||
motor_params=DATASHEET,
|
||||
)
|
||||
entity, sim = initialize_entity(create_entity_with_actuator(ROBOT_XML, cfg), device)
|
||||
V = torch.tensor([[10.0, -5.0]], device=device)
|
||||
_drive(entity, sim, device, effort_target=V)
|
||||
v_adr = entity.indexing.joint_v_adr
|
||||
expected = K * V[0] / R
|
||||
assert torch.allclose(sim.data.qfrc_actuator[0, v_adr], expected, atol=1e-4)
|
||||
|
||||
|
||||
def test_voltage_mode_voltage_limit_zero_is_noop(device):
|
||||
"""Docstring promises ``voltage_limit=0`` disables clamping. Verify against
|
||||
MuJoCo's ``dcmotor_voltage`` (which only clamps when ``Vmax > 0``)."""
|
||||
cfg = BuiltinDcMotorActuatorCfg(
|
||||
target_names_expr=("joint.*",),
|
||||
mode=DcMotorInputMode.VOLTAGE,
|
||||
motor_params=DATASHEET,
|
||||
voltage_limit=0.0,
|
||||
)
|
||||
entity, sim = initialize_entity(create_entity_with_actuator(ROBOT_XML, cfg), device)
|
||||
V = torch.tensor([[1000.0, 0.0]], device=device) # absurdly high voltage.
|
||||
_drive(entity, sim, device, effort_target=V)
|
||||
v_adr = entity.indexing.joint_v_adr
|
||||
expected = K * V[0] / R
|
||||
assert torch.allclose(sim.data.qfrc_actuator[0, v_adr], expected, atol=1e-2)
|
||||
|
||||
|
||||
def test_back_emf_reduces_torque_at_velocity(device):
|
||||
"""Same V, joint moving at omega: tau = K * (V - K * omega) / R."""
|
||||
cfg = BuiltinDcMotorActuatorCfg(
|
||||
target_names_expr=("joint.*",),
|
||||
mode=DcMotorInputMode.VOLTAGE,
|
||||
motor_params=DATASHEET,
|
||||
)
|
||||
entity, sim = initialize_entity(create_entity_with_actuator(ROBOT_XML, cfg), device)
|
||||
V = torch.tensor([[10.0, 0.0]], device=device)
|
||||
omega0 = torch.tensor([[2.0, 0.0]], device=device)
|
||||
_drive(entity, sim, device, effort_target=V, qd0=omega0)
|
||||
v_adr = entity.indexing.joint_v_adr
|
||||
expected = K * (V[0] - K * omega0[0]) / R
|
||||
assert torch.allclose(sim.data.qfrc_actuator[0, v_adr], expected, atol=1e-4)
|
||||
|
||||
|
||||
def test_position_mode_pid_at_rest(device):
|
||||
"""kd=0, no Vmax clamp: tau = K * kp * (target - q) / R."""
|
||||
# voltage_limit must be >0 (cfg invariant), pick it big enough not to clamp.
|
||||
entity, sim = initialize_entity(
|
||||
_make_entity(damping=0.0, voltage_limit=1000.0), device
|
||||
)
|
||||
pos = torch.tensor([[0.1, -0.05]], device=device)
|
||||
_drive(entity, sim, device, pos_target=pos)
|
||||
v_adr = entity.indexing.joint_v_adr
|
||||
expected = K * 5.0 * pos[0] / R
|
||||
assert torch.allclose(sim.data.qfrc_actuator[0, v_adr], expected, atol=1e-4)
|
||||
|
||||
|
||||
def test_position_mode_voltage_clamp(device):
|
||||
"""Huge position error -> PID voltage saturates at Vmax."""
|
||||
entity, sim = initialize_entity(
|
||||
_make_entity(stiffness=100.0, damping=0.0, voltage_limit=2.0),
|
||||
device,
|
||||
)
|
||||
# kp * err = 100 * 0.5 = 50 V, well above Vmax=2.
|
||||
pos = torch.tensor([[0.5, 0.0]], device=device)
|
||||
_drive(entity, sim, device, pos_target=pos)
|
||||
v_adr = entity.indexing.joint_v_adr
|
||||
qfrc = sim.data.qfrc_actuator[0, v_adr]
|
||||
expected_first = K * 2.0 / R # tau at clamped V.
|
||||
assert qfrc[0].item() == pytest.approx(expected_first, abs=1e-4)
|
||||
assert qfrc[1].item() == pytest.approx(0.0, abs=1e-4)
|
||||
|
||||
|
||||
def test_velocity_mode_pid(device):
|
||||
"""P-only velocity tracking: tau = K * kp * (target - qdot) / R."""
|
||||
entity, sim = initialize_entity(
|
||||
_make_entity(mode=DcMotorInputMode.VELOCITY, damping=0.0, voltage_limit=1000.0),
|
||||
device,
|
||||
)
|
||||
qd0 = torch.tensor([[1.0, 0.0]], device=device)
|
||||
vel_target = torch.tensor([[3.0, 0.0]], device=device)
|
||||
_drive(entity, sim, device, vel_target=vel_target, qd0=qd0)
|
||||
v_adr = entity.indexing.joint_v_adr
|
||||
# back-EMF subtracts K*omega; this is folded into the dcmotor bias.
|
||||
# voltage = kp*(target - qdot); tau = K*(voltage - K*omega)/R.
|
||||
voltage = 5.0 * (vel_target[0] - qd0[0])
|
||||
expected = K * (voltage - K * qd0[0]) / R
|
||||
assert torch.allclose(sim.data.qfrc_actuator[0, v_adr], expected, atol=1e-4)
|
||||
|
||||
|
||||
def test_effort_limit_clamps_torque(device):
|
||||
"""forcerange clamps the algebraic torque output."""
|
||||
entity, sim = initialize_entity(
|
||||
_make_entity(stiffness=100.0, damping=0.0, voltage_limit=1000.0, effort_limit=0.1),
|
||||
device,
|
||||
)
|
||||
m = sim.mj_model
|
||||
for i in range(2):
|
||||
assert m.actuator_forcelimited[i] == 1
|
||||
assert m.actuator_forcerange[i, 0] == pytest.approx(-0.1)
|
||||
assert m.actuator_forcerange[i, 1] == pytest.approx(0.1)
|
||||
|
||||
# Unclamped tau would be K * 100 * 0.5 / R ~= K*50/R, well above 0.1.
|
||||
pos = torch.tensor([[0.5, 0.0]], device=device)
|
||||
_drive(entity, sim, device, pos_target=pos)
|
||||
v_adr = entity.indexing.joint_v_adr
|
||||
qfrc = sim.data.qfrc_actuator[0, v_adr]
|
||||
assert qfrc[0].item() == pytest.approx(0.1, abs=1e-4)
|
||||
assert qfrc[1].item() == pytest.approx(0.0, abs=1e-4)
|
||||
|
||||
|
||||
# Cogging.
|
||||
|
||||
|
||||
def test_cogging_packed_into_biasprm(device):
|
||||
"""``cogging=(A, Np, phi)`` packs into ``biasprm[0:3]``."""
|
||||
cfg = BuiltinDcMotorActuatorCfg(
|
||||
target_names_expr=("joint.*",),
|
||||
mode=DcMotorInputMode.VOLTAGE,
|
||||
motor_params=DATASHEET,
|
||||
cogging=(0.5, 4.0, 0.1),
|
||||
)
|
||||
_, sim = initialize_entity(create_entity_with_actuator(ROBOT_XML, cfg), device)
|
||||
m = sim.mj_model
|
||||
for i in range(2):
|
||||
assert m.actuator_biasprm[i, 0] == pytest.approx(0.5)
|
||||
assert m.actuator_biasprm[i, 1] == pytest.approx(4.0)
|
||||
assert m.actuator_biasprm[i, 2] == pytest.approx(0.1)
|
||||
|
||||
|
||||
def test_cogging_contributes_torque(device):
|
||||
"""At ctrl=0 (no electromagnetic torque), qfrc_actuator equals the cogging
|
||||
term ``A * sin(Np * q + phi)`` evaluated at the joint angle."""
|
||||
A, Np, phi = 0.5, 4.0, 0.1
|
||||
cfg = BuiltinDcMotorActuatorCfg(
|
||||
target_names_expr=("joint.*",),
|
||||
mode=DcMotorInputMode.VOLTAGE,
|
||||
motor_params=DATASHEET,
|
||||
cogging=(A, Np, phi),
|
||||
)
|
||||
entity, sim = initialize_entity(create_entity_with_actuator(ROBOT_XML, cfg), device)
|
||||
q0, q1 = 0.3, -0.2
|
||||
_drive(
|
||||
entity,
|
||||
sim,
|
||||
device,
|
||||
q0=torch.tensor([[q0, q1]], device=device),
|
||||
effort_target=torch.zeros(1, 2, device=device),
|
||||
)
|
||||
v_adr = entity.indexing.joint_v_adr
|
||||
qfrc = sim.data.qfrc_actuator[0, v_adr]
|
||||
assert qfrc[0].item() == pytest.approx(A * math.sin(Np * q0 + phi), abs=1e-5)
|
||||
assert qfrc[1].item() == pytest.approx(A * math.sin(Np * q1 + phi), abs=1e-5)
|
||||
|
||||
|
||||
def test_cogging_bypasses_effort_limit(device):
|
||||
"""Cogging is added *after* the forcerange clamp (MuJoCo's intentional
|
||||
model: ``effort_limit`` bounds electromagnetic torque, cogging is
|
||||
mechanical). Total torque can exceed ``effort_limit`` by up to the
|
||||
cogging amplitude."""
|
||||
A, Np, phi = 0.5, 0.0, math.pi / 2 # sin(pi/2)=1, so cogging = A at any q.
|
||||
cfg = BuiltinDcMotorActuatorCfg(
|
||||
target_names_expr=("joint.*",),
|
||||
mode=DcMotorInputMode.VOLTAGE,
|
||||
motor_params=DATASHEET,
|
||||
cogging=(A, Np, phi),
|
||||
effort_limit=0.05, # An order of magnitude below A.
|
||||
)
|
||||
entity, sim = initialize_entity(create_entity_with_actuator(ROBOT_XML, cfg), device)
|
||||
# Pick a voltage large enough that the electromagnetic torque alone
|
||||
# would saturate forcerange at +/- 0.05.
|
||||
V = torch.tensor([[100.0, 0.0]], device=device)
|
||||
_drive(entity, sim, device, effort_target=V)
|
||||
v_adr = entity.indexing.joint_v_adr
|
||||
qfrc = sim.data.qfrc_actuator[0, v_adr]
|
||||
# joint1: electromagnetic clamped to +0.05, plus cogging A=0.5.
|
||||
assert qfrc[0].item() == pytest.approx(0.05 + A, abs=1e-5)
|
||||
# joint2: zero voltage, electromagnetic=0, only cogging.
|
||||
assert qfrc[1].item() == pytest.approx(A, abs=1e-5)
|
||||
|
||||
|
||||
# Optional stateful extensions (integral, slew, inductance, thermal, LuGre).
|
||||
# Each behavior check compares against a baseline with the feature disabled
|
||||
# so that removing the wiring in edit_spec causes the comparison to fail.
|
||||
|
||||
|
||||
def _step_n(entity, sim, device, n: int, *, pos_target=None, eff_target=None):
|
||||
zero = torch.zeros(1, 2, device=device)
|
||||
entity.write_joint_state_to_sim(position=zero, velocity=zero)
|
||||
for _ in range(n):
|
||||
entity.set_joint_position_target(pos_target if pos_target is not None else zero)
|
||||
entity.set_joint_velocity_target(zero)
|
||||
entity.set_joint_effort_target(eff_target if eff_target is not None else zero)
|
||||
entity.write_data_to_sim()
|
||||
sim.step()
|
||||
|
||||
|
||||
def _qfrc(entity, sim) -> torch.Tensor:
|
||||
return sim.data.qfrc_actuator[0, entity.indexing.joint_v_adr].clone()
|
||||
|
||||
|
||||
def test_integral_gain_ramps_torque(device):
|
||||
"""Integrator in position mode ramps torque over time even with ``kp``
|
||||
and ``kd`` near zero."""
|
||||
# stiffness must be > 0 (validation); choose tiny so ki dominates.
|
||||
base = dict(
|
||||
mode=DcMotorInputMode.POSITION, stiffness=1e-4, damping=0.0, voltage_limit=24.0
|
||||
)
|
||||
ent_off, sim_off = _make_initialized(device, **base, integral_gain=0.0)
|
||||
ent_on, sim_on = _make_initialized(device, **base, integral_gain=10.0)
|
||||
|
||||
target = torch.tensor([[0.5, 0.0]], device=device)
|
||||
for sim, ent in ((sim_off, ent_off), (sim_on, ent_on)):
|
||||
_step_n(ent, sim, device, n=20, pos_target=target)
|
||||
assert _qfrc(ent_on, sim_on)[0].abs() > 100.0 * _qfrc(ent_off, sim_off)[0].abs()
|
||||
|
||||
|
||||
def test_slew_rate_limits_voltage(device):
|
||||
"""``slew_rate`` rate-limits ``ctrl``: after one step, effective voltage
|
||||
is far below the requested input."""
|
||||
base = dict(
|
||||
mode=DcMotorInputMode.VOLTAGE, stiffness=0.0, damping=0.0, voltage_limit=0.0
|
||||
)
|
||||
ent_off, sim_off = _make_initialized(device, **base, slew_rate=0.0)
|
||||
ent_on, sim_on = _make_initialized(device, **base, slew_rate=10.0)
|
||||
|
||||
V = torch.tensor([[100.0, 0.0]], device=device)
|
||||
for sim, ent in ((sim_off, ent_off), (sim_on, ent_on)):
|
||||
_step_n(ent, sim, device, n=1, eff_target=V)
|
||||
assert _qfrc(ent_off, sim_off)[0] > 100.0 * _qfrc(ent_on, sim_on)[0]
|
||||
|
||||
|
||||
def test_inductance_lags_current(device):
|
||||
"""Large ``inductance`` (te >> dt) suppresses early-step torque."""
|
||||
base = dict(
|
||||
mode=DcMotorInputMode.VOLTAGE, stiffness=0.0, damping=0.0, voltage_limit=0.0
|
||||
)
|
||||
ent_off, sim_off = _make_initialized(device, **base, inductance=0.0)
|
||||
ent_on, sim_on = _make_initialized(device, **base, inductance=1.0)
|
||||
|
||||
V = torch.tensor([[10.0, 0.0]], device=device)
|
||||
for sim, ent in ((sim_off, ent_off), (sim_on, ent_on)):
|
||||
_step_n(ent, sim, device, n=2, eff_target=V)
|
||||
assert _qfrc(ent_off, sim_off)[0].abs() > 10.0 * _qfrc(ent_on, sim_on)[0].abs()
|
||||
|
||||
|
||||
def test_thermal_decays_torque(device):
|
||||
"""I^2R heating raises T, which raises effective resistance and decays
|
||||
torque over time."""
|
||||
# Params chosen for visible effect in a handful of steps without going
|
||||
# numerically unstable: small C (fast heating) and modest alpha.
|
||||
base = dict(
|
||||
mode=DcMotorInputMode.VOLTAGE, stiffness=0.0, damping=0.0, voltage_limit=0.0
|
||||
)
|
||||
ent_off, sim_off = _make_initialized(device, **base)
|
||||
ent_on, sim_on = _make_initialized(
|
||||
device, **base, thermal=(1.0, 0.1, 0.0, 0.01, 0.0, 0.0)
|
||||
)
|
||||
|
||||
V = torch.tensor([[100.0, 0.0]], device=device)
|
||||
for sim, ent in ((sim_off, ent_off), (sim_on, ent_on)):
|
||||
_step_n(ent, sim, device, n=5, eff_target=V)
|
||||
assert _qfrc(ent_on, sim_on)[0].abs() < 0.5 * _qfrc(ent_off, sim_off)[0].abs()
|
||||
|
||||
|
||||
def test_lugre_subtracts_friction(device):
|
||||
"""LuGre friction subtracts a velocity-dependent force after the
|
||||
``effort_limit`` clamp (mechanical, like cogging)."""
|
||||
# Static comparison at v>0, ctrl=0; avoids feedback between LuGre slowing
|
||||
# the joint and back-EMF easing off under sim.step().
|
||||
# no LuGre: qfrc = -K^2 * v / R (back-EMF only)
|
||||
# w/ LuGre: qfrc = -K^2 * v / R - sigma1*v - ...
|
||||
base = dict(
|
||||
mode=DcMotorInputMode.VOLTAGE, stiffness=0.0, damping=0.0, voltage_limit=0.0
|
||||
)
|
||||
ent_off, sim_off = _make_initialized(device, **base)
|
||||
ent_on, sim_on = _make_initialized(
|
||||
device, **base, lugre=(1e4, 100.0, 0.1, 0.15, 0.01)
|
||||
)
|
||||
|
||||
zero = torch.zeros(1, 2, device=device)
|
||||
v0 = torch.tensor([[1.0, 0.0]], device=device)
|
||||
for sim, ent in ((sim_off, ent_off), (sim_on, ent_on)):
|
||||
ent.write_joint_state_to_sim(position=zero, velocity=v0)
|
||||
ent.set_joint_position_target(zero)
|
||||
ent.set_joint_velocity_target(zero)
|
||||
ent.set_joint_effort_target(zero)
|
||||
ent.write_data_to_sim()
|
||||
sim.forward()
|
||||
assert abs(_qfrc(ent_on, sim_on)[0]) > 100.0 * abs(_qfrc(ent_off, sim_off)[0])
|
||||
|
||||
|
||||
# Config validation.
|
||||
|
||||
|
||||
def test_pid_mode_requires_gains():
|
||||
with pytest.raises(ValueError, match="stiffness"):
|
||||
BuiltinDcMotorActuatorCfg(
|
||||
target_names_expr=("j",),
|
||||
mode=DcMotorInputMode.POSITION,
|
||||
motor_params=DATASHEET,
|
||||
voltage_limit=1.0,
|
||||
)
|
||||
with pytest.raises(ValueError, match="voltage_limit"):
|
||||
BuiltinDcMotorActuatorCfg(
|
||||
target_names_expr=("j",),
|
||||
mode=DcMotorInputMode.POSITION,
|
||||
motor_params=DATASHEET,
|
||||
stiffness=1.0,
|
||||
)
|
||||
|
||||
|
||||
def test_voltage_mode_rejects_pid_gains():
|
||||
with pytest.raises(ValueError, match="VOLTAGE"):
|
||||
BuiltinDcMotorActuatorCfg(
|
||||
target_names_expr=("j",),
|
||||
mode=DcMotorInputMode.VOLTAGE,
|
||||
motor_params=DATASHEET,
|
||||
stiffness=1.0,
|
||||
)
|
||||
|
||||
|
||||
def test_site_rejected():
|
||||
with pytest.raises(ValueError, match="SITE"):
|
||||
BuiltinDcMotorActuatorCfg(
|
||||
target_names_expr=("j",),
|
||||
motor_params=DATASHEET,
|
||||
stiffness=1.0,
|
||||
voltage_limit=1.0,
|
||||
transmission_type=TransmissionType.SITE,
|
||||
)
|
||||
|
||||
|
||||
# Joint-level passthrough.
|
||||
|
||||
|
||||
def test_armature_applied(device):
|
||||
_, sim = initialize_entity(_make_entity(armature=0.7), device)
|
||||
m = sim.mj_model
|
||||
for jname in ("joint1", "joint2"):
|
||||
dof_id = m.jnt_dofadr[m.joint(jname).id]
|
||||
assert m.dof_armature[dof_id] == pytest.approx(0.7)
|
||||
|
||||
|
||||
# Domain randomization.
|
||||
|
||||
|
||||
def _scene_env(
|
||||
device,
|
||||
num_envs=2,
|
||||
mode: DcMotorInputMode = DcMotorInputMode.POSITION,
|
||||
):
|
||||
def spec_fn():
|
||||
spec = mujoco.MjSpec.from_string(ROBOT_XML)
|
||||
for a in list(spec.actuators):
|
||||
spec.delete(a)
|
||||
return spec
|
||||
|
||||
entity_cfg = EntityCfg(
|
||||
spec_fn=spec_fn,
|
||||
articulation=EntityArticulationInfoCfg(
|
||||
actuators=(
|
||||
BuiltinDcMotorActuatorCfg(
|
||||
target_names_expr=("joint.*",),
|
||||
mode=mode,
|
||||
motor_params=DATASHEET,
|
||||
stiffness=5.0 if mode != DcMotorInputMode.VOLTAGE else 0.0,
|
||||
damping=0.5 if mode != DcMotorInputMode.VOLTAGE else 0.0,
|
||||
voltage_limit=24.0 if mode != DcMotorInputMode.VOLTAGE else 0.0,
|
||||
effort_limit=50.0,
|
||||
),
|
||||
)
|
||||
),
|
||||
)
|
||||
scene_cfg = SceneCfg(num_envs=num_envs, entities={"robot": entity_cfg})
|
||||
scene = Scene(scene_cfg, device)
|
||||
model = scene.compile()
|
||||
sim = Simulation(num_envs=num_envs, cfg=SimulationCfg(), model=model, device=device)
|
||||
scene.initialize(model, sim.model, sim.data)
|
||||
|
||||
env = Mock()
|
||||
env.num_envs = num_envs
|
||||
env.device = device
|
||||
env.scene = {"robot": scene["robot"]}
|
||||
env.sim = sim
|
||||
return env
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"operation, kp_in, kd_in, kp_expected, kd_expected",
|
||||
[
|
||||
# scale: multiplies the configured defaults (kp=5.0, kd=0.5).
|
||||
("scale", 2.0, 3.0, 2.0 * 5.0, 3.0 * 0.5),
|
||||
# abs: writes the value directly.
|
||||
("abs", 10.0, 2.0, 10.0, 2.0),
|
||||
],
|
||||
)
|
||||
def test_dr_pd_gains_position_mode(
|
||||
device, operation, kp_in, kd_in, kp_expected, kd_expected
|
||||
):
|
||||
env = _scene_env(device)
|
||||
robot = env.scene["robot"]
|
||||
act = robot.actuators[0]
|
||||
assert isinstance(act, BuiltinDcMotorActuator)
|
||||
ctrl_ids = act.global_ctrl_ids
|
||||
env.sim.expand_model_fields(("actuator_gainprm", "actuator_biasprm"))
|
||||
|
||||
dr.pd_gains(
|
||||
env,
|
||||
env_ids=torch.tensor([0], device=device),
|
||||
kp_range=(kp_in, kp_in),
|
||||
kd_range=(kd_in, kd_in),
|
||||
asset_cfg=SceneEntityCfg("robot"),
|
||||
operation=operation,
|
||||
)
|
||||
|
||||
m = env.sim.model
|
||||
n = len(ctrl_ids)
|
||||
assert torch.allclose(
|
||||
m.actuator_gainprm[0, ctrl_ids, 4], torch.full((n,), kp_expected, device=device)
|
||||
)
|
||||
assert torch.allclose(
|
||||
m.actuator_gainprm[0, ctrl_ids, 6], torch.full((n,), kd_expected, device=device)
|
||||
)
|
||||
# Other env untouched (cfg defaults).
|
||||
assert torch.allclose(m.actuator_gainprm[1, ctrl_ids, 4], torch.tensor(5.0))
|
||||
assert torch.allclose(m.actuator_gainprm[1, ctrl_ids, 6], torch.tensor(0.5))
|
||||
|
||||
|
||||
def test_dr_pd_gains_voltage_mode_rejected(device):
|
||||
env = _scene_env(device, mode=DcMotorInputMode.VOLTAGE)
|
||||
env.sim.expand_model_fields(("actuator_gainprm", "actuator_biasprm"))
|
||||
with pytest.raises(ValueError, match="VOLTAGE"):
|
||||
dr.pd_gains(
|
||||
env,
|
||||
env_ids=torch.tensor([0], device=device),
|
||||
kp_range=(1.0, 1.0),
|
||||
kd_range=(1.0, 1.0),
|
||||
asset_cfg=SceneEntityCfg("robot"),
|
||||
)
|
||||
|
||||
|
||||
def test_dr_effort_limits_writes_forcerange(device):
|
||||
env = _scene_env(device)
|
||||
robot = env.scene["robot"]
|
||||
act = robot.actuators[0]
|
||||
ctrl_ids = act.global_ctrl_ids
|
||||
env.sim.expand_model_fields(
|
||||
("actuator_forcerange", "jnt_actfrcrange", "tendon_actfrcrange")
|
||||
)
|
||||
|
||||
dr.effort_limits(
|
||||
env,
|
||||
env_ids=torch.tensor([0], device=device),
|
||||
effort_limit_range=(123.0, 123.0),
|
||||
asset_cfg=SceneEntityCfg("robot"),
|
||||
operation="abs",
|
||||
)
|
||||
|
||||
m = env.sim.model
|
||||
n = len(ctrl_ids)
|
||||
assert torch.allclose(
|
||||
m.actuator_forcerange[0, ctrl_ids, 0],
|
||||
torch.full((n,), -123.0, device=device),
|
||||
)
|
||||
assert torch.allclose(
|
||||
m.actuator_forcerange[0, ctrl_ids, 1],
|
||||
torch.full((n,), 123.0, device=device),
|
||||
)
|
||||
# Env 1 keeps the configured default of 50.
|
||||
assert torch.allclose(m.actuator_forcerange[1, ctrl_ids, 1], torch.tensor(50.0))
|
||||
|
||||
|
||||
# Delay.
|
||||
|
||||
|
||||
def test_delay_position_mode(device):
|
||||
"""A 2-step lag should make position-mode torque reference step-0 target."""
|
||||
entity, sim = initialize_entity(
|
||||
_make_entity(
|
||||
stiffness=10.0,
|
||||
damping=0.0,
|
||||
voltage_limit=1000.0,
|
||||
delay_min_lag=2,
|
||||
delay_max_lag=2,
|
||||
),
|
||||
device,
|
||||
)
|
||||
zero = torch.zeros(1, 2, device=device)
|
||||
entity.write_joint_state_to_sim(position=zero, velocity=zero)
|
||||
targets = [
|
||||
torch.tensor([[0.1, 0.0]], device=device),
|
||||
torch.tensor([[0.3, 0.0]], device=device),
|
||||
torch.tensor([[0.5, 0.0]], device=device),
|
||||
]
|
||||
for p in targets:
|
||||
entity.set_joint_position_target(p)
|
||||
entity.set_joint_velocity_target(zero)
|
||||
entity.set_joint_effort_target(zero)
|
||||
entity.write_data_to_sim()
|
||||
sim.forward()
|
||||
|
||||
v_adr = entity.indexing.joint_v_adr
|
||||
# With lag=2 and three writes, the effective target is targets[0].
|
||||
expected = K * 10.0 * targets[0][0] / R
|
||||
assert torch.allclose(sim.data.qfrc_actuator[0, v_adr], expected, atol=1e-4)
|
||||
@@ -0,0 +1,468 @@
|
||||
"""Tests for BuiltinPdActuator.
|
||||
|
||||
Covers the unique surface of the actuator: paired <position>/<velocity>
|
||||
elements per target, joint/tendon-level actfrcrange sum-clamp, DR for both
|
||||
gains and effort limits, delay synchronization, and the ordering invariant
|
||||
that DR depends on.
|
||||
"""
|
||||
|
||||
from unittest.mock import Mock
|
||||
|
||||
import mujoco
|
||||
import pytest
|
||||
import torch
|
||||
from conftest import (
|
||||
create_entity_with_actuator,
|
||||
get_test_device,
|
||||
initialize_entity,
|
||||
load_fixture_xml,
|
||||
)
|
||||
|
||||
from mjlab.actuator import BuiltinPdActuator, BuiltinPdActuatorCfg
|
||||
from mjlab.actuator.actuator import TransmissionType
|
||||
from mjlab.entity import Entity, EntityArticulationInfoCfg, EntityCfg
|
||||
from mjlab.envs.mdp import dr
|
||||
from mjlab.managers.scene_entity_config import SceneEntityCfg
|
||||
from mjlab.scene import Scene, SceneCfg
|
||||
from mjlab.sim.sim import Simulation, SimulationCfg
|
||||
|
||||
ROBOT_XML = load_fixture_xml("floating_base_articulated")
|
||||
KP = 100.0
|
||||
KD = 10.0
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def device():
|
||||
return get_test_device()
|
||||
|
||||
|
||||
def _make_entity(
|
||||
*,
|
||||
effort_limit: float | None = 50.0,
|
||||
armature: float | None = None,
|
||||
delay_max_lag: int = 0,
|
||||
delay_min_lag: int = 0,
|
||||
delay_hold_prob: float = 0.0,
|
||||
) -> Entity:
|
||||
cfg = BuiltinPdActuatorCfg(
|
||||
target_names_expr=("joint.*",),
|
||||
stiffness=KP,
|
||||
damping=KD,
|
||||
effort_limit=effort_limit,
|
||||
armature=armature,
|
||||
delay_min_lag=delay_min_lag,
|
||||
delay_max_lag=delay_max_lag,
|
||||
delay_hold_prob=delay_hold_prob,
|
||||
)
|
||||
return create_entity_with_actuator(ROBOT_XML, cfg)
|
||||
|
||||
|
||||
def _at_rest_with_targets(
|
||||
entity: Entity,
|
||||
sim,
|
||||
device: str,
|
||||
pos_target: torch.Tensor,
|
||||
vel_target: torch.Tensor,
|
||||
) -> None:
|
||||
entity.write_joint_state_to_sim(
|
||||
position=torch.zeros(1, 2, device=device),
|
||||
velocity=torch.zeros(1, 2, device=device),
|
||||
)
|
||||
entity.set_joint_position_target(pos_target)
|
||||
entity.set_joint_velocity_target(vel_target)
|
||||
entity.set_joint_effort_target(torch.zeros(1, 2, device=device))
|
||||
entity.write_data_to_sim()
|
||||
sim.forward()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Structural invariants
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_two_ctrls_per_target_with_pos_then_vel_layout(device):
|
||||
"""Each target gets one <position> + one <velocity>, in halves."""
|
||||
entity, sim = initialize_entity(_make_entity(), device)
|
||||
act = entity.actuators[0]
|
||||
assert isinstance(act, BuiltinPdActuator)
|
||||
|
||||
n = act.num_targets
|
||||
assert n == len(act.target_names) == 2
|
||||
assert len(act.ctrl_ids) == 2 * n
|
||||
assert len(act.global_ctrl_ids) == 2 * n
|
||||
|
||||
names = [sim.mj_model.actuator(i).name for i in act.global_ctrl_ids.tolist()]
|
||||
assert names[:n] == [f"{name}_pd_pos" for name in act.target_names]
|
||||
assert names[n:] == [f"{name}_pd_vel" for name in act.target_names]
|
||||
|
||||
|
||||
def test_site_transmission_rejected():
|
||||
with pytest.raises(ValueError, match="SITE"):
|
||||
BuiltinPdActuatorCfg(
|
||||
target_names_expr=("x",),
|
||||
stiffness=1.0,
|
||||
damping=1.0,
|
||||
transmission_type=TransmissionType.SITE,
|
||||
)
|
||||
|
||||
|
||||
def test_armature_applied_once(device):
|
||||
"""Joint armature must come from the position element only; double-applying
|
||||
would silently double dof_armature."""
|
||||
_, sim = initialize_entity(_make_entity(armature=0.7), device)
|
||||
m = sim.mj_model
|
||||
for jname in ("joint1", "joint2"):
|
||||
dof_id = m.jnt_dofadr[m.joint(jname).id]
|
||||
assert m.dof_armature[dof_id] == pytest.approx(0.7)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Force computation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_position_only(device):
|
||||
"""Zero vel target: qfrc = kp * pos_target."""
|
||||
entity, sim = initialize_entity(_make_entity(effort_limit=None), device)
|
||||
pos = torch.tensor([[0.1, -0.05]], device=device)
|
||||
_at_rest_with_targets(entity, sim, device, pos, torch.zeros(1, 2, device=device))
|
||||
v_adr = entity.indexing.joint_v_adr
|
||||
assert torch.allclose(sim.data.qfrc_actuator[0, v_adr], KP * pos[0], atol=1e-4)
|
||||
|
||||
|
||||
def test_velocity_only(device):
|
||||
"""Zero pos target, joint at rest: qfrc = kd * vel_target."""
|
||||
entity, sim = initialize_entity(_make_entity(effort_limit=None), device)
|
||||
vel = torch.tensor([[0.3, -0.2]], device=device)
|
||||
_at_rest_with_targets(entity, sim, device, torch.zeros(1, 2, device=device), vel)
|
||||
v_adr = entity.indexing.joint_v_adr
|
||||
assert torch.allclose(sim.data.qfrc_actuator[0, v_adr], KD * vel[0], atol=1e-4)
|
||||
|
||||
|
||||
def test_pd_superposition(device):
|
||||
"""Both targets nonzero: qfrc = kp * pos_target + kd * vel_target."""
|
||||
entity, sim = initialize_entity(_make_entity(effort_limit=None), device)
|
||||
pos = torch.tensor([[0.1, -0.05]], device=device)
|
||||
vel = torch.tensor([[0.2, -0.1]], device=device)
|
||||
_at_rest_with_targets(entity, sim, device, pos, vel)
|
||||
v_adr = entity.indexing.joint_v_adr
|
||||
expected = KP * pos[0] + KD * vel[0]
|
||||
assert torch.allclose(sim.data.qfrc_actuator[0, v_adr], expected, atol=1e-4)
|
||||
|
||||
|
||||
def test_actfrcrange_sum_clamp(device):
|
||||
"""A pos error big enough to make kp*err exceed effort_limit must be
|
||||
clamped at the joint, not allowed to ride through the unbounded element."""
|
||||
entity, sim = initialize_entity(_make_entity(effort_limit=5.0), device)
|
||||
# kp * 10.0 = 1000, well over the 5.0 clamp.
|
||||
pos = torch.tensor([[10.0, 0.0]], device=device)
|
||||
_at_rest_with_targets(entity, sim, device, pos, torch.zeros(1, 2, device=device))
|
||||
v_adr = entity.indexing.joint_v_adr
|
||||
qfrc = sim.data.qfrc_actuator[0, v_adr]
|
||||
assert qfrc[0].item() == pytest.approx(5.0, abs=1e-4)
|
||||
assert qfrc[1].item() == pytest.approx(0.0, abs=1e-4)
|
||||
|
||||
|
||||
def test_effort_limit_none_leaves_joint_unlimited(device):
|
||||
"""effort_limit=None: jnt_actfrclimited stays 0 on the targeted joints."""
|
||||
_, sim = initialize_entity(_make_entity(effort_limit=None), device)
|
||||
m = sim.mj_model
|
||||
for jname in ("joint1", "joint2"):
|
||||
jid = m.joint(jname).id
|
||||
assert m.jnt_actfrclimited[jid] == 0
|
||||
|
||||
|
||||
def test_actuator_forcerange_not_set(device):
|
||||
"""We deliberately leave per-element forcerange unset; the limit lives on
|
||||
the joint. Inspection of actuator_force[i] thus shows the unclamped value."""
|
||||
entity, sim = initialize_entity(_make_entity(effort_limit=5.0), device)
|
||||
m = sim.mj_model
|
||||
for ctrl_id in entity.actuators[0].global_ctrl_ids.tolist():
|
||||
assert m.actuator_forcelimited[ctrl_id] == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Delay synchronization
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_delay_syncs_pos_and_vel(device):
|
||||
"""The shared delay buffer must lag pos and vel together."""
|
||||
entity, sim = initialize_entity(
|
||||
_make_entity(effort_limit=None, delay_min_lag=2, delay_max_lag=2),
|
||||
device,
|
||||
)
|
||||
pos_targets = [
|
||||
torch.tensor([[0.1, 0.0]], device=device),
|
||||
torch.tensor([[0.3, 0.0]], device=device),
|
||||
torch.tensor([[0.5, 0.0]], device=device),
|
||||
]
|
||||
vel_targets = [
|
||||
torch.tensor([[1.0, 0.0]], device=device),
|
||||
torch.tensor([[2.0, 0.0]], device=device),
|
||||
torch.tensor([[3.0, 0.0]], device=device),
|
||||
]
|
||||
entity.write_joint_state_to_sim(
|
||||
position=torch.zeros(1, 2, device=device),
|
||||
velocity=torch.zeros(1, 2, device=device),
|
||||
)
|
||||
for p, v in zip(pos_targets, vel_targets, strict=True):
|
||||
entity.set_joint_position_target(p)
|
||||
entity.set_joint_velocity_target(v)
|
||||
entity.set_joint_effort_target(torch.zeros(1, 2, device=device))
|
||||
entity.write_data_to_sim()
|
||||
sim.forward()
|
||||
|
||||
v_adr = entity.indexing.joint_v_adr
|
||||
# With lag=2, both halves should reference step-0 values.
|
||||
expected = KP * pos_targets[0][0] + KD * vel_targets[0][0]
|
||||
assert torch.allclose(sim.data.qfrc_actuator[0, v_adr], expected, atol=1e-4)
|
||||
|
||||
|
||||
def test_reset_clears_delay_buffer(device):
|
||||
entity, _ = initialize_entity(_make_entity(delay_min_lag=1, delay_max_lag=3), device)
|
||||
act = entity.actuators[0]
|
||||
assert act._delay_buffer is not None
|
||||
entity.set_joint_position_target(torch.full((1, 2), 0.5, device=device))
|
||||
entity.set_joint_velocity_target(torch.zeros(1, 2, device=device))
|
||||
entity.write_data_to_sim()
|
||||
|
||||
entity.reset(torch.tensor([0], device=device))
|
||||
assert act._delay_buffer.current_lags[0] == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Domain randomization
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _scene_env(device, transmission=TransmissionType.JOINT, num_envs=2):
|
||||
"""Build a real scene/sim with one BuiltinPd-driven entity for DR tests."""
|
||||
if transmission == TransmissionType.JOINT:
|
||||
xml = ROBOT_XML
|
||||
targets = ("joint.*",)
|
||||
else:
|
||||
xml = load_fixture_xml("tendon_finger")
|
||||
# tendon_finger ships with motor/position/velocity actuators; we need a
|
||||
# bare spec so BuiltinPd can attach to the tendon without name clashes.
|
||||
targets = ("finger_tendon",)
|
||||
|
||||
def spec_fn():
|
||||
spec = mujoco.MjSpec.from_string(xml)
|
||||
# Strip any pre-existing actuators so BuiltinPd's added elements own ctrl.
|
||||
for a in list(spec.actuators):
|
||||
spec.delete(a)
|
||||
return spec
|
||||
|
||||
entity_cfg = EntityCfg(
|
||||
spec_fn=spec_fn,
|
||||
articulation=EntityArticulationInfoCfg(
|
||||
actuators=(
|
||||
BuiltinPdActuatorCfg(
|
||||
target_names_expr=targets,
|
||||
stiffness=KP,
|
||||
damping=KD,
|
||||
effort_limit=50.0,
|
||||
transmission_type=transmission,
|
||||
),
|
||||
)
|
||||
),
|
||||
)
|
||||
scene_cfg = SceneCfg(num_envs=num_envs, entities={"robot": entity_cfg})
|
||||
scene = Scene(scene_cfg, device)
|
||||
model = scene.compile()
|
||||
sim = Simulation(num_envs=num_envs, cfg=SimulationCfg(), model=model, device=device)
|
||||
scene.initialize(model, sim.model, sim.data)
|
||||
|
||||
env = Mock()
|
||||
env.num_envs = num_envs
|
||||
env.device = device
|
||||
env.scene = {"robot": scene["robot"]}
|
||||
env.sim = sim
|
||||
return env
|
||||
|
||||
|
||||
def test_dr_pd_gains_scales_halves_independently(device):
|
||||
env = _scene_env(device)
|
||||
robot = env.scene["robot"]
|
||||
act = robot.actuators[0]
|
||||
assert isinstance(act, BuiltinPdActuator)
|
||||
n = act.num_targets
|
||||
pos_ids = act.global_ctrl_ids[:n]
|
||||
vel_ids = act.global_ctrl_ids[n:]
|
||||
|
||||
# Expand fields so DR can write per-env.
|
||||
env.sim.expand_model_fields(("actuator_gainprm", "actuator_biasprm"))
|
||||
|
||||
dr.pd_gains(
|
||||
env,
|
||||
env_ids=torch.tensor([0], device=device),
|
||||
kp_range=(2.0, 2.0),
|
||||
kd_range=(3.0, 3.0),
|
||||
asset_cfg=SceneEntityCfg("robot"),
|
||||
operation="scale",
|
||||
)
|
||||
|
||||
m = env.sim.model
|
||||
# Position half: gainprm[0] and biasprm[1] both scaled by kp=2, biasprm[2]
|
||||
# must stay zero (no kd injection).
|
||||
assert torch.allclose(
|
||||
m.actuator_gainprm[0, pos_ids, 0],
|
||||
torch.full((n,), 2.0 * KP, device=device),
|
||||
)
|
||||
assert torch.allclose(
|
||||
m.actuator_biasprm[0, pos_ids, 1],
|
||||
torch.full((n,), -2.0 * KP, device=device),
|
||||
)
|
||||
assert torch.allclose(
|
||||
m.actuator_biasprm[0, pos_ids, 2], torch.zeros(n, device=device)
|
||||
)
|
||||
# Velocity half: gainprm[0] and biasprm[2] both scaled by kd=3, biasprm[1]
|
||||
# stays zero (no kp injection).
|
||||
assert torch.allclose(
|
||||
m.actuator_gainprm[0, vel_ids, 0],
|
||||
torch.full((n,), 3.0 * KD, device=device),
|
||||
)
|
||||
assert torch.allclose(
|
||||
m.actuator_biasprm[0, vel_ids, 2],
|
||||
torch.full((n,), -3.0 * KD, device=device),
|
||||
)
|
||||
assert torch.allclose(
|
||||
m.actuator_biasprm[0, vel_ids, 1], torch.zeros(n, device=device)
|
||||
)
|
||||
# The other env must be untouched.
|
||||
assert torch.allclose(m.actuator_gainprm[1, pos_ids, 0], torch.tensor(KP))
|
||||
assert torch.allclose(m.actuator_gainprm[1, vel_ids, 0], torch.tensor(KD))
|
||||
|
||||
|
||||
def test_dr_pd_gains_abs_writes_correct_columns(device):
|
||||
env = _scene_env(device)
|
||||
robot = env.scene["robot"]
|
||||
act = robot.actuators[0]
|
||||
assert isinstance(act, BuiltinPdActuator)
|
||||
n = act.num_targets
|
||||
pos_ids = act.global_ctrl_ids[:n]
|
||||
vel_ids = act.global_ctrl_ids[n:]
|
||||
env.sim.expand_model_fields(("actuator_gainprm", "actuator_biasprm"))
|
||||
|
||||
dr.pd_gains(
|
||||
env,
|
||||
env_ids=torch.tensor([0, 1], device=device),
|
||||
kp_range=(200.0, 200.0),
|
||||
kd_range=(25.0, 25.0),
|
||||
asset_cfg=SceneEntityCfg("robot"),
|
||||
operation="abs",
|
||||
)
|
||||
|
||||
m = env.sim.model
|
||||
assert torch.allclose(
|
||||
m.actuator_gainprm[:, pos_ids, 0],
|
||||
torch.full((env.num_envs, n), 200.0, device=device),
|
||||
)
|
||||
assert torch.allclose(
|
||||
m.actuator_biasprm[:, pos_ids, 1],
|
||||
torch.full((env.num_envs, n), -200.0, device=device),
|
||||
)
|
||||
assert torch.allclose(
|
||||
m.actuator_biasprm[:, pos_ids, 2],
|
||||
torch.zeros(env.num_envs, n, device=device),
|
||||
)
|
||||
assert torch.allclose(
|
||||
m.actuator_gainprm[:, vel_ids, 0],
|
||||
torch.full((env.num_envs, n), 25.0, device=device),
|
||||
)
|
||||
assert torch.allclose(
|
||||
m.actuator_biasprm[:, vel_ids, 2],
|
||||
torch.full((env.num_envs, n), -25.0, device=device),
|
||||
)
|
||||
|
||||
|
||||
def test_dr_effort_limits_writes_jnt_actfrcrange(device):
|
||||
env = _scene_env(device, transmission=TransmissionType.JOINT)
|
||||
robot = env.scene["robot"]
|
||||
act = robot.actuators[0]
|
||||
assert isinstance(act, BuiltinPdActuator)
|
||||
env.sim.expand_model_fields(
|
||||
("actuator_forcerange", "jnt_actfrcrange", "tendon_actfrcrange")
|
||||
)
|
||||
|
||||
joint_ids = robot.indexing.joint_ids[act.target_ids]
|
||||
pre_forcerange = env.sim.model.actuator_forcerange.clone()
|
||||
|
||||
dr.effort_limits(
|
||||
env,
|
||||
env_ids=torch.tensor([0], device=device),
|
||||
effort_limit_range=(123.0, 123.0),
|
||||
asset_cfg=SceneEntityCfg("robot"),
|
||||
operation="abs",
|
||||
)
|
||||
|
||||
m = env.sim.model
|
||||
# The joint sum-clamp was rewritten on env 0 only.
|
||||
assert torch.allclose(
|
||||
m.jnt_actfrcrange[0, joint_ids],
|
||||
torch.tensor([[-123.0, 123.0]] * len(joint_ids), device=device),
|
||||
)
|
||||
assert torch.allclose(
|
||||
m.jnt_actfrcrange[1, joint_ids],
|
||||
torch.tensor([[-50.0, 50.0]] * len(joint_ids), device=device),
|
||||
)
|
||||
# Per-element actuator_forcerange must be untouched for BuiltinPd: that
|
||||
# field belongs to the existing single-element actuator semantic.
|
||||
assert torch.allclose(m.actuator_forcerange, pre_forcerange)
|
||||
|
||||
|
||||
def test_dr_effort_limits_scale_multiplies_default(device):
|
||||
"""``scale`` multiplies the configured ``effort_limit`` (50.0) by the sample."""
|
||||
env = _scene_env(device, transmission=TransmissionType.JOINT)
|
||||
robot = env.scene["robot"]
|
||||
act = robot.actuators[0]
|
||||
assert isinstance(act, BuiltinPdActuator)
|
||||
env.sim.expand_model_fields(
|
||||
("actuator_forcerange", "jnt_actfrcrange", "tendon_actfrcrange")
|
||||
)
|
||||
joint_ids = robot.indexing.joint_ids[act.target_ids]
|
||||
|
||||
dr.effort_limits(
|
||||
env,
|
||||
env_ids=torch.tensor([0], device=device),
|
||||
effort_limit_range=(2.0, 2.0),
|
||||
asset_cfg=SceneEntityCfg("robot"),
|
||||
operation="scale",
|
||||
)
|
||||
|
||||
m = env.sim.model
|
||||
# Default is [-50, 50], scaled by 2 -> [-100, 100].
|
||||
assert torch.allclose(
|
||||
m.jnt_actfrcrange[0, joint_ids],
|
||||
torch.tensor([[-100.0, 100.0]] * len(joint_ids), device=device),
|
||||
)
|
||||
|
||||
|
||||
def test_dr_effort_limits_writes_tendon_actfrcrange(device):
|
||||
env = _scene_env(device, transmission=TransmissionType.TENDON)
|
||||
robot = env.scene["robot"]
|
||||
act = robot.actuators[0]
|
||||
assert isinstance(act, BuiltinPdActuator)
|
||||
env.sim.expand_model_fields(
|
||||
("actuator_forcerange", "jnt_actfrcrange", "tendon_actfrcrange")
|
||||
)
|
||||
tendon_ids = robot.indexing.tendon_ids[act.target_ids]
|
||||
|
||||
dr.effort_limits(
|
||||
env,
|
||||
env_ids=torch.tensor([0], device=device),
|
||||
effort_limit_range=(77.0, 77.0),
|
||||
asset_cfg=SceneEntityCfg("robot"),
|
||||
operation="abs",
|
||||
)
|
||||
|
||||
m = env.sim.model
|
||||
assert torch.allclose(
|
||||
m.tendon_actfrcrange[0, tendon_ids],
|
||||
torch.tensor([[-77.0, 77.0]] * len(tendon_ids), device=device),
|
||||
)
|
||||
assert torch.allclose(
|
||||
m.tendon_actfrcrange[1, tendon_ids],
|
||||
torch.tensor([[-50.0, 50.0]] * len(tendon_ids), device=device),
|
||||
)
|
||||
@@ -1072,3 +1072,37 @@ def test_history_captures_impact_forces(device):
|
||||
assert torch.all(max_force_seen > steady_state_force * 1.5), (
|
||||
f"Peak force {max_force_seen} should be significantly above mg={steady_state_force}"
|
||||
)
|
||||
|
||||
|
||||
def test_global_frame_maxforce_rotation(device):
|
||||
"""A box at rest on a plane has its contact normals all vertical."""
|
||||
cfg = ContactSensorCfg(
|
||||
name="box_contact",
|
||||
primary=ContactMatch(mode="geom", pattern="box_geom", entity="box"),
|
||||
fields=("found", "force", "normal", "tangent"),
|
||||
reduce="maxforce",
|
||||
global_frame=True,
|
||||
)
|
||||
scene, sim = create_scene_with_sensor(FALLING_BOX_XML, "box", cfg, device)
|
||||
|
||||
root_state = torch.zeros((2, 13), device=sim.device)
|
||||
root_state[:, 2] = 0.11
|
||||
root_state[:, 3] = 1.0
|
||||
scene["box"].write_root_state_to_sim(root_state)
|
||||
for _ in range(150):
|
||||
sim.step()
|
||||
scene.update(dt=sim.cfg.mujoco.timestep)
|
||||
|
||||
sensor_force = scene["box_contact"].data.force[:, 0, :]
|
||||
|
||||
# On a flat plane the contact normal is vertical, so a correctly rotated
|
||||
# global-frame force should have its magnitude entirely on the z axis.
|
||||
assert torch.all(sensor_force[:, 0].abs() < 0.05), (
|
||||
f"sensor_force x-component should be ~0, got {sensor_force[:, 0].tolist()}"
|
||||
)
|
||||
assert torch.all(sensor_force[:, 1].abs() < 0.05), (
|
||||
f"sensor_force y-component should be ~0, got {sensor_force[:, 1].tolist()}"
|
||||
)
|
||||
assert torch.all(sensor_force[:, 2].abs() > 1.0), (
|
||||
f"sensor_force z-component should be non-trivial, got {sensor_force[:, 2].tolist()}"
|
||||
)
|
||||
|
||||
@@ -125,6 +125,72 @@ def test_delayed_ideal_applies_delay(device):
|
||||
assert torch.allclose(qfrc, expected_torque, atol=1e-4)
|
||||
|
||||
|
||||
def test_delayed_ideal_delays_velocity(device):
|
||||
"""Velocity targets share the same delay as position targets.
|
||||
|
||||
Regression test: the velocity reference used to bypass the delay buffer, so
|
||||
the damping term consumed the latest target instead of the delayed one.
|
||||
"""
|
||||
entity = create_entity_with_delayed_ideal(delay_min_lag=2, delay_max_lag=2)
|
||||
entity, sim = initialize_entity(entity, device)
|
||||
|
||||
joint_pos = torch.zeros(1, 2, device=device)
|
||||
joint_vel = torch.zeros(1, 2, device=device)
|
||||
entity.write_joint_state_to_sim(joint_pos, joint_vel)
|
||||
|
||||
# Only the velocity target varies; position and effort stay zero.
|
||||
vel_targets = [
|
||||
torch.tensor([[0.1, 0.2]], device=device),
|
||||
torch.tensor([[0.3, 0.4]], device=device),
|
||||
torch.tensor([[0.5, 0.6]], device=device),
|
||||
]
|
||||
|
||||
for vel_target in vel_targets:
|
||||
entity.set_joint_position_target(joint_pos)
|
||||
entity.set_joint_velocity_target(vel_target)
|
||||
entity.set_joint_effort_target(torch.zeros(1, 2, device=device))
|
||||
entity.write_data_to_sim()
|
||||
sim.forward()
|
||||
|
||||
joint_v_adr = entity.indexing.joint_v_adr
|
||||
qfrc = sim.data.qfrc_actuator[0, joint_v_adr]
|
||||
|
||||
# With lag=2, the damping term uses the velocity target from step 0:
|
||||
# kd * (delayed_vel_target - 0) = 10.0 * [0.1, 0.2].
|
||||
expected_torque = 10.0 * vel_targets[0][0]
|
||||
assert torch.allclose(qfrc, expected_torque, atol=1e-4)
|
||||
|
||||
|
||||
def test_delayed_ideal_delays_effort(device):
|
||||
"""Feedforward effort targets share the same delay as position targets."""
|
||||
entity = create_entity_with_delayed_ideal(delay_min_lag=2, delay_max_lag=2)
|
||||
entity, sim = initialize_entity(entity, device)
|
||||
|
||||
joint_pos = torch.zeros(1, 2, device=device)
|
||||
joint_vel = torch.zeros(1, 2, device=device)
|
||||
entity.write_joint_state_to_sim(joint_pos, joint_vel)
|
||||
|
||||
effort_targets = [
|
||||
torch.tensor([[1.0, 2.0]], device=device),
|
||||
torch.tensor([[3.0, 4.0]], device=device),
|
||||
torch.tensor([[5.0, 6.0]], device=device),
|
||||
]
|
||||
|
||||
for effort_target in effort_targets:
|
||||
entity.set_joint_position_target(joint_pos)
|
||||
entity.set_joint_velocity_target(joint_vel)
|
||||
entity.set_joint_effort_target(effort_target)
|
||||
entity.write_data_to_sim()
|
||||
sim.forward()
|
||||
|
||||
joint_v_adr = entity.indexing.joint_v_adr
|
||||
qfrc = sim.data.qfrc_actuator[0, joint_v_adr]
|
||||
|
||||
# With lag=2, the feedforward term uses the effort target from step 0.
|
||||
expected_torque = effort_targets[0][0]
|
||||
assert torch.allclose(qfrc, expected_torque, atol=1e-4)
|
||||
|
||||
|
||||
def test_delayed_actuator_reset(device):
|
||||
"""Test that reset clears the delay buffer."""
|
||||
entity = create_entity_with_delayed_builtin(delay_min_lag=1, delay_max_lag=3)
|
||||
|
||||
@@ -243,6 +243,27 @@ def test_unnamed_freejoint_gets_default_name():
|
||||
assert "floating_base_joint" in entity.all_joint_names
|
||||
|
||||
|
||||
def test_multiple_freejoints_raises():
|
||||
"""An entity with more than one freejoint is rejected at construction."""
|
||||
xml = """
|
||||
<mujoco>
|
||||
<worldbody>
|
||||
<body name="object_a" pos="0 0 1">
|
||||
<freejoint/>
|
||||
<geom type="box" size="0.1 0.1 0.1" mass="0.1"/>
|
||||
</body>
|
||||
<body name="object_b" pos="1 0 1">
|
||||
<freejoint/>
|
||||
<geom type="box" size="0.1 0.1 0.1" mass="0.1"/>
|
||||
</body>
|
||||
</worldbody>
|
||||
</mujoco>
|
||||
"""
|
||||
cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(xml))
|
||||
with pytest.raises(ValueError, match="2 freejoints"):
|
||||
Entity(cfg)
|
||||
|
||||
|
||||
def test_find_methods():
|
||||
"""Test find methods with exact and regex matches."""
|
||||
entity = create_floating_articulated_entity()
|
||||
|
||||
@@ -125,7 +125,9 @@ def test_dr_fields_registered_in_event_manager(device):
|
||||
assert "actuator_gainprm" in manager.domain_randomization_fields
|
||||
assert "actuator_biasprm" in manager.domain_randomization_fields
|
||||
assert "actuator_forcerange" in manager.domain_randomization_fields
|
||||
assert len(manager.domain_randomization_fields) == 5
|
||||
assert "jnt_actfrcrange" in manager.domain_randomization_fields
|
||||
assert "tendon_actfrcrange" in manager.domain_randomization_fields
|
||||
assert len(manager.domain_randomization_fields) == 7
|
||||
|
||||
|
||||
def test_recompute_level_ordering():
|
||||
@@ -418,6 +420,62 @@ def test_effort_limits_scale_no_accumulation(device):
|
||||
assert abs(actual_upper - 200.0) < 1e-5
|
||||
|
||||
|
||||
def test_pd_gains_accepts_operation_object(device):
|
||||
"""dr.scale / dr.abs Operation objects produce the same result as strings."""
|
||||
env_str, ideal_str = _make_pd_env(device)
|
||||
env_obj, ideal_obj = _make_pd_env(device)
|
||||
|
||||
ids = torch.tensor([0], device=device)
|
||||
kwargs = dict(
|
||||
kp_range=(1.5, 1.5), kd_range=(2.0, 2.0), asset_cfg=SceneEntityCfg("robot")
|
||||
)
|
||||
|
||||
torch.manual_seed(0)
|
||||
dr.pd_gains(env_str, ids, operation="scale", **kwargs)
|
||||
torch.manual_seed(0)
|
||||
dr.pd_gains(env_obj, ids, operation=dr.scale, **kwargs)
|
||||
|
||||
assert torch.allclose(
|
||||
env_str.sim.model.actuator_gainprm[0], env_obj.sim.model.actuator_gainprm[0]
|
||||
)
|
||||
assert torch.allclose(ideal_str.stiffness, ideal_obj.stiffness)
|
||||
|
||||
|
||||
def test_effort_limits_accepts_operation_object(device):
|
||||
"""dr.abs Operation object produces the same result as the string."""
|
||||
env_str, ideal_str = _make_effort_env(device)
|
||||
env_obj, ideal_obj = _make_effort_env(device)
|
||||
|
||||
ids = torch.tensor([0], device=device)
|
||||
kwargs = dict(effort_limit_range=(150.0, 150.0), asset_cfg=SceneEntityCfg("robot"))
|
||||
|
||||
dr.effort_limits(env_str, ids, operation="abs", **kwargs)
|
||||
dr.effort_limits(env_obj, ids, operation=dr.abs, **kwargs)
|
||||
|
||||
assert torch.allclose(
|
||||
env_str.sim.model.actuator_forcerange[0], env_obj.sim.model.actuator_forcerange[0]
|
||||
)
|
||||
assert torch.allclose(ideal_str.force_limit, ideal_obj.force_limit)
|
||||
|
||||
|
||||
def test_pd_gains_rejects_unsupported_operation(device):
|
||||
"""Operations other than scale/abs raise ValueError."""
|
||||
env, _ = _make_pd_env(device)
|
||||
ids = torch.tensor([0], device=device)
|
||||
|
||||
with pytest.raises(ValueError, match="only supports 'scale' and 'abs'"):
|
||||
dr.pd_gains(env, ids, kp_range=(1.0, 1.0), kd_range=(1.0, 1.0), operation=dr.add)
|
||||
|
||||
|
||||
def test_effort_limits_rejects_unsupported_operation(device):
|
||||
"""Operations other than scale/abs raise ValueError."""
|
||||
env, _ = _make_effort_env(device)
|
||||
ids = torch.tensor([0], device=device)
|
||||
|
||||
with pytest.raises(ValueError, match="only supports 'scale' and 'abs'"):
|
||||
dr.effort_limits(env, ids, effort_limit_range=(1.0, 1.0), operation=dr.add)
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# Section 3: Other events
|
||||
# ===========================================================================
|
||||
@@ -503,7 +561,9 @@ def test_step_mode_fires_every_call(device):
|
||||
assert call_count[0] == 5
|
||||
|
||||
|
||||
def _make_impulse_env(device, num_envs=2, num_bodies=1, body_ids=None):
|
||||
def _make_impulse_env(
|
||||
device, num_envs=2, num_bodies=1, body_ids=None, cooldown_s=(0.0, 0.0)
|
||||
):
|
||||
"""Create a mock env for apply_body_impulse tests."""
|
||||
if body_ids is None:
|
||||
body_ids = [0]
|
||||
@@ -523,7 +583,7 @@ def _make_impulse_env(device, num_envs=2, num_bodies=1, body_ids=None):
|
||||
|
||||
asset_cfg = SceneEntityCfg("robot", body_ids=body_ids)
|
||||
term_cfg = Mock()
|
||||
term_cfg.params = {"asset_cfg": asset_cfg}
|
||||
term_cfg.params = {"asset_cfg": asset_cfg, "cooldown_s": cooldown_s}
|
||||
impulse = events.apply_body_impulse(cfg=term_cfg, env=env)
|
||||
return env, mock_entity, asset_cfg, impulse
|
||||
|
||||
@@ -531,11 +591,13 @@ def _make_impulse_env(device, num_envs=2, num_bodies=1, body_ids=None):
|
||||
def test_apply_body_impulse_basic(device):
|
||||
"""Impulse is applied and cleared after duration expires."""
|
||||
env, mock_entity, asset_cfg, impulse = _make_impulse_env(
|
||||
device, num_envs=2, num_bodies=3, body_ids=[1]
|
||||
device, num_envs=2, num_bodies=3, body_ids=[1], cooldown_s=(10.0, 10.0)
|
||||
)
|
||||
|
||||
# First call: cooldown_s starts at 0 and gets decremented by dt,
|
||||
# so it becomes <= 0 and triggers.
|
||||
# Skip the initial cooldown so the first call triggers immediately;
|
||||
# the trigger/sustain/expire cycle is what's under test here.
|
||||
impulse._interval_time_left[:] = 0.0
|
||||
|
||||
impulse(
|
||||
env,
|
||||
None,
|
||||
@@ -643,6 +705,43 @@ def test_apply_body_impulse_reset_clears(device):
|
||||
assert env_ids_arg[0].item() == 0
|
||||
|
||||
|
||||
def test_apply_body_impulse_initial_cooldown(device):
|
||||
"""The first call after init/reset enters cooldown, not an immediate impulse.
|
||||
|
||||
Regression test for #973.
|
||||
"""
|
||||
env, mock_entity, asset_cfg, impulse = _make_impulse_env(
|
||||
device, num_envs=1, num_bodies=1, body_ids=[0], cooldown_s=(0.05, 0.05)
|
||||
)
|
||||
|
||||
def step():
|
||||
impulse(
|
||||
env,
|
||||
None,
|
||||
force_range=(10.0, 10.0),
|
||||
torque_range=(0.0, 0.0),
|
||||
duration_s=(1.0, 1.0),
|
||||
cooldown_s=(0.05, 0.05), # ~2.5 steps at dt=0.02
|
||||
asset_cfg=asset_cfg,
|
||||
)
|
||||
|
||||
# First two steps consume the sampled cooldown; impulse must not fire yet.
|
||||
step()
|
||||
assert not impulse._active.any()
|
||||
step()
|
||||
assert not impulse._active.any()
|
||||
|
||||
# Third step crosses the cooldown boundary and triggers.
|
||||
step()
|
||||
assert impulse._active.all()
|
||||
|
||||
# Reset re-enters cooldown: next step should not immediately re-trigger.
|
||||
impulse.reset(env_ids=torch.tensor([0], device=device))
|
||||
assert not impulse._active.any()
|
||||
step()
|
||||
assert not impulse._active.any()
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# Section 5: Recomputation integration
|
||||
# ===========================================================================
|
||||
|
||||
@@ -113,3 +113,16 @@ def test_select_gpus_cpu_mode_empty_cuda_visible_devices():
|
||||
selected, num = select_gpus([0])
|
||||
assert selected is None
|
||||
assert num == 0
|
||||
|
||||
|
||||
def test_select_gpus_mig_uuids():
|
||||
"""Handles MIG GPU UUIDs in CUDA_VISIBLE_DEVICES."""
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = "MIG-GPU-abc-123,MIG-GPU-def-456"
|
||||
|
||||
selected, num = select_gpus("all")
|
||||
assert selected == ["MIG-GPU-abc-123", "MIG-GPU-def-456"]
|
||||
assert num == 2
|
||||
|
||||
selected, num = select_gpus([0])
|
||||
assert selected == ["MIG-GPU-abc-123"]
|
||||
assert num == 1
|
||||
|
||||
@@ -1,961 +0,0 @@
|
||||
"""Tests for per-world mesh variant support."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, cast
|
||||
|
||||
import mujoco
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from mjlab.entity import EntityCfg, VariantCfg, VariantEntityCfg
|
||||
from mjlab.sim.mesh_variants import allocate_worlds, build_mesh_variant_model
|
||||
from mjlab.viewer.model_sync import (
|
||||
disable_model_sameframe_shortcuts,
|
||||
sync_model_fields,
|
||||
)
|
||||
|
||||
# Helpers: variant specs with visual + collision mesh geoms.
|
||||
|
||||
|
||||
def _sphere_2col_spec() -> mujoco.MjSpec:
|
||||
"""Sphere: 1 visual + 2 collision geoms."""
|
||||
spec = mujoco.MjSpec()
|
||||
mv = spec.add_mesh()
|
||||
mv.name = "visual"
|
||||
mv.make_sphere(subdivision=3)
|
||||
for i in range(2):
|
||||
mc = spec.add_mesh()
|
||||
mc.name = f"col_{i}"
|
||||
mc.make_sphere(subdivision=1)
|
||||
body = spec.worldbody.add_body()
|
||||
body.name = "prop"
|
||||
body.add_freejoint()
|
||||
gv = body.add_geom()
|
||||
gv.name = "visual"
|
||||
gv.type = mujoco.mjtGeom.mjGEOM_MESH
|
||||
gv.meshname = "visual"
|
||||
gv.contype = 0
|
||||
gv.conaffinity = 0
|
||||
for i in range(2):
|
||||
gc = body.add_geom()
|
||||
gc.name = f"col_{i}"
|
||||
gc.type = mujoco.mjtGeom.mjGEOM_MESH
|
||||
gc.meshname = f"col_{i}"
|
||||
return spec
|
||||
|
||||
|
||||
def _cone_4col_spec() -> mujoco.MjSpec:
|
||||
"""Cone: 1 visual + 4 collision geoms (more than sphere)."""
|
||||
spec = mujoco.MjSpec()
|
||||
mv = spec.add_mesh()
|
||||
mv.name = "visual"
|
||||
mv.make_cone(nedge=8, radius=0.05)
|
||||
for i in range(4):
|
||||
mc = spec.add_mesh()
|
||||
mc.name = f"col_{i}"
|
||||
mc.make_sphere(subdivision=1)
|
||||
body = spec.worldbody.add_body()
|
||||
body.name = "prop"
|
||||
body.add_freejoint()
|
||||
gv = body.add_geom()
|
||||
gv.name = "visual"
|
||||
gv.type = mujoco.mjtGeom.mjGEOM_MESH
|
||||
gv.meshname = "visual"
|
||||
gv.contype = 0
|
||||
gv.conaffinity = 0
|
||||
for i in range(4):
|
||||
gc = body.add_geom()
|
||||
gc.name = f"col_{i}"
|
||||
gc.type = mujoco.mjtGeom.mjGEOM_MESH
|
||||
gc.meshname = f"col_{i}"
|
||||
return spec
|
||||
|
||||
|
||||
def _simple_sphere_spec() -> mujoco.MjSpec:
|
||||
"""Single-geom sphere for simple tests."""
|
||||
spec = mujoco.MjSpec()
|
||||
m = spec.add_mesh()
|
||||
m.name = "sphere"
|
||||
m.make_sphere(subdivision=2)
|
||||
body = spec.worldbody.add_body()
|
||||
body.name = "prop"
|
||||
body.add_freejoint()
|
||||
g = body.add_geom()
|
||||
g.name = "visual"
|
||||
g.type = mujoco.mjtGeom.mjGEOM_MESH
|
||||
g.meshname = "sphere"
|
||||
return spec
|
||||
|
||||
|
||||
def _simple_cone_spec() -> mujoco.MjSpec:
|
||||
"""Single-geom cone for simple tests."""
|
||||
spec = mujoco.MjSpec()
|
||||
m = spec.add_mesh()
|
||||
m.name = "cone"
|
||||
m.make_cone(nedge=8, radius=0.05)
|
||||
body = spec.worldbody.add_body()
|
||||
body.name = "prop"
|
||||
body.add_freejoint()
|
||||
g = body.add_geom()
|
||||
g.name = "visual"
|
||||
g.type = mujoco.mjtGeom.mjGEOM_MESH
|
||||
g.meshname = "cone"
|
||||
return spec
|
||||
|
||||
|
||||
def _hinge_spec() -> mujoco.MjSpec:
|
||||
"""Object with a hinge joint (incompatible with freejoint variants)."""
|
||||
spec = mujoco.MjSpec()
|
||||
m = spec.add_mesh()
|
||||
m.name = "box"
|
||||
m.make_sphere(subdivision=1)
|
||||
body = spec.worldbody.add_body()
|
||||
body.name = "prop"
|
||||
j = body.add_joint()
|
||||
j.name = "hinge"
|
||||
j.type = mujoco.mjtJoint.mjJNT_HINGE
|
||||
g = body.add_geom()
|
||||
g.name = "visual"
|
||||
g.type = mujoco.mjtGeom.mjGEOM_MESH
|
||||
g.meshname = "box"
|
||||
return spec
|
||||
|
||||
|
||||
def _build_scene_with_variants(
|
||||
variant_a_fn, variant_b_fn, *, weight_a=0.5, weight_b=0.5
|
||||
):
|
||||
"""Build a scene spec + variant_info from two variant spec_fns."""
|
||||
cfg = VariantEntityCfg(
|
||||
variants={
|
||||
"a": VariantCfg(spec_fn=variant_a_fn, weight=weight_a),
|
||||
"b": VariantCfg(spec_fn=variant_b_fn, weight=weight_b),
|
||||
},
|
||||
)
|
||||
entity = cfg.build()
|
||||
assert entity.variant_metadata is not None
|
||||
scene_spec = mujoco.MjSpec()
|
||||
frame = scene_spec.worldbody.add_frame()
|
||||
scene_spec.attach(entity.spec, prefix="object/", frame=frame)
|
||||
return scene_spec, [("object/", entity.variant_metadata)]
|
||||
|
||||
|
||||
# allocate_worlds.
|
||||
|
||||
|
||||
def test_allocate_worlds_proportional():
|
||||
result = allocate_worlds((0.6, 0.4), 10)
|
||||
assert len(result) == 10
|
||||
assert result.count(0) == 6
|
||||
assert result.count(1) == 4
|
||||
|
||||
|
||||
def test_allocate_worlds_uniform():
|
||||
result = allocate_worlds((1.0, 1.0), 8)
|
||||
assert result.count(0) == 4
|
||||
assert result.count(1) == 4
|
||||
|
||||
|
||||
def test_allocate_worlds_single_variant():
|
||||
result = allocate_worlds((1.0,), 5)
|
||||
assert result == [0, 0, 0, 0, 0]
|
||||
|
||||
|
||||
def test_allocate_worlds_zero_weight_skips_variant():
|
||||
"""A zero-weight variant gets zero worlds; the rest split nworld."""
|
||||
result = allocate_worlds((1.0, 0.0, 1.0), 10)
|
||||
assert len(result) == 10
|
||||
assert result.count(1) == 0
|
||||
assert result.count(0) == 5
|
||||
assert result.count(2) == 5
|
||||
|
||||
|
||||
def test_allocate_worlds_rejects_negative_weight():
|
||||
with pytest.raises(ValueError, match="non-negative"):
|
||||
allocate_worlds((1.0, -0.1), 10)
|
||||
|
||||
|
||||
def test_allocate_worlds_rejects_all_zero():
|
||||
with pytest.raises(ValueError, match="positive sum"):
|
||||
allocate_worlds((0.0, 0.0), 10)
|
||||
|
||||
|
||||
def test_allocate_worlds_largest_remainder_sums_to_nworld():
|
||||
"""Largest-remainder rounding must always allocate exactly nworld worlds."""
|
||||
for nworld in (3, 7, 100, 1000):
|
||||
result = allocate_worlds((1.0, 1.0, 1.0), nworld)
|
||||
assert len(result) == nworld
|
||||
# Difference between any two variant counts is at most 1 (uniform).
|
||||
counts = [result.count(i) for i in range(3)]
|
||||
assert max(counts) - min(counts) <= 1
|
||||
|
||||
|
||||
# Entity merging.
|
||||
|
||||
|
||||
def test_entity_builds_with_variants():
|
||||
cfg = VariantEntityCfg(
|
||||
variants={
|
||||
"sphere": VariantCfg(spec_fn=_simple_sphere_spec, weight=0.5),
|
||||
"cone": VariantCfg(spec_fn=_simple_cone_spec, weight=0.5),
|
||||
},
|
||||
)
|
||||
entity = cfg.build()
|
||||
meta = entity.variant_metadata
|
||||
assert meta is not None
|
||||
assert meta.variant_names == ("sphere", "cone")
|
||||
assert meta.num_mesh_geoms == 1
|
||||
mesh_names = [m.name for m in entity.spec.meshes]
|
||||
assert any("sphere" in n for n in mesh_names)
|
||||
assert any("cone" in n for n in mesh_names)
|
||||
|
||||
|
||||
def test_multi_geom_body_padding():
|
||||
"""Sphere (3 geoms) + cone (5 geoms) -> body padded to 5 mesh geoms."""
|
||||
cfg = VariantEntityCfg(
|
||||
variants={
|
||||
"sphere": VariantCfg(spec_fn=_sphere_2col_spec, weight=0.5),
|
||||
"cone": VariantCfg(spec_fn=_cone_4col_spec, weight=0.5),
|
||||
},
|
||||
)
|
||||
entity = cfg.build()
|
||||
meta = entity.variant_metadata
|
||||
assert meta is not None
|
||||
assert meta.num_mesh_geoms == 5 # max(3, 5)
|
||||
# Sphere: 3 real + 2 padding (None).
|
||||
assert sum(1 for n in meta.variant_mesh_names[0] if n is None) == 2
|
||||
# Cone: 5 real, no padding.
|
||||
assert all(n is not None for n in meta.variant_mesh_names[1])
|
||||
|
||||
|
||||
# Validation.
|
||||
|
||||
|
||||
def test_mismatched_joint_structure_raises():
|
||||
cfg = VariantEntityCfg(
|
||||
variants={
|
||||
"sphere": VariantCfg(spec_fn=_simple_sphere_spec, weight=0.5),
|
||||
"hinge": VariantCfg(spec_fn=_hinge_spec, weight=0.5),
|
||||
},
|
||||
)
|
||||
with pytest.raises(ValueError, match="joint"):
|
||||
cfg.build()
|
||||
|
||||
|
||||
def test_single_variant_builds():
|
||||
"""A single variant degenerates cleanly; useful for templated variant sets."""
|
||||
cfg = VariantEntityCfg(
|
||||
variants={"only": VariantCfg(spec_fn=_simple_sphere_spec)},
|
||||
)
|
||||
entity = cfg.build()
|
||||
assert entity.variant_metadata is not None
|
||||
assert entity.variant_metadata.variant_names == ("only",)
|
||||
|
||||
|
||||
def test_empty_variants_raises():
|
||||
cfg = VariantEntityCfg(variants={})
|
||||
with pytest.raises(ValueError, match="at least one"):
|
||||
cfg.build()
|
||||
|
||||
|
||||
def _fixed_base_sphere_spec() -> mujoco.MjSpec:
|
||||
"""Fixed-base sphere variant (no free joint): currently unsupported."""
|
||||
spec = mujoco.MjSpec()
|
||||
m = spec.add_mesh(name="sphere")
|
||||
m.make_sphere(subdivision=2)
|
||||
body = spec.worldbody.add_body(name="prop")
|
||||
body.add_geom(type=mujoco.mjtGeom.mjGEOM_MESH, meshname="sphere")
|
||||
return spec
|
||||
|
||||
|
||||
def test_fixed_base_variants_rejected():
|
||||
"""Variants must be floating-base; fixed-base raises with a clear message."""
|
||||
cfg = VariantEntityCfg(
|
||||
variants={
|
||||
"a": VariantCfg(spec_fn=_fixed_base_sphere_spec, weight=0.5),
|
||||
"b": VariantCfg(spec_fn=_fixed_base_sphere_spec, weight=0.5),
|
||||
},
|
||||
)
|
||||
with pytest.raises(ValueError, match="floating-base"):
|
||||
cfg.build()
|
||||
|
||||
|
||||
def test_setting_spec_fn_on_variant_cfg_raises():
|
||||
"""VariantEntityCfg.spec_fn is unused; setting it should fail loudly."""
|
||||
with pytest.raises(ValueError, match="spec_fn cannot be set"):
|
||||
VariantEntityCfg(
|
||||
variants={"only": VariantCfg(spec_fn=_simple_sphere_spec)},
|
||||
spec_fn=_simple_sphere_spec,
|
||||
)
|
||||
|
||||
|
||||
def test_no_variants_unchanged():
|
||||
cfg = EntityCfg(spec_fn=_simple_sphere_spec)
|
||||
entity = cfg.build()
|
||||
assert entity.variant_metadata is None
|
||||
|
||||
|
||||
# build_mesh_variant_model: dataid and dependent fields.
|
||||
|
||||
|
||||
def test_dataid_assigned_per_world():
|
||||
"""Each world's geom_dataid points to its variant's meshes."""
|
||||
scene_spec, vi = _build_scene_with_variants(_simple_sphere_spec, _simple_cone_spec)
|
||||
result = build_mesh_variant_model(scene_spec, 4, vi)
|
||||
|
||||
dataid = result.wp_model.geom_dataid.numpy()
|
||||
assert dataid.shape == (4, result.mj_model.ngeom)
|
||||
assert dataid.ndim == 2
|
||||
|
||||
w2v = result.world_to_variant["object/"]
|
||||
assert w2v[0] == 0 # variant a (sphere)
|
||||
assert w2v[2] == 1 # variant b (cone)
|
||||
|
||||
# Sphere and cone worlds must have different dataid values.
|
||||
assert not np.array_equal(dataid[0], dataid[2])
|
||||
|
||||
|
||||
def test_padding_slots_get_disabled():
|
||||
"""Shorter variant's padding geom slots have dataid == -1."""
|
||||
scene_spec, vi = _build_scene_with_variants(_sphere_2col_spec, _cone_4col_spec)
|
||||
result = build_mesh_variant_model(scene_spec, 4, vi)
|
||||
|
||||
dataid = result.wp_model.geom_dataid.numpy()
|
||||
w2v = result.world_to_variant["object/"]
|
||||
|
||||
# Find a sphere world (variant 0, 3 mesh geoms -> 2 padding slots).
|
||||
sphere_world = int(np.where(w2v == 0)[0][0])
|
||||
# Find mesh geom columns (skip non-mesh geoms like worldbody).
|
||||
mesh_geom_ids = [
|
||||
gid
|
||||
for gid in range(result.mj_model.ngeom)
|
||||
if result.mj_model.geom_type[gid] == mujoco.mjtGeom.mjGEOM_MESH
|
||||
]
|
||||
sphere_dataid = dataid[sphere_world, mesh_geom_ids]
|
||||
# Last 2 mesh geom slots should be -1 (disabled padding).
|
||||
assert sphere_dataid[-1] == -1
|
||||
assert sphere_dataid[-2] == -1
|
||||
# Padding slots must still be collision-enabled in the template/warp model.
|
||||
# Short variants are disabled by per-world dataid=-1; long variants need the
|
||||
# same slots enabled so their extra hulls can collide.
|
||||
assert np.all(result.mj_model.geom_contype[mesh_geom_ids[-2:]] == 1)
|
||||
assert np.all(result.mj_model.geom_conaffinity[mesh_geom_ids[-2:]] == 1)
|
||||
assert np.all(result.wp_model.geom_contype.numpy()[mesh_geom_ids[-2:]] == 1)
|
||||
assert np.all(result.wp_model.geom_conaffinity.numpy()[mesh_geom_ids[-2:]] == 1)
|
||||
# First 3 should be valid (>= 0).
|
||||
assert all(d >= 0 for d in sphere_dataid[:3])
|
||||
|
||||
|
||||
def test_dependent_fields_match_individual_compilation():
|
||||
"""Per-world body_mass matches independently compiled variant models."""
|
||||
scene_spec, vi = _build_scene_with_variants(_simple_sphere_spec, _simple_cone_spec)
|
||||
result = build_mesh_variant_model(scene_spec, 4, vi)
|
||||
|
||||
# Compile each variant independently for reference values.
|
||||
sphere_model = _simple_sphere_spec().compile()
|
||||
cone_model = _simple_cone_spec().compile()
|
||||
|
||||
body_mass = result.wp_model.body_mass.numpy()
|
||||
w2v = result.world_to_variant["object/"]
|
||||
|
||||
sphere_w = int(np.where(w2v == 0)[0][0])
|
||||
cone_w = int(np.where(w2v == 1)[0][0])
|
||||
|
||||
# The object body is the last body in the scene.
|
||||
obj_body = result.mj_model.nbody - 1
|
||||
|
||||
# Mass should match individually compiled models.
|
||||
np.testing.assert_allclose(
|
||||
body_mass[sphere_w, obj_body],
|
||||
sphere_model.body_mass[-1],
|
||||
atol=1e-4,
|
||||
)
|
||||
np.testing.assert_allclose(
|
||||
body_mass[cone_w, obj_body],
|
||||
cone_model.body_mass[-1],
|
||||
atol=1e-4,
|
||||
)
|
||||
|
||||
# Sphere and cone should have different masses.
|
||||
assert not np.isclose(body_mass[sphere_w, obj_body], body_mass[cone_w, obj_body])
|
||||
|
||||
|
||||
def test_select_default_values_uses_per_world_variant_defaults():
|
||||
"""Per-world defaults are indexed by env first, then by entity."""
|
||||
from mjlab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg
|
||||
from mjlab.envs.mdp.dr._core import _select_default_values
|
||||
from mjlab.scene import SceneCfg
|
||||
from mjlab.terrains import TerrainEntityCfg
|
||||
|
||||
def _explicit_variant(
|
||||
mesh_name: str,
|
||||
mass: float,
|
||||
inertia: tuple[float, float, float],
|
||||
*,
|
||||
cone: bool = False,
|
||||
) -> mujoco.MjSpec:
|
||||
spec = mujoco.MjSpec()
|
||||
mesh = spec.add_mesh()
|
||||
mesh.name = mesh_name
|
||||
if cone:
|
||||
mesh.make_cone(nedge=8, radius=0.05)
|
||||
else:
|
||||
mesh.make_sphere(subdivision=1)
|
||||
body = spec.worldbody.add_body(name="prop")
|
||||
body.add_freejoint()
|
||||
body.explicitinertial = 1
|
||||
body.mass = mass
|
||||
body.ipos[:] = (0.0, 0.0, 0.0)
|
||||
body.inertia[:] = inertia
|
||||
body.iquat[:] = (1.0, 0.0, 0.0, 0.0)
|
||||
body.add_geom(
|
||||
name="visual",
|
||||
type=mujoco.mjtGeom.mjGEOM_MESH,
|
||||
meshname=mesh_name,
|
||||
contype=0,
|
||||
conaffinity=0,
|
||||
mass=0.0,
|
||||
)
|
||||
return spec
|
||||
|
||||
object_cfg = VariantEntityCfg(
|
||||
variants={
|
||||
"sphere": VariantCfg(
|
||||
lambda: _explicit_variant("sphere", 0.2, (1e-4, 2e-4, 3e-4)),
|
||||
weight=0.5,
|
||||
),
|
||||
"cone": VariantCfg(
|
||||
lambda: _explicit_variant("cone", 0.7, (4e-4, 5e-4, 6e-4), cone=True),
|
||||
weight=0.5,
|
||||
),
|
||||
},
|
||||
init_state=EntityCfg.InitialStateCfg(pos=(0.0, 0.0, 0.2)),
|
||||
)
|
||||
env_cfg = ManagerBasedRlEnvCfg(
|
||||
decimation=1,
|
||||
scene=SceneCfg(
|
||||
terrain=TerrainEntityCfg(terrain_type="plane"),
|
||||
num_envs=4,
|
||||
env_spacing=1.0,
|
||||
entities={"object": object_cfg},
|
||||
),
|
||||
)
|
||||
|
||||
env = ManagerBasedRlEnv(cfg=env_cfg, device="cpu")
|
||||
try:
|
||||
obj_body = int(env.scene["object"].indexing.root_body_id)
|
||||
env_ids = torch.arange(env.num_envs, device=env.device)
|
||||
body_ids = torch.tensor([obj_body], device=env.device)
|
||||
|
||||
for field in ("body_mass", "body_ipos", "body_inertia", "body_iquat"):
|
||||
selected = _select_default_values(env, field, env_ids, body_ids)
|
||||
torch.testing.assert_close(
|
||||
selected[:, 0],
|
||||
getattr(env.sim.model, field)[:, obj_body],
|
||||
)
|
||||
finally:
|
||||
env.close()
|
||||
|
||||
|
||||
def test_viser_builds_per_world_mesh_handles_for_variants():
|
||||
"""Viser dynamic meshes must not collapse all worlds onto env0's mesh."""
|
||||
from contextlib import nullcontext
|
||||
|
||||
from mjlab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg
|
||||
from mjlab.scene import SceneCfg
|
||||
from mjlab.terrains import TerrainEntityCfg
|
||||
from mjlab.viewer.viser.scene import MjlabViserScene, _PerWorldMeshGroup
|
||||
|
||||
class _Handle:
|
||||
def __init__(self, **kwargs):
|
||||
self.visible = kwargs.get("visible", True)
|
||||
self.batched_positions = kwargs.get("batched_positions", np.zeros((0, 3)))
|
||||
self.batched_wxyzs = kwargs.get("batched_wxyzs", np.zeros((0, 4)))
|
||||
self.batched_scales = kwargs.get("batched_scales")
|
||||
self.batched_colors = kwargs.get("batched_colors")
|
||||
self.batched_opacities = kwargs.get("batched_opacities")
|
||||
self.position = kwargs.get("position", np.zeros(3))
|
||||
self.wxyz = kwargs.get("wxyz", np.array([1.0, 0.0, 0.0, 0.0]))
|
||||
|
||||
def remove(self) -> None:
|
||||
pass
|
||||
|
||||
class _Scene:
|
||||
def __init__(self):
|
||||
self.batched: list[tuple[tuple, dict, _Handle]] = []
|
||||
|
||||
def configure_environment_map(self, **_kwargs) -> None:
|
||||
pass
|
||||
|
||||
def add_frame(self, *_args, **kwargs) -> _Handle:
|
||||
return _Handle(**kwargs)
|
||||
|
||||
def add_grid(self, *_args, **kwargs) -> _Handle:
|
||||
return _Handle(**kwargs)
|
||||
|
||||
def add_mesh_trimesh(self, *_args, **kwargs) -> _Handle:
|
||||
return _Handle(**kwargs)
|
||||
|
||||
def add_batched_meshes_trimesh(self, *args, **kwargs) -> _Handle:
|
||||
handle = _Handle(**kwargs)
|
||||
self.batched.append((args, kwargs, handle))
|
||||
return handle
|
||||
|
||||
def add_batched_meshes_simple(self, *args, **kwargs) -> _Handle:
|
||||
handle = _Handle(**kwargs)
|
||||
self.batched.append((args, kwargs, handle))
|
||||
return handle
|
||||
|
||||
class _Server:
|
||||
def __init__(self):
|
||||
self.scene = _Scene()
|
||||
|
||||
def atomic(self):
|
||||
return nullcontext()
|
||||
|
||||
def flush(self) -> None:
|
||||
pass
|
||||
|
||||
env_cfg = ManagerBasedRlEnvCfg(
|
||||
decimation=1,
|
||||
scene=SceneCfg(
|
||||
terrain=TerrainEntityCfg(terrain_type="plane"),
|
||||
num_envs=4,
|
||||
env_spacing=1.0,
|
||||
entities={
|
||||
"object": VariantEntityCfg(
|
||||
variants={
|
||||
"sphere": VariantCfg(_simple_sphere_spec, weight=0.5),
|
||||
"cone": VariantCfg(_simple_cone_spec, weight=0.5),
|
||||
},
|
||||
init_state=EntityCfg.InitialStateCfg(pos=(0.0, 0.0, 0.2)),
|
||||
)
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
env = ManagerBasedRlEnv(cfg=env_cfg, device="cpu")
|
||||
try:
|
||||
env.sim.expand_model_fields(("geom_rgba",))
|
||||
env.sim.model.geom_rgba[:, :, :3] = torch.linspace(
|
||||
0.2,
|
||||
0.9,
|
||||
env.num_envs,
|
||||
device=env.device,
|
||||
)[:, None, None]
|
||||
server = _Server()
|
||||
scene = MjlabViserScene(
|
||||
cast(Any, server),
|
||||
env.sim.mj_model,
|
||||
env.num_envs,
|
||||
sim_model=env.sim.model,
|
||||
expanded_fields=env.sim.expanded_fields,
|
||||
)
|
||||
groups = [mg for mg in scene._mesh_groups if isinstance(mg, _PerWorldMeshGroup)]
|
||||
|
||||
assert groups
|
||||
assert sum(len(mg.env_ids) for mg in groups) >= env.num_envs
|
||||
|
||||
body_xpos = env.sim.data.xpos.cpu().numpy()
|
||||
body_xmat = env.sim.data.xmat.cpu().numpy()
|
||||
mocap_pos = (
|
||||
env.sim.data.mocap_pos.cpu().numpy() if env.sim.mj_model.nmocap > 0 else None
|
||||
)
|
||||
mocap_quat = (
|
||||
env.sim.data.mocap_quat.cpu().numpy() if env.sim.mj_model.nmocap > 0 else None
|
||||
)
|
||||
scene.show_only_selected = True
|
||||
scene.update_from_arrays(body_xpos, body_xmat, mocap_pos, mocap_quat, env_idx=0)
|
||||
scene.update_from_arrays(body_xpos, body_xmat, mocap_pos, mocap_quat, env_idx=1)
|
||||
|
||||
assert any(mg.handle.visible for mg in groups)
|
||||
|
||||
handle_count = len(server.scene.batched)
|
||||
env.sim.model.geom_rgba[:, :, :3] = torch.linspace(
|
||||
0.9,
|
||||
0.2,
|
||||
env.num_envs,
|
||||
device=env.device,
|
||||
)[:, None, None]
|
||||
scene.update_from_arrays(body_xpos, body_xmat, mocap_pos, mocap_quat, env_idx=0)
|
||||
assert len(server.scene.batched) > handle_count
|
||||
finally:
|
||||
env.close()
|
||||
|
||||
|
||||
def test_viser_convex_hulls_are_per_variant():
|
||||
"""Convex-hull handles must differ across variants, not all show env0's hull."""
|
||||
from contextlib import nullcontext
|
||||
|
||||
from mjlab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg
|
||||
from mjlab.scene import SceneCfg
|
||||
from mjlab.terrains import TerrainEntityCfg
|
||||
from mjlab.viewer.viser.scene import MjlabViserScene, _PerWorldHullGroup
|
||||
|
||||
class _Handle:
|
||||
def __init__(self, **kwargs):
|
||||
self.visible = kwargs.get("visible", True)
|
||||
self.batched_positions = kwargs.get("batched_positions", np.zeros((0, 3)))
|
||||
self.batched_wxyzs = kwargs.get("batched_wxyzs", np.zeros((0, 4)))
|
||||
self.batched_scales = kwargs.get("batched_scales")
|
||||
self.batched_colors = kwargs.get("batched_colors")
|
||||
self.batched_opacities = kwargs.get("batched_opacities")
|
||||
self.position = kwargs.get("position", np.zeros(3))
|
||||
self.wxyz = kwargs.get("wxyz", np.array([1.0, 0.0, 0.0, 0.0]))
|
||||
self.vertices = kwargs.get("vertices")
|
||||
self.faces = kwargs.get("faces")
|
||||
|
||||
def remove(self) -> None:
|
||||
pass
|
||||
|
||||
class _Scene:
|
||||
def __init__(self):
|
||||
self.batched: list[tuple[tuple, dict, _Handle]] = []
|
||||
|
||||
def configure_environment_map(self, **_kwargs) -> None:
|
||||
pass
|
||||
|
||||
def add_frame(self, *_args, **kwargs) -> _Handle:
|
||||
return _Handle(**kwargs)
|
||||
|
||||
def add_grid(self, *_args, **kwargs) -> _Handle:
|
||||
return _Handle(**kwargs)
|
||||
|
||||
def add_mesh_trimesh(self, *_args, **kwargs) -> _Handle:
|
||||
return _Handle(**kwargs)
|
||||
|
||||
def add_batched_meshes_trimesh(self, *args, **kwargs) -> _Handle:
|
||||
handle = _Handle(**kwargs)
|
||||
self.batched.append((args, kwargs, handle))
|
||||
return handle
|
||||
|
||||
def add_batched_meshes_simple(self, path, vertices, faces, **kwargs) -> _Handle:
|
||||
# Capture the mesh identity so the test can compare hull shapes.
|
||||
kwargs = dict(kwargs)
|
||||
kwargs["vertices"] = np.asarray(vertices)
|
||||
kwargs["faces"] = np.asarray(faces)
|
||||
handle = _Handle(**kwargs)
|
||||
self.batched.append(((path,), kwargs, handle))
|
||||
return handle
|
||||
|
||||
class _Server:
|
||||
def __init__(self):
|
||||
self.scene = _Scene()
|
||||
|
||||
def atomic(self):
|
||||
return nullcontext()
|
||||
|
||||
def flush(self) -> None:
|
||||
pass
|
||||
|
||||
# Sphere and cone produce visibly different convex hulls.
|
||||
env_cfg = ManagerBasedRlEnvCfg(
|
||||
decimation=1,
|
||||
scene=SceneCfg(
|
||||
terrain=TerrainEntityCfg(terrain_type="plane"),
|
||||
num_envs=4,
|
||||
env_spacing=1.0,
|
||||
entities={
|
||||
"object": VariantEntityCfg(
|
||||
variants={
|
||||
"sphere": VariantCfg(_simple_sphere_spec, weight=0.5),
|
||||
"cone": VariantCfg(_simple_cone_spec, weight=0.5),
|
||||
},
|
||||
init_state=EntityCfg.InitialStateCfg(pos=(0.0, 0.0, 0.2)),
|
||||
)
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
env = ManagerBasedRlEnv(cfg=env_cfg, device="cpu")
|
||||
try:
|
||||
server = _Server()
|
||||
scene = MjlabViserScene(
|
||||
cast(Any, server),
|
||||
env.sim.mj_model,
|
||||
env.num_envs,
|
||||
sim_model=env.sim.model,
|
||||
expanded_fields=env.sim.expanded_fields,
|
||||
)
|
||||
groups: list[_PerWorldHullGroup] = list(scene._hull_per_world_groups)
|
||||
# Two distinct variants -> at least two hull handles on the same body.
|
||||
assert len(groups) >= 2, f"expected >=2 hull variants, got {len(groups)}"
|
||||
all_envs = np.concatenate([g.env_ids for g in groups])
|
||||
assert sorted(all_envs.tolist()) == list(range(env.num_envs))
|
||||
# Hulls must be shape-distinct, not all copies of env0's hull.
|
||||
shapes = {(g.handle.vertices.shape, g.handle.faces.shape) for g in groups}
|
||||
assert len(shapes) >= 2, (
|
||||
f"hull variants collapsed to one shape: {shapes} "
|
||||
"(all envs would share env0's hull)"
|
||||
)
|
||||
|
||||
body_xpos = env.sim.data.xpos.cpu().numpy()
|
||||
body_xmat = env.sim.data.xmat.cpu().numpy()
|
||||
scene.show_convex_hull = True
|
||||
scene.show_only_selected = True
|
||||
for target_env in range(env.num_envs):
|
||||
scene.update_from_arrays(body_xpos, body_xmat, env_idx=target_env)
|
||||
visible_groups = [g for g in groups if g.handle.visible]
|
||||
assert len(visible_groups) == 1
|
||||
assert target_env in visible_groups[0].env_ids
|
||||
assert visible_groups[0].handle.batched_positions.shape[0] == 1
|
||||
|
||||
scene.show_only_selected = False
|
||||
scene.update_from_arrays(body_xpos, body_xmat, env_idx=0)
|
||||
assert all(g.handle.visible for g in groups)
|
||||
finally:
|
||||
env.close()
|
||||
|
||||
|
||||
# DR consistency on variant scenes.
|
||||
|
||||
|
||||
def _explicit_mass_variant(
|
||||
mesh_name: str,
|
||||
mass: float,
|
||||
*,
|
||||
cone: bool = False,
|
||||
) -> mujoco.MjSpec:
|
||||
"""Build a single-geom freejoint variant with an explicit body mass."""
|
||||
spec = mujoco.MjSpec()
|
||||
mesh = spec.add_mesh()
|
||||
mesh.name = mesh_name
|
||||
if cone:
|
||||
mesh.make_cone(nedge=8, radius=0.05)
|
||||
else:
|
||||
mesh.make_sphere(subdivision=1)
|
||||
body = spec.worldbody.add_body(name="prop")
|
||||
body.add_freejoint()
|
||||
body.explicitinertial = 1
|
||||
body.mass = mass
|
||||
body.ipos[:] = (0.0, 0.0, 0.0)
|
||||
body.inertia[:] = (1e-4, 1e-4, 1e-4)
|
||||
body.iquat[:] = (1.0, 0.0, 0.0, 0.0)
|
||||
body.add_geom(
|
||||
name="visual",
|
||||
type=mujoco.mjtGeom.mjGEOM_MESH,
|
||||
meshname=mesh_name,
|
||||
contype=0,
|
||||
conaffinity=0,
|
||||
mass=0.0,
|
||||
)
|
||||
return spec
|
||||
|
||||
|
||||
def test_dr_body_mass_scale_preserves_variant_baseline():
|
||||
"""``dr.body_mass`` scale must use each variant's own baseline.
|
||||
|
||||
This is the load-bearing claim of ``_per_world_default_fields``: scaling
|
||||
body_mass on a variant scene by a per-env factor must produce
|
||||
``variant_default[env] * scale[env]``, not ``template_default * scale[env]``.
|
||||
"""
|
||||
from mjlab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg
|
||||
from mjlab.envs.mdp import dr
|
||||
from mjlab.managers.event_manager import EventTermCfg
|
||||
from mjlab.managers.scene_entity_config import SceneEntityCfg
|
||||
from mjlab.scene import SceneCfg
|
||||
from mjlab.terrains import TerrainEntityCfg
|
||||
|
||||
light_mass = 0.1
|
||||
heavy_mass = 1.0
|
||||
scale = 2.0
|
||||
|
||||
object_cfg = VariantEntityCfg(
|
||||
variants={
|
||||
"light": VariantCfg(
|
||||
lambda: _explicit_mass_variant("light", light_mass), weight=0.5
|
||||
),
|
||||
"heavy": VariantCfg(
|
||||
lambda: _explicit_mass_variant("heavy", heavy_mass, cone=True), weight=0.5
|
||||
),
|
||||
},
|
||||
init_state=EntityCfg.InitialStateCfg(pos=(0.0, 0.0, 0.2)),
|
||||
)
|
||||
env_cfg = ManagerBasedRlEnvCfg(
|
||||
decimation=1,
|
||||
scene=SceneCfg(
|
||||
terrain=TerrainEntityCfg(terrain_type="plane"),
|
||||
num_envs=4,
|
||||
env_spacing=1.0,
|
||||
entities={"object": object_cfg},
|
||||
),
|
||||
events={
|
||||
"scale_mass": EventTermCfg(
|
||||
func=dr.body_mass,
|
||||
mode="startup",
|
||||
params={
|
||||
"asset_cfg": SceneEntityCfg("object", body_names=("prop",)),
|
||||
"operation": "scale",
|
||||
"ranges": (scale, scale), # deterministic factor
|
||||
},
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
with pytest.warns(UserWarning, match="dr.body_mass only randomizes mass"):
|
||||
env = ManagerBasedRlEnv(cfg=env_cfg, device="cpu")
|
||||
try:
|
||||
obj_body = int(env.scene["object"].indexing.root_body_id)
|
||||
w2v = env.sim.world_to_variant["object"]
|
||||
actual = env.sim.model.body_mass[:, obj_body].cpu()
|
||||
|
||||
variant_baseline = torch.tensor([light_mass, heavy_mass], dtype=actual.dtype)
|
||||
expected = variant_baseline[w2v.cpu()] * scale
|
||||
torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5)
|
||||
|
||||
# Sanity: at least one env per variant, otherwise the test is vacuous.
|
||||
assert (w2v == 0).any() and (w2v == 1).any()
|
||||
finally:
|
||||
env.close()
|
||||
|
||||
|
||||
# Full env lifecycle.
|
||||
|
||||
|
||||
def test_env_step_with_variants():
|
||||
"""Build a full ManagerBasedRlEnv with variants; step without crashing."""
|
||||
from mjlab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg
|
||||
from mjlab.envs.mdp.events import reset_root_state_uniform
|
||||
from mjlab.managers.event_manager import EventTermCfg
|
||||
from mjlab.managers.scene_entity_config import SceneEntityCfg
|
||||
from mjlab.scene import SceneCfg
|
||||
from mjlab.terrains import TerrainEntityCfg
|
||||
|
||||
object_cfg = VariantEntityCfg(
|
||||
variants={
|
||||
"sphere": VariantCfg(_simple_sphere_spec, weight=0.5),
|
||||
"cone": VariantCfg(_simple_cone_spec, weight=0.5),
|
||||
},
|
||||
init_state=EntityCfg.InitialStateCfg(pos=(0.0, 0.0, 0.2)),
|
||||
)
|
||||
|
||||
env_cfg = ManagerBasedRlEnvCfg(
|
||||
decimation=2,
|
||||
scene=SceneCfg(
|
||||
terrain=TerrainEntityCfg(terrain_type="plane"),
|
||||
num_envs=4,
|
||||
env_spacing=1.0,
|
||||
entities={"object": object_cfg},
|
||||
),
|
||||
events={
|
||||
"reset": EventTermCfg(
|
||||
func=reset_root_state_uniform,
|
||||
mode="reset",
|
||||
params={
|
||||
"pose_range": {},
|
||||
"velocity_range": {},
|
||||
"asset_cfg": SceneEntityCfg("object"),
|
||||
},
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
env = ManagerBasedRlEnv(cfg=env_cfg, device="cpu")
|
||||
obs, _ = env.reset()
|
||||
actions = torch.zeros(env.num_envs, 0)
|
||||
for _ in range(10):
|
||||
obs, rew, term, trunc, info = env.step(actions)
|
||||
# No NaN in positions.
|
||||
qpos = env.sim.data.qpos[:].cpu().numpy()
|
||||
assert np.all(np.isfinite(qpos))
|
||||
env.close()
|
||||
|
||||
|
||||
# Viewer: sameframe shortcut fix.
|
||||
|
||||
|
||||
def _viewer_regression_sphere_spec() -> mujoco.MjSpec:
|
||||
spec = mujoco.MjSpec()
|
||||
m = spec.add_mesh()
|
||||
m.name = "sphere"
|
||||
m.make_sphere(subdivision=3)
|
||||
m.scale[:] = (0.05, 0.05, 0.05)
|
||||
body = spec.worldbody.add_body()
|
||||
body.name = "prop"
|
||||
body.add_freejoint()
|
||||
g = body.add_geom()
|
||||
g.name = "visual"
|
||||
g.type = mujoco.mjtGeom.mjGEOM_MESH
|
||||
g.meshname = "sphere"
|
||||
return spec
|
||||
|
||||
|
||||
def _viewer_regression_cone_spec() -> mujoco.MjSpec:
|
||||
spec = mujoco.MjSpec()
|
||||
m = spec.add_mesh()
|
||||
m.name = "cone"
|
||||
m.make_cone(nedge=16, radius=0.04)
|
||||
m.scale[:] = (0.05, 0.05, 0.05)
|
||||
body = spec.worldbody.add_body()
|
||||
body.name = "prop"
|
||||
body.add_freejoint()
|
||||
g = body.add_geom()
|
||||
g.name = "visual"
|
||||
g.type = mujoco.mjtGeom.mjGEOM_MESH
|
||||
g.meshname = "cone"
|
||||
return spec
|
||||
|
||||
|
||||
def test_sameframe_fix_makes_host_forward_match_variant():
|
||||
"""Clearing sameframe shortcuts aligns host mj_forward with variant."""
|
||||
base_model = _viewer_regression_sphere_spec().compile()
|
||||
cone_model = _viewer_regression_cone_spec().compile()
|
||||
|
||||
# Sync cone's kinematic fields onto sphere's model (like viewer does).
|
||||
for field in (
|
||||
"geom_size",
|
||||
"geom_pos",
|
||||
"geom_quat",
|
||||
"body_mass",
|
||||
"body_inertia",
|
||||
"body_ipos",
|
||||
"body_iquat",
|
||||
):
|
||||
getattr(base_model, field)[:] = getattr(cone_model, field)
|
||||
|
||||
base_data = mujoco.MjData(base_model)
|
||||
base_data.qpos[:] = cone_model.qpos0
|
||||
base_data.qpos[2] = 0.05
|
||||
mujoco.mj_forward(base_model, base_data)
|
||||
|
||||
cone_data = mujoco.MjData(cone_model)
|
||||
cone_data.qpos[:] = cone_model.qpos0
|
||||
cone_data.qpos[2] = 0.05
|
||||
mujoco.mj_forward(cone_model, cone_data)
|
||||
|
||||
# Before fix: positions differ due to stale sameframe flags.
|
||||
assert not np.allclose(base_data.geom_xpos, cone_data.geom_xpos)
|
||||
|
||||
# After fix: clearing sameframe makes them match.
|
||||
disable_model_sameframe_shortcuts(base_model)
|
||||
mujoco.mj_forward(base_model, base_data)
|
||||
np.testing.assert_allclose(base_data.geom_xpos, cone_data.geom_xpos, atol=1e-6)
|
||||
|
||||
|
||||
def test_sync_model_fields_copies_only_requested_env_fields():
|
||||
"""Viewer model sync copies explicit fields and leaves others unchanged."""
|
||||
model = _simple_sphere_spec().compile()
|
||||
|
||||
class _SimModel:
|
||||
geom_rgba = torch.tensor(
|
||||
[
|
||||
[[0.1, 0.2, 0.3, 0.4]],
|
||||
[[0.5, 0.6, 0.7, 0.8]],
|
||||
],
|
||||
dtype=torch.float32,
|
||||
)
|
||||
geom_pos = torch.tensor(
|
||||
[
|
||||
[[1.0, 2.0, 3.0]],
|
||||
[[4.0, 5.0, 6.0]],
|
||||
],
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
original_geom_pos = model.geom_pos.copy()
|
||||
|
||||
sync_model_fields(model, _SimModel(), {"geom_rgba"}, env_idx=1)
|
||||
|
||||
np.testing.assert_allclose(model.geom_rgba, [[0.5, 0.6, 0.7, 0.8]])
|
||||
np.testing.assert_allclose(model.geom_pos, original_geom_pos)
|
||||
@@ -0,0 +1,213 @@
|
||||
"""Tests for sensor-based projected gravity (framezaxis up-vector sensor).
|
||||
|
||||
The shipped robots expose a ``framezaxis`` sensor that outputs the world Z-axis in the
|
||||
IMU site frame; negating it gives projected gravity. These tests check the sensor (and
|
||||
the ``projected_gravity_from_sensor`` observation that wraps it) against an independent
|
||||
ground-truth computation, and verify that -- unlike the entity-data
|
||||
``projected_gravity_b`` -- it tracks the IMU site orientation, which is what makes IMU
|
||||
mounting domain randomization observable.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from typing import TYPE_CHECKING, cast
|
||||
|
||||
import mujoco
|
||||
import pytest
|
||||
import torch
|
||||
from conftest import get_test_device
|
||||
|
||||
from mjlab.entity import EntityCfg
|
||||
from mjlab.envs.mdp import dr
|
||||
from mjlab.envs.mdp.observations import projected_gravity_from_sensor
|
||||
from mjlab.managers.scene_entity_config import SceneEntityCfg
|
||||
from mjlab.scene import Scene, SceneCfg
|
||||
from mjlab.sim.sim import Simulation, SimulationCfg
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from mjlab.envs import ManagerBasedRlEnv
|
||||
|
||||
# Gravity points along world -Z; projected gravity is this expressed in a body frame.
|
||||
_GRAVITY_DIR_W = (0.0, 0.0, -1.0)
|
||||
|
||||
|
||||
def _quat_to_mat(q: tuple[float, float, float, float]) -> torch.Tensor:
|
||||
"""Rotation matrix from a (w, x, y, z) quaternion. Independent of MuJoCo/mjlab."""
|
||||
w, x, y, z = q
|
||||
return torch.tensor(
|
||||
[
|
||||
[1 - 2 * (y * y + z * z), 2 * (x * y - w * z), 2 * (x * z + w * y)],
|
||||
[2 * (x * y + w * z), 1 - 2 * (x * x + z * z), 2 * (y * z - w * x)],
|
||||
[2 * (x * z - w * y), 2 * (y * z + w * x), 1 - 2 * (x * x + y * y)],
|
||||
],
|
||||
dtype=torch.float64,
|
||||
)
|
||||
|
||||
|
||||
def _expected_projected_gravity(q: tuple[float, float, float, float]) -> torch.Tensor:
|
||||
"""Ground-truth projected gravity for a body with world orientation ``q``.
|
||||
|
||||
proj = R(q)^T @ g_world, computed from an explicit rotation matrix so it does not
|
||||
share a code path with the sensor or with ``projected_gravity_b``.
|
||||
"""
|
||||
g_w = torch.tensor(_GRAVITY_DIR_W, dtype=torch.float64)
|
||||
return _quat_to_mat(q).T @ g_w
|
||||
|
||||
|
||||
class Env:
|
||||
"""Minimal env stub for driving observation and dr functions in tests."""
|
||||
|
||||
def __init__(self, scene, sim, device):
|
||||
self.scene = scene
|
||||
self.sim = sim
|
||||
self.num_envs = scene.num_envs
|
||||
self.device = device
|
||||
|
||||
|
||||
def _make_env(scene, sim, device) -> ManagerBasedRlEnv:
|
||||
"""Build the env stub, typed as the real env for the functions under test."""
|
||||
return cast("ManagerBasedRlEnv", Env(scene, sim, device))
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def device():
|
||||
return get_test_device()
|
||||
|
||||
|
||||
def _robot_xml(site_euler: str = "0 0 0") -> str:
|
||||
"""Free-floating box with an IMU site and the framezaxis up-vector sensor."""
|
||||
return f"""
|
||||
<mujoco>
|
||||
<worldbody>
|
||||
<body name="base" pos="0 0 1">
|
||||
<freejoint name="free_joint"/>
|
||||
<geom name="base_geom" type="box" size="0.2 0.2 0.1" mass="5.0"/>
|
||||
<site name="imu" pos="0.05 0 0" euler="{site_euler}"/>
|
||||
</body>
|
||||
</worldbody>
|
||||
<sensor>
|
||||
<framezaxis name="imu_upvector" objtype="body" objname="world"
|
||||
reftype="site" refname="imu"/>
|
||||
</sensor>
|
||||
</mujoco>
|
||||
"""
|
||||
|
||||
|
||||
def _build(xml: str, device: str, num_envs: int = 2):
|
||||
entity_cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(xml))
|
||||
scene = Scene(
|
||||
SceneCfg(num_envs=num_envs, env_spacing=3.0, entities={"robot": entity_cfg}),
|
||||
device,
|
||||
)
|
||||
model = scene.compile()
|
||||
sim = Simulation(
|
||||
num_envs=num_envs, cfg=SimulationCfg(njmax=20), model=model, device=device
|
||||
)
|
||||
scene.initialize(sim.mj_model, sim.model, sim.data)
|
||||
return scene, sim
|
||||
|
||||
|
||||
def _set_root_quat(robot, q: tuple[float, float, float, float], device: str) -> None:
|
||||
root_state = robot.data.default_root_state.clone()
|
||||
root_state[:, 3:7] = torch.tensor(q, device=device, dtype=root_state.dtype)
|
||||
robot.write_root_state_to_sim(root_state)
|
||||
|
||||
|
||||
def test_sensor_matches_ground_truth_when_site_aligned(device):
|
||||
"""Sensor and entity both equal hand-computed projected gravity for a tilted base."""
|
||||
scene, sim = _build(_robot_xml(), device)
|
||||
robot = scene["robot"]
|
||||
|
||||
# Compose a 0.6 rad roll with a 0.3 rad pitch into a single root quaternion.
|
||||
ax = (math.cos(0.3), math.sin(0.3), 0.0, 0.0)
|
||||
ay = (math.cos(0.15), 0.0, math.sin(0.15), 0.0)
|
||||
q = (
|
||||
ax[0] * ay[0] - ax[1] * ay[1] - ax[2] * ay[2] - ax[3] * ay[3],
|
||||
ax[0] * ay[1] + ax[1] * ay[0] + ax[2] * ay[3] - ax[3] * ay[2],
|
||||
ax[0] * ay[2] - ax[1] * ay[3] + ax[2] * ay[0] + ax[3] * ay[1],
|
||||
ax[0] * ay[3] + ax[1] * ay[2] - ax[2] * ay[1] + ax[3] * ay[0],
|
||||
)
|
||||
_set_root_quat(robot, q, device)
|
||||
sim.forward()
|
||||
|
||||
expected = _expected_projected_gravity(q).to(device=device, dtype=torch.float32)
|
||||
# Guard against a vacuous pass: the tilt must actually move gravity off straight-down.
|
||||
straight_down = torch.tensor(_GRAVITY_DIR_W, device=device)
|
||||
assert (expected - straight_down).abs().max() > 0.3
|
||||
|
||||
sensor_grav = -scene["robot/imu_upvector"].data
|
||||
entity_grav = robot.data.projected_gravity_b
|
||||
torch.testing.assert_close(sensor_grav[0], expected, atol=1e-5, rtol=0)
|
||||
torch.testing.assert_close(entity_grav[0], expected, atol=1e-5, rtol=0)
|
||||
|
||||
|
||||
def test_observation_fn_tracks_site_orientation(device):
|
||||
"""The observation fn reflects IMU site tilt; the entity-data version does not.
|
||||
|
||||
With the base upright but the IMU site rolled 30 deg about x, projected gravity in the
|
||||
site frame is (0, -sin30, -cos30). The entity-data version stays straight-down because
|
||||
it uses the root body orientation and is blind to the site.
|
||||
"""
|
||||
scene_rot, sim_rot = _build(_robot_xml(site_euler="30 0 0"), device)
|
||||
scene_flat, sim_flat = _build(_robot_xml(site_euler="0 0 0"), device)
|
||||
sim_rot.forward()
|
||||
sim_flat.forward()
|
||||
|
||||
# Drive through the actual shipped observation function, not the raw sensor.
|
||||
env_rot = _make_env(scene_rot, sim_rot, device)
|
||||
env_flat = _make_env(scene_flat, sim_flat, device)
|
||||
grav_rot = projected_gravity_from_sensor(env_rot, "robot/imu_upvector")
|
||||
grav_flat = projected_gravity_from_sensor(env_flat, "robot/imu_upvector")
|
||||
|
||||
expected_rot = torch.tensor(
|
||||
[0.0, -math.sin(math.radians(30)), -math.cos(math.radians(30))], device=device
|
||||
)
|
||||
straight_down = torch.tensor(_GRAVITY_DIR_W, device=device)
|
||||
torch.testing.assert_close(grav_rot[0], expected_rot, atol=1e-5, rtol=0)
|
||||
torch.testing.assert_close(grav_flat[0], straight_down, atol=1e-5, rtol=0)
|
||||
|
||||
# The entity-data version is unchanged by the site rotation (so it cannot be used to
|
||||
# observe IMU mounting randomization), confirming why the sensor path is needed.
|
||||
entity_rot = scene_rot["robot"].data.projected_gravity_b
|
||||
torch.testing.assert_close(entity_rot[0], straight_down, atol=1e-5, rtol=0)
|
||||
|
||||
|
||||
@pytest.mark.filterwarnings(
|
||||
"ignore:Use of index_put_ on expanded tensors is deprecated:UserWarning"
|
||||
)
|
||||
def test_site_quat_randomization_changes_sensor(device):
|
||||
"""The full DR path: running ``dr.site_quat`` perturbs the gravity observation.
|
||||
|
||||
This is what the G1 example configs rely on -- randomizing the IMU site orientation
|
||||
must show up in the sensor-based projected gravity, per-environment.
|
||||
"""
|
||||
scene, sim = _build(_robot_xml(), device, num_envs=4)
|
||||
sim.expand_model_fields(("site_quat",))
|
||||
env = _make_env(scene, sim, device)
|
||||
|
||||
sim.forward()
|
||||
straight_down = torch.tensor(_GRAVITY_DIR_W, device=device)
|
||||
before = projected_gravity_from_sensor(env, "robot/imu_upvector").clone()
|
||||
# Upright base + identity site quat => straight-down gravity in every env.
|
||||
torch.testing.assert_close(before, straight_down.expand_as(before), atol=1e-5, rtol=0)
|
||||
|
||||
torch.manual_seed(0)
|
||||
dr.site_quat(
|
||||
env,
|
||||
env_ids=None,
|
||||
roll_range=(-0.3, 0.3),
|
||||
pitch_range=(-0.3, 0.3),
|
||||
yaw_range=(-0.3, 0.3),
|
||||
asset_cfg=SceneEntityCfg("robot", site_names=("imu",)),
|
||||
)
|
||||
sim.forward()
|
||||
after = projected_gravity_from_sensor(env, "robot/imu_upvector")
|
||||
|
||||
# Randomization moved the reading off straight-down and made it env-dependent.
|
||||
assert (after - straight_down).abs().max() > 0.05
|
||||
assert not torch.allclose(after, before, atol=1e-3)
|
||||
assert torch.unique(after, dim=0).shape[0] >= 2
|
||||
# The perturbation is a rotation, so gravity stays a unit vector.
|
||||
norms = torch.linalg.norm(after, dim=-1)
|
||||
torch.testing.assert_close(norms, torch.ones_like(norms), atol=1e-5, rtol=0)
|
||||
@@ -1,30 +0,0 @@
|
||||
"""Tests for mjlab.utils.random."""
|
||||
|
||||
import subprocess
|
||||
import sys
|
||||
import textwrap
|
||||
|
||||
|
||||
def test_seed_rng_cpu_device_does_not_initialize_warp_cuda() -> None:
|
||||
"""seed_rng(device="cpu") must not initialize Warp's CUDA runtime.
|
||||
|
||||
Runs in a subprocess so that Warp is guaranteed uninitialized before the
|
||||
call.
|
||||
"""
|
||||
script = textwrap.dedent("""
|
||||
import warp as wp
|
||||
from mjlab.utils.random import seed_rng
|
||||
|
||||
assert wp._src.context.runtime is None, "Warp must not be initialized yet"
|
||||
seed_rng(42, device="cpu")
|
||||
rt = wp._src.context.runtime
|
||||
if rt is not None:
|
||||
cuda = [d for d in wp.get_devices() if "cuda" in str(d)]
|
||||
assert not cuda, f"seed_rng(device='cpu') initialized CUDA devices {cuda}"
|
||||
""")
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", script], capture_output=True, text=True
|
||||
)
|
||||
assert result.returncode == 0, (
|
||||
f"subprocess failed:\nstdout={result.stdout}\nstderr={result.stderr}"
|
||||
)
|
||||
@@ -936,7 +936,7 @@ def test_multi_frame_body_exclusion(device):
|
||||
should skip body_b's own geom but HIT body_a's platform. Frame A's
|
||||
rays should skip body_a and hit the floor.
|
||||
"""
|
||||
xml = """
|
||||
body_a_xml = """
|
||||
<mujoco>
|
||||
<worldbody>
|
||||
<geom name="floor" type="plane" size="10 10 0.1" pos="0 0 0"/>
|
||||
@@ -945,6 +945,12 @@ def test_multi_frame_body_exclusion(device):
|
||||
<geom name="geom_a" type="box" size="2 2 0.1" mass="5.0"/>
|
||||
<site name="site_a" pos="0 0 0"/>
|
||||
</body>
|
||||
</worldbody>
|
||||
</mujoco>
|
||||
"""
|
||||
body_b_xml = """
|
||||
<mujoco>
|
||||
<worldbody>
|
||||
<body name="body_b" pos="0 0 3">
|
||||
<freejoint name="free_b"/>
|
||||
<geom name="geom_b" type="box" size="0.5 0.5 0.5" mass="5.0"/>
|
||||
@@ -957,15 +963,17 @@ def test_multi_frame_body_exclusion(device):
|
||||
cfg = RayCastSensorCfg(
|
||||
name="multi",
|
||||
frame=(
|
||||
ObjRef(type="site", name="site_a", entity="robot"),
|
||||
ObjRef(type="site", name="site_b", entity="robot"),
|
||||
ObjRef(type="site", name="site_a", entity="body_a"),
|
||||
ObjRef(type="site", name="site_b", entity="body_b"),
|
||||
),
|
||||
pattern=GridPatternCfg(size=(0.0, 0.0), resolution=0.1),
|
||||
max_distance=10.0,
|
||||
exclude_parent_body=True,
|
||||
)
|
||||
|
||||
scene, sim = make_scene_and_sim(device, xml, (cfg,))
|
||||
scene, sim = make_scene_and_sim(
|
||||
device, {"body_a": body_a_xml, "body_b": body_b_xml}, (cfg,)
|
||||
)
|
||||
sim.step()
|
||||
sim.sense()
|
||||
|
||||
|
||||
@@ -4,6 +4,7 @@ import ast
|
||||
import tempfile
|
||||
from dataclasses import asdict
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import mujoco
|
||||
import onnx
|
||||
@@ -521,3 +522,85 @@ def test_onnx_motion_model_clamps_out_of_bounds_time_step():
|
||||
_, joint_pos, *_ = model(x, time_step)
|
||||
|
||||
torch.testing.assert_close(joint_pos, motion.joint_pos[num_steps - 1 : num_steps])
|
||||
|
||||
|
||||
def _make_tracking_runner_shell(registry_name, logger_type, upload_model=True):
|
||||
"""Build a MotionTrackingOnPolicyRunner with all heavy parts mocked out."""
|
||||
from mjlab.tasks.tracking.rl.runner import MotionTrackingOnPolicyRunner
|
||||
|
||||
runner = MotionTrackingOnPolicyRunner.__new__(MotionTrackingOnPolicyRunner)
|
||||
runner.registry_name = registry_name
|
||||
runner.cfg = {"upload_model": upload_model}
|
||||
runner.logger = MagicMock()
|
||||
runner.logger.logger_type = logger_type
|
||||
|
||||
mock_motion_term = MagicMock()
|
||||
mock_motion_term.cfg.anchor_body_name = "pelvis"
|
||||
mock_motion_term.cfg.body_names = ["body1"]
|
||||
runner.env = MagicMock()
|
||||
runner.env.unwrapped.command_manager.get_term.return_value = mock_motion_term
|
||||
return runner
|
||||
|
||||
|
||||
@pytest.mark.parametrize("logger_type", ["wandb", "WandbLogWriter"])
|
||||
def test_tracking_runner_registers_artifact_for_wandb_logger_types(
|
||||
logger_type, monkeypatch, tmp_path
|
||||
):
|
||||
"""use_artifact is called for both legacy 'wandb' and current 'WandbLogWriter' logger types.
|
||||
|
||||
Regression test: rsl-rl-lib 5.4 renamed the WandB logger type from 'wandb'
|
||||
to 'WandbLogWriter'. If only 'wandb' is checked, use_artifact is silently
|
||||
skipped and the nightly report fails with 'No motion artifact found in the run.'
|
||||
"""
|
||||
from mjlab.rl.runner import MjlabOnPolicyRunner
|
||||
from mjlab.tasks.tracking.rl import runner as runner_mod
|
||||
|
||||
runner = _make_tracking_runner_shell("org/motions/motion:latest", logger_type)
|
||||
|
||||
monkeypatch.setattr(MjlabOnPolicyRunner, "save", lambda *a, **kw: None)
|
||||
monkeypatch.setattr(runner_mod, "get_base_metadata", lambda *a: {})
|
||||
monkeypatch.setattr(runner_mod, "attach_metadata_to_onnx", lambda *a: None)
|
||||
monkeypatch.setattr(
|
||||
runner.env.unwrapped.__class__,
|
||||
"export_policy_to_onnx",
|
||||
lambda *a, **kw: None,
|
||||
raising=False,
|
||||
)
|
||||
|
||||
checkpoint = tmp_path / "run-dir" / "model_100.pt"
|
||||
checkpoint.parent.mkdir()
|
||||
checkpoint.touch()
|
||||
|
||||
mock_run = MagicMock()
|
||||
mock_run.name = "test-run"
|
||||
|
||||
with patch.object(runner_mod, "wandb") as mock_wandb:
|
||||
mock_wandb.run = mock_run
|
||||
runner.export_policy_to_onnx = MagicMock()
|
||||
runner.save(str(checkpoint))
|
||||
|
||||
mock_run.use_artifact.assert_called_once_with("org/motions/motion:latest")
|
||||
|
||||
|
||||
def test_tracking_runner_does_not_register_artifact_for_tensorboard(
|
||||
monkeypatch, tmp_path
|
||||
):
|
||||
"""use_artifact is NOT called when using the tensorboard logger."""
|
||||
from mjlab.rl.runner import MjlabOnPolicyRunner
|
||||
from mjlab.tasks.tracking.rl import runner as runner_mod
|
||||
|
||||
runner = _make_tracking_runner_shell("org/motions/motion:latest", "tensorboard")
|
||||
|
||||
monkeypatch.setattr(MjlabOnPolicyRunner, "save", lambda *a, **kw: None)
|
||||
monkeypatch.setattr(runner_mod, "get_base_metadata", lambda *a: {})
|
||||
monkeypatch.setattr(runner_mod, "attach_metadata_to_onnx", lambda *a: None)
|
||||
|
||||
checkpoint = tmp_path / "run-dir" / "model_100.pt"
|
||||
checkpoint.parent.mkdir()
|
||||
checkpoint.touch()
|
||||
|
||||
with patch.object(runner_mod, "wandb") as mock_wandb:
|
||||
runner.export_policy_to_onnx = MagicMock()
|
||||
runner.save(str(checkpoint))
|
||||
|
||||
mock_wandb.run.use_artifact.assert_not_called()
|
||||
|
||||
@@ -2,8 +2,15 @@
|
||||
|
||||
import mujoco
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from mjlab.terrains.primitive_terrains import BoxSteppingStonesTerrainCfg
|
||||
from mjlab.terrains.config import ALL_TERRAIN_PRESETS
|
||||
from mjlab.terrains.primitive_terrains import (
|
||||
_MIN_BORDER_HEIGHT,
|
||||
BoxInvertedPyramidStairsTerrainCfg,
|
||||
BoxPyramidStairsTerrainCfg,
|
||||
BoxSteppingStonesTerrainCfg,
|
||||
)
|
||||
|
||||
_CFG = BoxSteppingStonesTerrainCfg(
|
||||
proportion=1.0,
|
||||
@@ -37,12 +44,10 @@ def _generate_stones(
|
||||
if geom is None:
|
||||
continue
|
||||
pos, size = geom.pos, geom.size
|
||||
# Skip platform, floor, and border geoms.
|
||||
is_platform = (
|
||||
np.isclose(pos[0], center)
|
||||
and np.isclose(pos[1], center)
|
||||
and np.isclose(size[0], cfg.platform_width / 2, atol=1e-4)
|
||||
)
|
||||
# Skip platform, floor, and border geoms. The platform is the geom centered
|
||||
# exactly at the patch center (its size is grid-snapped, not the configured
|
||||
# width, so it is identified by position alone).
|
||||
is_platform = np.isclose(pos[0], center) and np.isclose(pos[1], center)
|
||||
is_full_span = np.isclose(size[0], cfg.size[0] / 2) or np.isclose(
|
||||
size[1], cfg.size[1] / 2
|
||||
)
|
||||
@@ -74,3 +79,50 @@ def test_stone_size_decreases_with_difficulty():
|
||||
sizes[difficulty] = np.mean([hx + hy for _, _, hx, hy in stones])
|
||||
|
||||
assert sizes[0.0] > sizes[1.0]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"cfg_cls", [BoxPyramidStairsTerrainCfg, BoxInvertedPyramidStairsTerrainCfg]
|
||||
)
|
||||
def test_pyramid_stairs_border_present_at_zero_difficulty(cfg_cls):
|
||||
"""At difficulty 0 the step height collapses to 0, but the flat border frame
|
||||
must still be generated as solid, non-degenerate geometry (regression for the
|
||||
empty-boundary bug, issue #1033)."""
|
||||
cfg = cfg_cls(
|
||||
size=(8.0, 8.0),
|
||||
step_height_range=(0.0, 0.2),
|
||||
step_width=0.3,
|
||||
platform_width=3.0,
|
||||
border_width=1.0,
|
||||
)
|
||||
spec = mujoco.MjSpec()
|
||||
spec.worldbody.add_body(name="terrain")
|
||||
output = cfg.function(difficulty=0.0, spec=spec, rng=np.random.default_rng(0))
|
||||
|
||||
# The border frame sits below z=0 (top flush at ground level); inner step
|
||||
# boxes are centered at z=0. Identify the frame by its downward offset.
|
||||
border_geoms = [
|
||||
g.geom for g in output.geometries if g.geom is not None and g.geom.pos[2] < -1e-4
|
||||
]
|
||||
assert len(border_geoms) == 4, "Expected four border frame boxes."
|
||||
for geom in border_geoms:
|
||||
# Each frame box must be solid, not a degenerate zero-height geom, and its
|
||||
# top must be flush with the ground plane at z=0.
|
||||
assert geom.size[2] >= _MIN_BORDER_HEIGHT / 2 - 1e-9
|
||||
assert np.isclose(geom.pos[2] + geom.size[2], 0.0, atol=1e-6)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("preset_name", sorted(ALL_TERRAIN_PRESETS))
|
||||
@pytest.mark.parametrize("difficulty", [0.0, 1.0])
|
||||
def test_preset_compiles_across_difficulty(preset_name, difficulty):
|
||||
"""Every terrain preset must generate compilable MuJoCo geometry across the
|
||||
full difficulty range. Difficulty 0 is exercised explicitly because curriculum
|
||||
row 0 lands there deterministically, which previously produced degenerate
|
||||
geometry (zero-height hfields, NaN colors, missing borders)."""
|
||||
cfg = ALL_TERRAIN_PRESETS[preset_name](size=(8.0, 8.0))
|
||||
spec = mujoco.MjSpec()
|
||||
spec.worldbody.add_body(name="terrain")
|
||||
cfg.function(difficulty=difficulty, spec=spec, rng=np.random.default_rng(0))
|
||||
# Compiling validates geom/hfield sizes and rgba values (catches NaNs and
|
||||
# non-positive sizes that MuJoCo rejects).
|
||||
spec.compile()
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
"""Tests for motion tracking evaluation metrics."""
|
||||
|
||||
import math
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
@@ -34,11 +35,11 @@ def mock_command():
|
||||
|
||||
|
||||
def test_mpkpe_zero_when_positions_match(mock_command):
|
||||
"""Test MPKPE is zero when positions are identical."""
|
||||
"""Test MPKPE is zero when global positions are identical."""
|
||||
num_bodies = len(mock_command.cfg.body_names)
|
||||
positions = torch.rand(mock_command.num_envs, num_bodies, 3)
|
||||
|
||||
mock_command.body_pos_relative_w = positions.clone()
|
||||
mock_command.body_pos_w = positions.clone()
|
||||
mock_command.robot_body_pos_w = positions.clone()
|
||||
|
||||
mpkpe = compute_mpkpe(mock_command)
|
||||
@@ -48,10 +49,10 @@ def test_mpkpe_zero_when_positions_match(mock_command):
|
||||
|
||||
|
||||
def test_mpkpe_correct_error(mock_command):
|
||||
"""Test MPKPE computes correct mean error."""
|
||||
"""Test MPKPE computes the correct mean global error."""
|
||||
num_bodies = len(mock_command.cfg.body_names)
|
||||
|
||||
mock_command.body_pos_relative_w = torch.zeros(mock_command.num_envs, num_bodies, 3)
|
||||
mock_command.body_pos_w = torch.zeros(mock_command.num_envs, num_bodies, 3)
|
||||
mock_command.robot_body_pos_w = torch.zeros(mock_command.num_envs, num_bodies, 3)
|
||||
mock_command.robot_body_pos_w[:, :, 0] = 1.0 # 1 unit offset in x
|
||||
|
||||
@@ -60,58 +61,71 @@ def test_mpkpe_correct_error(mock_command):
|
||||
assert torch.allclose(mpkpe, torch.ones(mock_command.num_envs), atol=1e-6)
|
||||
|
||||
|
||||
def test_r_mpkpe_invariant_to_global_translation(mock_command):
|
||||
"""Test R-MPKPE is invariant to global translation."""
|
||||
def test_mpkpe_uses_global_reference(mock_command):
|
||||
"""MPKPE must read the global reference, not the drift-cancelled one.
|
||||
|
||||
Pins issue #1006: setting body_pos_relative_w to match the robot exactly
|
||||
would yield zero error if it were (incorrectly) used; the metric must
|
||||
instead follow body_pos_w.
|
||||
"""
|
||||
num_bodies = len(mock_command.cfg.body_names)
|
||||
robot_pos = torch.rand(mock_command.num_envs, num_bodies, 3)
|
||||
mock_command.robot_body_pos_w = robot_pos.clone()
|
||||
mock_command.body_pos_relative_w = robot_pos.clone() # zero error if misused
|
||||
mock_command.body_pos_w = robot_pos.clone()
|
||||
mock_command.body_pos_w[:, :, 0] += 1.0 # 1 unit of global drift
|
||||
|
||||
mock_command.anchor_pos_w = torch.zeros(mock_command.num_envs, 3)
|
||||
mock_command.body_pos_w = torch.rand(mock_command.num_envs, num_bodies, 3)
|
||||
mock_command.robot_anchor_pos_w = torch.zeros(mock_command.num_envs, 3)
|
||||
mock_command.robot_body_pos_w = mock_command.body_pos_w.clone()
|
||||
mpkpe = compute_mpkpe(mock_command)
|
||||
|
||||
r_mpkpe_1 = compute_root_relative_mpkpe(mock_command)
|
||||
|
||||
# Translate everything by large offset.
|
||||
offset = torch.tensor([100.0, 200.0, 300.0])
|
||||
mock_command.anchor_pos_w = offset.expand(mock_command.num_envs, 3).clone()
|
||||
mock_command.body_pos_w = mock_command.body_pos_w + offset
|
||||
mock_command.robot_anchor_pos_w = offset.expand(mock_command.num_envs, 3).clone()
|
||||
mock_command.robot_body_pos_w = mock_command.robot_body_pos_w + offset
|
||||
|
||||
r_mpkpe_2 = compute_root_relative_mpkpe(mock_command)
|
||||
|
||||
assert torch.allclose(r_mpkpe_1, r_mpkpe_2, atol=1e-5)
|
||||
assert torch.allclose(mpkpe, torch.ones(mock_command.num_envs), atol=1e-6)
|
||||
|
||||
|
||||
def test_r_mpkpe_detects_relative_error(mock_command):
|
||||
"""Test R-MPKPE detects errors in relative positions."""
|
||||
def test_r_mpkpe_zero_when_relative_positions_match(mock_command):
|
||||
"""R-MPKPE is zero when re-anchored positions are identical."""
|
||||
num_bodies = len(mock_command.cfg.body_names)
|
||||
positions = torch.rand(mock_command.num_envs, num_bodies, 3)
|
||||
|
||||
mock_command.anchor_pos_w = torch.zeros(mock_command.num_envs, 3)
|
||||
mock_command.body_pos_w = torch.zeros(mock_command.num_envs, num_bodies, 3)
|
||||
mock_command.body_pos_w[:, :, 0] = 1.0 # Bodies 1 unit from anchor
|
||||
mock_command.body_pos_relative_w = positions.clone()
|
||||
mock_command.robot_body_pos_w = positions.clone()
|
||||
|
||||
mock_command.robot_anchor_pos_w = torch.zeros(mock_command.num_envs, 3)
|
||||
mock_command.robot_body_pos_w = torch.zeros(mock_command.num_envs, num_bodies, 3)
|
||||
mock_command.robot_body_pos_w[:, :, 0] = 2.0 # Bodies 2 units from anchor
|
||||
r_mpkpe = compute_root_relative_mpkpe(mock_command)
|
||||
|
||||
assert r_mpkpe.shape == (mock_command.num_envs,)
|
||||
assert torch.allclose(r_mpkpe, torch.zeros(mock_command.num_envs), atol=1e-6)
|
||||
|
||||
|
||||
def test_r_mpkpe_uses_relative_reference(mock_command):
|
||||
"""R-MPKPE reads the re-anchored reference, not the global one.
|
||||
|
||||
Setting body_pos_w to match the robot exactly would yield zero error if
|
||||
it were (incorrectly) used; the metric must instead follow
|
||||
body_pos_relative_w.
|
||||
"""
|
||||
num_bodies = len(mock_command.cfg.body_names)
|
||||
robot_pos = torch.rand(mock_command.num_envs, num_bodies, 3)
|
||||
mock_command.robot_body_pos_w = robot_pos.clone()
|
||||
mock_command.body_pos_w = robot_pos.clone() # zero error if misused
|
||||
mock_command.body_pos_relative_w = robot_pos.clone()
|
||||
mock_command.body_pos_relative_w[:, :, 0] += 1.0 # 1 unit of local pose error
|
||||
|
||||
r_mpkpe = compute_root_relative_mpkpe(mock_command)
|
||||
|
||||
assert torch.allclose(r_mpkpe, torch.ones(mock_command.num_envs), atol=1e-6)
|
||||
|
||||
|
||||
def test_joint_velocity_error(mock_command):
|
||||
"""Test joint velocity error computes correct L2 norm."""
|
||||
def test_joint_velocity_error_rms(mock_command):
|
||||
"""Joint velocity error is the per-joint RMS of the velocity error."""
|
||||
num_joints = 3
|
||||
|
||||
mock_command.joint_vel = torch.zeros(mock_command.num_envs, num_joints)
|
||||
mock_command.robot_joint_vel = torch.zeros(mock_command.num_envs, num_joints)
|
||||
mock_command.robot_joint_vel[:, 0] = 3.0
|
||||
mock_command.robot_joint_vel[:, 1] = 4.0 # Error [3, 4, 0] has norm 5
|
||||
mock_command.robot_joint_vel[:, 1] = 4.0 # Error [3, 4, 0]
|
||||
|
||||
error = compute_joint_velocity_error(mock_command)
|
||||
|
||||
assert torch.allclose(error, torch.ones(mock_command.num_envs) * 5.0, atol=1e-6)
|
||||
expected = math.sqrt((3.0**2 + 4.0**2 + 0.0**2) / num_joints)
|
||||
assert torch.allclose(error, torch.ones(mock_command.num_envs) * expected, atol=1e-6)
|
||||
|
||||
|
||||
def test_ee_position_error_only_uses_specified_bodies(mock_command):
|
||||
@@ -153,3 +167,13 @@ def test_ee_orientation_error_detects_rotation(mock_command):
|
||||
# Error should be approximately pi/2 radians.
|
||||
expected = torch.ones(mock_command.num_envs) * (3.14159 / 2)
|
||||
assert torch.allclose(error, expected, atol=0.01)
|
||||
|
||||
|
||||
def test_ee_metrics_raise_on_unknown_body(mock_command):
|
||||
"""Unknown end-effector names raise instead of silently scoring zero."""
|
||||
num_bodies = len(mock_command.cfg.body_names)
|
||||
mock_command.body_pos_relative_w = torch.zeros(mock_command.num_envs, num_bodies, 3)
|
||||
mock_command.robot_body_pos_w = torch.zeros(mock_command.num_envs, num_bodies, 3)
|
||||
|
||||
with pytest.raises(ValueError, match="not tracked"):
|
||||
compute_ee_position_error(mock_command, ("nonexistent_body",))
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -4,7 +4,7 @@ import mujoco
|
||||
import pytest
|
||||
from conftest import get_test_device
|
||||
|
||||
from mjlab.actuator import XmlActuatorCfg
|
||||
from mjlab.actuator import XmlActuator, XmlActuatorCfg
|
||||
from mjlab.entity import Entity, EntityArticulationInfoCfg, EntityCfg
|
||||
from mjlab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg, mdp
|
||||
from mjlab.managers.observation_manager import ObservationGroupCfg, ObservationTermCfg
|
||||
@@ -160,6 +160,7 @@ def test_xml_actuator_explicit_command_field_bypasses_detection():
|
||||
entity.compile()
|
||||
|
||||
actuator = entity._actuators[0]
|
||||
assert isinstance(actuator, XmlActuator)
|
||||
assert actuator.command_field == "effort"
|
||||
assert actuator._target_names == ["joint1"]
|
||||
|
||||
|
||||
Generated
+34
-26
@@ -42,7 +42,7 @@ conflicts = [[
|
||||
|
||||
[manifest]
|
||||
constraints = [
|
||||
{ name = "gitpython", specifier = ">=3.1.47" },
|
||||
{ name = "gitpython", specifier = ">=3.1.49" },
|
||||
{ name = "lxml", specifier = ">=6.1.0" },
|
||||
]
|
||||
overrides = [{ name = "mujoco", specifier = ">=3.8.0.dev0", index = "https://py.mujoco.org/" }]
|
||||
@@ -823,14 +823,14 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "gitpython"
|
||||
version = "3.1.47"
|
||||
version = "3.1.50"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "gitdb" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/c1/bd/50db468e9b1310529a19fce651b3b0e753b5c07954d486cba31bbee9a5d5/gitpython-3.1.47.tar.gz", hash = "sha256:dba27f922bd2b42cb54c87a8ab3cb6beb6bf07f3d564e21ac848913a05a8a3cd", size = 216978, upload-time = "2026-04-22T02:44:44.059Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/33/f6/354ae6491228b5eb40e10d89c4d13c651fe1cf7556e35ebdded50cff57ce/gitpython-3.1.50.tar.gz", hash = "sha256:80da2d12504d52e1f998772dc5baf6e553f8d2fcfe1fcc226c9d9a2ee3372dcc", size = 219798, upload-time = "2026-05-06T04:01:26.571Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/f2/c5/a1bc0996af85757903cf2bf444a7824e68e0035ce63fb41d6f76f9def68b/gitpython-3.1.47-py3-none-any.whl", hash = "sha256:489f590edfd6d20571b2c0e72c6a6ac6915ee8b8cd04572330e3842207a78905", size = 209547, upload-time = "2026-04-22T02:44:41.271Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/20/7a/1c6e3562dfd8950adbb11ffbc65d21e7c89d01a6e4f137fa981056de25c5/gitpython-3.1.50-py3-none-any.whl", hash = "sha256:d352abe2908d07355014abdd21ddf798c2a961469239afec4962e9da884858f9", size = 212507, upload-time = "2026-05-06T04:01:23.799Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1647,7 +1647,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "mjlab"
|
||||
version = "1.3.0"
|
||||
version = "1.4.0"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "imageio-ffmpeg" },
|
||||
@@ -1658,6 +1658,8 @@ dependencies = [
|
||||
{ name = "onnxscript" },
|
||||
{ name = "prettytable" },
|
||||
{ name = "rsl-rl-lib" },
|
||||
{ name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or (extra == 'extra-5-mjlab-cpu' and extra == 'extra-5-mjlab-cu128')" },
|
||||
{ name = "scipy", version = "1.16.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' or (extra == 'extra-5-mjlab-cpu' and extra == 'extra-5-mjlab-cu128')" },
|
||||
{ name = "tensorboard" },
|
||||
{ name = "tensordict" },
|
||||
{ name = "torch", version = "2.9.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'darwin' or (extra == 'extra-5-mjlab-cpu' and extra == 'extra-5-mjlab-cu128') or (extra != 'extra-5-mjlab-cpu' and extra != 'extra-5-mjlab-cu128')" },
|
||||
@@ -1715,12 +1717,13 @@ docs = [
|
||||
requires-dist = [
|
||||
{ name = "imageio-ffmpeg" },
|
||||
{ name = "mediapy", specifier = ">=1.2.6" },
|
||||
{ name = "mjviser", git = "https://github.com/mujocolab/mjviser?rev=1bdfd6fe79066b847a5f430000fcfbb53ec31a6f" },
|
||||
{ name = "mujoco", specifier = ">=3.8.0", index = "https://py.mujoco.org/" },
|
||||
{ name = "mujoco-warp", git = "https://github.com/google-deepmind/mujoco_warp?rev=6f235d4" },
|
||||
{ name = "mjviser", specifier = ">=0.0.14" },
|
||||
{ name = "mujoco", specifier = "~=3.8.0", index = "https://py.mujoco.org/" },
|
||||
{ name = "mujoco-warp", git = "https://github.com/google-deepmind/mujoco_warp?rev=88b55fc2696960b927bc12584994bb8412b36558" },
|
||||
{ name = "onnxscript", specifier = ">=0.5.4" },
|
||||
{ name = "prettytable" },
|
||||
{ name = "rsl-rl-lib", specifier = "==5.2.0" },
|
||||
{ name = "rsl-rl-lib", specifier = "==5.4.0" },
|
||||
{ name = "scipy", specifier = ">=1.15" },
|
||||
{ name = "tensorboard", specifier = ">=2.20.0" },
|
||||
{ name = "tensordict" },
|
||||
{ name = "torch", specifier = ">=2.7.0" },
|
||||
@@ -1732,7 +1735,7 @@ requires-dist = [
|
||||
{ name = "tqdm" },
|
||||
{ name = "trimesh", specifier = ">=4.8.3" },
|
||||
{ name = "tyro", specifier = ">=1.0.1" },
|
||||
{ name = "viser", specifier = ">=1.0.26" },
|
||||
{ name = "viser", specifier = ">=1.0.27" },
|
||||
{ name = "wandb", specifier = ">=0.22.3" },
|
||||
{ name = "warp-lang", marker = "sys_platform != 'darwin'", specifier = ">=1.12.0", index = "https://pypi.nvidia.com/" },
|
||||
{ name = "warp-lang", marker = "sys_platform == 'darwin'", specifier = ">=1.12.0" },
|
||||
@@ -1767,8 +1770,8 @@ docs = [
|
||||
|
||||
[[package]]
|
||||
name = "mjviser"
|
||||
version = "0.0.13"
|
||||
source = { git = "https://github.com/mujocolab/mjviser?rev=1bdfd6fe79066b847a5f430000fcfbb53ec31a6f#1bdfd6fe79066b847a5f430000fcfbb53ec31a6f" }
|
||||
version = "0.0.14"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "mujoco" },
|
||||
{ name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or (extra == 'extra-5-mjlab-cpu' and extra == 'extra-5-mjlab-cu128')" },
|
||||
@@ -1777,6 +1780,10 @@ dependencies = [
|
||||
{ name = "trimesh" },
|
||||
{ name = "viser" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/1d/e4/eef89b279fb1811b5f120a99ac9284c32ff2ca4fad5e6f5c93035f72ba9a/mjviser-0.0.14.tar.gz", hash = "sha256:ebde2203dab89959a13ae549b4d3e5e5cf9eb69de11a1a2fd759cbe8f8c641f3", size = 29576, upload-time = "2026-05-07T03:35:13.128Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/9c/c2/4534d678ad1b3f7dee6fd83110112800287be21ba007d988abb9ddc8e0ac/mjviser-0.0.14-py3-none-any.whl", hash = "sha256:4b09f8e90506fc4a71d76fc628872147157947e9292428213ab78cfe137c9a26", size = 32296, upload-time = "2026-05-07T03:35:14.252Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ml-dtypes"
|
||||
@@ -1958,8 +1965,8 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "mujoco-warp"
|
||||
version = "3.8.0"
|
||||
source = { git = "https://github.com/google-deepmind/mujoco_warp?rev=6f235d4#6f235d46cb2ecf8f37c8f967f8dd9c87d0ca5807" }
|
||||
version = "3.8.0.2"
|
||||
source = { git = "https://github.com/google-deepmind/mujoco_warp?rev=88b55fc2696960b927bc12584994bb8412b36558#88b55fc2696960b927bc12584994bb8412b36558" }
|
||||
dependencies = [
|
||||
{ name = "absl-py" },
|
||||
{ name = "etils", extra = ["epath"] },
|
||||
@@ -2542,7 +2549,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "paramiko"
|
||||
version = "4.0.0"
|
||||
version = "5.0.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "bcrypt" },
|
||||
@@ -2550,9 +2557,9 @@ dependencies = [
|
||||
{ name = "invoke" },
|
||||
{ name = "pynacl" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/1f/e7/81fdcbc7f190cdb058cffc9431587eb289833bdd633e2002455ca9bb13d4/paramiko-4.0.0.tar.gz", hash = "sha256:6a25f07b380cc9c9a88d2b920ad37167ac4667f8d9886ccebd8f90f654b5d69f", size = 1630743, upload-time = "2025-08-04T01:02:03.711Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/62/93/dcc25d52f49022ae6175d15e6bd751f1acc99b98bc61fc55e5155a7be2e7/paramiko-5.0.0.tar.gz", hash = "sha256:36763b5b95c2a0dcfdf1abc48e48156ee425b21efe2f0e787c2dd5a95c0e5e79", size = 1548586, upload-time = "2026-05-09T18:28:52.256Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/a9/90/a744336f5af32c433bd09af7854599682a383b37cfd78f7de263de6ad6cb/paramiko-4.0.0-py3-none-any.whl", hash = "sha256:0e20e00ac666503bf0b4eda3b6d833465a2b7aff2e2b3d79a8bba5ef144ee3b9", size = 223932, upload-time = "2025-08-04T01:02:02.029Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/82/5b/eadf6d45de38d30ab603f49393b6cd2cbe7e233af8cf90197e32782b68a9/paramiko-5.0.0-py3-none-any.whl", hash = "sha256:b7044611c30140d9a75261653210e2002977b71a0497ff3ba0d98d7edbf62f7c", size = 208919, upload-time = "2026-05-09T18:28:50.295Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -3253,7 +3260,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "rsl-rl-lib"
|
||||
version = "5.2.0"
|
||||
version = "5.4.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "gitpython" },
|
||||
@@ -3261,6 +3268,7 @@ dependencies = [
|
||||
{ name = "numpy", version = "2.3.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' or (extra == 'extra-5-mjlab-cpu' and extra == 'extra-5-mjlab-cu128')" },
|
||||
{ name = "onnx" },
|
||||
{ name = "onnxscript" },
|
||||
{ name = "tensorboard" },
|
||||
{ name = "tensordict" },
|
||||
{ name = "torch", version = "2.9.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'darwin' or (extra == 'extra-5-mjlab-cpu' and extra == 'extra-5-mjlab-cu128') or (extra != 'extra-5-mjlab-cpu' and extra != 'extra-5-mjlab-cu128')" },
|
||||
{ name = "torch", version = "2.9.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "(sys_platform != 'darwin' and extra == 'extra-5-mjlab-cu128') or (extra == 'extra-5-mjlab-cpu' and extra == 'extra-5-mjlab-cu128')" },
|
||||
@@ -3268,9 +3276,9 @@ dependencies = [
|
||||
{ name = "torchvision", version = "0.24.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'darwin' or extra != 'extra-5-mjlab-cpu' or (extra == 'extra-5-mjlab-cpu' and extra == 'extra-5-mjlab-cu128')" },
|
||||
{ name = "torchvision", version = "0.25.0", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform != 'darwin' and extra == 'extra-5-mjlab-cpu') or (extra == 'extra-5-mjlab-cpu' and extra == 'extra-5-mjlab-cu128')" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/d7/a8/fb9aae0573a83dd510e228085f5e6ff9076a91f2bff6699270f07d31854c/rsl_rl_lib-5.2.0.tar.gz", hash = "sha256:cbb4eee96af9574495208381115d45d68ad1c0403710a2bc6512a2ee5bf57124", size = 60558, upload-time = "2026-04-23T12:40:54.259Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/0f/51/2d4c95b3642c0f659fbcddcd17fa0401903733250fa382f51f9b913bb85f/rsl_rl_lib-5.4.0.tar.gz", hash = "sha256:e1aa5cd5771f2d9a9e7a7ba5456b942ab588410cfd397b3c63e32d00a0744f0d", size = 65902, upload-time = "2026-05-27T10:42:39.656Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/7e/3a/3e8f39049cc5a8994bfdb2dd09d727f90a2781b9755adf59e49fa509500e/rsl_rl_lib-5.2.0-py3-none-any.whl", hash = "sha256:fc767059f329a184527dd10766c3382354f6deca4b049f23febc8140c3dc6029", size = 86451, upload-time = "2026-04-23T12:40:52.762Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/f8/dd/6797be77ce0cc12881f29cd9363b791f15276fde6a136a73443b6b6c2ce3/rsl_rl_lib-5.4.0-py3-none-any.whl", hash = "sha256:b30a0e59dac0ef7236f8f793c9bedfa8f2b8f8d29c2965140feeb1439153ebab", size = 92871, upload-time = "2026-05-27T10:42:38.415Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -4593,11 +4601,11 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "urllib3"
|
||||
version = "2.6.3"
|
||||
version = "2.7.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/c7/24/5f1b3bdffd70275f6661c76461e25f024d5a38a46f04aaca912426a2b1d3/urllib3-2.6.3.tar.gz", hash = "sha256:1b62b6884944a57dbe321509ab94fd4d3b307075e0c2eae991ac71ee15ad38ed", size = 435556, upload-time = "2026-01-07T16:24:43.925Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/53/0c/06f8b233b8fd13b9e5ee11424ef85419ba0d8ba0b3138bf360be2ff56953/urllib3-2.7.0.tar.gz", hash = "sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c", size = 433602, upload-time = "2026-05-07T16:13:18.596Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/39/08/aaaad47bc4e9dc8c725e68f9d04865dbcb2052843ff09c97b08904852d84/urllib3-2.6.3-py3-none-any.whl", hash = "sha256:bf272323e553dfb2e87d9bfd225ca7b0f467b919d7bbd355436d3fd37cb0acd4", size = 131584, upload-time = "2026-01-07T16:24:42.685Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/7f/3e/5db95bcf282c52709639744ca2a8b149baccf648e39c8cc87553df9eae0c/urllib3-2.7.0-py3-none-any.whl", hash = "sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897", size = 131087, upload-time = "2026-05-07T16:13:17.151Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -4672,7 +4680,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "viser"
|
||||
version = "1.0.26"
|
||||
version = "1.0.27"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "imageio" },
|
||||
@@ -4688,9 +4696,9 @@ dependencies = [
|
||||
{ name = "yourdfpy" },
|
||||
{ name = "zstandard" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/12/ce/82a0e50fae21f5e02fcc5d9aff2ab59dccb9c319b6c4cf528f2228049b05/viser-1.0.26.tar.gz", hash = "sha256:dc08c6f505e70324b0603bdddf9714c00ac828c259ee49abd8ad094bfc90c91c", size = 4828261, upload-time = "2026-03-30T11:43:19.513Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/fd/f5/48adb4e5e4234f48e96a1e7fc50cca6731280df0c279833e333963f9ea5c/viser-1.0.27.tar.gz", hash = "sha256:87e3239d6c1c2c003db93ac4072430ec790e336ffe7214781f035e54faebc0af", size = 4897986, upload-time = "2026-05-06T10:30:47.556Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/91/f7/762a2d5fab509d0c632b271e21e634462397cc02cca649771c3e9d2e0bcc/viser-1.0.26-py3-none-any.whl", hash = "sha256:03b177b4ef584f58f7b74fdf44cccb165b8a220ffd90728ef5c1e3d1b1fcf258", size = 4922888, upload-time = "2026-03-30T11:43:21.355Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/7c/ad/8ae712579e294b4395fb39f7d65524b51fc7b731eacce26af096b7e59b61/viser-1.0.27-py3-none-any.whl", hash = "sha256:8da5b7934416e6e2d3a7ebcf39fc840f21030b51eb63231e8cfef457bfb49031", size = 4998748, upload-time = "2026-05-06T10:30:49.965Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
||||
Reference in New Issue
Block a user