PaperlayerView exact revision on arXiv

Source-linked AI summary

Generalization in Reinforcement Learning with Selective Noise Injection and Information Bottleneck

Maximilian Igl, Kamil Ciosek, Yingzhen Li, Sebastian Tschiatschek, Cheng Zhang, Sam Devlin, Katja Hofmann

arXiv:1910.12911v1cs.LGcs.AIstat.ML

TL;DR

RL policies often overfit training environments, while RL’s model-dependent data generation and early low-data regime make supervised-learning regularization unreliable. The paper proposes SNI and IBAC to adapt noise-based regularization and feature compression to RL, with IBAC-SNI significantly outperforming prior state-of-the-art results on generalization benchmarks.

  • Problem

    RL policies must generalize beyond training environments, but regularization methods developed for supervised learning interact differently with RL’s model-dependent data generation and shifting, low-data regime.

  • Method

    The paper proposes Selective Noise Injection (SNI) to reduce harmful stochasticity and Information Bottleneck Actor Critic (IBAC) to encourage compressed features in RL.

  • Results

    IBAC-SNI significantly outperforms previous state-of-the-art results on the Multiroom and Coinrun generalization tasks.

  • Takeaways & Limitations

    The results support adapting regularization to RL’s data-generation and low-data characteristics, with SNI especially important on complex environments.

  • Takeaways & Limitations

    The evaluation focuses on VIB and Dropout, leaving SNI’s application to other regularization techniques for future work.

Abstract

from arXiv · show

The ability for policies to generalize to new environments is key to the broad application of RL agents. A promising approach to prevent an agent's policy from overfitting to a limited set of training environments is to apply regularization techniques originally developed for supervised learning. However, there are stark differences between supervised learning and RL. We discuss those differences and propose modifications to existing regularization techniques in order to better adapt them to RL. In particular, we focus on regularization techniques relying on the injection of noise into the learned function, a family that includes some of the most widely used approaches such as Dropout and Batch Normalization. To adapt them to RL, we propose Selective Noise Injection (SNI), which maintains the regularizing effect the injected noise has, while mitigating the adverse effects it has on the gradient quality. Furthermore, we demonstrate that the Information Bottleneck (IB) is a particularly well suited regularization technique for RL as it is effective in the low-data regime encountered early on in training RL agents. Combining the IB with SNI, we significantly outperform current state of the art results, including on the recently proposed generalization benchmark Coinrun.

1 Introduction

RL policies often overfit training environments, making generalization essential for autonomous agents. The paper adapts regularization to RL by addressing model-dependent data generation and non-stationary, low-data training.

  • Motivation: Generalization is essential for broad autonomous-agent applications because many RL agents overfit their training environments.Benchmarks commonly test generalization by procedurally generating or modifying video-game levels.
  • Motivation: Regularization methods from supervised learning, including weight decay, dropout, and batch normalization, can improve RL generalization.The paper argues that RL-specific differences must inform how these techniques are applied.
  • RL-specific challenges: RL training data depends on the learned model, so stochastic regularization can harm data quality and gradient estimation.Noise in the policy can prematurely end episodes, while noise can destabilize the critic and off-policy importance weights.
  • Proposed methods: SNI selectively injects noise only when it provides regularization, computing regularized-network outputs deterministically otherwise.The method is evaluated with Dropout and VIB but is intended to apply to most stochastic regularization techniques.
  • Proposed methods: IBAC adapts the Information Bottleneck to actor-critic RL by compressing input features to encourage robust generalization under shifting data distributions.The paper motivates IBAC by the limited diversity of states encountered early in RL training.
  • Evaluation: IBAC-SNI significantly outperforms previous state-of-the-art results on the Multiroom and Coinrun generalization tasks.The evaluation uses PPO and reports benefits from IBAC and SNI individually and in combination.

2 Background

