PaperlayerView exact revision on arXiv

Source-linked AI summary

When Does Muon Help Agentic Reinforcement Learning?

Kai Ruan, Jinghao Lin, Zihe Huang, Ziqi Zhou, Qianshan Wei, Xuan Wang, Hao Sun

arXiv:2607.16169v4cs.LGcs.AI

TL;DR

Whether Muon helps reinforcement-learning post-training remains unclear. This paper tests optimizer regimes across agentic tasks, objectives, and model scales, finding that fan-in Muon can sustain a larger stable effective step and improve late success under the evaluated recipe.

  • Problem

    Muon’s operating regime for sparse-reward, long-horizon agentic reinforcement learning remains unclear despite strong pre-training results.

  • Method

    The study compares fan-in Muon and AdamW across ALFWorld objectives, Qwen2.5 scales, AdamW rate controls, and update-magnitude-matched experiments under shared KL and clipping.

  • Results

    Fan-in Muon sustains a larger stable effective step and improves paired late success, while lower-rate Muon is inconsistent and tuned AdamW narrows the difference.

  • Takeaways & Limitations

    Optimizer selection should use recipe-level rate screening: aggressive Muon is most useful with optimization headroom, whereas saturation and AdamW tuning reduce its advantage.

  • Takeaways & Limitations

    The conclusions are limited to the shared regularization and fixed training budget, excluding broader checkpoint families and longer-horizon training.

Abstract

from arXiv · show

Muon is competitive with AdamW in large-scale pre-training, but its operating regime in reinforcement-learning post-training remains unclear. We map this regime on ALFWorld, a sparse-reward agentic benchmark, using three group-based objectives and Qwen2.5 models from 0.5B to 3B. Under a shared KL and clipping recipe, matched optimizer comparisons and AdamW rate controls trace the usable step-size range. AdamW responds non-monotonically to rate, whereas fan-in Muon remains stable at a more aggressive effective step: at $3 \times 10^{-5}$ it improves late success over an AdamW $10^{-6}$ baseline after correction across rate-metric tests. Its normalized-AUC effect is directionally positive but less uniform; the heuristic-matched lower-rate effect is less consistent, and tuned AdamW nearly matches high-rate Muon at 3B GraphGPO. High-rate Muon applies $3.53 \times$ AdamW's hidden-matrix update RMS; a full-budget RMS-matched control removes the late-success gain. Together, these results identify a recipe-level operating regime in which fan-in Muon supports a more aggressive stable effective step under shared KL and clipping: the margin is largest when optimization headroom remains and contracts near saturation, after AdamW tuning, or under magnitude matching. The scale-matched control ties this spectral effect to Muon's scale convention rather than establishing a universal optimizer ranking. Code is available at https://github.com/x66ccff/verl-muon.

Introduction

This study asks when fan-in Muon helps long-horizon, sparse-reward agentic RL, using controlled ALFWorld comparisons across Qwen2.5 scales and group-based objectives. It identifies a recipe-level regime where aggressive Muon can improve late success, while AdamW tuning, saturation, and RMS matching narrow or remove the difference.

  • Study design: Muon is evaluated on ALFWorld with Qwen2.5-Instruct agents at 0.5B, 1.5B, and 3B using GRPO, GiGPO, and GraphGPO.The objectives provide heterogeneous credit-assignment conditions, with a separate ablation testing optimizer-by-credit interaction.
  • Operating regime: AdamW remains rate-sensitive: 3×10^-6 is competitive but variable under GiGPO, while 10^-6 has the strongest tested GRPO mean.The matched comparisons use AdamW as the baseline, and separate rate controls assess when it is competitive.
  • Operating regime: High-rate Muon improves pooled late success over matched AdamW, with a directionally positive normalized-AUC effect after correction.The introduction frames these results as operating-regime evidence rather than a universal optimizer ranking.
  • Scale and saturation: At 1.5B, both Muon rates achieve higher AUC than AdamW, but at 3B GraphGPO, AdamW tuning nearly closes the AUC gap.In saturated GiGPO settings, high-rate Muon reaches the strong-policy regime earlier, indicating that visible gains depend on remaining optimization headroom.
  • Update magnitude: 3.53× AdamW’s hidden-matrix RMS characterizes the high-rate Muon update, while a full-budget RMS-matched control removes its standard recipe advantage.These measurements connect the observed difference to Muon’s scale convention and motivate treating spectral effects as diagnostics.

