Persistent Robot World Models: Stabilizing Multi-Step Rollouts via Reinforcement Learning

1. 문제 정의 (Problem Statement)

기존의 로봇 월드 모델은 주로 단기 예측에 최적화되어 있으며, 교사 강제 학습 (Teacher-forced training) 방식을 사용합니다. 즉, 모델은 과거의 실제 정답 데이터 (Ground Truth) 를 입력받아 다음 프레임을 예측하도록 훈련됩니다.

하지만 실제 로봇 제어나 장기 시뮬레이션에서는 자율적 (Autoregressive) 인 롤아웃이 필요합니다. 이는 모델이 예측한 프레임을 다시 다음 단계의 입력 (Context) 으로 사용하는 과정을 반복하는 것을 의미합니다. 이 과정에서 발생하는 주요 문제는 다음과 같습니다.

• 노출 편향 (Exposure Bias): 훈련 시에는 정답 데이터를 사용하지만, 추론 시에는 모델 자신의 불완전한 예측값을 입력으로 사용하게 됩니다.
• 오차 누적 (Error Compounding): 초기 단계의 미세한 예측 오차가 다음 단계의 입력으로 전달되며 기하급수적으로 증폭됩니다.
• 시각적 붕괴: 몇 초 만에 예측된 비디오에서 물체의 형태가 왜곡되거나 (예: 그릇이 뭉개짐), 로봇의 팔이 명령된 궤적에서 벗어나고, 전체 장면의 일관성이 무너집니다.

기존의 확산 모델 (Diffusion Models) 에는 이러한 장기적 일관성을 보장하기 위한 훈련 메커니즘이 부재했습니다.

2. 방법론 (Methodology)

저자들은 **강화학습 (Reinforcement Learning, RL)**을 도입하여 모델이 자신의 생성된 롤아웃 데이터에 대해 직접 학습하도록 하는 RL 사후 훈련 (RL Post-training) 프레임워크를 제안했습니다.

핵심 기술적 접근

1. 대조적 강화학습 (Contrastive RL) 적용:

   • 기존 확산 모델에 RL 을 적용할 때 발생하는 역전파 비용 문제를 해결하기 위해, **DiffusionNFT [42]**의 아이디어를 차용했습니다.
   • 확산 과정의 역전파를 거치지 않고, **여러 개의 후보 생성물 (Candidate Outputs)**을 생성하여 보상 (Reward) 으로 비교하는 방식을 사용합니다.
   • $x_0$ 예측 파라미터화 적응: 기존 연구는 속도 예측 (Velocity prediction) 기반 모델에 적용되었으나, 이 논문은 Ctrl-World와 같은 $x_0$(클린 데이터) 예측 기반 확산 모델에 해당 논리를 수학적으로 재도출하여 적용했습니다. 이론적 수렴 보장이 그대로 유지됨을 증명했습니다.

2. 자율적 롤아웃을 위한 훈련 프로토콜:

   • 공유 접두어 (Shared Prefix) 생성: 실제 정답 데이터에서 시작하여 모델이 자율적으로 $P$개의 프레임을 생성합니다. 이때 $P$는 무작위로 샘플링되어 초기 단계부터 심화된 오차 단계까지 다양한 상황을 학습하게 합니다.
   • 후보 분기 (Branching): 고정된 접두어 상태 (History Buffer) 에서 $K$개의 독립적인 후보 연속 시나리오를 생성합니다.
   • 상대적 평가: 생성된 $K$개의 후보를 정답 데이터와 비교하여 보상을 산출하고, 그룹 내 상대적 순위 (Group-normalized) 를 기반으로 모델을 업데이트합니다.

3. 다중 뷰 시각적 보상 (Multi-view Visual Rewards):

   • 단일 메트릭의 한계를 극복하기 위해 LPIPS(지각적 유사성), SSIM(구조적 유사성), PSNR(신호 대 잡음비) 을 결합한 복합 보상 함수를 설계했습니다.
   • 다중 카메라 뷰 통합: 로봇의 손목 카메라 (Wrist) 와 외부 카메라 (External) 등 모든 뷰에서 생성된 프레임을 평가하여, 특정 뷰에 치우치지 않는 전반적인 시각적 충실도를 높였습니다.
   • 클립 단위 평가: 긴 롤아웃 전체가 아닌, 짧은 클립 단위로 보상을 계산하여 학습 신호의 분산을 줄이고 구체적인 오류를 교정합니다.

