기술 요약: WAM-RL

문제 정의

최근의 월드-액션(World-Action, WA) 모델들은 비디오 기반 생성 모델을 통해 미래의 관측값과 행동을 결합하여 모델링함으로써 로봇 정책 학습에서 강력한 일반화 성능과 데이터 효율성을 입증해 왔습니다. 그러나 기존의 WA 모델은 전문가의 궤적(expert trajectories)을 이용한 지도 학습에 크게 의존합니다. 이러한 의존성은 두 가지 근본적인 한계를 초래합니다:

1. 분포적 제약(Distributional Constraints): 학습된 정책은 훈련 데이터의 서포트(support) 내로 제한되어, 데몬스트레이션 분포를 벗어나는 미세한 조작 기술을 습득하는 데 어려움을 겪습니다.
2. 지속적 개선의 부재: 모델이 실제 환경과의 상호작용을 통해 지속적으로 개선될 수 없으므로, 새로운 시나리오에 적응하거나 실행 중 오류를 수정하는 능력이 저해됩니다.

강화 학습(RL)은 자연스러운 해결책이지만, 이를 WA 모델에 적용하는 것은 쉽지 않습니다. 시각적 표현이 고정된 경우가 많은 표준 비전-언어-액션(VLA) 모델과 달리, WA 모델은 두 개의 밀접하게 결합된 구성 요소, 즉 월드 모델(미래를 예측하는 모델)과 액션 모델(예측을 행동으로 매핑하는 모델)로 구성됩니다. RL이나 온라인 미세 조정(fine-tuning)을 무분별하게 적용하면 월드 모델의 잠재 공간(latent space)에서 분포 변화(distribution shift)가 발생하여 불안정성을 초래하고, 액션 모델이 이 잠재 공간에 깊이 의존하고 있기 때문에 성능 저하를 일으킬 수 있습니다.

방법론: WAM-RL

저자들은 온라인 상호작용을 통해 월드 모델과 액션 모델의 공동 최적화를 가능하게 하는 프레임워크인 WAM-RL을 제안합니다. 핵심 통찰은 월드 모델이 예측 구조와 계획 능력을 제공하는 반면, 액션 모델은 주로 잠재적 예측을 실행 가능한 행동으로 변환하는 번역기 역할을 한다는 점입니다.

이 프레임워크는 두 부분으로 구성된 최적화 체계를 채택합니다:

1. KL 정규화를 이용한 온라인 비디오 SFT (월드 모델)

월드 모델을 정교화하기 위해, 저자들은 성공적인 상호작용 중에 수집된 궤적을 활용하여 온라인 자기 지도 미세 조정(SFT)을 수행합니다.

• 목적: 월드 모델이 표준 비디오 모델링 목적 함수(예: flow matching)를 사용하여 미래 관측값을 예측하도록 학습합니다.
• 안정화: 잠재 특징 분포가 급격하게 변화하여 액션 모델을 망가뜨리는 것을 방지하기 위해, **KL 정규화 항(KL regularization term)**을 도입합니다. 이는 업데이트된 월드 모델의 잠재 특징이 프리트레인된 모델의 고정된 복사본과 가깝게 유지되도록 제한합니다.
• 구현: 저자들은 지수 이동 평균(EMA)을 사용하여 현재 모델의 공분산을 추정하고 참조 공분산은 고정된 상태로 유지함으로써 잠재 분포를 가우시안으로 근사합니다. 이를 통해 액터(actor)가 기대하는 기하학적 구조를 보존하면서도 점진적인 적응을 허용합니다.

2. 재구성 기반 RL (액션 모델)

액션 모델은 실행된 궤적이 월드 모델의 예측을 충실히 실현하도록 하기 위해 RL을 사용하여 최적화됩니다.

• 보상 신호: 희소한 태스크 보상 대신, 저자들은 **재구성 기반의 조밀한 보상(reconstruction-based dense reward)**을 정의합니다. 이 보상은 월드 모델이 '상상'한 미래 관측값과 실제 환경에서 행동을 실행하여 얻은 실제 관측값 사이의 일관성(유사도)을 측정합니다.
• 유사도 메트릭: 저자들은 픽셀 수준의 MSE, 광학 흐름(optical flow) 일관성, 의미론적 특징 유사도(DINOv2, V-JEPA2)를 포함한 다양한 유사도 함수를 조사합니다.
• 최적화: 액터는 이러한 재구성 보상을 사용하는 정책 경사(policy gradient) 방법을 통해 업데이트되며, 이를 통해 월드 모델의 상상과 실제 실행 사이의 간극을 효과적으로 메웁니다.

