Latent Memory Palace: Reasoning for Control as Autoregressive Variational Inference

We cast reasoning for control as inference in a latent-variable model, where actions \(a\) are generated from observations \(o\) and latent traces \(z = (z_1, z_2, \ldots, z_{T(z)})\) drawn from a vocabulary \(\mathcal{V} \cup \{\texttt{eos}\}\). We maximize the evidence lower bound of the maximum likelihood objective:

\[ \log p(a \mid o) \;\ge\; \mathbb{E}_{z_{1:T(z)} \sim q_\theta(\cdot \mid o, a)} \big[\, \log p_\phi(a \mid o, z_{1:T(z)}) \,\big] - D_{\mathrm{KL}}\!\big(\, q_\theta(z_{1:T(z)} \mid o, a) \,\big\|\, p_\theta(z_{1:T(z)} \mid o) \,\big). \]

The learned prior \(p_\theta(z \mid o)\) and the action decoder \(p_\phi(a \mid o, z)\) together fully parameterize a policy: actions are generated by first sampling a reasoning trace \(z \sim p_\theta(z \mid o)\), then sampling \(a \sim p_\phi(a \mid o, z)\).

The objective above doesn't inherently incentivize adaptive computation. We introduce compression by decaying the Gaussian action decoder's variance with latent length. This imposes a length penalty and encourages the model to allocate latent steps parsimoniously.

We optimize the ELBO using latent-space reinforcement learning. Training alternates between two phases: (1) the posterior generates a buffer of latent trajectories; (2) the model optimizes a clipped surrogate objective (PPO) of the evidence lower bound, minimizing KL divergence to the posterior while maximizing action likelihood.

LMP uses a cross-attention encoder-decoder architecture. The encoder attends to observations and actions to autoregressively generate a variable-length sequence of latent tokens. The decoder attends to observations and latent tokens to predict action chunks. A EOS token terminates the reasoning chain, allowing the model to adaptively allocate its computation.

We train language-conditioned LMP-π and DP on the DROID dataset. Deployed zero-shot, LMP-π outperforms DP on block-bowl and marker-mug. We then finetune on two more tasks: peg-hole and clean-table. LMP-π demonstrates both precision and compositional generalization, notably picking up at least two objects with 95% success in clean-table.

On language-conditioned multitask LIBERO-90, LMP-π performs marginally better than DP across the board, but a closer look reveals significantly higher success on the bottom-10 tasks. This indicates that LMP-π organizes information with less cross-task interference. The per-task breakdown below shows LMP-π tapering off at a slower rate than DP.

We can instantiate LMP as a variable-length discrete action tokenizer, LMP-tok, by dropping the observation conditioning. When controlling for downstream autoregressive policies, LMP-tok significantly outperforms all baselines and is the only method achieving nontrivial performance on the RoboMimic tool-hang task.

We visualize the number of reasoning steps generated during rollouts. The model uses fewer reasoning steps during gripper movements, e.g., grasping and releasing in DROID tasks or hand-over in RoboMimic transport.

We manually slide a block through the open gripper while querying the policy for the number of reasoning steps at each state. The reasoning length shifts with the block's position relative to the gripper.

We compare LMP-π to two non-iterative latent-variable baselines: Gaussian VAE Policy and VQ-BeT. LMP-π consistently outperforms both on D3IL and RoboMimic, indicating that fixed-length latent representations lack the granularity needed for complex manipulation tasks.

We ablate compression strength by controlling the decoder variance schedule. A smaller σ_min penalizes prediction errors more heavily, yielding shorter latent traces. Across LIBERO tasks, increasing compression (smaller σ_min) improves performance, showing that parsimonious allocation of latent tokens is crucial for generalization.

On DROID checkpoints, we find that too little compression (0.1-0.08) causes the decoder to ignore the latents; too much (0.1-0.01) forces overly short traces and raises reconstruction error; we use a balanced schedule (0.2-0.02) in our experiments.