$\mathcal{B}$-Coder: Value-Based Deep Reinforcement Learning for Program Synthesis
Authors: Zishun Yu
Yunzhe Tao
Liyu Chen
Tao Sun
Hongxia Yang
Venue: ICLR 2024 spotlight (2024)
Keywords: Program Synthesis
Code Generation
Reinforcement Learning
Value-Based RL
Primary area: reinforcement learning

Abstract: Program synthesis aims to create accurate, executable programs from problem specifications, specifically from natural language descriptions in our context. 
Recent studies have leveraged the power of reinforcement learning (RL) in conjunction with large language models (LLMs), significantly enhancing code generation capabilities. The application of RL focuses on directly optimizing for functional correctness, offering an advantage over conventional supervised methods. 
Despite policy-based RL methods dominating the literature on RL for program synthesis, the nature of program synthesis tasks hints at a natural alignment with value-based methods.
This stems from the rich collection of off-policy programs, including those developed by human programmers and also historical samples, coupled with the straightforward verification of generated programs through automated unit testing, meaning rewards are easy to obtain.
Diverging from the dominant use of policy-based algorithms, our work explores the feasibility of value-based approaches, leading to the development of our $\mathcal{B}$-Coder (pronounced Bellman coder).
Yet, training value-based methods presents challenges due to the enormous search space inherent to program synthesis. 
To this end, we introduce an initialization protocol for RL agents utilizing pre-trained LMs and a conservative Bellman operator to reduce training complexities. 
Moreover, we demonstrate how to leverage the learned value functions as a dual strategy to post-process generated programs. 
Our empirical evaluations demonstrated $\mathcal{B}$-Coder's capability in achieving state-of-the-art performance when compared to policy-based methods. 
Remarkably, this achievement is reached with minimal reward engineering effort, highlighting the effectiveness of value-based RL, independent of reward designs.

