Title: Guided Adaptive Credit Assignment for Sample Efficient Policy Optimization

Abstract: Policy gradient methods have achieved remarkable successes in solving challenging reinforcement learning problems. However, it still often suffers from sparse reward tasks, which leads to poor sample efficiency during training. In this work, we propose a guided adaptive credit assignment method to do effectively credit assignment for policy gradient methods. Motivated by entropy regularized policy optimization, our method extends the previous credit assignment methods by introducing a more general credit assignment named guided adaptive credit assignment(GACA). The benefit of GACA is a principled way of utilizing off-policy samples. The effectiveness of proposed algorithm is demonstrated on the challenging WikiTableQuestions and WikiSQL benchmarks and an instruction following environment. The task is generating action sequences or program sequences from natural language questions or instructions, where only final binary success-failure execution feedback is available. Empirical studies show that our method significantly improves the sample efficiency of the state-of-the-art policy optimization approaches.

1. Introduction:

Deep reinforcement learning (RL) provides a general framework for solving challenging goal-oriented sequential decision-making problems, it has recently achieved remarkable successes in advancing the frontier of AI technologies (Silver et al., 2016; Mnih & Kavukcuoglu, 2013; Silver et al., 2017; Andrychowicz et al., 2017) . Policy gradient (PG) (Kakade, 2002; Mnih et al., 2016; Schulman et al., 2015) is one of the most successful model-free RL approaches that has been widely applied to high dimensional continuous control, vision-based robotics, playing video games, and program synthesis (Liang et al., 2018; Guu et al., 2017; Bunel et al., 2018) . Despite these successes, a key problem of policy gradient methods is that it often suffers from high sample complexity in sparse reward tasks. In sparse reward tasks, there is only a binary signal which indicate successful task completion but without carefully shaped reward function to properly guide the policy optimization. A naive yet effective solution to address this challenge is by exploring many diverse samples and re-labelling visited states as goal states during training (see e.g. Andrychowicz et al., 2017; Pong et al., 2019) . Regardless of the cost of generating large samples and the bias introduced during comparison, in many practical applications like program synthesis, it may not even be possible to compare between different states. A variety of credit assignment techniques have been proposed for policy gradient methods in settings where comparison of states is not available (See e.g. Liang et al. 2018 , Agarwal et al. 2019 , and Norouzi et al. 2016) .

In this work, we focus on entropy regularized reinforcement learning. Instead of directly optimizing the RL objective, which is hard in sparse reward tasks, we sort to optimize policy to approximate a learnable prior distribution called guiding prior distribution. By using so-called f -divergence (Csiszár et al., 2004; Liese & Vajda, 2006; Nowozin et al., 2016; Wang et al., 2018) which defines a broad class of divergence(e.g., KL and reverse KL divergence) that are sufficient to fully characterize the distributions under consideration, we construct a class of gradient estimator that allow us to generalize previous credit assignment methods. The neat property is that the gradient estimator can adaptively optimize policy based on divergence between itself and the prior distribution. It is natural to expect this more flexible gradient estimator provide an adaptive trade-off between different credit assignment methods, in addition, it also has a good property such that all off-policy samples are utilized to compute gradient, which can yield powerful credit assignment. Our approach tremendously extends the existing credit assignment used including REINFORCE (Sutton et al., 2000; Williams, 1992) , maximum marginal likelihood(MML) (Dempster et al., 1977; Guu et al., 2017) , MAPO (Liang et al., 2018) , iterative maximum likelihood(IML) (Liang et al., 2017; Abolafia et al., 2018) , and RAML (Norouzi et al., 2016) .

We evaluate our method on a variety of tasks, including the challenging WikiSQL (Zhong et al., 2017) and WikiTableQuestions (Pasupat & Liang, 2015) program synthesis benchmarks, and an instruction following navigation task TextWorld (Agarwal et al., 2019) . Our experiments show that GACA greatly improves the sample efficiency of the entire policy optimization, and leads to significant higher asymptotic performance over previous state-of-the-art methods.

2. Background:



2.1. Reinforcement Learning and Policy Optimization:

Reinforcement learning(RL) considers the problem of finding an optimal policy for an agent that interacts with an uncertain environment and collects reward per action. The goal of the agent is to maximize its cumulative reward. Formally, this problem can be viewed as a Markov decision process over the environment states s ∈ S and agent actions z ∈ Z , with the environment dynamics defined by the transition probability T (s |s, z) and reward function r(s t , z t ), which yields a reward immediately following the action z t performed in state s t . The agent's action z is selected by a conditional probability distribution π(z|s) called policy.

In policy gradient methods, we consider a set of candidate policies π θ (z|s) parameterized by θ and obtain the optimal policy by maximizing the expected cumulative reward or return

J(θ) = E s∼ρπ,z∼π(z|s) [r(s, z)] ,

where ρ π (s) = E ∞ t=1 γ t-1 Pr(s t = s) is the normalized discounted state visitation distribution with discount factor γ ∈ [0, 1).

2.2. Sparse Reward Reinforcement Learning and Credit Assignment:

Auto-regressive model is often used as a policy in many real world applications including program synthesis and combinational optimization (Liang et al., 2018; Guu et al., 2017) . In this work, we consider the following form of policy distribution.

EQUATION

where z <t = (z 1 , . . . , z t-1 ) denotes a prefix of the action sequence z, s 0 ∈ Z denotes some context information about the task, such as initial state or goal state (Andrychowicz et al., 2017) . And π θ (z|s 0 ) satisfy ∀z ∈ Z : π θ (z|s 0 ) ≥ 0 and E z∈Z π θ (z|s 0 ) = 1.

In environments where dense reward function is not available, only a small fraction of the agents' experiences will be useful to compute gradient to optimize policy, leading to substantial high sample complexity. Therefore, it is of great practical importance to develop algorithms which can learn from binary signal indicating successful task completion or other unshaped reward signal. Nicky English Tipperary

