Enhancing Policy Learning with World-Action Model

이 논문은 미래 시각 관측과 상태 전이를 유도하는 행동을 함께 추론하도록 역동역학 목표를 도입한 'World-Action Model(WAM)'을 제안하여, CALVIN 벤치마크의 조작 작업에서 기존 베이스라인 대비 행동 복제 성공률을 71.2% 로, 모델 기반 PPO 미세 조정 후 92.8% 로 크게 향상시키고 훈련 단계를 8.7 배 줄인 결과를 보여줍니다.

핵심

이 논문은 로봇이 더 똑똑하고 빠르게 배울 수 있도록 돕는 새로운 기술, **'WAM(월드-액션 모델)'**에 대해 설명합니다.

기존의 로봇 학습 방식과 이 새로운 방식의 차이를 이해하기 위해, **'로봇이 세상을 배우는 방법'**을 **'새로운 요리를 배우는 요리사'**에 비유해 보겠습니다.

1. 기존 방식 (DreamerV2): "요리 사진만 보는 요리사"

기존의 세계 모델 (World Model) 은 로봇에게 **"앞으로 어떤 장면이 펼쳐질지"**를 예측하도록 훈련시켰습니다.

• 비유: 요리사가 레시피를 보고 "이 재료를 넣으면 어떻게 생길지"만 상상하는 것과 같습니다.
• 문제점: 요리사는 "내가 칼을 어떻게 움직였더니 이렇게 된 거야?"라는 행동과 결과의 인과관계를 깊이 있게 이해하지 못합니다. 단순히 "이런 장면이 나올 거야"라고 예측하는 데만 집중하다 보니, 실제로 요리를 할 때 필요한 미세한 손동작 (행동) 에 대한 정보가 누락될 수 있습니다.

2. 새로운 방식 (WAM): "행동까지 예측하는 요리사"

이 논문에서 제안한 WAM은 기존 모델에 **'역동역학 (Inverse Dynamics)'**이라는 새로운 훈련 요소를 추가했습니다.

• 비유: 이제 요리사는 "앞으로 어떤 장면이 펼쳐질지"를 예측할 뿐만 아니라, **"그 장면이 만들어지려면 내가 어떤 행동을 했어야 했는지"**도 함께 추리하도록 훈련받습니다.
  • "아, 이 접시가 이렇게 움직였구나? 그럼 내가 손목을 이렇게 꺾었겠구나!"
  • "이 불이 켜졌네? 그럼 내가 스위치를 눌렀겠구나!"
• 핵심 아이디어: 로봇이 "무엇이 일어날까?"만 묻는 게 아니라, **"내가 무엇을 했기에 이런 일이 일어났을까?"**를 함께 생각하게 함으로써, 로봇의 뇌 (잠재 표현) 에 행동과 관련된 중요한 정보가 더 선명하게 새겨지도록 만든 것입니다.

3. 왜 이것이 중요한가요? (실제 효과)

이 방식은 로봇이 실제 물건을 조작할 때 (서랍을 열거나, 전구를 켜는 등) 놀라운 성과를 냈습니다.

• 더 빠른 학습: 기존 방식보다 약 8.7 배 적은 데이터로 같은 수준의 학습을 완료했습니다. 마치 요리사가 더 적은 재료로 더 맛있는 요리를 배운 것과 같습니다.
• 더 높은 성공률:
  • 초보 단계 (모방 학습): 기존 방식의 성공률 45.8% 에서 **61.7%**로 크게 향상되었습니다.
  • 숙련 단계 (추가 훈련): 추가 훈련을 거치면 성공률이 **92.8%**까지 치솟았습니다. (기존 방식은 79.8%)
  • 특히 서랍을 열거나 슬라이더를 움직이는 등 정교한 손동작이 필요한 작업에서 효과가 가장 컸습니다.

4. 요약: "행동과 결과를 함께 생각하는 로봇"

