Reinforcement Learning

​

Why Reinforcement Fine-tuning?

• Require high-quality labeled data
• May not directly optimize for desired behaviors
• Can be expensive and time-consuming
• Limited ability to handle complex objectives

• Using programmable reward functions
• Optimizing directly for desired outcomes
• Learning from trial and error
• Developing generalized strategies

​

How GRPO Works

​

1. Dual Model Architecture

• Frozen Reference Model: Maintains baseline performance
• Trainable Policy Model: Optimized through training

• Prevents catastrophic forgetting
• Maintains model stability
• Enables controlled optimization
• Preserves general capabilities

​

2. Training Process

1. Generation Phase:
   • Start with example prompts
   • Generate multiple completions
   • Use temperature sampling
   • Create response variations
2. Evaluation Phase:
   • Score completions with reward functions
   • Calculate relative advantages
   • Identify high/low performing outputs
3. Optimization Phase:
   • Update policy model weights using the advantages calculated for each completion
   • Maintain reference proximity by using the frozen reference model to generate completions
   • Track performance metrics by monitoring the reward function scores and the total reward

​

3. Implementation Details

• Uses LoRA for efficient updates
• Leverages LoRAX (multi-LoRA) infrastructure

​

When to Use GRPO

​

Ideal Use Cases

​

Implementing GRPO

​

1. Reward Function Design

• Clear Objectives: Define specific, unambiguous goals that align with your desired model behavior
• Measurable Outcomes: Create quantifiable metrics that can be consistently evaluated
• Comprehensive Coverage: Ensure the reward function captures all important aspects of the task
• Balanced Scoring: Weight different criteria appropriately to avoid overemphasizing any single aspect
• Robust Validation: Include error handling and edge cases to prevent reward manipulation

def format_reward(prompt: str, completion: str, example: dict) -> float:
    """Validate output format"""
    try:
        # Check format requirements
        if meets_criteria(completion):
            return 1.0
        return 0.0
    except Exception:
        return 0.0

def quality_reward(prompt: str, completion: str, example: dict) -> float:
    """Assess response quality"""
    score = 0.0
    # Add quality checks
    score += check_coherence(completion)
    score += check_relevance(prompt, completion)
    return score / 2.0 # To normalize the score to 0-1 range

​

2. Training Configuration

• Number of training steps: Controls total iterations (typically 1000-5000)
• Number of generations per step: Controls the number of completions generated for each prompt (typically 8 or 16)
• Beta: Controls the balance between exploration and exploitation (0.1-1.0).
• Sampling parameters:
  • temperature: Higher values increase diversity, lower values increase consistency
  • max_tokens: Maximum length of generated responses
  • top_p: Nucleus sampling threshold (0.1-1.0) (for temperature > 0)

from predibase import GRPOConfig, SamplingParams, RewardFunctionsConfig, RewardFunctionsRuntimeConfig

config = GRPOConfig(
    base_model="model-name",
    train_steps=1000,
    num_generations=16,
    beta=0.001, # 0.001 is a good starting point
    sampling_params=SamplingParams(
      temperature=0.9,
      max_tokens=1024,
    ),
    reward_fns=RewardFunctionsConfig(
        functions={
            "format": format_reward,
            "quality": quality_reward
        }
    )
)

​

3. Monitoring and Optimization

• Reward Function Trends: Monitor how individual reward functions evolve over time to identify learning patterns and potential issues
• Generation Diversity: Track the variety of responses being generated to ensure the model isn’t getting stuck in local optima
• Average Reward: Watch the mean reward across generations to gauge overall progress
• Reward Variance: Monitor the spread of rewards to understand model stability
• Convergence Rate: Track how quickly the model reaches stable performance
• Failure Modes: Document common types of incorrect or low-quality outputs by monitoring the logs and completions tabs.

• Fine-tune hyperparameters (e.g. beta, temperature, max_tokens)
• Monitor convergence (e.g. reward function trends, total reward)

​

Best Practices

​

1. Reward Function Design

• Start Simple: Begin with basic reward functions and incrementally add complexity as needed
• Comprehensive Testing: Test reward functions across diverse inputs and edge cases to ensure robust behavior
• Error Handling: Implement graceful error handling for unexpected inputs and edge cases
• Clear Documentation: Document all assumptions, requirements, and expected behaviors
• Graded Rewards: Use continuous scoring (0-1 range) instead of binary rewards to provide richer learning signals
• Multi-Objective Rewards: Combine format, correctness, and other relevant metrics for comprehensive feedback
• Validation: Include validation checks to ensure reward functions produce expected outputs
• Modularity: Design reward functions to be modular (one reward function per objective)

