.. _augmented-imitation-learning:

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.

Generating Demonstrations
~~~~~~~~~~~~~~~~~~~~~~~~~

We use the Isaac Lab Mimic feature that allows the generation of additional demonstrations automatically from a handful of annotated demonstrations.

.. note::
    This section assumes you already have an annotated dataset of collected demonstrations. If you don't, you can follow the instructions in :ref:`teleoperation-imitation-learning` to collect and annotate your own demonstrations.

In the following example, we will show you how to use Isaac Lab Mimic to generate additional demonstrations that can be used to train a visuomotor policy directly or can be augmented with visual variations using Cosmos (using the ``Isaac-Stack-Cube-Franka-IK-Rel-Visuomotor-Cosmos-Mimic-v0`` environment).

.. note::
    The ``Isaac-Stack-Cube-Franka-IK-Rel-Visuomotor-Cosmos-Mimic-v0`` environment is similar to the standard visuomotor environment (``Isaac-Stack-Cube-Franka-IK-Rel-Visuomotor-Mimic-v0``), but with the addition of segmentation masks, depth maps, and normal maps in the generated dataset. These additional modalities are required to get the best results from the visual augmentation done using Cosmos.

.. 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/mimic_dataset_1k.hdf5 \
    --task Isaac-Stack-Cube-Franka-IK-Rel-Visuomotor-Cosmos-Mimic-v0 \
    --rendering_mode performance

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 on a moderate laptop GPU.
On a more powerful desktop machine, use a larger number of environments for a significant speedup of this step.

Cosmos Augmentation
~~~~~~~~~~~~~~~~~~~

HDF5 to MP4 Conversion
^^^^^^^^^^^^^^^^^^^^^^

The ``hdf5_to_mp4.py`` script converts camera frames stored in HDF5 demonstration files to MP4 videos. It supports multiple camera modalities including RGB, segmentation, depth and normal maps. This conversion is necessary for visual augmentation using Cosmos as it only works with video files rather than HDF5 data.

.. rubric:: Required Arguments

.. list-table::
    :widths: 30 70
    :header-rows: 0

    * - ``--input_file``
      - Path to the input HDF5 file.
    * - ``--output_dir``
      - Directory to save the output MP4 files.

.. rubric:: Optional Arguments

.. list-table::
    :widths: 30 70
    :header-rows: 0

    * - ``--input_keys``
      - List of input keys to process from the HDF5 file. (default: ["table_cam", "wrist_cam", "table_cam_segmentation", "table_cam_normals", "table_cam_shaded_segmentation", "table_cam_depth"])
    * - ``--video_height``
      - Height of the output video in pixels. (default: 704)
    * - ``--video_width``
      - Width of the output video in pixels. (default: 1280)
    * - ``--framerate``
      - Frames per second for the output video. (default: 30)

.. note::
    The default input keys cover all camera modalities as per the naming convention followed in the ``Isaac-Stack-Cube-Franka-IK-Rel-Visuomotor-Cosmos-Mimic-v0`` environment. We include an additional modality "table_cam_shaded_segmentation" which is not a part of the generated modalities from simulation in the HDF5 data file. Instead, it is automatically generated by this script using a combination of the segmentation and normal maps to get a pseudo-textured segmentation video for better controlling the Cosmos augmentation.

.. note::
    We recommend using the default values given above for the output video height, width and framerate for the best results with Cosmos augmentation.

Example usage for the cube stacking task:

.. code:: bash

    python scripts/tools/hdf5_to_mp4.py \
    --input_file datasets/mimic_generated_dataset.hdf5 \
    --output_dir datasets/mimic_generated_dataset_mp4

Running Cosmos for Visual Augmentation
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

After converting the demonstrations to MP4 format, you can use a `Cosmos <https://github.com/NVIDIA/Cosmos?tab=readme-ov-file>`_ model to visually augment the videos. Follow the Cosmos documentation for details on the augmentation process. Visual augmentation can include changes to lighting, textures, backgrounds, and other visual elements while preserving the essential task-relevant features.

