NVIDIA Cosmos-H-Dreams: 실시간 생성 시뮬레이션을 수술 로봇에 적용하다

월드 파운데이션 모델(world foundation models)은 다른 경로를 제공합니다. 모든 객체와 물리적 상호작용을 수동으로 작성하는 대신, 동기화된 비디오와 로봇 운동학(robot kinematics)으로부터 시각적 역학(visual dynamics)을 직접 학습합니다. NVIDIA의 Cosmos-H-Surgical-Simulator는 초기 장면과 일련의 로봇 동작으로부터 미래의 수술 비디오를 생성함으로써 이 접근 방식을 입증했습니다. 이는 Open-H-Embodiment 생태계 전반에 걸쳐 물리적 시간보다 빠르게 평가하고 합성 데이터를 생성할 수 있게 했습니다.

오늘, 우리는 다음 단계인 Cosmos-H-Dreams를 소개합니다. 이것은 수술 로봇을 위한 실시간 액션 조건부(action-conditioned) 생성 시뮬레이터입니다. Cosmos-H-Dreams는 Cosmos-H-Surgical-Simulator의 역량을 인과적(causal)이며 몇 단계로 이루어진 학생 모델에 응축하여 NVIDIA의 가속 스트리밍 추론 라이브러리인 FlashDreams를 통해 제공합니다. 단일 NVIDIA RTX PRO 6000 GPU에서 실행되는 이 결과물은 사람이나 학습된 정책이 폐쇄 루프(closed loop)로 제어할 수 있는 인터랙티브 환경을 만듭니다.

Cosmos-H-Surgical-Simulator는 NVIDIA Cosmos-Predict2.5-2B를 기반으로 구축되고 Open-H-Embodiment 데이터셋으로 사후 훈련된 액션 조건부 월드 파운데이션 모델입니다. 수술 상황 프레임과 미래 로봇 궤적을 받으면, 해당 동작의 가능성이 높은 시각적 결과를 보여주는 비디오를 생성합니다. 이로 인해 오프라인 정책 평가 및 합성 데이터 생성에 유용합니다. 기록되거나 정책으로 생성된 궤적을 모델로 전송하면, 상응하는 롤아웃(rollout)이 생성될 수 있으며, 물리적 로봇에서 동작을 반복 실행할 필요 없이 그 결과를 검사하거나 점수화할 수 있습니다.

Cosmos-H-Dreams는 이 모델을 실시간 영역으로 이동시킵니다. Cosmos-H-Surgical-Simulator가 학습한 다중 체현(multi-embodiment) 수술 사전 지식(priors)에서 출발하여, da Vinci Research Kit (dVRK) 테이블탑 봉합에 맞게 모델을 특화하고 이를 장면을 자기회귀적(autoregressively)으로 생성하는 인과적 학생(causal student)에게 증류합니다. 출시된 이 모델은 초기 RGB 프레임과 로봇 운동학(robot kinematics)의 실시간 스트림을 입력받아, 다음 프레임 청크를 생성한 후 이어지는 동작 블록을 계속 진행합니다.

또한 저희는 CMR Surgical 및 Cambridge Consultants와 협력하여 Cosmos-H-Dreams가 Versius 외과의사 컨트롤러에 통합되는 다용성(versatility)도 시연했으며, 이를 통해 Versius 플랫폼에서 실시간 작동이 가능하도록 했습니다.

핵심 과제는 생성 비용을 줄이는 동시에 유용한 수술 역학(surgical dynamics)을 보존하는 것입니다. Cosmos-H-Dreams는 길고 자기회귀적인 롤아웃에 맞춰 설계된 교사-학생(teacher-to-student) 학습 파이프라인을 사용합니다.

양방향 교사는 통합된 44차원 동작 표현(action representation)을 사용하는 Cosmos-H-Surgical-Simulator Open-H 체크포인트에서 시작됩니다. 출시된 dVRK 테이블탑 모델의 경우, 상대적인 엔드 이펙터 변환(translation), 회전(rotation), 그리퍼 상태로 구성된 듀얼 암 dVRK 동작 콘텐츠가 이 공통 표현으로 매핑됩니다.

