Human decision-making ranges from the reflexive (e.g., walking) to the deliberate (e.g., playing chess), with variability in the time and effort devoted to each decision. This process is both iterative and adaptive, proceeding through a chain of logical steps that scales with the complexity of the decision. Modern large language models (LLMs) exhibit a similar pattern: a broad class of “reasoning” LLMs achieves substantially improved performance by generating intermediate tokens before the final answer. We ask: can iterative, adaptive computation benefit sequential decision making problems such as robotics? For robotic policies, we posit that the benefits are twofold: efficient decision making under task variability, and improved generalization. However, directly transferring methods from language models is unlikely to suffice, as language tokens may not capture the nuances required for spatial understanding and precise motions. We therefore propose that robotic policies reason iteratively in a latent space learned end-to-end for action prediction. This preserves the benefits of iterative, adaptive computation while allowing the intermediate representations to capture control-relevant information at the appropriate granularity.
Experiments
We conduct experiments to answer the following research questions: (1) Is LMP-a parameterized form a performant generative control policy? (2) Is LMP-the tokenization variant an effective action tokenizer? (3) Does LMP demonstrate adaptive allocation of test-time compute? (4) Which factors contribute to LMP’s performance?
Discussion
Since LMP is trained with sampling-based RL techniques, it is sensitive to hyperparameters. For example, the latent distribution is prone to collapsing without a sufficiently large rollout buffer or regularization. We leave the exploration of stable and scalable optimization of the proposed framework to future work.
Baseline Details
In this section, we describe the details of baseline methods. All policy baselines use an observation history of 1 and action chunk of 16. We train all baselines for the same number of gradient steps as our method. These are generative control policies which model the expert’s action distribution using a denoising diffusion process. Diffusion Policy is the single-task variant, and Large Behavior Model is the language-conditioned multi-task variant. We parameterize the noise prediction network as a bidirectional transformer with cross-attention observation conditioning. We use 100 diffusion steps during training and 10 diffusion steps with DDIM scheduler for inference.
Environment Details
The platform is a single-arm manipulation setup based on a 7-DoF Franka robot. The image observations consist of two (180, 320) images from a scene camera and a wrist camera. The low-dimensional observations consist of 6-dimensional absolute joint positions and a 1-dimensional absolute gripper position, represented as a continuous value in [0, 1]. The actions consist of 6-dimensional delta joint positions and a 1-dimensional absolute gripper position. The controller runs at 10 Hz. To ensure reliable comparisons, we mount an overhead camera and use a program to track 20 initial configurations for each task. The program alpha-blends the pre-captured snapshot and the live frame so that objects can be reset by aligning them.
We evaluate on two zero-shot tasks roughly within the training distribution of the DROID dataset, and two challenging tasks where we collect 100 demonstrations each and finetune the policies.
block-bowl (zero-shot) involves picking up a cube block and placing it in a bowl. The objects are randomly initialized across the workspace. An episode is considered successful if the gripper releases the block into the bowl. Language instruction: “Pick up the block and put it in the bowl.”
marker-mug (zero-shot) involves picking up a marker and dropping it in a mug. The objects are randomly initialized across the workspace. An episode is considered successful if the gripper releases the marker into the mug. Language instruction: “Pick up the marker and put it in the mug.”
peg-hole (finetuned) is a high-precision manipulation task which involves inserting a square peg into a base with a square hole. The base is randomly initialized in a square region near the center of the workspace, and the peg is randomly initialized outside the square region. An episode is considered successful if the gripper inserts the peg securely into the base. Language instruction: “Insert the peg into the hole.”
clean-table (finetuned) is a compositional manipulation task which involves placing 3 objects scattered around the workspace into a bowl. We collect demonstrations by randomly choosing 3 objects from a 7-object suite. We measure partial success of cleaning up at least 1, 2, and 3 objects. Language instruction: “Clean up the table.”
Distribution of Task Success Rates Across LIBERO-90
While both LMP-the parameterized form and DP achieve near-perfect success in a majority of tasks, LMP-the parameterized form tapers off at a slower rate, marking better generalization across task semantics and less cross-task interference.
Additional Qualitative Analysis of Test-Time Compute
We construct an analytical experiment in d3il-stacking to qualitatively understand the allocation of test-time compute. We keep the gripper in an open position and slide a block through the open gripper, while continuously querying the policy for the average number of latent steps. As the block moves, the number of latent steps reaches the lowest point the moment the block is right underneath the gripper. This qualitatively shows a correlation between reasoning steps and potential gripper movement.
We find that across a number of simulated and real-world tasks, the number of latent steps is low during grasping and releasing, and high during aligning and other movements. We did not observe as much latent step variation in simpler tasks such as Robomimic lift and can, or D3IL aligning and sorting where the end-effector is a peg.
Frequently Asked Questions
How does Latent Memory Palace differ from standard diffusion-based policies? LMP uses iterative latent reasoning in a learned latent space rather than denoising in action space, enabling adaptive test-time compute allocation.
What types of tasks benefit most from LMP’s adaptive computation? Complex tasks requiring alignment or precise movements see high latent step counts during critical phases, while simpler grasping tasks use fewer steps.
Is the LMP framework sensitive to hyperparameter choices? Yes, the latent distribution can collapse without a sufficiently large rollout buffer or regularization, making stable optimization an area for future work.
How is LMP evaluated in terms of generalization? It is tested on zero-shot tasks from the DROID dataset and finetuned on high-precision tasks like peg insertion, showing improved cross-task generalization over diffusion baselines.