3. 주요 기여 (Key Contributions)

1. 로봇 월드 모델을 위한 RL 사후 훈련 체계:

   • 정답 데이터가 아닌 모델 자신의 생성 데이터 (Autoregressive Rollouts) 를 기반으로 모델을 최적화하는 새로운 훈련 패러다임을 제시했습니다.
   • $x_0$ 예측 기반 확산 모델에 대조적 RL 목적 함수를 적용하고 그 수렴 보장을 이론적으로 입증했습니다.

2. 자율적 롤아웃을 위한 훈련 프로토콜 설계:

   • 고정된 프롬프트가 없는 비디오 생성 환경에서, 누적된 히스토리 버퍼를 공유 컨텍스트로 활용하여 여러 후보를 생성하고 비교하는 메커니즘을 고안했습니다.

3. 효율적인 다중 뷰 시각적 보상 함수:

   • 로봇 작업의 핵심인 물체 조작과 로봇 팔의 움직임을 정확히 포착하기 위해, 다양한 카메라 뷰와 지각적 메트릭을 결합한 보상 시스템을 개발했습니다.

4. DROID 데이터셋에서의 SOTA 달성:

   • DROID 데이터셋에서 기존 최강 베이스라인 (Ctrl-World, WPE, IRASim 등) 을 모든 지표에서 압도적으로 능가하는 성능을 입증했습니다.

4. 실험 결과 (Results)

**DROID 데이터셋 (Robot Manipulation Dataset)**을 기반으로 한 평가 결과입니다.

• 시각적 품질 지표 (14 단계 롤아웃, 약 11 초):

  • 외부 카메라: LPIPS 가 14.0% 감소 (0.091 → 0.070), PSNR 이 1.40dB 향상.
  • 손목 카메라 (Wrist): SSIM 이 9.1% 향상, PSNR 이 1.59dB 향상. 이는 미세한 물체 조작의 정확도가 크게 개선되었음을 의미합니다.
  • 비교: 기존 베이스라인 대비 모든 메트릭에서 우위를 점했으며, 특히 물체와 로봇 팔의 구조적 일관성이 크게 개선되었습니다.

• 페어 비교 및 인간 선호도:

  • 페어 비교: 검증 세트의 약 **98%**에서 베이스라인보다 우수한 결과를 보였습니다.
  • 블라인드 인간 평가: 인간 평가자들이 생성된 롤아웃 중 본 연구의 모델을 **80%**의 비율로 선호했습니다.

• 장기적 안정성:

  • 롤아웃 시간이 길어질수록 오차가 누적되는 현상이 기존 모델에 비해 현저히 억제되었으며, 물체의 형태가 붕괴되지 않고 구조적 무결성을 유지했습니다.

• 정책 평가 도구로서의 유효성:

  • 학습된 월드 모델을 사용하여 실제 로봇 정책의 성능을 시뮬레이션에서 평가했을 때, 실제 환경의 진행률과 높은 상관관계 (Pearson Correlation 0.822) 를 보였습니다.

5. 의의 및 결론 (Significance)

이 연구는 **로봇 학습을 위한 가상 환경 (World Model)**의 실용성을 크게 확장했습니다.

• 시뮬레이션의 신뢰성 확보: 기존 월드 모델이 가진 "단기 예측은 좋으나 장기 시뮬레이션은 불가능하다"는 한계를 극복하여, 로봇 정책의 학습과 평가를 위한 고충실도 (High-fidelity) 가상 환경을 제공할 수 있게 되었습니다.
• 데이터 효율성: 실제 로봇을 조작하는 데 드는 비용과 위험을 줄이고, 시뮬레이션 내에서 안전하게 정책을 반복적으로 훈련하고 평가할 수 있는 기반을 마련했습니다.
• 확장성: 제안된 RL 사후 훈련 프레임워크는 시각적 충실도뿐만 아니라 물리 법칙이나 작업 성공률과 같은 다른 보상 신호에도 적용 가능하여, 향후 더 정교한 로봇 에이전트 개발의 토대가 될 것입니다.

결론적으로, PersistWorld는 강화학습과 확산 모델을 결합하여 로봇이 복잡한 작업을 장기적으로 예측하고 시뮬레이션할 수 있는 새로운 표준을 제시한 획기적인 연구입니다.