Augmented Imitation Learning

This section describes how to use Isaac Lab’s imitation learning capabilities with the visual augmentation capabilities of 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.

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).

uv run --extra isaacsim,mimic python scripts/imitation_learning/isaaclab_mimic/generate_dataset.py \
--device cpu --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

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 CPU. On a more powerful desktop machine, use a larger number of environments for a significant speedup of this step. When running large camera-heavy jobs, switch to the RTX Minimal renderer for higher throughput, or override the task’s high-fidelity RTX settings directly if they are more expensive than your workflow requires.

Cosmos Augmentation

uv pip install opencv-python

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.

Required Arguments

Optional Arguments

Example usage for the cube stacking task:

uv run --extra mimic python scripts/tools/hdf5_to_mp4.py \
--input_file datasets/mimic_dataset_1k.hdf5 \
--output_dir datasets/mimic_dataset_1k_mp4

Running Cosmos for Visual Augmentation

After converting the demonstrations to MP4 format, you can use a Cosmos_ 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:

We provide an example augmentation output from Cosmos Transfer1 below:

We recommend using the 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. You can refer to the installation instructions, the checkpoint download instructions and this example for reference on how to use Transfer1 for this usecase. We further recommend the following settings to be used with the Transfer1 model for this task:

Hyperparameters

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.

Required Arguments

Optional Arguments

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.

Example command to use the Cosmos Transfer1 model for this usecase:

export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:=0}"
export CHECKPOINT_DIR="${CHECKPOINT_DIR:=./checkpoints}"
export NUM_GPU="${NUM_GPU:=1}"
PYTHONPATH=$(pwd) torchrun --nproc_per_node=$NUM_GPU --nnodes=1 --node_rank=0 cosmos_transfer1/diffusion/inference/transfer.py \
    --checkpoint_dir $CHECKPOINT_DIR \
    --video_save_folder outputs/cosmos_dataset_1k_mp4 \
    --controlnet_specs ./controlnet_specs/demo_0.json \
    --offload_text_encoder_model \
    --offload_guardrail_models \
    --num_gpus $NUM_GPU

Example ./controlnet_specs/demo_0.json json file to use with the above command:

{
    "prompt": "A robotic arm is picking up and stacking cubes inside a foggy industrial scrapyard at dawn, surrounded by piles of old robotic parts and twisted metal. The background includes large magnetic cranes, rusted conveyor belts, and flickering yellow floodlights struggling to penetrate the fog. The robot arm is bright teal with a glossy surface and silver stripes on the outer edges; the joints rotate smoothly and the pistons reflect a pale cyan hue. The robot arm is mounted on a table that is light oak wood with a natural grain pattern and a glossy varnish that reflects overhead lights softly; small burn marks dot one corner. The arm is connected to the base mounted on the table. The bottom cube is deep blue, the second cube is bright red, and the top cube is vivid green, maintaining their correct order after stacking. Sunlight pouring in from a large, open window bathes the table and robotic arm in a warm golden light. The shadows are soft, and the scene feels natural and inviting with a slight contrast between light and shadow.",
    "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.",
    "input_video_path" : "mimic_dataset_1k_mp4/demo_0_table_cam.mp4",
    "sigma_max": 50,
    "vis": {
        "input_control": "mimic_dataset_1k_mp4/demo_0_table_cam.mp4",
        "control_weight": 0.3
    },
    "edge": {
        "control_weight": 0.3
    },
    "depth": {
        "input_control": "mimic_dataset_1k_mp4/demo_0_table_cam_depth.mp4",
        "control_weight": 0.6
    },
    "seg": {
        "input_control": "mimic_dataset_1k_mp4/demo_0_table_cam_shaded_segmentation.mp4",
        "control_weight": 0.7
    }
}

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.

Required Arguments

Example usage for the cube stacking task:

uv run --extra mimic python scripts/tools/mp4_to_hdf5.py \
--input_file datasets/mimic_dataset_1k.hdf5 \
--videos_dir datasets/cosmos_dataset_1k_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 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.

Required Arguments

Optional Arguments

Example usage for the cube stacking task:

uv run --extra mimic python scripts/tools/merge_hdf5_datasets.py \
--input_files datasets/mimic_dataset_1k.hdf5 datasets/cosmos_dataset_1k.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 to train a policy. Any other framework or training method could be used.

To install the robomimic framework, use the following commands:

# install the dependencies
sudo apt install cmake build-essential
# resolve and verify Robomimic in the uv-managed environment
uv run --extra mimic python -c "import robomimic"

For a legacy environment, install the same dependencies with ./isaaclab.sh -i mimic.

Training an agent

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

uv run --extra isaacsim,mimic python scripts/imitation_learning/robomimic/train.py \
--task IsaacContrib-Stack-Cube-Franka-IK-Rel-Visuomotor-Cosmos --algo bc \
--dataset ./datasets/mimic_cosmos_dataset.hdf5 \
--name bc_rnn_image_franka_stack_mimic_cosmos

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:

Evaluation Settings

Required Arguments

Optional Arguments

Example usage for the cube stacking task:

uv run --extra isaacsim,mimic python scripts/imitation_learning/robomimic/robust_eval.py \
--task IsaacContrib-Stack-Cube-Franka-IK-Rel-Visuomotor-Cosmos \
--input_dir logs/robomimic/IsaacContrib-Stack-Cube-Franka-IK-Rel-Visuomotor-Cosmos/bc_rnn_image_franka_stack_mimic_cosmos/*/models \
--log_dir robust_results/bc_rnn_image_franka_stack_mimic_cosmos \
--log_file result \
--seeds 0 \
--num_rollouts 15

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:

Model Comparison

The above trained models’ checkpoints can be accessed here in case you wish to use the models directly.