TraceRL

A library for post-training large language models with reinforcement learning, built for agentic behavior rather than single-step reasoning.

Motivation

This project began out of a recurring frustration with the state of LLM-RL tooling. During a research project earlier this year, I found myself wanting a codebase built principally around agentic workflows — not single-turn reasoning — and unable to find one that matched how I wanted to think about the problem. Several strong frameworks exist for this purpose as of late 2025, but I built TraceRL because its underlying abstractions are, I believe, conceptually distinct from what is already out there.

Much of this year went into thinking through what an LLM-RL library's architecture should look like to support fast prototyping and genuine research iteration. When my last research project wrapped up in mid-November, I used the downtime to rewrite the codebase from the ground up and release it publicly.

Some of the design choices here will be familiar to anyone who has worked with other LLM-RL frameworks. Others depart from convention substantially. The sections below explain why.

Design Philosophy

TraceRL is a library, not a framework. Components are loosely coupled by design, so individual pieces can be removed or replaced without disturbing the rest of the system.

Agents and Environments are strictly separated.

• Environments are pure state-transition functions that may emit rewards on transition. They support multiple agents by design.
• Agents are LLMs with state, wrapped in a harness responsible for context management, parsing model outputs into the format an environment expects, and coordinating internal or external tools.
• The Interaction Protocol defines the agent–environment loop explicitly, which lets any agent harness pair with any compatible environment. An InteractionProtocol exposes a run() method that returns a list of Rollouts, one per agent perspective.

Rollouts capture both AgentSteps and EnvironmentSteps.

• An AgentStep records every model call, including internal tool loops, and carries a TokenTrace used for training.
• An EnvironmentStep records a state transition from env.step() and references the AgentSteps that produced the action.
• This distinction exists because training requires the full reasoning trace. A ReAct-style agent might invoke tools several times before producing a final action, and every one of those calls contributes tokens worth training on. Online batching concatenates all AgentSteps within a turn into a single training sample; see CONSIDERATIONS.md for details.

The Trainer is decoupled from the concept of an episode. The Trainer requests its next batch from a batch source via batch_source.get_next_batch(), which yields State-Action-Weight triples. The weight may represent a ground-truth reward, a rollout return, a group-relative advantage, or another signal entirely. Because the Trainer only deals in these triples, it has no notion of episodes or rollouts. A batch source can be a static dataset for offline methods like SFT, or a Rollout Engine connected to a live vLLM server for online RL.

An RL algorithm is a strategy that plugs into the Trainer. Each algorithm defines a Credit Assigner and a Loss — together, these two components are sufficient to express nearly any RL algorithm. The credit assigner determines how much credit a given state-action pair (prompt-completion pair, in LLM terms) receives; in GRPO, for instance, credit is computed relative to other rollouts in the group.

Batch composition is controlled at two levels.

• Intra-batch control governs how rollouts within a batch relate to one another. GRPO, for example, is implemented as an intra-batch mechanism that expands each rollout request into group_size copies before dispatching them to the rollout engine.
• Inter-batch control governs dependencies across batches over the course of training, covering concepts such as curriculum learning.

A Note on Maturity

This codebase is best understood as aspirational. It targets researchers who value clean abstractions and conceptual clarity over production hardening, and it does not aim to replace any existing framework. It has not been battle-tested at scale and should be treated as closer to an intellectual exercise than a production system. That said, extending it with new algorithms, agent harnesses, or environments — including with the help of tools like Claude Code or Codex — has been straightforward in practice.

Repository Layout

The codebase has two main components: the library itself, and a set of runnable scripts demonstrating end-to-end usage.

Path	Contents
src/tracerl/	Core library: agents and context management, environment interfaces and built-in environments, interaction protocols, inference clients, training and batching logic, evaluation, and distributed weight pushing.
examples/	Reference scripts that assemble the library's components into complete training and evaluation runs.
environments/	Standalone environments and configs, importable or runnable in isolation.
data/	Small datasets and artifacts used by select examples.
tests/	Unit and integration tests (pytest markers include integration and gpu).
scripts/	Standalone utilities, e.g. calibrate_micro_batch.py for micro-batch sizing; see scripts/README.md.

For details on truncation semantics — environment time limits versus protocol cutoffs versus model finish reasons — see CONSIDERATIONS.md.

Logging

Training statistics use canonical prefixes: train/, eval/, and perf/ (for example, train/loss, eval/accuracy, perf/gpu_mem_alloc_mb). The train/step and eval/step fields are used by loggers to annotate panels and runs. For Weights & Biases logging, set WANDB_PROJECT to your preferred project name; it defaults to TraceRL if unset.

Training Configuration

• rollouts_per_update: number of rollouts collected per trainer step. Must be divisible by group_size for GRPO-style grouping.
• max_seq_len: maximum token length permitted for any single sample. The trainer raises an error if this is exceeded.
• micro_token_budget: maximum padded tokens allowed per micro-batch. The collator buckets and splits each macro-batch to respect this limit.

Examples

Tic-Tac-Toe (examples/tic_tac_toe/) — a lightweight environment for iterating on the full stack without significant sampling cost.

• Online RL: train_tic_tac_toe.py performs LoRA fine-tuning with GRPO-style group-relative credit assignment.
• SFT: sft_tic_tac_toe.py is the offline counterpart, useful for bootstrapping a policy's format-following and basic competence before enabling online RL.
• Data and evaluation: generate_synth_data.py and eval_tic_tac_toe_vllm.py.

GSM8K (examples/gsm8k/) — a more conventional QA-shaped workload, with train_gsm8k.py and eval_gsm8k_vllm.py.

