Spark Policy Reward Co Evolution

Train Language Models with Co-Evolving Policy and Reward

Outcome

Build a unified LLM training system that eliminates separate reward models by recycling generated rollouts to train policy and reward judgment within a single model, reducing computational overhead while improving reasoning and reward accuracy.

Problem Context

Standard reinforcement learning for language models relies on two separate systems: a policy that generates responses and a reward model that evaluates them. This separation creates inefficiencies:

• Expensive human preference data collection for reward modeling
• Computational redundancy during training and inference (two models instead of one)
• Reward model quality often lags behind policy improvements, creating optimization misalignment
• Training data from policy improvements gets discarded rather than recycled

SPARK addresses these by creating a single model that simultaneously learns to reason better and judge answers more accurately, using the same rollouts for both objectives.

Core Concept

The central insight is co-evolution through data recycling: instead of separate systems, one model learns multiple complementary objectives from the same generated rollouts. Three interrelated training tasks create a feedback loop:

1. Pointwise scoring teaches the model to recognize when individual responses are correct
2. Pairwise comparison develops preference discrimination between response qualities
3. Reflection training enables the model to fix its own mistakes through self-correction

When the reward component improves, it produces better policy gradients. Better policies generate higher-quality rollouts. Better rollouts train more accurate reward judgment. This creates compounding gains without external infrastructure.

Architecture Overview

The training pipeline consists of four sequential phases per iteration:

• Rollout generation: Sample multiple candidate responses from the policy for each input
• On-policy evaluation: Assign advantage scores using verifiable ground truth (exact match, execution success, or reference comparison)
• Data construction: Transform rollouts into three complementary datasets for pointwise, pairwise, and reflection objectives
• Unified optimization: Backpropagate a combined loss that balances policy improvement with reward accuracy and drift regularization

The model architecture remains standard (transformer-based LLM or VLM). Training infrastructure differs only in objective composition, not in fundamental components.

Implementation

Phase 1: Rollout Generation and Evaluation

Generate multiple candidate responses per input and compute standardized advantage scores that will inform all downstream training objectives.

import torch
from torch.utils.data import DataLoader

def generate_rollouts(model, inputs, num_candidates=4, temperature=0.7):
    """Generate multiple candidate responses per input."""
    rollouts = []

    for input_text in inputs:
        candidates = []
        for _ in range(num_candidates):
            # Generate with sampling
            output = model.generate(
                input_text,
                temperature=temperature,
                max_length=256,
                do_sample=True
            )
            candidates.append(output)

        rollouts.append({
            'input': input_text,
            'candidates': candidates,
            'generated_at': 'current_policy'
        })

    return rollouts

def compute_advantages(rollouts, ground_truth_labels, baseline='mean'):
    """Assign advantage scores to candidates using verifiable rewards."""
    processed = []

    for rollout in rollouts:
        input_text = rollout['input']
        candidates = rollout['candidates']
        labels = ground_truth_labels[input_text]

        # Compute individual correctness scores
        scores = []
        for candidate in candidates:
            # Use exact match, execution, or reference-based comparison
            score = evaluate_correctness(candidate, labels)
            scores.append(score)

        # Standardize advantages
        if baseline == 'mean':
            baseline_val = sum(scores) / len(scores)
        elif baseline == 'min':
            baseline_val = min(scores)
        else:
            baseline_val = 0.0

        advantages = [s - baseline_val for s in scores]

        processed.append({
            'input': input_text,
            'candidates': candidates,
            'scores': scores,
            'advantages': advantages,
            'normalized_scores': [(s - min(scores)) / (max(scores) - min(scores) + 1e-8)
                                  for s in scores]
        })

    return processed

def evaluate_correctness(candidate, ground_truth):
    """Verify answer correctness using exact match or semantic similarity."""
    if isinstance(ground_truth, list):
        # Multiple acceptable answers
        return float(candidate.strip() in [g.strip() for g in ground_truth])
    else:
        # Single reference
        return float(candidate.strip() == ground_truth.strip())

Phase 2: Multi-Objective Data Construction

Transform single rollouts into three complementary datasets. Each serves a different learning objective, enabling the model to develop specialized subcomponents.

