Original: P OLICY O PTIMIZATION BY L OCAL I MPROVEMENT THROUGH S EARCH
Revised: POLICY OPTIMIZATION BY LOCAL IMPROVEMENT THROUGH SEARCH

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Imitation learning has emerged as a powerful strategy for learning initial policies that can be reﬁned with reinforcement learning techniques.
Revised: Imitation learning has emerged as a powerful strategy for learning initial policies that can be reﬁned with reinforcement learning techniques.

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: However, the learning signal in both cases arises from the expert actions.
Revised: However, the learning signal in both cases arises from the expert actions.

Original: On the other end of the spectrum, approaches rooted in Policy
Revised: On the other end of the spectrum, approaches rooted in Policy

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: In this paper, we propose an algorithm that ﬁnds a middle ground by using Monte Carlo Tree Search (MCTS) (Kocsis & Szepesv ´ ari, 2006) to perform local trajectory improvement over rollouts from the policy.
Revised: In this paper, we propose an algorithm that ﬁnds a middle ground by using Monte Carlo Tree Search (MCTS) (Kocsis & Szepesv ´ ari, 2006) to perform local trajectory improvement over rollouts from the policy.

Original: We provide theoretical justiﬁcation for both the proposed local trajectory search algorithm and for our use of MCTS as a local policy improvement operator.
Revised: We provide theoretical justiﬁcation for both the proposed local trajectory search algorithm and for our use of MCTS as a local policy improvement operator.

Original: We also show empirically that our method (Policy Optimization by
Revised: We also show empirically that our method (Policy Optimization by

Original: 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.
Revised: 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.

Original: Furthermore, the resulting policy outperforms strong baselines in both reinforcement learning and imitation learning.
Revised: Furthermore, the resulting policy outperforms strong baselines in both reinforcement learning and imitation learning.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: 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).
Revised: 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).

Original: These successes showcase the power of learning from direct interactions with environments.
Revised: These successes showcase the power of learning from direct interactions with environments.

Original: 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
Revised: 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

Original: Optimization (PPO) (Schulman et al., 2017) by playing 180 years worth of games against itself daily (OpenAI, 2018).
Revised: Optimization (PPO) (Schulman et al., 2017) by playing 180 years worth of games against itself daily (OpenAI, 2018).

Original: The issue with sample complexity can be mitigated by learning from expert behavior.
Revised: The issue with sample complexity can be mitigated by learning from expert behavior.

Original: 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 reﬁnement.
Revised: 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 reﬁnement.

Original: 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).
Revised: 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).

Original: Earlier approaches, such as behavioral cloning (BC), relied on training models to mimic expert behavior at various states in the demonstration data (Pomerleau, 1991).
Revised: Earlier approaches, such as behavioral cloning (BC), relied on training models to mimic expert behavior at various states in the demonstration data (Pomerleau, 1991).

Original: 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).
Revised: 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).

Original: In order to mitigate these issues, DAgger (Ross et al., 2011) uses an expert policy to provide supervision to the policy being learned.
Revised: In order to mitigate these issues, DAgger (Ross et al., 2011) uses an expert policy to provide supervision to the policy being learned.

Original: 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.
Revised: 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.

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: In this paper, we propose an algorithm that ﬁnds a middle ground by using Monte Carlo Tree Search (MCTS) (Kocsis & Szepesv´ari, 2006) to perform local trajectory improvement over states sampled from different time steps in the trajectory unrolled from the policy.
Revised: In this paper, we propose an algorithm that ﬁnds a middle ground by using Monte Carlo Tree Search (MCTS) (Kocsis & Szepesv´ari, 2006) to perform local trajectory improvement over states sampled from different time steps in the trajectory unrolled from the policy.

Original: This approach has the beneﬁt 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.
Revised: This approach has the beneﬁt 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.

Original: This approach stands in contrast to other works in interactive imitation learning that correct distribution mismatch by using one-step feedback.
Revised: This approach stands in contrast to other works in interactive imitation learning that correct distribution mismatch by using one-step feedback.

Original: We provide theoretical justiﬁcation for the advantage of a balanced local trajectory improvement and show that MCTS can serve as a policy improvement operator under certain conditions 1 .
Revised: We provide theoretical justiﬁcation for the advantage of a balanced local trajectory improvement and show that MCTS can serve as a policy improvement operator under certain conditions 1 .

Original: An added beneﬁt of our effort is computational – depth parallel MCTS on local trajectory segments is much faster than traditional sequential MCTS to generate demonstrations.
Revised: An added beneﬁt of our effort is computational – depth parallel MCTS on local trajectory segments is much faster than traditional sequential MCTS to generate demonstrations.

Original: We show that our proposed algorithm can easily be parallelized to enable more efﬁcient imitation learning from MCTS.
Revised: We show that our proposed algorithm can easily be parallelized to enable more efﬁcient imitation learning from MCTS.

Original: Notably, this level of parallelism is present at varying depths which differs from existing works on parallel MCTS (Chaslot et al., 2008).
Revised: Notably, this level of parallelism is present at varying depths which differs from existing works on parallel MCTS (Chaslot et al., 2008).

Original: In summary, our main contributions are:
Revised: In summary, our main contributions are:

Original: • A general interactive imitation learning algorithm that balances the expert feedback quality and the state distribution divergence.
Revised: • A general interactive imitation learning algorithm that balances the expert feedback quality and the state distribution divergence.

Original: Our proposed approach provides a ﬂexible local trajectory improvement based on MCTS.
Revised: Our proposed approach provides a ﬂexible local trajectory improvement based on MCTS.

Original: • Theoretical analysis on the general beneﬁt of local improvement and speciﬁc case study on using MCTS as a local improvement operator.
Revised: • Theoretical analysis on the general beneﬁt of local improvement and speciﬁc case study on using MCTS as a local improvement operator.

Original: • Strong empirical performance on a suite of high-dimensional continuous control problems based
Revised: • Strong empirical performance on a suite of high-dimensional continuous control problems based

Original: on both sample efﬁciency and training time.
Revised: on both sample efﬁciency and training time.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Imitation Learning.
Revised: Imitation Learning.

Original: Imitation learning (IL) refers to the problem of learning to perform a task from expert demonstrations.
Revised: Imitation learning (IL) refers to the problem of learning to perform a task from expert demonstrations.

Original: Behavioral cloning (Widrow & Smith, 1964) is one popular approach which maximizes the likelihood of expert actions under the agent policy (Pomerleau, 1989;
Revised: Behavioral cloning (Widrow & Smith, 1964) is one popular approach which maximizes the likelihood of expert actions under the agent policy (Pomerleau, 1989;

Original: Schaal, 1999;
Revised: Schaal, 1999;

Original: Muller et al., 2006;
Revised: Muller et al., 2006;

Original: Mlling et al., 2013; Bojarski et al., 2016;
Revised: Mlling et al., 2013; Bojarski et al., 2016;

Original: Giusti et al., 2016; Mahler & Goldberg, 2017; Wang et al., 2019; Bansal et al., 2018).
Revised: Giusti et al., 2016; Mahler & Goldberg, 2017; Wang et al., 2019; Bansal et al., 2018).