Rank:

x = "Which player ranked the most?" R(z) = I{Execute(z) == "Nicky English"} Figure 1 : An example of the program synthesis task, where an agent is presented with a context s 0 consists of a natural language question and a table, and is asked to generate a program z = (z 1 , z 2 , .., z n ). The agent receives a reward of 1 if execution of z on the relevant data table leads to the correct answer g (e.g., "Nicky English").

In Section 3, we will describe a method to efficiently utilize high-reward and zero-reward trajectories to address this challenge. We will evaluate the method on program synthesis and instructions following navigation, both are particular sparse reward tasks. Figure 1 shows an example of sparse reward program synthesis. The model needs to discover the programs that can generate the correct answer in a given context and generalizes over unseen context.

We consider goal-conditioned reinforcement learning from sparse rewards. This constitutes a modification to the reward function such that it depends on a goal g ∈ G, such that r(z, g, s) : S × Z × G → R. Every episode starts with sampling a state-goal pair from some distribution p(s 0 , g). Unlike the state, the goal stays fixed for the whole episode. At every time step, an action is chosen according to some policy π, which is expressed as a function of the state and the goal, π : S × G → Z . Therefore, we apply the following sparse reward function:

EQUATION

where g is a goal and F (z) denotes evaluating action sequence z on the task that controls when the goal is considered completed. The objective is given by

EQUATION

where H is the length of the trajectory. We can calculate the gradient of Equation 4 with REIN-FORCE (Williams, 1992) and estimate it using Monte Carlo samples.

EQUATION

Unfortunately, since the search space of programs is very large, most samples z have reward R(z) = 0, thus have no contribution to the gradient estimation in Equation 5. Besides, because the variance of score function estimators is very high, it is challenging to estimate the gradient in Equation 5 with a small number of successful programs. Previous method Liang et al. (2018) propose to estimate gradient as a combination of expectations inside and outside successful programs buffer, however it's still restricted to use successful programs only, and suffers from high sample complexity.

3. Method:

In this section, we fist introduce entropy regularized reinforcement learning and describe optimizing policy via minimizing a discrepancy between itself and a prior in Section 3.1, and then introduce learnable prior to guide policy optimization in Section 3.2, finally we introduce a class of flexible adaptive gradient estimator Section 3.3.

3.1. Entropy Regularized Reinforcement Learning.:

We consider a general entropy regularized objective (Ziebart et al., 2008) which favors stochastic policies by augmenting the objective with the relative entropy of the policy,

EQUATION

where λ is a regularization weight, H(π θ (z|s 0 )) is the entropy regularization term. Entropy based policy optimization is a general framework that has gained many successes in a variety of tasks (see e.g., Haarnoja et al., 2018; Teh et al., 2017) . Maximizing Equation 6 is equivalent to minimizing the Kullback-Leibler discrepancy between policy π θ (z|s 0 ) and an energy based prior distribution.

Lemma 1. Maximizing Equation 6 is equivalent to minimizing the following objective,

EQUATION

where V (s 0 ) = λ log z∼Z exp(R(z, g, s 0 )/λ) is a 'soft-version' of value function, serving as a normalization constant here. From Equation 7, we aim to approximate the distribution π(z) with a distribution from a family {π θ (z|s 0 ) : θ ∈ Θ}, where θ is the parameter that we want to optimize, and π θ (z|s 0 ) is represented as an autoregressive policy in Equation 1. In environments where only sparse reward function is available, only a small fraction of the agent's samples will be useful to compute gradient to optimize policy, thus Equation 6 often leads to a substantial sample complexity. Equation 7 seems would be a better objective since all of the agent's samples can contribute to the minimization of KL-divergence, however, for a given s 0 , the prior distribution is simply a binary value function over z which is not suitable. Intuitively, we would like π(z) weighs higher on 'almost success' action sequences z and weighs lower on 'far from success' action sequences z.

3.2. Guiding Prior Distribution.:

In this part, we will describe how to learn the prior distribution π(z) to guide policy optimization.

Proposition 1. Given a policy π θ (z|s 0 ), new guiding prior distribution π(z) that minimizes the discrepancy in Equation 7 is given by,

EQUATION

and the minimization of Equation 7 equals to mutual information between s 0 and z:

EQUATION

Proof. See Appendix C for details.

Proposition 1 indicates that alternatively optimizing π θ (z|s 0 ) and π(z) leads to a complex mixture distribution of π(z), increasing the expressive power of prior for credit assignment. Since Equation 8minimizes D KL (π θ (z|s 0 ) π(z)) and leads to a mutual information between s 0 and z, therefore the entropy regularized objective becomes the following mutual information regularized objective,

EQUATION

Equation 10 draws connection with rate distortion theory (Shannon, 1959; Cover & Thomas, 2012) , intuitively, the policy π θ (z|s 0 ) is encouraged to discard reward-irrelevant information in context s 0 subject to a limited channel capacity given by I(s 0 ; z). In the next section, we will present a class of gradient estimator that can adaptively update policy distribution to approximate the guiding prior.

3.3. Adaptive Gradient Estimation.:

While D KL (π θ (z|s 0 ) || π(z)) is the typical divergence measure widely used in variational inference and reinforcement learning (see e.g. Wainwright et al., 2008; Abdolmaleki et al., 2018; Hoffman et al., 2013) , it often leads to model collapse because of its mode-seeking property. Therefore, directly optimizing Equation 7often gives a suboptimal model π θ (z|s 0 ). It is therefore natural to consider alternative divergence measures. We approach this problem by minimizing the general f -divergence (Ali & Silvey, 1966; Morimoto, 1963) between π(z) and π θ (z|s 0 ). f -divergence includes a large spectrum of divergences (e.g., KL and reverse KL divergence) and is shown to be powerful in various settings (Nowozin et al., 2016; Wang et al., 2018; Ghasemipour et al., 2019) ,

