Action-Informed Belief via Latent Flow Matching for Memory-Intensive Robotic Tasks

Dawei Wang, Xingrui Yu, Zhenglin Wan, Ivor Tsang

Keywords: Flow Matching; Memory for Robotics; Partially Observable Markov Decision Process; Reinforcement Learning; Representation Learning

Abstract: Memory is essential for achieving fully autonomous robots in real-world settings, yet reliably remembering fine-grained information in complex and dynamic environments for robots remains an unsolved challenge. While RNN-based approaches combined with reinforcement learning can handle simple memory problems, they often struggle with memory-intensive robotics tasks where success hinges on subtle, fine-grained distinctions. Belief-based methods that leverage privileged information offer a solution for partially observable Markov decision processes (POMDPs). However, the current method focuses solely on reconstructing the belief representation, while ignoring the strong coupling between state variations and actions in the representation space. As a result, the structure of its latent representation is not well-suited to handling the high-dimensional observations and complex dynamics encountered in robotic environments. Motivated by this gap, we propose an action-informed belief reconstruction method. Specifically, we employ Flow Matching to model inverse dynamics as a structural constraint, explicitly steering the latent belief distribution to align with action requirements. By formulating this auxiliary objective within the latent space, our approach enforces strong causal coupling that better preserves fine-grained information and achieves strong performance on the memory-intensive robotic manipulation benchmarks.

Email Sharing: We authorize the sharing of all author emails with Program Chairs.

Data Release: We authorize the release of our submission and author names to the public in the event of acceptance.

Submission Number: 70