Published as a conference paper at ICLR 2024
B-CODER: VALUE-BASED DEEP REINFORCEMENT LEARNING FOR PROGRAM SYNTHESIS
Zishun Yu"
Yunzhe Tao, Liyu Chen, Tao Sun & Hongxia Yang ByteDance Inc.
Department of Computer Science University of Illinois Chicago Chicago, IL 60607 zyu32Quic.eds
Seattle, WA 98004
{yunzhe. tao, Liyu. tao sun, hx wx..yag@@byyeearcc.cc com
ABSTRACT
Program synthesis aims to create accurate, executable programs from problem specifications, specifically from natural language descriptions our context. Re- cent studies have leveraged the power of reinforcement learning (RL.) in conjunc- tion with large language models (LLMs), significantly enhancing code generation capabilities. The application of RL focuses on directly optimizing for functional correctness, offering an advantage over conventional supervised methods. De- spite policy-based RL methods dominating the literature RL for program syn- thesis, the nature of program synthesis tasks hints at a natural alignment with value-based methods. This stems from the rich collection of off-policy programs, including those developed by human programmers and also historical samples, coupled with the straightforward verification of generated programs through au- tomated unit testing, meaning rewards are easy to obtain. Diverging from the dominant use of policy-based algorithms, our work explores the feasibility of value-based approaches, leading to the development of our B-Coder (pronounced Bellman coder). Yet, training value-based methods presents challenges due to the enormous search space inherent to program synthesis. To this end, we introduce an initialization protocol for RL agents utilizing pre-trained LMs and a conserva- tive Bellman operator to reduce training complexities. Moreover, demonstrate how to leverage the learned value functions as dual strategy to erated programs. Our empirical evaluations demonstrated B-Coder's capability in achieving state-of-the-art performance when compared to policy-base methods. Remarkably, this achievement is reached with minimal reward engineering effort, highlighting the effectiveness of value-based RL, independent of reward designs.
gen-
INTRODUCTION
Program synthesis (or code generation) aims to create funetionally accurate executable programs from problem specifications, such as input-output (IO) examples (Summers, 1977; Gulwani et al., 2012), constraint- based (Osera & Zdancewic, 2015; Frankle et al., 2016) or natural language de- scriptions (Hendrycks et al., 2021; Austin et al., 2021), among others. The increasing attention to- wards this field can be attributed to its potential transforming the software development paradigm. Notably, Al-powered tools have shown evidence of boosting efficiency within the software industry.
Large language models (LLMs) (Brown et al., 2020; OpenAI, 2023; Anil et al., 2023; Chowdhery al., 2022; Rae al., 2021; Hoffmann al., 2022; Touvron et al., 2023) have garnered substantial interest and shown remarkable achievements. The scheme of pre-training on vast amounts of data has yielded notable successes in natural language generation. This trend extends its influence to program synthesis, where numerous specialized code LLMs (Li et al., 2023; 2022; Nijkamp et al., 2022: Zheng et al., 2023; Fried et al., 2022; Chen et al., 202 Wang et al., 2021; 2023; Xu al., 2023; Rozière et al., 2023) have been introduced to address challenges in program synthesis.
Unlike many free-form natural language generation tasks, where the quality of model's output is hard to assess, the correctness of synthesized programs can be verified through automated execution with
"This work done during internship at ByteDance Inc. Correspondence Zishun Yu.

---

Published as a conference paper at ICLR 2024
predefined unit tests. This allows for directly optimizing execution outcomes through reinforcement learning (RL), by formulating test outcomes as reward signals. Our discussion focuses on recent RL- based works (Le et al., 2022; Shojaee et al., 2023; Liu et al., 2023) that have achieved remarkable advancements in Python text-to-code generation, evaluated on the challenging benchmarks sourced from Codeforces programming contests (Hendrycks et al., 2021; Li et al., 2022) Notably, these works predominantly favor on-policy policy-based algorithms.
While (on-policy) policy-based methods are favored in existing program synthesis works, they are known to be sample inefficient (Nachum et al., 2017; Gu et al., 2016) due to their inability to use off-policy samples. In contrast, value-basee methods, using temporal difference learning, are known to be more sample-efficient (Gu et al., 2016; Nachum et al., 2017; Liu et al., 2020), as they solve a fixed-point iteration which does not explicitly require a specific data distribution, hence offering better compatibility with off-policy data. We defer the technical explanations on on/off-policy data and reasons for the different efficiency to Section 3.2, where we have notations and definitions ready.
In program synthesis, the primary sources of off-policy data include human programs and previously synthesized programs. Both are off-policy as they do not follow the sequence distribution induced by the current model. Current program synthesis works often directly use off-policy samples with on- policy methods. Unsurprisingly, Shojace et al. (2023) notices that an increase in off-policy synthetic programs may degrade performance. This occurs as off policy data lead to biased gradient estimates. Ideally, an objective should be to enhance or at least sustain performance as data volume grows.
To summarize, the reasons that suggest natural fit for value-based methods in program synthesis twofold: the availability of (inexpensive) rewards, similar to classical RL tasks like GÖ and Atari; and the principle compatibility with off-policy data for effectively leveraging human and historical data. However, value-based RL faces challenges such as difficulty in converging in large state-action spaces. To this end, we introduce B-Coder (Bellman coder), with our contributions being threefold:
We stabilize value-based RL for program synthesis by proposing an initialization protocol for Q-functions and a conservative Bellman operator to mitigate the training complexities.
We demonstrate how to leverage value functions dual strategy improve generation.
. B-Coder achieves strong empirical performance with minimal reward engineering, provid- ing further insights of RL algorithm design independent of reward function designs.
Paper structure. We introduce related works and notations in Section 2 and 3. Section 4 details our method and the rationale behind our design choices. Specifically, Sections 4.1, 4.2, and 4.3 address the challenges of value function training by: leveraging task structure, providing effective Q- function initialization, and a conservative operator for stable yet less ambitious updates, respectively. Section 4.5 shows additional benefit of value functions, and Section 5 shows our empirical results.
2 RELATED WORKS
Execution-guided program synthesis. The feasibility of verifying programs through test case out- comes has led to the line of execution-guided works (Chen al., 2018; Zohar & Wolf, 2018; Chen et al., 2021b). While these efforts leverage execution feedback, they do not directly optimize towards higher execution success rate due the inherent non-differentiability of execution outcomes.
RL for general sequence modeling, Supervised LM training, using next token predictions (NTP) or masked language modeling (Kenton & Toutanova, 2019), has recognized limitations. One prominent issue is the exposure bias: given that the training is done in a "teacher-
manner (Bengio et al., 2015: Ranzato et al., 2015), errors tend to accumulate during testing due to rego-segressive genera- tion. In contrast, prior works (Ranzato et al., 2015; Rennie et al., 2017) have demonstrated the effi- cacy of RL. in addressing exposure bias and optimizing non-differentiable metrics, e.g. BI.UU (Pap- ineni et al., 2002) and ROUGE (Lin, 2004), by leveraging automatic scoring as reward function.
RL for program synthesis. Supervised losses also fall short when assessing the functional accuracy of synthesized programs (Hendrycks et al., 2021; Chen et al., 2021a). As such, relying solely on su- pervised learning for program synthesis is not ideal. As RL provides a pathway to directly optimize non-differentiable objectives. plentiful work (Zhong et al., 2017; Simmons-Edle et al., 2018; Ellis et al., 2019; Wang et al., 2022) have studied enhancing code generation through RL. For the works most related to ours: CodeRL. (Le et al., 2022) adapted REIFFOCEE (Williams, 1992), a classic
2

---

Published as a conference paper at ICLR 2024
policy gradient (PG) algorithm, along with the baseline trick for variance reduction and supervise- trained reward model to alleviate the issue of sparse execution signals. In addition, they proposed critic sampling strategy to refine and repair program based the example unit tests feedback. PPOCoder (Shojaee et al., 2023) applied proximal policy gradient (Schulman et al., 2017, PPO) to fine-tune pre-trained LMs. In addition, they leverage the syntactic and semantic structure of code, such as syntax trees (Rabinovich et al., 2017) and data-flow graphs (Yasunaga & Liang, 2020), to improve reward function designs. RLTF (Liu et al., 2023) proposed an online training framework for program synthesis using policy gradient with heursitically desigged fine-grained rewards.
Additional discussions. Appendix D lists several RL applications, showing the analogies between program synthesis and tasks that benefit from value- based methods. In C, we extend the discussion on works that extend poliey-based methods to an off-policy setting. Such attempts often involve training a value function, further highlighting our motivation for starting with value-based methods.
3 PRELIMINARIES
One could formulate the program synthesis task as a to-suueeeeee generation task, where a model takes a problem description D as input and outputs a program W which aims to achieve the functionality specified by D. A generated program W = (wn...., wT) is composed by a sequence use constant T to denote the sequence length although it could
of tokens v0, E V. For brevity, be a variable in practice, and W to denote a program in general (both generated and ground truth). Let LM be an instance of LM, (( D).-) be the logits layer (language modelling head) output, and p(-/w D) be the probabilistic distribution over the vocabulary V (computed by passing {(.- through softmax), conditioned on a sequence Wet and D. Suppose W* is a ground truth program and Duain is the train set, conventionally LMs could be trained by minimizing the cross-entropy loss
Lee(p) = EEw logp(W"|D) = -EW. E, log p[ [22..
(I)
3.1 RL NOTATIONS
To make notations casier to interpret, we bridge program synthesis notations to standard RL ones. RL problems are typically formulated as Markov Decision Processes (MDPs) and an MDP M is often composed by 5-tuple M=(S, A, P,P, r) which are state space, action space, transition func- tion, reward function and discount factor, respectively. The discount factor y discounts future values emphasize the near futures, and we use 0=.999 (which slightly prefers more concise solution). A (stochastic) transition function P S x A + A(S) is a distribution over S conditioned on a state- action pair (s, a). In program synthesis, IP trivial as s = Sy o at, where denotes concatenation.
State and action. In code generation context, an action is token wt. Hence the action space A the vocabulary V. As the information used to generate token ú; is (w<r, D), the state is hence defined as ( D). For a given D, the state space S = yT For brevity, we will mainly use St, de rather than the we notations, and sometimes omit the time index t if it leads to no confusion. We will also use s', a' to denote 8; 1.41 whenever only the relative temporal position matters.
Policy, A policy æ : S - - ((A) assigns an action distribution A(A) to any state S meaning predicting a token wi based on current sequence Ver and the problem specification D. Prior works often define mo = Po and directly optimize LM parameters 0 with PG methods. We however define : f(0 D) to be function of 0 and other components D, see details in Section 4.
Reward function. A reward function r : S A + I de- termines reward of taking action a, at state St. We follow the reward design of Le et al. (2022) in equation 2. We may also use shorthand notation Ft := r(st, oz). Note that the reward is determined when the program W is completed at
r() ((8r,ar)=
+1.0, W passed all unit tests
0.3, if W failed any unit test
0.6, if W cannot be executed
1.0, if W cannot be compiled
T. Thus r 0 if + T otherwise defined equation 2.
(2)
Value functions. RL maximizes the discounted returns, value function Qº S x A- and the state value function V= S RR are defined recursively as:
"r The state-action
V'((s)= == E (s. := EE o M, = s, Ao a) E, (r(s + (s',
M, So = z] = -uaaa(()),: (r(s,a) + (s))
(3)
(4)
3

---

Published as a conference paper at ICLR 2024
where Q(s,a) In addition, the advantage function is (=ssa))=((6
3.2 VALUE-BASED RL AND DUELING DQN
Value-based algorithms especially the Q-learning family (Watkins & Dayan, 1992; Mnih et al., 2013; Van Hasselt et al., 2016; Bellemare et al., 2017) have achieved remarkable successes. A canonical framework of the Q-learning family iterates between policy evaluation and policy improvement:
policy evaluation (PE): Qs arg minq EplQx-1(s,a) (r 7Qx-1 [k-)
(5)
policy improvement (PI): Tk = arg max,
(6)
where 'D is an arbitrary dataset, the PE step estimates the previous policy πx-1 using the Bellman equation (Bellman, 1966), and the Pl step finds improved mi.. by maximizing Qs estimates.
In particular, we build our framework on top of Dueling DQN (Wang et al., 2016, DDQN). a nutshell, DDQN approximates (s)) and A(s, a) with separate heads, and run improvement and evaluation steps with Q(s, a) (s))+ + A(s, a). This bifurcation enables a robust estimation of Vss) without conflating with the actions, which subsequently ensures stable learning of A(s, a) given that it focuses solely on the relative values. As consequence, DDQN often exhibits enhanced stability in training dynamics and improved generalization. In addition to the prior mentioned advan- tages, DDQN enables us to leverage a task structure that ground truth programs should attain highest advantages, therefore reducing the searching space, which we will elaborate on Section 4.1.
Remarks on sample efficiency. We illustrate the inefficiency of policy-based methods using vanilla PG as an example. PG maximizes (()) == ETΣ, >' [T MM = Ew [E,,"'r) ?"r]). with gradient (() computed using the policy gradient theorem. This method requires training data W drawn from the distribution induced by current policy πu. hence called on-policy. Therefore, one should in principle generate new data and discard historical data every update, leading to undesired sample inefficiency. In contrast, policy evaluation as in equation works with arbitrary dataset D.
4 ALGORIHHMIC DESIGNS ACCELERATING VALUE-BASHE TRAINING
While value-based RL holds great promise, its training can be challenging due to the large action space A= V and the high- dimensional state space S= V" This leads to a notably large Q-table of size ((V(T). And the cardinality of policy space is - ( ), which grows doubly exponentially. Both challenges from large action spaces and high-dimensional state spaces are pivotal research topics in RL. The action space challenges are discussed by e.g. Dulac-Arnol et al. (2015); Tavakoli et al. (2018); Kalashnikov et al. (2018), while He et al. (2016); Nair et al. (2018), among others, considered the state spaces complexities. In partieular, Silver (2015); Duan et al. (2016) commented on that the potentially better training stability of policy-based methods in these scenarios.
Ours
Conservative Ours Conservstive Initiiaiizaioon
5000 10200 15100 2000: 20002 20200 35002
training iterations
Figure 1: Training curves on APPS train set. denotes B-Coder, re- moves our conservative operator, and v is B-Coder without both our operator and initialization.
To address the challenges inherent in training value-based RL for LMs, at high level, we developed B-Coder considering three key aspects: incorporation of task structure, initialization of Q-function, and backup using a conservative Bellman operator. Figure 1 previews the effectiveness of our al- gorithmic designs, which shows the training curve of different aalue-based RL algorithms on the APPS dataset. Due to aforementioned challenges, the performance of the vanilla DDQN continu- ously decreases even evaluated on the training set. In contrast, both the Q-function initialization and the conservative Bellman operator show benefits in stabilizing and accelerating the training process.
For notational convenience in subsequent sections, we begin with an overview of our notations and parameterizations, summarized in Figure 2. Figure 2(a) denotes a pre-trained encoder-decoder LM parameterized by fag. (where subscript ckpt denotes the fact it's a checkppint/consttan)) Figure 2(b) and (c) show the forward graphs of our two different training stages: (b) corresponds to a pre-training stage for o, to provide a good initialization for (c) the subsequent fine-tuning of d. Motivations and details are deferred Section 4.2 and 4.3, respectively. As we proceed to the rationale behind our
4

---

Published as a conference paper at ICLR 2024
designs, it is encouraged to maintain familiarity with app 6, and their corresponding products, especially the forward paths to Qo and Qs. to prevent confusion in the subsequent sections.
4.1 LEVERAGING TASK STRUCTURES
(a]
dkgss Im head
enc
der
Poass
As noted earlier, a key attribute of program synthesis task is the provision of human solutions, which are guaranteed to be cor- rect. As result, these solutions should at- tain the highest Q-values, even if the cor- rect solutions might not be unique. As such, for a ground truth program W' = (sh aj: .... Sip, o,), Q(si, a Z Q(st, a) holds for all a E V, hence A(st, d7) 2 A(st, a).
(b)
Ø
Ve
Ockpe
dec
Im head
ADcepp
Q+
(c)
Im head
Ep
Ap
*
dee
Qo
residual head
To enforce this structure, one could ensure A(W) < 0 and A(W*) x O, where we abuse the notation and by letting A(W) Σi= A(s,, at). I It ensures that W has ad- vantages that are roughly the highest. To this end, suppose g(-) is a general neural network, decompose Q as follows,
Ocipps
$
*
Va
Figure 2: (a) A forward graph of conventional enc-dec LMs, with a checkpoint Oaz,, where p a distribution over A and & denotes logits (b) Our forward graph for pre-training d; (c) Our forward graph for fine-tuning 0. * indicates frozen/constant component.
as. = g(s,a) max g(s, a) +V(s) A(s, a) + V(s).
non-positive advantage
It enforces our first condition that A(W) S 0. For the second condition A(W*) a O, we optimize an advantage function A by minimizing an auxiliary advantage loss function, namely Ladv.
Lad (A) = a6 )oom
E A(5 at)l
(8)
We also cap the Q-function with Rmax = 1, the maximum total rewards. See Appendix for details.
4.2 Q-FUNCTION INITIALIZATION
Despite the task structures introduced, training the Q-function from scratch remains extremely chal- lenging. While this is not problem for policy-based learning (given that directly fine-tune pre- trained LMs without requiring a Q-function at all), it presents significant challenges in value-based approaches because one often does not have a pre-trained Q-function. To this end, we show that one could initialize a Q-function from the logits output E(.,-) of a pre-trainee LMM
Initialization of Q via pre-trained models. Yu & Zhang (2023) considered the fine-tuning of RL agents after offline RL pre-training. Their main idea is to reconstruct a Q-function from the pre-trained policy, for fine-tuning. Drawing inspiration from their approach, one could similarly reconstruct/initiaiize Q-function using a pre-trained LM, akin to using a pre-trained policy.
This initialization was motivated by the mneryy-based policy line of works (Haarnoja et al., 2017; 2018 8)18) where a policy æ is the product of passing a Q-function through softmax transfer function. Analogously, in LMs, p- p-the the distribution over V is produced by passing logits é through softmax.
language modeling:
p(a]s) exp a)) EaeA exp (e(s,a))
energy-based x: (alss) exp (ss a)) EasA exp (s,
(10)
where a a temperature hyper-parameter. One could naturally set Q(s, a) = al(s, (s a) for initial- ization. Hence, with aforementioned dueling structure in equation 7 and our pre-defined parame- terization, one could set the advantage function as A a) := a alloss (s, a) max, to (s, a)}, leading to (s, a) == A (s, + Vs)).. See also our forward pass graph defined in Figure 2b. In a nutshell, this Qs funffnttionn produces policy TH identical to the output distribution Pray of LM
To ( softmax| Qg softmax (losp (8) max lo (8, a)+ a)+ Vo = Pa (a|s), (11)
5

---

Published as a conference paper at ICLR 2024
where Q(s) a))aaA and E(s) := [K(s, a.s,a)))....
Recalling equation 5 - 6, the Q-learning family can be viewed as iterations between policy evaluation and improvement. We now elaborate on how this Qo function initialization affects both steps.
Policy improvement. One could, informally, consider the operation of taking softmax with respect to Qo as seft policy improvement (Haarnoj et al., 2018) step with temperature a. Therefore, equation 11 can be interpreted as: running soft policy improvement alone with this initialized Q+ preserved the performance of pre- rre-trained LM offering good starting point of online fine-tuning.
Policy evaluation. Yet, this Qo function only captures relative values, since initialized only the advantages A0age the relative information as shown in equation 1.. V& can thereby be an arbitrary function. This would not affect the policy improvement step due to the translation invariance of the softmax function. However, during the policy evaluation step, see e.g. equation 5, the Bellman error can be heavily influenced by the V-values. When the V-values is the dominant source of error, the policy evaluation optimization could be largely driven by the state-only V-values. This can lead to a loss of the relative action values, that we intended to preserve in the previous step.
Pre-training of Vo. This can be addressed by adding a pre-training phase of (a(s), during which we we freeze the advantage function Ag and train Vi by minimizing the temporal difference error (or equivalently doing policy evaluation). In this stage, we optimize the following loss until convergence
(Va; ) + (s
Et=o [r+ + SG (1(+++11,(+)))
(5(s,, at)]². (12)
where SG is a stop gradient operator, (Q a')) follows standard semi mrnniggadient optimization, åe41 is a target action (details deferred to section 4.3), and Qp(s. a) = A0 (s. + (s)).
In summary, initialization steps ensures that, prior to fine-tuning 0, our Q meets two important conditions: it starts with the action distribution ar of a pre-trained LM and it begins with low temporal difference error (because the pre- pre-traine training of V in equation 12 directly minimizes it).
4.3 A CONSERVATIVE BELLMAN OPERATOR
With pre-trained state value function Va, are now ready to learn a good state-action value func- tion via fine-tuning. We parameterize Qo a) == Ao a)+ Vo() = alEo(s, a) -max, (o a)] VJ + Va. where we define Vi V + Ve. and we initialize 0 in a way such that Ey la and VJ" 0. It ensures that Qa Q+ initialization, good starting point for subsequent fine-tuning on 0. Technically speaking, setting Vo = VT + V& is not required, as one could finetune both 0 and 6. We however observed that finetuning residual head VT. with ø frozen, leads to better stability.
Although we avoid training Qa from scratch, optimizing Qo by Q-learning family algorithms can still be challenging. We attribute this to the characteristics of the Bellman optimality operator 3* that seeks to learn the optimal value function Q* and optimal policy x². which requires a good data cov- erage of the state-action space S x A (e.g. Jiang & Huang, 2020; Xie al., 202 a: Zhan et al., 2022). In program synthesis, however, such assumption can hardly be met due to the large state-action space and the high computational costs of Transformer inference. While conventional Q-learning family relies on the operator 13", recent works in RI., especially those considering limited data regime (e.g. Agarwal et al., 2020; Levine et al., 2020), often design "conservative" operators (c.g. Achiam et al., 2017; Kumar et al., 2020; Brandfonbrener et al., 2021) to address difficulties led by 1
Conservative Bellman operators. The concept behind conservative Bellman operators is to "aim low". Instead of learning the optimal Q" and n², these operators typically seeks to learn a policy that either surpasses a behavior policy (which is used to collect a RL dataset in offline RL literature, see e.g. Achiam et al., 2017: Brandfonbrener et al., 2021) or fine-tune a pre-existing policy (e.g. Xie et al., 2021b; Yu & Zhang, 2023). This is often achieved by introducing a regularizer that penalizes deviations from the behavior/pre-existin policy. In particular, as shown in equation 14, we define our conservative Bellman operator 189, which depends on a fixed, pre-defined policy q, as follows:
optimality B: (B* a) = r(s, a) + 7B [Q(s,, / a')], where à' = arg max, Q(s',a) a) conservative B: (B*Q)(s, a) = r(s,a) a) + / a')], where à' = arg max, q(als'). (14)
(13)
The intuition behind our operator 1 that we evaluate the action-value function Q of greed- ified policy q (als)= {a = arg max, q(a)s)), where I is the indicator function. The rationale
6

