Resolving State Ambiguity in Robot Manipulation via Adaptive Working Memory Recoding

TL;DR

This paper introduces PAM, an adaptive working memory policy for robotic manipulation that effectively resolves state ambiguity by utilizing a hierarchical memory system, enabling robust performance with a large history window and high inference speed.

Contribution

The paper proposes PAM, a novel visuomotor policy with adaptive working memory, supporting large history windows and efficient ambiguity resolution in robotic tasks.

Findings

Supports a 300-frame history window with high inference speed

Handles multiple state ambiguity scenarios simultaneously

Maintains stable training with a 10-second history window

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/