Rethinking Value Function Learning for Generalization in Reinforcement Learning
Anonymous Author(s)
AffiliationAddress email
Abstract
We focus on the problem of training RL agents on multiple training environmentsto improve the generalization performance and sample efficiency.
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 overfitting 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 significantly improvegeneralization performance and sample efficiency in the Procgen Benchmark.
Introduction
In recent years, deep reinforcement learning (RL) has achieved remarkable success in various taskssuch as robotic controls and games [28, 19, 33].
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 differentvisuals ( i.e. , observational generalization) even if they achieve high sample efficiency in their trainingenvironments [14, 41, 9].
Training RL agents on a finite set of environments and testing them on unseen environments is thestandard practice for evaluating observational generalization in RL [10].
In this framework, severalmethods have attempted to improve the generalization by adopting the regularization techniques thatoriginate from supervised learning or training robust state representations via self-supervised learning[9, 20, 31, 27].
However, these methods have mainly focused on developing new auxiliary objectiveson top of the existing RL algorithms intended for the conventional single-environment setting such as PPO [35].
Some recent works identify the interference between policy and value optimization arisingfrom the multiple training environments and propose new RL training schemes to obtain an accuratevalue function by decoupling the policy and value function training with separate networks [11, 30].
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 sufficientregularization.
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 find that common regularization methods that penalize large estimates of the value function, originally developed for preventing theoverfitting problem in the single-environment setting, are also beneficial to improving both train andtest performance in the multi-environment setting.
However, this benefit comes at the cost of earlyconvergence which hinders further performance enhancement.
Based on these observations, we introduce a new model-free algorithm named Delayed-Critic Policy Gradient (DCPG), where the value network is implicitly penalized by optimizing less frequently withmore data than the policy network.
We find that the value function with delayed updates suffers lessfrom the overfitting problem and improves the train and test performance.
We also demonstrate that itprovides better representations to the policy network with a shared encoder.
Furthermore, we proposea simple self-supervised task that implicitly learns the forward and inverse dynamics of environmentsusing a single discriminator on top of DCPG.
Our methods achieve state-of-the-art observationalgeneralization performance compared to prior model-free methods on the Procgen benchmark [10].
Preliminaries2.1 Observational Generalization in RL
We consider a collection of Markov Decision Processes (MDPs) denoted by M .
Each environmentm ∈ M is defined by a tuple ( S m , A , T m , r m , ρ m , γ ) , where S m ⊂ S is the state space with imageobservations, A is the action space shared across all environments, T m : S m × A → ∆( S m ) is thetransition function, r m : S m ×
A → R is the reward function, ρ m is the initial state distribution,and γ ∈ (0 , 1] is the discount factor.
We assume that there exists observational variation betweenenvironments, but the transition and reward functions share common structures across environments.
A policy π : S → ∆( A ) is trained using a fixed number of training environments M train = { m i } ni =1 .
The objective is to find the optimal policy π ∗ that maximizes the expected return on the unseen testenvironments M test [-environments E m,s t ,a t [ (cid:80) ∞ t =
-]=
M \ M train .
In this paper, we use the Procgen benchmark as a testbed for evaluating observational generalizationin RL [10].
The Procgen benchmark is a suite of 16 different games with high diversity comparableto the ALE benchmark [5].
Each game consists of procedurally generated environment instanceswith varying visuals ( e.g. , layout, background, and the position and spawn time of entities), which arealso called levels.
Each level is completely determined by an integer seed.
The standard evaluationprotocol on the Procgen benchmark is to train an agent on a fixed set of training levels and evaluateits performance on the unseen test levels sampled from the full distribution of levels [10].
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 [35].
PPO first 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 π designed to constrain the size of the policy update:
[Equation]
where ˆ A t is an estimate of the advantage function at timestep t .
The value network V ϕ is trained tominimize the following 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) [34].
In practice, the policy and value networks are optimized simultaneously with shared parameters ( i.e. ,θ = ϕ ), especially in image-based RL [13, 39].
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 beneficial 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 flexibility of PPO.
[Figure]
2.3 Phasic Policy Gradient
Phasic Policy Gradient (PPG) is an algorithm built upon PPO that significantly improves the observational generalization performance by addressing the problems of sharing parameters [11].
Specifically, PPG employs separate encoders for the policy and value networks as shown in Figure 1b and introduces an auxiliary value head V θ to distill useful features from the value function into the policynetwork.
For training, PPG alternates between policy and auxiliary phases.
During the policy phase,which is repeated N π times, the policy and value network are trained on newly-collected trajectoriesto optimize the policy and value objective, respectively, in the same way as PPO.
Then, all statesand value targets in the trajectories are stored in a buffer B .
During the auxiliary phase, the auxiliaryvalue head and the policy are jointly trained on all data in the buffer, which is N π times larger thanthe policy phase, to optimize the following auxiliary value objective J aux and policy regularizer C π :
[Equation]
That is, the value network is distilled into the policy encoder while preserving the original policy π θ old .
In addition, the value network is further trained on all data in the buffer to optimize the value objective.
It is claimed that the distillation of a better-trained value function with a separate network and theadditional training for a more accurate value function can be beneficial to improving generalizationperformance and sample efficiency [11].
Motivation
3.1 Overfitting in Value Function on Multiple Training Environments
We begin by investigating the difficulty of training an accurate value function on multiple trainingenvironments.
Obviously, a more accurate value network on given training environments leads tobetter training performance [37].
However, even in a simple setting where an agent is trained on asingle environment, it has been shown that value networks are prone to overfitting in the sense thatit does not extrapolate well to unseen states of the training environment [23, 12]. Also, it has beenclaimed that the value networks cannot fit the true value functions [21, 15].
Increasing the numberof training environments could further exacerbate the problem.
Intuitively, the value network mightrely more on memorization when fed with fewer training samples per environment given the samenumber of environment steps.
To corroborate this claim, we measure the stiffness (cosine similarity) of the value loss gradients 112 between states ( s, s ′ ) following the practice in Fort et al.
[16], Bengio et al.
[6], which is defined 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 a [Figure] BigFish
[Figure] StarPilot [Figure] value network with lower stiffness is more likely to memorize the train data [6].
Specifically, 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.
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 overfitting problem brings us to train the valuenetwork with sufficient 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 theoverfitting in value functions.
First, discount regularization (DR) is a method that trains a valuefunction with a lower discount factor γ ′
Second, activation regularization (AR) is a techniquethat trains a value network with L 2 penalty on the output of the value function [3]:
[Equation]
where α > 0 is the regularization coefficient.
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.
Specifically,we reduce the discount factor from γ = 0 .
999 to γ ′ = 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 reduceoverfitting in the multi-environment setting.
We also find 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 find 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 find a more flexible 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, namely Delayed-Critic Policy Gradient(DCPG), which deals with the overfitting problem of the value function in a simple and flexible [Figure] t u r n
BigFish
[Figure]
BigFish
[Figure]
fashion.
The key idea is that the value network should be optimized less frequently than the policynetwork to suppress the value estimates and reduce overfitting.
Also, it should be trained with moredata to prevent overfitting to small, recent training data, as mentioned in [25].
Algorithm Require: Policy network π θ , value network V θ , dynamics discriminator f θ 1: for phase = 1 , 2 , . . .
do 2:
Initialize empty buffer B 3: for iter = 1 , 2 , . .
, N π do ▷ Policy phase 4:
Sample trajectories τ using π θ and compute value function target ˆ R t for each state s t ∈ τ 5: for epoch = 1 , 2 , . . .
, E π do 6:
Optimize policy objective J π ( θ ) and value regularizer C V ( θ ) with τ 7: end for 8: Add ( s t , a t , ˆ R t ) to B 9: end for 10: for iter
4.1 Algorithm
DCPG proceeds by alternating the policy and auxiliary phases similar to PPG, but it uses a sharedencoder in the same way as PPO and does not require an auxiliary value head as shown in Figure 1c.
During the policy phase, the policy network is trained on newly-collected trajectories to optimize thepolicy objective.
However, the value network is constrained to preserve the outputs of the originalvalue function by optimizing the following value regularizer C V given by
[Equation]
in the sense that updating the value network frequently with a small batch size can further increaseapproximation errors for unseen states as observed in Section 3.1.
Then, all states and value targets inthe trajectories are stored in a buffer.
During the auxiliary phase, the value network is trained with a larger number of data in the buffer tooptimize the value objective, while the policy network is constrained to preserve its outputs by thepolicy regularizer.
Note that since the value network has the same encoder as the policy network,optimizing the value objective also plays the role of representation learning for the policy network.
Thus, we do not require an additional auxiliary head for feature distillation, unlike PPG.
The overallprocedure of DCPG can be found in Algorithm 1.
[Figure]
DCPG, BigFish
[Figure]
DCPG, Climber
[Figure]
DCPG, Dodgeball
True Predicted 4.2 Value Analysis of Delayed Critic Update The delayed critic update can operate as an implicit penalization of the value function.
Since thepolicy improvement is not reflected in the value network immediately, it encourages smaller valueestimates than the true values.
To validate this claim, we train DCPG agents using 200 training levelson the Procgen games and measure the true and predicted values for the initial states of the traininglevels.
Figure 4 shows the value estimates are kept lower than the true values at the early stage oftraining but recover the true values as training progresses.
More results can be found in Appendix C.
We also evaluate the effectiveness of the delayed critic update in mitigating the overfitting problem bymeasuring the stiffness of the value function. The blue lines in Figure 2 show that DCPG has higherstiffness compared to PPG, which implies that our value network is more robust against overfitting asaddressed in Section 3.1.
We conclude that the delayed update of the value network in DCPG servesas a good regularizer for the value network.
More results can be found in Appendix A.
4.3 Learning Forward and Inverse Dynamics with Single Discriminator
In the context of multi-task learning, learning an auxiliary task can serve as a regularizer that preventsthe overfitting in the main task [32]. Motivated by this, we consider learning an auxiliary task usingthe shared encoder with the value network concurrently.
Suppose the buffer B used in the auxiliaryphase contains a transition ( s t , a t , s t +1 ) .
Since the transition functions are assumed to have the sameunderlying structure, it is natural to train a neural network that models the forward dynamics ofenvironments in the feature space of the value encoder.
Concretely, a dynamics head f θ takes theembedding of the current state s t and the current action a t as inputs and predicts the embedding ofthe next state s t .
However, it is well known that a forward dynamics model using a neural networktends to overfit on a small number of training data and make poor predictions on unseen states [8].
Toaddress this, we consider learning the forward dynamics implicitly by training a discriminator thatdistinguishes 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 coefficient for the inverse dynamics.
We name the resulting algorithm Dynamics-aware Delayed-Critic Policy Gradient (DDCPG).
Thenetwork structure and algorithm for DDCPG are presented in Figure 1c and Algorithm 1, respectively(the difference with DCPG is marked in cyan). [Figure]
Experiments
[Figure]
5.1 Experimental Settings
We evaluate the observational generalization performance of our methods on the Procgen benchmark.
We use the “easy” difficulty mode, which most prior works have focused on.
We train agents on 200training levels generated by seeds from 1 to 200 for 25M environment steps and test them on held-outtest levels, following the practice in Cobbe et al.
We measure the average return of 100 testepisodes and report its mean and standard deviation 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 [31, 22, 11, 30].
UCB-DrAC is a data augmentation algorithmfor training policy and value networks to be robust against various types of transformations.
PLRis a sampling algorithm for the procedural generation that selects the next training level based onits future learning potential.
PPG is the previous state-of-the-art method on the Procgen benchmark.
DAAC is motivated by PPG and distills advantage function into the policy network instead.
For eachmethod, we use the same set of hyperparameters across all games, following the standard protocol inthe ALE and Procgen benchmarks [28, 10].
For our methods, we use the same hyperparameter settingas PPG for a fair comparison.
To compare the performance of each method with a single score, wecalculate the PPO-normalized score averaged over all Procgen games, which is computed by dividingthe average return of each method by PPO, following the practice in Jiang et al. [22], Raileanu et al.[31].
For the implementation details and hyperparameters, please refer to Appendix D.
5.2 Observational Generalization Performance on Procgen Benchmark
Table 1 shows the average test returns of each method on all 16 Procgen games.
DCPG significantlyoutperforms all the baselines, yielding a 24%p improvement on average compared to the previousstate-of-the-art PPG.
Adding dynamics learning to DCPG can further improve the test performancewith an 18%p increase in the PPO-normalized score.
We also provide the average training returnson all Procgen games in Appendix E, showing that our methods achieve better sample efficiencyin the multi-environment setting. Furthermore, we evaluate the test performance of our methodsusing Min-Max normalized score and report interquartile mean (IQM) scores as well as mean andmedian scores in Appendix G, following the recent practice in Agarwal et al.
We observe that theperformance improvements of our methods are statistically significant for all the metrics we consider.
[Figure]
a
[Figure]
e r
[Figure]
t u r n
[Figure]
[Figure]
[Table]
Figure 5 shows the test performance curves of each method on 6 Procgen games throughout training.
Our methods not only achieve superior final test performance but also acquire generalization abilitywith a fewer number of interactions between environments.
For example, our methods can reach thefinal performance of PPG using only 20% of total environment steps in BigFish and Climber and10% in Dodgeball and Miner.
More experimental results can be found in Appendix E.
5.3 Ablation Study for Delayed Critic Update
The delayed critic update not only alleviates the overfitting problem but also provides better representations for generalization to the policy network.
To validate this, we train DCPG agents with anadditional value network V ϕ using a separate encoder (separate DCPG).
The separate value network 240 V ϕ is aggressively optimized in the same way as PPG and is only used for computing advantagesfor policy optimization.
The original value network V θ is only used for learning representations forthe policy network with delayed updates.
Figure 6 shows that separate DCPG achieves better testperformance than PPG, which implies that the delayed critic update itself produces more generalizablerepresentations.
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 train DCPG agents with forward dynamics learning (DCPG+F),inverse dynamics learning (DCPG+I), and forward and inverse dynamics learning using separatediscriminators (DCPG+FI) and compare them with DDCPG.
Table 2 shows that inverse dynamics provides performance gain only when jointly trained with forward dynamics using a single discriminator.
We provide the detailed experimental settings and results in Appendix F.
[Table]
Related Works
Observational generalization in RL One major approach for improving observational generalization in RL is to apply regularization methods originally developed for supervised learning,such as batch normalization, weight regularization, and information bottleneck to model-free RLalgorithms [14, 9, 20].
Several works adopt data augmentation techniques that are commonly used incomputer vision to address the generalization problem [24, 38].
For example, UCB-DrAC proposesmethods that automatically find an effective augmentation and introduces new regularization terms tolearn robust state representations under various types of augmentations [31].
Recently, it has beenshown that MuZero, the state-of-the-art model-based RL algorithm, achieves better generalizationperformance than model-free approaches, though it requires a larger network architecture [4].
Another line of work uses self-supervised learning to learn invariant representations that can generalizeto unseen environments.
Some methods attempt to learn state representations that capture long-termbehavior proximity between states using bisimulation or policy similarity metrics to discard irrelevantinformation for generalization [40, 1, 27].
The idea of learning dynamics as a self-supervised task,originally developed for sample efficiency, has also been used to improve observational generalization[18, 36].
DIM learns a state representation that can predict the representations of successive timestepsusing self-supervised learning [26].
It has also been tested that predicting future state representationsconditioned on current state and action improves the generalization performance of MuZero [4]. Our work aims to further enhance dynamics learning by jointly modeling the forward and inversedynamics objectives using a single discriminator.
Most closely related to our work is PPG, which proposes a better training regime to improve bothgeneralization and sample efficiency in the multi-environment setting [11].
It trains the policy andvalue function on separate networks with different levels of sample reuse to avoid interference andobtain a better-trained value function.
DAAC also uses decoupled policy and value networks to traina more accurate value function [30].
However, it distills the advantage function into the policy, whichis less prone to overfitting to environment-specific features than the value function.
We focus onregularizing the value function while obtaining better representation, which distinguishes our workfrom the prior methods.
Learning value function with regularization Discount regularization is one common regularization method to prevent overfitting in value function [7].
A lower discount factor can lead to betterperformance in approximated dynamic programming when the approximation error is large [29].
It isalso helpful for mitigating overfitting when dealing with limited data in POMDPs [17].
Activationregularization penalizes the norm of the value function and has a regularization effect similar to thatof discount regularization [3].
It has been shown that activation regularization is equivalent to thediscount regularization in batch TD learning and also effective in online RL [3].
Conclusion
We have investigated the difficulty of learning an accurate value function in the multi-environmentsetting and shown that suppressing a value function with regularization can be beneficial to improvingthe 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 predicted values at the earlystage of training with delayed updates compared to the policy network.
We find that the delayed valueupdate not only prevents the memorization of training data but also produces more generalizablerepresentations that can extend to unseen test environments.
Furthermore, we propose a simple selfsupervised task that implicitly learns both forward and inverse dynamics using a single discriminator,which can be easily combined with DCPG.
Our methods exhibit state-of-the-art performance inobservation generalization
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes] We clearly reflect the main claims in Section 3 and
Section 4.
(b) Did you describe the limitations of your work?
[N/A]
(c) Did you discuss any potential negative societal impacts of your work?
[N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] We have read the ethics review guidelines and checked our paper has no
conflict.
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 experimental results (either in the supplemental material or as a URL)?
[Yes] We provide thecode for reproducing the main results in the supplementary material.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
[Yes] We provide the implementation details and hyperparameters of allbaselines 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 computationalresource 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 identifiable information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[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]