Integrate native NeRD and SuperMimic with Cartpole, ANYmal C and G1

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

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

scripts/learning/nerd/record.py

@@ -0,0 +1,89 @@
+# 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
+
+"""Collect a registered task's native NeRD transitions into an HDF5 dataset."""
+
+import argparse
+import json
+from pathlib import Path
+
+import gymnasium as gym
+import torch
+
+from isaaclab.app import add_launcher_args, launch_simulation
+from isaaclab.managers import DatasetExportMode
+from isaaclab.utils.assets import read_file
+from isaaclab.utils.io import dump_yaml
+
+import isaaclab_tasks  # noqa: F401
+from isaaclab_tasks.utils import resolve_task_config, setup_preset_cli
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--task", required=True)
+    parser.add_argument("--output", type=Path, required=True)
+    parser.add_argument("--num_envs", type=int, default=1024)
+    parser.add_argument("--steps", type=int, default=1500)
+    parser.add_argument("--action_scale", type=float, default=1.5, help="Uniform normalized action bound.")
+    parser.add_argument(
+        "--policy", help="Exported feedforward TorchScript policy path or URL, consuming policy observations."
+    )
+    parser.add_argument("--action_hold", type=int, default=1, help="Environment steps between action updates.")
+    parser.add_argument("--seed", type=int, default=42)
+    add_launcher_args(parser)
+    args, overrides = setup_preset_cli(parser)
+    if min(args.num_envs, args.steps, args.action_hold) < 1 or args.action_scale <= 0:
+        parser.error("World/step counts, action hold and action scale must be positive.")
+    path = args.output.resolve()
+    if path.suffix != ".hdf5":
+        parser.error("--output must end in .hdf5.")
+    if path.exists() or path.with_suffix(".json").exists():
+        parser.error("The output dataset already exists; choose another path.")
+    cfg, _ = resolve_task_config(args.task, "", overrides=overrides)
+    if cfg.recorders.transition is None:
+        parser.error("Choose a task with a configured NeRD transition recorder, such as a -Explore task.")
+    cfg.scene.num_envs = args.num_envs
+    cfg.sim.device = args.device
+    cfg.seed = args.seed
+    cfg.recorders.transition.record_to_episodes = True
+    cfg.recorders.dataset_export_mode = DatasetExportMode.EXPORT_ALL
+    cfg.recorders.dataset_export_dir_path = str(path.parent)
+    cfg.recorders.dataset_filename = path.stem
+    cfg.recorders.export_in_close = True
+    path.parent.mkdir(parents=True, exist_ok=True)
+    dump_yaml(str(path.with_suffix(".yaml")), cfg)
+    torch.manual_seed(args.seed)
+    with launch_simulation(cfg, args):
+        env = gym.make(args.task, cfg=cfg).unwrapped
+        try:
+            observations, _ = env.reset()
+            policy = torch.jit.load(read_file(args.policy), map_location=env.device).eval() if args.policy else None
+            capture = env.recorder_manager.get_term("transition")
+            metadata = {
+                "task": args.task,
+                "seed": args.seed,
+                "action_scale": args.action_scale,
+                "policy": args.policy,
+                "action_hold": args.action_hold,
+                "dynamics": capture.dynamics.descriptor,
+            }
+            path.with_suffix(".json").write_text(json.dumps(metadata, indent=2) + "\n")
+            with torch.inference_mode():
+                actions = torch.empty(env.num_envs, env.action_manager.total_action_dim, device=env.device)
+                for step in range(args.steps):
+                    if step % args.action_hold == 0:
+                        if policy is None:
+                            actions.uniform_(-args.action_scale, args.action_scale)
+                        else:
+                            actions.copy_(policy(observations["policy"]))
+                    observations, _, _, _, _ = env.step(actions)
+        finally:
+            env.close()
+    print(f"Recorded native transitions to {path}")
+
+
+if __name__ == "__main__":
+    main()

scripts/learning/nerd/train.py

