AxisRL: A Post-Training Framework for the Agentic Era

Why agentic RL changes the post-training system

LLM post-training workloads are moving beyond single-turn question answering and into complex agentic workflows. A model may now run inside a long-lived environment, observe state, call tools, read tool results, update context, and eventually receive a reward.

That changes what a post-training framework has to handle. The problem is no longer just “generate samples, then train on them.” The framework also has to track multi-turn interaction, tool calls, environment state updates, and reward collection.

In traditional single-turn post-training, rollout is close to batch generation. In agentic RL, rollout becomes a multi-step interaction. Assistant responses, tool calls, tool results, verifiers, rewards, and environment state updates all become part of the same trajectory. Each step can change the context seen by later steps, making sequence length, latency, and training cost much less predictable.

Environments and external harnesses are also more complex. Math and Search tasks can often be controlled directly by the training framework. Coding agents such as OpenHands and Codex, browser tasks, desktop workspaces, multi-process tool systems, and third-party platforms often have their own control flow. Re-implementing the full harness inside the training framework is not always practical, and it can introduce behavior drift. For these systems, the training framework is better off capturing the interaction between the model and the harness, together with the final reward.

At the same time, training behavior becomes harder to verify. Tokenization, chat templates, rollout and trainer execution paths, weight sync, MoE routing, context packing, and algorithm implementations can all introduce subtle differences. Many issues do not crash immediately. They show up later as training-curve drift, gradient spikes, or unstable rewards.

What is AxisRL?

AxisRL is a training framework for agentic LLM post-training. It connects SGLang rollout, Megatron training, and real-world agent workflows into a compact online RL loop.

Its core positioning is simple: a small core with a complete online RL loop.

• SGLang handles high-throughput rollout and model serving.
• Megatron handles large-scale model training, parallelism, checkpointing, and model execution.
• AxisRL handles agent workflow orchestration, weight sync, resource scheduling, data movement, training observability, and task-specific recipes.

This keeps AxisRL’s scope narrow. It does not replace the serving engine or the training engine. Instead, it focuses on the contracts between rollout, training, data movement, and debugging.

Source code and examples are available in the AxisRL GitHub repository.

Overall workflow

From a user’s perspective, the main AxisRL workflow has five steps:

1. Rollout actors run agent workflows according to a recipe.
2. SGLang workers serve model generation.
3. Environments, tools, verifiers, or external harnesses produce interaction records and rewards.
4. Megatron workers consume training samples and run PPO or GRPO-family training.
5. Updated weights are synchronized back to the rollout side for the next online RL iteration.

There are three important boundaries. First, rollout and training are decoupled through queues. The rollout side produces trainable samples, while the training side consumes batches according to its own parallelism strategy.

Second, the model-serving boundary is token-in, token-out. Raw text, chat messages, and tool schemas are tokenized before they enter model workers. This makes sample packing, routing replay, and context merge easier to keep consistent across rollout and trainer paths.

Third, the driver mostly handles scheduling and metadata. It does not act as the transport layer for large tensors. Heavy rollout-side payloads—such as MoE routing information, complex rollout artifacts, or future multimodal intermediate data—move through a handle-based data path and are read by trainers on demand.

Three design goals

AxisRL is organized around Flexibility, Efficiency, and Observability. These goals are not separate product surfaces; they are constraints on the same online training loop.

AxisRL treats observability as part of the system design. Tokenization, chat templates, weight sync, routing replay, and packing need tests, comparison tools, and reproducible debug inputs. This lowers regression risk while the system and its recipes evolve quickly.

Agent workflow integration: white-box and black-box

Agentic RL does not have a single integration pattern. AxisRL supports direct control for simple environments and interaction capture for complex external harnesses.

In white-box RL, AxisRL controls the agent loop. The user implements a gym-like environment: the model generates an action from the current observation, the environment executes the action, and then returns a new observation and reward. This works well when the environment is clear, the tool surface is limited, and the control flow can be expressed inside the training framework.

In black-box RL, AxisRL does not require the user to re-implement the full control flow of an external harness. The harness calls the model through an OpenAI-compatible API. AxisRL’s proxy captures model inputs, outputs, and necessary metadata, then combines them with the final reward to construct training samples.

This distinction matters for integration cost and behavior consistency. Many real agent systems derive much of their value from complex harnesses. Requiring those harnesses to be rewritten makes integration expensive and risks behavior drift. With both paths, simple environments can be integrated directly while complex systems enter the RL loop through black-box capture.

Recipes, run modes, and resource scheduling

AxisRL’s user entry point is a set of recipes, not a single global harness. Each recipe can define the dataset, rollout loop, environment, verifier, reward computation, metrics, and training configuration for a task.

For white-box RL, a recipe usually implements environment interaction logic. For black-box RL, it focuses more on launching the external harness, capturing model interactions, and collecting rewards. Different tasks can keep different control flows as long as they eventually produce a common training-sample format that reuses the same training path.