​

2. Training Process

• Start Small: Begin with a small number of training steps and gradually increase as needed
• Monitor Closely: Keep a close eye on reward function trends, generation diversity, and average reward to identify potential issues
• Wait For 40-50 Steps: Expect 40-50 steps before clear improvements
• Validate Frequently: Regularly validate reward functions to ensure they are working as expected
• Watch for Reward Hacking: Monitor for reward hacking by checking the logs and completions tabs

​

3. Reward Function Tips

​

Data Types in Example Dictionary

• Lists/dictionaries: Use ast.literal_eval()
• Numbers: Use float() or int()
• Booleans: Use bool()

​

Partial Credit Scoring

• Binary rewards (0 or 1) provide limited learning signal
• Partial credit creates smoother learning gradients
• Example: 0.5 points for correct approach but wrong final answer
• Helps identify incremental improvements

​

External Libraries

• LLM-as-judge using OpenAI client
• NLP libraries like spaCy
• Scientific computing with NumPy
• External APIs for validation

​

4. Monitoring and Debugging

​

Understanding Reward Graphs

• total_reward: Average reward across all reward functions, indicating overall model performance
• total_reward_std: Standard deviation of rewards, measuring performance stability and consistency
• Individual reward function trends: Track progress of each reward component separately
• Version-specific performance: Compare metrics across different model versions

• Steady upward trends in reward scores, showing continuous improvement
• Decreasing variance in rewards over time, indicating more stable performance
• Format rewards improving before task rewards, suggesting proper learning progression
• Consistent progress across model versions, demonstrating reliable training

​

Using Reward Logs

• Monitor training metrics and reward trends over time
• Debug reward function behavior using detailed logging and print statements
• Track and troubleshoot API response times and error rates
• Analyze periods of stagnant reward scores to identify potential issues

​

Analyzing Completions

• Side-by-side comparison of model outputs across training iterations
• Detailed reward score breakdowns for each completion
• Visual progress tracking of model improvements over time
• Detection of potential reward hacking attempts

​

Common Challenges

1. Training Issues:
   • Unstable or flat rewards: Monitor reward trends and adjust learning rate
   • Slow convergence: Consider increasing batch size or simplifying reward criteria
   • Quality regression: Check for reward hacking and adjust reward function weights
   • Reward hacking: Implement reward shaping and validation checks
2. Implementation Problems:
   • Complex reward calculations: Break down into smaller, testable components
   • Integration issues: Use proper error handling and logging
   • API timeouts: Implement retry logic and fallback mechanisms
3. Reward Function Issues:
   • Non-numeric returns: Ensure all reward functions return float values
   • Inconsistent scoring: Standardize scoring ranges and normalize outputs
   • Missing error handling: Add try-catch blocks and default values
   • Slow computation: Optimize calculations and cache results when possible

​

Frequently Asked Questions

​

How many training steps should I use?

• Dataset size and complexity
• Number of reward functions
• Base model size
• Task difficulty

​

What if my rewards are not improving?

1. Check reward function implementation
2. Verify training data quality
3. Consider simplifying reward criteria

​

Can I use external APIs in reward functions?

• API costs and rate limits (e.g., OpenAI, Anthropic, etc.):
  • Monitor usage to stay within your budget
• Robust error handling
  • Implement exponential backoff for retries
  • Add comprehensive logging for debugging
  • Handle various failure modes gracefully
• Fallback mechanisms
  • Maintain local alternatives for critical functions
  • Define sensible default scores for failures

​

What’s the recommended reward range?

• 0: Completely incorrect
• 0.1-0.9: Partial credit
• 1.0: Perfect response

• Provide clear learning signals
• Enable reward combination

​

Can I update reward functions during training?

1. Get current config using pb.adapters.get_config()
2. Update the reward function
3. Apply changes with pb.adapters.update_config()

# Get current config
cfg = pb.adapters.get_config("myrepo/1")

# Update existing or add new functions
cfg.reward_fns["my_reward"] = my_reward_function_v2

# Apply updates
pb.adapters.update_config("myrepo/1", cfg)

