Title: POLICY OPTIMIZATION BY LOCAL IMPROVE-MENT THROUGH SEARCH

Abstract: Imitation learning has emerged as a powerful strategy for learning initial policies that can be refined with reinforcement learning techniques. Most strategies in imitation learning, however, rely on per-step supervision either from expert demonstrations, referred to as behavioral cloning (Pomerleau, 1989; 1991) or from interactive expert policy queries such as DAgger (Ross et al., 2011) . These strategies differ on the state distribution at which the expert actions are collected -the former using the state distribution of the expert, the latter using the state distribution of the policy being trained. However, the learning signal in both cases arises from the expert actions. On the other end of the spectrum, approaches rooted in Policy Iteration, such as Dual Policy Iteration (Sun et al., 2018b) do not choose next step actions based on an expert, but instead use planning or search over the policy to choose an action distribution to train towards. However, this can be computationally expensive, and can also end up training the policy on a state distribution that is far from the current policy's induced distribution. In this paper, we propose an algorithm that finds a middle ground by using Monte Carlo Tree Search (MCTS) (Kocsis & Szepesvári, 2006) to perform local trajectory improvement over rollouts from the policy. We provide theoretical justification for both the proposed local trajectory search algorithm and for our use of MCTS as a local policy improvement operator. We also show empirically that our method (Policy Optimization by Local Improvement through Search or POLISh) is much faster than methods that plan globally, speeding up training by a factor of up to 14 in wall clock time. Furthermore, the resulting policy outperforms strong baselines in both reinforcement learning and imitation learning.

1. INTRODUCTION:

Reinforcement learning (RL) has seen a great deal of success in recent years, from playing games (Mnih et al., 2015; Silver et al., 2016) to robotic control (Gu et al., 2017; Singh et al., 2019) . These successes showcase the power of learning from direct interactions with environments. However, a well-known disadvantage of reinforcement learning approaches is the demand for a large number of samples for learning; for example, OpenAI Five learned to solve DOTA using Proximal Policy Optimization (PPO) (Schulman et al., 2017) by playing 180 years worth of games against itself daily (OpenAI, 2018) . The issue with sample complexity can be mitigated by learning from expert behavior. For example, AlphaStar (The Alpha Star Team, 2019) uses models that are pre-trained using supervised learning techniques on expert human demonstrations, before RL is used for refinement.

A variety of strategies have been developed in imitation learning to learn from expert behavior, where the expert can be a human (Stadie et al., 2017) or a pre-trained policy (Ho & Ermon, 2016) . Earlier approaches, such as behavioral cloning (BC), relied on training models to mimic expert behavior at various states in the demonstration data (Pomerleau, 1991) . However, these models suffer from the problem that even small differences between the learned policy and the expert behavior can lead to a snow-balling effect, where the state distribution diverges to a place where the behavior of the policy is now meaningless since it was not trained in that part of space (Ross & Bagnell, 2010) . In order to mitigate these issues, DAgger (Ross et al., 2011) uses an expert policy to provide supervision to the policy being learned. While these strategies differ on the state distribution at which the expert actions are optimized -for example BC uses the state distribution of the expert, DAgger-like approaches use the state distribution from the policy being trained -the supervision signal itself is the expert's behavior at each step.

On the other end of the spectrum, approaches rooted in policy iteration, such as Dual Policy Iteration (Sun et al., 2018b) do not mimic next step actions of a policy directly, but instead use planning or search over the policy to choose an action distribution to train towards (Silver et al., 2017) . However, this can be computationally expensive, and can also end up training the policy on a state distribution that is far from the current policy's induced distribution.

In this paper, we propose an algorithm that finds a middle ground by using Monte Carlo Tree Search (MCTS) (Kocsis & Szepesvári, 2006) to perform local trajectory improvement over states sampled from different time steps in the trajectory unrolled from the policy. This approach has the benefit of training the policy on a state distribution that is close to that induced by the policy itself, while using local search or planning over a smaller horizon to generate a good policy to train towards. This approach stands in contrast to other works in interactive imitation learning that correct distribution mismatch by using one-step feedback. We provide theoretical justification for the advantage of a balanced local trajectory improvement and show that MCTS can serve as a policy improvement operator under certain conditions1 .

An added benefit of our effort is computational -depth parallel MCTS on local trajectory segments is much faster than traditional sequential MCTS to generate demonstrations. We show that our proposed algorithm can easily be parallelized to enable more efficient imitation learning from MCTS. Notably, this level of parallelism is present at varying depths which differs from existing works on parallel MCTS (Chaslot et al., 2008) .