EQUATION

where f : R + → R is any twice-differentiable convex function. It can be shown by Jensen's inequality that D F (p || q) ≥ 0 for any p and q. Further, if

f (t) is strictly convex at t = 1, then D F (π(z) || π θ (z|s 0 )) = 0 implies π(z) = π θ (z|s 0 )

. We use stochastic optimization to minimizing Equation 11, then gradient of Equation 11 is given by: Lemma 2. Assume f is a differentiable convex function and log π θ (z|s 0 ) is differentiable w.r.t. θ.

For f-divergence defined in equation 11, we have

EQUATION

where

ρ f (t) = f (t)t -f (t).

Proof. See Appendix B for details or Wang et al. (2018) .

Equation 12 shows that the gradient of f -divergence between π θ (z|s 0 ) and π(z) can be specified through ρ f or f . In next section, we will describe how to adaptive choose ρ f or f based on the discrepancy between π θ (z|s 0 ) and π(z). Since the space of Z is enumerable and the environment is deterministic, the expectation over z ∼ π θ (z|s 0 ) can be efficiently computed through sampling in replay buffer. We proceed to describe how to estimate this gradient with samples.

3.4. Final Algorithm.:

Given Equation 12, it's natural to ask how to estimate the gradient, a naive way is simply store past trajectories in a replay buffer and sample random mini-batch from it to compute the gradient. However this approach suffers from the fact that a large fraction of sampled trajectories have zero-reward, which leads to high sample complexity. We propose to save high-reward trajectories and zero-reward trajectories into two separated replay buffers and estimate the following gradient:

Proposition 2. Given replay buffers B and C for saving high-reward and zero-reward trajectories, an unbiased and low variance estimation is given by,

∇ θ DF (π(z) || π θ (z|s 0 )) = w B E z∼π + θ (z|x) ρ f π θ (z|s 0 ) π(z) ∇ θ log π θ (z|s 0 ) + w C E z∼π - θ (z|x) ρ f π θ (z|s 0 ) π(z) ∇ θ log π θ (z|s 0 ) (13)

where w B and w B represent the total probability of trajectories in replay buffers B and C respectively, w B + w C = 1, and

π + θ (z | x) = π θ (z|s 0 )/w B if z ∈ B 0 if z ∈ C , π - θ (z | x) = 0 if z ∈ B π θ (z|s 0 )/w C if z ∈ C (14)

Proof. See Appendix D for details.

The gradient estimation uses high-reward trajectories thus π θ (z|s 0 ) will not forget them, the estimation also utilize zero-reward trajectories in the past, which improves sample efficiency. The corresponding framework is shown in Figure 2 . Note that different from MAPO where they also use a buffer to save successful programs, Equation 13differs in that all offpolicy samples can be used to estimate gradient, which leads to a higher sample efficiency. We follow Wang et al. (2018) in choosing f -divergence such that it achieve a trade-off between exploration and exploitation, specifically, let {z i } be drawn from buffers B and C and w i = π(z i |s 0 )/π(z i ), then we substitute ρ f (π θ (z|s 0 )/π(z)) with the inverse of approximate tail probability given by 1 n E n i=1 I(w i ≥ t). The benefit of doing this is policy distribution π θ (z|s 0 ) can adaptively coverage and approximate prior distribution π(z).

In practice, to overcome cold start problem in sparse reward policy optimization, we follow Liang et al. (2018) to clip w B to a given range such that w B = max(w B , w l ) and w B = min(w B , w u ) where w l ≤ w u . Note that Equation 13 generalizes previous work in credit assignment including REINFORCE, MML (Dempster et al., 1977; Berant et al., 2013; Guu et al., 2017) , MAPO (Liang et al., 2018) , RAML (Norouzi et al., 2016) , and IML (Liang et al., 2017; Abolafia et al., 2018) . It is natural to expect this more flexible gradient estimator provide an adaptive trade-off between different credit assignment methods and can yield powerful credit assignment. Due to page limit, we leave discussions and proofs around generalization in Appendix E. Combining Proportion 1 and Proportion 2 together, we summarize the main algorithm in Algorithm 1.

4. Experiment:

We first introduce the set up of experiments, then evaluate GACA on two sparse reward program synthesis benchmarks WikiTableQuestions and WikiSQL, and an instruction following sparse reward navigation task.

4.1. Experimental setup:

WikiTableQuestions (Pasupat & Liang, 2015) contains 2,108 tables and 18,496 question-answer pairs build from tables extracted from Wikipedia. WikiSQL (Zhong et al., 2017 ) is a recent Algorithm 1 Guided Adaptive Credit Assignment for Sample Efficient Policy Optimization Require: Training data p(s 0 , g), random initialized policy π θ (z|s 0 ), uniform initialized prior π(z), high-reward and zero-reward trajectory buffers B and C, and clipping thresholds w l and w u . repeat Sample initial states and goals {s 0 , g} from data distribution p(s 0 , g) Collect trajectories with π θ (z|s 0 ) given {s 0 , g} and push trajectories into replay buffers B and C according to their rewards. Draw {z i } from buffers B and C through stratified sampling, compute w B and w C Compute tail probability 1 n E n i=1 I(w i ≥ t), where w i = π(z i | x)/π(z i ) Update policy distribution π θ (z|s 0 ) with Equation 13by substituting ρ f (π θ (z|s 0 )/π(z)) with the inverse of tail probability Compute new guiding prior distribution π(z) = E s0,g∼p(s0,g) [π θ (z|s 0 )] until converge or early stop large scale dataset on learning natural language interfaces for databases. It contains 24,241 tables extracted from Wikipedia and 80,654 question-program pairs. It is annotated with programs (SQL). In both datasets, question-answers are split into train, validation, and test sets. We also evaluate GACA on an instruction following navigation environment TextWorld (Agarwal et al., 2019) . The task is an instruction following navigation in a maze of size N × N with K deadly traps distributed randomly over the maze. An agent is given need a language instruction which outlines an optimal path that the agent can take to reach the goal, the agent needs to generate a sequence of actions and the agent receives a reward of 1 if it succeeds in reaching the goal within a certain number of steps, otherwise 0. An example of this task is shown in Figure 5 . For details in experiments, refer to Appendix F.

