Updates Isaac Lab Mimic docs for V2 release (#299)

# Description

<!--
Thank you for your interest in sending a pull request. Please make sure
to check the contribution guidelines.

Link:
https://isaac-sim.github.io/IsaacLab/main/source/refs/contributing.html
-->

Documentation update for Isaac Lab Mimic for V2 release. Adds steps for
training visuomotor policy and dexmimic data generation.


## Type of change

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

- 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

- [ ] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./isaaclab.sh --format`
- [x] 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

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

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

@@ -24,6 +24,8 @@
     DirectMARLEnvCfg
     ManagerBasedRLMimicEnv
     MimicEnvCfg
+    SubTaskConfig
+    SubTaskConstraintConfig
     ViewerCfg
 
 Manager Based Environment
@@ -92,6 +94,18 @@ Mimic Environment
     :show-inheritance:
     :exclude-members: __init__, class_type
 
+.. autoclass:: SubTaskConfig
+    :members:
+    :inherited-members:
+    :show-inheritance:
+    :exclude-members: __init__, class_type
+
+.. autoclass:: SubTaskConstraintConfig
+    :members:
+    :inherited-members:
+    :show-inheritance:
+    :exclude-members: __init__, class_type
+
 Common
 ------
 

docs/source/overview/environments.rst

@@ -109,7 +109,9 @@ for the lift-cube environment:
     +--------------------+-------------------------+-----------------------------------------------------------------------------+
     | |lift-cube|        | |lift-cube-link|        | Pick a cube and bring it to a sampled target position with the Franka robot |
     +--------------------+-------------------------+-----------------------------------------------------------------------------+
-    | |stack-cube|       | |stack-cube-link|       | Stack three cubes (bottom to top: blue, red, green) with the Franka robot   |
+    | |stack-cube|       | |stack-cube-link|       | Stack three cubes (bottom to top: blue, red, green) with the Franka robot.  |
+    |                    |                         | Blueprint env used for the NVIDIA Isaac GR00T blueprint for synthetic       |
+    |                    | |stack-cube-bp-link|    | manipulation motion generation                                              |
     +--------------------+-------------------------+-----------------------------------------------------------------------------+
     | |cabi-franka|      | |cabi-franka-link|      | Grasp the handle of a cabinet's drawer and open it with the Franka robot    |
     |                    |                         |                                                                             |
@@ -127,6 +129,8 @@ for the lift-cube environment:
     +--------------------+-------------------------+-----------------------------------------------------------------------------+
     | |cube-shadow|      | |cube-shadow-vis-link|  | In-hand reorientation of a cube using Shadow hand using perceptive inputs   |
     +--------------------+-------------------------+-----------------------------------------------------------------------------+
+    | |gr1_pick_place|   | |gr1_pick_place-link|   | Pick up and place an object in a basket with a GR-1 humanoid robot          |
+    +--------------------+-------------------------+-----------------------------------------------------------------------------+
 
 .. |reach-franka| image:: ../_static/tasks/manipulation/franka_reach.jpg
 .. |reach-ur10| image:: ../_static/tasks/manipulation/ur10_reach.jpg
@@ -135,6 +139,7 @@ for the lift-cube environment:
 .. |cube-allegro| image:: ../_static/tasks/manipulation/allegro_cube.jpg
 .. |cube-shadow| image:: ../_static/tasks/manipulation/shadow_cube.jpg
 .. |stack-cube| image:: ../_static/tasks/manipulation/franka_stack.jpg
+.. |gr1_pick_place| image:: ../_static/tasks/manipulation/gr-1_pick_place.jpg
 
 .. |reach-franka-link| replace:: `Isaac-Reach-Franka-v0 <https://github.com/isaac-sim/IsaacLab/blob/main/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/config/franka/joint_pos_env_cfg.py>`__
 .. |reach-ur10-link| replace:: `Isaac-Reach-UR10-v0 <https://github.com/isaac-sim/IsaacLab/blob/main/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/config/ur_10/joint_pos_env_cfg.py>`__
@@ -146,6 +151,8 @@ for the lift-cube environment:
 .. |cube-allegro-link| replace:: `Isaac-Repose-Cube-Allegro-v0 <https://github.com/isaac-sim/IsaacLab/blob/main/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/config/allegro_hand/allegro_env_cfg.py>`__
 .. |allegro-direct-link| replace:: `Isaac-Repose-Cube-Allegro-Direct-v0 <https://github.com/isaac-sim/IsaacLab/blob/main/source/isaaclab_tasks/isaaclab_tasks/direct/allegro_hand/allegro_hand_env_cfg.py>`__
 .. |stack-cube-link| replace:: `Isaac-Stack-Cube-Franka-v0 <https://github.com/isaac-sim/IsaacLab/blob/main/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/franka/stack_joint_pos_env_cfg.py>`__
+.. |stack-cube-bp-link| replace:: `Isaac-Stack-Cube-Franka-Blueprint-v0 <https://github.com/isaac-sim/IsaacLab/blob/main/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/franka/stack_ik_rel_blueprint_env_cfg.py>`__
+.. |gr1_pick_place-link| replace:: `Isaac-PickPlace-GR1T2-Abs-v0 <https://github.com/isaac-sim/IsaacLab/blob/main/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/pick_place/pickplace_gr1t2_env_cfg.py>`__
 
 .. |cube-shadow-link| replace:: `Isaac-Repose-Cube-Shadow-Direct-v0 <https://github.com/isaac-sim/IsaacLab/blob/main/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/shadow_hand_env_cfg.py>`__
 .. |cube-shadow-ff-link| replace:: `Isaac-Repose-Cube-Shadow-OpenAI-FF-Direct-v0 <https://github.com/isaac-sim/IsaacLab/blob/main/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/shadow_hand_env_cfg.py>`__

docs/source/overview/teleop_imitation.rst

@@ -121,7 +121,8 @@ If, while performing a demonstration, a mistake is made, or the current demonstr
 Pre-recorded demonstrations
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-We provide a pre-recorded ``dataset.hdf5`` containing 10 human demonstrations for ``Isaac-Stack-Cube-Franka-IK-Rel-v0`` `here <https://omniverse-content-production.s3-us-west-2.amazonaws.com/Assets/Isaac/4.5/Isaac/IsaacLab/Mimic/dataset.hdf5>`_.
+We provide a pre-recorded ``dataset.hdf5`` containing 10 human demonstrations for ``Isaac-Stack-Cube-Franka-IK-Rel-v0``
+`here <https://omniverse-content-production.s3-us-west-2.amazonaws.com/Assets/Isaac/4.5/Isaac/IsaacLab/Mimic/dataset.hdf5>`_.
 This dataset may be downloaded and used in the remaining tutorial steps if you do not wish to collect your own demonstrations.
 
 .. note::
@@ -134,35 +135,75 @@ Additional demonstrations can be generated using Isaac Lab Mimic.
 
 Isaac Lab Mimic is a feature in Isaac Lab that allows generation of additional demonstrations automatically, allowing a policy to learn successfully even from just a handful of manual demonstrations.
 
+In the following example, we will show how to use Isaac Lab Mimic to generate additional demonstrations that can be used to train either a state-based policy
+(using the ``Isaac-Stack-Cube-Franka-IK-Rel-Mimic-v0`` environment) or visuomotor policy (using the ``Isaac-Stack-Cube-Franka-IK-Rel-Visuomotor-Mimic-v0`` environment).
+
+.. note::
+
+   All commands in the following sections must keep a consistent policy type. For example, if choosing to use a state-based policy, then all commands used should be from the "State-based policy" tab.
+
 In order to use Isaac Lab Mimic with the recorded dataset, first annotate the subtasks in the recording:
 
-.. code:: bash
+.. tabs::
+   .. tab:: State-based policy
+      .. code:: bash
+
+         ./isaaclab.sh -p scripts/imitation_learning/isaaclab_mimic/annotate_demos.py \
+         --device cuda --task Isaac-Stack-Cube-Franka-IK-Rel-Mimic-v0 --auto \
+         --input_file ./datasets/dataset.hdf5 --output_file ./datasets/annotated_dataset.hdf5
+
+   .. tab:: Visuomotor policy
+      .. code:: bash
+
+         ./isaaclab.sh -p scripts/imitation_learning/isaaclab_mimic/annotate_demos.py \
+         --device cuda --enable_cameras --task Isaac-Stack-Cube-Franka-IK-Rel-Visuomotor-Mimic-v0 --auto \
+         --input_file ./datasets/dataset.hdf5 --output_file ./datasets/annotated_dataset.hdf5
 
-   ./isaaclab.sh -p scripts/imitation_learning/isaaclab_mimic/annotate_demos.py --input_file ./datasets/dataset.hdf5 --output_file ./datasets/annotated_dataset.hdf5 --task Isaac-Stack-Cube-Franka-IK-Rel-Mimic-v0 --auto
 
 Then, use Isaac Lab Mimic to generate some additional demonstrations:
 
-.. code:: bash
+.. tabs::
+   .. tab:: State-based policy
+      .. code:: bash
 
-   ./isaaclab.sh -p scripts/imitation_learning/isaaclab_mimic/generate_dataset.py --input_file ./datasets/annotated_dataset.hdf5 --output_file ./datasets/generated_dataset_small.hdf5 --num_envs 10 --generation_num_trials 10
+         ./isaaclab.sh -p scripts/imitation_learning/isaaclab_mimic/generate_dataset.py \
+         --device cuda --num_envs 10 --generation_num_trials 10 \
+         --input_file ./datasets/annotated_dataset.hdf5 --output_file ./datasets/generated_dataset_small.hdf5
 
-.. note::
+   .. tab:: Visuomotor policy
+      .. code:: bash
 
-   The output_file of the ``annotate_demos.py`` script is the input_file to the ``generate_dataset.py`` script
+         ./isaaclab.sh -p scripts/imitation_learning/isaaclab_mimic/generate_dataset.py \
+         --device cuda --enable_cameras --num_envs 10 --generation_num_trials 10 \
+         --input_file ./datasets/annotated_dataset.hdf5 --output_file ./datasets/generated_dataset_small.hdf5
 
 .. note::
 
-   Isaac Lab is designed to work with manipulators with grippers. The gripper commands in the demonstrations are extracted separately and temporally replayed during the generation of additional demonstrations.
+   The output_file of the ``annotate_demos.py`` script is the input_file to the ``generate_dataset.py`` script
 
 Inspect the output of generated data (filename: ``generated_dataset_small.hdf5``), and if satisfactory, generate the full dataset:
 
-.. code:: bash
+.. tabs::
+   .. tab:: State-based policy
+      .. code:: bash
+
+         ./isaaclab.sh -p scripts/imitation_learning/isaaclab_mimic/generate_dataset.py \
+         --device cuda --headless --num_envs 10 --generation_num_trials 1000 \
+         --input_file ./datasets/annotated_dataset.hdf5 --output_file ./datasets/generated_dataset.hdf5
+
+   .. tab:: Visuomotor policy
+      .. code:: bash
+
+         ./isaaclab.sh -p scripts/imitation_learning/isaaclab_mimic/generate_dataset.py \
+         --device cuda --enable_cameras --headless --num_envs 10 --generation_num_trials 1000 \
+         --input_file ./datasets/annotated_dataset.hdf5 --output_file ./datasets/generated_dataset.hdf5
 
-   ./isaaclab.sh -p scripts/imitation_learning/isaaclab_mimic/generate_dataset.py --input_file ./datasets/annotated_dataset.hdf5 --output_file ./datasets/generated_dataset.hdf5 --num_envs 10 --generation_num_trials 1000 --headless
 
 The number of demonstrations can be increased or decreased, 1000 demonstrations have been shown to provide good training results for this task.
 
-Additionally, the number of environments in the ``--num_envs`` parameter can be adjusted to speed up data generation. The suggested number of 10 can be executed even on a laptop GPU. On a more powerful desktop machine, set it to 100 or higher for significant speedup of this step.
+Additionally, the number of environments in the ``--num_envs`` parameter can be adjusted to speed up data generation.
+The suggested number of 10 can be executed on a moderate laptop GPU.
+On a more powerful desktop machine, use a larger number of environments for a significant speedup of this step.
 
 Robomimic setup
 ^^^^^^^^^^^^^^^
@@ -181,23 +222,109 @@ To install the robomimic framework, use the following commands:
 Training an agent
 ^^^^^^^^^^^^^^^^^
 
-We can now train a BC agent for ``Isaac-Stack-Cube-Franka-IK-Rel-v0`` using the Mimic generated data:
+Using the Mimic generated data we can now train a state-based BC agent for ``Isaac-Stack-Cube-Franka-IK-Rel-v0``, or a visuomotor BC agent for ``Isaac-Stack-Cube-Franka-IK-Rel-Visuomotor-v0``:
 
-.. code:: bash
+.. tabs::
+   .. tab:: State-based policy
+      .. code:: bash
+
+         ./isaaclab.sh -p scripts/imitation_learning/robomimic/train.py \
+         --task Isaac-Stack-Cube-Franka-IK-Rel-v0 --algo bc \
+         --dataset ./datasets/generated_dataset.hdf5
 
-   ./isaaclab.sh -p scripts/imitation_learning/robomimic/train.py --task Isaac-Stack-Cube-Franka-IK-Rel-v0 --algo bc --dataset ./datasets/generated_dataset.hdf5
+   .. tab:: Image-based policy
+      .. code:: bash
 
-By default, the training script will save a model checkpoint every 100 epochs. The trained models and logs will be saved to logs/robomimic/Isaac-Stack-Cube-Franka-IK-Rel-v0/bc
+         ./isaaclab.sh -p scripts/imitation_learning/robomimic/train.py \
+         --task Isaac-Stack-Cube-Franka-IK-Rel-Visuomotor-v0 --algo bc \
+         --dataset ./datasets/generated_dataset.hdf5
+
+.. note::
+   By default the trained models and logs will be saved to ``IssacLab/logs/robomimic``.
 
 Visualizing results
 ^^^^^^^^^^^^^^^^^^^
 
-By inferencing using the generated model, we can visualize the results of the policy in the same environment:
+By inferencing using the generated model, we can visualize the results of the policy:
+
+.. tabs::
+   .. tab:: State-based policy
+      .. code:: bash
+
+         ./isaaclab.sh -p scripts/imitation_learning/robomimic/play.py \
+         --device cuda --task Isaac-Stack-Cube-Franka-IK-Rel-v0 --num_rollouts 50 \
+         --checkpoint /PATH/TO/desired_model_checkpoint.pth
+
+   .. tab:: Visuomotor policy
+      .. code:: bash
+
+         ./isaaclab.sh -p scripts/imitation_learning/robomimic/play.py \
+         --device cuda --enable_cameras --task Isaac-Stack-Cube-Franka-IK-Rel-Visuomotor-v0 --num_rollouts 50 \
+         --checkpoint /PATH/TO/desired_model_checkpoint.pth
+
+
+Demo: Data Generation and Policy Training for a Humanoid Robot
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. figure:: ../_static/tasks/manipulation/gr-1_pick_place.gif
+   :width: 100%
+   :align: center
+   :alt: GR-1 humanoid robot performing a pick and place task
+
+
+Isaac Lab Mimic supports data generation for robots with multiple end effectors. In the following demonstration, we will show how to generate data
+to train a Fourier GR-1 humanoid robot to perform a pick and place task.
+
+
+Generate the dataset
+^^^^^^^^^^^^^^^^^^^^
+
+Download the pre-recorded annotated dataset ``dataset_annotated_gr1.hdf5`` from
+`here <https://omniverse-content-production.s3-us-west-2.amazonaws.com/Assets/Isaac/4.5/Isaac/IsaacLab/Mimic/dataset_annotated_gr1.hdf5>`_.
+Place the file under ``IsaacLab/datasets`` and run the following command to generate a new dataset with 1000 demonstrations.
 
 .. code:: bash
 
-   ./isaaclab.sh -p scripts/imitation_learning/robomimic/play.py --task Isaac-Stack-Cube-Franka-IK-Rel-v0 --num_rollouts 50 --checkpoint /PATH/TO/desired_model_checkpoint.pth
+   ./isaaclab.sh -p scripts/imitation_learning/isaaclab_mimic/generate_dataset.py \
+   --device cuda --headless --num_envs 10 --generation_num_trials 1000 \
+   --input_file ./datasets/dataset_annotated_gr1.hdf5 --output_file ./datasets/generated_dataset_gr1.hdf5
+
+Train a policy
+^^^^^^^^^^^^^^
+
+Use Robomimic to train a policy for the generated dataset.
+
+.. code:: bash
 
+   ./isaaclab.sh -p scripts/imitation_learning/robomimic/train.py \
+   --task Isaac-PickPlace-GR1T2-Abs-v0 --algo bc \
+   --normalize_training_actions \
+   --dataset ./datasets/generated_dataset_gr1.hdf5
+
+The training script will normalize the actions in the dataset to the range [-1, 1].
+The normalization parameters are saved in the model directory under ``PATH_TO_MODEL_DIRECTORY/logs/normalization_params.txt``.
+Record the normalization parameters for later use in the visualization step.
+
+.. note::
+   By default the trained models and logs will be saved to ``IssacLab/logs/robomimic``.
+
+Visualize the results
+^^^^^^^^^^^^^^^^^^^^^
+
+Visualize the results of the trained policy by running the following command, using the normalization parameters recorded in the prior training step:
+
+.. code:: bash
+
+   ./isaaclab.sh -p scripts/imitation_learning/robomimic/play.py \
+   --device cuda \
+   --task Isaac-PickPlace-GR1T2-Abs-v0 \
+   --num_rollouts 50 \
+   --norm_factor_min <NORM_FACTOR_MIN> \
+   --norm_factor_max <NORM_FACTOR_MAX> \
+   --checkpoint /PATH/TO/desired_model_checkpoint.pth
+
+.. note::
+   Change the ``NORM_FACTOR`` in the above command with the values generated in the training step.
 
 Common Pitfalls when Generating Data
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -263,11 +390,13 @@ The config class :class:`~isaaclab_mimic.envs.FrankaCubeStackIKRelMimicEnvCfg` s
 
 The ``DataGenConfig`` member contains various parameters that influence how data is generated. It is initially sufficient to just set the ``name`` parameter, and revise the rest later.
 
-Subtasks are a list of ``SubTaskConfig`` objects, of which the most important members are:
+Subtasks are a list of :class:`~isaaclab.envs.SubTaskConfig` objects, of which the most important members are:
 
 * ``object_ref`` is the object that is being interacted with. This will be used to adjust motions relative to this object during data generation. Can be ``None`` if the current subtask does not involve any object.
 * ``subtask_term_signal`` is the ID of the signal indicating whether the subtask is active or not.
 
+For multi end-effector environments, subtask ordering between end-effectors can be enforced by specifying subtask constraints. These constraints are defined in the :class:`~isaaclab.envs.SubTaskConstraintConfig` class.
+
 Subtask annotation
 ^^^^^^^^^^^^^^^^^^
 

source/isaaclab/isaaclab/envs/mimic_env_cfg.py

@@ -77,7 +77,9 @@ class DataGenConfig:
 
 @configclass
 class SubTaskConfig:
-    """Configuration settings specific to the management of individual subtasks."""
+    """
+    Configuration settings for specifying subtasks used in Mimic environments.
+    """
 
     """Mandatory options that should be defined for every subtask."""
 
@@ -150,7 +152,9 @@ class SubTaskConstraintCoordinationScheme(enum.IntEnum):
 
 @configclass
 class SubTaskConstraintConfig:
-    """Configuration settings for subtask constraints."""
+    """
+    Configuration settings for specifying subtask constraints used in multi-eef Mimic environments.
+    """
 
     eef_subtask_constraint_tuple: list[tuple[str, int]] = (("", 0), ("", 0))
     """List of associated subtasks tuples in order.