[Paper] Gradient‑Based Planning을 위한 World Models에서 Train‑Test 격차 해소

Source: arXiv - 2512.09929v1

Overview

World‑model‑based reinforcement learning은 대규모 전문가 데이터셋으로 단일 동역학 예측기를 오프라인에서 학습한 뒤, 이를 여러 다운스트림 플래닝 작업에 재사용할 수 있다는 약속을 제공합니다. Gradient‑based 플래너는 빠르지만, 전통적인 gradient‑free 방법인 Cross‑Entropy Method (CEM)보다 성능이 낮은 경우가 많았습니다. 이 논문은 그 이유를 규명합니다—모델이 학습되는 방식(다음 상태 예측)과 테스트 시 사용되는 방식(행동 시퀀스 최적화) 사이에 불일치가 존재하기 때문이며, 그 격차를 메우는 구체적인 해결책을 제시합니다. 이를 통해 더 빠르고 정확도에서도 경쟁력 있는 gradient‑based 플래닝을 구현합니다.

Key Contributions

• Train‑test gap analysis: 다음 상태 예측 학습 목표와 추론 시 행동 최적화 사용 사례 사이의 불일치를 정형화합니다.
• Data synthesis tricks: imagined rollouts, action‑perturbation sampling 등 모델에 독립적인 간단한 증강 기법을 도입해, 플래닝 중에 모델이 마주하게 될 궤적을 모델이 미리 경험하도록 합니다.
• Improved gradient‑based planner: 합성 데이터를 활용하면, 표준 gradient descent가 CEM 성능을 맞추거나 초과하면서도 전체 연산량의 약 10 %만 사용합니다.
• Broad empirical validation: 블록 스태킹 같은 조작 환경과 미로 탐색 같은 네비게이션 환경 등 다양한 작업에서 일관된 향상을 보입니다.
• Open‑source implementation: 코드와 사전 학습된 모델을 제공해 실무자가 기술을 쉽게 도입할 수 있도록 합니다.

Methodology

1. Baseline world model: 전문가 궤적 대규모 코퍼스를 사용해 신경망 동역학 모델 (f_\theta(s_t, a_t) \rightarrow s_{t+1})을 학습합니다. 이때 일반적인 평균 제곱 오차(MSE)를 다음 상태 예측에 적용합니다.
2. Identify the gap: 테스트 시 플래너는 모델을 미분 가능한 시뮬레이터로 활용해 미래 보상에 정의된 손실을 역전파하여 후보 행동 시퀀스 (\mathbf{a}_{0:H})를 업데이트합니다. 그러나 모델은 자체의 불완전한 예측이 만든 상태 분포에 한 번도 노출되지 않았습니다.
3. Train‑time data synthesis:
   • Imagined rollouts: 실제 상태에서 시작해 현재 모델을 사용해 무작위 행동을 샘플링하며 합성 궤적을 생성합니다.
   • Action‑perturbation replay: 전문가 행동에 노이즈를 추가하고 재시뮬레이션해, 모델이 오프‑정책 행동에 강건하도록 합니다.
   • Reward‑aware sampling: 합성 샘플을 추정된 미래 보상으로 가중해, 플래닝 중에 모델이 탐색할 영역으로 편향시킵니다.
4. Joint training: 실제 전문가 데이터와 합성 샘플을 섞어 동역학 모델을 계속 학습합니다. 추가 손실 항은 필요 없으며, 동일한 다음 상태 예측 목표를 두 데이터 소스에 적용합니다.
5. Gradient‑based planning: 추론 시 무작위 행동 시퀀스를 초기화하고, 학습된 월드 모델을 사용해 예측 궤적을 계산한 뒤, 작업‑특정 보상을 평가하고, 보상 그래디언트를 역전파해 행동을 정제합니다(예: Adam 사용).

Results & Findings

• 합성 데이터로 학습된 모델은 성능 격차를 해소합니다: gradient‑based 플래너가 이제 모든 테스트 작업에서 CEM과 동등하거나 약간 앞섭니다.
• 연산량 절감이 크게 나타납니다—gradient descent가 CEM이 필요로 하는 반복 횟이 약 10 % 수준으로 수렴해 지연 시간과 에너지 소비가 감소합니다.
• Ablation 연구를 통해 imagined rollouts, perturbations, reward‑aware sampling 각각이 긍정적인 영향을 주며, 어느 하나라도 제거하면 성공률과 속도가 모두 저하됨을 확인했습니다.

Practical Implications

• Faster online planning: 로봇이나 에이전트가 밀리초 단위로 재플래닝할 수 있어, 움직이는 컨베이어에서의 pick‑and‑place와 같은 조작 작업이나 드론 장애물 회피와 같은 자율 네비게이션에서 실시간 대응이 가능해집니다.
• Reduced hardware requirements: Gradient‑based 플래너는 단순한 역전파만 필요하므로, 대규모 병렬 샘플링을 요구하는 CEM에 비해 일반 GPU나 온‑디바이스 가속기에서도 효율적으로 실행됩니다.
• Simplified pipelines: 개발자는 하나의 월드‑모델 학습 루프만 유지하고, 동일한 모델을 여러 다운스트림 작업에 재사용할 수 있어 별도 정책을 학습할 필요가 없습니다.
• Scalable to large datasets: 합성 롤아웃이 저비용이기 때문에, 자율주행 차량 플릿 로그와 같은 방대한 오프라인 데이터셋에도 잘 확장됩니다.
• Potential for hybrid systems: 모델‑프리 파인‑튜닝과 결합해 “양쪽 장점” 시스템을 만들 수 있습니다. 월드 모델이 강력한 사전 지식을 제공하고, gradient‑based 플래닝이 빠른 적응을 담당합니다.

Limitations & Future Work

• Model bias remains: 이 방법은 장기 롤아웃에서 발생하는 누적 오류를 완전히 없애지는 못합니다; 매우 깊은 플래닝 horizon에서는 여전히 성능 저하가 발생할 수 있습니다.
• Task‑specific reward design: Gradient‑based 플래닝은 미분 가능한 보상 신호가 필요합니다. 복잡하고 희소한 보상을 설계하는 것은 여전히 어려운 과제입니다.
• Limited to deterministic dynamics: 현재 접근법은 결정론적 월드 모델을 전제로 합니다. 확률적이거나 부분 관측 가능한 환경으로 확장하는 것은 아직 해결되지 않은 문제입니다.
• Future directions:
  • 불확실성 추정(예: 앙상블) 을 도입해 합성 과정에 활용.
  • 롤아웃 길이를 점진적으로 늘리는 커리큘럼‑스타일 합성 탐색.
  • 상태 추정이 추가적인 난이도를 제공하는 비전 기반 조작 등 고차원 인식 풍부 도메인에서의 테스트.

Bottom line: 학습 데이터 분포를 추론 시 실제 사용 방식과 일치시키는 것으로, 이 연구는 gradient‑based 플래닝의 속도 이점을 성능 저하 없이 활용할 수 있게 합니다—보다 민첩하고 데이터 효율적인 자율 시스템을 향한 흥미로운 진전입니다.

Authors

• Arjun Parthasarathy
• Nimit Kalra
• Rohun Agrawal
• Yann LeCun
• Oumayma Bounou
• Pavel Izmailov
• Micah Goldblum

Paper Information

• arXiv ID: 2512.09929v1
• Categories: cs.LG, cs.RO
• Published: December 10, 2025
• PDF: Download PDF