FSDP2 Math (examples/fsdp2_training/) — a multi-GPU template demonstrating FSDP2 wrapping, NCCL weight pushes to vLLM, and GRPO-style credit assignment via train_math_fsdp2.py.

Pipeline RL (examples/pipeline_rl/) — an actor/learner split over Redis for asynchronous sampling, via run_actor.py and run_trainer.py. Still experimental.

Rejection Sampling (examples/rejection_sampling.py) — generates rollouts, filters them, and writes training-ready JSONL for offline training.

Algorithm Presets

TraceRL ships several presets that pair a credit assignment strategy with a loss function:

Preset	Description
GRPO (make_grpo)	Group Relative Policy Optimization with token-level clipped surrogate loss. Default clipping range: (0.8, 1.2).
SAPO (make_sapo)	Uses a soft sigmoid gate in place of hard clipping, smoothly attenuating off-policy updates while preserving learning signal. Uses asymmetric temperatures (τ_neg > τ_pos) for stability.
GMPO (make_gmpo)	Uses the geometric mean of token-level importance ratios rather than the arithmetic mean, reducing sensitivity to outlier tokens. Wider default clipping range: (e^-0.4, e^0.4) ≈ (0.67, 1.49).
Dr. GRPO (make_dr_grpo)	An unbiased GRPO variant that omits standard-deviation normalization.
GSPO (make_gspo)	Sequence-level importance ratios computed via geometric mean.
CISPO (make_cispo)	Clipped importance-sampling weight optimization that preserves gradients from reflective reasoning tokens.
REINFORCE (make_reinforce)	Classic policy gradient with an importance sampling correction.
SFT (make_sft)	Supervised fine-tuning (behavioral cloning) for cold-start training.

Example usage:

from tracerl.training import make_gmpo, GRPORequestStrategy

# Create a GMPO algorithm instance with group size 4
algo = make_gmpo(group_size=4)

# Pair it with GRPO-style request expansion
request_strategy = GRPORequestStrategy(group_size=4)

Token-In Inference (Drift-Free)

TraceRL applies chat templates locally and sends pre-tokenized prompts directly to the vLLM /v1/completions endpoint. This keeps training aligned to the exact tokens sampled by the model, eliminating drift between training and inference.

Key details:

• Agents require a ChatTemplate instance (see tracerl.inference.HFChatTemplate).
• Use a single shared tokenizer per process to avoid redundant initialization.
• Tool calling relies on a text parser (for example, HermesToolParser) to extract tool calls from raw completions.
• Tools have two scopes, defined on ToolAgent: tools (executed internally by the agent, e.g. a calculator or code interpreter) and external_tools (returned to the protocol for external handling, e.g. delegation to sub-agents).
• To set explicit stop token IDs, use VLLMExtensions.extra_body_overrides, e.g. {"stop_token_ids": [...]}.

Migrating from the Chat Completions API

Earlier versions of this library used vLLM's /v1/chat/completions endpoint. The token-in API introduces the following breaking changes:

Previous API	Current API
ChatClient.complete(ChatCompletionRequest)	ChatClient.complete_tokens(TokenCompletionRequest)
ChatCompletionRequest(messages=...)	TokenCompletionRequest(prompt_token_ids=...)
Agent(client, model, ctx, parser)	Agent(client, model, ctx, parser, chat_template)
ToolRequest(tools, tool_choice="auto")	ToolRequest(tools) — tool_choice removed

The chat_template parameter is now required on all agents:

from transformers import AutoTokenizer
from tracerl.inference import HFChatTemplate, HermesToolParser

tokenizer = AutoTokenizer.from_pretrained("your-model")
chat_template = HFChatTemplate(tokenizer)

# For tool-calling agents:
chat_template = HFChatTemplate(tokenizer, tool_parser=HermesToolParser())

Requirements

• Python 3.12 or later
• PyTorch 2.8.0 or later, with CUDA, for training examples
• A vLLM server exposing the OpenAI-compatible API; NCCL is required for live weight pushes
• Redis, for the pipeline RL actor/trainer example

Installation

Using uv:

uv sync

To run the example scripts:

uv sync --extra examples

Roadmap

High Priority

• Add on-policy distillation, including a script for the MOPD variant introduced in MiMo-V2-Flash.
• Add a Gym-style registry for agent harnesses, environments, and interaction protocols, removing the need to construct them ad hoc and enabling general-purpose eval and training scripts.
• Implement proper FSDP2 wrapping throughout the training scripts.

Medium Priority

• Add Single Stream Policy Optimization, as described in this paper.
• Implement the importance-sampling findings from LLM Data Co.'s research, described here.

Hierarchical Agents and Delegation

• Implement a DelegatingProtocol for hierarchical agent architectures, where parent agents delegate subtasks to sub-agents via external tools and both parent and child rollouts are collected for training. The supporting infrastructure (external_tools and external_tool_handler) already exists; see CONSIDERATIONS.md. Inspired by Context-Folding, but implemented at the protocol level.

Environments and Agents

• Build an agent harness and environment for Pokémon. The harness and environment are currently fused in Claude Plays Pokemon; the goal is to disentangle them into a reusable agent harness paired with separate Pokémon-game environments.
• Add a Rustorio environment.

Low Priority

• Support --revision in scripts/push_to_hub.py for uploading specific checkpoints.
• Add a progress bar to the evaluation loop.

Package Restructuring (breaking — defer until everything else is complete)

• Split eval, training, and batch_gen into distinct modules.