Generative Control as Optimization: Time Unconditional Flow Matching for Adaptive and Robust Robotic Control

Zunzhe Zhang; Runhan Huang; Yicheng Liu; Shaoting Zhu; Linzhan Mou; Hang Zhao

cs.ROMar 18, 2026

Generative Control as Optimization: Time Unconditional Flow Matching for Adaptive and Robust Robotic Control

Authors: Zunzhe Zhang, Runhan Huang, Yicheng Liu, Shaoting Zhu, Linzhan Mou, Hang Zhao

Organizations: Institute for Interdisciplinary Information Sciences, Tsinghua University, Beijing, China · Galaxea Inc., Beijing, China · Department of Computer Science, Princeton University, New Jersey, US

Abstract

Diffusion models and flow matching have become a cornerstone of robotic imitation learning, yet they suffer from a structural inefficiency where inference is often bound to a fixed integration schedule that is agnostic to state complexity. This paradigm forces the policy to expend the same computational budget on trivial motions as it does on complex tasks. We introduce Generative Control as Optimization (GeCO), a time-unconditional framework that transforms action synthesis from trajectory integration into iterative optimization. GeCO learns a stationary velocity field in the action-sequence space where expert behaviors form stable attractors. Consequently, test-time inference becomes an adaptive process that allocates computation based on convergence--exiting early for simple states while refining longer for difficult ones. Furthermore, this stationary geometry yields an intrinsic, training-free safety signal, as the field norm at the optimized action serves as a robust out-of-distribution (OOD) detector, remaining low for in-distribution states while significantly increasing for anomalies. We validate GeCO on standard simulation benchmarks and demonstrate seamless scaling to pi0-series Vision-Language-Action (VLA) models. As a plug-and-play replacement for standard flow-matching heads, GeCO improves success rates and efficiency with an optimization-native mechanism for safe deployment. Video and code can be found at https://hrh6666.github.io/GeCO/

Figures & tables

Figure 1: Generative Control as Optimization (GeCO). (A) The Paradigm Shift: Unlike standard flow matching which relies on rigid, time-dependent integration schedules (top), GeCO learns a stationary velocity field where inference becomes an iterative optimization process toward stable attractors (bottom). (B) Adaptive Computation: This formulation enables the policy to dynamically allocate computational budget based on state complexity—exiting early for simple transit phases (Scenario 1) while performing deep refinement for precise manipulation (Scenario 2). (C) Intrinsic Safety: The stationary geometry provides a zero-shot safety mechanism. In-distribution (ID) states converge to low-energy equilibria ( $||f_{\theta}||\approx 0$ ), whereas out-of-distribution (OOD) anomalies exhibit persistently high field norms ( $||f_{\theta}||\gg 0$ ), enabling robust detection.

Panel of Figure 2 — Figure 2: Computation Follows Task Complexity. We visualize the spatial distribution of inference effort along a single rollout. The first three panels (a–c) are sampled from LIBERO-Spatial, and the last three panels (d–f) are from LIBERO-Object. The color of each line encodes the number of function evaluations (NFE) required for convergence at that state, ranging from blue (NFE = 1) to red (NFE = 20). This visualization illustrates how GeCO allocates more computation to challenging states while using fewer steps in easier regions.

Method	Max	Success Rate (%)					Efficiency
Method	Steps	Goal	Spat.	Obj.	Long	Avg.	NFE
Diff. Policy	100 (Fixed)	82.9	91.4	88.9	82.7	86.5	100.0
Rectified flow	20 (Fixed)	92.4	94.6	97.0	76.0	90.0	20.0
GeCO (Ours)	5	91.6	95.4	98.2	82.4	91.9	5.0
GeCO (Ours)	10	93.0	95.8	99.0	81.8	92.4	8.7
GeCO (Ours)	20	95.2	95.8	99.0	83.8	93.5	11.6

