Improves Isaac Lab Mimic tutorial documentation and flow (#5283)

# Description

This PR updates workflows and documentation for Isaac Lab Mimic to
improve ease of use and code legibility.
The changes include:

1. Add option for full sim buffer reset in HDF5 replay script when using
single envs. Copy and move a large chunk of the main() function into a
separate helper function.
2. Mark optional methods in ManagerBasedRLMimicEnv
3. Refactor Franka IK Rel envs to inherit directly from Stack Env base.
Eliminate illogical inheritance from Franka direct joint pose env.
4. Change idle action in pick place envs from torch tensor to standard
python list so allow for env serialization
5. Refactor Isaac Lab Mimic documentation for better clarity and flow.
6. Let uv override numpy <2 dep requirements to avoid downgrading numpy
for SRL usd-to-urdf-converter

## Type of change

<!-- As you go through the list, delete the ones that are not
applicable. -->

- New feature (non-breaking change which adds functionality)
- Documentation update

## Screenshots

Please attach before and after screenshots of the change if applicable.

<!--
Example:

| Before | After |
| ------ | ----- |
| _gif/png before_ | _gif/png after_ |

To upload images to a PR -- simply drag and drop an image while in edit
mode and it should upload the image directly. You can then paste that
source into the above before/after sections.
-->

## Checklist

- [x] I have read and understood the [contribution
guidelines](https://isaac-sim.github.io/IsaacLab/main/source/refs/contributing.html)
- [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./isaaclab.sh --format`
- [x] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [ ] I have added tests that prove my fix is effective or that my
feature works
- [x] I have updated the changelog and the corresponding version in the
extension's `config/extension.toml` file
- [x] I have added my name to the `CONTRIBUTORS.md` or my name already
exists there

<!--
As you go through the checklist above, you can mark something as done by
putting an x character in it

For example,
- [x] I have done this task
- [ ] I have not done this task
-->

Author: peterd-NV

docs/source/_static/mimic/franka_cube_stacking_env.jpg

@@ -5,6 +5,10 @@ Augmented Imitation Learning
 
 This section describes how to use Isaac Lab's imitation learning capabilities with the visual augmentation capabilities of `Cosmos <https://www.nvidia.com/en-us/ai/cosmos/>`_ models to generate demonstrations at scale to train visuomotor policies robust against visual variations.
 
+
+.. important::
+   The `Cosmos Transfer1 <https://github.com/nvidia-cosmos/cosmos-transfer1/tree/e4055e39ee9c53165e85275bdab84ed20909714a>`_ model used in this tutorial is `supported <https://huggingface.co/nvidia/Cosmos-Transfer1-7B#software-integration>`_ on Ampere and Hopper GPUs.
+
 Generating Demonstrations
 ~~~~~~~~~~~~~~~~~~~~~~~~~
 

docs/source/_static/mimic/franka_datagen.jpg

@@ -11,7 +11,7 @@ This page covers data generation and imitation learning workflows for humanoid r
 
 .. important::
 
-   Complete the tutorial in :ref:`Teleoperation and Imitation Learning with Isaac Lab Mimic <teleoperation-imitation-learning>`
+   Complete the tutorial in :ref:`Synthetic Data Generation and Imitation Learning with Isaac Lab Mimic <teleoperation-imitation-learning>`
    before proceeding with the following demonstrations to
    understand the data collection, annotation, and generation steps of Isaac Lab Mimic.
 
@@ -114,7 +114,7 @@ You can replay the collected demonstrations by running the following command:
 Annotate the demonstrations
 """""""""""""""""""""""""""
 
-Unlike the :ref:`Franka stacking task <generating-additional-demonstrations>`, the GR-1 pick and place task uses manual annotation to define subtasks.
+Unlike the :ref:`Franka stacking task <generate-additional-demonstrations>`, the GR-1 pick and place task uses manual annotation to define subtasks.
 
 The pick and place task has one subtask for the left arm (pick) and two subtasks for the right arm (idle, place).
 Annotations denote the end of a subtask. For the pick and place task, this means there are no annotations for the left arm and one annotation for the right arm (the end of the final subtask is always implicit).

docs/source/_static/mimic/mimic_env_architecture.jpg

@@ -154,6 +154,7 @@ explicit = false
 [tool.uv]
 index-strategy = "unsafe-best-match"
 prerelease = "allow"
+override-dependencies = ["numpy>=2"]
 
 [tool.uv.pip]
 index-strategy = "unsafe-best-match"

docs/source/overview/imitation-learning/augmented_imitation.rst

@@ -38,6 +38,15 @@
     default=False,
     help="Validate the replay success rate using the task environment termination criteria",
 )
+parser.add_argument(
+    "--reset_sim_buffer_each_episode",
+    action="store_true",
+    default=False,
+    help=(
+        "Before loading each episode's initial state, call env.sim.reset() to clear"
+        " simulation buffers. Only valid with --num_envs 1."
+    ),
+)
 
 # append AppLauncher cli args
 AppLauncher.add_app_launcher_args(parser)
@@ -106,86 +115,27 @@ def compare_states(state_from_dataset, runtime_state, runtime_env_index) -> (boo
     return states_matched, output_log
 
 
-def main():
-    """Replay episodes loaded from a file."""
-    global is_paused
-
-    # Load dataset
-    if not os.path.exists(args_cli.dataset_file):
-        raise FileNotFoundError(f"The dataset file {args_cli.dataset_file} does not exist.")
-    dataset_file_handler = HDF5DatasetFileHandler()
-    dataset_file_handler.open(args_cli.dataset_file)
-    env_name = dataset_file_handler.get_env_name()
-    episode_count = dataset_file_handler.get_num_episodes()
-
-    if episode_count == 0:
-        print("No episodes found in the dataset.")
-        exit()
-
-    episode_indices_to_replay = args_cli.select_episodes
-    if len(episode_indices_to_replay) == 0:
-        episode_indices_to_replay = list(range(episode_count))
-
-    if args_cli.task is not None:
-        env_name = args_cli.task.split(":")[-1]
-    if env_name is None:
-        raise ValueError("Task/env name was not specified nor found in the dataset.")
-
-    num_envs = args_cli.num_envs
-
-    env_cfg = parse_env_cfg(env_name, device=args_cli.device, num_envs=num_envs)
-
-    # extract success checking function to invoke in the main loop
-    success_term = None
-    if args_cli.validate_success_rate:
-        if hasattr(env_cfg.terminations, "success"):
-            success_term = env_cfg.terminations.success
-            env_cfg.terminations.success = None
-        else:
-            print(
-                "No success termination term was found in the environment."
-                " Will not be able to mark recorded demos as successful."
-            )
-
-    # Disable all recorders and terminations
-    env_cfg.recorders = {}
-    env_cfg.terminations = {}
-
-    # create environment from loaded config
-    env = gym.make(args_cli.task, cfg=env_cfg).unwrapped
-
-    teleop_interface = Se3Keyboard(Se3KeyboardCfg(pos_sensitivity=0.1, rot_sensitivity=0.1))
-    teleop_interface.add_callback("N", play_cb)
-    teleop_interface.add_callback("B", pause_cb)
-    print('Press "B" to pause and "N" to resume the replayed actions.')
-
-    # Determine if state validation should be conducted
-    state_validation_enabled = False
-    if args_cli.validate_states and num_envs == 1:
-        state_validation_enabled = True
-    elif args_cli.validate_states and num_envs > 1:
-        print("Warning: State validation is only supported with a single environment. Skipping state validation.")
-
-    # Get idle action (idle actions are applied to envs without next action)
-    if hasattr(env_cfg, "idle_action"):
-        idle_action = env_cfg.idle_action.repeat(num_envs, 1)
-    else:
-        idle_action = torch.zeros(env.action_space.shape)
+def replay_episodes_loop(  # noqa: C901
+    env,
+    dataset_file_handler: HDF5DatasetFileHandler,
+    episode_names: list[str],
+    episode_count: int,
+    episode_indices_to_replay: list[int],
+    num_envs: int,
+    success_term,
+    state_validation_enabled: bool,
+    idle_action: torch.Tensor,
+    reset_sim_buffer_each_episode: bool,
+) -> tuple[int, int, list[int]]:
+    """Run the replay loop until all selected episodes finish or the app exits.
 
-    # reset before starting
-    env.reset()
-    teleop_interface.reset()
-
-    # simulate environment -- run everything in inference mode
-    episode_names = list(dataset_file_handler.get_episode_names())
+    Returns:
+        Tuple of (replayed_episode_count, recorded_episode_count, failed_demo_ids).
+    """
     replayed_episode_count = 0
     recorded_episode_count = 0
-
-    # Track current episode indices for each environment
-    current_episode_indices = [None] * num_envs
-
-    # Track failed demo IDs
-    failed_demo_ids = []
+    current_episode_indices: list[int | None] = [None] * num_envs
+    failed_demo_ids: list[int] = []
 
     with contextlib.suppress(KeyboardInterrupt) and torch.inference_mode():
         while simulation_app.is_running() and not simulation_app.is_exiting():
@@ -195,7 +145,7 @@ def main():
             episode_ended = [False] * num_envs
             while has_next_action:
                 # initialize actions with idle action so those without next action will not move
-                actions = idle_action
+                actions = idle_action.clone()
                 has_next_action = False
                 for env_id in range(num_envs):
                     env_next_action = env_episode_data_map[env_id].get_next_action()
@@ -216,11 +166,9 @@ def main():
                                 )
                             else:
                                 # if not successful, add to failed demo IDs list
-                                if (
-                                    current_episode_indices[env_id] is not None
-                                    and current_episode_indices[env_id] not in failed_demo_ids
-                                ):
-                                    failed_demo_ids.append(current_episode_indices[env_id])
+                                cid = current_episode_indices[env_id]
+                                if cid is not None and cid not in failed_demo_ids:
+                                    failed_demo_ids.append(cid)
 
                             episode_ended[env_id] = True
 
@@ -243,6 +191,8 @@ def main():
                             env_episode_data_map[env_id] = episode_data
                             # Set initial state for the new episode
                             initial_state = episode_data.get_initial_state()
+                            if reset_sim_buffer_each_episode:
+                                env.sim.reset()
                             env.reset_to(initial_state, torch.tensor([env_id], device=env.device), is_relative=True)
                             # Get the first action for the new episode
                             env_next_action = env_episode_data_map[env_id].get_next_action()
@@ -275,6 +225,99 @@ def main():
                             print("\t- mismatched.")
                             print(comparison_log)
             break
+
+    return replayed_episode_count, recorded_episode_count, failed_demo_ids
+
+
+def main():
+    """Replay episodes loaded from a file."""
+    global is_paused
+
+    # Load dataset
+    if not os.path.exists(args_cli.dataset_file):
+        raise FileNotFoundError(f"The dataset file {args_cli.dataset_file} does not exist.")
+    dataset_file_handler = HDF5DatasetFileHandler()
+    dataset_file_handler.open(args_cli.dataset_file)
+    env_name = dataset_file_handler.get_env_name()
+    episode_count = dataset_file_handler.get_num_episodes()
+
+    if episode_count == 0:
+        print("No episodes found in the dataset.")
+        exit()
+
+    episode_indices_to_replay = list(args_cli.select_episodes)
+    if len(episode_indices_to_replay) == 0:
+        episode_indices_to_replay = list(range(episode_count))
+
+    if args_cli.task is not None:
+        env_name = args_cli.task.split(":")[-1]
+    if env_name is None:
+        raise ValueError("Task/env name was not specified nor found in the dataset.")
+
+    num_envs = args_cli.num_envs
+    if args_cli.reset_sim_buffer_each_episode and num_envs != 1:
+        raise ValueError(
+            "--reset_sim_buffer_each_episode is only supported with a single environment (--num_envs 1). "
+            f"Got num_envs={num_envs}. Use --num_envs 1 or disable --reset_sim_buffer_each_episode."
+        )
+
+    env_cfg = parse_env_cfg(env_name, device=args_cli.device, num_envs=num_envs)
+
+    # extract success checking function to invoke in the main loop
+    success_term = None
+    if args_cli.validate_success_rate:
+        if hasattr(env_cfg.terminations, "success"):
+            success_term = env_cfg.terminations.success
+            env_cfg.terminations.success = None
+        else:
+            print(
+                "No success termination term was found in the environment."
+                " Will not be able to mark recorded demos as successful."
+            )
+
+    # Disable all recorders and terminations
+    env_cfg.recorders = {}
+    env_cfg.terminations = {}
+
+    # create environment from loaded config
+    env = gym.make(args_cli.task, cfg=env_cfg).unwrapped
+
+    teleop_interface = Se3Keyboard(Se3KeyboardCfg(pos_sensitivity=0.1, rot_sensitivity=0.1))
+    teleop_interface.add_callback("N", play_cb)
+    teleop_interface.add_callback("B", pause_cb)
+    print('Press "B" to pause and "N" to resume the replayed actions.')
+
+    # Determine if state validation should be conducted
+    state_validation_enabled = False
+    if args_cli.validate_states and num_envs == 1:
+        state_validation_enabled = True
+    elif args_cli.validate_states and num_envs > 1:
+        print("Warning: State validation is only supported with a single environment. Skipping state validation.")
+
+    # Get idle action (idle actions are applied to envs without next action)
+    if hasattr(env_cfg, "idle_action"):
+        idle_action = torch.tensor(env_cfg.idle_action, device=env.unwrapped.device).repeat(num_envs, 1)
+    else:
+        idle_action = torch.zeros(env.action_space.shape)
+
+    # reset before starting
+    env.reset()
+    teleop_interface.reset()
+
+    episode_names = list(dataset_file_handler.get_episode_names())
+    replayed_episode_count, recorded_episode_count, failed_demo_ids = replay_episodes_loop(
+        env,
+        dataset_file_handler,
+        episode_names,
+        episode_count,
+        episode_indices_to_replay,
+        num_envs,
+        success_term,
+        state_validation_enabled,
+        idle_action,
+        args_cli.reset_sim_buffer_each_episode,
+    )
+
     # Close environment after replay in complete
     plural_trailing_s = "s" if replayed_episode_count > 1 else ""
     print(f"Finished replaying {replayed_episode_count} episode{plural_trailing_s}.")

docs/source/overview/imitation-learning/humanoids_imitation.rst

@@ -1,7 +1,7 @@
 [package]
 
 # Note: Semantic Versioning is used: https://semver.org/
-version = "4.6.14"
+version = "4.6.15"
 
 # Description
 title = "Isaac Lab framework for Robot Learning"

docs/source/overview/imitation-learning/teleop_imitation.rst

@@ -1,6 +1,17 @@
 Changelog
 ---------
 
+4.6.15 (2026-04-24)
+~~~~~~~~~~~~~~~~~~~
+
+Changed
+^^^^^^^
+
+* Marked :meth:`~isaaclab.envs.manager_based_rl_mimic_env.ManagerBasedRLMimicEnv.get_subtask_start_signals` and
+  :meth:`~isaaclab.envs.manager_based_rl_mimic_env.ManagerBasedRLMimicEnv.get_subtask_term_signals` with
+  ``@optional_method``.
+
+
 4.6.14 (2026-04-24)
 ~~~~~~~~~~~~~~~~~~~