PMTRM: Pseudo-Memory Temporal Re-encoding Module for Embodied Policy Learning

Changchuan Yang; Haoxuan Xu; Wenbo Chen; Shuai Ren; Jianlong Zheng; Huarui Zhang; Tianfu Li; Guanzhong Tian

cs.ROOct 8, 2026

PMTRM: Pseudo-Memory Temporal Re-encoding Module for Embodied Policy Learning

Authors: Changchuan Yang, Haoxuan Xu, Wenbo Chen, Shuai Ren, Jianlong Zheng, Huarui Zhang, Tianfu Li, Guanzhong Tian

Organizations: Zhejiang University · The Hong Kong University of Science and Technology · The Hong Kong University of Science and Technology (Guangzhou) · vivo Robotics Lab · South China Normal University

Abstract

Robotic manipulation often contains repeated motions whose local observations look similar at different phases. When these phases require different actions, a policy that relies mainly on the current observation may repeat completed motions or switch phases at the wrong time. To address this phase ambiguity, we present the Pseudo-Memory Temporal Re-encoding Module (PMTRM), a lightweight plug-in module with only 7.61M parameters that encodes a bounded history of executed states and actions into a latent sequence for existing policies. To help distinguish phases, a temporal heterogeneity objective penalizes positive similarity between distant positions in this sequence, while anchor and reconstruction losses preserve information needed for action prediction. The reconstruction decoder is used only during training, leaving the temporal re-encoder to supply history to the policy at inference. We train the module progressively on synthetic sequences and robot data, then jointly with the policy, using temporal masking to accommodate partial histories. This integration retains the original action head and action space and adds auxiliary losses to the original policy loss. Experiments with multiple policy backbones in simulation and on a real robot show improved task success on tasks with phase ambiguity, with little additional computation.

Figures & tables

Fig. 1: Conceptual comparison of temporal memory mechanisms for repetitive manipulation. (a) Representative memory-based approaches summarize history into recurrent states or memory tokens, or retrieve relevant information from an explicit memory bank. (b) PMTRM instead re-encodes a bounded queue of executed state–action history into a phase-discriminative latent sequence $Z$ .

Fig. 2: Overall framework of PMTRM. (a) The module employs a two-stage convolutional temporal re-encoding autoencoder to map synthetic or real executed state-action histories into pseudo-memory representations.(b) Stages 1–4 correspond to synthetic pretraining, open-source real-data fine-tuning, mirror-mask fine-tuning (Eq. ( 7 )), and joint policy training, respectively.(c) At deployment, $E_{2}$ is disabled and only $E_{1}$ encodes the executed state/action history queue (up to $T$ steps). The policy consumes this queue representation and predicts a separate future action chunk, which is not re-encoded until its actions are executed.

Dataset family	State/Action	Episodes	Frames
DROID/Bridge/RT-1 [ 11 , 12 , 13 ]	7–8/7	137,700	5.6M
LIBERO/CALVIN [ 14 , 15 ]	8–15/7	34,800	2.6M
BEHAVIOR-1K [ 16 ]	256/23	10,000	119.1M
TACO/Franka/TOTO [ 17 , 18 , 19 ]	7–13/7–15	4,200	0.6M
Austin/UT/Berkeley/Stanford [ 20 , 21 , 22 , 23 , 24 ]	7–15/7–8	3,800	1.8M
X-VLA soft-fold [ 25 ]	96/14	1,542	2.9M

TABLE I: Open-source robot manipulation datasets used in the fine-tuning setting. The corpus includes single-arm and bimanual data (UMI and BEHAVIOR-1K), deformable-object interactions (towel/dish washing and soft-fold), and multi-stage trajectories (LIBERO, CALVIN, BEHAVIOR-1K, and X-VLA); these tags describe pretraining diversity rather than a separate benchmark split.

Fig. 3: Representation analysis on a phase-aliasing cyclic trajectory. Top: temporal cosine-similarity matrices for a standard autoencoder and PMTRM (the same $[0,1]$ color scale is used; dotted boxes mark cycle-1/cycle-2 aliased-phase blocks measured by Distant Cos. in Table II , and dashed lines mark the cycle boundary). Bottom: PCA projections of two cycles; stars mark their starts. PMTRM suppresses off-diagonal similarities while retaining local continuity.

Setting	Avg. Success $\uparrow$	Recon. MSE $\downarrow$	Distant Cos. $\downarrow$
Single-stage AE + heterogeneity (same budget)	63.4	2.4	0.31
GRU + reconstruction (same policy interface)	64.6	3.1	0.29
Contrastive temporal loss (same budget)	67.2	3.0	0.25
PMTRM w/o synthetic pretraining	67.7	2.8	0.23
PMTRM, ACT MuJoCo Markov reach	91.6	1.9	0.30
PMTRM, $\beta=0$	65.1	7.8	0.18