We use the RGB, depth and shaded segmentation videos from the previous step as input to the Cosmos model as seen below:

.. figure:: https://download.isaacsim.omniverse.nvidia.com/isaaclab/images/cosmos_inputs.gif
   :width: 100%
   :align: center
   :alt: RGB, depth and segmentation control inputs to Cosmos

We provide an example augmentation output from `Cosmos Transfer1 <https://github.com/nvidia-cosmos/cosmos-transfer1>`_ below:

.. figure:: https://download.isaacsim.omniverse.nvidia.com/isaaclab/images/cosmos_output.gif
   :width: 100%
   :align: center
   :alt: Cosmos Transfer1 augmentation output

We recommend using the `Cosmos Transfer1 <https://github.com/nvidia-cosmos/cosmos-transfer1>`_ model for visual augmentation as we found it to produce the best results in the form of a highly diverse dataset with a wide range of visual variations. We further recommend the following settings to be used with the Transfer1 model for this task:

.. rubric:: Hyperparameters

.. list-table::
    :widths: 30 70
    :header-rows: 0

    * - ``negative_prompt``
      - "The video captures a game playing, with bad crappy graphics and cartoonish frames. It represents a recording of old outdated games. The images are very pixelated and of poor CG quality. There are many subtitles in the footage. Overall, the video is unrealistic and appears cg. Plane background."
    * - ``positive_prompt``
      - "realistic, photorealistic, high fidelity, varied lighting, varied background"
    * - ``sigma_max``
      - 50
    * - ``control_weight``
      - "0.3,0.3,0.6,0.7"
    * - ``hint_key``
      - "blur,canny,depth,segmentation"
    * - ``control_input_preset_strength``
      - "low"

Another crucial aspect to get good augmentations is the set of prompts used to control the Cosmos generation. We provide a script, ``cosmos_prompt_gen.py``, to construct prompts from a set of carefully chosen templates that handle various aspects of the augmentation process.

.. rubric:: Required Arguments

.. list-table::
    :widths: 30 70
    :header-rows: 0

    * - ``--templates_path``
      - Path to the file containing templates for the prompts.

.. rubric:: Optional Arguments

.. list-table::
    :widths: 30 70
    :header-rows: 0

    * - ``--num_prompts``
      - Number of prompts to generate (default: 1).
    * - ``--output_path``
      - Path to the output file to write generated prompts. (default: prompts.txt)

.. code:: bash

    python scripts/tools/cosmos/cosmos_prompt_gen.py \
    --templates_path scripts/tools/cosmos/transfer1_templates.json \
    --num_prompts 10 --output_path prompts.txt

In case you want to create your own prompts, we suggest you refer to the following guidelines:

1. Keep the prompts as detailed as possible. It is best to have some instruction on how the generation should handle each visible object/region of interest. For instance, the prompts that we provide cover explicit details for the table, lighting, background, robot arm, cubes, and the general setting.

2. Try to keep the augmentation instructions as realistic and coherent as possible. The more unrealistic or unconventional the prompt is, the worse the model does at retaining key features of the input control video(s).

3. Keep the augmentation instructions in-sync for each aspect. What we mean by this is that the augmentation for all the objects/regions of interest should be coherent and conventional with respect to each other. For example, it is better to have a prompt such as "The table is of old dark wood with faded polish and food stains and the background consists of a suburban home" instead of something like "The table is of old dark wood with faded polish and food stains and the background consists of a spaceship hurtling through space".

4. It is vital to include details on key aspects of the input control video(s) that should be retained or left unchanged. In our prompts, we very clearly mention that the cube colors should be left unchanged such that the bottom cube is blue, the middle is red and the top is green. Note that we not only mention what should be left unchanged but also give details on what form that aspect currently has.

MP4 to HDF5 Conversion
^^^^^^^^^^^^^^^^^^^^^^

The ``mp4_to_hdf5.py`` script converts the visually augmented MP4 videos back to HDF5 format for training. This step is crucial as it ensures the augmented visual data is in the correct format for training visuomotor policies in Isaac Lab and pairs the videos with the corresponding demonstration data from the original dataset.