def construct_training_data(processed_rollouts):
    """Create pointwise, pairwise, and reflection datasets from rollouts."""
    pointwise_data = []
    pairwise_data = []
    reflection_data = []

    for rollout in processed_rollouts:
        input_text = rollout['input']
        candidates = rollout['candidates']
        scores = rollout['normalized_scores']
        advantages = rollout['advantages']

        # POINTWISE: (input, response) -> correctness score
        for candidate, score in zip(candidates, scores):
            pointwise_data.append({
                'input': input_text,
                'response': candidate,
                'target_score': score,
                'task_type': 'pointwise'
            })

        # PAIRWISE: Compare two responses and predict preference
        for i in range(len(candidates)):
            for j in range(i + 1, len(candidates)):
                better_idx = i if scores[i] > scores[j] else j
                worse_idx = j if better_idx == i else i

                pairwise_data.append({
                    'input': input_text,
                    'response_a': candidates[better_idx],
                    'response_b': candidates[worse_idx],
                    'label': 0,  # response_a is better
                    'task_type': 'pairwise'
                })

        # REFLECTION: Train model to self-correct wrong answers
        for candidate, score in zip(candidates, scores):
            if score < 0.5:  # Incorrect response
                reflection_data.append({
                    'input': input_text,
                    'incorrect_response': candidate,
                    'task': 'Generate corrected response',
                    'task_type': 'reflection'
                })

    return {
        'pointwise': pointwise_data,
        'pairwise': pairwise_data,
        'reflection': reflection_data
    }

def prepare_batch(examples, tokenizer, max_length=512):
    """Tokenize and format batch for unified training."""
    batch = {'input_ids': [], 'attention_mask': [], 'labels': []}

    for example in examples:
        task_type = example.get('task_type', 'pointwise')

        if task_type == 'pointwise':
            # Format: "Question: {input}\nAnswer: {response}\nScore: {score}"
            text = f"Question: {example['input']}\nAnswer: {example['response']}\nScore:"
            target = str(int(example['target_score'] * 100))

        elif task_type == 'pairwise':
            # Format for preference prediction
            text = f"Compare:\nA: {example['response_a']}\nB: {example['response_b']}\nBetter:"
            target = "A"

        else:  # reflection
            text = f"Fix: {example['incorrect_response']}\nCorrected:"
            target = ""  # Will be generated

        encoded = tokenizer(
            text,
            truncation=True,
            max_length=max_length,
            padding='max_length',
            return_tensors='pt'
        )

        batch['input_ids'].append(encoded['input_ids'].squeeze())
        batch['attention_mask'].append(encoded['attention_mask'].squeeze())

    return {k: torch.stack(v) for k, v in batch.items()}

Phase 3: Unified Loss Computation

Combine multiple objectives into a single loss function that balances policy improvement, reward accuracy, and KL divergence regularization.

import torch.nn.functional as F

