🤖 machine learning

Reflective Prompted Policy Optimization: Trajectory-Grounded Revision and Salience Bias

This paper introduces Reflective Prompted Policy Optimization (R2PO), a two-stage LLM framework that enhances policy search by leveraging detailed trajectory-level evidence rather than scalar rewards to diagnose failures and guide revisions, while specifically addressing the "salience bias" failure mode to achieve faster, more stable, and near-optimal performance across diverse environments.

Original authors: Rahaf Abu Hara, Vaibbhav Murarri, Claudio Zito

Published 2026-05-12

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Rahaf Abu Hara, Vaibbhav Murarri, Claudio Zito

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Scorecard" vs. The "Story"

Imagine you are teaching a robot to play a video game. In traditional methods, you only get a scorecard at the end of the game.

The Scorecard says: "You got 45 points."
The Robot thinks: "Okay, I need to do something different to get 46."

But the scorecard doesn't tell you why you lost. Did you fall in a hole on step 3? Did you get stuck in a loop? Did you play perfectly for 19 rounds but crash spectacularly on the 20th? Without the story, the robot has to guess, leading to a lot of trial-and-error.

Recent AI methods tried to fix this by using a "Smart Coach" (a Large Language Model) to look at the score and suggest changes. But even this coach was only looking at the scorecard. It knew the score was low, but it didn't know what the robot actually did wrong.

The Solution: R2PO (The Two-Coach System)

The authors propose R2PO, a new system that uses two AI coaches working together to teach the robot. They treat the robot's actual gameplay footage (the "trajectory") as the most important evidence, not just the final score.

Coach 1: The Scout (Search-LLM)

Role: The Scout looks at the history of scores from previous attempts.
Action: It says, "Okay, based on the scores, let's try changing the robot's settings slightly in this direction." It proposes a new set of instructions for the robot.
Analogy: Think of the Scout as a hiker looking at a map. They guess which direction might lead to the summit based on where they've been before.

The Environment: The Test Run

The robot tries out the Scout's new instructions. It plays the game 20 times (called "rollouts"). This generates a lot of data: where it went, what it did, and how it failed.

Coach 2: The Critic (Critic-LLM)

Role: The Critic is the detective. It doesn't just look at the score; it watches the video footage of the robot's 20 attempts.
Action: It analyzes the footage to find the specific mistake. "Ah, I see! In 19 of those runs, the robot was great. But in one run, it fell into a hole because it turned left too early."
The Fix: Instead of guessing, the Critic suggests a tiny, targeted tweak to the robot's instructions to fix that specific problem.

The Safety Net: The Selection Step

Before the Critic's new instructions are saved, a "Referee" checks them.

If the new instructions work better, they are kept.
If the new instructions make things worse (even if the Critic thought they were a good idea), the old instructions are kept.
Analogy: This is like a chef tasting a new soup. If the new spice makes it taste worse, they throw the new spice away and stick with the original recipe.

The Hidden Trap: "Salience Bias"

The researchers discovered a funny but dangerous habit in the Critic Coach. They call it Salience Bias.

Imagine the robot plays 20 games.

19 games: It wins easily.
1 game: It crashes into a wall in a very dramatic, messy way.

If you show the Critic Coach all 20 videos, the Coach gets obsessed with that one dramatic crash. It thinks, "Oh no! The robot is terrible at avoiding walls!" and it changes the robot's brain to fix the crash.

The Result: The robot was actually doing great in the other 19 games. By trying to fix the one weird crash, the Critic accidentally breaks the robot's ability to play the other 19 games. The robot gets worse.

How R2PO Fixes This:
The R2PO system is smart about which video it shows the Critic.

The Median Video: Instead of showing the worst crash or the best win, it shows the "middle" video—the one that represents what the robot usually does.
The Stats Sheet: It gives the Critic a summary sheet: "95% of the time, the robot wins. The crash was a 5% fluke."
The "Don't Touch" Rule: If the robot is already doing really well (high score), the Critic is told: "Don't change anything unless you have a very, very good reason."

This stops the Critic from panicking over a single bad day and ruining a good strategy.

The Results: Faster, Smarter, and More Stable

The paper tested this system on 10 different environments (like balancing a pole, playing Pong, or navigating a maze).

Speed: R2PO learned much faster. For example, on the "CartPole" game (balancing a stick), it reached near-perfect scores in about 500 tries, while other methods needed 4,000 tries.
Stability: Other methods would find a great solution, then accidentally break it in the next step. R2PO found the solution and stayed there.
Efficiency: They used a relatively small, open-source AI model (20 billion parameters) to beat methods that used much larger, expensive, or proprietary models.

Summary

The paper argues that to teach AI effectively, you shouldn't just give it a score. You should show it the story of what happened. However, you have to be careful not to let the AI obsess over the one time it failed (Salience Bias). By showing the AI the "average" story and giving it a safety net, you can teach it to learn faster, fix specific problems, and avoid breaking what already works.

Technical Summary: Reflective Prompted Policy Optimization (R2PO)

Problem Statement

Existing Large Language Model (LLM)-based policy optimizers function as black-box searchers that rely exclusively on scalar reward feedback. While these methods can propose new policy parameters based on historical reward scores, they lack diagnostic insight into why a policy failed. An optimizer sees a score of 0.45 but cannot distinguish between a policy that consistently underperforms, one that gets stuck in a loop, or one that succeeds in 19 out of 20 rollouts but catastrophically fails in the 20th. This limitation forces LLM-based optimizers to rely on trial-and-error tuning rather than targeted correction, often resulting in unstable training and inefficient convergence.

