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

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.