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

TL;DR

This paper introduces Cobalt, a novel offline-online reinforcement learning method that improves multi-turn code generation by combining trajectory collection with contextual bandit training, outperforming existing baselines.

Contribution

Cobalt is a new method that leverages offline trajectories and online bandit learning to enhance multi-turn code generation in large language models.

Findings

Cobalt outperforms two multi-turn online RL baselines.

Significant improvements in Pass@1 scores on LiveCodeBench.

Analysis of in-context reward hacking behaviors.

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.