Robotics Research Hub
Latent Memory Palace: Reasoning for Control as Autoregressive Variational Inference

Latent Memory Palace: Reasoning for Control as Autoregressive Variational Inference

Chuning Zhu, Eva Xu, Jose Barreiros, Krishnan Srinivasan, Paarth Shah, Abhishek Gupta · N/A · 2026

reinforcement-learning control
PDF arXiv

Framework

N/A

License

N/A

Stars

N/A

Summary

Human decision-making is highly flexible -- some actions are taken immediately; others require longer deliberation.

Abstract Summary

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-$\texttt{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.

Key Points

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

Latent Memory Palace: Reasoning as Autoregressive Variational Inference

We consider the imitation learning setting, where the goal is to learn a control policy $\pi_{\theta}:\mathcal{O}\rightarrow\mathcal{A}$ from a dataset of expert demonstrations, $\mathcal{D}=\{(o_{i},a_{i})\}_{i=1}^{n}$, with observations $o_{i}$ and actions $a_{i}$. For the sake of practicality, $o_{i}$ may be a stacked history of observations with language instructions, and $a_{i}$ may be a chunk of actions.

Latent Memory Palace: Reasoning as Autoregressive Variational Inference — Latent Memory Palace: Reasoning for Control as Autoregressive Variational Inference
Figure 2: Training LMP-$\pi$ alternates between two stages. (1) Latent rollout: the posterior generates a buffer of latent trajectories. (2) Model update: the model optimizes a clipped-surrogate of the variational lower bound, computed from the posterior, the prior, and the decoder.

Experiments

We conduct experiments to answer the following research questions. (1) Is LMP-$\pi$ a performant generative control policy? (2) Is LMP- tok an effective action tokenizer? (3) Does LMP demonstrate adaptive allocation of test-time compute? (4) Which factors contribute to LMP’s performance? We evaluate our method on real-world and simulated visuomotor control benchmarks as shown in Fig. 4 .

(1) DROID [ 31 ] is a real-world manipulation platform with multitask training data.

Experiments — Latent Memory Palace: Reasoning for Control as Autoregressive Variational Inference
Figure 3: LMP-$\pi$ architecture. The model consists of a latent encoder and an action decoder. The encoder is a causal transformer cross-attending to observations and actions. The decoder is a bidirectional transformer cross-attending to observations and latents.

Sources and Demonstrations

Figures are reproduced from Zhu et al. See the full paper for experimental details and the project page, when the paper links one, for demonstrations and videos.

Share

Related Papers