How to Implement Custom Reward Functions for Domain-Specific Tasks in AReaL

AReaL implements custom reward functions as pure Python callables that receive prompts, completions, token IDs, and ground-truth answers, returning scalar scores that integrate directly into RL training loops via string import paths.

AReaL (Advanced Reasoning and Learning) is an open-source reinforcement learning framework designed for vision-language models and mathematical reasoning. To implement custom reward functions for domain-specific tasks in AReaL, you must adhere to a specific function contract and leverage the framework's dispatcher utilities. This guide walks through the exact implementation patterns found in the inclusionai/areal repository.

Understanding AReaL's Reward Architecture

The Reward Function Contract

Every reward function in AReaL must implement a fixed six-parameter signature. The trainer invokes this callable during the RL loop, passing raw model outputs and ground-truth data.

The exact signature required is:

def my_reward_fn(prompt, completions, prompt_ids, completion_ids, answer, **kwargs) -> float:
    """
    Compute reward score for domain-specific task.
    
    Args:
        prompt: Original input prompt string
        completions: Model-generated output strings
        prompt_ids: Token IDs of the prompt (tensor)
        completion_ids: Token IDs of the completion (tensor)
        answer: Ground-truth reference answer
        **kwargs: Optional domain-specific arguments
        
    Returns:
        Scalar reward value (float or int)
    """
    pass

The Dispatcher Pattern

AReaL uses get_custom_reward_fn in areal/reward/__init__.py (lines 11-20) to map string identifiers to concrete callable implementations:

# areal/reward/__init__.py (lines 11-20)

def get_custom_reward_fn(path: str, **kwargs):
    if "clevr_count_70k" in path:
        from .clevr_count_70k import clevr_count_70k_reward_fn
        return clevr_count_70k_reward_fn
    elif "geometry3k" in path:
        from .geometry3k import geometry3k_reward_fn
        return geometry3k_reward_fn
    else:
        raise ValueError(...)

This dispatcher enables configuration-driven training workflows where reward functions are specified via string import paths.

Utility Helpers for Numeric Verification

For math-heavy domains, AReaL provides MathVerifyWorker in areal/reward/__init__.py (lines 27-38). This utility wraps math_verify with configurable precision tolerance:

# areal/reward/__init__.py (lines 27-38)

class MathVerifyWorker:
    # wraps math_verify with configurable precision

    ...

def get_math_verify_worker() -> MathVerifyWorker:
    ...

Reuse this class to compare numeric predictions against ground truth without implementing custom tolerance logic.

Step-by-Step Implementation Guide

Step 1: Create the Reward Module

Create a new Python file within the areal/reward/ directory or any importable package in your Python path. Placing files in areal/reward/ keeps rewards version-controlled with the core framework.

touch areal/reward/my_domain_reward.py

Step 2: Define the Function Signature

Implement the required six-parameter signature exactly as specified in the AReaL contract:

def my_domain_reward_fn(prompt, completions, prompt_ids, completion_ids, answer, **kwargs) -> float:
    """Domain-specific reward implementation."""
    # Implementation logic here

    pass

Step 3: Extract and Parse Model Outputs

Raw LLM outputs often contain explanatory text, formatting markers, or multiple candidate answers. Use regular expressions or domain-specific parsers to isolate the structured prediction from the completions string:

import re

def _extract_number(text: str) -> str:
    """Return the last numeric token found in text."""
    nums = re.findall(r"-?\d*\.?\d+", text.replace(",", ""))
    return nums[-1] if nums else ""

Step 4: Compute the Scalar Score

Calculate a reward value between 0 and 1 (or any positive float) based on the similarity between prediction and ground truth. For numeric tasks, instantiate MathVerifyWorker via get_math_verify_worker():

from areal.reward import get_math_verify_worker

def my_domain_reward_fn(prompt, completions, prompt_ids, completion_ids, answer, **kwargs) -> float:
    pred = _extract_number(str(completions))
    gold = _extract_number(str(answer)) or str(answer)
    
    if not pred or not gold:
        return 0.0
    
    worker = get_math_verify_worker()
    return worker.verify(pred, gold)

Step 5: Register and Use the Reward

Optionally register your reward in areal/reward/__init__.py by adding the identifier to VALID_REWARD_FN or extending the get_custom_reward_fn dispatcher:

