​
Getting Started with Reward Functions

This guide will help you understand the basics of creating, testing, and deploying reward functions using the Reward Kit.

​
What is a Reward Function?

A reward function is a mechanism for evaluating the quality of model outputs in reinforcement learning from verifiable reward (RLVR) workflows. Reward functions help:

  • Evaluate model responses based on specific criteria.
  • Provide numerical scores that can be used to optimize models.
  • Offer explanations for why specific scores were assigned.

​
Installation

To get started with Reward Kit, install it via pip:

pip install reward-kit

For development, including running all examples and contributing to the codebase, install it in editable mode with development dependencies:

git clone https://github.com/fireworks-ai/reward-kit.git # Or your fork
cd reward-kit
pip install -e ".[dev]"

​
Authentication Setup

To use Reward Kit with the Fireworks AI platform, set up your authentication credentials:

# Set your API key
export FIREWORKS_API_KEY=your_api_key

​
Basic Reward Function Structure

Here’s a simple reward function that evaluates responses based on word count:

from reward_kit import reward_function, EvaluateResult, MetricResult
from typing import List, Dict, Optional

@reward_function
def word_count_reward(
    messages: List[Dict[str, str]],
    original_messages: Optional[List[Dict[str, str]]] = None,
    **kwargs
) -> EvaluateResult:
    """
    Evaluate a response based on its word count.

    Args:
        messages: List of conversation messages
        original_messages: Original messages (usually without the response being evaluated)
        **kwargs: Additional parameters

    Returns:
        EvaluateResult with score and metrics information
    """
    # Get the assistant's response (last message)
    if not messages or messages[-1].get("role") != "assistant":
        return EvaluateResult(
            score=0.0,
            reason="No assistant response found in messages.",
            metrics={"error": MetricResult(score=0.0, success=False, reason="No assistant response found")}
        )

    response = messages[-1].get("content", "")

    # Count words and calculate score
    word_count = len(response.split())
    score = min(word_count / 100.0, 1.0)  # Cap at 1.0
    success = word_count > 10 # Example: success if more than 10 words

    return EvaluateResult(
        score=score,
        reason=f"Overall word count evaluation: {word_count} words.",
        metrics={
            "word_count": MetricResult(
                score=score,
                success=success,
                reason=f"Word count: {word_count}"
            )
        }
    )

​
Testing and Evaluating

There are several ways to test your reward functions and run evaluations:

​
Programmatic Testing (for individual functions)

You can test your reward function directly in Python with sample conversations:

# Sample conversation
test_messages = [
    {"role": "user", "content": "What is machine learning?"},
    {"role": "assistant", "content": "Machine learning is a method of data analysis that automates analytical model building."}
]

# Test the reward function
result = word_count_reward(messages=test_messages)
print(f"Score: {result.score}")
print(f"Explanation: {result.metrics['word_count'].reason}")

For evaluating datasets or running complete examples, the primary method is the reward-kit run CLI command. This uses Hydra for configuration, allowing you to define your dataset, model, and reward logic in YAML files.

  1. Explore Examples: Check out the examples in the examples/ directory at the root of the repository. The main Examples README provides an overview and guidance on their structure. Each example (e.g., examples/math_example/) has its own README explaining how to run it.

  2. Run an Example:

    # Example: Running the math_example
python -m reward_kit.cli run \
  --config-path examples/math_example/conf \
  --config-name run_math_eval.yaml

    This command processes the dataset, generates model responses, applies reward functions, and saves detailed results.

​
Previewing Evaluation Outputs with reward-kit preview

After running an evaluation with reward-kit run, a preview_input_output_pairs.jsonl file is typically generated in the output directory. You can use reward-kit preview to inspect these pairs or re-evaluate them with different metrics:

# Preview outputs from a previous run
reward-kit preview \
  --samples ./outputs/YYYY-MM-DD/HH-MM-SS/preview_input_output_pairs.jsonl \
  --metrics-folders "your_metric_name=./path/to/your_metric_script"

Refer to the Evaluation Workflows guide for a more detailed lifecycle overview.

​
Deploying Your Reward Function

When you’re ready, deploy your reward function to use in training workflows:

# Deploy programmatically
evaluator_id = word_count_reward.deploy(
    name="word-count-evaluator",
    description="Evaluates responses based on word count"
)
print(f"Deployed with ID: {evaluator_id}")

Or using the CLI:

reward-kit deploy --id word-count-evaluator --metrics-folders "word_count=./path/to/metric" --force

​
Next Steps

Now that you have an overview of getting started:

  1. Dive deeper into Reward Function Anatomy.
  2. Understand the Core Data Types used in Reward Kit.
  3. Explore the Evaluation Workflows in more detail.
  4. Browse the Examples Overview and the main Examples README to find practical implementations.
  5. Follow our step-by-step tutorial for a hands-on walkthrough.