Source-linked AI summary
Variational Discriminator Bottleneck: Improving Imitation Learning, Inverse RL, and GANs by Constraining Information Flow
Xue Bin Peng, Angjoo Kanazawa, Sam Toyer, Pieter Abbeel, Sergey Levine
TL;DR
Adversarial training is unstable because overly accurate discriminators provide uninformative generator gradients. The paper constrains discriminator information flow with a variational bottleneck, improving performance across imitation learning, inverse reinforcement learning, and GAN image generation.
Problem
Adversarial models require a balance between generator and discriminator performance because highly accurate discriminators produce relatively uninformative gradients.
Method
VDB constrains mutual information between observations and the discriminator’s internal representation to modulate discriminator accuracy and maintain useful gradients.
Results
VDB substantially improves performance across adversarial imitation learning, inverse reinforcement learning, and GAN image generation, including learning dynamic skills from raw video demonstrations.
Takeaways & Limitations
The method provides a broadly applicable adaptive regularizer for adversarial learning across multiple tasks and architectures.
Abstract
from arXiv · showhide
Adversarial learning methods have been proposed for a wide range of applications, but the training of adversarial models can be notoriously unstable. Effectively balancing the performance of the generator and discriminator is critical, since a discriminator that achieves very high accuracy will produce relatively uninformative gradients. In this work, we propose a simple and general technique to constrain information flow in the discriminator by means of an information bottleneck. By enforcing a constraint on the mutual information between the observations and the discriminator's internal representation, we can effectively modulate the discriminator's accuracy and maintain useful and informative gradients. We demonstrate that our proposed variational discriminator bottleneck (VDB) leads to significant improvements across three distinct application areas for adversarial learning algorithms. Our primary evaluation studies the applicability of the VDB to imitation learning of dynamic continuous control skills, such as running. We show that our method can learn such skills directly from \emph{raw} video demonstrations, substantially outperforming prior adversarial imitation learning methods. The VDB can also be combined with adversarial inverse reinforcement learning to learn parsimonious reward functions that can be transferred and re-optimized in new settings. Finally, we demonstrate that VDB can train GANs more effectively for image generation, improving upon a number of prior stabilization methods.
1 INTRODUCTION
The paper introduces the variational discriminator bottleneck (VDB), a general regularization technique that constrains discriminator information flow to maintain useful generator gradients. It applies VDB across motion imitation, image generation, and transferable reward learning.
- Adversarial learning uses discriminators to supervise generators producing samples indistinguishable from data, but the approach spans GANs, inverse reinforcement learning, and imitation learning.
- VDB constrains mutual information between input observations and the discriminator’s internal representation, modulating discriminator accuracy and preserving informative gradients.
- Figure 1 presents applications in motion imitation, image generation, and transferable reward learning through adversarial inverse reinforcement learning.
- Adaptive information-constrained noise limits discriminator accuracy while compressing inputs to model the most discerning differences between generator and data distributions.
- The method substantially improves adversarial learning across acrobatic mocap imitation, raw-video continuous-control imitation, inverse reinforcement learning, and GAN image generation.
2 RELATED WORK
Adversarial methods support imitation and inverse reinforcement learning by replacing manually engineered rewards with discriminators, but policy quality can remain below methods using available designed rewards. The paper positions VDB as an improvement across these adversarial settings.
- GANs, inverse reinforcement learning, and imitation learning use adversarial discriminator–generator frameworks to classify real and generated samples or recover behaviors from demonstrations.
- Adversarial imitation can avoid explicit reward engineering, which is useful when reward functions are difficult to design.
- When manually designed rewards are available, policies trained with adversarial methods have remained below those produced by manually engineered rewards.
- The paper reports that VDB improves prior adversarial techniques and produces results comparable to state-of-the-art approaches using manually engineered reward functions.
- The information bottleneck regularizes internal representations by minimizing mutual information with inputs, encouraging compressed representations that ignore irrelevant distractors.
3 PRELIMINARIES
The preliminaries review variational information bottlenecks, which encode inputs into stochastic latent variables and constrain their mutual information with the original features. A variational bound makes this constraint practical in deep models.
- The VIB principle is instantiated in GANs, inverse reinforcement learning, and imitation learning, extending the approach beyond supervised learning.
- Variational information bottlenecks map features x to a latent distribution and train a classifier q(y|z) using the encoded samples z.
- The bottleneck imposes the constraint I(X, Z) ≤Ic on mutual information between original features X and latent representations Z.
- The marginal distribution p(z) can be difficult to compute, motivating an approximation r(z) for deriving a variational bound.
- Figure 2 depicts encoding samples into E(z|x), discriminator classification in latent space, and discriminator differences under varying KL bounds Ic.
- A coefficient β can subsume the information constraint into the optimization objective.
- Prior supervised-learning experiments found that VIB models can be less prone to overfitting and more robust to adversarial examples.
4 VARIATIONAL DISCRIMINATOR BOTTLENECK
The variational discriminator bottleneck (VDB) constrains information flow through a stochastic discriminator representation, adaptively controlling discriminator accuracy to preserve useful generator gradients. The framework extends to GANs, imitation learning, and inverse reinforcement learning, where it supports motion imitation and transferable reward learning.
- VDB framework: VDB inserts a stochastic encoder between each input and discriminator, constraining mutual information I(X, Z) between observations and their representation.The discriminator classifies encoded samples rather than the original inputs.
- Optimization: Dual gradient descent adaptively updates the Lagrange multiplier β to enforce the target mutual-information constraint I_c during training.In practice, one discriminator-and-encoder gradient step is followed by a β update.
- VDB framework: The information constraint limits discriminator accuracy and maintains non-degenerate generator gradients when the mutual-information budget is sufficiently small.Under additional simplifying assumptions, the constraint bounds gradient coefficients below; the result is presented in Appendix A.
- VAIL: VDB extends GAIL into VAIL by encoding policy states and using the discriminator output as the agent’s reward.The discriminator distinguishes target-policy states from agent-policy states, while the reward is r_t = −log (1 − D(µ_E(s))).
- VAIRL: VDB also extends AIRL into VAIRL by introducing stochastic encoders for the state representations used by the learned functions g and h.AIRL can recover a reusable reward under stated environmental restrictions when the learned state-only function corresponds to the expert reward up to a constant.
- Evaluation: VAIL outperforms other evaluated methods on all humanoid motion-imitation skills except policies trained with Peng et al.’s manually designed reward.Performance is measured by average joint rotation error in radians.
5 EXPERIMENTS
Experiments evaluate VDB-based adversarial learning in imitation learning, inverse reinforcement learning, and image generation. VAIL learns running directly from raw video, VAIRL improves reward transfer, and VGAN remains stable during image-generation training.
- VAIL: Variational Adversarial Imitation Learning: VAIL closely reproduces diverse humanoid skills from a single mocap demonstration, including dynamic flips and complex contacts.The evaluation uses a 34-degrees-of-freedom simulated humanoid character and compares VAIL with several adversarial and non-adversarial baselines.
- VAIL: Variational Adversarial Imitation Learning: VAIL consistently outperforms previous adversarial methods, while VAIL-GP achieves the best overall performance.The comparison measures average joint rotation error and includes handcrafted-reward and behavioral-cloning policies.
- VAIRL: Variational Adversarial Inverse Reinforcement Learning: VAIRL learns smoother reward functions than AIRL that enable more-reliable transfer to mirrored C-Maze and S-Maze environments.Performance is reported as mean return with standard deviation over five runs, alongside the expert’s mean return.
- VAIL: Variational Adversarial Imitation Learning: Only VAIL successfully learns the running gait from raw video; GAIL and pixel-wise reward training fail to learn it.VAIL’s discriminator attends to spatially coherent image patches, with gradients that tend to be significantly larger than GAIL’s.
- VAIL: Variational Adversarial Imitation Learning: Adaptive β updates enforcing Ic = 0.5 consistently achieve the best performance across tested constraint settings.Small fixed β values revert toward GAIL, whereas large fixed values improve early learning but converge to worse performance.
- VGAN: Variational Generative Adversarial Networks: VGAN remains stable on CIFAR-10 while vanilla GAN and instance noise are prone to diverging as training progresses.Image-generation performance is evaluated with Fréchet Inception Distance across CIFAR-10, CelebA, and CelebAHQ experiments.
6 CONCLUSION
The paper presents VDB as a broadly applicable regularization technique for adversarial learning, while identifying real-world video imitation and theoretical analysis as important future directions.
- 6 CONCLUSION: VDB yields significant improvements across adversarial learning tasks in imitation learning, inverse reinforcement learning, and image generation.The conclusion describes the technique as broadly applicable across domains.
- 6 CONCLUSION: Video-imitation results have primarily used synthetic scenes, motivating extension to imitation of real-world videos.The conclusion also calls for deeper theoretical analysis of convergence and stability conditions.
A ANALYSIS AND PROOFS
The analysis shows that constraining the encoder’s KL divergence keeps VDB generator gradients non-degenerate by preventing embedding coefficients from becoming arbitrarily small. The proof uses Gaussian encodings and simplifying assumptions, while noting that average rather than pointwise KL constraints are used in practice.
- Gradient guarantee: Theorem A.1 characterizes the VDB generator gradient under a Gaussian encoder and a KL constraint bounded by I_c.The gradient is analyzed for generator distribution G(x), data distribution p*(x), and encoder E(z|x) = N(μ_E(x), Σ).
- Assumptions: The proof assumes a diagonal, input-independent covariance for notational simplicity, although input-dependent variance is also analyzable.In practice, the method optimizes an average KL constraint rather than the pointwise constraint used in the theorem.
- Gradient guarantee: Without the bottleneck constraint, likelihood coefficients linking generated and data embeddings can become arbitrarily small, causing vanishing generator gradients.This occurs when real and fake samples are far apart in embedding space.
- Gradient guarantee: The KL constraint yields a continuous monotonic lower bound C(I_c) on those coefficients, with C(I_c) → δ > 0 as I_c → 0.The bound follows because the encoded means cannot be more than a constraint-dependent distance apart.
- Implementation: The VDB can be combined with gradient penalty by reparameterizing the encoder so gradients backpropagate through its stochastic representation.The gradient-penalty weight varies by application and is applied only to real samples.
C IMITATION LEARNING
The imitation experiments use phase-aware discriminators and evaluate their learning behavior across motion and video tasks. PFNN discriminators improve performance, while the VDB constrains discriminator accuracy to preserve informative training dynamics.
- Motion imitation: Motion imitation trains a 34-degree-of-freedom humanoid to mimic human mocap demonstrations using body-state features and a phase variable φ ∈ [0, 1].Policies are trained with PPO from discriminator rewards.
- Phase-functioned discriminators: PFNN parameters vary with phase by linearly interpolating between adjacent fixed parameter sets, using five parameter sets in the implementation.The interpolation weights are phase-dependent and correspond to uniformly spaced phase values.
- Evaluation: Figure 10 compares VAIL with other methods using average joint rotation error over three random seeds.The metric measures discrepancy between simulated-character and reference motions.
- Phase-functioned discriminators: PFNN discriminators significantly improve performance across evaluated motion tasks compared with fully connected alternatives, especially when phase information is included structurally.Policies without a phase input perform worst overall, while all methods perform well on simpler skills such as running.
- Discriminator behavior: With the KL constraint enforced, the VDB discriminator converges to approximately 80% accuracy instead of near-perfect accuracy reached by unconstrained VAIL.Dual gradient descent on β effectively enforces the information constraint I_c.
- Video imitation: Video imitation uses a simplified 2D biped to reduce depth ambiguity from monocular demonstrations rendered as sequences of video frames.The biped has 12 degrees of freedom.
D.1 EXPERIMENTAL SETUP
The inverse-reinforcement-learning experiments evaluate reward recovery and transfer in C- and S-shaped maze environments. VAIRL produces smoother, more ground-truth-like rewards, with gradient penalties strengthening this effect, particularly in the S-maze.
- Environments: The experiments use C-maze and larger S-maze environments with continuous two-dimensional actions and position observations but no velocity observations.The S-maze adds a wall between the agent and goal and permits stronger control forces.
- Environments: The C-maze ground-truth reward is r_t = −d_t − 10^-3||a_t||^2, combining distance-to-goal and action penalties.The action penalty is treated as zero in Figure 13.
- Demonstrations: Expert demonstrations are generated by training TRPO agents on the ground-truth training-environment reward and saving 107 trajectories from each final policy.The provided setup specifies additional training changes for the test S-maze.
- Reward recovery: VAIRL tends to recover smoother reward functions that more closely match ground truth than the evaluated IRL baselines.Figures 14 and 15 compare recovered rewards and trajectories on both maze tasks.
- Reward recovery: Adding a gradient penalty enhances reward-recovery similarity for both AIRL and VAIRL, especially when combined with the variational discriminator bottleneck in the S-maze.The setup uses gradient-penalty coefficients that differ across the C-maze and S-maze tasks.
E.1 EXPERIMENTAL SETUP:
The image-generation experiments apply VDB within convolutional GANs and evaluate established datasets and objectives. The implementation uses stochastic latent embeddings with adaptive constraint updates and preserves latent dimensionality.
- Training and evaluation: The experiments use the non-saturating GAN objective except for WGAN-GP, compute FID on samples of size 100002, and omit batch normalization.RMSprop and a fixed learning rate are used across experiments.
- VDB implementation: For convolutional GANs, VDB is implemented as a 1×1 convolution that outputs a Gaussian latent distribution with diagonal covariance.The bottleneck is applied in the final embedding space while preserving latent dimensionality.
- VDB implementation: All image experiments use adaptive β updates with dual stepsize α_β = 10^-5.The implementation also includes instance noise in the image-generation setup.
- Architectures: CIFAR-10 uses a ResNet-based architecture, while CelebA and CelebA-HQ use the architecture from Mescheder et al. (2018).The architecture details are referenced in the corresponding implementation tables.
E.2 RESULTS
The VDB improves GAN image-generation results on CIFAR-10, while offering little improvement on CelebA and supporting direct high-resolution CelebA-HQ training without progressive growing.
- CIFAR-10: The CIFAR-10 comparison evaluates VGAN-GP against WGAN-GP, instance noise, spectral normalization, and gradient penalty.Networks are reported at 750k iterations with Ic = 0.1 and gradient-penalty coefficient wGP = 10.
- CelebA: On CelebA, VGAN-GP does not substantially improve over the baselines, with FID scores of 7.64 for GAN, 7.76 for GP, and 7.25 for VGAN-GP.The passage attributes the limited improvement to an architecture already effectively tuned for the task, where vanilla GAN also trains well.
- CelebA-HQ: VGAN can be trained directly on CelebA-HQ at 1024 × 1024 resolution without progressive growing.The model uses Ic = 0.1, trains with VGAN-GP for 300k iterations, and fits batch size 8 on a single Tesla V100.
- CIFAR-10: FID 18.1: VGAN-GP outperforms GAN, instance noise, spectral normalization, gradient penalty, and WGAN-GP on CIFAR-10.The compared FID scores are 63.6, 30.7, 23.9, 22.6, and 19.9, respectively, while VGAN-GP achieves 18.1.
- Qualitative results: Figures visualize random VGAN samples on CelebA at 128 × 128 and CelebA-HQ at 1024 × 1024 after 300k iterations.The CelebA-HQ models are trained from scratch at full resolution without the progressive scheme described by Karras et al. (2017).