PaperlayerView exact revision on arXiv

Source-linked AI summary

Decoupling Value and Policy for Generalization in Reinforcement Learning

Roberta Raileanu, Rob Fergus

arXiv:2102.10330v2cs.LGcs.AI

TL;DR

Deep RL agents can struggle to generalize to unseen procedurally generated environments, and shared policy–value representations can require instance-specific features. IDAAC separates policy and value optimization and adds an auxiliary loss for instance-invariant policy features; it reports strong results on Procgen and DeepMind Control tasks with distractors.

  • Problem

    In procedurally generated environments, agents trained on relatively few levels must generalize to unseen instances, while shared policy–value representations can require instance-specific features that hurt policy generalization.

    Source: §1 · p. 1

  • Method

    IDAAC uses separate networks for policy and value, predicts advantage for policy training, and adds an auxiliary loss encouraging policy representations to be invariant to task instances.

    Source: §4.2 · p. 4, §6 · p. 9

  • Results

    IDAAC achieves state-of-the-art generalization on Procgen and outperforms strong RL algorithms on DeepMind Control tasks with distractors.

    Source: §6 · p. 9

  • Takeaways & Limitations

    The authors report that IDAAC combines good generalization with accurate value estimates and learns representations and predictions more robust to cosmetic observation changes that leave the underlying state unchanged.

    Source: §6 · p. 9

    The approach’s invariance bias targets the number of remaining episode steps and may not help all problems; the authors also lack a firm theoretical argument for why advantage prediction improves generalization.

    Source: §6 · p. 9

Abstract

from arXiv · show

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.

1. Introduction

The paper studies generalization to unseen procedurally generated levels and identifies shared policy–value representations as a source of overfitting. It proposes separate policy and value networks plus an auxiliary loss that encourages policy representations to be instance-invariant.

1. Introduction

  • The target setting is generalization from a relatively small training set to unseen levels that share a goal but differ in appearance, layout, and dynamics.These differences can include backgrounds and the locations, shapes, and colors of entities.Source: §1 · p. 1

1.1. Policy-Value Representation Asymmetry

  • Accurate value estimation can require instance-specific features that are unnecessary for learning the optimal policy, so sharing representations can lead to policies that generalize poorly.In Ninja, semantically identical initial observations can have different values because the levels have different lengths.Source: §1 · p. 1, §1.1 · pp. 1–2, Figure 1 · p. 2
  • IDAAC uses separate networks for the policy and value, with the policy network also predicting advantage, and adds a loss encouraging policy representations to be invariant to the task instance.The separate value network supplies estimates used to compute advantage targets.Source: §1.1 · p. 2
  • The paper reports state-of-the-art generalization on Procgen and stronger performance than popular methods on DeepMind Control tasks with distractors.It presents identifying shared policy–value representations as a source of overfitting and proposes decoupling and instance-invariance as remedies.Source: §1.1 · p. 2

2. Related Work

Prior work addresses deep RL generalization through regularization, data augmentation, representation learning, and other strategies. This paper distinguishes its focus on policy–value representation asymmetry and overfitting to spurious correlations.

Generalization in Deep RL

  • Prior deep RL generalization work has explored regularization, data augmentation, and representation learning, including state abstractions and feature alignment.Other approaches include distillation, surprise minimization, temporal mutual information, and automatic curricula.Source: §2 · pp. 2–3

Decoupling the Policy and Value Function

  • Although parameter sharing is common in pixel-based deep RL, some studies decouple policy and value optimization to improve sample efficiency.The paper instead focuses on generalization to unseen environments and identifies shared features as a source of overfitting to spurious correlations.Source: §2 · p. 3
  • The paper contrasts its approach with Cobbe et al. (2020): it avoids an auxiliary value-distillation phase, excludes value gradients from policy updates, and uses advantage and invariance losses for policy training.Prior advantage prediction in Q-learning does not learn policies directly.Source: §2 · p. 3

3. Background

The paper formalizes training across a distribution of partially observable Markov decision processes, with evaluation over the environment distribution. Procgen instantiates this setting using procedurally generated levels.

