Tune Gemma 3 1B in JAX with GRPO for reasoning (Part 5): Reward System Deep Dive

Tunix (Tune-In-JAX)

Reward System Deep Dive: Tuning, Testing & Debugging

This is Part 5 of an 8-part series. In Part 4, we loaded Gemma 3 1B with LoRA. Now we’ll take a deep dive into the reward system — the component that determines what your model actually learns.

The reward system is the most important part of GRPO training that you control. Get it right, and your model learns sophisticated reasoning, get it wrong, and you’ll spend hours training a model that produces nonsense — or worse, learns to game your rewards without actually improving.

This part provides examples of testing frameworks, weight tuning strategies, and debugging tools. The goal: validate your reward system before committing to a multi-hour training run.

Why Rewards Deserve a Deep Dive

In GRPO, reward functions are the teacher. They tell the model:

• What makes a good response → high reward, reinforce this behavior
• What makes a bad response → low reward, discourage this behavior
• How to improve incrementally → graduated scoring guides learning

Consider how GRPO ranks three responses to “What is 15% of 80?”:

Response A (Excellent):
<reasoning>
Step 1: Convert 15% to decimal: 0.15
Step 2: Multiply: 0.15 × 80 = 12
Therefore, 15% of 80 equals 12.
</reasoning>
<answer>12</answer>

→ Format: 1.0 | Coherence: 0.95 | Correctness: 1.0
→ Composite: 0.98 - STRONGLY REINFORCE
Response B (Wrong format):
The answer is 12.
→ Format: 0.0 | Coherence: 0.0 | Correctness: 0.7
→ Composite: 0.18 - WEAKLY REINFORCE (correct but wrong structure)
Response C (Wrong answer, good format):
<reasoning>15% of 80 = 80/15 = 5.3</reasoning>
<answer>5.3</answer>
→ Format: 1.0 | Coherence: 0.35 | Correctness: 0.0
→ Composite: 0.52 - MODERATE (good format, flawed reasoning)

GRPO compares these within a group and reinforces based on relative quality: A >> C >> B.

The key insight: rewards are normalized (0.0–1.0) and weighted. This provides interpretable scores and tunable component balance.

Our Four Reward Components (Recap)

In this tutorial, we designed a multi-component system in Part 3. Here’s how the pieces fit together:

                    ┌─────────────────┐
                    │   COMPOSITE     │
                    │    REWARD       │
                    └────────┬────────┘
                             │
       ┌──────────┬──────────┼──────────┬──────────┐
       ▼          ▼          ▼          ▼          │
  ┌────────┐ ┌────────┐ ┌────────┐ ┌────────┐      │
  │ FORMAT │ │COHERENCE││CORRECT-│ │EFFICIENCY│    │
  │  (25%) │ │  (20%) │ │NESS(55%)││  (0%)   │     │
  └────────┘ └────────┘ └────────┘ └────────┘      │
                                                   │
  Composite = Σ (weight_i × component_i) ◄─────────┘

Why These Weights?

• Correctness dominates (55%): For reasoning tasks, getting the right answer with good reasoning is the primary goal
• Format is foundational (25%): Without proper tags, we can’t even parse the response
• Coherence supports quality (20%): Rewards logical structure without dominating
• Efficiency disabled (0%): Length control is secondary; we don’t want to penalize thorough explanations

Testing Before Training: The RewardAnalyzer

Before committing hours to training, verify your reward system behaves correctly. We built a RewardAnalyzer class for exactly this purpose:

class RewardAnalyzer:
    """Analyze and visualize reward component contributions."""

    def analyze_response(self, response, ground_truth=None, prompt=""):
        """Get detailed breakdown of all component scores."""
        # Returns raw scores, weighted scores, and composite

    def print_breakdown(self, result):
        """Display formatted component breakdown with visual bars."""

    def compare_responses(self, responses, ground_truth, prompt):
        """Compare multiple responses side-by-side (simulates GRPO)."""

Usage:

analyzer = RewardAnalyzer(config)
result = analyzer.analyze_response(response, ground_truth="12", prompt="What is 15% of 80?")
analyzer.print_breakdown(result)

# Output:
# ─────────────────────────────────────────────────────
# Component       Raw      Weight   Weighted   Bar
# ─────────────────────────────────────────────────────
# Format          1.00     25%      0.250      ████████████████████
# Coherence       0.85     20%      0.170      █████████████████░░░
# Correctness     0.95     55%      0.523      ███████████████████░
# Efficiency      0.80     0%       0.000      ████████████████░░░░
# ─────────────────────────────────────────────────────
# COMPOSITE                         0.943

This visibility lets you catch problems before training.

Comprehensive Test Suite

We validate the reward system against diverse response patterns:

Test Categories

Sample Test: Format Compliance

# Missing all tags - should score LOW
no_tags = "Consciousness is difficult to define. Some say it's reducible to brain states."
result = analyzer.analyze_response(no_tags)
# → Format: 0.0, Composite: ~0.20

# Proper structure - should score HIGH  
with_tags = """<reasoning>
Consciousness presents what Chalmers calls the "hard problem"...
</reasoning>
<answer>Consciousness remains philosophically contested.</answer>"""
result = analyzer.analyze_response(with_tags)
# → Format: 1.0, Composite: ~0.85

Verification Checklist

Before training, confirm:

checks = [
    ("Excellent responses score > 0.85", best_score > 0.85),
    ("Missing tags penalized", no_tags_score < 0.40),
    ("Depth indicators increase score", deep_score > shallow_score),
    ("Too brief responses penalized", brief_score < 0.50),
]

