Improves Isaac Lab Mimic tutorial documentation and flow

[peterd-NV]

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

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

Screenshots

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

Checklist

• I have read and understood the contribution guidelines
• I have run the pre-commit checks with ./isaaclab.sh --format
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• I have updated the changelog and the corresponding version in the extension's config/extension.toml file
• I have added my name to the CONTRIBUTORS.md or my name already exists there

[isaaclab-review-bot]

🤖 Isaac Lab Review Bot

Summary

This PR decouples the Franka IK-relative and visuomotor stacking environment configs from the joint-position config inheritance chain, making each self-contained. It also converts idle_action fields from torch.tensor to plain Python lists across four humanoid pick-place configs (with a corresponding consumer fix in replay_demos.py), and includes a significant documentation overhaul of the Franka Mimic imitation learning tutorial.

The refactoring direction is sound — the IK-relative config had no business inheriting from the joint-position config just to get scene setup. However, there is an inconsistency between how the two refactored files handle the EventCfg, and a documentation typo that inverts the intended meaning.

Design Assessment

Acceptable, with one inconsistency to resolve.

Breaking the inheritance from stack_joint_pos_env_cfg.FrankaCubeStackEnvCfg is the right call. The IK-relative env should not depend on a joint-position env just for shared scene/cube/gripper configuration. Inlining these configs makes each environment self-contained and easier to reason about.

The idle_action change from torch.tensor to plain Python lists is also correct — config dataclasses should be serializable and device-agnostic; the tensor conversion now happens at the point of use in replay_demos.py with the correct device.

The code duplication between stack_ik_rel_env_cfg.py and stack_ik_rel_visuomotor_env_cfg.py (~100 lines of identical cube/ee_frame/gripper setup) is worth noting but acceptable for now — it mirrors the existing pattern in stack_joint_pos_env_cfg.py and keeps each config self-contained. A shared helper could reduce this in the future.

Findings

🟡 Warning: Visuomotor EventCfg still inherits init_franka_arm_pose from joint-pos parent — stack_ik_rel_visuomotor_env_cfg.py:47

The visuomotor file's EventCfg inherits from stack_joint_pos_env_cfg.EventCfg, which includes the init_franka_arm_pose event term (the one with the # FIXME comment). This means the visuomotor env will both:

1. Set initial joint pose via ArticulationCfg.InitialStateCfg(joint_pos=...) (new, correct approach)
2. Run set_default_joint_pose via the inherited init_franka_arm_pose event term on every reset (old approach)

The stack_ik_rel_env_cfg.py correctly avoids this by using a standalone EventCfg that doesn't inherit from the joint-pos parent. The visuomotor file should do the same for consistency — either use a standalone EventCfg (like the IK-rel file does) or explicitly override init_franka_arm_pose = None to suppress it.

This won't cause incorrect behavior (both set the same joint values), but it's redundant work on every reset and creates a confusing inconsistency between the two files that were refactored in the same PR.

🟡 Warning: Documentation typo inverts meaning — teleop_imitation.rst:232

The text reads:

"Do have extended pauses."

This should be:

"Do not have extended pauses."

The original text was "Do not pause." The rewrite accidentally dropped the negation, making the advice say the opposite of what's intended.

🔵 Suggestion: Consider inheriting visuomotor from IK-rel config — stack_ik_rel_visuomotor_env_cfg.py

Since the visuomotor env now needs the exact same robot/cube/ee_frame/gripper/event setup as the IK-rel env (all copy-pasted), consider having FrankaCubeStackVisuomotorEnvCfg inherit from the IK-rel FrankaCubeStackEnvCfg and only add the visuomotor-specific parts (cameras, observations, domain randomization events). This would eliminate ~100 lines of duplication while preserving the decoupling from the joint-position config.

Test Coverage

This is primarily a config refactoring and documentation PR. The environment configs are exercised by the existing Isaac Lab test suite (environment creation, stepping, reset). No new behavioral code was added that would require new tests.

The replay_demos.py change (torch.tensor(env_cfg.idle_action, device=...)) is a straightforward fix that follows from the idle_action type change. It would be covered by any existing replay demo integration test.

CI Status

• pre-commit: ✅ pass
• Check for Broken Links: ✅ pass
• Build Latest Docs: ⏳ pending
• license-check: ⏳ pending

Verdict

Minor fixes needed

The refactoring direction is correct and well-motivated. Two items to address:

1. The visuomotor EventCfg inheritance inconsistency (creates redundant init_franka_arm_pose event — moderate concern)
2. The "Do have extended pauses" documentation typo (easy fix)

[kellyguo11]

are there any breaking changes we should highlight? we can use the isaac lab 3.0 migration doc page to keep track of any changes users should be aware of.

[peterd-NV]

