AxisRL: an agentic RL post-training framework built on SGLang and Megatron

The problem AxisRL targets: multi-turn agent rollouts that silently diverge from training

Single-turn post-training has a simple contract. A prompt goes in, a completion comes out, a reward is computed, and the trainer consumes the result. Agentic RL breaks that contract. The README describes a model that may interact with a long-lived environment, call tools, observe tool results, update context, and only receive a reward after several turns. Between the rollout and the trainer sit environment state, tool calls, verifiers, reward collection, training sample construction and weight synchronization.

The README names the failure mode that follows: small differences in tokenization, chat templates, logprobs, routing, packing or weight versions can appear later as loss spikes, reward instability, or rollout-trainer mismatch. That is the specific problem AxisRL claims to solve. It is not a training algorithm library. It is the system layer around agentic post-training, and the README is explicit that SGLang and Megatron remain the core serving and training engines.

The intended user is therefore narrow. If you are fine-tuning a 7B model on a single node with single-turn data, this framework is heavier than the problem. The README also states it has been used in agent RL workflows with 300+ turn trajectories and training runs at hundreds-of-billions-parameter scale, which sets the scale of the audience it is written for.

How the rollout-trainer loop works, and why the driver stays thin

The README lays out a five-step loop. Rollout actors execute task-specific agent workflows. SGLang workers serve model generation. Environments, tools, verifiers or external harnesses produce interaction records and rewards. Megatron workers consume training samples and run PPO or GRPO-family training. Updated weights are synchronized back to the rollout side for the next iteration.

The design decision worth noticing is where the driver sits. AxisRL keeps the driver lightweight: it manages scheduling, lifecycle, metrics, phase transitions and metadata. Heavy payloads, such as routing replay data or future multimodal artifacts, move through a handle-based data path and are read by trainer workers on demand. That is a deliberate separation. A control plane that only moves handles does not stall on large tensors, and the trainer pulls what it needs when it needs it.

Around that loop the framework adds partial rollout, prefix-tree merge, MagiAttention, and off-policy stabilization tools named TIS, sequence masking and Icepop. Task logic lives in recipes rather than in the core. The README's flexibility row states that recipes carry control flow, tools, rewards, context management and resource needs, while white-box environments and black-box harness capture are both supported and heterogeneous components are managed through resource groups. Configurable policy objectives include PPO, GRPO/GRPO2, GSPO, TOPR and TIS.

Observability is treated as a first-class goal rather than a debugging afterthought. The README points to tests of critical boundaries plus mismatch analysis, routing replay checks and spike replay. Whether those tools are enough is something only a real run will tell you; the README does not publish a worked example of a mismatch being diagnosed.

Installing AxisRL from the Docker image and running a first GRPO job

The README recommends the project Docker image and states it includes SGLang, Megatron Core, MagiAttention, Ray, CUDA dependencies and the Python packages used by the current recipes. Pull the prebuilt image first.

docker pull leejunjie/sglang-mcore:cu130-sgl0.5.14-mcore0.18-magi

The repository also carries the Dockerfile used to build it, at docker/cuda/cu130-sgl0.5.14-mcore0.18-magi.Dockerfile, so you can rebuild rather than depend on the published tag. Inside the container, install the package in editable mode.

pip install -e .

setup.py declares python_requires=">=3.12" and a single install requirement, e2b>=2.34.0, which means the heavy dependencies are expected to come from the image rather than from pip resolution. The README offers an optional data download for the models and datasets referenced by the current recipes and tests.

python axrl/example/download_data.py

That bulk download can be large because it includes multi-billion-parameter models. For a narrower run the README suggests adjusting the recipe model and dataset paths instead of downloading everything. The smallest documented entry point is the GSM8K GRPO recipe, which assumes a GPU machine with enough resources for the default parallelism in that recipe.

AXRL_OUTPUT_DIR_NAME=grpo_gsm8k \
bash axis_recipe/grpo_gsm8k/run_train.sh \
  --online_rl_train.max_global_updates=4

Most config fields can be overridden from the command line with --path.to.field=value, which is what the max_global_updates override demonstrates. The README does not state what a successful short run prints, so treat the first launch as an environment check.

Search-R1 and the retrieval server: a second process you have to keep alive

The Search-R1 recipe is the clearest example of what agentic post-training costs operationally. It uses a retrieval server in addition to rollout and training workers, so a run is no longer one script. The README exports the port, starts the retriever, builds the config, and then launches training.

export AXRL_SEARCH_PORT=18000
bash axis_recipe/search_r1/start_retriever.sh
python axis_recipe/search_r1/search_r1_config.py

