VPWEM: Visuomotor Memory Architecture

• The paper introduces VPWEM, a framework that integrates a fixed-size sliding window (working memory) with a compact learned episodic memory for non-Markovian tasks.
• It employs a Transformer-based contextual memory compressor to summarize past observations, ensuring efficient and constant memory use per time step.
• Empirical results demonstrate over 20% improvement in success rates on memory-intensive tasks compared to state-of-the-art baselines.

Searching arXiv for the exact VPWEM paper and nearby usage. VPWEM denotes “Visuomotor Policy with Working and Episodic Memory”, a framework for imitation-learned robotic control designed for non-Markovian tasks in which the current observation does not contain all information needed for optimal action selection. The framework combines a fixed-size sliding window of recent observations as working memory with a compact learned summary of older observations as episodic memory, and instantiates this design on diffusion policies so that action generation conditions on both short-term and episode-wide information while maintaining nearly constant memory and computation per step (Lei et al., 5 Mar 2026).

1. Definition and problem setting

VPWEM is formulated in a POMDP,

$$\mathcal{M}=(\mathcal{S}, \mathcal{A}, \mathcal{P}, \mathcal{R}, \Omega, \mathcal{O}),$$

where the policy should in principle depend on the full history

$$\pi_\theta(a_t \mid h_t), \qquad h_t = o_{\le t} = (o_0,\dots,o_t).$$

In practice, most visuomotor policies use only a truncated recent history,

$\pi_\theta(a_t \mid o_{t-L:t}),$

which is often sufficient for nearly Markovian manipulation tasks but fails when task success depends on hidden past events, latent goals, or temporally distant cues (Lei et al., 5 Mar 2026).

The framework is motivated by two limitations of standard visuomotor and diffusion policies. First, they typically condition on a single frame or a very short history such as 2, 6, or 10 frames. Second, simply enlarging the context window is computationally costly because Transformer attention over length $L$ scales as $O(L^2)$, and longer histories can encourage reliance on nuisance correlations and “copycat” behavior. VPWEM addresses this by retaining explicit recent observations and compressing the rest of the episode into a fixed-size learned memory (Lei et al., 5 Mar 2026).

This design places VPWEM in the class of non-Markovian visuomotor policies rather than standard short-context behavior cloning systems. A plausible implication is that the framework is intended less as a replacement for local reactive control than as a memory augmentation for policies deployed in partially observable manipulation settings.

2. Memory structure: working memory and episodic memory

The central architectural distinction in VPWEM is between working memory and episodic memory. Working memory is a fixed-size sliding window over the most recent observation tokens,

$$w_t \triangleq o_{t-L:t} = \mathrm{concat}(o_{t-L+1},\dots,o_t)\in \mathbb{R}^{L\times D}.$$

This preserves fine-grained local temporal information for immediate control (Lei et al., 5 Mar 2026).

Episodic memory is a fixed number of learned memory tokens summarizing all older observations outside that window. Older observations,

$$o_{\le t-L} = \mathrm{concat}(o_0,o_1,\dots,o_{t-L}),$$

are not stored verbatim. Instead, when an observation token $o_\tau$ leaves the working-memory window, it is augmented with positional information,

$f_\tau = o_\tau + PE(\tau)\in\mathbb{R}^D,$

and inserted into an observation cache,

$$\mathcal{C}_f \leftarrow \mathcal{M}(\mathcal{C}_f \cup f_\tau),$$

where $$\pi_\theta(a_t \mid h_t), \qquad h_t = o_{\le t} = (o_0,\dots,o_t).$$0 is the cache management policy and the cache has maximum size $$\pi_\theta(a_t \mid h_t), \qquad h_t = o_{\le t} = (o_0,\dots,o_t).$$1. The resulting episodic memory output is

$$\pi_\theta(a_t \mid h_t), \qquad h_t = o_{\le t} = (o_0,\dots,o_t).$$2

The policy then conditions action generation on both $$\pi_\theta(a_t \mid h_t), \qquad h_t = o_{\le t} = (o_0,\dots,o_t).$$3 and $$\pi_\theta(a_t \mid h_t), \qquad h_t = o_{\le t} = (o_0,\dots,o_t).$$4 (Lei et al., 5 Mar 2026).

At the observation level, VPWEM processes multi-modal inputs including RGBD observations from $$\pi_\theta(a_t \mid h_t), \qquad h_t = o_{\le t} = (o_0,\dots,o_t).$$5 cameras and low-dimensional proprioceptive states. Camera observations are encoded into visual features $$\pi_\theta(a_t \mid h_t), \qquad h_t = o_{\le t} = (o_0,\dots,o_t).$$6, proprioception into $$\pi_\theta(a_t \mid h_t), \qquad h_t = o_{\le t} = (o_0,\dots,o_t).$$7, and these are concatenated and passed through an MLP to form a frame token $$\pi_\theta(a_t \mid h_t), \qquad h_t = o_{\le t} = (o_0,\dots,o_t).$$8 (Lei et al., 5 Mar 2026).