Table 1: Performance vs. Efficiency on LIBERO. GeCO achieves higher success rates by allowing adaptive refinement. By increasing the max step budget, GeCO refines actions for complex states, yet the average compute remains low compared to fixed-schedule baselines.

Method	Adjust Bottle		Beat Block Hammer		Blocks Ranking (Size)		Click Alarmclock		Click Bell		Avg.
Method	Easy	Hard	Easy	Hard	Easy	Hard	Easy	Hard	Easy	Hard	Easy	Hard
ACT ( Zhao et al., 2023 )	97.0	23.0	56.0	3.0	0.0	0.0	32.0	4.0	58.0	3.0	49.0	7.0
$\pi_{0}$ ( Black et al., 2024 )	90.0	56.0	43.0	21.0	7.0	1.0	63.0	11.0	44.0	3.0	49.0	18.0
$\pi_{0}$ with GeCO	86.0	66.0	52.0	36.0	0.0	0.0	72.0	16.0	42.0	21.0	50.0	28.0

Table 2: Success rate (%) of different VLA methods on the RoboTwin 2.0 benchmark. Each task is evaluated under two difficulty levels: Easy and Hard . For every method we report the success rate for each task and difficulty level, together with the average performance across all five tasks. Higher values indicate better task completion ability. The best result in each column is highlighted in bold .

Method	Track 1	Track 2	Track 3	Track 4	Track 6	Avg.
$\pi_{0}$	47.0	21.2	29.1	17.3	32.2	29.4
$\pi_{0.5}$	40.6	22.6	18.0	16.1	25.6	24.5
$\pi_{0}$ with GeCO	61.0	22.0	34.0	20.0	45.0	36.0

Table 3: VLA success rates (%) on VLABench tracks. We report track-wise success and the average over the listed tracks; higher is better.

Method	Goal	Spat.	Obj.	Long	Avg.
$\pi_{0}$ + FAST ( Pertsch et al., 2025 )	88.6	96.4	96.8	60.2	85.5
$\pi_{0}$ ( Black et al., 2024 )	95.8	96.8	98.8	85.2	94.2
$\pi_{0.5}$ ( Intelligence et al., 2025 )	98.0	98.8	98.2	92.4	96.9
Ours ( $\pi_{0}$ )	96.4	96.8	97.0	85.4	93.9
Ours ( $\pi_{0.5}$ )	96.4	97.4	98.6	92.4	96.2

Table 4: VLA success rates (%) on LIBERO benchmark. We report track-wise success and the average over the listed tracks; higher is better.

Method	AUROC ( $\uparrow$ )	TNR(%)	TPR (%)	Time Saved(%)
Baseline (Moving Avg)	0.53	14.4	27.8	10.3
Ours (Moving Avg)	0.93	82.4	89.7	42.4
Ours (Leaky Bucket)	0.93	84.0	90.0	44.3

Table 5: OOD Detection on LIBERO Suite Shift (Goal

\rightarrow

Spatial). We report planning-level AUROC for separating ID (Goal) and OOD (Spatial) using optimization-dynamics signals computed per planning call. At a fixed operating point, TPR is specificity on ID instances and TNR is recall on OOD instances. Time Saved is the average fraction of interaction time avoided on OOD episodes by early reporting, computed as

1-t_{\mathrm{report}}/T_{\mathrm{total}}

$\alpha$	Object	Goal	Spatial	Long	Avg
$\alpha=0.0$	0.984	0.946	0.934	0.770	0.909
$\alpha=0.2$	0.968	0.954	0.958	0.802	0.921
$\alpha=0.4$	0.982	0.934	0.956	0.782	0.914
$\alpha=0.6$	0.982	0.942	0.956	0.778	0.915
$\alpha=0.8$	0.978	0.940	0.960	0.810	0.922
$\alpha=1.0$	0.932	0.840	0.898	0.558	0.807

Table 6: Ablation Study on Truncation Parameter

\alpha

