Rethinking Value Function Learning for Generalization in Reinforcement Learning
Anonymous Author(s)
AfﬁliationAddress email
Abstract
We focus on the problem of training RL agents on multiple training environmentsto improve the generalization performance and sample efﬁciency.
In prior methods,policy and value function are separately optimized with the goal of avoidinginterference and obtaining a more accurate value function using a separate networkarchitecture.
We identify that the value function is more prone to overﬁtting trainingdata and an appropriate penalization of the value function is required for bettertraining and test performance in the multi-environment setting.
To this end, weintroduce Delayed-Critic Policy Gradient (DCPG), which implicitly penalizes thevalue estimates by training the value function less frequently with more trainingdata compared to the policy.
Furthermore, we propose a simple self-supervisedtask that implicitly learns the forward and inverse dynamics of environments usinga single discriminator, which can be jointly optimized with the value function.
Compared to the prior methods, our proposed algorithms signiﬁcantly improvegeneralization performance and sample efﬁciency in the Procgen Benchmark.
Introduction
In recent years, deep reinforcement learning (RL) has achieved remarkable success in various taskssuch as robotic controls and games [24, 15, 28].
To apply RL algorithms to more practical applicationssuch as autonomous vehicles or healthcare systems, they should be robust against the non-stationarityof the real world and capable of generalizing to unseen situations during their deployment.
However,current RL algorithms often fail to generalize to unseen environments with the same goal but differentcontexts even if they achieve high sample efﬁciency in their training environments [12, 36, 7].
Training RL agents on a ﬁnite set of environments and testing them on the full distribution of unseenenvironments is the standard practice for evaluating generalization in RL [8].
In this framework,several methods have attempted to improve generalization by adopting the regularization techniquesthat originate from supervised learning or training robust representations via self-supervised learning[7, 16, 27, 23].
However, these methods have mainly focused on developing new auxiliary objectiveson top of the existing actor-critic algorithms intended for the conventional single-environment settingsuch as PPO [30].
Some recent works identify the interference between policy and value optimizationarising from the multi-environment setting and propose new RL training schemes to obtain an accuratevalue function by decoupling the policy and value function training with separate networks [9, 26].
In this work, we argue that obtaining an accurate value function in the multi-environment setting ismore challenging than in the single-environment setting and it should be optimized with sufﬁcientregularization.
We identify that the value function trained on multiple environments is more prone tomemorizing the training data and cannot generalize to unvisited states of the training environments,which acts detrimentally to the training performance.
Moreover, we ﬁnd that common regularizationmethods that penalize large estimates of the value function, originally developed for preventing the overﬁtting problem in the single-environment setting, are also beneﬁcial to improving both train andtest performance in the multi-environment setting.
However, this beneﬁt comes at the cost of earlyconvergence which hinders further performance enhancement.
Based on these observations, we introduce Delayed-Critic Policy Gradient (DCPG), where the valuenetwork is implicitly penalized by optimizing less frequently with more data than the policy network.
We ﬁnd that the value function with delayed updates suffers less from the overﬁtting problem andimproves the train and test performance.
We also demonstrate that it provides better representationsfor generalization to the policy network with a shared encoder.
Furthermore, we propose a simple selfsupervised task that implicitly learns the forward and inverse dynamics of environments using a singlediscriminator on top of DCPG.
Our methods achieve state-of-the-art generalization performance andsample efﬁciency on the challenging Procgen benchmark [8].
Preliminaries2.1 Generalization in RL
We consider a collection of environments M formulated as Markov Decision Processes (MDPs).
Each environment m ∈ M is deﬁned by a tuple ( S m , A , T m , r m , ρ m , γ ) , where S m is the statespace, A is the action space shared across all environments, T c : S m × A → P ( S m ) is the transitionfunction, r m : S m ×
A → R is the reward function, ρ m is the initial state distribution, and γ ∈ (0 , 1]is the discount factor.
All states in M have equal dimensions and we denote the set of all possiblestates by S .
While the transition and reward functions are deﬁned dependent on an environment m ,we assume that they share some common structures.
A policy π : S → P ( A ) is trained using a ﬁxednumber of training environments M train = { m i } ni =1 .
The objective is to ﬁnd the optimal policy π ∗that maximizes the expected return on the unseen test {+environments E m,s t ,a t [ (cid:80) ∞ t =
+}
γ t r m ( s t , a t )] ,where m ∼ U ( M \ M train ) , s 0 | m ∼ ρ m , a t ∼ π ( · | s t ) , and s t +1 | m ∼ T m ( · | s t , a t ) .
In this paper, we use the Procgen benchmark as a testbed for evaluating generalization in RL [8]. The
Procgen benchmark is a suite of 16 games with high diversity comparable to the ALE benchmark[3].
Each game consists of procedurally generated environment instances with varying contexts ( e.g. ,layout, background, and the position and spawn time of entities), which are also called levels.
Eachlevel is completely determined by an integer seed.
The standard protocol for evaluating generalizationon the Procgen benchmark is to train an agent on a ﬁxed set of training levels and evaluate itsperformance on the unseen test levels sampled from the full distribution of levels [8].
2.2 Proximal Policy Optimization
Proximal Policy Optimization (PPO) is one of the most successful policy gradient algorithms thatlearns a policy π θ and value function V φ parameterized by neural networks [30].
PPO ﬁrst collectstrajectories τ using the old policy π θ old right before update.
Then, the policy network π θ is trainedwith the sampled trajectories for several epochs (often called sample reuse) to maximize the followingclipped surrogate policy objective J π :
[Equation]
where ˆ A t is an advantage estimate at timestep t .
The value network V φ is trained to minimize thefollowing value objective J V :
[Equation]
where ˆ R t = ˆ A t + V φ ( s t ) is a value function target.
It is used to compute the advantage estimates viageneralized advantage estimator (GAE) [29].
In practice, the policy and value networks are optimized simultaneously with shared parameters ( i.e. ,θ = φ ), especially in image-based RL [11, 34].
For example, they can be implemented using a sharedencoder with separate linear heads as shown in Figure 1a.
Sharing the parameters has an advantagein that features learned by each objective can be beneﬁcial to the other.
It also reduces the memorycost and accelerates the training time.
However, shared network design complicates the optimizationas a single network needs to be optimized over multiple objectives whose gradients may have varyingscales and directions.
This also constrains the policy and value networks to be trained with the samebatch of data to the same degree of sample reuse, which severely limits the ﬂexibility of PPO.
[Figure]
2.3 Phasic Policy Gradient
Phasic Policy Gradient (PPG) addresses the problems of sharing parameters and uses separatenetworks as shown in Figure 1b [9].
It also introduces an auxiliary value head V θ to distill usefulfeatures from the value function into the policy network.
PPG alternates between policy and auxiliaryphases.
During the policy phase, the policy network π θ and the value network V φ are trained on newlycollected trajectories to optimize the policy objective J π and the value objective J V , respectively.
Allstates s t and value targets ˆ R t in the trajectories are stored in a buffer B .
During the auxiliary phase,the auxiliary value head V θ and the policy π θ are jointly trained on all data in the buffer to optimizethe following auxiliary objective J aux and policy regularizer C π :
[Equation]
That is, the value function is distilled into the policy network while preserving the original policyπ θ old .
Additionally, the value network V φ is further trained on all data in the buffer to optimize thevalue objective J V .
This stems from the idea that the distillation of a better-trained value function witha separate network and the additional training for a more accurate value function can be beneﬁcial toimproving generalization performance and sample efﬁciency, as claimed in Cobbe et al.
Motivation
3.1 Overﬁtting in Value Function on Multiple Training Environments
We begin by investigating the difﬁculty of training an accurate value network on multiple trainingenvironments.
Obviously, a more accurate value network on given training environments leads tobetter training performance and sample efﬁciency [31].
However, even in a simple setting wherean agent is trained on a single environment, it has been shown that value networks are prone tooverﬁtting in the sense that it does not extrapolate well to unseen states of the training environment[19, 10].
Increasing the number of training environments could further exacerbate the problem.
Intuitively, the value network might rely more on memorization when fed with fewer training samplesper environment given the same number of environment steps.
To corroborate this claim, we measure the stiffness (cosine similarity) of the value loss gradients 109 between states ( s, s (cid:48) ) following the practice in Fort et al.
[13], Bengio et al.
It is deﬁned by
[Equation]
Note that low stiffness indicates that updating the parameters towards minimizing the value loss forone state has a negative impact on minimizing the value loss for other states.
It has been shown that avalue network with lower stiffness is more likely to memorize the train data [4].
Speciﬁcally, we train
PPG agents on the Procgen games while varying the number of training levels from 1 to 200 andcompute the average stiffness across all state pairs in a mini-batch of size 2 14 throughout training.
The detailed experimental settings can be found in Appendix A. [Figure] BigFish [Figure]
Climber
[Figure] StarPilot [Figure] The green lines in Figure 2 show that the stiffness of the value network decreases as the number oftraining environments increases as expected.
This indicates that the value function trained on multipleenvironments is more likely to memorize the training data and cannot accurately predict the valuesfor unseen states from the training environments.
This overﬁtting problem brings us to train the valuenetwork with sufﬁcient regularization.
3.2 Training Value Function with Explicit Regularization
Now, we examine the effectiveness of value function regularization in the multi-environment setting.
We introduce two common approaches taken from the single-environment setting to prevent theoverﬁtting in value functions.
First, discount regularization (DR) is a method that trains a valuefunction with a lower discount factor γ (cid:48)
Second, activation regularization (AR) is a techniquethat trains a value network with L 2 penalty on the output of the value function [2]:
[Equation]
where α > 0 is the regularization coefﬁcient.
To validate the effectiveness of the regularization in the multi-environment setting, we train PPGagents with each of these two methods using 200 training levels on the Procgen games.
Speciﬁcally,we reduce the discount factor from γ = 0 .
999 to γ (cid:48) = 0 .
995 for PPG+DR and use α = 0 .
for
PPG+AR.
We measure the average train and test returns to evaluate the training performance and itstransferability to unseen test environments.
As shown in Figure 3a, value function regularization improves both the train and test performance of
PPG to some extent.
This implies that explicitly suppressing the value function also helps to reduceoverﬁtting in the multi-environment setting.
We also ﬁnd that these regularization methods achievebetter train and test performance in all Procgen games on average.
More details on the experimentalsetting can be found in Appendix B.
Despite its effectiveness, explicit value function regularization can lead to a suboptimal solutionwhen the environment step increases.
Figure 3b shows the average of the true and predicted valuesmeasured at the initial states of training environments for PPG+DR and PPG+AR.
We ﬁnd that thevalue network trained with regularization reaches a plateau too quickly, which implies that excessiveregularization later hinders the value network from learning the correct value function.
This motivatesus to ﬁnd a more ﬂexible regularization that boosts training performance while allowing the valuenetwork to converge to true values.
Delayed-Critic Policy Gradient
In this section, we propose a new policy gradient algorithm that builds upon PPG, namely DelayedCritic Policy Gradient (DCPG), which deals with the overﬁtting problem of the value function in asimple and ﬂexible fashion.
The key idea is that the value network should be optimized less frequentlythan the policy network to suppress the value estimates and reduce overﬁtting. Also, it should betrained with more data to prevent overﬁtting to small, recent training data, as mentioned in [21].
[Figure]
BigFish, Train
[Figure]
Require: Policy network π θ , value network V θ , dynamics head f θ 1: for phase = 1 , 2 , . . .
do 2:
Initialize empty buffer B 3: for iter = 1 , 2 , . .
, N π do (cid:46) Policy phase 4:
Sample trajectories τ using π θ and compute value function target ˆ R t for each state s t ∈ τ 5: Add ( s t , a t , ˆ R t ) to B 6: for epoch = 1 , 2 , . .
, E π do 7:
Optimize policy objective J π ( θ ) and value regularizer C V ( θ ) with τ 8: end for 9: end for 10: for iter = 1 , 2 , . . .
, E aux do (cid:46) Auxiliary phase 11: Optimize value objective J V ( θ ) , dynamics objective J f ( θ ) , and policy regularizer C π ( θ ) with B 12: end for 13: end
DCPG proceeds by alternating the policy and auxiliary phases similar to PPG, but it uses a sharedencoder and does not require an auxiliary value head as shown in Figure 1c.
During the policy phase,new trajectories are collected and the policy network π θ is trained on the trajectories to optimize thepolicy objective J π .
However, the value network V θ is trained on the trajectories to optimize thevalue regularizer C V to preserve the outputs of the original value function V θ old , which is given by
[Equation]
in the sense that updating the value network frequently can further increase approximation errors forunseen states as observed in Section 3.1.
Finally, all states and value targets in the trajectories arestored in a buffer B .
During the auxiliary phase, the value network V θ and policy network π θ are trained with a largernumber of data in the buffer to optimize the value objective J V and the policy regularizer C π .
Notethat since the value network has the same encoder as the policy network, optimizing the valueobjective also plays the role of representation learning for the policy network.
Thus, we do not requirean additional auxiliary head for feature distillation, unlike PPG.
The overall procedure of DCPG canbe found in Algorithm 1.
4.2 Value Analysis of Delayed Value Update
The delayed value update can operate as an implicit penalization of the value function.
Since thepolicy improvement is not reﬂected in the value network immediately, it encourages smaller valueestimates than the true values.
To validate this claim, we train DCPG agents using 200 training levels [Figure] DCPG [Figure]
BigFish
[Figure] DCPG, Climber [Figure] DCPG, Dodgeball
[Figure]
on the Procgen games and measure the true and predicted values for the initial states of the traininglevels.
Figure 4 shows the outputs of the value network are kept lower than the true values at theearly stage of training but recover the true values as training progresses.
More results can be found in Appendix C.
We also evaluate the effectiveness of the delayed value update in mitigating the overﬁtting problemby measuring the stiffness of the value function while varying the number of training levels, The bluelines in Figure 2 show that DCPG has higher stiffness compared to PPG. It implies that our valuenetwork is more robust against overﬁtting as addressed in Section 3.1.
We observe that the delayedupdate of the value network in DCPG serves as a good regularizer for the value network.
More resultscan be found in Appendix A.
4.3 Learning Forward and Inverse Dynamics with Single Discriminator
To further improve generalization performance and sample efﬁciency, we utilize the buffer B to learnan auxiliary task.
Suppose the buffer contains a transition ( s t , a t , s t +1 ) ∈ B .
Since the transitionfunctions are assumed to have a shared structure, it is natural to train a neural network that modelsforward dynamics in feature spaces.
Concretely, given an encoder e θ , a dynamics head f θ takesthe current state embedding of e θ ( s t ) and the current action a t as inputs and predicts the next stateembedding of e θ ( s t +1 ) .
However, it is well known that a forward dynamics model using a neuralnetwork tends to overﬁt on a small number of train data and make poor predictions on unseen states[6].
To address this, we consider learning the forward dynamics implicitly by training a discriminatorthat distinguishes whether the next state is valid, given the current state and action:
[Equation]
Such a discriminator might discard information about the action to distinguish whether a transition isvalid.
For example, the discriminator can predict the valid transition by using only the proximity ofthe current and next states.
To avoid this, we train the discriminator to jointly distinguish whether thecurrent action is valid, given the current and next states ( i.e. , inverse dynamics):
[Equation]
Finally, we optimize the following objective during the auxiliary phase:
[Equation]
where ˆ s t +1 ∼ U
( B\{ s t +1 } ) , ˆ a t ∼ U
( A\{ a t } ) , and η > 0 is the coefﬁcient for the inverse dynamics.
We name the resulting algorithm Dynamics-aware Delayed-Critic Policy Gradient (DDCPG).
Thefull algorithm is described in Algorithm 1 (difference with DCPG is marked in blue).
Experiments
5.1 Experimental Settings
We evaluate the generalization performance and sample efﬁciency of our methods on all 16 games inthe Procgen benchmark.
We use the “easy” difﬁculty mode, which most prior works have focused on. [Figure] DAAC [Figure]
DCPG DDCPG We train agents on 200 training levels generated by seeds ranging from 1 to 200 for 25M environmentsteps and test them on held-out test levels, following the practice in Cobbe et al.
For evaluation,we measure the average return of 100 train and test episodes and report its mean and standarddeviation over 10 runs with different initialization.
We compare our methods with PPO and 4 other baselines that use PPO as the backbone algorithm:
UCB-DrAC, PLR, PPG, and DAAC [27, 18, 9, 26].
UCB-DrAC is a data augmentation methodfor training policy and value networks robust against image transformations.
PLR is a samplingalgorithm for procedurally generated environments that selects the next training level based on itsfuture learning potential.
PPG is the previous state-of-the-art method on the Procgen benchmark.
DAAC builds upon PPG and distills advantage function into the policy network.
For each method, weuse the same set of hyperparameters across all games, following the standard protocol in the ALE and Procgen benchmarks [24, 8].
For our methods, we use the same hyperparameter setting as PPG for afair comparison.
To compare the performance of each method with a single score, we calculate the PPO-normalized score averaged over all Procgen games, which is computed by dividing the averagereturn of each method by PPO, following the practice in Raileanu et al.
For the implementationdetails and hyperparameters, please refer to Appendix D.
5.2 Generalization Performance on Procgen Benchmark
Table 1 shows the average test returns of each method on all 16 Procgen games after training.
DCPGsigniﬁcantly outperforms all the baselines, yielding a 24%p improvement on average relative tothe previous state-of-the-art method.
Notably, the delayed value update without any auxiliary taskleads to better test performance in 8 of 16 games.
Furthermore, DDCPG improves upon DCPG withan 18%p increase in the normalized score. Adding dynamics learning to DCPG can boost the testperformance in 11 of 16 games, showing its effectiveness and consistency across various types ofgames.
We also provide the average training returns on all Procgen games in Appendix E, showingthat DCPG and DDCPG also have better sample efﬁciency in the multi-environment setting. Figure 5 shows the test performance curves of each method on 6 Procgen games throughout training. We ﬁnd that our methods not only achieve superior ﬁnal test performance but also exhibit fasterconvergence compared to the baselines.
It implies that our methods can acquire generalization ability
[Figure]
A v
[Figure]
r a g e r e t u r n
[Figure]
Dodgeball
[Figure]
[Table] with a fewer number of interactions between environments. For example
[Table]
our methods can reach theﬁnal performance of PPG using only 20% of total environment steps in BigFish and Climber and10% of total environment steps in Dodgeball and Miner.
We provide the train and test performancecurves on all Procgen games in Appendix E.
5.3 Ablation Study for Delayed Value Update
The delayed value update not only alleviates the overﬁtting problem but also provides better representations to the policy network that can generalize to unseen test environments.
To validate this, wetrain DCPG agents with an additional value network V φ using a separate encoder (separated DCPG).
The separate value network V φ is trained in the same way as PPG without delay and is only usedfor computing advantages for policy optimization.
The original value network V θ is only used forlearning representations for the policy network with delayed updates.
Figure 6 shows that separate DCPG achieves better test performance than PPG, which uses an aggressively optimized valuefunction to learn representations for the policy network.
This implies that the delayed value updateitself produces more generalizable representations than the value function without any regularization.
For the detailed experimental settings and results, please refer to Appendix F.
5.4 Ablation Study for Dynamics Learning
We conduct an ablation study to show the effectiveness of training forward and inverse dynamicsusing a single discriminator.
We compare the test performance of DDCPG to DCPG with forwarddynamics learning (DCPG+F), inverse dynamics learning (DCPG+I), and forward and inverse dynamics learning using separate discriminators (DCPG+FI).
Table 2 shows that inverse dynamicsprovides additional performance gain only when jointly trained with forward dynamics using a singlediscriminator.
We provide the detailed experimental settings and results in Appendix F.
[Table]
Related Works
Generalization in deep RL One major approach for improving generalization in deep RL is to applyregularization techniques originally developed for supervised learning to prevent overﬁtting, such asbatch normalization [17]. It has been shown that batch normalization and L 2 weight regularizationprovide considerable beneﬁts to generalization [12, 7].
SNI proposes a regularization techniquethat selectively injects noises into policy parameters and learns compressed state representationsusing VIB [16].
Several works employ data augmentation techniques that are commonly used incomputer vision to tackle the generalization problem [20, 33].
UCB-DrAC proposes methods thatautomatically ﬁnd an effective augmentation and introduces new regularization terms to learn robuststate representations under various augmentations [27].
Another line of work uses self-supervised learning techniques to learn invariant representationsthat can generalize to unseen environments.
DIM learns a state representation that can predict therepresentations of successive timesteps using self-supervised learning [22].
DBC trains encoderssuch that the distances in latent space are equivalent to the bisimulation metrics to discard irrelevantinformation for policy [35]. PSE learns state representations that capture long-term behavior proximitybetween states using contrastive learning [1]. CTRL proposes a new clustering algorithm that groupsrepresentations of behaviorally similar trajectories [23].
Most closely related to our work is PPG, which suggests a better training regime to improve bothgeneralization and sample efﬁciency in multi-environment settings [9].
It trains the policy and valuefunction on separate networks with different levels of sample reuse to avoid interference and obtain abetter-trained value function.
DAAC also uses decoupled policy and value networks to train a moreaccurate value function [26].
However, it distills the advantage function into the policy, which is lessprone to overﬁtting to environment-speciﬁc features than the value function.
We focus on regularizingthe value function while obtaining better representation, which distinguishes our work from aboveapproaches.
Learning value function with regularization Discount regularization is one common regularization method to prevent overﬁtting in value function, especially when training data is limited [5].
Alower discount factor can lead to better performance in approximated dynamic programming whenthe approximation error is large [25].
It can also be helpful for mitigating overﬁtting when dealingwith limited data in POMDPs [14].
Activation regularization penalizes the norm of the value functionand has regularization effect similar to that of discount regularization [2].
It has been shown thatactivation regularization is equivalent to the discount regularization in batch TD learning and alsoeffective in online RL [2].
Conclusion
We have investigated the difﬁculty of learning an accurate value function in a multi-environmentsetting and shown that suppressing a value function with appropriate regularization can be beneﬁcialto improving the train and test performance.
Based on this observation, we propose Delayed-Critic Policy Gradient (DCPG), where the value network is implicitly regularized to have lower predictedvalues at the early stage of training with delayed updates compared to the policy network.
Weﬁnd that the delayed value update not only prevents the memorization of training data in the multienvironment setting but also produces more generalizable representations that can extend to unseentest environments.
Furthermore, we propose a simple self-supervised task that implicitly learns bothforward and inverse dynamics using a single discriminator, which can be easily combined with DCPG.
Our methods exhibits state-of-the-
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’scontributions and scope?
[Yes] We clearly reﬂect the main claims in Section 3 and
Section 4.(b)
Did you describe the limitations of
[Table]
[N
[Table]
(d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?
[Yes] We have read the ethics review guidelines and checked our paper has no
conﬂict.
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[N/A]
(b) Did you include complete proofs of all theoretical results?
[N/A]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as a URL)?
[Yes] We provide the code for reproducing the main results in the supplementary material.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] We provide the implementation details and hyperparameters of all baselines and our methods in the supplementary material.
(c) Did you report error bars (e g., with respect to the random seed after running experi-
ments multiple times)?
[Yes] We conduct experiments over 10 training seeds and report
mean and standard deviation.
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes] We describe the computational resource in the supplementary material.
If you are using existing assets (e g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes] We cite the benchmark
used for our experiments.
(b) Did you mention the license of the assets?
[N/A]
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identiﬁableinformation or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
[Table]
(a) Did you include the full text of instructions given to participants and screenshots, ifapplicable?
[N/A](b)
Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]