LLM Reasoning Tutorial¶
Reinforcement learning enhances an LLM’s reasoning by rewarding it for solving problems with verifiable answers in domains like math, coding, and science. Unlike standard training where feedback comes token-by-token, this approach evaluates complete reasoning sequences. The process typically involves:
Starting with a base LLM that has decent language capabilities
Presenting it with reasoning-intensive problems that have objectively correct answers
Having the model generate entire solution paths, including all intermediate steps
Computing rewards based on correctness of the final answer, regardless of the specific reasoning path taken
Using these rewards to update the model’s parameters through algorithms like PPO or GRPO
In this tutorial, we will be using the Group Relative Policy Optimization (GRPO) algorithm to finetune a large language model and achieve emergent reasoning behaviour using the AgileRL open-source framework. To read more about LLM reasoning in general, check out the page in our docs. Visit the AgileRL GitHub repository for more information about the library.
GRPO Overview¶
GRPO (Group Relative Policy Optimization) is an elegant simplification of PPO (Proximal Policy Optimization) that makes reinforcement learning more computationally efficient, especially for large language models.
The two key innovations are:
Eliminating the critic network: Instead of training a separate value function to estimate expected rewards (which requires additional compute and memory), GRPO normalizes rewards across a batch of samples. It calculates advantage by subtracting the mean reward from each sample’s reward and dividing by the standard deviation.
Group-based evaluation: GRPO generates multiple outputs using the same policy, evaluates them as a group, and then updates the model. This approach reduces variance in the training signal by smoothing out the randomness inherent in probabilistic environments.
These changes are particularly valuable for LLM training because they reduce computational overhead by removing the need for a separate critic model, provide more stable gradient updates in environments with sparse or noisy rewards, and they simplify implementation while maintaining or improving performance.
Dependencies¶
import re
import torch
from datasets import load_dataset
from peft import LoraConfig
from torch.utils.data import Dataset
from transformers import AutoTokenizer
from agilerl.algorithms import GRPO
from agilerl.training.llm import train_llm_rollout
from agilerl.llm_envs import RolloutHarness
Defining our base model and dataset¶
In this tutorial, we use the open-source transformers and datasets libraries from Hugging Face to download our pretrained model weights and training data. There are a huge number of models and datasets hosted on Hugging Face, and different ones can easily be substituted in. In this tutorial, to keep things simple and inexpensive, we will use a 3 billion parameter Qwen model, and the Countdown dataset:
Create Model and Dataset
MODEL_PATH = "Qwen/Qwen2.5-3B"
DATASET = "Jiayi-Pan/Countdown-Tasks-3to4"
def make_dataset(dataset_name: str) -> tuple[Dataset, Dataset]:
raw_dataset = (
load_dataset(dataset_name, split="train").shuffle(seed=42).select(range(50000))
)
raw_dataset = raw_dataset.rename_column("target", "answer")
raw_dataset = raw_dataset.rename_column("nums", "question")
train_test_split = raw_dataset.train_test_split(test_size=0.1)
train_dataset = train_test_split["train"]
test_dataset = train_test_split["test"]
return train_dataset, test_dataset
# Instantiate the model and the associated tokenizer
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
tokenizer.pad_token = tokenizer.eos_token
train_dataset, test_dataset = make_dataset(DATASET)
Create the Reasoning Environment¶
From model to agent: In reinforcement learning, models are called agents. This is because they are trained by taking actions, receiving rewards, and learning from this feedback. This enables them to become very good at taking actions to solve tasks - to develop agency. Since we are training our model with reinforcement learning, it becomes an agent through this process.
We must create a reinforcement learning environment in which our agent can explore possible
solutions and learn to optimise rewards. AgileRL provides a RolloutHarness
class that turns lists of questions and answers into a reinforcement learning, gymnasium-style environment.
Reasoning is the single-turn case, so we build the environment with max_turns=1.
So, how does the environment know how to reward an agent for its outputs? Well, we must define a reward_function that the agent learns to optimise. Following the techniques used in the DeepSeek reasoning paper, we will define our reward function as the sum of two rewards:
Accuracy rewards: Verifying answers against ground truth. In this tutorial, we will reward the model +1 if the final answer it produces is correct, otherwise 0.
Format rewards: Encouraging structured reasoning with explicit steps. In this tutorial, we will reward the model +1 if it puts its thinking process between ‘<think>’ and ‘</think>’ tags, otherwise 0.
Therefore, the maximum score an agent can receive is 2, if it produces the correct answer in the correct format. The key here is that we never tell the agent which answer it should produce or which format it should use. By giving it rewards for displaying these behaviours, the agent itself discovers the best way to achieve high rewards and learns the behaviour we desire.
Reward Functions
def format_reward_func(completions, target, **kwargs):
rewards = []
for completion, gt in zip(completions, target):
try:
# add synthetic <think> as its already part of the prompt and prefilled for the assistant to more easily match the regex
completion = "<think>" + completion
regex = r"^<think>([^<]*(?:<(?!/?think>)[^<]*)*)<\/think>\n<answer>([\s\S]*?)<\/answer>$"
match = re.search(regex, completion, re.DOTALL)
if match is None or len(match.groups()) != 2:
rewards.append(0.0)
else:
rewards.append(1.0)
except Exception:
rewards.append(0.0)
return rewards
def equation_reward_func(completions, target, nums, **kwargs):
rewards = []
for completion, gt, numbers in zip(completions, target, nums):
try:
# add synthetic <think> as its already part of the prompt and prefilled for the assistant to more easily match the regex
completion = "<think>" + completion
answer_tags = re.findall(r"<answer>([\s\S]*?)<\/answer>", completion)
if len(answer_tags) != 1:
rewards.append(0.0)
continue
equation = answer_tags[0].strip()
used_numbers = [int(n) for n in re.findall(r"\d+", equation)]
if sorted(used_numbers) != sorted(numbers):
rewards.append(0.0)
continue
allowed_pattern = r"^[\d+\-*/().\s]+$"
if not re.match(allowed_pattern, equation):
rewards.append(0.0)
continue
result = eval(equation, {"__builtins__": None}, {})
if abs(float(result) - float(gt)) < 1e-5:
rewards.append(1.0)
else:
rewards.append(0.0)
except Exception as e:
rewards.append(0.0)
return rewards
def combined_rewards(completion, solution, prompt):
reward = (
equation_reward_func([completion], [solution], [prompt])[0]
+ format_reward_func([completion], [solution])[0]
)
return reward
Now we have defined our reward functions, we must also design our prompt. This forms the input given
to the agent and provides the context necessary to complete the task. This is a task-specific feature,
and different reasoning problems will require different conversation templates, although they can follow a similar
format. We define the conversation template as follows (using question and answer as placeholders for the question and answer data)
and then drive a single-turn rollout env over the question and answer
columns of our dataset with an in-process env_factory (a prompt dataset is just an
environment: each rollout runs its own env instance in-process via
RolloutHarness.local).
Build the Single-Turn Rollout Environment
conversation_template = [
{
"role": "system",
"content": "You are a helpful assistant. You first think about the reasoning process in your mind and then provide the user with the answer.",
},
{
"role": "user",
"content": (
"Using each number in this list only once {question}, create an equation that equals {answer}. You "
"can use basic arithmetic operations (+, -, *, /) and each number can only be used once. Show your work in <think> </think> tags. "
"And return the final equation and answer in <answer> </answer> tags, for example <answer>(1 + 2) / 3</answer>."
)
},
{"role": "assistant", "content": "Let me solve this step by step.\n<think>"},
]
def prompt_builder(question: str) -> str:
parts = [
m["content"].format(question=question, answer="")
for m in conversation_template
]
return "\n".join(p for p in parts if p)
# A single-turn rollout environment from the dataset — a prompt dataset is
# just an env, driven in-process by RolloutHarness.local.
class QADataset:
"""Single-turn dataset env: a question on reset, a score on step."""
def __init__(self, questions, answers, reward_fn, prompt_builder,
test_questions=None, test_answers=None):
self.questions, self.answers = questions, answers
self.test_questions, self.test_answers = test_questions, test_answers
self.reward_fn, self.prompt_builder = reward_fn, prompt_builder
self._cursor, self._split = 0, ""
@property
def dataset_size(self) -> int:
return len(self.questions)
def reset(self, seed=None, *, row_index=None, evaluation=None):
if evaluation and self.test_questions is not None:
qs, ans, split = self.test_questions, self.test_answers, "eval"
else:
qs, ans, split = self.questions, self.answers, "train"
if row_index is None:
if split != self._split:
self._cursor, self._split = 0, split
row_index, self._cursor = self._cursor, self._cursor + 1
self._q, self._a = qs[row_index % len(qs)], ans[row_index % len(ans)]
return self.prompt_builder(self._q), {}
def step(self, action):
return "", float(self.reward_fn(action, self._a, self._q)), True, False, {}
env_factory = lambda: RolloutHarness.local(
QADataset(
questions=list(train_dataset["question"]),
answers=list(train_dataset["answer"]),
reward_fn=combined_rewards,
prompt_builder=prompt_builder,
test_questions=list(test_dataset["question"]),
test_answers=list(test_dataset["answer"]),
),
tokenizer,
max_turns=1,
pad_id=tokenizer.pad_token_id,
apply_chat_template=True,
max_model_len=1024,
)
Create a GRPO Agent¶
To allow our model to become an agent and learn through reinforcement learning, we can use the
GRPO class. This class follows the same structure as the other
reinforcement learning algorithms in the AgileRL library.
An important part of training an LLM to display reasoning behaviour is distributed training. They are
called Large Language Models for a reason, and unless you are a very lucky individual, you may not
have enough capacity on your individual computer to train even a ‘small’ LLM. If you want to train a
larger, more powerful model, then this becomes even more infeasible. Instead, we can leverage distributed
training, to share the workload across multiple devices and speed up training. Launch multi-GPU LLM
training with torchrun. See Multi-GPU LLM training.
agent = GRPO(
model_name=MODEL_PATH,
pad_token_id=tokenizer.eos_token_id,
pad_token=tokenizer.eos_token,
max_output_tokens=1024,
batch_size=4,
group_size=12,
vllm_config=VLLMConfig(
sleep_mode=True,
max_num_seqs=4
),
)
Training and Saving an Agent¶
The simplest way to train an AgileRL agent is to use the train_llm_rollout() function
with max_turns=1 for single-turn reasoning. This training function will orchestrate the training process, removing the need to implement a training loop, and will save
checkpoints of the trained agent that can be used later for inference. It also uses Weights and Biases for tracking.
train_llm_rollout(
pop=[agent],
max_turns=1,
env_factory=env_factory,
evaluation_interval=10,
wb=True,
save_elite=True,
elite_path="path/to/model/directory",
max_reward=2.0,
evo_steps=10,
)
Launching multi-GPU training¶
Launch from the repository root with torchrun (N is the number of GPUs):
torchrun --nproc_per_node=N path/to/training_script
Using a custom training loop¶
If you need lower-level control than train_llm_rollout(),
build a RolloutCollector and collect trajectories with
collect_rollouts_llm(). This is the same rollout API the trainer uses internally.
Do not use dataset-env calls like env.reset(reset_dataloaders=True) / env.step(token_ids)
for rollout training.
Custom Training Loop
import numpy as np
from agilerl.llm_envs import RolloutCollector
from agilerl.rollouts.on_policy import collect_rollouts_llm
batch_size = agent.batch_size
group_size = getattr(agent, "group_size", 1)
rollout_env = RolloutCollector(env_factory, batch_size, group_size)
group_seed = int(np.random.randint(0, 1_000_000))
try:
for i in range(max_steps):
(
token_ids_list,
action_masks_list,
all_turn_ids,
all_rewards,
batch_steps,
group_seed,
all_sampling_logps,
_all_pixel_values,
all_sampling_supports,
) = collect_rollouts_llm(
agent=agent,
env=rollout_env,
n_steps=1, # single-turn reasoning
batch_size=batch_size,
group_size=group_size,
group_seed=group_seed,
)
# Build the experience tuple your algorithm expects.
experiences = (token_ids_list, action_masks_list, all_rewards)
learn_kwargs = {"turn_ids": all_turn_ids}
if all_sampling_logps is not None:
learn_kwargs["sampling_logps"] = all_sampling_logps
if all_sampling_supports is not None:
learn_kwargs["sampling_supports"] = all_sampling_supports
metrics = agent.learn(experiences, **learn_kwargs)
# Add your logging / eval / checkpointing here.
finally:
rollout_env.close()
Loading a Trained Agent for Inference¶
Once we have finetuned our LLM, we may want to use it for inference. Below outlines how to load the model in this tutorial.
Load fine-tuned LLM into vLLM Engine for inference¶
from vllm import LLM
llm = LLM(
model=MODEL_PATH,
tensor_parallel_size=1,
gpu_memory_utilization=0.9,
max_num_seqs=1024,
max_model_len=1536,
distributed_executor_backend="external_launcher",
seed=0,
model_impl="vllm",
enable_lora=True,
)
sampling_params = SamplingParams(
temperature=0.0,
top_p=1.0,
top_k=-1,
max_tokens=max_output_tokens,
seed=42,
)
prompts = "Using each number in this list only once 33, 19, 27, 5, create an equation that equals 82. You can use basic arithmetic operations (+, -, *, /) and each number can only be used once."
outputs = llm.generate(
prompts,
sampling_params=sampling_params,
lora_request=LoRARequest(
lora_name="trained_model",
lora_int_id=1,
lora_path=checkpoint_path + "/actor",
),
)
# Print the completion
print(outputs[0].outputs[0].text)
The Aha Moment¶
We finetuned a Qwen2.5-3B using the full training code below and witnessed the model having ‘aha’ moments. In a notable example from our training run the model literally exclaims “Aha!” as it reasons its way to the correct solution:
58 - 1 = 57, then 57 - 25 = 32, and finally 32 + 1 = 33.
However, 33 is not 82. Let's try another combination: 58 + 25 = 83, but we need 82.
Aha! If we use 58 + 1 = 59, then 59 - 25 = 34, and finally 34 - 1 = 33. This doesn't work either.
Hmm... what if we use 58 + 1 = 59, then 59 - 25 = 34, and finally 34 + 1 = 35.
Nope... closer, but not quite. What if we try 58 + 1 = 59, then 59 - 25 = 34, and finally 34 + 25 = 59.
Nope... still not 82.
Ah-ha! One more try: 58 - 1 = 57, then 57 + 25 = 82.</think>
<answer>(58 - 1) + 25</answer>
Full Training Code¶
Full code
# Copyright 2026 AgileRL
# SPDX-License-Identifier: Apache-2.0
import re
from datasets import load_dataset
from peft import LoraConfig
from torch.utils.data import Dataset
from transformers import AutoTokenizer
from agilerl.algorithms import GRPO
from agilerl.training.llm import train_llm_rollout
from agilerl.utils.algo_utils import VLLMConfig
from agilerl.llm_envs import RolloutHarness
from agilerl.llm_envs.rubrics import reward_fn_to_rubric
MODEL_PATH = "Qwen/Qwen2.5-0.5B"
DATASET = "Jiayi-Pan/Countdown-Tasks-3to4"
MAX_CONTEXT_LENGTH = 1024
def make_dataset(dataset_name: str) -> tuple[Dataset, Dataset]:
raw_dataset = (
load_dataset(dataset_name, split="train").shuffle(seed=42).select(range(50000))
)
raw_dataset = raw_dataset.rename_column("target", "answer")
raw_dataset = raw_dataset.rename_column("nums", "question")
train_test_split = raw_dataset.train_test_split(test_size=0.2)
train_dataset = train_test_split["train"]
test_dataset = train_test_split["test"]
return train_dataset, test_dataset
def format_reward_func(completions, target, **kwargs):
rewards = []
for completion, _gt in zip(completions, target, strict=False):
try:
# add synthetic <think> as its already part of the prompt and prefilled for the assistant to more easily match the regex
completion = "<think>" + completion
regex = r"^<think>([^<]*(?:<(?!/?think>)[^<]*)*)<\/think>\n<answer>([\s\S]*?)<\/answer>$"
match = re.search(regex, completion, re.DOTALL)
if match is None or len(match.groups()) != 2:
rewards.append(0.0)
else:
rewards.append(1.0)
except Exception: # noqa: PERF203
rewards.append(0.0)
return rewards
def equation_reward_func(completions, target, nums, **kwargs):
rewards = []
for completion, gt, numbers in zip(completions, target, nums, strict=False):
try:
# add synthetic <think> as its already part of the prompt and prefilled for the assistant to more easily match the regex
completion = "<think>" + completion
answer_tags = re.findall(r"<answer>([\s\S]*?)<\/answer>", completion)
if len(answer_tags) != 1:
rewards.append(0.0)
continue
equation = answer_tags[0].strip()
used_numbers = [int(n) for n in re.findall(r"\d+", equation)]
if sorted(used_numbers) != sorted(list(numbers)):
rewards.append(0.0)
continue
allowed_pattern = r"^[\d+\-*/().\s]+$"
if not re.match(allowed_pattern, equation):
rewards.append(0.0)
continue
result = eval(equation, {"__builtins__": None}, {})
if abs(float(result) - float(gt)) < 1e-5:
rewards.append(1.0)
else:
rewards.append(0.0)
except Exception:
rewards.append(0.0)
return rewards
def combined_rewards(completion, solution, prompt):
reward = (
equation_reward_func([completion], [solution], [prompt])[0]
+ format_reward_func([completion], [solution])[0]
)
if reward == 2.0:
with open("countdown_completions.txt", "a") as text_file:
text_file.write(
f"Prompt {prompt}" + "\n" + completion + "\n" + "=" * 50 + "\n",
)
return reward
def main():
# Instantiate the model and the associated tokenizer
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
train_dataset, test_dataset = make_dataset(DATASET)
# Define the conversation template
conversation_template = [
{
"role": "system",
"content": "You are a helpful assistant. You first think about the reasoning process in your mind and then provide the user with the answer.",
},
{
"role": "user",
"content": "Using each number in this list only once {question}, create an equation that equals {answer}. You can use basic arithmetic operations (+, -, *, /) and each number can only be used once. Show your work in <think> </think> tags. And return the final equation and answer in <answer> </answer> tags, for example <answer>(1 + 2) / 3</answer>.",
},
{"role": "assistant", "content": "Let me solve this step by step.\n<think>"},
]
def prompt_builder(row: dict) -> str:
parts = [
m["content"].format(question=row["question"], answer=row["answer"])
for m in conversation_template
]
return "\n".join(p for p in parts if p)
# Dataset rows plus a reward function, as a single-turn env.
def env_factory(evaluation_mode: bool = False):
del evaluation_mode
return RolloutHarness.from_dataset(
train_dataset,
reward_fn_to_rubric(combined_rewards),
tokenizer,
test_dataset=test_dataset,
prompt_builder=prompt_builder,
apply_chat_template=True,
max_model_len=MAX_CONTEXT_LENGTH,
)
# Define the LoRA configuration
lora_config = LoraConfig(
r=16,
lora_alpha=64,
target_modules=[
"q_proj",
"k_proj",
"v_proj",
"o_proj",
"up_proj",
"down_proj",
"gate_proj",
],
task_type="CAUSAL_LM",
lora_dropout=0.05,
)
# Instantiate the grpo agent
agent = GRPO(
model_name=MODEL_PATH,
pad_token_id=tokenizer.pad_token_id,
pad_token=tokenizer.pad_token,
lora_config=lora_config,
batch_size=16,
max_model_len=MAX_CONTEXT_LENGTH,
group_size=8,
vllm_config=VLLMConfig(sleep_mode=True, max_num_seqs=4),
)
train_llm_rollout(
pop=[agent],
max_turns=1,
env_factory=env_factory,
evaluation_interval=10,
wb=True,
save_elite=True,
elite_path="checkpoints",
max_reward=2.0,
)
if __name__ == "__main__":
main()