3. Background

  • Each environment m is a POMDP specified by state and observation spaces, actions, transitions, observations, rewards, and a discount factor.The training set is a fixed subset sampled from the POMDP distribution.Source: §3 · p. 3
  • The objective is to find a policy maximizing expected discounted reward over the distribution of POMDPs.Training access is restricted to a fixed set of sampled environments.Source: §3 · p. 3
  • Procgen provides 16 procedurally generated games, with each level corresponding to a POMDP sampled from its game’s distribution.Agents train on 200 fixed levels and are tested across the full distribution of levels.Source: §3 · p. 3

4. Invariant Decoupled Advantage Actor-Critic

IDAAC combines separate policy and value networks with advantage prediction and an adversarial loss that encourages policy features to omit episode-specific information. Its design aims to retain useful policy-learning gradients while reducing overfitting to instance-dependent cues.

4.1. Decoupling the Policy and Value Function

  • Separate networks alone can impair effective policy learning because value-loss gradients provide denser, less noisy features than sparse, high-variance policy gradients.The method therefore seeks another useful source of policy-network gradients without relying on value gradients.Source: §4.1 · pp. 3–4

4. Invariant Decoupled Advantage Actor-Critic

  • DAAC separates policy and value networks, while IDAAC adds a regularizer to the policy encoder that discourages episode-specific information.The value estimates are used to compute advantage targets.Source: §4 · p. 3, Figure 2 · p. 4

4.2. Using the Advantage instead of the Value Function

  • DAAC trains the policy network to predict advantage alongside optimizing the policy, using a separate value network to compute advantage targets.Advantage is a relative measure of an action’s expected additional return and may vary less with episode length than absolute state value.Source: §4.2 · p. 4
  • Advantage prediction provides useful policy-learning gradients and mitigates overfitting associated with using value gradients to update the policy network.The paper reports improved performance on test environments.Source: §4.2 · p. 5

4.3. Learning Instance-Invariant Features

  • IDAAC discourages policy features from carrying information about remaining episode steps by adversarially training an encoder and discriminator on the order of observations from one trajectory.The discriminator learns to classify which observation came first, while the encoder is trained to make that order uncertain; this loss is applied to the policy network, not the value network.Source: §4.3 · p. 5

5. Experiments

The experiments compare DAAC and IDAAC with established RL methods on Procgen and DeepMind Control tasks, then examine their ablations and the relationship between value estimation and generalization. Across these settings, the results support decoupling policy and value learning and using advantage gradients to improve generalization.

5. Experiments

  • The evaluation covers three DeepMind Control tasks with distractor backgrounds and Procgen’s 16 procedurally generated games, using Procgen’s easy setup of 200 training levels per game.Procgen agents are tested on the full distribution of levels.Source: §5 · p. 5

5.1. Generalization Performance on Procgen

  • DAAC and IDAAC outperform all baselines on aggregated Procgen train and test performance, with IDAAC further improving on DAAC and both showing superior test results.The approaches establish a new state-of-the-art on Procgen’s test distribution after 25M environment steps.Source: §5.1 · p. 6, Table 1 · p. 6

5.2. Ablations

  • DAAC outperforms PPO, DVAC, and AAC on both train and test environments across four Procgen games.AAC’s weaker generalization suggests that adding advantage prediction does not prevent overfitting when value-loss gradients still update the policy; DVAC’s weaker test performance shows DAAC’s gains are not merely due to a separate value network or a more accurate value estimate.Source: §5.2 · p. 7

5.3. DeepMind Control with Distractors

  • DAAC and IDAAC outperform PPO, UCB-DrAC, and PPG on three DeepMind Control tasks with synthetic or natural video backgrounds.The backgrounds are sampled from videos at the start of each episode, creating spurious correlations between backgrounds and rewards.Source: §5.3 · p. 8

5.4. Value Loss and Generalization

  • PPO agents trained on more Procgen levels generalize better to unseen levels despite having higher value loss, a positive correlation between value loss and generalization.The authors interpret this observation as consistent with shared policy-value representations overfitting.Source: §5.4 · p. 8

5.5. Advantage vs. Value During an Episode

  • On a Ninja level, PPO trained on 200 levels has accurate values but poor generalization and near-linear value predictions over episode steps; PPO trained on 10k levels generalizes better but has lower value accuracy.With DAAC, the advantage estimate lacks a clear dependence on episode step while the value estimate remains linear, without negatively affecting the policy because the networks are separate.Source: Table 2 · p. 9, §5.5 · p. 9

