# Scalable Bayesian Inverse Reinforcement Learning

Alex J. Chan & Mihaela van der Schaar
Department of Applied Mathematics and Theoretical Physics
University of Cambridge
Cambridge, UK
<EMAIL_ADDRESS>

###### Abstract

Bayesian inference over the reward presents an ideal solution to the ill-posed
nature of the inverse reinforcement learning problem. Unfortunately current
methods generally do not scale well beyond the small tabular setting due to
the need for an inner-loop MDP solver, and even non-Bayesian methods that do
themselves scale often require extensive interaction with the environment to
perform well, being inappropriate for high stakes or costly applications such
as healthcare. In this paper we introduce our method, Approximate Variational
Reward Imitation Learning (AVRIL), that addresses both of these issues by
jointly learning an approximate posterior distribution over the reward that
scales to arbitrarily complicated state spaces alongside an appropriate policy
in a completely offline manner through a variational approach to said latent
reward. Applying our method to real medical data alongside classic control
simulations, we demonstrate Bayesian reward inference in environments beyond
the scope of current methods, as well as task performance competitive with
focused offline imitation learning algorithms.

##  1 Introduction

For applications in complicated and high-stakes environments it can often mean
operating in the minimal possible setting - that is with no access to
knowledge of the environment dynamics nor intrinsic reward, nor even the
ability to interact and test policies. In this case learning and inference
must be done solely on the basis of logged trajectories from a competent
demonstrator showing only the states visited and the the action taken in each
case.

Clinical decision making is an important example of this, where there is great
interest in learning policies from medical professionals but is completely
impractical and unethical to deploy policies on patients mid-training.
Moreover this is an area where it is not only the policies, but also knowledge
of the demonstrator’s preferences and goals, that we are interested in. While
_imitation learning_ (IL) generally deals with the problem of producing
appropriate policies to match a demonstrator, with the added layer of
understanding motivations this would then usually be approached through
_inverse reinforcement learning_ (IRL). Here attempting to learn the assumed
underlying reward driving the demonstrator, before secondarily learning a
policy that is optimal with respect to the reward using some forward
_reinforcement learning_ (RL) technique. By composing the RL and IRL
procedures in order to perform IL we arrive at _apprenticeship learning_ (AL),
which introduces its own challenges, particularly in the offline setting.
Notably for any given set of demonstrations there are (infinitely) many
rewards for which the actions would be optimal (Ng et al., 2000). Max-margin
(Abbeel & Ng, 2004) and max-entropy (Ziebart et al., 2008) methods for
heuristically differentiating plausible rewards do so at the cost of
potentially dismissing the _true_ reward for not possessing desirable
qualities. On the other hand a _Bayesian_ approach to IRL (BIRL) is more
conceptually satisfying, taking a probabilistic view of the reward, we are
interested in the posterior distribution having seen the demonstrations
(Ramachandran & Amir, 2007), accounting for all possibilities. BIRL is not
without its own drawbacks though, as noted in Brown & Niekum (2019), making it
inappropriate for modern complicated environments: assuming linear rewards;
small, solvable environments; and repeated, inner-loop, calls to forward RL.

The main contribution then of this paper is a method for advancing BIRL beyond
these obstacles, allowing for approximate reward inference using an
arbitrarily flexible class of functions, in any environment, without costly
inner-loop operations, and importantly entirely offline. This leads to our
algorithm AVRIL, depicted in figure 1, which represents a framework for
jointly learning a variational posterior distribution over the reward
alongside an imitator policy in an _auto-encoder-esque_ manner. In what
follows we review the modern methods for offline IRL/IL (Section 2) with a
focus on the approach of _Bayesian_ IRL and the issues it faces when
confronted with challenging environments. We then address the above issues by
introducing our contributions (Section 3), and demonstrate the gains of our
algorithm in real medical data and simulated control environments, notably
that it is now possible to achieve Bayesian reward inference in such settings
(Section 4). Finally we wrap up with some concluding thoughts and directions
(Section 5). Code for AVRIL and our experiments is made available at
https://github.com/XanderJC/scalable-birl and
https://github.com/vanderschaarlab/mlforhealthlabpub.

Figure 1:  Overview. AVRIL is a framework for BIRL that works through an
approximation in the variational Bayesian framework, considering the reward to
be a latent representation of behaviour. A distribution over the reward, which
is amortised over the demonstration space, is learnt that then informs an
imitator $Q$-function policy. The dotted line represents a departure from a
traditional auto-encoder as the input, alongside the latent reward, informs
the decoder.

##  2 Approaching Apprenticeship and Imitation Offline

#### Preliminaries.

We consider the standard Markov decision process (MDP) environment, with
states $s\in\mathcal{S}$, actions $a\in\mathcal{A}$, transitions
$T\in\Delta(\mathcal{S})^{\mathcal{S}\times\mathcal{A}}$, rewards
$R\in\mathbb{R}^{\mathcal{S}\times\mathcal{A}}$111We define a state-action
reward here, as is usual in the literature. Extensions to a state-only reward
are simple, and indeed can be preferable, as we will see later., and discount
$\gamma\in[0,1]$. For a policy $\pi\in\Delta(\mathcal{A})^{\mathcal{S}}$ let
$\rho_{\pi}(s,a)=\mathbb{E}_{\pi,T}[\sum_{t=0}^{\infty}\gamma^{t}\mathbbm{1}_{\\{s_{t}=s,a_{t}=a\\}}]$
be the induced unique occupancy measure alongside the state-only occupancy
measure $\rho_{\pi}(s)=\sum_{a\in\mathcal{A}}\rho_{\pi}(s,a)$. Despite this
full environment model, the only information available to us is the
MDP$\backslash RT$, in that we have no access to either the underlying reward
or the transitions, with our lacking knowledge of the transitions being also
strong in the sense that further we are unable to simulate the environment to
sample them. The learning signal is then given by access to $m$-many
trajectories of some demonstrator assumed to be acting optimally w.r.t. the
MDP, following a policy $\pi_{D}$, making up a data set
$\mathcal{D}_{raw}=\\{(s_{1}^{(i)},a_{1}^{(i)},\ldots,s_{\tau^{(i)}}^{(i)},a_{\tau^{(i)}}^{(i)})\\}_{i=1}^{m}$
where $s_{t}^{(i)}$ is the state and $a_{t}^{(i)}$ is the action taken at step
$t$ during the $i$th demonstration, and $\tau^{(i)}$ is the (max) time horizon
of the $i$th demonstration. Given the Markov assumption though it is
sufficient and convenient to consider the demonstrations simply as a
collection of $n$-many state, action, next state, next action tuples such that
$\mathcal{D}=\\{(s_{i},a_{i},s^{\prime}_{i},a^{\prime}_{i})\\}_{i=1}^{n}$ with
$n=\sum_{i=1}^{m}(\tau^{(i)}-1)$.

#### Apprenticeship through rewards.

Typically AL proceeds by first inferring an appropriate reward function with
an IRL procedure (Ng et al., 2000; Ramachandran & Amir, 2007; Rothkopf &
Dimitrakakis, 2011; Ziebart et al., 2008) before running forward RL to obtain
an appropriate policy. This allows for easy mix-and-match procedures, swapping
in different standard RL and IRL methods depending on the situation. These
algorithms though depend on either knowledge of $T$ in order to solve exactly
or the ability to perform roll-outs in the environment, with little previous
work focusing on the entirely offline setting. One simple solution is through
attempting to learn the dynamics (Herman et al., 2016), though without a large
supply of diverse demonstrations or a small environment this becomes
impractical given imperfections in the model. Alternatively Klein et al.
(2011) and Lee et al. (2019) attempt off-policy feature matching through
least-squared temporal difference and deep neural networks to uncover
appropriate feature representations.

#### Implicit-reward policy learning.

Recent work has often forgone an explicit representation of the reward. Moving
within the maximum-entropy RL framework (Ziebart, 2010; Levine, 2018), Ho &
Ermon (2016) noted that the full procedure (RL $\circ$ IRL) can be interpreted
equivalently as the minimisation of some divergence between occupancy measures
of the imitator and demonstrator:

$\displaystyle\underset{\pi}{\operatorname*{arg\,min}}\\{\psi^{*}(\rho_{\pi}-\rho_{\pi_{D}})-H(\pi)\\},$
(1)

with $H(\pi)$ being the discounted causal entropy (Bloem & Bambos, 2014) of
the policy and $\psi^{*}$ the Fenchel conjugate of a chosen regulariser on the
form of the reward. These are typically optimised in an adversarial fashion
(Goodfellow et al., 2014) and given the focus on evaluating $\rho_{\pi}$ this
often requires extensive interaction with the environment, otherwise banking
on approximations over a replay buffer (Kostrikov et al., 2018) or a
reformulation of the divergence to allow for off-policy evaluation (Kostrikov
et al., 2019). Bear in mind that optimal policies within the maximum-entropy
framework are parameterised by a Boltzmann distribution:

$\displaystyle\pi(a|s)=\frac{\exp(Q(s,a))}{\sum_{b\in\mathcal{A}}\exp(Q(s,b))},$
(2)

with $Q(s,a)$ the _soft_ $Q$-function, defined recursively via the _soft_
Bellman-equation:

$\displaystyle Q(s,a)\triangleq
R(s,a)+\gamma\mathbb{E}_{s^{\prime}\sim\rho_{\pi}}\Big{[}\underset{a^{\prime}}{\text{soft
max }}Q(s^{\prime},a^{\prime}))\Big{]}.$ (3)

Then for a learnt parameterised policy given in terms of $Q$-values from a
function approximator $Q_{\theta}$, we can obtain an _implied_ reward given
by:

$\displaystyle
R_{Q_{\theta}}(s,a)=Q_{\theta}(s,a)-\gamma\mathbb{E}_{s^{\prime}\sim\rho_{\pi}}\Bigg{[}\log\Bigg{(}\sum_{a^{\prime}\in\mathcal{A}}\exp(Q_{\theta}(s^{\prime},a^{\prime}))\Bigg{)}\Bigg{]}.$
(4)

A number of algorithms make use of this fact with Piot et al. (2014) and Reddy
et al. (2019) working by essentially placing a sparsity prior on this implied
reward, encouraging it towards zero, and thus incorporating subsequent state
information. Alternatively Jarrett et al. (2020) show that even the simple
behavioural cloning (Bain & Sammut, 1995) is implicitly maximising some reward
with an approximation that the expectation over states is taken with respect
to the demonstrator, not the learnt policy. They then attempt to rectify part
of this approximation using the properties of the energy-based model implied
by the policy (Grathwohl et al., 2019).

The problem with learning an implicit reward in an offline setting is that it
remains just that, _implicit_ , only able to be evaluated at points seen in
the demonstrations, and even then only approximately. Thus even if their
consideration improves imitator policies performance they offer no real
improvement for interpretation.

###  2.1 Bayesian Inverse Reinforcement Learning

We are then resigned to directly reason about the underlying reward, bringing
us back to the question of IRL, and in particular BIRL for a principled
approach to reasoning under uncertainty. Given a prior over possible
functions, having seen some demonstrations, we calculate the posterior over
the function using a theoretically simple application of Bayes rule.
Ramachandran & Amir (2007) defines the likelihood of an action at a state as a
Boltzmann distribution with inverse temperature and respective state-action
values, yielding a probabilistic demonstrator policy given by:

$\displaystyle\pi_{D}(a|s,R)=\frac{\exp(\beta
Q_{R}^{\pi_{D}}(s,a))}{\sum_{b\in\mathcal{A}}\exp(\beta
Q_{R}^{\pi_{D}}(s,b))},$ (5)

where $\beta\in[0,\infty)$ represents the the confidence in the optimality of
the demonstrator. Note that despite similarities, moving forward we are no
longer within the maximum-entropy framework and $Q_{R}^{\pi}(s,a)$ now denotes
the traditional, not soft (as in equation 3), state-action value ($Q$-value)
function given a reward $R$ and policy $\pi$ such that
$Q_{R}^{\pi}(s,a)=\mathbb{E}_{\pi,T}[\sum_{t=0}^{\infty}\gamma^{t}R(s_{t})|s_{0}=s,a_{0}=a]$.
Unsurprisingly this yields an intractable posterior distribution leading to a
Markov chain Monte Carlo (MCMC) algorithm based on a random grid-walk to
sample from the posterior.

