[논문 리뷰] In-Context Planning with Latent Temporal Abstractions

Baiting Luo, Yunuo Zhang외 4명|arXiv (Cornell University)|2026. 02. 21.

Reinforcement Learning in Robotics인용 수 0

한 줄 요약

I-TAP은 학습된 이산 잠재 시간 추상 공간에서 계획을 수행함으로써 맥락 내 적응(in-context adaptation)과 온라인 계획을 통합하여 부분 관찰성 및 확률적 동역학 하에서 테스트 시 그레이디언트 업데이트 없이도 강건한 의사결정을 가능하게 한다.

ABSTRACT

Planning-based reinforcement learning for continuous control is bottlenecked by two practical issues: planning at primitive time scales leads to prohibitive branching and long horizons, while real environments are frequently partially observable and exhibit regime shifts that invalidate stationary, fully observed dynamics assumptions. We introduce I-TAP (In-Context Latent Temporal-Abstraction Planner), an offline RL framework that unifies in-context adaptation with online planning in a learned discrete temporal-abstraction space. From offline trajectories, I-TAP learns an observation-conditioned residual-quantization VAE that compresses each observation-macro-action segment into a coarse-to-fine stack of discrete residual tokens, and a temporal Transformer that autoregressively predicts these token stacks from a short recent history. The resulting sequence model acts simultaneously as a context-conditioned prior over abstract actions and a latent dynamics model. At test time, I-TAP performs Monte Carlo Tree Search directly in token space, using short histories for implicit adaptation without gradient update, and decodes selected token stacks into executable actions. Across deterministic MuJoCo, stochastic MuJoCo with per-episode latent dynamics regimes, and high-dimensional Adroit manipulation, including partially observable variants, I-TAP consistently matches or outperforms strong model-free and model-based offline baselines, demonstrating efficient and robust in-context planning under stochastic dynamics and partial observability.

연구 동기 및 목표

• horizon을 줄이고 가지를 줄이기 위해 잠재 시간 추상화를 도입하여 연속 제어를 위한 오프라인 RL의 계획 병목을 해결한다.
• 역사 기반 조건 priors를 통해 관찰되지 않은 잠재 동역학에 대한 맥락 내 적응을 가능하게 한다.
• 잠재 토큰 공간에서 작동하는 온라인 플래너(MCTS)를 통합하여 비최적 오프라인 데이터를 극복한다.
• 관찰 및 매크로-액션에 조건화된 잔차 양자화 VAE를 통해 이산 잠재 표현을 학습한다.
• 결정론적 및 확률적 MuJoCo 환경과 고차원 Adroit 조작에서 강건성을 입증한다.

제안 방법

• 관찰-매크로액션 시퀀스를 관찰 조건부 잔차 양자화 VAE(RQ-VAE)로 이산화하여 깊이-D 토큰 스택을 생성한다.
• 짧은 히스토리와 현재 관찰에 조건화된 토큰 스택을 자기회귀적으로 예측하도록 시간적 Transformer를 학습한다.
• 시간과 깊이에 걸쳐 분해하는 자기회귀 모델을 통해 맥락 조건부 잠재 코드 스택에 대한 잠재 사전(latent prior)을 학습한다.
• 맥락 가이드 사전을 사용하여 잠재 토큰 공간에서 직접 몬테카를로 트리 탐색(MCTS)을 수행하고 매크로액션을 선택한 다음 원시 행동으로 디코딩한다.
• 학습 안정화를 위해 학습 중 고분산 수익(G_t)을 마스킹하고, 맥락에서 G_t^(L)처럼 짧은 수익을 유지하여 체제 식별에 사용한다.
• 선택된 잠재 스택을 실행 가능한 행동 시퀀스로 디코딩하여 테스트 시 그레이디언트 업데이트 없이 환경과 상호 작용한다.

실험 결과

연구 질문

• RQ1 학습된 이산 잠재 시간 추상 공간에서의 계획이 부분 관찰성과 잠재 동역학하에서 계획의 효율성과 강건성을 향상시킬 수 있는가?
• RQ2 맥락 조건부 잠재 모델링과 온라인 계획이 확률적이고 고차원적인 과제 전반에서 오프라인 모델 프리 및 모델 기반 기준선보다 우수한가?
• RQ3 매크로액션 길이, 맥락 크기, 계획 수, 잔차 깊이가 성능 및 적응에 어떻게 영향을 미치는가?
• RQ4 히스토리를 통한 맥락 내 적응이 잠재 환경 매개변수를 추론하고 테스트 시 그레이디언트 업데이트 없이 계획을 이끄는 데 충분한가?

주요 결과

• I-TAP은 결정론적 및 확률적 MuJoCo 작업과 고차원 Adroit 조작에 대해 강력한 오프라인 RL 및 계획 기준과 동등하거나 이를 능가한다.
• 맥락 조건부 사전을 갖춘 잠재 토큰 위의 계획은 분기와 수평선을 줄여 강건한 맥락 내 적응을 가능하게 한다.
• 잔차 양자화 VAE는 깊이-D 코드가 있는 이산 표현을 통해 고차원 관찰 및 매크로액션에 대한 디코딩 충실도를 유지한다.
• 맥락 길이와 계획 수는 성능에 체계적인 영향을 미치며, 더 긴 맥락은 적응을 향상시키고 더 깊은 예측은 부분 관찰성 효과를 완화한다.
• I-TAP은 부분 관찰성(POMDP) 환경에서도 최상위 성능을 유지하며, 통합 계획과 맥락 내 적응이 없는 기준선보다 우수하다.