こんにちは！ABEJA で ABEJA Platform 開発を行っている坂井（@Yagami360）です。
以下の記事では、Cosmos を実際に動かしてみながら Cosmos により物理法則が考慮されたフォトリアリスティックな動画生成（合成データ生成）が可能であることを確かめました。
しかしながら Cosmos 単体では単に動画生成を行うだけなので、これだけではロボティクスモデル（VLAなど）用の学習用データセットとしては利用できなく、実際のロボティクス開発において Cosmos を活用できません。
そのため本記事では、Cosmos と NVIDIA Omniverse の Isaac Sim & Lab を組み合わせることにより、ロボティクスモデル用の学習用データセット生成（合成データ生成）として Cosmos を活用する方法を紹介します。
NVIDIA Omniverse や Isaac Sim & Lab については、以下の記事をご確認ください
また、以前の下記ブログで説明した シミュレーター上での手動遠隔操作や Isaac Mimic を使用したデータ生成手法もそのまま利用します
- NVIDIA Omniverse と Cosmos を活用したロボティクス向け合成データ生成のワークフロー
- Isaac Sim のシミュレーター上の遠隔操作による少数データセット作成
- Isaac Mimic による大量データセット自動生成
- Cosmos Transfer のための前処理を行なう
- Cosmos Transfer によるデータ拡張
- VLAモデル用学習用データセット形式への変換
- まとめ
- We Are Hiring!
NVIDIA Omniverse と Cosmos を活用したロボティクス向け合成データ生成のワークフロー
今回紹介するロボティクス用の学習用データセットとしての合成データ作成のワークフローは、以下の図のようになります。
- Isaac Sim のシミュレーター上で人手による遠隔操作で少数のデータセットを作成する
- Isaac Mimic を使用して AI モデルで大量エピソードのデータセットを自動生成する
- 上記大量生成した各エピソードのデータセットに対して、Cosmos Transfer でのフォトリアリスティックなデータ拡張を行い、データセットを更にかさ増しする
- VLAモデル用の学習用データセットに変換する
この方法の一部は、「Isaac GR00T Blueprint」として、以下の blueprint でのデモとして公開されている情報を参考にしています。
なお本ワークフローのうち、① と ② によるデータ生成の方法は既に以前のブログで紹介済みなので、ブログのリンク先を紹介するのにとどめます。
Isaac Sim のシミュレーター上の遠隔操作による少数データセット作成
以下ブログ記載の方法で、Isaac Sim のシミュレーター上での人手による手動遠隔操作で少数エピソード分の学習用データセットを作成します。
Isaac Mimic による大量データセット自動生成
以下ブログ記載の方法で、上記手動遠隔操作で作成した少数エピソードの学習用データセットから Isaac Mimic を使用して AI モデルで大量エピソードのデータセットを自動生成します。
Cosmos Transfer のための前処理を行なう
ここからが、本記事で新規に行なう部分になります。
Isaac Mimic を使用して大量エピソードの学習用データセットが得られますが、このデータセットは HDF5形式の学習用データセットでありまた内部にロボットからのカメラ画像も含まれていないので、このままでは Cosmos の世界基盤モデルに入力できません。
そのため、HDF5形式の学習用データセットを Isaac Sim のシミュレーター上のロボットでリプレイ（動かす）しながら、Cosmos に入力するためのカメラからの｛RGB動画・セグメンテーション動画｝を作成します
カメラを配置した環境のシミュレーターを起動する
上記 Isaac Mimic で使用していた環境（
Isaac-Stack-Cube-Franka-IK-Rel-Mimic-v0）にtable_high_camとtable_camの２つのカメラを配置した環境（Isaac-Stack-Cube-Franka-IK-Rel-Blueprint-Mimic-v0）を使用してシミュレーターを起動します。この際のカメラの設定は、Isaac Sim の機能を利用して法線マップ画像とセグメンテーション画像をレンダリングできる設定にされています。今回の例では、テーブル上のカメラ（
table_high_cam）のみ使用します。※ これら法線マップ画像とセグメンテーション画像は、Cosmos での生成動画の品質を向上させる動画作成ために使用します。
HDF5形式の学習用データセットでシミュレーター上のロボットをリプレイしながら、カメラから法線マップとセグメンテーションのフレーム画像を作成する
Isaac Mimic で作成したHDF5形式の学習用データセットでリプレイします。今回は、１エピソードのみ分リプレイします
リプレイ中に、シミュレーター上のカメラから各フレーム毎に｛法線マップ画像・セグメンテーション画像｝がリアルタイムに作成され、以下のようにある１エピソードに対しての各時間ステップ毎のフレーム画像リストが得られます。
フレーム画像リストから Cosmos に入力するための動画を作成する
上記作成したフレーム画像リストから Cosmos の世界基盤モデルに入力するための動画を作成します。
この際に、以下のようなオブジェクトのセグメンテーションや凹凸情報のみをもった動画を入力したほうが Cosmos での生成動画の品質が向上するため、セグメンテーション情報に法線マップ情報を適用してレンダリングした動画を作成します。
テーブル上部カメラ（table_high_cam）
Cosmos Transfer によるデータ拡張
上記前処理で得られたカメラから動画を入力として、Cosmos Transfer でのフォトリアリスティックなデータ拡張を行いデータセットを更にかさ増しします。
特にここまでのワークフローで作成したデータセットは、あくまでシミュレーター上で作成したデータセットですが、Cosmos の世界基盤モデルによるデータ拡張では、物理法則が考慮されたフォトリアリスティックな動画が得られるので現実世界でロボットを動かす際の汎化性能が向上することが期待できます。
以下ブログ記載の方法で Cosmos Transfer を動かすための環境構築を行ってください
必要ファイルを配置する
assets/Isaac-Stack-Cube-Franka-IK-Rel-Blueprint-Mimic-v0ディレクトリ以下に、上記作成したカメラからのセグメンテーション動画を配置してください。また推論時の設定用 JSON ファイルも配置してください設定用 JSON ファイル（
assets/Isaac-Stack-Cube-Franka-IK-Rel-Blueprint-Mimic-v0/cosmos_param_1.json）{ "prompt": "The scene depicts a robotic arm performing a precise block-stacking operation with glass cubes in a cluttered workshop. The workspace is centered on a wood table. On the table are three cubes. The robotic arm, featuring a sleek industrial design clad in smooth white plastic with metallic joints, extends from its base to manipulate the cubes. The entire scene is illuminated by professional studio lighting, creating soft shadows. A clutter of equipment can be seen in the background. The camera is fixed in place.", "input_video_path" : "assets/Isaac-Stack-Cube-Franka-IK-Rel-Blueprint-Mimic-v0/shaded_segmentation_table_high_cam_trial_0_tile_0.mp4", "seg": { "input_control": "assets/Isaac-Stack-Cube-Franka-IK-Rel-Blueprint-Mimic-v0/shaded_segmentation_table_high_cam_trial_0_tile_0.mp4", "control_weight": 1.0 } }
コンテナに接続する
docker run -it \ -v $(PWD)/cosmos-transfer1:/workspace \ -v $(PWD)/cosmos-transfer1/checkpoints:/workspace/checkpoints \ --gpus all \ nvcr.io/$USER/cosmos-transfer1:latestCosmos Transfer 世界基盤モデルの推論スクリプトを実行する
コンテナ内で以下のスクリプトを実行してください
export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:=0}" export CHECKPOINT_DIR="${CHECKPOINT_DIR:=./checkpoints}" export NUM_GPU="${NUM_GPU:=1}" # 推論スクリプトを実行する # GPUメモリに余裕があれば --offload_xxx 引数なしにすることで推論時間を短縮可能 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/Isaac-Stack-Cube-Franka-IK-Rel-Blueprint-Mimic-v0 \ --controlnet_specs assets/Isaac-Stack-Cube-Franka-IK-Rel-Blueprint-Mimic-v0/cosmos_param_1.json \ --offload_text_encoder_model \ --offload_guardrail_models \ --seed 42 \ --num_gpus $NUM_GPUスクリプト実行後に、以下のようなフォトリアリスティックにデータ拡張された動画が出力されました。
※ 動画の最後の方で画面が暗転してますが、単に Cosmos 世界基盤モデルの品質の問題なのか推論パラメータの問題なのかよくわからなかったです。
入力テキスト 入力動画 出力動画 The scene depicts a robotic arm performing a precise block-stacking operation with glass cubes in a cluttered workshop. The workspace is centered on a wood table. ... 
入力テキストを “The scene depicts a robotic arm performing a precise block-stacking operation with brightly colored plastic cubes in a manufacturing facility. The workspace is centered on a metal table. …” といった別パターンにしてみると、以下のような異なるキューブ素材やテーブル素材でデータ拡張された動画が出力されました。
入力テキスト 入力動画 出力動画 The scene depicts a robotic arm performing a precise block-stacking operation with plastic cubes in a manufacturing facility. The workspace is centered on a metal table. ... 
Isaac Sim 上や Isaac Lab Mimic で作成した学習用データセットは、あくまでシミュレーター上で作成したデータセットであり、カメラ画像（カメラ動画）が現実世界とは異なる簡易的なものでしたが、Cosmos で世界基盤モデルでの生成される合成データはこのように現実世界に近い画像や動画になっているので、現実世界でロボットを動かした際の汎化性能が向上することが期待できます。
そして、これらの操作を学習用データセットの全エピソードに対して繰り返し、以下のように各エピソードに対してデータ拡張を行ったデータセットを作成していきます。
※ 以下の生成動画には品質が悪い動画も混じってますが、品質の高い学習用データセットとして作成する場合はこういった動画はクレンジングする必要があります
| エピソード | カメラ画像 （データ拡張前） |
カメラ画像 （データ拡張後１） |
カメラ画像 （データ拡張後２） |
・・・ |
|---|---|---|---|---|
| 1 | |
|
|
... |
| 2 |  |
 |
 |