#### Issues in complex and unknown environments.

This original formulation, alongside extensions that consider maximum-a-
posteriori inference (Choi & Kim, 2011) and multiple rewards (Choi & Kim,
2012; Dimitrakakis & Rothkopf, 2011), suffer from three major drawbacks that
make them impractical for modern, complicated, and model-free task
environments.

  1. 1.

_The reward is a linear combination of state features._ Naturally this is a
very restrictive class of functions and assumes access to carefully hand-
crafted features of the state space.

  2. 2.

_The cardinality of the state-space is finite, $|\mathcal{S}|<\infty$._
Admittedly this can be relaxed in practical terms, although it does mean the
rapid-mixing bounds derived by Ramachandran & Amir (2007) do not hold at all
in the infinite case. For finite approximations they scale at
$\mathcal{O}(|\mathcal{S}|^{2})$, rapidly becoming vacuous and causing BIRL to
inherit the usual MCMC difficulties on assessing convergence and sequential
computation (Gamerman & Lopes, 2006).

  3. 3.

_The requirement of an inner-loop MDP solve._ Most importantly at every step a
new reward is sampled and the likelihood of the data must then be evaluated.
This requires calculating the $Q$-values of the policy with respect to the
reward, in other words running forward RL. While not an insurmountable problem
in the simple cases where everything is known and can be quickly solved with a
procedure guaranteed to converge correctly, this becomes an issue in the realm
where only deep function approximation works adequately (i.e. the non-tabular
setting). DQN training for example easily stretches into hours (Mnih et al.,
2013) and will have to be repeated thousands of times, making it completely
untenable.

We have seen that even in the most simple setting the problem of exact
Bayesian inference over the reward is intractable, and the above limitations
of the current MCMC methods are not trivial to overcome. Consequently very
little work has been done in the area and there still remain very open
challenges. Levine et al. (2011) addressed linearity through a Gaussian
process approach, allowing for a significantly more flexible and non-linear
representation though introducing issues of its own, namely the computational
complexity of inverting large matrices (Rasmussen, 2003). More recently Brown
& Niekum (2019) have presented the only current solution to the inner-loop
problem by introducing an alternative formulation of the likelihood, one based
on human recorded pairwise preferences over demonstrations that significantly
reduces the complexity of likelihood calculation but does necessitate that we
have these preferences available. This certainly can’t be assumed always
available and while very effective for the given task is not appropriate in
the general case. One of the key aspects of our contribution is that we deal
with all three of these issues while also not requiring any additional
information.

#### The usefulness of uncertainty.

On top of the philosophical consistency of Bayesian inference there are a
number of very good reasons for wanting a measure of uncertainty over any
uncovered reward that are not available from more traditional IRL algorithms.
First and foremost the (epistemic) uncertainty revealed by Bayesian inference
tells us a lot about what areas of the state-space we really cannot say
anything about because we haven’t seen any demonstrations there - potentially
informing future data collection if that is possible (Mindermann et al.,
2018). Additionally in the cases we are mostly concerned about (e.g. medicine)
we have to be very careful about letting algorithms pick actions in practice
and we are interested in performing _safe_ or _risk-averse_ imitation, for
which a degree of confidence over learnt rewards is necessary. Brown et al.
(2020) for example use a distribution over reward to optimise a conditional
value-at-risk instead of expected return so as to bound potential downsides.

##  3 Approximate Variational Reward Imitation Learning

#### A variational Bayesian approach.

In this section we detail our method, AVRIL, for efficiently learning an
imitator policy and performing reward inference simultaneously. Unlike the
previously mentioned methods, that take sampling or MAP-based approaches to
$p(R|\mathcal{D})$, we employ variational inference (Blei et al., 2017) to
reason about the posterior. Here we posit a surrogate distribution
$q_{\phi}(R)$, parameterised by $\phi$, and aim to minimise the Kullback-
Leibler (KL) divergence to the posterior, resulting in an optimisation
objective:

$\displaystyle\underset{\phi}{\min}\\{D_{\mathrm{KL}}(q_{\phi}(R)||p(R|\mathcal{D}))\\}.$
(6)

This divergence is still as troubling as the posterior to evaluate, leading to
an auxiliary objective function in the Evidence Lower BOund (ELBO):

$\displaystyle\mathcal{F}(\phi)=\mathbb{E}_{q_{\phi}}\big{[}\log
p(\mathcal{D}|R)\big{]}-D_{KL}\big{(}q_{\phi}(R)||p(R)\big{)},$ (7)

where it can be seen that maximisation over $\phi$ is equivalent to (6).
Generally we are agnostic towards the form of both the prior and variational
distribution, for simplicity we assume a Gaussian process prior with mean zero
and unit variance over $R$ alongside the variational posterior distribution
given by $q_{\phi}$ such that:

$\displaystyle q_{\phi}(R)=\mathcal{N}(R;\mu,\sigma^{2}),$ (8)

where $\mu,\sigma^{2}$ are the outputs of an encoder neural network taking $s$
as input and parameterised by $\phi$. Note that for the algorithm that we will
describe these choices are not a necessity and can be easily substituted for
more expressive distributions if appropriate. Maintaining the assumption of
Boltzmann rationality on the part of the demonstrator, our objective takes the
form:

$\displaystyle\mathcal{F}(\phi)=\mathbb{E}_{q_{\phi}}\Bigg{[}\sum_{(s,a)\in\mathcal{D}}\log\frac{\exp(\beta
Q_{R}^{\pi_{D}}(s,a))}{\sum_{b\in\mathcal{A}}\exp(\beta
Q_{R}^{\pi_{D}}(s,b))}\Bigg{]}-D_{KL}\big{(}q_{\phi}(R)||p(R)\big{)}.$ (9)

The most interesting (and problematic) part of this objective as ever centres
on the evaluation of $Q_{R}^{\pi_{D}}(s,a)$. Notice that what is really
required here is an expression of the $Q$-values as a smooth function of the
reward such that with samples of $R$ we could take gradients w.r.t. $\phi$. Of
course there is little hope of obtaining this simply, by itself it is a harder
problem than that of forward RL which only attempts to evaluate the $Q$-values
for a specific $R$ and already in complicated environments has to rely on
function approximation and limited guarantees.

A naive approach would be to sample $\hat{R}$ and then approximate the
$Q$-values with a second neural network, solving offline over the batched data
using a least-squared TD/$Q$-learning algorithm, as is the approach forced on
sampling based BIRL methods. It is in fact though doubly inappropriate for
this setting, not only does this require a solve as an inner-loop but
importantly differentiating through the solving operation is extremely
impractical, it requires backpropagating through a number of gradient updates
that are essentially unbounded as the complexity of the environment increases.

#### A further approximation.

This raises an important question - is it possible to jointly optimise a
policy and variational distribution only once instead of requiring a repeated
solve? This is theoretically suspect, the $Q$-values are defined on a singular
reward, constrained as
$R(s,a)=\mathbb{E}_{s^{\prime},a^{\prime}\sim\pi,T}[Q_{R}^{\pi}(s,a)-\gamma
Q_{R}^{\pi}(s^{\prime},a^{\prime})]$ so we cannot learn a particular standard
$Q$-function that reflects the entire distribution. But can we learn a policy
that reflects the expected reward using a second policy neural network
$Q_{\theta}$? We can’t simply optimise $\theta$ alongside $\phi$ to maximise
the ELBO though as that completely ignores the fact that the learnt policy is
intimately related to the distribution over the reward. Our solution to ensure
then that they behave as intended is by constraining $q_{\phi}$ and
$Q_{\theta}$ to be consistent with each other, specifically that the implied
reward of the policy is sufficiently likely under the variational posterior
(equivalently that the negative log-likelihood is sufficiently low). Thus we
arrive at a constrained optimisation objective given by:

$\displaystyle\underset{\phi,\theta}{\max}\,\sum_{(s,a)\in\mathcal{D}}\log\frac{\exp(\beta
Q_{\theta}(s,a))}{\sum_{b\in\mathcal{A}}\exp(\beta
Q_{\theta}(s,b))}-D_{KL}\big{(}q_{\phi}(R)||p(R)\big{)},$ (10)
$\displaystyle\text{subject to}\,\mathbb{E}_{\pi,T}[-\log
q_{\phi}(Q_{\theta}(s,a)-\gamma Q_{\theta}(s^{\prime},a^{\prime}))]<\epsilon.$

with $\epsilon$ reflecting the strength of the constraint. Rewriting (10) as a
Lagrangian under the KKT conditions (Karush, 1939; Kuhn & Tucker, 1951), and
given complimentary slackness, we obtain a practical objective function:

$\displaystyle\mathcal{F}(\phi,\theta,\mathcal{D})=\sum_{(s,a,s^{\prime},a^{\prime})\in\mathcal{D}}$
$\displaystyle\log\frac{\exp\beta
Q_{\theta}(s,a))}{\sum_{b\in\mathcal{A}}\exp(\beta
Q_{\theta}(s,b))}-D_{KL}\big{(}q_{\phi}(R(s,a))||p(R(s,a))\big{)}$
$\displaystyle+\lambda\log q_{\phi}(Q_{\theta}(s,a)-\gamma
Q_{\theta}(s^{\prime},a^{\prime})).$ (11)

Here the KL divergence between processes is approximated over a countable set,
and $\lambda$ is introduced to control the strength of constraint.

Result: Parameters $\phi$ of variational distribution and $\theta$ of policy
Q-function

Input: $\mathcal{D},S,A,\gamma,\lambda$, learning rate $\eta$, mini-batch size
$b$;

Initialise $\phi,\theta$ ;

$\triangleright$ Can concatenate into single vector $(\phi,\theta)$

while  _not converged_ do

Sample $\mathcal{D}_{mini}$ from $\mathcal{D}$ ;

$\mathcal{F}(\phi,\theta,\mathcal{D})=\mathbb{E}[\frac{n}{b}\mathcal{F}(\phi,\theta,\mathcal{D}_{mini})]$
;

$\triangleright$ MC estimate total loss

$(\phi^{\prime},\theta^{\prime})\leftarrow(\phi,\theta)+\eta\nabla_{\phi,\theta}\mathcal{F}(\phi,\theta,\mathcal{D})$
;

$\triangleright$ Gradient step for $\phi,\theta$

$\phi,\theta\leftarrow\phi^{\prime},\theta^{\prime}$

end while

Return: $\phi,\theta$

Algorithm 1 Approximate Variational Reward Imitation Learning (AVRIL)

#### On the implementation.

Optimisation is simple as both networks are maximising the same objective and
gradients can be easily obtained through backpropagation while being amenable
to mini-batching, allowing you to call your favourite gradient-based
stochastic optimisation scheme. We re-iterate though that AVRIL really
represents a framework for doing BIRL and not a specific model since
$Q_{\theta}$ and $q_{\phi}$ represent arbitrary function approximators. So far
we have presented both as neural networks, but this does not have to be the
case. Of course the advantage of them is their flexibility and ease of
training but they are still inherently black box. It is then perfectly
possible to swap in any particular function approximator if the task requires
it, using simple linear models for example may slightly hurt performance but
allow for more insight. Despite the specific focus on infinite state-spaces,
AVRIL can still even be applied in the tabular setting by simply representing
the policy and variational distribution with multi-dimensional tensors. Having
settled on their forms, equation (3) is calculated simply and the joint
gradient with respect to $\theta$ and $\phi$ is straight-forwardly returned
using any standard auto-diff package. The whole process is summarised in
Algorithm 1.

We can now see how AVRIL does not suffer the issues outlined in section 2.1.
Our form of $q_{\phi}(R)$ is flexible and easily accommodates a non-linear
form of the reward given a neural architecture - this also removes any
restriction on $\mathcal{S}$, or at least allows for any state space that is
commonly tackled within the IL/RL literature. Additionally we have a single
objective for which all parameters are maximised simultaneously - there are no
inner-loops, costly or otherwise, meaning training is faster than the MCMC
methods by a factor equal roughly to the number of samples they would require.

#### The generative model view.

Figure 2:  Graphical model for Bayesian IRL

