Batch Reinforcement Learning Through Continuation Method

1 INTRODUCTION . While RL is fundamentally an online learning paradigm , many practical applications of RL algorithms , e.g. , recommender systems [ 5 , 7 ] or autonomous driving [ 36 ] , fall under the batch RL setup . Under this setting , the agent is asked to learn its policy from a fixed set of interactions collected by a different ( and possibly unknown ) policy commonly referred to as the behavior policy , without the flexibility to gather new interactions . Realizing the interactive nature of online RL has been hindering its wider adoptions , researchers strive to bring these techniques offline [ 24 , 11 , 20 , 23 , 31 , 12 , 21 , 2 , 32 , 8 ] . We focus on policy optimization under batch RL setup . As pointed out in [ 3 , 26 ] , even with access to the exact gradient , the loss surface of the objective function maximizing the expected return is difficult to optimize , leading to slow convergence . Chen et al . [ 8 ] show that the objective function of expected return exhibits sub-optimal plateaus and exponentially many local optima in the worst case . Batch setup makes the learning even harder as it adds large variance to the gradient estimate , especially when the learned policy differs from the behavior policy used to generate the fixed trajectories . Recent works propose to constrain the size of the policy update [ 27 , 28 ] or the distance between the learned policy and the behavior policy [ 14 , 21 ] . The strength of that constraint is a critical hyperparameter that can be hard to tune [ 28 ] , as a loose constraint does not alleviate the distribution shift while a strict one results in conservative updates . Here we propose to address the challenges using continuation methods [ 35 , 6 , 17 ] . Continuation methods attempt to solve the global optimization problem by progressively solving a sequence of new objectives that can be optimized more efficiently and then trace back the solutions to the original one . We change the objective function of policy optimization by including an additional term penalizing the KL divergence between the parameterized policy ⇡✓ and the behavior policy . We then gradually decrease the weight of that penalty , eventually converging to optimizing the expected return . With this additional constraint , we benefit from more accurate policy evaluation in the early stage of training as the target policy is constrained to be close to the behavior policy . As training continues , we relax the constraint and allow for more aggressive improvement over the behavior policy as long as the policy evaluation is still stable and relatively reliable , i.e . with a small enough variance . By doing so , the proposed method exhaustively exploits the information in the collected trajectories while avoiding the overestimation of state-action pairs that lack support . The contributions of this paper are as follows : ( 1 ) We propose a soft policy iteration approach to batch RL through the continuation method . ( 2 ) We theoretically verify that in the tabular setting with exact gradients , maximizing KL regularized expected return leads to faster convergence than optimizing the expected return alone . Also , our method converges to the globally optimal policy if there are sufficient data samples for accurate value estimation . ( 3 ) We demonstrate the effectiveness of our method in reducing errors in value estimation using visualization ; ( 4 ) We empirically verify the advantages of our method over existing batch RL methods on various complex tasks . 2 RELATED WORK . Batch Reinforcement Learning . Off-policy reinforcement learning has been extensively studied [ 11 , 20 , 30 , 23 , 31 ] , with many works [ 12 , 21 , 2 ] focusing on variants of Q-learning . Fujimoto et al . [ 12 ] , Kumar et al . [ 21 ] investigated the extrapolation error in batch RL resulting from the mismatch of state-action visitation distribution between the fixed dataset and the current policy , and proposed to address it by constraining the action distribution of the current policy from deviating much from the training dataset distribution . Recent works [ 29 , 33 ] studied policy iteration under batch RL . The Q function is estimated in the policy evaluation step without special treatment while the policy updates are regularized to remain close to the prior policy with a fixed constraint . To further reduce uncertainty in Q learning , an ensemble of Q networks [ 21 , 29 ] and distributional Q-function [ 2 , 33 ] are introduced for the value estimation . [ 34 , 18 ] use the KL divergence between the the target policy and the behavior policy as a regularization term in the policy update and/or value estimation . The constraint is controlled by a fixed weight of the KL regularization or a fixed threshold for the KL divergence . While all of these works apply a fixed constraint determined by a sensitive hyperparameter to control the distance between the behavior/prior policy and the target policy , we focus on gradually relaxed constraints . Constrained Policy Updates . Several works [ 27 , 1 , 15 ] studied constrained policy updates in online settings . Kakade & Langford [ 19 ] show that large policy updates can be destructive , and propose a conservative policy iteration algorithm to find an approximately optimal policy . Schulman et al . [ 27 ] constrain the KL divergence between the old policy and new policy to guarantee policy improvement in each update . Grau-Moya et al . [ 15 ] force the policy to stay close to a learned prior distribution over actions , deriving a mutual-information regularization between state and action . Cheng et al . [ 9 ] propose to regularize in the function space . Again these methods focused on a fixed constraint while we are interested in continuing relaxing the constraint to maximize the expected return eventually . Also none of these methods have been extensively tested for batch RL with fixed training data . Continuation Method . Continuation method [ 35 ] is a global optimization technique . The main idea is to transform a nonlinear and highly non-convex objective function to a series of smoother and easier to optimize objective functions . The optimization procedure is successively applied to the new functions that are progressively more complex and closer to the original non-context problem , to trace their solutions back to the original objective function . Chapelle et al . [ 6 ] use the continuation method to optimize the objective function of semi-supervised SVMs and reach lower test error compared with algorithms directly minimizing the original objective . Hale et al . [ 17 ] apply the continuation method to l1-regularized problems and demonstrate better performance for compressed sensing problems . Inspired by prior works , we employ the continuation method to transform the objective of batch RL problems by adding regularization . We gradually decrease the regularization weight to trace the solution back to the original problem . 3 METHOD . In classical RL , an agent interacts with the environment while updating its policy . At each step t , the agent observes a state st 2 S , selects an action at 2 A according to its policy to receive a reward rt = r ( st , at ) : S ⇥ A ! R and transitions to the next state st+1 ⇠ P ( ·|st , at ) . The state value of a policy ⇡ at a state s is V ⇡ ( s ) = Es0=s , at⇠⇡ ( ·|st ) , st+1⇠P ( ·|st , at ) [ P1 t=0 tr ( st , at ) ] . 2 [ 0 , 1 ] is the discounting factor . At each step , the agent updates the policy ⇡ so that the expected return V ⇡ ( ⇢ ) = Es⇠⇢ [ V ⇡ ( s ) ] ( where ⇢ is the initial state distribution ) is maximized . In batch RL , the agent is not allowed to interact with the environment during policy learning . Instead it has access to a fixed set of trajectories sampled from the environment according to a behavior policy1 . A trajectory { ( s0 , a0 , r0 ) , ( s1 , a1 , r1 ) , · · · , ( sT , aT , rT ) } is generated by sampling s0 from the initial state distribution ⇢ , sampling the action at ⇠ ( ·|st ) at the state st and moving to st+1 ⇠ P ( ·|st , at ) for each step t 2 [ 0 , 1 , · · · , T ] . The length T can vary among trajectories . We then convert the generated trajectories to a dataset D = { ( si , ai , ri , s0i ) } Ni=1 , where s0i is the next state after si in a trajectory . The goal of batch RL is to learn a parameterized policy ⇡✓ with the provided dataset to maximize the expected return V ⇡ ( ⇢ ) . In Sec . 3.1 , we will first introduce a new objective function Ṽ ⇡ , ⌧ ( ⇢ ) , i.e . the expected return of policy ⇡ with KL regularization term and the regularization weight ⌧ . With exact gradients , Ṽ ⇡ , ⌧ ( ⇢ ) can be optimized more efficiently than the original objective V ⇡ ( ⇢ ) . With the 1If the behavior policy is not known in advance , it can be fitted from the data [ 30 , 7 ] . continuation method , solving a sequence of optimization problems for Ṽ ⇡ , ⌧ ( ⇢ ) with decaying value of ⌧ converges toward optimizing V ⇡ ( ⇢ ) and makes the optimization easier . In Sec . 3.2 , we derive soft policy iteration with KL regularization to optimize Ṽ ⇡ , ⌧ ( ⇢ ) , without the assumption of exact gradients . Finally , in Sec . 3.3 , we propose a practical batch RL algorithm with value estimation for target policy based on this theory . 3.1 OPTIMIZING EXPECTED RETURN WITH KL REGULARIZATION . In batch RL , the distribution of the trajectories generated by the behavior policy can be very different from that of the learned policy . We thus restrict the learned policy to stay close to the behavior policy via the regularization of KL divergence . Define the soft state value of a policy ⇡ at a state s as Ṽ ⇡ , ⌧ ( s ) = Es0=s , at⇠⇡ ( ·|st ) , st+1⇠P ( ·|st , at ) '' 1X t=0 t ✓ r ( st , at ) ⌧ log ⇡ ( at|st ) ( at|st ) ◆ # , ( 1 ) where the temperature parameter ⌧ controls the deviation from . The new objective function becomes Ṽ ⇡ , ⌧ ( ⇢ ) = Es⇠⇢ [ Ṽ ⇡ , ⌧ ( s ) ] . This KL regularized objective differs from the original objective V ⇡ ( ⇢ ) , which however can be recovered as ⌧ ! 0 . As pointed out in [ 3 ] , even with exact gradients , the objective function V ⇡ ( ⇢ ) is still difficult to optimize due to its highly non-smooth landscape . Mei et al . [ 26 ] further prove that , in a tabular setting with softmax parameterized policy and exact gradients , the vanilla policy gradient method ( i.e . directly updating the parameters of policy ⇡ to maximize V ⇡ ( ⇢ ) with gradient descent ) converges to the global optimal policy at a convergence rate O ( 1/t ) , while the entropy-regularized policy gradient enjoys a significantly faster linear convergence rate O ( e t ) . Motivated by this line of work , we investigate the convergence rate of optimizing Ṽ ⇡ , ⌧ ( ⇢ ) with the exact gradient descent and compare it with the vanilla policy gradient method . We study the smoothness and Łojasiewicz inequality for the function Ṽ ⇡ , ⌧ ( ⇢ ) to prove the convergence rate , similar to [ 26 ] . The detailed proofs of all following theorems are provided in the appendix . Theorem 1 . In the tabular setting with softmax parameterized policy ⇡✓ , maximizing Ṽ ⇡ , ⌧ ( ⇢ ) using policy gradient with the learning rate ⌘ = ( 1 ) 3 ( 8M+⌧ ( 4+8 logA ) ) , for all t > 1 , we have Ṽ ⇡ ⇤ ⌧ , ⌧ ( ⇢ ) Ṽ ⇡✓t , ⌧ ( ⇢ ) C · e C⌧ ( t 1 ) · M + ⌧ logA ( 1 ) 2 where ⇡⇤ ⌧ is the optimal policy maximizing Ṽ ⇡ , ⌧ ( ⇢ ) , M is the bound of the absolute value of r ( s , a ) + ⌧ log ( a|s ) , A is the size of action space , S is the size of state space , C⌧ / ( 1 ) 4 ( 8M/⌧+4+8 logA ) ·S , and C is a constant independent with t and ⌧ . Theorem 1 states that KL regularized expected return Ṽ ⇡ , ⌧ ( ⇢ ) can be optimized with a convergence rate O ( e t ) rather than the O ( 1/t ) , the convergence rate of vanilla policy gradient for expected return alone . The faster convergence inspires us to optimize Ṽ ⇡ , ⌧ ( ⇢ ) to reach policy ⇡⇤ ⌧ , then use ⇡⇤ ⌧ as initialization , gradually decrease the temperature ⌧ towards 0 , and eventually move from ⇡⇤ ⌧ to ⇡⇤ = argmax⇡ V ⇡ ( ⇢ ) . With a reasonable value of ⌧ , we enjoy a linear convergence rate toward ⇡⇤⌧ from the randomly initialized policy ⇡✓ . As ⌧ decreases , ⇡⇤⌧ gets closer to ⇡⇤ . The final optimization of V ⇡✓ ( ⇢ ) from ⇡⇤ ⌧ can be much faster than from a randomly initialized ⇡✓ . ( a ) 0 2000 4000 6000 8000 10000 iteration 0.0 0.2 0.4 0.6 0.8 1.0 V π θ i ( ρ ) vanilla continuation ( b ) Figure 1 : ( a ) A grid world with sparse rewards . ( b ) Learning curve of the value of learned policy ⇡✓i . We conduct a hyper-parameter search for the learning rate 5,1,0.5,0.1,0.05,0.01,0.005,0.001 and report the best performance for each method . We construct a toy example to illustrate this motivation . In the grid world ( Fig . 1a ) , the start state , annotated with ‘ S ’ , is in the center and the terminal states are marked in yellow . There are only two states with positive rewards ( 0.9 and 1 ) . There are four actions { up , down , left , right } . A badly initialized policy ⇡✓0 is shown as arrows in Fig . 1a ) . The initialization results in a poor policy , having high tendency to go right toward a terminal state with zero reward . The vanilla policy gradient method ( i.e . maximizing V ⇡ ( ⇢ ) with true gradient ) starting from this initial point takes more than 7000 iterations to escape a sub-optimal solution ( Fig . 1b ) . In contrast , we escape the sub-optimal solution much faster when ap- plying the continuation method to update the policy with the gradients of Ṽ ⇡ , ⌧ ( ⇢ ) , where the behavior policy ( ·|s ) = [ u1 , u2 , u3 , u4 ] with ui , i = 1 , 2 , 3 , 4 randomly sampled from U [ 0 , 1 ] and Algorithm 1 Soft Policy Iteration through Continuation Method 1 : Initialize : actor network ⇡✓ , ensemble critic network { Q ( 1 ) , Q ( 2 ) , · · · , Q ( K ) } , behavior pol- icy network , penalty coefficient ⌧ , decay rate , number of iterations I for each ⌧ 2 : Input : training dataset D = { ( si , ai , ri , s0i ) } Ni=0 3 : for update j = 0 , 1 , · · · do 4 : Sample batch of data { ( si , ai , ri ) } Bi=1 from D 5 : # Learn the behavior policy with behavior cloning objective 6 : Update to maximize 1 B P B i=1 log ( ai|si ) 7 : # Train the critic network 8 : Update ( k ) to minimize the temporal difference 1 B P B i=1 ri + V ( s0i ) Q ( k ) ( si , ai ) 2 9 : where V ( s ) = 1 K P K k=1 Ea⇠⇡✓ ( ·|s ) ( Q ( k ) ( s , a ) ) ⌧KL ( ⇡✓ ( ·|s ) | ( ·|s ) ) 10 : # Train the actor network 11 : Update ✓ to maximize 1 B P B i=1 h 1 K P K k=1 Ea⇠⇡✓ ( ·|si ) ( Q ( k ) ( si , a ) ) ⌧KL ( ⇡✓ ( ·|si ) | ( ·|si ) ) i 12 : # Decay the weight of KL regularization ⌧ for every I updates 13 : if j mod I = 0 then 14 : ⌧ ⌧ ⇤ 15 : end if 16 : end for normalized for each state s. In Fig . 1b , as we decrease ⌧ , the value of learned policy ⇡✓i for each iteration i quickly converges to the optimal value . In other words , optimizing a sequence of objective functions Ṽ ⇡ , ⌧ ( ⇢ ) can reach the optimal solution for V ⇡ ( ⇢ ) significantly faster .

Summary: The paper extends soft actor-critic (SAC) to the batch RL setting, replacing the policy entropy in the objective function with the KL divergence from the behavioral policy. The temperature parameter tau weighting the reward agains the KL term is annealed towards zero during the optimization process, which corresponds to starting with behavioral cloning for high values of tau and ending up with the standard reward maximization RL objective for tau=0. Theoretical analysis and experiments confirm the advantages of the proposed method.