In-Context Planning with Latent Temporal Abstractions

TL;DR

I-TAP addresses the bottlenecks of planning-based offline RL for continuous control under partial observability by learning a discrete latent temporal-abstraction space and an in-context sequence prior. It discretizes observation-macro-action sequences with an observation-conditioned Residual-Quantized VAE and models latent dynamics with a Transformer, enabling Monte Carlo Tree Search over latent tokens at decision time. The approach achieves robust adaptation across stochastic regimes and high-dimensional tasks, matching or surpassing strong baselines while requiring no gradient updates during testing. This facilitates efficient, context-aware planning in partially observable, stochastic environments, with scalable planning in latent spaces suitable for real-world control problems.

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.

Figures (8)

• Figure 1: Overview of I-TAP. Left: A residual-quantized VAE (RQ-VAE) discretizes continuous observation–action trajectories into a coarse-to-fine token stack. Right: Normalized return and per-decision latency as functions of planning horizon and context size on Stochastic MuJoCo, highlighting the importance of a properly sized context window for effective in-context planning under environmental stochasticity and partial observability.
• Figure 2: An overview of our RQ-VAE model that discretizes state-macro action sequences and temporal prior for I-TAP
• Figure 3: Macro-level MCTS overview. Each iteration uses P-UCT to select a macro-action, expands several candidates and their predicted outcomes in parallel, then backs up the resulting Q-estimates through the search tree to steer subsequent exploration.
• Figure 4: Ablation results across Adroit (expert) and MuJoCo Hopper. We plot $\Delta$ scores relative to I-TAP on Hopper medium– expert, which serves as the baseline (zero line).
• Figure 5: Test reconstruction losses across residual levels.