. We report the performance on the LIBERO benchmark across different values of

\alpha

Figure 3: GeCO policy execution for the Nut Assembly task. The robot performs high-precision alignment and rotational insertion.

Figure 4: GeCO policy execution for the Chemistry Tube Arrangement task. The policy adaptively handles the tight-tolerance insertion of fragile tubes.

Method	Nut Assembly			Tube Arrangement
	Success ( $\uparrow$ )	Avg. NFE ( $\downarrow$ )	Avg. Infer. Time ( $\downarrow$ )	Success ( $\uparrow$ )	Avg. NFE ( $\downarrow$ )	Avg. Infer. Time ( $\downarrow$ )
G0 Plus	30%	10.0	252.3	25%	10.0	254.6
G0 Plus + GeCO (Ours)	70%	5.06	172.4	80%	3.82	150.3

Table 7: Real-world performance on Galaxea R1 Lite. We report the success rate, average Number of Function Evaluations (NFE), and average inference time for each task.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Task	Track 1		Track 2		Track 3		Track 4		Track 6
Task	$\pi_{0}$	Ours	$\pi_{0}$	Ours	$\pi_{0}$	Ours	$\pi_{0}$	Ours	$\pi_{0}$	Ours
add_condiment	0.66	0.80	0.14	0.08	0.34	0.16	0.26	0.10	0.56	0.76
insert_flower	0.18	0.30	0.04	0.00	0.22	0.28	0.02	0.00	0.10	0.02
select_book	0.694	0.84	0.064	0.08	0.417	0.51	0.311	0.23	0.714	0.79
select_chemistry_tube	0.52	0.76	0.12	0.14	0.70	0.94	0.06	0.02	0.28	0.71
select_drink	0.52	0.78	0.224	0.30	0.08	0.10	0.10	0.04	0.44	0.66

Appendix

Table 8: Per-task success rates on VLABench primitive 10-task evaluation. We report success rates for the baseline

\pi_{0}

and

\pi_{0}

with GeCO (ours) across five tracks; higher is better. Track averages correspond to the averages over the 10 tasks in each track.

Figure 5: Task setups for the real-world robotic deployment, showing the configurations for both the Nut Assembly and the Chemistry Tube Arrangement tasks.

Panel of Figure 6 — Figure 6: Adaptive computation during real-world execution. We visualize the per-step Number of Function Evaluations (NFE) and total inference time over time for representative physical rollouts of Nut Assembly and Tube Arrangement . For Nut Assembly , GeCO uses an average of 3.50 NFEs and 166.7 ms total inference time per decision. For Tube Arrangement , GeCO uses an average of 3.73 NFEs and 147.0 ms per decision. Rather than relying on a fixed computation budget, GeCO automatically increases computation during bottleneck phases, where the NFE rises to 9–10, and exits early during easier free-space motions, where the NFE drops to as low as 2–3. The corresponding inference-time traces mirror this behavior, showing that GeCO adaptively allocates computation according to task difficulty and thereby improves real-world inference efficiency.

$\alpha$	Success Rate (%)
$\alpha$	Goal	Spat.	Obj.	Long	Avg.
$0.6$	85.8	98.0	91.6	48.6	81.0
$0.8$	96.4	97.4	98.6	92.4	96.2
$0.9$	95.8	97.8	94.0	81.2	92.2

Appendix

Table 9: Ablation Study on Truncation Parameter

\alpha

(VLA Architecture). We report the success rates (%) on the LIBERO benchmark across different values of

\alpha

to demonstrate its impact on policy performance.

$\tau$	LIBERO Object	LIBERO Goal	LIBERO Spatial	LIBERO Long	Avg	Avg. Steps
$0.4$	99.0	95.2	95.8	83.8	93.5	11.6
$0.6$	98.8	94.0	95.6	81.0	92.35	7.7
$0.8$	96.8	93.2	96.0	81.6	91.9	6.6