---

Published as a conference paper at ICLR 2024
behind greedification is that q" can be seen as q in greedy-decoding mode, which usually has bet- ter (one-shot) capability than sampling mode (although the latter has better generation diversity). Considering setting 9 Poasp the operator B seeks to learn policy и that outperforms Poaae"
We further comment some properties of 8y; proposition 4.1 shows 8 is contraction, meaning there is an unique fixed point. It leads to proposition 4.2, motivating our development of Section 4.5. Proposition 4.1. 8 is 7-contraction in &e norm.
Given our conservative Bellman operator, we could define our conservative temporal difference loss, (.+1)) a=)s (15)
(o(Qe:q) - 7 (s (e Dai En ir, SG
where 0 arg max, ((a(se+1). and Qu(s, = (Co(s,a) max. (o(8.u)) (s) + Vo(+).
4.4 IMPLEMENTATION AND OPTIMIZATION
Architecture and parameterization recap. Following (Le et al., 2022; Shojaee al., 2023; Liu et al., 2023), we choose T5 (Raffel et al., 2020) as our base architecture for ag Φ and #; and 0 is initialized with CodeRL checkpoint which is publicly available. Specifically, 0spis e and share a same encoder, and the encoder is frozen throughout, to reduce the amount of learnable parameters.
Two-stage training. As noted earlier, our training are composed with two stages: a pre-training stage of é, namely ø-stage, and fine-tuning stage of 0, namely 0-stage. A pseudo-algorithm could be found in Appendix A. In addition, further implementation details are deferred to Appendix H.
d-stage: Given our development of Section 4.2, we pre-train Vo function using stochastic gradient descent with VoLv (Va: (. with Ly defined in equation 12.
0-stage (fine-tuning): In this stage, we seek to optimize Qe to minimize our previously developed losses: Lady and LQ. defined in equation 8 and 15, respectively. In addition, it is also common practice to include cross-entrope loss L during fine- tuning. Therefore, we conclude our final loss function equation 17, and A is updated using stochastic gradient descent with Vo (Qo: Poay
Recall: Qo a) - As a) Vs(s)
(Lo(s, a) max, Co(s, (s) L (Qo: Paeam == La (Qa: am ) + Pady Lad (Ag) Bo Cae (To), where T. softmax (+Q (17)
(16)
4.5 A FREE REWARD MODEL
Reward modeling is crucial in language modeling and also in inverse RL (detailed discussions could be found in Appendix C). An intriguing finding from IRL, applicable to our framework, is that a trained Q-function can recover reward function without additional training. Analogously to Garg et al. (2021), an one-to-one correspondence between Q and reward holds with our conservative Bellman operator Ba, We define the inverse conservative Bellman operator T RSSAA -, S
(T9Q) a) = Q(s, a) - EE, Q (s', arg max, q(als')). Proposition 4.2. The inverse conservative Bellman operator T is bijection.
(18)
Proposition 4.2 shows that Qo uniquely corresponding a reward function ie == TuQo.! Given the definition of T we could recover a reward model ia with Qa without additional training:
a'a (s a) = (s a) Qo (s', arg max, Po (als)) = Qo(s, a) (Vo
(19)
We use the estimation fo a) == (s. aVo(s" in practice, with reasons deferred to Appendix F.
Candidates selection with ig. We leverage our reward model ig to do candidate programs selection, as example to highlight the additional benefits of value- based RL. We rank generated programs by the cumulative rewards Ro (W) := ΣT (s( at). predicted by our reward model 7g. to select the programs that are most likely to be correct. Specifically. for pass @1 metrics, we follow the evaluation protocol used in CodeT (Chen et al., 2022), work that considered program selection via automatic generated tests. This protocol computes pass@1 k by first generating programs and select a subset programs to evaluate pass @ k. In our case, we select the k-sized subset with top-k highest Rt(.)) from total m candidates. Our results Section 5 follow this evaluation protocol. We use and name our recovered reward model and the real reward function, respectively.

---

Published as a conference paper at ICLR 2024
Table 1: Empirical evaluation on APPS test set. 1, and # indicates results duplicated from Le et al. (2022), Shojaee al. (2023) and Liu et al. (2023), respectively. Bold number indicates the best result and underlined number means our result are the second best. Intro, inter and comp stand for introductory, interview and competition, respectively.
trainable
Pass@1
P
@ass 1000
Model
parameters
Intro
Inter
Comp
All
Intro
Inter
Comp
All
Intro
Inter
Comp
All
Codex
12B
4.14
0.14
0.02
0.92
9.65
0.51
0.09
2,25
25.02
3.70
3.23
7.87
AlphaCodef
IB
17.67
5.24
7.06
8.09
GPT31
175B
0.20
0.03
0.00
0.06
GPT2!
0.1B
1.00
0.33
0.00
0.40
2.70
0.73
0.00
1.02
GPT21
1.5B
1.30
0.70
0.00
0.68
3.60
1.03
0.00
1.34
25.00
9.27
8.80
12.32
GPT-Neo
2.7B
3.90
0.57
0.00
1.12
5.50
0.80
0.00
1.58
27.90
9.83
11,40
13.76
GPT-J'
6B
5.60
1,00
0.50
1.82
9.20
1.73
1,00
3,08
35.20
13.15
13.51
17.63
RL based methods-
without using example unit tests
CodeRL.
770M
6.20
1,50
0.30
2.20
9.39
1.90
0.42
3.10
35.30
13.33
13.60
17.78
PPOCoder!
770M
5.20
1.00
0.50
1.74
9.10
2.50
1.20
3.56
35.20
13.35
13.90
17.77
RLTF#:
770M
4.16
0.97
0.20
1.45
10.12
2.65
0.82
3.78
38,30
15.13
15,90
19.92
B-Coder
<770M/xraee
6.70
1.50
2.30
10.40
2.63
0.70
3.80
37.00
13.67
12.60
18.12
Remarks on ña. To further explain the motivation of rank- ing with Po. consider a realistic deployment setting where a fine-tuned model is deployed for end-user applications. Users often provide a language description of their needs but may not include test cases (which can also be challeng- ing for beginners casual users). Additionally, the model is usually required to offer a single best response instead of range of options. Therefore, the ability to rank programs without true rewards is desirable advantage.
compile ermrr
Derelity
0.15
0.3
To preview the effectiveness of fo. show the correlation between environmental reward r and our cumulative re- ward Rg. In Figure 3, green region corresponds to correct programs, and has the highest Rg on average. For incor- rect programs, those with compile and runtime errors have the lowest and the second lowest Rg. respectively. Pro- grams can be executed but fail some tests, have the second highest Rg. Hence, it concludes that Rg has evident positive correlation to the true reward r.
Oumulative Rewards
Figure 3: Kernel density estimation of R- evaluated on a collection of gen- erated programs. The x-axis represents the predicted reward given by Rg and the y-axis is its density. Color codes the true outcomes defined in equation 2.
EMPIRICAL EVALUATION
Sampling using Qa. Nucleus sampling (top-p sampling) (Holtzman al., 2019) with sampling temperature' (Ackley et al., 1985) has been one of the most important sampling techniques. It can also be easily implemented in our framework. One could simply consider Qa/a as logits and the sampling procedure would remain identical to standard LMS, see Appendix B for details.
APPS benchmark and baselines. In line with prior RL-based works (Le et al., 2022; Shojaee et al., 2023; Liu et al., 2023). we evaluate B-Coder on the challenging code contests benchmark APPS (Hendrycks et al., 2021). contains 5,000 training and 5,000 testing problems, with three dif- ficulty levels: introductory, interview and competition. We compare our B-Coder with pre-trained or supervise fine-tuned LLM baselines: GPT2 (Radford al., 2019), GPT3 (Brown et al., 2020), GPT- Neo (Black et al., 2021), GPT-J (Wang& & Komatsuzaki, 2021), Codex (Chen et al., 2021a) and Al- phaCode (Li et al., 2022); and RL fine-tuned baselines: CodeRL (Le et al., 2022), PPOCoder (Sho- jaee et al., 2023) and a concurrent work RLTF (Liu et al., 2023).
APPS: without example test outcomes. In the APPS dataset, each problem has several example unit tests (different from the hidden unit tests used for evaluation). These example tests are usually leveraged to refine generated samples. For example, CodeRL and RLTF considers a critic sampling
"Sampling temperature is different from temperature a in equation 10. They can he different values. Forr both & and 0-stage, our model trains decoder and heads, i.e. 5770M trainable params per stage.
8

---

Published as a conference paper at ICLR 2024
(CS) strategy that refines and repairs generated programs based the execution outcomes of the example tests. We start with experiments results in which example test outcomes are not used (hence CodeRL and RLTF results in Table are without CS). Table 1 shows that our B-Coder has overall the best pass@ l for k = (1, 5 and achieves second best place for k= 1000 (best result reported by the concurrent work RITTF). For Table 1 results,
nucleus sampling with a sampling temperature of 0.6. We set m to 256 for k= {1.5) 5) and m to 2500 for A = 1000, where m is parameter of our ranking protocol introduced in Section 4.5 (see Appendix for an ablation study on m).
APPS: using example test out- comes. Table lists the results using example tests. In addition to the CS strategy that uses ex- ample tests to refine/repair pro- grams, Li et al. (2022) and Chen et al. (2021a) consider filtered setting, in which programs fail- ing example tests excluded, and the pass @ ls is evaluated us-
Table 2: APPS results when using example test outcomes.
Model
Ou I
Pas
Intro
Comp
Intro
Comp
Codex! filtered
22.78
2.64
3.04
6.75
24.52
3.23
3.08
7.46
AlphaCode! fitered
14.36
5.63
4.58
7.17
CodeRL.
6.77
1.80
0.69
15.27
4.48
2.36
6.21
CodeRLI fittred
16.27
6.00
4.27
7.71
CodeRL estilltered
16.52
6,16
4.15
7.83
24.49
8.58
7.82
11.61
RLTTF
8.40
2,28
1.10
18.60
3.70
7.80
B-Coder littered
18.00
6.63
2.30
8.04
23.30
8.83
6.40
11.30
ing (a subset of) programs that pass example tests (which is also related to the k@@ metric (Li et al., 2022), the pass rate using l submissions from m. samples). We also test B-Coder in this filtered set- ting. Similarly, we first exclude programs that fail example tests. Suppose n out of m programs pass; we then follow our ranking protocol to get top-k out of n. programs for evaluation. B-Coder outper- forms baselines with either CS or filtered setting for k 51,5 The baseline, CodeRL+CS+fillered incorporated both strategies achieved a slight advantage over B-Coder for pass@5 5 while being sur- passed by B-Coder for pass @ It worth mentioning that CS is a alug-and-play component, which could also be combined with B-Coder, to further improve pass rate. For the results in Table 2, we use a temperature of 0.4 and ma set to 1000, matching the m used in Le et al. (2022).
Generalization ability. In addition, we test the generalization ability of our dual strat- egy, ranking with Rg. We study two aspects: generalization to other models and general- ization to different domains. To this end, we designed the following experiments, which confirmed its generalizability in positive.
Table 3: Generalization to CodeRL. Pass@k k eval- uated with top-k ranked programs, generated by CodeRL. indicates absolute improvement achieved by ranking, compared to un-ranked @
Pass@k
k Temp.
Intro
Inter
Comp
All
0.4
6.30
1.27
0.50 0.37
2.12 0.68
For the former, we generate (off-policy) pro- grams using CodeRL (with m = 256), and rank those programs by Rg. Table 3 shows our ranking strategy leads improvements in most cases, even though the programs to be ranked are not generated by B-Coder.
0.6
6.00
2.13
1.23 0.42
0.50 0.36
2.04 0.75
0.4
9.30
-0.2
2.10
0.01
0.70 0.15
3.26 0.00
0.6
10.20
0.58
2.57
0.41
0.80
0.16
3.74 0.39
Table 4: Zero-shot pass@k k on MBPP. absolute improvement achieved by ranking.
indicates
Temp.
k=1
k=5
k=10
k=80
For the latter, we test our dual strategy with another dataset MBPP (Austin al., 2021) (with m = 512). Table 4 shows consistent improvements for all temperatures and k.
0.7
20.13
6.61
37.04
5.61
44.45
4.63
64.00
1.41
0.8
18.89
6.99
36.59 7.21
44.46 6.59
65.20
4.28
0.9
17.32
7.34
35.04 8.58
43.15 8.22
63.20 4.33
CONCLUSION
In this work, we explore the feasibility of value-based RL algorithms for program synthesis task. We demonstrate how to stabilize and accelerate training through Q-function initialization and conser- vative updates. Moreover, our work is conducted with minimal reward engineering effort, thereby placing an emphasis on the perspective of algorithm designs. While policy- algorithms remain mainstream in the current program synthesis literature, the question of how to effectively lever- age off-policy programs, including historical synthetic samples, in principled way, might still be under-explored. We are convinced that value-basee RL offers a promising direction to address this question, and thereby to scale RL for code generation at large by (re)-using the extensive collection of off-policy programs. Our work could thus serve as an important initial step towards this direction.

---

Published as a conference paper at ICLR 2024
REFERENCES
Pieter Abbeel and Andrew Y Ng. Apprenticeship learning via inverse reinforcement learning. In Proceedings the twenty-first international conference on Machine learning, pp. 1, 2004.
Joshua Achiam, David Held, Aviv Tamar, and Pieter Abbeel. Constrained policy optimization. In International conference on machine learning, pp. 22-31. PMLR, 2017.
David H Ackley, Geoffrey E Hinton, and Terrence J Sejnowski. A learning algorithm for boltzmann machines. Cognitive science, 9(1 9()::147-19, 1985.
Rishabh Agarwal, Dale Schuurmans, and Mohammad Norouzi. An optimistic perspective offline reinforcement learning. In International Conference on Machine Learning, pp. 104-114. PMLR, 2020.
Rohan Anil, Andrew M Dai, Orhan Firat, Melvin Johnson, Dmitry Lepikhin, Alexandre Passos, Siamak Shakeri, Emanuel Taropa, Paige Bailey, Zhifeng Chen, al. Palm 2 technical report. arXiv preprint wwXivv2305..0440 2023.
Kai Arulkumaran, Antoine Cully, and Julian Togelius. Alphastar: An evolutionary computation perspective. In Proceedings of the genetic and evolutionary computation conference companion, pp. 314-315, 2019.
Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie Cai, Michael Terry Quoc Le, et al. Program synthesis with large language models. arXiv preprint vrXiv:2108007733 2021.
Yuntao Bai, Andy Jones, Kamal Ndousse, Amanda Askell, Anna Chen, Nova DasSarma, Dawn Drain, Stanislav Fort, Deep Ganguli, Tom Henighan, al. Training helpful and harmless assistant with reinforcement learning from human feedback. arXiv preprint arXiv:2204.08822 2022a.
Yuntao Bai, Saurav Kadavath, Sandipan Kundu, Amanda Askell, Jackson Kernion, Andy Jones, Anna Chen, Anna Goldie, Azalia Mirhoseini, Cameron McKinnon, et Constitutional ai: Harm- lessness from ai feedback. arXiv preprint arXiv:2212.080733 2022b.
Mare Bellemare, Will Dabney, and Rémi Munos. distributional perspective on reinforcement learning. In International conference machine learning, pp. 449-458. PMLR, 2017.
Richard Bellman. Dynamic programming. Science, 153(3731):34-37, 1966.
Samy Bengio, Oriol Vinyals, Navdeep Jaitly, and Noam Shazeer. Scheduled sampling for sequence prediction with recurrent neural networks. Advances in neural information processing systems, 28, 2015.
Sid Black, Leo Gao, Phil Wang, Connor Leahy, and Stella Biderman. GPT-Neo: Large Scale Autore- gressive Language Modeling with Mesh-Tensorflow, March 2021. URL https: / /oi org/ 10. 50.5221/zeee zenooo. 5297715. If you use this software, please cite it using these metadata.
David Brandfonbrener, Will Whitney, Rajesh Ranganath, and Joan Bruna. Offline without off- policy evaluation. Advances in neural information processing systems, 34:4933-4944 2021.
Noam Brown and Tuomas Sandholm. Superhuman ai for heads-up up no-limit poker: Libratus beats top professionals. Science, 359(6374):4184477 2018.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot leamners. Advances in neural information processing systems, 33:1877-1901, 2020.
Bei Chen, Fengji Zhang, Anh Nguyen, Daoguang Zan, Zeqi Lin. Jian Man-Guang Lou, and Weizhu Chen. Codet: Code generation with generated tests. In The Eleventh International Conference on Learning Representations, 2022.
10

---

Published as a conference paper at ICLR 2024
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. Evaluating large language models trained code. arXiv preprint arXiv:2107.03374, 2021a.
Xinyun Chen, Chang Liu, and Dawn Song. Execution-guideet neural program synthesis. In Interna- tional Conference on Learning Representations, 2018.
Xinyun Chen, Dawn Song, and Yuandong Tian. Latent execution for neural program synthesis beyond domain-specific languages. Advances in Neural Information Processing Systems, 34: 22196,22208, 2021b.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.000 02331, 2022.
Paul F Christiano, Jan Leike, Tom Brown, Miljan Martic, Shane Legg, and Dario Amodei. Deep reinforcement learning from human preferences. Advances in neural information processing sys- lems, 30, 2017.
Thomas Degris, Martha White, and Richard S Sutton. Off-policy actor-critic. arXiv preprint orXiv::205.48833 2012.
Shizhe Diao, Rui Pan, Hanze Dong, Ka Shun Shum, Jipeng Zhang, Wei Xiong, and Tong Zhang. Lmflow: An extensible toolkit for finetuning and inference of large foundation models. arXiv preprint arXiv:2306. Stti::2330.1222200 12420, 2023.
Yan Duan, Xi Chen, Rein Houthooft, John Schulman, and Pieter Abbeel. Benchmarking deep reinforcement learning for continuous control. In International conference on machine learning, pp. 1329-1338. PMLR, 2016.
Gabriel Dulac-Amold, Richard Evans, Hado van Hasselt, Peter Sunehag, Timothy Lillicrap, Jonathan Hunt, Timothy Mann, Theophane Weber, Thomas Degris, and Ben Coppin. Deep rein- forcement learning in large discrete action spaces. arXiv preprint 151i0:151 2015.
Kevin Ellis, Maxwell Nye, Yewen Pu, Felix Sosa, Josh Tenenbaum, and Armando Solar-Lezama. Write, execute, assess: Program synthesis with repl. Advances in Neural Information Processing Systems, 32, 2019.
Jonathan Frankle,
Peter-Michacl Osera, David Walker, and Steve Zdancewic. Example-directed synthesis: a type-theoretic interpretation. ACM Sigplan Notices, -1()):802-815,200
Daniel Fried, Armen Aghajanyan, Jessy Lin, Sida Wang, Eric Wallace, Freda Shi, Ruiqi Zhong, Scott Yih, Luke Zettlemoyer, and Mike Lewis. Incoder: A generative model for code infilling and synthesis. In The Eleventh International Conference Learning Representations, 2022.
Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor- critic methods. In International conference on machine learning, pp. 1587-1596. PMLR, 2018.
Divyansh Garg, Shuvam Chakraborty, Chris Cundy, Jiaming Song, and Stefano Ermon. Iq-learm: Inverse soft-q learning for imitation. Advances in Neural Information Processing Systems, 34: 4028 4039, 2021.
Shixiang Gu. Timothy Lillicrap. Zoubin Ghahramani, Richard Turner, and Sergey Levine. Q-prop: Samplc-efficient policy gradient with an off-policy critic. arXiv preprint 1611::6611.022477 2016.
Caglar Gulcehre, Tom Le Paine, Srivatsan Srinivasan, Ksenia Konyushkova, Lotte Weerts, Abhishek Sharma, Aditya Siddhant, Alex Ahern, Miaosen Wang, Chenjie Gu, et al. Reinforced self-training (rest) for language modeling. arXiv preprini arXiv:2308.008: 08998, 2023.
Sumit Gulwani, William R Harris, and Rishabh Singh. Spreadsheet data manipulation using exam- ples. Communications of the ACM, 55(8):77-155 2012.
Tuomas Haarojj,, Haoran Tang. Pieter Abbeel, and Sergey Levine. Reinforcement learning with deep energy-based policies. In International conference on machine learning, pp. 1352-1361. PMLR, 2017.
11

---

Published as a conference paper at ICLR 2024
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In International confer- on machine learning, pp. 1861-1870. PMLR, 2018
Frank S He, Yang Liu, Alexander G Schwing, and Jian Peng. Learning to play in day: Faster deep reinforcement learning by optimality tightening. International Conference on Learning Representations, 2016.
Dan Hendryeks, Steven Basart, Saurav Kadavath, Mantas Mazeika, Akul Arora, Ethan Guo, Collin Burns, Samir Puranik, Horace He, Dawn Song, et al. Measuring coding challenge competence with apps. arXiv preprint arXiv:2105. 09999998889 2021.
Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. Advances in neural information processing systems, 29, 2016.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Train- ing compute-optimal large language models. arXiv preprint arXiv:2203. 15556, 2022.
Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. The curious case of neural text degeneration. In International Conference on Learning Representations, 2019.
Ehsan Imani, Eric Graves, and Martha White. An off-policy policy gradient theorem using emphatic weightings. Advances in Neural Information Processing Systems, 31, 2018.
Alexis Jacq, Matthieu Geist, Ana Paiva, and Olivier Pietquin. Learning from learner. In Interna- tional Conference on Machine Learning, pp. 2990-2999. PMLR, 2019.
Nan Jiang and Jiawei Huang. Minimax value interval for off-policy evaluation and policy optimiza- tion. Advances in Neural Information Processing Systems, 33:2747-2755 2020.
Dmitry Kalashnikov, Alex Irpan, Peter Pastor, Julian Ibarz, Alexander Herzog, Eric Jang, Deirdre Quillen, Ethan Holly, Mrinal Kalakrishnan, Vincent Vanhoucke, et al. Scalable deep reinforce- ment learning for rision-based robotic manipulation. In Conference on Robot Learning, pp. 651- 673. PMLR, 2018
Alex Kendall, Jeffrey Hawke, David Janz, Przemyslaw Mazur, Daniele Reda, John-Mark Allen, Vinh-Dieu Lam, Alex Bewley, and Amar Shah. Learning to drive in day. In 2019 International Conference on Robotics and Automation (ICRA), pp. 8248-8288 8254. IEEE, 2019.
Jacob Devlin Ming-Wei Chang Kenton and Lee Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of NAACL-IILT, pp. 4171- 4186, 2019.
Vijay Konda and John Tsitsiklis. Actor-critic algorithms. Advances in neural information processing systems, 12, 1999.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems, 33:1 1191, 2020.
Hung Le, Yue Wang, Akhilesh Deepak Gotmare, Silvio Savarese, and Steven Chu Hong Hoi. Coderl: Mastering code generation through pretrained models and deep reinforcement learning. Advances in Neural Information Processing Systems, 0222114-21133 2022.
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tuto- rial, review, and perspectives on open problems. arXiv preprint arXiv:2005.0164) 2020.
Raymond Li, Loubna Ben Allal, Yangtian Zi, Niklas Muennighoff, Denis Kocetkov, Chenghao Mou, Mare Marone, Christopher Akiki, Jia Li, Jenny Chim, et al. Starcoder: may the source be with you! arxiv preprint 230.00330000666111 2023.
Yujia Li, David Choi, Junyoung Chung, Nate Kushman, Julian Schrittwieser. Rémi Leblond, Tom Eccles, James Keeling, Felix Gimeno, Agustin Dal Lago, et al. Competition- code generation with alphacode. Science, 37866624 2022.
12

---

Published as a conference paper at ICLR 2024
Chin- Lin. Rouge: A package for automatic evaluation of summaries. In Text summarization branches out, pp. 74 74,81,2200 2004.
Jiate Liu, Yiqin Zhu, Kaiwen Xiao, Qiang Fu, Xiao Han, Wei Yang, and Deheng Ye. Ritf: Rein- forcement learning from unit test feedback. arXiv preprint wrXiv:2307.04349 2023.
Yao Liu, Adith Swaminathan, Alekh Agarwal, and Emma Brunskill. Off-policy policy gradient with stationary distribution correction. In Uncertainty in artificial intelligence, pp. 1180-1190. PMLR, 2020.
lyy Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- ence on Learning Representations, 2018.
Ziyang Luo, Can Xu, Pu Zhao, Qingfeng Sun, Xiubo Geng, Wenxiang Hu, Chongyang Tao, Jing Ma, Qingwei Lin, and Daxin Jiang. Wizardcoder: Empowering code large language models with evol-instruct. arXiv preprint arXiv:2306.0850 2023.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wier- stra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint arXiv:1312.5602, 2013.
Matej Moravčík, Martin Schmid, Neil Burch, Viliam Lisy, Dustin Morrill, Nolan Bard, Trevor Davis, Kevin Waugh, Michael Johanson, and Michael Bowling. Deepstack: Expert-level artificial intelligence in heads-up no-limit poker. Science, 356(6337):508-513 2017.
Ofir Nachum, Mohammad Norouzi, Kelvin Xu, and Dale Schuurmans. Bridging the gap between value and policy based reinforcement learning. Advances in neural information processing sys- tems, 30, 2017.
Ashvin V Nair, Vitchyr Pong, Murtaza Dalal, Shikhar Bahl, Steven Lin, and Sergey Levine. Visual reinforcement learning with imagined goals. Advances in neural information processing systems, 31, 2018
Andrew Y Ng. Stuart Russell, et al. Algorithms for inverse reinforcement learning. In Icml, vol- ume 1, pp. 2. 2000.
Erik Nijkamp, Bo Pang, Hiroaki Hayashi, Lifu Tu, Huan Wang, Yingbo Zhou Silvio Savarese, and Caiming Xiong. Codegen: An open large language model for code with multi-turn program synthesis. In The Eleventh International Conference on Learning Representations, 2022.
R OpenAl. Gpt-4 technical report. arXiv, pp. 2303-08774, 2023.
Peter-Michael Osera and Steve Zdancewic. SIGPLAN Notices, 50(6):619-630 2015.
program synthesis. ACM
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems, 35: 27730-27744, 2022.
Kishore Papineni, Salim Roukos, Todd Ward, and Wei-Jing Zhu. Bleu: method for automatic evaluation of machine translation. In Proceedings of the 40th annual meeting of the Association for Computational Linguistics, pp. 311-318 2002.
Maxim Rabinovich, Mitchell Stern, and Dan Klein. Abstract syntax networks for code generation and semantic parsing. arXiv preprint arXiv: 1704.07535 2017.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog. 1(8):9, 2019.
Jack W Rae. Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, et al. Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint 12. 2021.
13

---

Published as a conference paper at ICLR 2024
Rafael Rafailov, Archit Sharma, Eric Mitchell, Stefano Ermon, Christopher D Manning, and Chelsea Finn. Direct preference optimization: Your language model is secretly a reward model. arXiv preprint arXiv:2305.18229, 2023.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu, Exploring the limits of transfer learning with a unified text-to-text transformer. The Journal of Machine Learning Research, 21(1):5485-551, 2020.
Marc' Aurelio Ranzato, Sumit Chopra, Michael Auli, and Wojciech Zaremba. Sequence level train- ing with recurrent neural networks. arXiv preprint arXiv:1511.6673 2015.
Steven Rennie, Etienne Marcheret, Youssef Mroueh, Jerret Ross, and Vaibhava Goel. Self-critical sequence training for image captioning. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 7008-7024. 2017.
Baptiste Rozière, Jonas Gehring, Fabian Gloeckle, Sten Sootla, Itai Gat, Xiaoqing Ellen Tan, Yossi Adi, Jingyu Liu, Tal Remez, Jérémy Rapin, al. Code Ilama: Open foundation models for code. arXiv preprint arXiv:2308. /et.3008.290000 /2900, 2023.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.0634, 2017.
Parshin Shojaee, Aneesh Jain, Sindhu Tipimeni, and Chandan K Reddy. Execution-basee code generation using deep reinforcement learning. arXiv preprint arXiv:2301.13811 2023.
David Silver. Lecture 7: Policy gradient. UCL Course on RL, 2015.
David Silver, Aja Huang, Chris Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Pannecrshelvam, Mare Lanctot, et al. Mastering the game of go with deep neural networks and tree search. nature, 529(7587):484 489, 2016.
Riley Simmons-Edie Edler, Anders Miltner, and Sebastian Seung. Program synthesis through reinforce- ment learning guided tree search. arXiv preprint (xXXix/800.029900 2018.
Nisan Stienmon, Long Jeffrey Wu, Daniel Ziegler, Ryan Lowe, Chelsea Voss, Alec Radford, Dario Amodei, and Paul F Christiano. Learning summarize with human feedback. Advances in Neural Information Processing Systems, 33:3008-3021, 2020.
Phillip D Summers. A methodology for lisp program construction from examples. Journal of the ACM (JACM), 24(1):161-175, 1977.
Arash Tavakoli, Fabio Pardo, and Petar Kormushev. Action branching architectures for deep rein- forcement learning. In Proceedings of the aaai conference on artificial intelligence, volume 32, 2018.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprini arXiv: 13971, 2023
Hado Van Hasselt, Arthur Guez, and David Silver. Deep reinforcement learning with double q- learning. In Proceedings of the AAAI conference artificial intelligence, volume 30, 2016.
Ben Wang and Aran Komatsuzaki. GPT-J-6B: Billion Parameter Autoregressive Language jax, May 2021.
Model. :tttpp://giiiee
com/ wwomm/ingofoooio z/
Xin Wang, Yasheng Wang, Yao Wan, Fei Mi, Yitong Li, Pingyi Zhou, Jin Liu, Hao Wu, Xin Jiang. and Qun Liu. Compilable neural code generation with compiler feedback. In Findings of the Association for Computational Linguistics: ACL 2022, pp. 9-19, 2022.
Yue Wang, Weishi Wang, Shafiq Joty, and Steven CH Hoi. Codet5: Identifier- unified pre- trained encoder-decoder models for code understanding and generation. Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pp. 8696-8708, 2021.
14

---

Published as a conference paper at ICLR 2024
Yue Wang, Hung Le, Akhilesh Deepak Gotmare, Nghi DQ Bui, Junnan Li, and Steven CH Hoi. Codet5+: Open code large language models for code understanding and generation. arXiv preprint arXhv:2300.07922, 2023.
Ziyu Wang, Tom Schaul, Matteo Hessel, Hado Hasselt, Mare Lanctot, and Nando Freitas. Dueling network architectures for deep reinforcement learning. In International conference on machine learning, pp. 1995-2003, PMLR, 2016.
Christopher JCH Watkins and Peter Dayan. Q-learning. Machine learning, 8:279-292, 1992.
Ronald J Williams. Simple statistical moadient-following algorithms for connectionist reinforcement learning. Machine learning, 8:229 2-229-256, 1992.
Tengyang Xie, Ching-An Cheng, Nan Jiang, Paul Mineiro, and Alekh Agarwal. Bellman-consistent pessimism for offline reinforcement learning. Advances in neural information processing systems, 34:6683-6694, 2021a.
Tengyang Xie, Nan Jiang, Huan Wang, Caiming Xiong, and Yu Bai. Policy finetuning: Bridg- ing sample-efficient offline and online reinforcement learning. Advances in neural information processing systems, 34:27395-27407, 2021b.
Can Xu, Qingfeng Sun, Kai Zheng, Xiubo Geng, Pu Zhao, Jiazhan Feng, Chongyang Tao, and Daxin Jiang. Wizardlm: Empowering large language models to follow complex instructions. arXiv preprint arXiv:2304. 12244, 2023.
Michihiro Yasunaga and Percy Liang. Graph-based, self-supervised program repair from diagnostic feedback. International Conference on Machine Learning, pp. 10799-10808 PMLR, 2020.
Zishun Yu and Xinhua Zhang. Actor-critic alignment for offline-to-online reinforcement learning. In Proceedings of the 40th International Conference Machine Learning, volume 202, pp. 40452- 40474, 2023.
Wenhao Zhan, Baihe Huang, Audrey Huang, Nan Jiang, and Jason Lee. Offline reinforcement learning with realizability and single-policy concentrability. In Conference Learning Theory, pp. 2730-2775. PMLR, 2022.
Qinkai Zheng, Xiao Xia, Xu Zou, Yuxiao Dong, Shan Wang, Yufci Xue. Zihan Wang, Lei Shen, Andi Wang, Yang Li, et al. Codegeex: A pre-trained model for code generation with multilingual evaluations on humaneval-x. arXiv preprint arXiv:2303. 17566, 2023.
Victor Zhong, Caiming Xiong, and Richard Socher. Seq2sql: Generating structured queries from natural language using reinforcement learning. arXiv preprint arXiv: 170.:01009000033 2017.
Brian Ziebart, Andrew L Maas, J Andrew Bagnell, Anind K Dey, et al. Maximum entropy inverse reinforcement learning. In Aaai, volume 8, pp. 1433-1438. Chicago, II., USA, 2008
Daniel Ziegler, Nisan Stiennon, Jeffrey Wu, Tom Brown, Alec Radford, Dario Amodei, Paul Christiano, and Geoffrey Irving. Fine-tuning language models from human preferences. arXiv preprint arXiv: wtt.//9999..9........ 2019.
Martin Zinkevich, Michael Johanson, Michael Bowling, and Carmelo Piccione. Regret minimization in games with incomplete information. Advances neural information processing systems, 20, 2007.
Amit Zohar and Lior Wolf. Automatic program synthesis of long programs with a learned garbage collector. Advances in neural information processing systems, 31, 2018.
15

---

Published as a conference paper at ICLR 2024
A PSEUDO-CODE FOR TRAINING
Algorithm Training Procedure with Φ- and O-stages
Require: 0ag o, and with shared frozen encoder
1: # pre-training stage, update & only
2: procedure PRETRAINUVVLLUE())
b d-stage
3:
for num_iters do
4:
Draw sample (s, a, r, 8') from dataset
5:
Compute logits a (x
6:
Compute state value (s(s)
7:
Compute loss Lv Lv(VV: klam))
> arguments omitted for brevity equation 12
8:
Gradient step with V VaCv( (oCv((en Voi (esp))
9: end for
end procedure
11: # fine-tuning stage, update 0 only
12:
procedure FINETUNEQVALUEGO
0-stage
13:
for num_iters do
14:
Draw sample 0, (,, from dataset
15:
Compute residual state-value Vi (s)
16:
Compute pre-trained state-value (g(8)
17:
Compute state-value (s)s) (s((s) Vs(s)
18:
Compute advantage Ag(s,) = (8(
max, (o(s,
Compute (s, aAp(s, Vo(s)
equation 16
20:
Compute To (-(s) = -sot=saof((( (Qa
Compute P0ak (-(s)
equation 14
22:
Compute Co( (Qe:
equation 15
23:
Compute Lee (Te) and Late AAa)
equation and
24:
LQ(Q0; po ) Bochce (TO) Compute fine-tune loss Lt Qo; poam)
Cadr
Gradient step with FoLn(Qo: p0
end for
27: end procedure
B PSEUDO-CODE FOR SAMPLING
Algorithm Sampling Procedure
Require: model parameters 0, o: SAMPLERp.(()) R > V that maps logits vector to token
with parametes (top-p sampling) and temperature
1:
2: procedure SAMPLEONFTTOKENNS
3: Obtain current state
4:
Compute logits vector (s(s) E p V
5:
Compute advantage vector As(s) Lg(s) max, La(s)[a)
6:
Compute Vo(s) = (s) + (s(s)
7:
Compute Q vector Qo(s) Ag(ss Vo(s)
8: Run SAMPLER,. (Qo(s)/a)
> sample with Qs(s)/a
9: end procedure
16

---

Published as a conference paper at ICLR 2024
C ADDITIONAL RELATED WORKS
Off-policy policy-based methods. One string of off-policy policy-based methods is based on im- portance ratio. Suppose the data is collected by a behavior policy 3. PG with off-policy data can be corrected by (H) = Es " (ael))) ent even though the data distribution is off-policy. However, computing the ratio Ty ( is not always feasible as the density function of off-policy data, such as human data, is often unknown. In addition, this correction can lead to high variance due to the product of ratios along trajectories.
[(a))s.) (E/= "r,,, log (ag(s))))) This allows unbiased gradi-
While vanilla importance- importance-weighted off-policy PG does not require the approximation of value func- tions, some advanced ratio-based methods often incorporate value functions, such (Imani et al., 2018; Liu et al., 2020). Another viable approach is the direct combination of value-based and policy- based methods, often referred to the actor-critic framework, e.g. (Konda & Tsitsiklis, 1999; De- gris et al., 2012), Although actor-critic methods are often conisdered as the third category, besides policy-based and value-based, we and some other works (Fujimoto al., 2018) lean towards cate- gorizing actor-critic to be more value-based, as the major difficulty lies in value function approxi- mations. Nevertheless, both directions of extending policy-based methods to an off off-policy setting, largely rely on the value functions. This emphasizes the motivation and significance of our work.
Reward modeling and beyond. Due to the successes of reinforcement learning from human/AI feedback (Christiano et al., 2017; Bai et al., 2022b). Reward modeling and RI fine-tuning with learned reward model has been a popular choice for post-SFT (supervised fine-tuning) refine ment (see e.g. Ziegler et al., 2019; Stiennon et al., 2020; Bai et al., 2022a; Ouyang et al., 2022). In particular, in program synthesis, Le et al. (2022) trains a classifier, that predicts unit test outcomes, as their reward model for RL fine-tuning. However, reward models train and their quality can heavily impact RL fine-tuning performance. Recent works (e.g. Rafailov et al., 2023; Diao et al., 2023) explore preference learning beyond conventional reward model.
sometimes be expensive to
Modeling reward function, on the other hand, has been a long- lant-nnsing topic in inverse RL or imi- tation learning (IRL or II., see e.g. Ng al., 2000; Abbeel & Ng, 2004; Ziebart et al., 2008; Ho & Ermon, 2016). While conventional IRI/ILL often iterates between reward model fitting and RL. training stages, recent IL works (Jacq et al., 2019; Garg et al., 2021) also explore beyond explicitly reward modeling to reduce training instability and optimization difficulty, led by the iterative opti- mization scheme. Specifically,, Garg et al. (2021) leverages the one-to-one correspondence between Q- Q-function and reward model, given the soft Bellman operator, to eliminate the reward fitting step.
Candidate selection in program synthesis. Existing works have shown one could improve program pass rate by filtering out programs that are likely to be incorrect. For instance, Chen et al. (2021a) filtered out programs that cannot pass example unit tests given in doc-strings, and Chen et al. (2022) filtered out programs that cannot pass generated unit tests. Furthermore, reward models are also often used to rank candidate programs (see e.g. Gulcehre et al., 2023; Touvron et al., 2023).
D A SPECTRUM OF RL APPLICATIONS
To conceptually demonstrate the differences between policy-based and value-based methods, and why program synthesis might be well-suited to value-based approaches, Figure 4 presents a spec- trum of RL applications. It could be observed that in scenarios where rewards are not expensive to evaluate or there's 's plenty of off-policy data (data not generated by the current policy/model) value-based methods tend to be preferred. Consider, for instance, InstructGPT (Ouyang et al., 2022) (pasicy-based) and AlphaGoo (Silver et al., 2016) (value-based). The former relies on human anno- tators (expensive) to label model-generated (on-policy) responses, while the latter obtains rewards from simulators (cheap), and leverages (1) human expert games (off-policy during training and (2) re-using historical games (off-policy) through experience reply.
Table 5 provides explanations our application plot of Figure 4. Applications in games typically find it easy to obtain rewards and make extensive use of off-policy data, e.g human games or historical replays. Conversely, InstructGPT obtains its rewards from preferences labeled by human annotators, with the data predominantly generated by the GPT model itself. The self-driving application notable has high cost of gathering rewards, due to the risks of real-world driving. While existing driving data
17

---

Published as a conference paper at ICLR 2024
Atari
extensive
Amount of off-policy data
Poker
StarCraft
InstructGPT
Image
value- based
captioning
Autonomous
driving
Code
limited
policy-
based
cheap
expensive
Costs to get rewards
Figure 4: A collection of RL applications. and represents value-based and policy-based RL, respectively. The x-axis shows the difficulty of obtaining rewards, while the y-axis measures the amount of off-policy data. Tasks that face significant hurdles in gathering rewards or have limited off-policy data typically lean towards policy- algorithms. Tasks where rewards are more readily obtained or that benefit from substantial collection of off-policy data favors value- based methods. See descriptions of each task Table
could be utilized, Kendall et al. (2019) specifically choose not to use pre-collected data, leading to their choice of policy-based algorithm.
In code generation, despite the availability of cheap rewards and the existing collection of off-policy programs, whether human writttn or historical synthetic programs, current literature leans towards policy-based methods. We believe that value-based methods could be promising direction, given their similarity to tasks with simulators.
Table 5: Summary of RL applications.
References
"Tpe of RL
Costs of Getting Rewards
Available Off-Policy Data
Atari
(Maih et al., 2013)
value
cheap: simulator
extensive: history/human games
(Silver al., 2016)
value
cheap: simulator
extensive: historyhumann games
(Moravčik al., 2017)
(Brown & Sandholm, 2018)
value
cheap:
extensive: history/human games
StarCraft
(Arulkumaran al., 2019)
cheap: simulator
extensive: history/human games
InstractGPT
(Ouyang al., 2022)
policy
expensive: human annotators
limited: mostly model-generated data
Image
Caption
(Ranzato al., 2015)
(Rennie al., 2017)
policy
cheap: automatic metrics
limited: mostly model generated data
Self-driving
(Kendall et al., 2019)
policy
expensive: driving in real-world
limited: mostly model--gnerrtedd data
(Ie al., 2022)
Code
Generation
(Shojace et al., 2023)
policy
cheap: unit lesting
extensive: collection of human programs
(Liu al., 2023)
E REWARD ENGINEERING COMPARISON
Table shows that ours has the least reward engineering effort. Note that our reward model Fg is directly derived from Qo, and is not used for training.
Table 7 shows the results when only basic reward function (defined in equation 2) used, under no example test outcomes setting. CodeRL. and RLTF results are duplicated from their reports.
'While Poker Al often uses counterfactual regret minimization (Zinkevich et al., 2007), which isn't strictly reinforcement learning, the shared principle of estimating action values allows us to categorize it under value- based methods.
18

---

Published as a conference paper at ICLR 2024
Table 6: Comparison of reward designs
Reward
Remark
Ours
CodeRL
RLTF
PPOCoder
Basic
equation 2
Reward Model
learned reward model
Fine-Grained
fine-gred grained error type & location of error ratio of passed tests
Adaptive
Syntactic Correctness
compilable
Synlactic Matching
syntactic similarity to ground truth
Semantic Matching
semantic similarity ground truth
Table 7: Performance with only basic reward (equation 2). and #t indicates results duplicated from Le et al. (2022) and Liu et al. (2023), respectively.
Pass@1 1
Pass@5
Model
Intro
Inter
Comp
All
Intro
Inter
Comp
All
CodeRL RLTFN
4.60
1.10
0.20
1.62
7.10
1.57
0.40
2.44
1.37
3.50
B-Coder
6.70
1.50
0.30
2.30
10.40
2.63
0.70
3.80
F ADVANTAGE OF APPROXIMATE VERSION OF 7
Recap that our recovered reward F is computed by
ca(s, a) = Qo(s, a) -E Qo (s', arg max, peaa (als')) Qo(s,a) - Qo(s,a)-7Vo($)) VVo(''.
(20)
Imagining scenario in which we sample/decode using trained Qe, the forward pass will compute a)) and (o(s) for each timestep, because of our dueling architecture. But Poag will not be evaluated during generation, because Paam is only used when computing Lo( :Pap)) Computing the exact version Qa a) >I,, Qo(s', arg max,, Palage (als') will require additional computation of pa during generation. In contrast, Q(s, a) and Vss) are already computed during generation, therefore it requires almost additional computation compute Fo(s,a).
G ADDITIONAL IMPLEMENTATION TRICKS
G.1 UPPER BOUND OF Q-FUNCTION
Given our reward design in equation 2, the cumulative reward is upper bounded by Rmax = 1. We enforce Q(s, a) < Rmax by transform the state value function V(s) SOFTABS (V(s)) + Rmax < Rmaxx where SOFTABS(2)) == function. Given A(s,o) a) < enforcing V(s) Rmax leads Q(8 a) S Rmax-
+
In 2 is soft absolute
G.2 RESIDUAL HEAD INITIALIZATION
In section 4.3, initialize 0 in way such that to Ea and Vg(s) The former be done by simply loading the checkpoint Oage Adding residual head Vy, that initialized output zeros, can be done with a simple trick. One can simply add two heads h and h22 let hy be trainable, and h2 be fixed for subsequent fine-tuning, setting ViJ h h2 achieves the desired functionality.
H TRAINING AND EVALUATION DETAILS
In supplement to implementation details in Section 4.4 and 5, give more low-level details here.
APPS dataset. addition to the train/test split details described in Section 5, APPS datast, on average, consists of 2 example unit tests, 21 hidden unit tests, and 23 ground truth programs. We
19

---

Published as a conference paper at ICLR 2024
follow the same procudure Hendrycks et al. (2021); Le et al. (2022) to construct prompts for both training and evaluation. Specifically, see Section of Hendrycks et al. (2021).
MBPP dataset. MBPP has 974 instances with a 374/90/500 train/val/test splits and, in addition, 10 problems reserved for few-shot learning. Because we only do zero-shot evaluation on MBPP, only the 500 test problems tests. In addition, these tests are usually not hidden. Therefore, prior works Le et al. (2022); Shojaee al. (2023); Liu et al. (2023) often explicitly incorporate the tests into prompt string. We follow WizardCoder (Luo et al., 2023) to construct our input format. Details could be found in this repo.
used for evaluation. Each problem of MBPP usually comes with three unit
Pre-trained model. We initialize our model with CodeRL. checkpoint publicty available at here, meaning we initialize aag 6, and 0 from it. Note that we freeze encoder for both ø-stage and 0-stage, therefore the encoder is shared during both training and generation. For both training and generation, we set the maximum length to 600 and 512 for source and target sequences, respectively.
Training data preparation. While we use Drai to represent our training dataset, yet we have not elaborated on how it is constructed. In general, we follow the protocol of prior RL-based works that combining all ground truth programs and a set of programs generated by the pre-trained model, for each problem D. Specifically, we generate 256 programs per problem using pre-trained checkpoint. Combined with ground truth programs, there are, on average, 278 programs per problem.
Mini-batch preparation. By prior definition, our dataset Drain now contains both ground truth programs and generated programs. Notably, the volume of generated programs significantly larger than that of the ground truth programs. This means that if one were to randomly sample from the dataset, generated programs would dominate the mini- batches. To address this, when preparing a mini-batch we sample Preat B ground truth programs and (1 Preat) x B generated programs, where B is batch size.
o--tage training. In the d-stage, we pre-train state-value function (s(s). We conduct our exper- iment with 4xA100-80G GPUs. Specifically, we use batch size of 16 for each GPU and gradient accumulation step of 4, resulting in total batch size of 256. For optimizer and scheduler, we use Adam W optimizer (Loshchilov & Hutter, 2018) with a constant learning rate of le-5 and a weight decay of 0.05. We train @ for 18k gradient steps.
0-stage training. In the 0-stage, we conduct our experiment with 8XA100-80G GPUs. Specificaly we use batch size of 16 for each GPU and gradient accumulation step of resulting total batch size of 128. For optimizer and scheduler, we use AdamW with a peak learning rate 3e-5, a weight decay of 0.05, and linear decay scheduler with no warmup. We train @ for 10k gradient steps.
Other hyper-parametery We set the ground truth data ratio Preat = 0.5 and the energy-based policy temperature (see equation 10) for all experiments. In 0-stage, we use Badv 0..1 and 3.e 0.5
I ABLATION ON m
comp
1.5
1.-
6.2
02
62
13
3.0
6.3
3.0
16.0
0.3
27 24
23
Figure 5: Ablation on our ranking strategy achieves consistent improvements under different
budgets m.
Table 5 conduct an ablation study on ranking budgets m, it can be observed that our ranking strategy achieves consistent improvements under different budgets m.
20

---

Published as a conference paper at ICLR 2024
J COMMENTS ON B PROPERTIES
J.1 PROPOSITION 4.1
Proof.
1B9Q1 19Q2|lo max/r((ss)) +1 r(s.a) -E, Q2(s'
5,a
(ā' = arg max, {(als)))
max R 1Q1(8' (2(6',0))
(21)
5,4
max
Eley 1Q a') Q2(),a')
(22)
5,4
< max yl max Q1(s', d') Q2(5', a')l
(23)
8,4
8'aa
7||Q Q2
(24)
L
J.2 PROPOSITION 4.2
Proof. The proof is similar to Lemma C.3. in Garg et al. (2021). To prove that TP is bijection, such that
suffices to show that for any r A -> I. there exists unique Q S x TPQ. Note that by proposition 4.1, there exists unique Q" Brr that satisfies Q" r(s, a) + y. QP (s', arg max, p(als')). Rearranging the terms gives r = TPQP. This completes the proof.
K DISCUSSION ON LIMITATIONS
Table 8: @ass 1 results are evalu- ated with greedy decoded programs, and pass @ (0,50, 100} are computed by sam- pled programs using temperature of 0.4.
Table Ranking with compared with filtering with real environmental reward function т, i.c. hidden tests. f-ranked results are duplicated from Table
F-ranked
@ass@
CodeRL
B-Coder
r-filtered
pass
pass@5
1
1.60
1.60
3.28
2.88
Intro
6.70
10.40
26.60
50
7.16
7.35
Inter
1.50
2.63
7.87
100
8.76
9.18
Comp
0.30
0.70
5.10
All
2.30
3.80
11.06
While being exploratory, our work admits certain limitations including: additional frozen parameters introduced, and we observe that raw performance (without ranking) is mixed compared to CodeRL (see Table 8) (which we believe is somewhat excusable as we use less reward designs). However, we remark the effectiveness of our overall framework including the dual strategy is non-trivial, especially with limited reward engineering.
It is also informative to show results filtered by the true environmental reward function instead of results ranked by our recovered reward function F. Although filtering with r requires using hidden tests, meaning it cannot be implemented in realistic settings, also discussions in Section 4.5. However, could serve as an upper limit for our ranking strategy and as a sanity check. (Roughly speaking, if f=r, the pass rate of f-ranking and r-fittering would be identical.) To this end, use the same set of candidate programs as those in Table 1, but apply the ground truth reward function filter candidates rather than using F for ranking. The corresponding results in Table show that, although T-ranking is effective, there remains a large room for improvement.
21


Reviews:
[{"review_id": "cD3oi8YwqH", "reviewer": "Reviewer_dB3F", "summary": "The task is to generate code (sequence generation) given a natural language description, and to have that code pass the associated unit tests.\n\nThis is framed as an RL problem, and unlike most prior work it does not use an algorithm from the policy gradient family, but uses a sort of Q learning approach.\n\nThe main focus is on how to overcome the very long and large action space you get by treating each token as an action, and the sparse reward at the end. The main idea is to make good use of pre-trained language models that have been fine-tuned to generate code for given ground truth (description, code solution) pairs. \n\nUsing this is non-trivial however, and the paper takes a creative an impressive approach. \n\nThe main ideas are known but very nicely combined. Some highlights:\n- the pretrained model is used to initialise a sensible Q function before the RL training\n- decomposing Q(s,a) as state value + state-action advantage (dueling dqn)\n- conservative bellman updates to stabilise the rl training loop.\n\nA drawback of the approach is that it takes multiple steps, uses various tricky pre-training and residual fitting steps, and could be difficult to apply in practice. But this is not suprising since using Q learning here is challenging, and it is impressive the authors managed to make it work so well.\n\nDetailed comments:\n\n2.1 Policy: should say \"assigns an action *from the set* $\\Delta(\\mathcal A)$\" and generally this notation here needs checking.\n\nIn figure 2, why is zero training iterations the best? anyway this is a nice ablation / teaser for the paper.\n\nFig 3 is very cool indeed.\n\nGeneral question - since you sample a softened Q as your policy, you could actually apply policy gradients as well. What are the connections here? Could you combine your scheme with a policy gradient type of update?", "questions": "See the summary.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "See the summary.", "weaknesses": "See the summary.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Wx9Om7UI3i", "reviewer": "Reviewer_R5vH", "summary": "this work proposes a value-function based training/inference procedure for program synthesis. rather than training a model (a policy) that ingest context (e.g. given as natural language comments) and generates tokens of code, this work leverages a value function Q(s,a) = V(s) + A(s,a), where V and A are implemented as neural networks.\n\nIn order to make Q learning work, this paper proposes two \"tricks\", one being a good Q-value initialization, the other being a conservative bellman update. Without these tricks, the training procedure would either not converge at all or improve very slowly. \n\nOverall, this paper insists that they want to tackle the task of program generation by training a value function, and overcome a lot of challenges to make it work.", "questions": "## time to take sample?\n\nThe \"go to\" solution of program synthesis in the presence of given input-output examples (test cases) is to simply generate a vast number of programs, test all of them against spec through execution, and pick the ones that passed the test cases.\n\nThus, suppose I am a practitioner, the metric that I really care about is, given a synthesis problem, complete with natural language descriptions and a set of input-output test cases, and a time budget of 1 minute, can my synthesizer find a correct program?\n\nThis allows different methods to be compared head to head. A \"dumb\" enumerator -- for instance, a bigram generator -- might actually be reasonable of it can generate vast numbers of programs quickly, and a \"smart\" enumerator might actually be non-performative due to the time it takes to sample a full program to even check.\n\nThus, I would like to ask for the times it take to sample a full program on average, when compared against some of the other policy based methods.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "## quality: good\nthe overall quality of the paper is good, in that it even worked at all, and the lack of adhoc tricks.\n\n### that it even worked at all\n\nThe biggest strength of this paper is that it has shown that code generation _can be_ done with a value network, despite the large number of state x action pairs needs to be considered, the unavailbility of a good initial Q function, and the volatility of the Q function \"blowing up\". Evaluation shown that the proposed method is on par with the policy based methods in terms of pass@k metrics.\n\nThe knowledge that value-based approach for code generation can be done _at all_ is a great merit, and cannot be under stated.\n\n### lack of ad-hoc \"tricks\"\n\nI find all the \"tricks\" presented in this work well founded and easy to understand. The reward hacking part is minimal and frankly very reasonable. Thus I am convinced this approach would work on different program synthesis domains just as well.", "weaknesses": "## significance\n\nlack of evidence on why we need to do this to begin with\n\nThe following statement from the conclusion section is a very good motivation \" it is recognized being sample-inefficient (see e.g. Nachum et al., 2017), meaning poorly using off-policy data, even data previously generated\". Does this work actually demonstrate this with an experiment? I believe it should, to show (in a table) that the proposed method can \"learn more with less\", with x axis showing the amount of training data, and y-axis some performance metric. In case I have misread the paper and this is indeed one of the experiments, it should be top and foremost contribution this paper needs to claim. Currently this paper reads akin to \"we did it because we can\" but lacks the justification -- in the form of an experiment -- that this is something valuable to begin with.\n\nI could think of many other benefits of having a value function, for instance, sometimes people might want a top-k ranked programs, which could be tricky to do based on only policy, but very intuitive with a value based model.\n\nIf the authors are able to justify their work with some experimental evidence, showing the advantages of a value-based modeling, this work will be that much better.\n\n## clarity\nthis paper would benefit from a Figure1 that outline the overall \"workflow\" of the method. During reading of the paper, the phrase \"which is yet another difficulty one must overcome to train a value function .... we introduce technique X to tackle this problem.\" The number of tricks quickly adds up, making the paper difficult to follow.\n\nI feel the paper is overboard with definition and philosophy of \"what is policy-based\" and \"what is value-based / RL\". I believe these things have a very technical definition that the community can all agree with. The clean story could've been \"off-policy learning through value iteration for generation would be good, as it use less data. here's how to do it, here's an experiment showing how well we can do as a function of amount of human data required\".\n\nDefining program synthesis as \"generating program from natural language\" is imprecise, as there are other works of \"program synthesis\" that generates program from input-output examples alone, or those that work with constraint solvers that generates programs from correctness specifications. A better wording could benefit here, as I instinctively reacted poorly by the opening definition of program synthesis, thinking it is way too narrow.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "v4xdti95hK", "reviewer": "Reviewer_F9ov", "summary": "This paper derives a value-based approach for fine-tuning LLM's for code generation, motivated by the success of value-based methods in domains with similar constraints on data and reward modeling. The authors build on top of the Dueling DQN family of algorithms and design architectures and pretraining schemes for the Q/A/V functions that carefully incorporate the pretrained LLM policy in a way that provides mathematical guarantees that certain properties are met, such as ranking human solutions highly, leading to a policy function that is identical to the original LLM policy, and minimizing initial Bellman error. They then use a conservative Bellman operator based off greedy decoding, and show that they can derive a reward model from an inverse conservative Bellman operator without any additional training needed. Their approach is called Bellman-Coder. Experiments show that Bellman-Coder performs comparably to related work, and is often the best. In addition, the reward model generalizes to new domains and models.", "questions": "In figure 2, why don't the approaches train for the same number of iterations? In the same vein, what do the dots signify on the dashed lines, and why do their distances change? \n\nis it true that the reward model coming for free in section 4 is possible for any Bellman operator? what are the preconditions necessary for this? I ask because the \"further comments\" right before Proposition 3.1 made me think that the conservativeness of the Bellman operator defined was needed, but it sounds like in section 4 that it works for any Bellman operator. (I do not have much expertise in this area)\n\nThe related work paragraph in section 1 only discusses work on supervised LLM training and RL alternatives for sequence generation. given the amount of highly related work discussed in other sections of the introduction, it seems like this paragraph should be renamed to something more precise. \n\nAlphaGo is described has having lots of off-policy data, but how does AlphaZero relate to the taxonomy of Figure 1? Isn't AlphaZero pretty much the same approach, but with all on-policy data?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- The paper is very well-written. The authors do an excellent job of discussing related work, motivating the design of a value-based approach for code generation, explaining the challenges behind designing a value-based approach, and describing how their design tackles these challenges. The technical sections are presented with care to make sure readers don't get lost in symbols, and I found them easy to understand (at least at a medium-to-high level) despite not being a RL expert. The experiments are presented well. Everything about the writing and clarity is excellent.\n- The paper's approach is original and very high quality. It does appear that no prior work has leveraged value-based methods for code generation to the degree of success the authors present. Making progress on value-based methods for code generation seems like a significant research direction to pursue, given the importance of code generation as a primary application area of AI and the promise of value-based methods in other areas (as motivated by the authors). The only uncertainty of the significance might be due to the limited performance gains, but this is easily understandable given this paper seems to be the first to attempt it, and the uphill battle of designing such a method given the comparable ease of using a LLM as an off-the-shelf policy for policy-based methods. The authors took the harder route, and seem to have managed to get it to work comparably.\n- I'd like to emphasize how high quality the mathematical derivations seem. The authors seem to have a very valuable expertise in RL and have put it to good use in making this method work. The ablations in figure 2 attest to the effectiveness of their formulation, with some caveats.", "weaknesses": "The primary weakness of the paper is that the results are not much better than prior work. However, given how difficult the perceived approach seems to be, I think this is okay. \n\nI am not very up-to-date with the different techniques and evaluation metrics for RL code generation, so I can't comment too much on whether the evaluation is missing anything. But taking a quick look at related work, the authors seem to use the same benchmarks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qLd4lwPOZX", "reviewer": "Reviewer_3d3T", "summary": "In this paper, the authors introduce a value-based method for code generation with pretrained LLMs. In their opinion, their method fills a vacuum, due to the fact that all present LLM-based code-generation systems are policy-based in order to easily leverage pre-trained LLMs. The authors get around this difficulty via proposing initialisation and semi-gradient policy-improvement schemes in which the LLMs logits are visualised as Q-values. In addition to this, they propose a conservative Bellman operator in which the \"best action\" is not chosen based on the trained Q-values, but on the logits of the initial pre-trained LLM, which in their view stabilises training and alleviates the usual convergence problems of value-based RL.\nThe authors test their framework on a standard code-generation benchmark (the APPS benchmark) and use a relatively straightforward reward design. They report performance superior to, or competitive with, with an exhaustive set of baselines.", "questions": "Main question and concern: why do the authors use a filtering procedure at evaluation time based on a different reward model $\\tilde{r}(s,a)$ than the environmental reward $r(s,a)$ they use for training? Given that the real reward $r$ can be easily and cheaply computed (a point which the authors themselves make in the introduction), I cannot see the need for such a procedure, and the authors do not elaborate on it. The authors should either:\n - Justify this choice.\n - Recompute their performance numbers with filtering based on $r$ at evaluation time.\n - Re-train their method with the modelled reward $\\tilde{r}$ from the outset.\n\nOther concerns/questions/recommendations:\n- Could the mixed performance numbers be due to the conservative Bellman operator, which keeps the policy close to the initial checkpoint? Did the authors consider alternative definitions for $q(a|s)$, or to update it at some point during training?\n - Value-based methods, including DQN, can suffer from the so-called \"Deadly Triad\" of deep RL. Can the authors elaborate on how they avoid this pitfall?\n- It seems that the authors add a value head $V^r_\\theta(s)$ (besides its LM head) to their LLM. How is that done, and can it be done with a generic LLM?\n- In figure 2, it seems that the pass@1 metric is actually at its best at very beginning of training and that it never recovers afterwards. Is the figure only meant to be illustrative, or are those the real training curves?\n- Can the authors outline their paper's content and structure at the end of the intro?\n- In figure 3, it would be helpful to explicitly define what $\\ell$ and $p$ stand for.\n- Why is the Nucleus Sampling outlined at the end of section 3? In my opinion it would make more sense to have in section 5, since it appears to be part of the evaluation protocol. Or is it also used during training?\n- It would be helpful if the authors defined (e.g. in a caption) what \"Intro\", \"Inter\", \"Comp\", and \"All\", mean in tables 1 and 2. While the text does mention it, it would be helpful to readers to be more explicit.\n- Why do the authors not consider $k=1000$ in table 2, as they do in table 1?\n- The study on generalisation of the reward model (notwithstanding my concerns on its use, detailed above) would better be moved to an appendix.\n\n# Post rebuttal edit:\nThe authors have exhaustively addressed my concerns and revised their paper as I suggested. I shall therefore raise my score to 8.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The authors are correct in that value-based methods for code-generation were absent from the literature, and that the gap should be filled.\n- Their solution for doing so is quite novel and elegant, especially the idea of visualising the LLM's logits as Q-values, and use them for both initialisation and policy improvement.\n- They consider an exhaustive set of baselines.\n- They show performance superior to, or competitive with, baselines, and on a programming benchmark notorious for its high difficulty.", "weaknesses": "- While the framework is novel, its comparison with baselines gives mixed results: in some cases, it does not actually come out on top. This contradicts the authors' claim that value-based method should be more suited to code generation: all of the baselines are policy-based.\n- The authors filter programs at evaluation time using a different reward than the one used during training. This makes it hard to interpret their performance numbers. Could it be the filtering that is mostly responsible for them?\n- The authors' method needs to be applied to a pre-trained LLM, which is used not only for initialisation but also for training. This is only possible with LLMs whose checkpoints are publicly available. In the case of the largest LLMs available today, even with a checkpoint available, considerable computational resources would be needed to apply the method.\n- The paper has a bit of an idiosyncratic structure, with e.g. a whole section dedicated only to the reward filtering at evaluation time, a very small Related Work section, and no outline of the paper's contents in the introduction, which makes it a bit hard to read.\n- Some more proof-reading is needed as typos can be found here and there.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (spotlight)