Original: 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;
Revised: 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;

Original: Fu et al., 2018;
Revised: Fu et al., 2018;

Original: Ho & Ermon, 2016).
Revised: Ho & Ermon, 2016).

Original: In this work we focus on imitation learning via behavioral cloning.
Revised: In this work we focus on imitation learning via behavioral cloning.

Original: 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´e et al., 2009).
Revised: 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´e et al., 2009).

Original: To explain distributional shift, let us assume we train an agent by performing supervised learning on the actions an expert has taken.
Revised: To explain distributional shift, let us assume we train an agent by performing supervised learning on the actions an expert has taken.

Original: A small error in the ﬁrst time step of training may bring the agent to a state that the expert has rarely visited and is therefore not as well modeled.
Revised: A small error in the ﬁrst time step of training may bring the agent to a state that the expert has rarely visited and is therefore not as well modeled.

Original: Over time, this error compounds, leading the agent to states far from expert behavior and diverging further from the expert demonstrations.
Revised: Over time, this error compounds, leading the agent to states far from expert behavior and diverging further from the expert demonstrations.

Original: The longer the episode is, the more likely the agent is to deviate from expert demonstrations.
Revised: The longer the episode is, the more likely the agent is to deviate from expert demonstrations.

Original: Various solutions to the problem of distributional shift have been proposed (Daum´e et al., 2009;
Revised: Various solutions to the problem of distributional shift have been proposed (Daum´e et al., 2009;

Original: Ross & Bagnell, 2010; Ross et al., 2011; Ho & Ermon, 2016; Laskey et al., 2017; Bansal et al., 2018; de Haan et al., 2019).
Revised: Ross & Bagnell, 2010; Ross et al., 2011; Ho & Ermon, 2016; Laskey et al., 2017; Bansal et al., 2018; de Haan et al., 2019).

Original: Some of these approaches reduce the effects of distributional shift by iteratively querying the expert (Daum´e et al., 2009; Ross et al., 2011).
Revised: Some of these approaches reduce the effects of distributional shift by iteratively querying the expert (Daum´e et al., 2009; Ross et al., 2011).

Original: 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.
Revised: 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.

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: We use MCTS (Kocsis & Szepesv´ari, 2006) to perform the local trajectory improvements.
Revised: We use MCTS (Kocsis & Szepesv´ari, 2006) to perform the local trajectory improvements.

Original: We show the promise of our approach with both theoretical and empirical results.
Revised: We show the promise of our approach with both theoretical and empirical results.

Original: Monte Carlo Tree Search.
Revised: Monte Carlo Tree Search.

Original: We refer the reader to (Browne et al., 2012) for a comprehensive survey on MCTS.
Revised: We refer the reader to (Browne et al., 2012) for a comprehensive survey on MCTS.

Original: 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).
Revised: 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).

Original: MCTS is a best-ﬁrst tree search method which uses a policy to explore the most promising actions ﬁrst.
Revised: MCTS is a best-ﬁrst tree search method which uses a policy to explore the most promising actions ﬁrst.

Original: Through repeated simulation, MCTS builds a tree whose nodes represent states and whose branches correspond to actions that can be taken from those states.
Revised: Through repeated simulation, MCTS builds a tree whose nodes represent states and whose branches correspond to actions that can be taken from those states.

Original: The objective is to maximize total return, so after all simulations have been completed, a ﬁnal trajectory is generated by traversing the most visited sequence of nodes.
Revised: The objective is to maximize total return, so after all simulations have been completed, a ﬁnal trajectory is generated by traversing the most visited sequence of nodes.

Original:  B ACKGROUND & P RELIMINARIES
Revised:  B ACKGROUND & P RELIMINARIES

Original: Markov Decision Processes.
Revised: Markov Decision Processes.

Original: We consider policy learning for Markov Decision Processes (MDPs) represented as a tuple ( S , A , P , r, γ, D ) .
Revised: We consider policy learning for Markov Decision Processes (MDPs) represented as a tuple ( S , A , P , r, γ, D ) .

Original: Let S denote the state space,
Revised: Let S denote the state space,

Original: A the action space, P ( s (cid:48) | s, a ) the (probabilistic) state dynamics, r ( s, a ) the reward function, γ the discount factor and D the initial state distribution.
Revised: A the action space, P ( s (cid:48) | s, a ) the (probabilistic) state dynamics, r ( s, a ) the reward function, γ the discount factor and D the initial state distribution.

Original: Policy Learning.
Revised: Policy Learning.

Original: A stochastic policy π maps a state s ∈ S to a distribution over the actions A , denoted by π ( s ) .
Revised: A stochastic policy π maps a state s ∈ S to a distribution over the actions A , denoted by π ( s ) .

Original: At each state, an action a is sampled from π ( s ) with probability π ( a | s ) and a reward of r ( s, a ) is received by π .
Revised: At each state, an action a is sampled from π ( s ) with probability π ( a | s ) and a reward of r ( s, a ) is received by π .

Original: The goal is to learn a policy that maximizes the accumulated discounted rewards J D ( π ) =
Revised: The goal is to learn a policy that maximizes the accumulated discounted rewards J D ( π ) =

Original: E τ ∼ π
Revised: E τ ∼ π

Original: [ (cid:80) ∞ i =

Revised: [ (cid:80) ∞ i =

Original: γ i r ( s i , a i )] .
Revised: γ i r ( s i , a i )] .

Original: We omit the dependency of the initial state distribution D when the context is clear.
Revised: We omit the dependency of the initial state distribution D when the context is clear.

Original: A few useful quantities related to a policy are the value function V π , the state-action value function Q π and the advantage function A π deﬁned as follows:
Revised: A few useful quantities related to a policy are the value function V π , the state-action value function Q π and the advantage function A π deﬁned as follows:

Original: [Equation]
Revised: [Equation]

Original: As the policy takes a sequence of actions, its performance has a strong connection to the state distribution induced by its actions.
Revised: As the policy takes a sequence of actions, its performance has a strong connection to the state distribution induced by its actions.

Original: We deﬁne quantities related to the state distributions induced by policies at different time steps.
Revised: We deﬁne quantities related to the state distributions induced by policies at different time steps.

Original: Let d tπ denote the state distribution obtained by following π for t steps.
Revised: Let d tπ denote the state distribution obtained by following π for t steps.

Original: We use d π = (1 − γ ) (cid:80) ∞ t =

Revised: We use d π = (1 − γ ) (cid:80) ∞ t =

Original: γ t d t π to denote the accumulated discounted state distribution.
Revised: γ t d t π to denote the accumulated discounted state distribution.

