verl-muon: When Does Muon Help Agentic Reinforcement Learning?

This repository accompanies "When Does Muon Help Agentic Reinforcement Learning?". Its primary artifact is a working Muon implementation for verl-style agentic reinforcement learning. The experiment launchers and analysis tools are included as supporting material, but the main entry point is the optimizer itself.

Start here: MuonWithAdamW, the Muon implementation used in our paper

What Is Muon?

Muon is an optimizer designed for matrix-valued neural-network parameters. It first forms a momentum update, then uses Newton-Schulz iterations to approximately orthogonalize that update before applying a shape-aware scale. This changes the geometry and effective magnitude of updates to hidden weight matrices, which can support more aggressive learning rates than a conventional AdamW recipe.

Our implementation is a hybrid optimizer:

• hidden matrix parameters use Muon;
• embeddings, normalization parameters, the language-model head, and parameters with fewer than two dimensions fall back to AdamW;
• the implementation is local and does not require an external Muon package.

This routing makes Muon practical for transformer policy training without forcing every parameter type through the same matrix update.

What Did We Find?

We studied ALFWorld agent training with Qwen2.5 models from 0.5B to 3B under multiple group-based policy objectives. The main finding is not that Muon always beats AdamW, but that Muon opens a useful higher-update operating regime:

• fan-in-scaled Muon remained stable at aggressive effective update magnitudes;
• Muon at a learning rate of 3e-5 improved late validation success over the 1e-6 AdamW baseline across rate-metric tests, although normalized-AUC gains were less uniform;
• the high-rate Muon recipe applied about 3.53x the hidden-matrix update RMS of AdamW, and matching that RMS removed the late-success advantage;
• the advantage narrowed when performance approached saturation or when AdamW was better tuned—for example, tuned AdamW nearly matched high-rate Muon on the 3B GraphGPO setting.

The practical conclusion is therefore recipe-level: Muon is most useful when the task still has optimization headroom, and its learning rate and update magnitude should be treated as part of the RL recipe rather than as a universal optimizer ranking.

Muon Across Scales and Objectives

Figure 4. Configuration-matched comparisons across model scales and objectives. Muon reaches strong validation success earlier in several settings, while the margin narrows as performance approaches saturation.

Muon Implementation Map

• _muon_update: momentum, Newton-Schulz orthogonalization, and fan-in scaling
• _use_muon_update: Muon-versus-AdamW parameter routing
• MuonWithAdamW: the complete hybrid optimizer
• get_optimizer_from_config: optimizer construction from verl config
• FSDP handling that preserves full 2D parameters
• Actor optimizer wiring with named parameters
• Optimizer configuration fields
• Muon and AdamW recipes evaluated in the paper

Integrate Muon into an Existing verl Repository

The integration has two code-level requirements, one configuration change, and one verification step. Port the linked sections selectively so that you preserve any upstream verl changes in your own checkout.

1. Add the hybrid optimizer

Port the Muon update and parameter routing, MuonWithAdamW, and the type: muon factory branch. The implementation between these sections also contains optional update diagnostics; either port their helper classes or remove the diagnostic hooks in your minimal version. Pass named parameters to the factory: parameter names are what keep embeddings, normalization layers, and the language-model head on AdamW.

2. Preserve matrix shapes under FSDP

Muon must see full 2D weight matrices. Our FSDP1 integration enables original parameters and switches the actor to ShardingStrategy.NO_SHARD when local Muon is selected; see fsdp_workers.py and the actor construction.

This is an important memory tradeoff: NO_SHARD replicates actor parameters instead of fully sharding them. Do not copy this behavior blindly into a setup that depends on FSDP memory savings. A fully sharded or FSDP2 integration needs a different mechanism that still reconstructs or exposes each complete matrix to the Muon update.

3. Add the optimizer fields

Add the Muon fields to the actor optimizer config and keep the critic on AdamW or Adam unless you intentionally want to study a different recipe:

actor_rollout_ref:
  actor:
    optim:
      type: muon
      lr: 1e-6
      muon_lr: 1e-5
      momentum: 0.95
      ns_steps: 5
      nesterov: true
      muon_adjust_lr_fn: original
      weight_decay: 0.01

critic:
  optim:
    type: adam

Here, lr is the AdamW fallback learning rate and muon_lr is the learning rate for hidden matrix parameters. We evaluated Muon learning rates of 1e-5 and 3e-5; these are paper settings, not universal defaults, so tune them for your model, task, batch construction, KL control, and training horizon.

4. Verify the port

At startup, the optimizer should report nonzero counts for both Muon and AdamW parameters. A 0 matrix parameters error normally means that FSDP is exposing flattened or sharded 1D views rather than full matrices. During training, verify the resolved Muon learning rate and, when comparing optimizers, monitor applied update magnitude rather than relying on the configured learning rate alone.

Paper and Experiment Materials

The rest of the repository supports auditing and reproducing the paper:

• REPRODUCE.md maps figures and tables to launch scripts;
• INSTALL.md describes the reference environment;
• METRICS.md defines the reported curve aggregations;
• scripts/experiments/ contains paper-aligned launchers;
• analysis/aggregate_metrics.py recomputes late-window and normalized-AUC summaries;
• THIRD_PARTY_NOTICES.md records third-party source and environment dependencies.

Full experiment runs require GPUs, model weights, and the corresponding ALFWorld or WebShop assets. These materials are secondary to the reusable Muon optimizer implementation linked above.

Citation

If you find this work useful, please cite our paper:

@misc{ruan2026doesmuonhelpagentic,
  title         = {When Does Muon Help Agentic Reinforcement Learning?},
  author        = {Kai Ruan and Jinghao Lin and Zihe Huang and Ziqi Zhou and Qianshan Wei and Xuan Wang and Hao Sun},
  year          = {2026},
  eprint        = {2607.16169},
  archivePrefix = {arXiv},
  primaryClass  = {cs.LG},
  url           = {https://arxiv.org/abs/2607.16169}
}

Acknowledgements

This repository builds on verl and verl-agent. It also uses the official implementations of GiGPO (code) and GraphGPO (code). We sincerely thank their authors and contributors for making these codebases and methods publicly available.