TABLE II: Ablation of PMTRM design choices on the memory-intensive evaluation (ACT MuJoCo Button Press-and-Release, LIBERO Drawer Open-and-Close, Pour Bottle Twice, and Shell Game). Avg. success follows the metric in Table V ; all variants use the same demonstrations, seeds, and

T=128

Fig. 4: Left: The MuJoCo tasks Transfer Cube , Bimanual Insertion , Transfer Plus , and Stack Two Blocks . Right: The real-world tasks Stack Block , Store Items , Potato Placement , and Bowl Stacking .

Fig. 5: Left: The simulation tasks Button Press-and-Release (MuJoCo) and Drawer Open-and-Close (LIBERO). Right: The real-world tasks Pour Bottle Twice and Shell Game . These tasks require the policy to distinguish repetitive motions and execution phases using historical context rather than only the current observation; the key phase transitions are press versus release, open versus close, first versus second pour, and target tracking after swaps.

	Transfer Cube (Sim)			Bimanual Insertion (Sim)			Transfer Plus (Sim)			Stack Two Blocks (Sim)
	Touch	Lift	Transfer	Grasp	Contact	Insert	Lift	Stack	Finish	Stack	Lift	Finish
DP (DDPM, CNN) [ 5 ]	$95.1$ $\pm$ $1.2$	$92.4$ $\pm$ $1.3$	$90.2$ $\pm$ $1.9$	$77.4$ $\pm$ $2.7$	$69.0$ $\pm$ $3.1$	$63.2$ $\pm$ $3.3$	$62.1$ $\pm$ $3.5$	$52.9$ $\pm$ $3.9$	$52.9$ $\pm$ $3.9$	$82.0$ $\pm$ $2.7$	$63.5$ $\pm$ $3.2$	$46.7$ $\pm$ $3.6$
DP-PMTRM	$97.0$ $\pm$ $1.0$	$94.1$ $\pm$ $1.1$	$92.8$ $\pm$ $1.5$	$81.0$ $\pm$ $2.3$	$73.6$ $\pm$ $2.7$	$67.8$ $\pm$ $3.0$	$66.8$ $\pm$ $3.2$	$58.4$ $\pm$ $3.5$	$58.4$ $\pm$ $3.5$	$85.1$ $\pm$ $2.4$	$68.0$ $\pm$ $3.0$	$53.2$ $\pm$ $3.4$
ACT [ 4 ]	$98.4$ $\pm$ $0.9$	$96.2$ $\pm$ $1.0$	$94.8$ $\pm$ $2.1$	$81.3$ $\pm$ $2.6$	$73.5$ $\pm$ $3.0$	$68.1$ $\pm$ $3.2$	$66.7$ $\pm$ $3.6$	$57.5$ $\pm$ $4.0$	$57.5$ $\pm$ $4.0$	$85.8$ $\pm$ $2.8$	$67.3$ $\pm$ $3.4$	$50.4$ $\pm$ $3.8$
ACT-PMTRM	$99.1$ $\pm$ $0.5$	$97.5$ $\pm$ $0.8$	$96.4$ $\pm$ $1.5$	$85.7$ $\pm$ $2.2$	$79.4$ $\pm$ $2.6$	$73.8$ $\pm$ $2.9$	$72.6$ $\pm$ $3.1$	$63.8$ $\pm$ $3.5$	$63.8$ $\pm$ $3.5$	$89.3$ $\pm$ $2.3$	$72.8$ $\pm$ $3.0$	$58.7$ $\pm$ $3.4$
HACT-VQ [ 28 ]	$98.5$ $\pm$ $0.9$	$97.6$ $\pm$ $1.2$	$96.2$ $\pm$ $1.4$	$87.4$ $\pm$ $2.4$	$82.2$ $\pm$ $2.7$	$76.3$ $\pm$ $2.8$	$79.2$ $\pm$ $2.2$	$68.5$ $\pm$ $3.1$	$68.5$ $\pm$ $3.1$	$90.6$ $\pm$ $2.3$	$76.1$ $\pm$ $2.8$	$55.8$ $\pm$ $3.2$

TABLE III: Simulation evaluation of PMTRM with different action policy backbones. Success rate (%)

\uparrow

is reported as mean

\pm

standard deviation over 10 random seeds; the rightmost subcolumn of each task is end-to-end completion, and all PMTRM variants use

T=128

	Stack Block			Store Items			Potato Placement			Bowl Stacking
	Grasp	Lift	Stack	First	Second	Finish	Grasp	Lift	Place	First	Second	Finish
DP (DDPM, CNN) [ 5 ]	0.82	0.70	0.62	0.79	0.61	0.54	0.91	0.83	0.74	0.77	0.59	0.50
DP-PMTRM	0.88	0.78	0.71	0.86	0.86	0.64	0.95	0.89	0.82	0.84	0.76	0.61
ACT [ 4 ]	0.85	0.74	0.69	0.86	0.65	0.58	0.93	0.86	0.80	0.83	0.62	0.55
ACT-PMTRM	0.91	0.84	0.78	0.91	0.78	0.70	0.97	0.92	0.87	0.89	0.75	0.67
GR00T N 1.7 [ 30 ]	0.86	0.76	0.70	0.85	0.68	0.60	0.94	0.88	0.81	0.83	0.65	0.57

