[논문 리뷰] Chain of World: 잠재 공간에서 생각하는 로봇을 위한 세계 모델

TL;DR

비전-언어-행동(Vision-Language-Action, VLA) 모델은 로봇이 시각 정보와 언어 지시를 이해하고 행동하게 만드는 AI입니다. 하지만 기존 모델들은 세상의 물리적 법칙이나 시간적 인과관계를 깊이 이해하지 못하는 한계가 있었습니다. **CoWVLA(Chain-of-World VLA)**는 이러한 문제를 해결하기 위해, 세계 모델(World Model)의 시간적 추론 능력과 잠재 공간 계획(Latent Space Planning)의 효율성을 결합한 새로운 패러다임을 제시합니다. CoWVLA는 비디오를 정적인 **'구조(Structure)'**와 동적인 **'모션(Motion)'**으로 분리하여 세상을 이해하고, 이를 통해 더 효율적이고 일반화 성능이 뛰어난 로봇 제어를 실현합니다.

연구 배경 및 동기

로봇이 "컵을 들어 선반에 올려놓으세요"와 같은 지시를 수행하려면, 단순히 현재 이미지를 보는 것을 넘어 앞으로 일어날 일들을 예측하고 계획할 수 있어야 합니다. 기존 VLA 모델들은 이러한 시간적 추론 능력이 부족했습니다.

• 기존 세계 모델: 미래의 모든 픽셀을 예측하려다 보니 계산량이 너무 많고, 정작 중요한 '움직임'의 본질을 놓치기 쉬웠습니다.
• 잠재 행동 모델: 행동을 추상적인 잠재 공간에서 표현하여 효율성을 높였지만, 세상의 물리 법칙을 일관되게 모델링하는 능력이 부족했습니다.

CoWVLA는 이 두 가지 접근법의 장점을 결합합니다. 픽셀 단위의 예측 대신, 추상화된 '모션'의 연쇄(Chain)를 잠재 공간에서 예측합니다. 이 "생각의 사슬"을 통해 로봇은 복잡한 작업을 더 효율적으로 계획하고 수행할 수 있습니다. 특히 비디오를 '구조'(배경, 물체 모양 등)와 '모션'(움직임의 패턴)으로 분리함으로써, 주방이 바뀌거나 컵 모양이 달라져도 '컵을 옮기는' 동작의 본질을 파악하여 높은 일반화 성능을 달성합니다.

관련 연구

로봇 제어 분야는 강화 학습, 세계 모델 등 다양한 아이디어를 바탕으로 발전해왔습니다.

CoWVLA는 이러한 선행 연구들의 어깨 위에서, 세계 모델의 예측 능력과 잠재 공간의 효율성을 구조-모션 분리라는 새로운 아이디어로 융합하여 기존의 한계를 돌파합니다.

핵심 기여

1. CoWVLA 모델 제안: 세계 모델의 시간적 추론과 잠재 공간 계획의 효율성을 결합한 새로운 VLA 패러다임을 제시합니다.
2. 구조-모션 분리(Disentanglement): 비디오를 정적인 '구조'($z_s$)와 동적인 '모션'($z_m$)으로 분리하여, 환경 변화에 강인한 일반화 성능을 확보했습니다.
3. 잠재 공간에서의 연쇄적 사고 (Chain of Thought in Latent Space): 언어 지시를 바탕으로 미래에 필요한 '잠재 모션'의 연쇄(Chain of World)를 예측하여 복잡하고 긴 작업을 효율적으로 계획합니다.
4. SOTA 성능 달성: LIBERO, SimplerEnv 등 까다로운 로봇 제어 벤치마크에서 기존 모델들을 능가하는 성능과 뛰어난 데이터 효율성을 입증했습니다.

제안 방법론

CoWVLA는 2단계 학습 전략을 따릅니다. 첫 번째 단계에서는 세상의 물리적 원리를 배우고, 두 번째 단계에서는 그 원리를 바탕으로 실제 로봇 행동을 학습합니다.

1단계: 세계 모델 사전 학습 (세상의 원리 배우기)

이 단계에서는 대규모 비디오 데이터를 사용하여 세상이 어떻게 작동하는지에 대한 일반적인 원리를 학습합니다. 핵심은 **Video VAE(Variational Autoencoder)**를 사용하여 비디오를 '구조'와 '모션'으로 분리하는 것입니다.

• 구조 잠재 변수 ($z_s$): 장면의 정적인 요소(배경, 객체 모양, 색상 등)를 인코딩합니다. 비디오의 첫 프레임에서 추출됩니다.
• 모션 잠재 변수 ($z_m$): 시간의 흐름에 따른 변화(객체의 움직임, 상태 변화 등)를 인코딩합니다. 비디오의 전체 프레임 시퀀스에서 추출됩니다.

