Imagine-then-Plan: Agent Learning from Adaptive Lookahead with World Models

1. 문제 정의 (Problem)

최근 대규모 언어 모델 (LLM) 기반 에이전트는 다양한 디지털 및 물리적 환경에서 추론과 상호작용이 가능해졌으나, 여전히 **'얕은 그라운딩 (Shallow Grounding)'**이라는 한계에 직면해 있습니다.

• 현재의 한계: 기존 에이전트는 주로 즉각적인 관찰과 과거의 상호작용 기록에 기반하여 결정을 내립니다. 이는 미래의 상태를 예측하거나 행동의 장기적인 결과를 시뮬레이션할 수 없음을 의미합니다.
• 결과: 에이전트는 환경의 인과 관계를 깊이 이해하지 못해, 실행 후 되돌릴 수 없는 오류나 상태 충돌을 발견하게 되며, 이는 복잡한 장기 작업 (Long-horizon tasks) 에서 치명적인 실패로 이어집니다.
• 기존 방법의 부족: 기존 세계 모델 (World Models) 기반 접근법들은 주로 단일 단계 검증이나 고정된 시간 범위 (Fixed-horizon) 의 시뮬레이션에 의존합니다. 이는 복잡한 작업의 장기적 의존성을 놓치거나, 사소한 행동에 대해 불필요하게 깊은 시뮬레이션을 수행하여 계산 비용을 낭비하는 비효율성을 보입니다.

2. 제안 방법: Imagine-then-Plan (ITP) (Methodology)

이 논문은 에이전트가 학습된 세계 모델을 통해 미래 상태를 능동적으로 상상하고, 이를 기반으로 계획을 수립하는 Imagine-then-Plan (ITP) 프레임워크를 제안합니다.

핵심 개념: POIMDP (Partially Observable and Imaginable MDP)

기존의 부분 관측 마르코프 결정 과정 (POMDP) 을 확장하여, 에이전트가 **관측 가능한 현재 (Concrete Present)**와 **상상 가능한 미래 (Foresighted Future)**를 모두 고려하는 새로운 의사결정 과정을 정의합니다.

• 에이전트의 정책은 현재 상태뿐만 아니라, 세계 모델을 통해 생성된 '상상된 궤적 (Imagined Trajectories)'에 조건을 두어 최적의 행동을 선택합니다.

적응형 Lookahead 메커니즘 (Adaptive Lookahead)

ITP 의 가장 중요한 혁신은 작업의 진행 상황과 최종 목표 간의 균형을 고려하여 상상하는 시간 범위 (Lookahead Horizon, $K_t$) 를 동적으로 조절하는 것입니다.

• 동적 조절: 중요한 결정이 필요한 단계에서는 깊은 미래 시뮬레이션을 수행하고, 단순한 단계에서는 시뮬레이션을 줄여 계산 효율성을 극대화합니다.
• 자기 수정 (Self-Correction): 상상된 궤적을 통해 목표 달성 여부, 잠재적 병목 현상, 충돌 등을 사전에 감지하여 실제 환경에 행동하기 전에 계획을 수정합니다.

구현 변형 (Variants)

ITP 는 두 가지 변형으로 구현됩니다:

1. ITP-I (In-Imagination Learning, Training-free):
   • 추가 학습 없이 추론 시간 (Inference time) 에 작동합니다.
   • 고정된 정책과 세계 모델을 사용하며, '상상 - 반성 (Reflect)' 단계를 거쳐 행동을 선택합니다.
   • LLM 에이전트의 플러그 앤 플레이 (Plug-and-play) 향상 도구로 활용됩니다.
2. ITP-R (Reinforced Training):
   • 적응형 Lookahead 선택과 행동 정책을 함께 최적화하기 위해 강화학습을 적용합니다.
   • 가짜 라벨링 (Pseudo-Labeling): 전문가 궤적을 기반으로 세계 모델을 통해 다양한 Lookahead 길이의 시뮬레이션을 수행하고, 최적의 길이를 가짜 라벨로 생성합니다.
   • 온라인 최적화: Actor-Critic 알고리즘 (A2C) 을 사용하여 행동 정책과 Lookahead 예측기 (K-head) 를 동시에 학습하며, 계산 비용과 작업 성공률 사이의 균형을 위한 보상 함수를 설계합니다.