Ultimately a policy represents a generative model for the behavioural data we
see while executing in an environment. Ho & Ermon (2016) explicitly make use
of this fact by casting the problem in the GAN framework (Goodfellow et al.,
2014). Our method is more analogous to a VAE (Kingma & Welling, 2013), though
to be clear not exactly, where given the graphical model in figure 2 the
reward can be seen as a latent representation of the policy. Our approach
takes the seen data and amortises the inference, _encoding_ over the state
space. While the policy does not act as a _decoder_ in precisely taking the
encoded reward and outputting a policy, it does take the reward and state
information and translate into actions and therefore behaviour. This approach
has its advantages, first in meaningful interpretation of the latent reward
(which is non-existent in adversarial methods), and secondly that we forgo the
practical difficulties of alternating _min-max_ optimisation (Kodali et al.,
2017) while maintaining a generative view of the policy.

#### Temporal consistency through reward regularisation.

Considering only the first term of (3) yields the standard behavioural cloning
setup (where the logits output of a classification network can be interpreted
as the $Q$-values) as it removes the reward from the equation and just focuses
on matching actions to states. AVRIL can then be seen as a policy-learning
method regularised by the need for the reward implied by the logits to be
consistent. Note that this does not induce any necessary bias since the logits
normally contain an extra degree of freedom allowing them to arbitrarily shift
by some scale factor. This factor is now explicitly constrained by giving the
logits additional meaning in that they represent $Q$-values. This places great
importance on the KL term, since every parameterisation of a policy will have
an associated implied reward, the KL regularises these to be not so far from
the prior and preventing the reward from overfitting to the policy and
becoming pointless. It also is able to double as a regularising term in a
similar manor to previous reward-regularisation methods (Piot et al., 2014;
Reddy et al., 2019) depending on the chosen prior, encouraging the reward to
be close to zero:

Proposition 1 (Reward Regularisation) Assume that the constraint in (10) is
satisfied in that
$\mathbb{E}_{q_{\phi}}[R(s,a)]=\mathbb{E}_{\pi,T}[Q_{\theta}(s,a)-\gamma
Q_{\theta}(s^{\prime},a^{\prime})]$, then given a standard normal prior
$p(R)=\mathcal{N}(R;0,1)$ the KL divergence yields a sparsity regulator on the
implied reward:

$\displaystyle\mathcal{L}_{reg}=\sum_{(s,a,s^{\prime},a^{\prime})\in\mathcal{D}}\frac{1}{2}\big{(}Q_{\theta}(s,a)-\gamma
Q_{\theta}(s^{\prime},a^{\prime})\big{)}^{2}+g(\mathrm{Var}_{q_{\phi}}[R(s,a)]).$
(12)

_Proof._ Appendix. $\square$ This follows immediately from the fact that the
divergence evaluates as
$D_{KL}\big{(}q_{\phi}(R(s,a))||p(R(s,a))\big{)}=\frac{1}{2}(-\log(\mathrm{Var}_{q_{\phi}}[R(s,a)])-1+\mathrm{Var}_{q_{\phi}}[R(s,a)]+\mathbb{E}_{q_{\phi}}[R(s,a)]^{2})$

This then allows AVRIL to inherit the benefit of these methods while also
explicitly learning a reward that can be queried at any point. We are also
allowed the choice of whether it is state-only or state-action. This has so
far been arbitrary, but it is important to consider that a state-only reward
is a necessary and sufficient condition for a reward that is fully
disentangled from the dynamics (Fu et al., 2018). Thus by learning such a
reward and given the final term of (3) that directly connects one-step rewards
in terms of the policy, this forces the policy (not the reward) to account for
the dynamics of the system ensuring temporal consistency in a way that BC for
example simply can’t. Alternatively using a state-action reward means that
inevitably some of the temporal information leaks out of the policy and into
the reward - ultimately to the detriment of the policy but potentially
allowing for a more interpretable (or useful) form of reward depending on the
task at hand.

##  4 Experiments

#### Experimental setup.

We are primarily concerned with the case of medical environments, which is
exactly where the issue of learning without interaction is most crucial, you
just cannot let a policy sample treatments for a patient to try to learn more
about the dynamics. It is also where a level of interpretability in what has
been learnt is important, since the consequence of actions are potentially
very impactful on human lives. As such we focus our evaluation on learning on
a real-life healthcare problem, with demonstrations taken from the Medical
Information Mart for Intensive Care (MIMIC-III) dataset (Johnson et al.,
2016). The data contains trajectories of patients in intensive care recording
their condition and theraputic interventions at one day intervals. We evaluate
the ability of the methods to learn a medical policy in both the two and four
action setting - specifically whether the patient should be placed on a
ventilator, and the decision for ventilation in combination with antibiotic
treatment. These represent the two most common, and important, clinical
interventions recorded in the data. Without a recorded notion of reward,
performance is measured with respect to _action matching_ against a held out
test set of demonstrations with cross-validation.

Alongside the healthcare data and for the purposes of demonstrating
generalisability, we provide additional results on standard environments of
varying complexity in the RL literature, the standard control problems of:
CartPole, a classic control environment aiming to swing up and balance a
pendulum; Acrobot, which aims to maintain a sequence of joints above a given
height; and LunarLander, guiding a landing module to a safe touchdown on the
moon surface. In these settings given sufficient demonstration data all
benchmarks are very much capable of reaching demonstrator level performance,
so we test the algorithms on their ability to handle sample complexity in the
low data regime by testing their performance when given access to a select
number of trajectories which we adjust, replicating the setup in Jarrett et
al. (2020). With access to a simulation through the OpenAI gym (Brockman et
al., 2016), we measure performance by deploying the learnt policies live and
calculating their average return over 300 episodes.

Figure 3:  Control environments performance. We plot the average returns received by the policies when deployed live in the environment against the number of trajectories seen during training. Figure 4: (Top) A state-action reward is learnt and plotted for an otherwise average patient as their blood oxygen level changes. (Bottom) The associated weights given a state-only reward as a linear function of the state-space. Table 1: Healthcare performance. Comparison of methods on the MIMIC-III dataset. Performance of the policy is evaluated on the quality of _action matching_ against a held out test set of demonstrations. We report the accuracy (ACC), area under the receiving operator characteristic curve (AUC) and average precision score (APS). | Ventilator | Ventilator + Antibiotics
---|---|---
Metric | ACC | AUC | APS | ACC | AUC | APS
BC | $0.873\pm 0.007$ | $0.916\pm 0.002$ | $0.904\pm 0.003$ | $0.700\pm 0.009$ | $0.864\pm 0.003$ | $0.665\pm 0.009$
VDICE | $0.879\pm 0.002$ | $0.915\pm 0.002$ | $0.904\pm 0.003$ | $0.710\pm 0.005$ | $0.863\pm 0.002$ | $0.675\pm 0.004$
RCAL | $0.870\pm 0.012$ | $0.916\pm 0.003$ | $0.904\pm 0.005$ | $0.702\pm 0.008$ | $0.865\pm 0.004$ | $0.669\pm 0.006$
DSFN | $0.869\pm 0.005$ | $0.905\pm 0.003$ | $0.885\pm 0.001$ | $0.683\pm 0.007$ | $0.856\pm 0.002$ | $0.670\pm 0.004$
EDM | $0.882\pm 0.011$ | $\mathbf{0.920\pm 0.002}$ | $0.909\pm 0.003$ | $0.716\pm 0.008$ | $0.873\pm 0.002$ | $0.682\pm 0.004$
A-RL | $0.875\pm 0.010$ | $0.904\pm 0.002$ | $0.927\pm 0.002$ | $0.718\pm 0.010$ | $0.864\pm 0.002$ | $0.665\pm 0.005$
AVRIL | $\mathbf{0.891\pm 0.002}$ | $0.917\pm 0.001$ | $\mathbf{0.940\pm 0.001}$ | $\mathbf{0.754\pm 0.001}$ | $\mathbf{0.884\pm 0.000}$ | $\mathbf{0.708\pm 0.002}$

#### Benchmarks.

We test our method (AVRIL) against a number of benchmarks from the offline
IRL/IL setting: Deep Successor Feature Network (DSFN) (Lee et al., 2019), an
offline adaptation of max-margin IRL that generalises past the linear methods
using a deep network with least-squares temporal-difference learning, the only
other method that produces both a reward and policy; Reward-regularized
Classification for Apprenticeship Learning (RCAL) (Piot et al., 2014), where
an explicit regulariser on the sparsity of the implied reward is introduced in
order to account for the dynamics information; ValueDICE (VDICE) (Kostrikov et
al., 2019), an adversarial imitation learning, adapted for the offline setting
by removing the replay regularisation; Energy-based Distribution Matching
(EDM) (Jarrett et al., 2020), the state-of-the-art in offline imitation
learning; and finally the standard example of Behavioural Cloning (BC). To
provide evidence that we are indeed learning an appropriate reward we show an
ablation of our method on the MIMIC data: we take the reward learnt by AVRIL
and use it as the ‘true’ reward used to train a $Q$-network offline to learn a
policy (A-RL). Note that we have not included previous BIRL methods for the
reasons explained in section 2.1, training a network just once in these
environments takes in the order of minutes and repeating this sequentially
thousands of times is just not practical. For aid in comparison all methods
share the same network architecture of two hidden layers of 64 units with ELU
activation functions and are trained using Adam (Kingma & Ba, 2014) with
learning rates individually tuned. Further details on experimental setup and
the implementation of benchmarks can be found in the appendix.

#### Evaluation.

We see for all tasks AVRIL learns an appropriate policy that performs strongly
across the board, being competitive in all cases and in places beating out all
of the other benchmarks. The results for our healthcare example are given in
table 1, with AVRIL performing very strongly, having the highest accuracy and
precision score in both tasks. The results for the control environments are
shown in figure 3. AVRIL performs competitively and is easily capable of
reaching demonstrator level performance in the samples given for these tasks,
though not always as quickly as some of the dedicated offline IL methods.

#### Reward insight.

Remember though that task performance is not exactly our goal. Rather the key
aspect of AVRIL is the inference over the unseen reward in order to gain
information about the preferences of the agent that other black-box policy
methods can’t. In the previous experiments our reward encoder was a neural
network for maximum flexibility and we can see from the performance of A-RL we
learn a representation of the reward that can be used to relearn in the
environment very effectively, albeit not quite to the same standard of AVRIL.
Note this also reflects an original motivation for AVRIL in that offpolicy RL
on top of a learnt reward suffers. In figure 4 we explore how to gain more
insight from the learnt reward using different parameterisations of the
reward. The top graph shows how a learnt state-action reward changes as a
function of blood-oxygen level for an otherwise healthy patient, and it can be
seen that as it drops below average the reward for ventilating the patient
becomes much higher (note this is average for patients _in_ the ICU, not
across the general population). While this is intuitive we still have to query
a neural network repeatedly over the state space to gain insight, the bottom
graph of figure 4 presents then a simpler but perhaps more useful
representation. In this case we learn a state-only reward as before but as a
linear model. This is not as strong a constraint on the policy since that is
still free to be non-linear as a neural network but simultaneously allows us
the insight of what our model considers high value in the environment as we
plot the relative model coefficients for each covariate. We can see here for
example that the biggest impact on the overall estimated quality of a state is
given by blood pressure, well known as an important indicator of health
(Hepworth et al., 1994), strongly impacted by trauma and infection.

#### Gridworld ground-truth comparison

While environments like MIMIC are the main focus of this work they do not lend
them selves to inspection of the uncovered reward as the ground truth simply
is not available to us. We thus demonstrate on a toy gridworld environment, in
order to clearly see the effect of learning a posterior distribution over the
reward. In this (finite) example both the encoder and decoder are represented
by tensors but otherwise the procedure remains the same. Figure 5 plots scaled
heat-maps of: a) the ground truth reward; b) the relative state occupancy of
the expert demonstrations, obtained using value-iteration; c) the reward
posterior mean; and d) the reward standard deviation. The interesting thing to
note is that the standard deviation of the learnt reward essentially resembles
the complement of the state occupancy - revealing the epistemic uncertainty
around that part of the state-space given we haven’t seen any demonstrations
there.

