Muon Optimizer Boosts Agentic Reinforcement Learning Performance

What Changed

The landscape of reinforcement learning (RL) post-training is undergoing a shift as researchers begin to evaluate non-standard optimizers in agentic contexts. Historically, AdamW has been the default optimizer for both pre-training and fine-tuning large language models. However, the emergence of the Muon optimizer—which has demonstrated competitive performance in large-scale pre-training—has prompted investigations into its efficacy for RL. A recent study, 'When Does Muon Help Agentic Reinforcement Learning?', provides the first empirical evidence that Muon can significantly outperform AdamW in sparse-reward agentic environments, specifically when applied to the hidden weight matrices of policy models.

This shift is critical because sparse-reward environments, such as those found in ALFWorld, often present significant optimization challenges for standard gradient-based methods. The research suggests that the choice of optimizer is not merely a hyperparameter to be tuned, but a structural decision that interacts deeply with the policy optimizer and the advantage estimation strategy. By moving away from the ubiquity of AdamW in these specific sub-tasks, developers may unlock higher success rates and faster convergence in agentic workflows.

Technical Details

The study evaluates vanilla Muon against AdamW using the Qwen2.5-0.5B-Instruct model, focusing on the ALFWorld benchmark. The researchers employed three distinct policy optimization frameworks: Group-in-Group Policy Optimization (GiGPO), Group Relative Policy Optimization (GRPO), and GraphGPO. The core innovation lies in the selective application of Muon: rather than applying it globally, the researchers found that applying Muon specifically to hidden weight matrices yields the most substantial improvements.

In the context of GiGPO, the optimizer's impact is stark. When Muon is utilized, the final-window validation success rates show a marked improvement compared to high-rate AdamW controls, which often fail to maintain post-update success. The research highlights that the efficacy of Muon is highly dependent on the interplay between the advantage estimator and the learning rate. For instance, at a learning rate of 3e-5, Muon provides a clear advantage over GRPO baselines. At a more conservative 1e-5 learning rate, the synergy between GraphGPO and Muon becomes even more pronounced, leading to higher normalized validation AUC and significantly faster convergence metrics.

The findings suggest that Muon’s mechanism—which likely differs from AdamW’s adaptive moment estimation in how it handles curvature and weight updates—is particularly well-suited to the non-stationary nature of RL policy updates. By effectively navigating the loss landscape of sparse-reward tasks, Muon allows the agent to maintain stable performance where AdamW might otherwise diverge or stagnate.

Benchmark Analysis

The empirical results demonstrate clear performance gains across different optimization frameworks. In the GiGPO framework, applying Muon to hidden weight matrices increased the final-window validation success from 0.290 to 0.546. In the GRPO framework at a 3e-5 learning rate, Muon improved success from 0.161 to 0.268. Finally, using GraphGPO at a 1e-5 learning rate, the model reached a success rate of 0.901, while the normalized validation AUC improved from 0.399 to 0.556. Furthermore, the Muon-optimized agent reached 0.5 and 0.75 success thresholds 30 and 60 updates earlier than the baseline, respectively.

Developer Implications

For AI/ML engineers working on agentic RL, these results suggest that the 'default' optimizer choice should be re-evaluated. If you are currently struggling with convergence in sparse-reward environments, switching to Muon for hidden weight matrices may provide an immediate performance boost. However, this is not a 'plug-and-play' solution. The study emphasizes that the success of Muon is contingent upon the learning rate and the specific policy optimizer being used.

Engineers should adopt a more holistic approach to hyperparameter tuning, treating the optimizer, the advantage estimator, and the learning rate as a coupled system. When implementing Muon, it is advisable to start with the hidden weight matrices and conduct sensitivity analyses across different learning rates. Given that the research is based on single-seed comparisons, developers should also exercise caution and perform their own multi-seed validations to ensure that these gains are robust across different task variations and random initializations. The potential for faster training and higher success rates makes this a high-value area for further experimentation in production agentic pipelines.

Bottom Line

The study of Muon in agentic RL marks an important step toward optimizing the post-training phase of agentic models. By demonstrating that Muon can significantly outperform AdamW in sparse-reward tasks, the researchers have opened a new avenue for improving agent reliability and training efficiency. While multi-seed and cross-task validation remain necessary to confirm the generalizability of these results, the current evidence strongly suggests that Muon is a powerful tool for developers looking to push the boundaries of agentic performance. The future of RL optimization will likely require a more nuanced understanding of how these optimizers interact with the specific structural requirements of agentic policies.