Original: With these quantities, we can rewrite J ( π ) =
Revised: With these quantities, we can rewrite J ( π ) =

Original: (cid:80) ∞ t =

Revised: (cid:80) ∞ t =

Original: E s ∼ d tπ ,a ∼ π ( s )
Revised: E s ∼ d tπ ,a ∼ π ( s )

Original: [ γ t r ( s, a )]
Revised: [ γ t r ( s, a )]

Original: =
Revised: =

Original: E s ∼ d π ,a ∼ π ( s )
Revised: E s ∼ d π ,a ∼ π ( s )

Original: [ r ( s, a )] .
Revised: [ r ( s, a )] .

Original: We will make use of the following relationship between the two policies in our analysis (Kakade &
Revised: We will make use of the following relationship between the two policies in our analysis (Kakade &

Original: Langford, 2002):
Revised: Langford, 2002):

Original: [Equation]
Revised: [Equation]

Original: Throughout the paper, we use π ∗ to denote the expert policy.
Revised: Throughout the paper, we use π ∗ to denote the expert policy.

Original:  T HE POLISH A LGORITHM
Revised:  T HE POLISH A LGORITHM

Original: We ﬁrst present a general framework for imitation learning from local trajectory improvements in
Revised: We ﬁrst present a general framework for imitation learning from local trajectory improvements in

Original: Section 4.1.
Revised: Section 4.1.

Original: We then describe a method based on MCTS to perform the proposed local trajectory optimization in Section 4.2.
Revised: We then describe a method based on MCTS to perform the proposed local trajectory optimization in Section 4.2.

Original: 4. M AIN A LGORITHM
Revised: 4. M AIN A LGORITHM

Original: [Figure]
Revised: [Figure]

Original: Our main algorithm builds on the observation that behavioral cloning (Pomerleau, 1991) and
Revised: Our main algorithm builds on the observation that behavioral cloning (Pomerleau, 1991) and

Original: DAgger (Ross et al., 2011) are two extremes along a spectrum of possible algorithms using experts to generate trajectory improvement.
Revised: DAgger (Ross et al., 2011) are two extremes along a spectrum of possible algorithms using experts to generate trajectory improvement.

Original: This is illustrated in Figure 1.
Revised: This is illustrated in Figure 1.

Original: The solid nodes represent the states visited by a current policy π during a rollout of length T .
Revised: The solid nodes represent the states visited by a current policy π during a rollout of length T .

Original: DAgger seeks to collect expert feedback on every state along τ , while behavioral cloning only imitates a single complete trajectory from the expert.
Revised: DAgger seeks to collect expert feedback on every state along τ , while behavioral cloning only imitates a single complete trajectory from the expert.

Original: DAgger collects 1-step demonstrations from every state while behavioral cloning collects T -step from a single step.
Revised: DAgger collects 1-step demonstrations from every state while behavioral cloning collects T -step from a single step.

Original: 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.
Revised: 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.

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

Original: For each iteration, we ﬁrst collect trajectories from a current policy
Revised: For each iteration, we ﬁrst collect trajectories from a current policy

Original: π i (Line 3).
Revised: π i (Line 3).

Original: Then, we collect every k states from each trajectory to use as starting states for running the expert policy (Line 7, 8, 9).
Revised: Then, we collect every k states from each trajectory to use as starting states for running the expert policy (Line 7, 8, 9).

Original: From each state collected, we rollout π ∗ for t steps to generate a partial trajectory (Line 12, 13).
Revised: From each state collected, we rollout π ∗ for t steps to generate a partial trajectory (Line 12, 13).

Original: Once all the data has been generated, we optimize for the behavioral cloning loss L ( D, π ) = 

Revised: Once all the data has been generated, we optimize for the behavioral cloning loss L ( D, π ) =

Original: | D | (cid:80) ( s,a ∗ )
Revised: | D | (cid:80) ( s,a ∗ )

Original: ∈ D ( I ( π ( s ) (cid:54) = a ∗ )) , i.e., the supervised learning loss from data collected in D , to generate a new policy and repeat the process.
Revised: ∈ D ( I ( π ( s ) (cid:54) = a ∗ )) , i.e., the supervised learning loss from data collected in D , to generate a new policy and repeat the process.

Original: 4. MCTS AS THE EXPERT POLICY
Revised: 4. MCTS AS THE EXPERT POLICY

Original: General imitation learning requires having access to an expert policy.
Revised: General imitation learning requires having access to an expert policy.

Original: We follow a series of recent works that explore applying MCTS to provide feedback for policy improvement (Guo et al., 2014;
Revised: We follow a series of recent works that explore applying MCTS to provide feedback for policy improvement (Guo et al., 2014;

Original: Silver et al., 2017; Anthony et al., 2017).
Revised: Silver et al., 2017; Anthony et al., 2017).

Original: Speciﬁcally, for every iteration, π ∗ becomes the policy obtained by running MCTS with π i by following a variant of MCTS that uses a slightly modiﬁed
Revised: Speciﬁcally, for every iteration, π ∗ becomes the policy obtained by running MCTS with π i by following a variant of MCTS that uses a slightly modiﬁed

Original: UCT rule to select leaf node in MCTS: UCT ( s, a ) =
Revised: UCT rule to select leaf node in MCTS: UCT ( s, a ) =

Original: r ( s,a ) n ( s,a ) + c · π ( a | s ) (cid:113) 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 .
Revised: r ( s,a ) n ( s,a ) + c · π ( a | s ) (cid:113) 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 .

Original: 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.
Revised: 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.

Original: This approach removes the need for a pre-deﬁned expert policy and can be effective if MCTS can serve to improve the policy.
Revised: This approach removes the need for a pre-deﬁned expert policy and can be effective if MCTS can serve to improve the policy.

Original: We provide a sufﬁcient condition that justiﬁes this design choice.
Revised: We provide a sufﬁcient condition that justiﬁes this design choice.

Original:  T HEORETICAL A NALYSIS
Revised:  T HEORETICAL A NALYSIS

Original: In this section, we provide theoretical justiﬁcation for considering local trajectory improvement in relation to both behavioral cloning and DAgger.
Revised: In this section, we provide theoretical justiﬁcation for considering local trajectory improvement in relation to both behavioral cloning and DAgger.

Original: We then prove that, under certain conditions, MCTS 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).
Revised: We then prove that, under certain conditions, MCTS 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).

Original: We provide our proofs in the
Revised: We provide our proofs in the

Original: Appendix.
Revised: Appendix.

Original: 5. M OTIVATION FOR L OCAL T RAJECTORY I MPROVEMENT
Revised: 5. M OTIVATION FOR L OCAL T RAJECTORY I MPROVEMENT

Original: Imitation learning only makes sense if the policy we are imitating is better than our current policy.
Revised: Imitation learning only makes sense if the policy we are imitating is better than our current policy.

