Observations Meet Actions: Learning Control-Sufficient Representations for Robust Policy Generalization
This work addresses robust policy generalization for reinforcement learning agents in varying environments, representing a novel method for a known bottleneck.
The paper tackled the problem of robust policy generalization in reinforcement learning by introducing a dual inference-control framework that distinguishes observation sufficiency from control sufficiency, resulting in BCPO, which matches or surpasses baselines on continuous-control benchmarks with shifting parameters while using fewer samples and maintaining performance outside the training regime.
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.