In-Context Planning with Latent Temporal Abstractions

Baiting Luo; Yunuo Zhang; Nathaniel S. Keplinger; Samir Gupta; Abhishek Dubey; Ayan Mukhopadhyay

cs.LGFeb 21, 2026

In-Context Planning with Latent Temporal Abstractions

Authors: Baiting Luo, Yunuo Zhang, Nathaniel S. Keplinger, Samir Gupta, Abhishek Dubey, Ayan Mukhopadhyay

Organizations: Vanderbilt University · William & Mary

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 & 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 $\pm$ 3.55	71.49 $\pm$ 3.46	37.31 $\pm$ 3.66	37.99 $\pm$ 2.71	61.87 $\pm$ 2.56	68.03 $\pm$ 3.94	44.83 $\pm$ 2.58
	Hopper (2.5)	104.45 $\pm$ 2.66	93.40 $\pm$ 3.65	40.86 $\pm$ 5.42	52.19 $\pm$ 8.37	73.86 $\pm$ 2.68	106.17 $\pm$ 2.16	60.61 $\pm$ 3.46
	Hopper (0)	111.71 $\pm$ 0.08	106.74 $\pm$ 2.24	85.55 $\pm$ 3.83	57.40 $\pm$ 6.06	101.6 $\pm$ 1.85	105.4	91.5
	Walker2D (12)	93.50 $\pm$ 3.15	92.75 $\pm$ 1.34	91.09 $\pm$ 2.78	32.38 $\pm$ 4.55	52.42 $\pm$ 1.27	83.18 $\pm$ 3.70	68.61 $\pm$ 3.33
	Walker2D (7)	97.29 $\pm$ 1.68	93.48 $\pm$ 1.20	91.40 $\pm$ 1.42	56.48 $\pm$ 7.51	64.67 $\pm$ 1.00	91.44 $\pm$ 1.44	86.66 $\pm$ 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 $\pm$ 3.17	62.90 $\pm$ 3.50	58.96 $\pm$ 0.21
	Hopper (3.75)	85.07 $\pm$ 3.57	80.51 $\pm$ 3.86	64.97 $\pm$ 0.25
	Hopper (1.25)	106.91 $\pm$ 2.17	94.65 $\pm$ 3.03	87.41 $\pm$ 0.24
	Walker2D (15)	88.26 $\pm$ 3.67	87.21 $\pm$ 3.42	44.07 $\pm$ 0.26
	Walker2D (9.5)	94.29 $\pm$ 1.97	92.31 $\pm$ 1.81	56.23 $\pm$ 0.22
	Walker2D (3.5)	98.80 $\pm$ 1.00	96.53 $\pm$ 0.62	65.07 $\pm$ 0.18

Table 2: Normalised results (mean

\pm

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 $\pm$ 8.19	60.68 $\pm$ 7.88	46.44 $\pm$ 7.54
Cloned	Hammer	4.38 $\pm$ 1.28	2.43 $\pm$ 0.29	1.32 $\pm$ 0.12
Cloned	Door	14.17 $\pm$ 1.34	13.22 $\pm$ 1.34	13.45 $\pm$ 1.43
Cloned	Relocate	0.08 $\pm$ 0.02	0.15 $\pm$ 0.13	-0.23 $\pm$ 0.01
Expert	Pen	133.81 $\pm$ 5.23	126.60 $\pm$ 5.60	127.40 $\pm$ 7.70
Expert	Hammer	128.37 $\pm$ 0.21	127.16 $\pm$ 0.29	127.60 $\pm$ 1.70

Table 3: Adroit robotic hand control results.

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).

Panel of Figure 5 — Figure 5: Test reconstruction losses across residual levels.

Panel of Figure 6 — Figure 6: Average decision time vs. context length with 16 action candidates (left) and 32 action candidates (right).

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Environment	Hyper-parameter	Value
All	learning rate	$1\times 10^{-4}$
All	batch size	512
All	dropout probability	0.1
All	number of attention heads	4
All	macro action length $L$	3
All	embedding size (latent code)	512

Appendix

Table 4: List of Hyper-parameters

Model-Based Planning

Model-Based Planning uses explicit or learned models of how actions change states to simulate futures, compare plans, and choose actions.

See all 199 papers

LRC-JEPA: Disentangling Dynamics and Residual Context for Efficient World Models
Sep 28, 2026Luzhe Huang, Lei Chu, Jingyi Liang +1Disentangled Representation Learning World Models for Robotics
Independent researcher
Reinforced Planning with Latent World Models
Aug 19, 2026Armin Sommer, Jannik SchillingWorld Model-Based Planning Robotic RL
Pantheon Industries
Latent Memory Palace: Reasoning for Control as Autoregressive Variational Inference
Jul 9, 2026Chuning Zhu, Eva Xu, Jose Barreiros +3RL Control Latent Action Learning
University of Washington · Toyota Research Institute

Sep 28, 2026J/K move · Enter open · S save

LRC-JEPA: Disentangling Dynamics and Residual Context for Efficient World Models

Luzhe Huang, Lei Chu, Jingyi Liang, Yuhuan Zhao

Independent researcher

Disentangled Representation Learning World Models for Robotics World Model-Based Planning Action-Conditioned World Models Model-Based Planning Latent World Models+1

Open paper AskPDF arXiv

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

\mathbf{z}

and learned-query residual-context embeddings

\mathbf{u}

. Only

\mathbf{z}

is propagated by the dynamics model and used for planning, while

\mathbf{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.

Abstract

Model-Based Planning

Explore similar work

LRC-JEPA: Disentangling Dynamics and Residual Context for Efficient World Models

Reinforced Planning with Latent World Models

Latent Memory Palace: Reasoning for Control as Autoregressive Variational Inference