Training GRPO on GSM8K in Arena¶

This tutorial walks through the full Arena workflow for LLM reasoning: registering a dataset, writing a reward function, submitting a training job, monitoring progress, and deploying the fine-tuned model. We fine-tune Qwen/Qwen2.5-0.5B-Instruct with GRPO on GSM8K, a dataset of grade-school math word problems, as our running example.

It is the language-model counterpart to Training PPO on a Custom Environment in Arena. Where that tutorial validates a custom Gymnasium environment, here we register a Hugging Face dataset and supply a reward function instead.

Prerequisites¶

Installation¶

Arena is provided by the separate agilerl-arena distribution (agilerl.arena.* in the shared namespace). It pulls in lightweight client dependencies (httpx, rich, etc.) rather than core training stacks. Install it directly, or through the AgileRL extra:

pip install agilerl-arena
# or
pip install "agilerl[arena]"

Authentication¶

All Arena operations require authentication. You can authenticate in one of two ways:

API Key: set the ARENA_API_KEY environment variable:
```
export ARENA_API_KEY="arena_pat..."
```
Note

Personal access tokens can be found in the Arena dashboard under Profile Management → CLI API Key.
Device login (interactive): run the login command, which opens a browser for OAuth authentication:
Python
from agilerl.arena import ArenaClient client = ArenaClient() client.login()
CLI
arena login
Credentials are persisted locally so you only need to authenticate once per machine.

The Task¶

GRPO (Group Relative Policy Optimization) rewards a model for solving problems that have a verifiable answer. Instead of training a separate critic to estimate value, GRPO samples a group of completions for each prompt and scores each one relative to the group mean. This removes the critic network and gives a stable training signal, which makes it a good fit for reasoning tasks.

Our task is GSM8K. Each row is a question (a math word problem) and an answer whose final line holds the solution after a #### marker:

question: Natalia sold clips to 48 of her friends in April, and then she sold half as many
          clips in May. How many clips did she sell altogether in April and May?
answer:   Natalia sold 48/2 = 24 clips in May.
          Altogether she sold 48+24 = 72 clips.
          #### 72

We reward the model for producing the correct final number, and for wrapping its reasoning and answer in the format we ask for.

Register the Dataset¶

Before training, we register the dataset on Arena. We import GSM8K straight from Hugging Face and map its columns onto the keys Arena’s reasoning datasets expect (question and answer). GSM8K already uses those names, so the mapping is one-to-one.

Python

client.create_dataset(
    name="gsm8k",
    category="reasoning",
    hf_dataset_name="openai/gsm8k",
    hf_config="main",
    hf_split="train",
    column_mapping={"question": "question", "answer": "answer"},
)

CLI

arena datasets create gsm8k \
    --category reasoning \
    --hf-dataset openai/gsm8k \
    --hf-config main \
    --hf-split train \
    --column-mapping '{"question": "question", "answer": "answer"}'

Datasets fall into one of three categories, and each expects its own column mapping keys:

reasoning: question and answer pairs for algorithms like GRPO, as here. Expects {"question": "<column>", "answer": "<column>"}.
preference: a prompt with a chosen and a rejected completion, for algorithms like DPO. Expects {"prompt": "<column>", "chosen": "<column>", "rejected": "<column>"}.
sft: supervised fine-tuning pairs mapping a prompt to a target completion. Expects {"prompt": "<column>", "target": "<column>"}.

GRPO trains on reasoning datasets.

You can browse datasets already registered for your organization, or search Hugging Face, with arena datasets list (add --search to query Hugging Face).

The Reward Function¶

A reasoning job needs a reward function: a Python file that scores each completion the model produces. Arena calls a function named reward with three arguments, question, answer, and completion, and expects a single float back. The higher the reward, the better the completion.

Our reward has two parts. A correctness reward parses the gold answer after ####, extracts the final number from the model’s <answer> block, and returns 1.0 when they match. A format reward returns 1.0 when the completion wraps its reasoning and answer in <think>...</think> and <answer>...</answer> tags. The most a completion can score is 2.0. We never tell the model which answer or format to produce; it discovers both from the reward.

reward.py

import re


def reward(question: str, answer: str, completion: str) -> float:
    """Combined GSM8K reward: answer correctness plus output format."""
    return correctness_reward(answer, completion) + format_reward(completion)


def extract_gold_answer(answer: str) -> str | None:
    """Return the final numeric answer from a GSM8K answer (the value after ``####``)."""
    match = re.search(r"####\s*(-?[\d,]+)", answer)
    return match.group(1).replace(",", "").strip() if match else None


def extract_model_answer(completion: str) -> str | None:
    """Return the final number inside the model's ``<answer>...</answer>`` block."""
    block = re.search(r"<answer>([\s\S]*?)</answer>", completion)
    if block is None:
        return None
    number = re.search(r"-?[\d,]+", block.group(1))
    return number.group(0).replace(",", "").strip() if number else None


def correctness_reward(answer: str, completion: str) -> float:
    """Reward 1.0 when the model's final answer matches the gold answer, else 0.0."""
    gold = extract_gold_answer(answer)
    predicted = extract_model_answer(completion)
    if gold is None or predicted is None:
        return 0.0
    return 1.0 if predicted == gold else 0.0


