Why Generalization in RL is Difficult: Epistemic POMDPs and Implicit Partial Observability

TL;DR

The paper reframes RL generalization as a problem of navigating implicit partial observability induced by epistemic uncertainty over unseen MDPs, formalizing this as an Epistemic POMDP. It argues that standard MDP-focused RL methods fail to generalize optimally under test-time uncertainty and provides a Bayes-optimal, memory-based perspective on policy behavior. To address this, it introduces a tractable approximation called Linked Ensembles for the Epistemic POMDP (LEEP), which enforces coordination among an ensemble of policies through a KL-disagreement penalty and an optimistic linking rule. Empirically, LEEP yields significant improvements in generalization on ProcGen benchmarks, reducing train-test gaps compared to traditional PPO baselines. The work highlights the need to model epistemic uncertainty during training to achieve robust RL generalization and points to future directions in scalable posterior modeling and principled POMDP optimization.

Abstract

Generalization is a central challenge for the deployment of reinforcement learning (RL) systems in the real world. In this paper, we show that the sequential structure of the RL problem necessitates new approaches to generalization beyond the well-studied techniques used in supervised learning. While supervised learning methods can generalize effectively without explicitly accounting for epistemic uncertainty, we show that, perhaps surprisingly, this is not the case in RL. We show that generalization to unseen test conditions from a limited number of training conditions induces implicit partial observability, effectively turning even fully-observed MDPs into POMDPs. Informed by this observation, we recast the problem of generalization in RL as solving the induced partially observed Markov decision process, which we call the epistemic POMDP. We demonstrate the failure modes of algorithms that do not appropriately handle this partial observability, and suggest a simple ensemble-based technique for approximately solving the partially observed problem. Empirically, we demonstrate that our simple algorithm derived from the epistemic POMDP achieves significant gains in generalization over current methods on the Procgen benchmark suite.

Figures (9)

• Figure 1: Visualization of the robotic zookeeper example. Standard RL algorithms learn the classifier strategy, since it is optimal in every training zoo, but this strategy is sub-optimal for generalization because peeking generalizes better than the classifier at test-time. This failure occurs due to the following disconnect: while the task is fully-observed since the image uniquely specifies the location of the otter habitat, to an agent that has limited training data, the location is implicitly partially observed at test-time because of the agent's epistemic uncertainty about the parameters of the image classifier.
• Figure 2: Sequential Classification RL Problem. In this task, an agent must keep guessing the label for an image until it gets it correct. To avoid low test return, policies should change actions if the label guessed was incorrect, but standard RL methods fail to do so, instead guessing the same incorrect label repeatedly.
• Figure 3: DQN on RL FashionMNIST. DQN achieves lower test performance than simple variants that leverage the structure of the RL problem.
• Figure 4: Test set return for LEEP and PPO throughout training in four Procgen environments (averaged across 5 random seeds). LEEP achieves higher test returns than PPO on three tasks (Maze, Heist and Dodgeball) and matches test return on Bigfish while having less variance across seeds.
• Figure 5: (top) Performance of LEEP and PPO with only 50 training levels on Maze. (bottom) Ablations of LEEP in Maze.

Theorems & Definitions (20)

• Proposition 5.1
• Remark 5.1
• Remark 5.2: Failure of MDP-Optimal Policies, Propositions \ref{['prop:deterministic-dominated-random']}, \ref{['prop:suboptimal-actions-everywhere']}
• Remark 5.3: Failure of Generic Stochasticity, Proposition \ref{['prop:stochastic-dominated-random']}
• Proposition 6.1
• Proposition 6.2
• Proposition B.1
• proof
• Proposition B.1
• proof