diff --git a/docs/index.rst b/docs/index.rst index 78e4932e35b3..377e1d7e1295 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -168,6 +168,7 @@ Table of Contents source/experimental-features/ray source/experimental-features/visuo_tactile_sensor source/experimental-features/rlinf_vla_posttraining + source/experimental-features/supermimic source/features/imitation-learning/index .. toctree:: diff --git a/docs/licenses/dependencies/supermimic-license.txt b/docs/licenses/dependencies/supermimic-license.txt new file mode 100644 index 000000000000..46e96afd1585 --- /dev/null +++ b/docs/licenses/dependencies/supermimic-license.txt @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 ETH Zurich, Chong Zhang + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/docs/source/_static/css/environment-browser.js b/docs/source/_static/css/environment-browser.js index edb297ec0bf1..8182504b0a85 100644 --- a/docs/source/_static/css/environment-browser.js +++ b/docs/source/_static/css/environment-browser.js @@ -91,6 +91,7 @@ ["IsaacContrib-Lift-Cube-Franka-IK-Abs", "", "", "", "", "environment-previews/franka_lift.jpg"], ["IsaacContrib-Lift-Cube-Franka-IK-Rel", "", "", "", "", "environment-previews/franka_lift.jpg"], ["IsaacContrib-Lift-Cube-OpenArm", "rl_games,rsl_rl", "", "", "", "environment-previews/openarm_uni_lift.jpg"], + ["IsaacContrib-Motion-Tracking-G1", "supermimic", "newton_mjwarp", "", "", "", false, {}, {}, "", null, ["||", "newton_mjwarp||"]], ["IsaacContrib-Multitask-Manipulation", "rsl_rl", "", "", "", "environment-previews/multitask_manipulation.jpg"], ["IsaacContrib-Navigation-3DObstacles-ARL-Robot-1", "rl_games,rsl_rl,skrl", "", "", "", "tasks/drone_arl/arl_robot_1_navigation.jpg"], ["IsaacContrib-Navigation-Flat-AnymalC", "rsl_rl,skrl", "", "", "", "environment-previews/anymal_c_nav.jpg"], @@ -285,6 +286,7 @@ ["skrl", "skrl", "skrl"], ["sb3", "sb3", "Stable-Baselines3"], ["rlinf", "rlinf", "RLinf"], + ["supermimic", "supermimic", "SuperMimic"], ]], ]; diff --git a/docs/source/api/lab/isaaclab.managers.rst b/docs/source/api/lab/isaaclab.managers.rst index d378386b86de..f83a4bd85b8e 100644 --- a/docs/source/api/lab/isaaclab.managers.rst +++ b/docs/source/api/lab/isaaclab.managers.rst @@ -231,6 +231,28 @@ Curriculum Manager Recorder Manager ---------------- +Recorder terms select which tensors to capture at the existing pre/post step and reset callbacks. +For an online consumer, a term can keep its latest batch in preallocated tensors and return +``(None, None)`` from its callbacks. The training root obtains that term with +``env.recorder_manager.get_term(name)`` and copies the batch into its own replay storage. +The term owns the lifetime of its published tensors; consumers must copy them before the next capture. + +In an RL environment, post-step callbacks run after reward and termination computation and before +automatic reset. A term reading physical tensors directly should set +``requires_post_step_observations = False``. This avoids an extra observation computation and its +noise draws. Existing terms default to ``True``. Callback dispatch is independent of this setting. + +Returning a field path and a batched tensor (or nested dictionary of tensors) retains samples in +the manager. Each field has a time-major allocation and independent per-world insertion indices, +so reset and per-step fields can have different sample rates. ``buffer_initial_capacity`` controls +the first allocation per field: choose it with the payload size and expected episode length in mind, +especially for images. Storage grows without dropping samples and is reused across resets. + +``get_episode(env_id)`` returns an owned snapshot whose leaves are tensors with a leading sample +dimension. Reading a snapshot can synchronize device indices and copy data. Changes to that snapshot +do not change live recording; use ``add_to_episodes`` and ``set_success_to_episodes`` for live writes. +File export materializes these snapshots; the hot capture path does not construct per-world episodes. + .. autoclass:: RecorderManager :members: :inherited-members: diff --git a/docs/source/experimental-features/supermimic.rst b/docs/source/experimental-features/supermimic.rst new file mode 100644 index 000000000000..e4bc7ba2b65b --- /dev/null +++ b/docs/source/experimental-features/supermimic.rst @@ -0,0 +1,205 @@ +.. _supermimic: + +SuperMimic motion tracking +========================== + +The experimental ``IsaacContrib-Motion-Tracking-G1`` task integrates the +`SuperMimic workflow `__ into Isaac Lab. +Use the standard reinforcement-learning entrypoints with ``--rl_library supermimic``. +The physical teacher runs Newton/MuJoCo Warp with CUDA graphs enabled. A learned world model +fits its recorded transitions, and the policy improves through differentiable imagined rollouts. + +This integration supports the MLP world model and the optional one-frame NeRD architecture. +It includes stationary padding of reference goals beyond recorded clip endpoints: padded poses +hold their position and orientation, with zero velocity. Recorded endpoint velocities remain unchanged. +The integration does not establish the paper's reported zero-failure result or an exact reproduction +of its motion selection. Evaluation reports failures over the actual prepared corpus. + +Ownership +--------- + +The task is an ordinary :class:`~isaaclab.envs.ManagerBasedRLEnv`. Its configuration lives in +``isaaclab_tasks.contrib.g1_tracking.g1_tracking_env_cfg`` and owns the scene, actuators, observations, +commands, rewards, terminations and recorder terms. The matching 29-DoF G1 asset variant and its +actuator settings are task-specific; substituting another G1 asset changes the dynamics. + +The remaining task modules have separate responsibilities: + +* ``mdp`` implements the native manager terms, including transition capture and adaptive reference sampling. +* ``tracking`` defines the state layout and the differentiable observation, reward and termination + expressions shared by real and imagined transitions. +* ``motion`` loads and resamples the reference corpus and constructs future goals. +* ``kinematics`` supplies the matching differentiable link geometry and joint order. + +``isaaclab_rl.supermimic`` owns the learner configuration, networks, normalizers, replay ring and +optimization. The native ``train_supermimic`` and ``play_supermimic`` entrypoint backends construct +the environment once and own stepping, checkpoint I/O, logging and evaluation. No external +SuperMimic package or launcher is required. The task does not import the learner, W&B or PhysicsNeMo. + +Transition capture and storage +------------------------------ + +``env.recorders.transition`` is disabled in the task's default configuration. SuperMimic training +enables ``TransitionRecorderCfg`` and reads its ``latest`` tensor dictionary after each control step. +This term captures the state, action, prior action, randomized context, reference index, reward, +episode step and termination flags. It captures the true successor before an automatic reset. +Its preallocated tensors are overwritten on the next step; a consumer that retains transitions +must copy them into its own storage. The learner does this in its time/world replay ring. + +The term uses the normal recorder lifecycle. It does not replace environment stepping or own replay. +By default, ``record_to_episodes=False`` avoids retaining per-world episodes. Offline consumers can +enable that setting and configure ``env.recorders.dataset_export_mode`` and the export directory +to use the same capture with the manager's batched episode storage and HDF5 exporter. +Other tasks can define their own :class:`~isaaclab.managers.RecorderTerm` schema at task-config level; +the core manager has no dependency on G1 or a learning algorithm. + +Prepare the robot and motions +----------------------------- + +Use a CUDA-capable Isaac Lab installation with the Newton backend and its standalone MJCF importer. +Start with the 40-motion retargeted LAFAN1 corpus. The preparation command downloads the robot +description and motions at a pinned dataset revision and records file hashes. The conversion +command is Isaac Lab's existing MJCF-to-USD tool: + +.. code-block:: bash + + uv run python scripts/tools/prepare_g1_tracking.py --output /data/g1_tracking + uv run python scripts/tools/convert_mjcf.py \ + /data/g1_tracking/g1_description/g1_29dof.xml \ + /data/g1_tracking/g1_usd/g1_29dof.usd + +The converter writes the composed asset to +``/data/g1_tracking/g1_usd/g1_29dof/g1_29dof.usda``. To reuse existing local inputs, supply +``--source_dir /path/to/prepared_inputs`` to the first command and choose an empty output directory. +That directory must contain ``g1_description/g1_29dof_rev_1_0.urdf``, its meshes and ``lafan1_g1``. +Local reuse records the local source and hashes, rather than claiming a downloaded revision. +The task defaults to the same layout under ``/data/g1_tracking``; the examples below +override those paths for an external data volume. + +For BONES-SEED, obtain the ``g1.tar.gz`` archive from the ``isaac/bones-seed-g1:1`` dataset and run: + +.. code-block:: bash + + uv run python scripts/tools/prepare_bones.py \ + --archive /data/bones-seed-g1/g1.tar.gz \ + --mjcf /data/g1_tracking/g1_description/g1_29dof.xml \ + --output /data/bones_prepared --workers 16 + +The converter reads named CSV columns, converts centimetres/degrees to metres/radians and resamples +120 Hz poses to 50 Hz. It writes retained poses under ``motions/``, per-clip decisions in +``motions.jsonl`` and filter settings and hashes in ``manifest.json``. The geometric and name filters +are explicit independent choices; the output is not claimed to match the authors' 129,514-clip subset. +Robot and motion data remain external inputs subject to their source licenses. BONES motion data +is attributed to Bones Studio. + +Train and resume +---------------- + +Train the MLP baseline through the unified entrypoint: + +.. code-block:: bash + + uv run python scripts/reinforcement_learning/train.py \ + --rl_library supermimic --task IsaacContrib-Motion-Tracking-G1 \ + --num_envs 4096 --log_dir logs/supermimic/g1_mlp \ + env.scene.robot.spawn.usd_path=/data/g1_tracking/g1_usd/g1_29dof/g1_29dof.usda \ + env.commands.motion.kinematics_path=/data/g1_tracking/g1_description/g1_29dof.xml \ + env.commands.motion.motion_dir=/data/g1_tracking/lafan1_g1 + +The default agent uses BF16 neural projections, FP32 geometry and compiled tensor rollouts. +Task settings use ``env.*`` overrides; optimizer, replay and runner settings use ``agent.*``. +The analytical imagined objective requires the task's observation and termination functions. +Observation history, stateless policy noise and task reward parameters remain configurable. +The runner rejects incompatible observation functions or layouts instead of silently optimizing +a different imagined task. + +For a short integration smoke test, use ``--num_envs 32 --max_iterations 2`` and append: + +.. code-block:: bash + + agent.buffer_steps=16 agent.warmup_steps=4 agent.world_model_window=2 \ + agent.world_model_batch_size=32 agent.world_model_batches=1 agent.policy_batch_size=32 \ + agent.horizon=4 agent.critic_epochs=1 agent.critic_minibatches=1 \ + agent.world_model_only_iterations=0 agent.log_interval=1 \ + env.observations.policy.history_length=3 env.commands.motion.goal_frames=2 + +These reduced settings test the pipeline; they are not a learning benchmark. For the prepared +BONES corpus, replace the motion directory with ``/data/bones_prepared/motions`` and set +``env.commands.motion.source_fps=50``. For a nine-hour budget, add +``--max_iterations 250000 agent.max_hours=9``. The budget is cumulative when resuming and is +checked after each update. Initial environment setup and replay warmup are outside that budget. + +Optional NeRD world model +~~~~~~~~~~~~~~~~~~~~~~~~~ + +Set ``agent.world_model=nerd-one-frame`` to select PhysicsNeMo Robotics' ``ParallelNeRD`` with +four attention spans of one. The policy, critic, task, replay and optimization objective are shared +with the MLP configuration. NeRD is used inside the learner; the physical teacher still uses +MuJoCo Warp. This option does not deploy a learned collision or constraint solver into Newton. + +This experimental option requires a compatible PhysicsNeMo Robotics installation exposing +``physicsnemo_robotics.models.parallel_nerd.ParallelNeRD``. It is imported only when selected; +the default MLP workflow has no PhysicsNeMo dependency. The integration was validated with +PhysicsNeMo Robotics ``0.1.0.dev0`` at revision ``a81160ef`` and PhysicsNeMo ``2.1.1`` in the +development environment. A clean, released package combination is not established here; +the experimental dependency must be supplied separately. + +The offline NeRD workflow below uses that same network implementation with temporal attention. +Its generalized-coordinate state and fixed offline normalization differ from SuperMimic's +task-specific state and online statistics. Those representations are intentionally distinct. + +Logging and preemption +~~~~~~~~~~~~~~~~~~~~~~ + +Every run writes resolved configurations, a run manifest and ``train_log.csv``. To enable W&B, +authenticate its SDK and append ``agent.wandb_project= agent.wandb_entity=``. +An optional stable ``agent.wandb_run_id=`` lets resumed runs continue the same history. +The runner records the actual run identity in ``wandb-run.json`` and the checkpoint. + +``last.pt`` contains the world model, policy, critic, their optimizers, normalizers, adaptive sampler +statistics and semantic configurations. It is replaced atomically after the first update, at the +configured interval, after 120 seconds by default, and on completion. When W&B is enabled, each +save is uploaded as a version of ``run--checkpoint`` with the ``latest`` alias; the runner waits +for upload before reporting ``durable_checkpoint_iteration``. Milestones also save +``policy_.pt`` for playback. SIGINT/SIGTERM request a save after the current update. +An immediate kill can lose progress since the last completed checkpoint. + +To resume, repeat the original task and learning overrides, add ``--checkpoint /path/to/last.pt`` +and increase ``--max_iterations``. Logging, save intervals and budget settings can change. +Physical state and replay are rebuilt; model and optimizer state and sampler statistics resume. +This is a training continuation, not a bitwise continuation of the simulator trajectory. +The runner accepts its native checkpoint metadata; legacy external-run checkpoints require +explicit conversion before use. + +Evaluate and record video +------------------------- + +Playback restores the native checkpoint's task settings and evaluates every clip exactly once: + +.. code-block:: bash + + uv run python scripts/reinforcement_learning/play.py \ + --rl_library supermimic --checkpoint logs/supermimic/g1_mlp/policy_014000.pt \ + --num_envs 64 --output logs/supermimic/g1_mlp/evaluation.json + +The output includes the failure count, failed names, per-clip returns and lengths, and both +clip-weighted and frame-weighted normalized returns. Failed tails remain in the full-clip +denominator. Use these physical evaluation metrics to assess policy quality; dynamics loss +measures prediction fit and can improve without improving tracking success. + +Use ``--motion_dir`` to select another corpus in the same pose schema and source frame rate. +``--asset_path`` and ``--kinematics_path`` relocate the matching assets. Add the standard +``--video --video_length 600`` options to record playback with Isaac Lab's native renderer. +For an interpretable video, select a small motion directory and a small ``--num_envs`` count. + +Validation scope +---------------- + +The native tests cover terminal capture, observation/kinematics agreement, episode export, +stationary goal padding, replay boundaries, complete evaluation batches and gradients through +compiled MLP/NeRD dynamics. Asset-dependent tests use ``ISAACLAB_G1_TRACKING_DATA`` pointing to +the prepared directory above. The dependency-direction tests prevent reintroducing an external +training forwarder or moving learner ownership into the task. + +The source adaptation preserves the upstream SuperMimic MIT notice in +``docs/licenses/dependencies/supermimic-license.txt``. diff --git a/scripts/imitation_learning/isaaclab_mimic/annotate_demos.py b/scripts/imitation_learning/isaaclab_mimic/annotate_demos.py index 1659510f7412..5a232f90fbc2 100644 --- a/scripts/imitation_learning/isaaclab_mimic/annotate_demos.py +++ b/scripts/imitation_learning/isaaclab_mimic/annotate_demos.py @@ -515,7 +515,6 @@ def annotate_episode_in_manual_mode( print(f'\tThe episode will be replayed again for re-marking subtask signals for the eef "{eef_name}".\n') - annotated_episode = env.recorder_manager.get_episode(0) for ( subtask_term_signal_name, subtask_term_signal_action_index, @@ -523,7 +522,9 @@ def annotate_episode_in_manual_mode( # subtask termination signal is false until subtask is complete, and true afterwards subtask_signals = torch.ones(len(episode.data["actions"]), dtype=torch.bool) subtask_signals[:subtask_term_signal_action_index] = False - annotated_episode.add(f"obs/datagen_info/subtask_term_signals/{subtask_term_signal_name}", subtask_signals) + env.recorder_manager.add_to_episodes( + f"obs/datagen_info/subtask_term_signals/{subtask_term_signal_name}", subtask_signals.unsqueeze(0), [0] + ) if args_cli.annotate_subtask_start_signals: for ( @@ -532,8 +533,8 @@ def annotate_episode_in_manual_mode( ) in subtask_start_signal_action_indices.items(): subtask_signals = torch.ones(len(episode.data["actions"]), dtype=torch.bool) subtask_signals[:subtask_start_signal_action_index] = False - annotated_episode.add( - f"obs/datagen_info/subtask_start_signals/{subtask_start_signal_name}", subtask_signals + env.recorder_manager.add_to_episodes( + f"obs/datagen_info/subtask_start_signals/{subtask_start_signal_name}", subtask_signals.unsqueeze(0), [0] ) return True diff --git a/scripts/tools/prepare_bones.py b/scripts/tools/prepare_bones.py new file mode 100644 index 000000000000..ff04d30d1004 --- /dev/null +++ b/scripts/tools/prepare_bones.py @@ -0,0 +1,247 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Portions derived from SuperMimic (MIT), Copyright (c) 2026 ETH Zurich, Chong Zhang. +# See docs/licenses/dependencies/supermimic-license.txt. + +"""Convert the public BONES-SEED archive into metric G1 poses and a versioned selection manifest. + +The name exclusions adapt SONIC's published filter, retaining floor acrobatics. The geometric +checks follow SuperMimic Appendix E; the unreported stationary gates are explicit below. +This independent selection is not claimed to match the authors' exact 129514-file manifest. +""" + +import argparse +import csv +import hashlib +import io +import json +import multiprocessing as mp +import tarfile +from collections import Counter, deque +from pathlib import Path + +import mujoco +import numpy as np +import torch +from scipy.spatial.transform import Rotation + +from isaaclab_tasks.contrib.g1_tracking.kinematics import G1_JOINT_NAMES +from isaaclab_tasks.contrib.g1_tracking.motion import CONTROL_DT, resample_motion + +# Derived from NVlabs/GR00T-WholeBodyControl filter_and_copy_bones_data.py, +# blob f09061012a1f1d1da3ef35d614440629138659cd. Floor acrobatics remain included. +EXCLUDED_NAMES = ( + "bed", + "bike", + "chair", + "climb", + "com_up_50cm", + "sitting", + "step_on", + "seat", + "table", + "sit_", + "ladder", + "crutch", + "_ride_", + "scooter", + "stepdown", + "box_hspu", + "50cm_box_", + "on_box", + "fall_from", + "on_1m", + "form_box", + "off_1m", + "230m", + "jump_over_obstacle_", + "lift_crate_come_up_", + "stair", + "box_jump", + "monkey_jump", + "box_dips", + "walking_on_edge", + "push_obstacle", + "jump_off_50cm", + "look_over_fence", +) +FILTER = { + "elevated_height_m": 0.20, + "elevated_seconds": 0.6, + "penetration_m": -0.12, + "sustained_penetration_m": -0.08, + "penetration_fraction": 0.8, + "floating_clearance_m": 0.02, + "floating_seconds": 0.5, + "standing_pelvis_m": 0.55, + "upright_cosine": 0.9, + "stationary_max_link_speed_m_s": 0.2, + "stationary_gate_provenance": "Independent choices; these gates are not specified numerically in the paper.", +} + + +def initialize_worker(mjcf_path, output, source_fps): + """Own one MuJoCo model/data pair per conversion process; never share mutable simulator data.""" + global MODEL, DATA, FLOOR, FEET, OUTPUT, SOURCE_FPS + torch.set_num_threads(1) + spec = mujoco.MjSpec.from_file(mjcf_path) + spec.worldbody.add_geom(name="filter_floor", type=mujoco.mjtGeom.mjGEOM_PLANE, size=[0, 0, 0.1]) + MODEL = spec.compile() + DATA = mujoco.MjData(MODEL) + FLOOR = MODEL.geom("filter_floor").id + FEET = [] + for side in ("left", "right"): + body = MODEL.body(f"{side}_ankle_roll_link").id + geoms = np.flatnonzero((MODEL.geom_bodyid == body) & ((MODEL.geom_contype | MODEL.geom_conaffinity) != 0)) + if len(geoms) == 0: + raise ValueError(f"No collision geometry on {side} foot.") + FEET.append(geoms) + joint_names = [MODEL.joint(j).name for j in range(MODEL.njnt) if MODEL.jnt_type[j] != mujoco.mjtJoint.mjJNT_FREE] + if joint_names != G1_JOINT_NAMES: + raise ValueError("Filtering model joint order differs from the pose schema.") + OUTPUT, SOURCE_FPS = Path(output), source_fps + + +def convert_csv(data): + """BONES header-defined cm/degree/Euler poses to metric xyzw/radian poses in task joint order.""" + header = next(csv.reader(io.StringIO(data.decode("utf-8-sig")))) + expected = [ + "Frame", + *[f"root_translate{x}" for x in "XYZ"], + *[f"root_rotate{x}" for x in "XYZ"], + *[f"{name}_dof" for name in G1_JOINT_NAMES], + ] + if len(set(header)) != len(header) or set(header) != set(expected): + raise ValueError(f"Unexpected BONES CSV fields: {header}") + raw = np.loadtxt(io.BytesIO(data), delimiter=",", skiprows=1, ndmin=2) + raw = raw[:, [header.index(name) for name in expected]] + if len(raw) < 4 or not np.isfinite(raw).all() or not np.allclose(np.diff(raw[:, 0]), 1, atol=0, rtol=0): + raise ValueError("Invalid pose values or nonconsecutive source frame indices.") + quat = Rotation.from_euler("xyz", raw[:, 4:7], degrees=True).as_quat() + return np.concatenate((raw[:, 1:4] / 100, quat, np.deg2rad(raw[:, 7:])), axis=1).astype(np.float32) + + +def sustained(mask, seconds): + n = int(round(seconds / CONTROL_DT)) + return len(mask) >= n and bool((np.convolve(mask.astype(np.int32), np.ones(n, dtype=np.int32), "valid") == n).any()) + + +def process_motion(name, data): + record = {"source": name, "source_sha256": hashlib.sha256(data).hexdigest()} + raw = convert_csv(data) + p, quat, q = resample_motion(torch.from_numpy(raw), SOURCE_FPS) + poses = torch.cat((p, quat, q), dim=-1).numpy() + n = len(poses) + zmin, clearance, positions = np.empty(n), np.empty(n), [] + for i, pose in enumerate(poses): + DATA.qpos[:] = np.concatenate((pose[:3], pose[[6, 3, 4, 5]], pose[7:])) + mujoco.mj_kinematics(MODEL, DATA) + zmin[i] = DATA.xpos[1:, 2].min() + clearance[i] = min(mujoco.mj_geomDistance(MODEL, DATA, int(g), FLOOR, 10, None) for gs in FEET for g in gs) + positions.append(DATA.xpos[1:].copy()) + max_speed = np.linalg.norm(np.gradient(np.stack(positions), CONTROL_DT, axis=0), axis=-1).max(axis=-1) + up = Rotation.from_quat(poses[:, 3:7]).as_matrix()[:, 2, 2] + floating = ( + (clearance > FILTER["floating_clearance_m"]) + & (poses[:, 2] > FILTER["standing_pelvis_m"]) + & (up > FILTER["upright_cosine"]) + & (max_speed < FILTER["stationary_max_link_speed_m_s"]) + ) + reason = None + if sustained(zmin > FILTER["elevated_height_m"], FILTER["elevated_seconds"]): + reason = "elevated_support" + elif (zmin < FILTER["penetration_m"]).any() or (zmin < FILTER["sustained_penetration_m"]).mean() > FILTER[ + "penetration_fraction" + ]: + reason = "ground_penetration" + elif sustained(floating, FILTER["floating_seconds"]): + reason = "floating_feet" + record.update(source_frames=len(raw), frames=n, reason=reason) + if reason is None: + path = Path(name).with_suffix(".npy") + target = OUTPUT / "motions" / path + target.parent.mkdir(parents=True, exist_ok=True) + np.save(target, poses, allow_pickle=False) + record.update(path=path.as_posix(), pose_sha256=hashlib.sha256(poses.tobytes()).hexdigest()) + return record + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--archive", required=True) + parser.add_argument("--mjcf", required=True, help="Matching 29-DoF G1 MJCF from prepare_g1_tracking.py.") + parser.add_argument("--output", required=True) + parser.add_argument("--workers", type=int, default=16) + parser.add_argument("--source_fps", type=float, default=120) + args = parser.parse_args() + output = Path(args.output) + if output.exists() and any(output.iterdir()): + parser.error("Output must be empty; do not mix dataset versions.") + output.mkdir(parents=True, exist_ok=True) + counts, kept_frames, total = Counter(), 0, 0 + manifest_path = output / "motions.jsonl" + with ( + manifest_path.open("w") as log, + mp.get_context("spawn").Pool( + args.workers, initializer=initialize_worker, initargs=(args.mjcf, str(output), args.source_fps) + ) as pool, + ): + pending = deque() + + def collect_one(): + nonlocal kept_frames + result = pending.popleft().get() + reason = result["reason"] or "kept" + counts[reason] += 1 + if reason == "kept": + kept_frames += result["frames"] + log.write(json.dumps(result) + "\n") + if sum(counts.values()) % 5000 == 0: + log.flush() + print("PROGRESS", dict(counts), "kept_frames", kept_frames, flush=True) + + with tarfile.open(args.archive, "r|gz") as archive: + for member in archive: + if not member.isfile() or not member.name.endswith(".csv"): + continue + name = Path(member.name) + if name.is_absolute() or ".." in name.parts: + raise ValueError(f"Unsafe archive path: {name}") + total += 1 + keyword = next((key for key in EXCLUDED_NAMES if key in name.stem.lower()), None) + if keyword: + counts["name"] += 1 + log.write(json.dumps({"source": member.name, "reason": "name", "keyword": keyword}) + "\n") + continue + pending.append(pool.apply_async(process_motion, (member.name, archive.extractfile(member).read()))) + if len(pending) >= 2 * args.workers: + collect_one() + while pending: + collect_one() + summary = { + "source": "isaac/bones-seed-g1:1", + "attribution": "Motion Data by Bones Studio", + "raw_clips": total, + "counts": dict(counts), + "kept_frames": kept_frames, + "kept_hours": (kept_frames - counts["kept"]) * CONTROL_DT / 3600, + "source_fps": args.source_fps, + "prepared_fps": 1 / CONTROL_DT, + "layout": "root_xyz_m,root_quaternion_xyzw,29_joints_rad", + "joint_names": G1_JOINT_NAMES, + "name_exclusions": EXCLUDED_NAMES, + "geometric_filter": FILTER, + "exact_paper_subset": False, + "selection_sha256": hashlib.sha256(manifest_path.read_bytes()).hexdigest(), + } + (output / "manifest.json").write_text(json.dumps(summary, indent=2) + "\n") + print("PREPARATION_COMPLETE", json.dumps(summary), flush=True) + if not total or not counts["kept"]: + raise ValueError("No motions retained; inspect the selection manifest.") + + +if __name__ == "__main__": + main() diff --git a/scripts/tools/prepare_g1_tracking.py b/scripts/tools/prepare_g1_tracking.py new file mode 100644 index 000000000000..7aabd90b680e --- /dev/null +++ b/scripts/tools/prepare_g1_tracking.py @@ -0,0 +1,96 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Prepare pinned G1 tracking inputs; convert the resulting MJCF with scripts/tools/convert_mjcf.py. + +Downloads the retargeted LAFAN1 motions and matching Unitree robot description, retaining the source +revision in a manifest. The datasets retain their original licenses; this script does not redistribute them. +""" + +import argparse +import hashlib +import json +import shutil +import tempfile +import urllib.request +import xml.etree.ElementTree as ET +from pathlib import Path + +import mujoco + +REVISION = "ce1572906efe6157840e8474d5a0d7aa87481e74" +REPOSITORY = "lvhaidong/LAFAN1_Retargeting_Dataset" + + +def main() -> None: + """Download the inputs or reuse a local copy, then compile a free-floating 29-DoF G1 MJCF.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument( + "--source_dir", type=Path, help="Existing prepared directory with g1_description and lafan1_g1." + ) + args = parser.parse_args() + output = args.output.resolve() + if output.exists() and any(output.iterdir()): + parser.error("Output must be empty to keep preparation reproducible.") + output.mkdir(parents=True, exist_ok=True) + if args.source_dir: + source = {"local_directory": str(args.source_dir.resolve())} + for name in ("g1_description", "lafan1_g1"): + shutil.copytree(args.source_dir / name, output / name) + else: + source = {"repository": REPOSITORY, "revision": REVISION} + paths = ["robot_description/g1/g1_29dof_rev_1_0.urdf"] + for directory in ("g1", "robot_description/g1/meshes"): + with urllib.request.urlopen( + f"https://huggingface.co/api/datasets/{REPOSITORY}/tree/{REVISION}/{directory}?limit=1000" + ) as response: + paths.extend(entry["path"] for entry in json.load(response) if entry["type"] == "file") + for name in paths: + relative = Path(name) + if relative.parts[0] == "g1": + target = output / "lafan1_g1" / relative.name + else: + target = output / "g1_description" / relative.relative_to("robot_description/g1") + target.parent.mkdir(parents=True, exist_ok=True) + urllib.request.urlretrieve( + f"https://huggingface.co/datasets/{REPOSITORY}/resolve/{REVISION}/{name}", target + ) + description = output / "g1_description" + robot = ET.parse(description / "g1_29dof_rev_1_0.urdf") + root = robot.getroot() + for element in root.findall("mujoco"): + root.remove(element) + compiler = ET.SubElement(ET.SubElement(root, "mujoco"), "compiler") + compiler.attrib.update( + meshdir=".", balanceinertia="true", discardvisual="false", fusestatic="false", strippath="false" + ) + with tempfile.NamedTemporaryFile(dir=description, suffix=".urdf") as temporary: + robot.write(temporary.name) + model = mujoco.MjModel.from_xml_path(temporary.name) + mjcf = description / "g1_29dof.xml" + mujoco.mj_saveLastXML(str(mjcf), model) + compiled = ET.parse(mjcf) + pelvis = compiled.find(".//body[@name='pelvis']") + if pelvis is None: + raise ValueError("The pinned robot must retain its pelvis body.") + pelvis.insert(0, ET.Element("freejoint", name="floating_base_joint")) + compiled.write(mjcf) + model = mujoco.MjModel.from_xml_path(str(mjcf)) + if model.nq != 36 or model.nv != 35: + raise ValueError("Expected a free-floating 29-DoF robot.") + source["sha256"] = {} + for path in sorted(output.rglob("*")): + if path.is_file(): + with path.open("rb") as stream: + source["sha256"][path.relative_to(output).as_posix()] = hashlib.file_digest( + stream, "sha256" + ).hexdigest() + (output / "source.json").write_text(json.dumps(source, indent=2) + "\n") + print(f"Prepared {mjcf}; convert it with scripts/tools/convert_mjcf.py before training.") + + +if __name__ == "__main__": + main() diff --git a/scripts/tools/test/test_prepare_bones.py b/scripts/tools/test/test_prepare_bones.py new file mode 100644 index 000000000000..8692dbb50059 --- /dev/null +++ b/scripts/tools/test/test_prepare_bones.py @@ -0,0 +1,54 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Independent unit, frame-rate and physical-support contracts for BONES preparation.""" + +import io +import os +from pathlib import Path + +import numpy as np +import pytest +from scipy.spatial.transform import Rotation + +from isaaclab_tasks.contrib.g1_tracking.kinematics import G1_JOINT_NAMES + +from scripts.tools import prepare_bones + + +@pytest.fixture +def elevated_clip(): + columns = [ + "Frame", + *[f"root_translate{x}" for x in "XYZ"], + *[f"root_rotate{x}" for x in "XYZ"], + *[f"{name}_dof" for name in G1_JOINT_NAMES], + ] + values = np.zeros((121, len(columns))) + values[:, 0], values[:, 3], values[:, 6], values[:, 7] = np.arange(121), 300, 90, 45 + order = np.arange(len(columns))[::-1] + stream = io.StringIO() + np.savetxt(stream, values[:, order], delimiter=",", header=",".join(columns[i] for i in order), comments="") + return stream.getvalue().encode() + + +def test_conversion_uses_header_order_and_named_units(elevated_clip): + poses = prepare_bones.convert_csv(elevated_clip) + np.testing.assert_allclose(poses[:, :3], np.tile([0, 0, 3], (121, 1))) + np.testing.assert_allclose(poses[:, 7], np.pi / 4) + np.testing.assert_allclose(Rotation.from_quat(poses[0, 3:7]).apply([1, 0, 0]), [0, 1, 0], atol=1e-6) + with pytest.raises(ValueError, match="fields"): + prepare_bones.convert_csv(elevated_clip.replace(b"root_translateX", b"root_translateY")) + + +def test_filter_resamples_and_rejects_missing_support(elevated_clip, tmp_path): + data_dir = os.environ.get("ISAACLAB_G1_TRACKING_DATA") + if not data_dir: + pytest.skip("Set ISAACLAB_G1_TRACKING_DATA to the prepared G1 tracking directory.") + mjcf = Path(data_dir) / "g1_description/g1_29dof.xml" + prepare_bones.initialize_worker(str(mjcf), str(tmp_path), 120) + result = prepare_bones.process_motion("session/elevated.csv", elevated_clip) + assert result["reason"] == "elevated_support" and result["frames"] == 51 + assert not list(tmp_path.rglob("*.npy")) diff --git a/source/isaaclab/changelog.d/quaternion-zero-gradient.fixed.rst b/source/isaaclab/changelog.d/quaternion-zero-gradient.fixed.rst new file mode 100644 index 000000000000..4ab770ab120a --- /dev/null +++ b/source/isaaclab/changelog.d/quaternion-zero-gradient.fixed.rst @@ -0,0 +1,2 @@ +* Fixed NaN gradients in ``axis_angle_from_quat`` at identity rotations by keeping the + unselected exact-form division finite during the Taylor approximation. diff --git a/source/isaaclab/changelog.d/recorder-capture.changed.rst b/source/isaaclab/changelog.d/recorder-capture.changed.rst new file mode 100644 index 000000000000..d0eebccaa3cd --- /dev/null +++ b/source/isaaclab/changelog.d/recorder-capture.changed.rst @@ -0,0 +1,9 @@ +* Replaced per-world recorder episode containers with batched field allocations and + independent insertion indices, reusing storage across resets and growing without dropping samples. +* Added :meth:`~isaaclab.managers.RecorderManager.get_term` and the term capability + :attr:`~isaaclab.managers.RecorderTerm.requires_post_step_observations` so physical-state + capture can avoid extra observation computation and noise draws. +* **Breaking:** Changed :meth:`~isaaclab.managers.RecorderManager.get_episode` to return an owned + snapshot with tensor leaves. To update live recording, use ``add_to_episodes`` and + ``set_success_to_episodes``; do not mutate the snapshot. Set ``buffer_initial_capacity`` + to control the initial allocation per field. diff --git a/source/isaaclab/changelog.d/recorder-capture.major b/source/isaaclab/changelog.d/recorder-capture.major new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab/changelog.d/recorder-inference-lifetime.fixed.rst b/source/isaaclab/changelog.d/recorder-inference-lifetime.fixed.rst new file mode 100644 index 000000000000..1302fe5c27d3 --- /dev/null +++ b/source/isaaclab/changelog.d/recorder-inference-lifetime.fixed.rst @@ -0,0 +1,2 @@ +* Kept retained recorder allocations writable after leaving PyTorch inference mode, including + grown buffers and schema changes, so subsequent reset and export can finalize those episodes. diff --git a/source/isaaclab/changelog.d/supermimic.added.rst b/source/isaaclab/changelog.d/supermimic.added.rst new file mode 100644 index 000000000000..3babbcc46a52 --- /dev/null +++ b/source/isaaclab/changelog.d/supermimic.added.rst @@ -0,0 +1 @@ +* Added SuperMimic to the training-library choices reported by the task-listing CLI. diff --git a/source/isaaclab/isaaclab/cli/commands/list_envs.py b/source/isaaclab/isaaclab/cli/commands/list_envs.py index fdba9c8edf1a..d2ea258a99d0 100644 --- a/source/isaaclab/isaaclab/cli/commands/list_envs.py +++ b/source/isaaclab/isaaclab/cli/commands/list_envs.py @@ -73,7 +73,7 @@ def command_list_envs(args: list[str] | None = None) -> None: for index, spec in enumerate(task_specs, start=1): libraries = [ library - for library in ("rl_games", "rlinf", "rsl_rl", "sb3", "skrl", "torchrl") + for library in ("rl_games", "rlinf", "rsl_rl", "sb3", "skrl", "supermimic", "torchrl") if spec.kwargs.get(f"{library}_cfg_entry_point") is not None ] row = [index, spec.id, spec.entry_point, spec.kwargs["env_cfg_entry_point"], ", ".join(libraries) or "(none)"] diff --git a/source/isaaclab/isaaclab/envs/manager_based_rl_env.py b/source/isaaclab/isaaclab/envs/manager_based_rl_env.py index 859cb837cd5b..122fb5087a6d 100644 --- a/source/isaaclab/isaaclab/envs/manager_based_rl_env.py +++ b/source/isaaclab/isaaclab/envs/manager_based_rl_env.py @@ -237,8 +237,8 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: self.reward_buf = self.reward_manager.compute(dt=self.step_dt) if len(self.recorder_manager.active_terms) > 0: - # update observations for recording if needed - self.obs_buf = self.observation_manager.compute() + if self.recorder_manager.requires_post_step_observations: + self.obs_buf = self.observation_manager.compute() self.recorder_manager.record_post_step() # -- reset envs that terminated/timed-out and log the episode information diff --git a/source/isaaclab/isaaclab/envs/mdp/recorders/recorders.py b/source/isaaclab/isaaclab/envs/mdp/recorders/recorders.py index a9315bbca63d..44252e55a030 100644 --- a/source/isaaclab/isaaclab/envs/mdp/recorders/recorders.py +++ b/source/isaaclab/isaaclab/envs/mdp/recorders/recorders.py @@ -14,6 +14,8 @@ class InitialStateRecorder(RecorderTerm): """Recorder term that records the initial state of the environment after reset.""" + requires_post_step_observations = False + def record_post_reset(self, env_ids: Sequence[int] | None): def extract_env_ids_values(value): nonlocal env_ids @@ -27,6 +29,8 @@ def extract_env_ids_values(value): class PostStepStatesRecorder(RecorderTerm): """Recorder term that records the state of the environment at the end of each step.""" + requires_post_step_observations = False + def record_post_step(self): return "states", self._env.scene.get_state(is_relative=True) @@ -34,6 +38,8 @@ def record_post_step(self): class PreStepActionsRecorder(RecorderTerm): """Recorder term that records the actions in the beginning of each step.""" + requires_post_step_observations = False + def record_pre_step(self): return "actions", self._env.action_manager.action @@ -41,6 +47,8 @@ def record_pre_step(self): class PreStepFlatPolicyObservationsRecorder(RecorderTerm): """Recorder term that records the policy group observations in each step.""" + requires_post_step_observations = False + def record_pre_step(self): return "obs", self._env.obs_buf["policy"] @@ -48,6 +56,8 @@ def record_pre_step(self): class PostStepProcessedActionsRecorder(RecorderTerm): """Recorder term that records processed actions at the end of each step.""" + requires_post_step_observations = False + def record_post_step(self): processed_actions = None diff --git a/source/isaaclab/isaaclab/managers/recorder_manager.py b/source/isaaclab/isaaclab/managers/recorder_manager.py index f4894840afb9..3ac51533f963 100644 --- a/source/isaaclab/isaaclab/managers/recorder_manager.py +++ b/source/isaaclab/isaaclab/managers/recorder_manager.py @@ -54,6 +54,14 @@ class RecorderManagerBaseCfg: export_in_close: bool = False """Whether to export episodes in the close call.""" + buffer_initial_capacity: int = 64 + """Initial samples per world and field. Must be positive. + + Fields allocate on first use and double their capacity when necessary, retaining all samples. + Resets reuse storage. Set this to the expected episode length (including decimation callbacks) + to avoid growth during recording. Each field has independent write indices. + """ + dataset_compression: bool = True """Enable dataset compression.""" @@ -71,6 +79,13 @@ class RecorderTerm(ManagerTermBase): * Post-step recording: This callback is invoked at the end of `env.step()` when all the managers are processed. """ + requires_post_step_observations: bool = True + """Whether the RL environment must compute observations before this term's post-step callback. + + Terms that only read physical tensors should set this to False. The default preserves the + observation access of existing terms; disabling it avoids extra noise draws and computation. + """ + def __init__(self, cfg: RecorderTermCfg, env: ManagerBasedEnv): """Initialize the recorder term. @@ -163,6 +178,8 @@ def __init__(self, cfg: object, env: ManagerBasedEnv): """ self._term_names: list[str] = [] self._terms: dict[str, RecorderTerm] = {} + self._dataset_file_handler = None + self._failed_episode_dataset_file_handler = None # Do nothing if cfg is None or an empty dict if not cfg: @@ -177,30 +194,28 @@ def __init__(self, cfg: object, env: ManagerBasedEnv): if not isinstance(cfg, RecorderManagerBaseCfg): raise TypeError("Configuration for the recorder manager is not of type RecorderManagerBaseCfg.") - # create episode data buffer indexed by environment id - self._episodes: dict[int, EpisodeData] = dict() - for env_id in range(env.num_envs): - self._episodes[env_id] = EpisodeData() + if cfg.buffer_initial_capacity < 1: + raise ValueError("buffer_initial_capacity must be positive.") + self._fields: list[_RecorderField] = [] + self._env_ids = torch.arange(env.num_envs, device=env.device) + # -1 denotes an episode whose success has not been set. + self._success = torch.full((env.num_envs,), -1, dtype=torch.int8, device=env.device) + self._exported_counts = torch.zeros(2, dtype=torch.long, device=env.device) env_name = getattr(env.cfg, "env_name", None) - self._dataset_file_handler = None if cfg.dataset_export_mode != DatasetExportMode.EXPORT_NONE: self._dataset_file_handler = cfg.dataset_file_handler_class_type() self._dataset_file_handler.create( os.path.join(cfg.dataset_export_dir_path, cfg.dataset_filename), env_name=env_name ) - self._failed_episode_dataset_file_handler = None if cfg.dataset_export_mode == DatasetExportMode.EXPORT_SUCCEEDED_FAILED_IN_SEPARATE_FILES: self._failed_episode_dataset_file_handler = cfg.dataset_file_handler_class_type() self._failed_episode_dataset_file_handler.create( os.path.join(cfg.dataset_export_dir_path, f"{cfg.dataset_filename}_failed"), env_name=env_name ) - self._exported_successful_episode_count = {} - self._exported_failed_episode_count = {} - def __str__(self) -> str: """Returns: A string representation for recorder manager.""" msg = f" contains {len(self._term_names)} active terms.\n" @@ -232,36 +247,31 @@ def active_terms(self) -> list[str]: return self._term_names @property - def exported_successful_episode_count(self, env_id=None) -> int: - """Number of successful episodes. + def requires_post_step_observations(self) -> bool: + """Whether any active recorder term requires post-step observations.""" + return any(term.requires_post_step_observations for term in self._terms.values()) - Args: - env_id: The environment id. Defaults to None, in which case all environments are considered. + @property + def exported_successful_episode_count(self) -> int: + """Number of successful episodes. Returns: The number of successful episodes. """ - if not hasattr(self, "_exported_successful_episode_count"): + if not hasattr(self, "_exported_counts"): return 0 - if env_id is not None: - return self._exported_successful_episode_count.get(env_id, 0) - return sum(self._exported_successful_episode_count.values()) + return int(self._exported_counts[1].item()) @property - def exported_failed_episode_count(self, env_id=None) -> int: + def exported_failed_episode_count(self) -> int: """Number of failed episodes. - Args: - env_id: The environment id. Defaults to None, in which case all environments are considered. - Returns: The number of failed episodes. """ - if not hasattr(self, "_exported_failed_episode_count"): + if not hasattr(self, "_exported_counts"): return 0 - if env_id is not None: - return self._exported_failed_episode_count.get(env_id, 0) - return sum(self._exported_failed_episode_count.values()) + return int(self._exported_counts[0].item()) """ Operations. @@ -280,96 +290,81 @@ def reset(self, env_ids: Sequence[int] | slice | None = None) -> dict[str, torch # Do nothing if no active recorder terms are provided if len(self.active_terms) == 0: return {} - # resolve environment ids - if env_ids is None: - env_ids = slice(None) - if isinstance(env_ids, slice): - env_ids = list(range(self._env.num_envs)[env_ids]) - if isinstance(env_ids, torch.Tensor): - env_ids = env_ids.tolist() - for term in self._terms.values(): term.reset(env_ids=env_ids) - - for env_id in env_ids: - self._episodes[env_id] = EpisodeData() - - # nothing to log here + self._reset_buffers(env_ids) return {} - def get_episode(self, env_id: int) -> EpisodeData: - """Returns the episode data for the given environment id. + def get_term(self, name: str) -> RecorderTerm: + """Return the configured recorder term with the given name.""" + return self._terms[name] - Args: - env_id: The environment id. + def get_episode(self, env_id: int) -> EpisodeData: + """Return an owned snapshot of a world's current episode. - Returns: - The episode data for the given environment id. + Leaves are tensors with a leading sample dimension, matching the exported dataset. + Reading synchronizes write indices and copies the recorded samples. Later recording, + reset, or snapshot mutation cannot affect the other copy. Add live data through + :meth:`add_to_episodes`, and set success through :meth:`set_success_to_episodes`. """ - return self._episodes.get(env_id, EpisodeData()) + episode = EpisodeData() + if not self.active_terms or not 0 <= env_id < self._env.num_envs: + return episode + for field in self._fields: + length = int(field.lengths[env_id].item()) + if length: + data = episode.data + keys = field.name.split("/") + for key in keys[:-1]: + data = data.setdefault(key, {}) + data[keys[-1]] = field.data[:length, env_id].clone() + success = int(self._success[env_id].item()) + episode.success = None if success < 0 else bool(success) + return episode def add_to_episodes(self, key: str, value: torch.Tensor | dict, env_ids: Sequence[int] | slice | None = None): - """Adds the given key-value pair to the episodes for the given environment ids. + """Append a batch to preallocated storage, advancing selected worlds' write indices. Args: - key: The key of the given value to be added to the episodes. The key can contain nested keys - separated by '/'. For example, "obs/joint_pos" would add the given value under ['obs']['policy'] - in the underlying dictionary in the episode data. - value: The value to be added to the episodes. The value can be a tensor or a nested dictionary of tensors. - The shape of a tensor in the value is (env_ids, ...). - env_ids: The environment ids. Defaults to None, in which case all environments are considered. + key: Field path, with nested keys separated by '/'. + value: Tensor of shape (selected worlds, ...) or nested dictionary of tensors. + Warp arrays are also accepted. Storage copies the values and preserves dtype/device. + A field's shape, dtype, and device must remain constant while it contains samples. + env_ids: Unique world indices in the same order as the value batch. None selects all worlds. """ - # Do nothing if no active recorder terms are provided - if len(self.active_terms) == 0: + if not self.active_terms or key is None: return - - # resolve environment ids - if key is None: - return - if env_ids is None: - env_ids = slice(None) - if isinstance(env_ids, slice): - env_ids = list(range(self._env.num_envs)[env_ids]) - if isinstance(env_ids, torch.Tensor): - env_ids = env_ids.tolist() - if isinstance(value, dict): for sub_key, sub_value in value.items(): - if isinstance(sub_value, wp.array): - sub_value = wp.to_torch(sub_value) self.add_to_episodes(f"{key}/{sub_key}", sub_value, env_ids) return - if isinstance(value, wp.array): value = wp.to_torch(value) - value = value.clone() # Clone once for all envs - for value_index, env_id in enumerate(env_ids): - if env_id not in self._episodes: - self._episodes[env_id] = EpisodeData() - self._episodes[env_id].env_id = env_id - self._episodes[env_id].add(key, value[value_index], clone=False) + num_selected = ( + self._env.num_envs + if env_ids is None + else (len(range(self._env.num_envs)[env_ids]) if isinstance(env_ids, slice) else len(env_ids)) + ) + if value.ndim < 1 or value.shape[0] != num_selected: + raise ValueError(f"Field '{key}' requires {num_selected} rows, got {tuple(value.shape)}.") + if not num_selected: + return + field = next((field for field in self._fields if field.name == key), None) + if field is None: + # Reject ambiguous leaf/subtree paths before allocating any storage. + if any(key.startswith(f"{f.name}/") or f.name.startswith(f"{key}/") for f in self._fields): + raise ValueError(f"Recorder field '{key}' conflicts with an existing field path.") + field = _RecorderField(key, value, self._env.num_envs, self.cfg.buffer_initial_capacity) + self._fields.append(field) + field.append(value, env_ids) def set_success_to_episodes(self, env_ids: Sequence[int] | slice | None, success_values: torch.Tensor): - """Sets the task success values to the episodes for the given environment ids. - - Args: - env_ids: The environment ids. Defaults to None, in which case all environments are considered. - success_values: The task success values to be set to the episodes. The shape of the tensor is (env_ids, 1). - """ - # Do nothing if no active recorder terms are provided - if len(self.active_terms) == 0: + """Set selected worlds' success flags from a tensor of shape (selected worlds,) or (selected worlds, 1).""" + if not self.active_terms: return - - # resolve environment ids if env_ids is None: env_ids = slice(None) - if isinstance(env_ids, slice): - env_ids = list(range(self._env.num_envs)[env_ids]) - if isinstance(env_ids, torch.Tensor): - env_ids = env_ids.tolist() - - for value_index, env_id in enumerate(env_ids): - self._episodes[env_id].success = success_values[value_index].item() + self._success[env_ids] = success_values.reshape(-1).to(device=self._success.device, dtype=torch.int8) def record_pre_step(self) -> None: """Trigger recorder terms for pre-step functions.""" @@ -412,11 +407,9 @@ def record_pre_reset(self, env_ids: Sequence[int] | slice | None, force_export_o return if env_ids is None: - env_ids = slice(None) - if isinstance(env_ids, slice): - env_ids = list(range(self._env.num_envs)[env_ids]) - if isinstance(env_ids, torch.Tensor): - env_ids = env_ids.tolist() + env_ids = self._env_ids + elif isinstance(env_ids, slice): + env_ids = self._env_ids[env_ids] for term in self._terms.values(): key, value = term.record_pre_reset(env_ids) @@ -444,7 +437,7 @@ def record_post_reset(self, env_ids: Sequence[int] | slice | None) -> None: return if isinstance(env_ids, slice): - env_ids = list(range(self._env.num_envs)[env_ids]) + env_ids = self._env_ids[env_ids] for term in self._terms.values(): key, value = term.record_post_reset(env_ids) @@ -485,11 +478,11 @@ def export_episodes( return if env_ids is None: - env_ids = slice(None) - if isinstance(env_ids, slice): - env_ids = list(range(self._env.num_envs)[env_ids]) - if isinstance(env_ids, torch.Tensor): - env_ids = env_ids.tolist() + env_ids = self._env_ids + elif isinstance(env_ids, slice): + env_ids = self._env_ids[env_ids] + else: + env_ids = torch.as_tensor(env_ids, dtype=torch.long, device=self._env_ids.device) # Handle demo_ids processing if demo_ids is not None: @@ -502,53 +495,43 @@ def export_episodes( duplicates = [x for i, x in enumerate(demo_ids) if demo_ids.index(x) != i] raise ValueError(f"demo_ids must be unique. Found duplicates: {list(set(duplicates))}") - # Export episode data through dataset exporter - need_to_flush = False + present = torch.zeros(self._env.num_envs, dtype=torch.bool, device=self._env_ids.device) + for field in self._fields: + present |= (field.lengths > 0).to(present.device) + succeeded = self._success == 1 - if any(env_id in self._episodes and not self._episodes[env_id].is_empty() for env_id in env_ids): + # Only file export materializes per-world snapshots. EXPORT_NONE stays batched. + need_to_flush = False + mode = self.cfg.dataset_export_mode + if mode != DatasetExportMode.EXPORT_NONE and present[env_ids].any().item(): ep_meta = self.get_ep_meta() - if self._dataset_file_handler is not None: - self._dataset_file_handler.add_env_args(ep_meta) - if self._failed_episode_dataset_file_handler is not None: - self._failed_episode_dataset_file_handler.add_env_args(ep_meta) - - for i, env_id in enumerate(env_ids): - if env_id in self._episodes and not self._episodes[env_id].is_empty(): - self._episodes[env_id].pre_export() - - episode_succeeded = self._episodes[env_id].success - target_dataset_file_handler = None - if (self.cfg.dataset_export_mode == DatasetExportMode.EXPORT_ALL) or ( - self.cfg.dataset_export_mode == DatasetExportMode.EXPORT_SUCCEEDED_ONLY and episode_succeeded - ): - target_dataset_file_handler = self._dataset_file_handler - elif self.cfg.dataset_export_mode == DatasetExportMode.EXPORT_SUCCEEDED_FAILED_IN_SEPARATE_FILES: - if episode_succeeded: - target_dataset_file_handler = self._dataset_file_handler - else: - target_dataset_file_handler = self._failed_episode_dataset_file_handler - if target_dataset_file_handler is not None: - # Use corresponding demo_id if provided, otherwise None - current_demo_id = demo_ids[i] if demo_ids is not None else None - target_dataset_file_handler.write_episode( - self._episodes[env_id], current_demo_id, self.cfg.dataset_compression - ) - need_to_flush = True - # Update episode count - if episode_succeeded: - self._exported_successful_episode_count[env_id] = ( - self._exported_successful_episode_count.get(env_id, 0) + 1 - ) - else: - self._exported_failed_episode_count[env_id] = self._exported_failed_episode_count.get(env_id, 0) + 1 - # Reset the episode buffer for the given environment after export - self._episodes[env_id] = EpisodeData() - + for handler in (self._dataset_file_handler, self._failed_episode_dataset_file_handler): + if handler is not None: + handler.add_env_args(ep_meta) + selected = env_ids.tolist() + selected_present = present[env_ids].tolist() + selected_success = succeeded[env_ids].tolist() + for i, env_id in enumerate(selected): + if not selected_present[i]: + continue + handler = self._dataset_file_handler + if mode == DatasetExportMode.EXPORT_SUCCEEDED_FAILED_IN_SEPARATE_FILES: + if not selected_success[i]: + handler = self._failed_episode_dataset_file_handler + elif mode == DatasetExportMode.EXPORT_SUCCEEDED_ONLY and not selected_success[i]: + continue + episode = self.get_episode(env_id) + demo_id = demo_ids[i] if demo_ids is not None else None + handler.write_episode(episode, demo_id, self.cfg.dataset_compression) + need_to_flush = True + + self._exported_counts[0] += (present & ~succeeded)[env_ids].sum() + self._exported_counts[1] += (present & succeeded)[env_ids].sum() + self._reset_buffers(env_ids) if need_to_flush: - if self._dataset_file_handler is not None: - self._dataset_file_handler.flush() - if self._failed_episode_dataset_file_handler is not None: - self._failed_episode_dataset_file_handler.flush() + for handler in (self._dataset_file_handler, self._failed_episode_dataset_file_handler): + if handler is not None: + handler.flush() def close(self): """Closes the recorder manager by exporting any remaining data to file as well as properly @@ -570,6 +553,12 @@ def close(self): Helper functions. """ + def _reset_buffers(self, env_ids): + """Discard selected samples by resetting indices; retain allocations for the next episodes.""" + for field in self._fields: + field.reset(env_ids) + self._success[slice(None) if env_ids is None else env_ids] = -1 + def _prepare_terms(self): """Prepares a list of recorder terms.""" # check if config is dict already @@ -587,6 +576,7 @@ def _prepare_terms(self): "export_in_record_pre_reset", "export_in_close", "dataset_compression", + "buffer_initial_capacity", ]: continue # check if term config is None @@ -605,3 +595,65 @@ def _prepare_terms(self): # add term name and parameters self._term_names.append(term_name) self._terms[term_name] = term + + +class _RecorderField: + """One field's time-major allocation and independent per-world insertion indices. + + No per-world Python containers or tensor views are created during append. The host upper + bound avoids reading device indices in steady state; only capacity exhaustion checks synchronize. + """ + + def __init__(self, name: str, value: torch.Tensor, num_envs: int, capacity: int): + self.name = name + # Retained storage outlives the caller's inference context: reset/export may run outside it. + with torch.inference_mode(False): + self.data = torch.empty((capacity, num_envs, *value.shape[1:]), dtype=value.dtype, device=value.device) + self.lengths = torch.zeros(num_envs, dtype=torch.long, device=value.device) + self.worlds = torch.arange(num_envs, device=value.device) + self.upper_bound = 0 + + @torch.no_grad() + def append(self, value: torch.Tensor, env_ids): + if value.shape[1:] != self.data.shape[2:] or value.dtype != self.data.dtype or value.device != self.data.device: + if self.lengths.any().item(): + raise ValueError(f"Recorder field '{self.name}' changed shape, dtype, or device with active samples.") + with torch.inference_mode(False): + self.data = torch.empty( + (*self.data.shape[:2], *value.shape[1:]), dtype=value.dtype, device=value.device + ) + self.lengths = self.lengths.to(value.device) + self.worlds = self.worlds.to(value.device) + self.upper_bound = 0 + if env_ids is None: + worlds, indices = self.worlds, self.lengths + else: + if not isinstance(env_ids, slice): + env_ids = torch.as_tensor(env_ids, dtype=torch.long, device=self.data.device) + worlds = self.worlds[env_ids] + indices = self.lengths[worlds] + capacity = self.data.shape[0] + if self.upper_bound >= capacity: + self.upper_bound = int(self.lengths.max().item()) + if self.upper_bound >= capacity: + with torch.inference_mode(False): + data = self.data.new_empty((2 * capacity, *self.data.shape[1:])) + data[:capacity].copy_(self.data) + self.data = data + self.data[indices, worlds] = value + if env_ids is None: + self.lengths.add_(1) + else: + self.lengths[worlds] = indices + 1 + self.upper_bound += 1 + + def reset(self, env_ids): + if env_ids is None or (isinstance(env_ids, slice) and env_ids == slice(None)): + self.lengths.zero_() + self.upper_bound = 0 + else: + if not isinstance(env_ids, slice): + env_ids = torch.as_tensor(env_ids, dtype=torch.long, device=self.data.device) + self.lengths[env_ids] = 0 + if not isinstance(env_ids, slice) and len(env_ids) == len(self.lengths): + self.upper_bound = 0 diff --git a/source/isaaclab/isaaclab/utils/math.py b/source/isaaclab/isaaclab/utils/math.py index c3366a3d8f81..032c815d7807 100644 --- a/source/isaaclab/isaaclab/utils/math.py +++ b/source/isaaclab/isaaclab/utils/math.py @@ -495,7 +495,7 @@ def axis_angle_from_quat(quat: torch.Tensor, eps: float = 1.0e-6) -> torch.Tenso angle = 2.0 * half_angle # check whether to apply Taylor approximation sin_half_angles_over_angles = torch.where( - angle.abs() > eps, torch.sin(half_angle) / angle, 0.5 - angle * angle / 48 + angle.abs() > eps, torch.sin(half_angle) / angle.clamp_min(eps), 0.5 - angle * angle / 48 ) return quat[..., :3] / sin_half_angles_over_angles.unsqueeze(-1) diff --git a/source/isaaclab/test/envs/test_manager_based_rl_env_unit.py b/source/isaaclab/test/envs/test_manager_based_rl_env_unit.py index e8042d488292..df28cf38669c 100644 --- a/source/isaaclab/test/envs/test_manager_based_rl_env_unit.py +++ b/source/isaaclab/test/envs/test_manager_based_rl_env_unit.py @@ -12,8 +12,20 @@ import gymnasium as gym import numpy as np import pytest +import torch from isaaclab.envs import ManagerBasedRLEnv +from isaaclab.managers import ( + DatasetExportMode, + ObservationGroupCfg, + ObservationManager, + ObservationTermCfg, + RecorderManager, + RecorderManagerBaseCfg, + RecorderTerm, + RecorderTermCfg, +) +from isaaclab.utils.noise import UniformNoiseCfg pytestmark = pytest.mark.unit @@ -102,3 +114,86 @@ def test_obs_space_follows_clip_constraint(): assert term_space.shape == expected_shapes[term_name] assert np.all(term_space.low == low) assert np.all(term_space.high == high) + + +def _state_observation(env): + return env.scene.state + + +class _PhysicalRecorder(RecorderTerm): + requires_post_step_observations = False + + def record_post_step(self): + self.latest = self._env.scene.state.clone() + return None, None + + +class _ObservationRecorder(RecorderTerm): + def record_post_step(self): + self.latest = self._env.obs_buf["policy"].clone() + return None, None + + +def _make_stepping_env(recorder_type=None): + """Exercise the real step/observation/recorder owners with a minimal deterministic physics leaf.""" + env = object.__new__(ManagerBasedRLEnv) + env._is_closed = True + env.cfg = SimpleNamespace(decimation=1, sim=SimpleNamespace(dt=0.02, render_interval=1), compute_final_obs=False) + env.scene = SimpleNamespace(num_envs=3, state=torch.zeros(3, 2)) + env.scene.write_data_to_sim = lambda: None + env.scene.update = lambda dt: None + env.sim = SimpleNamespace( + device="cpu", + is_playing=lambda: True, + is_rendering=False, + step=lambda render: env.scene.state.add_(1), + consume_reset_request=lambda: False, + ) + env._physics_handles_decimation = False + env._sim_step_counter = env.common_step_counter = 0 + env.episode_length_buf = torch.zeros(3, dtype=torch.long) + env.render_enabled = False + env.video_recorders = [] + env.extras = {} + env.action_manager = SimpleNamespace(process_action=lambda action: None, apply_action=lambda: None) + mask = torch.zeros(3, dtype=torch.bool) + env.termination_manager = SimpleNamespace(compute=lambda: mask, terminated=mask, time_outs=mask) + env.reward_manager = SimpleNamespace(compute=lambda dt: torch.zeros(3)) + env.command_manager = SimpleNamespace(compute=lambda dt: None) + env.event_manager = SimpleNamespace(active_terms=[]) + group = ObservationGroupCfg(history_length=3, flatten_history_dim=False, enable_corruption=True) + group.state = ObservationTermCfg(func=_state_observation, noise=UniformNoiseCfg(n_min=-0.1, n_max=0.1)) + env.observation_manager = ObservationManager({"policy": group}, env) + cfg = RecorderManagerBaseCfg(dataset_export_mode=DatasetExportMode.EXPORT_NONE) + if recorder_type is not None: + cfg.capture = RecorderTermCfg(class_type=recorder_type) + env.recorder_manager = RecorderManager(cfg, env) + env.obs_buf = env.observation_manager.compute(update_history=True) + return env + + +def test_raw_capture_preserves_observation_noise_history_and_rng(): + """Raw capture dispatch must not consume extra noise or advance the real policy history.""" + outputs, states = [], [] + for recorder_type in (None, _PhysicalRecorder, _ObservationRecorder): + torch.manual_seed(17) + env = _make_stepping_env(recorder_type) + history = [env.step(torch.empty(3, 0))[0]["policy"].clone() for _ in range(4)] + outputs.append(torch.stack(history)) + states.append(torch.get_rng_state()) + if recorder_type is not None: + term = env.recorder_manager.get_term("capture") + assert not env.recorder_manager._fields + if recorder_type is _PhysicalRecorder: + assert torch.equal(term.latest, torch.full((3, 2), 4.0)) + else: + # The default-compatible path really supplies observations to its callback. + assert term.latest.shape == (3, 3, 2) + env.recorder_manager.close() + assert torch.equal(outputs[0], outputs[1]) + assert torch.equal(states[0], states[1]) + assert not torch.equal(outputs[0], outputs[2]) + assert not torch.equal(states[0], states[2]) + # After four steps the three history entries contain steps 2, 3 and 4, each noised once. + expected = torch.tensor([2.0, 3.0, 4.0])[None, :, None] + assert torch.all((outputs[1][-1] - expected).abs() <= 0.1) diff --git a/source/isaaclab/test/managers/test_recorder_manager.py b/source/isaaclab/test/managers/test_recorder_manager.py index c07dfba99d10..43e05f8aa320 100644 --- a/source/isaaclab/test/managers/test_recorder_manager.py +++ b/source/isaaclab/test/managers/test_recorder_manager.py @@ -4,6 +4,8 @@ # SPDX-License-Identifier: BSD-3-Clause from __future__ import annotations +import ast +import inspect import os import shutil import tempfile @@ -11,11 +13,14 @@ from collections.abc import Sequence from typing import TYPE_CHECKING, cast +import h5py import pytest import torch +import warp as wp from isaaclab.managers import DatasetExportMode, RecorderManager, RecorderManagerBaseCfg, RecorderTerm, RecorderTermCfg from isaaclab.utils import configclass +from isaaclab.utils.datasets import EpisodeData pytestmark = pytest.mark.unit @@ -164,8 +169,8 @@ def test_record(dataset_dir): # check the recorded data for env_id in range(env.num_envs): episode = recorder_manager.get_episode(env_id) - assert torch.stack(episode.data["record_pre_step"]).shape == (2, 4) - assert torch.stack(episode.data["record_post_step"]).shape == (2, 5) + assert episode.data["record_pre_step"].shape == (2, 4) + assert episode.data["record_post_step"].shape == (2, 5) # Trigger pre-reset callbacks which then export and clean the episode data recorder_manager.record_pre_reset(env_ids=slice(None)) @@ -176,8 +181,138 @@ def test_record(dataset_dir): recorder_manager.record_post_reset(env_ids=slice(None)) for env_id in range(env.num_envs): episode = recorder_manager.get_episode(env_id) - assert torch.stack(episode.data["record_post_reset"]).shape == (1, 3) + assert episode.data["record_post_reset"].shape == (1, 3) recorder_manager.reset(slice(1, None, 2)) for env_id in range(env.num_envs): assert recorder_manager.get_episode(env_id).is_empty() == (env_id % 2 == 1) recorder_manager.close() + + +def test_batched_episode_indices_and_ownership(monkeypatch): + """Asynchronous resets, field cadence, growth and snapshots must preserve exact samples.""" + cfg = DummyRecorderManagerCfg(dataset_export_mode=DatasetExportMode.EXPORT_NONE, buffer_initial_capacity=2) + manager = RecorderManager(cfg, create_dummy_env(num_envs=3)) + manager.add_to_episodes("unused", torch.empty(0, 7), []) + assert not manager._fields + with pytest.raises(ValueError, match="requires 3 rows"): + manager.add_to_episodes("invalid", torch.empty(2, 7)) + assert not manager._fields + source = torch.tensor([[1.0, -0.0], [2.0, float("nan")], [3.0, 4.0]]) + # Policies commonly collect under inference_mode, then reset/export after leaving it. + with torch.inference_mode(): + manager.add_to_episodes("obs", {"state": source}) + wp.init() + manager.add_to_episodes("initial_state", wp.from_torch(torch.tensor([[10], [20], [30]], dtype=torch.int32))) + source.fill_(99) + manager.add_to_episodes("obs/state", source) + snapshot = manager.get_episode(1) + assert snapshot.data["obs"]["state"][0, 0] == 2 + assert torch.isnan(snapshot.data["obs"]["state"][0, 1]) + assert torch.signbit(manager.get_episode(0).data["obs"]["state"][0, 1]) + assert snapshot.data["initial_state"].dtype == torch.int32 + assert snapshot.success is None + + manager.reset(torch.tensor([1])) + with torch.inference_mode(): + manager.add_to_episodes("obs/state", torch.tensor([[42.0, 43.0], [52.0, 53.0]]), [2, 1]) + manager.add_to_episodes("obs/state", torch.tensor([[62.0, 63.0]]), slice(0, 1)) + assert torch.equal(manager.get_episode(1).data["obs"]["state"], torch.tensor([[52.0, 53.0]])) + assert manager.get_episode(1).data.keys() == {"obs"} + assert torch.equal(manager.get_episode(2).data["obs"]["state"][:, 0], torch.tensor([3.0, 99.0, 42.0])) + assert torch.equal(manager.get_episode(0).data["obs"]["state"][:, 0], torch.tensor([1.0, 99.0, 62.0])) + assert snapshot.data["obs"]["state"].shape == (2, 2) + snapshot = manager.get_episode(1) + snapshot.data["obs"]["state"].fill_(-100) + assert manager.get_episode(1).data["obs"]["state"][0, 0] == 52 + + with pytest.raises(ValueError, match="changed shape"): + manager.add_to_episodes("obs/state", torch.zeros(3, 5)) + with pytest.raises(ValueError, match="requires 3 rows"): + manager.add_to_episodes("obs/state", torch.zeros(2, 2)) + with pytest.raises(ValueError, match="conflicts"): + manager.add_to_episodes("obs", torch.zeros(3, 2)) + manager.add_to_episodes("obs/state", torch.empty(0, 2), []) + assert manager.get_episode(0).data["obs"]["state"].shape == (3, 2) + + manager.reset() + # A new schema is allowed once every world has released this field, e.g. manual annotation. + with torch.inference_mode(): + manager.add_to_episodes("obs/state", torch.zeros(3, 5, dtype=torch.bool)) + assert manager.get_episode(0).data["obs"]["state"].dtype == torch.bool + manager.reset() + + def forbidden(*args, **kwargs): + raise AssertionError("Steady-state recording must reuse storage without episode objects or host index reads") + + # Architectural/performance gate: preallocated insertion and reset need neither episode objects, + # per-world container appends, nor device-to-host scalar reads. Export-none also stays batched. + with monkeypatch.context() as patch: + patch.setattr(EpisodeData, "__init__", forbidden) + patch.setattr(EpisodeData, "add", forbidden) + patch.setattr(torch.Tensor, "item", forbidden) + patch.setattr(torch, "empty", forbidden) + manager.add_to_episodes("obs/state", torch.zeros(3, 5, dtype=torch.bool)) + manager.set_success_to_episodes([0, 2], torch.tensor([[True], [False]])) + manager.export_episodes([0, 2]) + manager.reset([1]) + assert manager.exported_successful_episode_count == 1 + assert manager.exported_failed_episode_count == 1 + assert all(manager.get_episode(i).is_empty() for i in range(3)) + manager.close() + + +@pytest.mark.parametrize("mode", list(DatasetExportMode)) +def test_export_batched_episodes(dataset_dir, mode): + """Export modes, custom IDs, nested tensor schema and counters survive partial finalization.""" + cfg = DummyRecorderManagerCfg(dataset_export_dir_path=dataset_dir, dataset_export_mode=mode, export_in_close=True) + manager = RecorderManager(cfg, create_dummy_env(num_envs=3)) + values = torch.tensor([[1.0, -0.0], [2.0, float("nan")], [3.0, 4.0]]) + manager.add_to_episodes("actions", values) + manager.add_to_episodes("obs", {"flag": torch.tensor([True, False, True])}) + manager.set_success_to_episodes(None, torch.tensor([[True], [False], [True]])) + with pytest.raises(ValueError, match="must match"): + manager.export_episodes([2, 1], [30]) + with pytest.raises(ValueError, match="must be unique"): + manager.export_episodes([2, 1], [30, 30]) + manager.export_episodes(torch.tensor([2, 1]), [30, 10]) + assert manager.get_episode(1).is_empty() + assert manager.get_episode(2).is_empty() + assert manager.get_episode(0).success is True + manager.export_episodes(slice(0, 1), [20]) + manager.export_episodes() # Empty episodes do not affect counters. + assert manager.exported_successful_episode_count == 2 + assert manager.exported_failed_episode_count == 1 + manager.close() + if mode == DatasetExportMode.EXPORT_NONE: + assert os.listdir(dataset_dir) == [] + return + with h5py.File(os.path.join(dataset_dir, "dataset.hdf5")) as file: + expected = {"demo_30", "demo_20"} + if mode == DatasetExportMode.EXPORT_ALL: + expected.add("demo_10") + assert set(file["data"]) == expected + assert file["data/demo_20/actions"][()].tobytes() == values[0:1].numpy().tobytes() + assert bool(file["data/demo_30/obs/flag"][0]) + assert bool(file["data/demo_20"].attrs["success"]) + if mode == DatasetExportMode.EXPORT_SUCCEEDED_FAILED_IN_SEPARATE_FILES: + with h5py.File(os.path.join(dataset_dir, "dataset_failed.hdf5")) as file: + assert set(file["data"]) == {"demo_10"} + assert file["data/demo_10/actions"][()].tobytes() == values[1:2].numpy().tobytes() + assert not bool(file["data/demo_10"].attrs["success"]) + + +def test_recorder_storage_boundary(): + """Do not reintroduce per-world episode ownership or Python world loops in the append path.""" + import isaaclab.managers.recorder_manager as module + + source = ast.parse(inspect.getsource(module)) + assert not any(isinstance(node, ast.Attribute) and node.attr == "_episodes" for node in ast.walk(source)) + forbidden = ("physicsnemo", "supermimic", "isaaclab_physicsnemo") + for node in ast.walk(source): + if isinstance(node, ast.Import): + assert not any(alias.name.startswith(forbidden) for alias in node.names) + elif isinstance(node, ast.ImportFrom): + assert not (node.module or "").startswith(forbidden) + field = next(node for node in source.body if isinstance(node, ast.ClassDef) and node.name == "_RecorderField") + append = next(node for node in field.body if isinstance(node, ast.FunctionDef) and node.name == "append") + assert not any(isinstance(node, (ast.For, ast.While, ast.comprehension)) for node in ast.walk(append)) diff --git a/source/isaaclab/test/utils/test_math.py b/source/isaaclab/test/utils/test_math.py index d9d217931b56..1085b8f6b4b0 100644 --- a/source/isaaclab/test/utils/test_math.py +++ b/source/isaaclab/test/utils/test_math.py @@ -174,6 +174,12 @@ def test_axis_angle_from_quat_approximation(device): # Assert that the computed values are close to the expected values torch.testing.assert_close(axis_angle_computed, axis_angle_expected) + # The unselected exact formula must not introduce NaN gradients through 0/0 at identity. + identities = torch.tensor([[0.0, 0.0, 0.0, 1.0], [0.0, 0.0, 0.0, -1.0]], device=device, requires_grad=True) + math_utils.axis_angle_from_quat(identities).sum().backward() + expected_gradient = torch.tensor([[2.0, 2.0, 2.0, 0.0], [-2.0, -2.0, -2.0, 0.0]], device=device) + torch.testing.assert_close(identities.grad, expected_gradient) + @pytest.mark.parametrize("device", test_devices()) def test_quat_error_magnitude(device): diff --git a/source/isaaclab_experimental/changelog.d/recorder-capture.changed.rst b/source/isaaclab_experimental/changelog.d/recorder-capture.changed.rst new file mode 100644 index 000000000000..99e5426a4d83 --- /dev/null +++ b/source/isaaclab_experimental/changelog.d/recorder-capture.changed.rst @@ -0,0 +1,2 @@ +* Honored recorder terms' post-step observation requirements in the Warp RL environment, + allowing physical-state capture without extra observation computation. diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py index be522e073a2b..b35fb410a88c 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py @@ -321,12 +321,12 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: ) if len(self.recorder_manager.active_terms) > 0: - # update observations for recording if needed - self._manager_call_switch.call_stage( - stage="ObservationManager_compute_no_history", - warp_call={"fn": self.observation_manager.compute, "kwargs": {"return_cloned_output": False}}, - timer=DEBUG_TIMER_STEP, - ) + if self.recorder_manager.requires_post_step_observations: + self._manager_call_switch.call_stage( + stage="ObservationManager_compute_no_history", + warp_call={"fn": self.observation_manager.compute, "kwargs": {"return_cloned_output": False}}, + timer=DEBUG_TIMER_STEP, + ) self.recorder_manager.record_post_step() # -- reset envs that terminated/timed-out and log the episode information diff --git a/source/isaaclab_rl/changelog.d/supermimic.added.rst b/source/isaaclab_rl/changelog.d/supermimic.added.rst new file mode 100644 index 000000000000..1928560251bb --- /dev/null +++ b/source/isaaclab_rl/changelog.d/supermimic.added.rst @@ -0,0 +1,2 @@ +* Added native SuperMimic training and playback through the unified reinforcement-learning entrypoints, + with MLP and optional one-frame NeRD dynamics, learner-owned replay, resumable checkpoints and W&B uploads. diff --git a/source/isaaclab_rl/changelog.d/supermimic.minor b/source/isaaclab_rl/changelog.d/supermimic.minor new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_rl/isaaclab_rl/entrypoints/api.py b/source/isaaclab_rl/isaaclab_rl/entrypoints/api.py index 8db3c29ed195..2652462473e9 100644 --- a/source/isaaclab_rl/isaaclab_rl/entrypoints/api.py +++ b/source/isaaclab_rl/isaaclab_rl/entrypoints/api.py @@ -12,7 +12,7 @@ from .dispatch import run_play_cli, run_random_agent_cli, run_train_cli, run_zero_agent_cli -BackendName = Literal["rl_games", "rlinf", "rsl_rl", "sb3", "skrl", "torchrl"] +BackendName = Literal["rl_games", "rlinf", "rsl_rl", "sb3", "skrl", "supermimic", "torchrl"] """Backends of the unified entrypoints.""" MULTI_GPU_BACKENDS: tuple[BackendName, ...] = ("rl_games", "rsl_rl", "skrl") diff --git a/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/play_supermimic.py b/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/play_supermimic.py new file mode 100644 index 000000000000..3b955bcaf470 --- /dev/null +++ b/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/play_supermimic.py @@ -0,0 +1,150 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Portions derived from SuperMimic (MIT), Copyright (c) 2026 ETH Zurich, Chong Zhang. +# See docs/licenses/dependencies/supermimic-license.txt. + +"""Native SuperMimic policy playback and complete motion-set evaluation.""" + +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path + +import torch + +from isaaclab.app import add_launcher_args, launch_simulation +from isaaclab.utils import update_from_dict + +import isaaclab_tasks # noqa: F401 +from isaaclab_tasks.contrib.g1_tracking.g1_tracking_env_cfg import TransitionRecorderCfg +from isaaclab_tasks.contrib.g1_tracking.tracking import MAX_REWARD_RATE, OBS_DIM +from isaaclab_tasks.utils import resolve_task_config, setup_preset_cli + +from ...supermimic.networks import Policy +from ..common import ( + add_common_play_args, + apply_env_overrides, + apply_video_recording, + create_isaaclab_env, + pre_launch_video_config, + scoped_torch_backend_flags, +) + + +@torch.no_grad() +def evaluate(env, policy: Policy, real_time: bool = False) -> dict: + """Score each clip once, retaining failed tails in the full-reference-length denominator.""" + native_env = env.unwrapped + command = native_env.command_manager.get_term("motion") + capture = native_env.recorder_manager.get_term("transition") + motion = command.motion + clip_count = len(motion.names) + returns = torch.zeros(clip_count, device=native_env.device) + steps = torch.zeros(clip_count, dtype=torch.long, device=native_env.device) + failed = torch.zeros(clip_count, dtype=torch.bool, device=native_env.device) + for offset in range(0, clip_count, native_env.num_envs): + count = min(native_env.num_envs, clip_count - offset) + command.cfg.evaluation_clip_offset = offset + observations, _ = env.reset() + finished = torch.arange(native_env.num_envs, device=native_env.device) >= count + while not bool(finished.all()): + started = time.perf_counter() + action = policy(observations["policy"], observations["goal"]) + observations, _, _, _, _ = env.step(action) + transition = capture.latest + alive = ~finished + returns[offset : offset + count] += torch.where(alive[:count], transition["r"][:count], 0.0) + steps[offset : offset + count] += alive[:count] + failed[offset : offset + count] |= alive[:count] & transition["term"][:count] + finished |= transition["done"] + if real_time: + time.sleep(max(0.0, native_env.step_dt - (time.perf_counter() - started))) + maximum_step = MAX_REWARD_RATE * native_env.step_dt * native_env.reward_manager.get_term_cfg("tracking").weight + denominators = (motion.clip_len - 1).float() * maximum_step + normalized = returns / denominators + return { + "norm_return": normalized.mean().item(), + "norm_return_framew": (returns.sum() / denominators.sum()).item(), + "n_failed": int(failed.sum()), + "failed": [motion.names[i] for i in failed.nonzero().flatten().tolist()], + "per_clip": dict(zip(motion.names, normalized.tolist(), strict=True)), + "per_clip_steps": dict(zip(motion.names, steps.tolist(), strict=True)), + } + + +def run(argv: list[str]) -> None: + """Restore a native checkpoint's task settings and evaluate its policy with physical simulation.""" + parser = argparse.ArgumentParser(description=__doc__) + add_common_play_args( + parser, agent_default="supermimic_cfg_entry_point", agent_help="SuperMimic agent configuration." + ) + parser.add_argument("--checkpoint", required=True, help="Full or policy-only native SuperMimic checkpoint.") + parser.add_argument("--output", help="Evaluation JSON; defaults to _eval.json.") + parser.add_argument("--motion_dir", help="Evaluate a different directory of reference motions.") + parser.add_argument("--asset_path", help="Relocated USD asset path.") + parser.add_argument("--kinematics_path", help="Relocated matching MJCF asset path.") + add_launcher_args(parser) + args, overrides = setup_preset_cli(parser, argv) + if overrides or args.train_env_cfg or args.frontend != "torch": + parser.error( + "Evaluation restores the saved task and uses its play mode; use explicit relocation/device/seed options." + ) + checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=True) + task = args.task or checkpoint["task"] + if task != checkpoint["task"]: + parser.error("The selected task must match the checkpoint.") + env_cfg, _ = resolve_task_config(task, None, overrides=[]) + env_cfg.recorders.transition = TransitionRecorderCfg() + update_from_dict(env_cfg, checkpoint["env_cfg"]) + apply_env_overrides(args, env_cfg) + if args.motion_dir: + env_cfg.commands.motion.motion_dir = args.motion_dir + if args.asset_path: + env_cfg.scene.robot.spawn.usd_path = args.asset_path + if args.kinematics_path: + env_cfg.commands.motion.kinematics_path = args.kinematics_path + if args.seed is not None: + env_cfg.seed = args.seed + count = sum(path.suffix in (".csv", ".npy") for path in Path(env_cfg.commands.motion.motion_dir).rglob("*")) + if count < 1 or env_cfg.scene.num_envs < 1: + parser.error("Evaluation needs reference motions and a positive world count.") + env_cfg.scene.num_envs = min(env_cfg.scene.num_envs, count) + env_cfg.play_mode() + output = Path(args.output or Path(args.checkpoint).with_suffix("").as_posix() + "_eval.json") + output.parent.mkdir(parents=True, exist_ok=True) + pre_launch_video_config(env_cfg, args) + apply_video_recording(env_cfg, str(output.parent), args, subdir="play", checkpoint_path=args.checkpoint) + torch.manual_seed(env_cfg.seed) + with ( + scoped_torch_backend_flags( + cuda_matmul_allow_tf32=False, cudnn_allow_tf32=True, cudnn_deterministic=False, cudnn_benchmark=False + ), + launch_simulation(env_cfg, args), + ): + env = create_isaaclab_env(task, env_cfg, args, convert_marl_to_single_agent=False) + try: + dtype = torch.bfloat16 if checkpoint["agent_cfg"]["amp"] == "bf16" else None + policy = Policy( + OBS_DIM, + 73 * env_cfg.commands.motion.goal_frames, + hist=env_cfg.observations.policy.history_length, + amp_dtype=dtype, + ).to(env.unwrapped.device) + policy.load_state_dict(checkpoint["pi"]) + policy.eval() + result = evaluate(env, policy, args.real_time) + result.update( + checkpoint=str(Path(args.checkpoint).resolve()), iteration=checkpoint["iteration"], seed=env_cfg.seed + ) + output.write_text(json.dumps(result, indent=2) + "\n") + print( + f"Completed {len(result['per_clip'])} clips: {result['n_failed']} failures, " + f"normalized return {result['norm_return']:.6f}; {output}" + ) + finally: + env.close() diff --git a/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/train_supermimic.py b/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/train_supermimic.py new file mode 100644 index 000000000000..a4472e78008d --- /dev/null +++ b/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/train_supermimic.py @@ -0,0 +1,451 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""SuperMimic training backend of the unified reinforcement learning entrypoint.""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import signal +import time +from contextlib import nullcontext +from datetime import datetime +from pathlib import Path + +import torch + +from isaaclab.app import add_launcher_args, launch_simulation +from isaaclab.utils import string_to_callable, to_dict +from isaaclab.utils.noise import NoiseCfg + +import isaaclab_tasks # noqa: F401 +from isaaclab_tasks.contrib.g1_tracking.g1_tracking_env_cfg import ( + G1TrackingEnvCfg, + ObservationsCfg, + TerminationsCfg, + TransitionRecorderCfg, +) +from isaaclab_tasks.contrib.g1_tracking.tracking import MAX_REWARD_RATE, OBS_DIM +from isaaclab_tasks.utils import resolve_task_config, setup_preset_cli + +from ...supermimic.learner import Learner +from ...supermimic.networks import WM_IN, WM_OUT, WorldModel +from ..common import ( + add_common_train_args, + apply_env_overrides, + apply_video_recording, + create_isaaclab_env, + dump_train_configs, + enable_cameras_for_video, + pre_launch_video_config, + resolve_seed, + scoped_torch_backend_flags, + wrap_sensor_capture, + write_run_manifest, +) + + +def run(argv: list[str]) -> None: + """Fit dynamics from real transitions and improve a policy through imagined tracking rollouts.""" + parser = argparse.ArgumentParser(description=__doc__) + add_common_train_args( + parser, + agent_default="supermimic_cfg_entry_point", + agent_help="SuperMimic agent configuration.", + include_distributed=False, + ) + parser.add_argument("--checkpoint", help="Full native checkpoint to resume; physical state and replay restart.") + parser.add_argument("--log_dir", help="Output directory; defaults to logs/supermimic//.") + add_launcher_args(parser) + args, overrides = setup_preset_cli(parser, argv) + if args.frontend != "torch": + parser.error("SuperMimic currently supports the native Torch environment frontend.") + env_cfg, agent_cfg = resolve_task_config(args.task, args.agent, overrides=overrides) + _configure_training(args, env_cfg, agent_cfg) + + output = Path( + args.log_dir + or Path("logs/supermimic") / agent_cfg.experiment_name / datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + ) + output.mkdir(parents=True, exist_ok=True) + resumed = torch.load(args.checkpoint, map_location="cpu", weights_only=True) if args.checkpoint else None + # Capture semantic configuration before adding run-specific log paths and video outputs. + saved_env_cfg, saved_agent_cfg = to_dict(env_cfg), to_dict(agent_cfg) + if resumed: + if resumed["task"] != args.task or resumed["env_cfg"] != saved_env_cfg: + raise ValueError("Resume requires the checkpoint's task and environment configuration.") + mutable = { + "max_iterations", + "max_hours", + "save_interval", + "save_interval_seconds", + "log_interval", + "wandb_project", + "wandb_entity", + "wandb_run_id", + "wandb_name", + } + changed = { + key for key, value in saved_agent_cfg.items() if key not in mutable and resumed["agent_cfg"][key] != value + } + if changed or resumed["iteration"] >= agent_cfg.max_iterations: + raise ValueError(f"Resume requires unchanged learning settings and a later iteration limit: {changed}.") + identity = resumed.get("wandb") if resumed else None + if identity and agent_cfg.wandb_run_id is not None and agent_cfg.wandb_run_id != identity["run_id"]: + identity = None + if identity: + for key in ("project", "entity", "run_id"): + value = getattr(agent_cfg, f"wandb_{key}") + if value is not None and value != identity[key]: + raise ValueError(f"Resuming a W&B run must preserve its {key}.") + setattr(agent_cfg, f"wandb_{key}", identity[key]) + logging = nullcontext(None) + if agent_cfg.wandb_project: + import wandb + + logging = wandb.init( + project=agent_cfg.wandb_project, + entity=agent_cfg.wandb_entity, + id=agent_cfg.wandb_run_id, + name=agent_cfg.wandb_name, + resume="allow", + dir=str(output), + config={"task": args.task, "agent": saved_agent_cfg, "env": saved_env_cfg}, + allow_val_change=True, + ) + env_cfg.log_dir = str(output) + enable_cameras_for_video(args) + pre_launch_video_config(env_cfg, args) + apply_video_recording(env_cfg, str(output), args) + write_run_manifest(str(output), library="supermimic", task=args.task, metadata={"agent": args.agent}) + dump_train_configs(str(output), env_cfg, agent_cfg) + torch.manual_seed(agent_cfg.seed) + with ( + logging as run, + scoped_torch_backend_flags( + cuda_matmul_allow_tf32=False, cudnn_allow_tf32=True, cudnn_deterministic=False, cudnn_benchmark=False + ), + launch_simulation(env_cfg, args), + ): + if run: + run.define_metric("iteration") + run.define_metric("*", step_metric="iteration") + identity = {"entity": run.entity, "project": run.project, "run_id": run.id, "url": run.url} + (output / "wandb-run.json").write_text(json.dumps(identity, indent=2) + "\n") + env = create_isaaclab_env(args.task, env_cfg, args, convert_marl_to_single_agent=False) + env = wrap_sensor_capture(env, str(output), args) + metadata = { + "task": args.task, + "agent_cfg": saved_agent_cfg, + "env_cfg": saved_env_cfg, + "wandb": identity, + "state_layout": "root-link-xyzw-v-w-q-qd", + } + try: + _train(env, agent_cfg, output, resumed, run, metadata) + finally: + env.close() + + +def _configure_training(args, env_cfg, agent_cfg): + """Apply native overrides and validate that real and imagined task semantics agree.""" + if not isinstance(env_cfg, G1TrackingEnvCfg): + raise ValueError("SuperMimic currently requires the G1 motion-tracking task's analytical objective.") + apply_env_overrides(args, env_cfg) + if args.seed is not None: + agent_cfg.seed = resolve_seed(args.seed) + if args.max_iterations is not None: + agent_cfg.max_iterations = args.max_iterations + env_cfg.seed = agent_cfg.seed + if env_cfg.recorders.transition is None: + env_cfg.recorders.transition = TransitionRecorderCfg() + if min(agent_cfg.buffer_steps, agent_cfg.warmup_steps) < agent_cfg.world_model_window: + raise ValueError("Replay capacity and warmup must each cover a full world-model window.") + if ( + min( + agent_cfg.max_iterations, + agent_cfg.horizon, + agent_cfg.world_model_window, + agent_cfg.world_model_batches, + agent_cfg.world_model_batch_size, + agent_cfg.policy_batch_size, + agent_cfg.critic_epochs, + agent_cfg.critic_minibatches, + agent_cfg.learning_rate_decay_iterations, + agent_cfg.save_interval, + agent_cfg.log_interval, + env_cfg.observations.policy.history_length, + env_cfg.commands.motion.goal_frames, + ) + < 1 + ): + raise ValueError("Learning counts, observation history and goal lengths must be positive.") + if agent_cfg.critic_minibatches > agent_cfg.horizon * agent_cfg.policy_batch_size: + raise ValueError("Critic minibatches must contain samples.") + if not 0 <= env_cfg.commands.motion.gamma < 1 or not 0 <= agent_cfg.gae_lambda <= 1: + raise ValueError("Discount must be in [0, 1) and GAE lambda in [0, 1].") + if not math.isfinite(agent_cfg.action_noise) or agent_cfg.action_noise < 0: + raise ValueError("Action noise must be finite and nonnegative.") + if agent_cfg.world_model_only_iterations < 0: + raise ValueError("World-model-only iterations must be nonnegative.") + for value in (agent_cfg.update_to_data_ratio, agent_cfg.learning_rate, agent_cfg.save_interval_seconds): + if not math.isfinite(value) or value <= 0: + raise ValueError("Learning rates, update/data ratio and checkpoint intervals must be finite and positive.") + if agent_cfg.max_hours is not None and (not math.isfinite(agent_cfg.max_hours) or agent_cfg.max_hours <= 0): + raise ValueError("The training-hour budget must be finite and positive.") + if agent_cfg.amp not in ("none", "bf16") or agent_cfg.world_model not in ("mlp", "nerd-one-frame"): + raise ValueError("Select amp=none/bf16 and world_model=mlp/nerd-one-frame.") + if env_cfg.sim.dt * env_cfg.decimation != 0.02: + raise ValueError("The reference motions and world model require a 20 ms control period.") + # The same analytical observations are used inside imagined rollouts. Reject changes that + # preserve tensor width but alter their meaning; native history and stateless noise remain configurable. + template = ObservationsCfg() + for name in ("policy", "goal"): + expected, actual = to_dict(getattr(template, name)), to_dict(getattr(env_cfg.observations, name)) + if name == "policy": + for group in (expected, actual): + group.pop("history_length") + group.pop("enable_corruption") + for term in group.values(): + if isinstance(term, dict): + term.pop("noise", None) + if actual != expected or tuple(actual) != tuple(expected): + raise ValueError(f"Analytical {name} observations require the task's declared functions and layout.") + if to_dict(env_cfg.terminations) != to_dict(TerminationsCfg()): + raise ValueError("Analytical termination requires the task's declared failure, clip-end and timeout terms.") + + +def _train(env, agent_cfg, output, resumed, run, metadata): + """Own the learning schedule, checkpoint publication and signal handling for one environment.""" + if run: + import wandb + + previous_handlers = {} + try: + native_env = env.unwrapped + command = native_env.command_manager.get_term("motion") + capture = native_env.recorder_manager.get_term("transition") + observations, _ = env.reset() + from isaaclab_tasks.contrib.g1_tracking.mdp import tracking_reward + + reward_cfg = native_env.reward_manager.get_term_cfg("tracking") + reward_func = string_to_callable(reward_cfg.func) if isinstance(reward_cfg.func, str) else reward_cfg.func + if ( + native_env.reward_manager.active_terms != ["tracking"] + or reward_func is not tracking_reward + or reward_cfg.weight <= 0 + ): + raise ValueError("The imagined objective requires the task's single tracking reward with positive weight.") + noise_terms, offset = [], 0 + group = native_env.cfg.observations.policy + for name, shape in zip( + native_env.observation_manager.active_terms["policy"], + native_env.observation_manager.group_obs_term_dim["policy"], + strict=True, + ): + term = getattr(group, name) + noise = term.noise if group.enable_corruption else None + if noise is not None and not isinstance(noise, NoiseCfg): + raise ValueError("Analytical observations require stateless policy noise.") + fn = ( + None if noise is None else string_to_callable(noise.func) if isinstance(noise.func, str) else noise.func + ) + noise_terms.append((slice(offset, offset + shape[-1]), fn, noise)) + offset += shape[-1] + if offset != OBS_DIM or int(command.motion.clip_len.max()) <= agent_cfg.horizon: + raise ValueError("Observation width or motion lengths are incompatible with the learner.") + network = None + if agent_cfg.world_model == "nerd-one-frame": + from physicsnemo_robotics.models.parallel_nerd import ParallelNeRD + + network = ParallelNeRD(WM_IN, WM_OUT, temporal_attention_spans=(1, 1, 1, 1)) + torch.nn.init.zeros_(network.head_output.weight) + torch.nn.init.zeros_(network.head_output.bias) + dtype = torch.bfloat16 if agent_cfg.amp == "bf16" else None + world_model = WorldModel(network=network, amp_dtype=dtype).to(native_env.device) + learner = Learner( + agent_cfg, + command.mdp, + capture.latest, + tuple(noise_terms), + native_env.step_dt * reward_cfg.weight, + world_model, + goal_frames=command.cfg.goal_frames, + history=group.history_length, + gamma=command.cfg.gamma, + ) + first_iteration, prior_wall_time, sim_steps = 0, 0.0, 0 + if resumed: + for name in ("wm", "pi", "critic", "opt_wm", "opt_pi", "opt_v"): + getattr(learner, name).load_state_dict(resumed[name]) + command.bin_return.copy_(resumed["bin_return"]) + command.bin_weight.copy_(resumed["bin_weight"]) + first_iteration, prior_wall_time, sim_steps = ( + resumed["iteration"], + resumed["wall_time"], + resumed["sim_steps"], + ) + del resumed + stop_signal = 0 + + def request_stop(signum, frame): + nonlocal stop_signal + stop_signal = signum + + previous_handlers = {sig: signal.signal(sig, request_stop) for sig in (signal.SIGTERM, signal.SIGINT)} + stats = {name: [] for name in ("reward", "length", "return", "failed", "done")} + + @torch.no_grad() + def collect(count): + nonlocal observations, sim_steps + for _ in range(count): + observations, _, _, _, extras = env.step(learner.act(observations)) + transition = capture.latest + learner.observe(transition) + done = transition["done"] + done_count = done.sum() + stats["reward"].append(transition["r"].mean() / (MAX_REWARD_RATE * learner.reward_scale)) + stats["length"].append(((transition["ep_step"] + 1) * done).sum()) + stats["return"].append(extras["log"].get("Metrics/motion/return", 0.0) * done_count) + stats["failed"].append(transition["term"].sum()) + stats["done"].append(done_count) + sim_steps += native_env.num_envs + + collect(agent_cfg.warmup_steps) + training_start = log_start = checkpoint_time = time.time() + step_accumulator = 0.0 + aggregate = {name: [] for name in ("wm", "wm_last", "pi", "v", "imagined_reward", "gradient", "survival")} + columns = ( + "iteration", + "wall_h", + "sim_steps", + "wm_loss", + "wm_last_error", + "pi_loss", + "v_loss", + "imagined_reward", + "actor_gradient", + "survival", + "sim_reward", + "episode_length", + "episode_return", + "failure_rate", + "iterations_per_second", + "learning_rate", + ) + log_path = output / "train_log.csv" + if first_iteration and log_path.exists(): + lines = log_path.read_text().splitlines() + lines = lines[:1] + [line for line in lines[1:] if int(float(line.split(",")[0])) <= first_iteration] + log_path.write_text("\n".join(lines) + "\n") + else: + log_path.write_text(",".join(columns) + "\n") + with log_path.open("a") as log: + for iteration in range(first_iteration, agent_cfg.max_iterations): + fraction = min(iteration / agent_cfg.learning_rate_decay_iterations, 1.0) + lr = agent_cfg.learning_rate * (0.03 + 0.97 * 0.5 * (1 + math.cos(math.pi * fraction))) + for optimizer in (learner.opt_wm, learner.opt_pi, learner.opt_v): + for group_cfg in optimizer.param_groups: + group_cfg["lr"] = lr + step_accumulator += 1.0 / agent_cfg.update_to_data_ratio + count = int(step_accumulator) + step_accumulator -= count + collect(count) + for _ in range(agent_cfg.world_model_batches): + loss, error = learner.wm_update() + aggregate["wm"].append(loss) + aggregate["wm_last"].append(error) + if iteration >= agent_cfg.world_model_only_iterations: + results = learner.policy_update() + for name, value in zip( + ("pi", "v", "imagined_reward", "gradient", "survival"), results, strict=True + ): + aggregate[name].append(value) + if (iteration + 1) % agent_cfg.log_interval == 0: + means = [float(torch.stack(values).mean()) if values else 0.0 for values in aggregate.values()] + totals = { + name: float(torch.stack(values).sum()) if values else 0.0 for name, values in stats.items() + } + done_count = max(totals["done"], 1.0) + wall_h = (prior_wall_time + time.time() - training_start) / 3600 + row = [ + iteration + 1, + wall_h, + sim_steps, + *means, + totals["reward"] / max(len(stats["reward"]), 1), + totals["length"] / done_count, + totals["return"] / done_count, + totals["failed"] / done_count, + agent_cfg.log_interval / (time.time() - log_start), + lr, + ] + if not all(math.isfinite(value) for value in row): + raise FloatingPointError(f"Nonfinite training statistics at update {iteration + 1}: {row}") + log.write(",".join(f"{value:.6g}" for value in row) + "\n") + log.flush() + if run: + run.log(dict(zip(columns, row, strict=True))) + print( + f"Iteration {iteration + 1}: dynamics={means[0]:.4f}, policy={means[2]:.4f}, " + f"value={means[3]:.4f}, hours={wall_h:.3f}", + flush=True, + ) + log_start = time.time() + for values in (*stats.values(), *aggregate.values()): + values.clear() + elapsed = prior_wall_time + time.time() - training_start + stop = agent_cfg.max_hours is not None and elapsed >= agent_cfg.max_hours * 3600 + milestone = ( + (iteration + 1) % agent_cfg.save_interval == 0 or iteration + 1 == agent_cfg.max_iterations or stop + ) + if ( + milestone + or stop_signal + or iteration == first_iteration + or time.time() - checkpoint_time >= agent_cfg.save_interval_seconds + ): + checkpoint = { + **metadata, + **{ + name: getattr(learner, name).state_dict() + for name in ("wm", "pi", "critic", "opt_wm", "opt_pi", "opt_v") + }, + "iteration": iteration + 1, + "wall_time": elapsed, + "sim_steps": sim_steps, + "bin_return": command.bin_return, + "bin_weight": command.bin_weight, + "saved_at": datetime.now().isoformat(), + } + torch.save(checkpoint, output / "last.pt.tmp") + os.replace(output / "last.pt.tmp", output / "last.pt") + if milestone: + policy_checkpoint = { + key: value + for key, value in checkpoint.items() + if key not in ("wm", "critic", "opt_wm", "opt_pi", "opt_v") + } + path = output / f"policy_{iteration + 1:06d}.pt" + torch.save(policy_checkpoint, path.with_suffix(".tmp")) + os.replace(path.with_suffix(".tmp"), path) + if run: + artifact = wandb.Artifact( + f"run-{run.id}-checkpoint", type="model", metadata={"iteration": iteration + 1} + ) + artifact.add_file(str(output / "last.pt"), name="last.pt") + run.log_artifact(artifact, aliases=["latest"]).wait(timeout=180) + run.summary["durable_checkpoint_iteration"] = iteration + 1 + checkpoint_time = time.time() + if stop_signal: + raise SystemExit(128 + stop_signal) + if stop: + break + finally: + for sig, handler in previous_handlers.items(): + signal.signal(sig, handler) diff --git a/source/isaaclab_rl/isaaclab_rl/entrypoints/dispatch.py b/source/isaaclab_rl/isaaclab_rl/entrypoints/dispatch.py index aff2b314193e..b50d2378cf69 100644 --- a/source/isaaclab_rl/isaaclab_rl/entrypoints/dispatch.py +++ b/source/isaaclab_rl/isaaclab_rl/entrypoints/dispatch.py @@ -33,6 +33,7 @@ "rsl_rl": "isaaclab_rl.entrypoints.backends.train_rsl_rl", "sb3": "isaaclab_rl.entrypoints.backends.train_sb3", "skrl": "isaaclab_rl.entrypoints.backends.train_skrl", + "supermimic": "isaaclab_rl.entrypoints.backends.train_supermimic", "torchrl": "isaaclab_rl.entrypoints.backends.train_torchrl", }, "play": { @@ -41,6 +42,7 @@ "rsl_rl": "isaaclab_rl.entrypoints.backends.play_rsl_rl", "sb3": "isaaclab_rl.entrypoints.backends.play_sb3", "skrl": "isaaclab_rl.entrypoints.backends.play_skrl", + "supermimic": "isaaclab_rl.entrypoints.backends.play_supermimic", "torchrl": "isaaclab_rl.entrypoints.backends.play_torchrl", }, } diff --git a/source/isaaclab_rl/isaaclab_rl/supermimic/__init__.py b/source/isaaclab_rl/isaaclab_rl/supermimic/__init__.py new file mode 100644 index 000000000000..d0e155aa7944 --- /dev/null +++ b/source/isaaclab_rl/isaaclab_rl/supermimic/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""SuperMimic dynamics fitting and differentiable policy learning.""" diff --git a/source/isaaclab_rl/isaaclab_rl/supermimic/learner.py b/source/isaaclab_rl/isaaclab_rl/supermimic/learner.py new file mode 100644 index 000000000000..4adee693220d --- /dev/null +++ b/source/isaaclab_rl/isaaclab_rl/supermimic/learner.py @@ -0,0 +1,290 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Portions derived from SuperMimic (MIT), Copyright (c) 2026 ETH Zurich, Chong Zhang. +# See docs/licenses/dependencies/supermimic-license.txt. + +"""Replay ownership, supervised dynamics fitting and differentiable actor/critic updates.""" + +from __future__ import annotations + +import torch +import torch.nn as nn + +from isaaclab_tasks.contrib.g1_tracking.tracking import MAX_REWARD_RATE, OBS_DIM, unflatten_state + +from .networks import CRITIC_EXTRA, INV_DIM, Critic, Policy, WorldModel, critic_features, normalize_xi, state_difference +from .supermimic_cfg import SuperMimicCfg + + +class ReplayBuffer: + """Learner-owned time/world ring; sampled model windows cannot cross physical discontinuities.""" + + def __init__(self, capacity: int, transition: dict[str, torch.Tensor]): + if capacity < 1: + raise ValueError("Replay capacity must be positive.") + self.capacity, self.size, self.pointer = (capacity, 0, 0) + self.num_envs, self.device = (transition["s"].shape[0], transition["s"].device) + fields = ("s", "s1", "a", "a_prev", "xi", "f", "ep_step", "done") + self.data = {key: transition[key].new_empty((capacity, *transition[key].shape)) for key in fields} + + @torch.no_grad() + def add(self, transition): + for key, data in self.data.items(): + data[self.pointer].copy_(transition[key]) + self.pointer = (self.pointer + 1) % self.capacity + self.size = min(self.size + 1, self.capacity) + + def _row(self, age): + return (self.pointer - self.size + age) % self.capacity + + def sample_windows(self, batch_size, length): + if self.size < length: + raise ValueError("Replay does not yet contain a full model window.") + start = torch.randint(self.size - length + 1, (batch_size,), device=self.device) + env = torch.randint(self.num_envs, (batch_size, 1), device=self.device) + rows = self._row(start[:, None] + torch.arange(length, device=self.device)) + out = {key: self.data[key][rows, env] for key in ("s", "s1", "a", "xi", "done")} + steps, frames = (self.data["ep_step"][rows, env], self.data["f"][rows, env]) + connected = ( + ~out["done"][:, :-1] + & (steps[:, 1:] == steps[:, :-1] + 1) + & (frames[:, 1:] == frames[:, :-1] + 1) + & (out["s"][:, 1:] == out["s1"][:, :-1]).all(-1) + ) + out["mask"] = torch.cat((torch.ones(batch_size, 1, device=self.device), connected.float()), -1).cumprod(-1) + return out + + def sample_starts(self, batch_size, history_length): + if not self.size: + raise ValueError("Cannot sample an empty replay buffer.") + age = torch.randint(self.size, (batch_size,), device=self.device) + env = torch.randint(self.num_envs, (batch_size,), device=self.device) + rows = self._row(age) + steps = self.data["ep_step"][rows, env].long() + back = torch.arange(history_length - 1, -1, -1, device=self.device) + back = torch.minimum(torch.minimum(back[None], steps[:, None]), age[:, None]) + history_rows = self._row(age[:, None] - back) + fields = ("s", "a_prev", "xi", "f") + start = {key: self.data[key][rows, env] for key in fields} + history = {key: self.data[key][history_rows, env[:, None]] for key in fields} + return (start, history) + + +class Learner: + """Own models, normalization, replay and optimization, independent of simulator stepping.""" + + def __init__( + self, + cfg: SuperMimicCfg, + mdp, + transition, + noise_terms, + reward_scale, + world_model: WorldModel, + *, + goal_frames: int, + history: int, + gamma: float, + ): + self.cfg, self.mdp, self.motion = (cfg, mdp, mdp.motion) + self.gamma = gamma + self.device = transition["s"].device + self.noise_terms, self.reward_scale = (noise_terms, reward_scale) + if not 0 <= gamma < 1: + raise ValueError("Discount must lie in [0, 1).") + self.goal_frames, self.history, self.horizon = (goal_frames, history, cfg.horizon) + dtype = torch.bfloat16 if cfg.amp == "bf16" else None + self.wm = world_model + self.pi = Policy(OBS_DIM, 73 * goal_frames, hist=history, amp_dtype=dtype).to(self.device) + self.critic = Critic(INV_DIM + CRITIC_EXTRA + 73 * goal_frames, amp_dtype=dtype).to(self.device) + self.opt_wm = torch.optim.Adam(self.wm.parameters(), lr=cfg.learning_rate) + self.opt_pi = torch.optim.Adam(self.pi.parameters(), lr=cfg.learning_rate, betas=(0.9, 0.95)) + self.opt_v = torch.optim.Adam(self.critic.parameters(), lr=cfg.learning_rate) + self.buffer = ReplayBuffer(cfg.buffer_steps, transition) + self.error_floor = torch.cat( + ( + torch.full((3,), 0.003), + torch.full((3,), 0.005), + torch.full((3,), 0.02), + torch.full((3,), 0.05), + torch.full((29,), 0.003), + torch.full((29,), 0.05), + ) + ).to(self.device) + self.rollout = torch.compile(self._wm_rollout, disable=not cfg.compile) + self.error = torch.compile(self._wm_err, disable=not cfg.compile) + self.imagine = torch.compile(self._imag_step, disable=not cfg.compile) + self.reward_termination = torch.compile(self._reward_term, disable=not cfg.compile) + self.critic_input_compiled = torch.compile(self._critic_input, disable=not cfg.compile) + + def add_noise(self, observation): + return torch.cat( + [ + observation[..., region] if fn is None else fn(observation[..., region], cfg) + for region, fn, cfg in self.noise_terms + ], + -1, + ) + + @torch.no_grad() + def act(self, observations): + self.pi.obs_norm.update(observations["policy"][:, -1]) + self.pi.goal_norm.update(observations["goal"]) + action = self.pi(observations["policy"], observations["goal"]) + return (action + self.cfg.action_noise * torch.randn_like(action)).clamp(-20, 20) + + @torch.no_grad() + def observe(self, transition): + self.buffer.add(transition) + state, successor = (unflatten_state(transition["s"]), unflatten_state(transition["s1"])) + self.wm.in_norm.update(self.wm.inputs(state, transition["a"], normalize_xi(transition["xi"]))) + self.wm.out_norm.update(state_difference(state, successor)) + + def _wm_rollout(self, s, a, xi): + S = unflatten_state(s[:, 0]) + preds = [] + for j in range(self.cfg.world_model_window): + S = self.wm.step(S, a[:, j], normalize_xi(xi[:, j])) + preds.append(S) + return tuple(torch.stack([P[i] for P in preds], 1) for i in range(6)) + + def _wm_err(self, S_hat, s1, mask): + std = torch.maximum(self.wm.out_norm.std(), self.error_floor) + e = (state_difference(unflatten_state(s1), S_hat) / std) ** 2 + e = e.mean(-1) + return ((e * mask).sum() / s1.shape[0], e[:, -1]) + + def _imag_step(self, S, o_hist, a_prev, f, xi, xi_n, g): + a = self.pi(o_hist, g) + S1 = self.wm.step(S, a, xi_n) + f1 = f + 1 + o1 = self.add_noise(self.mdp.obs(S1, a, xi, f1)) + o_hist = torch.cat((o_hist[:, 1:], o1.unsqueeze(1)), 1) + return (S1, a, o_hist, f1, self.motion.goal(f1, self.goal_frames)) + + def _reward_term(self, S, f, a, a_prev): + return (self.reward_scale * self.mdp.reward(S, f, a, a_prev), self.mdp.soft_termination(S, f)) + + def _critic_input(self, S, f, xi_n, g): + return torch.cat((critic_features(self.mdp, S, f, xi_n), g), -1) + + def wm_update(self): + w = self.buffer.sample_windows(self.cfg.world_model_batch_size, self.cfg.world_model_window) + S_hat = self.rollout(w["s"], w["a"], w["xi"]) + loss, e_last = self.error(S_hat, w["s1"], w["mask"]) + self.opt_wm.zero_grad(set_to_none=True) + loss.backward() + nn.utils.clip_grad_norm_(self.wm.parameters(), 1.0, error_if_nonfinite=True) + self.opt_wm.step() + m = w["mask"][:, -1] + return (loss.detach() / self.cfg.world_model_window, (e_last.detach() * m).sum() / m.sum().clamp(min=1)) + + def sample_valid_starts(self, B: int): + starts, hists, n = ([], [], 0) + attempts = 0 + while n < B: + attempts += 1 + if attempts > 32: + raise RuntimeError("Replay has too few valid imagination starts for the configured horizon.") + s, h = self.buffer.sample_starts(2 * B, self.history) + f = s["f"] + ok = f + self.horizon <= self.motion.clip_end[self.motion.clip_id[f]] + idx = ok.nonzero(as_tuple=False).squeeze(-1)[: B - n] + starts.append({k: v[idx] for k, v in s.items()}) + hists.append({k: v[idx] for k, v in h.items()}) + n += len(idx) + return tuple( + {key: torch.cat([sample[key] for sample in samples]) for key in samples[0]} for samples in (starts, hists) + ) + + def policy_update(self): + s, h = self.sample_valid_starts(self.cfg.policy_batch_size) + S = unflatten_state(s["s"]) + xi, f, a_prev = (s["xi"], s["f"], s["a_prev"]) + xi_n = normalize_xi(xi) + with torch.no_grad(): + o_hist = self.mdp.obs(unflatten_state(h["s"]), h["a_prev"], h["xi"], h["f"]) + o_hist = self.add_noise(o_hist) + self.wm.requires_grad_(False) + self.critic.requires_grad_(False) + B = self.cfg.policy_batch_size + g0 = self.motion.goal(f, self.goal_frames) + g = g0 + S0 = S + Ss, As, Aps, Fs, Gs = ([], [], [], [], []) + for _ in range(self.horizon): + Aps.append(a_prev) + S, a_prev, o_hist, f, g = self.imagine(S, o_hist, a_prev, f, xi, xi_n, g) + Ss.append(S) + As.append(a_prev) + Fs.append(f) + Gs.append(g) + + def flat(xs): + return torch.stack(xs).flatten(0, 1) + + S_all = tuple(flat([St[i] for St in Ss]) for i in range(6)) + f_all = flat(Fs) + r, d = self.reward_termination(S_all, f_all, flat(As), flat(Aps)) + r, d = (r.view(self.horizon, B), d.view(self.horizon, B)) + alive = torch.cumprod(1 - d, 0) + c_t = torch.cat((torch.ones(1, B, device=self.device), alive[:-1]), 0) + disc = self.gamma ** torch.arange(self.horizon, device=self.device, dtype=torch.float32).unsqueeze(-1) + J = (disc * c_t * r).sum(0) + c = alive[-1] + r_sum = (c_t * r).detach().sum(0) + rews, dones = (r.detach(), d.detach()) + terminal_value = self.critic(self._critic_input(S, f, xi_n, g)) + J = J + self.gamma**self.horizon * c * terminal_value + loss = -J.mean() + self.opt_pi.zero_grad(set_to_none=True) + loss.backward() + gn = nn.utils.clip_grad_norm_(self.pi.parameters(), 1.0) + nn.utils.clip_grad_value_(self.pi.parameters(), 0.1) + if torch.isfinite(gn): + self.opt_pi.step() + self.wm.requires_grad_(True) + self.critic.requires_grad_(True) + with torch.no_grad(): + S_c = tuple( + torch.cat((S0[i].unsqueeze(0), torch.stack([St[i].detach() for St in Ss])), 0).flatten(0, 1) + for i in range(6) + ) + f_c = torch.cat((s["f"].unsqueeze(0), torch.stack(Fs)), 0).flatten() + g_c = torch.cat((g0.unsqueeze(0), torch.stack(Gs)), 0).flatten(0, 1) + xi_c = xi_n.unsqueeze(0).expand(self.horizon + 1, -1, -1).flatten(0, 1) + X = self.critic_input_compiled(S_c, f_c, xi_c, g_c).view(self.horizon + 1, B, -1) + self.critic.norm.update(X) + with torch.no_grad(): + V = self.critic(X) + R = V[-1] + targets = [] + for t in range(self.horizon - 1, -1, -1): + R = rews[t] + self.gamma * (1 - dones[t]) * ( + (1 - self.cfg.gae_lambda) * V[t + 1] + self.cfg.gae_lambda * R + ) + targets.append(R) + targets = torch.stack(targets[::-1]) + Xf, Yf = (X[:-1].reshape(-1, X.shape[-1]), targets.reshape(-1)) + n = Xf.shape[0] + mb = n // self.cfg.critic_minibatches + for _ in range(self.cfg.critic_epochs): + perm = torch.randperm(n, device=self.device) + for i in range(self.cfg.critic_minibatches): + idx = perm[i * mb : (i + 1) * mb] + prediction = self.critic(Xf[idx]) + v_loss = ((prediction - Yf[idx]) ** 2).mean() + self.opt_v.zero_grad(set_to_none=True) + v_loss.backward() + nn.utils.clip_grad_norm_(self.critic.parameters(), 1.0, error_if_nonfinite=True) + self.opt_v.step() + return ( + loss.detach(), + v_loss.detach(), + (r_sum / self.horizon / (MAX_REWARD_RATE * self.reward_scale)).mean(), + gn.detach(), + c.mean().detach(), + ) diff --git a/source/isaaclab_rl/isaaclab_rl/supermimic/networks.py b/source/isaaclab_rl/isaaclab_rl/supermimic/networks.py new file mode 100644 index 000000000000..8dabc2a4b11c --- /dev/null +++ b/source/isaaclab_rl/isaaclab_rl/supermimic/networks.py @@ -0,0 +1,194 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Portions derived from SuperMimic (MIT), Copyright (c) 2026 ETH Zurich, Chong Zhang. +# See docs/licenses/dependencies/supermimic-license.txt. + +"""World model, policy (motion encoder + shared decoder), and critic (paper Table A.4).""" + +from __future__ import annotations + +import torch +import torch.nn as nn + +from isaaclab_tasks.contrib.g1_tracking.kinematics import mat_to_6d, so3_exp, so3_log, yaw_mat +from isaaclab_tasks.contrib.g1_tracking.tracking import NUM_J, XI_DIM + +INV_DIM = 1 + 6 + 3 + 3 + 29 + 29 +CRITIC_EXTRA = 29 + 14 * 3 + 3 + 6 + XI_DIM + + +def invariant_features(S) -> torch.Tensor: + """Heading- and position-invariant state features: height, tilt (6D), body-frame velocities, joints.""" + p, R, v, w, q, qd = S + H = yaw_mat(R) + Rt = R.transpose(-1, -2) + return torch.cat( + ( + p[..., 2:3], + mat_to_6d(H.transpose(-1, -2) @ R), + (Rt @ v.unsqueeze(-1)).squeeze(-1), + (Rt @ w.unsqueeze(-1)).squeeze(-1), + q, + qd, + ), + -1, + ) + + +def normalize_xi(xi: torch.Tensor) -> torch.Tensor: + return torch.cat(((xi[..., :1] - 0.75) / 0.45, xi[..., 1:30] / 0.01, xi[..., 30:33] / 0.05), -1) + + +def mlp(inp: int, hidden: list[int], out: int, act=nn.ELU, layer_norm: bool = False) -> nn.Sequential: + layers, d = [], inp + for h in hidden: + layers.append(nn.Linear(d, h)) + if layer_norm: + layers.append(nn.LayerNorm(h)) + layers.append(act()) + d = h + layers.append(nn.Linear(d, out)) + return nn.Sequential(*layers) + + +class RunningNorm(nn.Module): + """Running mean/std normalizer; updated from data without gradients, frozen after ``max_count``.""" + + def __init__(self, dim: int, max_count: float = 2e8, eps: float = 1e-2): + super().__init__() + self.register_buffer("mean", torch.zeros(dim)) + self.register_buffer("var", torch.ones(dim)) + self.register_buffer("count", torch.zeros(())) + self.max_count, self.eps = max_count, eps + + @torch.no_grad() + def update(self, x: torch.Tensor): + x = x.reshape(-1, x.shape[-1]).float() + if self.count >= self.max_count: + return + n = x.shape[0] + m, v = x.mean(0), x.var(0, unbiased=False) + tot = self.count + n + delta = m - self.mean + self.mean += delta * n / tot + self.var = (self.var * self.count + v * n + delta**2 * self.count * n / tot) / tot + self.count = tot + + def forward(self, x): + return (x - self.mean) / torch.sqrt(self.var + self.eps**2) + + def std(self): + return torch.sqrt(self.var + self.eps**2) + + +WM_IN = INV_DIM + NUM_J + XI_DIM +WM_OUT = 3 + 3 + 3 + 3 + NUM_J + NUM_J + + +def run_network(net: nn.Module, x: torch.Tensor, dtype: torch.dtype | None) -> torch.Tensor: + if dtype is None: + return net(x) + with torch.autocast("cuda", dtype=dtype): + return net(x).float() + + +def state_difference(S, S1): + """Frame-local increments from S to S1 (the regression targets).""" + p, R, v, w, q, qd = S + p1, R1, v1, w1, q1, qd1 = S1 + H = yaw_mat(R) + Rt = R.transpose(-1, -2) + return torch.cat( + ( + (H.transpose(-1, -2) @ (p1 - p).unsqueeze(-1)).squeeze(-1), + so3_log(Rt @ R1), + (Rt @ (v1 - v).unsqueeze(-1)).squeeze(-1), + (Rt @ (w1 - w).unsqueeze(-1)).squeeze(-1), + q1 - q, + qd1 - qd, + ), + -1, + ) + + +class WorldModel(nn.Module): + """Deterministic dynamics s_{t+1} = f(s_t, a_t, xi) predicting frame-local state increments.""" + + def __init__(self, hidden=(2048, 2048, 2048), amp_dtype=None, network: nn.Module | None = None): + super().__init__() + self.amp_dtype = amp_dtype + self.net = network if network is not None else mlp(WM_IN, list(hidden), WM_OUT, act=nn.SiLU, layer_norm=True) + if network is None: + nn.init.zeros_(self.net[-1].weight) + nn.init.zeros_(self.net[-1].bias) + self.in_norm = RunningNorm(WM_IN) + self.out_norm = RunningNorm(WM_OUT, eps=1e-4) + + def inputs(self, S, a, xi_n): + return torch.cat((invariant_features(S), a, xi_n), -1) + + def step(self, S, a, xi_n): + p, R, v, w, q, qd = S + tokens = self.in_norm(self.inputs(S, a, xi_n)).unsqueeze(-2) + y = run_network(self.net, tokens, self.amp_dtype).squeeze(-2) + d = y * self.out_norm.std() + self.out_norm.mean + H = yaw_mat(R) + dp, dth, dv, dw, dq, dqd = d.split((3, 3, 3, 3, NUM_J, NUM_J), -1) + p1 = p + (H @ dp.unsqueeze(-1)).squeeze(-1) + R1 = R @ so3_exp(dth) + v1 = v + (R @ dv.unsqueeze(-1)).squeeze(-1) + w1 = w + (R @ dw.unsqueeze(-1)).squeeze(-1) + return (p1, R1, v1, w1, q + dq, qd + dqd) + + +class Policy(nn.Module): + """a = dec(enc(g), o_{t-9:t}); deterministic, latent z in R^32 bounded by tanh.""" + + def __init__(self, obs_dim: int, goal_dim: int, hist: int = 10, z_dim: int = 32, amp_dtype=None): + super().__init__() + self.hist = hist + self.amp_dtype = amp_dtype + self.enc = mlp(goal_dim, [2048, 1024, 512, 256], z_dim) + self.dec = mlp(z_dim + obs_dim * hist, [2048, 1024, 512, 256], NUM_J) + with torch.no_grad(): + self.dec[-1].weight.mul_(0.01) + self.dec[-1].bias.zero_() + self.obs_norm = RunningNorm(obs_dim) + self.goal_norm = RunningNorm(goal_dim) + + def forward(self, obs_hist: torch.Tensor, goal: torch.Tensor) -> torch.Tensor: + """obs_hist: (B, hist, obs_dim) noisy observations, newest last.""" + z = torch.tanh(run_network(self.enc, self.goal_norm(goal), self.amp_dtype)) + o = self.obs_norm(obs_hist).flatten(-2) + return run_network(self.dec, torch.cat((z, o), -1), self.amp_dtype) + + +class Critic(nn.Module): + def __init__(self, in_dim: int, amp_dtype=None): + super().__init__() + self.amp_dtype = amp_dtype + self.net = mlp(in_dim, [256, 128], 1) + self.norm = RunningNorm(in_dim) + + def forward(self, x): + return run_network(self.net, self.norm(x), self.amp_dtype).squeeze(-1) + + +def critic_features(mdp, S, f, xi) -> torch.Tensor: + """Critic input: invariant state, goal, sim context, and robot-vs-reference errors.""" + m = mdp.motion + p, R, v, w, q, qd = S + H = yaw_mat(R) + Ht = H.transpose(-1, -2) + rp, rR = m.root_p[f], m.root_R[f] + bp = mdp.kin.forward(p, R, v, w, q, qd)[0] + dR = yaw_mat(R @ rR.transpose(-1, -2)) + anchor = torch.cat((p[..., :2], rp[..., 2:3]), -1) + ref_bp = anchor.unsqueeze(-2) + ((dR.unsqueeze(-3) @ (m.body_p[f] - rp.unsqueeze(-2)).unsqueeze(-1)).squeeze(-1)) + e_body = (Ht.unsqueeze(-3) @ (ref_bp - bp).unsqueeze(-1)).squeeze(-1).flatten(-2) + e_root = (Ht @ (rp - p).unsqueeze(-1)).squeeze(-1) + rel = mat_to_6d(R.transpose(-1, -2) @ rR) + return torch.cat((invariant_features(S), m.q[f] - q, e_body, e_root, rel, xi), -1) diff --git a/source/isaaclab_rl/isaaclab_rl/supermimic/supermimic_cfg.py b/source/isaaclab_rl/isaaclab_rl/supermimic/supermimic_cfg.py new file mode 100644 index 000000000000..8181ab7161f5 --- /dev/null +++ b/source/isaaclab_rl/isaaclab_rl/supermimic/supermimic_cfg.py @@ -0,0 +1,51 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Learning and checkpoint settings for the native SuperMimic workflow.""" + +from isaaclab.utils import configclass + + +@configclass +class SuperMimicCfg: + """SuperMimic settings; task timing, observations and reference goals belong to the environment.""" + + seed: int = 0 + experiment_name: str = "g1_tracking" + max_iterations: int = 14_000 + buffer_steps: int = 1000 + """Replay capacity in control steps per world.""" + update_to_data_ratio: float = 0.4 + """Learning iterations per physical control step.""" + horizon: int = 32 + """Differentiable imagined rollout length.""" + world_model_window: int = 8 + world_model_batches: int = 10 + world_model_batch_size: int = 2048 + policy_batch_size: int = 4096 + learning_rate: float = 2e-4 + gae_lambda: float = 0.95 + action_noise: float = 0.1 + critic_epochs: int = 5 + critic_minibatches: int = 4 + warmup_steps: int = 64 + world_model_only_iterations: int = 500 + learning_rate_decay_iterations: int = 250_000 + world_model: str = "mlp" + """Dynamics architecture: ``mlp`` or ``nerd-one-frame`` (optional PhysicsNeMo Robotics).""" + amp: str = "bf16" + """Neural projection precision: ``bf16`` or ``none``. Geometry always uses float32.""" + compile: bool = True + """Compile the tensor rollout, objective and regression functions with Torch.""" + save_interval: int = 2000 + save_interval_seconds: float = 120.0 + """Maximum time between checkpoint requests, checked at iteration boundaries.""" + log_interval: int = 50 + max_hours: float | None = None + """Cumulative training-hour budget, including prior progress when resuming.""" + wandb_project: str | None = None + wandb_entity: str | None = None + wandb_run_id: str | None = None + wandb_name: str | None = None diff --git a/source/isaaclab_rl/test/test_entrypoints.py b/source/isaaclab_rl/test/test_entrypoints.py index 4f1f98ae4c13..92357c887ce5 100644 --- a/source/isaaclab_rl/test/test_entrypoints.py +++ b/source/isaaclab_rl/test/test_entrypoints.py @@ -46,6 +46,11 @@ ("import isaaclab_rl.rl_games.pbt", ["isaaclab_rl.rl_games.pbt.pbt", "rl_games", "torch"]), ("import isaaclab_rl.rsl_rl", ["isaaclab_rl.rsl_rl.vecenv_wrapper", "rsl_rl", "torch"]), ("import isaaclab_rl.utils", ["torch"]), + ( + "import isaaclab_tasks.contrib.g1_tracking", + ["isaaclab_rl.supermimic", "physicsnemo_robotics", "wandb", "supermimic"], + ), + ("import isaaclab_rl.supermimic.networks", ["physicsnemo_robotics", "wandb", "supermimic"]), # resolves through isaaclab_rl.entrypoints, so it also covers importing that package directly ( "from isaaclab_rl import run_play_cli", diff --git a/source/isaaclab_rl/test/test_supermimic.py b/source/isaaclab_rl/test/test_supermimic.py new file mode 100644 index 000000000000..92f761e57f7a --- /dev/null +++ b/source/isaaclab_rl/test/test_supermimic.py @@ -0,0 +1,231 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Portions derived from SuperMimic (MIT), Copyright (c) 2026 ETH Zurich, Chong Zhang. +# See docs/licenses/dependencies/supermimic-license.txt. + +"""Contracts for native SuperMimic ownership, replay and evaluation.""" + +import ast +from pathlib import Path + +import numpy as np +import pytest +import torch + +from isaaclab.paths import ISAACLAB_ROOT +from isaaclab.test.utils import DeviceScope, test_devices + +from isaaclab_rl.supermimic.learner import ReplayBuffer + + +@pytest.mark.parametrize("architecture", ("mlp", "nerd-one-frame")) +@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) +def test_compiled_world_model_propagates_finite_action_gradients(architecture, device): + """Imagined control needs action gradients through frozen dynamics, including identity rotations.""" + from isaaclab_rl.supermimic.networks import WM_IN, WM_OUT, WorldModel, state_difference + + torch.manual_seed(0) + network = None + if architecture == "nerd-one-frame": + dependency = pytest.importorskip("physicsnemo_robotics.models.parallel_nerd") + network = dependency.ParallelNeRD(WM_IN, WM_OUT, temporal_attention_spans=(1, 1, 1, 1)) + torch.nn.init.normal_(network.head_output.weight, std=0.001) + world = WorldModel(hidden=(32, 32), network=network).to(device) + if network is None: + torch.nn.init.normal_(world.net[-1].weight, std=0.001) + world.requires_grad_(False) + state = ( + torch.zeros(4, 3, device=device), + torch.eye(3, device=device).expand(4, -1, -1), + torch.zeros(4, 3, device=device), + torch.zeros(4, 3, device=device), + torch.zeros(4, 29, device=device), + torch.zeros(4, 29, device=device), + ) + action = torch.zeros(4, 29, device=device, requires_grad=True) + context = torch.zeros(4, 33, device=device) + successor = torch.compile(world.step)(state, action, context) + loss = (state_difference(state, successor) - 0.1).square().mean() + gradient = torch.autograd.grad(loss, action)[0] + assert torch.isfinite(loss) and gradient.isfinite().all() and gradient.abs().sum() > 0 + assert all(parameter.grad is None for parameter in world.parameters()) + + +def test_native_ownership_has_no_external_training_package(): + """Reject external forwarders, duplicate task stepping and task-to-learner dependency inversions.""" + task = Path(ISAACLAB_ROOT) / "source/isaaclab_tasks/isaaclab_tasks/contrib/g1_tracking" + learning = Path(ISAACLAB_ROOT) / "source/isaaclab_rl/isaaclab_rl/supermimic" + backends = learning.parent / "entrypoints/backends" + roots = (backends / "train_supermimic.py", backends / "play_supermimic.py") + for path in (*task.rglob("*.py"), *learning.rglob("*.py"), *roots): + tree = ast.parse(path.read_text()) + imports = [] + for node in ast.walk(tree): + if isinstance(node, ast.Import): + imports.extend(alias.name for alias in node.names) + elif isinstance(node, ast.ImportFrom) and node.level == 0: + imports.append(node.module or "") + assert not any(name == "supermimic" or name.startswith("supermimic.") for name in imports), path + symbols = {node.name for node in ast.walk(tree) if isinstance(node, (ast.ClassDef, ast.FunctionDef))} + assert symbols.isdisjoint({"TrackingSim", "TransitionCapture"}), path + if path.is_relative_to(task): + assert not any(name.startswith(("isaaclab_rl", "physicsnemo", "wandb")) for name in imports), path + assert "step" not in symbols and path.name not in ("train.py", "learner.py"), path + elif path.is_relative_to(learning): + assert path.name not in ("train.py", "evaluate.py", "sim.py", "robot.py"), path + assert not any(name.startswith(("argparse", "gymnasium", "wandb", "isaaclab.envs")) for name in imports), ( + path + ) + + +def test_native_task_registration_and_optional_capture(): + from isaaclab_tasks.utils import resolve_task_config + + env_cfg, agent_cfg = resolve_task_config( + "IsaacContrib-Motion-Tracking-G1", "supermimic_cfg_entry_point", overrides=[] + ) + assert env_cfg.recorders.transition is None + assert type(env_cfg).__module__.startswith("isaaclab_tasks.") + assert type(agent_cfg).__module__.startswith("isaaclab_rl.") + + +def test_replay_ring_masks_boundaries_and_owns_samples(): + transition = { + "s": torch.zeros(1, 71), + "s1": torch.zeros(1, 71), + "a": torch.zeros(1, 29), + "a_prev": torch.zeros(1, 29), + "xi": torch.zeros(1, 33), + "r": torch.zeros(1), + "f": torch.zeros(1, dtype=torch.long), + "ep_step": torch.zeros(1, dtype=torch.long), + "done": torch.zeros(1, dtype=torch.bool), + } + ring = ReplayBuffer(4, transition) + assert "r" not in ring.data # This learner evaluates analytical rewards; capture still supplies them to others. + for i in range(6): + transition["s"].fill_(i) + transition["s1"].fill_(i + 1) + transition["f"].fill_(i) + transition["ep_step"].fill_(i) + transition["done"].fill_(i == 3) + ring.add(transition) + transition["s"].fill_(99) + window = ring.sample_windows(4, 4) + torch.testing.assert_close(window["s"][..., 0], torch.tensor([[2.0, 3.0, 4.0, 5.0]]).expand(4, -1)) + torch.testing.assert_close(window["mask"], torch.tensor([[1.0, 1.0, 0.0, 0.0]]).expand(4, -1)) + # A state intervention or manual reset is a boundary even when no sampled done flag was set. + ring.data["done"].zero_() + ring.data["s"][ring._row(1), 0, 0] = -1 + torch.testing.assert_close(ring.sample_windows(4, 4)["mask"], torch.tensor([[1.0, 0.0, 0.0, 0.0]]).expand(4, -1)) + ring.data["ep_step"].zero_() + start, history = ring.sample_starts(32, 10) + torch.testing.assert_close(history["s"], start["s"][:, None].expand(-1, 10, -1)) + + +def test_evaluation_batches_cover_last_partial_batch_and_failed_tails(): + from types import SimpleNamespace + + from isaaclab_rl.entrypoints.backends.play_supermimic import evaluate + + lengths = torch.tensor([3, 4, 2, 5, 3, 4, 2]) + motion = SimpleNamespace(names=[f"clip{i}" for i in range(7)], clip_len=lengths) + command = SimpleNamespace(motion=motion, cfg=SimpleNamespace(evaluation_clip_offset=0)) + capture = SimpleNamespace(latest={}) + + class Env: + num_envs, device, step_dt = 3, "cpu", 0.02 + command_manager = SimpleNamespace(get_term=lambda name: command) + recorder_manager = SimpleNamespace(get_term=lambda name: capture) + reward_manager = SimpleNamespace(get_term_cfg=lambda name: SimpleNamespace(weight=1.0)) + + @property + def unwrapped(self): + return self + + def reset(self): + self.ids = (torch.arange(3) + command.cfg.evaluation_clip_offset) % 7 + self.elapsed = 0 + return {"policy": None, "goal": None}, {} + + def step(self, action): + self.elapsed += 1 + term = (self.ids == 3) & (self.elapsed == 1) + capture.latest = { + "r": (self.ids + 1) * 0.01, + "done": term | (self.elapsed >= lengths[self.ids] - 1), + "term": term, + } + return {"policy": None, "goal": None}, None, None, None, None + + result = evaluate(Env(), lambda *args: None) + assert result["failed"] == ["clip3"] + assert list(result["per_clip_steps"].values()) == [2, 3, 1, 1, 2, 3, 1] + np.testing.assert_allclose(list(result["per_clip"].values()), [0.1, 0.2, 0.3, 0.1, 0.5, 0.6, 0.7], rtol=1e-6) + + +def test_training_root_rejects_different_analytical_task_semantics(tmp_path, monkeypatch): + from isaaclab.envs import mdp + + from isaaclab_rl.entrypoints.backends import train_supermimic as train + from isaaclab_rl.supermimic.supermimic_cfg import SuperMimicCfg + + from isaaclab_tasks.contrib.g1_tracking.g1_tracking_env_cfg import G1TrackingEnvCfg + + class ReachedLauncher(Exception): + pass + + def launch(*args): + raise ReachedLauncher + + monkeypatch.setattr(train, "launch_simulation", launch) + invalid = ( + (("policy", "joint_vel", "func"), mdp.joint_pos), # Same width, different physical quantity. + (("policy", "joint_vel", "params"), {"asset_cfg": "different_asset"}), + (("policy", "joint_vel", "delay_max_lag"), 1), + (("policy", "flatten_history_dim"), True), + (("goal", "reference", "scale"), 2.0), + ) + argv = ["--task", "IsaacContrib-Motion-Tracking-G1", "--log_dir", str(tmp_path)] + for path, value in invalid: + cfg = G1TrackingEnvCfg() + target = cfg.observations + for name in path[:-1]: + target = getattr(target, name) + setattr(target, path[-1], value) + monkeypatch.setattr(train, "resolve_task_config", lambda *args, **kwargs: (cfg, SuperMimicCfg())) + with pytest.raises(ValueError, match="declared functions and layout"): + train.run(argv) + cfg = G1TrackingEnvCfg() + term = cfg.observations.policy.__dict__.pop("joint_vel") + cfg.observations.policy.joint_vel = term + monkeypatch.setattr(train, "resolve_task_config", lambda *args, **kwargs: (cfg, SuperMimicCfg())) + with pytest.raises(ValueError, match="declared functions and layout"): + train.run(argv) + for name, attribute, value in ( + ("failed", "func", mdp.time_out), + ("failed", "time_out", True), + ("clip_end", None, None), + ): + cfg = G1TrackingEnvCfg() + if attribute is None: + setattr(cfg.terminations, name, value) + else: + setattr(getattr(cfg.terminations, name), attribute, value) + monkeypatch.setattr(train, "resolve_task_config", lambda *args, **kwargs: (cfg, SuperMimicCfg())) + with pytest.raises(ValueError, match="Analytical termination"): + train.run(argv) + cfg = G1TrackingEnvCfg() + cfg.observations.policy.enable_corruption = False + cfg.observations.policy.joint_vel.noise.n_max = 0.25 + cfg.commands.motion.objective.term_height = 0.3 + cfg.episode_length_s = 4.0 + monkeypatch.setattr(train, "resolve_task_config", lambda *args, **kwargs: (cfg, SuperMimicCfg())) + with pytest.raises(ReachedLauncher): + train.run(argv) # Supported task parameters reach construction without a simulator. + cfg.observations.policy.history_length = 0 + with pytest.raises(ValueError, match="must be positive"): + train.run(argv) diff --git a/source/isaaclab_tasks/changelog.d/g1-tracking.added.rst b/source/isaaclab_tasks/changelog.d/g1-tracking.added.rst new file mode 100644 index 000000000000..3851fd1d62ec --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/g1-tracking.added.rst @@ -0,0 +1,2 @@ +* Added the contributed G1 motion-tracking task with native manager terms, optional transition capture, + differentiable tracking objectives and stationary goal padding beyond recorded clip endpoints. diff --git a/source/isaaclab_tasks/changelog.d/g1-tracking.minor b/source/isaaclab_tasks/changelog.d/g1-tracking.minor new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/g1_tracking/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/g1_tracking/__init__.py new file mode 100644 index 000000000000..513fa5cf146d --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/g1_tracking/__init__.py @@ -0,0 +1,19 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""G1 reference-motion tracking with native environment managers.""" + +import gymnasium as gym + +gym.register( + id="IsaacContrib-Motion-Tracking-G1", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.g1_tracking_env_cfg:G1TrackingEnvCfg", + "default_agent": "supermimic", + "supermimic_cfg_entry_point": "isaaclab_rl.supermimic.supermimic_cfg:SuperMimicCfg", + }, +) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/g1_tracking/g1_tracking_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/g1_tracking/g1_tracking_env_cfg.py new file mode 100644 index 000000000000..905194202ab9 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/g1_tracking/g1_tracking_env_cfg.py @@ -0,0 +1,250 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Portions derived from SuperMimic (MIT), Copyright (c) 2026 ETH Zurich, Chong Zhang. +# See docs/licenses/dependencies/supermimic-license.txt. + +"""The G1 tracking task's native scene, manager terms and Newton physics preset.""" + +from pathlib import Path + +from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg +from isaaclab_newton.sim.schemas import NewtonArticulationCfg +from isaaclab_physx.sim.schemas import PhysxRigidBodyCfg + +import isaaclab.sim as sim_utils +from isaaclab.actuators import ImplicitActuatorCfg +from isaaclab.assets import ArticulationCfg, AssetBaseCfg +from isaaclab.envs import ManagerBasedRLEnvCfg, mdp +from isaaclab.managers import ( + CommandTermCfg, + DatasetExportMode, + EventTermCfg, + ObservationGroupCfg, + ObservationTermCfg, + RecorderManagerBaseCfg, + RecorderTermCfg, + RewardTermCfg, + TerminationTermCfg, +) +from isaaclab.paths import ISAACLAB_ROOT +from isaaclab.scene import InteractiveSceneCfg +from isaaclab.utils import configclass, replace +from isaaclab.utils.noise import UniformNoiseCfg + +from isaaclab_tasks.utils import PresetCfg + +from .motion import CONTROL_DT +from .tracking import RewardCfg + +G1_DATA_DIR = Path(ISAACLAB_ROOT) / "data/g1_tracking" + +ARMATURE_5020 = 0.003609725 +ARMATURE_7520_14 = 0.010177520 +ARMATURE_7520_22 = 0.025101925 +ARMATURE_4010 = 0.00425 + +NATURAL_FREQ = 10 * 2.0 * 3.1415926535 # 10 Hz +DAMPING_RATIO = 2.0 + + +# Per-joint (effort limit, velocity limit, armature) by motor type. +_M5020 = (25.0, 37.0, ARMATURE_5020) +_M4010 = (5.0, 22.0, ARMATURE_4010) +JOINT_MOTORS = { + "hip_pitch": (88.0, 32.0, ARMATURE_7520_14), + "hip_roll": (139.0, 20.0, ARMATURE_7520_22), + "hip_yaw": (88.0, 32.0, ARMATURE_7520_14), + "knee": (139.0, 20.0, ARMATURE_7520_22), + "ankle_pitch": (50.0, 37.0, 2.0 * ARMATURE_5020), + "ankle_roll": (50.0, 37.0, 2.0 * ARMATURE_5020), + "waist_yaw": (88.0, 32.0, ARMATURE_7520_14), + "waist_roll": (50.0, 37.0, 2.0 * ARMATURE_5020), + "waist_pitch": (50.0, 37.0, 2.0 * ARMATURE_5020), + "shoulder_pitch": _M5020, + "shoulder_roll": _M5020, + "shoulder_yaw": _M5020, + "elbow": _M5020, + "wrist_roll": _M5020, + "wrist_pitch": _M4010, + "wrist_yaw": _M4010, +} + + +DEFAULT_JOINT_POS = { + ".*_hip_pitch_joint": -0.312, + ".*_knee_joint": 0.669, + ".*_ankle_pitch_joint": -0.363, + ".*_elbow_joint": 0.6, + "left_shoulder_roll_joint": 0.2, + "left_shoulder_pitch_joint": 0.2, + "right_shoulder_roll_joint": -0.2, + "right_shoulder_pitch_joint": 0.2, +} + + +G1_CFG = ArticulationCfg( + # Prepare the matching 29-DoF MJCF and convert it to USD before constructing this task. + spawn=sim_utils.UsdFileCfg( + usd_path=str(G1_DATA_DIR / "g1_usd/g1_29dof/g1_29dof.usda"), + variants={"Physics": "mujoco"}, + activate_contact_sensors=False, + rigid_props=PhysxRigidBodyCfg(disable_gravity=False, linear_damping=0.0, angular_damping=0.0), + articulation_props=NewtonArticulationCfg(self_collision_enabled=False), + ), + init_state=ArticulationCfg.InitialStateCfg( + pos=(0.0, 0.0, 0.76), joint_pos=DEFAULT_JOINT_POS, joint_vel={".*": 0.0} + ), + soft_joint_pos_limit_factor=0.9, + actuators={ + "joints": ImplicitActuatorCfg( + joint_names_expr=[".*"], + joint_effort_limit={f".*{name}_joint": motor[0] for name, motor in JOINT_MOTORS.items()}, + joint_velocity_limit={f".*{name}_joint": motor[1] for name, motor in JOINT_MOTORS.items()}, + stiffness={f".*{name}_joint": motor[2] * NATURAL_FREQ**2 for name, motor in JOINT_MOTORS.items()}, + damping={ + f".*{name}_joint": 2 * DAMPING_RATIO * motor[2] * NATURAL_FREQ for name, motor in JOINT_MOTORS.items() + }, + armature={f".*{name}_joint": motor[2] for name, motor in JOINT_MOTORS.items()}, + ) + }, +) + + +@configclass +class TrackingPhysicsCfg(PresetCfg): + newton_mjwarp = NewtonCfg( + solver_cfg=MJWarpSolverCfg( + njmax=500, nconmax=128, cone="pyramidal", integrator="implicitfast", use_mujoco_contacts=True + ), + num_substeps=4, + use_cuda_graph=True, + ) + default = newton_mjwarp + + +@configclass +class TrackingSceneCfg(InteractiveSceneCfg): + ground = AssetBaseCfg( + prim_path="/World/ground", + spawn=sim_utils.GroundPlaneCfg( + physics_material=sim_utils.RigidBodyMaterialBaseCfg(static_friction=0.3, dynamic_friction=0.3) + ), + ) + robot = replace(G1_CFG, prim_path="{ENV_REGEX_NS}/Robot") + + +@configclass +class MotionCommandCfg(CommandTermCfg): + class_type = "isaaclab_tasks.contrib.g1_tracking.mdp:MotionCommand" + resampling_time_range = (1.0e9, 1.0e9) + motion_dir: str = str(G1_DATA_DIR / "lafan1_g1") + """Directory containing metric/xyzw/radian 36-channel motion CSV or NumPy files.""" + kinematics_path: str = str(G1_DATA_DIR / "g1_description/g1_29dof.xml") + """MJCF matching the scene asset, used for differentiable link kinematics.""" + source_fps: float = 30.0 + goal_frames: int = 10 + gamma: float = 0.99 + evaluation: bool = False + evaluation_clip_offset: int = 0 + reset_noise: bool = True + objective: RewardCfg = RewardCfg() + + +@configclass +class CommandsCfg: + motion = MotionCommandCfg() + + +@configclass +class ActionsCfg: + joint_pos = mdp.JointPositionActionCfg( + class_type="isaaclab_tasks.contrib.g1_tracking.mdp:ContextJointPositionAction", + asset_name="robot", + joint_names=[".*"], + scale={ + name: 0.25 * G1_CFG.actuators["joints"].joint_effort_limit[name] / stiffness + for name, stiffness in G1_CFG.actuators["joints"].stiffness.items() + }, + ) + + +@configclass +class ObservationsCfg: + @configclass + class PolicyCfg(ObservationGroupCfg): + joint_pos = ObservationTermCfg( + func="isaaclab_tasks.contrib.g1_tracking.mdp:joint_pos_rel", noise=UniformNoiseCfg(n_min=-0.01, n_max=0.01) + ) + joint_vel = ObservationTermCfg(func=mdp.joint_vel, noise=UniformNoiseCfg(n_min=-0.5, n_max=0.5)) + angular_velocity = ObservationTermCfg(func=mdp.base_ang_vel, noise=UniformNoiseCfg(n_min=-0.2, n_max=0.2)) + gravity = ObservationTermCfg(func=mdp.projected_gravity, noise=UniformNoiseCfg(n_min=-0.05, n_max=0.05)) + previous_action = ObservationTermCfg(func=mdp.last_action) + orientation = ObservationTermCfg( + func="isaaclab_tasks.contrib.g1_tracking.mdp:relative_orientation", + noise=UniformNoiseCfg(n_min=-0.05, n_max=0.05), + ) + history_length = 10 + flatten_history_dim = False + concatenate_terms = True + enable_corruption = True + + @configclass + class GoalCfg(ObservationGroupCfg): + reference = ObservationTermCfg(func=mdp.generated_commands, params={"command_name": "motion"}) + + policy = PolicyCfg() + goal = GoalCfg() + + +@configclass +class RewardsCfg: + tracking = RewardTermCfg(func="isaaclab_tasks.contrib.g1_tracking.mdp:tracking_reward", weight=1.0) + + +@configclass +class TerminationsCfg: + failed = TerminationTermCfg(func="isaaclab_tasks.contrib.g1_tracking.mdp:tracking_failure") + clip_end = TerminationTermCfg(func="isaaclab_tasks.contrib.g1_tracking.mdp:clip_end", time_out=True) + time_out = TerminationTermCfg(func=mdp.time_out, time_out=True) + + +@configclass +class EventsCfg: + context = EventTermCfg(func="isaaclab_tasks.contrib.g1_tracking.mdp:randomize_context", mode="startup") + + +@configclass +class TransitionRecorderCfg(RecorderTermCfg): + class_type = "isaaclab_tasks.contrib.g1_tracking.mdp:TransitionRecorder" + record_to_episodes: bool = False + + +@configclass +class RecordersCfg(RecorderManagerBaseCfg): + transition: TransitionRecorderCfg | None = None + dataset_export_mode = DatasetExportMode.EXPORT_NONE + + +@configclass +class G1TrackingEnvCfg(ManagerBasedRLEnvCfg): + scene = TrackingSceneCfg(num_envs=4096, env_spacing=0.0) + commands = CommandsCfg() + actions = ActionsCfg() + observations = ObservationsCfg() + rewards = RewardsCfg() + terminations = TerminationsCfg() + events = EventsCfg() + recorders = RecordersCfg() + sim = sim_utils.SimulationCfg(dt=CONTROL_DT, render_interval=1, physics=TrackingPhysicsCfg()) + decimation = 1 + episode_length_s = 10.0 + + def play_mode(self): + self.commands.motion.evaluation = True + self.observations.policy.enable_corruption = False + self.episode_length_s = 1.0e9 + self.scene.ground.spawn.physics_material.static_friction = 1.0 + self.scene.ground.spawn.physics_material.dynamic_friction = 1.0 diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/g1_tracking/kinematics.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/g1_tracking/kinematics.py new file mode 100644 index 000000000000..9988b66f3fdc --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/g1_tracking/kinematics.py @@ -0,0 +1,164 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Portions derived from SuperMimic (MIT), Copyright (c) 2026 ETH Zurich, Chong Zhang. +# See docs/licenses/dependencies/supermimic-license.txt. + +"""Differentiable G1 geometry shared by reference motion, real states and imagined dynamics. + +Native math owns quaternion conversion (xyzw). The smooth SO(3) exponential/logarithm +and regularized heading projection retain the derivatives required by the learned rollout. +""" + +from __future__ import annotations + +import mujoco +import numpy as np +import torch + +from isaaclab.utils.math import axis_angle_from_quat, matrix_from_quat, quat_from_angle_axis, quat_from_matrix + +# Joint order of the LAFAN1 retarget CSVs (Unitree SDK order). +G1_JOINT_NAMES = [ + "left_hip_pitch_joint", "left_hip_roll_joint", "left_hip_yaw_joint", "left_knee_joint", + "left_ankle_pitch_joint", "left_ankle_roll_joint", + "right_hip_pitch_joint", "right_hip_roll_joint", "right_hip_yaw_joint", "right_knee_joint", + "right_ankle_pitch_joint", "right_ankle_roll_joint", + "waist_yaw_joint", "waist_roll_joint", "waist_pitch_joint", + "left_shoulder_pitch_joint", "left_shoulder_roll_joint", "left_shoulder_yaw_joint", "left_elbow_joint", + "left_wrist_roll_joint", "left_wrist_pitch_joint", "left_wrist_yaw_joint", + "right_shoulder_pitch_joint", "right_shoulder_roll_joint", "right_shoulder_yaw_joint", "right_elbow_joint", + "right_wrist_roll_joint", "right_wrist_pitch_joint", "right_wrist_yaw_joint", +] # fmt: skip + +# BeyondMimic's tracked bodies; the anchor/root is the pelvis. +TRACKED_BODIES = [ + "pelvis", "left_hip_roll_link", "left_knee_link", "left_ankle_roll_link", + "right_hip_roll_link", "right_knee_link", "right_ankle_roll_link", "torso_link", + "left_shoulder_roll_link", "left_elbow_link", "left_wrist_yaw_link", + "right_shoulder_roll_link", "right_elbow_link", "right_wrist_yaw_link", +] # fmt: skip + + +def so3_exp(v: torch.Tensor) -> torch.Tensor: + """Rotation vector (..., 3) -> matrix, smooth at zero.""" + theta2 = (v * v).sum(-1, keepdim=True) + theta = torch.sqrt(theta2 + 1e-12) + small = theta2 < 1e-8 + a = torch.where(small, 1 - theta2 / 6, torch.sin(theta) / theta) + b = torch.where(small, 0.5 - theta2 / 24, (1 - torch.cos(theta)) / (theta2 + 1e-12)) + x, y, z = v.unbind(-1) + zero = torch.zeros_like(x) + K = torch.stack((zero, -z, y, z, zero, -x, -y, x, zero), -1).reshape(v.shape[:-1] + (3, 3)) + eye = torch.eye(3, device=v.device, dtype=v.dtype).expand(K.shape) + return eye + a[..., None] * K + b[..., None] * (K @ K) + + +def so3_log(R: torch.Tensor) -> torch.Tensor: + """Matrix to rotation vector with a finite zero-angle derivative for model fitting.""" + return axis_angle_from_quat(quat_from_matrix(R)) + + +def rot_angle_sq(Ra: torch.Tensor, Rb: torch.Tensor) -> torch.Tensor: + """Squared geodesic angle between rotations, differentiable (via acos of trace, clamped).""" + tr = (Ra * Rb).sum((-1, -2)) # trace(Ra^T Rb) + cos = torch.clamp((tr - 1) / 2, -1 + 1e-6, 1 - 1e-6) + return torch.acos(cos) ** 2 + + +def yaw_mat(R: torch.Tensor) -> torch.Tensor: + """Heading (yaw-only) rotation of R, from the projected x axis.""" + fx, fy = R[..., 0, 0], R[..., 1, 0] + n = torch.sqrt(fx * fx + fy * fy + 1e-9) + c, s = fx / n, fy / n + o, one = torch.zeros_like(c), torch.ones_like(c) + return torch.stack((c, -s, o, s, c, o, o, o, one), -1).reshape(R.shape[:-2] + (3, 3)) + + +def mat_to_6d(R: torch.Tensor) -> torch.Tensor: + """First two columns, flattened (6D continuous rotation representation).""" + return R[..., :, :2].reshape(R.shape[:-2] + (6,)) + + +class G1Kinematics: + def __init__(self, mjcf_path: str, device: str, tracked: list[str] = TRACKED_BODIES): + m = mujoco.MjModel.from_xml_path(mjcf_path) + self.device = device + jname_to_idx = {n: i for i, n in enumerate(G1_JOINT_NAMES)} + # moving bodies in tree order (mujoco orders parents before children) + bodies, parent, joint_idx, axes, offs_p, offs_R = [], [], [], [], [], [] + mj_to_local = {} + for b in range(1, m.nbody): + name = m.body(b).name + if m.body_jntnum[b] == 0: + continue # fixed links are merged visually only; none are tracked + jid = m.body_jntadr[b] + mj_to_local[b] = len(bodies) + bodies.append(name) + if m.jnt_type[jid] == mujoco.mjtJoint.mjJNT_FREE: + parent.append(-1) + joint_idx.append(-1) + else: + pb = m.body_parentid[b] + if pb not in mj_to_local: + raise RuntimeError(f"fixed parent for {name}") + parent.append(mj_to_local[pb]) + joint_idx.append(jname_to_idx[m.joint(jid).name]) + axes.append(m.jnt_axis[jid].copy()) + offs_p.append(m.body_pos[b].copy()) + offs_R.append(m.body_quat[b].copy()) + self.body_names = bodies + self.parent = parent + self.joint_idx = joint_idx + self.axes = torch.tensor(np.array(axes), dtype=torch.float32, device=device) + self.off_p = torch.tensor(np.array(offs_p), dtype=torch.float32, device=device) + self.off_R = matrix_from_quat( + torch.tensor(np.array(offs_R)[:, [1, 2, 3, 0]], dtype=torch.float32, device=device) + ) + self.tracked_idx = [bodies.index(n) for n in tracked] + self.num_bodies = len(bodies) + # joint limits (in G1_JOINT_NAMES order) + lo, hi = np.zeros(29), np.zeros(29) + for j in range(m.njnt): + n = m.joint(j).name + if n in jname_to_idx: + lo[jname_to_idx[n]], hi[jname_to_idx[n]] = m.jnt_range[j] + self.q_lo = torch.tensor(lo, dtype=torch.float32, device=device) + self.q_hi = torch.tensor(hi, dtype=torch.float32, device=device) + + def forward(self, p, R, v, w, q, qd, all_bodies: bool = False): + """Tracked-body world positions, rotations, linear (link-origin) and angular velocities. + + Args: + p, R, v, w: root (pelvis) position (B,3), rotation (B,3,3), world lin/ang velocity (B,3). + q, qd: joint positions/velocities (B,29) in G1_JOINT_NAMES order. + """ + P, Rs, V, W = ( + [None] * self.num_bodies, + [None] * self.num_bodies, + [None] * self.num_bodies, + [None] * self.num_bodies, + ) + for i in range(self.num_bodies): + par = self.parent[i] + if par < 0: + P[i], Rs[i], V[i], W[i] = p, R, v, w + continue + j = self.joint_idx[i] + Rp = Rs[par] + pi = P[par] + (Rp @ self.off_p[i].unsqueeze(-1)).squeeze(-1) + R_off = Rp @ self.off_R[i] + Ri = R_off @ matrix_from_quat(quat_from_angle_axis(q[..., j], self.axes[i])) + axis_w = (R_off @ self.axes[i].unsqueeze(-1)).squeeze(-1) + wi = W[par] + axis_w * qd[..., j : j + 1] + vi = V[par] + torch.cross(W[par], pi - P[par], dim=-1) + P[i], Rs[i], V[i], W[i] = pi, Ri, vi, wi + idx = range(self.num_bodies) if all_bodies else self.tracked_idx + return ( + torch.stack([P[i] for i in idx], -2), + torch.stack([Rs[i] for i in idx], -3), + torch.stack([V[i] for i in idx], -2), + torch.stack([W[i] for i in idx], -2), + ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/g1_tracking/mdp.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/g1_tracking/mdp.py new file mode 100644 index 000000000000..bedef9e29510 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/g1_tracking/mdp.py @@ -0,0 +1,278 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Portions derived from SuperMimic (MIT), Copyright (c) 2026 ETH Zurich, Chong Zhang. +# See docs/licenses/dependencies/supermimic-license.txt. + +"""Native manager terms for reference tracking and physical transition capture.""" + +from __future__ import annotations + +import torch +import warp as wp +from newton import ModelFlags + +from isaaclab.envs.mdp.actions import JointPositionAction +from isaaclab.managers import CommandTerm, RecorderTerm +from isaaclab.utils.math import matrix_from_quat, quat_from_matrix + +from .kinematics import G1_JOINT_NAMES, G1Kinematics, mat_to_6d, so3_exp +from .motion import BIN_FRAMES, MotionLib +from .tracking import MAX_REWARD_RATE, STATE_DIM, TrackingObjective + + +class MotionCommand(CommandTerm): + """Own reference sampling, physical context and the task's analytical objectives.""" + + def __init__(self, cfg, env): + super().__init__(cfg, env) + self.robot = env.scene["robot"] + if list(self.robot.joint_names) != G1_JOINT_NAMES: + raise ValueError(f"The tracking asset must preserve the motion joint order: {self.robot.joint_names}") + self.kin = G1Kinematics(cfg.kinematics_path, self.device) + self.motion = MotionLib(cfg.motion_dir, self.kin, self.device, cfg.source_fps) + self.mdp = TrackingObjective( + self.kin, self.motion, self.robot.data.default_joint_pos.torch[0].clone(), cfg.objective + ) + self.physics = env.sim.backend_manager + self.robot_view = self.robot.root_view + self.torso_id = self.robot.find_bodies("torso_link")[0][0] + self.env_ids = torch.arange(self.num_envs, device=self.device) + self.xi = torch.zeros(self.num_envs, 33, device=self.device) + self.start_frame = torch.zeros(self.num_envs, dtype=torch.long, device=self.device) + self.start_bin = torch.zeros_like(self.start_frame) + self._accounted_step = torch.zeros_like(self.start_frame) + self.episode_return = torch.zeros(self.num_envs, device=self.device) + self.discounted_return = torch.zeros_like(self.episode_return) + self.bin_return = torch.zeros(self.motion.num_bins, device=self.device) + self.bin_weight = torch.full_like(self.bin_return, 1.0 / self.motion.num_bins) + self._command = torch.zeros(self.num_envs, 73 * cfg.goal_frames, device=self.device) + self.metrics = { + "return": torch.zeros_like(self.episode_return), + "length": torch.zeros_like(self.episode_return), + } + + @property + def command(self): + return self._command + + @property + def frame(self): + return self.start_frame + self._env.episode_length_buf + + def state(self): + data = self.robot.data + pose, velocity = data.root_link_pose_w.torch, data.root_link_vel_w.torch + return ( + pose[:, :3] - self._env.scene.env_origins, + matrix_from_quat(pose[:, 3:7]), + velocity[:, :3], + velocity[:, 3:], + data.joint_pos.torch, + data.joint_vel.torch, + ) + + def _account_reward(self, env_ids): + steps = self._env.episode_length_buf[env_ids] + pending = steps > self._accounted_step[env_ids] + if not hasattr(self._env, "reward_buf"): + return pending + reward = torch.where(pending, self._env.reward_buf[env_ids], 0.0) + self.episode_return[env_ids] += reward + self.discounted_return[env_ids] += self.cfg.gamma ** (steps.float() - 1) * reward + self._accounted_step[env_ids] = steps + return pending + + def reset(self, env_ids=None): + ids = self.env_ids if env_ids is None else self.env_ids[env_ids] + pending = self._account_reward(ids) + # Public reset after command.compute discards a partial episode without updating its bin twice. + if hasattr(self._env, "reset_terminated") and not self.cfg.evaluation: + boundary = pending & (self._env.reset_terminated[ids] | self._env.reset_time_outs[ids]) + ended = ids[boundary] + if len(ended): + steps = self._env.episode_length_buf[ended].float() + bootstrap = torch.where( + self._env.reset_terminated[ended], + 0.0, + self.cfg.gamma**steps * (self.episode_return[ended] / steps) / (1 - self.cfg.gamma), + ) + returns = self.discounted_return[ended] + bootstrap + sums, counts = torch.zeros_like(self.bin_return), torch.zeros_like(self.bin_return) + sums.index_add_(0, self.start_bin[ended], returns) + counts.index_add_(0, self.start_bin[ended], torch.ones_like(returns)) + self.bin_return.copy_( + torch.where(counts > 0, 0.8 * self.bin_return + 0.2 * sums / counts.clamp_min(1), self.bin_return) + ) + scale = MAX_REWARD_RATE * self._env.step_dt * self._env.reward_manager.get_term_cfg("tracking").weight + self.metrics["return"][ids] = self.episode_return[ids] / ( + self._env.episode_length_buf[ids].clamp_min(1) * scale + ) + self.metrics["length"][ids] = self._env.episode_length_buf[ids].float() + return super().reset(ids) + + def _update_metrics(self): + self._account_reward(self.env_ids) + + def _update_command(self): + self._command.copy_(self.motion.goal(self.frame, self.cfg.goal_frames)) + + def _resample_command(self, env_ids): + ids = self.env_ids[env_ids] + n, motion = len(ids), self.motion + if self.cfg.evaluation: + frames = motion.clip_start[(ids + self.cfg.evaluation_clip_offset) % len(motion.names)] + bins = motion.frame_bin[frames] + else: + maximum_return = ( + MAX_REWARD_RATE + * self._env.step_dt + * self._env.reward_manager.get_term_cfg("tracking").weight + / (1 - self.cfg.gamma) + ) + soft = torch.softmax(-25.0 * self.bin_return / maximum_return, 0) + self.bin_weight.mul_(0.999).add_(0.001 * soft) + probabilities = 0.5 / motion.num_bins + 0.4 * soft + 0.1 * self.bin_weight / self.bin_weight.sum() + bins = torch.multinomial(probabilities, n, replacement=True) + start = motion.bin_start[bins] + end = torch.minimum(start + BIN_FRAMES, motion.clip_end[motion.bin_clip[bins]]) + frames = start + (torch.rand(n, device=self.device) * (end - start).float()).long() + + p, R, v, w, q, qd = (value.clone() for value in motion.ref_state(frames)) + if self.cfg.reset_noise and not self.cfg.evaluation: + p += torch.empty_like(p).uniform_(-1, 1) * p.new_tensor((0.05, 0.05, 0.01)) + rotation = torch.empty_like(p).uniform_(-1, 1) * p.new_tensor((0.1, 0.1, 0.2)) + R = so3_exp(rotation) @ R + v += torch.empty_like(v).uniform_(-1, 1) * v.new_tensor((0.5, 0.5, 0.2)) + w += torch.empty_like(w).uniform_(-1, 1) * w.new_tensor((0.52, 0.52, 0.78)) + q += torch.empty_like(q).uniform_(-0.1, 0.1) + q = q.clamp(self.kin.q_lo, self.kin.q_hi) + pose = torch.cat((p + self._env.scene.env_origins[ids], quat_from_matrix(R)), -1) + self.robot.write_root_link_pose_to_sim_index(root_pose=pose, env_ids=ids) + self.robot.write_root_link_velocity_to_sim_index(root_velocity=torch.cat((v, w), -1), env_ids=ids) + self.robot.write_joint_state_to_sim_index(position=q, velocity=qd, env_ids=ids) + self.robot.actuators.target_command.set_position_index(value=q, env_ids=ids) + self.start_frame[ids], self.start_bin[ids] = frames, bins + self._accounted_step[ids] = 0 + self.episode_return[ids] = 0 + self.discounted_return[ids] = 0 + if not self.cfg.evaluation: + self.xi[ids, 1:30] = torch.empty(n, 29, device=self.device).uniform_(-0.01, 0.01) + # The native environment clears episode_length_buf after manager resets. + self._command[ids] = motion.goal(frames, self.cfg.goal_frames) + + +class ContextJointPositionAction(JointPositionAction): + def __init__(self, cfg, env): + super().__init__(cfg, env) + self.motion = env.command_manager.get_term("motion") + + def process_actions(self, actions): + super().process_actions(actions) + self._processed_actions.add_(self.motion.xi[:, 1:30]) + + +def randomize_context(env, env_ids): + command = env.command_manager.get_term("motion") + if command.cfg.evaluation: + command.xi[:, 0] = 1.0 + else: + command.xi[:, 0].uniform_(0.3, 1.2) + command.xi[:, 30:33] = torch.empty(env.num_envs, 3, device=env.device).uniform_(-1, 1) * command.xi.new_tensor( + (0.025, 0.05, 0.05) + ) + default_com = command.robot.data.body_com_pos_b.torch[:, command.torso_id].clone() + command.robot.set_coms_index( + coms=(default_com + command.xi[:, 30:33]).unsqueeze(1), body_ids=[command.torso_id] + ) + model = command.physics.get_model() + shape = command.robot_view.get_attribute("shape_material_mu", model).shape + values = command.xi[:, 0].reshape(env.num_envs, 1, 1).expand(shape).contiguous() + command.robot_view.set_attribute("shape_material_mu", model, wp.from_torch(values)) + command.physics.add_model_change(ModelFlags.SHAPE_PROPERTIES) + + +def joint_pos_rel(env): + command = env.command_manager.get_term("motion") + return command.robot.data.joint_pos.torch - command.mdp.default_q - command.xi[:, 1:30] + + +def relative_orientation(env): + command = env.command_manager.get_term("motion") + R = matrix_from_quat(command.robot.data.root_link_pose_w.torch[:, 3:7]) + return mat_to_6d(R.transpose(-1, -2) @ command.motion.root_R[command.frame]) + + +def tracking_reward(env): + command = env.command_manager.get_term("motion") + return command.mdp.reward(command.state(), command.frame, env.action_manager.action, env.action_manager.prev_action) + + +def tracking_failure(env): + command = env.command_manager.get_term("motion") + return command.mdp.termination(command.state(), command.frame) + + +def clip_end(env): + command = env.command_manager.get_term("motion") + return command.frame >= command.motion.clip_end[command.motion.clip_id[command.start_frame]] + + +class TransitionRecorder(RecorderTerm): + """Publish actual pre/post tensors until the next step; optionally retain them as episodes. + + An exported episode's end is also a boundary, including explicit resets that happen after a + transition was retained. Online consumers additionally check adjacent state/step continuity. + """ + + requires_post_step_observations = False + + def __init__(self, cfg, env): + super().__init__(cfg, env) + self.motion = env.command_manager.get_term("motion") + self.robot = self.motion.robot + self.latest = { + "s": torch.empty(self.num_envs, STATE_DIM, device=self.device), + "s1": torch.empty(self.num_envs, STATE_DIM, device=self.device), + "a": torch.empty(self.num_envs, 29, device=self.device), + "a_prev": torch.empty(self.num_envs, 29, device=self.device), + "xi": torch.empty_like(self.motion.xi), + "f": torch.empty(self.num_envs, dtype=torch.long, device=self.device), + "ep_step": torch.empty(self.num_envs, dtype=torch.int32, device=self.device), + "r": torch.empty(self.num_envs, device=self.device), + "done": torch.empty(self.num_envs, dtype=torch.bool, device=self.device), + "term": torch.empty(self.num_envs, dtype=torch.bool, device=self.device), + } + self._has_transition = False + + def _capture_state(self, target): + data = self.robot.data + target[:, :7].copy_(data.root_link_pose_w.torch) + target[:, :3].sub_(self._env.scene.env_origins) + target[:, 7:13].copy_(data.root_link_vel_w.torch) + target[:, 13:42].copy_(data.joint_pos.torch) + target[:, 42:71].copy_(data.joint_vel.torch) + + def record_pre_step(self): + self._capture_state(self.latest["s"]) + self.latest["a"].copy_(self._env.action_manager.action) + self.latest["a_prev"].copy_(self._env.action_manager.prev_action) + self.latest["xi"].copy_(self.motion.xi) + self.latest["f"].copy_(self.motion.frame) + self.latest["ep_step"].copy_(self._env.episode_length_buf) + self._has_transition = True + return None, None + + def record_post_step(self): + self._capture_state(self.latest["s1"]) + self.latest["r"].copy_(self._env.reward_buf) + self.latest["done"].copy_(self._env.reset_buf) + self.latest["term"].copy_(self._env.reset_terminated) + return ("transition", self.latest) if self.cfg.record_to_episodes else (None, None) + + def record_pre_reset(self, env_ids): + if self._has_transition: + self.latest["done"][slice(None) if env_ids is None else env_ids] = True + return None, None diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/g1_tracking/motion.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/g1_tracking/motion.py new file mode 100644 index 000000000000..eec9b12a6d95 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/g1_tracking/motion.py @@ -0,0 +1,165 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Portions derived from SuperMimic (MIT), Copyright (c) 2026 ETH Zurich, Chong Zhang. +# See docs/licenses/dependencies/supermimic-license.txt. + +"""G1 reference motions resampled to 50 Hz, with bounded-memory FK construction.""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import torch +import torch.nn.functional as F + +from isaaclab.utils.math import matrix_from_quat + +from .kinematics import G1Kinematics, mat_to_6d, so3_log, yaw_mat + +CONTROL_DT = 0.02 +BIN_FRAMES = 50 # 1 s bins for return-based reset sampling + + +def resample_motion(raw: torch.Tensor, src_fps: float): + """Resample metric/xyzw/radian G1 poses to the control rate using shortest-arc SLERP.""" + p, quat, j = raw[:, :3], raw[:, 3:7], raw[:, 7:36] + # resample to 50 Hz + T_src = raw.shape[0] + dur = (T_src - 1) / src_fps + n = int(np.floor(dur / CONTROL_DT)) + 1 + if n < 2: + raise ValueError("Motion is shorter than one control transition.") + t = torch.arange(n, dtype=torch.float32) * CONTROL_DT * src_fps + i0 = torch.clamp(t.floor().long(), max=T_src - 2) + a = (t - i0.float()).clamp(0, 1) + p = p[i0] * (1 - a[:, None]) + p[i0 + 1] * a[:, None] + j = j[i0] * (1 - a[:, None]) + j[i0 + 1] * a[:, None] + # Batched shortest-arc SLERP; the native helper accepts a single interpolation pair. + q0, q1 = quat[i0], quat[i0 + 1] + dot = (q0 * q1).sum(-1, keepdim=True) + q1 = torch.where(dot < 0, -q1, q1) + theta = torch.acos(dot.abs().clamp(max=1.0)) + sine, weight = torch.sin(theta), a[:, None] + w0 = torch.where(sine < 1e-6, 1 - weight, torch.sin((1 - weight) * theta) / (sine + 1e-12)) + w1 = torch.where(sine < 1e-6, weight, torch.sin(weight * theta) / (sine + 1e-12)) + quat = F.normalize(w0 * q0 + w1 * q1, dim=-1) + return p, quat, j + + +class MotionLib: + def __init__(self, path: str, kin: G1Kinematics, device: str, src_fps: float = 30.0): + if not np.isfinite(src_fps) or src_fps <= 0: + raise ValueError("Source frame rate must be finite and positive.") + files = sorted((*Path(path).rglob("*.csv"), *Path(path).rglob("*.npy"))) + if not files: + raise ValueError(f"No motion CSV/NumPy files found in {path}.") + self.device = device + self.names = [f.relative_to(path).with_suffix("").as_posix() for f in files] + if len(set(self.names)) != len(self.names): + raise ValueError("Motion names must be unique across CSV and NumPy files.") + ps, qs, js = [], [], [] + lens = [] + for f in files: + raw = torch.tensor( + np.load(f, allow_pickle=False) if f.suffix == ".npy" else np.loadtxt(f, delimiter=","), + dtype=torch.float32, + ) + if raw.ndim != 2 or raw.shape[0] < 2 or raw.shape[1] != 36: + raise ValueError(f"Motion {f} must contain at least two 36-channel frames.") + if not torch.isfinite(raw).all() or (raw[:, 3:7].norm(dim=-1) < 1e-6).any(): + raise ValueError(f"Motion {f} contains nonfinite values or invalid quaternions.") + if src_fps == 1.0 / CONTROL_DT: + p, quat, j = raw[:, :3], F.normalize(raw[:, 3:7], dim=-1), raw[:, 7:] + else: + p, quat, j = resample_motion(raw, src_fps) + n = len(p) + ps.append(p), qs.append(quat), js.append(j) + lens.append(n) + self.clip_len = torch.tensor(lens, device=device) + self.clip_start = torch.cumsum( + torch.cat((torch.zeros(1, dtype=torch.long), torch.tensor(lens[:-1], dtype=torch.long))), 0 + ).to(device) # dtype: one clip gives an empty (float) tensor + self.clip_end = self.clip_start + self.clip_len - 1 # last valid frame (inclusive) + self.num_frames = int(sum(lens)) + self.clip_id = torch.cat([torch.full((n,), i, dtype=torch.long) for i, n in enumerate(lens)]).to(device) + + root_p, root_R, root_v, root_w, jq, jqd = [], [], [], [], [], [] + for p, quat, j in zip(ps, qs, js): + R = matrix_from_quat(quat) + v = torch.gradient(p, spacing=CONTROL_DT, dim=0)[0] + qd = torch.gradient(j, spacing=CONTROL_DT, dim=0)[0] + # angular velocity (world) by central differences of rotations + w = torch.zeros_like(p) + w[1:-1] = so3_log(R[2:] @ R[:-2].transpose(-1, -2)) / (2 * CONTROL_DT) + w[0] = so3_log(R[1] @ R[0].T) / CONTROL_DT + w[-1] = so3_log(R[-1] @ R[-2].T) / CONTROL_DT + root_p.append(p), root_R.append(R), root_v.append(v), root_w.append(w), jq.append(j), jqd.append(qd) + cat = lambda xs: torch.cat(xs).to(device) # noqa: E731 + self.root_p, self.root_R, self.root_v, self.root_w = cat(root_p), cat(root_R), cat(root_v), cat(root_w) + self.q, self.qd = cat(jq), cat(jqd) + with torch.no_grad(): + for s in range(0, self.num_frames, 65536): + e = min(s + 65536, self.num_frames) + out = kin.forward( + self.root_p[s:e], self.root_R[s:e], self.root_v[s:e], self.root_w[s:e], self.q[s:e], self.qd[s:e] + ) + if s == 0: + self.body_p, self.body_R, self.body_v, self.body_w = ( + value.new_empty((self.num_frames, *value.shape[1:])) for value in out + ) + for destination, value in zip((self.body_p, self.body_R, self.body_v, self.body_w), out, strict=True): + destination[s:e].copy_(value) + # per-frame static goal parts: q, qd, root velocities in the root frame + Rt = self.root_R.transpose(-1, -2) + self.root_v_local = (Rt @ self.root_v.unsqueeze(-1)).squeeze(-1) + self.root_w_local = (Rt @ self.root_w.unsqueeze(-1)).squeeze(-1) + self.root_gz = self.root_R[:, 2, 2] # -(projected gravity z) of the reference root: g_b = R^T (0,0,-1) + # 1 s bins + bin_clip, bin_start, frame_bins = [], [], [] + frame_offset = 0 + for c in range(len(lens)): + nb = (lens[c] - 2) // BIN_FRAMES + 1 + frame_bins.append((torch.arange(lens[c]) // BIN_FRAMES).clamp_max(nb - 1) + len(bin_start)) + for b in range(nb): + bin_clip.append(c) + bin_start.append(frame_offset + b * BIN_FRAMES) + frame_offset += lens[c] + self.bin_clip = torch.tensor(bin_clip, device=device) + self.bin_start = torch.tensor(bin_start, device=device) + self.num_bins = len(bin_start) + self.frame_bin = torch.cat(frame_bins).to(device) + + def future_frames(self, f: torch.Tensor, K: int) -> torch.Tensor: + """Frames f+1..f+K clamped to the clip's last frame, shape (B, K).""" + end = self.clip_end[self.clip_id[f]] + fut = f.unsqueeze(-1) + torch.arange(1, K + 1, device=f.device) + return torch.minimum(fut, end.unsqueeze(-1)) + + def goal(self, f: torch.Tensor, K: int) -> torch.Tensor: + """Goal features from the reference only: K future frames relative to the current reference frame.""" + fut = self.future_frames(f, K) # (B,K) + # Beyond the recorded clip, the padded pose is held still rather than continuing at its endpoint velocity. + recorded = (fut == f.unsqueeze(-1) + torch.arange(1, K + 1, device=f.device)).unsqueeze(-1) + H = yaw_mat(self.root_R[f]) # (B,3,3) + Ht = H.transpose(-1, -2).unsqueeze(1) + dp = (Ht @ (self.root_p[fut] - self.root_p[f].unsqueeze(1)).unsqueeze(-1)).squeeze(-1) + rot = mat_to_6d(Ht @ self.root_R[fut]) + feats = torch.cat( + ( + dp, + rot, + self.root_v_local[fut] * recorded, + self.root_w_local[fut] * recorded, + self.q[fut], + self.qd[fut] * (0.1 * recorded), + ), + -1, + ) # (B,K,73) + return feats.reshape(f.shape[0], -1) + + def ref_state(self, f: torch.Tensor): + return self.root_p[f], self.root_R[f], self.root_v[f], self.root_w[f], self.q[f], self.qd[f] diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/g1_tracking/tracking.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/g1_tracking/tracking.py new file mode 100644 index 000000000000..859fd256b95d --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/g1_tracking/tracking.py @@ -0,0 +1,128 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Portions derived from SuperMimic (MIT), Copyright (c) 2026 ETH Zurich, Chong Zhang. +# See docs/licenses/dependencies/supermimic-license.txt. + +"""Analytic MDP heads shared by the simulator and the world model (paper Appendix A). + +State tuple ``S = (p, R, v, w, q, qd)``: pelvis position, rotation matrix, world linear/angular +velocity of the pelvis link frame, joint positions and velocities. +""" + +from __future__ import annotations + +import torch + +from isaaclab.utils import configclass +from isaaclab.utils.math import matrix_from_quat + +from .kinematics import G1Kinematics, mat_to_6d, rot_angle_sq, yaw_mat +from .motion import MotionLib + +STATE_DIM = 71 +OBS_DIM = 29 + 29 + 3 + 3 + 29 + 6 +XI_DIM = 33 +NUM_J = 29 + +MAX_REWARD_RATE = 5.0 + + +@configclass +class RewardCfg: + body_pos_std: float = 0.3 + body_ori_std: float = 0.4 + body_linvel_std: float = 1.0 + body_angvel_std: float = 3.14 + root_pos_std: float = 0.3 + root_ori_std: float = 0.4 + w_action_rate: float = -0.1 + w_joint_limit: float = -2.0 + soft_limit_factor: float = 0.9 + term_height: float = 0.25 + term_tilt: float = 0.8 + + +def unflatten_state(x: torch.Tensor): + return x[..., 0:3], matrix_from_quat(x[..., 3:7]), x[..., 7:10], x[..., 10:13], x[..., 13:42], x[..., 42:71] + + +class TrackingObjective: + def __init__(self, kin: G1Kinematics, motion: MotionLib, default_q: torch.Tensor, cfg: RewardCfg | None = None): + self.kin, self.motion, self.cfg = kin, motion, cfg if cfg is not None else RewardCfg() + cfg = self.cfg + self.default_q = default_q + mid, half = (kin.q_lo + kin.q_hi) / 2, (kin.q_hi - kin.q_lo) / 2 + self.soft_lo, self.soft_hi = mid - half * cfg.soft_limit_factor, mid + half * cfg.soft_limit_factor + + # ---------------- observation ---------------- + def obs(self, S, a_prev, xi, f) -> torch.Tensor: + """Noise-free proprioceptive observation (+ reference-relative root orientation).""" + p, R, v, w, q, qd = S + Rt = R.transpose(-1, -2) + dq0 = xi[..., 1:30] + g_b = -R[..., 2, :] # R^T (0,0,-1) + rel = mat_to_6d(Rt @ self.motion.root_R[f]) + return torch.cat((q - self.default_q - dq0, qd, (Rt @ w.unsqueeze(-1)).squeeze(-1), g_b, a_prev, rel), -1) + + # ---------------- reward / termination ---------------- + def reward(self, S, f, a, a_prev, terms: bool = False): + """Reward rate comparing S with reference frame f; the caller applies its step duration.""" + c, m = self.cfg, self.motion + p, R, v, w, q, qd = S + bp, bR, bv, bw = self.kin.forward(p, R, v, w, q, qd) + rp, rR = m.root_p[f], m.root_R[f] + # reference bodies re-anchored to the robot's root (xy + yaw), as in BeyondMimic + dR = yaw_mat(R @ rR.transpose(-1, -2)) + anchor = torch.cat((p[..., :2], rp[..., 2:3]), -1) + ref_bp = anchor.unsqueeze(-2) + ( + (dR.unsqueeze(-3) @ (m.body_p[f] - rp.unsqueeze(-2)).unsqueeze(-1)).squeeze(-1) + ) + ref_bR = dR.unsqueeze(-3) @ m.body_R[f] + e_pos = ((ref_bp - bp) ** 2).sum(-1).mean(-1) + e_ori = rot_angle_sq(ref_bR, bR).mean(-1) + e_lv = ((m.body_v[f] - bv) ** 2).sum(-1).mean(-1) + e_av = ((m.body_w[f] - bw) ** 2).sum(-1).mean(-1) + e_rp = ((rp - p) ** 2).sum(-1) + e_ro = rot_angle_sq(rR, R) + r_track = ( + torch.exp(-e_pos / c.body_pos_std**2) + + torch.exp(-e_ori / c.body_ori_std**2) + + torch.exp(-e_lv / c.body_linvel_std**2) + + torch.exp(-e_av / c.body_angvel_std**2) + + 0.5 * torch.exp(-e_rp / c.root_pos_std**2) + + 0.5 * torch.exp(-e_ro / c.root_ori_std**2) + ) + r_rate = ((a - a_prev) ** 2).sum(-1) + r_lim = (torch.relu(q - self.soft_hi) + torch.relu(self.soft_lo - q)).sum(-1) + r = r_track + c.w_action_rate * r_rate + c.w_joint_limit * r_lim + if terms: + return r, { + "track": r_track, + "e_body_pos": e_pos, + "e_body_ori": e_ori, + "e_root_pos": e_rp, + "action_rate": r_rate, + "joint_limit": r_lim, + } + return r + + def termination(self, S, f) -> torch.Tensor: + p, R = S[0], S[1] + rp, rR = self.motion.root_p[f], self.motion.root_R[f] + e_z = (rp[..., 2] - p[..., 2]).abs() + e_tilt = (rR[..., 2, 2] - R[..., 2, 2]).abs() + return (e_z > self.cfg.term_height) | (e_tilt > self.cfg.term_tilt) + + def soft_termination(self, S, f, tz: float = 0.02, tt: float = 0.04) -> torch.Tensor: + """Differentiable termination probability for imagined rollouts.""" + p, R = S[0], S[1] + rp, rR = self.motion.root_p[f], self.motion.root_R[f] + e_z = (rp[..., 2] - p[..., 2]).abs() + e_tilt = (rR[..., 2, 2] - R[..., 2, 2]).abs() + alive = (1 - torch.sigmoid((e_z - self.cfg.term_height) / tz)) * ( + 1 - torch.sigmoid((e_tilt - self.cfg.term_tilt) / tt) + ) + return 1 - alive diff --git a/source/isaaclab_tasks/test/test_g1_tracking.py b/source/isaaclab_tasks/test/test_g1_tracking.py new file mode 100644 index 000000000000..2d9072c5e254 --- /dev/null +++ b/source/isaaclab_tasks/test/test_g1_tracking.py @@ -0,0 +1,212 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Portions derived from SuperMimic (MIT), Copyright (c) 2026 ETH Zurich, Chong Zhang. +# See docs/licenses/dependencies/supermimic-license.txt. + +"""G1 capture and motion contracts, including stationary padding beyond recorded clip ends.""" + +import os +from pathlib import Path +from types import SimpleNamespace + +import gymnasium as gym +import h5py +import mujoco +import numpy as np +import pytest +import torch +import warp as wp + +from isaaclab.app import launch_simulation +from isaaclab.managers import DatasetExportMode, RewardTermCfg, TerminationTermCfg +from isaaclab.test.utils import DeviceScope, test_devices +from isaaclab.utils.math import matrix_from_quat +from isaaclab.utils.string import resolve_matching_names_values + +from isaaclab_tasks.contrib.g1_tracking.g1_tracking_env_cfg import TransitionRecorderCfg +from isaaclab_tasks.contrib.g1_tracking.kinematics import TRACKED_BODIES +from isaaclab_tasks.contrib.g1_tracking.motion import MotionLib +from isaaclab_tasks.utils import resolve_task_config + + +def test_motion_rates_and_stationary_endpoint_padding(tmp_path): + raw = np.zeros((121, 36), dtype=np.float32) + raw[:, 2], raw[:, 6] = 0.8, 1.0 + raw[:, 0] = np.linspace(0, 0.2, len(raw)) + roll = np.linspace(0, 0.2, len(raw)) + raw[:, 3], raw[:, 6] = np.sin(roll / 2), np.cos(roll / 2) + raw[:, 7] = np.linspace(0, 0.1, len(raw)) + for directory in ("a", "b"): + (tmp_path / directory).mkdir() + np.savetxt(tmp_path / directory / "walk.csv", raw, delimiter=",") + # Kinematics is an independent dependency here; this test checks temporal pose/rate semantics. + kin = SimpleNamespace(forward=lambda p, R, v, w, q, qd: (p[:, None], R[:, None], v[:, None], w[:, None])) + motion = MotionLib(str(tmp_path), kin, "cpu", src_fps=120) + assert motion.names == ["a/walk", "b/walk"] + assert motion.clip_len.tolist() == [51, 51] and motion.num_bins == 2 + torch.testing.assert_close(motion.root_v, torch.tensor([0.2, 0, 0]).expand(102, -1), atol=2e-6, rtol=1e-5) + assert motion.future_frames(torch.tensor([49, 100]), 3).tolist() == [[50, 50, 50], [101, 101, 101]] + goals = motion.goal(torch.tensor([47, 49, 50, 100]), 3).reshape(4, 3, 73) + # Recorded endpoint rates remain factual; only slots beyond the clip describe a stationary held pose. + expected_velocity = torch.tensor([[0.2, 0.2, 0.2], [0.2, 0.0, 0.0], [0.0, 0.0, 0.0], [0.2, 0.0, 0.0]]) + torch.testing.assert_close(goals[:, :, 9], expected_velocity, atol=2e-6, rtol=1e-5) + torch.testing.assert_close(goals[:, :, 12], expected_velocity, atol=1e-5, rtol=1e-5) + torch.testing.assert_close(goals[:, :, 44], expected_velocity * 0.05, atol=2e-6, rtol=1e-5) + torch.testing.assert_close(goals[1:, :, 15], torch.full((3, 3), 0.1), atol=2e-6, rtol=1e-5) + with pytest.raises(ValueError, match="frame rate"): + MotionLib(str(tmp_path), kin, "cpu", src_fps=0) + + +def never(env): + return torch.zeros(env.num_envs, dtype=torch.bool, device=env.device) + + +@pytest.fixture(scope="module", params=test_devices(DeviceScope.CUDA)) +def task(tmp_path_factory, request): + torch.backends.cuda.matmul.allow_tf32 = False + output = tmp_path_factory.mktemp("native_task") + cfg, _ = resolve_task_config( + "IsaacContrib-Motion-Tracking-G1", None, overrides=["env.scene.num_envs=40", "env.seed=123"] + ) + cfg.sim.device = request.param + data_dir = os.environ.get("ISAACLAB_G1_TRACKING_DATA") + if not data_dir: + pytest.skip("Set ISAACLAB_G1_TRACKING_DATA to the prepared G1 tracking directory.") + cfg.scene.robot.spawn.usd_path = str(Path(data_dir) / "g1_usd/g1_29dof/g1_29dof.usda") + cfg.commands.motion.kinematics_path = str(Path(data_dir) / "g1_description/g1_29dof.xml") + cfg.commands.motion.motion_dir = str(Path(data_dir) / "lafan1_g1") + cfg.recorders.transition = TransitionRecorderCfg() + cfg.observations.policy.enable_corruption = False + cfg.observations.policy.history_length = 3 + cfg.commands.motion.goal_frames = 2 + cfg.commands.motion.reset_noise = False + cfg.recorders.transition.record_to_episodes = True + cfg.recorders.dataset_export_mode = DatasetExportMode.EXPORT_ALL + cfg.recorders.dataset_export_dir_path = str(output) + cfg.recorders.buffer_initial_capacity = 2 + for name in ("failed", "clip_end", "time_out"): + getattr(cfg.terminations, name).func = never + with launch_simulation(cfg): + env = gym.make("IsaacContrib-Motion-Tracking-G1", cfg=cfg).unwrapped + try: + yield env, output + finally: + env.close() + + +def test_real_capture_observations_context_and_offline_export(task): + env, output = task + command = env.command_manager.get_term("motion") + recorder = env.recorder_manager.get_term("transition") + observations, _ = env.reset() + assert env.cfg.sim.physics.use_cuda_graph + robot = command.robot + model = mujoco.MjModel.from_xml_path(command.cfg.kinematics_path) + expected_com = torch.tensor(model.body("torso_link").ipos, device=env.device, dtype=torch.float32) + torch.testing.assert_close( + robot.data.body_com_pos_b.torch[:, command.torso_id], expected_com + command.xi[:, 30:33] + ) + mu = wp.to_torch(command.robot_view.get_attribute("shape_material_mu", command.physics.get_model())) + torch.testing.assert_close(mu, command.xi[:, :1, None].expand_as(mu)) + previous_action = torch.zeros(env.num_envs, 29, device=env.device) + records = [] + for step in range(3): + state = torch.cat( + ( + robot.data.root_link_pose_w.torch, + robot.data.root_link_vel_w.torch, + robot.data.joint_pos.torch, + robot.data.joint_vel.torch, + ), + -1, + ).clone() + frame = command.frame.clone() + action = torch.full_like(previous_action, 0.01 * (step + 1)) + observations, reward, terminated, truncated, _ = env.step(action) + transition = recorder.latest + assert not (terminated | truncated).any() + torch.testing.assert_close(transition["s"], state, rtol=0, atol=0) + torch.testing.assert_close(transition["a_prev"], previous_action, rtol=0, atol=0) + torch.testing.assert_close(transition["a"], action, rtol=0, atol=0) + torch.testing.assert_close(transition["r"], reward, rtol=0, atol=0) + torch.testing.assert_close(transition["f"], frame, rtol=0, atol=0) + analytical = command.mdp.obs(command.state(), action, command.xi, command.frame) + torch.testing.assert_close(observations["policy"][:, -1], analytical, rtol=1e-5, atol=2e-6) + torch.testing.assert_close(observations["goal"], command.motion.goal(command.frame, 2)) + action_term = env.action_manager.get_term("joint_pos") + _, _, values = resolve_matching_names_values(env.cfg.actions.joint_pos.scale, robot.joint_names) + scale = torch.tensor(values, device=env.device) + expected_target = command.mdp.default_q + command.xi[:, 1:30] + scale * action + torch.testing.assert_close(action_term.processed_actions, expected_target) + records.append({key: value.clone() for key, value in transition.items()}) + previous_action = action + body_ids = [robot.body_names.index(name) for name in TRACKED_BODIES] + positions, rotations, velocities, angular = command.kin.forward(*command.state()) + torch.testing.assert_close(positions, robot.data.body_link_pos_w.torch[:, body_ids], atol=5e-5, rtol=1e-5) + torch.testing.assert_close( + rotations, matrix_from_quat(robot.data.body_link_quat_w.torch[:, body_ids]), atol=5e-5, rtol=1e-5 + ) + torch.testing.assert_close(velocities, robot.data.body_link_lin_vel_w.torch[:, body_ids], atol=1e-4, rtol=1e-4) + torch.testing.assert_close(angular, robot.data.body_link_ang_vel_w.torch[:, body_ids], atol=1e-4, rtol=1e-4) + snapshot = env.recorder_manager.get_episode(0).data["transition"] + for key in records[0]: + torch.testing.assert_close(snapshot[key], torch.stack([record[key][0] for record in records]), rtol=0, atol=0) + env.recorder_manager.export_episodes([0]) + with h5py.File(output / "dataset.hdf5", "r") as dataset: + saved = dataset["data/demo_0/transition"] + for key in snapshot: + torch.testing.assert_close(torch.from_numpy(saved[key][:]), snapshot[key].cpu(), rtol=0, atol=0) + + +def test_terminal_capture_and_return_accounting(task): + env, _ = task + command = env.command_manager.get_term("motion") + capture = env.recorder_manager.get_term("transition") + reward_cfg = env.reward_manager.get_term_cfg("tracking") + term_cfgs = {name: env.termination_manager.get_term_cfg(name) for name in ("failed", "clip_end", "time_out")} + gamma = command.cfg.gamma + try: + env.reward_manager.set_term_cfg( + "tracking", RewardTermCfg(func=lambda env: env.episode_length_buf / env.step_dt, weight=1.0) + ) + env.termination_manager.set_term_cfg( + "failed", + TerminationTermCfg( + func=lambda env: ( + (env.episode_length_buf == 3) & (torch.arange(env.num_envs, device=env.device) % 2 == 0) + ) + ), + ) + env.termination_manager.set_term_cfg( + "time_out", + TerminationTermCfg( + func=lambda env: ( + (env.episode_length_buf == 3) & (torch.arange(env.num_envs, device=env.device) % 2 == 1) + ), + time_out=True, + ), + ) + command.cfg.gamma = 0.9 + env.reset() + command.bin_return.zero_() + command.start_bin.zero_() + for _ in range(3): + env.step(torch.zeros(env.num_envs, 29, device=env.device)) + assert capture.latest["done"].all() + assert capture.latest["term"].sum() == env.num_envs // 2 + assert (capture.latest["ep_step"] == 2).all() + assert (env.episode_length_buf == 0).all() + assert (capture.latest["s1"][:, :3] != command.state()[0]).any() + # Discounted 1,2,3 = 5.23; timeout adds .9**3 * mean(1,2,3)/.1 = 14.58. + torch.testing.assert_close(command.bin_return[0], torch.tensor(0.2 * (5.23 + 19.81) / 2, device=env.device)) + before = command.bin_return.clone() + env.reset() + torch.testing.assert_close(command.bin_return, before, rtol=0, atol=0) + finally: + command.cfg.gamma = gamma + env.reward_manager.set_term_cfg("tracking", reward_cfg) + for name, cfg in term_cfgs.items(): + env.termination_manager.set_term_cfg(name, cfg) diff --git a/tools/environ_docs.py b/tools/environ_docs.py index b1c0141e48d3..858f39d4a2a9 100644 --- a/tools/environ_docs.py +++ b/tools/environ_docs.py @@ -58,7 +58,7 @@ ) # RL libraries listed in a stable order across generated docs. -_RL_LIBRARY_ORDER = ("rl_games", "rsl_rl", "skrl", "sb3", "rlinf", "torchrl") +_RL_LIBRARY_ORDER = ("rl_games", "rsl_rl", "skrl", "sb3", "rlinf", "supermimic", "torchrl") _ALGORITHM_LABELS = {"": "PPO", "ppo": "PPO", "amp": "AMP", "ippo": "IPPO", "mappo": "MAPPO"} # Gym IDs excluded from the training list. The ``-Eval`` suffix marks dedicated @@ -223,7 +223,9 @@ def parse_rl_libraries_from_kwargs(kwargs: dict) -> dict[str, list[str]]: agents: dict[str, set[str]] = collections.defaultdict(set) for value, library, algo_suffix in entries: - algorithm = _ALGORITHM_LABELS.get(algo_suffix.lower()) + algorithm = ( + "MBRL" if library == "supermimic" and not algo_suffix else _ALGORITHM_LABELS.get(algo_suffix.lower()) + ) if algorithm is None: # Registry suffixes can name recipes or preset variants, not algorithms. continue diff --git a/tools/test/test_environ_docs.py b/tools/test/test_environ_docs.py index 48a82681e63d..667b3951b9d6 100644 --- a/tools/test/test_environ_docs.py +++ b/tools/test/test_environ_docs.py @@ -109,11 +109,13 @@ def test_parse_rl_libraries_from_kwargs_handles_multi_agent_and_amp(): "skrl_mappo_cfg_entry_point": "agents:skrl_mappo_cfg.yaml", "skrl_amp_cfg_entry_point": "agents:skrl_amp_cfg.yaml", "rsl_rl_with_symmetry_cfg_entry_point": "agents.rsl_rl_ppo_cfg:RunnerCfg", + "supermimic_cfg_entry_point": "isaaclab_rl.supermimic.supermimic_cfg:SuperMimicCfg", "rl_games_cfg_entry_point_ignored": "nope", } agents = parse_rl_libraries_from_kwargs(kwargs) assert agents["rl_games"] == ["PPO"] assert agents["skrl"] == ["PPO", "AMP", "IPPO", "MAPPO"] + assert agents["supermimic"] == ["MBRL"] def test_parse_rl_libraries_does_not_label_canonical_amp_alias_as_ppo():