Original: We formalize below what “better” means in this context:
Revised: We formalize below what “better” means in this context:

Original: Assumption 1.
Revised: Assumption 1.

Original: Assume that π ∗ is an expert policy such that for every state
Revised: Assume that π ∗ is an expert policy such that for every state

Original: [Equation]
Revised: [Equation]

Original: π ∗ ( s )
Revised: π ∗ ( s )

Original: [ Q π ( s, a )] ≥
Revised: [ Q π ( s, a )] ≥

Original: E a ∼ π ( s ) [ Q π ( s, a )] and E a ∼ π ∗ ( s ) ,s (cid:48) ∼P ( ·| s,a ) [ V kπ ∗ ( s (cid:48) )] ≥
Revised: E a ∼ π ( s ) [ Q π ( s, a )] and E a ∼ π ∗ ( s ) ,s (cid:48) ∼P ( ·| s,a ) [ V kπ ∗ ( s (cid:48) )] ≥

Original: E a ∼ π ( s ) ,s (cid:48) ∼P ( ·| s,a ) [ V kπ ∗ ( s (cid:48) )] for every k ≥ 1 , where V kπ denotes the accumulated discounted k -step returns.
Revised: E a ∼ π ( s ) ,s (cid:48) ∼P ( ·| s,a ) [ V kπ ∗ ( s (cid:48) )] for every k ≥ 1 , where V kπ denotes the accumulated discounted k -step returns.

Original: The ﬁrst assumption ensures that J ( π ∗ ) ≥ J ( π ) according to Equation 1.
Revised: The ﬁrst assumption ensures that J ( π ∗ ) ≥ J ( π ) according to Equation 1.

Original: 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.
Revised: 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.

Original: This is typically true if π is a greedy policy so it pays less attention to future rewards.
Revised: This is typically true if π is a greedy policy so it pays less attention to future rewards.

Original: 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.
Revised: 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.

Original: A major issue with (long-horizon) behavioral cloning is the potential for cascading errors if a learned policy makes a mistake early on.
Revised: A major issue with (long-horizon) behavioral cloning is the potential for cascading errors if a learned policy makes a mistake early on.

Original: This compounding behavior is captured in the following theorem:
Revised: This compounding behavior is captured in the following theorem:

Original: Theorem 1.
Revised: Theorem 1.

Original: (Theorem 2.1 in (Ross & Bagnell, 2010), rewritten in terms of rewards instead of costs)
Revised: (Theorem 2.1 in (Ross & Bagnell, 2010), rewritten in terms of rewards instead of costs)

Original: Let π be a policy such that (cid:15) =
Revised: Let π be a policy such that (cid:15) =

Original: E s ∼ d π ∗
Revised: E s ∼ d π ∗

Original: e π ( s )] .
Revised: e π ( s )] .

Original: Then J ( π ) ≥ J ( π ∗ ) − T 2 (cid:15) , where T is the task horizon and e π ( s ) =
Revised: Then J ( π ) ≥ J ( π ∗ ) − T 2 (cid:15) , where T is the task horizon and e π ( s ) =

Original: E a ∼ π ( s ) ,a ∗∼ π ∗ ( s )
Revised: E a ∼ π ( s ) ,a ∗∼ π ∗ ( s )

Original: [ I ( a (cid:54) = a ∗ )] .
Revised: [ I ( a (cid:54) = a ∗ )] .

Original: The cascading error effect manifests as a quadratic function in the task horizon.
Revised: The cascading error effect manifests as a quadratic function in the task horizon.

Original: We can do a more careful analysis that motivates the design of an algorithm that clones behaviors at a shorter time scale.
Revised: We can do a more careful analysis that motivates the design of an algorithm that clones behaviors at a shorter time scale.

Original: Deﬁne J t
Revised: Deﬁne J t

Original: i : t i +1 π
Revised: i : t i +1 π

Original: ( π (cid:48) ) to be the rewards obtained by following π (cid:48) from time step t i to t
Revised: ( π (cid:48) ) to be the rewards obtained by following π (cid:48) from time step t i to t

Original: i +1 from the start distribution d t i − 1 π .
Revised: i +1 from the start distribution d t i − 1 π .

Original: Assume we divide the time horizon into equal parts with length t so that we have T/t segments with t
Revised: Assume we divide the time horizon into equal parts with length t so that we have T/t segments with t

Original: i = it .
Revised: i = it .

Original: [Equation]
Revised: [Equation]

Original: where (cid:15) i is the error rate computed between steps t i and t i +1 .
Revised: where (cid:15) i is the error rate computed between steps t i and t i +1 .

Original: Assuming (cid:15) i = (cid:15) holds across different segments, the expression above simpliﬁes to
Revised: Assuming (cid:15) i = (cid:15) holds across different segments, the expression above simpliﬁes to

Original: [Equation]
Revised: [Equation]

Original: The objective is to maximize J ( π ) and this is done by optimizing the lower bound on the right hand side via imitating π ∗ .
Revised: The objective is to maximize J ( π ) and this is done by optimizing the lower bound on the right hand side via imitating π ∗ .

Original: Next, we consider how the two terms on the right hand side interact.
Revised: Next, we consider how the two terms on the right hand side interact.

Original: Proposition 1.
Revised: Proposition 1.

Original: Under Assumption 1, (cid:80) T/t
Revised: Under Assumption 1, (cid:80) T/

Original: [Equation]
Revised: [Equation]

Original: =1 J
Revised: 1 γ t (

Original: [Equation]
Revised: [Equation]

Original: t
Revised: J

Original: [Equation]
Revised: [Equation]

Original: π [Equation] t .
Revised: is a multiple of t .

Original: Furthermore, it is maximized at t = T .
Revised: Furthermore, it is maximized at t = T .

Original: This proposition says that by multiplying the length of a segment, the ﬁrst term grows.
Revised: This proposition says that by multiplying the length of a segment, the ﬁrst term grows.

Original: Since the second term, Tt(cid:15) , is  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 π .
Revised: Proposition 2. 1 − γ T 1 − γ t t 2 (cid:15) 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 π .

Original: There is empirical evidence, (Silver et al., 2017; Anthony et al., 2017), that ﬁtting to the expert action distribution is a more robust approach than ﬁtting to a single action that the expert takes.
Revised: There is empirical evidence, (Silver et al., 2017; Anthony et al., 2017), that ﬁtting to the expert action distribution is a more robust approach than ﬁtting to a single action that the expert takes.

Original: In
Revised: In

Original: Proposition 2, we show that we can relate (cid:15) to the KL-divergence between the action distributions of the policy and the expert.
Revised: Proposition 2, we show that we can relate (cid:15) to the KL-divergence between the action distributions of the policy and the expert.

Original: Proposition 2.
Revised: Proposition 3.

Original: (cid:15) =
Revised: (cid:15) =

Original: E s ∼ d π ∗
Revised: [Equation] e [Equation]

Original: e π ( s )] ≤ E s ∼ d π∗
Revised: [Equation] [ D [Equation] π

Original: [ D KL ( π ( s ) ||
Revised: [Equation] s ) ||

Original: π ∗( s ))
Revised: [Equation] ∗

Original:  + H ( π ( s ))] , where D KL is the
Revised: [Equation] s )) + H [Equation]