6. Discussion

IDAAC addresses policy–value representation asymmetry with decoupled optimization and instance-invariant policy representations, while reporting robust representations and strong generalization. The discussion limits the expected value of step-count invariance, and the accompanying PPO appendix describes its clipped policy objective.

6. Discussion

  • IDAAC combines decoupled policy–value optimization with an instance-invariance loss and reports strong generalization and robustness to cosmetic observation changes.The discussion also reports accurate value estimates alongside good generalization.Source: §6 · p. 9
  • Step-count invariance may not help all problems; the authors expect the most gains with partial observability, goal states, and varying episode lengths.They identify investigating other auxiliary losses as a direction for future work.Source: §6 · p. 9

A. PPO

  • PPO alternates between collecting environment data and optimizing an actor-critic objective with stochastic gradient ascent.The policy learns a policy and value function for an MDP.Source: A. PPO · p. 13
  • PPO’s policy objective uses importance sampling, and its importance weight estimates the advantage function.The behavior policy generates training states and actions, while the optimized policy generates the target state-action distribution.Source: A. PPO · p. 13
  • PPO uses a clipped surrogate objective that constrains the policy probability ratio to [1 − ϵ, 1 + ϵ] and discourages extreme policy updates.The objective takes the minimum of the original and clipped versions.Source: A. PPO · p. 13

B. DAAC and IDAAC

DAAC and IDAAC use separate policy and value networks, with IDAAC adding an instance-invariance loss. Their optimization alternates between policy and value updates, with independently controlled update schedules.

B. DAAC and IDAAC

  • Separate networks allow different epoch counts and value updates every N_π policy updates; less frequent value updates can improve stability.The authors also report that the value network permits greater sample reuse than the policy network; Table 3 lists the hyperparameters used.Source: B. DAAC and IDAAC · p. 14, Table 3 · p. 15
  • Each iteration collects trajectories and computes value and advantage targets, which support policy-gradient and advantage-prediction objectives.The algorithm descriptions include separate policy and value update loops.Source: B. DAAC and IDAAC · p. 14
  • IDAAC adds an instance-invariance loss to the decoupled policy and value networks used by DAAC.The IDAAC policy update includes an invariance loss and a discriminator loss.Source: B. DAAC and IDAAC · p. 14

C. Hyperparameters

The experiments build on a PPO implementation and ResNet architecture, use a hyperparameter search to select the reported settings, and train Procgen baselines under a shared setup.

C. Hyperparameters

  • The methods use a PPO implementation and the ResNet architecture used for strong Procgen results, with Cobbe et al. (2019) easy-mode hyperparameters unless noted otherwise.The implementation is based on Kostrikov (2018)’s PPO implementation.Source: C. Hyperparameters · p. 15
  • The selected settings were E_π = 1, E_V = 9, N_π = 1, α_a = 0.25, and α_i = 0.001.These values were selected as best overall in the reported hyperparameter search; experiments used Adam.Source: C. Hyperparameters · p. 15
  • Baselines were trained for 25M frames on Procgen easy mode using the same 200 training levels and full-distribution testing setup as the proposed methods.The authors used reported best hyperparameters for Mixreg, PLR, UCB-DrAC, and PPG, and released or recommended settings for other baselines.Source: C. Hyperparameters · p. 15

D. Procgen Results

On Procgen, IDAAC and DAAC outperform other methods on most test games and are comparable on most remaining games. DAAC also outperforms its ablations across train and test environments for most games.

D. Procgen Results

  • IDAAC and DAAC outperform other methods on most Procgen test games and are comparable on most remaining games.Figure 7 reports test performance for IDAAC, DAAC, PPG, UCB-DrAC, and PPO across Procgen games.Source: Figure 7 · p. 16, D. Procgen Results · p. 16
  • IDAAC outperforms the other methods on most Procgen training games and is significantly better than PPO.Figure 8 presents the training performance comparison.Source: Figure 8 · p. 17
  • DAAC is substantially better than DVAC and AAC on most train and test games, and comparable on the remaining games, including CaveFlyer and Maze.Figures 9 and 10 compare PPO, DAAC, and the two ablations across Procgen games.Source: D. Procgen Results · p. 16, Figure 9 · p. 18, Figure 10 · p. 19
  • Tables 4 and 5 report final train and test scores for Procgen games after 25M environment steps, with means and standard deviations over 10 runs.The tables include the methods and baselines’ scores for each game.Source: D. Procgen Results · p. 20, Table 4 · p. 20, Table 5 · p. 20

