Prl Process Reward Learning

---
name: prl-process-reward-learning
title: "PRL: Process Reward Learning Improves LLMs' Reasoning Ability and Broadens the Reasoning Boundary"
version: 0.0.2
engine: skillxiv-v0.0.2-claude-opus-4.6
license: MIT
url: "https://arxiv.org/abs/2601.10201"
keywords: [reasoning, process-rewards, RL, trajectory-supervision, entropy-regularization]
description: "Improves LLM reasoning by decomposing RL objectives into intermediate process rewards assigned to reasoning steps, improving both final accuracy and reasoning capacity without expensive Monte Carlo Tree Search."
---

## Overview

Enhance LLM reasoning capabilities by providing reward signals at intermediate steps of the reasoning process rather than only at the end. This approach decomposes entropy-regularized RL objectives into granular process-level rewards that guide step-by-step improvement.

## When to Use

- For complex reasoning tasks (math, logic, science, coding)
- When you want to improve pass@k performance
- For multi-step problem-solving where intermediate correctness matters
- When computational efficiency is important (vs. MCTS or expensive sampling)

## When NOT to Use

- For single-step tasks without intermediate reasoning
- When intermediate ground truth is unavailable
- For real-time applications where reward computation adds latency
- For tasks where final-outcome-only feedback is sufficient

## Key Technical Components

### Process Reward Decomposition

Break down the global RL objective into step-level rewards.

```python
# Process reward decomposition
class ProcessRewardLearner:
    def __init__(self, reference_model):
        self.reference_model = reference_model
        self.process_rewards = {}

    def decompose_objective(self, trajectory, final_reward):
        """Decompose end-reward into process rewards"""
        # Original objective: max R(trajectory) - KL(policy || reference)
        # Decomposed: sum of process rewards for each step

        trajectory_length = len(trajectory)
        base_reward = final_reward / trajectory_length

        process_rewards = []
        for step_idx, step in enumerate(trajectory):
            # Assign step-level reward
            step_reward = self.compute_step_reward(
                step,
                step_idx,
                trajectory,
                final_reward
            )
            process_rewards.append(step_reward)

        return process_rewards

    def compute_step_reward(self, step, idx, full_trajectory, final_reward):
        """Compute reward for individual reasoning step"""
        # Reward based on:
        # 1. Step correctness (if available)
        # 2. Progress toward solution
        # 3. KL penalty vs reference model

        step_logprob = step["log_probability"]
        reference_logprob = self.reference_model.get_logprob(step["text"])

        kl_penalty = step_logprob - reference_logprob
        progress_bonus = self.estimate_progress(idx, full_trajectory)

        return progress_bonus - 0.1 * kl_penalty
```

### Step-Level Correctness Annotation

When available, use ground truth to label intermediate steps.

```python
# Step correctness annotation
def annotate_step_correctness(trajectory, problem, solution_steps):
    """Mark which reasoning steps are correct"""
    annotations = []
    for i, step in enumerate(trajectory):
        if i < len(solution_steps):
            is_correct = matches_solution_step(step, solution_steps[i])
        else:
            is_correct = None  # Unknown

        annotations.append({
            "step_index": i,
            "text": step["text"],
            "is_correct": is_correct,
            "confidence": compute_confidence(step, solution_steps[i])
        })
    return annotations
```

### Entropy Regularization Integration

Incorporate KL divergence penalty to maintain exploration.

```python
# Entropy-regularized process rewards
class EntropyRegularizedPRL:
    def __init__(self, beta=0.1, reference_model=None):
        self.beta = beta  # Entropy regularization coefficient
        self.reference_model = reference_model

    def compute_regularized_reward(self, step_text, reference_logprob):
        """Add entropy regularization to step reward"""
        # KL(policy || reference) = E[log(p) - log(q)]
        policy_logprob = self.policy_model(step_text)
        kl_divergence = policy_logprob - reference_logprob

        # Entropy regularization penalizes KL
        regularized_reward = -self.beta * kl_divergence

        return regularized_reward

    def batch_compute_process_rewards(self, trajectories):
        """Compute process rewards for batch of trajectories"""
        all_rewards = []
        for trajectory in trajectories:
            rewards = []
            for step in trajectory:
                ref_logprob = self.reference_model(step)
                reward = self.compute_regularized_reward(step, ref_logprob)
                rewards.append(reward)
            all_rewards.append(rewards)
        return all_rewards
```

### Advantage Estimation

Compute advantages for policy gradient updates.