이 연구의 핵심은 **"로봇에게 미래를 상상할 때, '내가 무엇을 했는지'도 함께 생각하게 하라"**는 것입니다.

• 기존: "앞으로 저렇게 될 거야." (수동적 관찰)
• WAM: "내가 저렇게 했기 때문에 저렇게 될 거야." (능동적 이해)

이 작은 변화가 로봇의 뇌를 훨씬 더 똑똑하게 만들었고, 적은 노력으로 복잡한 작업을 척척 해내는 결과를 가져왔습니다. 마치 요리사가 단순히 레시피를 외우는 게 아니라, 손끝의 감각까지 이해하게 되어 더 훌륭한 요리를 만들어내는 것과 같습니다.

기술 요약

1. 문제 정의 (Problem Statement)

기존의 세계 모델 (예: DreamerV2, DreamerV3 등) 은 주로 관측값 (이미지) 예측에만 초점을 맞추고 훈련됩니다. 즉, 과거의 관측과 행동을 입력받아 미래의 관측을 예측하는 데 최적화되어 있지만, 행동 (Action) 자체를 명시적으로 모델링하지는 않습니다.

이러한 비대칭적인 훈련 방식은 다음과 같은 문제를 야기합니다:

• 행동 관련 정보의 손실: 잠재 공간 (Latent Space) 의 표현 (Representation) 이 픽셀 재구성 (Pixel Reconstruction) 에만 최적화되므로, 환경이 에이전트의 행동에 어떻게 반응하는지에 대한 미세한 구조적 정보 (Action-relevant structure) 가 누락될 수 있습니다.
• 하류 제어의 비효율성: 이렇게 학습된 잠재 상태 ($z_t$) 가 하류의 확산 정책 (Diffusion Policy) 에 입력되지만, 행동 예측에 필요한 정보가 부족하여 정책 학습의 수렴 속도와 최종 성능이 제한받습니다.

2. 방법론 (Methodology)

저자들은 아키텍처를 완전히 재설계하는 대신, 기존 DreamerV2의 훈련 목적 함수에 역동역학 (Inverse Dynamics) 헤드를 추가하여 World-Action Model (WAM) 을 제안했습니다.

A. 아키텍처 및 훈련 목표

• 기반 구조: DreamerV2 의 RSSM (Recurrent State-Space Model) 아키텍처를 기반으로 합니다.
• 역동역학 헤드 (Inverse Dynamics Head): 인코더 임베딩 ($e_t, e_{t+1}$) 을 입력받아 해당 상태 전이를 일으킨 행동 ($\hat{a}_t$) 을 예측하는 추가 헤드를 도입합니다.
  • 기존 방식: $o_{t+1} = f(o_t, a_t)$ (관측 예측)
  • WAM 방식: $o_{t+1} = f(o_t, a_t)$ 및 $a_t = g(e_t, e_{t+1})$ (행동 예측)
• 연쇄 효과 (Cascading Effect): 행동 예측 헤드는 인코더 ($e_t$) 에 직접적으로 작용합니다. 이는 KL 발산 손실을 통해 후방 (Prior) 으로 전파되고, 최종적으로 확산 정책이 사용하는 잠재 특징 ($f_t$) 에 행동 관련 정보가 풍부하게 포함되도록 합니다.
• 손실 함수:
  $$L_{WAM} = \lambda_{KL} L_{KL} + \lambda_{img} L_{recon} + \lambda_{act} L_{action}$$
  여기서 $L_{action}$은 인코더가 행동 관련 구조를 학습하도록 정규화하는 역할을 합니다.

B. 정책 학습 파이프라인

WAM 은 두 단계의 정책 학습을 지원합니다:

