Scalable Online Planning via Reinforcement Learning Fine-Tuning

1 Introduction . Lookahead search has been a key component of successful AI systems in sequential decision-making problems . For example , in order to achieve superhuman performance in go , chess and shogi , AlphaZero leveraged Monte Carlo tree search ( MCTS ) [ 38 ] . MuZero extended this even further to Atari games , again using MCTS [ 32 ] . Without MCTS , AlphaZero performs below a top human level , and more generally no superhuman Go bot has yet been developed that does not use some form of MCTS . Similarly , search algorithms were a critical component of AI successes in backgammon [ 45 ] , chess [ 10 ] , poker [ 27 , 7 , 8 ] , and Hanabi [ 22 ] . However , even though different search algorithms were used in each domain , all of them were tabular search algorithms , i.e. , a distinct policy was computed for each state encountered during search , without any function approximation to generalize between similar states . While tabular search has achieved great success , particularly in perfect-information deterministic environments , its applicability is clearly limited . For example , in the popular partially observable stochastic AI benchmark game Hanabi [ 5 ] , one-step lookahead search involves a search over about 500 possible next states . However , searching over all two-step joint policies would require a search over 2020 states , which is clearly intractable for tabular search . Additionally , unlike perfectinformation deterministic games where it is only necessary to search over a tiny fraction of the next several moves , partial observability and stochasticity make it impossible to limit the search to a tiny subset of all possible states . Fortunately , many of these states are extremely similar , so a search algorithm can in theory benefit by generalizing between similar states . This is the motivation for our non-tabular search algorithm . ∗Equal Contribution . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . In this paper we take inspiration from related research in continuous control environments that use non-tabular planning algorithms to improve performance at inference time [ 47 , 25 , 1 ] . These methods leverage finite-horizon model-based rollouts . Specifically , we replace tabular search with finetuning of the policy network at inference time . We show that with this approach we are able to achieve state-of-the-art performance in Hanabi . Specifically , our method is able to search multiple moves ahead and discover joint deviations , which in general is intractable using tabular search . We also show the generality of our approach by showing that it outperforms Monte Carlo tree search in deterministic and stochastic versions of the Atari game Ms. Pacman . 2 Background . 2.1 MDPs , POMDPs , and Dec-POMDPs . We consider Markov decision processes ( MDPs ) [ 42 ] , where the agent observes the state st ∈ S at time t , performs an action according to their policy at ∼ π ( st ) and receive reward rt = r ( st , at ) . The environment transit to the next state following transition probability st+1 ∼ P ( st+1|st , at ) . POMDPs [ 39 ] extend MDPs to the partially observable setting where the agent can not observe the true underlying state but instead receive information through a observation function ot = Ω ( st ) . Due to the partial observability , the agent policy π often needs to take into account the entire actionobservation history ( AOH ) denoted as τ ′t = { o0 , a0 , r0 , . . . ot } with τt = { s0 , a0 , r0 , . . . , st } representing the true underlying trajectory . Dec-POMDPs [ 29 , 5 ] extend POMDPs to the cooperative multi-agent setting . At each time step t , each agent i receives their individual observation from the state oit = Ω i ( st ) and selects action ait . The joint action of N player is a tuple at = ( a 1 t , a 1 t , ... , a N t ) and can be observed by all players . We denote the individual policy for each player as πi and joint policy as π . We use τ it = { oi0 , a0 , r0 , . . . oit } to represent the AOH from the perspective of agent i . 2.2 Tabular Search in Dec-POMDPs . SPARTA [ 22 ] is a tabular search algorithm previously applied in Dec-POMDPs that is proven to never hurt performance and in practice greatly improves performance . Specifically , SPARTA assumes that all agents in a game agree on a joint blueprint policy πbp and that one or more agents may choose to deviate from πbp at test time . We assume a two player setting in the following discussion . There are two versions of SPARTA . In single-agent SPARTA , only one agent may deviate from the blueprint . We denote the search agent as agent i and the other agent as−i . Agent i maintains their private belief over the trajectories they are in given their AOH , Bi ( τt ) = P ( τt|τ it ) . At action selection time , agent i compute the Monte Carlo estimation for the value of each action assuming both players will follow blueprint policy πbp thereafter , Q ( τ it , at ) = Eτt∼Bi ( τt ) , τT∼P ( τT |τt , at , πbp ) Rt ( τT ) , ( 1 ) where Rt ( τT ) = ∑T t′=t rt′ is the forward looking return from t to termination step T . Agent i will pick argmaxat Q ( τ i t , at ) instead of the blueprint action abp if maxat Q ( τ i t , at ) −Q ( τ it , abp ) > ϵ . In multi-agent SPARTA , both agents are allowed to search at test time . However , since agent i ’ s belief over states depends on agent −i ’ s policy , it is essential that agent i replicate the search policy of agent −i . Since agent i does not know the private state of agent −i , agent i must therefore compute agent −i ’ s policy for every possible AOH τ it agent i may have seen , a process referred to as range-search . Once a correct belief is computed , agent −i runs their own search . The rangesearch can be prohibitively expensive to compute in many cases because players might have millions of potential AOH . For example , Hanabi might have up to about 10 million AOHs . To mitigate this problem , the authors introduce a way to execute multi-agent search only when the range of AOHs is small enough that it is computationally feasible to do so , and the algorithm uses single-agent search otherwise . Even so , SPARTA multi-agent search is limited to one-ply search – each agent calculates the expected value of only the next action . The computational complexity grows exponentially with each additional ply . In contrast , our non-tabular multi-agent search technique makes it possible to conduct 2-ply , and deeper , search . 2.3 Reinforcement Learning . Our goal is to learn a stationary policy π ( a|s ) such that the expected non-discounted return J ( π ) = Es0∼ρ ( s ) Eπ ∑∞ t=0 r ( st , at ) is maximized , where ρ ( s ) is the initial state distribution . For any policy π we define its value Vπ ( st ) = Eπ ( ∑∞ l=0 r ( st+l , at+l ) |st ) , its Q-value Qπ ( st , at ) = Eπ ( ∑∞ l=0 r ( st+l , at+l ) |st , at ) and its advantage Aπ ( st , at ) = Qπ ( st , at ) − Vπ ( st ) . In reinforcement learning ( RL ) , the transition probability p and the reward r are unknown , and the policy is learned by sampling trajectories τ = ( ( s0 , a0 ) , ... , ( sT , aT ) ) from the environment . In Q-learning , we use the sampled trajectories to learn the Q-function of the optimal policy , given by the Bellman equation : Q ( s , a ) = Es′ ( r ( s , a ) + γmaxa′ Q ( s′ , a′ ) ) . The learned policy is greedy with respect to Q : π ( a|s ) = argmaxa Q ( s , a ) . In DQN , the Q-value is approximated with a neural network Qθ that is trained by taking gradients of the mean squared Bellman error : ∇θÊ ( Qθ ( s , a ) − ( r ( s , a ) + γmax a′ Qθ′ ( s ′ , a′ ) ) ) 2 ( 2 ) where θ′ is a fixed copy of θ and Ê denotes the empirical average over a finite number of batches . In policy gradient ( PG ) , we directly learn a parametrized policy πθ by performing gradient descent using an estimator of J : ∇θÊ ( log πθ ( a|s ) Â ( s , a ) ) ( 3 ) where Â is an estimate of the advantage of πθ . 2.4 Decision-Time Planning . To make better decisions at test time , we can specialize the learned policy , which we refer to as the blueprint policy ( also known as the prior policy ) , to the current state by using it as prior in an online planning algorithm . One of the most popular planning algorithms is Monte Carlo tree search ( MCTS ) , which has famously enabled superhuman performance in go , chess , and many other settings [ 36 , 37 , 38 , 32 ] . MCTS builds a tree of potential future states starting from the current state . Each node s keeps track of its action-state visitation counts N ( s , a ) and an estimate of its Q-values Q ( s , a ) , refined each time the node is visited . The prior policy is used to select the next action to take during the tree traversal . MCTS don ’ t scale well in environments with large branching factors caused by a high amount of stochasticity or partial observability . For example , MCTS becomes very expensive as the size of the action space increases because it needs to build an explicit tree to keep statistics about state-action pairs , and nodes need to be visited several times for the method to be effective . Furthermore , MCTS can in the worst case expand very deeply a branch that it misidentifies to be optimal due to a lack of exploration – exacerbated in high dimensions– which leads to a worst-case sample complexity worse than uniform sampling [ 28 ] . 3 Reinforcement Learning Fine-Tuning . To alleviate the inefficiency of tabular search methods like MCTS in large state and action spaces , we propose to formulate online planning as a RL problem that should be quickly solvable given a sufficiently optimal blueprint policy . We do so by using RL as a multiple-step policy improvement operator [ 12 ] . More specifically , we bias the initial state distribution towards the current state s∗ and we reduce the horizon of the problem , leading to the following truncated objective : max π̃ ( Es0∼ρs∗ ( s ) Ea0 : H−1 , s1 : H∼π̃ H−1∑ t=0 ( r ( st , at ) ) + EaH : ∞ , sH+1 : ∞∼π ∞∑ t=H ( r ( st , at ) ) ) ( 4 ) where ρs∗ ( s ) = 1 ( s = s∗ ) and π is the blueprint policy . If π is optimal , then π′ = π is a solution to the problem . We call this procedure RL Fine-Tuning or RL Search and we use the two terms interchangeably in the remaining . The blueprint π is either directly parameterized by θ or is greedy with respect to a Q-value parameterized by θ . In both cases , we improve π for the next H steps by following the gradient of the truncated objective with respect to θ . We present the two approaches in more detail in the following subsections .

Summary: This paper proposes to replace tabular search methods like MCTS with fine tuning of neural policies for online planning.   The motivation is that fine tuning the neural policies online may be more efficient than doing MCTS which has no generalization between states during the search.   The key algorithmic contribution is the biased RL objective used for more efficient online fine tuning of the neural policy.