The paper formalizes training across a distribution of MDPs, introduces PPO as its actor-critic baseline, and reviews regularization methods used or extended for RL.

  • RL formulation: Training environments are modeled as MDPs sampled from a distribution q(m), with experiments using either unlimited access or a fixed training set.Each MDP includes state and action spaces, transitions, rewards, and an initial-state distribution.
  • RL formulation: The learning goal is to find a parameterized policy πθ(a|s) that maximizes discounted expected reward.The policy maps states to action distributions and is optimized over the environment distribution and dynamics.
  • PPO: PPO collects trajectories with a rollout policy and optimizes a clipped surrogate objective using an advantage estimate over multiple epochs.The clipping provides an efficient approximate trust-region procedure on collected data.
  • PPO: PPO’s value loss minimizes a clipped upper bound on squared prediction error using a bootstrapped target and previous value function.The overall objective combines the policy and value-function terms.
  • Regularization techniques: Weight decay reduces parameter magnitudes by adding an L2 penalty to the optimization objective.The resulting update includes multiplicative weight shrinkage in addition to minimizing the original loss.
  • Regularization techniques: Data augmentation changes or distorts inputs; this work uses a cutout variant that fills random rectangular image regions with random colors.The method is intended to improve generalization.
  • Regularization techniques: Batch Normalization estimates mini-batch activation statistics, whereas Dropout randomly zeros activations during training to prevent neuron co-adaptation.The architecture applies Dropout to the last hidden layer.
  • Regularization techniques: VIB regularizes a model through a stochastic latent variable Z, minimizing I(X,Z) while maximizing I(Z,Y).Its variational formulation uses an encoder, decoder, approximate latent marginal, and hyperparameter β.

3 The Problem of Using Stochastic Regularization in RL

Stochastic regularization interacts adversely with RL because the learned policy helps generate training data and influences gradient estimation. The paper identifies several noise-induced failure modes motivating SNI.

  • Implication: These interactions explain why stochastic optimization methods can be less effective or detrimental when combined with other RL regularization techniques.The paper uses these observations to motivate selective rather than indiscriminate noise injection.
  • RL gradient structure: Unlike supervised learning, RL’s learned policy participates in data generation, so policy gradients require corrections for the rollout policy.Actor-critic methods also use a learned critic as a baseline and for bootstrapping.
  • RL gradient structure: The rollout policy may differ from the learned policy, but it should remain similar to keep the off-policy correction variance low.The correction depends on the ratio between learned-policy and rollout-policy probabilities.
  • Noise-induced effects: Noise in the rollout policy can produce undesirable actions that prematurely end episodes and degrade the quality of observed data.This affects the trajectories used for subsequent learning.
  • Noise-induced effects: Noise increases the variance of the off-policy correction term, especially when the learned policy has not compensated for the injected stochasticity.This can reduce the quality of policy-gradient estimates.
  • Noise-induced effects: Noise in the critic increases variance in both policy and critic gradient updates through the computation of Vθ.The critic’s role in actor-critic optimization makes its stochasticity affect both update pathways.

4 Method

The method adapts noise-based regularization to reinforcement learning through Selective Noise Injection and introduces Information Bottleneck Actor Critic to learn compressed, potentially more robust features. IBAC-SNI combines both ideas by suspending noise for selected computations while retaining regularization during policy updates.

  • Selective Noise Injection: Selective Noise Injection mitigates adverse effects of stochastic regularization by suspending noise during selected parts of RL gradient computation.The method uses noise-suspended critics and rollout policies, then mixes gradients from noise-injected and noise-suspended policies.
  • Selective Noise Injection: SNI computes the policy gradient as a mixture while sharing the noise-suspended rollout policy across both terms.This lowers off-policy importance-weight variance while retaining the regularizing effect of noise-injected updates and reusing collected data.
  • Information Bottleneck Actor Critic: Early RL training has limited state diversity and high target-value variance, making spurious correlations and robust feature learning difficult.Observed states initially cluster around starting states, while the policy and critic have not sufficiently converged.
  • Information Bottleneck Actor Critic: IBAC regularizes a stochastic latent representation by minimizing βI(S, Z) while maximizing I(Z, A), encouraging compressed features that retain action-relevant information.The approach follows the Variational Information Bottleneck and adds a KL regularization term to a modified policy-gradient objective.
  • Information Bottleneck Actor Critic: IBAC uses the latent mode z = µθ(s) to compute noise-suspended policies and critics, while conditioning both on compressed features.Although IBAC encourages feature compression, it also introduces stochasticity; SNI is applied to address that stochasticity.
  • Information Bottleneck Actor Critic: IBAC-SNI combines Information Bottleneck Actor Critic with Selective Noise Injection, while the study focuses on VIB and Dropout as the most promising techniques without SNI.Applications of SNI to other regularization techniques are left for future work.

