Variational Discriminator Bottleneck: Improving Imitation Learning, Inverse RL, and GANs by Constraining Information Flow

1 INTRODUCTION . Adversarial learning methods provide a promising approach to modeling distributions over highdimensional data with complex internal correlation structures . These methods generally use a discriminator to supervise the training of a generator in order to produce samples that are indistinguishable from the data . A particular instantiation is generative adversarial networks , which can be used for high-fidelity generation of images ( Goodfellow et al. , 2014 ; Karras et al. , 2017 ) and other highdimensional data ( Vondrick et al. , 2016 ; Xie et al. , 2018 ; Donahue et al. , 2018 ) . Adversarial methods can also be used to learn reward functions in the framework of inverse reinforcement learning ( Finn et al. , 2016a ; Fu et al. , 2017 ) , or to directly imitate demonstrations ( Ho & Ermon , 2016 ) . However , they suffer from major optimization challenges , one of which is balancing the performance of the generator and discriminator . A discriminator that achieves very high accuracy can produce relatively uninformative gradients , but a weak discriminator can also hamper the generator ’ s ability to learn . These challenges have led to widespread interest in a variety of stabilization methods for adversarial learning algorithms ( Arjovsky et al. , 2017 ; Kodali et al. , 2017 ; Berthelot et al. , 2017 ) . In this work , we propose a simple regularization technique for adversarial learning , which constrains the information flow from the inputs to the discriminator using a variational approximation to the information bottleneck . By enforcing a constraint on the mutual information between the input observations and the discriminator ’ s internal representation , we can encourage the discriminator to learn a representation that has heavy overlap between the data and the generator ’ s distribution , thereby effectively modulating the discriminator ’ s accuracy and maintaining useful and informative 1xbpeng.github.io/projects/VDB/ gradients for the generator . Our approach to stabilizing adversarial learning can be viewed as an adaptive variant of instance noise ( Salimans et al. , 2016 ; Sønderby et al. , 2016 ; Arjovsky & Bottou , 2017 ) . However , we show that the adaptive nature of this method is critical . Constraining the mutual information between the discriminator ’ s internal representation and the input allows the regularizer to directly limit the discriminator ’ s accuracy , which automates the choice of noise magnitude and applies this noise to a compressed representation of the input that is specifically optimized to model the most discerning differences between the generator and data distributions . The main contribution of this work is the variational discriminator bottleneck ( VDB ) , an adaptive stochastic regularization method for adversarial learning that substantially improves performance across a range of different application domains , examples of which are available in Figure 1 . Our method can be easily applied to a variety of tasks and architectures . First , we evaluate our method on a suite of challenging imitation tasks , including learning highly acrobatic skills from mocap data with a simulated humanoid character . Our method also enables characters to learn dynamic continuous control skills directly from raw video demonstrations , and drastically improves upon previous work that uses adversarial imitation learning . We further evaluate the effectiveness of the technique for inverse reinforcement learning , which recovers a reward function from demonstrations in order to train future policies . Finally , we apply our framework to image generation using generative adversarial networks , where employing VDB improves the performance in many cases . 2 RELATED WORK . Recent years have seen an explosion of adversarial learning techniques , spurred by the success of generative adversarial networks ( GANs ) ( Goodfellow et al. , 2014 ) . A GAN framework is commonly composed of a discriminator and a generator , where the discriminator ’ s objective is to classify samples as real or fake , while the generator ’ s objective is to produce samples that fool the discriminator . Similar frameworks have also been proposed for inverse reinforcement learning ( IRL ) ( Finn et al. , 2016b ) and imitation learning ( Ho & Ermon , 2016 ) . The training of adversarial models can be extremely unstable , with one of the most prevalent challenges being balancing the interplay between the discriminator and the generator ( Berthelot et al. , 2017 ) . The discriminator can often overpower the generator , easily differentiating between real and fake samples , thus providing the generator with uninformative gradients for improvement ( Che et al. , 2016 ) . Alternative loss functions have been proposed to mitigate this problem ( Mao et al. , 2016 ; Zhao et al. , 2016 ; Arjovsky et al. , 2017 ) . Regularizers have been incorporated to improve stability and convergence , such as gradient penalties ( Kodali et al. , 2017 ; Gulrajani et al. , 2017a ; Mescheder et al. , 2018 ) , reconstruction loss ( Che et al. , 2016 ) , and a myriad of other heuristics ( Sønderby et al. , 2016 ; Salimans et al. , 2016 ; Arjovsky & Bottou , 2017 ; Berthelot et al. , 2017 ) . Task-specific architectural designs can also substantially improve performance ( Radford et al. , 2015 ; Karras et al. , 2017 ) . Similarly , our method also aims to regularize the discriminator in order to improve the feedback provided to the generator . But instead of explicit regularization of gradients or architecture-specific constraints , we apply a general information bottleneck to the discriminator , which previous works have shown to encourage networks to ignore irrelevant cues ( Achille & Soatto , 2017 ) . We hypothesize that this then allows the generator to focus on improving the most discerning differences between real and fake samples . Adversarial techniques have also been applied to inverse reinforcement learning ( Fu et al. , 2017 ) , where a reward function is recovered from demonstrations , which can then be used to train policies to reproduce a desired skill . Finn et al . ( 2016a ) showed an equivalence between maximum entropy IRL and GANs . Similar techniques have been developed for adversarial imitation learning ( Ho & Ermon , 2016 ; Merel et al. , 2017 ) , where agents learn to imitate demonstrations without explicitly recovering a reward function . One advantage of adversarial methods is that by leveraging a discriminator in place of a reward function , they can be applied to imitate skills where reward functions can be difficult to engineer . However , the performance of policies trained through adversarial methods still falls short of those produced by manually designed reward functions , when such reward functions are available ( Rajeswaran et al. , 2017 ; Peng et al. , 2018 ) . We show that our method can significantly improve upon previous works that use adversarial techniques , and produces results of comparable quality to those from state-of-the-art approaches that utilize manually engineered reward functions . Our variational discriminator bottleneck is based on the information bottleneck ( Tishby & Zaslavsky , 2015 ) , a technique for regularizing internal representations to minimize the mutual information with the input . Intuitively , a compressed representation can improve generalization by ignoring irrelevant distractors present in the original input . The information bottleneck can be instantiated in practical deep models by leveraging a variational bound and the reparameterization trick , inspired by a similar approach in variational autoencoders ( VAE ) ( Kingma & Welling , 2013 ) . The resulting variational information bottleneck approximates this compression effect in deep networks ( Alemi et al. , 2016 ; Achille & Soatto , 2017 ) . A similar bottleneck has also been applied to learn disentangled representations ( Higgins et al. , 2017 ) . Building on the success of VAEs and GANs , a number of efforts have been made to combine the two . Makhzani et al . ( 2016 ) used adversarial discriminators during the training of VAEs to encourage the marginal distribution of the latent encoding to be similar to the prior distribution , similar techniques include Mescheder et al . ( 2017 ) and Chen et al . ( 2018 ) . Conversely , Larsen et al . ( 2016 ) modeled the generator of a GAN using a VAE . Zhao et al . ( 2016 ) used an autoencoder instead of a VAE to model the discriminator , but does not enforce an information bottleneck on the encoding . While instance noise is widely used in modern architectures ( Salimans et al. , 2016 ; Sønderby et al. , 2016 ; Arjovsky & Bottou , 2017 ) , we show that explicitly enforcing an information bottleneck leads to improved performance over simply adding noise for a variety of applications . 3 PRELIMINARIES . In this section , we provide a review of the variational information bottleneck proposed by Alemi et al . ( 2016 ) in the context of supervised learning . Our variational discriminator bottleneck is based on the same principle , and can be instantiated in the context of GANs , inverse RL , and imitation learning . Given a dataset { xi , yi } , with features xi and labels yi , the standard maximum likelihood estimate q ( yi|xi ) can be determined according to min q Ex , y∼p ( x , y ) [ −log q ( y|x ) ] . ( 1 ) Unfortunately , this estimate is prone to overfitting , and the resulting model can often exploit idiosyncrasies in the data ( Krizhevsky et al. , 2012 ; Srivastava et al. , 2014 ) . Alemi et al . ( 2016 ) proposed regularizing the model using an information bottleneck to encourage the model to focus only on the most discriminative features . The bottleneck can be incorporated by first introducing an encoder E ( z|x ) that maps the features x to a latent distribution over Z , and then enforcing an upper bound Ic on the mutual information between the encoding and the original features I ( X , Z ) . This results in the following regularized objective J ( q , E ) J ( q , E ) = min q , E Ex , y∼p ( x , y ) [ Ez∼E ( z|x ) [ −log q ( y|z ) ] ] s.t . I ( X , Z ) ≤ Ic . ( 2 ) Note that the model q ( y|z ) now maps samples from the latent distribution z to the label y . The mutual information is defined according to I ( X , Z ) = ∫ p ( x , z ) log p ( x , z ) p ( x ) p ( z ) dx dz = ∫ p ( x ) E ( z|x ) logE ( z|x ) p ( z ) dx dz , ( 3 ) where p ( x ) is the distribution given by the dataset . Computing the marginal distribution p ( z ) = ∫ E ( z|x ) p ( x ) dx can be challenging . Instead , a variational lower bound can be obtained by using an approximation r ( z ) of the marginal . Since KL [ p ( z ) ||r ( z ) ] ≥ 0 , ∫ p ( z ) log p ( z ) dz ≥∫ p ( z ) log r ( z ) dz , an upper bound on I ( X , Z ) can be obtained via the KL divergence , I ( X , Z ) ≤ ∫ p ( x ) E ( z|x ) logE ( z|x ) r ( z ) dx dz = Ex∼p ( x ) [ KL [ E ( z|x ) ||r ( z ) ] ] . ( 4 ) This provides an upper bound on the regularized objective J̃ ( q , E ) ≥ J ( q , E ) , J̃ ( q , E ) = min q , E Ex , y∼p ( x , y ) [ Ez∼E ( z|x ) [ −log q ( y|z ) ] ] s.t . Ex∼p ( x ) [ KL [ E ( z|x ) ||r ( z ) ] ] ≤ Ic . ( 5 ) To solve this problem , the constraint can be subsumed into the objective with a coefficient β min q , E Ex , y∼p ( x , y ) [ Ez∼E ( z|x ) [ −log q ( y|z ) ] ] + β ( Ex∼p ( x ) [ KL [ E ( z|x ) ||r ( z ) ] ] − Ic ) . ( 6 ) Alemi et al . ( 2016 ) evaluated the method on supervised learning tasks , and showed that models trained with a VIB can be less prone to overfitting and more robust to adversarial examples .

Summary: The authors propose to apply the Deep Variational Information Bottleneck (VIB) method of [1] on discriminator networks in various adversarial-learning-based scenarios. They propose a way to adaptively update the value for the bêta hyper-parameter to respect the constraint on I(X,Z). Their technique is shown to stabilize/allow training when P_g and P_data do not overlap, similarly to WGAN and gradient-penalty based approaches, by essentially pushing their representation distributions (p_z) to overlap with the mutual information bottleneck. It can also be considered as an adaptive version of instance noise, which serves the same goal. The method is evaluated on different adversarial learning setup (imitation learning, inverse reinforcement learning and GANs), where it compares positively to most related methods. Best results for ‘classical’ adversarial learning for image generation are however obtained when combining the proposed VIB with gradient penalty (which outperforms by itself the VGAN in this case).