.. rubric:: Required Arguments

.. list-table::
    :widths: 30 70
    :header-rows: 0

    * - ``--input_file``
      - Path to the input HDF5 file containing the original demonstrations.
    * - ``--videos_dir``
      - Directory containing the visually augmented MP4 videos.
    * - ``--output_file``
      - Path to save the new HDF5 file with augmented videos.

.. note::
    The input HDF5 file is used to preserve the non-visual data (such as robot states and actions) while replacing the visual data with the augmented versions.

.. important::
    The visually augmented MP4 files must follow the naming convention ``demo_{demo_id}_*.mp4``, where:

    - ``demo_id`` matches the demonstration ID from the original MP4 file

    - ``*`` signifies that the file name can be as per user preference starting from this point

    This naming convention is required for the script to correctly pair the augmented videos with their corresponding demonstrations.

Example usage for the cube stacking task:

.. code:: bash

    python scripts/tools/mp4_to_hdf5.py \
    --input_file datasets/mimic_generated_dataset.hdf5 \
    --videos_dir datasets/cosmos_dataset_mp4 \
    --output_file datasets/cosmos_dataset_1k.hdf5

Pre-generated Dataset
^^^^^^^^^^^^^^^^^^^^^

We provide a pre-generated dataset in HDF5 format containing visually augmented demonstrations for the cube stacking task. This dataset can be used if you do not wish to run Cosmos locally to generate your own augmented data. The dataset is available on `Hugging Face <https://huggingface.co/datasets/nvidia/PhysicalAI-Robotics-Manipulation-Augmented>`_ and contains both (as separate dataset files), original and augmented demonstrations, that can be used for training visuomotor policies.

Merging Datasets
^^^^^^^^^^^^^^^^

The ``merge_hdf5_datasets.py`` script combines multiple HDF5 datasets into a single file. This is useful when you want to combine the original demonstrations with the augmented ones to create a larger, more diverse training dataset.

.. rubric:: Required Arguments

.. list-table::
    :widths: 30 70
    :header-rows: 0

    * - ``--input_files``
      - A list of paths to HDF5 files to merge.

.. rubric:: Optional Arguments

.. list-table::
    :widths: 30 70
    :header-rows: 0

    * - ``--output_file``
      - File path to merged output. (default: merged_dataset.hdf5)

.. tip::
    Merging datasets can help improve policy robustness by exposing the model to both original and augmented visual conditions during training.

Example usage for the cube stacking task:

.. code:: bash

    python scripts/tools/merge_hdf5_datasets.py \
    --input_files datasets/mimic_generated_dataset.hdf5 datasets/cosmos_dataset.hdf5 \
    --output_file datasets/mimic_cosmos_dataset.hdf5

Model Training and Evaluation
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

Robomimic Setup
^^^^^^^^^^^^^^^

As an example, we will train a BC agent implemented in `Robomimic <https://robomimic.github.io/>`__ to train a policy. Any other framework or training method could be used.

To install the robomimic framework, use the following commands:

.. code:: bash

   # install the dependencies
   sudo apt install cmake build-essential
   # install python module (for robomimic)
   ./isaaclab.sh -i robomimic

Training an agent
^^^^^^^^^^^^^^^^^

Using the generated data, we can now train a visuomotor BC agent for ``Isaac-Stack-Cube-Franka-IK-Rel-Visuomotor-v0``:

.. code:: bash

    ./isaaclab.sh -p scripts/imitation_learning/robomimic/train.py \
    --task Isaac-Stack-Cube-Franka-IK-Rel-Visuomotor-v0 --algo bc \
    --dataset ./datasets/mimic_cosmos_dataset.hdf5

.. note::
   By default the trained models and logs will be saved to ``IssacLab/logs/robomimic``.

Evaluation
^^^^^^^^^^

The ``robust_eval.py`` script evaluates trained visuomotor policies in simulation. This evaluation helps assess how well the policy generalizes to different visual variations and whether the visually augmented data has improved the policy's robustness.

Below is an explanation of the different settings used for evaluation:

.. rubric:: Evaluation Settings

