Robust Losses for Learning Value Functions

1 INTRODUCTION . Learning value functions from off-policy data remains an open challenge due to high-variance samples and the inability to optimize the objective of interest . Progress towards this goal has been made over years of algorithm development , by reducing the variance for temporal difference ( TD ) algorithms ( Precup et al. , 2000 ; Munos et al. , 2016 ; Mahmood et al. , 2017 ) ; following approximate gradients of a proxy objective—the mean squared Bellman error ( MSBE ) —which upper bounds our objective of interest ( Dai et al. , 2017 ; 2018 ; Feng et al. , 2019 ) ; and following an approximate gradient of a projected version of the MSBE ( Sutton et al. , 2009 ; Maei et al. , 2009 ; Mahadevan et al. , 2014 ; Liu et al. , 2016 ; Ghiassian et al. , 2020 ) . For many years , however , there was only a limited set of choices—mainly vanilla TD algorithms . Because it was unclear precisely what objective TD was optimizing ( Baird , 1995 ; Antos et al. , 2008 ) , it was difficult to extend the algorithm . The known alternative—optimizing the MSBE—suffers from an issue of double sampling in the absence of a simulator ( Baird , 1995 ; Scherrer , 2010 ) . Driving the recent innovations are two key advances for objectives in RL . The first is the formalization of the objective underlying TD , called the mean squared projected Bellman error ( MSPBE ) ( Antos et al. , 2008 ; Sutton et al. , 2009 ) , which projects the Bellman error into the space spanned by the function approximator . Many algorithms are built on the originally introduced variants , GTD2 and TD with gradient corrections ( TDC ) ( Sutton et al. , 2009 ) . These algorithms , however , are limited to linear function approximation because the MSPBE is defined only for the linear case . The second advance is the introduction of a conjugate form for the MSBE to handle the double sampling problem ( Dai et al. , 2017 ) . By transforming the MSBE using biconjugate functions , the double sampling problem instead became a better understood saddlepoint optimization problem . The SBEED algorithm ( Dai et al. , 2018 ) later extended the conjugate MSBE to control using a smoothed Bellman optimality operator and parameterizing both the policy and value function estimates . By transforming the MSBE into an objective for which we can obtain unbiased sample gradients , the conjugate MSBE allowed a natural extension to nonlinear function approximation . A natural next step is to use these advances to revisit defining and optimizing statistically robust losses for value functions . The MSBE and MSPBE are built on squared errors , which are known to be sensitive to outliers . In the RL setting , this translates into overemphasizing states for which the Bellman residual is high , at the cost of obtaining accurate estimates in other states . For example , consider the CliffWorld domain ( Sutton and Barto , 2018 ) where the agent must navigate alongside a cliff to reach a rewarding goal state . Should the agent step into the cliff , the agent is sent back to the start state with a large negative reward . In order to represent the optimal policy , the agent must learn that actions which lead to the cliff yield negative return , while actions that lead towards the goal yield positive return ( or at least less negative ) . Perfectly representing the exact magnitude of negative returns for stepping in the cliff is not useful , as it is sufficient to only know that these actions are more negative than actions leading towards the goal . Further , because the agent bootstraps off these value estimates in the states near the goal , squared errors tend to magnify these inaccuracies across all states visited in the episode . In general , the ability to replace squared errors with alternatives , like absolute errors or the Huber loss ( Huber , 1964 ) , provides another dimension to improve our algorithms and , to the best of our knowledge , has yet to be explored . Issues with squared errors have been noted in the RL literature and addressed heuristically in control using clipping . DQN uses clipped TD errors by default to avoid large magnitude updates ( Mnih et al. , 2015 ) , likewise , Dueling DQN ( Wang et al. , 2016 ) directly clips its gradients . There is a close relationship between clipping TD errors and using a Huber loss in the linear function approximation setting . However , even for linear Q-learning or TD-learning , clipping TD errors does not correspond to minimizing a Huber loss for the Bellman error . Rather , it takes a non-gradient update and applies clipping to that update to mimic the Huber loss which results in a bias similar to residual TD methods ( Baird , 1995 ; Sutton and Barto , 2018 ) . In this work , we develop practical algorithms that can use absolute and Huber errors instead of squared errors for the BE . To do so , we develop a unified perspective of the mean absolute Bellman Error ( MABE ) , mean squared Bellman Error ( MSBE ) , and mean Huber Bellman Error ( MHBE ) which smoothly interpolates between these other two . We rely on the insight that the MSBE can be reformulated into a saddlepoint problem with the introduction of an auxiliary learned variable ; using the same strategy , we derive a biconjugate form for the MHBE and MABE amenable to simple gradient-descent techniques . The resulting approach is a simple modification of existing gradient TD algorithms , using a different update for the auxiliary variable , making it straightforward to use either of these losses . We show that the MHBE and MABE can significantly change the solution quality , improving accuracy in terms of the mean squared value error , as well as the mean absolute value error . We then show that gradient algorithms for the MHBE tend to perform similar to those for the MSBE , but provide significant robustness improvements in certain cases , particularly under bad state aliasing . Finally , we show that in control with nonlinear function approximation , gradient-based algorithms minimizing the MHBE often outperform those using squared losses or those using non-gradient updates . 2 PROBLEM FORMULATION . We model the agent ’ s interactions with its environment as a Markov Decision Process ( MDP ) , ( S , A , P , R , γ ) . At each time-step t , the agent observes states St ∈ S , selects an action At ∈ A according to policy π : S → ∆ ( A ) , transitions to the next state St+1 ∈ S according to transition function P : S × A → ∆ ( S ) , and receives a scalar reward signal Rt+1 and discount γt+1 ∈ [ 0 , 1 ] . The discount depends on the state , and encodes termination when γt+1 = 0 ( White , 2017 ) . For the prediction setting , the agent ’ s goal is to estimate the value function vπ for a given policy . The value function can be defined recursively , using the Bellman operator ( T v ) ( s ) def= E [ Rt+1 + γt+1v ( St+1 ) | St = s ] where the expectation is taken with respect to the policy π and transition dynamics P . The true values vπ are the fixed point for this operator : T vπ = vπ . Our goal is to approximate vπ , with vθ ∈ F for some ( parameterized ) function space F . Typically , the quality of this approximation is evaluated under the value error , either mean squared value error ( MSVE ) or mean absolute value error ( MAVE ) MSVE ( vθ ) def = ∑ s∈S d ( s ) ( vθ ( s ) − vπ ( s ) ) 2 MAVE ( vθ ) def = ∑ s∈S d ( s ) |V ( s ) − vπ ( s ) | where d : S → [ 0 , 1 ] is typically the visitation distribution under a behavior policy.1 1For exposition , we assume discrete states and actions throughout this paper . The connection to continuous state-actions is straightforward and we will explicitly call out where this connection is less obvious . One objective used to learn approximation vθ is the mean squared Bellman error ( MSBE ) MSBE ( θ ) def= ∑ s∈S d ( s ) ( ( T vθ ) ( s ) − vθ ( s ) ) 2 = ∑ s∈S d ( s ) E [ δ ( θ ) | S = s ] 2 ( 1 ) where δ ( θ ) def= Rt+1 + γt+1vθ ( St+1 ) − vθ ( St ) . If vπ ∈ F , then there exists a θ such that vθ = vπ and MSBE ( θ ) = 0 . Otherwise , if vπ /∈ F , then this objective trades-off Bellman error across states . The trade-off in errors across states is impacted by both the weighting d as well as the fact that a squared error is used . The function approximator focuses on states with high weighting d , which is sensible . However , by using a squared error , it heavily emphasizes states with higher error which may not be desirable . In the next section , we develop an approach to use robust losses—the absolute error and the Huber error—in place of this squared error . The same approach as above can also be used for control to approximate the optimal action-values q∗ . These values can similarly be defined using a Bellman optimality operator ( T ∗q ) ( s , a ) def= E [ Rt+1 + γt+1 max a′∈A q ( St+1 , a ′ ) | St = s , At = a ] with T ∗q∗ = q∗ . The corresponding mean squared Bellman error for learning approximate qθ is∑ s∈S , a∈A d ( s , a ) [ ( T ∗qθ ) ( s , a ) − qθ ( s , a ) ] 2 where we overload d to be a state-action weighting , which typically corresponds to the state-action visitation frequency under a policy . 3 ROBUST BELLMAN ERRORS . In this section , we provide reformulations of the MABE and the MHBE which avoid the double sampling issue by using their biconjugates . For some intuition on how these objectives differ from the MSVE and MSBE , we visualize them in Figure 1 . To sample this gradient requires a sample of δ ( θ ) for the first expectation and an independent sample of δ ( θ ) for the second expectation . Otherwise , using the same sample , we estimate the gradient of E [ δ ( θ ) 2 | S = s ] instead of E [ δ ( θ ) | S = s ] 2 . Due to the chain rule , both the MABE and MHBE will suffer from the same issue as the MSBE . The strategy is to reformulate the objectives using conjugates , introducing an auxiliary variable to estimate part of the gradient . For a real-valued function f : R → R , the conjugate is f∗ ( h ) def= supx∈R xh − f ( x ) . This function f∗ also has a conjugate , f∗∗ , which is called the biconjugate of f . Further , for any function f that is proper , convex , and lower semi-continuous , the biconjugate f∗∗ ( x ) = f ( x ) for all x by the Fenchel-Moreau theorem ( Fenchel , 1949 ; Moreau , 1970 ) . Because the three functions we want to reformulate—the absolute , huber , and square functions—are all proper , convex , and lower semi-continuous , this equivalence allows us to reformulate these losses using biconjugates to avoid the double sampling issue without changing the solutions to these losses . To rewrite existing results in our notation and provide some intuition , we first write the reformulation for the MSBE . The conjugate of the squared error f ( x ) = 12x 2 is f∗ ( h ) = maxx∈R hx− 12x 2 , which is in fact again the squared error : f∗ ( h ) = 12h 2 ( this result is well known , but for completeness we include the proof in Appendix A ) . The biconjugate is f∗∗ ( x ) = maxh∈R xh− 12h 2 and f ( x ) = f∗∗ ( x ) . We can use this to get that , for x = E [ δ ( θ ) | S = s ] , E [ δ ( θ ) | S = s ] 2 = max h∈R 2E [ δ ( θ ) | S = s ] h− h2 . If we have function space Fall—the set of all possible functions h : S → R—then we get that MSBE ( θ ) = ∑ s∈S d ( s ) max h∈R ( 2E [ δ ( θ ) | S = s ] h− h2 ) = max h∈Fall ∑ s∈S d ( s ) ( 2E [ δ ( θ ) | S = s ] h ( s ) − h ( s ) 2 ) where the maximization comes out of the sum using the interchangeability property ( Shapiro et al. , 2014 ; Dai et al. , 2017 ) and h ( s ) is a function that allows us to independently pick a maximizer for every state in the summation . If we have the maximizing function h∗ ( s ) , it is straightforward to sample the gradient . Because h∗ ( s ) itself is not directly a function of θ , then the gradient is ∇θ ∑ s∈S d ( s ) ( 2E [ δ ( θ ) | St=s ] h∗ ( s ) −h∗ ( s ) 2 ) = ∑ s∈S 2d ( s ) h∗ ( s ) E [ γ∇θv ( S′ ) −∇θv ( s ) ) | St=s ] . Drawing samples S ∼ d ( · ) , A ∼ π ( ·|S ) , and S′ ∼ P ( ·|S , A ) , we can easily compute a stochastic sample of the gradient . In practice , we simply optimize the resulting saddlepoint problem with a minimization over θ and maximization over h. Note the optimal h∗ ( s ) = E [ δ ( θ ) | St = s ] . We can use the same procedure for the Huber error and the absolute error . We derive the biconjugate form for the Huber error in the following proposition . Though it is a relatively straightforward result to obtain , to the best of our knowledge , it is new and so worth providing formally . Proposition 3.1 The biconjugate of the Huber function is p∗∗τ ( x ) = max h∈ [ −τ , τ ] xh− 1 2 h2 . ( 2 ) Due to space restrictions , we provide a complete proof in Appendix A . The absolute value has biconjugate maxh∈ [ −1,1 ] xh . As in the squared error case , this is a well known result but we include the proof for completeness in Appendix A . Notice again the constrained optimization problem for this biconjugate , as in the case of the Huber biconjugate function . We can now provide the forms for MABE and MHBE : MABE ( θ ) def= max h∈Fsign ∑ s∈S d ( s ) h ( s ) E [ δ ( θ ) |S = s ] MHBE ( θ ) def= max h∈Fclipτ ∑ s∈S d ( s ) ( 2h ( s ) E [ δ ( θ ) |S = s ] − h ( s ) 2 ) Fsign is the set of all functions hsign : S→ { −1 , 1 } andFclipτ the set of all functions hclipτ : S→ [ −τ , τ ] . For all of these reformulations , in practice we will have parameterized functions h , and so only approximate the objective . For example , for the MSBE , we may use linear functions Fh = { h | h ( s ) = θh > x ( s ) , θh ∈ Rd } where x : S → Rd is a d-dimensional feature generating function . For the MHBE , we might use Fh = { h | h ( s ) = clipτ ( θh > x ( s ) ) , θh ∈ Rd } . For the conjugate reformulation of the MSBE , limiting vθ and h ( s ) to both be linear functions of the same features results in the mean squared Projected Bellman error , as noted by Dai et al . ( 2017 ) . There are no equivalent existing projected Bellman errors for these new absolute and Huber variants . We note , however , that when the function space is constrained to only functions of our features x ( s ) , these conjugate reformulations do not suffer from the same identifiability issue raised for the MSBE ( Sutton and Barto , 2018 ; Scherrer , 2010 ) as we show in Appendix B . The final form of the MHBE highlights a connection to a recently proposed gradient TD algorithm that seemed to significantly improve stability over vanilla gradient TD methods . The algorithm , called TD with Regularized Corrections ( TDRC ) , adds a regularizer to the parameters for h ( Ghiassian et al. , 2020 ) . This regularizer has the effect of constraining h ( s ) to be closer to zero , and so could be seen as providing some of the same robustness benefits as a Huber loss . Of course , the correspondence is by no means exact . Regularizing with ` 2 on the parameters is different than restricting between −1 and 1 . Further , as shown in that work , TDRC does not alter the fixed-point of its mean squared objective . Using the Huber error in place of the squared error , however , will likely alter the fixed-point . Finally , we show that the MHBE bounds the MAVE . It is well-known that absolute errors are hard to optimize and so the Huber error acts as a smooth approximation to the absolute error . Theorem 3.2 first shows that minimizing the MABE acts as a proxy for minimizing the MAVE . Then , by bounding the absolute value error with the Huber error , we get that the MHBE acts as an approximation of the MABE implying that minimizing the MHBE likewise minimizes the MAVE . Because of this approximation , the MHBE has an irreducible error term which is controlled by the Huber parameter τ . We provide a complete proof of Theorem 3.2 in Appendix A.1 . Theorem 3.2 Let τcap = min ( τ , 1 ) , then for arbitrary v ∈ Rd and 0 < < τ2cap we have ‖vπ − v‖1 ≤ ‖ ( I − Pπ , γ ) −1‖1‖T v − v‖1 ≤ ‖ ( I − Pπ , γ ) −1‖1 d∑ s=0 ( √ 2 pτ ( T vs − vs ) + √ 2 ) .

Summary: This paper starts with the premise that squared error minimization, despite its wide use, might not be the most effective option for learning value functions. The authors hypothesize that this might be because of squared error's emphasis on outlier states where the bellman error is large at the expense of accuracy on other states. To address this, they consider the absolute value and huber errors alongside the squared error objective and propose a saddle-point reformulation for these objectives that requires learning an auxiliary learned state function which, essentially, attempts to predict the residual error at each state.  Based on this motivation, starting with Section 4 the authors then make a connection of the proposed robust loss framework involving a learned auxiliary state function with prior work on algorithms for improving TD learning, namely, GTD2 and TDC.  Empirical analysis covers both prediction tasks as well as control tasks. In the prediction experiments, evaluations are conducted over carefully designed synthetic problems with linear function approximation in order to highlight particular challenges for each objective. These demonstrate that the Huber objective can often improve the prediction accuracy over squared error over a wide range of step sizes (learning rates).  For the empirical investigation on non-linear control problems, the authors consider a Huber-extension of the previously proposed QRC algorithm (which itself extends the TDC update to be considered in conjunction with the DQN target definition) and demonstrate that QRC-huber can improve on baselines in certain environments while being competitive elsewhere. These algorithms are also validated on a mini version of the Atari domain called Minatar.