ReCal: Reward Calibration for RL-based LLM Routing

2026-06-10

AI Summary

This work addresses key challenges in reinforcement learning–based large language model routing for heterogeneous tasks, including incomparable reward signals, ambiguous credit assignment, and optimization bias. To mitigate these issues, the authors propose a hierarchical reward decomposition mechanism coupled with a variance-aware reweighting strategy, which enables component-level advantage estimation and cross-dataset reward normalization to effectively calibrate reward signals. The approach significantly enhances the training stability and decision quality of routing policies. Empirical evaluation across seven benchmark datasets demonstrates consistent and substantial improvements over existing baselines, validating both the effectiveness and generalization capability of the proposed method.

Abstract

Large language model (LLM) routing has emerged as an effective paradigm for leveraging the complementary strengths of multiple LLMs through dynamic model and reasoning-strategy selection. Recent reinforcement learning (RL)-based routing methods further improve routing quality by optimizing routing policies from interaction feedback. However, they still struggle to provide informative and comparable learning signals under heterogeneous tasks with varying difficulty. In practice, multiple objectives (e.g., correctness, format behavior) are aggregated into a single scalar reward, leading to ambiguous credit assignment and conflicting optimization signals. Moreover, reward signals exhibit significant variability across instances, where some instances produce higher or more variable rewards, introducing optimization bias that favors trivial samples over informative ones. To address these issues, we propose \textbf{ReCal}, a \textbf{\underline{Re}}ward \textbf{\underline{Cal}}ibration framework for RL-based LLM routing. We first introduce a hierarchical reward decomposition mechanism with component-wise advantage estimation. We further propose a distribution-aware optimization strategy that calibrates optimization variability through variance-aware reweighting and per-dataset normalization. Experiments on seven datasets demonstrate that ReCal consistently improves routing performance, and training stability over baselines. Code is available at https://anonymous.4open.science/r/ReCal.

Problem

• LLM routing
• reward calibration
• reinforcement learning
• credit assignment
• optimization bias

Innovation

• Reward Calibration
• LLM Routing
• Reinforcement Learning
• Hierarchical Reward Decomposition
• Distribution-aware Optimization