def format_reward(completion: str) -> float:
    """Reward 1.0 when the completion follows ``<think>...</think>\\n<answer>...</answer>``."""
    text = "<think>" + completion
    regex = r"^<think>[\s\S]*?</think>\n<answer>[\s\S]*?</answer>$"
    return 1.0 if re.match(regex, text.strip()) else 0.0

The reward function runs on Arena’s infrastructure during training. Arena also runs it once when you submit the job, as a quick check that it imports and returns a number.

Creating a Project¶

Before submitting a training job, we need a project to submit it to. Language-model runs live in an LLM-based project, so we pass that flag when creating it.

Python

client.create_project(name="GSM8K Tutorial", description=None, llm_based=True)

CLI

arena projects create "GSM8K Tutorial" --llm-based

Tip

You can set a default project to work on by using the arena projects set-default <project-name> CLI command.

Submit a Training Job¶

We define the training configuration in a YAML manifest (gsm8k_grpo.yaml). The environment section references the dataset by the name we registered it under. The network section names the base model and its LoRA adapter, so training only updates the adapter weights.

gsm8k_grpo.yaml

---
algorithm:
    name: GRPO
    batch_size: 6
    beta: 0.001
    lr: 0.000005
    clip_coef: 0.2
    max_grad_norm: 0.1
    update_epochs: 1
    group_size: 8
    temperature: 0.9
    use_vllm: false
    calc_position_embeddings: true
    reduce_memory_peak: false

environment:
    name: gsm8k

training:
    max_steps: 100_000
    evo_steps: 5
    pop_size: 2
    eval_loop: 1
    evaluation_interval: 5

mutation:
    mutation_sd: 0.1
    rand_seed: 42
    probabilities:
        no_mut: 0.1
        arch_mut: 0.0
        new_layer: 0.0
        params_mut: 0.0
        act_mut: 0.0
        rl_hp_mut: 0.6
    rl_hp_selection:
        lr:
            min: 0.0000001
            max: 0.00001
        beta:
            min: 0.0001
            max: 0.01

tournament_selection:
    tournament_size: 2
    elitism: true

network:
    pretrained_model_name_or_path: Qwen/Qwen2.5-0.5B-Instruct
    max_context_length: 512
    lora_config:
        lora_r: 16
        lora_alpha: 64
        target_modules:
            - q_proj
            - k_proj
            - v_proj
            - o_proj
        lora_dropout: 0.05
        task_type: CAUSAL_LM

Unlike the PPO tutorial, there is no observation space for Arena to infer an encoder from. The model architecture comes entirely from pretrained_model_name_or_path and the lora_config.

We submit the manifest together with the reward function. Passing --reward-file tells Arena this is a reasoning job, and it validates the reward function against a sample completion before the run starts.

Python

result = client.submit_experiment(
    manifest="gsm8k_grpo.yaml",
    reward_file="reward.py",
    resource_id="arena-medium",
    num_nodes=2,
    project="GSM8K Tutorial",
    experiment_name="gsm8k-grpo-v1",
)

CLI

arena experiments submit gsm8k_grpo.yaml \
    --reward-file reward.py \
    --resource-id arena-medium \
    --num-nodes 2 \
    --project 'GSM8K Tutorial' \
    --experiment-name gsm8k-grpo-v1

Tip

You can view all of the available resources to train on by running the CLI command arena resources list.

Monitor Training¶

You can monitor training progress directly from the Arena dashboard, or download metrics programmatically:

Python

# Download training metrics
client.download_experiment_metrics(
    experiment_name="gsm8k-grpo-v1",
    output_path="metrics.csv",
)

# List available checkpoints
checkpoints = client.list_checkpoints(
    experiment_name="gsm8k-grpo-v1"
)
print(checkpoints)

CLI

# Download metrics
arena experiments metrics gsm8k-grpo-v1 --output-file metrics.csv

# List checkpoints
arena experiments checkpoints gsm8k-grpo-v1

Deploy the Trained Model¶

Once training is complete, deploy the best checkpoint to an inference endpoint:

Python

# Deploy the best checkpoint
client.deploy_agent(experiment_name="gsm8k-grpo-v1")

# Or deploy a specific checkpoint
client.deploy_agent(
    experiment_name="gsm8k-grpo-v1",
    checkpoint="step_500",
)

CLI

# Deploy the best checkpoint
arena agent deploy gsm8k-grpo-v1

# Deploy a specific checkpoint
arena agent deploy gsm8k-grpo-v1 --checkpoint step_500

Interact with the Deployed Model¶

After deployment, open the model by name and send it a prompt to receive a completion. The deployment name matches the experiment you deployed; you can also list your deployments with arena agent list.

Python

from agilerl.arena import ArenaClient

prompt = "Weng earns $12 an hour for babysitting. Yesterday, she just did 50 minutes of babysitting. How much did she earn?"

with ArenaClient() as client:
    with client.open_inference_agent("gsm8k-grpo-v1") as agent:
        # Generate a full completion
        result = agent.generate(prompt)
        print(result.results[0].completion)

        # Or stream tokens as they are produced
        for chunk in agent.generate_stream(prompt):
            print(chunk, end="")

CLI

# Always streams tokens as they are generated
arena agent generate gsm8k-grpo-v1 \
    --prompt "Weng earns \$12 an hour for babysitting. Yesterday, she just did 50 minutes of babysitting. How much did she earn?"

Tip

Set a default agent with arena agent run gsm8k-grpo-v1 and you can drop the deployment name from later commands. arena agent generate --prompt "..." then uses the active agent.