MemCtrl: Using MLLMs as Active Memory Controllers on Embodied Agents

TL;DR

MemCtrl tackles memory bottlenecks in embodied AI by equipping small Multimodal LLMs with a trainable memory head $\mu$ that actively gates what observations are stored. It supports two training modes—offline supervised learning from an expert and online RL—allowing $\mu$ to selectively prune memory, thereby improving decision-making under tight memory constraints. Evaluations on EmbodiedBench across two backbones show an average performance gain of about $16\%$, with larger improvements on long-horizon and complex instructions and fewer invalid actions due to better memory management. The approach is modular, transferable across backbones, and reduces memory footprint, making it practical for real-time embodied agents operating on edge devices.

Abstract

Foundation models rely on in-context learning for personalized decision making. The limited size of this context window necessitates memory compression and retrieval systems like RAG. These systems however often treat memory as large offline storage spaces, which is unfavorable for embodied agents that are expected to operate under strict memory and compute constraints, online. In this work, we propose MemCtrl, a novel framework that uses Multimodal Large Language Models (MLLMs) for pruning memory online. MemCtrl augments MLLMs with a trainable memory head μthat acts as a gate to determine which observations or reflections to retain, update, or discard during exploration. We evaluate with training two types of μ, 1) via an offline expert, and 2) via online RL, and observe significant improvement in overall embodied task completion ability on μ-augmented MLLMs. In particular, on augmenting two low performing MLLMs with MemCtrl on multiple subsets of the EmbodiedBench benchmark, we observe that μ-augmented MLLMs show an improvement of around 16% on average, with over 20% on specific instruction subsets. Finally, we present a qualitative analysis on the memory fragments collected by μ, noting the superior performance of μaugmented MLLMs on long and complex instruction types.

Figures (6)

• Figure 1: Overview: We present MemCtrl, a novel memory filterting scheme to improve decision making performance on small MLLMs tackling embodied tasks. Our approach proposes a trainable memory head (green box labeled "Memory") that learns to actively filter out redundant observations on-the-go. This form of active filtering alleviates issues with inefficient retrieval from stored observations, while also enabling scalability as a detachable memory head.
• Figure 2: Comparison with Prior Work: We present MemCtrl, a novel approach to train "memory heads" to filter observations on the go. Prior work either used the entirety of stored observations as context (left) or filtered them via a variety of Retrieval Augmented Generation (RAG) based schemes ( red arrows), both of which assume the parsing of large amounts of data offline. MemCtrl introduces transferrable heads to use on MLLM backbone ( green arrows) to actively filter observations.
• Figure 3: MemCtrl: We experiment with $3$ augmentations. The simple case acts as a non-trained baseline, where the MLLM is directly queried about storage. In the offline supervised case, $\mu$ is first pretrained using expert answers from a high performing, expert MLLM (GPT-4o here). This trained binary classifier then acts as a head on top of the MLLM backbone. In the Online RL case, we train the memory head online as a policy network. We use a sparse reward on task success and a dense reward on action success. Note that MemCtrl is trained as a detachable head that takes the visuolingual MLLM embeddings as input.
• Figure 4: Long Horizon performance on EB-Habitat: We notice that on long horizon tasks, the expert tends to end the task early by hastily assuming that it is done (finishing after placing one plate instead of all plates). Memory heads highlight unique performance improvements, with $\mu_{\text{RL}}$ exhibiting a more exploratory nature by continuing to place new objects at the right counter, and $\mu_{\text{Exp.}}$ being more exploitative by repeating the same activity over and over, with a single plate. Note: Grayed out images indicate discarded memories.
• Figure 5: Base Performance on EB-Habitat: Here, we compare the performance of GPT-4o vs Qwen2.5-VL-7B-Ins, with various memory augmentations. While GPT-4V gives superior zero-shot performance, it is a very large model that is not easily finetunable. In this instance it also takes more steps to complete the task. On the other hand, $\mu$ boosts the performance of a significantly weaker model, and in this scenario, even doing it quicker with $\mu_{RL}$. The expert head fails here however, similar to the complete memory, causing the episode to end early. Note: Grayed out images indicate discarded memories.