[ArXiv 2025] RewardFlow: Topology-Aware Reward Propagation for Agentic RL
RewardFlow: Topology-Aware Reward Propagation on State Graphs for Agentic RL with Large Language Models
RewardFlow is a novel, lightweight reinforcement learning framework that improves Large Language Model (LLM) agents by estimating dense, state-level rewards through state graph construction and topology-aware propagation. It achieves SOTA results on benchmarks like ALFWorld and Sokoban without requiring expensive external reward models.
TL;DR
Training LLM agents is notoriously difficult due to sparse rewards: you only know if the agent succeeded at the very end of a 20-step task. RewardFlow solves this by building a "State Graph" from sampled trajectories. It automatically propagates the terminal success signal backward through the graph's topology, creating dense, state-level rewards without any help from human annotators or expensive secondary reward models.
- Achievement: SOTA on ALFWorld, WebShop, and Sokoban.
- Efficiency: Adds <1% training time overhead.
- Core Insight: States closer to success on the graph are inherently more valuable.
The Credit Assignment Crisis in Agentic RL
Most agentic environments (like ALFWorld or Sokoban) provide a binary reward: 1 for success, 0 for everything else. For an agent that took 50 actions to reach a goal, which of those actions were actually responsible for the success?
Traditional RL (PPO, GRPO) often spreads the terminal reward uniformly across all tokens or steps. This is "blunt force" optimization; it rewards lucky accidents and irrelevant moves just as much as pivotal decisions. While Process Reward Models (PRMs) have emerged to score intermediate steps, they usually depend on:
- Human labels: Extremely slow and expensive.
- LLM-as-a-Judge: Prone to hallucinations and computationally heavy.
Methodology: Mining the State Graph
RewardFlow operates on a simple but profound observation: Intermediate states that appear frequently in successful paths are the "bottlenecks" of reasoning.
1. State Graph Construction
The system takes multiple sampled trajectories for a single task and merges identical states into nodes.
- Normalization: It uses embeddings and rule-based checks (e.g., noting if an item is "cleaned") to ensure different textual descriptions of the same semantic state are merged.
- Pruning: It removes "hallucinated" actions that the LLM proposed but the environment rejected.
2. Topology-Aware Propagation
Once the graph is built, RewardFlow uses Inverse BFS (Breadth-First Search). It starts at the "Success" node and flows reward backward.
- A state receives a reward , where is the shortest path to success.
- The Action Reward is then the "Potential Difference": .
This rewards actions that move the agent topologically closer to the goal.
3. Synergistic Advantages
RewardFlow doesn't throw away the global signal. It combines:
- Local Advantage: How much better was this action compared to other actions taken from the same state?
- Global Advantage: How much better was this trajectory compared to the average?
Experimental Results: Dominating the Baselines
RewardFlow was tested against heavyweights like GRPO (the engine behind DeepSeek-R1) and RLOO.
Quantitative Performance
In the visual puzzle task Sokoban, RewardFlow achieved a staggering 62.4% success rate using Qwen2.5-7B-VL, compared to just 34.4% for the next best method (GiGPO).
Robustness & Efficiency
- Generalization: It maintained performance on Out-of-Distribution (OOD) tasks where other RL methods collapsed.
- Speed: As shown in the time breakdown, the graph construction logic is practically "free" compared to the time spent on LLM inference.
Critical Analysis & Takeaways
Why does it work? RewardFlow essentially "rediscovers" the value function of the environment through exploration. By aggregating multiple paths, it creates a more stable, objective baseline for what a "good" state looks like, compared to methods that look at trajectories in isolation.
Limitations: Its efficacy depends on state distinguishability. In environments where the state is "hidden" or purely internal (e.g., pure mathematical chain-of-thought without external environment feedback), building a graph becomes difficult.
Future Outlook: The logical next step is scaling this to open-ended web environments and multi-agent systems, where the state graph could potentially be shared across different agents to accelerate collective learning.
Summary: RewardFlow is a masterclass in "work smarter, not harder." By using basic graph theory to structure existing LLM rollouts, it bypasses the need for massive human annotation and delivers a far more precise training signal.