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

Motivation

• Reward-maximizing post-training can collapse policies to one high-reward mode, yielding similar images or dominant action strategies and weakening VLA task and scene generalization.Reference KL and transition entropy do not consistently retain distinct complete outputs or strategies; proxy rewards can rise while perceptual quality or diversity declines.Source: §1 · p. 1

Unified RL post-training framework

• Uni-TMPO uses groupwise forward KL to match reward-induced and policy-induced distributions over complete trajectories, applying one objective to diffusion and flow policies in T2I and VLA.The paper frames this as a shared approach because both domains assign probabilities to complete trajectories through stochastic diffusion or flow transitions.Source: §1 · pp. 1–2

Progress-conditioned coarse-to-fine T2I sampling

• Progress-conditioned coarse-to-fine T2I sampling explores global structures early and refines local details later, achieving higher reward and diversity than TreeGRPO while reducing iteration time by 10.2%.The comparison uses the same trajectory-matching objective and settings.Source: §1 · p. 2

Feedback-conditioned action-chunk VLA sampling

• Feedback-conditioned VLA sampling builds action-chunk trajectories under changing observations and a shared language instruction, using transition log probabilities to provide trajectory probabilities.Executing a chunk prefix yields the observation for the next policy query.Source: §1 · p. 2

Extensive evaluation

• Uni-TMPO achieves the best T2I reward–diversity–efficiency trade-off and exceeds the strongest reward-maximization baselines on VLA ID results and task and scene OOD generalization.Its average gains are 8.0 percentage points on task OOD generalization and 4.8 points on scene OOD generalization; real-robot evaluation shows multiple strategies support completion when the higher-reward target is unavailable.Source: §1 · p. 2

2.1 T2I Post-Training and Generative Diversity

• T2I post-training methods improve rewards but often narrow high-reward image modes; proxy rewards can rise as perceptual quality and diversity decline, while diversity methods may require predefined reference distributions.Uni-TMPO instead derives its groupwise target from scalar rewards already used in conventional RL post-training.Source: §2.1 · p. 3

2.2 VLA Post-Training and OOD Generalization

• VLA reinforcement learning can improve task success while reducing a policy to one behavior, and losing alternatives is especially harmful when tasks or scenes change.Prior work motivates evaluation on held-out tasks, scene transfer, and controlled execution changes.Source: §2.2 · p. 3

2.3 Trajectory, Flow, and Distribution Matching

• Unlike reward optimization, constrained improvement, or diffusion-chain flow fitting, Uni-TMPO directly matches complete-trajectory probabilities to terminal-reward targets using exact groupwise forward KL.It shares the objective across T2I denoising and VLA action-chunk trajectories while adapting trajectory construction to each domain.Source: §2.3 · p. 3

3.1 Stochastic Diffusion and Flow Trajectories

• T2I trajectories record a complete denoising rollout from Gaussian image noise, while VLA trajectories combine repeated observation-conditioned action-chunk queries with executed prefixes.Both domains provide a context, complete trajectory, outcome reward, and log probabilities of stochastic policy transitions; only stochastic policy decisions contribute to the differentiable trajectory score.Source: §3.1 · p. 3

3.2 Reward Maximization with Multiple Valid Trajectories

• Expected-reward maximization does not distinguish assigning all probability to one maximally rewarded trajectory from distributing probability across several equally rewarded trajectories.Distinct high-reward trajectories can represent different image modes or physical strategies, while relative advantages do not specify a target distribution over them.Source: §3.2 · p. 4

3.3 Boltzmann Reward Distributions

• Boltzmann weighting converts scalar rewards into a trajectory distribution whose inverse temperature β controls preference strength, retaining several useful trajectories at small or moderate values and favoring top rewards at large values.The target is constructed within each sampled group of trajectories sharing a context.Source: §3.3 · p. 4

4 UNIFIED TRAJECTORY-MATCHING FRAMEWORK

• Uni-TMPO matches a reward-defined target distribution to trajectory probabilities using a groupwise forward-KL objective for stochastic diffusion and flow policies.T2I and VLA use domain-specific trajectory sampling and score computation.Source: §4 · p. 4

4.1 Trajectory Groups

• Each group shares a context: a prompt in T2I or an instruction and initialization identifier in VLA, with trajectories scored from recorded stochastic transition log densities.A trajectory is a complete image denoising process for T2I; VLA trajectories store the physical rollout and stochastic action-flow decisions, while environment transitions are excluded from scores.Source: §4.1 · p. 4

4.2 Reward-Induced Target Distribution

• Standardized rewards form a groupwise Boltzmann target in which higher-reward trajectories receive more probability, with inverse temperature β controlling preference strength.Standardization removes within-group reward offsets and reduces sensitivity to evaluator scales.Source: §4.2 · p. 4

4.3 Groupwise Trajectory Probabilities

• Normalizing trajectory scores within each group yields relative policy probabilities; comparing them with the target reveals high-reward trajectories assigned too little probability.Shared terms across all trajectories cancel, while terms shared by only a subset still affect relative probabilities.Source: §4.3 · p. 4

4.4 Forward-KL Distribution Matching

• Groups with constant rewards provide no relative preference and are excluded; forward-KL updates move each sampled trajectory toward its target probability and propagate corrections through recorded stochastic decisions.Under-target trajectories gain relative probability, while over-target trajectories lose it.Source: §4.4 · pp. 4–5

4.5 On-Policy Update

• The on-policy update uses fresh trajectory groups, detached rewards, and valid groups; VLA gradients pass through the trainable action expert while the vision-language backbone stays frozen.A reference-policy KL, when used, remains separate from the trajectory-matching loss.Source: §4.5 · p. 5