Related Work

Prior work frames Muon as a matrix-aware optimizer with scalable pre-training variants, while highlighting learning-rate and optimizer-continuity issues in post-training. Related agent-training research instead focuses on group-based credit assignment and improving optimization signals across long-horizon trajectories.

  • Muon and matrix-aware optimizers: Muon orthogonalizes each 2D hidden-layer momentum matrix with Newton–Schulz iterations and commonly uses Adam-family optimization for embeddings, norms, and output heads.Liu et al. (2025) reported scalability to LLM pre-training at roughly half AdamW’s FLOP budget, while MuonClip targets attention-logit instability.
  • Muon and matrix-aware optimizers: Concurrent variants modify Muon through adaptive second moments, operator-geometry budgets, row-magnitude control, or mixtures of orthogonalized and momentum updates.Their evidence concerns pre-training or supervised generalization, so they are treated as complementary designs rather than evaluated baselines.
  • Muon and matrix-aware optimizers: Update-RMS matching, distributed implementations, adaptive scaling, and convergence analyses extend matrix-aware optimization beyond the original Muon formulation.These developments include SOAP and Muon extensions for multi-billion-parameter pre-training, NAMO and Nor-Muon scaling, and condition-number-independent convergence results for simplified Muon.
  • Muon in post-training and optimizer mismatch: Switching Adam-pretrained models to Muon is learning-rate sensitive and can worsen the learning–forgetting tradeoff, whereas optimizer continuity is more favorable.Muon-pretrained models have also been successfully instruction-tuned with Muon.
  • Group-based RL for LLM agents: Group-based RL methods estimate advantages from within-group rollout statistics, with multi-turn agents making long-horizon credit assignment a central difficulty.GiGPO adds step-level relative advantages by regrouping actions from repeated anchor states across trajectories, while HGPO addresses context inconsistency through history-conditioned groups but is not evaluated here.
  • Group-based RL for LLM agents: Recent agent-training methods improve optimizer inputs through token-level distillation, counterfactual or failed-trajectory credit, hierarchical trajectory signals, and state or inverse-dynamics objectives.These methods operate along complementary axes to the fixed group-based estimator sweep.

Background

The background defines three group-relative credit objectives for sparse-reward agent trajectories and explains Muon’s fan-in-scaled polar-factor update. It motivates pairing AdamW at 10^-6 with Muon at 10^-5 and 3 × 10^-5 while separating fixed-recipe comparisons from rate-sensitivity analysis.

  • Credit structure: GRPO normalizes trajectory returns, GiGPO contrasts actions from repeated anchor states, and GraphGPO derives transition credit from rollout-graph distance to the goal.All three use a group-relative policy loss while distributing sparse terminal reward differently.
  • Muon: Muon replaces momentum-gradient singular-value magnitudes with an approximate polar factor and applies a fan-in-scaled update to hidden weight matrices.Non-matrix parameters remain updated with AdamW.
  • Learning-rate comparability: 10× is the practical learning-rate conversion heuristic used to pair AdamW at 10^-6 with Muon at 10^-5, while 3 × 10^-5 probes a more aggressive fan-in recipe.The higher Muon rate is not presented as rate-equivalent to AdamW 10^-6.
  • Learning-rate comparability: Fixed rates across estimators and objectives separate paired optimizer comparisons from post-hoc rate selection, while a separate AdamW sweep measures rate sensitivity.Matched diagnostics and a full-budget control assess applied-update scale and an RMS-matched Muon convention.

Experiments