are there any breaking changes we should highlight? we can use the isaac lab 3.0 migration doc page to keep track of any changes users should be aware of.

No breaking changes here (beyond the CLI changes that applies to all of Lab e.g. --viz). Everything still works the way it did before but just a bit more refined now in terms of docs.

[isaaclab-review-bot]

🤖 Isaac Lab Review Bot

Summary

This PR is primarily a documentation overhaul for the Isaac Lab Mimic tutorial, plus several code quality improvements: refactoring Franka stack environment configs to eliminate inheritance from joint-pos configs, converting idle_action from tensors to plain Python lists for serializability, adding a --reset_sim_buffer_each_episode flag to replay_demos.py, and marking optional methods in ManagerBasedRLMimicEnv. The changes are largely beneficial for clarity and maintainability.

Architecture Impact

The refactoring of stack_ik_rel_env_cfg.py and stack_ik_rel_visuomotor_env_cfg.py to inherit directly from StackEnvCfg instead of stack_joint_pos_env_cfg.FrankaCubeStackEnvCfg is a breaking change for any external code that relied on the previous inheritance chain. However, this makes the configs self-contained and eliminates confusing inheritance from a joint-position config when using IK controllers. The replay_demos.py refactor extracts the main loop into replay_episodes_loop(), improving testability.

Implementation Verdict

Minor fixes needed — The code changes are sound with one critical bug in replay_demos.py and a few minor issues.

Test Coverage

The PR does not include new unit tests for the refactored environment configs or the new --reset_sim_buffer_each_episode flag. The existing mimic tests (test_generate_dataset_franka_state.py, etc.) exercise the annotation and generation pipelines but don't directly test replay_demos.py. Test timeout increases from 600→5000 seconds are substantial and may warrant investigation into why tests are slow rather than just increasing timeouts.

CI Status

Most checks are pending; pre-commit and build wheel passed. Cannot fully assess until installation and docker image builds complete.

Findings

🔴 Critical: scripts/tools/replay_demos.py:290 — Missing device specification for idle_action tensor

idle_action = torch.zeros(env.action_space.shape)

When env_cfg lacks idle_action, this creates a tensor on CPU without specifying the device. This will cause a device mismatch error when cloning and assigning to actions (line 148) if the environment runs on GPU. Should be:

idle_action = torch.zeros(env.action_space.shape, device=env.unwrapped.device)

🟡 Warning: scripts/tools/replay_demos.py:245-247 — List mutation in argument affects caller

episode_indices_to_replay = list(args_cli.select_episodes)

While this now correctly creates a copy of the CLI argument, the list is passed to replay_episodes_loop() and mutated via .pop(0) at line 178. This is fine since it's a local copy, but .pop(0) on a list is O(n) and could be slow for large episode counts. Consider using collections.deque or reversing and using .pop().

🟡 Warning: source/isaaclab_tasks/.../stack_ik_rel_env_cfg.py:42 — Typo in yaw range tuple

"pose_range": {"x": (0.4, 0.6), "y": (-0.10, 0.10), "z": (0.0203, 0.0203), "yaw": (-1.0, 1, 0)},

The yaw tuple has 3 elements (-1.0, 1, 0) instead of the expected 2 elements (-1.0, 1.0). This appears copied from the original stack_joint_pos_env_cfg.py:53 — verify if randomize_object_pose handles 3-element tuples or if this is a latent bug.

🟡 Warning: pyproject.toml:157 — Aggressive numpy override may break dependencies

override-dependencies = ["numpy>=2"]

Forcing numpy>=2 globally may break packages that genuinely require numpy<2. The PR description mentions this is to "avoid downgrading numpy for SRL usd-to-urdf-converter" but this override affects all dependencies, not just that package. Consider a more targeted solution.

🔵 Improvement: source/isaaclab_tasks/.../stack_ik_rel_env_cfg.py + stack_ik_rel_visuomotor_env_cfg.py — Significant code duplication
Both files now contain nearly identical EventCfg classes, cube configurations, and ee_frame setup (~100 lines duplicated). Consider extracting common Franka stack configuration into a shared base or mixin to reduce maintenance burden.

🔵 Improvement: docs/source/overview/imitation-learning/teleop_imitation.rst:667 — Documentation references truncated file
The diff shows the RST file was truncated at 30K chars. Ensure the full file renders correctly and cross-references (like ``:ref:`generate-additional-demonstrations```) are all valid.

🔵 Improvement: source/isaaclab/isaaclab/envs/manager_based_rl_mimic_env.py:17-19 — @optional_method decorator defined but not exported
The optional_method decorator is defined locally but not added to __all__ or documented. If this is intended to be a public API marker, it should be properly exported from the module.

[isaaclab-review-bot]

🤖 Isaac Lab Review Bot

Summary