for check_name, passed in checks:
    print(f"{'✅' if passed else '❌'} {check_name}")

If any check fails, adjust your reward functions or weights before training based on your specifications.

*💡 *Full test suite: The complete testing framework with 15+ test cases across all categories is on GitHub.

Weight Tuning Strategies

Default weights work well for general reasoning, but you may need adjustments for specific domains.

When to Adjust

Sensitivity Analysis

We can see how composite scores change as we vary the correctness weight:

Correctness Weight    Coherence    Composite    Notes
────────────────────────────────────────────────────
     20%               35%          0.72        Too low for reasoning
     40%               15%          0.78        
     55%               10%          0.82        ← Our choice
     70%               -5%          N/A         Coherence goes negative

The 55% correctness weight balances answer accuracy with reasoning quality.

Domain-Specific Configurations Examples

# Ethics/Moral Philosophy: Higher coherence for argument validity
ethics_weights = {'format': 0.20, 'coherence': 0.25, 'correctness': 0.45, 'efficiency': 0.10}

# Pure Epistemology: Maximum reasoning depth
epistemology_weights = {'format': 0.20, 'coherence': 0.10, 'correctness': 0.60, 'efficiency': 0.10}

# Logic/Formal Reasoning: Formal validity matters
logic_weights = {'format': 0.20, 'coherence': 0.35, 'correctness': 0.35, 'efficiency': 0.10}

Simulating GRPO Rankings

GRPO learns by comparing responses within groups. We can simulate this process:

prompt = "Can we ever truly know another person's mind?"

responses = [
    excellent_multi_perspective_response,  # Rich depth indicators
    good_but_one_sided_response,           # Logical but limited view
    shallow_with_good_format,              # Tags present, no substance
    no_structure_at_all,                   # Plain text, no tags
]
results = analyzer.compare_responses(responses, ground_truth=None, prompt=prompt)
# GRPO Advantage Calculation:
# Response 1: Score 0.92, Advantage +1.2 → STRONGLY REINFORCE
# Response 2: Score 0.71, Advantage +0.3 → MILD REINFORCE  
# Response 3: Score 0.45, Advantage -0.4 → MILD DISCOURAGE
# Response 4: Score 0.18, Advantage -1.1 → STRONGLY DISCOURAGE

If rankings match expectations, GRPO will push the model toward sophisticated reasoning. If they don’t, fix your rewards before training.

Debugging During Training

Training issues often trace back to reward problems. The RewardDebugger helps diagnose:

class RewardDebugger:
    def diagnose_response(self, response, ground_truth=None):
        """Detailed diagnosis of why a response got its score."""
        # Shows: tag detection, content extraction, indicator counts

    def batch_statistics(self, responses, ground_truths):
        """Compute statistics across a batch."""
        # Returns: mean, std, min, max, zero-count per component

    def print_batch_report(self, stats):
        """Print formatted report with warnings."""

Common Issues and Fixes

Batch Statistics Warning Signs

stats = debugger.batch_statistics(responses)
# Warnings:
# ⚠️  Low composite variance (std < 0.05) - may indicate reward collapse
# ⚠️  High zero rate for format - check tag detection
# ⚠️  Low format scores (mean < 0.3) - model not learning structure

Key Principles for Reward Design

1. Test Before You Train

Run your test suite on diverse examples, verify rankings match expectations and fix issues before committing hours to training.

2. Components Should Be Complementary

Each component should measure something distinct as in our example:

• Format: Structure compliance
• Coherence: Logical flow
• Correctness: Answer + reasoning quality
• Efficiency: Length appropriateness

Overlapping components can cause double-counting or conflicting signals.

3. Weights Reflect Priorities

Your weights encode what you care about, if correctness is 55%, you’re saying “right answers with good reasoning matter most.” Adjust based on your actual goals.

4. Monitor During Training

Use the debugger to check batch statistics periodically, watch for reward collapse or component imbalance and adjust if needed.

5. Iterate

Reward design is empirical, your first configuration probably isn’t optimal. Run short training experiments, evaluate outputs, and refine.

Part 5 Complete ✅

We’ve built a comprehensive reward validation framework:

1. RewardAnalyzer : Detailed component breakdown for any response
2. Test Suite : 15+ test cases covering format, depth, coherence, edge cases
3. Weight Tuning :Sensitivity analysis and domain-specific configurations
4. GRPO Simulation : Verify rankings before training
5. RewardDebugger :Diagnose issues during training

Key Takeaways

• Test rewards extensively before training : Hours of training can be wasted on misconfigured rewards
• Verify rankings match expectations : GRPO learns from relative quality within groups
• Monitor batch statistics : Catch reward collapse or component imbalance early
• Iterate on weights : Default weights are starting points, not final answers
• Debugging tools save time : Diagnose problems before they compound

What’s Next

In Part 6, we’ll bring everything together into the training pipeline. We’ll configure the Tunix GRPO learner, set up checkpointing, and start the actual training run.

→ Continue to Part 6: Training Pipeline

Series Navigation:

• Part 1: Introduction & Environment Setup
• Part 2: Understanding GRPO Algorithm
• Part 3: Data Loading & Preprocessing
• Part 4: Model Loading & LoRA Setup
• Part 5: Reward Function Design ← You are here
• Part 6: Training Pipeline
• Part 7: Evaluation & Results
• Part 8: Inference & Next Steps

All code is available on GitHub. Found this helpful? Follow for the complete series.