.. list-table::
    :widths: 30 70
    :header-rows: 0

    * - ``Vanilla``
      - Exact same setting as that used during Mimic data generation.
    * - ``Light Intensity``
      - Light intensity/brightness is varied, all other aspects remain the same.
    * - ``Light Color``
      - Light color is varied, all other aspects remain the same.
    * - ``Light Texture (Background)``
      - Light texture/background is varied, all other aspects remain the same.
    * - ``Table Texture``
      - Table's visual texture is varied, all other aspects remain the same.
    * - ``Robot Arm Texture``
      - Robot arm's visual texture is varied, all other aspects remain the same.

.. rubric:: Required Arguments

.. list-table::
    :widths: 30 70
    :header-rows: 0

    * - ``--task``
      - Name of the environment.
    * - ``--input_dir``
      - Directory containing the model checkpoints to evaluate.

.. rubric:: Optional Arguments

.. list-table::
    :widths: 30 70
    :header-rows: 0

    * - ``--horizon``
      - Step horizon of each rollout. (default: 400)
    * - ``--num_rollouts``
      - Number of rollouts per model per setting. (default: 15)
    * - ``--num_seeds``
      - Number of random seeds to evaluate. (default: 3)
    * - ``--seeds``
      - List of specific seeds to use instead of random ones.
    * - ``--log_dir``
      - Directory to write results to. (default: /tmp/policy_evaluation_results)
    * - ``--log_file``
      - Name of the output file. (default: results)
    * - ``--norm_factor_min``
      - Minimum value of the action space normalization factor.
    * - ``--norm_factor_max``
      - Maximum value of the action space normalization factor.
    * - ``--disable_fabric``
      - Whether to disable fabric and use USD I/O operations.
    * - ``--enable_pinocchio``
      - Whether to enable Pinocchio for IK controllers.

.. note::
    The evaluation results will help you understand if the visual augmentation has improved the policy's performance and robustness. Compare these results with evaluations on the original dataset to measure the impact of augmentation.

Example usage for the cube stacking task:

.. code:: bash

    ./isaaclab.sh -p scripts/imitation_learning/robomimic/robust_eval.py \
    --task Isaac-Stack-Cube-Franka-IK-Rel-Visuomotor-v0 \
    --input_dir logs/robomimic/Isaac-Stack-Cube-Franka-IK-Rel-Visuomotor-v0/bc_rnn_image_franka_stack_mimic_cosmos_table_only/*/models \
    --log_dir robust_results/bc_rnn_image_franka_stack_mimic_cosmos_table_only \
    --log_file result \
    --enable_cameras \
    --seeds 0 \
    --num_rollouts 15 \
    --rendering_mode performance

We use the above script to compare models trained with 1000 Mimic-generated demonstrations, 2000 Mimic-generated demonstrations and 2000 Cosmos-Mimic-generated demonstrations (1000 original mimic + 1000 Cosmos augmented) respectively. We use the same seeds (0, 1000 and 5000) for all three models and provide the metrics (averaged across best checkpoints for each seed) below:

.. rubric:: Model Comparison

.. list-table::
    :widths: 25 25 25 25
    :header-rows: 0

    * - **Evaluation Setting**
      - **Mimic 1k Baseline**
      - **Mimic 2k Baseline**
      - **Cosmos-Mimic 2k**
    * - ``Vanilla``
      - 62%
      - 96.6%
      - 86.6%
    * - ``Light Intensity``
      - 11.1%
      - 20%
      - 62.2%
    * - ``Light Color``
      - 24.6%
      - 30%
      - 77.7%
    * - ``Light Texture (Background)``
      - 16.6%
      - 20%
      - 68.8%
    * - ``Table Texture``
      - 0%
      - 0%
      - 20%
    * - ``Robot Arm Texture``
      - 0%
      - 0%
      - 4.4%

The above trained models' checkpoints can be accessed `here <https://huggingface.co/datasets/nvidia/PhysicalAI-Robotics-Manipulation-Augmented/tree/main/robomimic_bc_rnn_visuomotor_models>`_ in case you wish to use the models directly.
