Cosmos-H-Dreams, RTX PRO 6000 1대로 160fps 실시간 구현

NVIDIA RTX PRO 6000 GPU 1대로 초당 160프레임(fps)의 실시간 상호작용 환경을 구현한 Cosmos-H-Dreams가 공개됐다. 기존 수술 시뮬레이션은 생성 속도가 느려 조작과 화면 출력 사이에 시차가 발생했으나, 이번 모델은 이를 해결해 사용자가 실시간으로 로봇을 제어하며 반응을 확인할 수 있다. 고가의 서버 팜 없이 단일 워크스테이션 수준의 하드웨어만으로 고정밀 시뮬레이션을 구동할 수 있어, 학습된 정책이나 사람이 폐루프 제어(closed loop, 출력값이 다시 입력으로 들어가는 제어 방식)를 통해 환경을 직접 조작하는 접근성이 높아졌다.

Cosmos-H-Dreams는 로봇의 동작 명령에 따라 다음 장면을 생성하는 액션 조건부 생성 시뮬레이터(action-conditioned generative simulator)로 작동한다. Cosmos-H-Surgical-Simulator의 능력을 이전 상태를 바탕으로 다음 상태를 순차적으로 예측하는 경량 모델인 인과적 학생 모델(causal student model)로 증류하여 실시간성을 확보했다. 시스템은 초기 RGB 이미지 프레임과 로봇 관절의 각도 및 위치 정보인 로봇 키네마틱스(robot kinematics) 스트림을 입력받아 다음 프레임 묶음을 생성하며, 이어지는 액션 블록에 맞춰 장면을 연속적으로 업데이트한다.

이 모델은 dVRK(da Vinci Research Kit)의 테이블탑 봉합 작업과 Versius 수술 제어기 플랫폼에서 구동 가능하다. NVIDIA는 CMR Surgical 및 Cambridge Consultants와 협력하여 Versius 제어기를 통합함으로써 실제 수술 플랫폼에서의 작동 가능성을 검증했다. 표준 Cosmos-H-Surgical-Simulator가 초당 약 10프레임 수준의 추론 속도를 보였던 것과 대비해, Cosmos-H-Dreams는 단일 GPU 환경에서 약 16배 빠른 160fps를 달성했다.

교사-학생 증류와 self-forcing으로 해결한 누적 오차

학습 구조는 Cosmos-H-Surgical-Simulator라는 교사 모델의 능력을 Causal Student라는 학생 모델로 전이시키는 증류 방식을 취한다. 이때 로봇의 움직임을 정의하는 44차원 통합 액션 표현(unified action representation)을 사용한다. dVRK 테이블탑 모델의 경우 말단 장치의 상대적 이동과 회전, 그리퍼 상태로 구성된 듀얼 암 액션 콘텐츠를 이 통합 표현으로 매핑하여 처리한다.

학습 데이터로는 JHU dVRK 테이블탑 혼합 데이터셋을 활용한다. 여기에는 성공적인 시연뿐만 아니라 바늘을 떨어뜨리거나 매듭짓기에 실패한 사례 등 다양한 실패 및 분포 외 에피소드(out-of-distribution episodes)가 포함된다. 시뮬레이터가 잘못된 동작에 따른 인과적 결과까지 정확히 재현해야 수술 정책 평가 도구로서의 정확도를 높일 수 있기 때문이다.

장기 예측 시 발생하는 불안정성을 잡기 위해 모델이 한 번에 예측하는 시간 범위인 템포럴 호라이즌(temporal horizon)을 12프레임에서 시작해 최종 72프레임까지 단계적으로 확대했다. 각 단계에서 이전 단계의 학습된 가중치를 초기값으로 사용하여 모델이 급격한 변화 없이 안정적으로 긴 시간의 흐름을 학습하도록 설계했다.

생성 과정에서 발생하는 누적 오차는 self-forcing distillation(셀프 포싱 증류)으로 해결했다. 학생 모델이 스스로 생성한 컨텍스트를 다시 입력값으로 넣어 교사 모델의 데이터 분포와 일치하도록 강제하는 방식이다. 이를 통해 반복 구조에서 작은 오차가 증폭되어 화면이 뭉개지는 현상을 막고, 장기 실행 환경에서도 도구 끝단의 위치나 장면 구조가 유지되는 성능을 확보했다.

10fps에서 160fps로, FlashDreams가 만든 추론 가속

이러한 학습 구조로 안정성을 확보한 후, NVIDIA는 추론 속도를 극대화하기 위한 최적화 단계를 적용했다. 생성 효율을 높이기 위해 latent frame(잠재 프레임)당 최소 2단계의 디노이징(denoising) 과정만으로 프레임을 생성한다. 학생 모델은 교사 모델의 수술 도메인 사전 지식을 유지하면서 스트리밍에 필요한 인과적 구조를 결합해 연산 횟수를 줄였다.