주요 기여

1. 최초의 WA 모델용 RL 프레임워크: WAM-RL은 월드 모델과 액션 모델 모두의 공동 최적화를 가능하게 함으로써, 월드-액션 패러다임에 강화 학습을 통합한 최초의 프레임워크입니다.
2. 안정적인 학습 전략: 저자들은 온라인 비디오 SFT와 KL 정규화를 결합한 월드 모델 학습법과 재구성 기반 RL을 사용하는 액터 학습법을 결합한 새로운 학습 전략을 제안합니다. 이 설계는 잠재 공간을 불안정하게 만들지 않으면서 두 구성 요소의 공진화(co-evolution)를 촉진합니다.
3. 장기 궤적(Long-Horizon) 작업에 대한 경험적 통찰: 본 연구는 액터 단독 최적화의 한계에 대한 중요한 통찰을 제공합니다. 저자들은 액터만을 최적화할 경우 단기 궤적 작업에서는 개선이 나타나지만, 누적된 예측 오류로 인해 장기 궤적 작업에서는 유의미한 이득을 얻지 못함을 보여줍니다. 복잡한 장기 성능을 위해서는 공동 최적화가 필수적임을 입증했습니다.

실험 결과

이 방법은 LIBERO(객체 중심 조작) 및 RLBench(다단계 기술) 벤치마크에서 평가되었습니다.

• 성능 향상: WAM-RL은 프리트레인 베이스라인(Base)과 액터 전용 RL($\pi$RL)을 모두 일관되게 능가했습니다.
  • LIBERO-Object에서 성공률은 68%(Base)에서 82%(WAM-RL)로 향상되었으며, 이는 액터 전용 방식(78%)을 크게 앞지른 결과입니다.
  • **RLBench (Water Plants)**에서 WAM-RL은 성공률을 19%에서 22%로 개선한 반면, 액터 전용 방식($\pi$RL)은 베이스 모델(18%)보다 개선되지 못했습니다.
• 행동 분석:
  • 회복 행동(Recovery Behaviors): 월드 모델의 온라인 적응을 통해, 즉각적인 실패를 예측하는 대신 더 현실적인 회복 행동(예: 실패한 잡기를 다시 시도함)을 생성할 수 있게 되었습니다.
  • 보상에 대한 절제 연구(Ablation on Rewards): 연구 결과, 광학 흐름(optical flow)이 성공과 실패 사이에서 가장 변별력 있는 보상 신호를 제공했지만, 픽셀 수준의 MSE가 가장 좋은 다운스트림 성능을 보였습니다. 저자들은 이것이 픽셀 MSE가 월드 모델의 훈련 목적 함수와 더 잘 일치하며, 분포 외(out-of-distribution) 행동에 대해 더 강한 규제를 가하기 때문이라고 가설을 세웠습니다.
  • 비디오 SFT에 대한 절제 연구: 온라인 비디오 SFT 없이 훈련된 모델은 초기 실수를 회복하지 못하고 빠르게 분포 외 행동으로 표류했습니다. 반면, 비디오 SFT를 포함한 모델은 실패를 예측하고 단일 오픈 루프 청크 내에서 교정 동작을 생성하는 법을 배웠습니다.

의의 및 주장

본 논문은 월드-액션 모델의 효과성이 근본적으로 월드 모델의 품질에 의해 제한된다고 주장합니다. 액터 단독 최적화는 단기 궤적 작업을 개선할 수는 있지만, 장기 궤적 설정에서의 누적된 예측 오류를 해결할 수는 없습니다.

WAM-RL의 의의는 월드 모델과 액션 모델을 공동으로 개선하는 것이 견고한 장기 의사결정에 필수적임을 입증했다는 데 있습니다. 온라인 상호작용을 통해 월드 모델이 실패 역학과 회복 전략을 학습하게 함으로써, 시스템은 편차를 예측하고 수정하여 더욱 견고한 실행을 이끌어낼 수 있습니다.

저자들은 KL 정규화가 월드 모델의 적응 범위를 제한하여 프리트레인 분포 이상의 능력 확장을 제한할 수 있다는 점을 언급하며 한계를 겸허히 인정했습니다. 또한, 재구성 기반 보상(종종 프리트레인된 표현에 기반함)에 대한 의존도가 성공과 실패 사이의 대비를 제한할 수 있으므로, 향후 연구에서 더 표현력이 풍부하고 태스크 인지적인 보상 공식이 필요함을 시사했습니다.