FlowMAP: Flow Matching for Generalizable Agent Planning

Agent planning faces dynamic heterogeneity—nonstationary observations, dynamics, and objectives with sparse, delayed rewards—which dominant methods largely ignore, leading to poor generalization under environment shifts. We propose Flow-Matching for Agent Planning (FlowMAP), which formulates planning as a continuous-time flow-matching problem by learning a planning-time velocity field that transports an initial meta-state distribution toward a task-conditioned target. FlowMAP introduces Value-Transport Flow Matching to provide a distribution-level planning objective that steers transport toward high-value regions in the meta-state distribution, mitigating error accumulation under environmental shifts. To enforce alignment between meta-state distribution transport and action--environment interaction, FlowMAP further proposes Flow--Policy Co-Training, which jointly optimizes the planning flow and policy so that the flow transport directly regularizes the policy-induced meta-distribution dynamics. Across diverse agent planning benchmarks, FlowMAP consistently outperforms strong baselines, yielding improvements in planning generalization.

AI agents are increasingly expected to make decisions in complex environments, such as games, simulators, robots, or interactive digital tasks. A major challenge is that these environments often change: the layout may vary, objects may move, rewards may be delayed, and the same action may lead to different outcomes in different situations. Existing planning methods often focus on predicting one future path or improving actions step by step, which can become unreliable when the environment changes. This paper introduces FlowMAP, a new way to help agents plan more robustly. Instead of only asking which single action or trajectory is best, FlowMAP looks at how the agent’s internal understanding of the environment changes over time. It learns to guide this internal state toward more promising regions that are likely to lead to better long-term outcomes. Experiments on several challenging benchmarks show that this approach improves planning and generalization compared with strong existing methods. More broadly, the work suggests that future AI agents may become more reliable by controlling how their internal beliefs evolve, rather than relying only on short-term action updates or predicted rollouts.