Experiments on ALFWorld and WebShop compare AdamW with lower- and higher-rate Muon across Qwen2.5 scales and group-based objectives under matched training conditions. High-rate Muon generally improves late success and, with remaining optimization headroom, normalized AUC, while AdamW is rate-sensitive and tuning or saturation narrows the gap.

  • Setup: ALFWorld experiments compare AdamW at 10−6 with Muon at 10−5 or 3 × 10−5 across three estimators and Qwen2.5 models from 0.5B upward.Matched cells hold model, data, rollouts, schedules, initialization, splits, evaluation, and non-optimizer settings fixed; comparisons generally use five seeds.
  • Optimizer comparisons: High-rate Muon improves late success, but its normalized-AUC effect is less uniform; Holm-adjusted p-values are .004 and .022, respectively.Across matched 0.5B comparisons, lower-rate Muon is not consistently signed, while AdamW responds non-monotonically and rates of 5 × 10−6 and above lose nearly all post-update success.
  • Scale and headroom: 0.5B late success rises from 0.290 to 0.546 with high-rate Muon in the GiGPO scale sweep.At 1.5B and 3B, late success approaches saturation, and the corresponding normalized-AUC increases are 0.635 to 0.708 and 0.707 to 0.805.
  • Scale and headroom: At 3B GraphGPO, high-rate Muon reaches 0.75 success first, while tuned AdamW at 3 × 10−6 attains nearly the same AUC.This boundary case shows the optimizer margin contracting after AdamW tuning.
  • Transfer: On WebShop, high-rate Muon raises partial-task-score AUC from 0.078 to 0.548 and late score from 0.197 to 0.854.Lower-rate Muon reaches 0.256 AUC and 0.554 late score; exact completion is reported separately.

GiGPO: Qwen2.5-0.5B

For GiGPO, configuration-matched comparisons extend across model scales, while the optimizer-level analysis shows that Muon changes applied-update spectral shape and that RMS matching removes the task-level gain. The reported success and success-AUC ranges are 0.008 to 0.689 and 0.005 to 0.305, respectively.

  • GiGPO: Qwen2.5-0.5B: Configuration-matched comparisons cover GiGPO at 0.5B, 1.5B, and 3B, with saturated GiGPO comparisons emphasizing earlier progress.GRPO and GraphGPO retain late-success headroom at 1.5B.
  • GiGPO: Qwen2.5-0.5B: 0.008 to 0.689 success and 0.005 to 0.305 success AUC are the reported GiGPO ranges.Appendix C.6 provides the trajectories and complete summary.
  • Optimizer diagnostics: 0.015 for AdamW’s applied update and 0.585–0.635 for Muon show that Muon produces a substantially flatter applied-update spectrum in matched 1.5B GRPO.The raw statistic remains low for both optimizers, and the statistic is averaged equally across matrices over training.
  • Optimizer diagnostics: RMS-matched scaling does not preserve the task-level gain, identifying the scaling rule as part of the effective recipe without isolating spectral shape.The convention changes both global magnitude and layerwise scale allocation.

Discussion

Fan-in Muon is most useful when optimization headroom remains and its aggressive update is stable under shared KL and clipping. The evidence supports recipe-level rate screening rather than a universal optimizer conversion rule.

  • Operating regime: Fan-in Muon improves late success when optimization headroom remains, while near saturation its advantage appears as earlier progress.The lower, approximately rate-equivalent Muon setting is inconsistent, and tuned AdamW closes the AUC gap at 3B GraphGPO.
  • Rate sensitivity: Under shared KL and clipping, Muon sustains an effective step that AdamW rate increases do not in the tested sweeps.At 3×10−6, AdamW is competitive but variable; higher rates fail under the shared recipe.
  • Practical selection: 10−5 is the conservative Muon candidate, whereas 3 × 10−5 is the high-headroom candidate under the evaluated regularization.The high-rate recipe should be selected only when early validation improves without instability in response clipping, entropy, or valid-action diagnostics.
  • Scope of evidence: The scale results characterize the full optimizer recipe rather than model size in isolation because Muon’s hidden-matrix parameter coverage increases with model scale.Matched 0.5B ALFWorld comparisons provide the primary statistical evidence; scale and WebShop experiments are extension cases.
  • Mechanism and limits: The full-budget RMS-matched control falls below standard Muon and AdamW, showing that Muon’s scaling convention is part of the effective recipe.An AdamW-like global magnitude does not preserve the gain because the accompanying layerwise allocation change prevents isolating the spectral effect.