@@ -0,0 +1,161 @@
+# 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
+
+"""Fit temporal NeRD to native HDF5 episodes, with resumable checkpoints and optional W&B uploads."""
+
+import argparse
+import json
+import os
+import signal
+import time
+from contextlib import nullcontext
+from pathlib import Path
+
+import torch
+import torch.distributed as dist
+
+from isaaclab.utils import to_dict
+
+from isaaclab_rl.nerd import NeRDTrainCfg, OfflineNeRD
+
+import isaaclab_tasks  # noqa: F401
+from isaaclab_tasks.utils import load_cfg_from_registry
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--dataset", type=Path, nargs="+", required=True)
+    parser.add_argument("--output", type=Path, required=True)
+    parser.add_argument("--config", type=Path, help="JSON overrides of NeRDTrainCfg, including network geometry.")
+    parser.add_argument("--max_updates", type=int)
+    parser.add_argument("--batch_size", type=int)
+    parser.add_argument("--device", default="cuda:0")
+    parser.add_argument("--checkpoint", type=Path, help="Resume a native checkpoint.")
+    parser.add_argument("--save_interval", type=int, default=1000)
+    parser.add_argument("--save_seconds", type=float, default=120.0)
+    parser.add_argument("--log_interval", type=int, default=100)
+    parser.add_argument("--wandb_project")
+    parser.add_argument("--wandb_entity")
+    parser.add_argument("--wandb_name")
+    args = parser.parse_args()
+    if min(args.save_interval, args.save_seconds, args.log_interval) <= 0:
+        parser.error("Checkpoint and logging intervals must be positive.")
+    distributed = int(os.environ.get("WORLD_SIZE", "1")) > 1
+    if distributed:
+        if args.device.startswith("cuda"):
+            args.device = f"cuda:{os.environ['LOCAL_RANK']}"
+            torch.cuda.set_device(args.device)
+        dist.init_process_group(backend="nccl" if args.device.startswith("cuda") else "gloo")
+    rank = dist.get_rank() if distributed else 0
+    task = json.loads(args.dataset[0].with_suffix(".json").read_text())["task"]
+    task_cfg = load_cfg_from_registry(task, "nerd_cfg_entry_point")
+    settings = task_cfg if isinstance(task_cfg, dict) else to_dict(task_cfg)
+    if args.config:
+        settings.update(json.loads(args.config.read_text()))
+    cfg = NeRDTrainCfg(**settings)
+    if args.max_updates is not None:
+        cfg.max_updates = args.max_updates
+    if args.batch_size is not None:
+        cfg.batch_size = args.batch_size
+    args.output.mkdir(parents=True, exist_ok=True)
+    target = args.output / "last.pt"
+    if target.exists() and not args.checkpoint:
+        parser.error("Output contains last.pt; use --checkpoint to resume or choose another output.")
+    saved = torch.load(args.checkpoint, map_location="cpu", weights_only=True) if args.checkpoint else None
+    settings = to_dict(cfg)
+    if saved:
+        previous = saved["training_cfg"]
+        if any(previous[key] != value for key, value in settings.items() if key not in ("max_updates", "compile")):
+            parser.error("Resume requires unchanged optimization settings, split and model geometry.")
+    learner = OfflineNeRD(args.dataset, cfg, args.device, checkpoint=saved)
+    if learner.updates >= cfg.max_updates:
+        parser.error("The requested update limit has already been reached.")
+    logging = nullcontext(None)
+    if args.wandb_project and rank == 0:
+        import wandb
+
+        identity = saved.get("wandb") if saved else None
+        resume_identity = (
+            identity
+            and identity["project"] == args.wandb_project
+            and (args.wandb_entity is None or identity["entity"] == args.wandb_entity)
+        )
+        logging = wandb.init(
+            project=args.wandb_project,
+            entity=identity["entity"] if resume_identity else args.wandb_entity,
+            name=args.wandb_name,
+            id=identity["id"] if resume_identity else None,
+            resume="allow",
+            config=settings,
+            allow_val_change=True,
+            dir=str(args.output),
+        )
+    stopped = False
+
+    def stop(signum, frame):
+        nonlocal stopped
+        stopped = True
+
+    handlers = {signum: signal.signal(signum, stop) for signum in (signal.SIGINT, signal.SIGTERM)}
+    try:
+        metrics_file = (args.output / "metrics.jsonl").open("a") if rank == 0 else nullcontext(None)
+        with logging as run, metrics_file as metrics:
+            if run:
+                run.define_metric("update")
+                run.define_metric("*", step_metric="update")
+            last_save = time.monotonic()
+            while learner.updates < cfg.max_updates:
+                values = learner.update()
+                checkpoint_due = (
+                    learner.updates % args.save_interval == 0 or time.monotonic() - last_save >= args.save_seconds
+                )
+                finished = learner.updates == cfg.max_updates or stopped
+                if distributed:
+                    decisions = torch.tensor([checkpoint_due, finished], dtype=torch.int32, device=args.device)
+                    dist.all_reduce(decisions, op=dist.ReduceOp.MAX)
+                    checkpoint_due, finished = decisions.tolist()
+                if learner.updates % args.log_interval == 0 or finished or checkpoint_due:
+                    values |= learner.validate()
+                    values["update"] = learner.updates
+                    if metrics:
+                        metrics.write(json.dumps(values) + "\n")
+                        metrics.flush()
+                        print(json.dumps(values), flush=True)
+                    if run:
+                        run.log(values)
+                if checkpoint_due or finished:
+                    checkpoint = learner.checkpoint()
+                    if rank == 0:
+                        checkpoint["training_cfg"] = settings
+                        checkpoint["wandb"] = (
+                            {"id": run.id, "project": run.project, "entity": run.entity} if run else None
+                        )
+                        temporary = target.with_suffix(".tmp")
+                        with temporary.open("wb") as stream:
+                            torch.save(checkpoint, stream)
+                            stream.flush()
+                            os.fsync(stream.fileno())
+                        os.replace(temporary, target)
+                    if run:
+                        artifact = wandb.Artifact(f"nerd-{run.id}", type="model", metadata={"update": learner.updates})
+                        artifact.add_file(str(target), name="last.pt")
+                        logged = run.log_artifact(artifact, aliases=["latest", f"update-{learner.updates}"])
+                        if run.settings.mode != "offline":
+                            logged.wait(timeout=180)
+                        run.log({"update": learner.updates, "checkpoint/durable_update": learner.updates})
+                    last_save = time.monotonic()
+                if finished:
+                    break
+    finally:
+        for signum, handler in handlers.items():
+            signal.signal(signum, handler)
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    finally:
+        if dist.is_initialized():
+            dist.destroy_process_group()