Revised: π

Original:  KL-divergence between two distributions and H is the entropy of a distribution.
Revised: [Equation] s ))] , where D KL is the KL-divergence between two distributions and H is the entropy of a distribution.

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

Original: 5. M ONTE C ARLO T REE S EARCH AS A L OCAL T RAJECTORY I MPROVEMENT O PERATOR
Revised: 5. M ONTE C ARLO T REE S EARCH AS A L OCAL T RAJECTORY I MPROVEMENT O PERATOR

Original: Now we turn our attention to π ∗ .
Revised: Now we turn our attention to π ∗ .

Original: 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).
Revised: 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).

Original: In either case, π ∗ is considered ﬁxed and independent of the current learned policy.
Revised: In either case, π ∗ is considered ﬁxed and independent of the current learned policy.

Original: 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.
Revised: 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.

Original: This approach circumvents the requirement of having a pre-deﬁned expert and has the nice property of bootstrapping from the existing policy (Sun et al., 2018b).
Revised: This approach circumvents the requirement of having a pre-deﬁned expert and has the nice property of bootstrapping from the existing policy (Sun et al., 2018b).

Original: However, it is not clear whether MCTS can always produce better trajectories for demonstrations.
Revised: However, it is not clear whether MCTS can always produce better trajectories for demonstrations.

Original: Another complication is the modiﬁed
Revised: Another complication is the modiﬁed

Original: 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.
Revised: 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.

Original: Let V e π be a value network that estimates the true value function V π for π .
Revised: Let V eπ be a value network that estimates the true value function V π for π .

Original: During MCTS, whenever a value estimation for a leaf node is needed, the value network is used.
Revised: During MCTS, whenever a value estimation for a leaf node is needed, the value network is used.

Original: We can derive a Q -function estimation Q eπ for π by using the following equation Q eπ ( s, a ) =
Revised: We can derive a Q -function estimation Q eπ for π by using the following equation Q eπ ( s, a ) =

Original: E [ r ( s, a ) + γV eπ
Revised: E [ r ( s, a ) + γV eπ

Original: ( s (cid:48) )] ; 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.
Revised: ( s (cid:48) )] ; 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.

Original: The following theorem, for which we provide a proof in the appendix, demonstrates that MCTS can effectively generate local trajectory improvement.
Revised: The following theorem, for which we provide a proof in the appendix, demonstrates that MCTS can effectively generate local trajectory improvement.

Original: Theorem 2.
Revised: Theorem 2.

Original: For a state s ∈ S , let (cid:15) ≤ (1 − π ( a ∗ | s ))
Revised: For a state s ∈ S , let (cid:15) ≤ (1 − π ( a ∗ | s ))

Original: min a (cid:54) =
Revised: min a (cid:54) =

Original: a ∗ ( Q eπ ( s, a ∗ ) − Q eπ ( s, a )) .
Revised: a ∗ ( Q eπ ( s, a ∗ ) − Q eπ ( s, a )) .

Original: After performing n MCTS searches according to the modiﬁed 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 ) .
Revised: After performing n MCTS searches according to the modiﬁed 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 ) .

Original: If | V eπ ( s ) −
Revised: If | V eπ ( s ) −

Original: V π ( s ) | ≤ (cid:15) 2 γ , as n → ∞ , E a ∼ π n ( s )
Revised: V π ( s ) | ≤ (cid:15) 2 γ , as n → ∞ , E a ∼ π n ( s )

Original: [ Q π ( s, a )] ≥ E a ∼ π ( s )
Revised: [ Q π ( s, a )] ≥ E a ∼ π ( s )

Original: [ Q π ( s, a )] , that is, π n ( s ) provides a onestep improvement over π . 
Revised: [ Q π ( s, a )] , that is, π n ( s ) provides a onestep improvement over π . [Figure] [Figure]

Original:  E XPERIMENT
Revised:  E XPERIMENT

Original: 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 [Figure]
Revised: 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 signiﬁcant speedup with a parallel implementation that is simple and efﬁcient

Original: [Figure]
Revised: [Figure]

Original: 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 signiﬁcant speedup with a parallel implementation that is simple and efﬁcient.

Original: 6. E XPERIMENT S ETUP
Revised: 6. E XPERIMENT S ETUP

Original: 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).
Revised: 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).

Original: For each environment, we use the standard maximum trajectory length of 1000.
Revised: For each environment, we use the standard maximum trajectory length of 1000.

Original: For each trajectory, there are multiple MCTS rollouts, each of which is distributed on a separate machine.
Revised: For each trajectory, there are multiple MCTS rollouts, each of which is distributed on a separate machine.

Original: This implementation takes advantage of the parallelizable nature of Lines 11, 12 in Algorithm 1.
Revised: This implementation takes advantage of the parallelizable nature of Lines 11, 12 in Algorithm 1.

Original: For instance, with a trajectory length of 1000 and a segment length of t = 32 , we can achieve maximum parallelization by using 32 machines.
Revised: For instance, with a trajectory length of 1000 and a segment length of t = 32 , we can achieve maximum parallelization by using 32 machines.

Original: Because the modiﬁed 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.
Revised: Because the modiﬁed 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.

Original: Thus, we pre-train a policy with 1000 iterations of PPO (Schulman et al., 2017).
Revised: Thus, we pre-train a policy with 1000 iterations of PPO (Schulman et al., 2017).

Original: In the policy optimization step (Line 15 in Algorithm 1), we use an augmented version of L ( D, π ) =
Revised: In the policy optimization step (Line 15 in Algorithm 1), we use an augmented version of L ( D, π ) =

Original: D KL ( π || π ∗ )
Revised: D KL ( π || π ∗ )

Original: + H ( π ) + L V
Revised: + H ( π ) + L V

Original: ( D, π ) .
Revised: ( D, π ) .

Original: The ﬁrst two terms, as shown in Proposition 2, quantify the disagreement between π and π ∗ , which we want to minimize.
Revised: The ﬁrst two terms, as shown in Proposition 3, quantify the disagreement between π and π ∗ , which we want to minimize.

Original: Note that in practice, we do not use the 0-1 loss presented in Section 1, since ﬁtting a distribution target leads to more stable learning as noted in (Anthony et al., 2017).
Revised: Note that in practice, we do not use the 0-1 loss presented in Section 1, since ﬁtting a distribution target leads to more stable learning as noted in (Anthony et al., 2017).

