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 & tables
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.
Model-Based
Model-Free
Dataset Type
Env (Noise)
I-TAP
L-MAP
TAP
1R2R
DT
CQL
IQL
Medium-Expert
Hopper (5)
82.87 ± 3.55
71.49 ± 3.46
37.31 ± 3.66
37.99 ± 2.71
61.87 ± 2.56
68.03 ± 3.94
44.83 ± 2.58
Hopper (2.5)
104.45 ± 2.66
93.40 ± 3.65
40.86 ± 5.42
52.19 ± 8.37
73.86 ± 2.68
106.17 ± 2.16
60.61 ± 3.46
Hopper (0)
111.71 ± 0.08
106.74 ± 2.24
85.55 ± 3.83
57.40 ± 6.06
101.6 ± 1.85
105.4
91.5
Walker2D (12)
93.50 ± 3.15
92.75 ± 1.34
91.09 ± 2.78
32.38 ± 4.55
52.42 ± 1.27
83.18 ± 3.70
68.61 ± 3.33
Walker2D (7)
97.29 ± 1.68
93.48 ± 1.20
91.40 ± 1.42
56.48 ± 7.51
64.67 ± 1.00
91.44 ± 1.44
86.66 ± 1.84
Table 1: Normalised results for high-noise (5), moderate-noise (2.5), deterministic (0) environments. Bold numbers indicate the best scores in each row.
Dataset Type
Env (Noise)
I-TAP
L-MAP
DT
Medium-Expert
Hopper (7.5)
70.43 ± 3.17
62.90 ± 3.50
58.96 ± 0.21
Hopper (3.75)
85.07 ± 3.57
80.51 ± 3.86
64.97 ± 0.25
Hopper (1.25)
106.91 ± 2.17
94.65 ± 3.03
87.41 ± 0.24
Walker2D (15)
88.26 ± 3.67
87.21 ± 3.42
44.07 ± 0.26
Walker2D (9.5)
94.29 ± 1.97
92.31 ± 1.81
56.23 ± 0.22
Walker2D (3.5)
98.80 ± 1.00
96.53 ± 0.62
65.07 ± 0.18
Table 2: Normalised results (mean ± std). Noise parameter values are shown in parentheses. Bold indicates the best score in each row.
Dataset Type
Env
I-TAP
L-MAP
TAP
Cloned
Pen
85.44 ± 8.19
60.68 ± 7.88
46.44 ± 7.54
Cloned
Hammer
4.38 ± 1.28
2.43 ± 0.29
1.32 ± 0.12
Cloned
Door
14.17 ± 1.34
13.22 ± 1.34
13.45 ± 1.43
Cloned
Relocate
0.08 ± 0.02
0.15 ± 0.13
-0.23 ± 0.01
Expert
Pen
133.81 ± 5.23
126.60 ± 5.60
127.40 ± 7.70
Expert
Hammer
128.37 ± 0.21
127.16 ± 0.29
127.60 ± 1.70
Table 3: Adroit robotic hand control results.
Figure 4: Ablation results across Adroit (expert) and MuJoCo Hopper. We plot Δ scores relative to I-TAP on Hopper medium–expert, which serves as the baseline (zero line).
Figure 5: Test reconstruction losses across residual levels.
Figure 6: Average decision time vs. context length with 16 action candidates (left) and 32 action candidates (right).
Compact JEPA world models enable efficient latent-space planning, but low-dimensional representation trained under reward-free self-supervision must encode both action-conditioned dynamics and predictable visual context. This competition can entangle controllable state with high-rank nuisance appearance and degrade planning as scenes become more complex. We introduce LRC-JEPA, a lightweight end-to-end world model that routes information into a compact predictive latent z and learned-query residual-context embeddings u. Only z is propagated by the dynamics model and used for planning, while u captures temporally persistent information for cross-attention reconstruction; a differentiable residual connection encourages the latent to retain complementary dynamic content. Under explicit assumptions, we show that the resulting representation is sufficient, minimal, nuisance-invariant, and disentangled. Across four simulated control environments, LRC-JEPA improves average planning success over a parameter-matched JEPA baseline by 9 percentage points and matches or exceeds substantially larger pretrained models. On the real-world Bridge-v2 set, its 5.5M-parameter active encoder outperforms DINO-WM (22.1M) and V-JEPA2 (303.9M) encoders while also enabling faster planning. Physical-state probes, reconstruction interventions, and ablations confirm the effectiveness of LRC-JEPA's representation disentanglement.
Humans solve complex problems by constructing plans and mentally simulating their outcomes with an internal model of the world. Machine learning has made substantial progress in learning world models that predict the consequences of action sequences, yet the procedures used to plan with these models remain largely hand-designed. Most planners rely on fixed search or optimization rules; approaches that learn aspects of search typically imitate a predefined optimizer or use planning to inform an amortized policy, rather improving multi-step plans. We introduce Reinforced Planning, a method that learns the plan-update itself by reinforcing update rules that produce better plans, using gradients propagated through a differentiable world model. We instantiate Reinforced Planning in RP1, which learns a critic over imagined outcomes via temporal-difference learning and a neural plan-improvement operator trained via imagined rollouts with a pretrained world model. RP1 can be trained fully offline without environment interaction; environment episodes are used only for checkpoint selection. Across visual navigation, arm reaching, and robotic manipulation on two world-model backbones, RP1 matches or exceeds existing planners, achieving near-perfect success in several settings while using 1,000× fewer world-model rollouts than the strongest alternative (CEM) and planning up to 67× faster under concurrent planners inference.
Human decision-making is highly flexible -- some actions are taken immediately; others require longer deliberation. Language models have exhibited a similar capacity for adaptive "reasoning." However, transferring this capability to continuous control policies has been challenging, as directly reasoning in language space may lack the granularity for spatial understanding and precise motions. In this work, we show that reasoning for control policies can emerge by organizing information in an autoregressive latent space reminiscent of a memory palace, where retrieval is iterative and adaptive. Our method, Latent Memory Palace (LMP), formulates reasoning as variational inference with an autoregressive latent distribution. We derive a latent-space reinforcement learning technique to tractably optimize its variational lower bound. The resulting policy, LMP-π, achieves strong empirical performance in simulation and real-world domains while exhibiting interpretable, adaptive allocation of test-time compute. We further show that the same framework yields a variable-length action tokenizer, LMP-tok, which significantly improves the performance of downstream autoregressive policies. Together, these results present a new perspective on latent reasoning for control through the lens of variational inference.
Chuning Zhu, Eva Xu, Jose Barreiros +3
University of Washington · Toyota Research Institute