Learning Robust Visuomotor Policies via Dynamics Alignment

Dohyeok Lee; Jung Min Lee; Munkyung Kim; Seokhun Ju; Jin Woo Koo; Kyungjae Lee; Dohyeong Kim; Taehyun Cho

cs.ROOct 31, 2025

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,

Organizations: Department of Electrical and Computer Engineering, Seoul National University, Seoul 08826, Korea · Department of Statistics, Korea University, Seoul 02841, Korea

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.

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 $\pm 12.3$	100/82.3 $\pm 9.3$	12/6.0 $\pm 5.6$	2/0.0 $\pm 0.0$	48/35.1 $\pm 2.7$
LPB	47/46.0 $\pm 0.0$	100/ 100.0 $\pm 0.0$	22/18.7 $\pm 3.1$	12/8.3 $\pm 1.2$	45/43.3 $\pm 0.5$
UVA	7/4.3 $\pm 0.6$	78/75.3 $\pm 0.6$	45/44.0 $\pm 0.0$	3/0.7 $\pm 0.6$	33/31.1 $\pm 0.1$
$\pi_{0}$	42/34.0 $\pm 1.0$	100/98.3 $\pm 0.6$	10/2.0 $\pm 1.0$	56 / 41.3 $\pm 3.1$	52/43.9 $\pm 0.9$
DAP	94 / 88.7 $\pm 0.6$	100/98.3 $\pm 0.6$	57 / 55.7 $\pm 2.3$	40/32.3 $\pm 4.0$	73 / 68.8 $\pm 1.6$

Table 1: Success rates (%) on ManiSkill3 under the 100-demonstration setting, reported as best checkpoint / mean over the last three checkpoints

\pm

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 $\pm 0.6$	73.7 $\pm 1.5$
DrawTriangle-v1	55.7 $\pm 2.3$	22.3 $\pm 1.2$
StackCube-v1	32.3 $\pm 4.0$	15.0 $\pm 1.0$

Table 2: Simulation dynamics ablation. Success rates (%), mean over the last three checkpoints

\pm

standard deviation.

Figure 5: Alignment error

R^{(\tau)}(\hat{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$ $\pm 0.09$	$+5.09$ $\pm 0.09$	$+6.29$ $\pm 0.10$
Video	$-0.06$ $\pm 0.10$	$+0.01$ $\pm 0.10$	$-0.00$ $\pm 0.10$

Table 3:

\Delta_{\text{PSNR}}

(dB) on PickCube-v1, PSNR under the true action minus mean PSNR under 10 random actions,

\pm

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
$\pi_{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 $\tau$ 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 $\tau$ (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 $\tau$ .

	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(o_{t},a_{t})$ (left) learns from mixed expert and random data using a flow matching VAE-DiT architecture with cross-attention conditioning. The policy $\pi(a_{t}|o_{t},o_{t+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 $\tau$ to endpoint estimates for $o_{t+1}$ and $a_{t}$ , 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

\pi_{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

Visuomotor policy learning trains robots to turn camera observations, often with language and proprioception, into actions for manipulation and navigation.

See all 279 papers

Potential-Guided Flow Matching for Vision-Language-Action Policy Improvement
Jun 3, 2026Yunpeng Mei, Jiakai He, Hongjie Cao +12Flow Matching Visuomotor Policy Learning
1Beijing Institute of Technology · 2Tsinghua University · 3The Hong Kong University of Science and Technology (Guangzhou)
ChronoFlow-Policy: Unifying Past-Current-Future Interaction Flow in Visuomotor Policy Learning
Jun 30, 2026Bokai Lin, Yifu Xu, Xinyu Zhan +6Long-Horizon Robotic Manipulation Visuomotor Policy Learning
Shanghai Jiao Tong University, China · Shanghai Innovation Institute, China · Noematrix, China
SeFA-Policy: Fast and Accurate Visuomotor Policy Learning with Selective Flow Alignment
Nov 11, 2025Rong Xue, Jiageng Mao, Mingtong Zhang +1Rectified Flow Visuomotor Policy Learning
University of Southern California

Jun 3, 2026J/K move · Enter open · S save

Potential-Guided Flow Matching for Vision-Language-Action Policy Improvement

Yunpeng Mei, Jiakai He, Hongjie Cao, Chenyu Wang+11

1Beijing Institute of Technology · 2Tsinghua University · 3The Hong Kong University of Science and Technology (Guangzhou)

Flow Matching Visuomotor Policy Learning Vision-Language-Action Models Flow-Based Generative Modeling Flow-Based VLA Models Flow Matching for Action Generation

Open paper AskPDF arXiv

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.

Abstract

Visuomotor Policy Learning

Explore similar work

Potential-Guided Flow Matching for Vision-Language-Action Policy Improvement

ChronoFlow-Policy: Unifying Past-Current-Future Interaction Flow in Visuomotor Policy Learning

SeFA-Policy: Fast and Accurate Visuomotor Policy Learning with Selective Flow Alignment