def compute_unified_loss(
    model,
    batch,
    ref_model=None,
    loss_weights=None,
    beta_kl=0.1
):
    """
    Compute combined loss across pointwise, pairwise, and reflection objectives.

    Args:
        model: LLM with unified policy+reward head
        batch: Mixed batch of [pointwise, pairwise, reflection] examples
        ref_model: Reference model for KL divergence (optional)
        loss_weights: Dict with 'pointwise', 'pairwise', 'reflection' keys
        beta_kl: KL divergence regularization coefficient
    """
    if loss_weights is None:
        loss_weights = {'pointwise': 0.4, 'pairwise': 0.3, 'reflection': 0.3}

    total_loss = 0.0
    losses = {}

    # POINTWISE LOSS: Regression on correctness scores
    pointwise_mask = batch.get('pointwise_mask', torch.zeros_like(batch['input_ids']))
    if pointwise_mask.any():
        logits = model(batch['input_ids'], attention_mask=batch['attention_mask']).logits

        # Extract score predictions (final token)
        score_preds = logits[:, -1, :2]  # Assume 0-100 range mapped to 2 dimensions
        targets = batch.get('target_scores', torch.zeros(logits.shape[0]))

        pointwise_loss = F.smooth_l1_loss(score_preds.float(), targets.unsqueeze(1).float())
        losses['pointwise'] = pointwise_loss
        total_loss += loss_weights['pointwise'] * pointwise_loss

    # PAIRWISE LOSS: Classification of preference
    pairwise_mask = batch.get('pairwise_mask', torch.zeros_like(batch['input_ids']))
    if pairwise_mask.any():
        logits = model(batch['input_ids'], attention_mask=batch['attention_mask']).logits

        # Binary classification: is response_a better than response_b?
        pairwise_logits = logits[:, -1, :2]
        pairwise_targets = batch.get('pairwise_labels', torch.zeros(logits.shape[0]))

        pairwise_loss = F.cross_entropy(pairwise_logits, pairwise_targets.long())
        losses['pairwise'] = pairwise_loss
        total_loss += loss_weights['pairwise'] * pairwise_loss

    # REFLECTION LOSS: Generation quality for self-correction
    reflection_mask = batch.get('reflection_mask', torch.zeros_like(batch['input_ids']))
    if reflection_mask.any():
        logits = model(batch['input_ids'], attention_mask=batch['attention_mask']).logits

        # Reconstruction loss for generating corrections
        reflection_targets = batch.get('target_corrections', torch.full_like(batch['input_ids'], -100))

        shift_logits = logits[..., :-1, :].contiguous()
        shift_labels = batch['input_ids'][..., 1:].contiguous()

        reflection_loss = F.cross_entropy(
            shift_logits.view(-1, shift_logits.size(-1)),
            shift_labels.view(-1)
        )
        losses['reflection'] = reflection_loss
        total_loss += loss_weights['reflection'] * reflection_loss

    # KL DIVERGENCE REGULARIZATION: Prevent policy drift
    if ref_model is not None:
        ref_logits = ref_model(batch['input_ids'], attention_mask=batch['attention_mask']).logits
        current_logits = model(batch['input_ids'], attention_mask=batch['attention_mask']).logits

        kl_loss = F.kl_div(
            F.log_softmax(current_logits, dim=-1),
            F.softmax(ref_logits.detach(), dim=-1),
            reduction='batchmean'
        )
        losses['kl_div'] = kl_loss
        total_loss += beta_kl * kl_loss

    return total_loss, losses

def training_step(
    model,
    batch,
    optimizer,
    ref_model=None,
    loss_weights=None,
    beta_kl=0.1
):
    """Perform one training step with unified objectives."""
    model.train()

    loss, component_losses = compute_unified_loss(
        model,
        batch,
        ref_model=ref_model,
        loss_weights=loss_weights,
        beta_kl=beta_kl
    )

    optimizer.zero_grad()
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    optimizer.step()

    return loss.item(), {k: v.item() for k, v in component_losses.items()}

Phase 4: Complete Training Loop

Orchestrate all phases into a complete training loop that cycles between on-policy rollout collection and model optimization.

