Unified Trajectory Matching Policy Optimization: Diverse T2I Generation and VLA Generalization

Zhiyuan Ma; Jiaming Li; Lingzhen Li; Yu Liu; Xuekai Zhu; Dingkang Liang; Kaiyan Zhang; Jianjun Li

cs.CVSep 28, 2026

Unified Trajectory Matching Policy Optimization: Diverse T2I Generation and VLA Generalization

Authors: Zhiyuan Ma, Jiaming Li, Lingzhen Li, Yu Liu, Xuekai Zhu, Dingkang Liang, Kaiyan Zhang, Jianjun Li, +2 more

Organizations: School of Computer Science and Technology, Huazhong University of Science and Technology, Wuhan, China

Abstract

Reward-maximizing reinforcement learning (RL) is widely used to post-train stochastic diffusion and flow policies for text-to-image (T2I) generation. However, reward-maximizing RL causes policy mode collapse even under reference KL or entropy regularization, reducing the policy to a single high-reward mode. In T2I, this produces similar images and reward hacking. When extended to vision-language-action (VLA) models, the same collapse removes alternative successful strategies and weakens task and scene generalization. To address this limitation, we introduce Unified Trajectory Matching Policy Optimization (Uni-TMPO), a unified RL post-training framework for diffusion and flow policies. First, Uni-TMPO converts standardized rewards into a target distribution within each trajectory group and derives the policy distribution from trajectory log probabilities. Then, forward Kullback-Leibler optimization matches the two distributions instead of maximizing expected reward. A progress-conditioned coarse-to-fine scheduler efficiently constructs T2I trajectories. Within the unified framework, feedback-conditioned sampling uses updated observations to construct VLA trajectories. Extensive experiments show that Uni-TMPO achieves higher T2I rewards and VLA ID success rates than the strongest baselines. More importantly, it achieves the best T2I reward-diversity-efficiency trade-off and VLA generalization to held-out tasks and scenes, while real-robot evaluation demonstrates the value of multiple action strategies when the higher-reward target is blocked.

Figures & tables

Fig. 1: Overview of Uni-TMPO for T2I generation and VLA manipulation. Progress-conditioned coarse-to-fine sampling constructs T2I trajectory groups for each prompt, while feedback-conditioned action-chunk sampling constructs VLA trajectory groups from a shared initialization. Unlike scalar reward maximization, Uni-TMPO matches reward-induced and policy-induced allocations within each group, preserving diverse image modes and action strategies. The right panels summarize T2I performance and VLA OOD generalization.

Fig. 2: Five-mode distribution matching with a three-layer MLP. The left panels show the five-peak reward landscape and shared initialization. The remaining panels show intermediate and final samples and final mode mass under Uni-TMPO (top) and Reward Maximization (bottom). Uni-TMPO matches the target and preserves all five modes, whereas Reward Maximization collapses to the highest-reward mode.

Method	Time (s) $\downarrow$	GenEval $\uparrow$	OCR $\uparrow$	PickScore $\uparrow$	HPS $\uparrow$	ImgRwd $\uparrow$	LGMD $\uparrow$	Cos.Div. $\uparrow$
Compositional Image Generation (GenEval)
FLUX.1-dev	–	0.647	–	22.301	0.301	1.099	$-0.031$	0.211
DAG-DB	187.5	0.889	–	21.998	0.291	1.071	0.097	0.237
DGFS-SubTB	178.6	0.917	–	22.210	0.298	1.107	0.113	0.241
Flow-GRPO	160.8	0.946	–	22.113	0.289	1.074	$-0.089$	0.198
TreeGRPO	126.2	0.936	–	21.524	0.281	1.083	$-0.281$	0.184

TABLE I: Comparison of FLUX.1-dev T2I post-training across compositional image generation, visual text rendering, and human preference alignment. The task-specific training reward is shown in parentheses after each task. Best and second-best results are bolded and underlined.

Fig. 3: T2I reward–diversity–efficiency trade-off. Metrics are normalized across the post-training methods within each task. Larger values indicate better results on every axis; the Time axis is reversed because shorter iteration time is better. Uni-TMPO provides the strongest overall trade-off across compositional generation, visual text rendering, and human preference alignment.

Fig. 4: T2I diversity under matched prompts. Each panel compares three stochastic samples per method under the same prompt. Compared with Flow-GRPO, Uni-TMPO produces broader variation in object appearance, aesthetic style, and spatial layout. Table I reports quantitative results.

Backbone	Method	LIBERO					MetaWorld-MT50	CALVIN-D
		Spatial	Object	Goal	Long	Avg.	Avg.	Len-5
$\pi_{0}$	SFT	65.3	64.4	49.8	51.2	57.6	50.8	57.5
	Flow-SDE	98.4	99.4	96.2	90.2	96.1	78.1	61.7
	Flow-Noise	99.0	99.2	98.2	93.8	97.6	85.8	59.9
	Uni-TMPO (ours)	99.2	99.6	98.8	93.6	97.8	88.6	63.9