​

What if my model is “reward hacking”?

1. Add validation checks for input/output quality
2. Implement penalties for suspicious patterns
3. Update reward functions to discourage exploitation
4. Use multiple reward criteria to prevent single-metric optimization
5. Review edge cases and boundary conditions
6. Use partial credit scoring for better learning signals

​

Do the rewards need to be binary?

​

What if I need to use an external library in my reward function?

from predibase import RewardFunctionsConfig, RewardFunctionsRuntimeConfig

def my_reward_function(prompt, completion, example) -> float:
    import my_pkg
    return my_pkg.score(...)

cfg = RewardFunctionsConfig(
    runtime=RewardFunctionsRuntimeConfig(
        packages=[
            "mypkg", # Add any packages you need here
        ]
    ),
    functions={
        "my_reward": my_reward_function,
    },

​

How do I pass additional columns to my reward functions?

​

What data types are supported in the example dictionary?

• Lists, dictionaries, sets and other data structures can be converted using ast.literal_eval()
• Numbers can be converted using float() or int()
• Booleans can be converted using bool()

​

How do I interpret my reward graphs?

1. total_reward: Shows the average total reward across all reward functions combined. An upward trend indicates the model is learning to optimize for the defined rewards.
2. total_reward_std: Shows the standard deviation of the total reward, indicating how consistent the model’s performance is across different examples. Over time, this should decrease as the model more consistently gets higher rewards for its completions.
3. Individual reward function graphs: Each reward function you define has a graph showing how well the model is learning that specific objective. For example, you might see separate graphs for format_reward_func, correctness_reward_func, etc. The value graphed is the average reward for that function across all completions at the given training step.

• Look for an overall upward trend (except for the standard deviation graph, which should trend downwards). Progress may not be linear.
• It’s normal to need 40-50 training steps before seeing clear signs of improvement.
• Format-related rewards (like proper JSON structure) often improve first, before more complex task-specific rewards.
• High variance (spiky graphs) is common early in training but should generally decrease over time.
• If a specific reward function’s graph remains flat or decreases, you may need to adjust its implementation by updating the reward function. It’s a good idea to check the Reward Logs tab first to see if there is an obvious issue with it before working towards improving it.

​

How do I use the reward logs tab?

1. Track training progress by showing which reward functions are executing, along with step and epoch numbers.
2. Print out the reward score assigned to the completion before returning the reward score itself.
3. Debug issues in your reward function implementations by inspecting the logged output.
4. Monitor performance issues like API timeouts or errors.
5. Investigate why reward scores might be flatlining or behaving unexpectedly.

​

How do I use the completions tab?

1. Select a Prompt: Use the prompt selection table to choose any prompt from your training set. The table displays:
   • Prompt Index: A unique identifier for each prompt.
   • Prompt Text: The actual text of the prompt shown to the model.
   • Prompt Length: Number of tokens in the prompt.
   • Additional Columns: Any extra metadata or context provided with the prompt.
   • Epoch: The training epochs where this prompt was used.
   • Total Completions: Number of completions generated for this prompt.
2. Compare Completions: This section allows you to compare two completions side-by-side.
   • By default the right column shows the best completion generated at the most recent epoch. The left column allows you to select any other completion from the same prompt (from any epoch).
   • A slider allows you to pick the epoch for the left column, and arrows allow you to switch between completions at that epoch.
   • Each column displays the completion text and all individual reward scores (and the total reward).
   • You can also easily compare any two completions (not just against the best one from the current epoch).
3. View all completions: Shows information about every completion generated for the selected prompt.
   • Epoch: The training epoch when the completion was generated.
   • Completion: The full text output from the model.
   • Individual reward scores: Separate columns for each reward function (e.g. format_reward_func_v1).
   • Total: The combined reward score across all functions.
   • Length: Number of tokens in the completion.
   • View: Allows you to display this completion in the “Compare Completions” section.

• Verifying that higher reward scores correspond to better quality outputs.
• Detecting reward hacking (where the model games the reward functions).
• Understanding how responses evolve throughout training.
• Identifying areas needing improved reward functions.
• Detecting breakthrough moments in the training process.
• Validating proper output formatting and structure.

​

How do I control the length of my GRPO training run?

​

Any additional tips?

​

Next Steps

• Explore Use Cases
• Learn about Production Deployment
• Reinforcement learning documentation