This separation of memory roles suggests a specific inductive bias: recent observations are treated as high-fidelity control context, whereas distant history is treated as content to be compressed for long-horizon retrieval.

3. Contextual memory compressor

The technical core of VPWEM is a Transformer-based contextual memory compressor that recursively converts out-of-window observations into a fixed number of episodic memory tokens. The compressor uses trainable query tokens $$\pi_\theta(a_t \mid h_t), \qquad h_t = o_{\le t} = (o_0,\dots,o_t).$$9, an observation cache $\pi_\theta(a_t \mid o_{t-L:t}),$0, and one summary-token cache $\pi_\theta(a_t \mid o_{t-L:t}),$1 for each compressor block (Lei et al., 5 Mar 2026).

The initial query state is

$\pi_\theta(a_t \mid o_{t-L:t}),$2

The observation cache and summary cache are concatenated as

$\pi_\theta(a_t \mid o_{t-L:t}),$3

$\pi_\theta(a_t \mid o_{t-L:t}),$4

The compressor then applies, for each layer $\pi_\theta(a_t \mid o_{t-L:t}),$5,

$\pi_\theta(a_t \mid o_{t-L:t}),$6

$\pi_\theta(a_t \mid o_{t-L:t}),$7

$\pi_\theta(a_t \mid o_{t-L:t}),$8

and outputs

$\pi_\theta(a_t \mid o_{t-L:t}),$9

The operator $L$0 refers to standard Transformer attention (Lei et al., 5 Mar 2026).

The first attention stage performs self-attention over cached summary tokens; the second performs cross-attention over cached historical observations. Operationally, the compressor updates online as tokens leave the working-memory window: the token is positionalized, inserted into the observation cache, processed by the compressor, and the resulting layerwise query tokens are stored back into summary caches,

$L$1

Because working-memory length $L$2, episodic-memory size $L$3, and cache size $L$4 are fixed, the paper characterizes memory and compute per step as nearly constant (Lei et al., 5 Mar 2026).

An important implementation detail is that cached $L$5 and $L$6 are detached from the computational graph. This ensures that historical information is propagated only via summary tokens, reduces memory consumption, and is described as mitigating overfitting (Lei et al., 5 Mar 2026).

4. Diffusion-policy instantiation and learning objective

VPWEM is instantiated on diffusion policies, specifically DP and MaIL, so that action generation is conditioned on both short-term working memory and long-term episodic memory. Conceptually, the policy is written as

$L$7

or, for chunked actions,

$L$8

The diffusion backbone uses a Transformer decoder architecture and applies a cross-attention mask for conditioned action-chunk generation (Lei et al., 5 Mar 2026).

Given a clean action chunk $L$9, the forward diffusion process is

$O(L^2)$0

where $O(L^2)$1. The denoiser predicts injected noise conditioned on noisy action, working memory, episodic memory, and diffusion timestep,

$O(L^2)$2

Training minimizes the DDPM-style denoising loss

$O(L^2)$3

The compressor is trained end-to-end jointly with this objective; it is not pretrained separately (Lei et al., 5 Mar 2026).

At inference time, VPWEM encodes the newest observation, updates the working-memory queue, inserts any evicted token into the observation cache, updates the summary caches by running the compressor, obtains episodic memory $O(L^2)$4, and then performs reverse diffusion conditioned on $O(L^2)$5 and $O(L^2)$6. The reverse step is given in the paper as the standard DDPM update,

$O(L^2)$7

with $O(L^2)$8 (Lei et al., 5 Mar 2026).

For DP-VPWEM, the reported settings are: DDPM with 50 sampling steps, action chunk length $O(L^2)$9, temperature $$w_t \triangleq o_{t-L:t} = \mathrm{concat}(o_{t-L+1},\dots,o_t)\in \mathbb{R}^{L\times D}.$$0, 8 layers, embedding dimension $$w_t \triangleq o_{t-L:t} = \mathrm{concat}(o_{t-L+1},\dots,o_t)\in \mathbb{R}^{L\times D}.$$1, compressor depth $$w_t \triangleq o_{t-L:t} = \mathrm{concat}(o_{t-L+1},\dots,o_t)\in \mathbb{R}^{L\times D}.$$2, short-term tokens $$w_t \triangleq o_{t-L:t} = \mathrm{concat}(o_{t-L+1},\dots,o_t)\in \mathbb{R}^{L\times D}.$$3, long-term tokens $$w_t \triangleq o_{t-L:t} = \mathrm{concat}(o_{t-L+1},\dots,o_t)\in \mathbb{R}^{L\times D}.$$4, maximum cache size $$w_t \triangleq o_{t-L:t} = \mathrm{concat}(o_{t-L+1},\dots,o_t)\in \mathbb{R}^{L\times D}.$$5, dropout ratio of memory tokens $$w_t \triangleq o_{t-L:t} = \mathrm{concat}(o_{t-L+1},\dots,o_t)\in \mathbb{R}^{L\times D}.$$6, batch size $$w_t \triangleq o_{t-L:t} = \mathrm{concat}(o_{t-L+1},\dots,o_t)\in \mathbb{R}^{L\times D}.$$7, learning rate $$w_t \triangleq o_{t-L:t} = \mathrm{concat}(o_{t-L+1},\dots,o_t)\in \mathbb{R}^{L\times D}.$$8, EMA rate $$w_t \triangleq o_{t-L:t} = \mathrm{concat}(o_{t-L+1},\dots,o_t)\in \mathbb{R}^{L\times D}.$$9, and AdamW with $$o_{\le t-L} = \mathrm{concat}(o_0,o_1,\dots,o_{t-L}),$$0, $$o_{\le t-L} = \mathrm{concat}(o_0,o_1,\dots,o_{t-L}),$$1 (Lei et al., 5 Mar 2026).