Original: The third term is a value loss similar to PPO (Schulman et al., 2017), as we need to reﬁne the value network for better MCTS performance.
Revised: The third term is a value loss similar to PPO (Schulman et al., 2017), as we need to reﬁne the value network for better MCTS performance.

Original: 6. M AIN R ESULTS
Revised: 6. M AIN R ESULTS

Original: Figure 2 compares POLISH with different segment lengths: t = 1000 (behavioral cloning), t =and DAgger (Ross et al., 2011), which can be regarded as t = 1 with dataset aggregation across iterations, as well as the vanilla PPO.
Revised: Figure 2 compares POLISH with different segment lengths: t = 1000 (behavioral cloning), t =and DAgger (Ross et al., 2011), which can be regarded as t = 1 with dataset aggregation across iterations, as well as the vanilla PPO.

Original: Overall, MCTS-driven imitation learning achieves a higher return compared to PPO in all three environments, and by a signiﬁcant margin with the best choice of segment length in particular.
Revised: Overall, MCTS-driven imitation learning achieves a higher return compared to PPO in all three environments, and by a signiﬁcant margin with the best choice of segment length in particular.

Original: 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
Revised: 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

Original: Walker2d which are complex environments.
Revised: Walker2d which are complex environments.

Original: This provides empirical support for Equation 3, which states that there exists a sweet spot for segment length in MCTS rollouts.
Revised: This provides empirical support for Equation 3, which states that there exists a sweet spot for segment length in MCTS rollouts.

Original: However, in the simple
Revised: However, in the simple

Original: Hopper environment, it is inconclusive which segment length leads to the best return.
Revised: Hopper environment, it is inconclusive which segment length leads to the best return.

Original: 6. E MPIRICAL E VIDENCE OF T HEORY BEHIND POLISH
Revised: 6. E MPIRICAL E VIDENCE OF T HEORY BEHIND POLISH

Original: 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.
Revised: 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.

Original: 3a shows how the ﬁrst term in Equation 3 changes as a function of segment length, and 3b depicts an approximation of the second term.
Revised: 3a shows how the ﬁrst term in Equation 3 changes as a function of segment length, and 3b depicts an approximation of the second term.

Original: For the Ant environment, we show both metrics for a select range of segment lengths of (1, 4, 32, 256, 1000).
Revised: For the Ant environment, we show both metrics for a select range of segment lengths of (1, 4, 32, 256, 1000).

Original: As shown in our theoretical analysis, the MCTS return improvement (over the policy) generally increases with segment length, but so does the KL-divergence.
Revised: As shown in our theoretical analysis, the MCTS return improvement (over the policy) generally increases with segment length, but so does the KL-divergence.

Original: 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.
Revised: 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.

Original: 6. P ARALLEL I MPLEMENTATION S PEEDUP
Revised: 6. P ARALLEL I MPLEMENTATION S PEEDUP

Original: Lastly, we show that we are able to achieve a signiﬁcant speedup with a parallel implementation.
Revised: Lastly, we show that we are able to achieve a signiﬁcant speedup with a parallel implementation.

Original: Our algorithm, by design, is easy to parallelize over distributed compute (Lines 11, 12 in Algorithm 1).
Revised: Our algorithm, by design, is easy to parallelize over distributed compute (Lines 11, 12 in Algorithm 1).

Original: Every MCTS rollout from an initial state is independent, and thus can be conducted on separate workers.
Revised: Every MCTS rollout from an initial state is independent, and thus can be conducted on separate workers.

Original: The overhead of a parallel implementation, which arises from model synchronization across different workers and data transmission over the network, is likely insigniﬁcant compared to the workload of a MCTS rollout, especially for complex environments.
Revised: The overhead of a parallel implementation, which arises from model synchronization across different workers and data transmission over the network, is likely insigniﬁcant compared to the workload of a MCTS rollout, especially for complex environments.

Original: 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 .
Revised: 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 .

Original: Even with a small segment length, e.g., t = 32 , which implies a high number of parallel tasks (=32) over a ﬁxed trajectory length (=1000), the actual runtime is close to ideal.
Revised: Even with a small segment length, e.g., t = 32 , which implies a high number of parallel tasks (=32) over a ﬁxed trajectory length (=1000), the actual runtime is close to ideal.

Original: This suggests that we can easily achieve O(10)X speedup for OpenAI Gym-like environments.
Revised: This suggests that we can easily achieve O(10)X speedup for OpenAI Gym-like environments.

Original: The overhead outweighs the beneﬁt only when segment length is very small.
Revised: The overhead outweighs the beneﬁt only when segment length is very small.

Original:  C ONCLUSION AND FUTURE WORK
Revised:  C ONCLUSION AND FUTURE WORK

Original: In this work, we propose POLISH, a general imitation learning algorithm that provides ﬂexible local trajectory improvement based on MCTS.
Revised: In this work, we propose POLISH, a general imitation learning algorithm that provides ﬂexible local trajectory improvement based on MCTS.

Original: We provide theoretical analysis that sheds light on the beneﬁt of local trajectory improvement as well as a speciﬁc case study on MCTS.
Revised: We provide theoretical analysis that sheds light on the beneﬁt of local trajectory improvement as well as a speciﬁc case study on MCTS.

Original: 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.
Revised: 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.

Original: Model based methods have been proposed as a way to improve the sample efﬁciency 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).
Revised: Model based methods have been proposed as a way to improve the sample efﬁciency 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).

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: In particular, we can learn a model of the environment and use that model to perform MCTS.
Revised: In particular, we can learn a model of the environment and use that model to perform MCTS.

Original: 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.
Revised: 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.

Original: In this work, we focused on an expert feedback collection method with a ﬁxed local search horizon.
Revised: In this work, we focused on an expert feedback collection method with a ﬁxed local search horizon.

Original: 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.
Revised: 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.

Original: A P ROOFS FROM S ECTIONProof of Proposition
Revised: A P ROOFS FROM S ECTIONProof of Proposition

Original: Proof.
Revised: Proof.

Original: Let t (cid:48) =
Revised: Let t (cid:48) =

Original: kt for some positive integer k .
Revised: kt for some positive integer k .

Original: Since Algorithm 1 collects initial state every t and t (cid:48) steps respectively.
Revised: Since Algorithm 1 collects initial state every t and t (cid:48) steps respectively.

Original: The set of starting states collected with the interval t (cid:48) is a subset of those collected with an interval of t .
Revised: The set of starting states collected with the interval t (cid:48) is a subset of those collected with an interval of t .

Original: We will show that for the ﬁrst segment of length t (cid:48) starting from some state s 
Revised: We will show that for the ﬁrst segment of length t (cid:48) starting from some state s , J 0: t (cid:48) π ( π ∗ ) ≥ (cid:80) ki =1 γ t ( i − 1)

Original: J 0:
Revised: J t

Original: [Equation] ( π ∗ ) .
Revised: i : t i +1 π ( π ∗ ) .

Original: Similar proof holds for the other segments as well.
Revised: Similar proof holds for the other segments as well.

