C ROSS -S TATE S ELF -C ONSTRAINT FOR FEATURE GENERALIZATION IN DEEP REINFORCEMENT LEARNING
Anonymous authors
Paper under double-blind review
A BSTRACT
Representation learning on visualized input is an important yet challenging task for deep reinforcement learning
RL)
The feature space learned from visualized input not only dominates the agent’s generalization ability in new environments but also affect the data efﬁciency during training.
To help the RL agent learn more general and discriminative representation among various states, we present cross-state self-constraint(CSSC), a novel technique that regularizes the representation feature space by comparing representation similarity across different pairs of state.
Based on the implicit feedback between state and action from the agent’s experience, this constraint helps reinforce the general feature recognition during the learning process and thus enhance the generalization to unseen environment.
We test our proposed method on the OpenAI ProcGen benchmark and see signiﬁcant improvement on generalization performance across most of Procgen games.
I NTRODUCTION
Deep Reinforcement learning has achieved tremendous success on mastering video games(Mnih et al., 2015) and the game of GO(Silver et al., 2017).
While training agent by using deep reinforcement learning algorithms, we usually assume that the agent could extract appropriate and effective features from different states and take actions accordingly.
However, as more and more research works(Zhang et al.
(2018), Song et al.
, Dabney et al.
(2020)) have pointed out, even welltrained RL agents that learns from visualized input tend to memorizing spurious pattern rather than understanding the essential generic features of a given state.
For example, an agent might pay more attention to the change of irrelevant background rather than noticing the obstacles or enemies(Song et al., 2019).
To improve generalization in the new environment, various kinds of regularization method like dropout(Farebrother et al., 2018) and data augmentation(Laskin et al., 2020) has been proposed and tested in combination with reinforcement learning.
Conventional methods like dropout and batch-norm has been proven to be effective in supervised-learning, and for self-supervised learning like RL we see multiple related applications across various environments.
Data augmentation like random crop(Laskin et al., 2020) or random convolution(Lee et al., 2019) have also been proposed recently and provide considerable generalization enhancement to the unseen levels of various tested environment(Tassa et al.
, Cobbe et al. , Cobbe et al.
The agent is acting on multiple augmented views of the same input and learn from these prior injected data.
However, modifying state information(injecting prior to the data) may be risky or even detrimental for representation learning because vital features may be altered or lost (ex: ﬂipping state image might change the corresponding behavioral meaning, cropping the input image might lose critical features like the enemy position in the game).
To avoid losing informative features of the visualized input, we choose a different approach.
As a human learner, we rarely depend on multiple augmented views of the same input to discriminate important or ﬁctitious features.
Instead, human learners try to recognize general patterns across multiple states and act accordingly.
In other words, if the same action(or behavior) has been conducted by a well-trained agent in two different states, we would infer that the agent has conceived similar feature patterns in these states.
For example, if one car stops for ninety seconds at two different intersections, we would guess that the car might be stopped by the red light at both places(Figure 1b).
From here we get the intuition about the relation between action(or behavior) and representation feature space learned by the agent: for a RL agent acting rationally across sequence of states, its behavior would be similar if it perceives similar critical patterns while acting differently if it perceives different patterns.
Based on this intuition, we designed a novel constraint that performs regularization directly in the representation feature space of the learning agent.
Our hypothesis for this constraint is simple: states in which the rational agent behaves identically should share more representation resemblance than states the agent behave differently.
We test this novel constraint in combination with Rainbow(Hessel et al., 2017), the state-of-theart Q-learning method that combines numerous improvements with the original DQN model(Mnih et al., 2015).
Inspired by the pair-wise structure used in BPR-opt(Rendle et al., 2012), we design the self-constraint based on the agent’s behavior and utilize implicit feedback between positive and negative state pairs in the replay buffer.
One thing worth noticing is that no change is needed for the underlying RL algorithm or model to adapt the proposed method, and it can be easily applied to other models with minimal effort(Figure 1a).
[Figure]
To measure the improvement on generalization to unseen environment levels, We test CSSC across 16 games on the OpenAI ProcGen benchmark(Cobbe et al., 2020).
We see signiﬁcant improvement on generalization capability for most ProcGen games in comparison with base Rainbow model.
We highlight the main contributions of CSSC as follows:
• Directly optimize feature generalization across various input states • Requires no additional modiﬁcation on the input data and the base model • Enhance generalization and data efﬁciency on existing RL methods
R ELATED W ORK 2. A UXILIARY T ASK FOR R EPRESENTATION L EARNING IN RL
Representation Learning has been a vital part for deep RL algorithms.
Even though the main goal of RL is to ﬁnd the optimal value function, Mccallum & Ballard (1996) and Li et al.
(2006) show that a representation specialized to this function may not be suitable for the sequence of value functions leading to it.
On the other hand, Dabney et al.
(2020) argue that the path toward optimal value function might be hindered by overﬁtting the representation to any intermediate value function during training.
One popular and effective method to address this challenge is by using auxiliary tasks (Jaderberg et al.
(2016), Bellemare et al.
Dabney et al.
(2020) also propose new auxiliary tasks by considering the value-improvement path holistically.
2. R EGULARIZATION ON RL 2.2. A DDING STOCHASTICITY TO RL
Conventional practices like stochastic policy(Hausknecht & Stone, 2015), random starts(Mnih et al., 2015), sticky actions(Hausknecht & Stone, 2015) and frame skipping are wildly used for popular tasks like Atari(Machado et al., 2017).
By adding stochasticity to the environments during training and testing, we would prevent simple algorithms like trajectory tree (Kearns et al., 1999) and brute (Machado et al., 2017) from optimizing over open-loop sequences of actions without even considering the input states.
However, as Zhang et al.
(2018) has pointed out, injecting stochasticity to the maze-like environments does not necessarily prevent the deep RL agents from overﬁtting.
2.2. C ONVENTIONAL R EGULATIZATION ON RL
For deep supervised learning, regularizations has long been a major issue as generalization from training to test set is the major concern.
Common regularization practices like L2 regularization, dropout, batch-normalization from supervised learning has also been tested in the reinforcement learning scenario.
Cobbe et al.
(2018) and Farebrother et al.
(2018) show that dropout and L regularization improve generalization to test environments in Coinrun(Cobbe et al., 2018) and Atari 2600(Bellemare et al., 2013) respectively.
2.2. D ATA A UGMENTATION ON RL
For supervised learning related to images or computer vision, data augmentation is a popular way to enhance performance by injecting useful priors on the training data.
In recent years, we see numerous data augmentation methods for RL that improve generalization or data efﬁciency.
Lee et al.
(2019) introduce a randomized convolution network that randomly perturbs input features.
Srinivas et al.
(2020) utilize contrastive learning and show considerable performance gains in terms of sample-efﬁciency.
Kostrikov et al.
(2020) uses augmented data and weighted Q-functions to achieve state-of-the-art data-efﬁciency on the DMControl(Tassa et al., 2018).
Laskin et al.
(2020) investigate ten different data augmentation on RL and point out that random crop is the most effective on DMControl(Tassa et al., 2018)
B ACKGROUND
CSSC is a general framework of cross-state representation regularization for RL.
In principle, one can apply CSSC to other variation of DQN or policy-based models for discrete action-space environments.
In this work we pick Rainbow (Hessel et al., 2017) and PPO (Schulman et al., 2017) as our base model to show that CSSC is compatible with the original Nature DQN(Mnih et al., 2015), its multiple improvements and policy-based algorithms.
In the following subsections we will review Rainbow DQN and introduce the concept of implicit feedback in Bayesian Personalized Ranking (BPR) proposed by Rendle et al.
3. R AINBOW
In combination with a convolutional neural network as visualized input encoder, Deep Q network(Mnih et al., 2015) demonstrates that it is possible to use neural network as a function approximator that maps raw pixels to Q values.
Since then, multiple improvements such as Double
Q Learning(van Hasselt et al., 2015), Dueling Network Architecture(Wang et al., 2015), Prioritized
Experience Replay(Schaul et al., 2015) and Noisy Networks(Fortunato et al., 2017) have been proposed.
In addition, Bellemare et al.
(2017) proposed the method of predicting a distribution over possible value support through the C51 algorithm.
By combining all the above techniques into a single off-policy algorithm, Rainbow DQN showcases the state-of-the-art sample-efﬁciency on Atari
Benchmarks.
The resulting loss function for Rainbow DQN is as follows:
[Equation]
where Φ z is the projection onto ﬁxed support z .
d ( n ) t and d t are the n-step target return distribution and the model-predicted return distribution respectively.
3. BPR- OPT
As opposed to explicit feedback like user ratings, implicit feedback in recommendation systems focuses on interactions like click or view between users and items.
In implicit feedback system only positive observations are available, and non-observed user-item pairs -e.g.
a user has not view an item yet - are considered as negative observations.
Instead of rating prediction, BPR-opt(Rendle et al., 2012) is designed for direct ranking optimization based on the implicit feedback between users and items.
BPR extends the user preference from observed interaction pairs to non-observed data by ranking the preference of positive observation and negatived observation across the training data.
The formulation of the maximum posterior estimator for the personalized ranking optimization (BPR-opt) can be written as follows:
[Equation]
where Θ represents the parameter vector of the base model.
D s represent the batch samples from training data.
p (Θ | > u ) is the posterior probability conditioned on latent preference structure > u for user u , λ
Θ is the model speciﬁc regularization parameter, σ is the sigmoid function, and ˆ x uij (Θ) is the real-valued function model Θ which captures the preference relationship between user u , item i and item j .
Rendle et al.
(2012) use ˆ x uij :
=
ˆ x ui − ˆ x uj to indicate that the user u prefers item i over item j .
M ETHOD
The main idea of CSSC is to extend the similarity ranking across various state representation based on our hypothesis: representations motivating identical behavior should share more similarity than those motivating different behavior.
We name this method as ”self-constraint” to emphasize that the state pairing is decided by the agent’s behavior rather than pre-deﬁned ﬁxed labels.
In the following subsection below, we introduce the following concepts to provide further explanation and insight: i
) deﬁnition of the behavior of an agent at a given state (ii) implicit feedback between state and
behavior
(iii) cross-state self-constraint as an auxiliary loss for representation
regularization
4. B EHAVIOR D EFINITION
We describe a typical Markov Decision Process(MDP) as ( X , A , R , P , γ ) with state space as X , action space as A , reward function as R , state transition function as P and discount factor as γ .
The agent would take an action a i at state x
i and then being transited to x i +1 by the environment with a given step reward r i .
Here we deﬁne the behavior set B ni =
{ ( a i , a i +1 , ..., a i + n − 1 ) | a ∈ A } as a set of action series of length n taken by the agent since state x i .
For b ni ∈ B ni and b nj ∈ B nj , b n i =
b n j only if a i + p = a j + p for 0 ≤ p ≤ n − 1 .
With this deﬁnition in mind, we can infer that if the agent conduct the behavior b 2 i =
( left, fire ) at state x
i , it would move left in state x
i and ﬁre in the next state x i +1 .
In the following paragraph we coin names like unigram, bigram and trigram for behaviors of length one, two and three respectively.
4. I MPLICIT F EEDBACK IN R EINFORCEMENT L EARNING
In the traditional deep Q-learning scenario, each state representation is highly correlated to neighboring states by the Bellman equation during the training process.
However, observational overﬁtting could still happen as long as the representation of neighboring states can ﬁt to speciﬁc sub-optimal version of Q-function indeﬁnitely as mentioned in Dabney et al.
Besides the base RL algorithm, we can view the interaction between state and behavior as an implicit feedback system by regarding states sharing same action as positive observation.
The general feature among these positive observation would be reinforced if we compare their similarity with negative observation -e.g.
states sharing different action.
Through this “ranking” procedure the pair-wise relationship would be extended across non-neighboring states and therefore preventing observational overﬁtting to speciﬁc Q-function.
4. I MPLEMENTATION OF C ROSS -S TATE S ELF -C ONSTRAINT IN COMBINATION WITH R AINBOW AND PPO
For each sample of state triples ( x p , x q , x r ) ∈ X with ( b np , b nq , b nr ) ∈ B n and b np
= b nq (cid:54) =
b nr , we decompose the estimator ˆ x pqr and deﬁne it as:
[Equation]
where e θ is the encoder function with weight θ that maps the pixel input to 1-D array feature vector with shape like [ element 0, element 1, ..., element n ] as the state representation.
We directly perform inner-product on e θ ( x p ) and e θ ( x q ) to calculate the representation similarity ˆ x pq between x p and x q .
Please note that the encoder function is the same part of the original base model used for feature extraction on visualized input.
For more details about neural network structures used in this paper, please refer to Figure 7 and Figure 8. To sample these state triplets without modifying the base DQN or PPO algorithm, we collect these state triples in the same batch of transitions used to calculate Bellman loss or policy loss.
For policy-based algorithms and DQN without prioritized replay (Schaul et al., 2015), we pair up the state triples for each transition in the sample and calculate the CSSC loss as follows:
[Equation]
where D s represents the sample batch of size N D s from the replay buffer and σ is the sigmoid function.
Then, we train the base algorithm in combination with the auxiliary CSSC loss as follows:
[Equation]
where β cssc is the hyperparameter to control the contribution of CSSC during training.
For pseudocode of PPO with unigram CSSC, please refer to Algorithm 2 in the appendix.
In the case of Rainbow with CSSC, we design the loss function in combination with the importance sampling (IP) weight as follows:
[Equation]
where (cid:76) and (cid:12) are element-wise add and multiplication respectively.
For pseudocode of Rainbow with unigram CSSC, please refer to Algorithm 1 in the appendix.
We ﬁnd that setting β cssc to 0.01 works quite well in most cases, so we stick with this setting for all the experiment conducted in the next section.
E XPERIMENT
Our primary goal for CSSC is to enhance the generalization capability of RL algorithm to unseen levels that share similar mechanism.
Fortunately, OpenAI ProcGen(Cobbe et al., 2020) presents a collection of game-like environments where the training and testing environments differs in visual appearance and structure.
Therefore, we evaluate CSSC in three different ways:(i) generalization improvement on 12 easy-mode games and 8 hard-mode games with Rainbow on OpenAI ProcGen (ii) generalization improvement on 16 easy-mode PPO improvement sentation
[Table] [Figure] [Table]
5. E NVIRONMENT AND M ODEL S ETTING
For the following experiments, we use the IMPALA CNN architecture recommended by Cobbe et al.
(2020) for the Rainbow and PPO model on ProcGen benchmark.
For the Gym Atari benchmark, we stay with the Nature CNN as mentioned in Hessel et al.
(2017) for the Rainbow model.
All the experiments in this paper are conducted in the single-agent setting.
For all hyperparameter used on ProcGen and Atari benchmark, please refer to Table 3 in the Appendix.
We use “unigram CSSC” or “bigram CSSC” to indicate the behavior length used to pair up the state triples for CSSC.
5. G ENERALIZATION ON P ROCGEN IN E ASY MODE WITH R AINBOW
As Cobbe et al.
(2020) have suggested, we conduct the training of 25 million timesteps ontraining levels and evaluate the generalization improvement on the full distribution of testing levels across 12 ProcGen games.
We test both the unigram and bigram CSSC in comparison with vanilla
Rainbow and see signiﬁcant improvement on testing performance.
To evaluate the mean normalized score for bigram and unigram CSSC across 12 games, we follow the normalization method used in
Cobbe et al.
and show the result in Figure 2a and Figure 12a respectively.
We also display the learning curve of bigram-CSSC in Figure 3 and unigram-CSSC in Figure 13.
In the following we summarize the main ﬁndings below:
• In supervised learning, posing regularization or constraint to the model would usually improve its testing performance on unseen data at the expense of hurting training performance.
However, we see that the proposed self-constraint improves both training and testing performance across most of the 12 game tested.
• For games like Starpilot, Chaser and Bigﬁsh, we see that the testing performance is even better than the training performance of vanilla Rainbow.
• For the mean normalized score in Table 1, Bigram-CSSC and Unigram-CSSC bring 85% and 64% improvement respectively on the testing performance in comparison with base Rainbow model.
5. G ENERALIZATION ON P ROCGEN IN H ARD MODE WITH R AINBOW
To further explore the generalization capability in more complicated and challenging tasks we examine the generalization improvement brought by CSSC on 8 ProcGen games.
We use 500 training levels for 30 million timesteps of training across 3 seeds and evaluate generalization on the full distribution of testing levels.
We test the bigram CSSC in comparison with the normal version of Rainbow across three different seeds and see signiﬁcant improvement on testing performance.
The learning curve of bigram CSSC is summarized in Figure 4.
To evaluate the mean normalized score for bigram CSSC across 8 games in hard mode, we also follow the normalization method used in Cobbe et al.
and show the result in Figure 2b.
In the following we summarize the main ﬁndings below:
• As shown in Figure4, CSSC substantially improves both training and testing performance in Bigﬁsh, Dodgeball, Starpilot, Fruitbot and Bossﬁght.
• In particular, we see nearly 4x performance jump in Bigﬁsh at 30M timestep in comparison with the base Rainbow model.
The mean episodic return at 30M timestep of bigram CSSC is even higher than that of PPO at 200M timestep as shown in Figure 4 of Cobbe et al.
• In the case of Heist (a puzzle-solving task in maze-like layout), coinrun(a 2-D scroll platform game) and Plunder(a challenging shooting game), the gain from CSSC is not as obvious because all these games require careful manipulation and planning.
Add constraint that modiﬁes state representation directly would be risky or even detrimental for the learning process of the base RL algorithm.
5. G ENERALIZATION ON P ROCGEN IN E ASY MODE WITH PPO
We test the CSSC with PPO on all the 16 games of the Procgen benchmark in Easy mode. We conduct the training of 25 million timesteps across 3 seeds on 200 training levels and evaluate the generalization improvement on the full distribution of testing levels across all 16 ProcGen games. From the normalized learning curve shown in Figure 5a we can tell that CSSC helps reduce the gap between training and testing performance. For detailed ﬁnal scores and learning curves, please refer to Table 7 and Figure 14, 15 in the Appendix.
We take Gym Atari benchmark (Brockman et al., 2016) as our second set of environments to measure the effectiveness of CSSC. Even though this classic benchmark does not explicitly split the training levels out of the training levels in all environments, we still see signiﬁcant improvement on 23 out of 38 games tested as shown in Figure 5b. For training hyperparameters and scoreboard on Gym
Atari benchmark, please refer to Table 3 and 8 in the Appendix.
5. V ISUALIZATION OF R EPRESENTATION E MBEDDING
To give a more tangible explanation about the effectiveness of CSSC, we plot the representation distribution in the embedding feature space across 4096 states from the replay buffer.
We ﬁrst perform dimension reduction on these state embeddings using principal decomposition analysis and label each embedding point with the index of corresponding action.
Here we show the embedding space of both Rainbow and bigram-CSSC models trained to play the Bigﬁsh game.
In ﬁgure 6 we can see that states of the same action are more clustered in bigram-CSSC than that of the vanilla Rainbow model.
We believe this is the possible reason behind the signiﬁcant improvement on generalization as shown in Figure 4.
For representation visualization in other ProcGen games, please refer to A.4.
C ONCLUSION
In this paper we propose a novel regularization on state representation learning based on the connection between agent behavior and the visual input.
Our hypothesis is derived from the observation that the behavior of a rational agent would have certain relationship with general cross-state features or patterns.
We see signiﬁcant improvement on generalization brought by the proposed cross-state self-constraint(CSSC) on most of the games in OpenAI ProcGen benchmark.
The connection between behavior and conceived visual input can be considered as some kind of “motivation” that acts as a decisive factor behind the learned policy.
It worth further study to better understand the derivation and transformation of agent’s ”motivation” during the Markov Decision Process, and we believe the concept proposed in this work would facilitate more research in this direction. [Figure] [Figure] A.
[Table]
H YPERPARAMTERS AND M ODEL SETTING FOR O PEN AI P ROC G EN AND
YM B ENCHMARK
[Table]
[Table]
A.
[Table]
S COREBOARDS ON O PEN AI P ROC G EN
[Table]
[Table]
[Table]
2.96 ± 0.
3.17 ± 0.
2.94 ± 0.
Starpilot 27.19 ± 2.
26.19 ± 1.
28.41 ± 2.
19.82 ± 0.
22.12 ± 1.
22.28 ± 1.
[Figure]
[Figure]
[Figure]
V ISUALIZATION OF R EPRESENTATION E
ON
A
G ENERALIZATION ON
ROC G EN I N E ASY
ODE FOR
AINBOW UNIGRAM -CSSC A
G ENERALIZATION ON
ROC G EN IN E ASY
ODE FOR PPO WITH CSSC A. P SEUDOCODE FOR CSSC EXPERIMENT A. E VALUATION OF BIGRAM CSSC WITH R AINBOW ON G YM A TARI B ENCHMARK
transition j t =
( s t , a t , r n t , d t , s t n in H with p t = max t p i 7:
P p
Σ p
P
j state e
Θ Compute L =
= d distribution projection function z 12:
Update priority p for each in the batch 13: p , s r ) , p (cid:54) =
r with j p , j r ∈ J and behavior b 1 p (cid:54) =
b 1 r for each transition j p in J 15:
Compute (cid:92) L CSSC =
[ ˆ s 0 p 0 q 0 r 0 , ..., ˆ s kp k q k r k
{+| ˆ s jp j q j r j := − 1 . 
+}
× σ ( e θ ( s p j ) · e θ ( s q j ) − e θ ( s p j ) · e θ ( s r j ))] 16:
Compute (cid:92) L total =
( (cid:92) L Rainbow (cid:76) ( β CSSC · (cid:92) L CSSC ))
(cid:12) (cid:99) W where (cid:76) and (cid:12) are elementwise add and multiplication respectively 17:
Fit Θ and θ with Adam optimizer to minimize (cid:92) L total .mean() 18: end if 19: end for
[Figure] do 4:
Collect set trajectories D
{T i } policy environment 5: Compute advantages estimates ˆ A t using current value function V φ 6: Store transition j t = ( s t , a t , r t , d t , ˆ A t , log ( π ( θ k ))) H 7: if do 9:
Predict state feature vector e prob transitions
Compute ratio = exp(log prob
Compute policy loss = batch.advantages clamp
Compute L policy = policy loss.mean() 13:
Sample positive pair ( s p , s q ) , p (cid:54) =
q with j p , j q ∈ J and behavior b 1 p =
b 1 q for each transition j p in batch transitions J 14:
Sample negative pair ( s p , s r ) , p (cid:54) =
r with j p , j r ∈ J and behavior b 1 p (cid:54) =
b 1 r for each transition j p in batch transitions J 15:
{+Compute cssc loss = [ ˆ s 0 p 0 q 0 r 0 , ..., ˆ s kp k q k r k | ˆ s jp j q j r j := − 1 . 
 × σ ( e θ ( s p j ) · e θ ( s q j ) − e θ ( s p j ) · e θ ( s r j ))] 16:+}
Compute L CSSC =
.mean
Fit Θ , θ and φ with Adam optimizer to minimize L total =
L Rainbow + β cssc ·L CSSC 18: end for 19: end if 20: end for
[Table]