增强模仿学习

本节描述了如何使用 Isaac Lab 的模仿学习功能，结合 Cosmos 模型的视觉增强功能，以生成大规模演示以训练对抗视觉变化的视觉动作策略。

生成演示

我们使用 Isaac Lab Mimic 功能，该功能允许从少数标注的演示中自动生成额外演示。

备注

本节假设您已经有一个收集到的演示数据集。如果没有，您可以按照 ​​Isaac Lab Mimic 中的遥操作与模仿学习​ 中的说明来收集和标注您自己的演示。

在以下示例中，我们将向您展示如何使用 Isaac Lab Mimic 生成额外的演示，这些演示可直接用于训练视觉运动策略，或者可以使用 Cosmos 进行视觉变化增强（使用 Isaac-Stack-Cube-Franka-IK-Rel-Visuomotor-Cosmos-Mimic-v0 环境）。

备注

Isaac-Stack-Cube-Franka-IK-Rel-Visuomotor-Cosmos-Mimic-v0 环境类似于标准的视觉动作环境（ Isaac-Stack-Cube-Franka-IK-Rel-Visuomotor-Mimic-v0 ），但在生成的数据集中添加了分割蒙版、深度图和法线图。这些额外的模态是使用 Cosmos 进行视觉增强时获得最佳结果所必需的。

./isaaclab.sh -p scripts/imitation_learning/isaaclab_mimic/generate_dataset.py \
--device cpu --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

演示次数可以增加或减少，已经证明展示1000次可以为这个任务提供良好的训练结果。

另外，可以调整 --num_envs 参数中环境的数量以加快数据生成。建议的数量为10，可以在一般笔记本电脑CPU上运行。在性能更强的台式机上，可以使用更多环境以显著加快此步骤。

Cosmos 增强

HDF5 到 MP4 转换

hdf5_to_mp4.py 脚本将存储在 HDF5 示范文件中的摄像头帧转换为 MP4 视频。它支持多种摄像头模式，包括 RGB、分割、深度和法线图。使用 Cosmos 进行视觉增强时，此转换是必要的，因为它仅适用于视频文件而不是 HDF5 数据。

必需参数

--input_file	输入HDF5文件的路径。
--output_dir	保存输出MP4文件的目录。

可选参数

--input_keys	从 HDF5 文件中处理的输入键列表。 (默认值: [“table_cam”, “wrist_cam”, “table_cam_segmentation”, “table_cam_normals”, “table_cam_shaded_segmentation”, “table_cam_depth”])
--video_height	输出视频的高度，以像素为单位。 (默认值: 704)
--video_width	输出视频的宽度（以像素为单位）。（默认值: 1280）
--framerate	输出视频的帧率。 (默认值: 30)

备注

默认输入键覆盖所有相机模式，按照 Isaac-Stack-Cube-Franka-IK-Rel-Visuomotor-Cosmos-Mimic-v0 环境中遵循的命名约定。我们包括一个额外的模式 “table_cam_shaded_segmentation” ，这不是从HDF5数据文件中通过仿真生成的模式的一部分。相反，它是由该脚本自动生成的，使用分割和法线贴图的组合来获取一个伪纹理分割视频，以更好地控制Cosmos增强。

备注

我们建议使用上面给出的默认值来获得与Cosmos增强最佳效果的输出视频高度、宽度和帧速率。

立方体堆叠任务的示例用法:

python scripts/tools/hdf5_to_mp4.py \
--input_file datasets/mimic_dataset_1k.hdf5 \
--output_dir datasets/mimic_dataset_1k_mp4

运行 Cosmos 用于视觉增强

将示范转换为 MP4 格式后，您可以使用 Cosmos 模型对视频进行视觉增强。请参考 Cosmos 文档以了解增强过程的详细信息。视觉增强可以包括对光照、纹理、背景和其他视觉元素的更改，同时保留基本的任务相关特征。

我们将上一步的 RGB、深度和阴影分割视频作为输入，输入到 Cosmos 模型中，如下所示:

我们在下面提供了来自 Cosmos Transfer1 的示例增强输出:

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 here, the checkpoint download instructions here 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:

超参数

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.”
sigma_max	50
control_weight	“0.3,0.3,0.6,0.7”
hint_key	“blur,canny,depth,segmentation”

另一个关键方面是获取良好的增强的关键是用于控制Cosmos生成的prompts集。我们提供了一个脚本， cosmos_prompt_gen.py ，用于从一组精心选择的模板中构建prompts，处理增强过程的各个方面。

必需参数

--templates_path

包含prompts模板的文件路径。

可选参数

--num_prompts	生成prompts数量（默认值: 1）。
--output_path	生成prompts写入的输出文件路径。 (默认值: prompts.txt)

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

如果您想要创建自己的prompts，则建议您参考以下准则:

1. 尽量详细地保留prompts。最好对生成过程中每个可见对象/感兴趣区域都有一些指导。例如，我们提供的prompts涵盖了桌子、灯光、背景、机械臂、立方体和整体布景的细节。

2. 尽量使增强指令尽可能地真实和连贯。模型在保留输入控制视频的关键特征方面做得越不真实或不传统，它的表现就越糟糕。

3. 保持每个方面的增强指令同步。我们的意思是，所有对象/感兴趣区域的增强应该在彼此之间是连贯和传统的。例如，最好是有一个prompts，例如 “The table is of old dark wood with faded polish and food stains and the background consists of a suburban home” ，而不是像 “The table is of old dark wood with faded polish and food stains and the background consists of a spaceship hurtling through space” 。