E. DeepMind Control Suite Experiments

On DeepMind Control tasks with natural and synthetic video distractors, DAAC and IDAAC outperform PPO, PPG, and UCB-DrAC. Evaluation uses held-out environments and averages return over training.

E. DeepMind Control Suite Experiments

F. Value Loss and Generalization

Across Procgen games, value loss is positively correlated with both test score and the number of training levels. This runs counter to the usual single-environment relationship and is consistent with the paper's account of overfitting under shared representations.

F. Value Loss and Generalization

  • For most Procgen games, value loss is positively correlated with test score and the number of training levels.The figures compare PPO agents trained on 200, 500, 1000, 2000, 5000, and 10000 levels.Source: F. Value Loss and Generalization · p. 22, Figure 12 · p. 22, Figure 13 · p. 23
  • Unlike the usual single-environment intuition, higher value loss here accompanies better test performance and more training levels.The authors say this supports their claim that accurate value estimation can encourage representations to overfit when policy and value share a network.Source: F. Value Loss and Generalization · p. 22

G. Advantage vs. Value During an Episode

The analysis contrasts value predictions, which can reflect level-specific or timestep information, with advantages that are less dependent on those idiosyncrasies. It also examines how value variance changes across training levels.

G. Advantage vs. Value During an Episode

  • For two Ninja levels, true and estimated values differ while the true and estimated advantages are approximately the same.The comparison uses initial observations that are semantically identical; the advantage shown is for the noop action, with similar conclusions for other actions.Source: G. Advantage vs. Value During an Episode · p. 25, Figure 15 · p. 25
  • With 200 training levels, PPO and PPG values rise nearly linearly with episode step, while DAAC's advantage lacks that dependence.A PPO model trained on 10k levels also lacks the linear trend; DAAC's value estimate still trends linearly, but separate policy and value networks prevent that trend from negatively affecting its policy.Source: G. Advantage vs. Value During an Episode · p. 25, Figure 16 · p. 25, Figure 17 · p. 26

H. Value Variance

  • Across Climber, Jumper, and Ninja, predicted-value variance for initial observations generally decreases as the number of training levels increases.The 200 observations used to compute the variance were part of every agent's training set; the authors relate lower variance to better generalization.Source: H. Value Variance · p. 27, Figure 20 · p. 27

I. Robustness to Spurious Correlations

The robustness analysis varies visual backgrounds while holding the underlying state fixed, comparing representation, output, and policy changes. DAAC and IDAAC show greater robustness than PPO and PPG in the reported comparisons, with qualifications for JSD and PPO-10k.

I. Robustness to Spurious Correlations

  • The comparison changes backgrounds without changing the underlying state, measuring feature L1/L2 norms, output differences, and policy JSD.The experiment uses 1,000 observations from 20 training levels, creates 10 background variants per observation, and reports means and standard deviations.Source: §I · p. 27
  • For Ninja, Jumper, and Climber, DAAC and IDAAC representations are more robust to background changes than PPO and PPG trained on the same number of levels.The figures report feature norms, value or advantage differences, and policy JSD; the methods' auxiliary-output differences are also smaller than PPO and PPG's.Source: §I · p. 27, Figure 21 · p. 28, Figure 22 · p. 28
  • JSD may not reflect semantic policy differences when distinct actions have the same environmental effect, making the policy-robustness analysis inconclusive.In Ninja, for example, two actions move the agent to the right.Source: §I · p. 28
  • PPO-10k is an upper-bound comparison rather than a directly fair baseline because it trains on 10,000 levels while DAAC and IDAAC train on 200.In some cases PPO-10k has lower feature-norm and value-or-advantage differences; the authors note that 10k-level training generally reaches the test distribution for most Procgen games.Source: §I · p. 28
Loading 2102.10330v2…