Stability and Generalisation in Batch Reinforcement Learning

1 INTRODUCTION . A central aim of machine learning theory is to provide wort-case bounds on overfitting . However , in reinforcement learning ( RL ) , which typically employs online learning on an unlimited stream of data and does not minimise an empirical risk , generalisation is more difficult to characterise than in the supervised learning case . This said , overfitting is still observed in RL settings , with various definitions of data scarcity . This work looks to characterise generalisation and overfitting in batch , off-policy RL . Batch RL has been of particular recent interest , for reasons of data efficiency , policy safety , or as a component of a more complex online algorithm . Despite this , overfitting has never been directly characterised in this context . Fortunately , batch RL makes use of a limited data regime close to the ERM formulation , allowing us to contribute such an analysis . Algorithmic stability ( Bousquet & Elisseeff , 2002 ) describes the change in the model learned by an algorithm when the data employed is changed at a single point . It is a useful property , as it bounds the generalisation gap–the difference between error on our dataset , and the true population error . As such , if we can bound the stability of a batch RL algorithm , we can also make guarantees about how prone it is to overfitting : how well the error on our dataset represents the error across the whole MDP . Specifically , this work formulates an off-distribution notion of uniform stability ( Bousquet & Elisseeff , 2002 ) , suitable for batch policy evaluation algorithms . We show the relationship between this variant of stability and an off-policy generalisation gap . This allows us to characterise overfitting in batch RL via stability , answering the question : when can we expect a batch RL algorithm to generalise well ? This framework is then applied to a partially fitted variant of temporal difference ( TD ) learning , that bridges the gap between fitted and traditional TD methods . Importantly , no realisability assumptions about the ability of the function class to fit values are needed . The iterative , partially fitted batch method employed here means that our result characterises the use of target networks ( Mnih et al. , 2015 ) , common in deep RL algorithms . Our bound expresses a bias-variance trade-off in the number of steps taken between updates of the target network–as more steps are taken , less data is incorporated into the update dynamics , leading to a higher variance , but lower bias procedure , as the value targets are also less correlated with the value parameters . 2 RELATED WORK . Our work builds directly on the results of Hardt et al . ( 2016 ) , which establishes stability bounds for stochastic gradient methods in non-convex optimisation settings . That work in turn builds on algorithmic stability , a topic in learning theory that characterises the generalisation gap in a manner agnostic to model capacity ( Rogers & Wagner , 1978 ; Devroye & Wagner , 1979 ; Bousquet & Elisseeff , 2002 ) . Bounds for batch RL have seen significant recent interest , especially for linearly parameterised value estimators . Existing bounds in this setting study generalisation by directly bounding the suboptimality of value function estimators , rather than the generalisation gap . Most relevant here is the work of Duan & Wang ( 2020 ) , where tight upper bounds on finite time regret are given for linear value function approximators applied to fitted off-policy evaluation , but restricted to the finite-horizon RL problem , where data can be separated into specific trajectories , and under a function approximation class that is closed under Bellman backups . Fan et al . ( 2020 ) provides generalisation bounds in deep RL with use of a target network but these are capacity dependent and do not characterise partial fitting . Other key bounds for batch fitted methods include those of Murphy ( 2005 ) ; Antos et al . ( 2007 ) ; Munos & Szepesvári ( 2008 ) ; Tosatto et al . ( 2017 ) , though settings and assumptions differ from those used here . Zhang et al . ( 2021 ) also characterise the convergence properties of a target network used with linear function approximation . Furthermore , exponential worst-case lower bounds have been constructed for batch RL . Chen & Jiang ( 2019 ) show that concentration assumptions are needed for polynomial-time learning without strong assumptions on the function approximation class . Zanette ( 2021 ) ; Wang et al . ( 2020 ) ; Amortila et al . ( 2020 ) provide exponential lower bounds for batch RL with only realizability and coverage assumptions . Wang et al . ( 2021 ) investigate these lower bounds empirically for fitted Q-iteration , and provides upper bounds using either realizability , or strong concentration assumptions . Direct characterisation of overfitting in RL has seen more interest recently , though the vast majority of work thus far is empirical . To our knowlege , Wang et al . ( 2019 ) provide the only theoretical characterisation of the generalisation gap beyond this work , though it applies in the online setting , depends on the model complexity , and makes the restrictive assumption that the reparameterization trick ( Kingma & Welling , 2013 ) can be applied to the random variables of the MDP . Francois-Lavet et al . ( 2019 ) characterise overfitting asymptotically for batch RL as applied to partially observable Markov decision processes with finite observation , state , and action spaces . Empirical characterisation of overfitting in RL across several settings can be found in Zhang et al . ( 2018a ) , Zhang et al . ( 2018b ) ; Cobbe et al . ( 2019 ) , and Packer et al . ( 2019 ) . 3 PRELIMINARIES . We consider the infinite horizon discounted RL setting . The agent interacts with an environment , formalised as a Markov Decision Process , M with ( potentially infinite ) state space S , a finite action space A , transition kernel P : S × A × S → [ 0 , 1 ] , bounded stochastic reward kernel R : S × A → P ( −rmax , rmax ) , where P ( a , b ) is a probability distribution with bounded support on the interval [ a , b ] and scalar discount factor γ ∈ [ 0 , 1 ) . The agent ’ s behaviour is determined by a policy that maps a state to a distribution over actions : pi : S × A → [ 0 , 1 ] . We seek to optimise ( in the control case ) , or estimate ( in the policy evaluation case ) the expected discounted sum of future rewards starting from a given state . This quantity is given by the state value function , V ( S ) = ∑ A∈A π ( S , A ) Q ( S , A ) , with Q : S × A → [ −rmax/ ( 1 − γ ) , rmax/ ( 1 − γ ) ] , the action value function , given recursively through the Bellman equation : Q ( S , A ) = r̄ ( S , A ) + γES′∼P ( S , A ) , A′∼π ( S′ ) [ Q ( S′ , A′ ) ] , ( 1 ) where r̄ ( S , A ) is the mean reward as sampled from the reward kernel . The Bellman operator T π projects functions forwards by one step through the dynamics of the MDP : T π ( Q ) ( S , A ) = r̄ ( S , A ) + γES′∼P ( S , A ) [ EA′∼π ( S′ ) [ Q ( S′ , A′ ) ] ] . ( 2 ) T is a γ-contractive mapping and thus has a fixed point . The fixed point of this operator corresponds to the true value of the policy in question ( Puterman , 2014 ) . We make use of a feature mapping to encode states and actions , φ : S×A : → Rd . When estimating MDP values , we employ parameterised value functions with a linear approximation : Qw ( φ ( S , A ) ) = φ > ( S , A ) w . 3.1 BATCH RL . In batch RL , rather than interacting directly with the MDP , the algorithm is given a fixed dataset , with samples drawn from the transition and reward kernels of the MDP . The agent uses these samples to estimate either the value of a fixed target policy π , or to learn an optimal policy , π∗ . Usually , the state-action distribution induced by the target policy differs from that of the dataset , so batch RL is generally viewed as an extreme case of off-policy RL , in which additional data can not be generated from any policy . In this paper , we focus on off-policy evaluation , leaving control to future work . Fitted temporal difference methods are a class of batch RL algorithms that iteratively fix a target estimator for the value function , g : S × A → R , and then minimise the mean squared empirical Bellman error ( MSBE ) of a separate estimator , fw : S ×A → R , over a batch of data : 1 D D∑ i=1 min w ( fw ( si , ai ) − T̂ g ( si , ai ) ) 2 , ( 3 ) where T̂ is the pointwise empirical Bellman operator : T̂ π ( Q ) ( S , A ) = r + γ ∑ A′ π ( A′|S′ ) Q ( S′ , A′ ) , ( 4 ) and r and S′ are sampled by taking action A in state S. Once the minimum is attained , g is updated to match fw . This procedure repeats until convergence . In practice , particularly when using neural networks , the minimiser of ( 3 ) can not be found in closed form . Furthermore , it may be too computationally expensive , or undesirable from the perspective of learning dynamics , to reach this optimum . Instead , partial optimisation is performed : a few updates of an iterative optimisation procedure often suffice . We summarise this process using linear function approximation and expected SARSA as our off-policy estimator of the Bellman error in Algorithm 1 . Algorithm 1 Partially Fitted Linear Expected SARSA Input : D ∼ Dµ , learning rate α , target policy π Output : value parameters w 1 : initialize value parameters w0 2 : for epoch i = 0 , ... , N do 3 : fix w̄ = wiK 4 : sample minibatch of size K , U = { ( Sj , Aj , Rj , S ′ j ) } 5 : for update k = 1 , ... , K do 6 : t = iK + k 7 : δu ( w ) = ( φ ( Sk , Ak ) > w −Rk − γ ∑ A′ π ( S ′ k , A ′ ) φ ( S′k , A ′ ) > w̄ ) 8 : wt = wt−1 + αφ ( Su , Au ) δu ( wt−1 ) 9 : end for 10 : project w = argmin||w′||≤wmax ||w − w ′|| 11 : end for 12 : return wt Algorithm 1 also includes a projection step for the purposes of the theory in Section 4 . After each epoch i , the value parameters , w are projected to a closed ball of fixed radius . This prevents the step sizes from growing too quickly , as our update rule is non-Lipschitz for unbounded parameter spaces . By varying the hyperparameters of Algorithm 1 , we move between traditional and fitted temporal difference algorithms . If the minimiser of ( 3 ) exists , setting K large enough to allow for convergence recovers fitted expected SARSA , while setting K = 1 recovers expected SARSA applied to the uniform distribution over data . Update Rules and Parameterisation One challenge that arises when treating the stability of Algorithm 1 comes from the fact that we need to reason over the behaviour of two separate sets of parameters–both the value and target parameters affect the update rule . In order to manage this , we propose a unified parameterisation for which we maintain a single set of parameters , ŵ which combines the value parameters and the frozen ones : ŵ : = [ w w̄ ] With this setup we need to use two update rules that operate on the same set of parameters . The parameter fitting update is given by : G ( ŵ ) s , a , r , s′ : = ŵ − α [ φ ( s , a ) 0 ] ( φ ( s , a ) > w − r − γ ∑ a′ π ( a′|s′ ) φ ( s′ , a′ ) > w̄ ) ( 5 ) Where 0 is the zero vector in d dimensions . This is identical to performing the fitted TD update on w and does not alter the value of w̄ . When we need to update the frozen parameters , we simply use an assignment : Ĝ ( ŵ ) = [ w w ] ( 6 ) This setup allows our update rules to depend only on the data being used to perform that update , rather than having an implicit dependency on previous updates , due to the fact that there are two sets of parameters being maintained .

Summary: The paper aims to study the effect of stability on generalization of a particular off-policy policy evaluation algorithm. Specifically, the authors study a version of Fitted Q-evaluation where the iterative procedure is instead thought of as a gradient update and is performed by only partially fitting the new Q-function to the current target Q function. The authors look at a linear function approximation setting and connect the recent results on stability of stochastic gradient methods (Hardt et al, 2016) with their method and show stability results, and connect that with overfitting of the policy evaluation scheme. The authors discuss the effect of the number of updates performed in each epoch during the partially fitted linear-update scheme and further show empirical simulations.