$\pi_{0.5}$	SFT	84.6	95.4	84.6	43.9	77.1	43.8	61.3

TABLE II: In-distribution VLA performance of

\pi_{0}

and

\pi_{0.5}

. Baseline configurations follow

\pi_{\mathrm{RL}}

[ 4 ] . LIBERO and MetaWorld-MT50 report average task success, while CALVIN-D reports five-subtask sequence completion (Len-5). All values are percentages. Bold values indicate the best result for each backbone.

Fig. 5: Task and scene OOD generalization after VLA post-training. MetaWorld ML45 uses 45 tasks for online RL and evaluates five unseen tasks. CALVIN uses Scenes ABC for supervised initialization and online RL, then evaluates five-subtask completion (Len-5) in Scene D over 1,000 sequences. The lower panels compare SFT and post-RL performance for $\pi_{0}$ and $\pi_{0.5}$ .

Fig. 6: WallDetour experiment. LEFT is shorter than RIGHT, although both reach the target under the same $0.9\times\mathrm{Success}+0.1\times\mathrm{PathEfficiency}$ reward. Reward Maximization concentrates successful rollouts on LEFT ( $0.95/0.05$ ), whereas Uni-TMPO retains both routes ( $0.55/0.45$ ), with normalized binary route entropies of $0.29$ and $0.99$ . When LEFT is blocked at evaluation, only Uni-TMPO reaches the target via RIGHT.

Fig. 7: Real-robot dual-target placement. The robot is instructed to place the red block on either yellow target. The left target receives slightly higher reward because it is closer. In the unblocked setting, both Reward Maximization rollouts select the left target, whereas Uni-TMPO succeeds at both targets. After the left target is blocked, Reward Maximization fails, while Uni-TMPO succeeds through the alternative right-target strategy.

Explore similar work

TMPO: Trajectory Matching Policy Optimization for Diverse and Efficient Diffusion Alignment
May 9, 2026Jiaming Li, Chenyu Zhu, Nanxi Yi +9Diffusion Alignment Diffusion-Based Reinforcement Learning Methods
1MAIR Lab, Huazhong University of Science and Technology · 2Kuaishou Technology · 3Nanyang Technological University +2
STAR: SpatioTemporal Adaptive Reward Allocation for Text-to-Image RL Post-Training
Jun 16, 2026Jinjie Shen, Wei Deng, Xian Hu +2Reinforcement Learning Post-Training
Finite Difference Flow Optimization for RL Post-Training of Text-to-Image Models
Mar 13, 2026David McAllister, Miika Aittala, Tero Karras +4Text-To-Image Diffusion Models Reinforcement Learning Post-Training
UC Berkeley · NVIDIA

May 9, 2026, cs.LGJ/K move · Enter open · S save

TMPO: Trajectory Matching Policy Optimization for Diverse and Efficient Diffusion Alignment

Jiaming Li, Chenyu Zhu, Nanxi Yi, Youjun Bao+8

1MAIR Lab, Huazhong University of Science and Technology · 2Kuaishou Technology · 3Nanyang Technological University+2

Diffusion Alignment Diffusion-Based Reinforcement Learning Methods Reward Functions Diversity Boltzmann Distributions

Open paper AskPDF arXiv

Reinforcement learning (RL) has shown extraordinary potential in aligning diffusion models to downstream tasks, yet most of them still suffer from significant reward hacking, which degrades generative diversity and quality by inducing visual mode collapse and amplifying unreliable rewards. We identify the root cause as the mode-seeking nature of these methods, which maximize expected reward without effectively constraining probability distribution over acceptable trajectories, causing concentration on a few high-reward paths. In contrast, we propose Trajectory Matching Policy Optimization (TMPO), which replaces scalar reward maximization with trajectory-level reward distribution matching. Specifically, TMPO introduces a Softmax Trajectory Balance (Softmax-TB) objective to match the policy probabilities of K trajectories to a reward-induced Boltzmann distribution. We prove that this objective inherits the mode-covering property of forward KL divergence, preserving coverage over all acceptable trajectories while optimizing reward. To further reduce multi-trajectory training time on large-scale flow-matching models, TMPO incorporates Dynamic Stochastic Tree Sampling, where trajectories share denoising prefixes and branch at dynamically scheduled steps, reducing redundant computation while improving training effectiveness. Extensive results across diverse alignment tasks such as human preference, compositional generation and text rendering show that TMPO improves generative diversity over state-of-the-art methods by 9.1%, and achieves competitive performance in all downstream and efficiency metrics, attaining the optimal trade-off between reward and diversity.

Abstract

Explore similar work

TMPO: Trajectory Matching Policy Optimization for Diverse and Efficient Diffusion Alignment

STAR: SpatioTemporal Adaptive Reward Allocation for Text-to-Image RL Post-Training

Finite Difference Flow Optimization for RL Post-Training of Text-to-Image Models