```python
# Advantage estimation from process rewards
def compute_advantages(process_rewards, discount_factor=0.99):
    """Convert process rewards to advantages for PG update"""
    advantages = []
    cumulative_return = 0

    # Reverse traversal for discount computation
    for reward in reversed(process_rewards):
        cumulative_return = reward + discount_factor * cumulative_return
        advantages.insert(0, cumulative_return)

    # Normalize for stability
    advantages = (advantages - np.mean(advantages)) / (np.std(advantages) + 1e-8)

    return advantages
```

### Training Loop Integration

Incorporate process rewards into standard policy gradient training.

```python
# PRL training loop
class PRLTrainer:
    def __init__(self, policy_model, reference_model, beta=0.1):
        self.policy = policy_model
        self.reference = reference_model
        self.beta = beta

    def train_step(self, problem, trajectories, final_rewards):
        """Single training step using process rewards"""
        losses = []

        for trajectory, final_reward in zip(trajectories, final_rewards):
            # 1. Decompose end-reward into process rewards
            process_rewards = self.decompose_objective(trajectory, final_reward)

            # 2. Add entropy regularization
            regularized_rewards = self.add_entropy_penalty(
                trajectory,
                process_rewards
            )

            # 3. Compute advantages
            advantages = compute_advantages(regularized_rewards)

            # 4. Policy gradient loss
            for step, advantage in zip(trajectory, advantages):
                step_logprob = self.policy(step)
                loss = -step_logprob * advantage  # PG loss

                losses.append(loss)

        # Optimization step
        total_loss = sum(losses) / len(losses)
        self.policy.backward(total_loss)
        self.policy.optimize()

        return total_loss.item()

    def decompose_objective(self, trajectory, final_reward):
        """Decompose final reward to process level"""
        return ProcessRewardLearner().decompose_objective(trajectory, final_reward)

    def add_entropy_penalty(self, trajectory, process_rewards):
        """Add entropy regularization to rewards"""
        regularized = []
        for step, reward in zip(trajectory, process_rewards):
            ref_logprob = self.reference(step)
            kl_penalty = self.compute_kl(step, ref_logprob)
            regularized.append(reward - self.beta * kl_penalty)
        return regularized

    def compute_kl(self, step, ref_logprob):
        """Compute KL divergence for step"""
        policy_logprob = self.policy(step)
        return policy_logprob - ref_logprob
```

### Performance Measurement

Track improvements in both pass@1 and pass@k metrics.

```python
# Performance tracking
class PerformanceTracker:
    def __init__(self):
        self.pass_at_1 = []
        self.pass_at_k = []
        self.reasoning_breadth = []  # Coverage of solution approaches

    def evaluate(self, model, test_set, k=5):
        """Measure reasoning improvements"""
        # Pass@1: single attempt success
        pass_1 = sum(
            1 for problem in test_set
            if model.solve(problem, attempts=1)
        ) / len(test_set)

        # Pass@k: success within k attempts
        pass_k = sum(
            1 for problem in test_set
            if any(model.solve(problem, attempts=1) for _ in range(k))
        ) / len(test_set)

        # Reasoning breadth: diversity of approaches explored
        breadth = self.measure_approach_diversity(model, test_set)

        self.pass_at_1.append(pass_1)
        self.pass_at_k.append(pass_k)
        self.reasoning_breadth.append(breadth)

        return {"pass@1": pass_1, "pass@k": pass_k, "breadth": breadth}

    def measure_approach_diversity(self, model, test_set):
        """Count distinct reasoning approaches discovered"""
        approaches = set()
        for problem in test_set:
            for attempt in range(5):
                trajectory = model.generate_trajectory(problem)
                approach = self.classify_approach(trajectory)
                approaches.add(approach)
        return len(approaches)
```

## Performance Characteristics

- Pass@1 improvement: +2-5% vs baseline
- Pass@k improvement: Broader reasoning capability expansion
- Computational cost: No expensive MCTS required
- Training efficiency: Faster convergence than outcome-only RL

## Key Advantages

1. **Granular feedback**: Every step receives guidance
2. **Efficiency**: Avoids computationally expensive tree search
3. **Interpretability**: Track which steps improve reasoning
4. **Stability**: Process rewards provide denser learning signal

## References

- Process-level rewards enable fine-grained policy guidance
- Entropy regularization maintains exploration without collapse
- Equivalent to reward maximization + KL penalty decomposed to steps
- More efficient than outcome-only RL for reasoning tasks