Integrate native SuperMimic training and G1 tracking

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::

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.

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"],
         ]],
     ];
 

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:

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 <https://github.com/zita-ch/supermimic_vibed>`__ 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 ``<IsaacLab>/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=<project> agent.wandb_entity=<entity>``.
+An optional stable ``agent.wandb_run_id=<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-<id>-checkpoint`` with the ``latest`` alias; the runner waits
+for upload before reporting ``durable_checkpoint_iteration``. Milestones also save
+``policy_<iteration>.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``.

scripts/imitation_learning/isaaclab_mimic/annotate_demos.py

@@ -515,15 +515,16 @@ 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,
     ) in subtask_term_signal_action_indices.items():
         # 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