AXRL_OUTPUT_DIR_NAME=search_r1 \
python -u axis_recipe/search_r1/train_search_r1.py \
  --config_path=axis_recipe/search_r1/search-r1-config.yaml \
  --online_rl_train.max_global_updates=4

There is also a combined script that sets the port inline for the default full recipe.

AXRL_SEARCH_PORT=18000 \
bash axis_recipe/search_r1/run_train.sh

Note what this implies. The retrieval server is a separate lifecycle from the trainer, and the README does not document how the retriever is health-checked, restarted or torn down when training exits. If a retrieval server dies mid-run, the failure surfaces as a rollout problem rather than as an infrastructure problem. That is a real operational gap in the documentation, not a hypothetical one.

Where AxisRL is the wrong tool, and what the black-box path does not yet cover

The README is honest about one thing that many project pages bury: the black-box RL recipe with OpenHands and E2B is described as still a work in progress. It demonstrates the black-box harness integration path, but the README text is truncated mid-sentence at that point, so the scope of what works is not stated. Anyone whose plan depends on capturing a closed harness through the OpenAI-compatible proxy should treat that as unverified until they read the recipe and tests in the repository.

The second limitation is the one implied by the whole architecture. AxisRL assumes SGLang for rollout and Megatron for training. If your stack is vLLM plus DeepSpeed, adopting AxisRL means adopting two engines, not one framework. The README does not describe a supported path for other serving or training backends.

The third is resource shape. Recipes assume a GPU machine with enough resources for their default parallelism. Partial rollout and prefix-tree merge exist precisely because long-tail trajectories and tool latency leave resources idle, but those mechanisms do not remove the floor. The README's own guidance to adjust recipe model and dataset paths instead of downloading everything is a hint that the defaults are sized for large machines.

Finally, there are no retrieved releases. The repository has no release history to pin to, so reproducibility depends on a commit plus the Docker image tag, and the image tag is the only version identifier the README gives.

Alternatives and the actual difference in approach

The closest comparison is verl. It also targets RL post-training for language models and also separates rollout from training, but its documented ecosystem centers on FSDP and Megatron training backends with vLLM and SGLang rollout, and its recipes are primarily single-turn or tool-calling RL over standard datasets. AxisRL's README frames the framework around real agent workflows, resource groups for heterogeneous components, and a handle-based data path for large payloads, with mismatch analysis and spike replay as named features. The difference is where the engineering effort went: verl spends it on backend coverage, AxisRL spends it on the contracts between a long-lived agent environment and the trainer.

OpenRLHF is the other common starting point. Its design is simpler and its dependency surface smaller, which makes it easier to install and reason about. It does not document routing replay or rollout-trainer mismatch tooling, so the class of bug AxisRL is built to catch is largely outside its scope. If your rollouts are single-turn and your loss curves are stable, OpenRLHF's simplicity is worth more than AxisRL's instrumentation.

The honest framing is that these are not interchangeable. AxisRL is a bet that agentic post-training failures are mostly system-boundary failures, and it invests accordingly.

Licence, maintenance and upgrade cost

AxisRL is Apache-2.0, which permits commercial use, modification and redistribution with the usual requirements around notices and patent terms. The repository also ships third_party content and depends on SGLang, Megatron Core, MagiAttention, Ray and CUDA components, each under its own licence. Apache-2.0 on AxisRL says nothing about those, so a licence review has to cover the image contents as well as this repository.

On maintenance, the last push to the default branch was on 2026-08-03. The repository is not archived. The README advertises PRs welcome, and the repository carries .pre-commit-config.yaml, an AGENTS.md file and a skills/ directory alongside tests/ and benchmark/, which suggests an active contribution process rather than a frozen snapshot.

Upgrade cost is the part to weigh before adopting. The recommended install is a pinned image tag, cu130-sgl0.5.14-mcore0.18-magi, which couples your AxisRL version to a specific SGLang and Megatron Core pair. Moving to a newer SGLang or Megatron means rebuilding that image and re-validating the rollout-trainer boundaries the framework exists to protect. Because there are no retrieved releases, there is no changelog to read between image tags. Budget for a rebuild and a boundary re-test on every engine upgrade, not just a pip install.

Editorial conclusion

AxisRL fits teams that already run SGLang and Megatron at multi-node scale and need agentic rollout, weight synchronization and mismatch debugging in one loop; it does not fit anyone without that GPU footprint, and the black-box OpenHands/E2B recipe is still a work in progress, so verify that path before planning around it.

Frequently asked questions

Official sources

1. Issues
2. License: Apache-2.0
3. Project website
4. README
5. XYZ-AI-Lab/axrl on GitHub