4.2. Comparing GACA with baselines:

Firstly, we compare GACA with several baseline methods that are special cases of GACA, to show the effectiveness of guiding prior and adaptive gradient estimation. We briefly introduce each baseline method here and leave the detailed discussion and proof of generalization in Appendix E. REINFORCE: REINFORCE maximizes expected return, we use on-policy samples to estimate the gradient ∇ θ J RL = E s0,g∼p(s0,g) E z∼π θ (z|s0), [∇ θ log π θ (z|s 0 )r(z, s 0 , g)].

IML: Iterative maximize likelihood (Liang et al., 2017; Abolafia et al., 2018) uniformly maximizes the likelihood of all the high-reward trajectories in past experience, the gradient is given by ∇ θ J IML = E s0,g∼p(s0,g) z∈B ∇ θ log π θ (z|s 0 )r(z, s 0 , g). RAML: Reward Augmented Maximum Likelihood (Norouzi et al., 2016 ) is a more general variant of IML, which weights off-policy samples with an energy based prior distribution in Equation 7, J RAML = E s0,g∼p(s0,g) E z∼Z π(z) log π θ (z|s 0 )r(z, s 0 , g), where π(z) = exp 1 λ (r(z, s 0 , g) -V (x)) . MML: Maximize Marginal Likelihood (Dempster et al., 1977; Berant et al., 2013) maximizes the marginal probability of the replay buffer B. The gradient of J MML is given by ∇ θ J MML = E s0,g∼p(s0,g) z∈B π θ (z|s0) ẑ∈B π θ (ẑ|s0) r(z, s 0 , g)∇ θ log π θ (z|s 0 ). MAPO/MAPOX: Memory Augmented Policy Optimization (Liang et al., 2018 ) is a recent method for reusing high-reward trajectories, it maximizes the expected reward and estimate the gradient with off-policy high-reward trajectories:

∇ θ J MAPO = (1 - α)E s0,g∼p(s0,g) E z∼π(z|x) [∇ θ log π θ (z|s 0 )r(z, s 0 , g)] + αE z∼B [∇ θ log π θ (z|s 0 )r(z, s 0 , g)],

where α is a weight equals to the total probability of high-reward trajectory z in buffer B. MAPOX (Agarwal et al., 2019) We first compare GACA with its various special cases on sparse reward program synthesis benchmarks, results are shown in Figure 3 and Table 1 , the comparison shows that noticeably improves upon previous methods in terms of both sample efficiency and asymptotic performance significantly by performing better credit assignment. The ablation study shows that both adaptive gradient estimation(AG) and guiding prior(GP) improve performance. The improvement over MAPO demonstrates that reusing all off-policy trajectories can greatly improve policy, the improvement over IML and MAPOX shows that encouraging exploration with reverse KL-divergence is not enough because it's simply impossible to explore such a large state space, while guiding prior and adaptive gradient estimation provide an efficient way of exploration and exploitation. We also analyzed a trained model qualitatively on program synthesis tasks and see that it can generate fairly complex programs, see Appendix G for some examples of generated programs. Our experiments follow the settings of MAPO (Liang et al., 2018) and MeRL (Agarwal et al., 2019) , refer to Appendix F for more details in experiments.

4.3. Comparing GACA with state-of-the-art:

We present the results on sparse reward program synthesis in Table 3 and Table 2 

5. Related Work:

Credit assignment is a critical part of various sequential decision making methods. Guu et al. (2017) builds connection between REINFORCE and MML by proposing hybrid approaches to take advantages of both MML and REINFORCE. Entropy based policy optimization is widely used in reinforcement learning (Ziebart et al., 2008; Schulman et al., 2017) , recently entropy based off-policy policy optimization is also proposed to approximate optimal policy distribution by minimizing the Kullback-Leibler(KL) divergence between policy and optimal distribution (Haarnoja et al., 2018) , Norouzi et al. (2016) considers an alternative direction of the KL divergence, where samples from exponential payoff distribution are used to estimate gradient. Recent work Grau-Moya et al. (2019) also propose to learn the prior distribution for Q-learning and show that this leads to a mutual information regularization. Experience replay is widely used in sparse reward reinforcement learning in order to exploit past high reward trajectories (Gangwani et al., 2019; Liang et al., 2018; Oh et al., 2018; Abolafia et al., 2018) . Andrychowicz et al. (2017) proposes to re-label visited states as goal states during training. More recent progress includes meta-learning the reward(such as discount factor) (Xu et al., 2018) . Weber et al. 2019 provides a comprehensive review of credit assignment methods in stochastic computation graph. Recently, there are a surge of interest in applying policy optimization in program synthesis through sparse supervision (Krishnamurthy et al., 2017; Guu et al., 2017; Liang et al., 2017; 2018; Agarwal et al., 2019) . GACA differs from previous methods by enabling reusing off-policy samples through learned prior and generalized gradient estimation.

6. Conclusion:

We developed the Guided Adaptive Credit Assignment(GACA), a new and general credit assignment method for obtaining sample efficiency of policy optimization in sparse reward setting. Our method generalizes several previous approaches. We demonstrated its practical advantages over existing methods, including MML, IML, REINFORCE, etc, on several challenging sparse reward tasks. In the future, we will investigate how to extend GACA to stochastic environments and apply it to robot learning from binary reward feedback. We would also like to point out that our method can be useful in other challenging tasks with deterministic environments such as combinational optimization and structural prediction where credit assignment from binary feedback remains a major challenge.

