NVIDIA Cosmos-H-Dreams: 手術ロボット工学のためのリアルタイム生成シミュレーション

TL;DR

NVIDIAは、手術ロボット工学のためのリアルタイム・アクション条件付き生成シミュレータ「Cosmos-H-Dreams」をリリースしました。Cosmos-H-Surgical-Simulatorを因果的（causal）な生徒モデルに蒸留し、FlashDreams推論エンジンを活用することで、NVIDIAは単一のGPU上で160 FPS（フレーム/秒）のインタラクティブなクローズドループ手術シーンのシミュレーションを実現しました。

ワールドモデルからインタラクティブ・シミュレータへの移行

Cosmos-H-Dreamsは、NVIDIA Cosmos-Predict2.5-2Bに基づきOpen-H-Embodimentデータセットでトレーニングされたアクション条件付きワールド基盤モデルである、従来のCosmos-H-Surgical-Simulatorを進化させたものです。元のシミュレータは主にオフラインのポリシー評価や合成データの生成に使用されていましたが、Cosmos-H-Dreamsはこれらの機能をリアルタイムの領域へと移行させます。

具体的には、このモデルは da Vinci Research Kit (dVRK) のテーブルトップ縫合に特化しています。初期のRGBフレームとロボットの運動学（kinematics）のライブストリームを入力として受け取り、次のフレームのシーケンスを自己回帰的に生成します。また、NVIDIAは、CMR SurgicalおよびCambridge ConsultantsのVersiusサージャンコントローラーと統合することで、システムの汎用性も実証しています。

リアルタイム性能を実現する蒸留パイプライン

手術のダイナミクスを損なうことなくリアルタイムの速度を達成するために、Cosmos-H-Dreamsは、長い自己回帰的なロールアウト（rollouts）向けに設計された「教師から生徒への（teacher-to-student）」トレーニングパイプラインを採用しています。

手術用教師モデル (The Surgical Teacher)

プロセスは、Cosmos-H-Surgical-SimulatorのOpen-Hチェックポイントに基づく双方向教師モデルから始まります。この教師モデルは、統合された44次元のアクション表現を使用します。dVRKテーブルトップモデルの場合、両腕のアクション内容（相対的なエンドエフェクタの移動、回転、およびグリッパーの状態）がこの表現にマッピングされます。教師モデルは、JHU dVRKテーブルトップ混合データセットで微調整されており、これには成功したデモンストレーションと失敗（例：針の落下や投げ損じ）の両方が含まれており、シミュレータが不適切なアクションの結果を正確に再現できるようにしています。

長いロールアウト中の安定性を維持するため、教師モデルのタイムホライゾンはトレーニング中に12フレームから72フレームへと段階的に拡大されます。

因果的ウォームアップと自己強制蒸留 (Causal Warmup and Self-Forcing Distillation)

蒸留プロセスは、主に2つの段階で行われます：

1. Causal Warmup: 因果的生徒モデルは教師モデルから初期化され、事前計算およびキャッシュされたデノイジング・トラジェトリを模倣するようにトレーニングされます。これにより、生徒モデルは因果的アテンション（causal attention）とストリーミングKVキャッシュを使用することを学習します。
2. Self-Forcing Distillation: 自己回帰モデルが自身の不完全な出力に条件付けされる際に発生する誤差の累積を防ぐため、NVIDIAはself-forcing蒸留を使用しています。生徒モデルが自身のコンテキストを生成し、凍結された教師モデルが分布一致の監督（supervision）を提供することで、これらのロールアウトを現実的なビデオへと導きます。

その結果として得られる生徒モデルは、潜在フレームあたりわずか2回のデノイジングステップで済む、few-step拡散（diffusion）をサポートしています。

FlashDreams: リアルタイム推論エンジン

Cosmos-H-Dreamsは、自己回帰的なワールドモデルおよびビデオモデルのための高速化推論ライブラリであるFlashDreamsを介して提供されます。FlashDreamsは、以下のものを含む、レイテンシを削減するためのいくつかの最適化を実装しています：

• Streaming KV cache
• CUDA Graph capturing
• Model compilation

これらの最適化により、推論速度は単一のNVIDIA RTX PRO 6000 GPU上で、標準的なCosmos-H-Surgical-Simulatorの約10 FPSから約160 FPSへと向上しました。これにより、以下のいくつかのインタラクションモードが可能になります：

• Browser Client: WebRTCを介して配信されるキーボードコマンドと生成フレーム。
• Meta Quest Client: WebXRを介してロボットのアクションにマッピングされたトラッキングコントローラーの動き。
• Closed-Loop Policy: 観測とアクションがリアルタイムで交換される、学習済み手術ポリシーとの統合。

手術用Physical AIの将来の方向性

NVIDIAは、Cosmos-H-Dreamsをクローズドループ手術Physical AIの基盤と見なしています。今後の開発は、視覚的な品質を超えて、以下を含む新しいベンチマークを通じて機能的な正確性を評価することに焦点を当てています：

• ツールチップの到達範囲とポーズの正確性
• グリッパーサイクルの忠実度
• アイドル時の安定性
• 反事実的なアクションの多様性
• 長期的なドリフト
• シミュレーションと実際のポリシー結果の一致

潜在的なダウンストリームアプリケーションには、レイテンシを考慮した遠隔手術、インタラクティブな手術リハーサル、術前計画、および術中意思決定支援が含まれます。NVIDIAは、Cosmos-H-Dreamsは研究開発プラットフォームであり、診断システム、術中イメージングの代替、または物理的なロボットコントローラーとして意図されたものではないことを注記しています。