Conclusion

Fan-in Muon supports a larger stable effective step in the evaluated agent-training recipes, improving paired late success where larger AdamW rates are unreliable. The advantage is recipe-dependent: lower-rate Muon is inconsistent, AdamW tuning narrows the difference, and RMS matching removes the gain while spectral flattening remains observed.

  • Fan-in Muon remains stable at an aggressive effective step where larger AdamW rates are unreliable and improves paired late success.The finding holds under shared regularization settings.
  • The lower Muon setting is inconsistent, while AdamW tuning narrows the difference.
  • RMS-matched scaling removes Muon's gain, although spectral flattening remains directly observed.This ties the observed advantage to the complete optimizer recipe rather than the optimizer name alone.

A. Experimental Hyperparameters

The experiments use a shared training setup alongside estimator-specific settings. Additional implementation details are reported separately in Appendix B.

  • Experimental Hyperparameters: The reported comparisons use a shared training setup.Table 3 summarizes the common configuration across experiments.
  • Experimental Hyperparameters: The comparisons also specify settings separately for each estimator.These estimator-specific settings are included in the summarized experimental configuration.
  • Experimental Hyperparameters: Implementation details are reported separately in Appendix B.Table 3 identifies Appendix B as the location for those details.

B. Implementation and Update-Scale Details · C. Additional Experimental Results · C.1 Credit-Control Interaction

The implementation assigns hidden matrices to fan-in Muon and other parameters to AdamW, producing substantially larger applied updates than AdamW under the standard recipe. Additional controls show that update magnitude and credit-term interactions materially shape Muon’s observed gains, while normalized-AUC effects are more variable.

  • B. Implementation and Update-Scale Details: 72%, 85%, and 90% of optimized parameters are hidden matrices assigned to Muon in the 0.5B, 1.5B, and 3B models, respectively.Embeddings, normalization parameters, tied embedding/LM-head weights, and other non-matrix parameters use AdamW.
  • B. Implementation and Update-Scale Details: Muon uses Keller–Jordan fan-in scaling, multiplying updates by max(1, dout/din), unlike constant-update-RMS scalable Muon conventions.The implementation uses 0.95 Nesterov momentum and five bfloat16 Newton–Schulz iterations.
  • B. Implementation and Update-Scale Details: 3.53× AdamW’s hidden-matrix RMS and 3.20× its total RMS are applied by standard high-rate Muon in the matched five-update probe.RMS matching reduces these ratios to 0.80 and 0.83, while Newton–Schulz raises median optimizer-step latency about 1.8× and cumulative five-update time by at most 3.1%.
  • B. Implementation and Update-Scale Details: RMS matching removes the late-success advantage of standard fan-in Muon in the full-budget 0.5B GiGPO comparison.The control tests applied-update magnitude and implementation convention rather than providing a pure spectral-shape ablation.
  • C. Additional Experimental Results: Late success averages evaluations at updates 175–200, whereas normalized AUC is trapezoidal area over updates 0–200 divided by 200.Heavy trajectories are five-point trailing means computed within each curve; faint trajectories show raw checkpoints or traces.
  • C. Additional Experimental Results: Matched comparisons hold model, estimator, data, rollout budget, schedule, initialization, task split, evaluation set, and non-optimizer settings fixed across five seeds per cell.These controls isolate optimizer and credit effects in the additional experiments.
  • C.1 Credit-Control Interaction: Enabling the GiGPO step term raises late success for both AdamW and Muon, while Muon’s late-success effect is positive in all ω = 1 and two of three ω = 0 comparisons.Muon’s normalized-AUC effects are more variable, indicating that optimizer–credit interactions differ across comparisons.

C.2 Robustness Across Matched Comparisons … E. Learning-Rate Controls and Stress Diagnostics