Methodology: R2PO Framework

The authors propose Reflective Prompted Policy Optimization (R2PO), a two-stage LLM framework designed to augment scalar rewards with trajectory-level behavioral evidence. R2PO operates over compact policy classes (linear policies for continuous control and tabular policies for discrete tasks) and consists of the following components:

Search-LLM (Global Optimizer): This stage functions similarly to prior LLM optimizers (e.g., ProPS). Given a history of evaluated policies and their scalar mean rewards, the Search-LLM proposes an initial candidate parameter vector ( $\theta_{init}$ ).
Environment Evaluation: The candidate policy is executed in the environment for $K$ independent rollouts, generating a set of trajectories $T$ containing states, actions, and rewards.
Critic-LLM (Local Reviser): This stage introduces the core innovation. Instead of receiving only the scalar mean reward, the Critic-LLM inspects a "calibrated trajectory evidence package" derived from the $K$ rollouts. It analyzes what went wrong and proposes a targeted parameter revision ( $\theta_{rev}$ ).
Selection Mechanism: Both the original candidate ( $\theta_{init}$ ) and the revised candidate ( $\theta_{rev}$ ) are re-evaluated. The system retains whichever achieves the higher mean reward, filtering out destructive edits or LLM hallucinations.

Calibrated Trajectory Evidence

To prevent the Critic-LLM from overfitting to outliers, R2PO constructs the evidence package using three specific components:

Median Trajectory: A single rollout whose return is closest to the median of the $K$ returns, representing typical behavior rather than best- or worst-case outliers.
Aggregate Rollout Statistics: A structured report summarizing mean, min, and max rewards, episode lengths, and success/failure rates across all $K$ rollouts. This conveys failure frequency without foregrounding a specific failed trace.
Revision Rule: An explicit instruction to the Critic-LLM to abstain from major changes if the initial mean reward exceeds a threshold ( $\tau_c$ ) and statistics do not indicate a consistent failure pattern. This prevents "over-correction" of already strong policies.

Key Findings and Analysis

1. Salience Bias

The paper identifies a dominant failure mode termed salience bias. When presented with multiple trajectories (e.g., best, worst, and median), the Critic-LLM tends to fixate on the worst-case trajectory as representative, even when the majority of rollouts indicate successful behavior. In a three-trajectory variant on the CartPole environment, this bias accounted for 76.6% of performance regressions. The LLM would diagnose a problem based on a single unlucky failure and propose edits that destroyed the policy's performance in the other 19 successful rollouts.

2. Ablation Studies

Systematic ablations across ten environments demonstrate that R2PO's success relies on the combination of three factors, none of which is sufficient alone:

Two-Stage Architecture: Separating global proposal from local, behavior-grounded revision is essential. A single rich prompt (CriticOnly) or a sequential scalar-only search (ActorSecondPass) fails to match R2PO's performance.
Trajectory-Grounded Revision: Using actual rollout evidence (median trajectory + statistics) significantly outperforms scalar-only feedback.
Selection Mechanism: A "keep-best" selection step is critical. Without it (AlwaysCritic), the high variance of trajectory-conditioned edits leads to frequent regressions.

3. Empirical Results

Using a relatively small, open-weight 20B-parameter model (gpt-oss:20b), R2PO was evaluated on ten environments (including CartPole, InvertedPendulum, Pong, and custom Maze/Nim tasks) against strong baselines (ProPS, ProPS+, and best-performing Stable-Baselines3 algorithms like PPO, SAC, TRPO).

Performance: R2PO achieved the highest mean reward on 9 of 10 environments and tied on the 10th (MountainCarContinuous). It reached the maximum possible reward on CartPole (500), InvertedPendulum (1000), Pong (3.0), and Nim (1.0).
Efficiency: R2PO reached near-optimal performance substantially earlier than baselines (e.g., near-maximum CartPole reward within ~500 episodes, compared to ~4,000 for ProPS+ and SB3).
Stability: R2PO maintained strong policies once found, exhibiting a much smaller gap between mean and best rewards compared to prior LLM methods, which often showed high variance and isolated spikes.
Resource Efficiency: Even when compared to ProPS results obtained with double the interaction budget (8,000 episodes) and a larger proprietary model (GPT-4o), R2PO (trained on 4,000 episodes with a 20B open model) matched or exceeded the baselines on 9 of 10 environments.

Significance and Claims

The paper claims that treating trajectories as first-class in-context evidence, rather than reducing them to scalar returns, fundamentally changes how LLMs can search policy spaces. The key contributions are:

Diagnostic Precision: R2PO enables LLMs to learn faster and diagnose failures more precisely by leveraging behavioral evidence, moving beyond trial-and-error tuning.
Mitigation of LLM Biases: The framework successfully mitigates the "salience bias" inherent in LLM reasoning over multiple evidence sources by using aggregate statistics and median trajectories, rather than raw worst-case traces.
Accessibility: The results demonstrate that these benefits are not limited to large, proprietary models; a 20B open-weight model is sufficient to outperform deep RL baselines and prior LLM-based methods when the architecture is correctly designed.

The authors note limitations, specifically that the method currently focuses on compact policy classes (linear and tabular) and requires coarse manual calibration of the revision threshold $\tau_c$ . Future work is proposed to extend R2PO to high-dimensional neural policies and to develop adaptive calibration mechanisms.