От офлайн-оценки к интерактивному управлению
Ранние модели, такие как Cosmos-H-Surgical-Simulator (на базе Cosmos-Predict2.5-2B), позволяли генерировать видео хирургических сцен для оценки политик и создания синтетических данных, но работали медленно. Cosmos-H-Dreams решает проблему задержки, превращая мир-модель в причинно-следственную (causal) студент-модель с малым числом шагов. Она принимает начальный RGB-кадр и поток кинематики робота, генерируя следующие кадры в замкнутом цикле.
Технология дистилляции: от учителя к студенту
Ключевое достижение — сохранение динамики хирургических тканей при резком снижении вычислительных затрат. Процесс обучения включает три этапа:
- Учитель (Teacher): Бидирекционная модель, дообученная на датасете JHU dVRK (включая сценарии провалов: падение иглы, неудачные узлы) для стабильности на горизонте до 72 кадров.
- Причинная инициализация (Causal Warmup): Студент обучается имитировать кэшированные траектории учителя, используя причинное внимание и потоковый кэш ключей/значений (KV cache).
- Самопринуждающая дистилляция (Self-Forcing Distillation): Студент генерирует rollout на основе своих же предыдущих (ошибочных) выводов, а замороженный учитель корректирует распределение, предотвращая накопление ошибок.
Производительность и инфраструктура
Модель обслуживается через библиотеку FlashDreams, которая оптимизирует инференс с помощью CUDA Graph, потокового KV-кэша и компиляции модели. Это позволяет перейти от ~10 FPS у базовой модели к интерактивным ~160 FPS на одном GPU NVIDIA RTX PRO 6000.
| Параметр | Cosmos-H-Surgical-Simulator | Cosmos-H-Dreams |
|---|---|---|
| Архитектура | World Foundation Model (Cosmos-Predict2.5-2B) | Causal Student Model (Few-step diffusion) |
| Режим работы | Офлайн (генерация по траектории) | Реал-тайм (интерактивный, замкнутый цикл) |
| Скорость инференса | ~10 FPS | ~160 FPS |
| Железо | Не указано (требует больше ресурсов) | 1x NVIDIA RTX PRO 6000 |
| Интеграция | Open-H-Embodiment | dVRK, Versius (CMR Surgical), WebRTC/WebXR |
Практическое применение и будущее
Система уже интегрирована с контроллером хирургической системы Versius (CMR Surgical) в сотрудничестве с Cambridge Consultants. Пользователи могут взаимодействовать через браузер (WebRTC) или шлем Meta Quest (WebXR). Следующий шаг — создание замкнутых бенчмарков для оценки не только визуального качества, но и физической точности: достижения кончика инструмента, стабильности захвата и переноса политик с симуляции на реального робота.
Источник: Hugging Face blog ↗