# Reward-Aware Trajectory Shaping for Few-step Visual Generation

arXiv: https://arxiv.org/abs/2604.14910 · HTML: https://mlsanity.com/paper/2604.14910/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion and multimodal model alignment, one of 10 directions in LLM alignment / RLHF / preference optimisation.
- Builds on 28 of its 55 references that are in the map.
- First posted on arXiv 16 Apr 2026.

## Placed in

- LLM alignment / RLHF / preference optimisation / [Diffusion and multimodal model alignment](https://mlsanity.com/alignment/diffusion-and-multimodal-model-alignment/)

## Builds on (most cited first, 28 of 28)

- [Deep reinforcement learning from human preferences](https://mlsanity.com/paper/1706.03741/) (2017)
- [HunyuanVideo: A Systematic Framework For Large Video Generative Models](https://mlsanity.com/paper/2412.03603/) (2024)
- [Training Diffusion Models with Reinforcement Learning](https://mlsanity.com/paper/2305.13301/) (2023)
- [Pick-a-Pic: An Open Dataset of User Preferences for Text-to-Image Generation](https://mlsanity.com/paper/2305.01569/) (2023)
- [Flow-GRPO: Training Flow Matching Models via Online RL](https://mlsanity.com/paper/2505.05470/) (2025)
- [Sequence Level Training with Recurrent Neural Networks](https://mlsanity.com/paper/1511.06732/) (2015)
- [Consistency Trajectory Models: Learning Probability Flow ODE Trajectory of Diffusion](https://mlsanity.com/paper/2310.02279/) (2023)
- [DanceGRPO: Unleashing GRPO on Visual Generation](https://mlsanity.com/paper/2505.07818/) (2025)
- [Knowledge Distillation in Iterative Generative Models for Improved Sampling Speed](https://mlsanity.com/paper/2101.02388/) (2021)
- [Improving Video Generation with Human Feedback](https://mlsanity.com/paper/2501.13918/) (2025)
- [Simplifying, Stabilizing and Scaling Continuous-Time Consistency Models](https://mlsanity.com/paper/2410.11081/) (2024)
- [Motion-I2V: Consistent and Controllable Image-to-Video Generation with Explicit Motion Modeling](https://mlsanity.com/paper/2401.15977/) (2024)
- [Hyper-SD: Trajectory Segmented Consistency Model for Efficient Image Synthesis](https://mlsanity.com/paper/2404.13686/) (2024)
- [VisionReward: Fine-Grained Multi-Dimensional Human Preference Learning for Image and Video Generation](https://mlsanity.com/paper/2412.21059/) (2024)
- [Video Diffusion Alignment via Reward Gradients](https://mlsanity.com/paper/2407.08737/) (2024)
- [TempFlow-GRPO: When Timing Matters for GRPO in Flow Models](https://mlsanity.com/paper/2508.04324/) (2025)
- [InstructVideo: Instructing Video Diffusion Models with Human Feedback](https://mlsanity.com/paper/2312.12490/) (2023)
- [Flash Diffusion: Accelerating Any Conditional Diffusion Model for Few Steps Image Generation](https://mlsanity.com/paper/2406.02347/) (2024)
- [Goodhart's Law in Reinforcement Learning](https://mlsanity.com/paper/2310.09144/) (2023)
- [SenseFlow: Scaling Distribution Matching for Flow-based Text-to-Image Distillation](https://mlsanity.com/paper/2506.00523/) (2025)
- [Fine-Grained Preference Optimization Improves Spatial Reasoning in VLMs](https://mlsanity.com/paper/2506.21656/) (2025)
- [RayFlow: Instance-Aware Diffusion Acceleration via Adaptive Flow Trajectories](https://mlsanity.com/paper/2503.07699/) (2025)
- [Integrating Reinforcement Learning with Visual Generative Models: Foundations and Advances](https://mlsanity.com/paper/2508.10316/) (2025)
- [Generative Video Compression: Towards 0.01% Compression Rate for Video Transmission](https://mlsanity.com/paper/2512.24300/) (2025)
- [VAST 1.0: A Unified Framework for Controllable and Consistent Video Generation](https://mlsanity.com/paper/2412.16677/) (2024)
- [TeleBoost: A Systematic Alignment Framework for High-Fidelity, Controllable, and Robust Video Generation](https://mlsanity.com/paper/2602.07595/) (2026)
- [Rethinking Reward Signals in Video GRPO: When Scores Become Targets](https://mlsanity.com/paper/2511.19356/) (2025)
- [Learning What to Trust: Bayesian Prior-Guided Optimization for Visual Generation](https://mlsanity.com/paper/2511.18919/) (2025)

## Built on by (most cited first, 2 of 2)

- [Physics-Grounded Fluid Video Generation with a Simulation Dataset and Dual-Stream Optical-Flow Supervision](https://mlsanity.com/paper/2607.25321/) (2026)
- [CatPortrait: Taming Video Generative Models for Portrait Animation via In-context Conditioning](https://mlsanity.com/paper/2610.14957/) (2026)

## Abstract (arXiv)

Achieving high-fidelity generation in extremely few sampling steps has long been a central goal of generative modeling. Existing approaches largely rely on distillation-based frameworks to compress the original multi-step denoising process into a few-step generator. However, such methods inherently constrain the student to imitate a stronger multi-step teacher, imposing the teacher as an upper bound on student performance. We argue that introducing preference alignment awareness enables the student to optimize toward reward-preferred generation quality, potentially surpassing the teacher instead of being restricted to rigid teacher imitation. To this end, we propose Reward-Aware Trajectory Shaping (RATS), a lightweight framework for preference-aligned few-step generation. Specifically, teacher and student latent trajectories are aligned at key denoising stages through horizon matching, while a reward-aware gate is introduced to adaptively regulate teacher guidance based on their relative reward performance. Trajectory shaping is strengthened when the teacher achieves higher rewards, and relaxed when the student matches or surpasses the teacher, thereby enabling continued reward-driven improvement. By seamlessly integrating trajectory distillation, reward-aware gating, and preference alignment, RATS effectively transfers preference-relevant knowledge from high-step generators without incurring additional test-time computational overhead. Experimental results demonstrate that RATS substantially improves the efficiency--quality trade-off in few-step visual generation, significantly narrowing the gap between few-step students and stronger multi-step generators.
