Bridging Online and Offline RL: Contextual Bandit Learning for Multi-Turn Code Generation

TL;DR

Cobalt reframes multi-turn code generation as a one-step recoverable problem and uses offline trajectories to train a contextual bandit that optimizes single-turn completions given trajectory context, reducing online training cost. It theoretically bounds the gap between the online multi-turn RL objective and the stepwise objective under KL regularization, showing a scalable $O(T\sqrt{\eta})$ regret. Empirically, Cobalt improves Pass@1 on LiveCodeBench and TACO-Dev compared to online baselines and maintains generalization to longer horizons, while perturbation-based data augmentation mitigates in-context reward hacking. These results indicate that contextual bandit learning with offline trajectories is a promising, efficient paradigm for self-improving LLMs in iterative decision tasks like multi-turn code generation.

Abstract

Recently, there have been significant research interests in training large language models (LLMs) with reinforcement learning (RL) on real-world tasks, such as multi-turn code generation. While online RL tends to perform better than offline RL, its higher training cost and instability hinders wide adoption. In this paper, we build on the observation that multi-turn code generation can be formulated as a one-step recoverable Markov decision process and propose contextual bandit learning with offline trajectories (Cobalt), a new method that combines the benefits of online and offline RL. Cobalt first collects code generation trajectories using a reference LLM and divides them into partial trajectories as contextual prompts. Then, during online bandit learning, the LLM is trained to complete each partial trajectory prompt through single-step code generation. Cobalt outperforms two multi-turn online RL baselines based on GRPO and VeRPO, and substantially improves R1-Distill 8B and Qwen3 8B by up to 9.0 and 6.2 absolute Pass@1 scores on LiveCodeBench. Also, we analyze LLMs' in-context reward hacking behaviors and augment Cobalt training with perturbed trajectories to mitigate this issue. Overall, our results demonstrate Cobalt as a promising solution for iterative decision-making tasks like multi-turn code generation. Our code and data are available at https://github.com/OSU-NLP-Group/cobalt.

Figures (11)

• Figure 1: A multi-turn code generation example with four public test cases (two correct and two incorrect) and 16 hidden test cases. The LLM starts with the coding problem and generates a program. It passes all benign public tests and hidden tests, while failing the two perturbed tests. As a result, one of the failed test case is returned as feedback. All other test cases and their pass rates are hidden. The LLM mistakenly follows the feedback and changes its program, which now passes the perturbed test case but only one hidden test case.
• Figure 2: Illustration of Cobalt. We first use a reference LLM $\pi_{\mathrm{ref}}$ to collect trajectories offline and divide them into partial trajectories. During online bandit learning, we sample single-step program completions $a_{t} \sim \pi(\cdot | s_t)$ using each partial trajectory as the contextual state, where $s_t = (o_0, a_0, ..., o_t)$. Then, we optimize the LLM under $J_{\mathrm{step}}$ (§ \ref{['sec:online_rl']}) to maximize $R(s_t, a_t)$.
• Figure 3: TACO-Dev.
• Figure 4: LiveCodeBench.
• Figure 5: Self-improvement results (numbers in Appendix \ref{['sec:seq_revision_tables']}). Cobalt improves both models' performance by a large margin and generalizes to unseen horizon ($t \ge 4$) at test time.

Theorems & Definitions (2)

• Definition 2.1
• Theorem 3.1