# In areal/reward/__init__.py

VALID_REWARD_FN = ["clevr_count_70k", "geometry3k", "my_domain"]

Reference the reward in your training workflow by passing the fully-qualified import path to the reward_fn parameter:

# examples/vlm/my_domain_grpo.py

workflow_kwargs = dict(
    reward_fn="areal.reward.my_domain_reward.my_domain_reward_fn",
    gconfig=config.gconfig,
    tokenizer=config.tokenizer_path,
    processor=config.tokenizer_path,
    enable_thinking=False,
)

Complete Code Example

Here is the full implementation for a numeric domain-specific task, following the patterns established in areal/reward/geometry3k.py:

# areal/reward/my_domain_reward.py

import re
from areal.reward import get_math_verify_worker

def _extract_number(text: str) -> str:
    """Return the last numeric token found in *text*."""
    nums = re.findall(r"-?\d*\.?\d+", text.replace(",", ""))
    return nums[-1] if nums else ""

def my_domain_reward_fn(
    prompt,
    completions,
    prompt_ids,
    completion_ids,
    answer,
    **kwargs,
) -> float:
    """
    Custom reward for a domain‑specific regression task.

    - *prompt*: original prompt (unused here).
    - *completions*: model output string.
    - *answer*: ground‑truth answer (string or numeric).
    - Returns a float in [0, 1].
    """
    # 1️⃣ Extract model prediction

    pred = _extract_number(str(completions))
    # 2️⃣ Normalize ground‑truth

    gold = _extract_number(str(answer)) or str(answer)

    if not pred or not gold:
        return 0.0

    # 3️⃣ Use the shared MathVerifyWorker for tolerant numeric comparison

    worker = get_math_verify_worker()
    return worker.verify(pred, gold)

Integration with Training Workflows

AReaL dynamically imports reward functions at runtime using the string path provided in the configuration. As demonstrated in examples/vlm/geometry3k_grpo.py (lines 51-53), the reward_fn parameter accepts the fully-qualified import path:

workflow_kwargs = dict(
    reward_fn="examples.vlm.geometry3k_grpo.geometry3k_reward_fn",
    gconfig=config.gconfig,
    tokenizer=config.tokenizer_path,
    processor=config.tokenizer_path,
    enable_thinking=False,
)

This pattern decouples reward logic from the training loop, allowing you to iterate on domain-specific evaluation metrics without modifying core AReaL infrastructure.

Summary

• AReaL reward functions are pure Python callables with a fixed six-parameter signature (prompt, completions, prompt_ids, completion_ids, answer, **kwargs) returning scalar floats.
• The get_custom_reward_fn dispatcher in areal/reward/__init__.py maps string identifiers to concrete implementations, enabling configuration-driven training.
• MathVerifyWorker provides reusable numeric verification for math-heavy domains with configurable precision tolerance.
• Implementation requires creating a Python module, defining the signature, parsing completions, computing scores, and passing the fully-qualified import path to the trainer's reward_fn argument.

Frequently Asked Questions

What is the exact function signature required for AReaL reward functions?

AReaL requires reward functions to accept six positional arguments—prompt, completions, prompt_ids, completion_ids, answer—plus optional **kwargs, and return a scalar float or integer. This contract is strictly enforced when the trainer calls your function via the dispatcher in areal/reward/__init__.py.

Can I reuse existing verification logic for mathematical tasks?

Yes. The MathVerifyWorker class defined in areal/reward/__init__.py (lines 27-38) wraps math_verify with configurable precision tolerance. Import it via from areal.reward import get_math_verify_worker to compare numeric predictions against ground truth without writing custom tolerance logic.

How do I register a new reward function so AReaL can find it by name?

Add your reward's short identifier to the VALID_REWARD_FN list in areal/reward/__init__.py, or implement a new conditional branch in get_custom_reward_fn that imports your module. Alternatively, pass the fully-qualified Python import path (e.g., "areal.reward.my_domain_reward.my_domain_reward_fn") directly to the reward_fn argument in your workflow configuration.

Where should I place my custom reward function files?

Place new reward modules inside the areal/reward/ directory to keep them version-controlled with the core framework, or place them in any Python package that is on the system path (such as examples/vlm/). AReaL dynamically imports the function at runtime using the string path you provide, so the location only needs to be importable by the Python interpreter running the training job.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}