A Proof of Lemma 1:

Proof. To derive Lemma 1, consider the KL divergence between π θ (z|s 0 ) and π(z) = exp 1 λ (r(z, g, s 0 ) -V (s 0 )) , where V (s 0 ) = λ log z∼Z exp(r(z, g, s 0 )/λ) is a 'soft-version' of value function, serving as a normalization constant here.

D KL (π θ (z|s 0 ) π(z)) = E z∼π θ (z|s0) [log π θ (z|s 0 ) -log π(z)] = E z∼π θ (z|s0) log π θ (z|s 0 ) -r(z, g, s 0 )/λ + log V (s 0 ) = E z∼π θ (z|s0) [log π θ (z|s 0 ) -r(z, g, s 0 )/λ] + log V (s 0 ),

Rearranging,

E z∼π θ (z|s0) [r(z, g, s 0 )] + λH(π θ (z|s 0 )) = -λD KL (π θ (z|s 0 ) π(z)) + λ log V (s 0 ), thus maximizing left hand side E z∼π θ (z|s0) [r(z, g, s 0 )] + λH(π θ (z|s 0 )) is equivalent to minimizing D KL (π θ (z|s 0 ) π(z)).

B Proof of Lemma 2:

Proof. To derive Lemma 2, consider that ∇ θ π θ (z|s 0 ) = π θ (z|s 0 )∇ θ log π θ (z|s 0 ), then we have

∇ θ D f (π(z) || π θ (z|s 0 )) = E π θ (z|s0) ∇ θ f π(z) π θ (z|s 0 ) + f π(z) π θ (z|s 0 ) ∇ θ log π θ (z|s 0 ) = E π θ (z|s0) f π(z) π θ (z|s 0 ) ∇ θ π(z) π θ (z|s 0 ) + f π(z) π θ (z|s 0 ) ∇ θ log π θ (z|s 0 ) = E π θ (z|s0) -f π(z) π θ (z|s 0 ) π(z) π θ (z|s 0 ) ∇ θ log π θ (z|s 0 ) + f π(z) π θ (z|s 0 ) ∇ θ log π θ (z|s 0 ) = -E π θ (z|s0) ρ f π(z) π θ (z|s 0 ) log π θ (z|s 0 ) , where ρ f (t) = f (t)t -f (t).

For convex function f , we have

f (t) ≥ 0, which implies ρ f (t) = f (t)t ≥ 0 on t ∈ R + , thus ρ f is a monotonically increasing function on R + . If ρ t is strictly increasing at t = 1, we have f is strictly convex at t = 1, which guarantees D F (p || q) = 0 imply p = q.

C Proof of Proposition 1

Let p(s 0 ) and p(z) denote the distribution of s 0 , and z respectively, for notation simplicity, we omit g in the following derivation and simply use s 0 ∈ p(s 0 ) to represent s 0 , g ∈ p(s 0 , g), and denote

π(z) = E s0∼p(s0) [π θ (z|s 0 )], then we have D KL (p(s 0 )π θ (z|s 0 ) || p(s 0 )p(z)) -D KL (p(s 0 )π θ (z|s 0 ) || p(s 0 )π(z)) = E s0∼p(s0) E z∼Z [p(s 0 )π θ (z|s 0 ) log p(s 0 )π θ (z|s 0 ) p(s 0 )p(z) ] -E s0∼p(s0) E z∼Z [p(s 0 )π θ (z|s 0 ) log p(s 0 )p(z|s 0 ) p(s 0 )π(z) ] = E s0∼p(s0) E z∼Z [p(s 0 )π θ (z|s 0 ) log π(z) p(z) ] = E s0∼p(s0) E z∼Z [π(z) log π(z) p(z) ] = E s0∼p(s0) [D KL (π(z) || p(z))] ≥ 0, thus π(z) = E s0∼p(s0) [π θ (z|s 0 )] = arg min p(z) D KL (p(s 0 )π θ (z|s 0 ) || p(s 0 )p(z)). Substituting π(z) in E s0∼p(s0) [D KL (p(s 0 )π θ (z|s 0 ) || p(s 0 )p(z))], we have D KL (p(s 0 )π θ (z|s 0 ) || p(s 0 )π(z)) = E s0∼p(s0) [p(s 0 , z) log p(s 0 , z) p(s 0 )p(z) ] = I(s 0 ; z) Thus E s0∼p(s0) [π θ (z|s 0 )

] is the solution of the minimization objective, and D KL (p(s 0 )π θ (z|s 0 ) || p(s 0 )π(z)) equals mutual information between state and action.

D Proof of Proposition 2:

Proof. To prove Equation 13 is an unbiased estimation of Equation 12, note that we can either enumerate replay buffers B and C when the size of buffers are small or approximate sampling from both buffers according to the specified ratio. In any case, this gives us a stratified sampling estimator of Equation 12, which is unbiased and low variance.

E Proof of generalization of previous credit assignment methods:

In this section, we discuss the connection between GACA and each credit assignment method, we will show that GACA is a unified form of existing credit assignment method. Firstly, we summarize existing method in Table 4 . Then we describe each method and give a proof of how to reduce GACA to it.

E.1 REINFORCE::

REINFORCE maximizes the expected reward and estimate the gradient with on-policy samples J RL = E s0,g∼p(s0,g) E z∼π θ (z|s0) r(z, s 0 , g), the gradient of REINFORCE objective is given by ∇ θ J RL = E s0,g∼p(s0,g) E z∼π θ (z|s0) ∇ θ log π θ (z|s 0 )r(z, s 0 , g). Apart from high variance issue in REINFORCE, it also suffers from sparse reward because reward r(z, s 0 , g) is low for most trajectories z. In contrast, GACA utilizes off-policy samples and still maintain unbiased gradient estimate. GACA reduces to REINFORCE by simply choosing ρ f as constant 1.

