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 a new planning method, aiming at improving performance in a large and realistic environment. The proposed method performs as a planning module and can be integrated with different types of online learning agents. Different from tabular search algorithms, the method formulates the planning step as a reinforcement learning problem and updates the policy or action-value function by bootstrapping. Thus, the method reduces the horizon of search and gains learning efficiency in large search spaces and continuous control problems.