5. Empirical evaluation and reported performance

VPWEM is evaluated on three benchmark families: MIKASA, MoMaRT, and Robomimic. MIKASA includes the memory-intensive tabletop tasks ShellGameTouch-v0 and RememberColor3-v0. MoMaRT includes five long-horizon mobile-manipulation tasks: Table Setup from Dishwasher, Table Setup from Dresser, Table Cleanup to Dishwasher, Table Cleanup to Sink, and Unload Dishwasher to Dresser. Robomimic includes mostly near-Markovian tasks such as Square and Transport, together with Long-Horizon Square, where the robot must remember the initial block configuration (Lei et al., 5 Mar 2026).

The baseline set includes RNN policy, Diffusion Policy (DP), DP-PTP, and MaIL. On MIKASA, the paper also reports comparisons against VLA baselines including Octo-small, OpenVLA, $$o_{\le t-L} = \mathrm{concat}(o_0,o_1,\dots,o_{t-L}),$$2, SpatialVLA, CronusVLA, and MemoryVLA (Lei et al., 5 Mar 2026).

The paper states that VPWEM “outperforms state-of-the-art baselines including diffusion policies and vision-language-action (VLA) models by more than 20% on the memory-intensive manipulation tasks in MIKASA” and achieves “an average 5% improvement on the mobile manipulation benchmark MoMaRT” (Lei et al., 5 Mar 2026). On Robomimic, VPWEM performs on par with baselines, which the paper interprets as indicating no significant degradation in almost-Markovian settings (Lei et al., 5 Mar 2026).

A compute–performance comparison is reported on MoMaRT unload-suboptimal, contrasting long-context DP-PTP variants with DP-VPWEM (L2S2):

Method	Context	Model Size (M)	GPU Mem (MB)	Train Time (s/step)	Inference Time (s)	Success Rate (%)
DP-PTP	4	50.67	656	0.02	0.17	46.1
DP-PTP	8	50.67	684	0.02	0.18	45.0
DP-PTP	16	50.68	794	0.08	0.27	49.1
DP-PTP	32	50.69	942	0.08	0.35	49.6
DP-PTP	64	50.70	1258	0.17	0.44	49.6
DP-PTP	128	50.73	1792	0.18	0.72	40.4
DP-VPWEM	L2S2	52.98	734	0.09	0.22	58.3

This result is used in the paper to support the claim that simply increasing context length causes substantial growth in memory and inference cost, whereas VPWEM attains the best reported success rate with bounded memory structures (Lei et al., 5 Mar 2026).

6. Ablations, limitations, and significance

The ablation studies are performed on Long-Horizon Square. They show that even L1S1 achieves 65% success rate, already outperforming the listed baselines. Increasing either short-term or long-term token count generally helps at first, but too many tokens can hurt, so the main experiments use L2S2. Increasing compressor depth improves performance, and increasing cache size helps until performance saturates beyond about 8. For cache management, FIFO, K-Means, and AdjSim perform similarly, whereas Random performs worse; the main model therefore uses FIFO (Lei et al., 5 Mar 2026).

The paper frames several limitations and future directions. It suggests extending VPWEM to more base policies, adding auxiliary objectives such as reconstruction, and deploying on real robots (Lei et al., 5 Mar 2026). A plausible implication is that the current formulation is benchmark-centered and primarily validated in simulation-oriented or standard offline imitation-learning settings rather than in broad real-world deployment.

In methodological terms, VPWEM’s significance lies in treating long-horizon robotic memory not as unbounded context accumulation but as a structured combination of explicit short-term history and learned compressed long-term state. The framework’s contribution is therefore architectural as much as empirical: it proposes a concrete mechanism by which diffusion-based policies can access temporally distant information without scaling attention over the full trajectory (Lei et al., 5 Mar 2026).