Policy Optimization by Local Improvement through Search

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 & Szepesvá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 ( Daumé 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 . 1We 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 Various solutions to the problem of distributional shift have been proposed ( Daumé 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 ( Daumé 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 & Szepesvá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 andD 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 JD ( π ) = Eτ∼π [ ∑∞ i=0 γ ir ( si , ai ) ] . 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 γir ( si , ai ) |s0 = s ] , Qπ ( s , a ) = Eτ∼π [ ∞∑ i=0 γir ( si , ai ) |s0 = s , a0 = 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 dtπ denote the state distribution obtained by following π for t steps . We use dπ = ( 1−γ ) ∑∞ t=0 γ tdtπ to denote the accumulated discounted state distribution . With these quantities , we can rewrite J ( π ) = ∑∞ t=0 Es∼dtπ , a∼π ( s ) [ γ tr ( s , a ) ] = Es∼dπ , a∼π ( s ) [ r ( s , a ) ] . We will make use of the following relationship between the two policies in our analysis ( Kakade & Langford , 2002 ) : J ( π′ ) = J ( π ) + Es∼dπ [ Ea∼π′ ( s ) [ Aπ ( s , a ) ] ] ( 1 ) Throughout the paper , we use π∗ to denote the expert policy .

Summary: This paper proposes POLISH, a reinforcement learning learning algorithm based on imitating partial trajectories produced by an MCTS procedure. The intuition behind this idea is that behavioral cloning suffers from distribution shift over time, and using MCTS allows imitation learning to be done on states closer to the policy's state distribution, which the authors justify using techniques similar to DAgger. The authors evaluate this method on continuous OpenAI Gym tasks, and show that it consistently beats a PPO baseline.