In summary, our main contributions are:

• A general interactive imitation learning algorithm that balances the expert feedback quality and the state distribution divergence. Our proposed approach provides a flexible local trajectory improvement based on MCTS.

• Theoretical analysis on the general benefit of local improvement and specific case study on using MCTS as a local improvement operator.

• Strong empirical performance on a suite of high-dimensional continuous control problems based on both sample efficiency and training time.

2. RELATED WORK:

Imitation Learning. Imitation learning (IL) refers to the problem of learning to perform a task from expert demonstrations. Behavioral cloning (Widrow & Smith, 1964) is one popular approach which maximizes the likelihood of expert actions under the agent policy (Pomerleau, 1989; Schaal, 1999; Muller et al., 2006; Mlling et al., 2013; Bojarski et al., 2016; Giusti et al., 2016; Mahler & Goldberg, 2017; Wang et al., 2019; Bansal et al., 2018) . Inverse Reinforcement Learning is another popular form of IL where a reward function is extracted from expert demonstrations and then a policy is trained to maximize that reward (Ziebart et al., 2008; Finn et al., 2016; Fu et al., 2018; Ho & Ermon, 2016) .

In this work we focus on imitation learning via behavioral cloning. Despite its success in canonical problems, such as Go (Silver et al., 2017) and Starcraft (The Alpha Star Team, 2019), behavioral cloning suffers from many challenges, most notably distributional shift (Daumé et al., 2009) .

To explain distributional shift, let us assume we train an agent by performing supervised learning on the actions an expert has taken. A small error in the first time step of training may bring the agent to a state that the expert has rarely visited and is therefore not as well modeled. Over time, this error compounds, leading the agent to states far from expert behavior and diverging further from the expert demonstrations. The longer the episode is, the more likely the agent is to deviate from expert demonstrations.

Various solutions to the problem of distributional shift have been proposed (Daumé et al., 2009; Ross & Bagnell, 2010; Ross et al., 2011; Ho & Ermon, 2016; Laskey et al., 2017; Bansal et al., 2018; de Haan et al., 2019) . Some of these approaches reduce the effects of distributional shift by iteratively querying the expert (Daumé et al., 2009; Ross et al., 2011) . DAgger (Ross et al., 2011) , one of the most widely used of these solutions, queries the expert on each of the states that are visited by the policy and uses expert demonstrations to improve the policy.

At the other extreme, we have Policy Iteration approaches, such as Dual Policy Iteration (Sun et al., 2018b) , AlphaZero (Silver et al., 2017) and ExIT (Anthony et al., 2017) that do not directly mimic the actions of an expert, but instead plan or search over the policy to choose an action distribution to train towards (Silver et al., 2017) . These methods can be computationally expensive and their long planning horizon can cause the state distribution to diverge far from the current policy's induced distribution.

In this paper, we propose a new algorithm, named POLISh, that strikes a middle ground by performing multi-step improvements over states sampled from different time steps in the trajectories generated by the policy. We use MCTS (Kocsis & Szepesvári, 2006) to perform the local trajectory improvements. We show the promise of our approach with both theoretical and empirical results.

Monte Carlo Tree Search. We refer the reader to (Browne et al., 2012) for a comprehensive survey on MCTS. To generate the proposed locally optimized trajectories, we follow recent work that explores using MCTS to provide feedback for policy improvement (Guo et al., 2014; Silver et al., 2017; Anthony et al., 2017) . MCTS is a best-first tree search method which uses a policy to explore the most promising actions first. Through repeated simulation, MCTS builds a tree whose nodes represent states and whose branches correspond to actions that can be taken from those states. The objective is to maximize total return, so after all simulations have been completed, a final trajectory is generated by traversing the most visited sequence of nodes.

3. BACKGROUND & PRELIMINARIES:

Markov Decision Processes. We consider policy learning for Markov Decision Processes (MDPs) represented as a tuple (S, A, P, r, γ, D). Let S denote the state space, A the action space, P(s |s, a) the (probabilistic) state dynamics, r(s, a) the reward function, γ the discount factor and D the initial state distribution.

Policy Learning. A stochastic policy π maps a state s ∈ S to a distribution over the actions A, denoted by π(s). At each state, an action a is sampled from π(s) with probability π(a|s) and a reward of r(s, a) is received by π. The goal is to learn a policy that maximizes the accumulated discounted rewards