5.1 Progress-Conditioned Coarse-to-Fine T2I Sampling

• A complete T2I image rollout records stochastic transitions, whose log densities are summed to score the trajectory.Leaves have equal numbers of stochastic transitions, so their scores need no length normalization.Source: §5.1 · p. 5
• Shared prefixes avoid repeating common denoising steps, while a progress-conditioned scheduler branches earlier for global exploration and later for appearance refinement.The scheduler changes which complete trajectories enter the group, not the matching loss.Source: §5.1 · pp. 5–6

5.2 Feedback-Conditioned Action-Chunk VLA Sampling

• VLA rollouts generate action chunks from current observations, execute a prefix, and use the resulting observation to condition the next chunk.Trajectory scores aggregate stochastic action-flow transition log probabilities across policy queries.Source: §5.2 · p. 6

5.3 Shared Objective Across Domains

• The same groupwise forward-KL objective aligns policy and reward-induced target distributions for both T2I and VLA trajectories.Source: §5.3 · p. 8

• The experiments evaluate the same trajectory-matching objective under both T2I and VLA trajectory structures.Source: §6 · p. 8
• The T2I study measures reward, diversity, and training efficiency.Source: §6 · p. 8
• The VLA study evaluates in-distribution performance and task and scene out-of-distribution generalization.Source: §6 · p. 8

6.1.1 Experimental Setup

• The evaluation uses FLUX.1-dev with LoRA across GenEval, OCR accuracy, and PickScore; diversity is measured by LGMD and DINOv2 cosine diversity on matched prompt groups.Baselines include distribution matching, reward maximization, and diversity regularization; comparisons share prompts, seeds, group size, decoding steps, and image counts.Source: §6.1.1 · p. 8, §6.1.2 · p. 8

6.1.3 Main Results

• Uni-TMPO directly matches complete trajectories to terminal-reward targets and achieves higher primary metrics with 42–49% lower iteration time than DGFS in the three-transition FLUX setting.DAG-DB and DGFS-SubTB retain positive LGMD, supporting distribution matching as a way to preserve visual modes.Source: §6.1.3 · pp. 8–9
• Across three protocols, Uni-TMPO provides the best reward–diversity–efficiency trade-off, reaching 0.954 GenEval and 0.951 OCR accuracy while improving both diversity measures.It obtains the highest ImageReward and PickScore and the second-highest HPS in preference optimization.Source: §6.1.3 · pp. 8–9
• The scheduler comparison holds the Uni-TMPO objective and other settings fixed while comparing the coarse-to-fine scheduler with fixed positions and TreeGRPO’s scheduler.Source: §6.1.3 · p. 9
• Uni-TMPO is best or tied for best in all six GenEval categories, spanning object composition, counting, color, position, and attribute binding.Source: §6.1.3 · p. 9
• Under matched prompts, Uni-TMPO varies appearance, aesthetics, and spatial layout while Flow-GRPO produces more similar samples; the variations remain consistent with the prompt.Source: §6.1.3 · p. 9

6.1.4 Multi-Reward Evaluation

• With OCR and PickScore weights scheduled from 3:1 to 1:3, Uni-TMPO reaches 0.932 OCR accuracy and 23.897 PickScore, versus Flow-GRPO’s 0.928 and 23.081.Uni-TMPO also records LGMD 0.155 and cosine diversity 0.248, versus −0.049 and 0.209 for Flow-GRPO; fixed 1:1 weighting gives the same ordering.Source: §6.1.4 · p. 9

6.2.1 Training and Evaluation Setup

• The study uses SFT-initialized OpenPI policies and compares Uni-TMPO with πRL Flow-SDE and Flow-Noise variants under matched environments, resets, interaction budgets, and horizons.Evaluation covers LIBERO, MetaWorld-MT50, CALVIN, task and scene OOD transfer, and alternative strategies.Source: §6.2.1 · p. 10, §6.2.2 · p. 10

6.2.2 In-Distribution Performance

• With π0, Uni-TMPO reaches 97.8% on LIBERO, 88.6% on MT50, and 63.9% on CALVIN-D, exceeding the strongest πRL baseline by 0.2, 2.8, and 2.2 points.With π0.5, results are 98.6%, 72.8%, and 89.2%, gains of 0.3, 2.1, and 2.2 points.Source: §6.2.2 · pp. 10–11

6.2.3 Out-of-Distribution Generalization

• On MetaWorld ML45, Uni-TMPO reaches 69.9% with π0 and 59.3% with π0.5, exceeding the corresponding πRL baselines by 4.4 and 11.5 points; on CALVIN ABC→D, it obtains 61.9% and 82.1% Len-5 success.On CALVIN, π0 improves 4.4 points over SFT and finishes 6.6 points higher; π0.5 improves 20.8 points over SFT and exceeds flow baselines by 3.0 points.Source: §6.2.3 · p. 11
• In WallDetour, Uni-TMPO retains LEFT/RIGHT route mass of 0.55/0.45 and route entropy 0.99; after LEFT is blocked, it completes via RIGHT while Reward Maximization fails.Source: §6.2.3 · p. 11
• In real-robot placement, Uni-TMPO succeeds at both targets and completes through the right-target strategy after the higher-reward left target is blocked.Source: §6.2.3 · p. 11

7 CONCLUSION

• Uni-TMPO matches reward-induced and policy-induced probabilities over trajectories, sharing a groupwise forward-KL update across T2I and VLA despite their different sampling procedures.The conclusion reports the best T2I reward–diversity–efficiency trade-off, best VLA ID performance, and task and scene OOD generalization.Source: §7 · p. 11
• Future work may extend RL post-training to video, audio, and 3D generation, molecular and protein design, and diffusion-based planning.Source: §7 · p. 11