Source-linked AI summary
Resolving State Ambiguity in Robot Manipulation via Adaptive Working Memory Recoding
Qingda Hu, Ziheng Qiu, Zijun Xu, Kaizhao Zhang, Xizhou Bu, Zuolei Sun, Bo Zhang, Jieru Zhao, Zhongxue Gan, Wenchao Ding
TL;DR
State ambiguity can require different actions in similar situations; resolving it may require up to 7 seconds of history, but long histories raise inference cost and challenge stable, affordable training. PAM uses a two-stage visuomotor policy with adaptive working memory that recodes current multimodal inputs and routes historical context into compact features. Across tasks requiring diverse contextual cues, PAM averages 0.91 success, 49.2% above LongDP (0.61) and 102.2% above MTIL (0.45).
Problem
State ambiguity can require different actions in similar situations; resolving it may require up to 7 seconds of history, but long histories raise inference cost and challenge stable, affordable training.
Source: §I · p. 1Method
PAM uses a two-stage visuomotor policy with adaptive working memory that recodes current multimodal inputs and routes historical context into compact features.
Source: §I · pp. 1–2Results
Across tasks requiring diverse contextual cues, PAM averages 0.91 success, 49.2% above LongDP (0.61) and 102.2% above MTIL (0.45).
Source: B. Main Results in Real-World Tasks · p. 6Takeaways & Limitations
PAM provides a practical long-history policy with up to 300 frames and real-time inference at 20 Hz, and its design principles facilitate extension to larger models.
Source: §I · p. 2, §V · pp. 7–8
Abstract
from arXiv · showhide
State ambiguity is common in robotic manipulation. Identical observations may correspond to multiple valid behavior trajectories. The visuomotor policy must correctly extract the appropriate types and levels of information from the history to identify the current task phase. However, naively extending the history window is computationally expensive and may cause severe overfitting. Inspired by the continuous nature of human reasoning and the recoding of working memory, we introduce PAM, a novel visuomotor Policy equipped with Adaptive working Memory. With minimal additional training cost in a two-stage manner, PAM supports a 300-frame history window while maintaining high inference speed. Specifically, a hierarchical frame feature extractor yields two distinct representations for motion primitives and temporal disambiguation. For compact representation, a context router with range-specific queries is employed to produce compact context features across multiple history lengths. And an auxiliary objective of reconstructing historical information is introduced to ensure that the context router acts as an effective bottleneck. We meticulously design 7 tasks and verify that PAM can handle multiple scenarios of state ambiguity simultaneously. With a history window of approximately 10 seconds, PAM still supports stable training and maintains inference speeds above 20Hz. Project website: https://tinda24.github.io/pam/
I. INTRODUCTION
State ambiguity means identical observations can fit multiple valid trajectories, but resolving it requires temporal cues that short observation windows may miss. PAM uses adaptive working memory to represent task context over longer histories, targeting this need while avoiding the cost and overfitting risks of naively extending the window.
I. INTRODUCTION
- State ambiguity arises when identical observations correspond to multiple valid trajectories, leaving short observations unable to provide crucial temporal-disambiguation cues.The introduction describes five scenarios: bidirectional overlap, fixed-duration static posture, sequential subtasks, repetitive cycles, and dynamic interaction response.Source: §I · p. 1
- Longer histories can resolve ambiguity, but naive extension is computationally expensive, may overfit, and makes stable, affordable training more challenging.Some tasks, including Wipe the Table Twice, may require information from up to seven seconds earlier.Source: §I · p. 1
- PAM uses adaptive working memory to recode multimodal inputs into motion primitives and context features, then routes historical context into compact features for action prediction.The auxiliary objective reconstructs historical information to encourage the context router to act as an effective bottleneck.Source: §I · pp. 1–2
- PAM supports a history window of up to 300 frames while preserving real-time inference at 20 Hz and revealing which historical segments and modalities inform disambiguation.Its inference process maintains adaptive working memory and encodes only the current observation at each step.Source: §I · p. 2, Figure 1 · p. 2
- Across seven designed ambiguity tasks, PAM achieves the best performance among methods whose history windows suffice for those tasks and performs competitively on Libero-Long.The introduction also states that adaptive recoding does not require additional annotations.Source: §I · p. 2
II. RELATED WORK
History-aware visuomotor policies encode or retrieve information from past observations in several ways, with different capacity, adaptability, and cost trade-offs. Temporally dependent inference methods maintain evolving state, but the cited approaches still face training-cost and stability challenges.
A. History-Aware Visuomotor Policy
- Manual historical representations add explicit past-motion or object-tracking inputs, but can have limited capacity and substantial annotation costs.Examples include rendered spatial traces and object-centric tracking.Source: A. History-Aware Visuomotor Policy · p. 2
- Memory-bank methods preserve keyframes through dynamic sampling or retrieval, fusion, and consolidation, but require complex storage and retrieval designs.The passage notes that encoding duration and repetition remains difficult to specify with rules.Source: A. History-Aware Visuomotor Policy · p. 2
- Multi-frame feature methods use image sequences or concatenated perceptual features, while sliding-window features can lack adaptability, confuse features, and increase costs with longer histories.The related-work passage characterizes PAM as enhancing adaptability while mitigating negative effects of longer windows.Source: A. History-Aware Visuomotor Policy · p. 2
B. Temporally Dependent Inference
- Temporally dependent approaches include memory-based LSTMs and trajectory-history hidden states, but these methods remain challenged by training cost and stability.The overview figure characterizes PAM as using motion primitives from the current frame and compact context features from extended history.Source: B. Temporally Dependent Inference · p. 3, Figure 2 · p. 3
III. METHOD
PAM reformulates policy inference as temporally dependent: it encodes each current observation while carrying compact context features forward as working memory. A context query recodes multimodal information for that memory, addressing the redundancy and weak semantic grounding of alternative feature sources.
A. Overview of PAM
- PAM predicts future action chunks from a multi-task dataset of language instructions, observations, joint states, and executed actions.Its maximum history window is 300 frames; longer-context scenarios can be handled by the robot system's high-level task planner.Source: A. Overview of PAM · p. 3
- Unlike static policies that treat inference steps independently, PAM carries context features forward as working memory while encoding only the current observation at each step.The context features also support an auxiliary objective that predicts historical image embeddings.Source: A. Overview of PAM · p. 3
- A context query recodes information for working memory because concatenated motion primitives are redundant and action-head intermediate states lack capacity and clear semantic grounding.Causal attention lets the query jointly attend to low-level perceptual cues and high-level decision-relevant features.Source: A. Overview of PAM · p. 3
B. Model Architecture
PAM separates action-relevant motion primitives from context features that summarize task progression and environmental changes. A context router compresses historical context, while an auxiliary objective predicts historical information from it.
B. Model Architecture
- The frame extractor encodes image, language, and proprioceptive inputs, with separate query tokens for action primitives and context features.The context query encodes contextual features used as adaptive working memory.Source: B. Model Architecture · p. 4
- Range-specific context queries summarize history at multiple lengths, producing compact features that capture both short- and long-horizon information.A validity mask handles inference steps with less history than the maximum window.Source: B. Model Architecture · p. 4
- The action head conditions on motion primitives and compact context features to predict action chunks using flow matching.The action head predicts a velocity field that transports an initial action state toward the target action.Source: B. Model Architecture · p. 4
- Four conditioning variants were considered, and staggered cross-attention beginning with motion primitives was the default.The alternatives were concatenation, staggered cross-attention beginning with context features, and ControlNet-style multi-condition injection.Source: B. Model Architecture · p. 4
- The auxiliary objective predicts past image features or trajectories using context features; historical image embeddings were more effective in the authors’ early experiments.A decoder-only transformer predicts past image embeddings, with historical observation patches encoded by a VAE-based autoencoder.Source: B. Model Architecture · p. 4
C. Training Recipe
PAM uses two training stages to learn motion primitives and then temporal disambiguation, while retaining parallel sampling. The overall loss combines flow-matching and auxiliary objectives.
C. Training Recipe
- PAM is trained in two stages with efficient parallel sampling in both, avoiding the instability and high training cost of sequentially training temporally dependent inference models.The authors state that this strategy scales to larger models.Source: C. Training Recipe · p. 4
- Stage 1 trains the action head on motion primitives while context-feature parameters remain inactive, stabilizing multimodal frame features for the next stage.For tasks without pronounced state ambiguity, the policy can already complete them effectively after this stage.Source: C. Training Recipe · p. 4
- In Stage 2, the extractor is frozen and its key–value states are cached, allowing efficient parallel sampling while learning temporal disambiguation.The auxiliary-loss weight λaux is set to 0.4.Source: C. Training Recipe · p. 4
IV. EXPERIMENTS
The experiments evaluate PAM on real-world state-ambiguity tasks, long-horizon Libero-Long tasks, and ablations of its modules and parameter configurations.
IV. EXPERIMENTS
- The experiments first assess PAM’s ability to resolve state ambiguity in real-world tasks.The real-world evaluations are covered in Sections IV-A and IV-B.Source: §IV · p. 4
- The evaluation then tests PAM on long-horizon tasks in Libero-Long.This evaluation is covered in Section IV-C.Source: §IV · p. 4
- Ablations examine the efficacy of each module and how parameter configurations affect overall performance.These studies are covered in Section IV-D.Source: §IV · p. 4
A. Real-World Experimental Setting
The real-world setup evaluates a seven-task suite designed around diverse state ambiguities, using repeated trials and success rates averaged across task stages. PAM is compared with two history-aware baselines.
A. Real-World Experimental Setting
- The real-world experiments use an Agilex Piper arm with a parallel gripper, fixed global and wrist RGB cameras at 320x240, and joint-angle and gripper-state actions.For each task, 50 expert trajectories were collected at 30 Hz using master-slave teleoperation.Source: A. Real-World Experimental Setting · pp. 4–5
- The seven-task suite includes Wipe Two Times, where about 210 frames separate repeated identical wiping trajectories, and Guessing Game, with about 100 frames of occlusion before action.Guessing Game relies on passive observation, whereas Wipe Two Times requires tracking progress across repeated active wiping.Source: A. Real-World Experimental Setting · p. 5
- The real-world baselines are MTIL and LongDP; PAM is a unified multi-task policy, while both baselines use separate single-task models.MTIL encodes trajectory history into a compact latent state, and LongDP reconstructs past action tokens and combines multi-frame features for decoding.Source: A. Real-World Experimental Setting · p. 5
- Each policy is evaluated over 20 trials per task, with overall success averaged across the task’s internal stages; the figure defines task success as mean subtask completion.PAM and MTIL run at 20 Hz, while LongDP’s maximum feasible frequency is 5 Hz.Source: A. Real-World Experimental Setting · p. 5, Figure 3 · p. 5
B. Main Results in Real-World Tasks
PAM performs strongly across real-world tasks requiring diverse contextual cues, while adaptive recoding supports compact history representations and interpretable decisions. Its auxiliary prediction head has task-dependent effects, with overall performance largely consistent when it is removed.
B. Main Results in Real-World Tasks
- PAM averages 0.91 success, a 49.2% gain over LongDP (0.61) and a 102.2% gain over MTIL (0.45) across diverse real-world tasks.PAM achieves success rates of 0.95 on Hold the Pot Lid and 0.92 on Wipe the Table Twice.Source: Table I · p. 6, B. Main Results in Real-World Tasks · p. 6
- PAM's attention maps reveal which historical segments and sensory modalities the policy uses to resolve ambiguity, including visual block-position cues and joint-state cues for posture maintenance.In Wipe the Table Twice, visual observations and joint states both provide effective contextual cues.Source: Figure 4 · p. 6, B. Main Results in Real-World Tasks · p. 6
- LongDP's raw multi-frame conditioning often confuses multi-stage tasks and adds computational overhead, while MTIL's propagated hidden state is susceptible to distortion and erratic early actions.LongDP performs strongly on some tasks, including Sponge and Square and Wipe One Time.Source: B. Main Results in Real-World Tasks · p. 6
- PAM uses adaptive recoding to distill compact, informative temporal representations, and performs seven tasks simultaneously with shared parameters.The paper attributes the task performance to learning highly discriminative history representations.Source: B. Main Results in Real-World Tasks · p. 6
- The auxiliary prediction head leaves overall performance largely consistent, improves Guessing Game success, and can reduce motion smoothness when its weight is too large.The head is disabled at inference, leaving inference cost unchanged.Source: B. Main Results in Real-World Tasks · p. 6
C. Evaluation on Libero-Long Long-Horizon Tasks
On Libero-Long, PAM is evaluated on long-horizon manipulation tasks and reaches an average success rate matching π0 while exceeding the other reported baselines. The evaluation uses a 50-frame history window because these tasks do not exhibit pronounced state ambiguity.
C. Evaluation on Libero-Long Long-Horizon Tasks
- PAM uses a 50-frame history window for these tasks because they do not exhibit pronounced state ambiguity; each task is evaluated over 100 trials across five random seeds.The official dataset provides 50 expert trajectories per task.Source: C. Evaluation on Libero-Long… · p. 7
- PAM reaches an 84.7% average success rate on Libero-Long, matching π0 and significantly exceeding other baselines.The evaluation covers 10 long-horizon tasks.Source: Table II · p. 7, C. Evaluation on Libero-Long… · p. 7
D. Ablation Studies
The ablations examine the visual encoder, condition injection, context router, and context-feature sampling interval. Results support distinct design choices for representing visual cues, sequencing action conditions, compacting context, and balancing sampling cost against information density.
D. Ablation Studies
- The ablations target four factors: visual encoder, multi-condition injection, context-router variants, and context-feature sampling interval.The adopted configuration is the one highlighted in the ablation table and corresponds to the setup in Section IV-B.Source: D. Ablation Studies · p. 7, Table III · p. 7
- ViT-Base scores 0.91, while ViT-Small scores 0.72 in the visual-encoder ablation.The paper reports stereotypical behavior with a lightweight DINOv2 encoder and identifies low pixel occupancy of the red cube as a risk for contextual-cue representation.Source: D. Ablation Studies · p. 7
- Later Inject Context Features scores 0.91, versus Parallel Inject 0.89, First Inject Context Features 0.72, and Concatenate 0.69.The authors report that early motion-primitive injection is essential for guiding action-manifold transition.Source: D. Ablation Studies · p. 7
- Replacing the context router with average pooling performs worst; too few queries fail to integrate cues across frames, while too many reduce compactness and burden action decoding.The router's staged queries attend to different history lengths.Source: D. Ablation Studies · p. 7
- Interval sampling trades computational cost against information density: overly small intervals are redundant, while overly large intervals may miss critical cues and increase sensitivity to inference-time distribution shifts.The selected interval reflects minimal changes between adjacent frames and visual persistence.Source: D. Ablation Studies · p. 7
V. CONCLUSION
PAM supports a 300-frame history window and reports strong performance on tasks with diverse contextual cues and long horizons. Its design also supports efficient parallel sampling and extension to larger models.
V. CONCLUSION
- PAM combines a 300-frame history window with strong performance on tasks requiring diverse contextual cues and on long-horizon tasks.The conclusion also describes PAM as versatile.Source: §V · pp. 7–8
- A compatible training scheme and model architectures support efficient parallel sampling and an affordable long history window while maintaining high inference speed.Source: §V · p. 7
- PAM’s simple architecture and design principles facilitate extension to larger models.Source: §V · p. 8