DRED: Zero-Shot Transfer in Reinforcement Learning via Data-Regularised Environment Design

TL;DR

This work addresses how the procedure of sampling training environment instances affects zero-shot generalisation in reinforcement learning. It formalises the link between level sampling and the generalisation gap through mutual information, showing that adaptive sampling acts as an information bottleneck on $\mathrm{I}(L;\pi)$, and introduces data-regularised environment design (DRED) to augment training levels while keeping the augmented distribution aligned with the target CMDP. DRED combines a generative model of level parameters with an adaptive replay mechanism, and a VAE-based variant (VAE-DRED) demonstrates improved ZSG and robustness to edge cases, while controlling distributional shift (ShiftGap). The approach yields substantial gains on Procgen and Minigrid benchmarks, offering a scalable framework for grounding level-generation in practical RL deployments and reducing start-up data requirements. Overall, the paper provides both theoretical insights and a practical, empirically validated method for enhancing zero-shot transfer in RL via data-regularised environment design.

Abstract

Autonomous agents trained using deep reinforcement learning (RL) often lack the ability to successfully generalise to new environments, even when these environments share characteristics with the ones they have encountered during training. In this work, we investigate how the sampling of individual environment instances, or levels, affects the zero-shot generalisation (ZSG) ability of RL agents. We discover that, for deep actor-critic architectures sharing their base layers, prioritising levels according to their value loss minimises the mutual information between the agent's internal representation and the set of training levels in the generated training data. This provides a novel theoretical justification for the regularisation achieved by certain adaptive sampling strategies. We then turn our attention to unsupervised environment design (UED) methods, which assume control over level generation. We find that existing UED methods can significantly shift the training distribution, which translates to low ZSG performance. To prevent both overfitting and distributional shift, we introduce data-regularised environment design (DRED). DRED generates levels using a generative model trained to approximate the ground truth distribution of an initial set of level parameters. Through its grounding, DRED achieves significant improvements in ZSG over adaptive level sampling strategies and UED methods. Our code and experimental data are available at https://github.com/uoe-agents/dred.

Figures (20)

• Figure 1: The agent (yellow) must navigate to the goal (green) but cannot pass through walls (grey) and only observes tiles directly adjacent to itself (highlighted yellow). An agent trained over levels (a)-(c) will transfer zero-shot to level (d) if it has learnt a behavior adapted to the task semantics of following blue tiles to the goal location.
• Figure 2: Aggregated $\text{I}(L;b)$, $\text{GenGap}$, train and test scores of different sampling strategies over 5 seeds across all Procgen games, using the rliable library rliable. Coloured boxes indicate the 95% confidence interval. For each game, the train and test scores and the $\text{GenGap}$ are normalised by the mean score of $(P_S = \mathcal{U})$ over the test set. Per-game scores, mutual information and classifier accuracy are reported in \ref{['app:results_procgen']}, and \ref{['app:procgen']} provides an extended description of our experimental setup.
• Figure 3: Investigating the effect of increasing $|L|$ when adaptive sampling induces distributional shift. Top row reports the test and train performance; the $\text{GenGap}$ is on the bottom left; on the bottom right is compared, for $(S=S^V)$, the $\text{ShiftGap}$ and the score differential with $(P_S=\mathcal{U})$ over the train set. Vertical bars indicate standard error over 5 seeds.
• Figure 4: Overview of the data-regularised environment design framework.
• Figure 5: Aggregate final performance and mean solved rate on $X_\text{test}$, an evaluation set of 2048 levels sampled from $p({\mathbf{x}})$ and held-out during training (top), and on 448 in-context edge cases (bottom). Example layouts from each evaluation set are plotted on the left. The coloured boxes indicate a 99% confidence interval and the black horizontal bars indicate standard error across 5 training seeds. We refer the reader to \ref{['app:cmdp_levelset']} for additional details on our evaluation sets.

Theorems & Definitions (4)

• Theorem 2.1
• Lemma 3.1
• Lemma 3.2
• Lemma 1.1