Add optional NeRD force and joint-wrench reconstruction

docs/index.rst

@@ -165,6 +165,7 @@ Table of Contents
 
    source/experimental-features/bleeding-edge
    source/concepts/warp_environments
+   source/experimental-features/nerd
    source/experimental-features/ray
    source/experimental-features/visuo_tactile_sensor
    source/experimental-features/rlinf_vla_posttraining

docs/source/api/lab/isaaclab.managers.rst

@@ -228,9 +228,40 @@ Curriculum Manager
     :members:
     :exclude-members: __init__
 
+.. _recorder-capture:
+
 Recorder Manager
 ----------------
 
+Tasks configure recorder terms to capture the state, controls and other fields their consumers need.
+Terms run at pre/post step and reset callbacks; reset callbacks receive slices or indices directly.
+Consumers own subsequent processing, replay and learning.
+
+Returning a field path and a batched tensor (or nested dictionary of tensors) retains samples.
+For online use without retention, publish the batch on the term and return ``(None, None)``.
+Consumers access it through ``env.recorder_manager.get_term(name)`` and copy data before the next capture.
+``EXPORT_NONE`` disables file export only; disabling the term removes capture itself.
+
+Fields may have different sample rates. ``buffer_initial_capacity`` sets the initial samples per world
+and field; choose it with payload size and episode length in mind. Storage grows without dropping
+samples and is reused across resets.
+
+**Observation migration:** RL post-step callbacks run after rewards and terminations, before reset
+and observation updates. ``env.obs_buf`` still contains the preceding policy input. Read physical
+state from the scene or call ``env.observation_manager.compute_group("recording", update_history=False)``
+for current observations. Use a separate noiseless group without history or delay to avoid stale
+measurements and additional policy noise draws. History is not advanced, but noise and modifiers still
+run. Non-RL environments update their normal observations before post-step capture.
+
+**Storage migration:** ``get_episode(env_id)`` returns an owned tensor snapshot with a leading sample
+dimension; reading copies data and may synchronize the device. Modify live recording through
+``add_to_episodes`` and ``set_success_to_episodes``, not the snapshot. ``EpisodeData.add()`` and
+``pre_export()`` are removed: collect through the manager, or assign completed tensors to
+``episode.data``. HDF5 loading and playback keep the same tensor format.
+
+The unused ``RecorderTerm.close(file_path)`` hook is removed. Dataset handlers finalize files;
+applications explicitly close any additional resources they own.
+
 .. 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()

source/isaaclab/changelog.d/recorder-capture.changed.rst

@@ -0,0 +1,21 @@
+* 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` for direct access to captured batches.
+* Preserve slices through recorder reset callbacks; terms accept slices or indices directly.
+* **Breaking:** Removed the implicit observation computation before RL post-step recorder callbacks.
+  Physical capture no longer consumes extra observation noise. Custom terms needing observations
+  of the completed transition must call ``env.observation_manager.compute_group(name, update_history=False)``
+  explicitly; ``env.obs_buf`` still contains the preceding policy input. Use a noiseless group without
+  history or delay for current measurements that do not consume policy 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.
+* **Breaking:** Removed ``EpisodeData.add()`` and ``pre_export()``. Collection now belongs to
+  ``RecorderManager.add_to_episodes()``; assign nested tensors with leading sample dimensions
+  to ``EpisodeData.data`` when constructing an episode directly. Playback and HDF5 loading remain supported.
+* **Breaking:** Removed the unused ``RecorderTerm.close(file_path)`` hook. Dataset handlers own
+  file finalization; applications must explicitly close any additional resources they own.
+
+* Documented independent capture, retention and export choices, with task-selected schemas and
+  consumer-owned processing. Core recorder storage remains independent of task and learning packages.

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.

source/isaaclab/isaaclab/envs/manager_based_rl_env.py

@@ -236,10 +236,7 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn:
         self.reset_time_outs = self.termination_manager.time_outs
         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()
-            self.recorder_manager.record_post_step()
+        self.recorder_manager.record_post_step()
 
         # -- reset envs that terminated/timed-out and log the episode information
         reset_env_ids = self.reset_buf.nonzero(as_tuple=False).squeeze(-1).int()

source/isaaclab/isaaclab/envs/mdp/recorders/recorders.py

@@ -14,7 +14,7 @@
 class InitialStateRecorder(RecorderTerm):
     """Recorder term that records the initial state of the environment after reset."""
 
-    def record_post_reset(self, env_ids: Sequence[int] | None):
+    def record_post_reset(self, env_ids: Sequence[int] | slice | None):
         def extract_env_ids_values(value):
             nonlocal env_ids
             if isinstance(value, dict):