# How the GRPO RL Trainer Executes Unit Test Rewards in OlmOCR

> Discover how the GRPO RL trainer in OlmOCR executes unit test rewards. Learn how pass rates drive policy updates for vision-language models.

- Repository: [Ai2/olmocr](https://github.com/allenai/olmocr)
- Tags: internals
- Published: 2026-07-06

---

**The OlmOCR training pipeline uses the Group-Relative Policy Optimization (GRPO) algorithm to fine-tune vision-language models by executing unit tests against generated completions, converting pass rates into scalar rewards that drive policy updates.**

The allenai/olmocr repository implements a novel reinforcement learning approach for document understanding by treating OCR quality as a unit testing problem. At the heart of this system is the **GRPO RL trainer**, which evaluates model outputs against predefined test suites and uses the results to compute low-variance policy gradients. This article examines the implementation details found in the source code, specifically how unit test rewards are calculated and applied during training.

## GRPO Trainer Initialization and Reward Configuration

The training orchestration begins in [`olmocr/train/grpo_train.py`](https://github.com/allenai/olmocr/blob/main/olmocr/train/grpo_train.py), where the `GRPOTrainer` from the 🤗 TRL library is instantiated with a custom list of reward functions. Before trainer initialization, the code constructs the `reward_funcs` list that will evaluate each model completion:

```python

# olmocr/train/grpo_train.py

logger.info("Initializing GRPO trainer")
trainer = GRPOTrainer(
    model=model,
    args=grpo_config,
    processing_class=processor,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    reward_funcs=reward_funcs,
)

```

The `reward_funcs` list is populated immediately prior to this instantiation. By default, it includes the benchmark reward that executes unit tests defined in OlmOCR bench JSONL files:

```python

# olmocr/train/grpo_train.py

reward_funcs.append(olmocr_bench_reward)               # micro‑average

# also add a macro‑averaged variant when requested

def olmocr_bench_reward_macroavg(...):
    return olmocr_bench_reward(..., macro_average=True)
reward_funcs.append(olmocr_bench_reward_macroavg)

```

According to the source code at lines 1429‑1434, these functions are appended to the list only when the corresponding command-line flags (`--reward_bench` and `--reward_bench_macroavg`) are provided.

## The Unit Test Reward Function

The core evaluation logic resides in the `olmocr_bench_reward` function defined at line 1093 of [`olmocr/train/grpo_train.py`](https://github.com/allenai/olmocr/blob/main/olmocr/train/grpo_train.py). This function receives batch inputs including the PDF path, JSONL test file location, and specific test IDs to evaluate:

```python

# olmocr/train/grpo_train.py

def olmocr_bench_reward(
    prompts,
    completions,
    completion_ids,
    pdf_path,
    jsonl_file,
    test_ids,
    macro_average: bool = False,
    **kwargs,
):
    """Run unit tests for each completion and return a reward."""
    # … parallel evaluation via ThreadPoolExecutor …

    # … per‑completion reward = overall pass rate (micro) or

    #  average of per‑type pass rates (macro) …

```

The reward computation follows a three-stage pipeline:

- **Test Loading** – `load_specific_tests_cached` (LRU-cached at line 5253) reads only the specific test IDs required for the current completion, avoiding redundant I/O.
- **Parallel Evaluation** – `evaluate_single_completion` executes each test via `test.run(completion)` inside a `ThreadPoolExecutor`, enabling concurrent validation of multiple assertions.
- **Reward Calculation** – The function returns a **float in [0, 1]** representing either the *micro-average* (`passed / total`) or the *macro-average* (mean pass-rate across test types like `present`, `absent`, and `order`).

Both reward modes provide dense supervision signals, but the macro-average variant prevents dominant test categories from masking performance on rare assertion types.

## Data Flow During Training

The GRPO RL trainer orchestrates a specific execution flow for each training step:

1. **Dataset Sampling** – `OlmOCRBenchDataset` (defined at line 373) yields batches containing `pdf_path`, `jsonl_file`, and `test_ids` for each PDF page.
2. **Generation** – The trainer requests `num_completions` outputs from the model, using VLLM when `--vllm_mode` is enabled for accelerated inference.
3. **Reward Computation** – `olmocr_bench_reward` receives the completion texts and metadata, executes the unit test suite, and returns a scalar reward per completion.
4. **Policy Update** – The `GRPOTrainer` computes **group-relative advantages**: each completion’s advantage is calculated relative to the mean reward of its generation batch, inherently reducing variance compared to traditional policy gradients.
5. **Statistics Logging** – `DetailedRewardLoggingCallback` aggregates pass-rates by test category and source file for monitoring.

This flow repeats for `--num_iterations` epochs, with the β (beta) parameter controlling the KL divergence penalty against the reference model.

## Monitoring Reward Statistics

While training progresses, the `DetailedRewardLoggingCallback` class (lines 41‑48) provides granular visibility into unit test performance:

```python

# olmocr/train/grpo_train.py

class DetailedRewardLoggingCallback(TrainerCallback):
    def on_log(...):
        detailed_reward_logger.log_to_wandb(state.global_step)
        detailed_reward_logger.clear()

```

This callback aggregates statistics **by test type** (e.g., presence checks versus ordering constraints) and **by JSONL file**, then reports metrics to Weights & Biases or stdout. Researchers can monitor which specific assertion categories are improving during RL fine-tuning, enabling targeted debugging of model capabilities.

## Command-Line Configuration

To launch training with unit test rewards, use the following pattern:

```bash
python -m olmocr.train.grpo_train \
    --train_bench_data_folder /path/to/bench_data \
    --output_dir ./output \
    --reward_bench ./bench_data/test1.jsonl \
    --beta 0.1 \
    --num_iterations 5 \
    --num_completions 8 \
    --loss_type "policy" \
    --scale_rewards true \
    --reward_bench_macroavg ./bench_data/test1.jsonl \
    --vllm_mode "colocate"

```

Key parameters include:
- `--reward_bench` – Activates the micro-average benchmark reward.
- `--reward_bench_macroavg` – Adds the macro-average variant for balanced test type supervision.
- `--beta` – Sets the KL penalty coefficient passed to `GRPOConfig` (line 1498).
- `--vllm_mode` – Enables colocated VLLM serving for high-throughput generation during GRPO rollouts.

## Summary

- The **GRPO RL trainer** in [`olmocr/train/grpo_train.py`](https://github.com/allenai/olmocr/blob/main/olmocr/train/grpo_train.py) wraps the TRL library's `GRPOTrainer` to fine-tune vision-language models using unit test pass rates as rewards.
- The **`olmocr_bench_reward`** function executes tests from JSONL files in parallel, returning micro-averaged (overall) or macro-averaged (per-type) rewards in the range [0, 1].
- **Group-relative advantages** computed by the GRPO algorithm reduce gradient variance by comparing each completion against the batch mean.
- **`DetailedRewardLoggingCallback`** tracks per-category pass rates, enabling fine-grained analysis of which OCR capabilities improve during training.
- Configuration is handled via command-line flags that populate `GRPOConfig` and select specific reward functions from the benchmark suite.

## Frequently Asked Questions

### What is the difference between micro-average and macro-average rewards in OlmOCR?

**Micro-average** calculates the overall pass rate across all unit tests (`passed / total`), while **macro-average** computes the mean of per-type pass rates (e.g., averaging the pass rates for `present`, `absent`, and `order` test types separately). The micro-average reflects raw accuracy, whereas the macro-average prevents large test categories from dominating the reward signal.

### How does the GRPO trainer reduce variance during policy updates?

The GRPO algorithm computes **group-relative advantages** by subtracting the mean reward of the current generation batch from each completion's individual reward. This relative advantage estimation removes the need for a separate value network and inherently reduces variance compared to traditional advantage actor-critic methods.

### Can I use custom unit test suites with the OlmOCR GRPO trainer?

Yes. The trainer loads tests from JSONL files specified via the `--reward_bench` and `--reward_bench_macroavg` arguments. Each JSONL file contains test definitions with assertions (present, absent, order, etc.) that `olmocr_bench_reward` evaluates against model completions. Custom suites must follow the bench data format expected by `load_specific_tests_cached`.

### What is the role of VLLM in the GRPO training pipeline?

When `--vllm_mode` is set to `"colocate"`, the trainer uses VLLM for accelerated generation of the `num_completions` samples required per GRPO step. This significantly increases throughput during the rollout phase compared to standard Hugging Face generation, especially when computing rewards for multiple completions per prompt.