5 Experiments

The experiments test whether IBAC and SNI improve feature learning and generalization across a synthetic low-data setting, Multiroom, and Coinrun. IBAC-SNI performs especially well on difficult layouts and significantly outperforms the new Coinrun baseline while SNI reduces instability associated with stochastic regularization.

  • 5.1 Learning Features in the Low-Data Regime: The VIB finds more general features than other regularizers in low-data settings and when competing features differ only slightly in generality.All models fit the training data; test performance measures reliance on the shared feature g_c rather than the less-general feature f_c.
  • 5.2 Multiroom: IBAC-SNI solves 43% of two-room layouts and 21% of three-room layouts, whereas other fully trained agents solve less than 3% and 0%, respectively.Layouts are randomly generated each episode, so multi-room settings require generalization across rarely observed states.
  • 5.3 Coinrun: IBAC agents learn faster and find more general features, translating into higher final performance on previously unseen Coinrun environments.Coinrun contains high-dimensional observations and substantial level variety.
  • 5.3 Coinrun: Almost all previously proposed regularizers reduce performance relative to the weight-decay and data-augmentation baseline, with BatchNorm performing worst.The authors suggest stochastic regularization may add instability when its effect overlaps with existing regularization.
  • 5.3 Coinrun: SNI mitigates the adverse effects of stochastic regularization, and IBAC-SNI significantly outperforms the new state-of-the-art baseline.For IBAC, λ = 0.5 performs better than λ = 1.
  • 5.3 Coinrun: The deterministic formulation reduces approximate KL divergence, while Dropout has higher stochastic KL divergence than IBAC.The KL divergence is used as a proxy for importance-weight variance in the off-policy correction.

6 Related Work

Related work distinguishes robustness to uncertainty from feature robustness. This paper focuses on feature robustness in high-dimensional observations, extending regularization-based approaches by accounting for differences between RL and supervised learning and introducing VIB as an RL regularizer.

  • Robustness to uncertainty: Robustness to uncertainty concerns unknown MDP effects on transition dynamics or rewards, often addressed through domain randomization or adaptive policies.This setting is common in robotics and control with domain shifts.
  • Feature robustness: Feature robustness targets generalization across unseen states by learning better features from high-dimensional observations such as images.Procedurally generated-level benchmarks commonly evaluate this form of generalization.
  • Feature robustness: Feature-robustness methods include diverse observation data, relational architectural biases, and supervised-learning regularizers such as weight decay, dropout, and batch normalization.Prior work applied these methods especially when training environments were limited.
  • Relation to prior work: This paper adapts regularization to RL by accounting for model-generated data and introduces VIB as a regularization technique not considered by closely related work.Earlier RL-VIB studies conditioned action selection on both encoded and raw observations and explored other RL settings.

7 Conclusion

The paper identifies two RL-specific challenges for regularization: model-dependent data generation and a noisy, low-data regime early in training. It addresses them with Selective Noise Injection and Information Bottleneck Actor Critic, whose combination achieves state-of-the-art performance on challenging generalization benchmarks.

  • 7 Conclusion: RL regularization must account for model-generated data and the noisy, low-data regime that complicates learning general features.Stochastic regularization can reduce data quality, while early training may provide limited state diversity.
  • 7 Conclusion: Selective Noise Injection selectively injects noise to avoid reduced data quality and higher gradient variance through a noisy critic.SNI is designed to mitigate adverse effects of stochastic regularization in RL.
  • 7 Conclusion: IBAC-SNI achieves state-of-the-art performance on challenging generalization benchmarks by combining compressed-feature learning with selective noise injection.The paper reports this outcome across the evaluated challenging benchmarks.

