# How to Perform GRPO Reinforcement Training with olmOCR: A Complete Guide

> Learn GRPO reinforcement training with olmOCR. Fine-tune vision-language models using GRPOTrainer and olmocr_bench_reward for unit-test pass rate rewards. Get the complete guide.

- Repository: [Ai2/olmocr](https://github.com/allenai/olmocr)
- Tags: how-to-guide
- Published: 2026-07-02

---

**GRPO reinforcement training with olmOCR fine-tunes vision-language models using the `GRPOTrainer` from the `trl` library, rewarding completions based on unit-test pass rates calculated by the `olmocr_bench_reward` function.**

OlmOCR 2 includes a dedicated reinforcement learning step called **GRPO (Group Relative Policy Optimization)** to improve PDF text extraction accuracy. This training pipeline, implemented in the `allenai/olmocr` repository, leverages the 🤗 Trainer infrastructure and specialized benchmark datasets to optimize model outputs through policy-gradient updates based on task-specific rewards.

## Prerequisites and Data Structure

Before launching GRPO reinforcement training with olmOCR, you must prepare your environment and organize your benchmark data according to the expected schema.

### Bench Data Layout

The training script expects a folder containing two essential components:

- `pdfs/` – A subdirectory with the raw PDF files referenced by your test suite.
- `*.jsonl` – One or more JSONL files defining unit tests for specific PDF pages, where each line contains at least `pdf`, `page`, `id`, and `type` fields.

You can filter which JSONL files to load using the `--jsonl_filter` argument (accepts a regex pattern). The `OlmOCRBenchDataset` class in [`olmocr/train/dataloader.py`](https://github.com/allenai/olmocr/blob/main/olmocr/train/dataloader.py) scans this folder, loads each PDF page once, and aggregates all test IDs belonging to that specific page.

### Required Dependencies

Install the exact dependency versions used by the olmOCR team:

```bash
pip install "olmocr[train]"
pip install trl==0.22.2 transformers==4.55.2

```

These versions ensure compatibility with the `GRPOTrainer` and the custom reward callbacks implemented in the repository.

## Core Components of the GRPO Pipeline

Understanding the three primary components helps you customize the training process for specific document understanding tasks.

### Dataset Construction (`OlmOCRBenchDataset`)

Located in [`olmocr/train/dataloader.py`](https://github.com/allenai/olmocr/blob/main/olmocr/train/dataloader.py), the `OlmOCRBenchDataset` class handles the heavy lifting of data preparation:

1. **PDF Rendering** – Converts each PDF page to a base64-encoded PNG using `render_pdf_to_base64png`.
2. **Prompt Assembly** – Builds a two-part message using `build_no_anchoring_v4_yaml_prompt` for the text instruction and appends the image token.
3. **Metadata Aggregation** – Returns a dictionary containing `prompt`, `pdf_path`, `jsonl_file`, `test_ids`, and the rendered image for the trainer.

### Reward Computation (`olmocr_bench_reward`)

The reward function, defined in [`olmocr/train/reward.py`](https://github.com/allenai/olmocr/blob/main/olmocr/train/reward.py), evaluates model completions during training:

- **Selective Loading** – Uses `load_specific_tests_cached` to load only the tests required for the current completion rather than the entire benchmark.
- **Pass-Rate Calculation** – By default, computes a micro-averaged pass rate (`overall_passes / total_tests`). Enable macro-averaging with the `--macro_average` flag to average the pass rates across different test types.
- **Detailed Logging** – The `DetailedRewardLogger` tracks per-type statistics (e.g., table extraction vs. math equation recognition) and logs them to Weights & Biases.

### GRPO Trainer Configuration

The main training entry point in [`olmocr/train/grpo_train.py`](https://github.com/allenai/olmocr/blob/main/olmocr/train/grpo_train.py) instantiates `GRPOTrainer` from `trl` with a custom `GRPOConfig`. The trainer:
- Samples completions from the policy model.
- Calls `olmocr_bench_reward` to compute scalar rewards.
- Calculates the group-relative advantage and performs policy-gradient updates.

## Local Training Execution

To run GRPO training on a local GPU or single-node setup, invoke the training module directly with your prepared benchmark data:

```bash

# Prepare your bench folder structure

# bench_data/

# ├── pdfs/

# │   └── research_paper.pdf

# └── validation_tests.jsonl

python -m olmocr.train.grpo_train \
  --train_bench_data_folder ./bench_data \
  --output_dir ./grpo_checkpoints \
  --model_name Qwen/Qwen2.5-VL-7B-Instruct \
  --per_device_train_batch_size 2 \
  --learning_rate 5e-5 \
  --num_iterations 5000 \
  --logging_steps 100 \
  --save_steps 500 \
  --macro_average

```

**Key arguments explained:**
- `--train_bench_data_folder` – Path to the directory containing your `pdfs/` folder and JSONL test definitions.
- `--model_name` – Supports any vision-language model compatible with the Qwen-VL API, including `Qwen2_5_VLForConditionalGeneration` checkpoints.
- `--macro_average` – Changes the reward calculation from micro-averaged (total passes / total tests) to macro-averaged (average of per-type pass rates).

## Distributed Training on Beaker

For large-scale experiments, use the provided Beaker wrapper script [`scripts/train/grpotrainer-beaker.sh`](https://github.com/allenai/olmocr/blob/main/scripts/train/grpotrainer-beaker.sh) to automate environment setup and distributed execution:

```bash
scripts/train/grpotrainer-beaker.sh \
  --model_name s3://ai2-oe-data/models/olmocr-2-7b-fp8 \
  --train_bench_data_folder /data/olmOCR-bench/bench_data \
  --output_dir /weka/oe-training-default/jakep/olmocr-grpo-checkpoints \
  --num_iterations 8000 \
  --logging_steps 200 \
  --save_steps 1000 \
  --macro_average \
  --preemptible

```

Behind the scenes, this script:
1. Builds a Docker image tagged `olmocr-grpo-<VERSION>-<GIT_HASH>`.
2. Syncs benchmark data from S3 (specifically `s3://ai2-oe-data/jakep/olmocr/olmOCR-bench-snapshot-...`).
3. Downloads model checkpoints from S3 if provided with an S3 path.
4. Submits the experiment to the `ai2/olmocr` Beaker workspace with proper distributed rank detection via `get_rank` and `is_main_process`.

The script also configures `S3SyncCallback` to automatically upload checkpoints to your specified S3 bucket after every `save_steps` interval.

## Monitoring and Logging

During GRPO reinforcement training with olmOCR, the `DetailedRewardLoggingCallback` pushes comprehensive metrics to Weights & Biases after each `logging_steps` interval. Monitor these key metrics:

- `bench_reward/overall_pass_rate` – The primary optimization target.
- `bench_reward/total_completions` – Number of completions evaluated.
- `bench_reward/{type}/pass_rate` – Per-test-type pass rates (e.g., `table`, `math`).
- `bench_reward/{jsonl_filename}/pass_rate` – Performance broken down by specific benchmark file.

These statistics appear in the WandB run associated with your experiment, allowing you to identify which document types require additional training iterations.

## Summary

- GRPO training in olmOCR uses `GRPOTrainer` from `trl==0.22.2` to optimize vision-language models on PDF comprehension tasks.
- The `OlmOCRBenchDataset` in [`olmocr/train/dataloader.py`](https://github.com/allenai/olmocr/blob/main/olmocr/train/dataloader.py) prepares data by rendering PDF pages and aggregating test IDs from JSONL files.
- Rewards are computed by `olmocr_bench_reward` in [`olmocr/train/reward.py`](https://github.com/allenai/olmocr/blob/main/olmocr/train/reward.py) using either micro-averaged or macro-averaged pass rates.
- Run locally via `python -m olmocr.train.grpo_train` or at scale using [`scripts/train/grpotrainer-beaker.sh`](https://github.com/allenai/olmocr/blob/main/scripts/train/grpotrainer-beaker.sh) for distributed Beaker experiments.
- The `DetailedRewardLogger` provides granular WandB logging of per-type and per-file performance statistics.

## Frequently Asked Questions

### What is GRPO and why does olmOCR use it?

GRPO (Group Relative Policy Optimization) is a reinforcement learning algorithm that optimizes language models by comparing generated completions within a group and computing relative advantages. OlmOCR uses GRPO instead of traditional supervised fine-tuning because it allows the model to learn from binary reward signals (test pass/fail) rather than needing perfect gold-standard outputs for every PDF page, enabling optimization on complex extraction tasks where multiple valid output formats exist.

### What data format does olmOCR expect for GRPO training?

The training pipeline requires a benchmark data folder containing a `pdfs/` subdirectory with your PDF files and one or more `.jsonl` files defining unit tests. Each JSONL line must include `pdf` (filename), `page` (number), `id` (unique test identifier), and `type` (test category) fields. The `OlmOCRBenchDataset` scans these files to build training prompts that combine rendered page images with instructions generated by `build_no_anchoring_v4_yaml_prompt`.

### How does the reward function calculate pass rates?

The `olmocr_bench_reward` function evaluates completions against cached unit tests loaded via `load_specific_tests_cached`. By default, it calculates a **micro-averaged** reward equal to `overall_passes / total_tests`. If you specify `--macro_average`, the reward becomes the average of per-type pass rates, giving equal weight to rare test types and common ones. The function supports detailed logging through `DetailedRewardLogger` to track performance across different document elements like tables or mathematical expressions.

### Can I run GRPO training without Beaker?

Yes, you can run GRPO training entirely on local hardware or standard cloud VMs by installing the required dependencies (`trl==0.22.2`, `transformers==4.55.2`) and invoking `python -m olmocr.train.grpo_train` with the appropriate `--train_bench_data_folder` and `--output_dir` arguments. The Beaker wrapper script ([`grpotrainer-beaker.sh`](https://github.com/allenai/olmocr/blob/main/grpotrainer-beaker.sh)) is optional and primarily automates Docker builds, S3 data synchronization, and distributed rank management for AI2's internal infrastructure.