Leveraging weights signals -- Predicting and improving generalizability in reinforcement learning

TL;DR

The paper tackles the generalization challenge in reinforcement learning by proposing a weight-based predictor that estimates an agent's ability to generalize to unseen environments from its neural network weights. It constructs datasets of agent weights and corresponding generalization scores, then develops two predictor architectures (CNN and DNN) to estimate generalization from weights, ultimately integrating the predictor into the PPO loss with a generalization-term to guide training. Empirical results show the DNN-based predictor provides robust signals across Minigrid and CoinRun, and that incorporating this predictor into PPO (Upgraded PPO) yields agents with better generalization than standard PPO, often earlier in training. The work demonstrates that weight-structure signals can be leveraged to improve RL generalization without external data augmentation, and it assesses the practical trade-offs between predictor types and computational considerations. The key contributions include: (i) assembling a dataset of agent weights with measured generalization, (ii) implementing CNN- and DNN-based predictors to estimate generalization from weights, (iii) augmenting the PPO objective with a generalization term guided by the predictor, and (iv) demonstrating superior generalization performance on unseen environments in Minigrid and CoinRun.

Abstract

Generalizability of Reinforcement Learning (RL) agents (ability to perform on environments different from the ones they have been trained on) is a key problem as agents have the tendency to overfit to their training environments. In order to address this problem and offer a solution to increase the generalizability of RL agents, we introduce a new methodology to predict the generalizability score of RL agents based on the internal weights of the agent's neural networks. Using this prediction capability, we propose some changes in the Proximal Policy Optimization (PPO) loss function to boost the generalization score of the agents trained with this upgraded version. Experimental results demonstrate that our improved PPO algorithm yields agents with stronger generalizability compared to the original version.

Figures (15)

• Figure 1: In our approach, we first train a dataset of agents on different environments to serve as training inputs for the generalizability predictor. The generalizability predictor is then used to modify the loss function of the training algorithm to increase the generalization capacity of the agents trained with it.
• Figure 2: From Agents to prediction of generalizability with a CNN, presenting the different steps from training the predictor and using it to predict the generalization score of a given agent based on its weights
• Figure 3: From Agents to prediction of generalizability with DNN,presenting the different steps from training the predictor and using it to predict the generalization score of a given agent based on its weights
• Figure 4: Modification of the agent training algorithm using the trained predictor as loss function component to improve trained agent generalization capability (numbers represent order of actions
• Figure 5: Minigrid SimpleCrossing and Coinrun environments used for training the predictors and testing the improved PPO algorithm