그런 다음 교사는 JHU dVRK 테이블탑 혼합물에 대해 미세 조정되는데, 여기에는 성공적인 시연뿐만 아니라 바늘 낙하(needle drops), 놓친 투척(missed throws), 그리고 불성공한 매듭 묶기(unsuccessful knot ties)와 같은 실패 및 분포 외(out-of-distribution) 에피소드가 포함됩니다. 이러한 실패 사례들은 중요합니다: 정책을 평가하도록 의도된 시뮬레이터는 이상적인 시연뿐만 아니라 잘못된 동작의 결과를 재현해야 하기 때문입니다.

긴 롤아웃 동안 안정성을 개선하기 위해, 학습 과정은 교사의 시간적 지평(temporal horizon)을 점진적으로 증가시킵니다. 저희는 12프레임 지평에서 훈련을 시작하여 이 숫자를 72프레임까지 점진적으로 늘립니다. 각 지평 증분 시마다, 사전 훈련된 가중치로 워밍업된 모델을 초기화합니다.

교사(teacher)의 디노이징 궤적은 먼저 사전 계산되어 캐시됩니다. 인과적 학생(causal student)은 교사로부터 초기화되고, 이 캐시된 궤적을 모방하도록 학습합니다. 이 워밍업 단계는 학생이 자체 생성한 히스토리로부터 학습하기 전에 인과적 어텐션(causal attention) 및 스트리밍 키/값 캐시(streaming key/value cache)를 사용하여 작동하는 방법을 가르칩니다.

자기회귀 모델(Autoregressive models)은 익숙한 문제에 직면합니다. 훈련 중에는 깨끗하고 정답인 컨텍스트(ground-truth context)를 볼 수 있지만, 배포 시에는 자신의 불완전한 출력에 조건화되어야 합니다. 따라서 작은 오류가 시간이 지남에 따라 누적될 수 있습니다.

Cosmos-H-Dreams는 **자가 강제 증류(self-forcing distillation)**를 통해 이러한 불일치(mismatch)를 해결합니다. 훈련 중, 학생은 자체 생성한 컨텍스트를 사용하여 전진(roll forward)합니다. 고정된 교사로부터의 분포 매칭 감독(Distribution-matching supervision)이 이 자가 생성된 롤아웃을 현실적인 수술 비디오로 안내합니다. 이는 학생을 상호작용적 추론(interactive inference) 중에 마주칠 동일한 조건에 대비시킵니다.

결과 모델은 전체 교사가 사용했던 다단계 프로세스 대신, 잠재 프레임당 단 2단계의 디노이징 단계만으로 몇 단계 전파(few-step diffusion)를 지원합니다. 이는 교사의 수술 사전 지식(surgical priors)과 스트리밍에 필요한 인과적 구조를 결합한 것입니다.

모델 증류는 실시간 스토리의 일부일 뿐입니다. Cosmos-H-Dreams는 자기회귀 월드 및 비디오 모델을 위한 가속화된 추론 라이브러리인 FlashDreams를 통해 제공됩니다.

FlashDreams는 스트리밍 KV 캐시(streaming KV cache), CUDA 그래프 캡처(CUDA Graph capturing), 또는 모델 컴파일과 같은 여러 보완적인 최적화를 통해 증류된 학생을 저지연 스트리밍 시스템으로 변환합니다.

이러한 기술들이 결합되어, 증류된 수술 월드 모델은 표준 Cosmos-H-Surgical-Simulator 추론의 초당 약 10프레임 수준에서 상호작용적 작동(~초당 160프레임)으로 향상되었으며, 이는 단일 NVIDIA RTX PRO 6000에서 가능합니다.

Cosmos-H-Dreams는 또한 생성을 상호작용으로 전환하는 인간-기계 인터페이스를 제공합니다. 브라우저 클라이언트는 키보드 명령을 전송하고 WebRTC를 통해 생성된 프레임을 수신할 수 있습니다. Meta Quest 클라이언트는 추적된 컨트롤러 움직임을 로봇 동작에 매핑하고 WebXR을 통해 합성된 장면을 표시할 수 있습니다. 동일한 모델은 학습된 수술 정책에도 연결될 수 있으며, 생성된 관측값과 예측된 동작이 폐쇄 루프 내에서 교환됩니다.