Figure 5:  Gridworld example. Scaled heat-maps of: the ground truth reward;
the relative state occupancy of the expert demonstrations; the reward
posterior mean; and reward standard deviation.

##  5 Conclusions

We have presented a novel algorithm, Approximate Variational Reward Imitation
Learning, for addressing the scalability issues that prevent current Bayesian
IRL methods being used in large and unknown environments. We show that this
performs strongly on real and toy data for learning imitation policies
completely offline and importantly recovers a reward that is both effective
for retraining policies but also offers useful insight into the preferences of
the demonstrator. Of course this still represents an approximation, and there
is room for further, more exact methods or else guarantees on the maximum
divergence. We have focused on simply obtaining the appropriate uncertainty
over reward as well as imitation in high stakes environments - in these
settings it is crucial that learnt policies avoid catastrophic failure and so
how exactly to use the uncertainty in order to achieve truly _safe_ imitation
(or indeed better-that-demonstrator apprenticeship) is increasingly of
interest.

## Acknowledgements

AJC would like to acknowledge and thank Microsoft Research for its support
through its PhD Scholarship Program with the EPSRC. This work was additionally
supported by the Office of Naval Research (ONR) and the NSF (Grant number:
1722516).

We would like to thank all of the anonymous reviewers on OpenReview, alongside
the many members of the van der Schaar lab, for their input, comments, and
suggestions at various stages that have ultimately improved the manuscript.

## References

  * Abbeel & Ng (2004) Pieter Abbeel and Andrew Y Ng.  Apprenticeship learning via inverse reinforcement learning.  In _Proceedings of the twenty-first international conference on Machine learning_ , pp. 1, 2004.
  * Bain & Sammut (1995) Michael Bain and Claude Sammut.  A framework for behavioural cloning.  In _Machine Intelligence 15_ , pp. 103–129, 1995.
  * Blei et al. (2017) David M Blei, Alp Kucukelbir, and Jon D McAuliffe.  Variational inference: A review for statisticians.  _Journal of the American statistical Association_ , 112(518):859–877, 2017.
  * Bloem & Bambos (2014) Michael Bloem and Nicholas Bambos.  Infinite time horizon maximum causal entropy inverse reinforcement learning.  In _53rd IEEE Conference on Decision and Control_ , pp. 4911–4916. IEEE, 2014.
  * Brockman et al. (2016) Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba.  Openai gym, 2016.
  * Brown et al. (2020) Daniel Brown, Scott Niekum, and Marek Petrik.  Bayesian robust optimization for imitation learning.  _Advances in Neural Information Processing Systems_ , 33, 2020.
  * Brown & Niekum (2019) Daniel S Brown and Scott Niekum.  Deep bayesian reward learning from preferences.  _arXiv preprint arXiv:1912.04472_ , 2019.
  * Choi & Kim (2011) Jaedeug Choi and Kee-Eung Kim.  Map inference for bayesian inverse reinforcement learning.  In _Advances in Neural Information Processing Systems_ , pp. 1989–1997, 2011.
  * Choi & Kim (2012) Jaedeug Choi and Kee-Eung Kim.  Nonparametric bayesian inverse reinforcement learning for multiple reward functions.  In _Advances in Neural Information Processing Systems_ , pp. 305–313, 2012.
  * Dimitrakakis & Rothkopf (2011) Christos Dimitrakakis and Constantin A Rothkopf.  Bayesian multitask inverse reinforcement learning.  In _European workshop on reinforcement learning_ , pp. 273–284. Springer, 2011.
  * Fu et al. (2018) Justin Fu, Katie Luo, and Sergey Levine.  Learning robust rewards with adverserial inverse reinforcement learning.  In _International Conference on Learning Representations_ , 2018.
  * Gamerman & Lopes (2006) Dani Gamerman and Hedibert F Lopes.  _Markov chain Monte Carlo: stochastic simulation for Bayesian inference_.  CRC Press, 2006.
  * Goodfellow et al. (2014) Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio.  Generative adversarial nets.  In _Advances in neural information processing systems_ , pp. 2672–2680, 2014.
  * Grathwohl et al. (2019) Will Grathwohl, Kuan-Chieh Wang, Joern-Henrik Jacobsen, David Duvenaud, Mohammad Norouzi, and Kevin Swersky.  Your classifier is secretly an energy based model and you should treat it like one.  In _International Conference on Learning Representations_ , 2019.
  * Hepworth et al. (1994) Joseph T Hepworth, Sherry Garrett Hendrickson, and Jean Lopez.  Time series analysis of physiological response during icu visitation.  _Western journal of nursing research_ , 16(6):704–717, 1994.
  * Herman et al. (2016) Michael Herman, Tobias Gindele, Jörg Wagner, Felix Schmitt, and Wolfram Burgard.  Inverse reinforcement learning with simultaneous estimation of rewards and dynamics.  In _Artificial Intelligence and Statistics_ , pp. 102–110. PMLR, 2016.
  * Hill et al. (2018) Ashley Hill, Antonin Raffin, Maximilian Ernestus, Adam Gleave, Anssi Kanervisto, Rene Traore, Prafulla Dhariwal, Christopher Hesse, Oleg Klimov, Alex Nichol, Matthias Plappert, Alec Radford, John Schulman, Szymon Sidor, and Yuhuai Wu.  Stable baselines.  https://github.com/hill-a/stable-baselines, 2018.
  * Ho & Ermon (2016) Jonathan Ho and Stefano Ermon.  Generative adversarial imitation learning.  In _Advances in neural information processing systems_ , pp. 4565–4573, 2016.
  * Jarrett et al. (2020) Daniel Jarrett, Ioana Bica, and Mihaela van der Schaar.  Strictly batch imitation learning by energy-based distribution matching.  _Advances in Neural Information Processing Systems_ , 33, 2020.
  * Johnson et al. (2016) Alistair EW Johnson, Tom J Pollard, Lu Shen, H Lehman Li-Wei, Mengling Feng, Mohammad Ghassemi, Benjamin Moody, Peter Szolovits, Leo Anthony Celi, and Roger G Mark.  Mimic-iii, a freely accessible critical care database.  _Scientific data_ , 3(1):1–9, 2016.
  * Karush (1939) William Karush.  Minima of functions of several variables with inequalities as side constraints.  _M. Sc. Dissertation. Dept. of Mathematics, Univ. of Chicago_ , 1939\.
  * Kingma & Ba (2014) Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * Kingma & Welling (2013) Diederik P Kingma and Max Welling.  Auto-encoding variational bayes.  _arXiv preprint arXiv:1312.6114_ , 2013.
  * Klein et al. (2011) Edouard Klein, Matthieu Geist, and Olivier Pietquin.  Batch, off-policy and model-free apprenticeship learning.  In _European Workshop on Reinforcement Learning_ , pp. 285–296. Springer, 2011.
  * Kodali et al. (2017) Naveen Kodali, Jacob Abernethy, James Hays, and Zsolt Kira.  On convergence and stability of gans.  _arXiv preprint arXiv:1705.07215_ , 2017.
  * Kostrikov et al. (2018) Ilya Kostrikov, Kumar Krishna Agrawal, Debidatta Dwibedi, Sergey Levine, and Jonathan Tompson.  Discriminator-actor-critic: Addressing sample inefficiency and reward bias in adversarial imitation learning.  In _International Conference on Learning Representations_ , 2018.
  * Kostrikov et al. (2019) Ilya Kostrikov, Ofir Nachum, and Jonathan Tompson.  Imitation learning via off-policy distribution matching.  _International Conference on Learning Representations (ICLR)_ , 2019\.
  * Kuhn & Tucker (1951) HW Kuhn and AW Tucker.  Nonlinear programming.  In _Proceedings of the Second Berkeley Symposium on Mathematical Statistics and Probability_. The Regents of the University of California, 1951\.
  * Lee et al. (2019) Donghun Lee, Srivatsan Srinivasan, and Finale Doshi-Velez.  Truly batch apprenticeship learning with deep successor features.  _International Joint Conference on Artificial Intelligence (IJCAI)_ , 2019.
  * Levine (2018) Sergey Levine.  Reinforcement learning and control as probabilistic inference: Tutorial and review.  _arXiv preprint arXiv:1805.00909_ , 2018.
  * Levine et al. (2011) Sergey Levine, Zoran Popovic, and Vladlen Koltun.  Nonlinear inverse reinforcement learning with gaussian processes.  In _Advances in Neural Information Processing Systems_ , pp. 19–27, 2011.
  * Mindermann et al. (2018) Sören Mindermann, Rohin Shah, Adam Gleave, and Dylan Hadfield-Menell.  Active inverse reward design.  _arXiv preprint arXiv:1809.03060_ , 2018.
  * Mnih et al. (2013) Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller.  Playing atari with deep reinforcement learning.  _arXiv preprint arXiv:1312.5602_ , 2013.
  * Ng et al. (2000) Andrew Y Ng, Stuart J Russell, et al.  Algorithms for inverse reinforcement learning.  In _Icml_ , volume 1, pp. 2, 2000.
  * Piot et al. (2014) Bilal Piot, Matthieu Geist, and Olivier Pietquin.  Boosted and reward-regularized classification for apprenticeship learning.  In _Proceedings of the 2014 international conference on Autonomous agents and multi-agent systems_ , pp. 1249–1256. International Foundation for Autonomous Agents and Multiagent Systems, 2014.
  * Raffin (2018) Antonin Raffin.  Rl baselines zoo.  https://github.com/araffin/rl-baselines-zoo, 2018.
  * Ramachandran & Amir (2007) Deepak Ramachandran and Eyal Amir.  Bayesian inverse reinforcement learning.  In _IJCAI_ , volume 7, pp. 2586–2591, 2007.
  * Rasmussen (2003) Carl Edward Rasmussen.  Gaussian processes in machine learning.  In _Summer School on Machine Learning_ , pp. 63–71. Springer, 2003\.
  * Reddy et al. (2019) Siddharth Reddy, Anca D Dragan, and Sergey Levine.  Sqil: Imitation learning via reinforcement learning with sparse rewards.  In _International Conference on Learning Representations_ , 2019.
  * Rothkopf & Dimitrakakis (2011) Constantin A Rothkopf and Christos Dimitrakakis.  Preference elicitation and inverse reinforcement learning.  In _Joint European conference on machine learning and knowledge discovery in databases_ , pp. 34–48. Springer, 2011.
  * Schulman et al. (2017) John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov.  Proximal policy optimization algorithms.  _arXiv preprint arXiv:1707.06347_ , 2017.
  * Ziebart (2010) Brian D Ziebart.  _Modeling Purposeful Adaptive Behavior with the Principle of Maximum Causal Entropy_.  PhD thesis, University of Washington, 2010.
  * Ziebart et al. (2008) Brian D Ziebart, Andrew L Maas, J Andrew Bagnell, and Anind K Dey.  Maximum entropy inverse reinforcement learning.  In _Aaai_ , volume 8, pp. 1433–1438. Chicago, IL, USA, 2008.

##  Appendix A Experimental Setup

#### Expert Demonstrators.

Demonstrations are produced by running pre-trained and hyperparmeter-optimised
agents taken from the RL Baselines Zoo (Raffin, 2018) in OpenAI Stable
Baselines (Hill et al., 2018). For Acrobot and LunarLander these are DQNs
(Mnih et al., 2013), while CartPole uses PPO2 (Schulman et al., 2017).
Trajectories were then sub-sampled for every 20th step in Acrobot and
CartPole, and every 5th step in LunarLander.

#### Testing setup.

For control environments algorithms were presented with (1,3,7,10,15)
trajectories uniformly sampled from a pool of 1000 expert trajectories. Each
algorithm was then trained until convergence and tested by performing 300 live
roll-outs in the simulated environment and recording the average accumulated
reward received in each episode. This whole process was then repeated 10
times, consequently with different initialisations and seen trajectories.

#### Implementations.

All methods are neural network based and so in experiments they share the same
architecture of 2 hidden layers of 64 units each connected by exponential
linear unit (ELU) activation functions.

Publicly available code was used in the implementations of a number of the
benchmarks, specifically:

  * •

VDICE (Kostrikov et al., 2019):
https://github.com/google-research/google-research/tree/master/value_dice

  * •

