On Pathologies in KL-Regularized Reinforcement Learning from Expert Demonstrations

1 Introduction . Reinforcement learning ( RL ) [ 15 , 24 , 46 , 47 ] is a powerful paradigm for learning complex behaviors . Unfortunately , many modern reinforcement learning algorithms require agents to carry out millions of interactions with their environment to learn desirable behaviors , making them of limited use for a wide range of practical applications that can not be simulated [ 8 , 28 ] . This limitation has motivated the study of algorithms that can incorporate pre-collected offline data into the training process either fully offline or with online exploration to improve sample efficiency , performance , and reliability [ 2 , 6 , 16 , 23 , 52 , 53 ] . An important and well-motivated subset of these methods consists of approaches for efficiently incorporating expert demonstrations into the learning process [ 5 , 11 , 18 , 42 ] . Reinforcement learning with Kullback-Leibler ( KL ) regularization is a particularly successful approach for doing so [ 3 , 27 , 29 , 31 , 44 , 51 ] . In KL-regularized reinforcement learning , the standard reinforcement learning objective is augmented by a Kullback-Leibler divergence term that penalizes dissimilarity between the online policy and a behavioral reference policy derived from expert demonstrations . The resulting regularized objective pulls the agent ’ s online policy towards the behavioral reference policy while also allowing it to improve upon the behavioral reference policy by exploring and interacting with the environment . Recent advances that leverage explicit or implicit KL-regularized objectives , such as BRAC [ 51 ] , ABM [ 44 ] , and AWAC [ 27 ] , have shown that KLregularized reinforcement learning from expert demonstrations is able to significantly improve the sample efficiency of online training and reliably solve challenging environments previously unsolved by standard deep reinforcement learning algorithms . ∗Equal contribution . † Corresponding author : tim.rudner @ cs.ox.ac.uk . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Contributions . In this paper , we show that despite some empirical success , KL-regularized reinforcement learning from expert demonstrations can suffer from previously unrecognized pathologies that lead to instability and sub-optimality in online learning . To summarize , our core contributions are as follows : • We illustrate empirically that commonly used classes of parametric behavioral policies experi- ence a collapse in predictive variance about states away from the expert demonstrations . • We demonstrate theoretically and empirically that KL-regularized reinforcement learning al- gorithms can suffer from pathological training dynamics in online learning when regularized against behavioral policies that exhibit such a collapse in predictive variance . • We show that the pathology can be remedied by non-parametric behavioral policies , whose predictive variances are well-calibrated and guaranteed not to collapse about previously unseen states , and that fixing the pathology results in online policies that significantly outperform stateof-the-art approaches on a range of challenging locomotion and dexterous hand manipulation tasks . The left panel of Figure 1 shows an example of the collapse in predictive variance away from the expert trajectories in parametric behavioral policies . In contrast , the right panel of Figure 1 shows the predictive variance of a non-parametric behavioral policy , which—unlike in the case of the parametric policy—increases off the expert trajectories . By avoiding the pathology , we obtain a stable and reliable approach to sample-efficient reinforcement learning , applicable to a wide range of reinforcement learning algorithms that leverage KL-regularized objectives.2 2 Background . We consider the standard reinforcement learning setting where an agent interacts with a discounted Markov Decision Process ( MDP ) [ 46 ] given by a 5-tuple ( S , A , p , r , γ ) , where S and A are the state and action spaces , p ( · | st , at ) are the transition dynamics , r ( st , at ) is the reward function , and γ is a discount factor . ρπ ( τt ) denotes the state–action trajectory distribution from time t induced by a policy π ( · | st ) . The discounted return from time step t is given by R ( τt ) = ∑∞ k=t γ kr ( sk , ak ) for t ∈ N0 . The standard reinforcement learning objective to be maximized is the expected discounted return Jπ ( τ0 ) = Eρπ ( τ0 ) [ R ( τ0 ) ] under the policy trajectory distribution . 2.1 Improving and Accelerating Online Training via Behavioral Cloning . We consider settings where we have a set of expert demonstrations without reward , D0 = { ( sn , an ) } Nn=1 = { S̄ , Ā } , which we would like to use to speed up and improve online learn- 2Code and visualizations of our results can be found at https : //sites.google.com/view/nppac . ing [ 5 , 42 ] . A standard approach for turning expert trajectories into a policy is behavioral cloning [ 1 , 4 ] which involves learning a mapping from states in the expert demonstrations to their corresponding actions , that is , π0 : S → A . As such , behavioral cloning does not assume or require access to a reward function and only involves learning a mapping from states to action in a supervised fashion . Since expert demonstrations are costly to obtain and often only available in small number , behavioral cloning alone is typically insufficient for agents to learn good policies in complex environments and has to be complemented by a method that enables the learner to build on the cloned behavior by interacting with the environment . A particularly successful and popular class of algorithms used for incorporating behavioral policies into online training is KL-regularized reinforcement learning [ 10 , 37 , 43 , 48 ] . 2.2 KL-Regularized Objectives in Reinforcement Learning . KL-regularized reinforcement learning modifies the standard reinforcement learning objective by augmenting the return with a negative KL divergence term from the learned policy π to a reference policy π0 , given a temperature parameter α . The resulting discounted return from time step t ∈ N0 is then given by R̃ ( τt ) = ∞∑ k=t γk [ r ( sk , ak ) − αDKL ( π ( · | sk ) ‖ π0 ( · | sk ) ) ] ( 1 ) and the reinforcement learning objective becomes J̃π ( τ0 ) = Eρπ ( τ0 ) [ R̃ ( τ0 ) ] . When the reference policy π0 is given by a uniform distribution , we recover the entropy-regularized reinforcement learning objective used in Soft Actor–Critic ( SAC ) [ 13 ] up to an additive constant . Under a uniform reference policy π0 , the resulting objective encourages exploration , while also choosing high-reward actions . In contrast , when π0 is non-uniform , the agent is discouraged to explore areas of the state space S where the variance of π0 ( · | s ) is low ( i.e. , more certain ) and encouraged to explore areas of the state space where the variance of π0 ( · | s ) is high . The KLregularized reinforcement learning objective can be optimized via policy–gradient and actor–critic algorithms . 2.3 KL-Regularized Actor–Critic . An optimal policy π that maximizes the expected KL-augmented discounted return J̃π can be learned by directly optimizing the policy gradient ∇πJ̃π . However , this policy gradient estimator exhibits high variance , which can lead to unstable learning . Actor–critic algorithms [ 7 , 17 , 32 , 38 ] attempt to reduce this variance by making use of the state value function V π ( st ) = Eρπ ( τt ) [ R̃ ( τt ) | st ] or the state–action value function Qπ ( st , at ) = Eρπ ( τt ) [ R̃ ( τt ) | st , at ] to stabilize training . Given a reference policy π0 ( at | st ) , the state value function can be shown to satisfy the modified Bellman equation V π ( st ) =̇ Eat∼π ( ·|st ) [ Q π ( st , at ) ] − αDKL ( π ( · | st ) ||π0 ( · | st ) ) with a recursively defined Q-function Qπ ( st , at ) =̇ r ( st , at ) + γ Est+1∼p ( ·|st , at ) [ V π ( st+1 ) ] . Instead of directly optimizing the objective function J̃π via the policy gradient , actor–critic methods alternate between policy evaluation and policy improvement [ 7 , 13 ] : Policy Evaluation . During the policy evaluation step , Qπθ ( s , a ) , parameterized by parameters θ , is trained by minimizing the Bellman residual JQ ( θ ) =̇ E ( st , at ) ∼D [ ( Qθ ( st , at ) − ( r ( st , at ) + γEst+1∼p ( ·|st , at ) [ Vθ̄ ( st+1 ) ] ) ) 2 ] , ( 2 ) where D is a replay buffer and θ̄ is a stabilizing moving average of parameters . Policy Improvement . In the policy improvement step , the policy πφ , parameterized by parameters φ , is updated towards the exponential of the KL-augmented Q-function , Jπ ( φ ) =̇ Est∼D [ αDKL ( πφ ( · | st ) ‖ π0 ( · | st ) ) ] − Est∼D [ Eat∼πφ ( ·|st ) [ Qθ ( st , at ) ] ] , ( 3 ) with states sampled from a replay buffer D and actions sampled from the parameterized online policy πφ . The following sections will focus on the policy improvement objective and how certain types of references policies can lead to pathologies when optimizing Jπ ( φ ) with respect to φ . 3 Identifying the Pathology . In this section , we investigate the effect of KL-regularization on the training dynamics . To do so , we first consider the properties of the KL divergence to identify a potential failure mode for KL-regularized reinforcement learning . Next , we consider parametric Gaussian behavioral reference policies commonly used in practice for continuous control tasks [ 13 , 51 ] and show that for Gaussian behavioral reference policies with small predictive variance , the policy improvement objective suffers from exploding gradients with respect to the policy parameters φ . We confirm that this failure occurs empirically and demonstrate that it results in slow , unstable , and suboptimal online learning . Lastly , we show that various regularization techniques used for estimating behavioral policies are unable to prevent this failure and also lead to suboptimal online policies . 3.1 When Are KL-Regularized Reinforcement Learning Objectives Meaningful ? . We start by considering the properties of the KL divergence and discuss how these properties can lead to potential failure modes in KL-regularized objectives . A well-known property of KL-regularized objectives in the variational inference literature is the occurrence of singularities when the support of one distribution is not contained in the support of the other . To illustrate this problem , we consider the case of Gaussian behavioral and online policies commonly used in practice . Mathematically , the KL divergence between two full Gaussian distributions is always finite and well-defined . Hence , we might hope KL-regularized reinforcement learning with Gaussian behavioral and online policies to be unaffected by the failure mode described above . However , the support of a Gaussian online policy πφ ( · | st ) will not be contained in the support of a behavioral reference policy π0 ( · | st ) as the predictive variance σ20 ( st ) tends to zero , and hence DKL ( πφ ( · | st ) ‖ π0 ( · | st ) ) → ∞ as σ20 ( st ) → 0 . In other words , as the variance of a behavioral reference policy tends to zero and the behavioral distribution becomes degenerate , the KL divergence blows up to infinity [ 25 ] . While in practice , Gaussian behavioral policy would not operate in the limit of zero variance , the functional form of the KL divergence between ( univariate ) Gaussians , DKL ( πφ ( · | st ) ‖ π0 ( · | st ) ) ∝ log σ0 ( st ) σφ ( st ) + σ2φ ( st ) + ( µφ ( st ) − µ0 ( st ) ) 2 2σ20 ( st ) , implies a continuous , quadratic increase in the magnitude of the divergence as σ0 ( st ) decreases , further exacerbated by a large difference in predictive means , |µφ ( st ) − µ0 ( st ) | . As a result , for Gaussian behavioral reference policies π0 ( · | st ) that assign very low probability to sets of points in sample space far away from the distribution ’ s mean µ0 ( st ) , computing the KL divergence can result in divergence values so large to cause numerical instabilities and arithmetic overflow . Hence , even for a suitably chosen behavioral reference policy class , vanishingly small behavioral reference policy predictive variances can cause the KL divergence to ‘ blow up ’ and cause numerical issues at evaluation points far away from states in the expert demonstrations . One way to address this failure mode may be to lower-bound the output of the variance network ( e.g. , by adding a small constant bias ) . However , placing a floor on the predictive variance of the behavioral reference policy is not sufficient to encourage effective learning . While it would prevent the KL divergence from blowing up , it would also lead to poor gradient signals , as well-calibrated predictive variance estimates that increase on states far away from the expert trajectories are necessary to keep the KL penalty from pulling the predictive mean of the online policy towards poor behavioral reference policy predictive means on states off the expert trajectories . Another possible solution could be to use heavy-tailed behavioral reference policies distributions , for example , Laplace distributions , to avoid pathological training dynamics . However , in Appendix B.3 we show that Laplace behavioral reference policies also suffer from pathological training dynamics , albeit less severely . In the following sections , we explain how an explosion in DKL ( πφ ( · | st ) ‖ π0 ( · | st ) ) caused by small σ20 ( st ) affects the gradients of Jπ ( φ ) in KL-regularized RL and discuss of how and why σ 2 0 ( st ) may tend to zero in practice .

Summary: The authors identify a potential failure in KL-regularized reinforcement learning in which small predictive variance of the target policy may lead to exploding gradients, possibly destabilizing gradient-based learning algorithms. The authors discuss an uncertainty collapse of parametric models under Maximum Likelihood estimation in OOD regions of the expert, resulting in the aforementioned pathology. Finally, they suggest to use non-parametric models of the behavior policy (in which uncertainty collapse is not present), showing improved performance on several continuous control benchmarks.