This PR refactors Franka cube-stacking environment configurations to eliminate illogical inheritance chains (IK-rel envs no longer inherit from joint-position env), adds documentation improvements for the Isaac Lab Mimic tutorial, fixes a serialization issue with idle_action (changed from torch.tensor to plain Python list), and adds a new --reset_sim_buffer_each_episode option to the replay script. The changes are architecturally sound with one critical bug in the replay script.

Architecture Impact

• Environment Configs: stack_ik_rel_env_cfg.py and stack_ik_rel_visuomotor_env_cfg.py now inherit directly from StackEnvCfg instead of FrankaCubeStackEnvCfg (joint-position variant), eliminating confusing inheritance where IK-based envs inherited joint-position action configurations only to override them.
• Serialization Fix: Changing idle_action from torch.tensor to Python list enables proper environment serialization (pickle/JSON) which is required for distributed training and checkpointing.
• Replay Script: The new --reset_sim_buffer_each_episode flag and refactored loop affect anyone using replay_demos.py for demonstration validation.

Implementation Verdict

Minor fixes needed

Test Coverage

The PR modifies test timeouts and expected annotation counts but does not add new tests for:

• The --reset_sim_buffer_each_episode feature in replay_demos.py
• The refactored environment configurations

The existing mimic tests (test_generate_dataset_franka_state.py) exercise the modified environments indirectly. The expected successful annotations changed from 10 to 9 (test_generate_dataset_franka_state.py:23) — this should be verified as intentional.

CI Status

All jobs are pending. Pre-commit passed.

Findings

🔴 Critical: scripts/tools/replay_demos.py:293 — idle_action device mismatch

idle_action = torch.zeros(env.action_space.shape)

When env_cfg.idle_action is not defined, the fallback creates a tensor on CPU (default device) rather than on env.unwrapped.device. This will cause a device mismatch error when actions = idle_action.clone() at line 148 is used with actions that get passed to env.step(). Should be:

idle_action = torch.zeros(env.action_space.shape, device=env.unwrapped.device)

🟡 Warning: scripts/tools/replay_demos.py:244-246 — Mutating input list

episode_indices_to_replay = list(args_cli.select_episodes)
...
next_episode_index = episode_indices_to_replay.pop(0)

While list() creates a copy of the argparse result, the function replay_episodes_loop receives this list and mutates it via pop(0). The list is passed by reference, so after replay_episodes_loop returns, the list in main() will be empty. This isn't a bug per se (the list isn't used after), but pop(0) on a list is O(n) — consider using collections.deque or iterating with an index for large episode lists.

🟡 Warning: test_generate_dataset_franka_state.py:23 — Changed expected annotation count without explanation

EXPECTED_SUCCESSFUL_ANNOTATIONS = 9  # was 10

The expected successful annotations dropped from 10 to 9. If this is due to dataset changes or annotation logic changes, it should be documented. If it's masking a regression, this is problematic.

🔵 Improvement: stack_ik_rel_env_cfg.py — Significant code duplication
Lines 42-189 in stack_ik_rel_env_cfg.py are nearly identical to lines 48-206 in stack_ik_rel_visuomotor_env_cfg.py. The cube definitions, frame transformer, gripper settings, and event configs are duplicated. Consider extracting a shared mixin or base class FrankaCubeStackBaseEnvCfg containing the common Franka + cube setup, then having both IK-rel and visuomotor configs inherit from it.

🔵 Improvement: manager_based_rl_mimic_env.py:18-20 — Decorator has no runtime effect

def optional_method(func):
    """Decorator to mark a method as optional."""
    func.__is_optional__ = True
    return func

This decorator only sets a metadata attribute. Nothing in the codebase appears to check __is_optional__. If this is intended for documentation purposes only, consider using typing.runtime_checkable protocols or documenting the expected usage pattern. If introspection is planned, document where it will be checked.

🔵 Improvement: pyproject.toml:157 — Override dependency comment needed

override-dependencies = ["numpy>=2"]

This overrides the SRL usd-to-urdf-converter's numpy<2 requirement. A comment explaining why this is safe (or what breaks if removed) would help future maintainers understand the tradeoff.

🔵 Improvement: replay_demos.py:175-176 — Type narrowing could be cleaner

cid = current_episode_indices[env_id]
if cid is not None and cid not in failed_demo_ids:

The type annotation current_episode_indices: list[int | None] is good, but the temporary variable cid is only needed because the list element is used twice. This is fine but could use a comment explaining the None-check is for type narrowing.

[isaaclab-review-bot]

🤖 Isaac Lab Review Bot

Summary

This PR improves Isaac Lab Mimic documentation flow, refactors Franka IK-relative environment configs to inherit directly from StackEnvCfg (eliminating diamond inheritance), changes idle_action from torch.tensor to Python lists for serialization compatibility, and adds a --reset_sim_buffer_each_episode flag to the replay script. The changes are largely correct but contain one critical bug in the idle action handling.