DSFN (Lee et al., 2019):
https://github.com/dtak/batch-apprenticeship-learning

  * •

EDM (Jarrett et al., 2020):
https://github.com/wgrathwohl/JEM

Note that VDICE was originally designed for continuous actions with a Normal
distribution output which we adapt for the experiments by replacing with a
Gumbel-softmax.

##  Appendix B Proofs

#### Proof of proposition 1.

Assuming the constraint is satisfied, we are maximising the following
objective:

$\displaystyle\mathcal{F}(\phi,\theta)=\sum_{(s,a,s^{\prime},a^{\prime})\in\mathcal{D}}$
$\displaystyle\log\frac{\exp\beta
Q_{\theta}(s,a))}{\sum_{b\in\mathcal{A}}\exp(\beta
Q_{\theta}(s,b))}-D_{KL}\big{(}q_{\phi}(R(s,a))||p(R(s,a))\big{)}$ (13)

Which is equivalent to minimising the negative value

$\displaystyle\mathcal{F}(\phi,\theta)=\sum_{(s,a,s^{\prime},a^{\prime})\in\mathcal{D}}$
$\displaystyle\underbrace{-\log\frac{\exp\beta
Q_{\theta}(s,a))}{\sum_{b\in\mathcal{A}}\exp(\beta
Q_{\theta}(s,b))}}_{\mathcal{L}_{BC}}+\underbrace{D_{KL}\big{(}q_{\phi}(R(s,a))||p(R(s,a))\big{)}}_{\mathcal{L}_{reg}},$
(14)

with the first term $\mathcal{L}_{BC}$ being the negative log-likelihood of
the data and the classic behavioural cloning objective. Now given a standard
Gaussian prior then the KL divergence of a Gaussian with mean $\mu$ and
variance $\sigma^{2}$ from the prior is given by
$\frac{1}{2}(-\log(\sigma^{2})+\sigma^{2}-1+\mu^{2})$ (Kingma & Welling,
2013). Then given our prior $p(R(s,a))=\mathcal{N}(R;0,1)$, the KL evaluates
as:

$\displaystyle\mathcal{L}_{reg}$
$\displaystyle=\sum_{(s,a,s^{\prime},a^{\prime})\in\mathcal{D}}D_{KL}\big{(}q_{\phi}(R(s,a))||p(R(s,a))\big{)}$
(15)
$\displaystyle=\sum_{(s,a,s^{\prime},a^{\prime})\in\mathcal{D}}\frac{1}{2}(-\log(\mathrm{Var}_{q_{\phi}}[R(s,a)])+\mathrm{Var}_{q_{\phi}}[R(s,a)]-1+\mathbb{E}_{q_{\phi}}[R(s,a)]^{2})$
(16)
$\displaystyle=\sum_{(s,a,s^{\prime},a^{\prime})\in\mathcal{D}}\frac{1}{2}(\mathbb{E}_{q_{\phi}}[R(s,a)]^{2})+g(\mathrm{Var}_{q_{\phi}}[R(s,a)])$
(17)
$\displaystyle=\sum_{(s,a,s^{\prime},a^{\prime})\in\mathcal{D}}\frac{1}{2}\big{(}Q_{\theta}(s,a)-\gamma
Q_{\theta}(s^{\prime},a^{\prime})\big{)}^{2}+g(\mathrm{Var}_{q_{\phi}}[R(s,a)])$
(18)

Since by assumption
$\mathbb{E}_{q_{\phi}}[R(s,a)]=\mathbb{E}_{\pi,T}[Q_{\theta}(s,a)-\gamma
Q_{\theta}(s^{\prime},a^{\prime})]$ with the expectation approximated over
samples in the data and considering a definition of the function $g$ to be
$g(\mathrm{Var}_{q_{\phi}}[R(s,a)])=\frac{1}{2}(-\log(\mathrm{Var}_{q_{\phi}}[R(s,a)])+\mathrm{Var}_{q_{\phi}}[R(s,a)]-1)$.
$\square$


#  MapCoder: Multi-Agent Code Generation
for Competitive Problem Solving

