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

> Learn to implement custom reward functions for domain-specific tasks in AReaL. Integrate Python callables directly into RL training loops for tailored AI development.

- Repository: [inclusionAI/areal](https://github.com/inclusionai/areal)
- Tags: how-to-guide
- Published: 2026-03-04

---

**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:

```python
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`](https://github.com/inclusionai/areal/blob/main/areal/reward/__init__.py) (lines 11-20) to map string identifiers to concrete callable implementations:

```python

# 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`](https://github.com/inclusionai/areal/blob/main/areal/reward/__init__.py) (lines 27-38). This utility wraps `math_verify` with configurable precision tolerance:

```python

# 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.

```bash
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:

```python
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:

```python
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()`:

```python
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`](https://github.com/inclusionai/areal/blob/main/areal/reward/__init__.py) by adding the identifier to `VALID_REWARD_FN` or extending the `get_custom_reward_fn` dispatcher:

```python

# 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:

```python

# 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`](https://github.com/inclusionai/areal/blob/main/areal/reward/geometry3k.py):

```python

# 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`](https://github.com/inclusionai/areal/blob/main/examples/vlm/geometry3k_grpo.py) (lines 51-53), the `reward_fn` parameter accepts the fully-qualified import path:

```python
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`](https://github.com/inclusionai/areal/blob/main/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`](https://github.com/inclusionai/areal/blob/main/areal/reward/__init__.py).

### Can I reuse existing verification logic for mathematical tasks?

Yes. The `MathVerifyWorker` class defined in [`areal/reward/__init__.py`](https://github.com/inclusionai/areal/blob/main/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`](https://github.com/inclusionai/areal/blob/main/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.