Appendix

Table 10: Ablation on the Early-Stop Threshold

\tau

. We report the success rates (%) and the average number of inference steps across the LIBERO benchmark.

Step-size schedule	LIBERO Object	LIBERO Goal	LIBERO Spatial	LIBERO Long	Avg
$0.015\times 20$	97.0	92.0	94.8	77.2	90.3
$0.02\times 20$	98.0	94.8	95.4	84.0	93.1
$0.025\times 20$	98.4	95.0	94.2	80.4	92.0
$0.035\times 20$	98.0	95.0	95.0	77.4	91.4
Current schedule	99.0	95.2	95.8	83.8	93.5

Appendix

Table 11: Ablation on Fixed Step-Size Schedules. We compare our current adaptive schedule with several fixed step-size schedules (constant step size

\times

total steps) on the LIBERO benchmark.

Explore similar work

CoLA-Flow Policy: Temporally Coherent Imitation Learning via Continuous Latent Action Flow Matching for Robotic Manipulation
Jan 30, 2026Wu Songwei, Jiang Zhiduo, Sun Wandong +4Flow Matching Latent Action Learning
Harbin Institute of Technology, Harbin 150001, China · Honor Device Co., Ltd., Shenzhen, China
KoopmanFlow: Spectrally Decoupled Generative Control Policy via Koopman Structural Bias
Mar 14, 2026Chengsi Yao, Ge Wang, Kai Kang +9Robotic Manipulation Flow-Based Generative Modeling
ISing AI · Future Network of Intelligence Institution, The Chinese University of Hong Kong, Shenzhen · Guangdong Polytechnic Normal University +1
One-Step Flow Policy: Self-Distillation for Fast Visuomotor Policies
Mar 12, 2026Shaolong Li, Lichao Sun, Yongchao ChenVisuomotor Policy Learning One-Step Generative Modeling
CSE Department, University of Michigan · Lehigh University · College of AI, Tsinghua University

Jan 30, 2026J/K move · Enter open · S save

CoLA-Flow Policy: Temporally Coherent Imitation Learning via Continuous Latent Action Flow Matching for Robotic Manipulation

Wu Songwei, Jiang Zhiduo, Sun Wandong, Xie Guanghu+3

Harbin Institute of Technology, Harbin 150001, China · Honor Device Co., Ltd., Shenzhen, China

Flow Matching Latent Action Learning Long-Horizon Robotic Manipulation Robot Imitation Learning

Open paper AskPDF arXiv

Learning long-horizon robotic manipulation requires jointly achieving expressive behavior modeling, real-time inference, and stable execution, which remains challenging for existing generative policies. Diffusion-based approaches offer strong modeling capacity but incur high inference latency, while flow matching enables fast, near-single-step generation yet often suffers from unstable execution when operating directly in the raw action space. We propose Continuous Latent Action Flow Policy (CoLA-Flow Policy), a trajectory-level imitation learning framework that performs flow matching in a continuous latent action space. By encoding action sequences into temporally coherent latent trajectories and learning an explicit latent-space flow, CoLA-Flow Policy decouples global motion structure from low-level control noise, enabling smooth and reliable long-horizon execution. The framework further integrates geometry-aware point cloud conditioning and execution-time multimodal modulation, using visual cues as a representative modality to enhance real-world robustness. Experiments in simulation and on real robots show that CoLA-Flow Policy achieves near-single-step inference, improves trajectory smoothness by up to 93.7% and task success by up to 25 percentage points over raw action-space flow baselines, while remaining significantly faster than diffusion-based policies.

Abstract

Explore similar work

CoLA-Flow Policy: Temporally Coherent Imitation Learning via Continuous Latent Action Flow Matching for Robotic Manipulation

KoopmanFlow: Spectrally Decoupled Generative Control Policy via Koopman Structural Bias

One-Step Flow Policy: Self-Distillation for Fast Visuomotor Policies