1. 행동 복제 (Behavioral Cloning, BC): 고정된 WAM 의 잠재 공간에서 전문가 시연 데이터를 기반으로 확산 정책 (Diffusion Policy) 을 사전 훈련합니다.
2. 모델 기반 강화 학습 (Model-based PPO): 훈련된 WAM 을 시뮬레이터로 사용하여, 실제 물리적 상호작용 없이 잠재 공간 내에서 PPO(Proximal Policy Optimization) 를 통해 정책을 미세 조정 (Fine-tuning) 합니다.

3. 주요 기여 (Key Contributions)

1. WAM 제안: DreamerV2 를 가볍게 확장하여 역동역학 헤드를 추가함으로써, 잠재 표현을 행동 관련 구조에 맞게 정규화했습니다.
2. 생성 품질 향상: CALVIN 벤치마크에서 WAM 은 DreamerV2 대비 더 적은 훈련 단계 (약 8.7 배 감소) 로 PSNR, SSIM, LPIPS, FVD 등 모든 영상 예측 지표에서 더 높은 성능을 달성했습니다.
3. 하류 정책 학습 개선: 행동 정규화된 표현이 하류의 확산 정책 학습을 크게 향상시켜, 행동 복제 (BC) 와 모델 기반 RL 미세 조정 모두에서 기존 베이스라인 (DiWA) 을 압도하는 성능을 보였습니다.

4. 실험 결과 (Results)

실험은 CALVIN 벤치마크의 8 가지 조작 작업 (Manipulation Tasks) 에서 수행되었습니다.

• 세계 모델 생성 품질:

  • WAM 은 DreamerV2 대비 더 사실적인 미래 상태 예측을 보여주었으며, 특히 객체의 형태 왜곡과 색상 드리프트가 적었습니다.
  • 훈련 단계 수: WAM 은 약 23 만 단계로 훈련되었으며, 이는 베이스라인 (200 만 단계) 대비 8.7 배 적은 데이터 효율성을 보입니다.

• 행동 복제 (BC) 성능:

  • WAM 기반 정책은 평균 성공률 **61.7%**를 기록하여, 베이스라인 DiWA (45.8%) 보다 15.9%p 향상되었습니다.
  • 특히 서랍 열기/닫기, 슬라이더 이동 등 정밀한 위치 제어가 필요한 작업에서 큰 개선을 보였습니다.

• PPO 미세 조정 (Fine-tuning) 성능:

  • 모델 기반 PPO 를 적용한 후 WAM 은 평균 성공률 **92.8%**를 달성했습니다.
  • 이는 베이스라인 DiWA (79.8%) 대비 13.0%p 높은 성능이며, '전구 켜기'와 'LED 끄기' 작업에서는 100% 성공률을 기록했습니다.
  • 동일한 정책 아키텍처와 훈련 절차를 사용했음에도 불구하고, WAM 은 베이스라인이 달성하는 성능에 도달하는 데 필요한 환경 상호작용 (Training Steps) 을 크게 줄였습니다.

5. 의의 및 결론 (Significance)

이 연구는 세계 모델이 단순히 관측을 예측하는 것을 넘어, 행동을 예측하는 역동역학 정보를 포함할 때 학습된 표현의 질이 어떻게 향상되는지를 입증했습니다.

• 핵심 통찰: 상태 전이를 일으킨 행동을 예측하도록 강제하면, 인코더가 환경의 동역학에 대한 인과적 구조 (Causal Structure) 를 포착하게 되어, 하류 제어 에이전트가 더 정확한 계획을 수립할 수 있게 됩니다.
• 실용적 가치: 정책 아키텍처나 훈련 프로세스를 변경하지 않고도, 세계 모델의 훈련 목적만 보완함으로써 로봇 조작 작업의 성공률을 획기적으로 높이고 데이터 효율성을 개선할 수 있음을 보여주었습니다. 이는 실제 로봇 시스템에 적용 시 시뮬레이션과 실제 환경 간의 격차 (Sim-to-Real) 를 줄이고 학습 비용을 절감하는 데 중요한 기여를 합니다.