Latest / Embodied AI 101 / LaMem-VLA: Latent-Memory-Native VLA Framework

Transcript

LaMem-VLA: Latent-Memory-Native VLA Framework · Embodied AI 101

• LaMem-VLA: Latent-Memory-Native VLA Framework. Vision-language-action (VLA) models – large multimodal policies that map camera inputs and language instructions to robot actions – have recently achieved impressive results in robotic manipulation. Most existing VLA models, however, process each decision in a “Markovian” way, predicting the next action primarily from the current camera view and command. This simplification works well for short subtasks, but it overlooks the historical context of a multi-step job. In practice, a robot may need to remember what it did moments or minutes ago – for example, which drawers were already opened, or where it last saw a tool – in order to decide what to do next. Without this temporal awareness, VLA policies struggle when tasks require remembering key events or progress over long horizons. To address this, Qu et al. propose LaMem-VLA, a “latent-memory-native” framework that weaves explicit memory of past experience directly into the reasoning of a VLA model. Unlike prior memory-augmented VLAs that merely retrieve past frames or use an external memory bank as an extra context for the policy, LaMem-VLA treats history as part of the same continuous token space where the model already processes vision and language. In concrete terms, LaMem-VLA encodes important snippets of history into latent memory tokens, and then prepends these tokens to the current observation and instruction before passing everything through the vision-language transformer. This way, the model’s self-attention process naturally integrates historical cues along with the new input, letting memory participate in the decision process on equal footing. The result is a VLA policy that not only sees the present, but also remembers what just happened – all within the same latent space. In simulated benchmarks, this memory-augmented approach notably outperforms prior VLA models on long-horizon manipulation tasks. In this episode we’ll unpack the key ideas of LaMem-VLA. We’ll start by reviewing why memory matters for robotics (and how earlier VLAs tried to add it), then dive into LaMem-VLA’s architecture and its novel memory mechanisms. We’ll go through its four main components – memory Curator, Seeker, Condenser, and Weaver – and see how they build and use separate short-term and long-term memory vaults. Finally, we’ll survey the experimental results on two benchmarks (SimplerEnv-Bridge and LIBERO) and discuss what these findings mean for the future of embodied AI. Motivation: Why Memory Matters in VLA Models. Imagine a robot learning to cook. Its chef mentor tells it to “finish setting the table and then pour water.” The VLA model controlling the robot sees an image of the table, hears the language command, and plans actions. But which objects are already placed, and which still need handling? If the robot only reacts to the current visible scene (as a strict Markov process would), it might forget even moments-old context: perhaps it has already put plates out and wiped the table, or maybe it opened a water jug a step ago. These details, which can be crucial to decide the next correct motion, require remembering past observations and actions. The robotics literature recognizes that manipulation tasks are often non-Markovian – the current state is not sufficient to decide the next move; one must remember the recent sequence of events. In cognitive terms, one can think of a hierarchy: a robot needs a short-term memory (like human working memory) to recall what just happened in the last few seconds, and a long-term memory or episodic cache to preserve the semantic context of the entire task (e.g. how many subtasks completed). Recent VLA models, however, largely ignore this. They simply encode the present image and instruction into tokens and roll out an action sequence via a transformer or diffusion policy. This works if tasks are simple and short, but for long, multi-stage tasks it inherently biases the policy to focus only on immediate observations. For example, if the robot must pick up a carrot and then put it in a bowl, once it drops the carrot it may no longer see it – but it still needs to recall that the carrot is the thing to pick up. Some recent works have attempted to give VLAs a memory. One approach (pioneered in the “MemoryVLA” framework ) is to maintain a working memory and a memory bank that stores past perceptual and semantic tokens. When making a decision, the model retrieves relevant entries from this bank and conditions the policy on them. Another group of methods simply expands the context window of the model: they append several past frames of video or recent steps to the input, so the policy is aware of short-term context. These methods do improve performance, but they have drawbacks. Temporal context expansion quickly becomes computationally expensive and still can’t capture truly long-horizon information beyond the fixed window. Explicit memory banks add content, but they typically connect to the policy as an external bag of tokens. That means the retrieved memory is fused in only at the final stage of action decoding, rather than being an integral part of the VLA’s reasoning process. In other words, past experience is treated as a separate adjunct, not as genuine context in the model’s own latent space. LaMem-VLA asks: What if we truly made memory tokens native to the model’s latent space? Since modern vision-language backbones already represent everything (visual frames, instructions, and action queries) as continuous token embeddings, we can conceivably also represent memory as a stream of tokens. If we keep those memory tokens in the same embedding space, we can let the transformer attention mix memory with current observations. This means the model’s own inference process can attend to and incorporate history as if it were just more of the input, rather than having to treat it as a separate feature. Crucially, this design can encapsulate unlimited history in a compact form (by condensing it into tokens) without blowing up the context length – enabling long-term memory at roughly constant computational cost. Qu et al. articulate exactly this vision. They posit two complementary forms of memory: a short-term visual memory storing recent perceptual details, and a long-term semantic memory encoding the progress and high-level task context. Neither memory mode is just the raw pixels; both are ultimately transformed into latent vectors. To do this coherently, they introduce a unified framework where:. A Curator observes the robot’s ongoing experience and writes it into two memory vaults (short-term and long-term). A Seeker retrieves relevant items from these vaults based on the current multimodal context. A Condenser takes the retrieved evidence and distills it into fixed-size latent tokens. A Weaver then weaves those latent memory tokens into the model’s input sequence along with the current image and instruction. In short, history is stored, retrieved, and consumed entirely in the model’s own embedding space, letting the action generation process be guided by memory in a seamless way. The results, as we will see, show clear benefits on long-horizon tasks. The LaMem-VLA Architecture. At a high level, LaMem-VLA is still a diffusion-based VLA policy. It builds on a powerful pretrained vision-language backbone and a diffusion-based action generator. But it wraps around this backbone a memory subsystem, made of the four modules mentioned above. Let’s break down each part. Latent Memory Vaults: Short-Term and Long-Term. The Curator module is responsible for building and updating the robot’s memory as it operations unfold. It sees each timestep’s data (the current camera image and the resulting action representations) and decides what to stash into memory. Memory is organized into two vaults:. Short-Term Memory Vault (Mshort): This holds visual perceptual evidence from the current episode. Intuitively, it captures what was in view (object positions, transient changes) during the recent seconds. To construct it, the Curator uses a learnable compression module on the current visual tokens Xt (the image observations encoded by the vision tower). This compressor (a small transformer or SE-bottleneck) outputs a set of compact visual tokens, then mean-pools them to form a retrieval key ks. The corresponding value vs is the un-pooled set of short-term tokens; together they form a key-value pair ms = (ks, vs). Each time step’s ms gets appended to Mshort. Crucially, the Curator also keeps this vault a bounded size (capacity L, e.g. 16 entries). When new entries arrive beyond capacity, it merges redundant ones. It compares the keys of adjacent short-term entries (via cosine similarity); the most similar pair is averaged together (both key and value) into one, eliminating duplication. This simple consolidation keeps the memory focused on distinct visual events. Long-Term Memory Vault (Mlong): This stores semantic and action-progression cues across the task. Instead of compressing the visual input, the Curator here writes the raw action-query hidden states Haction (from the diffusion policy’s transformer) into the vault. Each Haction essentially encodes what action was just taken at that step, so appending them builds up a timeline of what the robot did. There is no separate compression or key-value for long-term entries; each entry is simply one time step’s action embedding. Over time this vault accumulates a trajectory of how the task has progressed. Like the short-term vault, it also prunes itself if needed: once it exceeds capacity, it merges the two most redundant adjacent entries by averaging their embeddings. In sum, Mshort tracks what the robot saw just recently, while Mlong tracks what the robot has done or progressed over the longer run. Both vaults are living stores, updated step by step, as shown in Figure 2 of the paper (not reproduced here). Importantly, both vaults store latent vectors, not raw images: the short-term vault’s values vs are already in the vision-language embedding space, and the long-term vault’s entries are the same (they come from the VLM backbone’s action-query outputs). This setup ensures that everything in memory can eventually be treated as contextual tokens in the same space. Memory Seeker and Condenser: Retrieval and Tokenization. When it’s time to generate the next action, the model needs to pull out relevant information from memory. This is the job of the Seeker. The Seeker first computes a context query from the current state: it takes the latest visual tokens Xt (encoding the current camera view) and the instruction tokens I, and feeds them through the VLM to get hidden states. On top of those, it appends a small block of learnable query slots (initialized parameters) and runs a tiny transformer (called the query builder) so that these new query slots attend to the current context. After a couple of transformer layers, the appended slots become a refined context vector Qt. By mean-pooling Qt, the method obtains a single vector qt that represents the summary of “what the robot is focusing on now”. This query qt is used to sift through both memory vaults:. From the short-term vault, each entry’s key ks is compared to qt via cosine similarity. The top K most similar keys are selected. Suppose the indices are Is, then the Seeker grabs the corresponding short-term values vs and concatenates them into a set Zshort of size up to K×(number of tokens per value). These vectors are the recent visual evidence deemed relevant to the current context. From the long-term vault, the process is similar. Each stored action embedding ml (mean-pooled) is compared to qt (again by cosine), and the top K entries are taken. These K long-term vectors (each same dimension) are concatenated into Zlong. They represent high-level cues about what part of the task was most similar to now. Crucially, the Seeker retrieves sets Zshort and Zlong without inserting them directly into the policy. Instead, these raw sets may be large and partially redundant, so the next module Condenser compresses them into fixed-size tokens. The Memory Condenser has the job of synthesizing a small bundle of latent tokens from the retrieved evidence. It uses two tiny transformer “memory formers”, one for each type of evidence. For short-term memory, it prepares a fixed, learnable array of Ls token vectors Ts (e.g. 8 slots), and similarly Ll tokens Tl for long-term (e.g. 4 slots). Then it concatenates [Qt; Zshort; Ts] and processes them through transformer layers to produce updated slots Mshort (size Ls × C), and [Qt; Zlong; Tl] to produce Mlong (size Ll × C). These Mshort and Mlong are the final condensed memory tokens: they each represent a summary of the retrieved evidence, tuned to the current query context. Importantly, their lengths Ls and Ll are small and fixed, so this step caps the amount of memory entering the policy. Intuitively, the Seeker + Condenser pipeline works as follows: We ask memory, “what in the recent past seems relevant now?” (via retrieval). Then “let’s distill that set of evidences into a compact clue for the policy” (via the transformer slots). The result is two small latent token groups – one for visual detail and one for long-term semantics. Memory Weaver: Integrating Memory into VLA Reasoning. Finally, the Weaver takes these latent tokens and stitches them into the model’s input so that the policy attends to them. At timestep t, the model already planned to create an input sequence whose elements are: short-term memory tokens, long-term memory tokens, the current visual tokens Xt, the instruction tokens I, and some number of learnable action query tokens Qaction. The Weaver now inserts the condensed memory: it prefixes Mshort and Mlong to the front of the sequence. Technically, it also feeds them with learnable “source” embeddings (bs, bl) to let the model know they are memory tokens, but otherwise they are just treated as normal input tokens. So the final input sequence St looks like:. where all tokens are vectors in the same C-dimensional embedding space. The notation Mshort + 1Ls bs means that each of the Ls short-memory tokens gets added (broadcast) the source embedding bs. Similarly for long memory. Then the whole sequence is fed into the frozen vision-language decoder (a large transformer). The crucial point is that during the VLM’s self-attention, the memory tokens attend to and from the current image and instruction tokens, just as if they were part of the scene context. This yields updated action query embeddings Zaction at the end of the transformer, which now are grounded in memory as well as the present. Put another way, the VLA backbone processes memory tokens and current tokens together: the latent memory is “woven” into the fabric of the scene understanding. After this, a diffusion action decoder takes the memory-informed action queries Zaction and generates the actual continuous control actions. In sum, the memory tokens have shaped the policy’s reasoning implicitly, without any brute-force concatenation of raw history beyond a fixed-size latent glimpse. Example Flow of Information. Let’s trace one decision step to make this concrete. At time t, the robot has current observation ot and instruction I. The backbone encodes ot into visual tokens Xt and I into text tokens. Past steps have been feeding entries into Mshort and Mlong. Now the Seeker forms a query qt from Xt and I, retrieves some top 8 entries from each vault, yielding Zshort and Zlong. The Condenser maps these into 8 short-term memory slots Mshort and 4 long-term slots Mlong. The Weaver then assembles the input [Mshort; Mlong; Xt; I; Qaction], the model processes everything, and outputs refined action tokens Zaction. A diffusion policy then decodes Zaction into the next chunk of physical joint actions for the robot. Implementation Details. In implementation, LaMem-VLA uses a 7B-parameter Prismatic vision-language model as its backbone. (Prismatic is a recent large VLM pretrained on broad real-robot data.) The observation ot is a single third-person RGB image, resized to 224×224, which yields visual tokens via the frozen vision tower. Those tokens are passed through LLaMA-7B to get context-aware embeddings. The action queries Qaction are additional learnable tokens appended to the sequence, whose final hidden states serve as the base for action decoding. For action generation, LaMem-VLA uses a diffusion-based transformer policy with ~300M parameters, trained to denoise next-chunk motions. Key hyperparameters include vault capacity L=16 for each vault, and retrieval of K=8 units from each vault. The short-term condenser outputs Ls=8 tokens, and the long-term condenser outputs Ll=4 tokens. The query builder and memory formers are shallow two-layer transformers. During training, the whole model is finetuned end-to-end on robot episode data. Once trained, inference is quite efficient: the static backbone is run once per step, memory operations add a small overhead for retrieval and transformer layers, and then the diffusion sampler generates actions in ~10 denoising steps. (The authors note that VLA systems like this often orchestrate planning at a slower cadence than the low-level controller, so throughput is acceptable.). Experimental Evaluation. Qu et al. evaluate LaMem-VLA in simulation on two benchmark suites: SimplerEnv-Bridge and LIBERO. Both are common benchmarks for embodied instruction-following. Importantly, LaMem-VLA is compared against strong baselines and various state-of-the-art VLA models, to see whether its latent memory really pays off. SimplerEnv-Bridge. SimplerEnv-Bridge is a set of kitchen-table tasks designed to test real-to-sim generalization. The tasks include subtasks like “Put Spoon on Towel,” “Put Carrot on Plate,” “Stack Cube,” and “Put Eggplant in Basket” (Table 1 in the paper). Models are trained on some dataset of demonstrations and then run in simulation with a WidowX robot. Success is measured by accurate placement of objects. Table 1 of the paper reports the success rates of many methods on these four tasks. Some recent baselines:. CogACT (an earlier 2024 baseline) got only around 57% average success (mostly failing on Stack Cube). π0 (a 2025 VLA) had ~69% average. SemanticVLA (CVPR’26) got 65%. MemoryVLA (the ICLR’26 memory-enhanced model) scored about 71.9% average (75% on Spoon and Carrot tasks, 37.5% on Stack). In contrast, LaMem-VLA hits 73.9% average success. Notably, it beats CogACT (57.3%) by ~16.6 points and outperforms π0 (69.2%) and SemanticVLA (65.1%) substantially. Its per-task scores were 83.3% on Spoon, 75.0% on Carrot, 41.7% on Stack Cube, and 95.8% on Eggplant. The gains are clearest on the hardest task “Stack Cube,” where all models struggle, but LaMem-VLA achieves 41.7% (versus, say, 29.2% for SpatialVLA and 25% for CogACT). These results indicate that injecting memory in LaMem-VLA gives a robust boost. The tasks where an object might be occluded or where successive steps must be coordinated benefit from memory of object locations. For example, if the spoon was moved earlier, a memory token can remind the policy of its last seen spot. The authors emphasize that such historical context “provides effective memory for action prediction, improving manipulation performance across diverse tasks”. In other words, LaMem-VLA is better able to resolve ambiguities that a purely Markovian policy would miss. LIBERO. The LIBERO benchmark is a more comprehensive suite of household tasks. It has five “suites”: Spatial, Object, Goal, and two long-horizon sets (Long-10 and Long-90). In LIBERO, a Franka robot follows open vocabulary instructions across many environment variations. The first three suites each have 10 tasks of moderate length, while Long-10 has 10 longer tasks and Long-90 has 90 very long tasks. In the LIBERO experiments, models are trained on 50 demonstrations per task (the OpenVLA protocol). The evaluation is reported as success percentages in each suite. Table 2 of the paper (partially shown in [12]) compares many methods. Leaders before LaMem-VLA included: CogACT (with proprio/state inputs) at ~93.2% avg, π0 (with extras) at ~94.2%, and MemoryVLA (ICLR’26) at 96.5%. SemanticVLA (CVPR’26) was already at 97.0%. LaMem-VLA pushes the envelope further: it reaches an overall 97.6% success, the best among methods without extra sensors. Breaking it down, LaMem-VLA scored 98.8% on Spatial, 99.0% on Object, 97.2% on Goal, 95.8% on Long-10, and 97.0% on Long-90. It even exceeds the starred π0 and SemanticVLA on Spatial and Object (98-99%). On the long-horizon tasks, LaMem-VLA’s 97.0% average easily beats the prior best 95.6% of MemoryVLA. Put succinctly, LaMem-VLA yields the highest scores on LIBERO: “It achieves an average of 97.6% across five suites, improving 1.1 points over the best memory-augmented baseline and 4.4 points over a strong VLA policy.” Even without wrist cams or proprioceptive inputs (that some starred models used), LaMem-VLA wins overall. The gain is especially notable on Long-10 and Long-90: the model’s success is a few points higher than others in those hard suites. The authors conclude that weaving memory into the model helps it keep track of long-term progress and resolve temporal ambiguities. For instance, if an object was picked up several steps ago, a long-term memory token might remind the policy of that fact even if the object is now out of view. Ablation Studies. The paper also contains ablation experiments to isolate what makes LaMem-VLA work. Several key findings stood out:. Short-term vs. Long-term memory importance: Table 3 in the paper shows the effect of removing each memory component. Without any memory tokens (essentially turning memory off), success drops drastically (to ~57% on SimplerEnv and 92% on LIBERO-90). Including only short-term or only long-term memory improves performance somewhat (mid-60s Simpler, ~95% LIBERO), but neither is as good as having both. The full dual-memory model hits 73.9% Simpler and 97.0% LIBERO. This confirms the complementary roles of the two vaults: short-term captures visual details (e.g. current object positions), long-term carries global task context. Removing either stream weakens the reasoning, and using both yields the best results. Memory integration vs. policy-side: In another variant (Table 4), the authors compare three ways of using memory with a fixed backbone: (a) Baseline with no memory, (b) policy-side memory (like MemoryVLA, where retrieved tokens are concatenated after the backbone), (c) LaMem-VLA style latent memory. They find the memory-free baseline gets 57.3%/92.1%. Adding memory externally (policy-side) bumps that to 71.9%/94.8%, proving memory helps even if not woven in. Feeding raw retrieved tokens also helps (69.8%/95.1%). Critically, LaMem-VLA’s latent-native integration lifts it to 73.9%/97.0%. The authors interpret this as showing that embedding memory in the same space yields extra benefit beyond just providing context externally. The memory tokens allow the internal reasoning layers to jointly attend to memory, vision, and language, and this yields a further edge. Varying retrieval budget: Table 5 and Figure 3 (in the paper) vary the number K of memory units retrieved and the number of latent tokens Ls, Ll. They find that too few retrieved units (K=2) is insufficient, while K=8 performs best. Pulling too many (K=12) slightly hurts, likely due to more noise/redundancy. Similarly, increasing the number of latent slots helps up to a point (8 short, 4 long was the chosen setting). These experiments suggest the design choices (retrieve 8, condense to 8+4 slots) were sensible trade-offs. Overall, these ablations make clear that both memory vaults matter, and that the latent-memory-native approach is superior to older memory schemes. The authors report that the performance gain indeed comes from this integrated memory, not just the raw fact of more data. Significance and Discussion. LaMem-VLA makes a compelling case that enabling VLA models to explicitly remember history – in a way that is native to their transformer reasoning – can materially improve long-horizon robot tasks. The key novelty is the end-to-end design of memory: historical information is not a hacky add-on but becomes part of the model’s fabric. By encoding both recent percepts and long-term context as latent tokens, the model can flexibly draw on both detailed observations and broader task knowledge. In the broader context of embodied AI, this work joins a growing trend of neural memory systems for sequential decision-making. People often draw inspiration from the brain’s memory systems: indeed, the authors explicitly compare their short-term vs long-term vaults to working memory vs hippocampal memory in humans. LaMem-VLA’s memory curator-seeker-condense-weave pipeline is a concrete instantiation of a dual-memory system in the robotic VLA setting. It combines ideas from retrieval-based memory (storing experience, retrieving relevant parts) with ideas from transformer context expansion (creating learned tokens that auto-regress). Compared to previous VLA memory efforts, LaMem-VLA stands out by where the memory lives. For instance, MemoryVLA (Shi et al., ICLR’26) also splits memory into low-level vs high-level, but it still uses a separate passive memory bank. In MemoryVLA, the policy eventually ingests memory as an extra context vector in the diffusion step. Suited for its tasks, MemoryVLA achieved 71.9% on SimplerEnv-Bridge. LaMem-VLA pushes that further, partly by letting the attention layers themselves see the memory tokens as peers to the observation tokens, rather than as an afterthought. Similarly, other recent works like EventVLA (Ganlin et al., NeurIPS’26) focus on when to write events into memory, but they also feed memory externally. By contrast, LaMem-VLA uses a learnable condensation of events (through its memory formers) to fit memory directly into the input of the VLM. A pragmatic benefit of the LaMem-VLA approach is that it can retrofit to any VLA model that uses a transformer-like backbone. The memory tokens have a fixed small size, so the overall context length overhead is bounded regardless of task duration. This contrasts with simply concatenating entire histories or frames, which would blow up the sequence length. That said, there are limitations. The current results are all in simulation. The authors note they plan to test this on real robots in future work – real-world validation is still needed. Also, the memory system does add computational cost (extra transformer layers, retrieval computation, etc.), which could matter in time-sensitive applications. Finally, the system still relies on good retrieval: if qt fails to find the right clue in memory (e.g. in a completely novel situation), the advantages may wane. But the ablation on K shows they tried to tune this and found a sweet spot. In the grander scheme, LaMem-VLA is part of a fascinating wave of papers exploring memory in LLMs and RL agents. In foundation models, ideas like recurrent memory in GPT or external “long-context” storage have been gaining attention. For embodied agents, the stakes are high: a robot that literally forgets half of what it was doing is crippled. By bridging vision-language models with structured memory, this work contributes a valuable building block. One could imagine extending it further: perhaps future versions might store more abstract episodic memories (key frames) or even form hierarchies (stack memories of memories). But already, this work demonstrates a clear principle: let the model’s own latent space be the attic where it stashes past experience. To sum up, LaMem-VLA shows that a thoughtfully designed memory pipeline – curator, seeker, condenser, weaver – can lift a VLA agent’s performance substantially on long and complex tasks. The approach marries the strengths of pretrained multimodal transformers with classical notions of memory vaults. For practitioners, the takeaway is that when building vision-language policies, it may pay to let the model remember: not by piling on more inputs, but by weaving memory tokens into the transformer’s view. As embodied AI aims for truly long-horizon autonomy, mechanisms like LaMem-VLA will likely become an essential part of the toolbox.