Architecture Impact

• Environment Config Refactoring: stack_ik_rel_env_cfg.py and stack_ik_rel_visuomotor_env_cfg.py now inherit directly from StackEnvCfg instead of through stack_joint_pos_env_cfg.FrankaCubeStackEnvCfg. This eliminates confusing indirect inheritance and duplicates some configuration (cubes, ee_frame, events) explicitly in each file. The duplication is intentional for clarity.
• Serialization Fix: Changing idle_action from torch.tensor to plain Python lists enables environment config serialization (pickle/JSON), which is required for the Mimic workflow.
• replay_demos.py: The extraction of the main loop into replay_episodes_loop() improves testability and the new --reset_sim_buffer_each_episode flag addresses physics buffer issues in single-env replay.

Implementation Verdict

Minor fixes needed — One bug in idle action device handling, and a minor inconsistency in test assertions.

Test Coverage

• The PR modifies test utilities (mimic_test_utils.py) and test files but does not add new tests for the refactored environment configs.
• The existing CI tests (test_generate_dataset_franka_state.py, test_generate_dataset_skillgen.py) validate the Mimic workflow end-to-end.
• Test timeouts were increased substantially (3000→10000s) which is appropriate given the workload.
• Missing: No explicit tests for the --reset_sim_buffer_each_episode flag behavior.

CI Status

Most CI checks are pending/skipped. Pre-commit and wheel build passed, which validates formatting and basic package structure.

Findings

🔴 Critical: scripts/tools/replay_demos.py:287 — Device mismatch when creating idle_action tensor

idle_action = torch.tensor(env_cfg.idle_action, device=env.unwrapped.device).repeat(num_envs, 1)

The env.unwrapped.device is accessed, but env is already unwrapped on line 271 (env = gym.make(...).unwrapped). This works but is redundant. More critically, if env_cfg.idle_action is missing the else branch creates:

idle_action = torch.zeros(env.action_space.shape)

This tensor has no explicit device and will default to CPU, while actions passed to env.step() likely need to be on GPU. This causes a device mismatch. Should be:

idle_action = torch.zeros(env.action_space.shape, device=env.device)

🟡 Warning: source/isaaclab_mimic/test/test_generate_dataset_franka_state.py:111 — Weakened assertion reduces test reliability
The assertion was changed from checking for exactly 10 successful annotations to just > 0:

assert successful_count > 0, "No successful annotations found."

While this makes the test less brittle, it now passes even if only 1/10 annotations succeed, which may mask data quality issues. Consider a middle ground like >= 5.

🟡 Warning: source/isaaclab_tasks/.../stack_ik_rel_env_cfg.py:49 — Typo in yaw range tuple

"yaw": (-1.0, 1, 0)

This is a 3-element tuple where a 2-element range is expected (min, max). The same pattern exists in stack_ik_rel_visuomotor_env_cfg.py:70 and stack_joint_pos_env_cfg.py:56. This appears to be pre-existing code copied during the refactor, but should be fixed to (-1.0, 1.0).

🔵 Improvement: scripts/tools/replay_demos.py:234-241 — Redundant variable initialization

episode_indices_to_replay = list(args_cli.select_episodes)
if len(episode_indices_to_replay) == 0:
    episode_indices_to_replay = list(range(episode_count))

Then this same list is passed to replay_episodes_loop() where it's mutated via pop(0). The original code used args_cli.select_episodes directly which was also a reference. The conversion to list() is correct for isolation, but consider documenting that the function mutates this list.

🔵 Improvement: source/isaaclab_tasks/.../stack_ik_rel_env_cfg.py and stack_ik_rel_visuomotor_env_cfg.py — Significant code duplication
Both files now contain nearly identical ~150 lines defining EventCfg, cube rigid objects, and FrameTransformerCfg. While the PR author intentionally chose duplication over inheritance for clarity, consider extracting shared constants (like _FRANKA_STACK_IK_REL_INIT_JOINT_POS and cube properties) into a common module to reduce maintenance burden.

🔵 Improvement: docs/source/overview/imitation-learning/teleop_imitation.rst:667 — Misplaced section
The "Common Pitfalls when Generating Data" section was moved from earlier in the document to after "Creating Your Own Isaac Lab Mimic Compatible Environments". This breaks the logical flow since pitfalls are relevant to users before they create custom environments. Consider keeping pitfalls in the main tutorial flow.

🔵 Improvement: pyproject.toml:157 — Forced numpy>=2 override may cause issues

override-dependencies = ["numpy>=2"]

This overrides the numpy<2 constraint from srl-usd-to-urdf-converter. If that package genuinely requires numpy<2 for API compatibility, forcing numpy>=2 could cause runtime failures. Verify the SRL package works with numpy 2.x before merging.