J D (π) = E τ ∼π [ ∞ i=0 γ i r(s i , a i )].

We omit the dependency of the initial state distribution D when the context is clear. A few useful quantities related to a policy are the value function V π , the state-action value function Q π and the advantage function A π defined as follows:

V π (s) = E τ ∼π [ ∞ i=0 γ i r(s i , a i )|s 0 = s], Q π (s, a) = E τ ∼π [ ∞ i=0 γ i r(s i , a i )|s 0 = s, a 0 = a], A π (s, a) = Q π (s, a) -V π (s)

As the policy takes a sequence of actions, its performance has a strong connection to the state distribution induced by its actions. We define quantities related to the state distributions induced by policies at different time steps. Let d t π denote the state distribution obtained by following π for t steps. We use

d π = (1-γ) ∞ t=0 γ t d t

π to denote the accumulated discounted state distribution. With these quantities, we can rewrite , a) ]. We will make use of the following relationship between the two policies in our analysis (Kakade & Langford, 2002) :

J(π) = ∞ t=0 E s∼d t π ,a∼π(s) [γ t r(s, a)] = E s∼dπ,a∼π(s) [r(s

EQUATION

Throughout the paper, we use π * to denote the expert policy.

4. THE POLISH ALGORITHM:

We first present a general framework for imitation learning from local trajectory improvements in Section 4.1. We then describe a method based on MCTS to perform the proposed local trajectory optimization in Section 4.2. Our main algorithm builds on the observation that behavioral cloning (Pomerleau, 1991) and DAgger (Ross et al., 2011) are two extremes along a spectrum of possible algorithms using experts to generate trajectory improvement. This is illustrated in Figure 1 . The solid nodes represent the states visited by a current policy π during a rollout of length T . DAgger seeks to collect expert feedback on every state along τ , while behavioral cloning only imitates a single complete trajectory from the expert. DAgger collects 1-step demonstrations from every state while behavioral cloning collects T -step from a single step. It is easy to see that there is a spectrum of algorithms that lie in between, parametrized by the time horizon t to collect expert demonstrations and the frequency with which to choose the start state for the expert policy.

4.1. MAIN ALGORITHM:

We present our main algorithm in Algorithm 1 that collects locally improved partial trajectories, as shown in Figure 1 and trains a policy to imitate the expert on the collected trajectories.

For each iteration, we first collect trajectories from a current policy π i (Line 3). Then, we collect every k states from each trajectory to use as starting states for running the expert policy (Line 7, 8, 9). From each state collected, we rollout π * for t steps to generate a partial trajectory (Line 12, 13) . Once all the data has been generated, we optimize for the behavioral cloning loss L(D, π) = 1 |D| (s,a * )∈D (I(π(s) = a * )), i.e., the supervised learning loss from data collected in D, to generate a new policy and repeat the process.

4.2. MCTS AS THE EXPERT POLICY:

General imitation learning requires having access to an expert policy. We follow a series of recent works that explore applying MCTS to provide feedback for policy improvement (Guo et al., 2014; Silver et al., 2017; Anthony et al., 2017) . Specifically, for every iteration, π * becomes the policy obtained by running MCTS with π i by following a variant of MCTS that uses a slightly modified UCT rule to select leaf node in MCTS: UCT(s, a) = r(s,a) n(s,a) + c • π(a|s) log n(s) n(s,a) , where n(s) is the number of times the node s has been visited so far, n(s, a) is the number of times action a is selected at node s, and r(s, a) is the sum of all rewards obtained in simulations by taking a in s. At a leaf node, Monte Carlo rollouts for value estimation (Browne et al., 2012) is replaced by an estimate obtained from a value network trained with the policy. This approach removes the need for a pre-defined expert policy and can be effective if MCTS can serve to improve the policy. We provide a sufficient condition that justifies this design choice.

Algorithm 1 POLISH (Policy Optimization by Local Improvement through Search):

1: Input: an initial policy π 1 , the expert policy π * , the segment length t, number of iterations N , a policy class Π.

2: for 1 ≤ i ≤ N do 3: Sample {τ j } n j=1 from π i 4: S ← ∅ 5: D ← ∅ 6:

for each trajectory

τ j = (s 0 , a 0 , s 1 , • • • , s T -1 , a T -1 , s T ) do 7: k ← T /t -1 8: S j ← {s 0 , s t , s 2t , • • • , s kt } 9:

S ← S ∪ S j 10:

end for 11:

for each s ∈ S do 12:

Run π * for t steps from the state s to generate a partial trajectory

τ s = (s, a 0 , s 1 , • • • , a t-1 , s t ) 13: D ← D ∪ {(s, a 0 ), (s 1 , a 1 ), • • • , (s t-1 , a t-1 )} 14:

end for 15:

π i+1 ← arg min π∈Π L(D, π) 16: end for 17: return π N +1

can serve as a method for improving local trajectories, thus providing a theoretical foundation for previous empirical works (Silver et al., 2017; Anthony et al., 2017) . We provide our proofs in the Appendix.

5.1. MOTIVATION FOR LOCAL TRAJECTORY IMPROVEMENT:

Imitation learning only makes sense if the policy we are imitating is better than our current policy. We formalize below what "better" means in this context: Assumption 1. Assume that π * is an expert policy such that for every state s ∈ S,

E a∼π * (s) [Q π (s, a)] ≥ E a∼π(s) [Q π (s, a)] and E a∼π * (s),s ∼P(•|s,a) [V k π * (s )] ≥ E a∼π(s),s ∼P(•|s,a) [V k π * (s )] for every k ≥ 1, where V k π denotes the accumulated discounted k-step returns.

The first assumption ensures that J(π * ) ≥ J(π) according to Equation 1. The second assumption ensures that one-step deviation from π * according to π will result in lower-valued states if we were to follow π * from that point on. This is typically true if π is a greedy policy so it pays less attention to future rewards.

To motivate the idea of using local trajectory improvement instead of complete trajectory imitation learning, we extend the analysis in (Ross & Bagnell, 2010 ) on behavioral cloning. A major issue with (long-horizon) behavioral cloning is the potential for cascading errors if a learned policy makes a mistake early on. This compounding behavior is captured in the following theorem:

Theorem 1. (Theorem 2.1 in (Ross & Bagnell, 2010) , rewritten in terms of rewards instead of costs) Let π be a policy such that = E s∼d π * [e π (s)]. Then J(π) ≥ J(π * ) -T 2 , where T is the task horizon and e π (s) = E a∼π(s),a * ∼π * (s) [I(a = a * )].

The cascading error effect manifests as a quadratic function in the task horizon. We can do a more careful analysis that motivates the design of an algorithm that clones behaviors at a shorter time scale. Define J ti:ti+1 π (π ) to be the rewards obtained by following π from time step t i to t i+1 from the start distribution d ti-1 π . Assume we divide the time horizon into equal parts with length t so that we have T /t segments with t i = it.

EQUATION

)

where i is the error rate computed between steps t i and t i+1 . Assuming i = holds across different segments, the expression above simplifies to

J(π) ≥ T /t i=1 γ t(i-1) J ti:ti+1 π (π * ) - T /t i=1 γ t(i-1) t 2 = T /t i=1 γ t(i-1) J ti:ti+1 π (π * ) - 1 -γ T 1 -γ t t 2 (3)

The objective is to maximize J(π) and this is done by optimizing the lower bound on the right hand side via imitating π * . Next, we consider how the two terms on the right hand side interact.

Proposition 1. Under Assumption 1,

T /t i=1 γ t(i-1) J ti:ti+1 π (π * ) ≤ T /t i=1 γ t (i-1) J t i :t i+1 π (π * ) if t is a multiple of t. Furthermore, it is maximized at t = T .

This proposition says that by multiplying the length of a segment, the first term grows.

Proposition 2. 1-γ T 1-γ t t 2 is an increasing function of t > 0 for γ ∈ (0, 1).

So the second term is also an increasing function in t, there can be a balance point where the right hand side is maximized, giving the tightest lower bound on the performance of policy π.

There is empirical evidence, (Silver et al., 2017; Anthony et al., 2017) , that fitting to the expert action distribution is a more robust approach than fitting to a single action that the expert takes. In Proposition 2, we show that we can relate to the KL-divergence between the action distributions of the policy and the expert.

Proposition 3. = E s∼d π * [e π (s)] ≤ E s∼d π * [D KL (π(s)||π * (s)) + H(π(s))]

, where D KL is the KL-divergence between two distributions and H is the entropy of a distribution.

As a result, by optimizing the objective presented in Section 4.1, we are improving the lower bound presented in Equation 3.

5.2. MONTE CARLO TREE SEARCH AS A LOCAL TRAJECTORY IMPROVEMENT OPERATOR:

Now we turn our attention to π * . In traditional imitation learning, the expert policy π * can be a human expert (Ross et al., 2011; Stadie et al., 2017) or a pre-trained policy (Ho & Ermon, 2016) . In either case, π * is considered fixed and independent of the current learned policy. Recently, several works (Guo et al., 2014; Silver et al., 2016; Anthony et al., 2017) explored using MCTS to plan ahead when a generative model of the underlying MDP is available (Kearns et al., 2002; Kakade et al., 2003) and thus produce better action decisions compared with directly following a policy. This approach circumvents the requirement of having a pre-defined expert and has the nice property of bootstrapping from the existing policy (Sun et al., 2018b) . However, it is not clear whether MCTS can always produce better trajectories for demonstrations. Another complication is the modified MCTS procedure, whose changes include that the UCT algorithm weights the exploration term by the policy prior probability and actual Monte Carlo rollout estimation of a state value is replaced by a value network.

Let V e π be a value network that estimates the true value function V π for π. During MCTS, whenever a value estimation for a leaf node is needed, the value network is used. We can derive a Q-function estimation Q e π for π by using the following equation Q e π (s, a) = E[r(s, a) + γV e π (s )]; that is, we use the sampled rewards and value estimation for the next state to form an estimation for the Qfunction, where the expectation is taken over the stochasticity of the environment transitions. The following theorem, for which we provide a proof in the appendix, demonstrates that MCTS can effectively generate local trajectory improvement.

Theorem 2. For a state s ∈ S, let ≤ (1 -π(a * |s)) min a =a * (Q e π (s, a * ) -Q e π (s, a)). After performing n MCTS searches according to the modified MCTS algorithm described in Section 4.2, we use π n (s) to denote the action distribution at state s, which is proportional to n(s, a).

If |V e π (s)- V π (s)| ≤ 2γ , as n → ∞, E a∼πn(s) [Q π (s, a)] ≥ E a∼π(s) [Q π (s, a)],

that is, π n (s) provides a onestep improvement over π. 

6. EXPERIMENT:

The goals of our experimental evaluation are threefold: (1) to verify that there exists a balance point for optimized length for local trajectory improvement, (2) to show that there is a trade-off between reward improvement by MCTS, and the divergence of action distributions between the policy and the expert (MCTS), (3) to demonstrate that the POLISH algorithm, by design, enables a significant speedup with a parallel implementation that is simple and efficient.

6.1. EXPERIMENT SETUP:

We evaluate our algorithm on three environments (Ant, Walker2d, and Hopper) from OpenAI Gym (Brockman et al., 2016) with the MuJoCo physical engine (Todorov et al., 2012) . For each environment, we use the standard maximum trajectory length of 1000. For each trajectory, there are multiple MCTS rollouts, each of which is distributed on a separate machine. This implementation takes advantage of the parallelizable nature of Lines 11, 12 in Algorithm 1. For instance, with a trajectory length of 1000 and a segment length of t = 32, we can achieve maximum parallelization by using 32 machines. Because the modified MCTS described in Section 4.2 relies on a policy during search, we need a reasonably performant policy in order for MCTS to be able to improve trajectories locally as shown in Theorem 2. Thus, we pre-train a policy with 1000 iterations of PPO (Schulman et al., 2017) . In the policy optimization step (Line 15 in Algorithm 1), we use an augmented version of L(D, π) = D KL (π||π * ) + H(π) + L V (D, π). The first two terms, as shown in Proposition 3, quantify the disagreement between π and π * , which we want to minimize. Note that in practice, we do not use the 0-1 loss presented in Section 1, since fitting a distribution target leads to more stable learning as noted in (Anthony et al., 2017) . The third term is a value loss similar to PPO (Schulman et al., 2017) , as we need to refine the value network for better MCTS performance.

6.2. MAIN RESULTS:

Figure 2 compares POLISH with different segment lengths: t = 1000 (behavioral cloning), t = 32 and DAgger (Ross et al., 2011) , which can be regarded as t = 1 with dataset aggregation across iterations, as well as the vanilla PPO. Overall, MCTS-driven imitation learning achieves a higher return compared to PPO in all three environments, and by a significant margin with the best choice of segment length in particular. Regarding the choice of segment lengths, we show that our choice of t = 32 achieves a higher return than the two other extremes (t = 1 and t = 1000), in Ant and Walker2d which are complex environments. This provides empirical support for Equation 3, which states that there exists a sweet spot for segment length in MCTS rollouts. However, in the simple Hopper environment, it is inconclusive which segment length leads to the best return.

6.3. EMPIRICAL EVIDENCE OF THEORY BEHIND POLISH:

To further understand the impact of segment length on the return, we show the reward improvement after running MCTS with the current policy in Figure 3a , and the KL-divergence of the action distributions between the MCTS and the policy in Figure 3b . 3a shows how the first term in Equation 3 changes as a function of segment length, and 3b depicts an approximation of the second term. For the Ant environment, we show both metrics for a select range of segment lengths of (1, 4, 32, 256, 1000) . As shown in our theoretical analysis, the MCTS return improvement (over the policy) generally increases with segment length, but so does the KL-divergence. Thus, there is a sweet spot for the segment length, where both the MCTS return improvement is high and the policy can effectively learn to imitate the MCTS expert.

6.4. PARALLEL IMPLEMENTATION SPEEDUP:

Lastly, we show that we are able to achieve a significant speedup with a parallel implementation. Our algorithm, by design, is easy to parallelize over distributed compute (Lines 11, 12 in Algorithm 1). Every MCTS rollout from an initial state is independent, and thus can be conducted on separate workers. The overhead of a parallel implementation, which arises from model synchronization across different workers and data transmission over the network, is likely insignificant compared to the workload of a MCTS rollout, especially for complex environments. Figure 3c shows the actual runtime of the parallel version of MCTS rollouts over a trajectory, compared to the ideal lower-bound which is the runtime for one single MCTS rollout at a segment length t. Even with a small segment length, e.g., t = 32, which implies a high number of parallel tasks (=32) over a fixed trajectory length (=1000), the actual runtime is close to ideal. This suggests that we can easily achieve speedup for OpenAI Gym-like environments. The overhead outweighs the benefit only when segment length is very small.

7. CONCLUSION AND FUTURE WORK:

In this work, we propose POLISH, a general imitation learning algorithm that provides flexible local trajectory improvement based on MCTS. We provide theoretical analysis that sheds light on the benefit of local trajectory improvement as well as a specific case study on MCTS. We validate the proposed approach by demonstrating strong empirical performance on a suite of high-dimensional continuous control problems, providing empirical support for our theory.

Model based methods have been proposed as a way to improve the sample efficiency of policy optimization (Watter et al., 2015; Finn & Levine, 2016; Hafner et al., 2018; Piergiovanni et al., 2018; Clavera et al., 2018; Sun et al., 2018b; Janner et al., 2019) . However, model based policy optimization methods are vulnerable to model inaccuracy, which can lead to compounding errors for long horizon planning (Cheng et al., 2018) . Nevertheless, the model can still be accurate enough for planning shorter trajectories, meaning that we can leverage it for the proposed POLISH algorithm, as it relies only on local trajectory improvement. In particular, we can learn a model of the environment and use that model to perform MCTS. As long as the learned model is reasonably accurate on short time horizons (Sun et al., 2018b) , we can expect that it will be able to improve the policy.

In this work, we focused on an expert feedback collection method with a fixed local search horizon.

One future direction would be to design a more adaptive approach where the length of local search horizons are learned, striking a balance between feedback quality and state distribution shift.

Brian D Ziebart, Andrew Maas, J Andrew Bagnell, and Anind K Dey. Maximum entropy inverse reinforcement learning. pp. 14331438, 2008.

A PROOFS FROM SECTION 5

Proof of Proposition 1

Proof. Let t = kt for some positive integer k. Since Algorithm 1 collects initial state every t and t steps, respectively. The set of starting states collected with the interval t is a subset of those collected with an interval of t. We will show that for the first segment of length t starting from some state s, J 0:t π (π * ) ≥ k i=1 γ t(i-1) J ti:ti+1 π (π * ). Similar proof holds for the other segments as well.

EQUATION

The inequality holds term-by-term because by Assumption 1, following π instead of π * generates state distributions that induce a lower valued states.

To show that the sum is maximized at t = T , consider any t < T .

EQUATION

It follows that the sum is maximized at t = T .

Proof of Proposition 2:

Proof. Let f (t) = t 2 1-γ t . To prove this proposition, it suffices to show that f (t) is an increasing function for t > 0. First we take the derivative and get

f (t) = 2t(1-γ t )+t 2 γ t ln γ (1-γ t ) 2 .

Let g(t) = 2t(1 -γ t ) + t 2 γ t ln γ be the numerator term. It is sufficient to show that g(t) > 0 for t > 0. Taking the derivative for g(t), we obtain g (t) = 2 -2γ t -2tγ t ln γ + ln γ(2tγ t + t 2 γ t ln γ) = 2 -2γ t + (ln γ) 2 t 2 γ t > 0 the last inequality follows from γ ∈ (0, 1). So g(t) is increasing, i.e., g(t) > g(0) = 0. It follows that f (t) > 0 and f (t) is an increasing function of t.

Proof of Proposition 3:

EQUATION

Proof of Theorem 2:

We first provide an analysis that extends theoretical guarantees of UCT (Kocsis et al., 2006; Kocsis & Szepesvári, 2006) to the setting with the modified UCT rule as in 4.2.

As usual, we connect UCT with the Upper-Confidence Bound algorithm in multi-armed bandits problem (Lai & Robbins, 1985; Agrawal, 1995; Auer et al., 2002) by framing the UCT as a nonstationary bandit problem. Specifically, we assume there are K arms. At time step t, arm i has a reward distribution X it which lies in the interval [0, 1] for simplicity. Unlike in traditional multiarmed bandits, we allow X it to vary as a function of t. Following the standard procedure for analyzing UCT-type algorithms in (Kocsis et al., 2006; Kocsis & Szepesvári, 2006) , we assume that the reward distributions concentrate. Assumption 2. For each

1 ≤ i ≤ K, X in = 1 n n t=1 X it converges. Let µ in = E[X in

] and µ i = lim n→∞ µ in . Let δ in be the bias with respect to µ i , i.e., δ in = µ i -µ in , assume that there exists a constant C p > 0 and an integer N p such that for n ≥ N p , for any δ > 0, ∆ n (δ) = C p n ln(1/δ), the following two tail probability inequalities are satisfied

EQUATION

The UCB algorithm balances the exploration-exploitation problem by selecting arms with the highest upper confidence bound. We consider a related variant of UCB that weights the exploration term by some prior probabilities p i for arm i. A central quantity analyzed in multi-armed bandit problems is the regret which connects to the number of times some sub-optimal arm is chosen. Let T i (n) be a random variable representing the total number of times that arm i is chosen in the first n rounds. We assume a unique optimal arm exists and its related quantities will be indexed by a * , i.e., p * , µ * , T * (n), X * t . We first upper bound the expected number of choices of a sub-optimal arm in terms of its sub-optimality ∆ i = µ * -µ i and the prior probability.

Since µ in converges to µ i as n → ∞, δ in converges to zero, so for any > 0, there exists an index N 0 ( ) such that if t ≥ N 0 ( ), then |δ it | ≤ ∆i 2 for every arm i. Theorem 3. Under Assumption 2, if we use c t,s = 2C p ln t s for the modified UCB algorithm. Fix

> 0, E[T i (n)] ≤ 16p 2 i ln n (1-) 2 ∆ 2 i + N 0 ( ) + N p + π 2 6 (e p 2 i + e p 2 * ).

Proof. We follow the proof in (Kocsis et al., 2006) 

EQUATION

We observe that X * s + p * c t,s ≤ X i,s + p i c t,s implies that at least one of the following is true

EQUATION

)

