hokindeng/object-permanence

解決する問題

このプロジェクトは、ワールドモデルにおける「物体の永続性」の課題に対処します。これは、AIが物体が視界から隠れている（遮蔽されている）場合でも存在し続けることを理解する能力です。初期ビデオクリップに基づいてシーンの将来の状態を予測するための合成データを生成し、モデルを訓練する包括的なフレームワークを提供します。

仕組み

システムは2つの主要コンポーネントに分かれています：

1. データファクトリー：ビデオからビデオへのサンプルを作成する150個のBlenderレンダリングタスクジェネレーターのセット。これらのサンプルは、入力ビデオ（分割点までのシミュレーション）、ターゲットビデオ（分割点以降のシミュレーション）、テキストプロンプト、およびグラウンドトゥルースの軌跡で構成されます。
2. トレーニングスタック：AWS Trainium2向けに最適化されたNVIDIA Cosmos3-NanoのネイティブPyTorch実装。テンソル並列性とFSDP2を使用して、生成された合成データでモデルを訓練し、遮蔽された物体に関する予測能力を向上させます。

対象読者

主に、ワールドモデル、ビデオ生成、機械学習における空間推論に取り組むAI研究者や開発者を対象としています。

ハイライト

• 大規模な合成データセット：高品質なトレーニングデータを作成するための150個の認知基盤タスクジェネレーターを含みます。
• 高性能トレーニング：ネイティブPyTorchを使用してAWS Trainium2向けに最適化され、XLAグラフトレーシングを回避します。
• 詳細なグラウンドトゥルース：正確な評価のために、シーン内のすべてのオブジェクトのフレームごとの軌跡を提供します。
• 統合パイプライン：Blenderによるデータ生成から大規模モデルトレーニングまでの完全なワークフローをカバーします。