Disentangling Recognition and Decision Regrets in Image-Based Reinforcement Learning

TL;DR

The paper tackles observational overfitting in image-based RL by introducing recognition regret and decision regret to separate errors due to feature extraction from those due to downstream decision-making. It formalizes environments, policies, and regrets, and then demonstrates the framework with a worked three-state example plus empirical studies in Minigrid and Pong, showing how suboptimal recognition or decision policies produce distinct regret signatures. It extends the analysis to generalization, defining generalization regret and decomposing it into recognition and decision components to diagnose over-specific vs under-specific representations. The findings offer a practical diagnostic toolkit for model selection and architectural adjustments to improve generalization in RL tasks with visual inputs, with public code for replication.

Abstract

In image-based reinforcement learning (RL), policies usually operate in two steps: first extracting lower-dimensional features from raw images (the "recognition" step), and then taking actions based on the extracted features (the "decision" step). Extracting features that are spuriously correlated with performance or irrelevant for decision-making can lead to poor generalization performance, known as observational overfitting in image-based RL. In such cases, it can be hard to quantify how much of the error can be attributed to poor feature extraction vs. poor decision-making. To disentangle the two sources of error, we introduce the notions of recognition regret and decision regret. Using these notions, we characterize and disambiguate the two distinct causes behind observational overfitting: over-specific representations, which include features that are not needed for optimal decision-making (leading to high decision regret), vs. under-specific representations, which only include a limited set of features that were spuriously correlated with performance during training (leading to high recognition regret). Finally, we provide illustrative examples of observational overfitting due to both over-specific and under-specific representations in maze environments and the Atari game Pong.

Figures (6)

• Figure 1: Block diagram of an agent consisting of a recognition policy (encoder) followed by a decision policy (output head).
• Figure 2: States (left), observations (center), and representations (right) for our 3-state example at $t\!=\!14$.
• Figure 3: Recognition and Decision Regrets under Randomized Actions and Masked Observations. Both scenarios lead to regret. For random actions, it is mostly decision regret, indicating: Even when forced to occasionally take random actions, agents still learn useful representations of their environment (to be able to optimize their non-random actions). For masked observations, it is recognition regret, indicating: Partially masked observations fundamentally limit the agent's ability to optimize their actions regardless of their decision policy.
• Figure 4: Possible Maze Configurations. The agent is the red triangle, and the goal is the green square. For training, the key-door color is determined based on the door location, but for testing, the key-door color is always gray (making it both an irrelevant and a spurious feature).
• Figure 5: Regret and Generalization Regrets of Identity in Maze Environments. Reducing the number of latent channels prevents over-specific representations and lowers $\widehat{G\!R}{}^{\text{dec}}$. However, this also results in the recognition policy overfitting to the key-door color, ultimately leading to no improvements in terms of the overall generalization regret.