A Dropout with SNI

Applying SNI to Dropout requires suspending the noise when computing the policy. The implementation samples one dropout mask and keeps it fixed for all gradient updates using the collected data.

  • A Dropout with SNI: SNI applied to Dropout requires deciding how to suspend the noise when computing the policy.The paper denotes the resulting policy as ¯πθ.
  • A Dropout with SNI: The implementation samples one dropout mask and keeps it fixed for all gradient updates using the collected data.This choice was empirically better than applying no mask and scaling activations.
  • A Dropout with SNI: The fixed-mask implementation follows the approach used in the cited prior work.The passage attributes this implementation choice to reference [11].

B Supervised Classification Task

The supervised classification task generates inputs by embedding each class label through two pattern families, f^c and g^c, with only g^c shared across training and testing. Hyperparameter sweeps indicate stable qualitative results.

  • B Supervised Classification Task: Each data point encodes its class label twice, using the pattern families f^c and g^c.The two encodings use ω_f and ω_g possible patterns, respectively.
  • B Supervised Classification Task: Increasing ω makes an encoding less general because each pattern applies to fewer data points.The experiment uses n_g = 3, σ_ε = 1, d_x = 100, d_g = 20, and n_c = 5.
  • B Supervised Classification Task: The f^c and g^c patterns differ in dimensionality, with f^c vectors of dimension d_x and g^c vectors of smaller dimension d_g < d_x.The patterns are generated from random functions and evaluated at d_x points.
  • B Supervised Classification Task: The g^c encoding replaces a randomly located part of the noisy f^c-based vector, using one of n_g fixed possible locations shared across training and testing.The replacement location is randomly drawn per data point but restricted to fixed experiment-wide locations.
  • B Supervised Classification Task: The test-set loss shows the same qualitative results across multiple hyperparameters, indicating stability over a wide hyperparameter range.Figure 5 reports test-set error for multiple hyperparameter settings.

C Multiroom

The Multiroom environment represents observations as 11 × 11 × 3 grids and uses a three-layer CNN followed by a 64-unit hidden layer. The experiments use PPO with specified optimization and regularization settings.

  • C Multiroom: Observations have space 11 × 11 × 3, with channels encoding object type, object features, and door properties.The spatial grid contains 11 × 11 locations.
  • C Multiroom: The agent uses a three-layer CNN with 16, 32, and 32 filters, followed by a 64-unit hidden layer.IBAC or Dropout is applied to the hidden layer when enabled.
  • C Multiroom: Training uses PPO with Adam, γ = 0.99, λ_GAE = 0.95, entropy coefficient λ_H = 0.01, value-loss coefficient λ_V = 0.5, and gradient clipping at 0.5.The learning rate is 7e-4.

D Coinrun

On Coinrun, the experiments compare Dropout variants and IBAC with or without Selective Noise Injection under different regularization settings. SNI improves performance, while Dropout combined with weight decay and data augmentation remains substantially below the baseline.

  • Experimental setup: The Coinrun experiments use the Impala architecture, default policy-gradient hyperparameters, and the benchmark authors’ codebase.Dropout and IBAC were applied to the last hidden layer, with best reported settings of p_d = 0.2 for Dropout and 1e−4 for IBAC and weight decay.
  • Dropout implementations: λ = 1 learns fastest across the evaluated Dropout implementations.The authors attribute this possibly to high importance-weight variance in the stochastic SNI term when λ < 1.
  • Dropout implementations: All Dropout implementations converge to roughly the same value, significantly below the baseline agent, with weight decay and data augmentation.This result indicates that Dropout is not suitable for combination with those regularizers in this setting.
  • IBAC and SNI: SNI improves test performance for both IBAC and Dropout without weight decay and data augmentation.The comparison evaluates each method with and without SNI on the test environments.
  • IBAC and SNI: IBAC does not prevent overfitting by itself, although it leads to faster learning.Test performance decreases with longer training, suggesting that IBAC alone ultimately overfits the test-set environments.
Loading 1910.12911v1…