Stabilize RL training on reasoning tasks by performing entropy-guided rollouts from uncertain decision points, avoiding policy collapse and premature convergence. Increases fully correct trajectories on math reasoning while maintaining stable entropy throughout training.
Stabilize RL training on reasoning tasks by performing entropy-guided rollouts from uncertain decision points, avoiding policy collapse and premature convergence. Increases fully correct trajectories on math reasoning while maintaining stable entropy throughout training.
Entropy-Guided Exploration: Structured Uncertainty-Driven Rollouts for Stable Reasoning
Reinforcement learning on mathematical reasoning with verifiable rewards (RLVR) is notoriously unstable—models converge prematurely to shallow tactics, ignore long chains of reasoning, or collapse into reward-gaming strategies. Traditional RL policies optimize based on final trajectory rewards without understanding where the model is uncertain about its reasoning path. FR3E solves this through two-phase structured exploration: first, generate base trajectories and identify high-entropy decision points (where the model is most uncertain); second, launch multiple exploratory rollouts from those uncertain states, gathering intermediate reward signals that guide the policy toward more robust reasoning.
When fine-tuning reasoning models on AIME, competition math, or other verifiable reasoning benchmarks, entropy-guided exploration prevents the catastrophic policy drift common in standard RL. By targeting exploration at moments of genuine uncertainty rather than applying uniform exploration pressure, the model learns to question weak assumptions rather than just generate longer outputs.
Core Concept
FR3E operates in two synchronized phases. The First Return phase generates trajectories and computes token-level entropy across the sequence, identifying semantic decision points where entropy is high (model is uncertain). These positions become anchors for structured exploration. The Entropy-Eliciting Explore phase launches multiple diverse rollouts from each high-entropy state, simulating what would happen if the model had chosen differently at that critical juncture. Intermediate rewards from these partial trajectories inform the policy: if rollouts from a state tend to succeed, that state's decisions were good; if they fail, reconsider. Asymmetric clipping encourages the policy to explore beyond symmetric PPO bounds, and adaptive advantage modulation scales learning signals based on marginal value improvements between consecutive states, preventing both exploration collapse and reward-gaming.
Architecture Overview
Base Trajectory Generator: Produces initial reasoning chains conditioned on problem statements
Token-wise Entropy Computation: Identifies decision points where model uncertainty is highest
Semantic Segmentation: Groups entropy peaks into meaningful reasoning blocks
Rejection Sampling Filter: Removes degenerate rollouts where all branches yield identical rewards
Multi-Rollout Explorer: Generates diverse branches from high-entropy states via temperature sampling
Intermediate Reward Estimator: Assigns partial credit based on continuation success
Asymmetric Policy Gradient: Clip-Higher mechanism encouraging exploration without collapse
Implementation
This example demonstrates entropy computation and identification of critical reasoning decision points in trajectories.
# Entropy-guided exploration for reasoning trajectoriesimport torch
import torch.nn.functional as F
from collections import defaultdict
classEntropyGuidedExplorer:
def__init__(self, model, tokenizer, temperature=0.8):
self.model = model
self.tokenizer = tokenizer
self.temperature = temperature
defgenerate_base_trajectory(self, problem_statement, max_length=512):
"""Generate reasoning trajectory and compute token-level entropy."""# Tokenize problem
input_ids = self.tokenizer(problem_statement, return_tensors='pt')['input_ids']
# Generate trajectory with logits
trajectories = []
entropies = []
token_ids = []
with torch.no_grad():
for _ inrange(max_length):
outputs = self.model(input_ids)
logits = outputs.logits[:, -1, :] # Last token logits# Compute entropy at this position
probs = F.softmax(logits, dim=-1)
entropy = -torch.sum(probs * torch.log(probs + 1e-10), dim=-1)
entropies.append(entropy.item())
# Sample next token
probs_normalized = probs / self.temperature ifself.temperature > 0else probs
next_token = torch.multinomial(probs_normalized, num_samples=1)
token_ids.append(next_token.item())
# Append to input for next iteration
input_ids = torch.cat([input_ids, next_token], dim=-1)
# Check for end-of-sequenceif next_token.item() == self.tokenizer.eos_token_id:
break
trajectory_text = self.tokenizer.decode(token_ids)
return trajectory_text, torch.tensor(entropies)
defidentify_critical_decision_points(self, trajectory_text, entropies, threshold_percentile=75):
"""Identify high-entropy positions as semantic decision points."""# Compute threshold
entropy_threshold = torch.quantile(entropies, threshold_percentile / 100.0)
# Find positions with high entropy
high_entropy_positions = torch.where(entropies > entropy_threshold)[0]
# Group into semantic blocks (consecutive positions form one decision point)
decision_blocks = []
current_block = [high_entropy_positions[0].item()]
for i inrange(1, len(high_entropy_positions)):
pos = high_entropy_positions[i].item()
if pos - current_block[-1] <= 3: # Within 3 tokens
current_block.append(pos)
else:
decision_blocks.append(current_block)
current_block = [pos]
if current_block:
decision_blocks.append(current_block)
# Convert positions to text segments
tokens = self.tokenizer.tokenize(trajectory_text)
decision_points = []
for block in decision_blocks:
start_token = tokens[min(block)]
avg_entropy = torch.mean(entropies[block])
decision_points.append({
'position': block[0],
'segment': start_token,
'entropy': avg_entropy.item(),
'token_range': (block[0], block[-1])
})
return decision_points
This example shows the entropy-eliciting explore phase: launching diverse rollouts from critical decision points.
defexplore_from_decision_point(self, problem_statement, trajectory_prefix, decision_point_pos, num_rollouts=5):
"""Generate diverse rollouts from a critical decision point."""# Encode trajectory up to decision point
prefix_ids = self.tokenizer(
problem_statement + trajectory_prefix[:decision_point_pos],
return_tensors='pt'
)['input_ids']
rollout_trajectories = []
rollout_rewards = []
for _ inrange(num_rollouts):
# Generate diverse continuation with higher temperaturewith torch.no_grad():
outputs = self.model.generate(
prefix_ids,
max_length=512,
temperature=self.temperature * 1.5, # Higher temperature for diversity
top_p=0.95,
do_sample=True
)
rollout_text = self.tokenizer.decode(outputs[0])
rollout_trajectories.append(rollout_text)
return rollout_trajectories
defestimate_intermediate_reward(self, trajectory_continuation, verifier):
"""Estimate reward for partial trajectory using verifier."""try:
# Extract intermediate answers from trajectory
answers = self.extract_intermediate_answers(trajectory_continuation)
ifnot answers:
return0.0# Check last extracted answer
final_answer = answers[-1]
reward = verifier.check_answer(final_answer)
return reward
except:
return0.0# Failed parsing = no rewarddefadaptive_advantage_modulation(self, states, state_values):
"""Scale learning signals based on marginal value improvements."""
advantages = []
for i inrange(len(state_values) - 1):
current_value = state_values[i]
next_value = state_values[i + 1]
# Marginal improvement: how much does value increase at next state?
marginal_improvement = max(0, next_value - current_value)
# Scale advantage by improvement (small improvements get lower weight)
advantage = marginal_improvement * 10.0# Scaling factor
advantages.append(advantage)
return torch.tensor(advantages)
This example demonstrates the complete FR3E training loop with asymmetric clipping and adaptive advantage scaling.
When to use: Apply FR3E when training models on verifiable reasoning tasks—mathematics, logic puzzles, code correctness verification. Use when RL training is unstable and models collapse to shallow strategies. Ideal for problems where intermediate reasoning steps can be verified, not just final answers.
When NOT to use: Skip for tasks without verifiable intermediate steps (creative writing, open-ended reasoning). Avoid if computational budget is severely limited—multi-rollout exploration adds 4-5× overhead. Don't use for simple tasks where standard RL succeeds. Skip if your problem domain has sparse rewards and rarely verifiable intermediate results.
Common pitfalls: Setting entropy threshold too low (high percentile) includes all tokens, destroying decision point specificity. Too high discards useful exploration signals. Not using rejection sampling causes learning from degenerate rollouts where all branches have identical rewards. Forgetting asymmetric clipping negates exploration benefits—symmetric PPO prevents high advantages. Over-scaling with adaptive modulation can amplify noise. Not maintaining entropy bonus during RL causes policy collapse to deterministic tokens. Forgetting to verify that intermediate reward signals actually correlate with final correctness.
Reference
FR3E Team. (2025). First Return, Entropy-Eliciting Explore: Stable Reasoning in LLMs. arXiv preprint arXiv:2507.07017. https://arxiv.org/abs/2507.07017