def train_spark(
    model,
    train_dataloader,
    val_inputs,
    val_labels,
    num_epochs=3,
    rollout_steps_per_epoch=4,
    batch_size=16,
    learning_rate=1e-5,
    loss_weights=None,
    beta_kl=0.1
):
    """
    Train model using SPARK framework.

    Args:
        model: LLM/VLM to train as unified policy+reward
        train_dataloader: Iterator over training inputs
        val_inputs: Validation input examples
        val_labels: Ground truth labels for validation
        num_epochs: Total training epochs
        rollout_steps_per_epoch: Rollout generation iterations per epoch
        batch_size: Training batch size
        learning_rate: Optimizer learning rate
        loss_weights: Balance between pointwise/pairwise/reflection objectives
        beta_kl: KL divergence regularization strength
    """
    if loss_weights is None:
        loss_weights = {'pointwise': 0.4, 'pairwise': 0.3, 'reflection': 0.3}

    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
    ref_model = copy.deepcopy(model)  # For KL regularization

    training_history = []

    for epoch in range(num_epochs):
        epoch_losses = {'pointwise': [], 'pairwise': [], 'reflection': [], 'kl_div': []}

        for step, input_batch in enumerate(train_dataloader):
            # Phase 1: Generate on-policy rollouts
            rollouts = generate_rollouts(model, input_batch, num_candidates=4)
            processed = compute_advantages(rollouts, val_labels)

            # Phase 2: Construct multi-objective datasets
            training_data = construct_training_data(processed)

            # Phase 3 & 4: Unified optimization over multiple sub-steps
            for _ in range(rollout_steps_per_epoch):
                # Sample from all three objectives
                pointwise_batch = torch.utils.data.random_split(
                    training_data['pointwise'],
                    [min(batch_size, len(training_data['pointwise'])),
                     max(0, len(training_data['pointwise']) - batch_size)]
                )[0]

                pairwise_batch = torch.utils.data.random_split(
                    training_data['pairwise'],
                    [min(batch_size, len(training_data['pairwise'])),
                     max(0, len(training_data['pairwise']) - batch_size)]
                )[0]

                reflection_batch = torch.utils.data.random_split(
                    training_data['reflection'],
                    [min(batch_size // 2, len(training_data['reflection'])),
                     max(0, len(training_data['reflection']) - batch_size // 2)]
                )[0]

                # Combine and prepare
                combined = list(pointwise_batch) + list(pairwise_batch) + list(reflection_batch)
                prepared = prepare_batch(combined, model.tokenizer)

                # Training step
                loss, losses = training_step(
                    model,
                    prepared,
                    optimizer,
                    ref_model=ref_model,
                    loss_weights=loss_weights,
                    beta_kl=beta_kl
                )

                for task_type in epoch_losses:
                    if task_type in losses:
                        epoch_losses[task_type].append(losses[task_type])

            if (step + 1) % 10 == 0:
                avg_losses = {k: sum(v) / len(v) if v else 0 for k, v in epoch_losses.items()}
                print(f"Epoch {epoch}, Step {step}: {avg_losses}")

        # Periodic checkpoint and reference model update
        if (epoch + 1) % 2 == 0:
            ref_model = copy.deepcopy(model)
            torch.save(model.state_dict(), f'checkpoint_epoch_{epoch}.pt')

        training_history.append(epoch_losses)

    return model, training_history

Practical Guidance

Hyperparameter Configuration

When to Use SPARK

• Objective tasks with verifiable rewards: Math, coding, information retrieval where ground truth is checkable
• Scaling reasoning without external reward models: When human preference annotation is bottlenecked
• Reducing inference cost: Single unified model vs. separate policy + reward infrastructure
• On-policy training preferred: When distribution drift is a concern and fresh rollouts are feasible
• Limited preference data: When collecting human feedback is expensive or unavailable
• Multi-scale reasoning: Combining pointwise judgment with pairwise preference and self-reflection

When NOT to Use SPARK

• Subjective tasks without ground truth: Content that lacks verifiable correctness signals (creative writing, aesthetic judgments)
• Preference-based RLHF only: If your primary goal is aligning with human subjective preferences and you have preference data, DPO or IPO may be more direct
• Fully offline data: SPARK requires on-policy rollout generation; fully offline scenarios suit standard SFT or offline RL
• Extreme efficiency constraints: Rollout generation per epoch adds computational cost compared to single forward pass SFT
• Real-time deployment requirements: Multi-candidate generation and reflection add latency at inference
• Noisy reward signals: Pointwise and pairwise objectives suffer if ground truth evaluation is unreliable or partially labeled
• Low-resource settings: KL regularization with reference models requires memory for two model copies

Common Pitfalls

1. Unbalanced loss weights across objectives: If pointwise weight dominates, the model optimizes for score prediction at the expense of generation quality. Start with equal weights and adjust after 1–2 epochs of training.

2. Insufficient rollout diversity: Using low temperature (< 0.5) during generation reduces candidate diversity. Meaningless comparisons occur when all candidates are similar. Maintain temperature >= 0.7.

3. Ignoring baseline selection: Using a fixed baseline (e.g., 0.0) instead of per-batch standardization inflates advantages early in training. Always subtract batch mean or min from scores.

4. Reference model staleness: If the reference model for KL divergence isn't updated periodically, policy drift regularization becomes ineffective. Refresh every 2–3 epochs or when validation accuracy plateaus.

5. Over-weighting reflection loss on clean data: If your evaluation function is perfect (100% of rollouts labeled correctly), reflection training wastes capacity. Scale reflection weight down when error rate < 20%.

6. Batch size inconsistency across objectives: Combining pointwise (which scales to 4N candidates) with pairwise (scales to 6N comparisons) in fixed batch_size can cause memory spikes. Use dynamic batching or separate sub-batches.

7. Mixing on-policy and off-policy data: Rollouts generated from old policies have different advantage scales. Always regenerate rollouts before each training epoch rather than recycling old data.

Reference

Synergistic Policy And Reward Co-Evolving Framework for Large Language Models

• arXiv: https://arxiv.org/abs/2509.22624
• Submitted: September 26, 2025
• Code and resources available via official project repository