Md. Ashraful Islam1, Mohammed Eunus Ali1, Md Rizwan Parvez2
1Bangladesh University of Engineering and Technology
2Qatar Computing Research Institute (QCRI)
{mdashrafulpramanic<EMAIL_ADDRESS><EMAIL_ADDRESS>
###### Abstract

Code synthesis, which requires a deep understanding of complex natural
language (NL) problem descriptions, generation of code instructions for
complex algorithms and data structures, and the successful execution of
comprehensive unit tests, presents a significant challenge. Thus, while large
language models (LLMs) demonstrate impressive proficiency in natural language
processing (NLP), their performance in code generation tasks remains limited.
In this paper, we introduce a new approach to code generation tasks leveraging
the multi-agent prompting that uniquely replicates the full cycle of program
synthesis as observed in human developers. Our framework, MapCoder, consists
of four LLM agents specifically designed to emulate the stages of this cycle:
recalling relevant examples, planning, code generation, and debugging. After
conducting thorough experiments, with multiple LLMs ablations and analyses
across eight challenging competitive problem-solving and program synthesis
benchmarks—MapCoder showcases remarkable code generation capabilities,
achieving their new state-of-the-art (pass@1) results—(HumanEval 93.9%, MBPP
83.1%, APPS 22.0%, CodeContests 28.5%, and xCodeEval 45.3%). Moreover, our
method consistently delivers superior performance across various programming
languages and varying problem difficulties. We open-source our framework at
https://github.com/Md-Ashraful-Pramanik/MapCoder.

MapCoder: Multi-Agent Code Generation
for Competitive Problem Solving

Md. Ashraful Islam1, Mohammed Eunus Ali1, Md Rizwan Parvez2 1Bangladesh
University of Engineering and Technology 2Qatar Computing Research Institute
(QCRI) {mdashrafulpramanic<EMAIL_ADDRESS><EMAIL_ADDRESS>

##  1 Introduction

Computer Programming has emerged as an ubiquitous problem-solving tool that
brings tremendous benefits to every aspects of our life Li et al. (2022a);
Parvez et al. (2018); Knuth (1992). To maximize programmers’ productivity, and
enhance accessibility, automation in program synthesis is paramount. With the
growth of LLMs, significant advancements have been made in program
synthesis—driving us in an era where we can generate fully executable code,
requiring no human intervention Chowdhery et al. (2022); Nijkamp et al.
(2022).

Despite LLMs’ initial success and the scaling up of model size and data, many
of these models still struggle to perform well on complex problem-solving
tasks, especially in competitive programming problems Austin et al. (2021). To
mitigate this gap, in this paper, we introduce MapCoder: a Multi-Agent
Prompting Based Code Generation approach that can seamlessly synthesize
solutions for competition-level programming problems.

Competitive programming or competition-level code generation, often regarded
as the pinnacle of problem-solving, is an challenging task. It requires a deep
comprehension of NL problem descriptions, multi-step complex reasoning beyond
mere memorization, excellence in algorithms and data structures, and the
capability to generate substantial code that produces desired outputs aligned
with comprehensive test cases Khan et al. (2023).

Figure 1: Overview of MapCoder (top). It starts with a retrieval agent that
generates relevant examples itself, followed by planning, coding, and
iterative debugging agents. Our dynamic traversal (bottom) considers the
confidence of the generated plans as their reward scores and leverages them to
guide the code generation accordingly.

Early approaches utilizing LLMs for code generation employ a direct prompting
approach, where LLMs generate code directly from problem descriptions and
sample I/O Chen et al. (2021a). Recent methods like chain-of-thought (Wei et
al., 2022a) advocates modular or pseudo code-based generation to enhance
planning and reduce errors, while retrieval-based approaches such as Parvez et
al. (2021) leverage relevant problems and solutions to guide LLMs’ code
generations. However, gains in such approaches remains limited in such a
complex task like code generation where LLMs’ generated code often fails to
pass the test cases and they do not feature bug-fixing schema Ridnik et al.
(2024).

A promising solution to the above challenge is self-reflection Shinn et al.
(2023); Chen et al. (2022), which iteratively evaluates the generated code
against test cases, reflects on mistakes and modifies accordingly. However,
such approaches have limitations too. Firstly, while previous studies indicate
that superior problem-solving capabilities are attained when using in-context
exemplars (Shum et al., 2023; Zhang et al., 2022; Wei et al., 2022a) or plans
(Jiang et al., 2023b), these approaches, during both code generation and
debugging, only leverage the problem description itself in a zero-shot manner.
Consequently, their gains can be limited.

To confront the above challenge, we develop MapCoder augmenting the generation
procedure with possible auxiliary supervision. We draw inspiration from human
programmers, and how they use various signals/feedback while programming. The
human problem-solving cycle involves recalling past solutions, planning, code
writing, and debugging. MapCoder imitates these steps using LLM agents -
retrieval, planning, coding, and debugging. In contrast to relying on human
annotated examples, or external code retrieval models, we empower our
retrieval agent to autonomously retrieve relevant problems itself Yasunaga et
al. (2023). Moreover, we design a novel structured pipeline schema that
intelligently cascades the LLM agents and incorporates a dynamic iteration
protocol to enhance the generation procedure at every step. Figure 1 shows an
overview of our approach, MapCoder .

Additionally, existing iterative self-reflection methods rely on extra test
cases generated by LLM agents (e.g., AgentCoder Huang et al. (2023), LATS Zhou
et al. (2023), self-reflection Shinn et al. (2023)) or external tools,
compounding the challenges. Test case generation is equally challenging as
code generation Pacheco et al. (2007), and incorrect test cases can lead to
erroneous code. Blindly editing code based on these test cases can undermine
problem-solving capabilities. For instance, while self-reflection boosts
GPT-4’s performance on the HumanEval dataset, it drops by 3% on the MBPP
dataset Shinn et al. (2023). Upon identification, to validate this, on the
HumanEval dataset itself, we replace their GPT-4 with ChatGPT, and see that
model performance drops by 26.3%. Therefore, our debugging agent performs unit
tests and bug fixing using only the sample I/O, without any artifact-more
plausible for real-world widespread adoption.

We evaluate MapCoder on seven popular programming synthesis benchmarks
including both basic programming like HumanEval, MBPP and challenging
competitive program-solving benchmarks like APPS, CodeContests and xCodeEval.
With multiple different LLMs including ChatGPT, GPT-4, and Gemini Pro, our
approach significantly enhances their problem-solving capabilities -
consistently achieving new SOTA performances, outperforming strong baselines
like Reflexion (Shinn et al., 2023), and AlphaCodium (Ridnik et al., 2024).
Moreover, our method consistently delivers superior performance across various
programming languages and varying problem difficulties. Furthermore, with
detailed ablation studies, we analyze MapCoder to provide more insights.

##  2 Related Work

Program Synthesis: Program synthesis has a long standing history in AI systems
(Manna and Waldinger, 1971). A large number of prior research attempted to
address it via search/data flow approaches Li et al. (2022a); Parisotto and
Salakhutdinov (2017); Polozov and Gulwani (2015); Gulwani (2011). LMs, prior
to LLMs, attempt to generate code by fine-tuning (i.e., training) neural
language models (Wang et al., 2021; Ahmad et al., 2021; Feng et al., 2020;
Parvez et al., 2018; Yin and Neubig, 2017; Hellendoorn and Devanbu, 2017;
Rabinovich et al., 2017; Hindle et al., 2016), conversational intents or data
flow features (Andreas et al., 2020; Yu et al., 2019).
Large Language Models: Various LLMs have been developed for Code synthesis (Li
et al., 2022b; Fried et al., 2022; Chen et al., 2021b; Austin et al., 2021;
Nijkamp et al., 2022; Allal et al., 2023). Recent open source LLMs include
Llama-2 (Touvron et al., 2023), CodeLlama-2 (Roziere et al., 2023), Mistral
(Jiang et al., 2023a) Deepseek Coder Guo et al. (2024), MoTCoder Li et al.
(2023) that are capable of solving many basic programming tasks.

---
Table 1: Features in code generation prompt techniques.

Prompting LLMs: As indicated in Section 1, LLM prompting can be summarized
into three categories: retrieval Yasunaga et al. (2023); Parvez et al. (2023,
2021); planning (Wei et al., 2022b; Jiang et al., 2023b); debugging (Ridnik et
al., 2024; Chen et al., 2023, 2022; Le et al., 2022) apart from the direct
code generation approaches. In contrast, we combine all these paradigms and
bridge their gaps (See Table 1). Among others, in different contexts of
generic problem-solving, Tree-of-thoughts Yao et al. (2023), and Cumulative
reasoning Zhang et al. (2023) approaches consider a tree traversal approach to
explore different sub-steps towards a solution while our code generation
approach mirrors the human programming cycle through various LLM agents.
Notably, our traversal does not rely on sub-steps toward the solution but
instead utilizes different forms of complete solutions.

##  3 MapCoder

Our goal is to develop a multi-agent code generation approach for competitive
problem-solving. In order to do so, our framework, MapCoder, replicates the
human programming cycle through four LLM agents - retrieval, plan, code, and
debug. We devise a pipeline sequence for MapCoder, intelligently cascading the
agents in a structured way and enhancing each agent’s capability by augmenting
in-context learning signals from previous agents in the pipeline. However, not
all the agent responses/outputs are equally useful. Therefore, additionally,
MapCoder features an adaptive agent traversal schema to interact among
corresponding agents dynamically, iteratively enhancing the generated code by,
for example, fixing bugs, while maximizing the usage of the LLM agents. In
this section, we first discuss the agents (as per the pipeline), their
prompts, and interactions, followed by the dynamic agent traversal protocol in
MapCoder towards code generation for competitive problem-solving.

###  3.1 Retrieval Agent

Our first agent, the _Retrieval Agent_ , recalls past relevant problem-solving
instances, akin to human memory. It finds $k$ (user-defined) similar problems
without manual crafting or external retrieval models. Instead, we leverage the
LLM agent itself, instructing it to generate such problems. Our prompt extends
the analogical prompting principles Yasunaga et al. (2023), generating
examples and their solutions simultaneously, along with additional metadata
(e.g., problem description, code, and plan) to provide the following agents as
auxiliary data. We adopt a specific sequence of instructions, which is crucial
for the prompt’s effectiveness. In particular, initially, we instruct the LLM
to produce similar and distinct problems and their solutions, facilitating
problem planning reverse-engineering. Then, we prompt the LLM to generate
solution code step-by-step, allowing post-processing to form the corresponding
plan. Finally, we direct the LLM to generate relevant algorithms and provide
instructional tutorials, enabling the agent to reflect on underlying
algorithms and generate algorithmically similar examples.

###  3.2 Planning Agent

The second agent, the _Planning Agent_ , aims to create a step-by-step plan
for the original problem. Our _Planning Agent_ uses examples and their plans
obtained from the retrieval agent to generate plans for the original problem.
A straightforward approach would be to utilize all examples collectively to
generate a single target plan. However, not all retrieved examples hold equal
utility. Concatenating examples in a random order may compromise the LLM’s
ability to generate accurate planning. For instance, Xu et al. (2023)
demonstrated that even repeating more relevant information (e.g., query)
towards the end of the in-context input aids LLM reasoning more effectively
than including relatively less relevant contexts. A similar conclusion of
"separating noisy in-context data" can also be drawn from the state-of-the-art
retrieval augmented generation approaches like Wang et al. (2023). Therefore,
we generate a distinct target plan for each retrieved example. Additionally,
multiple plans offer diverse pathways to success.

To help the generation steps in the following agents with the utility
information for each plan, our designed prompt for the planning agent asks the
LLM to generate both plans and a confidence score. Figure 2 shows our prompt
got this agent.

Figure 2: Prompt for _Planning Agent_.

###  3.3 Coding Agent

Next is the _Coding Agent_. It takes the problem description, and a plan from
the _Planning Agent_ as input and translates the corresponding planning into
code to solve the problem. During the traversing of agents, _Coding Agent_
takes the original problem and one particular plan from the _Planning Agent_ ,
generates the code, and test on sample I/O. If the initial code fails, the
agent transfers it to the next agent for debugging. Otherwise, predicts that
as the final solution.

###  3.4 Debugging Agent

Finally, the _Debugging Agent_ utilizes sample I/O from the problem
description to rectify bugs in the generated code. Similar to humans cross-
checking their plan while fixing bugs, our pipeline supplements the _Debugging
Agent_ with plans from the _Planning Agent_. This plan-derived debugging
significantly enhances bug fixing in MapCoder, underscoring the pivotal roles
played by both the _Debugging Agent_ and the _Planning Agent_ in the
generation process. We verify this in Section 6. For each plan, this process
is repeated $t$ times. The prompt for this step is illustrated in Figure 3.
Note that, different from Reflexion Shinn et al. (2023) and AlphaCodium Ridnik
et al. (2024), our _Debugging Agent_ does not require any additional test case
generation in the pipeline.

Figure 3: Prompt for _Debugging Agent_. Figure 4: Example problem and solution
generation using Direct, CoT, Reflexion, and MapCoder prompts. MapCoder
explores high-utility plans first and uniquely features a plan-derived
debugging for enhanced bug fixing.

###  3.5 Dynamic Agent Traversal

The dynamic traversal in MapCoder begins with the _Planning Agent_ , which
outputs the plans for the original problem with confidence scores. These plans
are sorted, and the highest-scoring one is sent to the Coding Agent. The
Coding Agent translates the plan into code, tested with sample I/Os. If all
pass, the code is returned; otherwise, it’s passed to _Debugging Agent_. They
attempt to rectify the code iteratively up to $t$ times. If successful, the
code is returned; otherwise, responsibility shifts back to the _Planning
Agent_ for the next highest confidence plan. This iterative process continues
for $k$ iterations, reflecting a programmer’s approach. We summarize our agent
traversal in Algorithm A in Appendix. Our algorithm’s complexity is $O(kt)$.
An example illustrating MapCoder’s problem-solving compared to Direct, Chain-
of-thought, and Reflexion approaches is in Figure 4. All detailed prompts for
each agent are in Appendix B.

##  4 Experimental Setup

###  4.1 Datasets

For extensive evaluation, we have used eight benchmark datasets: five from
basic programming and three from complex competitive programming domains. Five
basic programming datasets are: HumanEval Chen et al. (2021a), HumanEval-ET
Dong et al. (2023a), EvalPlus Liu et al. (2023), MBPP) Austin et al. (2021),
and MBPP-ET Dong et al. (2023a). HumanEval-ET, EvalPlus extend HumanEval and
MBPP-ET comprehends MBPP by incorporating more test cases. The problem set
size of HumanEval and MBPP (and their extensions) are 164 and 397,
respectively. Due to the absence of sample I/O in MBPP and MBPP-ET, our
approach for code moderation involves randomly removing one test-case from
MBPP-ET for each problem and provide this test-case as a sample I/O for the
problem. Importantly, this removed test-case is carefully selected to ensure
mutual exclusivity from the hidden test sets in MBPP and MBPP-ET. Three
competitive programming datasets are: Automated Programming Progress Standard
(APPS), xCodeEval Khan et al. (2023), and CodeContest, where we have used 150,
106, and 156 problems, respectively, in our experiments.

###  4.2 Baselines

We have compared MapCoder with several baselines and state-of-the-art
approaches. Direct Prompting instructs language models to generate code
without explicit guidance, relying on their inherent capabilities of LLM.
Chain of Thought Prompting (CoT) Wei et al. (2022b) breaks down problems into
step-by-step solutions, enabling effective tackling of complex tasks. Self-
Planning Prompting Jiang et al. (2023b) divides the code generation task into
planning and implementation phases. Analogical Reasoning Prompting Yasunaga et
al. (2023) instructs models to recall relevant problems from training data.
Reflexion Shinn et al. (2023) provides verbal feedback to enhance solutions
based on unit test results. Self-collaboration Dong et al. (2023b) proposes a
framework where different LLMs act as analyst, coder, and tester to
cooperatively generate code for complex tasks, achieving better performance
than directly using a single LLM. AlphaCodium Ridnik et al. (2024) iteratively
refines code based on AI-generated input-output tests.

---
Table 2: Pass@1 results for different approaches. The results of the yellow
and blue colored cells are obtained from Jiang et al. (2023b) and Shinn et al.
(2023), respectively. The results of the Self-collaboration Dong et al.
(2023b) paper are collected from their paper. The green texts indicate the
state-of-the-art results, and the red text is gain over Direct Prompting
approach.

###  4.3 Foundation Models, Evaluation Metric, $k$, and $t$

With $k=t=5$ in HumanEval, and $k=t=3$ for others, we evaluate all the
datasets using ChatGPT (gpt-3.5-turbo-1106), GPT-4 (gpt-4-1106-preview) from
OpenAI and Gemini Pro from Google. We have also evaluated our method using an
open-source LLM, Mistral-7B-instruct. We have used the Pass@k evaluation
metric, where the model is considered successful if at least one of the $k$
generated solutions is correct.

##  5 Results

In this section, we evaluate the code generation capabilities of our
framework, MapCoder, for competitive problem solving. Our experimental results
are reported in Table 2. Overall, MapCoder shows a tremendous excellence in
code generation, significantly outperforms all baselines, and achieves new
state-of-the-art results in all benchmarks. In general the scales with GPT-4
are higher than ChatGPT.

###  5.1 Performance on basic code generation

The highest scale of performance (Pass@1) scores are observed in simple
program synthesis tasks like HumanEval, MBPP in Table 2. Though with the
simpler problem (non-contests) datasets such as HumanEval, HumanEval-ET, the
current state-of-the-art method, Reflexion Shinn et al. (2023) perform
reasonably well, this approach does not generalize across varying datasets
depicting a wide variety of problems. Self-reflection techniques enhance
GPT-4’s performance on HumanEval but result in a 3% decrease on the MBPP
dataset. Similarly, with ChatGPT, there’s a notable 26.3% drop in performance
where in several cases their AI generated test cases are incorrect. We observe
that 8% of failures in HumanEval and 15% in MBPP is caused by their AI
generates incorrect test cases while our approach is independent of AI test
cases, and consistently improves code generations in general. Consequently,
even in HumanEval, with GPT-4, our Pass@1 surpasses Reflexion by $\sim$3%. On
top, in all four simple programming datasets, MapCoder enhances the Direct
prompting significantly with a maximum of 88% on HumanEvalET by ChatGPT.

Figure 5: The number of correct answers wrt algorithm types (tags) and
difficulty levels (xCodeEval dataset).

###  5.2 Performance on competitive problem solving

The significance of MapCoder shines through clearly when evaluated in
competitive problem-solving contexts. Across datasets such as APPS, xCodeEval,
and CodeContests, MapCoder demonstrates substantial enhancements over Direct
prompting methods, with improvements of 41.3%, 52.6%, and 132.8% for ChatGPT,
and 73.7%, 41.2%, and 135.1% for GPT4, respectively. Notably, the most
challenging datasets are APPS and CodeContest, where MapCoder’s performance
stands out prominently. We deliberately compare against strong baselines on
these datasets, regardless of whether they are prompt-based or not.
Importantly, on CodeContest our Pass@1 results match the Pass@5 scores of the
concurrent state-of-the-art model AlphaCodium (Ridnik et al., 2024): 28.5% vs.
their 29% (see Table 3). Furthermore, our Pass@5 results demonstrate an
additional improvement of 12.8%. On APPS, MapCoder consistently surpasses the
Pass@1 scores of all baseline prompts for both ChatGPT and GPT-4.

---
Table 3: Pass@5 results on CodeContest dataset. AlphCodium result are from
Ridnik et al. (2024). The green cells indicate the SoTA and the red text
indicates improvement w.r.t Direct approach.

###  5.3 Performance with Varying Difficulty Levels

The APPS dataset comprises problems categorized into three difficulty levels:
(i) Introductory, (ii) Interview, and (iii) Competition. Figure 6 illustrates
the performance of various competitive approaches for these three categories.
The results reveal that our MapCoder excels across all problem categories,
with highest gain in competitive problem-solving indicating its superior code
generation capabilities in general, and on top, remarkable effectiveness in
competitive problem-solving. In order to gather more understanding on what
algorithm problems it’s capable of solving and in fact much difficulty level
it can solve, we have also conducted a comparison between MapCoder and the
Direct approach, considering the difficulty levels111Difficulty levels in
xCodeEval dataset represents an integer number, a higher value means more
difficult problem and tags222Tags in xCodeEval dataset represents algorithm
type that can be used to solve the problem i.e., greedy, dp, brute-force,
constructive, and so on. present in the xCodeEval dataset. The results of this
comparison are depicted in Figure 5. This comparison showcases that MapCoder
is effective across various algorithm types and exhibits superior performance
even in higher difficulty levels, compared to the Direct approach. However,
beyond (mid-level: difficulties>1000), its gains are still limited.

Figure 6: Performance vs problem types (APPS).

###  5.4 Performance Across Different LLMs

To show the robustness of MapCoder across various LLMs, we evaluate MapCoder
using Gemini Pro, a different family of SoTA LLM in Table 4. We also evaluate
MapCoder using an open-source LLM Mistral-7B instruct in Table 5. As expected,
our method shows performance gains over other baseline approaches in equitable
trends on both simple (HumanEval) and contest-level problems (CodeContest).

---
Table 4: Pass@1 results with using Gemini Pro. The red text is gain over
Direct Prompting approach.
---
Table 5: Pass@1 results with using Mistral-7B-instruct. The red text is gain
over Direct Prompting approach.

###  5.5 Performance Across Different Programming Languages

Furthermore, we evaluate model performances using MapCoder across different
programming languages. We utilize the xCodeEval dataset, which features
multiple languages. Figure 7 shows that consistent proficiency across
different programming languages is achieved by MapCoder with respect to
baselines.

Figure 7: The number of correct answers wrt different programming languages
(xCodeEval dataset).

##  6 Ablations Studies and Analyses

We present the ablation study of the MapCoder on HumanEval dataset as the
problems are simpler and easy to diagnose by us humans.

###  6.1 Impact of Different Agents

We have also conducted a study by excluding certain agents from our MapCoder,
which helps us investigate each agent’s impact in our whole pipeline. As
expected, the results (Table 6) show that every agent has its role in the
pipeline as turning off any agent decreases the performance of MapCoder.
Furthermore, we observe that the Debugging Agent has the most significant
impact on the pipeline, as evidenced by a performance drop of 17.5% when
excluding this agent exclusively, and an avg performance drop of 24.83% in all
cases. The _Planning agent_ has the second best important with avg drop of
16.7% in all cases. In Table 6), we perform an ablation study of our multi-
agent framework investigate each agent’s impact in our whole pipeline.

