Learning Robust Visuomotor Policies via Dynamics Alignment

Abstract

Generative BC policies model expert action distributions well, but they do not explicitly reason about what scene a candidate action would produce, so they cannot tell whether it reaches a task-relevant state. We address this by coupling action generation with an action-conditioned dynamics model that predicts the future observations induced by actions. However, naively placing such a dynamics model inside a generative policy creates a circular dependency. Future observations are needed to guide action generation, yet those observations can only be predicted from the action samples being generated. We introduce Dynamics-Aligned Flow Matching Policy (DAP), where a flow matching policy and an action-conditioned flow matching dynamics model co-sample actions and future observations through flow extrapolation. At each step the two flows correct each other, and the samples converge to a self-consistent action and future pair on which the policy and dynamics model agree. Across four simulated tasks and three real-world tasks, DAP achieves the highest average success rate, with relative improvements of 56.5% in simulation and 33.3% in the real world over the strongest baseline in each setting. Moreover, when the future-observation estimate is corrupted at inference time, DAP recovers a self-consistent action-future pair through iterative refinement, whereas video-conditioned and post-hoc dynamics baselines, which use a fixed future surrogate, degrade or fail.