IPD: Boosting Sequential Policy with Imaginary Planning Distillation in Offline Reinforcement Learning

이 논문은 오프라인 강화학습에서 결정 트랜스포머 기반 정책의 한계를 극복하기 위해, 불확실성 추정과 준최적 가치 함수를 활용한 세계 모델을 구축하고 모델 예측 제어 (MPC) 를 통해 이상적인 롤아웃을 생성하여 데이터를 증강한 후 이를 지식 증류하는 '상상적 계획 증류 (IPD)' 프레임워크를 제안하고 D4RL 벤치마크에서 기존 최첨단 방법들보다 우수한 성능을 입증했습니다.

핵심

🚀 IPD: 실패한 경험을 '상상력'으로 보완하는 AI 의 새로운 비법

이 논문은 **'오프라인 강화학습 (Offline RL)'**이라는 분야에서 AI 가 더 똑똑하게 행동할 수 있도록 돕는 새로운 방법, **IPD(Imaginary Planning Distillation, 상상력 기반 계획 증류)**를 소개합니다.

쉽게 말해, **"AI 가 직접 실수를 반복하며 배우는 대신, 이미 가진 실패 기록을 분석하고 '만약에'라는 상상을 통해 최고의 시나리오를 만들어낸 뒤, 그걸로 다시 배우게 한다"**는 아이디어입니다.

---

🎬 1. 배경: 왜 새로운 방법이 필요할까?

상황:
전통적인 AI 는 로봇을 조종하거나 게임을 할 때, 직접 부딪히고 넘어지며 (온라인 학습) 배웁니다. 하지만 현실 세계에서는 실수가 치명적일 수 있죠 (예: 자율주행차가 추돌하거나, 로봇이 사람을 다치게 함). 그래서 우리는 **"이미 쌓아둔 실패와 성공의 기록 (데이터)"**만 보고 배우는 오프라인 학습을 사용합니다.

문제점:
기존의 오프라인 학습 AI 들은 두 가지 큰 한계가 있습니다.

1. 나쁜 습관 고수: 기록된 데이터에 '실수'가 많으면, AI 는 그 실수를 그대로 따라 합니다.
2. 상상력 부족: "이 상황에서 내가 조금만 다르게 행동했다면 더 잘했을 텐데?"라는 **상상 (계획)**을 하지 못합니다. 단순히 기록된 대로만 따라 하느라, 최적의 행동을 찾아내지 못합니다.

비유:
마치 나쁜 요리 레시피 책만 보고 요리하는 셰프를 상상해보세요. 그 책에는 실패한 요리법도 섞여 있고, "소금 좀 더 넣으면 맛있었을 텐데"라는 생각은 전혀 없습니다. 이 셰프는 아무리 연습해도 맛있는 요리를 못 만들죠.

---

💡 2. IPD 의 핵심 아이디어: "상상력 (Imaginary Planning)"

IPD 는 이 문제를 해결하기 위해 **"상상력"**을 도입합니다. AI 가 직접 실험하지 않고도, 머릿속으로 **"만약에 내가 최선의 선택을 했다면 어떻게 되었을까?"**를 시뮬레이션하는 것입니다.

이 과정은 크게 3 단계로 나뉩니다.

1 단계: '현실 감각'과 '가치 판단'을 익히기

먼저 AI 는 기존 데이터로 두 가지를 배웁니다.

• 세계 모델 (World Model): "내가 이렇게 행동하면 세상이 어떻게 변할까?"를 예측하는 능력. (예: 공을 차면 어디로 날아가는지)
• 가치 함수 (Value Function): "이 상태가 얼마나 좋은 상태일까?"를 판단하는 능력. (예: 이 위치에서 공을 차면 득점 확률이 높음)

비유:
요리사가 레시피 책 (데이터) 을 보며, "이 재료를 섞으면 맛이 변할 거야 (세계 모델)"라고 예측하고, "이 조합은 실패할 확률이 높아 (가치 함수)"라고 판단하는 능력을 기르는 것입니다.

2 단계: 실패 구간을 찾아 '상상'으로 채우기

이제 AI 는 데이터 속의 **나쁜 부분 (실패한 구간)**을 찾아냅니다.

• "여기서 실수했구나."
• "만약 내가 여기서 최고의 요리사처럼 행동했다면 어땠을까?"라고 **상상 (MPC, 모델 예측 제어)**을 합니다.
• AI 는 머릿속으로 "최고의 행동"을 시뮬레이션하여, 새로운, 더 좋은 데이터를 만들어냅니다.
• 중요: AI 는 "내가 이 시뮬레이션을 얼마나 잘 믿을 수 있을까?"를 계산해서, 믿을 수 없는 엉터리 상상은 버리고 믿을 수 있는 좋은 상상만 데이터에 추가합니다.

