Learning Robust Visuomotor Policies via Dynamics Alignment
Authors: Dohyeok Lee, Jung Min Lee, Munkyung Kim, Seokhun Ju, Jin Woo Koo, Kyungjae Lee, Dohyeong Kim, Taehyun Cho, Tapomayukh Bhattacharjee, Jungwoo Lee
Organizations: Department of Electrical and Computer Engineering, Seoul National University, Seoul 08826, Korea · Department of Statistics, Korea University, Seoul 02841, Korea
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.
Figures & tables
Figure 2: Training and inference pipeline for DAP. (1) During training, DAP learns an action-conditioned dynamics flow and a future-conditioned policy flow. (2) During inference, the two frozen flows exchange extrapolated action and future-observation estimates, iteratively co-refining mutually aligned samples.
Figure 3: Benchmark simulation and real-world manipulation tasks used in our experiments.
Tasks
PickCube-v1
PushCube-v1
DrawTriangle-v1
StackCube-v1
Average
# of Demos
100
100
100
100
–
DP
78/52.0 ±12.3
100/82.3 ±9.3
12/6.0 ±5.6
2/0.0 ±0.0
48/35.1 ±2.7
LPB
47/46.0 ±0.0
100/ 100.0 ±0.0
22/18.7 ±3.1
12/8.3 ±1.2
45/43.3 ±0.5
UVA
7/4.3 ±0.6
78/75.3 ±0.6
45/44.0 ±0.0
3/0.7 ±0.6
33/31.1 ±0.1
π0
42/34.0 ±1.0
100/98.3 ±0.6
10/2.0 ±1.0
56 / 41.3 ±3.1
52/43.9 ±0.9
DAP
94 / 88.7 ±0.6
100/98.3 ±0.6
57 / 55.7 ±2.3
40/32.3 ±4.0
73 / 68.8 ±1.6
Table 1: Success rates (%) on ManiSkill3 under the 100-demonstration setting, reported as best checkpoint / mean over the last three checkpoints ± standard deviation (gray). Best is bolded .
Figure 4: Robustness, prediction-quality, and data-efficiency analyses on PickCube-v1. Error bars denote checkpoint standard deviations in (a,b) and validation-set confidence intervals in (c). No error bars in (d).
Tasks
DAP
DAP (w/o Dyn.)
PickCube-v1
88.7 ±0.6
73.7 ±1.5
DrawTriangle-v1
55.7 ±2.3
22.3 ±1.2
StackCube-v1
32.3 ±4.0
15.0 ±1.0
Table 2: Simulation dynamics ablation. Success rates (%), mean over the last three checkpoints ± standard deviation.
Figure 5: Alignment error R(τ)(o^) of DAP over flow time on the four simulation tasks, with 95% bootstrap confidence intervals over 200 held-out states.
# of Demos
100
200
500
Dynamics
+4.75±0.09
+5.09±0.09
+6.29±0.10
Video
−0.06±0.10
+0.01±0.10
−0.00±0.10
Table 3: ΔPSNR (dB) on PickCube-v1, PSNR under the true action minus mean PSNR under 10 random actions, ± 95% bootstrap confidence interval over 1600 test samples.
Tasks
Pick and Place
Towel Folding
Rim-Grasp Cup Placement
Average
# of Demos
59
59
99
–
DP
80
70
0
50.0
UVA
0
30
0
10.0
π0
0
20
0
6.7
DAP
90
80
30
66.7
DAP (w/o Dyn.)
40
40
10
30.0
Table 4: Real-world task success rates (%) and dynamics-ablation results.
Appendix figures & tables11 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 6: Top: Quantitative evaluation of extrapolated samples during flow generation. The peak signal-to-noise ratio (PSNR) values of the extrapolated future observation at each flow timestep τ are provided (Cam. 1 Ext. PSNR and Cam. 2 Ext. PSNR) with PSNR of final prediction (Cam. 1 PSNR and Cam. 2 PSNR). Likewise, we provide the mean squared error (MSE) values of extrapolated action for each τ (Act. Ext. MSE) with MSE of final prediction (Act. MSE). Evaluation is done on the entire validation dataset. Bottom: Qualitative samples of extrapolated future observation for each flow timestep τ .
PP
TF
RCP
Average
Exp.
Rand.
Exp.
Rand.
Exp.
Rand.
Exp.
Rand.
DAP
24.12
21.63
24.67
22.76
25.94
24.12
24.91
22.84
UVA
24.57
18.13
26.03
20.11
26.51
19.06
25.70
19.10
Appendix
Table 5: Per-task PSNR comparison between DAP and UVA on real-world data, reported separately on expert trajectories (Exp.) and random trajectories (Rand.) for Pick and Place (PP), Towel Folding (TF), Rim-Grasp Cup Placement (RCP), and their average.
Figure 7: Mean future-observation PSNR and success rate of the PickCube-v1 DAP policy as a function of classifier-free guidance (CFG) scale. Left: dynamics action-guidance sweep. Right: policy future-observation-guidance sweep.
Figure 8: Rotation-based task-difficulty summary for Pick and Place, Towel Folding, and Rim-Grasp Cup Placement. Left: per-episode cumulative end-effector rotation shown as violin plots with jittered episode samples and mean markers. Cumulative rotation is computed as the sum of geodesic angle changes between consecutive end-effector orientations. Right: density-normalized histogram of absolute end-effector rotation angles aggregated over all timesteps.
Task
Description
Success criterion
Action dim
Max steps
Camera views
PickCube-v1
Grasp-and-place task that moves a cube to a target goal position.
The cube position is within goal_thresh of the goal position (default Euclidean threshold: 0.025m).
7
300
base + hand
PushCube-v1
Non-prehensile pushing task that moves a cube to a goal region in front of the robot.
The cube’s xy position is within goal_radius of the target xy position (default Euclidean threshold: 0.1m), and the cube remains on the table.
4
300
base + hand
DrawTriangle-v1
Drawing task that traces a goal triangle on a canvas using a stick.
The drawn points are within a Euclidean distance of 0.05m from points on the goal triangle.
3
300
base
StackCube-v1
Pick up a red cube, stack it on top of a green cube, and release it without the cube falling. Both cubes have randomized z-axis rotations and non-colliding xy positions on the table.
The red cube is on top of the green cube within half a cube size, is static, and is not being grasped by the robot.
7
300
base + hand
Appendix
Table 8: Simulation task definitions, success criteria, evaluation horizons, and task interfaces. All methods are evaluated with the same maximum episode steps for each task.
Figure 9: The dynamics model f(ot,at) (left) learns from mixed expert and random data using a flow matching VAE-DiT architecture with cross-attention conditioning. The policy π(at∣ot,ot+1) generates actions through flow matching with a cross-attention U-Net architecture. Flow extrapolation (blue arrows) enables mutual correction: it maps flow samples at timestep τ to endpoint estimates for ot+1 and at , aligning action generation with dynamics predictions. This iterative parallel generation ensures consistent and robust behaviors while maintaining real-time performance.
PickCube-v1
PushCube-v1
DrawTriangle-v1
StackCube-v1
Hyperparameters
batch_size
64
policy_lr_peak
1.00e-04
policy_betas
(0.9, 0.95)
policy_weight_decay
0.02 (selective)
dynamics_lr_peak
1.00e-04
Appendix
Table 9: Hyperparameters and evaluation settings used for DAP on PickCube-v1, PushCube-v1, DrawTriangle-v1, and StackCube-v1. All values are shared across the four tasks except the past-action conditioning pathway.
PickCube-v1
PushCube-v1
DrawTriangle-v1
StackCube-v1
Hyperparameters
batch_size
256
lr_peak
1.00e-04
betas
(0.95, 0.999)
weight_decay
1.00e-06
lr_scheduler
cosine
Appendix
Table 10: Hyperparameters and evaluation settings used for the ManiSkill DP baseline on PickCube-v1, PushCube-v1, DrawTriangle-v1, and StackCube-v1. All values are shared across the four tasks.
PickCube-v1
PushCube-v1
DrawTriangle-v1
StackCube-v1
Hyperparameters
policy batch_size
64
policy num_epochs
500
dynamics batch_size
96
dynamics num_epochs
100
horizon
16
Appendix
Table 11: Hyperparameters and evaluation settings used for the LPB baseline on PickCube-v1, PushCube-v1, DrawTriangle-v1, and StackCube-v1. Values are shared across tasks except for the calibrated threshold and resulting LPB fire rate.
PickCube-v1
PushCube-v1
DrawTriangle-v1
StackCube-v1
Hyperparameters (shared across stages unless noted)
optimizer
AdamW
betas
(0.9, 0.95)
weight_decay
0.02
lr_peak
1.00e-04
lr_scheduler
cosine
Appendix
Table 12: Hyperparameters and evaluation settings used for UVA on PickCube-v1, PushCube-v1, DrawTriangle-v1, and StackCube-v1. All values are shared across the four tasks. Stage 1 (video pretraining) is followed by Stage 2 (joint video+action fine-tuning from the Stage 1 checkpoint).
PickCube-v1
PushCube-v1
DrawTriangle-v1
StackCube-v1
Hyperparameters
batch_size
16
8
lr_peak
1.00e-05
5.00e-06
warmup_steps
5000
1000
decay_steps
100k
100k
decay_lr
0
0
Appendix
Table 13: Training hyperparameters and evaluation settings used for the π0 baseline on PickCube-v1, PushCube-v1, DrawTriangle-v1, and StackCube-v1. All values are grouped by the PickCube/PushCube and DrawTriangle/StackCube task pairs.
Visuomotor policy learning trains robots to turn camera observations, often with language and proprioception, into actions for manipulation and navigation.
Large vision-language-action (VLA) policies are increasingly trained as conditional generative models over action chunks. Yet deployment produces mixed-quality experience-successful demonstrations, partial completions, recoverable mistakes, and failures-that is difficult to use with standard imitation. Full behavior cloning (BC) imitates failures, filtered BC discards useful sub-trajectories, and offline reinforcement learning adds a large critic. We introduce ForesightFlow, a self-guided flow-matching policy that augments each generated action chunk with a learned success-potential trajectory. The same flow proposes and scores candidate actions, enabling best-of-K inference without an external critic. The key issue is that policy improvement and value calibration require different supervision: advantage weighting should emphasize high-quality actions, but applying the same weights to potential coordinates suppresses failure gradients and creates overconfident scores. We address this with decoupled advantage-weighted flow matching, applying exponentiated advantage weights only to action velocities while training potential velocities uniformly. We further derive a one-step boundary estimator for conditional flow matching, allowing advantage computation with a single stop-gradient forward pass. Across five BEHAVIOR-1K simulation tasks and five real-world bimanual tasks, ForesightFlow improves over imitation baselines, matches the strongest separate-critic baseline in simulation success, improves real-world success, and reduces training compute by 38%. Ablations show that decoupling prevents value hallucination, the one-step estimator preserves candidate-ranking fidelity, and self-guided sampling improves long-horizon execution.
Yunpeng Mei, Jiakai He, Hongjie Cao +12
1Beijing Institute of Technology · 2Tsinghua University · 3The Hong Kong University of Science and Technology (Guangzhou)
Visual signals play a crucial role in policy learning by enabling models to capture object motion and interaction dynamics. Just as humans reason about actions using both past experience and anticipated outcomes, effective policies should integrate past interactions with future predictions. However, existing visuomotor policies typically model either historical context or future dynamics in isolation, lacking a unified temporal representation of interaction dynamics. In this work, we introduce ChronoFlow, a temporally unified representation that captures past, current, and future interaction dynamics through sparse 3D keypoints of both objects and the gripper. Based on this representation, we propose ChronoFlow-Policy, a diffusion-based visuomotor policy that jointly learns ChronoFlow and action sequences through a co-training objective. Experiments on 14 simulated tasks and 5 real-world manipulation tasks demonstrate that ChronoFlow-Policy consistently outperforms strong diffusion-policy baselines and improves robustness in long-horizon and non-Markovian manipulation scenarios. Our project page is available at https://the-kamisato-sii.github.io/ChronoFlow-Policy-project-page/.
Bokai Lin, Yifu Xu, Xinyu Zhan +6
Shanghai Jiao Tong University, China · Shanghai Innovation Institute, China · Noematrix, China
Developing efficient and accurate visuomotor policies poses a central challenge in robotic imitation learning. While recent rectified flow approaches have advanced visuomotor policy learning, they suffer from a key limitation: After iterative distillation, generated actions may deviate from the ground-truth actions corresponding to the current visual observation, leading to accumulated error as the reflow process repeats and unstable task execution. We present Selective Flow Alignment (SeFA), an efficient and accurate visuomotor policy learning framework. SeFA resolves this challenge by a selective flow alignment strategy, which leverages expert demonstrations to selectively correct generated actions and restore consistency with observations, while preserving multimodality. This design introduces a consistency correction mechanism that ensures generated actions remain observation-aligned without sacrificing the efficiency of one-step flow inference. Extensive experiments across both simulated and real-world manipulation tasks show that SeFA Policy surpasses state-of-the-art diffusion-based and flow-based policies, achieving superior accuracy and robustness while reducing inference latency by over 98%. By unifying rectified flow efficiency with observation-consistent action generation, SeFA provides a scalable and dependable solution for real-time visuomotor policy learning. Code is available on https://github.com/RongXueZoe/SeFA.