Since s ≥ A(n, ) and by the construction of A(n, ), A(n, ) ≥ A 0 (n, ) and A(n, ) ≥ N 0 ( ), so the inequality (19) does not hold. Thus, either (17) or (18) must hold. So I(X * s + p * c t,s ≤ X i,s + p i c t,s ) ≤ I(X * s + p * c t,s ≤ µ * ) + I(X i,s ≥ µ i + p i c t,s ). Now we can utilize inequalities (4) and ( 5) to derive upper bound on the two probabilities, so I(X * s + p * c t,s ≤ µ * ) = P(X * s ≤ µ * -p * c t,s ) ≤ e p 2 * /t 4 (20)

I(X i,s ≥ µ i + p i c t,s ) = P(X i,s ≥ µ i + p i c t,s ) ≤ e p 2 i /t 4 (21)

Plugging those quantities into (8) and use the relaxation that A(n, ) ≤ A 0 (n, ) + N 0 ( ) + N p , we obtain the desired result.

It follows that the asymptotic probability of choosing a sub-optimal arm goes to zero, the same guarantee in the original UCT algorithm, the rate of convergence is multiplied by the the square of the prior probability p 2 i . This means that if the prior probability is informative that it assigns small probabilities to sub-optimal arms, we can identify the optimal arm faster.

