[Paper Review] Decoupling Value and Policy for Generalization in Reinforcement Learning

Roberta Răileanu, Rob Fergus|arXiv (Cornell University)|Feb 20, 2021

Reinforcement Learning in Robotics52 references4 citations

TL;DR

This paper proposes IDAAC, a novel reinforcement learning method that decouples policy and value function learning using separate networks and enforces representation invariance to task-irrelevant visual features. By addressing the policy-value representation asymmetry—where value estimation requires instance-specific details not needed for policy generalization—IDAAC achieves state-of-the-art generalization on Procgen and DeepMind Control benchmarks, especially with limited training levels.

ABSTRACT

Standard deep reinforcement learning algorithms use a shared representation for the policy and value function, especially when training directly from images. However, we argue that more information is needed to accurately estimate the value function than to learn the optimal policy. Consequently, the use of a shared representation for the policy and value function can lead to overfitting. To alleviate this problem, we propose two approaches which are combined to create IDAAC: Invariant Decoupled Advantage Actor-Critic. First, IDAAC decouples the optimization of the policy and value function, using separate networks to model them. Second, it introduces an auxiliary loss which encourages the representation to be invariant to task-irrelevant properties of the environment. IDAAC shows good generalization to unseen environments, achieving a new state-of-the-art on the Procgen benchmark and outperforming popular methods on DeepMind Control tasks with distractors. Our implementation is available at https://github.com/rraileanu/idaac.

Motivation & Objective

• To address poor generalization in deep reinforcement learning when training on limited instances of procedurally generated environments.
• To identify and resolve the policy-value representation asymmetry, where value estimation requires instance-specific features not needed for optimal policy learning.
• To improve generalization by decoupling policy and value function optimization and enforcing invariance to task-irrelevant visual variations.
• To achieve state-of-the-art performance on Procgen and DeepMind Control benchmarks with minimal training data.

Proposed method

• IDAAC decouples the policy and value function by using separate neural networks for each, avoiding shared representation that may encode spurious correlations.
• It introduces an auxiliary loss that encourages the shared representation to be invariant to task-irrelevant visual changes, such as background color or pattern.
• The method uses contrastive learning to minimize representation changes under visual perturbations (e.g., background swaps) while preserving task-relevant state information.
• The policy and value networks are trained jointly with standard advantage actor-critic objectives, but with decoupled backpropagation to reduce interference.
• The approach is evaluated on Procgen and DeepMind Control environments with distractors, using metrics like feature norm, value/advantage difference, and policy divergence under background changes.
• The model is trained on only 200 levels, making it suitable for low-data generalization scenarios.

Experimental results

Research questions

• RQ1 Does sharing a single representation between policy and value function lead to overfitting in procedurally generated environments?
• RQ2 Can decoupling policy and value learning improve generalization to unseen levels?
• RQ3 Does enforcing invariance to irrelevant visual features (e.g., background) improve robustness and generalization?
• RQ4 How does IDAAC compare to SOTA methods like PPO and PPG in low-data and distractor-rich settings?
• RQ5 Can representation invariance be effectively learned via an auxiliary contrastive loss without harming policy performance?

Key findings

• IDAAC achieves state-of-the-art performance on the Procgen benchmark, outperforming prior methods despite training on only 200 levels.
• The method shows significantly improved robustness to background changes: representation norms and value/advantage differences are lower than in PPO and PPG, indicating reduced sensitivity to irrelevant visual features.
• IDAAC learns representations that are invariant to background changes, as shown by low L1 and L2 norm differences across 10 background variants of the same observation.
• The policy robustness, measured via Jensen-Shannon divergence, is comparable or better than baselines, even though JSD is not a perfect metric due to action equivalence in some environments.
• IDAAC generalizes better than PPO-200 and PPG, even though PPO-10k (trained on 10k levels) shows better robustness metrics—demonstrating that IDAAC achieves high generalization with far less data.
• The auxiliary invariance loss effectively reduces memorization of level-specific features in value estimation, without degrading policy performance.