RLFR: Extending Reinforcement Learning for LLMs with Flow Environment

Abstract: Reinforcement Learning with Verifiable Rewards (RLVR) has recently emerged as a promising framework for improving reasoning abilities in LLMs. However, policy optimized with binary verification prone to overlook potential valuable exploration in reasoning trajectory. In view of heavy annotation cost of golden Process Reward Models (PRMs), recent works attempt using auxiliary signals for reward shaping of process tokens, involving entropy and likelihood collected from logit space. In this work, we offer a novel perspective on shaping RLVR with flow rewards derived from latent space, and propose RLFR, where the flow fields of model latents are constructed from either off-policy high-quality data and on-policy rejection sampling data, and the velocity deviations of policy latents within it are quantified to serve as a reward signal. RLFR first demonstrates that a well-established flow field can be a sound environment for reward signal collection, highlighting the expressive latent space is much underexplored. Moreover, RLFR is able to compress any off-policy expert data as reference for constituting reward signals, and we show that the efficient context dependence compressed within the hidden states are utilized, rather than individual token-level denotation for context comprehending. Experiments on both language and multimodal reasoning benchmarks demonstrate the reliability of flow rewards, and suggesting a promising paradigm for reward shaping with auxiliary signals.

• The paper introduces a novel approach for credit assignment by using velocity deviations in latent space as dense, context-dependent reward signals.
• The methodology integrates off-policy expert data with on-policy rejection sampling, leading to improved reasoning and multimodal performance.
• Experimental results show RLFR outperforms traditional binary verification methods, achieving gains of 1.5% to 5.3% across diverse benchmarks.

RLFR: Extending Reinforcement Learning for LLMs with Flow Environment

Introduction

The RLFR framework introduces a novel approach to reward shaping in reinforcement learning for LLMs, leveraging the latent space via flow environments. Traditional RLVR (Reinforcement Learning with Verifiable Rewards) relies on binary outcome verification, which, while robust against reward hacking, suffers from coarse granularity and limited exploration in reasoning trajectories. RLFR addresses these limitations by constructing flow fields in the latent space, using velocity deviations as dense reward signals. This paradigm enables more nuanced credit assignment and facilitates the integration of off-policy expert data and on-policy rejection sampling for reference distribution construction.

Motivation and Theoretical Foundations

RLVR's binary verification mechanism often penalizes trajectories that contain valuable intermediate reasoning steps but ultimately yield incorrect answers. Process Reward Models (PRMs) offer stepwise rewards but are bottlenecked by annotation costs and misalignment between training and online trajectories. Recent works have explored auxiliary signals from logit space, such as entropy and likelihood, but these are susceptible to self-policy rewarding risks and reward hacking.

RLFR proposes to utilize the latent space, which is shown to be highly expressive and underexplored for reward signal collection. Flow Matching (FM) is employed to learn a velocity field that characterizes the reference distribution of high-quality reasoning trajectories. The reward is then computed as the deviation of policy latents from this flow field, penalizing large deviations and encouraging alignment with expert trajectories.

Figure 1: RLFR leverages flow fields in latent space for reward shaping, overcoming the limitations of binary verification and logit-based auxiliary signals.

The connection between velocity deviation and probability likelihood is formalized, establishing that minimizing velocity deviation maximizes the evidence lower bound (ELBO) of the log-likelihood under the reference distribution. This theoretical insight justifies the use of flow-based rewards as a surrogate for likelihood-based evaluation.

Latent Space Analysis

Empirical analysis demonstrates that the latent representations of tail tokens in reasoning trajectories are significantly more expressive than those of head tokens. These latents encode context-dependent information, enabling more reliable identification of trajectory quality compared to logit space or textual content.

Figure 2: Distribution of trajectory tokens in latent and logit space, highlighting the superior expressiveness of latent space for reward signal extraction.

Further, the latent space exhibits consistent signals across different layer percentiles, reinforcing its suitability for reward design.

Figure 3: Consistent expressiveness of latent space across layer percentiles in Qwen2.5-Base-7B, supporting robust reward signal collection.

RLFR Framework and Implementation

RLFR extends RLVR by shaping the advantage term for each token using flow returns derived from velocity deviations in latent space. The flow network is pre-trained on off-policy expert data and updated online using rejection sampling of high-quality trajectories. The reward for each token is computed as:

$\hat{A}_{k} = \sum_{s=k}^{||}\gamma^{s-k}r^{\bm{v}_{\phi}_{s}} + \hat{A}_{o}$

where $r^{\bm{v_\phi}_{k}}$ is the normalized flow reward, filtered to discard noisy fluctuations. The flow network conditions on subsequent token latents to enhance context dependence.

Algorithmically, RLFR proceeds as follows:

1. Offline Start: Pre-train the flow network on expert data to establish the reference distribution.
2. Online Optimization: Generate rollouts, optimize policy with RL (e.g., GRPO), and update the flow network using rejection-sampled high-quality trajectories.
3. Reward Calculation: For each token, compute velocity deviation in latent space and shape the advantage accordingly.

Experimental Results

RLFR demonstrates consistent improvements over RLVR and entropy-based shaping methods across both language and multimodal reasoning benchmarks. On Qwen2.5-Math-7B, RLFR surpasses RLVR by 1.5% average score; on Llama3.1-8B, the improvement is 5.3%. In multimodal settings, RLFR achieves notable gains on MathVision and MathVerse, and generalizes well to logic benchmarks.

Figure 4: RLFR achieves superior performance on both language and multimodal reasoning benchmarks compared to RLVR and logit-based shaping.

Training dynamics indicate that RLFR yields stable improvements without degeneration, with policy entropy stabilizing at a higher level, indicative of enhanced exploration.

Figure 5: Training logs of RLVR and RLFR on Qwen2.5VL-3B, showing accelerated and stable policy optimization with RLFR.

Ablation Studies

Ablations reveal that both offline flow pretraining and online rejection sampling are critical for optimal performance. Filtering out noisy flow reward fluctuations is essential for stable training. Correctness-based metrics for rejection sampling outperform entropy-based or composite metrics.

Timesteps used for velocity prediction significantly impact reward quality; larger timesteps with less noise are preferred, and debiasing weighting further improves performance.

Reward Behavior Analysis

Case studies show that flow rewards preferentially encourage tokens that execute the question, while suppressing empty or connection tokens. High-entropy tokens in logit space correspond to larger velocity deviations in latent space, reflecting ambiguity in hidden states. Flow rewards rely on context dependence compressed within hidden states, rather than individual token-level denotation.

Figure 6: Case study of reward tokens during training, illustrating the preference for contextually meaningful tokens in flow reward assignment.

Practical Implications and Future Directions

RLFR provides a scalable and annotation-efficient framework for reward shaping in LLM reinforcement learning. By leveraging latent space flow environments, it enables dense, context-dependent reward signals that facilitate exploration and robust policy optimization. The integration of off-policy expert data and online rejection sampling allows for flexible and adaptive reference distribution construction.

Future research may explore scaling the flow environment to larger models and datasets, investigating the potential of latent signals for test-time scaling, and extending the framework to domains beyond mathematical and multimodal reasoning.

Conclusion

RLFR advances the state of reward shaping in LLM reinforcement learning by harnessing the expressiveness of latent space through flow environments. The framework demonstrates reliable and consistent performance gains across diverse benchmarks, substantiated by theoretical analysis and empirical validation. RLFR's approach to leveraging velocity deviations in latent space for reward signal collection represents a promising direction for future research in scalable and robust LLM optimization.