모델은 초기 프레임($o_0$)과 언어 지시($l$)를 입력받아, 앞으로 일어나야 할 일련의 **'잠재적 모션 체인' ($z_m^{1...T}$)**을 예측합니다. 이 예측된 모션 체인과 초기 프레임의 구조($z_s$)를 결합하여 최종 프레임의 모습을 시각적으로 상상해내며 세계 모델의 예측 능력을 강화합니다.

이 과정의 손실 함수는 다음과 같이 구성됩니다.

• $L_{rec}$ (재구성 손실): 인코딩된 $z_s$, $z_m$으로부터 원본 영상을 얼마나 잘 복원하는지 측정합니다.
• $L_{p}$ (지각 손실): 복원된 영상이 사람의 눈에 얼마나 자연스럽게 보이는지 평가합니다.
• $L_{adv}$ (적대적 손실): 생성된 영상이 실제 영상과 구분할 수 없을 만큼 사실적이도록 유도합니다.
• $L_{KL}$ (KL 발산 손실): 잠재 공간이 통계적으로 잘 정돈된 정규 분포를 따르도록 규제합니다.

2단계: 행동 정책 공동 미세 조정 (로봇 제어법 배우기)

사전 학습으로 세상의 원리를 깨우친 모델에게 실제 로봇 제어 방법을 가르치는 단계입니다. 이 단계에서는 로봇 시연 데이터(관측, 행동, 언어 지시)를 사용합니다.

1. 잠재 계획 생성: 사전 학습된 세계 모델을 사용하여, 초기 관측과 언어 지시로부터 전체 작업에 대한 **잠재 모션 계획($z_m^{plan}$)**을 생성합니다.
2. 행동 예측: 현재 관측($o_t$), 언어 지시($l$), 그리고 **생성된 전체 잠재 계획($z_m^{plan}$)**을 입력으로 받는 정책 모델(Policy Model)이 구체적인 로봇 행동($a_t$)을 예측합니다.

이 과정의 손실 함수는 다음과 같습니다.

• $L_{action}$ (행동 예측 손실): 모델이 예측한 행동이 전문가의 실제 행동과 일치하도록 학습합니다. (주로 Cross-Entropy 손실 사용)
• $L_{latent\_dynamics}$ (잠재 동역학 손실): 미세 조정 중에도 잠재 표현이 사전 학습된 세계 모델의 물리 법칙을 일관되게 따르도록 규제하여, 학습 안정성을 높입니다.

개념적 코드 예시

# --- 1단계: 세계 모델 사전 학습 ---
world_model = WorldVAE()
motion_planner = MotionPlanner() # Transformer 기반 예측 모델

# 대규모 비디오 데이터셋으로 학습
for video, instruction in large_video_dataset:
    # VAE를 통해 구조(z_s)와 실제 모션(z_m_gt) 분리
    z_s, z_m_gt = world_model.encode(video)
    
    # VAE가 영상을 잘 복원하도록 학습 (L_VAE)
    world_model.train_on_reconstruction(video, z_s, z_m_gt)
    
    # 모션 플래너가 실제 모션을 잘 예측하도록 학습
    z_m_pred = motion_planner.predict(initial_frame=video[0], instruction=instruction)
    motion_planner.train_on_prediction(z_m_pred, z_m_gt)

# --- 2단계: 정책 미세 조정 ---
policy_model = ActionPolicy()

# 로봇 시연 데이터셋으로 학습
for observations, actions, instruction in robot_demonstration_dataset:
    # 1. 사전 학습된 모델로 잠재 계획 생성
    z_m_plan = motion_planner.predict(initial_frame=observations[0], instruction=instruction)
    
    # 2. 현재 상태와 전체 계획을 바탕으로 행동 예측
    for t, obs_t in enumerate(observations):
        # 정책 모델은 현재 관측뿐 아니라 미래 계획 전체를 참고함
        predicted_action = policy_model.act(obs_t, instruction, z_m_plan)
        
        # 실제 행동을 모방하도록 학습 (L_action)
        policy_model.train_on_imitation(predicted_action, actions[t])

실험 설정

CoWVLA의 성능은 복잡하고 긴 순서의 로봇 조작 작업을 포함하는 LIBERO와 SimplerEnv 벤치마크에서 평가되었습니다. 이 벤치마크들은 다양한 객체와 환경에서 모델의 일반화 및 문제 해결 능력을 테스트하기에 적합합니다. 실험은 NVIDIA A800 GPU 16~32대를 사용하여 진행되었으며, 대규모 로봇 데이터셋인 OXE를 활용하여 학습되었습니다.

실험 결과 분석