비유:
나쁜 레시피를 보고, "아, 여기서 소금을 덜 넣었구나. 만약 내가 소금을 적절히 넣었다면 어땠을까?"라고 상상합니다. 그리고 "내 상상력이 이 정도라면 믿을 만해"라고 판단한 뒤, 그 상상 속의 완벽한 레시피를 실제 레시피 책에 추가합니다.

3 단계: 상상한 지혜를 '학습'으로 바꾸기 (Distillation)

마지막으로, AI 는 실제 데이터 + 상상한 좋은 데이터를 섞어서 다시 학습합니다.

• 이때 단순히 따라 하는 게 아니라, "어떤 행동이 더 높은 점수를 줄까?"를 계산하며 학습합니다.
• 결과적으로 AI 는 나쁜 레시피만 보던 상태에서, 상상력을 통해 완성된 최고의 레시피를 배우게 되어, 훨씬 더 똑똑해집니다.

비유:
이제 셰프는 실패한 레시피와, 상상력을 통해 보완된 완벽한 레시피를 모두 보고 요리합니다. 그 결과, 처음엔 실패하던 요리도 이제는 완벽하게 만들어냅니다.

---

🌟 3. IPD 의 놀라운 성과

이론만 좋은 게 아닙니다. 연구진은 D4RL이라는 유명한 AI 테스트 벤치마크에서 이 방법을 시험해 보았습니다.

• 결과: 기존에 가장 잘하던 방법들보다 훨씬 더 높은 점수를 받았습니다.
• 특징: 데이터가 부족하거나 실패 사례가 많은 상황에서도, 상상력을 통해 부족한 부분을 채워주어 안정적이고 강력한 성능을 발휘했습니다.

---

📝 4. 한 줄 요약

IPD 는 AI 가 "실패한 기록"만 보고 배울 때 겪는 한계를, "머릿속으로 완벽한 시나리오를 상상하고 그걸로 다시 배우는" 방식을 통해 해결해줍니다.

마치 나쁜 경험담을 읽다가, **"만약에 내가 이렇게 했다면?"**이라는 상상을 통해 최고의 경험을 만들어내고, 그걸로 자신을 업그레이드하는 것과 같습니다. AI 가 더 안전하고 똑똑하게 현실 세계의 문제를 해결할 수 있는 새로운 길을 열어준 셈입니다.

기술 요약

1. 문제 정의 (Problem)

오프라인 강화학습 (Offline RL) 은 환경과의 상호작용 없이 고정된 데이터셋만으로 정책을 학습할 수 있어 안전성과 비용 효율성이 높지만, 다음과 같은 근본적인 한계를 가지고 있습니다.

• 정적 데이터셋의 한계: 기존 데이터가 하위 최적 (suboptimal) 인 경우, 학습된 정책이 이를 그대로 모방하게 되어 최적의 행동을 도출하기 어렵습니다.
• Decision Transformer 의 구조적 결함: 최근 주목받는 Decision Transformer 와 같은 시퀀스 기반 정책 모델은 강력한 시퀀스 모델링 능력을 가지고 있지만, 동적 계획법 (Dynamic Programming) 기반의 RL 메커니즘이 부재합니다. 이로 인해 하위 최적의 궤적들을 연결하여 최적 정책을 구성하는 데 어려움을 겪습니다.
• Return-to-Go (RTG) 의 불안정성: 기존 Transformer 기반 모델은 수동으로 설정된 목표 반환값 (Return-to-Go) 에 의존하는데, 이는 추론 시 불안정성을 초래하고 최적의 의사결정을 방해할 수 있습니다.

2. 제안 방법: IPD (Imaginary Planning Distillation)

저자들은 상상적 계획 증류 (Imaginary Planning Distillation, IPD) 라는 새로운 프레임워크를 제안합니다. 이는 오프라인 데이터 생성, 지도 학습, 온라인 추론 전 과정에 '가상 계획 (Imaginary Planning)'을 통합하여 Transformer 기반 정책을 강화합니다.

IPD 는 크게 4 단계로 구성됩니다:

가. 준최적 가치 함수 학습 (Offline Quasi-Optimal Value Function Learning)

• 오프라인 Q-러닝의 분포 외 (OOD) 상태 - 행동 쌍에 대한 가치 과대평가 문제를 해결하기 위해 Implicit Q-Learning (IQL) 원리를 따릅니다.
• Huber-expectile 회귀를 사용하여 이상치에 강건한 준최적 가치 함수 ($V_\psi$) 와 Q 함수 ($Q_\theta$) 를 학습합니다.
• 이를 통해 데이터셋 내에서 가장 유망한 행동을 식별할 수 있는 기준을 마련합니다.

나. 불확실성 측정 가능한 월드 모델 학습 (World Model with Uncertainty Measure)

