Demystifying Diffusion Policies: Action Memorization and Simple Lookup Table Alternatives

2025-05-09

AI Summary

This paper investigates why diffusion-based policies achieve strong performance in few-shot robotic manipulation, revealing that their efficacy stems from implicitly constructing an Action Lookup Table (ALT), relying on action memorization rather than generalization. Method: We propose an explicit, interpretable, lightweight ALT policy with out-of-distribution (OOD) detection capability. It employs a contrastively trained image encoder to learn joint action-observation embeddings; actions are retrieved via nearest-neighbor search over training trajectories; and OOD detection is performed online using a learned embedding distance threshold. Contribution/Results: Our method matches the performance of diffusion policies on small-scale datasets while accelerating inference by 294× and reducing memory footprint to just 1.15% of theirs. It enables real-time, closed-loop control on resource-constrained platforms—offering transparency, efficiency, and robustness unattainable with implicit generative policies.

Technology Category

• Intelligent Robots: Manipulation
• Search and Optimization: Learning to Search
• Computer Vision: Diffusion Models for Vision

Application Category

• Responsible Web: Machine-in-the-loop, human agency and autonomy
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

Abstract

Diffusion policies have demonstrated remarkable dexterity and robustness in intricate, high-dimensional robot manipulation tasks, while training from a small number of demonstrations. However, the reason for this performance remains a mystery. In this paper, we offer a surprising hypothesis: diffusion policies essentially memorize an action lookup table -- and this is beneficial. We posit that, at runtime, diffusion policies find the closest training image to the test image in a latent space, and recall the associated training action sequence, offering reactivity without the need for action generalization. This is effective in the sparse data regime, where there is not enough data density for the model to learn action generalization. We support this claim with systematic empirical evidence. Even when conditioned on wildly out of distribution (OOD) images of cats and dogs, the Diffusion Policy still outputs an action sequence from the training data. With this insight, we propose a simple policy, the Action Lookup Table (ALT), as a lightweight alternative to the Diffusion Policy. Our ALT policy uses a contrastive image encoder as a hash function to index the closest corresponding training action sequence, explicitly performing the computation that the Diffusion Policy implicitly learns. We show empirically that for relatively small datasets, ALT matches the performance of a diffusion model, while requiring only 0.0034 of the inference time and 0.0085 of the memory footprint, allowing for much faster closed-loop inference with resource constrained robots. We also train our ALT policy to give an explicit OOD flag when the distance between the runtime image is too far in the latent space from the training images, giving a simple but effective runtime monitor. More information can be found at: https://stanfordmsl.github.io/alt/.

Problem

• Understanding why diffusion policies excel in robot manipulation tasks
• Investigating if diffusion policies memorize action lookup tables
• Proposing a lightweight alternative to diffusion policies for efficiency

Innovation

• Diffusion policies memorize action lookup tables
• Action Lookup Table (ALT) replaces diffusion policies
• ALT uses contrastive image encoder for indexing