rollout-only · train-only · rollout + train · eval-only · mismatch analysis

These modes make it possible to inspect environments and rewards before training, train from existing samples, run the full online loop, evaluate behavior independently, or fix data and configuration for end-to-end rollout/trainer comparison.

For resource scheduling, SGLang workers, Megatron workers, verifiers, reward models, teacher models, and black-box adapters can be managed through Ray actors and resource groups. Rollout, training, verification, and external services do not need to share the same resource type; the system can place components according to the workload.

Efficiency at Scale

One core efficiency problem in multi-turn agent training is the GPU bubble. Agent trajectories make the long tail much longer.

One trajectory may finish quickly. Another may require dozens or hundreds of tool calls. Some verifiers return immediately, while external harnesses can have long-tail latency. Under a strong synchronization boundary, both training and rollout wait for the slowest samples.

In this release, AxisRL first supports partial rollout to reduce this waiting. Completed or partially completed samples can move to the trainer before every trajectory has finished, allowing the training side to consume available data earlier and wait less on long-tail workflows.

AxisRL can also work with TIS, sequence masking, and Icepop to perform importance correction and filtering at token and sequence levels, improving training stability under off-policy settings.

Thin control plane, handle-based data plane

Many RL systems encounter a hidden bottleneck as they scale: the driver process becomes a large-data relay. Rollout-side data flows through the driver and is then distributed to multiple trainer ranks. This is usually acceptable for metadata, but large payloads such as MoE routing information, multimodal artifacts, or complex context intermediates can turn it into a network, serialization, and CPU bottleneck.

AxisRL keeps the control plane thin. The driver handles scheduling, lifecycle management, metrics, phase transitions, and sample metadata. Large payloads go through a handle-based data plane. The rollout side places heavy data in local storage, an object store, or a tensor store; the training sample carries only a handle and necessary metadata. A Megatron worker reads the payload on demand for its own batch, with communication overlapped with training compute where possible.

R3 as a representative payload

In MoE post-training, Rollout Routing Replay (R3) reduces expert-routing mismatch between rollout and trainer, making KL and loss more stable. Routing information itself can be large. If it follows every sample through a centralized driver, system cost rises quickly.

AxisRL puts routing payloads on the handle-based data path. This makes R3 part of rollout/trainer end-to-end consistency without turning the driver into a heavy-data bottleneck. Future multimodal inputs, intermediate features, video artifacts, and other large rollout-side data can reuse the same mechanism.

Context management and MagiAttention

The context of an agentic workflow is often not a simple linear sequence.

A search agent may execute many tool calls. To control context length, the system may keep only recent tool results while compressing, hiding, or replacing older results with placeholders. The context visible to the assistant at turn N is not necessarily a strict prefix of the context at turn N+1. The full trajectory looks more like a tree with shared prefixes.

If every turn is expanded into an independent training sample, the samples contain many repeated tokens and repeat the same attention computation. AxisRL uses prefix-tree merge to find shared token prefixes across turns, then uses MagiAttention to express tree-like attention. Shared parts can be computed with less duplication while each turn still sees only the context it was allowed to see during rollout.

This serves both flexibility and efficiency. Recipes can define task-specific retention, compression, and replacement policies; the trainer can consume these contexts through prefix-tree merge and complex attention masks without recomputing the same prefix for every turn.

Observability: correctness first, mismatch analysis, and spike replay

Agentic RL systems often fail without an immediate crash. Problems appear as abnormal training curves, KL spikes, unstable rewards, growing rollout/trainer mismatch, or gradient spikes. AxisRL turns correctness into an engineering practice built around tests, reproducible inputs, and analysis tools.

Tests are part of the architecture

AxisRL puts important training behavior under test, including tokenization, chat templates, weight sync, checkpointing, routing replay, RolloutTrace packing, prefix-tree merge, MagiAttention forward behavior, the OpenAI proxy, and mismatch analysis.

These tests are not just pre-release checks. They define critical boundaries between rollout, training, context management, and debug tooling. They also provide clearer safety boundaries when coding agents participate in implementation, refactoring, and experimentation.

Mismatch analysis

Mismatch means that rollout and training produce different outputs for the same input tokens because of implementation differences. Logprob divergence is a common example and can be an important source of training instability.

AxisRL compares token-level differences across backends, configurations, and routing-replay settings. The analysis helps distinguish global drift from a few outliers or from problems concentrated in a sequence type, token range, or context layout.

Spike replay

Spike replay targets occasional gradient or loss spikes. AxisRL can save weights, optimizer state, data, and relevant routing information before a suspicious update. Users can reload the same snapshot, reproduce the spike on the same data, and inspect which samples, tokens, parameters, or routing patterns contributed.

This changes debugging from “wait for the next random spike” into repeatable experiments on the same spike. At large post-training scale, that reproducibility is often more useful than observing an anomaly once.

Roadmap

AxisRL is still moving quickly. The next steps focus on expanding real workloads while keeping the core small and understandable.