Original: [Equation]
Revised: [Equation]

Original: [Equation] [Equation] The inequality holds term-by-term because by Assumption 1, following π instead of π ∗ generates state distributions that induce a lower valued states.
Revised: The inequality holds term-by-term because by Assumption 1, following π instead of π ∗ generates state distributions that induce a lower valued states.

Original: To show that the sum is maximized at t =
Revised: To show that the sum is maximized at t =

Original: T , consider any t < T .
Revised: T , consider any t < T .

Original: [Equation]
Revised: [Equation]

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

Original: Proof of Proposition2:
Revised: Proof of Proposition

Original: Proof.
Revised: Proof.

Revised: Let f ( t ) =

Revised: t 2 1 − γ t .

Revised: To prove this proposition, it sufﬁces to show that f ( t ) is an increasing function for t > 0 .

Revised: First we take the derivative and get f (cid:48) ( t ) =

Original: [Equation]
Revised: 2 t (1 − γ t )+ t 2 γ t ln γ (1 − γ t ). Let g ( t ) = 2 t (1 − γ t ) + t 2 γ t ln γ be the numerator term.

Revised: It is sufﬁcient to show that g ( t ) > 0 for t > 0 .

Revised: Taking the derivative for g ( t ) , we obtain

Revised: [Equation]

Revised: >the last inequality follows from γ ∈ (0 , 1) .

Revised: So g ( t ) is increasing, i.e., g ( t ) > g (0)

Revised: = 0 .

Revised: It follows that f (cid:48) ( t ) > 0 and f ( t ) is an increasing function of t .

Revised: Proof of Proposition 3:

Revised: Proof.

Revised: [Equation]

Original: Proof of Theorem 2:
Revised: Proof of Theorem 2:

Original: We ﬁrst provide an analysis that extends theoretical guarantees of UCT (Kocsis et al., 2006; Kocsis & Szepesv ´ ari, 2006) to the setting with the modiﬁed UCT rule as in 4.2.
Revised: We ﬁrst provide an analysis that extends theoretical guarantees of UCT (Kocsis et al., 2006; Kocsis & Szepesv´ari, 2006) to the setting with the modiﬁed UCT rule as in 4.2.

Original: As usual, we connect UCT with the Upper-Conﬁdence 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.
Revised: As usual, we connect UCT with the Upper-Conﬁdence 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.

Original: Speciﬁcally, we assume there are K arms.
Revised: Speciﬁcally, we assume there are K arms.

Original: At time step t , arm i has a reward distribution X it which lies in the interval [0 , 1] for simplicity.
Revised: At time step t , arm i has a reward distribution X it which lies in the interval [0 , 1] for simplicity.

Original: Unlike in traditional multiarmed bandits, we allow X it to vary as a function of t .
Revised: Unlike in traditional multiarmed bandits, we allow X it to vary as a function of t .

Original: Following the standard procedure for analyzing UCT-type algorithms in (Kocsis et al., 2006; Kocsis & Szepesv ´ ari, 2006), we assume that the reward distributions concentrate.
Revised: Following the standard procedure for analyzing UCT-type algorithms in (Kocsis et al., 2006; Kocsis & Szepesv´ari, 2006), we assume that the reward distributions concentrate.

Original: Assumption 2.
Revised: Assumption 2.

Original: For each 1 ≤ i ≤ K , X in = 1 n (cid:80) nt =1 X it converges.
Revised: For each 1 ≤ i ≤ K , X in = 1 n (cid:80) nt =1 X it converges.

Original: Let µ in = E
Revised: Let µ in = E

Original: [ X in ] and µ i = lim n →∞ µ in .
Revised: [ X in ] and µ i = lim n →∞ µ in .

Original: 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 ( δ ) =
Revised: 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 ( δ ) =

Original: C p (cid:112) n ln(1 /δ ) , the following two tail probability inequalities are satisﬁed
Revised: C p (cid:112) n ln(1 /δ ) , the following two tail probability inequalities are satisﬁed

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: The UCB algorithm balances the exploration-exploitation problem by selecting arms with the highest upper conﬁdence bound.
Revised: The UCB algorithm balances the exploration-exploitation problem by selecting arms with the highest upper conﬁdence bound.

Original: We consider a related variant of UCB that weights the exploration term by some prior probabilities p i for arm i .
Revised: We consider a related variant of UCB that weights the exploration term by some prior probabilities p i for arm i .

Original: 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.
Revised: 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.

Original: Let
Revised: Let

Original: T i ( n ) be a random variable representing the total number of times that arm i is chosen in the ﬁrst n rounds.
Revised: T i ( n ) be a random variable representing the total number of times that arm i is chosen in the ﬁrst n rounds.

Original: We assume a unique optimal arm exists and its related quantities will be indexed by a ∗ , i.e., p ∗ , µ ∗ , T ∗ ( n ) , X ∗ t .
Revised: We assume a unique optimal arm exists and its related quantities will be indexed by a ∗ , i.e., p ∗ , µ ∗ , T ∗ ( n ) , X ∗ t .

Original: We ﬁrst upper bound the expected number of choices of a sub-optimal arm in terms of its sub-optimality ∆ i = µ ∗ − µ i and the prior probability.
Revised: We ﬁrst upper bound the expected number of choices of a sub-optimal arm in terms of its sub-optimality ∆ i = µ ∗ − µ i and the prior probability.

Original: Since µ in converges to µ i as n → ∞ , δ in converges to zero, so for any (cid:15) > 0 , there exists an index
Revised: Since µ in converges to µ i as n → ∞ , δ in converges to zero, so for any (cid:15) > 0 , there exists an index

Original: N 0 ( (cid:15) )
Revised: N 0 ( (cid:15) )

Original: such that if t ≥ N 0 ( (cid:15) ) , then | δ it | ≤ (cid:15) ∆ i 2 for every arm i .
Revised: such that if t ≥ N 0 ( (cid:15) ) , then | δ it | ≤ (cid:15) ∆ i 2 for every arm i .

Original: Theorem 3.
Revised: Theorem 3.

Original: Under Assumption 2, if we use c t,s = 2 C p (cid:113) ln ts for the modiﬁed UCB algorithm.
Revised: Under Assumption 2, if we use c t,s = 2 C p (cid:113) ln ts for the modiﬁed UCB algorithm.

Original: Fix 
Revised: Fix (cid:15) [Equation] [Equation]

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: We follow the proof in (Kocsis et al., 2006).
Revised: We follow the proof in (Kocsis et al., 2006).

Original: Deﬁne A 0 ( n, (cid:15) )
Revised: Deﬁne A 0 ( n, (cid:15) )

Original: =
Revised: =

Original: min { [Equation] | c t,s ≤ (1 − (cid:15) )∆ i 2 p i } .
Revised: min { s | c t,s ≤ (1 − (cid:15) ) ∆ i 2 p i } .