Next, we consider the more drastic change to the MCTS procedure by replacing the Monte Carlo rollout step with an estimation, usually from a value network learned with a policy (Silver et al., 2017; Anthony et al., 2017) . An accurate value estimation can effectively reduce the amount of lookahead necessary to identify an optimal action (Ng, 2003; Sun et al., 2018a) . On the other hand, an inaccurate value estimation can mislead the MCTS process. We provide a sufficient condition under which MCTS is guaranteed to improve the policy. It follows that, compared with the true Q-function Q π (s, a). 

:

We note that other techniques such as Generative Adversarial Imitation Learning (GAIL) have been developed that use example demonstrations in a modified objective which is still trained by interacting with the environment(Ho & Ermon

, 2016). These approaches directly operate on the induced state distribution itself; we do not consider these approaches here

THEORETICAL ANALYSISIn this section, we provide theoretical justification for considering local trajectory improvement in relation to both behavioral cloning and DAgger. We then prove that, under certain conditions, MCTS

Figures and Tables:

Figure 1: DAgger(Ross et al., 2011) collects 1step feedback at every state. Behavioral cloning collects T -step feedback only from the initial state. Our method collects t-step feedback at the frequency of every t states.

Figure 2: Average returns for continuous control tasks with different t values and the PPO baseline. Experiment results are across 5 random seeded runs. Shaded area indicates ±1 standard deviation.