... |
| … | … | ... | ... | ... |
VLAモデル用学習用データセット形式への変換
ここまでの方法で、学習用データセットの各エピソードに対してのシミュレーター上のカメラ動画に対する Comsos でのデータ拡張を行えましたが、このままでは以前としてロボティクスの VLA モデルの学習用データセットとしては利用できません。
VLA モデルの学習用データセットとして利用するには、各エピソードと時間ステップに対しての｛ロボットの状態・ロボットの行動・観測情報（カメラ画像）・制御テキスト｝のペアデータにする必要があります。
そのため Cosmos のデータ拡張で得られた生成動画と Isaac Mimic で作成したHDF5形式の学習用データセットを元に、以下のような学習用データセットの形式に変換します。（※ 但し、このような変換スクリプトは現時点で NVIDIA Omniverse や Cosmos で提供されていないようなので自作する必要ありです）
こうして得られた学習用データセットでVLA モデルをはじめて学習可能になり、このデータセットでモデルを学習すれば合成データで学習した VLA モデルを開発することができます。
まとめ
前回記事での NVIDIA Cosmos 単体利用では、動画データとしての合成データ生成はできても、実際に VLA などのロボティクスモデルで学習可能な合成データ生成はできない状態でしたが、今回記事のように NVIDIA Omniverse と Cosmos を双方向に利用すれば、ロボティクス VLA モデルの学習用データセットとしての合成データ生成も確かに可能であることがわかりました。
一方で、Cosmos の世界基盤モデルでの推論には、A100 などの Ampere 世代以降の GPU が必要で、GPUメモリも50GB 以上とかなり高価なマシンが必要になり、中々そのような環境をもっている人はいないよなという所感でした。
また１エピソードあたりの推論時間もかなり長く（今回の自分の環境では１パターンのデータ拡張だけで30分程度かかりました）、大量エピソードの学習用データセットに対して推論を行おうとすると膨大な時間がかかる（厳密に測定してませんが月レベルでかかる印象）ので、その点が大きな問題になると思いました。
実際に本記事ではこのデータ拡張処理に時間がかかりすぎる点が問題になり、Cosmos Transfer でデータ拡張した大量エピソードの学習用データセットで VLA モデルを学習させて、Cosmos Transfer によるデータ拡張でどの程度汎化性能が向上させれるのか？といったところまでの検証ができなかったです。
また以下のように生成動画の品質が悪いケース（キューブの数が不自然に増えたり途中で色が変わったりするなど）もかなりありました。
推論パラメータの調整で解決できるのかもしれませんが、品質のよい学習用データセットとして作成するには、結局人手による目視での品質評価が必要になって現状ではまだ全てを自動化できない印象がありました。
ここらへんの品質問題は、Cosmos Reason での生成動画の評価機能も触ってみて判断しようかと思いました。
We Are Hiring!
株式会社ABEJAでは共に働く仲間を募集しています！
ロボティクスやLLMに興味ある方々！機械学習プロダクトに関わるフロントエンド開発やバックエンド開発に興味ある方々！ こちらの採用ページから是非ご応募くださいませ！