Optimization gradient

REINFORCE ∇ θ J RL = E s0,g∼p(s0,g) E z∼π θ (z|s0) [∇ θ log π θ (z|s 0 )r(z, s 0 , g)] MML ∇ θ J MML = E s0,g∼p(s0,g) z∈B π θ (z|s 0 ) ẑ∈B π θ (ẑ|s 0 )

r(z, s 0 , g)∇ θ log π θ (z|s 0 ) IML ∇ θ J IML = E s0,g∼p(s0,g) z∈B ∇ θ log π θ (z|s 0 )r(z, s 0 , g) 4 : The optimization gradient for various credit assignment algorithms. GACA generalizes existing methods by f -divergence based adaptive gradient estimation.

MAPO, MAPOX:

∇ θ J MAPO = E s0,g∼p(s0,g) [(1 -α)E z∼π θ (z|s0) [∇ θ log π θ (z|s 0 )r(z, s 0 , g)] + α z∼B [∇ θ log π θ (z|s 0 )r(z, s 0 , g)]], α ∈ [0, 1] RAML ∇ θ J RAML = E s0,g∼p(s0,g) E z∼Z [log π θ (z|s 0 )r(z, s 0 , g)π(z)] GACA ∇ θ DF (π(z) || π θ (z|s 0 )) = w B E z∼π + θ (z|x) ρ f (π θ (z|s 0 )/π(z)) ∇ θ log π θ (z|s 0 ) + w C E z∼π - θ (z|x) ρ f (π θ (z|s 0 )/π(z)) ∇ θ log π θ (z|s 0 ) Table

E.2 MML::

Maximize Marginal Likelihood(MML) (Dempster et al., 1977; Berant et al., 2013) maximizes the marginal probability of the replay buffer B, the objective of MML is given by J MML = E s0,g∼p(s0,g) log z∈B π θ (z|s 0 )r(z, s 0 , g). The gradient of J MML has the form:

EQUATION

Taking a step in the direction of J MML up-weights the probability of high-reward trajectory z and thus attempts to up-weight each reward-earning trajectory. More discussion of this objective can be found in (Guu et al., 2017; Liang et al., 2018) .

Choosing w l = 1 in Equation 13, and clearly there exists a monotonically increasing function ρ f satisfy ρ f ( π θ (z|s0) π(z) ) = π θ (z|s0)

ẑ∈B π θ (ẑ|s0) . Choosing such ρ f , GACA reduces to MML.

E.3 IML::

Iterative maximize likelihood(IML) (Liang et al., 2017; Abolafia et al., 2018) uniformly maximizes the likelihood of all the high-reward trajectories in past experience. The objective is given by

J IML = E s0,g∼p(s0,g) E z∼B [log π θ (z|s 0 )r(z, s 0 , g)].

The gradient of IML is given by

EQUATION

Choosing ρ f = 1 and w B = 1 in Equation 13, GACA reduces to IML. For each given s 0 , g, IML can be expressed as optimizing policy distribution by minimizing the reverse KL divergence between the parameterized policy distribution and an optimal policy distribution, i.e., D KL (π ||π) where π is the optimal distribution. It's well-known that reverse KL-divergence promotes model-covering, thus IML seeks to explore diverse samples thus will have a higher chance of collecting high-reward trajectories. Recent work MAPOX (Agarwal et al., 2019) exploits this property of IML by running IML to collect diverse samples for training.

E.4 MAPO, MAPOX::

Memory Augmented Policy Optimization(MAPO) (Liang et al., 2018 ) is a recent method for reusing high-reward trajectories, it maximizes the expected reward and estimate the gradient with off-policy high-reward trajectories. The gradient of MAPO is

EQUATION

)

where α is a weight equals to the total probability of high-reward trajectory z in buffer B.

MAPOX (Agarwal et al., 2019) improves MAPO by running MAPO on trajectories collected with IML for exploration. As shown previously, IML can be viewed as minimizing 'reverse' KL divergence between policy distribution π θ (z|s 0 ) and the prior distribution, thus IML promotes exploration.

When choosing ρ f ( π(z) π θ (z|s0) ) = log( π(z) π θ (z|s0) ) -1 and setting w B = 1, GACA reduces to MAPO.

E.5 RAML::

Reward Augmented Maximum Likelihood(RAML) (Norouzi et al., 2016 ) is a more general variant of IML, which weights off-policy samples with an energy based prior,

EQUATION

where π(z) = exp 1 λ (r(z, s 0 , g) -V (x)) is the energy based prior distribution defined in Equation 7. Similar to IML, for each given s 0 , g, RAML can be expressed as optimizing the policy distribution by minimizing the KL divergence between the parameterized policy distribution and an energy based optimal policy distribution defined as exp 1 λ (r(z, s 0 , g) -V (x)) . The gradient estimation is over possible trajectories z ∼ Z, only few of them are high-reward and most trajectories cannot guide the policy learn good behavior, thus RAML suffers from high sample complexity. Choosing ρ f ( π(z) π θ (z|s0) ) = π(z) π θ (z|s0) and w B = 1 in Equation 13, GACA reduces to RAML.

F Experiments Details:

For WikiTableQuestions, we follow the construction in Pasupat & Liang (2015) for converting a table into a directed graph that can be queried. The rows and cells are converted to graph nodes while column names become labeled directed edges. Each batch includes samples from 25 examples For WikiSQL, we follow the setting in Liang et al. (2018) for choosing the sampling batch size. Our model use a seq2seq model as π θ (z|s 0 ), and two key-variable memory as high-reward buffer B and zero-reward buffer C, and associated with a domain specific language interpreter (Liang et al., 2017) .