Figure 3: Tradeoff between t and improved rewards from MCTS, KL-divergence, time to generate data from MCTS in the Ant environment.

|Q e π (s, a) -Q π (s, a)| = |E[r(s, a) + γV e π (s )] -E[r(s, a) + γV e π (s )]| = γ|E[]V e π (s ) -V π (s prove the theorem, we need the following lemma stating improvement when computed with respect to the estimated Q-function. Lemma 1. Let a * = arg max a Q e π (s, a) for a state s and ≤(1 -π(a * |s)) min a =a * Q e π (s, a * ) -Q e π (s, a). Then as n → ∞, E a∼πn(s) [Q e π (s, a)] ≥ E a∼π(s) [Q e π (s, a)] + .Proof. As n → ∞, π n (s) becomes a delta distribution with probability mass 1 on the a * , soE a∼πn(s) [Q e π (s, a)] -E a∼π(s) [Q e π (s, a)] = Q e π (s, a * ) -E a∼π(s) [Q e π (s, a)] (23) = (1 -π(a * |s))Q e π (s, a * ) -E a∼π(s),a =a * [Q e π (s, a)](24)≥ (1 -π(a * |s)) min a =a * Q e π (s, a * ) -Q e π (s, a) prove the main theorem, E a∼πn(s) [Q π (s, a)] ≥ E a∼πn(s) [Q e π (s, a) -E a∼π(s) [Q π (s, a)] (31)

Peer review (iclr2020):
The paper was confusing and difficult to read. I think it is trying to devise a methodology that improves upon a given expert policy, but I am not confident whether if this is the objective. The main contribution of the paper is the algorithm POLISH on page 5. However, it was unclear how MCTS with UCT is used with the expert policy in the algorithm. The first mention of the 0-1 loss objective is in section 4.1, and Algorithm 1 on page 5 claims to minimize this loss on line 15. However, in the experiment, the loss function is then switched to another L(D, pi) = D_{KL}(pi || pi*) + H(pi) + Lv(D, pi). Why the switch? And what is the definition of the third term? By examining the experimental results, Algorithm 1 seems to outperform PPO. Since the expert policy is known and given to Algorithm 1, I speculate Algorithm 1 is only replicating the expert policy, which would have outperformed PPO that has to learn from scratch. Thus, the comparison does not seem fair. It would be interesting to see how POLISH compares with the expert policy. Other comments: * Page 1 in Introduction: “However, these models suffer from the problem that even small difference between the learned policy and the expert behavior can lead to a snow-balling effect, where the state distribution diverges to a place where the behaviour of the policy is now meaningless since it was not trained that part of space”. Do you mean that the algorithm diverges instead of the state distribution diverges? The agent may incur errors in a space that has not been observed, but it should be able to learn eventually. So, what is causing divergence in states that have not yet encountered? * Page 4 in The POLISH Algorithm Main Algorithm: the 0-1 loss function is defined to be “L(D, pi) = 1/|D| sum_{s,a*}in D (I(pi(s) 
eq a*)), where a* is the expert policy’s selected action. I think it makes more sense to write “L(D, pi)* as 1/|D| sum_{s,a*}in D I ( a 
eq a* : a ~ pi(s)). Also, in RL, we don*t say that the policy receives a reward, but rather the agent receives a reward.