Cosmos-H-Dreams는 테이블탑 봉합에 대한 사전 훈련 체크포인트를 포함하고 있지만, 시스템은 사용자의 특정 구현체(embodiment)로 확장되도록 설계되었습니다. 자체 데이터셋을 위한 실시간 학생 모델을 훈련하려면, 단계별 가이드에서 교사 미세 조정(teacher fine-tuning) 및 자기 강제 증류(self-forcing distillation)에 대한 완벽한 레시피를 제공합니다.

Cosmos-H-Dreams는 수술 시뮬레이션의 새로운 영역을 열었습니다: 실제 로봇 데이터로부터 학습되어 거주할 만큼 충분히 반응적인 환경입니다.

당장의 다음 단계는 시각적 품질 이상을 평가하는 것입니다. 유용한 수술 시뮬레이터는 동작에 올바르게 반응하고, 긴 롤아웃(rollouts) 동안 기구 및 장면 구조를 보존하며, 물리적 로봇으로 전이될 수 있는 결론을 지원해야 합니다. 이는 새로운 계열의 폐쇄 루프 벤치마크를 동기 부여합니다: 툴팁 도달 범위 및 자세 정확도, 그리퍼-사이클 충실도(fidelity), 유휴 안정성(idle stability), 반사실적 동작 다양성(counterfactual action diversity), 장거리 드리프트(long-horizon drift), 그리고 시뮬레이션된 정책과 실제 정책 결과 간의 일치 여부입니다.

실시간 월드 모델은 수술 정책 개발에서도 능동적인 파트너가 될 수 있습니다. 이들은 필요에 따라 희귀한 실패 사례를 생성하고, 강화 학습(RL) 또는 모방 학습을 위한 확장 가능한 환경을 제공하며, 모든 실험을 희소한 로봇 하드웨어에 묶지 않고 새로운 정책의 신속한 평가를 가능하게 합니다.

더 나아가, 실시간 시뮬레이션은 지연 시간 인식 원격 수술(latency-aware telesurgery)과 같이 세계 모델이 보다 안정적인 디스플레이를 유지하는 데 도움을 주는 일련의 다운스트림 애플리케이션이나, 상호작용적 수술 리허설, 시술 계획, 그리고 수술 중 의사 결정 지원 등의 기능을 가능하게 합니다. Cosmos-H-Dreams는 진단 시스템이나 수술 중 이미징 대체재가 아니며, 물리적인 수술 로봇의 컨트롤러도 아닙니다. 그럼에도 불구하고 이 플랫폼은 이러한 가능성을 안전하게 탐색할 수 있는 기반을 제공합니다.

모델 충실도(model fidelity), 시간적 안정성(temporal stability), 하드웨어 효율성이 계속 개선됨에 따라, 실시간 생성 시뮬레이션은 외과 의사 교육, 합성 데이터 생성, 정책 훈련(policy training), 그리고 정책 평가(policy evaluation)를 하나의 공유된 물리 AI 환경 내에서 연결하는 데 도움을 줄 수 있습니다.

Cosmos-H-Dreams의 모델, 데이터, 및 런타임에 대해 살펴보세요:

Cosmos-H-Dreams 코드 및 예시: GitHub 저장소 Cosmos-H-Dreams 모델: Hugging Face 체크포인트 Cosmos-H-Surgical-Simulator: Hugging Face 모델 및 GitHub 저장소 자체 데이터에 맞게 조정하기: 교사 미세 조정(teacher fine-tuning) 및 자체 강제 증류(self-forcing distillation)를 위한 단계별 레시피 Open-H-Embodiment: Hugging Face 데이터셋 FlashDreams: GitHub 저장소 NVIDIA Cosmos-Predict2.5: GitHub 저장소 Cosmos-Surg-dVRK: World Foundation Model 기반 수술 로봇 정책 학습의 자동 온라인 평가 (arXiv 논문)

Cosmos-H-Dreams는 행동 조건부(action-conditioned) 수술 세계 모델링을 실시간 루프에 통합하여, 사람과 정책이 연습하고, 탐색하며, 데이터를 생성하고, 다음 상황을 평가할 수 있는 새로운 환경을 조성합니다.