Table 1 shows a comparison GACA with various baselines on two challenging sparse reward program synthesis tasks, where we also present an ablation study of each technique in GACA. Specifically, we studied the performance of GACA w/o AG which represents GACA without adaptive gradient estimation(Section 3.3), and GACA w/o GP which represents GACA without guiding prior(Section 3.2).

In detail, GACA w/o AG means use standard KL-divergence to calculate gradient in Eq 13 instead of f-divergence, while GACA w/o GP means don't learn the prior policy distribution as in 8, but fix prior policy distribution to the energy based distribution as in 7. Our code is based on open source implementation of MAPO (Liang et al., 2018) which implements a distributed actor-learner architecture Espeholt et al. (2018) to accelerate sampling through distributed actors. We also use open source code from MeRL (Agarwal et al., 2019) , tail-adapted variational inference (Wang et al., 2018) .

Our experiments follow the settings of MAPO (Liang et al., 2018) and MeRL (Agarwal et al., 2019) . An agent is presented with a sequence of (Left, Right, Up, Down) instructions. Given the input text, the agent on the blue dot need to perform a sequence of actions, and only receives a reward of 1 if it reaches the goal at the orange star.

We port their code to PyTorch (Paszke et al., 2017) and implement GACA on top of them to conduct experiments. We will release the code later. Gradients are estimated and periodically updated through a central learner (Espeholt et al., 2018) . For TextWorld1 , we use a set of 300 randomly generated environments with training and validation splits of 80% and 20% respectively following Agarwal et al. (2019) . The agent is evaluated on 300 unseen test environments from the same distribution. An example of TextWorld is shown in Figure 5 . We used the Adam Optimizer (Kingma & Ba, 2015) for WikiSQL, WikiTABLE, and TextWorld. We performed hyper-parameter sweeps via random search over the interval 10 -4 , 10 -2 for learning rate. All the hyperparameters are tuned on the evaluation set.

G Qualitative Results:

In order to evaluate the qualitative quality of the proposed method, we compare GACA with the recent state-of-the-art MAPO on WIKITABLE-QUESTIONS. Figure 5 shows examples of generated programs from natural language queries using model trained with GACA or MAPO, the difference between generated programs show that sometimes GACA is capable of generating correct programs that capture the meaning of the natural language queries while MAPO generates either wrong answer programs or spurious programs. Comment Query nu-1147: Which opponent of the kansas city chiefs in 1987 saw a total of more than 70,000 in attendance? MAPO: r0 = (argmax all_rows r.attendance-number); rres = (hop r0 r.opponent-string) GACA: r0 = (filter> all_rows [70,000] r.attendancenumber); rres = (hop r0 r.opponent-string)

The program generated by MAPO leads to wrong answer, because it wrongly assume that the one with largest number of attendance number is the opponent of kansas city. Instead GACA captures the semantic meaning of the query and generates correct program.

