공동 지도 학습을 거친 모델은 동역학 연쇄(CoD) 시퀀스를 생성하고, 이를 조건으로 미래 시각 상태를 예측합니다. 그러나 교사 강제 방식으로 주어진 CoD만을 이용한 예측 학습으로는 원하는 미래 상태에 도달하는 데 필요한 중간 동역학을 계획하도록 시각-언어 모델(VLM)을 훈련할 수 없습니다. 이에 ThinkWorld는 상상한 미래와 관측된 미래가 얼마나 일치하는지에 따라 샘플링된 CoD에 보상을 부여하여, CoD가 암시하는 세계의 변화를 VLM이 추론하도록 유도합니다. 이 보상은 관측된 전이만 필요하므로 VLM은 대안적인 CoD를 탐색하고, 행동 라벨 없이 또는 후보 동작을 실행하지 않고도 사람과 로봇 영상으로부터 학습할 수 있습니다.
세계 상상을 통한 동역학 연쇄 추론의 확장
세계 상상 기반 보상 설계
각 후보 CoD에 대해 ThinkWorld는 미래 시각 표현을 예측한 뒤 잠재 공간에서 관측된 미래와 비교합니다. 보상에는 현재 프레임과 미래 프레임 사이에 변화가 큰 패치일수록 DINO 특징의 일치도에 더 큰 가중치를 주는 움직임 강화 특징 유사도 점수를 사용합니다.
이 점수는 장면 전체의 균일한 기여도를 유지해 전역 일관성을 보존하는 동시에, 중복되는 배경 영역이 보상을 좌우하지 않도록 합니다. 움직임 가중치 산출에는 관측된 프레임만 사용하므로 점수 계산에 행동 주석이 필요하지 않습니다.
세계 상상 강화 학습을 통한 CoD 추론 최적화
ThinkWorld는 현재 장면에서 원하는 미래 상태에 이르는 중간 동역학을 포착하는 CoD를 학습하고자 합니다. 이 추론 과정을 훈련하기 위해 동일한 시각적 맥락에서 여러 대안 CoD를 샘플링하고, 상상한 미래가 관측된 미래와 가장 잘 일치하는 CoD에 보상을 부여합니다. 이러한 보상을 통해 ThinkWorld는 현재 장면에서 관측된 미래로 이어질 수 있는 대안 CoD를 탐색하며, 수동적으로 수집된 사람 및 로봇 영상으로 행동 추론을 유도할 수 있습니다. 이때 행동 라벨이나 후보 동작의 온라인 실행은 필요하지 않습니다.
VLM은 정규화된 세계 상상 보상을 사용한 그룹 상대 정책 최적화(GRPO)로 업데이트되며, 지도 학습 리플레이를 통해 지도 미세 조정(SFT) 과정에서 학습한 움직임 기반 접지를 유지합니다.
상상한 미래와 관측된 미래를 비교하는 이 단계는 연속 제어에 나중에 사용되는 표현에 행동 관련 동역학을 인코딩하도록 VLM을 유도합니다. 강화 학습 중에는 그 밖의 모든 구성 요소가 고정되며, 배포 시에는 CoD 생성이나 미래 예측이 필요하지 않습니다. 최적화 세부 사항은 부록 A.6에 제시되어 있습니다.
구현 세부 사항
실제 로봇 실험 설정
실제 로봇 평가는 양팔 YAM 시스템을 사용하며, 카메라 3대(전역 상단 시점 카메라 1대와 각 팔의 손목 시점 카메라 1대씩)를 장착합니다. 학습과 추론에는 RGB 관측값만 사용합니다. 이후 신체 구성에 맞게 조정된 ThinkWorld-YAM 모델을 효율적인 작업별 미세 조정의 초기 모델로 사용합니다. 추가 구현 및 학습 세부 사항은 부록에 제시되어 있습니다.
분포 변화 상황에서의 실제 환경 일반화
실제 환경 일반화는 세 가지 조작 작업, 즉 컵을 걸이에 걸기, 색상 맞추기, 물체 분류를 대상으로 대응되는 분포 내(ID) 및 분포 외(OOD) 설정에서 평가합니다. OOD 설정에서는 물체의 색상이나 범주를 바꾸며, 일부 경우에는 작업별 학습 중 관측되지 않은 궤적을 요구합니다.
각 작업에 대해 최소 여섯 가지 물체 구성에서 시연 데이터를 수집하고, 각 설정마다 5회 시행하여 여러 설정을 ID 및 OOD 조건 모두에서 평가합니다. 이 작업을 수행하려면 정책이 색상과 범주 등 작업과 관련된 물체 속성을 식별하고 그에 맞춰 움직임을 조정해야 합니다.
결론
ThinkWorld는 세계 상상을 활용한 강화 학습으로 전이 가능한 동역학 연쇄를 학습하는 시각-언어-행동(VLA) 학습 프레임워크입니다. 행동에 관한 생각을 그 행동이 유발하는 세계의 변화에 기반해 학습함으로써, ThinkWorld는 사람과 로봇의 신체 구성을 아우르는 공통 표현을 학습하고 명시적인 행동 주석이 없는 영상까지 행동 추론을 확장합니다.
실험을 통해 사람에서 로봇으로의 효과적인 전이와 미관측 작업 및 환경에 대한 강한 일반화 성능을 확인했습니다. 이는 관측된 세계 동역학으로부터 전이 가능한 행동 추론을 학습하는 방식의 효과를 입증하며, 다양한 사람 및 로봇 경험을 활용해 로봇 학습을 확장할 가능성을 보여줍니다.
모델 구성
연속 동작은 MolmoAct2로 초기화한 행동 전문가가 생성합니다. 이 전문가는 은닉 차원이 768인 비인과 Transformer 블록 36개로 구성되며 FP32로 동작합니다. 최대 동작 텐서 크기를 제공하며, 2차원 유효성 마스크를 사용해 각 예시를 해당 데이터셋 고유의 시퀀스 길이와 동작 채널로 제한합니다.
전문가는 입력 토큰에서만 계층별 Qwen 키/값 특징을 받습니다. 생성된 CoD 토큰과 미래 질의 토큰은 전문가의 직접적인 맥락에 포함되지 않습니다.
입력 및 출력 구성
모든 미디어는 단일 사용자 턴에서 텍스트 메타데이터보다 먼저 배치됩니다. 시간적 맥락을 사용할 수 있는 경우, 메인 시점의 과거 영상, 현재 메인 시점 영상, 데이터셋에 정의된 순서에 따른 나머지 현재 시점 영상 순으로 입력합니다. 과거 영상에는 최대 16개의 메인 시점 프레임이 포함되며, 2 Hz로 샘플링해 가장 오래된 프레임부터 최신 프레임 순으로 배열합니다. 현재 메인 이미지가 과거 영상에 중복 포함되지는 않습니다.
전체 VLA 입출력 템플릿은 논문에 제시되어 있습니다. 대괄호로 묶인 줄은 선택 사항이며, 텍스트 부분은 여러 줄로 나뉘지 않고 하나의 정규화된 문자열로 렌더링됩니다.
실제 로봇 설정
BimanualYAM 로봇 시스템은 실험실 내 데이터셋 수집과 학습된 정책 평가에 사용됩니다. 이 시스템은 YAM 팔 2개와 RealSense D405 카메라 3대로 구성됩니다. 카메라 한 대는 상단 시점을 제공하고, 나머지 두 대는 각 손목에 장착됩니다. 정책은 RGB 관측값만 사용합니다.
정책 추론은 단일 NVIDIA RTX 3090에서 실행됩니다. 추론 지연 시간의 차이를 고려해 각 정책은 동일한 15 Hz 제어 속도로 호출되고 동작이 실행되며, 나머지 시스템 구성 요소는 30 Hz로 동작합니다. 예측된 동작 청크마다 앞 절반을 실행한 뒤, 최신 관측을 사용해 정책이 다시 계획을 세웁니다.
추가 PCA 및 어텐션 분석
추가 정성적 시각화는 본문 분석을 확장합니다. 미래 표현 PCA와 어텐션 시각화에는 Qwen 레이어 24–27을 사용합니다. PCA와 어텐션은 기술적 분석을 위한 것으로, 인과 관계의 증거로 사용되지 않습니다.
미래 예측 표현의 확장 PCA 시각화에서는 다섯 가지 CoD 조건을 비교합니다. 원본 CoD, 셔플된 CoD, 빈 CoD, 다른 에피소드에서 가져온 기증자 CoD, 그리고 CoD가 없는 조건입니다. 이 비교는 CoD 조건에 따라 예측된 미래 표현이 어떻게 달라지는지 보여줍니다.
확장된 CoD-입력 관측 어텐션 맵에는 미래 예측 토큰에서 입력 관측으로 향하는 어텐션도 표시됩니다. CoD 어텐션은 동적이고 작업과 관련된 영역에 집중되는 경향이 있는 반면, 미래 예측 어텐션은 장면 전반에 더 넓게 분포합니다. 이러한 패턴은 CoD가 행동에 필요한 국소적 변화에 초점을 맞추고 미래 예측은 더 넓은 장면 맥락을 반영함을 시사합니다.
추론 효율성
정상 상태의 배치 크기 1 정책 지연 시간은 CUDA 그래프 캡처를 비활성화하고 정적 형상 그래프 설정 없이 즉시 실행 방식의 추론을 사용해 측정했습니다. 모든 모델은 동일한 NVIDIA H100 80GB GPU에서 측정했으며, 워밍업 호출 5회 후 CUDA 동기화 타이밍 호출 50회를 수행했습니다. 각 모델은 배포 시 고유한 입력 규약을 사용하므로, 측정값은 동일 입력을 적용한 아키텍처 제거 실험이 아니라 실제 정책 지연 시간을 반영합니다.
ThinkWorld는 학습 중에만 추론 과정을 사용하며, 이를 통해 VLM이 행동 중심 표현을 학습하도록 유도합니다. 배포된 정책은 추론 시 이 과정을 실행하지 않으므로 정책 호출마다 추가 실행 계산이 발생하지 않습니다. 이러한 분리는 학습 중 추론이 강한 작업 성능에 기여하면서도 낮은 추론 지연 시간을 유지할 수 있는 이유를 설명합니다.
자주 묻는 질문
ThinkWorld는 세계 상상을 통해 무엇을 학습하나요? 현재 장면에서 원하는 미래 상태로 이어지는 중간 동역학을 표현하는 전이 가능한 동역학 연쇄를 학습합니다.
ThinkWorld는 영상에서 학습할 때 행동 라벨이 필요한가요? 아니요. 세계 상상 보상은 상상한 미래와 관측된 미래를 비교하므로 행동 주석 없이 사람과 로봇 영상에서 학습할 수 있습니다.
실제 환경 평가에는 어떤 로봇 시스템을 사용하나요? 팔 2개와 RealSense D405 카메라 3대를 갖춘 양팔 YAM 시스템을 사용합니다.
배포된 정책은 세계 상상 추론 과정을 실행하나요? 아니요. 추론 과정은 학습 중에 사용되며, 배포 시에는 CoD 생성이나 미래 예측이 필요하지 않습니다.