---

Table 6: Pass@1 results for different versions of MapCoder (by using ChatGPT
on HumanEval dataset).

###  6.2 Qualitative Example

To verify the above numerical significance, and to understand how our method
enhance the code generation, we have performed a qualitative analysis to find
the underlying reason for the superior performance of MapCoder over other
competitive prompting approaches. An example problem and the output with the
explanation of Direct, CoT, Reflexion, and MapCoder prompting is shown in
Figure 4. This example demonstrates how the _Debugging Agent_ fixes the bugs
leveraging the plan as a guide from the _Planning Agent_. This verifies the
impact of these two most significant agents. We present more detailed examples
in Appendix.

###  6.3 Impact of $k$ and $t$

MapCoder involves two hyper-parameters: the number of self-retrieved
exemplars, $k$, and the number of debugging attempts, $t$. Our findings (Table
7) reveal that higher $k$, $t$ is proportionate performance gain at the
expense of time.

---

Table 7: Pass@1 results by varying $k$ and $t$.
---
Table 8: Average number of API calls, thousands of tokens used, required time
in minutes to get the API response.

###  6.4 Impact of Number of Sample I/Os

Given the limited number of sample I/Os in the HumanEval dataset (average of
2.82 per problem), we supplemented it with an additional 5 sample I/Os from
the HumanEval-ET dataset. Experiments with this augmented set showed an 1.5%
performance gain.

###  6.5 Error Analysis and Challenges

Although MapCoder demonstrates strong performance compared to other methods,
it faces challenges in certain algorithmic domains. For example, Figure 5
illustrates MapCoder’s reduced performance on more difficult problems
requiring precise problem understanding and concrete planning—capabilities
still lacking in LLMs. In the xCodeEval dataset (see Figure 5), it solves a
limited number of problems in categories like Combinatorics, Constructive,
Number Theory, Divide and Conquer, and Dynamic Programming (DP). Manual
inspection of five DP category problems reveals occasional misinterpretation
of problems, attempts to solve using greedy or brute-force approaches, and
struggles with accurate DP table construction when recognizing the need for a
DP solution.

##  7 Conclusion and Future Work

In this paper, we introduce MapCoder, a novel framework for effective code
generation in complex problem-solving tasks, leveraging the multi-agent
prompting capabilities of LLMs. MapCoder captures the complete problem-solving
cycle by employing four agents - retrieval, planning, coding, and debugging -
which dynamically interact to produce high-quality outputs. Evaluation across
major benchmarks, including basic and competitive programming datasets,
demonstrates MapCoder’s consistent outperformance of well-established
baselines and SoTA approaches across various metrics. Future work aims to
extend this approach to other domains like question answering and mathematical
reasoning, expanding its scope and impact.

##  8 Limitations

Among the limitations of our work, firstly, MapCoder generates a large number
of tokens, which may pose challenges in resource-constrained environments.
Table 8 shows the number of average API calls and token consumption with the
default $k$ and $t$ (i.e., with respect to the reported performance) while
Table 7) shows how $k$, $t$ can be adjusted to proportionate the performance
gain at the expense of time/token. We have not addressed the problem of
minimizing tokens/API-calls in this paper and leave it for future works.
Secondly, our method currently relies on sample input-output (I/O) pairs for
bug fixing. Although sample I/Os provide valuable insights for LLMs’ code
generation, their limited number may not always capture the full spectrum of
possible test cases. Consequently, enhancing the quality of additional test
case generation could reduce our reliance on sample I/Os and further improve
the robustness of our approach. Additionally, future exploration of open-
source code generation models, such as CodeLLaMa, LLaMa3, Mixtral 8x7B could
offer valuable insights and potential enhancements to our approach. Another
important concern is that while running machine-generated code, it is
advisable to run it inside a sandbox to avoid any potential risks.

