Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

The paper tackles state ambiguity in robotic manipulation by relaxing the Markov assumption through a long, adaptively recoded working-memory mechanism. PAM combines a causal frame feature extractor, a context router with history-length queries, and a flow-matching action head, trained in two stages to achieve temporal disambiguation with a 300-frame window at 20 Hz. Empirical results show strong real-world performance across seven tasks and robust long-horizon generalization on Libero-Long, outperforming baselines and offering interpretable attention maps. This approach provides a practical, scalable path to robust manipulation under temporal uncertainty, with potential for extension to larger models and richer memory representations.

Abstract

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/

Resolving State Ambiguity in Robot Manipulation via Adaptive Working Memory Recoding

TL;DR

The paper tackles state ambiguity in robotic manipulation by relaxing the Markov assumption through a long, adaptively recoded working-memory mechanism. PAM combines a causal frame feature extractor, a context router with history-length queries, and a flow-matching action head, trained in two stages to achieve temporal disambiguation with a 300-frame window at 20 Hz. Empirical results show strong real-world performance across seven tasks and robust long-horizon generalization on Libero-Long, outperforming baselines and offering interpretable attention maps. This approach provides a practical, scalable path to robust manipulation under temporal uncertainty, with potential for extension to larger models and richer memory representations.

Abstract

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/
Paper Structure (14 sections, 5 equations, 4 figures, 3 tables)

This paper contains 14 sections, 5 equations, 4 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. History-Aware Visuomotor Policy
  4. Temporally Dependent Inference
  5. Method
  6. Overview of PAM
  7. Model Architecture
  8. Training Recipe
  9. Experiments
  10. Real-World Experimental Setting
  11. Main Results in Real-World Tasks
  12. Evaluation on Libero-Long Long-Horizon Tasks
  13. Ablation Studies
  14. Conclusion

Figures (4)

  • Figure 1: Left: We illustrate a representative example of state ambiguity and elaborate on its various scenarios. State ambiguity is common in robotic manipulation, and the relevant state transitions do not follow the Markov assumption; therefore, a long history window is essential. Right: Existing methods and our method for encoding long history. Analogous to human continuous reasoning, PAM’s inference process is temporally dependent, extending the history window by maintaining a long-term adaptive working memory. The model only needs to encode the current observation at each inference step.
  • Figure 2: Overview of PAM. PAM uses two distinct types of features to guide action generation: motion primitives extracted from the current frame and compact context features drawn from the extended history window. The context features serve as working memory. (a) illustrates the frame feature extractor used to obtain these features. The extractor employs a query-based mechanism to recode multimodal inputs, enabling adaptive working memory recoding. (b) illustrates the context router, which receives the the context features and utilizes a set of query tokens spanning different history lengths to produce compact context features. PAM is trained in a two-stage manner. As shown in the figure, different subsets of model parameters are progressively activated across the two stages.
  • Figure 3: Real-world Tasks. We carefully designed a set of real-world tasks as shown in the figure, with the primary types of state ambiguity indicated in the top-right corner. Each task comprises multiple subtasks, and to provide a more fine-grained evaluation, the task success rate is computed as the average completion rate across its subtasks.
  • Figure 4: PAM exhibits strong interpretability. Left: The attention maps of the context router reveal which history portions PAM references to resolve state ambiguity, accurately identifying key frames from preceding task stages in Wipe the Table Twice. Right: The attention maps of extractor indicate which modalities PAM uses for working memory encoding. In Guessing Game, the context query extracts historical cues from visual observation of the block's position, while the motion primitives attends to joint states for posture maintenace. In Wipe the Table Twice, both visual observations and joint states provide effective contextual cues, demonstrating the effectiveness of our adaptive working memory encoding.