Behavior Predictive Representations for Generalization in Reinforcement Learning
Abstract
Deep reinforcement learning (RL) agents trained on a few environments, often struggle to generalize on unseen environments, even when such environments are semantically equivalent to training environments.
Such agents learn representations that overﬁt the characteristics of the training environments.
We posit that generalization can be improved by assigning similar representations to scenarios with similar sequences of long-term optimal behavior.
To do so, we propose behavior predictive representations (BPR) that capture long-term optimal behavior.
BPR trains an agent to predict latent state representations multiple steps into the future such that these representations can predict the optimal behavior at the future steps.
We demonstrate that BPR provides large gains on a jumping task from pixels, a problem designed to test generalization.
Introduction
Deep reinforcement learning (RL) agents, even when trained on diverse environments with similar high level goals but different dynamics and visual appearances, often struggle to generalize on unseen environments, even when such environments are semantically equivalent to training environments [Farebrother et al., 2018, Cobbe et al., 2020, Agarwal et al., 2021a, Packer et al., 2018].
Such agents learn state representations from high-dimensional observations that typically overﬁt to the peculiarities of training environments [Song et al., 2019, Raileanu and Fergus, 2021] rather than capturing generalizable skills which can be transferred to unseen environments.
Such overﬁtting hinders the real-world applicability of RL, making generalization in RL an important challenge.
To improve generalization using better representations, we revisit predictive representations [Littman et al., 2001, Rafols et al., 2005] that describe the environment in terms of predictions about future observations, such as representations that encode the underlying environments dynamics.
While learning such temporally predictive representations has been shown to improve sample efﬁciency [Oord et al., 2018, Schwarzer et al., 2021] within a training environment, it is unclear whether such representations would improve performance in unseen environments.
More recently, Agarwal et al.
[2021a] enhance generalization by learning similar state representations for observations with similar long-term optimal behavior.
Inspired by their ﬁndings, we posit that predictive representations that capture long-term optimal behavior might be better suited for generalization.
We expect such behavior predictive representations to generalize as two observations, possibly across different environments, are assigned similar representations if they exhibit similar sequences of optimal behaviors, irrespective of their differences in obtained rewards, visual appearances, or even the underlying dynamics.
For learning behavior predictive representations (BPR), we train the agent to predict latent state representations multiple steps into the future such that these representations can predict the optimal behavior at the future steps (Figure 1).
BPR can be viewed as a representation learning approach where the agent predicts the optimal behavior at future states resulting from following a sequence of actions from a given state.
We show the efﬁcacy of BPR on the jumping task on pixels and show it improves generalization upon existing methods including PSEs [Agarwal et al., 2021a] and
SPR [Schwarzer et al., 2021].
We also provide ablations demonstrating the effect of predicting suboptimal policies as well as the horizon for predicting future behavior.
Preliminaries
We describe an environment as a Markov decision process (MDP) that corresponds to a tuple
M = ( S, A, P, R, γ ) where S is the state space,
A is the action space, P : S ×
A × S → [0 , 1] is the state transition function, R : S ×
A → R is the reward function and γ ∈ [0 , 1] is the discount factor.
A policy π ( ·| s ) maps a state s ∈ S to a distribution over the action space A .
A trajectory is deﬁned as the sequence of states, actions and corresponding rewards i.e. s 0 , a 0 , r 1 , s 1 , · · · .
The goal of a reinforcement learning agent is to maximize the cumulative expected return E p ( τ ) [ (cid:80) t γ t r ( s t , a t )] where p ( τ ) =
p ( s 0 ) (cid:81) t p ( s t +1 | s t , a t ) π ( a t
| s t ) .
Behavior Predictive Representations
In this work, we aim to learn a policy that can generalize across related environments.
Speciﬁcally, we train an agent using a ﬁnite number of environments (or tasks) sampled from a distribution of environments.
The performance of this agent is evaluated using unseen environments sampled from the same distribution.
For example, consider the generalization problem in a jumping task from pixels [Tachet des Combes et al., 2018], where an agent needs to jump over an obstacle (Figure 2).
Standard deep RL agents trained on a small number of training tasks with different obstacle positions struggle to generalize to unseen obstacle positions [Agarwal et al., 2021a].
Inspired from the recent success of representation learning to improve generalization [Agarwal et al., 2021a, Raileanu and Fergus, 2021, Zhang et al., 2020], we also focus on learning better representations to improve generalization.
We posit that learning better representations requires understanding which states are similar in terms of their long-term optimal behavior.
To do so, we aim to learn latent state representations that not only capture the behaviour at the current state but will also be able to predict the behaviour at future states, which we call behavior predictive representations (BPR).
Since BPR simply uses an auxiliary objective L BPR , it can be easily combined with any RL or imitation learning setup, as shown in Figure 1.
To describe the auxiliary loss L BPR for predicting the long-term optimal behavior, we deﬁne some notation ﬁrst.
Let s t be the state at the time step t and z t be the corresponding latent representation learned by the policy network f .
The policy π predicts the action distribution given the latent representation z t .
We use an encoder network f to generate these latent representations from states as, z t = f
( s t ) .
A transition function h : S ∗ ×
A → S ∗ learns the state dynamics and predicts the representations at the next step, ˆ z t +
= h ( s t , a t ) .
We predict the representations for K future steps by iteratively applying the transition function.
While such latent state dynamics are typically learned by minimizing the mean squared loss between ˆ z t + k and z t + k , we instead use these predicted representations to predict the optimal action distributions in the future steps t + 1 to t + k .
To do so, the agent minimizes the cross entropy between the predicted action distribution and optimal action distributions at these steps.
Speciﬁcally, given access to the optimal policy π ∗ on training environments, the auxiliary loss L BPR is given by:
[Equation]
where L CE ( π 1 ( ·| s ) , π 2 ( ·| s ))
= − (cid:80)
a ∈
A π 1 ( a | s ) log π 2 ( a | s ) .
In the general RL setting, where we do not have access to the optimal policy, we propose to use the learned policy on training environments for specifying the future behavior in the objective L BPR .
Speciﬁcally, the target distribution for the auxiliary cross entropy loss, L BPR , comes from the same policy network that is being trained ( f in Figure 1).
To provide some stability from the continuous changes in the learned policy, we use a separate target policy network that is periodically updated with the learned policy network parameters, analogous to deep Q-learning [Mnih et al., 2013] and self-supervised learning [Grill et al., 2020].
So, the target representations ˆ z t + k are derived form the learned transition function h while the action distribution π learned ( z t + k ) comes from the target policy network.
For this setting, the auxiliary loss ˆ L BPR is given by,
[Equation]
The ﬁnal training objective is the combination of both of these loss functions.
Let L RL be the traditional model-free RL (or imitaion learning) objective.
Then, the combined loss function for learning behavior predictive representations is L =
L RL + λ BPR L BPR , where λ BPR is the weighting coefﬁcient for the auxiliary loss.
Experiments
We ﬁrst thoroughly investigate behavior predictive representations (BPR) on the jumping task [Tachet des Combes et al., 2018, Agarwal et al., 2021a] that captures whether agents can learn the correct invariances for generalization directly from image inputs.
4.1 Jumping Task From Pixels
Task Description .
The task consists of an agent trying to jump over an obstacle using two actions: right and jump .
Different tasks consist in shifting the ﬂoor height and/or the obstacle position (Fig- [Figure] ure 2).
To generalize, the agent needs to be invariant to the ﬂoor height while jump based on the obstacle position
Problem Setup .
Following Agarwal et al.
[2021a], we use three different conﬁgurations (Figure 3), each consisting of 18 seen (training) and 268 unseen (test) tasks, to test generalization in regimes without and with data augmentation using RandConv [Lee et al., 2020].
As discussed by Agarwal et al.
[2021a], the different grids conﬁgurations capture different types of generalization: the “wide” grid tests generalization via “interpolation”, the “narrow” grid tests out-of-distribution generalization via “extrapolation”, and the random grid instances evaluate generalization similar to supervised learning where train and test samples are drawn i.i.d.
from the same distribution.
Refer to Agarwal et al.
[2021a] for more more experimental details.
Baselines .
We compare the efﬁcacy of our method with a number of techniques that have been used to achieve generalization including regularization such as (cid:96) 2 -regularization and dropout [Farebrother et al., 2018] and data augmentation [Lee et al., 2020].
Policy Similarity Embeddings (PSEs)
[Agarwal et al., 2021a] are the state-of-the-art generalization method on the jumping task.
PSEs form an important baseline for BPR as PSEs also use the future behaviour as a similarity metric between states.
Speciﬁcally, PSEs learn contrastive metric embeddings using a policy similarity metric d (Equation 3) that uses policy to measure the long term behavior similarity between among states.
[Equation]
Self-predictive representations (SPR)
[Schwarzer et al., 2021] is another relevant baseline which has been shown to improve sample-efﬁciency on training environments on the Atari 100k benchmark [Kaiser et al., 2019, Agarwal et al., 2021b].
SPR’s objective is that the agent learns to predict its own latent representations at future steps.
Similar to BPR, it uses a transition function to iteratively generate these latent representations for the future steps.
However, while BPR optimizes the latent representations to predict future behavior, SPR tries to maximize the similarity between the predicted latent representations ˆ z t +1 :
ˆ z t + K with the true future state representations z t +1 :
z t + K .
To do so,
SPR uses a self-supervised learning objective [Grill et al., 2020] as the auxiliary loss,
[Equation]
where g o , g m and q are online projection network, target projection network and prediction networks respectively.
SPR linearly combines the auxiliary objective, L SPR , with the RL objective.
Results .
Table 1 summarizes the performance of BPR and all the baselines with and without data augmentation.
Without data augmentation, with only 18 training environments, BPR generalizes quite well in all the three grid conﬁgurations, signiﬁcantly outperforming regularization and PSEs by a large margin.
These results exhibit that BPR is effective even without data augmentation.
Data augmentation complements all the methods and boosts generalization performance.
Comparing
RandConv + BPR to RandConv, we see that BPR is much more effective on top of RandConv.
Moreover, when used in conjunction with data augmentation, BPR performs comparably to the current state-of-the-art method PSEs.
Compared to BPR, SPR degrades the generalization performance signiﬁcantly and even performs poorly than simply using RandConv.
We hypothesize [Table] that the self-supervised learning objective in SPR might be exacerbating the overﬁtting in learned representations by trying to predict the spurious features captured by the learned representations on training environments.
4.2 Effect of Policy Suboptimality on BPR
On the jumping task, we use the optimal policy on training environments to learn BPR.
To understand the dependence of BPR on the optimal policy, we utilize -suboptimal policies to the auxiliary loss during training.
Speciﬁcally, the prediction at future steps predicts the optimal action with probability 1 − (cid:15) and suboptimal action with probability (cid:15) .
We plot the performance of BPR on the “wide" conﬁguration for the degree of suboptimality speciﬁed by (cid:15) , starting with the optimal policy ( (cid:15) = 0 ) to a uniform random policy ( (cid:15) = 0 . 5 ).
It can be seen from Figure 4 that as (cid:15) increases, the performance decreases which is expected.
For small values of (cid:15) , i.e., < = 0 . 2 , the performance decreases slightly but for larger values, the performance decreases sharply.
This means that BPR is tolerant to certain levels of suboptimality.
4.3 Effect of look ahead on BPR
The lookahead of the agent is the number of future steps K for which the optimal action is to be predicted by latent representations.
Greater the value of K , latent representations are required to predict actions on further in the future and possibly improve their generalizability.
But it will be difﬁcult for the latent representation to predict actions for steps that are far from the current step.
Thus the performance drops for larger values of K .
Figure 5 shows the plot of the performance of BPR on the "wide" conﬁguration with increase the value of K .
The performance of BPR is good for even low values of K and it remains similar for
K = 1 to 7 .
But it increases slightly from K = 1 to 3 and then decreases thereafter.
As a result, we use K = 3 for all the results that we have reported so far.
Conclusion and Future Work
In this paper we introduced Behavior Predictive Representations for improving generalization in reinforcement Learning.
We show that predicting optimal actions at future steps is more beneﬁcial than dynamics prediction or predicting future latent state representations.
As seen from the results,
BPR also performs well without data augmentations.
We plan to extend our work to more complex environments designed for testing generalization such as the Procgen benchmark [Cobbe et al., 2019] and Distracting DM Control Suite [Stone et al., 2021].
In these environments, we do not have access to the optimal policies and thus we plan to build BPR on top of RL objectives and show the effectiveness of BPR in such settings