Original: Since c t,s = 2 C p (cid:113) ln ts , A 0 ( n, (cid:15) )
Revised: Since c t,s = 2 C p (cid:113) ln ts , A 0 ( n, (cid:15) )

Original: = (cid:100) 16 C 2 p p 2 i ln n (1 − (cid:15) )
Revised: = (cid:100) 16 C 2 p p 2 i ln n (1 − (cid:15) )

Original: 2 ∆ 2 i (cid:101) .
Revised: 2 ∆ 2 i (cid:101) .

Original: Let A ( n, (cid:15) )
Revised: Let A ( n, (cid:15) )

Original: =
Revised: =

Original: max( A 0 ( n, (cid:15) ) , N 0 ( (cid:15) ) , N p ) .
Revised: max( A 0 ( n, (cid:15) ) , N 0 ( (cid:15) ) , N p ) .

Original: Then
Revised: Then

Revised: T i ( n ) = 1 + n (cid:88) t = K +

Revised: I ( I t

Revised: = i

Original: [Equation]
Revised: [Equation]

Revised: (14) ≤ A ( n

Original: [Equation]
Revised: [Equation]

Revised: (cid:15)

Original: [Equation]
Revised: [Equation]

Original: We observe that X ∗ s + p ∗ c t,s ≤ X i
Revised: + n (cid:88) t

Original: [Equation]
Revised: [Equation]

Original:  (cid:48)
Revised: [Equation] [Equation] We observe that X ∗ s + p ∗ c t,s ≤ X i,s (cid:48)

Original: + p
Revised: + p

Original: i
Revised: i

Original: c t,s (cid:48) implies that at least one of the following is true
Revised: c t,s (cid:48) implies that at least one of the following is true

Original: [Equation]
Revised: [Equation]

Original: Since s (cid:48) ≥ A ( n, (cid:15) ) and by the construction of A ( n, (cid:15) )
Revised: Since s (cid:48) ≥ A ( n, (cid:15) ) and by the construction of A ( n, (cid:15) )

Original: , A ( n, (cid:15) ) ≥ A 0 ( n, (cid:15) ) and A ( n, (cid:15) ) ≥ N 0 ( (cid:15) ) , so the inequality (19) does not hold.
Revised: , A ( n, (cid:15) ) ≥ A 0 ( n, (cid:15) ) and A ( n, (cid:15) ) ≥ N 0 ( (cid:15) ) , so the inequality (19) does not hold.

Original: Thus, either (17) or (18) must hold.
Revised: Thus, either (17) or (18) must hold.

Original: So I ( X ∗ s + p ∗ c t,s ≤
Revised: So I ( X ∗ s + p ∗ c t,s ≤

Original: X i,s (cid:48)
Revised: X i,s (cid:48)

Original: + p
Revised: + p

Original: i
Revised: i

Original: c t,s (cid:48) )
Revised: c t,s (cid:48) )

Original: ≤
Revised: ≤

Original: I ( X ∗ s + p ∗ c t,s ≤ µ ∗ )
Revised: I ( X ∗ s + p ∗ c t,s ≤ µ ∗ )

Original: + I
Revised: + I

Original: ( X i,s (cid:48) ≥ µ i + p
Revised: ( X i,s (cid:48) ≥ µ i + p

Original: i
Revised: i

Original: c t,s (cid:48) ) .
Revised: c t,s (cid:48) ) .

Original: Now we can utilize inequalities (4) and (5) to derive upper bound on the two probabilities, so
Revised: Now we can utilize inequalities (4) and (5) to derive upper bound on the two probabilities, so

Original: [Equation]
Revised: [Equation]

Original: Plugging those quantities into (8) and use the relaxation that A ( n, (cid:15) )
Revised: Plugging those quantities into (8) and use the relaxation that A ( n, (cid:15) )

Original: ≤
Revised: ≤

Original: A [Equation] n, [Equation]
Revised: A 0 ( n, (cid:15) )

Original: ) + N 0 ( [Equation]
Revised: + N 0 ( (cid:15) )

Original: ) + N p , we obtain the desired result.
Revised: + N p , we obtain the desired result.

Original: 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 .
Revised: 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 .

Original: 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.
Revised: 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.

Original: 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).
Revised: 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).

Original: An accurate value estimation can effectively reduce the amount of lookahead necessary to identify an optimal action (Ng, 2003; Sun et al., 2018a).
Revised: An accurate value estimation can effectively reduce the amount of lookahead necessary to identify an optimal action (Ng, 2003; Sun et al., 2018a).

Original: On the other hand, an inaccurate value estimation can mislead the MCTS process.
Revised: On the other hand, an inaccurate value estimation can mislead the MCTS process.

Original: We provide a sufﬁcient condition under which MCTS is guaranteed to improve the policy.
Revised: We provide a sufﬁcient condition under which MCTS is guaranteed to improve the policy.

Original: It follows that, compared with the true
Revised: It follows that, compared with the true

Original: Q-function Q π ( s, a ) .
Revised: Q-function Q π ( s, a ) .

Original: [Equation]
Revised: [Equation]

Original: Before we ﬁnally prove the main theorem, we need the following lemma stating improvement when computed with respect to the estimated Q -function.
Revised: Before we ﬁnally prove the main theorem, we need the following lemma stating improvement when computed with respect to the estimated Q -function.

Original: Lemma 1.
Revised: Lemma 1.

Original: Let a ∗ = arg max a Q eπ ( s, a ) for a state s and (cid:15) ≤ (1 − π ( a ∗ | s )) min a (cid:54) =
Revised: Let a ∗ = arg max a Q eπ ( s, a ) for a state s and (cid:15) ≤ (1 − π ( a ∗ | s )) min a (cid:54) =

Original: a ∗ Q eπ ( s, a ∗ ) −
Revised: a ∗ Q eπ ( s, a ∗ ) −

Original: Q eπ ( s, a ) .
Revised: Q eπ ( s, a ) .

Original: Then as n → ∞ , E a ∼ π n ( s ) [ Q eπ ( s, a )] ≥ E a ∼ π ( s )
Revised: Then as n → ∞ , E a ∼ π n ( s ) [ Q eπ ( s, a )] ≥ E a ∼ π ( s )

Original: [ Q eπ ( s, a )] + (cid:15) .
Revised: [ Q eπ ( s, a )] + (cid:15) .

Original: Proof.
Revised: Proof.

Original: As n → ∞ , π n ( s ) becomes a delta distribution with probability mass 1 on the a ∗ , so
Revised: As n → ∞ , π n ( s ) becomes a delta distribution with probability mass 1 on the a ∗ , so

Original: [Equation]
Revised: [Equation]

Original: Finally
Revised: Finally

Original: [Equation]
Revised: [Equation]

Original: we prove the main theorem,
Revised: we prove the main theorem,

Original: [Equation]
Revised: [Equation]