TABLE IV: Real-world evaluation of PMTRM with different action policy backbones. Success rates are reported in

[0.0,1.0]

with

T=128

for PMTRM; the rightmost subcolumn of each task is final completion. Final-completion entries use ten Bernoulli trials, while intermediate stage rates pool step outcomes; for the former, a 0.70 estimate has Wilson 95% CI

[0.397,0.892]

rather than a standard-deviation interpretation. The real-world benchmark includes four physical manipulation tasks: Stack Block , Store Items , Potato Placement , and Bowl Stacking .

	Button Press-and-Release (Sim)		Drawer Open-and-Close (Sim)		Open and Close Drawer (Real)		Wipe Table (Real)
	TSR	CSR	TSR	CSR	First Open	Finish	First Wipe	Finish
HiF-VLA [ 31 ]	65.8	59.2	68.6	60.4	50.0	30.0	30.0	20.0
MemoryVLA [ 1 ]	64.1	57.5	72.8	64.6	50.0	30.0	40.0	20.0
MemER [ 32 ]	70.6	64.4	79.3	71.5	60.0	40.0	60.0	40.0
PrediMem [ 33 ]	76.4	70.6	84.8	78.2	70.0	60.0	70.0	50.0
PrediMem-PMTRM	79.2	73.8	87.4	81.6	80.0	70.0	80.0	60.0

TABLE V: Evaluation on memory-intensive and repetitive manipulation tasks. For simulation tasks, we report Task Success Rate (TSR, %) and Cumulative Success Rate (CSR, %). For real-world tasks, we report success rate (%) over 10 randomized closed-loop rollouts; First Open and First Wipe denote successful completion of the first repeated manipulation, while Finish requires completion of the full two-stage sequence. For reference, a 70% success rate over 10 Bernoulli trials has a Wilson 95% confidence interval of approximately

[39.7,89.2]\%

. PMTRM uses

T=128

Explore similar work

Action-Effect Memory Pretraining for Robot Manipulation
Jun 10, 2026Yijing Zhou, Qiwei Liang, Sitong Zhuang +5Representation Learning Robotic Manipulation
Hong Kong University of Science and Technology (Guangzhou) · Shenzhen University
ReCAT: Remember, Count, and Time: Structured Recurrent Memory for Robot Manipulation
Sep 28, 2026Pankhuri Vanjani, Mostafa Hatab, Can Mizrakli +4Language-Conditioned Robot Manipulation Recurrent Neural Networks
Intuitive Robots Lab, Karlsruhe Institute of Technology (KIT), Germany · University of Edinburgh, UK · NVIDIA +1
DRAM: Delta-rule Recurrent Associative Memory for Robot Manipulation Policies
Sep 26, 2026Xinyu Zhao, Yixiang Shan, Tao Yang +5Long-Horizon Robotic Manipulation Delta-Rule Linear Attention
Simplexity Robotics

Jun 10, 2026, cs.ROJ/K move · Enter open · S save

Action-Effect Memory Pretraining for Robot Manipulation

Yijing Zhou, Qiwei Liang, Sitong Zhuang, Jiaxi Li+4

Hong Kong University of Science and Technology (Guangzhou) · Shenzhen University

Representation Learning Robotic Manipulation Dynamic Robotic Manipulation Memory-Augmented Agents Robot Learning

Open paper AskPDF arXiv

We present AEM, an Action-Effect Memory pretraining framework for robot manipulation that learns compact temporal representations from vision-action history. Unlike prior robot representation pretraining methods that mainly focus on single-frame visual encoding, AEM targets the temporal nature of manipulation, where the current observation alone is often insufficient under partial observability. AEM models manipulation as an action-driven interaction process by interleaving visual and action features and applying masked modeling to recover missing content from incomplete histories, thereby learning action-conditioned state evolution. The Mamba-encoded output of the final vision token is used as a compact history representation, serving as the global context for decoding and downstream control. This design preserves a single-vector temporal bottleneck while keeping inference efficient. We evaluate AEM with Diffusion Policy and Flow Policy. AEM consistently improves manipulation performance in both simulation and real-world settings, outperforming baselines across clean scenes, cluttered and random scenes, and non-Markovian tasks. Ablation studies further show that history-aware pretraining surpasses single-frame pretraining and direct frame stacking, while reducing inference latency and computational cost.

Abstract

Explore similar work

Action-Effect Memory Pretraining for Robot Manipulation

ReCAT: Remember, Count, and Time: Structured Recurrent Memory for Robot Manipulation

DRAM: Delta-rule Recurrent Associative Memory for Robot Manipulation Policies