Technical Summary: ASTOR – Utility-Guided Multi-Task Reinforcement Learning for Code LLMs

Problem Statement

Reinforcement learning (RL) with verifiable rewards has proven effective for post-training Large Language Models (LLMs) on coding tasks. However, the current paradigm relies on training separate, task-specific RL specialists (e.g., one model for code generation, another for unit test generation). This approach incurs memory and compute costs that scale linearly with the number of tasks. Conversely, relying on a single specialist leads to sharp performance degradation when the model is applied outside its specific training domain.

While Multi-Task Reinforcement Learning (MTRL) offers a unified solution by training a single model across all tasks, existing methods suffer from two critical limitations due to their uniform treatment of heterogeneous coding tasks:

1. Data Scheduling ("What to learn"): Current methods sample from a mixed task pool or follow a fixed curriculum, ignoring inter-task effects. They fail to account for how training on one task (e.g., code generation) might positively or negatively impact another (e.g., unit test generation).
2. Policy Optimization ("How to learn"): Existing approaches apply a uniform update strategy and a single KL regularization coefficient across all tasks. This is suboptimal because coding tasks have divergent optimization requirements; syntactically strict tasks (like unit test generation) require conservative updates to maintain executability, while semantically flexible tasks (like commit message generation) benefit from broader exploration.

Methodology: ASTOR

The authors propose ASTOR (multi-tASk code reinforcement learning framework via uTility-driven coORdination), a framework that resolves these challenges through two coupled modules centered on a unified signal: Task Utility. Task utility captures both a task's learning potential and its cross-task synergy.

1. Hierarchical Utility-Routed Data Scheduling

This module addresses "what to learn" by dynamically allocating the training budget based on estimated task utilities.

• Dual-Level Utility Estimation:
  • Task-Level: Utility is defined by two components:
    • Learning Potential: Measured by the variance of rollout rewards. High variance indicates a heterogeneous set of outcomes (successes and failures), suggesting actionable optimization signals.
    • Cross-Task Synergy: Measured by the gradient cosine similarity between the current task and other tasks. Positive similarity indicates constructive alignment, while negative similarity suggests conflict.
  • Prompt-Level: Within a task, utility is determined by the current learning potential (reward variance) and historical progress (improvement in mean reward).
• Hierarchical Allocation:
  • Task Quota Allocation: The global training batch size is partitioned among tasks using a temperature-scaled softmax over the aggregated task utilities. This ensures tasks with high potential and synergy receive more training data.
  • Prompt Prioritization: Within each task's allocated quota, prompts are selected via weighted sampling based on their prompt-level utility, prioritizing samples where the model is actively improving.

2. Adaptive Utility-Calibrated Policy Optimization

This module addresses "how to learn" by dynamically calibrating the KL regularization coefficient ($\beta_k$) for each task based on its current utility.

• Mechanism: The KL coefficient for task $k$ is decomposed into a task-specific base ($\beta_{base,k}$) and a dynamic multiplier conditioned on task utility:
  $$\beta_k(t) = \beta_{base,k} \cdot (1 + \lambda \cdot U(t)(T_k))$$
• Adaptation Logic:
  • Tasks with high utility (strong learning signals) receive a larger $\beta_k$, tightening the KL constraint to prevent overshooting when the gradient signal is concentrated.
  • Tasks with low utility receive a smaller coefficient, allowing more aggressive updates where the policy has room to improve.
• Reward System: The framework supports four heterogeneous tasks spanning the software development workflow: Code I/O Prediction (execution correctness), Code Generation (pass rate), Unit Test Generation (coverage), and Commit Message Generation (textual similarity). Each utilizes a verifiable, execution-based reward signal.

Key Contributions

1. First Study of MTRL for Code LLMs: The paper formulates multi-task RL for code around two coupled decisions ("what to learn" and "how to learn") and introduces task utility as a unified signal governing both.
2. ASTOR Framework: A coordination framework featuring:
   • Hierarchical Utility-Routed Data Scheduling: Allocates budget across tasks and prioritizes prompts based on learning potential and synergy.
   • Adaptive Utility-Calibrated Policy Optimization: Dynamically adjusts per-task KL regularization to match the training state of heterogeneous tasks.
3. Empirical Validation: Extensive experiments demonstrating that a single ASTOR model outperforms both task-specific specialists and existing MTRL baselines.

Experimental Results

The authors evaluated ASTOR on two base models (Qwen2.5-Coder-7B and Qwen3-8B) across four representative coding tasks.

• vs. Single-Task Learning (STL): ASTOR achieves strong multi-task performance, rivaling or exceeding task-specific specialists. On Qwen2.5-Coder-7B, it surpasses the best single-task model by 9.0% in overall average performance. On Qwen3-8B, the improvement is 9.5%.
• vs. MTRL Baselines: ASTOR significantly outperforms three baselines: Joint Learning (uniform mixing), Curriculum Learning (fixed sequential order), and Model Merging (post-hoc integration of specialists).
  • Improvements over the strongest MTRL baseline (Joint Learning) range from 7.5% to 19.2% across models and metrics.
  • Specifically, on Qwen2.5-Coder-7B, ASTOR outperforms Joint Learning by 12.8%.
• Ablation Studies: Removing either the data scheduling module (task allocation or prompt prioritization) or the policy optimization module (dynamic or task-specific regularization) leads to significant performance degradation, confirming that both components are essential for coordinating heterogeneous tasks.

Significance and Claims

The paper claims that ASTOR validates the effectiveness of utility-driven coordination for multi-task code RL. By treating tasks heterogeneously—allocating resources based on learning potential and synergy, and constraining updates based on task-specific needs—a single model can achieve diverse capabilities without the linear scaling costs of deploying multiple specialists.

The authors position this work as a step toward efficient, unified AI-assisted coding systems that can handle the full software development workflow (understanding, implementation, testing, and documentation) within a single model, reducing deployment barriers while maintaining or improving performance across all stages. The work acknowledges limitations regarding task/language coverage and the manual tuning of base KL coefficients, suggesting these as areas for future work.