4. 在输入控制视频的关键方面，必须包括应保留或保持不变的细节。在我们的prompts中，我们非常明确地提到立方体的颜色应保持不变，使得底部立方体是蓝色，中间是红色，顶部是绿色。请注意，我们不仅提到了应该保持不变的内容，还详细说明了该方面当前的形式。

示例命令，用于在此用例中使用 Cosmos Transfer1 模型:

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

示例 ./controlnet_specs/demo_0.json json 文件可用于上述命令:

{
    "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 到 HDF5 转换

mp4_to_hdf5.py 脚本将经过视觉增强的MP4视频转换回HDF5格式，用于训练。这一步非常关键，因为它确保了增强的视觉数据以正确格式用于在Isaac Lab中训练视觉动作策略，并将视频与原始数据集中相应的演示数据配对。

必需参数

--input_file	包含原始演示的输入 HDF5 文件的路径。
--videos_dir	包含视觉增强的MP4视频的目录。
--output_file	保存新的带有增强视频的HDF5文件的路径。

备注

输入 HDF5 文件用于保留非视觉数据（例如机器人状态和动作），同时用增强版本替换视觉数据。

重要

视觉增强的 MP4 文件必须遵循命名约定 demo_{demo_id}_*.mp4 ，其中:

• demo_id 匹配原始 MP4 文件中的演示ID

• * 表示文件名可以根据用户的偏好从此处开始。

这种命名惯例是为了让脚本能够正确地将增强视频与相应的演示配对。

立方体堆叠任务的示例用法:

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

预生成的数据集

我们提供了一个预先生成的数据集，其中包含了用于立方体堆叠任务的视觉增强演示，格式为HDF5。 如果您不希望在本地运行Cosmos来生成自己的增强数据，可以使用此数据集。 该数据集位于 Hugging Face ，包含原始和增强演示（作为单独的数据集文件），可用于训练视觉动作策略。

合并数据集

merge_hdf5_datasets.py 脚本将多个 HDF5 数据集合并为单个文件。当您希望将原始演示与增强演示合并以创建更大、更多样化的训练数据集时，这将非常有用。

必需参数

--input_files

将要合并的HDF5文件路径列表。

可选参数

--output_file

合并输出的文件路径。（默认值: merged_dataset.hdf5）

小技巧

合并数据集可以通过在训练过程中使模型接触原始和增强的视觉条件，从而帮助提高政策的稳健性。

立方体堆叠任务的示例用法:

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

模型训练和评估

Robomimic 设置

作为一个示例，我们将训练一个在 Robomimic 中实现的 BC 智能体来训练一个策略。可以使用任何其他框架或训练方法。

要安装robomimic框架，请使用以下命令:

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

训练一个智能体

使用生成的数据，现在我们可以为 Isaac-Stack-Cube-Franka-IK-Rel-Visuomotor-Cosmos-v0 训练一个视觉-运动 BC 智能体。

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

备注

默认情况下，训练模型和日志将保存到 IssacLab/logs/robomimic 。

评估

robust_eval.py 脚本在仿真环境中评估经过训练的视觉动作策略。这种评估有助于评估策略在不同视觉变化下的泛化能力，以及视觉增强数据是否提高了策略的稳健性。

以下是用于评估的不同设置的解释:

评估设置

Vanilla	与生成Mimic数据时使用的完全相同设置。
Light Intensity	光强/亮度变化，其他方面保持不变。
Light Color	浅色变化，其他方面保持不变。
Light Texture (Background)	轻质纹理/背景多样化，所有其他方面保持不变。
Table Texture	桌子的视觉质感多种多样，其他方面保持不变。
Robot Arm Texture	机器人手臂的视觉纹理是多样的，其他所有方面保持不变。

必需参数

--task	环境名称。
--input_dir	包含要评估的模型检查点的目录。

可选参数

--start_epoch	开始评估的检查点纪元。(默认值: 100)
--horizon	每次执行的步骤范围。 (默认值: 400)
--num_rollouts	每个模型每个设置的 rollout 次数。 (默认值: 15)
--num_seeds	要评估的随机种子数量。 （默认值: 3）
--seeds	替代随机种子使用的特定种子列表。
--log_dir	写入结果的目录。 (默认: /tmp/policy_evaluation_results)
--log_file	输出文件的名称。 (默认: results)
--norm_factor_min	动作空间归一化因子的最小值。
--norm_factor_max	动作空间归一化因子的最大值。
--disable_fabric	是否禁用 fabric 并使用 USD I/O 操作。
--enable_pinocchio	是否启用 IK 控制器的 Pinocchio。

备注

评估结果将帮助您了解视觉增强是否提高了策略的性能和鲁棒性。将这些结果与在原始数据集上的评估结果进行比较，以衡量增强的影响。

立方体堆叠任务的示例用法:

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

备注

这个脚本可能需要一天甚至更长的时间才能运行（取决于使用的硬件）。这种行为是正常的。

我们使用上述脚本来比较分别经过1000次Mimic生成的演示、2000次Mimic生成的演示和2000次Cosmos-Mimic生成的演示（1000个原始mimic + 1000个Cosmos增强）训练的模型。我们对这三个模型使用相同的种子（0、1000和5000），并提供以下度量指标（在每个种子的最佳检查点上进行平均）:

模型比较

评估设置	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%

以上训练模型的检查点可以通过 这里 进行访问，以便直接使用模型。