Query nu-1167: Who was the first oldest living president? MAPO: r0 = (first all_rows); rans = (hop r0 r.president) GACA: r0 = (filter_str_contain_any all_rows [oldest living president'] r.became_oldest_living_president-string ); rres =( hop r0 r.president-string )

MAPO gets correct answer by chance because it is spurious program, MAPO wrongly assumes the data is ordered, in contrast, GACA generates programs that capture underlying semantic information.

Query nu-3733: At least how many more people attended gamestorm 15 than gamestrom 13? MAPO: r0 = (filter_str_contain_any all_rows [13] r.datesstring); rres = ( hop r0 r.attendance-number ) GACA: r0 = (filter_str_contain_any all_rows [13] r.datesstring); r1 = (filter_str_contain_any all_rows [13] r.datesstring); (diff r0 r1 r.attendance-number)

MAPO gets correct answer by chance because it is spurious program, MAPO wrongly assumes the data is ordered, in contrast, GACA generates programs that capture underlying semantic information.

Table 5 : Example of generated programs from models trained using MAPO and GACA on Wik-iTableQuestions.

:

https://github.com/google-research/google-research/tree/master/meta_reward_learning/ textworld

Figures and Tables:

Figure 2: Overview of GACA: it consists of multiple actors for sampling and storing high reward episodes into buffer B and zero reward episodes into buffer C, gradient is estimated at central learner periodically using samples from both B, C based on Equation 13.

Figure 3: Comparing GACA and baselines on benchmarks. The plot is average of 5 runs with a bar of one standard deviation. Examples of generated programs are shown in Figure 5. The policy needs to generate correct programs from binary feedback, making these two benchmarks very challenging.We also evaluate GACA on an instruction following navigation environment TextWorld(Agarwal et al., 2019). The task is an instruction following navigation in a maze of size N × N with K deadly traps distributed randomly over the maze. An agent is given need a language instruction which outlines an optimal path that the agent can take to reach the goal, the agent needs to generate a sequence of actions and the agent receives a reward of 1 if it succeeds in reaching the goal within a certain number of steps, otherwise 0. An example of this task is shown in Figure5. For details in experiments, refer to Appendix F.

Figure 5: Instruction following navigation in maze.An agent is presented with a sequence of (Left, Right, Up, Down) instructions. Given the input text, the agent on the blue dot need to perform a sequence of actions, and only receives a reward of 1 if it reaches the goal at the orange star.

improves MAPO by running MAPO on data collected by IML. Comparison to various credit assignment methods on WikiTableQuestions and WikiSQL benchmarks. We report mean accuracy, and its standard deviation on test sets based on 5 runs.

Table content: <table><tr><td>Method</td><td>Dataset</td><td>Val</td><td colspan="2">WikiTableQuestions Test Improvement</td><td>Val</td><td>WikiSQL Test</td><td>Improvement</td></tr><tr><td colspan="2">REINFORCE</td><td>&lt; 10</td><td>&lt; 10</td><td/><td>&lt; 10</td><td>&lt; 10</td></tr><tr><td>IML</td><td/><td colspan="2">35.4 ± 0.6 36.8 ± 0.4</td><td>+7.8</td><td colspan="2">69.3 ± 0.5 70.1 ± 0.2</td><td>+6.2</td></tr><tr><td>RAML</td><td/><td colspan="2">35.4 ± 0.7 35.9 ± 0.6</td><td>+8.7</td><td colspan="2">57.5 ± 0.3 61.4 ± 0.3</td><td>+14.9</td></tr><tr><td>MML</td><td/><td colspan="2">37.8 ± 0.7 39.7 ± 0.3</td><td>+4.9</td><td colspan="2">68.7 ± 0.2 70.7 ± 0.1</td><td>+5.6</td></tr><tr><td>MAPO</td><td/><td colspan="2">42.3 ± 0.1 42.8 ± 0.3</td><td>+1.8</td><td colspan="2">71.9 ± 0.6 72.5 ± 0.1</td><td>+3.8</td></tr><tr><td>MAPOX</td><td/><td colspan="2">42.5 ± 0.4 43.6 ± 0.4</td><td>+1.1</td><td colspan="2">74.4 ± 0.5 74.9 ± 0.3</td><td>+1.4</td></tr><tr><td colspan="2">GACA w/o GP</td><td colspan="2">43.8 ± 0.4 44.2 ± 0.2</td><td/><td colspan="2">74.6 ± 0.3 75.3 ± 0.2</td></tr><tr><td colspan="2">GACA w/o AG</td><td colspan="2">44.1 ± 0.5 44.1 ± 0.3</td><td/><td colspan="2">74.9 ± 0.3 75.2 ± 0.1</td></tr><tr><td>GACA</td><td/><td colspan="2">45.1 ± 0.3 44.6 ± 0.2</td><td/><td colspan="2">75.7 ± 0.3 76.3 ± 0.1</td></tr><tr><td>Method</td><td>Val.</td><td>Test</td><td/><td/><td/><td/></tr><tr><td colspan="4">Oracle 95.7(±1.3) 92.6(±1.0)</td><td/><td/><td/></tr><tr><td colspan="4">MAPO 73.1(±2.1) 68.5(±2.6)</td><td/><td/><td/></tr><tr><td colspan="4">MeRL 75.3(±1.6) 72.3(±2.2)</td><td/><td/><td/></tr><tr><td colspan="4">BoRL 83.0(±3.6) 74.5(±2.5)</td><td/><td/><td/></tr><tr><td colspan="4">GACA 87.3(±4.1) 80.1(±2.8)</td><td/><td/><td/></tr><tr><td colspan="4">Figure 4: Evaluation of GACA on</td><td/><td/><td/></tr><tr><td colspan="4">TextWorld with state-of-the-art.</td><td/><td/><td/></tr></table>

. The results of TextWorld are shown in Table4. GACA outperforms most recent state-of-the-art methods BoRL and MeRL proposed inAgarwal et al. (2019) by a large margin. The results demonstrate the efficacy of the proposed credit assignment compared to previous SOTA credit assignment methods. We would like to point out that GACA is a general method and can be combined with these techniques to further boost performance. Evaluation of GACA with state-of-the-art on WikiTableQuestions.

Table content: <table><tr><td>Method</td><td colspan="2">E.S. Val. Test</td><td>Method</td><td>E.S. Val. Test</td></tr><tr><td>MAPO</td><td colspan="2">1 42.3 42.8 +5.8</td><td>MAPO</td><td>1 71.9 72.5 +6.3</td></tr><tr><td>MeRL</td><td colspan="2">1 44.1 43.2 +5.4</td><td>MeRL</td><td>1 74.9 74.8 +4.0</td></tr><tr><td>BoRL</td><td colspan="2">1 42.9 43.8 +4.8</td><td>BoRL</td><td>1 74.6 74.2 +4.6</td></tr><tr><td colspan="2">MAPO (ensemble) 10</td><td>-46.3 +2.3</td><td colspan="2">MAPO (ensemble) 10</td><td>-74.9 +3.9</td></tr><tr><td colspan="2">MeRL (ensemble) 10</td><td>-46.9 +1.7</td><td colspan="2">MeRL (ensemble) 10</td><td>-77.1 +1.7</td></tr><tr><td>GACA</td><td colspan="2">1 45.1 44.6</td><td>GACA</td><td>1 75.7 76.3</td></tr><tr><td colspan="2">GACA (ensemble) 10</td><td>-48.6</td><td colspan="2">GACA (ensemble) 10</td><td>-78.8</td></tr></table>

Evaluation of GACA with state-of-the-art on WikiSQL.

Table content: <table/>

Peer review (iclr2020):
The authors formulate the credit assignment method as minimizing the divergence between policy function and a learned prior distribution. Then they apply f-divergence optimization to avoid the model collapse in this framework. Empirical experiments are conducted on the program synthesis benchmark with sparse rewards. The main contribution of this paper is applying f-divergence optimization on the program synthesis task for credit assignment. + One of my concerns is that the experiment section is in a limited domain to argue it is a broad algorithm for credit assignment. The paper will be stronger if the comparison is applied in a distant domain like goal-based robot learning etc. With some experiments on a different domain, the paper will be more convincing. + The improvement/margin in program synthesis task needed to be explained well, is the margin significant enough? + The paper could discuss more on related papers on program synthesis in the related work section as the main experiment is in this work. + The authors claim that the two-buffer estimation is better and lead to better gradient estimation, but it is not demonstrated empirically or theoretically. It could be better if the ablation study is conducted in the experiment. Or the author could provide a theoretical analysis of why equation (13) is better. Moreover, the investigation of different choices of and is necessary. + Another study needed is the investigation of different divergences; the work will be stronger if a KL divergence version is compared. Otherwise, it is not clear how much the f-divergence will contribute to the performance.