Schedule-and-Calibrate: Utility-Guided Multi-Task Reinforcement Learning for Code LLMs

TL;DR

This paper introduces ASTOR, a utility-guided multi-task reinforcement learning framework for code LLMs that dynamically allocates training resources based on task utility, leading to significant performance improvements.

Contribution

ASTOR's novel utility-driven scheduling and calibration modules enable more effective multi-task training for code LLMs, outperforming existing methods.

Findings

ASTOR improves performance on four coding tasks by 9.0%-9.5%.

ASTOR surpasses the strongest MTRL baseline by 7.5%-12.8%.

ASTOR outperforms task-specific specialists in experiments.

Abstract

Reinforcement learning (RL) with verifiable rewards has proven effective at post-training LLMs for coding, yet deploying separate task-specific specialists incurs costs that scale with the number of tasks, motivating a unified multi-task RL (MTRL) approach. However, existing MTRL methods treat all coding tasks uniformly, relying on fixed data curricula under a shared optimization strategy, ultimately limiting the effectiveness of multi-task training. To address these limitations, we propose ASTOR, a multi-tASk code reinforcement learning framework via uTility-driven coORdination. Centered on task utility, a signal capturing each task learning potential and cross-task synergy, ASTOR comprises two coupled modules: 1) Hierarchical Utility-Routed Data Scheduling module hierarchically allocates training budget and prioritizes informative prompts, steering training toward the most valuable data and 2) Adaptive Utility-Calibrated Policy Optimization module dynamically scales per-task KL regularization, matching update constraints to each tasks current training state. Experiments on two widely-used LLMs across four representative coding tasks demonstrate that ASTOR consistently improves a single model across all tasks, outperforming the best task-specific specialist by 9.0%-9.5% and surpassing the strongest MTRL baseline by 7.5%-12.8%.