Observations Meet Actions: Learning Control-Sufficient Representations for Robust Policy Generalization

TL;DR

This paper introduces a new framework for reinforcement learning that learns control-sufficient representations to improve policy generalization across varying contexts, using a novel objective and algorithm called BCPO.

Contribution

It formalizes the distinction between observation and control sufficiency, and proposes a variational information-bottleneck approach with BCPO for robust policy learning.

Findings

BCPO outperforms baselines on continuous control benchmarks with shifting parameters.

The approach requires fewer samples while maintaining or improving performance.

Framework unifies theory, diagnostics, and practice for context-based RL.

Abstract

Capturing latent variations ("contexts") is key to deploying reinforcement-learning (RL) agents beyond their training regime. We recast context-based RL as a dual inference-control problem and formally characterize two properties and their hierarchy: observation sufficiency (preserving all predictive information) and control sufficiency (retaining decision-making relevant information). Exploiting this dichotomy, we derive a contextual evidence lower bound(ELBO)-style objective that cleanly separates representation learning from policy learning and optimizes it with Bottlenecked Contextual Policy Optimization (BCPO), an algorithm that places a variational information-bottleneck encoder in front of any off-policy policy learner. On standard continuous-control benchmarks with shifting physical parameters, BCPO matches or surpasses other baselines while using fewer samples and retaining performance far outside the training regime. The framework unifies theory, diagnostics, and practice for context-based RL.