Latent Memory Palace: Reasoning for Control as Autoregressive Variational Inference

Abstract: Human decision-making is highly flexible -- some actions are taken immediately; others require longer deliberation. LLMs 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.

• The paper presents a novel autoregressive variational inference framework that organizes latent reasoning for adaptive control in robotic tasks.
• It demonstrates significant performance gains and improved generalization compared to diffusion-based policies across diverse manipulation benchmarks.
• The model combines a causal transformer encoder with a bidirectional decoder to enable interpretable latent allocation and efficient action tokenization.

Latent Memory Palace: Reasoning for Control as Autoregressive Variational Inference

Motivation and Conceptual Framework

The "Latent Memory Palace: Reasoning for Control as Autoregressive Variational Inference" (2607.08724) constructs a bridge between the iterative and adaptive reasoning observed in human cognition and the demands of sequential decision-making in robotic and control tasks. Existing LLMs leverage intermediate token generation for improved reasoning—a paradigm not directly transferable to robot control due to granularity and semantic mismatches between language and motor actions. The paper introduces Latent Memory Palace (LMP), which organizes reasoning as variational inference in a variable-length, discrete, autoregressive latent space. This latent space serves as an internal deliberation path, terminated via EOS, enabling adaptive allocation of computational resources depending on the complexity and ambiguity of the task.

Figure 1: LMP enables adaptive, iterative reasoning via variable-length autoregressive latent distributions, yielding policies that dynamically allocate test-time computation.

Mathematical Formulation and Learning Protocol

LMP parameterizes the control policy $\pi_\theta$ via an autoregressive latent variable model: actions $a$ are not directly estimated from observations $o$ but emerge from a latent trace $z$ generated iteratively and adaptively. The learning objective maximizes a variational lower bound (ELBO):

$$\log p(a|o) \geq \mathbb{E}_{z \sim q_\theta(z|o,a)}[\log p_\phi(a|o,z)] - D_{KL}(q_\theta(z|o,a) \| p_\theta(z|o))$$

where both $q_\theta$ (posterior) and $p_\theta$ (prior) are autoregressive distributions over variable-length sequences of tokens, terminated by EOS. Compression is enforced by scheduling decoding variance $\sigma(T(z)) = \gamma^{T(z)}\sigma_0$, introducing a length penalty that biases the model toward parsimony: additional latent steps are only generated when the predictive precision justifies their cost.

Optimization is performed via RL-like policy-gradient methods applied to latent token trajectories, using a trust-region PPO surrogate loss to stabilize training and mitigate posterior collapse.

Figure 2: LMP training alternates between latent rollout (posterior-sampled latent trajectories) and clipped surrogate model updates to optimize the variational bound.

Model Architecture

The latent encoder is a causal transformer cross-attending to both observations and actions, while the action decoder is a bidirectional transformer cross-attending to observations and latents. During training, latent sequences are padded/truncated to a fixed maximum length for computational tractability. For inference, the policy adaptively decides the length of latent reasoning—directly impacting computational cost and action expressivity.

Figure 3: LMP architecture utilizes cross-attention for both latent encoding and action decoding, tightly coupling latent reasoning to observation dynamics.

Action Tokenization and Sequential Compression

By removing observation conditioning, LMP instantiates "tok", a variable-length sequential action tokenizer compressing high-dimensional continuous action spaces into discrete latent codes. Unlike prior tokenizers (e.g., FAST, VQ-VAE), tok allows any partial latent sequence to decode into valid actions, facilitating flexible downstream autoregressive policy training and improving generalization across task domains.

Experimental Evaluation

Benchmarks and Tasks

LMP was evaluated on multiple real (DROID) and simulated (LIBERO, D3IL, RoboMimic) manipulator domains with diverse language-conditioned multitask and single-task scenarios, including high-precision insertion ("peg-hole"), compositional cleaning ("clean-table"), and zero-shot transfer (block-bowl, marker-mug).

Figure 4: Visualization of diverse simulated and real-world manipulation domains used for evaluation.

Empirical Results

LMP consistently outperformed state-of-the-art diffusion-based policies (DP), particularly in generalization and compositional tasks. On DROID, LMP achieved 95%+ success rate in compositional cleaning and significant gains in zero-shot transfer and precision tasks. On LIBERO-90, LMP demonstrated higher minimum task success rates across tail tasks, evidencing reduced cross-task interference.

Figure 5: LMP improves per-task success rates and tail performance across LIBERO-90 tasks compared to diffusion policies.

Tok outperformed VQ-VAE, FAST, and OAT in action tokenization, notably achieving non-trivial results on high-precision tasks (RoboMimic tool-hang), which prior tokenizers failed to solve.

Analysis of Adaptive Computation

LMP policies dynamically modulate latent trace length in response to action uncertainty and control demands. Fewer latent steps are allocated during gripper operations with irreducible timing variance, while aligning and trajectory phases receive more computation. Quantitative analysis showed a negative correlation between latent steps and action uncertainty estimated via KNN variances.

Figure 6: LMP exhibits interpretable adaptive allocation—gripper events elicit fewer latent steps, while manipulation phases lengthen reasoning.

Compression strength (variance schedule) directly impacts policy performance: excessive compression increases reconstruction error, while insufficient compression causes latents to be ignored. Empirically, a balanced schedule yields optimal utilization and generalization.

Figure 7: Analysis of compression—balance between reconstruction error and latent utilization is achieved via schedule tuning.

Qualitative rollouts with latent step truncation showed that LMP retains mode-seeking behavior and does not degenerate into mode averaging—deterministic policies with low latent steps still decode valid actions, possibly attributed to RL optimization in latent space.

Figure 8: Action token allocation in downstream policies mirrors LMP reasoning behavior when tok is trained with a length penalty.

Implications and Future Directions

Theoretically, LMP repositions latent-space variational inference as a framework capable of supporting both iterative and adaptive computation in control contexts—emulating flexible cognitive reasoning in embodied agents. Practically, LMP provides a scalable, interpretable mechanism for compute allocation in robotic systems, with implications for energy-efficient inference, transfer across tasks, and deployment in real-world environments.

Potential extensions include:

• Transitioning to continuous latent chains (diffusion-like processes)
• Episodic persistence of latents for inter-step memory
• Meta-RL fine-tuning for further optimization of test-time compute [snell2024scalingllmtesttimecompute]
• Multimodal and hierarchical reasoning, coupling vision-language-action with latent space adaptation

Limitations

Training stability remains a concern, with LMP sensitive to buffer size, regularization, and hyperparameters. Posterior collapse and latent utilization are nontrivial to balance, suggesting a need for future research into more robust, scalable variational inference and RL integration for latent-space models.

Conclusion

LMP advances adaptive, interpretable reasoning in control policies by leveraging autoregressive variational inference in latent spaces. Empirical results show gains in generalization and compositional behavior, and its tokenizer variant tok sets a new standard for discrete action representation. The framework opens up practical and theoretical avenues for sophisticated, scalable reasoning in robotic sequential decision-making.