3. 주요 기여 (Key Contributions)

1. POIMDP 개념 정립: 과거 상호작용과 상상된 미래를 통합한 새로운 에이전트 의사결정 프레임워크를 수학적으로 정의했습니다.
2. 적응형 Lookahead 메커니즘: 고정된 시뮬레이션 범위의 한계를 극복하고, 작업의 복잡도와 단계에 따라 동적으로 상상 깊이를 조절하는 방식을 도입했습니다.
3. 성능 검증: 학습이 없는 변형 (ITP-I) 과 강화학습 변형 (ITP-R) 을 통해 다양한 벤치마크에서 기존 방법론을 압도하는 성능을 입증했습니다.

4. 실험 결과 (Results)

저자는 ALFWorld(가정 환경), ScienceWorld(과학 실험), WebShop(웹 쇼핑), StableToolBench(도구 사용) 등 4 가지 주요 에이전트 벤치마크에서 실험을 수행했습니다.

• 성능 향상:
  • ITP-I: 추가 학습 없이도 ReAct, RAP 등 강력한 프롬프팅 기반 베이스라인보다 성공률 (Success Rate) 을 크게 향상시켰습니다.
  • ITP-R: 기존 학습 기반 방법론 (SFT, WKM, IWM 등) 을 모두 능가하는 최고 성능을 기록했습니다. (예: ALFWorld 에서 Qwen3-8B 기반 시 88.57% 성공률).
• 적응형 Lookahead의 효과:
  • 고정된 Lookahead vs. 적응형: 고정된 깊이 ($k$) 를 사용하는 경우, 깊이가 깊어질수록 계산 비용은 급증하지만 성공률은 오히려 감소하는 경향이 있었습니다. 반면, ITP 의 적응형 방식은 낮은 계산 비용으로 높은 성공률을 달성했습니다.
  • 무작위 Lookahead vs. 적응형: 상태에 기반한 적응형 선택이 단순히 깊이를 무작위로 변경하는 것보다 훨씬 효율적이었습니다.
• 세계 모델의 신뢰성:
  • 시뮬레이션 깊이가 깊어질수록 세계 모델의 예측 정확도 (Fact-F1) 는 감소하지만, 에이전트의 성공률은 특정 깊이 (예: $k=3$) 에서 정점을 보입니다. 이는 ITP 가 모델의 예측 신뢰 범위 내에서 최적의 깊이를 동적으로 찾음을 의미합니다.

5. 의의 및 결론 (Significance)

이 논문은 자율 에이전트 학습에 있어 세계 모델을 통한 '의식적인 (Deliberative)' 사고의 중요성을 강조합니다.

• 패러다임 전환: 단순한 반응형 (Reactive) 에이전트에서, 미래를 시뮬레이션하고 잠재적 오류를 사전에 교정하는 계획형 에이전트로의 전환을 가능하게 합니다.
• 효율성과 성능의 균형: 불필요한 계산 비용을 줄이면서도 복잡한 장기 작업의 성공률을 극대화하는 적응형 Lookahead 메커니즘은 향후 복잡한 에이전트 시스템 설계에 중요한 통찰을 제공합니다.
• 범용성: 다양한 LLM 아키텍처 (Qwen, Llama, DeepSeek 등) 와 다양한 작업 도메인에서 효과적으로 작동함을 입증하여, 모델에 구애받지 않는 강력한 프레임워크임을 보여줍니다.

요약하자면, ITP는 에이전트가 "생각한 후 행동 (Imagine-then-Plan)"할 수 있도록 하여, 실제 환경과의 비효율적인 상호작용을 줄이고 복잡한 문제 해결 능력을 획기적으로 향상시킨 혁신적인 접근법입니다.