• 가상 롤아웃 (Imaginary Rollouts) 을 생성하기 위해 불확실성 (Uncertainty) 을 정량화할 수 있는 월드 모델을 학습합니다.
• 앙상블 (Ensemble) 기반의 확률적 모델을 사용하여 환경의 무작위성 (Aleatoric uncertainty) 과 데이터 부족으로 인한 지식의 불확실성 (Epistemic uncertainty) 을 모두 모델링합니다.
• 계산 효율성을 위해 기하학적 Jensen-Shannon (GJS) 발산을 사용하여 모델 간 불일치를 측정하고, 이를 통해 신뢰할 수 있는 상태 - 행동 쌍만 선별합니다.

다. 상상적 계획을 통한 데이터 증강 (Data Augmentation with Imaginary Planning)

• 학습된 가치 함수를 사용하여 오프라인 데이터셋 내의 하위 최적 상태 (Suboptimal States) 를 식별합니다.
• 식별된 상태에 대해 모델 예측 제어 (MPC) 를 적용하여, 학습된 월드 모델과 준최적 가치 함수를 기반으로 더 높은 보상을 기대할 수 있는 가상 최적 궤적 (Imaginary Optimal Trajectories) 을 생성합니다.
• 생성된 궤적은 월드 모델의 불확실성 임계값을 통과한 경우에만 증강 데이터셋에 포함됩니다.

라. 상상적 계획 증류 (Imaginary Planning Distillation)

• 증강된 데이터셋과 가치 함수의 지식을 Transformer 정책으로 증류합니다.
• 손실 함수 (Loss Function):
  1. 시퀀스 모델링 항: 증강된 데이터셋의 고품질 행동을 모방합니다.
  2. Q-값 정규화 항: 가치 함수가 제공하는 행동 기울기 ($\nabla Q$) 를 활용하여 정책이 더 높은 Q-값을 갖는 행동을 하도록 유도합니다.
• 동적 Return-to-Go: 추론 시 수동으로 설정된 RTG 대신, 학습된 준최적 가치 함수 ($V_\psi$) 를 조건부 입력으로 사용하여 동적이고 안정적인 목표 반환값을 예측합니다.

3. 주요 기여 (Key Contributions)

1. 새로운 프레임워크 (IPD): 지도 학습 (시퀀스 모델링) 과 상상적 계획 (가상 롤아웃) 을 원활하게 통합한 오프라인 RL 프레임워크를 제안했습니다.
2. 계획 원리의 통합: Transformer 기반 정책에 암묵적 동적 계획법 (가치 함수) 과 명시적 모델 예측 제어 (MPC) 를 모두 통합하여, 하위 최적 데이터셋의 한계를 극복하고 최적 궤적 생성 능력을 향상시켰습니다.
3. 성능 검증 및 스케일링 법칙: D4RL 벤치마크를 통한 광범위한 실험을 통해 기존 SOTA 방법들 (Q-learning 기반 및 Transformer 기반) 보다 우수한 성능을 입증했습니다. 또한, 생성된 가상 데이터의 양과 정책 성능 간의 스케일링 법칙 (Scaling Law) 을 발견하여 데이터 증강의 효과를 규명했습니다.

4. 실험 결과 (Results)

• 벤치마크: D4RL (Gym, Kitchen, Adroit 도메인 포함 10 개 작업) 에서 평가 수행.
• 성능: IPD 는 CQL, IQL, Decision Transformer (DT), QDT, Reinformer 등 다양한 최신 오프라인 RL 방법론을 압도적으로 능가했습니다.
  • 특히 하위 최적 데이터가 많은 환경 (예: walker-medium-replay, kitchen-partial) 에서 큰 성능 향상을 보였습니다.
  • Adroit 도메인 (희소 보상 환경) 에서도 인간 데모의 한계를 극복하고 우수한 성능을 달성했습니다.
• Ablation Study:
  • MPC vs Greedy Q-learning: MPC 를 활용한 데이터 생성이 단순 탐욕적 (Greedy) Q-러닝 기반 생성보다 훨씬 우수한 성능을 보임을 확인했습니다.
  • 가치 함수 가이드: 수동 RTG 대신 학습된 준최적 가치 함수를 사용할 경우, 추론 시 성능 변동성 (Variance) 이 크게 감소하고 안정성이 향상됨을 입증했습니다.

5. 의의 및 결론 (Significance)

IPD 는 오프라인 강화학습의 핵심 난제인 '하위 최적 데이터의 한계' 와 'Transformer 모델의 계획 능력 부재' 를 동시에 해결하는 획기적인 접근법입니다.

• 실용성: 실제 환경과의 상호작용 없이도, 학습된 월드 모델을 통해 '상상'된 최적 경로를 학습 데이터에 추가함으로써 정책의 성능을 극대화할 수 있습니다.
• 안정성: 수동 튜닝이 필요한 RTG 를 학습된 가치 함수로 대체함으로써 추론 과정의 안정성을 크게 개선했습니다.
• 미래 전망: 이 연구는 오프라인 RL 분야에서 모델 기반 (Model-based) 접근법과 시퀀스 모델링을 결합한 새로운 패러다임을 제시하며, 복잡한 의사결정 작업에서의 정책 학습 효율성을 높이는 중요한 발걸음이 됩니다.