XYZ-AI-Lab/axrl

AxisRL is an agentic RL post-training framework built on SGLang rollout, Megatron training, and real-world agent workflows.

What it solves

AxisRL is designed for agentic reinforcement learning (RL) post-training, where models must interact with environments over multiple turns, call tools, and receive rewards based on long-term outcomes. It solves the coordination challenges of managing multi-turn rollouts, environment states, tool calls, and weight synchronization while ensuring that the training process remains stable and observable.

How it works

AxisRL acts as a system layer that connects high-throughput rollout engines (SGLang) with large-scale distributed training engines (Megatron). It manages the loop where rollout actors execute agent workflows, SGLang workers generate text, environments produce rewards, and Megatron workers perform policy optimization (such as PPO or GRPO).

Key technical mechanisms include:

  • White-box and Black-box Integration: Supports internal agent loops (white-box) or capturing interactions via an OpenAI-compatible proxy for external harnesses (black-box).
  • Partial Rollout: Allows completed samples to be sent to the trainer immediately to reduce idle time caused by long-tail trajectory latency.
  • Handle-Based Data Plane: Moves heavy payloads (like MoE routing data) via handles rather than a central driver to prevent bottlenecks.
  • Context Management: Uses prefix-tree merge and MagiAttention to optimize repeated attention compute in multi-turn contexts.
  • Consistency Tools: Includes mismatch analysis and "spike replay" to debug divergence between rollout and trainer execution paths.

Who it’s for

It is intended for researchers and engineers training large-scale agentic LLMs (hundreds-of-billions of parameters) that require complex, multi-turn interactions and high-throughput RL post-training.

Highlights

  • Built on SGLang for rollout and Megatron for distributed training.
  • Supports multiple optimization objectives including PPO, GRPO/GRPO2, GSPO, TOPR, and TIS.
  • Handles trajectories with 300+ turns.
  • Features Rollout Routing Replay (R3) to stabilize MoE expert routing during training.
  • Provides a reproducible debugging suite for loss spikes and token-level mismatches.

Related

  • Project
    THUDM/slime

    An LLM post-training framework for RL scaling that integrates Megatron and SGLang to provide high-performance training and flexible data generation.

  • Project
    radixark/miles

    An enterprise-grade reinforcement learning framework for large-scale model post-training that integrates SGLang for rollouts and Megatron-LM for scalable training.

  • Project
    areal-project/AReaL

    A large-scale asynchronous reinforcement learning system designed for training high-performance reasoning and agentic AI models.

  • Project
    langfengQ/verl-agent

    An extension of veRL for training LLM agents via reinforcement learning, featuring a step-independent rollout mechanism to efficiently handle long-horizon, multi-turn tasks.

  • Project
    Tencent-Hunyuan/UniRL

    UniRL is a reinforcement learning framework that provides a unified post-training loop for multimodal models, including image and video diffusion, autoregressive VLMs, and LLMs.