Across matched comparisons and scale extensions, Muon’s advantage is robust in learning speed and selected late-window outcomes, while optimizer separation appears primarily in policy movement rather than gradient magnitude. The evidence supports a conditional, rate- and scale-dependent benefit rather than a universal optimizer ranking.

  • C.2 Robustness Across Matched Comparisons: p = .0018 remains significant after removing the largest high-rate late-success difference, with Holm-adjusted padj = .0073.The result preserves the late-success direction under a sensitivity analysis that removes the most influential matched group.
  • C.2 Robustness Across Matched Comparisons: High-rate differences remain significant under estimator-stratified permutation tests for both metrics after discarding pairing assignments.The exhaustive enumeration contains 16,003,008 assignments per test and evaluates effect magnitudes rather than paired directions.
  • C.3 Task-Level Results: Muon’s differences span multiple task categories rather than concentrating in one task family.Figure 7 summarizes late-window performance across all three estimators and both Muon learning rates.
  • C.4 Optimizer Dynamics Across Estimators: Policy movement separates estimators more clearly than gradient magnitude, with KL trajectories varying while gradient norms remain comparable.Stable GRPO training through step 200 shows Muon moving farther from the reference policy on the same gradient-norm scale.
  • C.5 Scale Extensions: 0.75 success arrives in 30 updates for high-rate Muon, versus 65 for default AdamW and 40 for tuned AdamW at 3B GraphGPO.High-rate Muon’s normalized AUC is 0.856, versus 0.853 for tuned AdamW, illustrating a smaller margin after tuning.
  • C.6 WebShop Transfer Case: 0.766 peak exact success and 0.688 final success for high-rate Muon exceed AdamW’s final 0.023 in the matched 0.5B WebShop comparison.High-rate Muon also crosses continuous task score 0.5 by update 75 and remains above it through training.
  • C.7 Applied-Update Spectrum at 1.5B: 0.814 and 0.813 late-window success, with normalized AUC 0.529 and 0.574, follow Muon rates 10^-5 and 3 × 10^-5 from an AdamW baseline of 0.598 and 0.398.At 1.5B GRPO, the spectral diagnostics compare optimizer transformations on the same probed hidden-weight matrices.
  • D. Spectral–Credit Compatibility Hypothesis: Muon’s flattened update is most useful when promoted weak directions carry reliable policy-improvement signal, summarized by the reliability quantity Rtail_a.This compatibility view explains why spectral flattening can help conditionally rather than implying universal optimizer superiority.

E.1 AdamW Learning-Rate Controls · E.2 High-Rate Stress Diagnostics

AdamW’s usable rate range is narrow under the shared regularization: an intermediate rate increases dispersion without improving average performance, while higher rates become unstable. At 3 × 10^-5, Muon retains valid behavior and improves success where AdamW fails, under the same KL and clipping recipe.

  • E.1 AdamW Learning-Rate Controls: 10^-6 remains the strongest AdamW rate under GRPO at 0.5B, while 10^-5 produces zero late success.GiGPO also includes a competitive 3 × 10^-6 control and a 3 × 10^-5 stress setting.
  • E.1 AdamW Learning-Rate Controls: 3 × 10^-6 AdamW spans late success from 0.026 to 0.628, with a mean below the 10^-6 reference.The strongest individual intermediate-rate control therefore does not translate into better average performance.
  • E.1 AdamW Learning-Rate Controls: Higher AdamW rates fall outside the stable range under the shared regularization, while the intermediate rate mainly increases dispersion rather than average performance.This control pattern indicates non-monotonic rate behavior rather than a uniformly beneficial increase.
  • E.2 High-Rate Stress Diagnostics: 3 × 10^-5 causes AdamW to lose valid behavior, whereas Muon at 3 × 10^-5 retains valid behavior and improves success.The comparison uses the shared KL and clipping recipe, with the KL coefficient fixed.
  • E.2 High-Rate Stress Diagnostics: 10^-5 produces late KL excursions for AdamW under the shared KL and clipping recipe.These excursions occur before the more severe 3 × 10^-5 loss of valid behavior.
  • E.2 High-Rate Stress Diagnostics: The high-rate diagnostics report validation success and valid-action ratio alongside KL trajectories for GiGPO.Figure 14 distinguishes raw traces from five-point trailing means and plots KL on a logarithmic axis.
Loading 2607.16169v4…