## References

  * Ahmad et al. (2021) Wasi Uddin Ahmad, Saikat Chakraborty, Baishakhi Ray, and Kai-Wei Chang. 2021.  Unified pre-training for program understanding and generation.  _arXiv preprint arXiv:2103.06333_.
  * Allal et al. (2023) Loubna Ben Allal, Raymond Li, Denis Kocetkov, Chenghao Mou, Christopher Akiki, Carlos Munoz Ferrandis, Niklas Muennighoff, Mayank Mishra, Alex Gu, Manan Dey, et al. 2023.  Santacoder: don’t reach for the stars!  _arXiv preprint arXiv:2301.03988_.
  * Andreas et al. (2020) Jacob Andreas, John Bufe, David Burkett, Charles Chen, Josh Clausman, Jean Crawford, Kate Crim, Jordan DeLoach, Leah Dorner, Jason Eisner, Hao Fang, Alan Guo, David Hall, Kristin Hayes, Kellie Hill, Diana Ho, Wendy Iwaszuk, Smriti Jha, Dan Klein, Jayant Krishnamurthy, Theo Lanman, Percy Liang, Christopher H. Lin, Ilya Lintsbakh, Andy McGovern, Aleksandr Nisnevich, Adam Pauls, Dmitrij Petters, Brent Read, Dan Roth, Subhro Roy, Jesse Rusak, Beth Short, Div Slomin, Ben Snyder, Stephon Striplin, Yu Su, Zachary Tellman, Sam Thomson, Andrei Vorobev, Izabela Witoszko, Jason Wolfe, Abby Wray, Yuchen Zhang, and Alexander Zotov. 2020.  Task-oriented dialogue as dataflow synthesis.  _Transactions of the Association for Computational Linguistics_ , 8:556–571.
  * Austin et al. (2021) Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie Cai, Michael Terry, Quoc Le, et al. 2021.  Program synthesis with large language models.  _arXiv preprint arXiv:2108.07732_.
  * Chen et al. (2022) Bei Chen, Fengji Zhang, Anh Nguyen, Daoguang Zan, Zeqi Lin, Jian-Guang Lou, and Weizhu Chen. 2022.  Codet: Code generation with generated tests.  _arXiv preprint arXiv:2207.10397_.
  * Chen et al. (2021a) Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, Alex Ray, Raul Puri, Gretchen Krueger, Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, Nick Ryder, Mikhail Pavlov, Alethea Power, Lukasz Kaiser, Mohammad Bavarian, Clemens Winter, Philippe Tillet, Felipe Petroski Such, Dave Cummings, Matthias Plappert, Fotios Chantzis, Elizabeth Barnes, Ariel Herbert-Voss, William Hebgen Guss, Alex Nichol, Alex Paino, Nikolas Tezak, Jie Tang, Igor Babuschkin, Suchir Balaji, Shantanu Jain, William Saunders, Christopher Hesse, Andrew N. Carr, Jan Leike, Josh Achiam, Vedant Misra, Evan Morikawa, Alec Radford, Matthew Knight, Miles Brundage, Mira Murati, Katie Mayer, Peter Welinder, Bob McGrew, Dario Amodei, Sam McCandlish, Ilya Sutskever, and Wojciech Zaremba. 2021a.  Evaluating large language models trained on code.
  * Chen et al. (2021b) 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. 2021b.  Evaluating large language models trained on code.  _arXiv preprint arXiv:2107.03374_.
  * Chen et al. (2023) Xinyun Chen, Maxwell Lin, Nathanael Schärli, and Denny Zhou. 2023.  Teaching large language models to self-debug.  _arXiv preprint arXiv:2304.05128_.
  * Chowdhery et al. (2022) Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. 2022.  Palm: Scaling language modeling with pathways.  _arXiv preprint arXiv:2204.02311_.
  * Dong et al. (2023a) Yihong Dong, Jiazheng Ding, Xue Jiang, Zhuo Li, Ge Li, and Zhi Jin. 2023a.  Codescore: Evaluating code generation by learning code execution.  _arXiv preprint arXiv:2301.09043_.
  * Dong et al. (2023b) Yihong Dong, Xue Jiang, Zhi Jin, and Ge Li. 2023b.  Self-collaboration code generation via chatgpt.
  * Feng et al. (2020) Zhangyin Feng, Daya Guo, Duyu Tang, Nan Duan, Xiaocheng Feng, Ming Gong, Linjun Shou, Bing Qin, Ting Liu, Daxin Jiang, et al. 2020.  Codebert: A pre-trained model for programming and natural languages.  In _Findings of the Association for Computational Linguistics: EMNLP 2020_ , pages 1536–1547.
  * Fried et al. (2022) Daniel Fried, Armen Aghajanyan, Jessy Lin, Sida Wang, Eric Wallace, Freda Shi, Ruiqi Zhong, Wen-tau Yih, Luke Zettlemoyer, and Mike Lewis. 2022.  Incoder: A generative model for code infilling and synthesis.  _arXiv preprint arXiv:2204.05999_.
  * Gulwani (2011) Sumit Gulwani. 2011.  Automating string processing in spreadsheets using input-output examples.  _ACM Sigplan Notices_ , 46(1):317–330.
  * Guo et al. (2024) Daya Guo, Qihao Zhu, Dejian Yang, Zhenda Xie, Kai Dong, Wentao Zhang, Guanting Chen, Xiao Bi, Y Wu, YK Li, et al. 2024.  Deepseek-coder: When the large language model meets programming–the rise of code intelligence.  _arXiv preprint arXiv:2401.14196_.
  * Hellendoorn and Devanbu (2017) Vincent J. Hellendoorn and Premkumar Devanbu. 2017.  Are deep neural networks the best choice for modeling source code? In _Proceedings of the 2017 11th Joint Meeting on Foundations of Software Engineering_ , ESEC/FSE 2017, pages 763–773, New York, NY, USA. ACM.
  * Hindle et al. (2016) Abram Hindle, Earl T. Barr, Mark Gabel, Zhendong Su, and Premkumar Devanbu. 2016.  On the naturalness of software.  _Commun. ACM_ , 59(5):122–131.
  * Huang et al. (2023) Dong Huang, Qingwen Bu, Jie M Zhang, Michael Luck, and Heming Cui. 2023.  Agentcoder: Multi-agent-based code generation with iterative testing and optimisation.  _arXiv preprint arXiv:2312.13010_.
  * Jiang et al. (2023a) Albert Q. Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, Lélio Renard Lavaud, Marie-Anne Lachaux, Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed. 2023a.  Mistral 7b.
  * Jiang et al. (2023b) Xue Jiang, Yihong Dong, Lecheng Wang, Qiwei Shang, and Ge Li. 2023b.  Self-planning code generation with large language model.  _arXiv preprint arXiv:2303.06689_.
  * Khan et al. (2023) Mohammad Abdullah Matin Khan, M Saiful Bari, Xuan Long Do, Weishi Wang, Md Rizwan Parvez, and Shafiq Joty. 2023.  xcodeeval: A large scale multilingual multitask benchmark for code understanding, generation, translation and retrieval.  _arXiv preprint arXiv:2303.03004_.
  * Knuth (1992) Donald E Knuth. 1992.  Literate programming.  _CSLI Lecture Notes, Stanford, CA: Center for the Study of Language and Information (CSLI), 1992_.
  * Le et al. (2022) Hung Le, Yue Wang, Akhilesh Deepak Gotmare, Silvio Savarese, and Steven Chu Hong Hoi. 2022.  Coderl: Mastering code generation through pretrained models and deep reinforcement learning.  _Advances in Neural Information Processing Systems_ , 35:21314–21328.
  * Li et al. (2023) Jingyao Li, Pengguang Chen, and Jiaya Jia. 2023.  Motcoder: Elevating large language models with modular of thought for challenging programming tasks.  _arXiv preprint arXiv:2312.15960_.
  * Li et al. (2022a) Yujia Li, David Choi, Junyoung Chung, Nate Kushman, Julian Schrittwieser, Rémi Leblond, Tom Eccles, James Keeling, Felix Gimeno, Agustin Dal Lago, et al. 2022a.  Competition-level code generation with alphacode.  _Science_ , 378(6624):1092–1097.
  * Li et al. (2022b) Yujia Li, David Choi, Junyoung Chung, Nate Kushman, Julian Schrittwieser, Rémi Leblond, Tom Eccles, James Keeling, Felix Gimeno, Agustin Dal Lago, Thomas Hubert, Peter Choy, Cyprien de Masson d’Autume, Igor Babuschkin, Xinyun Chen, Po-Sen Huang, Johannes Welbl, Sven Gowal, Alexey Cherepanov, James Molloy, Daniel Mankowitz, Esme Sutherland Robson, Pushmeet Kohli, Nando de Freitas, Koray Kavukcuoglu, and Oriol Vinyals. 2022b.  Competition-level code generation with alphacode.
  * Liu et al. (2023) Jiawei Liu, Chunqiu Steven Xia, Yuyao Wang, and Lingming Zhang. 2023.  Is your code generated by chatGPT really correct? rigorous evaluation of large language models for code generation.  In _Thirty-seventh Conference on Neural Information Processing Systems_.
  * Manna and Waldinger (1971) Zohar Manna and Richard J. Waldinger. 1971.  Toward automatic program synthesis.  _Commun. ACM_ , 14(3):151–165.
  * Nijkamp et al. (2022) Erik Nijkamp, Bo Pang, Hiroaki Hayashi, Lifu Tu, Huan Wang, Yingbo Zhou, Silvio Savarese, and Caiming Xiong. 2022.  Codegen: An open large language model for code with multi-turn program synthesis.  _arXiv preprint arXiv:2203.13474_.
  * Pacheco et al. (2007) Carlos Pacheco, Shuvendu K Lahiri, Michael D Ernst, and Thomas Ball. 2007.  Feedback-directed random test generation.  In _29th International Conference on Software Engineering (ICSE’07)_ , pages 75–84. IEEE.
  * Parisotto and Salakhutdinov (2017) Emilio Parisotto and Ruslan Salakhutdinov. 2017.  Neural map: Structured memory for deep reinforcement learning.  _arXiv preprint arXiv:1702.08360_.
  * Parvez et al. (2021) Md Rizwan Parvez, Wasi Uddin Ahmad, Saikat Chakraborty, Baishakhi Ray, and Kai-Wei Chang. 2021.  Retrieval augmented code generation and summarization.  _arXiv preprint arXiv:2108.11601_.
  * Parvez et al. (2018) Md Rizwan Parvez, Saikat Chakraborty, Baishakhi Ray, and Kai-Wei Chang. 2018.  Building language models for text with named entities.  In _Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_ , pages 2373–2383, Melbourne, Australia. Association for Computational Linguistics.
  * Parvez et al. (2023) Md Rizwan Parvez, Jianfeng Chi, Wasi Uddin Ahmad, Yuan Tian, and Kai-Wei Chang. 2023.  Retrieval enhanced data augmentation for question answering on privacy policies.  In _Proceedings of the 17th Conference of the European Chapter of the Association for Computational Linguistics_ , pages 201–210, Dubrovnik, Croatia. Association for Computational Linguistics.
  * Polozov and Gulwani (2015) Oleksandr Polozov and Sumit Gulwani. 2015.  Flashmeta: A framework for inductive program synthesis.  In _Proceedings of the 2015 ACM SIGPLAN International Conference on Object-Oriented Programming, Systems, Languages, and Applications_ , pages 107–126.
  * Rabinovich et al. (2017) Maxim Rabinovich, Mitchell Stern, and Dan Klein. 2017.  Abstract syntax networks for code generation and semantic parsing.  _CoRR_ , abs/1704.07535.
  * Ridnik et al. (2024) Tal Ridnik, Dedy Kredo, and Itamar Friedman. 2024.  Code generation with alphacodium: From prompt engineering to flow engineering.  _arXiv preprint arXiv:2401.08500_.
  * Roziere et al. (2023) Baptiste Roziere, Jonas Gehring, Fabian Gloeckle, Sten Sootla, Itai Gat, Xiaoqing Ellen Tan, Yossi Adi, Jingyu Liu, Tal Remez, Jérémy Rapin, et al. 2023.  Code llama: Open foundation models for code.  _arXiv preprint arXiv:2308.12950_.
  * Shinn et al. (2023) Noah Shinn, Federico Cassano, Ashwin Gopinath, Karthik R Narasimhan, and Shunyu Yao. 2023.  Reflexion: Language agents with verbal reinforcement learning.  In _Thirty-seventh Conference on Neural Information Processing Systems_.
  * Shum et al. (2023) Kashun Shum, Shizhe Diao, and Tong Zhang. 2023.  Automatic prompt augmentation and selection with chain-of-thought from labeled data.  In _Findings of the Association for Computational Linguistics: EMNLP 2023_ , pages 12113–12139, Singapore. Association for Computational Linguistics.
  * Touvron et al. (2023) Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. 2023.  Llama 2: Open foundation and fine-tuned chat models.  _arXiv preprint arXiv:2307.09288_.
  * Wang et al. (2021) Yue Wang, Weishi Wang, Shafiq Joty, and Steven CH Hoi. 2021.  Codet5: Identifier-aware unified pre-trained encoder-decoder models for code understanding and generation.  In _EMNLP_ , pages 8696–8708.
  * Wang et al. (2023) Zhiruo Wang, Jun Araki, Zhengbao Jiang, Md Rizwan Parvez, and Graham Neubig. 2023.  Learning to filter context for retrieval-augmented generation.  _arXiv preprint arXiv:2311.08377_.
  * Wei et al. (2022a) Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al. 2022a.  Chain-of-thought prompting elicits reasoning in large language models.  _Advances in Neural Information Processing Systems_ , 35:24824–24837.
  * Wei et al. (2022b) Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al. 2022b.  Chain-of-thought prompting elicits reasoning in large language models.  _Advances in Neural Information Processing Systems_ , 35:24824–24837.
  * Xu et al. (2023) Xiaohan Xu, Chongyang Tao, Tao Shen, Can Xu, Hongbo Xu, Guodong Long, and Jian guang Lou. 2023.  Re-reading improves reasoning in language models.
  * Yao et al. (2023) Shunyu Yao, Dian Yu, Jeffrey Zhao, Izhak Shafran, Thomas L Griffiths, Yuan Cao, and Karthik Narasimhan. 2023.  Tree of thoughts: Deliberate problem solving with large language models.  _arXiv preprint arXiv:2305.10601_.
  * Yasunaga et al. (2023) Michihiro Yasunaga, Xinyun Chen, Yujia Li, Panupong Pasupat, Jure Leskovec, Percy Liang, Ed H Chi, and Denny Zhou. 2023.  Large language models as analogical reasoners.  _arXiv preprint arXiv:2310.01714_.
  * Yin and Neubig (2017) Pengcheng Yin and Graham Neubig. 2017.  A syntactic neural model for general-purpose code generation.  _CoRR_ , abs/1704.01696.
  * Yu et al. (2019) Tao Yu, Rui Zhang, Heyang Er, Suyi Li, Eric Xue, Bo Pang, Xi Victoria Lin, Yi Chern Tan, Tianze Shi, Zihan Li, Youxuan Jiang, Michihiro Yasunaga, Sungrok Shim, Tao Chen, Alexander Fabbri, Zifan Li, Luyao Chen, Yuwen Zhang, Shreya Dixit, Vincent Zhang, Caiming Xiong, Richard Socher, Walter Lasecki, and Dragomir Radev. 2019.  CoSQL: A conversational text-to-SQL challenge towards cross-domain natural language interfaces to databases.  In _Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP)_ , pages 1962–1979, Hong Kong, China. Association for Computational Linguistics.
  * Zhang et al. (2023) Yifan Zhang, Jingqin Yang, Yang Yuan, and Andrew Chi-Chih Yao. 2023.  Cumulative reasoning with large language models.  _arXiv preprint arXiv:2308.04371_.
  * Zhang et al. (2022) Zhuosheng Zhang, Aston Zhang, Mu Li, and Alex Smola. 2022.  Automatic chain of thought prompting in large language models.  _arXiv preprint arXiv:2210.03493_.
  * Zhou et al. (2023) Andy Zhou, Kai Yan, Michal Shlapentokh-Rothman, Haohan Wang, and Yu-Xiong Wang. 2023.  Language agent tree search unifies reasoning acting and planning in language models.  _arXiv preprint arXiv:2310.04406_.

Appendix

##  Appendix A Algorithm of MapCoder

Algorithm 1 shows the pseudo-code of our prompting technique.

Algorithm 1 MapCoder

1:$k\leftarrow$ number of self-retrieved exemplars

2:$t\leftarrow$ number of debugging attempts

3:

4:$exemplars\leftarrow$ RetrivalAgent($k$)

5:

6:$plans\leftarrow$ empty array of size $k$

7:for $example$ in $exemplars$ do

8: $plans[i]\leftarrow$ PlanningAgent($example$)

9:end for

10:

11:$plans\leftarrow$ SortByConfidence($plans$)

12:

13:for $i\leftarrow 1$ to $k$ do

14: $code\leftarrow$ CodingAgent($code$, $plan[i]$)

15: $passed,log\leftarrow$ test($code$, $sample\\_io$)

16: if $passed$ then

17: Return $code$

18: else

19: for $j\leftarrow 1$ to $t$ do

20: $code\leftarrow$ DebuggingAgent($code$, $log$)

21: $passed,log\leftarrow$ test($code$, $sample\\_io$)

22: if $passed$ then

23: Return $code$

24: end if

25: end for

26: end if

27:end for

28:Return $code$

##  Appendix B Details Promptings of MapCoder

The detailed prompting of the Retrieval Agent, Planning Agent, Coding Agent,
and Debugging Agent are shown in Figure 8, 9, and 10 respectively. Note that
we adopt a specific sequence of instructions in the prompt for Retrieval Agent
which is a crucial design choice.

##  Appendix C Example Problem

Two complete examples of how MapCoder works by showing all the prompts and
responses for all four agents is given in this link.

Figure 8: Prompt for self-retrieval Agent.  Figure 9: Prompt for Planning
Agent. The example problems that are mentioned in this figure will come from
the Retrieval Agent. Figure 10:  Prompt for Coding and Debugging Agent.