COBALT: Contextual Bandit Learning with Offline Trajectories for Multi-Turn Code Generation

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Today's paper: "COBALT: Contextual Bandit Learning with Offline Trajectories for Multi-Turn Code Generation".

Jane: The gist The COBALT method proposes contextual bandit learning to bridge online and offline reinforcement learning for multi-turn code generation,

Tom: First, who's behind it and why it matters.

Paper summary: Tom: So folks, today we're looking at a paper called "COBALT: Contextual Bandit Learning with Offline Trajectories for Multi-Turn Code Generation." It tackles how to train large language models to write code that takes multiple turns.

Jane: That’s right. The main idea here is bridging the gap between online reinforcement learning, which is great but costly, and offline RL, which uses existing data but can be limited. COBALT proposes a way to use both together efficiently for this coding task.

Lu: Basically, they look at multi-turn code generation as a one-step recoverable Markov decision process and suggest contextual bandit learning with offline trajectories to train these models.

Meng: One quick thing I want to say is that the paper claims this approach reuses offline multi-turn trajectories to efficiently train LLMs for multi-turn code generation.

Lalam: I think what they’re doing is treating the generation process as a sequence of decisions where each decision needs to be optimized step by step rather than planning the whole thing at once.

Tom: Exactly, and why this matters is because it makes training less expensive and more stable than the online RL methods we usually see used for these complex tasks.

Jane: They start by using a reference LLM to collect code generation trajectories, which they then divide into partial trajectories that act as contextual prompts.

Meng: So they are not training the model end-to-end across the whole sequence at once, but focusing on completing each of these smaller pieces through single-step code generation.

Tom: And during this online bandit learning phase, the LLM is trained to complete that specific partial trajectory prompt using just one step of code generation to maximize its reward.

Lu: The theoretical underpinning they use is the one-step recoverability property, which implies that the advantage function of an optimal policy has a uniform bound where it’s between negative one and zero.

Jane: So because of this, they argue that a suboptimal action doesn't cause huge problems down the line, which inspires them to train LLMs to greedily optimize each action.

Tom: That’s how they reduce the problem from sequential reinforcement learning to contextual bandit learning, which is a simpler structure for optimization.

Meng: The methodology has three main stages: offline trajectory collection, online bandit learning, and then inference for testing.

Lalam: In the first stage, they use a reference LLM to sample sixteen independent trajectories with a horizon of no more than three turns. They keep only those containing at least one correct program and then divide them up into partial trajectories for the online learning part.

Jane: The second stage is the online contextual bandit learning, where the objective is defined as Jstep(pi) equals one over T times the expected reward given a policy pi ref t h x at R(st, at) pi of a t given st.

Paper summary: Tom: So they are training the LLM to complete each partial trajectory prompt through single-step code generation based on that context and reward structure.

Lu: The theoretical analysis shows that this stepwise objective Jstep approximates the online multi-turn RL objective J with at most O(T times root of eta performance difference under KL regularization.

Jane: That means the performance gap between the stepwise optimal policy and the full multi-turn objective is bounded by O of T times root of eta, which simplifies to a linear error bound O of T under KL regularization.

Tom: That’s a big deal because it shows that contextual bandit learning can actually be well suited for multi-turn code generation tasks.

Meng: They also looked at how LLMs can exploit noisy or under-specified feedback during self-improvement, which is called in-context reward hacking.

Jane: To handle that, they augment the COBALT training with perturbed trajectories to make the models more robust against misleading execution feedback.

Lu: The perturbation test case involves exchanging the expected outputs of two public tests to create impossible test cases, and these cases are used as feedback at each turn during training.

Tom: The results on LiveCodeBench were pretty strong; COBALT improved R1-Distill 8B by up to nine point zero absolute Pass@one points and Qwen3 8B by up to six point two points.

Jane: And in terms of comparison, COBALT matches the performance of two multi-turn online RL baselines, GRPO and VeRPO, in the multi-turn setting.

Meng: On a practical side, they also noted that the training efficiency is good; taking only about sixteen point nine seconds on a training example with four GPUs compared to about twenty-six point seven seconds for VeRPO-MT.

Lalam: For me, this means we have a method that’s both effective and faster to train, which makes iterating on these complex coding tasks much more feasible for development teams.

Tom: So, to wrap up the paper "COBALT: Contextual Bandit Learning with Offline Trajectories for Multi-Turn Code Generation," it establishes a promising solution by combining offline data generation with online stepwise exploration and optimization.

Jane: It shows that contextual bandit learning is a viable framework for improving LLMs in multi-turn code generation, especially when dealing with the cost and instability of purely online reinforcement learning.

Lu: The implication here is that this approach suggests that contextual bandit learning might be useful for training LLMs on other iterative decision-making tasks like mathematical reasoning or deep research.

Meng: From an engineering standpoint, the perturbation-based data augmentation is a scalable way to improve robustness against incorrect execution feedback in these models.

Lalam: I think the main point is that COBALT improves both training efficiency and learning efficacy simultaneously for this specific task.

Conclusion: Tom: So we’re wrapping up this look at COBALT, which is "Contextual Bandit Learning with Offline Trajectories for Multi-Turn Code Generation."

Jane: It really boils down to taking those big online reinforcement learning problems and making them much more manageable by using offline data smartly.

Lu: The authors are showing that they can use a reference LLM to gather trajectories and then train the target model step-by-step on those pieces.

Meng: So instead of trying to solve the whole multi-turn code generation problem at once, you’re training it to just complete the next single step based on what happened before.

Tom: Exactly, and this approach uses contextual bandit learning because it simplifies the decision process down to picking the best single action at each moment.

Jane: The results show that this method actually beats some of the established online methods for multi-turn coding tasks, like GRPO and VeRPO.

Lu: And they even showed how you can make these models more resilient to getting bad feedback during self-correction by using those perturbed trajectories.

Meng: It seems like a really practical way to improve the stability of these complex code generation systems when you’re actually trying to deploy them.

Tom: The implication here is that we might be able to train powerful LLMs for coding tasks much more efficiently without needing massive amounts of expensive online interaction data.

Jane: COBALT suggests that offline data, when structured correctly, can be a really powerful tool for iterative decision-making processes like coding.

Lu: It opens up possibilities because it points toward using this kind of contextual bandit approach for other complex reasoning tasks besides just writing code.