Smaller Abstract State Spaces Enable Cross-Scale Generalization in Reinforcement Learning

2026-05-18

AI Summary

This work addresses the limited out-of-distribution (OOD) generalization of reinforcement learning agents across tasks with varying scales, a capability that remains far from human-like abstraction transfer. Building upon partially observable Markov decision processes (POMDPs), the study introduces an abstraction function to assess equivalence among experiences and proposes a successor-weighted model reduction method to construct a more compact abstract state space. It establishes the first theoretical framework for OOD generalization in reinforcement learning within the POMDP setting by extending state abstraction beyond fully observable environments. Through a decomposition of generalization error, the paper theoretically demonstrates that reducing the size of the abstract state space effectively lowers both approximation and estimation errors, thereby enhancing generalization performance on cross-scale tasks.

Abstract

While humans readily generalize abstract concepts to more complex or larger tasks, building Reinforcement Learning (RL) systems with this ability remains elusive. Here, we present the first theoretical model of how such Out-of-Distribution (OOD) generalization can be achieved in RL agents. Our approach considers Partially Observable Markov Decision Processes (POMDPs) and assumes that an intelligent agent uses an abstraction function to determine which experiences can be treated as equivalent and which must be distinguished. First, we extend the existing state abstraction framework and proof techniques to POMDPs. Then, we define a successor-weighted model reduction, a model reduction variant that enables compression into smaller abstract spaces than prior definitions allow. We derive a bound on the agent's OOD test performance, thereby defining the conditions under which OOD generalization is achievable. This bound decomposes an agent's performance loss into approximation and estimation errors, revealing how reducing an agent's abstract state space size improves test performance and OOD generalization. Our analysis suggests that constraining an agent to operate over a small, finite set of abstract states is necessary for achieving generalization to more complex tasks. Our results motivate further research into learning RL architectures that scale across tasks of varying complexity levels.

Problem

• out-of-distribution generalization
• reinforcement learning
• state abstraction
• POMDPs
• cross-scale generalization

Innovation

• state abstraction
• out-of-distribution generalization
• POMDP
• model reduction
• reinforcement learning