CoWVLA는 기존의 순수 세계 모델이나 잠재 행동 기반 모델들을 대부분의 태스크에서 큰 차이로 능가했습니다.

특히 CoWVLA의 진가는 도메인 일반화(Cross-Domain Generalization) 능력에서 드러났습니다. 학습 환경과 다른 새로운 환경(예: 다른 배경, 다른 모양의 도구)에 투입되었을 때, 다른 모델들은 성능이 급격히 저하된 반면 CoWVLA는 구조와 모션을 분리한 덕분에 훨씬 강건한 성능을 유지했습니다.

데이터 효율성 측면에서도 CoWVLA는 적은 양의 데이터와 계산 자원으로도 UniVLA와 같은 초대형 모델에 필적하는 높은 성능을 달성했습니다.

주요 결과 (vs. SOTA Baseline LAVA)

또한, **Ablation Study(요소 제거 연구)**를 통해 CoWVLA의 핵심 요소(구조-모션 분리, 세계 모델 사전 학습 등)를 하나씩 제거하며 성능 변화를 측정한 결과, 각 요소가 최종 성능에 결정적인 기여를 한다는 것을 입증했습니다.

비판적 평가

CoWVLA는 분명한 강점을 가지고 있지만, 몇 가지 한계점과 고려사항도 존재합니다.

강점:

• 뛰어난 일반화 성능: 구조와 모션을 분리하는 접근법은 환경 변화에 매우 강인합니다.
• 효율적인 장기 계획: 잠재 공간에서 모션 체인을 예측함으로써 복잡한 작업을 효율적으로 계획합니다.
• 데이터 효율성: 대규모 비디오 사전 학습 덕분에 상대적으로 적은 로봇 시연 데이터로도 높은 성능을 낼 수 있습니다.

한계점:

• 복잡성: 2단계 학습 과정과 여러 모델 구성 요소로 인해 전체 파이프라인이 복잡하고 학습 시간이 길어질 수 있습니다.
• 분리(Disentanglement)의 한계: 구조와 모션이 매우 복잡하게 얽혀있는 특정 작업(예: 점토로 모양 만들기)에서는 두 요소를 완벽하게 분리하기 어려울 수 있습니다.
• 재현성: 대규모 비디오 데이터셋과 상당한 양의 컴퓨팅 자원(GPU)이 필요하여 개인 연구자가 재현하기에는 장벽이 있습니다.

향후 연구 방향

CoWVLA는 다양한 확장 가능성을 제시합니다.

• 다양한 로봇 작업으로의 확장: 현재의 팔 조작 작업을 넘어, 이동 로봇이나 다리가 있는 로봇 등 다른 유형의 로봇 제어에 적용할 수 있습니다.
• 실제 환경 테스트: 시뮬레이션 환경을 넘어, 예측 불가능한 변수가 많은 실제 환경에서의 강건성을 검증하고 개선하는 연구가 필요합니다.
• 모델 경량화: 더 적은 자원으로도 CoWVLA를 구동할 수 있도록 모델의 효율성을 높이는 연구는 실용성을 높이는 데 중요합니다.

실무 적용 가이드

CoWVLA를 실제 산업 현장이나 로봇 서비스에 적용하려면 다음 사항을 고려해야 합니다.

1. 데이터 확보: 모델의 성능은 데이터의 양과 질에 크게 의존합니다. 사전 학습을 위한 다양한 비디오 데이터와, 특정 작업을 위한 고품질 로봇 시연 데이터를 확보하는 것이 가장 중요합니다.
2. 컴퓨팅 자원: 모델을 학습시키기 위해 고성능 GPU 클러스터가 필요합니다.
3. 작업 특화 미세 조정: 일반적인 세계 모델을 사전 학습한 후, 목표로 하는 특정 작업 환경과 로봇에 맞춰 정책 모델을 신중하게 미세 조정해야 합니다. 구조-모션 분리가 잘 작동하는지 확인하는 과정이 필요합니다.

결론

CoWVLA는 '생각하는' 로봇을 향한 중요한 진일보를 보여줍니다. 단순히 현재를 보고 반응하는 것을 넘어, 세상의 작동 원리를 이해하고 미래의 움직임을 잠재 공간에서 계획하는 'Chain of World' 개념을 도입했습니다. 비디오를 '구조'와 '모션'으로 분리하는 혁신적인 아이디어를 통해, 기존 VLA 모델들의 한계였던 일반화 성능과 장기 계획 능력을 크게 향상시켰습니다. CoWVLA는 앞으로 더 지능적이고 유연한 로봇을 만드는 데 핵심적인 기술적 토대를 제공할 것입니다.

참고 자료

• 논문 링크: arXiv:2403.03195
• 프로젝트 웹사이트: CoWVLA Project
• 코드 저장소: GitHub Repository