모델의 구조적 경량화와 더불어 FlashDreams라는 NVIDIA 가속 스트리밍 추론 라이브러리를 통해 시스템 전체를 최적화했다. FlashDreams는 자기회귀(autoregressive) 기반의 월드 모델과 비디오 모델의 추론 속도를 높이는 전용 라이브러리다. 하드웨어 자원을 효율적으로 사용하기 위해 streaming KV cache(스트리밍 키-값 캐시)를 적용해 메모리 접근 횟수를 줄였다.

추가적으로 CUDA Graph capturing(쿠다 그래프 캡처링)과 모델 컴파일을 통해 실행 경로를 최적화했다. 모델 컴파일은 고수준의 모델 정의를 특정 GPU 하드웨어에 최적화된 기계어 형태로 변환하여 실행 효율을 높이는 과정이다. 이러한 소프트웨어적 가속 기법들은 증류된 학생 모델의 가벼운 구조와 결합하여 저지연 스트리밍 시스템을 완성한다.

WebRTC와 WebXR을 통한 폐루프 정책 검증 환경

브라우저와 VR 기기를 사용하는 개발자는 WebRTC와 WebXR을 통해 생성된 시뮬레이션 영상에 실시간으로 개입할 수 있다. WebRTC를 사용하는 브라우저 클라이언트는 키보드 명령을 전송하고 생성된 프레임을 수신한다. Meta Quest 클라이언트는 컨트롤러의 모션을 로봇의 액션으로 매핑하여 WebXR을 통해 합성된 장면을 시각화한다. 사용자는 VR 컨트롤러로 로봇의 관절 각도나 위치를 조정하고, 모델이 생성한 시각적 피드백을 즉각적으로 확인한다.

학습된 수술 정책(policy)을 모델과 연결하면 관측값과 액션을 주고받는 폐루프(closed-loop) 제어가 가능하다. 모델이 생성한 시각적 관측값을 정책 모델이 입력받아 다음 행동을 결정하고, 이 행동이 다시 월드 모델의 입력으로 들어가는 순환 구조다. 이를 통해 물리 로봇을 구동하지 않고도 가상 환경 내에서 관측과 행동의 피드백 루프를 완성함으로써, 하드웨어 손상 위험 없이 AI 정책의 성능을 검증하고 수정할 수 있다.

실시간 월드 모델은 실제 환경에서 수집하기 어려운 희귀한 실패 사례를 인위적으로 생성하는 도구로 활용된다. 바늘을 떨어뜨리거나 조직을 잘못 건드리는 엣지 케이스(edge case)를 생성하여 AI가 대응 방법을 학습하게 한다. 이는 강화 학습(reinforcement learning)과 모방 학습(imitation learning)을 위한 확장 가능한 환경을 제공하며, 고가의 수술 로봇 하드웨어 점유 없이 정책을 빠르게 평가하여 개발 주기를 단축한다.

커스텀 로봇 데이터 적용을 위한 증류 레시피와 R&D 경계

테이블탑 봉합용 사전 학습 체크포인트가 제공되며, 연구자는 자체 데이터셋을 활용해 교사 모델 미세 조정과 self-forcing 증류 단계를 거쳐 실시간 학생 모델을 구축할 수 있다. 모델의 성능 검증은 툴팁 도달 거리, 포즈 정확도, 그리퍼 사이클 충실도 같은 정밀 지표로 수행한다. 특히 시간이 흐를수록 예측 영상이 물리 법칙에서 벗어나는 장기 롤아웃 드리프트(long rollout drift) 현상을 억제하는 것이 중요하다. 연구자는 제공된 가이드를 통해 로봇 하드웨어 특성에 맞는 검증 수치를 설정하고 정책의 물리적 제약 조건 작동 여부를 확인한다.

Cosmos-H-Dreams는 철저히 R&D 플랫폼으로 정의되며 실제 의료 현장에서의 사용 범위는 엄격히 제한된다. 이 모델을 환자 상태를 판단하는 진단 시스템이나 수술 중 실시간 영상을 대체하는 용도로 사용할 수 없다. 또한 실제 수술 로봇의 움직임을 직접 제어하는 컨트롤러로 연결하는 행위 역시 금지된다. 시뮬레이션 기반의 월드 모델이 생성하는 영상은 물리적 근사치이며, 이를 실제 수술 제어의 직접적인 근거로 삼는 것은 안전상 허용되지 않는다.

물리 로봇 없이도 학습된 정책을 폐루프(closed-loop, 입력이 다시 출력에 영향을 주는 구조) 환경에서 빠르게 검증하여 실제 배치 전의 리스크를 최소화하는 것이 이 시스템의 실질적인 활용 기준이다.