Source-linked AI summary
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
TL;DR
Phase-ambiguous manipulation requires history when similar observations occur at different phases calling for different actions. PMTRM encodes bounded executed state-action history and combines temporal discrimination with information-preserving objectives. It improves phase disambiguation, with reported inference overhead of 6.2% for ACT.
Problem
Locally similar observations can occur at different execution phases that require different actions, making history important for phase-aware action selection.
Source: §I · p. 1Method
PMTRM re-encodes bounded executed state-action history into a latent sequence, using temporal heterogeneity with anchor and reconstruction losses to preserve action-relevant information.
Source: §I · pp. 1–2Results
PMTRM improves phase disambiguation, with 0.42 ms per ACT decision and 6.2% overhead over the 6.8 ms base policy.
Source: G. Computational Cost and Scaling · p. 5, §V · p. 8Takeaways & Limitations
The bounded-history module targets phase disambiguation without retrieval, while explicit memories remain preferable for unbounded semantic recall.
Source: §V · p. 8When relevant events lie outside the bounded queue, increasing T beyond 128 gives little benefit; external retrieval, spatial memory, or recurrent full-history modules are appropriate complements.
Source: C. Experiment 3: Memory and Repetitive… · p. 8
Abstract
from arXiv · showhide
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.
I. INTRODUCTION
Phase ambiguity arises when locally similar observations occur at different execution phases that call for different actions. PMTRM encodes bounded executed state-action history for policy use, pairing temporal discrimination with information-preserving objectives and progressive training.
I. INTRODUCTION
- Observations can look similar across phases that require different actions, so insufficient history may repeat a completed motion or trigger a premature phase switch.Examples include an unchanged-looking button after a press and motion along the same path in opposite directions.Source: §I · p. 1
- Existing recurrent states, history tokens, and memory-bank retrieval provide context but can require dedicated update or retrieval mechanisms and policy-specific fusion layers.The paper targets phase information available in bounded recent history while seeking limited overhead and compatibility across policy backbones.Source: §I · p. 1
- PMTRM re-encodes bounded executed state-action history into a latent sequence for existing policies.The module has 7.61M parameters and regularizes its policy-facing latent for distant-phase discriminability while keeping it reconstructable.Source: §I · pp. 1–2
- Temporal heterogeneity discourages similarity between distant positions, while anchor and reconstruction losses preserve action-relevant information without phase labels.The objectives use temporal distance rather than phase annotations.Source: §I · pp. 1–2
- Training progresses from synthetic sequences through robot data and mirror-mask fine-tuning to joint policy training, retaining the original action head and action space at inference.The paper evaluates multiple policy backbones on phase-ambiguous tasks in simulation and on a real robot, including analyses of the objectives and mirror-mask training.Source: §I · p. 2
II. RELATED WORK
The related work contrasts memory mechanisms and policy-side conditioning with PMTRM’s bounded latent sequence, which directly regularizes temporal discriminability while preserving action-relevant information.
A. Temporal Memory and Pseudo-Memory
- History-aware policies use recurrent states, working-memory tokens, or retrieval, while cyclic manipulation also requires tracking execution progress.PMTRM instead regularizes a bounded state-action history sequence to distinguish distant temporal positions.Source: A. Temporal Memory and Pseudo-Memory · p. 2
B. Policy-Side Temporal Conditioning
- PMTRM exposes the full bounded latent sequence to an unchanged action head, using its decoder only during training rather than compressing history to one summary.It is suited to recent phase aliasing; explicit memory remains preferable for semantic retrieval or context beyond the queue.Source: B. Policy-Side Temporal Conditioning · p. 2
C. Temporal Discriminability
- Unlike prior temporal structuring through correspondence, alignment, or future prediction, PMTRM directly regularizes policy-consumed history to separate distant positions while preserving action-relevant information.This motivates its temporal heterogeneity objective.Source: C. Temporal Discriminability · p. 2
III. METHOD
- The method is organized around the history representation, its temporal heterogeneity and information-preservation objectives, and progressive training with downstream policy integration.Source: §III · p. 2
A. Pseudo-Memory Temporal Re-encoding Module
PMTRM re-encodes a masked, bounded history sequence and exposes the resulting sequence representation through a policy interface, with the decoder reserved for training.
A. Pseudo-Memory Temporal Re-encoding Module
- The input is a padded state-action trajectory sequence with temporal and feature masks for valid steps and channels.The feature mask is episode-constant, while the temporal mask varies with the history prefix.Source: A. Pseudo-Memory Temporal Re-encoding… · p. 2
- The policy-facing encoder maps the masked sequence to Z, while the decoder reconstructs the trajectory during training only.Visual backbones align frame features to the same queue entries; their projected features are concatenated with Z for backbone-specific input adapters.Source: A. Pseudo-Memory Temporal Re-encoding… · p. 2
- Z preserves the input sequence shape rather than compressing coordinates, and is exposed to the policy and heterogeneity loss while E2 preserves action-relevant information.Source: A. Pseudo-Memory Temporal Re-encoding… · p. 2
- The four-stage framework trains PMTRM through synthetic pretraining, real-data fine-tuning, mirror-mask fine-tuning, and joint policy training.At deployment only E1 encodes the executed history queue, and the predicted future action chunk is not re-encoded until execution.Source: Figure 2 · p. 3
B. Masked Reconstruction and Latent Anchor Loss
The latent anchor loss keeps the re-encoded sequence near the normalized input trajectory without forcing identity, while reconstruction preserves information needed for executable action prediction.
B. Masked Reconstruction and Latent Anchor Loss
- The latent anchor loss defines a small residual trust region around the normalized input and prevents heterogeneity from destroying action semantics.Its masked mean-squared-error and cosine-distance terms have a deliberately small total-objective weight, so the anchor does not dominate representation learning.Source: B. Masked Reconstruction and Latent… · p. 3
- The reconstruction term preserves information needed for executable action prediction.Source: B. Masked Reconstruction and Latent… · p. 3
C. Temporal Heterogeneity Objective
The temporal heterogeneity objective discourages distant positions in a trajectory from becoming overly similar, while leaving local smoothness unconstrained. Anchor and reconstruction objectives preserve action-relevant information, and the heterogeneity term is skipped when too few valid distant pairs remain.
C. Temporal Heterogeneity Objective
- The objective targets overly similar latent representations at distant time steps, especially in trajectories with repeated or locally similar motions.It penalizes positive similarity only for distant pairs; local smoothness is unconstrained while anchor and reconstruction objectives protect action semantics.Source: C. Temporal Heterogeneity Objective · p. 3
- With K = 32 anchors and δ = 8, the objective excludes local pairs and invalid temporal positions from its comparisons.Anchors are resampled per batch; the margin and anchor-count settings were selected on synthetic pretraining and checked on Button Press-and-Release and Stack Block.Source: C. Temporal Heterogeneity Objective · p. 3
- For short or heavily masked windows, the heterogeneity term is skipped when fewer than two distant valid pairs remain.The remaining batches contribute only anchor and reconstruction gradients, avoiding unstable large gradients when history is too short.Source: C. Temporal Heterogeneity Objective · p. 4
- The selected K and δ settings change final success by at most 1.4 points, and the same settings are fixed across backbones.At least eight active anchor–time pairs occur in 93.4% of synthetic windows, 89.7% of real-data windows, and 86.1% of mirror crops.Source: C. Temporal Heterogeneity Objective · p. 4
D. Overall Training Objective
PMTRM combines temporal heterogeneity, latent anchoring, and reconstruction, then trains progressively on synthetic sequences and robot trajectories before joint policy training. The schedule includes mirror-and-mask representation augmentation and adapts the module to the selected backbone.
D. Overall Training Objective
- The final objective combines temporal heterogeneity, latent anchoring, and reconstruction, with implementation weights αh = 1.5, β = 0.05, and γ = 8.0.Typical normalized batch losses are 0.10–0.16 for heterogeneity, 0.03–0.06 for anchoring, and 0.006–0.012 for reconstruction.Source: D. Overall Training Objective · p. 4
E. Progressive Pretraining and Fine-tuning
- Synthetic pretraining uses mixed periodic and nonperiodic windows with random lengths and active dimensions, padded and masked to the model shape.The stage uses 50,000 windows for 100 epochs; skipping it lowers Button Press-and-Release CSR by 3.2 points and cross-embodiment mean completion by 2.0 points.Source: E. Progressive Pretraining and… · p. 4
- Robot-data fine-tuning uses normalized, padded state-action trajectories, while the dataset corpus spans varied embodiments, object interactions, and multi-stage tasks.Channels are ordered as pose, gripper, then action; datasets are sampled with per-epoch caps so no corpus dominates gradients.Source: E. Progressive Pretraining and… · p. 4, Table I · p. 4
- Mirror-and-mask training creates a virtual reversed sequence for representation augmentation, then samples a contiguous window of at most 128 entries.The reversed stream is not executed or used as a policy target; reversal is applied to normalized vectors without sign-flipping channels.Source: E. Progressive Pretraining and… · p. 4
- Joint training adapts PMTRM to the policy backbone; the full four-stage schedule reaches 70.9% average memory-task success.Omitting Stage 1, Stage 2, or Stage 3 yields 67.7%, 69.1%, or 68.4%, respectively; omitting Stage 3 costs 2.5 points while keeping the same inference graph.Source: E. Progressive Pretraining and… · p. 4
F. Integration with Downstream Embodied Policies
PMTRM supplies encoded preceding state-action history to existing ACT, Diffusion Policy, and VLA backbones without changing their action interfaces. Its auxiliary objective is added to the original policy loss, and inference retains only the policy-facing encoder.
F. Integration with Downstream Embodied Policies
- PMTRM supplies padded, encoded preceding states and actions to ACT, Diffusion Policy, and VLA policies while leaving their action heads and action spaces unchanged.The integration preserves continuous actions, discrete tokens, action chunks, decoders, and normalization.Source: F. Integration with Downstream… · p. 4
- Joint optimization adds PMTRM's auxiliary objective to each backbone's original policy loss without task-specific loss redesign.Policy-specific objectives, schedules, and architectural settings follow each backbone's default configuration.Source: F. Integration with Downstream… · p. 4
G. Computational Cost and Scaling
- At inference, the training module's 7.61M parameters reduce to a 3.81M-parameter policy-facing encoder, with linear cost in history length.The reconstruction decoder is discarded; an episode-level queue supplies preceding executed states and actions at each decision step.Source: G. Computational Cost and Scaling · pp. 4–5
- At T = 128, PMTRM adds 0.42 ms per ACT decision on an RTX 4090, or 6.2% over the 6.8 ms base policy.The added latency is 2.7 ms on an i7-13700K CPU; GPU latency is 0.16, 0.27, 0.42, and 0.83 ms for T = 32, 64, 128, and 256.Source: G. Computational Cost and Scaling · p. 5
IV. EXPERIMENTS
The experiments examine representation quality and objective trade-offs, including frozen-representation phase probes. PMTRM separates distant phases while preserving local continuity, and its probe and ablation results quantify that behavior and its sensitivity to design choices.
A. Experiment 1: Loss Dynamics and Objective Trade-offs
- 98.1% phase accuracy from frozen Z compares with 71.4% from X and 76.8% from a single-stage AE.On real Drawer Open-and-Close, Z predicts open/close at 94.6% versus 79.2% from raw X; on Pour Bottle Twice, it predicts pour count at 86.0% versus 61.5%.Source: A. Experiment 1: Loss Dynamics and… · p. 5
- 70.9% at 0.42 ms for PMTRM, versus 68.8/69.4/69.7/69.2% at 0.61/0.48/0.46/0.58 ms for short-attention, diagonal-SSM, Mamba, and action-memory controls.The controls share the corpus, masks, demonstrations, seeds, and T = 128.Source: A. Experiment 1: Loss Dynamics and… · p. 5
- For Button CSR / Drawer CSR / Pour Finish / Shell Finish, β = 0.01, 0.05, and 0.10 yield 61/73/70/60, 65/79/80/70, and 64/77/80/60%, respectively.A symmetric negative-cosine penalty reduces local-continuity probe accuracy by 2.3 points and increases MSE by 0.6; the one-sided hinge is retained.Source: A. Experiment 1: Loss Dynamics and… · p. 5
- With ten real rollouts, a 70% estimate has a Wilson 95% CI of [39.7, 89.2]%; real gains are directional, while multi-seed simulation is higher-confidence.Distant Cos. is diagnostic rather than a value to minimize blindly; around 0.2–0.25 separates aliased phases while preserving continuity, whereas Markov tasks can remain strong near 0.30.Source: A. Experiment 1: Loss Dynamics and… · p. 5
B. Experiment 2: Compatibility with Conventional Action Policies
PMTRM is evaluated with conventional action-policy backbones on simulated and real-robot manipulation tasks. Results include improvements on later-stage metrics and fewer approach oscillations in Stack Blocks.
B. Experiment 2: Compatibility with Conventional Action Policies
- Removing proprioceptive/action history lowers ACT-PMTRM simulation success by 5.4 points, while removing visual features lowers it by 0.8 points.The authors interpret these results as showing that PMTRM supplies phase information complementary to visual grounding.Source: B. Experiment 2: Compatibility with… · p. 5
- Methods use identical demonstrations and splits, and the action decoders remain unchanged; simulation uses ten fixed seeds and real trials use ten randomized poses per method.Stage 1–3 provide a shared initialization, followed by independent PMTRM fine-tuning for each backbone.Source: B. Experiment 2: Compatibility with… · p. 5
- The evaluations cover transfer, insertion, and stacking tasks alongside repeated-phase tasks such as press/release, drawer open/close, twice-pouring, and Shell Game.The repeated-phase tasks require distinguishing execution stages using historical context.Source: Figure 4 · p. 6, Figure 5 · p. 6
- Gains are largest in later-stage Transfer, Insert, and Finish metrics; in Stack Blocks, sign changes fall from 4.7 to 2.9 per rollout and successful grasps rise from 69% to 78%.The Stack Blocks changes are consistent with fewer approach oscillations.Source: B. Experiment 2: Compatibility with… · p. 6
C. Experiment 3: Memory and Repetitive Manipulation
On memory-intensive repetitive tasks, adding PMTRM improves simulation CSR and physical-task completion, with gains that level off as history extends beyond the relevant interaction horizon. The experiments also show a bounded-history scope: long sequences can place needed phase or progress information outside the retained window.
C. Experiment 3: Memory and Repetitive Manipulation
- Evaluation uses held-out scene seeds and no test trajectories for tuning; demonstrations are episode-split, with LIBERO drawer variants additionally split by scene seed.Simulation and real-robot demonstrations and rollout counts are specified in the experimental protocol.Source: C. Experiment 3: Memory and Repetitive… · p. 6
- The four tasks test repeated-action phase distinctions, progress tracking, and history-dependent execution, including opposite actions under similar observations and tracking previously wiped regions.The task set includes simulation and physical interaction settings.Source: C. Experiment 3: Memory and Repetitive… · p. 6
- CSR improves by 2.4 and 1.6 points on simulation tasks, while final completion improves by 10 points on each physical task with the PrediMem+PMTRM adapter.PMTRM receives executed state/action history while PrediMem retains predictive tokens.Source: C. Experiment 3: Memory and Repetitive… · p. 6
- On Button Press-and-Release, CSR is 74% at T = 128 and remains 74% at T = 256, while GPU overhead rises from 0.42 to 0.83 ms.The reported practical rule is to cover one to two task-phase cycles and extend history only while success improves faster than latency.Source: C. Experiment 3: Memory and Repetitive… · p. 7
- Long or interrupted drawer sequences can push the needed opening or closing transition outside retained history, while extended wiping requires progress or coverage information beyond that window.Increasing T beyond 128 gives little benefit once the relevant event lies outside the queue.Source: C. Experiment 3: Memory and Repetitive… · pp. 7–8
V. CONCLUSION
PMTRM is a lightweight bounded-history representation module for robotic manipulation. It improves phase disambiguation without retrieval or extra policy tokens, while explicit memories remain preferable for unbounded semantic recall.
V. CONCLUSION
- PMTRM improves phase disambiguation without retrieval or extra policy tokens, but explicit memories remain preferable for unbounded semantic recall.The module is trained and evaluated with imitation objectives; in RL, distribution shift and reward-driven phase changes may require adjusting the heterogeneity objective.Source: §V · p. 8