L OCALLY C ONSTRAINED R EPRESENTATIONS IN R EINFORCEMENT L EARNING
A BSTRACT
The success of Reinforcement Learning (RL) heavily relies on the ability to learn robust representations from the observations of the environment.
In most cases, the representations learned purely by the reinforcement learning loss can differ vastly across states depending on how the value functions change.
However, the representations learned need not be very specific to the task at hand.
Relying only on the RL objective may yield representations that vary greatly across successive time steps.
In addition, since the RL loss has a changing target, the representations learned would depend on how good the current values/policies are.
Thus, disentangling the representations from the main task would allow them to focus more on capturing transition dynamics which can improve generalization.
To this end, we propose locally constrained representations, where an auxiliary loss forces the state representations to be predictable by the representations of the neighbouring states.
This encourages the representations to be driven not only by the value/policy learning but also self-supervised learning, which constrains the representations from changing in an unexplained manner.
We evaluate the proposed method on several known benchmarks and observe strong performance.
Especially in continuous control tasks, our experiments show a significant advantage over a strong baseline.
I NTRODUCTION
Representation Learning is a crucial problem in machine learning, particularly in understanding how complex inputs can be represented in a compact, well-defined manner, while retaining useful information of the input.
This has been studied extensively for vision and natural language processing tasks but not as much for Reinforcement Learning (RL).
RL has shown great success in games like
Atari Mnih et al.
, Go Silver et al.
(2016; 2018) and Chess Silver et al.
However, its application in the real-world setting is limited.
One main reason for this problem is the lack of an appropriate input feature space.
In games, the input space is rather well-defined, e.g. it can be the state of the game for board games or it can be the image of the game screen in computer games
Mnih et al.
However, it has been shown that learning policies directly from high-dimensional inputs can be sample
inefficient Lake et al.
Additionally, in real-world scenarios, we often have inputs from different sensors that the agent needs to fuse into a compact representation, which can be challenging.
Therefore, representation learning is a key to enabling reinforcement learning agents to solve real-world problems.
Most of the work on representation learning can be divided into the following categories depending on the learning methods: supervised Watter et al.
and self-supervised Srinivas et al.
(2020) learning, both of which can benefit from data augmentation Laskin et al.
representation learning with additional auxiliary tasks Jaderberg et al.
(2016); and using state metrics Ferns et al.
Our approach falls into self-supervised learning category.
The objective encourages state representations to be predictable as a linear combination of the representations of nearby states.
In this paper, we first provide a brief background of the supervised and self-supervised methods that exists in the literature2.
We then explain the proposed method and how it fits into the general landscape of representation learning methods for RL.
Following that, we detail the implementation
B ACKGROUND & R ELATED W ORKS
In Reinforcement Learning (RL), the agent-environment interface is represented by Markov Decision Processes (MDPs), defined by < S , A , R , P > , where S , A , R and P represent the state space, the action space, the reward function and the transition function of the environment, respectively.
RL is a sequential decision making problem where the goal is to find an optimal policy (a mapping from states to actions).
An optimal policy is a policy that maximizes the total discounted return with discount factor, γ ∈ [0 , 1] , obtained when starting from a particular state and following that policy.
Thus, mathematically, an optimal policy, π ∗ is:
[Equation]
where S t ∈ S , A t ∈ A and R t ∈ R are the states, actions and rewards at time t .
The key problem we intend to address here is how to find a good representation of the agent state in conjunction with the main RL problem.
In this section, we will discuss the current approaches for representation learning using supervised and self-supervised methods and describe where our algorithm fits in the literature space.
Reconstruction Loss and Learning Forward Models: The representation is learned by two losses, the original RL loss, that is, the Mean Squared TD Error (MSTDE) and the Auxiliary Loss.
The auxiliary target ˆ Y t depends on the representation learning algorithm used.
For example, it can be trying to predict the next reward or the next state Chung et al.
or even reconstruct the current state.
Additionally, it can also be General Value Functions Sutton et al.
(2011) that learn subtasks.
The most popular approach is to design a low-dimensional representation that is capable of reconstructing the current state.
The commonly used method is Auto Encoders Hinton & Salakhutdinov (2006) that can learn well from high-dimensional observations, such as images Mattner et al.
; van
Hoof et al.
; Watter et al.
Recently, Gelada et al.
and Schwarzer et al.
(2021) developed forward models that can predict the next state in the latent space without the need for reconstruction.
Additionally, these methods can predict multiple steps in the future instead of a one-step prediction.
A large portion of the work on forward models was mainly developed for building world models Ha & Schmidhuber (2018).
These were used for learning the complete dynamics of the world, so that the agent can perform model-based reinforcement learning without the need for interaction with the world.
Recently, Hafner et al.
(2018), Hafner et al.
and Hafner et al.
(2020) have developed model-based RL approaches using forward models in the latent space.
Data Augmentations & Contrastive Learning: Data augmentation methods are commonly used for image input data.
These methods randomly apply different transformations, such as cropping, rotation, or jitter, to generate additional training data, which helps the agent to be robust to small variations in the input space.
Laskin et al.
(2020) and Yarats et al.
(2021b) showed that augmenting the states with such modulations can improve the performance of RL without dedicated representation learning objectives.
Contrastive learning is a self-supervised learning algorithm that takes advantage of similarity constraints between data for representation learning
van den Oord et al.
; Chen et al.
The most popular integration of contrastive learning comes with data augmentation, where positive pairs are modulations of the same image, whereas positive and negative pairs are modulations from different images.
The main goal of the algorithm Srinivas et al.
(2020) is to increase the similarity between positive pairs while keeping negative pairs far apart.
Other recent approaches include prototypical representations Yarats et al.
(2021a), predictive information Lee et al.
and augmented temporal contrast Stooke et al.
(2020), Sermanet et al.
(2017) to associate pairs within a short time.
These methods attempt to make the representations of one state closer to the representations of neighboring states.
Constraining State Representations by Prior Knowledge:
Often, the agents can leverage prior knowledge of the environment to constrain the representations.
A very common method would be to assume that the “interesting” features of the state vary slowly, and thus setting a constraint on the state representation can be beneficial Kompella et al.
(2011); Jonschkowski et al.
Jonschkowski et al.
(2017) also used another prior that the positions of important objects (example, the agent) would vary and this forces the representations to change.
Using both slow and variable features creates a trade-off and prevents the representations from becoming constant.
Jonschkowski & Brock (2015) introduced two priors for robotics which are proportionality and repeatability.
Proportionality introduces a constraint where the representations of two states would vary by the same amount for the same actions taken in the two states.
Repeatability adds a constraint that forces representations of similar states to vary similarly when the same action is applied to them.
Our proposed Locally Constrained Representations (LCR) algorithm falls in this category.
LCR imposes the constraint that state representations should be linearly predictable by representations of surrounding states.
Thus, it is different from the contrastive learning approaches that attempt to pull together representations of neighboring states, as LCR merely introduces a linear explainability constraint to the representations.
Thus, two neighboring representations can be far apart from each other, even though they are linearly related.
In the following section, we describe the intuition behind this approach, formulate an algorithm and provide details on its implementation.
L OCALLY C ONSTRAINED R EPRESENTATIONS 3. I NTUITION
Our main motivation for locally constraining state representations comes from the assumption that in most Reinforcement Learning environments the main objects of interest rarely change rapidly with each action.
Under this assumption, the dynamics should be relatively simple in local policy space, where the states are a few actions apart.
In order to exploit this for representation learning, we introduce the soft constraint that the representation of a state should be closely approximated by a linear combination of neighboring state representations.
This constraint encourages the state representation to remain close to a manifold on which local dynamics can be expressed in a simple function of linear form.
Without this constraint, the representations can change arbitrarily from one state to the next as the representations become intertwined with the local reward signal from the environment.
This might result in over-fitting as the representation would fail to consider the dynamics of the environment.
Even predicting a single future step Chung et al.
(2018) in a self-supervised learning loss can take advantage of the smoothness of the state space.
Predicting multiple steps Gelada et al.
Schwarzer et al.
(2021) in the future can improve the learned representation by helping it capture features that do not change much (example: slow-changing features, such as objects).
Similarly in
LCR, since we use multiple neighboring states for prediction, it can help to capture such features.
3. A LGORITHM
The core idea of this method is to impose a constraint on the representation of a state and to encourage it to be linearly predictable by the representations of nearby states.
This learning formulation is inspired by locally linear embeddings Roweis & Saul (2000), where the embeddings are constrained to be linearly representable by the nearest data points.
Here we define the neighbourhood to be a window of K states in the sequence of transitions that are chosen by the current behaviour policy.
In this way, the representation of a state takes into consideration neighboring states instead of only focusing on the RL objective, which can improve generalization.
However, we do not want features that are completely unrelated to the main task, so we only impose a soft constraint on the representation.
Algorithm 1: Function LCR( sequence length, steps, batch ) : t ←for update in 1,2,3 . . .
updates do
Observe state s
Take Action a
Train policy/value network by base RL algorithm if t % batch =
=
then // Run LCR on the last batch
All States = GetLastBatch ( t, batch )
Initialize W ∼ Uniform(0,1)
N ← Set of all Neighbours in Sequence for state in All States do
N [ T ] =
GetSequence ( sequence length, state ) end for ϕ nearest ( T ) =
f ( N )
Loss =
( f ( All States ) − Wϕ nearest ( T ))for step in 1 , 2 , 3 . . .
steps do
Run Gradient Descent on LCR Loss end for end if end for [Figure]
We need to keep a batch of state transitions of size B , on the basis of which we constrain the representations.
Since RL is online, we need additional memory to keep track of the samples. If
[Figure] the main RL algorithm already uses an experience replay buffer Lin (2004); Mnih et al. , we can directly take the last B samples from the buffer.
More training samples mean better statespace coverage and better generalization.
So higher batch sizes are often better and this comes as a trade-off with respect to compute and memory requirements.
Once we have a batch, B , of states, we obtain K neighbours of each sample from the trajectory, assuming total sequence length is K + 1 , including the sample itself.
We then add a constraint to encourage the representation of the sample to be predictable by the representations of the neighbors.
Let s T be a sample drawn from the buffer of length
B .
Also, let s T − K/ 2 , . . .
, s T − 1 , s T +1 , . . .
, s T + K/ 2 be the nearest neighbors.
Let
[Equation]
be the D -dimensional representation of state s T , where f is a parameterized function.
For all of our experiments, we use a neural network to implement f .
Also, we define
[Equation]
where ϕ nearest ( T ) ∈ R K
× D is the matrix of neighbouring state representations stacked row-wise.
We define the loss
[Equation]
where Wϕ nearest ( T ) is the predicted representation of the state by linear combination of the neighbours with coefficients W ∈ R 1 × K .
We learn W and f by minimizing this loss using gradient descent steps.
After every new batch of samples of size B , this optimization is done in an interleaved manner with the standard RL training.
Since this is added as an auxiliary loss for representation learning, it can be applied on top of any RL algorithm as an auxiliary loss.
We want a soft constraint on the representation and thus we do not solve this optimization completely but take only a few gradient descent steps while keeping W positive (by clipping all negative weights to 0) to limit the solution space.
This is because we are only solving this on a limited batch size with a short state-space coverage, so taking too many steps can overfit to the samples in the batch.
Figure 1 (a) shows a diagram of the method.
R ESULTS
In this section, we test our algorithm in 14 different environments (2 from MiniGrid environments
Chevalier-Boisvert et al.
(2018) and 8 environments from Mujoco Todorov et al.
, andenvironments from Arcade Learning Environment (Atari) Bellemare et al.
The reason for choosing most of these environments is that there are well-defined learnable dynamics and the states change slowly in a given local policy space.
For diversity, we also analyze the performance of LCR in simpler environments such as CartPole and Acrobot.
[Figure] [Figure]
R e w a r d
[Figure]
4. I MPLEMENTATION D ETAILS AND B ASELINES
The description of the environments and the detailed hyper-parameters are mentioned in the supplementary material along with the code for all the environments for reproducibility of our results.
We provide ablation studies of the respective hyper-parameters introduced by LCR, both in the main paper and in the supplementary material, and discuss a suitable thumb rule so that extensive sweeps are not required.
For our baselines, we intentionally avoided choosing other auxiliary representation learning methods.
Our main aim in this paper is to highlight how this strategy of localizing representations can actually help in generalization in representation learning and all our experiments henceforth highlight that.
Furthermore, the experiments that directly perform similar constraints on the environment
Jonschkowski et al.
(2017), were mostly designed for robotics, while Locally Constrained Representations (LCR) is a more general applicable method.
Since this is an additional constraint that is applied in the form of an auxiliary loss, it can be easily added on top of any aforementioned representation learning approach.
4. M INI G RID E NVIRONMENTS
LCR helps in learning representations which prevents over-fitting in environments by constraining the representations locally.
So, given a stochastic environment, is the representation learned by LCR more generizable?
We try to answer this question by choosing two environments from the MiniGrid suite
), where the goals change randomly after every episode.
In both environments, the agent has to navigate a grid and reach a goal.
The agent receives a reward of +1 on reaching the goal and − 0 .
01 for each step that it takes.
We designed the environment RandomGoal , which is an extension of the EmptyRoom environment from MiniGrid Chevalier-Boisvert et al.
(2018) with the goal position changing randomly at the end of every episode.
FourRooms environment is the same as the generic one from MiniGrid Chevalier-Boisvert et al.
(2018), but with reduced size to alleviate the problem of exploration which we are not studying in this paper.
We run DQN Mnih et al.
and DQN with LCR in both environments with fully observable input.
Since the output is a 2D matrix, both architectures have 3 convolutional layers, with the first 2 followed by a max-pooling layer.
The output of the last convolutional layer is “flattened”, that is, reshaped into a vector.
This entire network corresponds to f θ in Fig.
1 (a).
The output of this network is then passed to a value network which consists of two dense layers and which outputs the
Q values.
The entire of flatten layer is representation layer ( ϕ in Fig.
1 (a)) on which LCR is applied.
Exploration is handled by the decaying ϵ -greedy strategy, which is shown by the dotted black line.
Figure 2 demonstrates that LCR is able to perform better because it does not overfit to one goal location.
This problem can happen with learning representations with the RL loss only.
As a result,
DQN (red) performs poorly in both settings, with complete failure in the FourRooms environment.
Adding LCR on top of DQN however, improves the performance significantly, reaching near-optimal policies in both cases.
Note, that the final rewards are noisy because the optimal reward changes with the location of the goal.
Constrained Representations: Figure 3 additionally shows the tSNE van der Maaten & Hinton (2008) plots in 2 dimensions for the representations learned by DQN and DQN with LCR after training for 5000 episodes.
The plots represent the representations of the states ϕ , in 20 trajectories with the same random policy (so that they all visit the same states equally).
It can be seen that the spread of different state representations is much smaller for LCR even though both of the algorithms visit the exact same states.
This indeed confirms that LCR learns a much more constrained representation.
This is especially true for the RandomGoal environment, where most of the states would be easily represented by neighboring states because the grid is empty.
On the other-hand, the representations learned by DQN without LCR would be biased because value estimation is the only learning signal and as a result the representations are somewhat saturated which can lead to over-fitting.
4. M UJOCO
For our second set of experiments, we take Mujoco , a very popular suite Todorov et al.
Brockman et al.
(2016) with well-defined physics, and observe the performance of LCR in easy and difficult tasks.
For the experiments on Mujoco, we used Soft Actor Critic (SAC) Haarnoja et al.
(2018) implementation of Ding et al.
(2020), which is known to have good performance on these tasks.
For these experiments, we used the default architecture from Ding et al.
(2020) with 2 hidden layers followed by an output layer.
The last hidden layer is used as the representation layer ϕ on which LCR was implemented.
For LCR, we used the exact set of hyper-parameters that as the
MiniGrid environment: a batch size B of 5000, 1 / 10 th of the SAC learning rate, sequence length of 11 and gradient steps for LCR optimization as 100.
Figure 4 shows the learning curves on Mujoco .
For relatively difficult tasks, such as HalfCheetahv2, Ant-v2, Humanoid-v2 and HumanoidStandup-v2, we can see that adding LCR on top of SAC dramatically improves performance, and it actually learns a lot faster.
This is probably due to the fact that the dynamics of Mujoco is well defined, and, as such, LCR encourages the representations to consider them better as opposed to learning representations with just the policy and value losses in SAC.
For the easier environments however, it is to be noted that adding LCR does not hurt performance at all, rather in some cases like Pusher-v2 and Striker-v2 it is marginally better than
SAC.
4. A TARI
Although the Arcade Learning Environment does not have continuous dynamics like Mujoco, we tested Rainbow Hessel et al.
(2017) with and without LCR in 4 Atari environments.
For Rainbow, we used the original configuration of Hessel et al.
(2017) with 3 convolutional layers followed by
NoisyLinear layers Fortunato et al.
The flattened output of the convolutional layer is chosen
[Figure]
as the representation layer ϕ for LCR.
We use a batch size of 1000 for LCR, because the state space (images) is much more high dimensional, and to avoid overfitting, we reduced the gradient steps to 20 and the learning rate to 1 / 10 th of the original RL learning rate.
Figure 5 shows the learning curves of LCR compared with the baseline.
Adding LCR significantly improves the performance and sample efficiency in 3 out of 4 environments.
In MsPacman LCR does not show improvement over the baseline but also does not seem to hurt.
4. A BLATION S TUDIES
LCR introduces 4 new hyper-parameters: the size of the batches, B over which we apply LCR, the number of steps of gradient descent, the learning rate, and the neighbourhood size K .
Figure 6 (a) illustrates the parameter sensitivity on the MiniGrid domains w.r.t.
the number of steps of gradient descent.
In both environments, we find a similar plot.
With very few steps, LCR does not work as well, and the same is true for higher number of steps, where the representations are constrained heavily to focus on the linear representations and thus collapses.
Similarly, Figure 6 (b) demonstrates the sensitivity to the number of nearest neighbours.
On the other hand, this depends on the environments.
For example, in the RandomGoal environment, we found higher values of sequence length does equally well if not better, because the grid is empty and having representations to be linearly represented by the neighbours does not affect performance.
However, with higher sequence length for FourRooms , LCR does not perform as well because the representations are forces to be representable from a large portion of the state space, which violates the neighbours being local .
Thus, we see a drop in performance for higher K .
Figure 6 (c) highlights the sensitivity to the learning rate.
For most of our experiments, we found that too small of a learning rate does not help learning at all compared to DQN, because the updates to the representations are not big enough to have a sufficient effect on the actions taken.
On the other hand, a high value of learning rate can be detrimental, as the representations become too constrained.
For our experiments, setting the learning rate to be 1 / 10 th of the learning rate used for the main RL agent works well.
The fourth hyper-parameter, batch size, has no such trade-offs.
The higher the batch size, the less is the chance of over-fitting on the main task, and thus better the performance.
Now, it might seem like we need to run sweeps across all these hyper-parameters to find the best one, and that would be tedious.
However, as mentioned in our experiments we found that normally setting sequence length to 11, number of steps to 100 and learning rate 1 / 10 th of the RL learning rate worked really well.
Thus, as a thumb rule, we set these parameters for all our experiments.
Detailed ablation studies across all the hyper-parameters are provided in the supplementary material.
4. A DDING LCR TO LOW DIMENSIONAL INPUTS
In very simple RL environments, separate representation learning is normally not required.
Since we have the assumption of representations being linearly predictable by the neighbours, it might happen that the representations for simple tasks get constrained to the point that the agent does not learn anything.
To test the performance of LCR with low-dimensional state space, we run LCR on two very simple environments, CartPole and Acrobot Brockman et al.
For these environments we used simple DQN Mnih et al.
(2015) with 2 hidden layers followed by the value layer.
The output of the second hidden layer is the representation.
From Figure 7, we can see that LCR does not impact the performance of the DQN over various sequence lengths.
In fact, for some sequence lengths, DQN with LCR outperforms DQN.
C ONCLUSION AND F UTURE W ORK
LCR adds an auxiliary loss to constrain the state representations in a local policy space.
This can improve generalization and robustness of the representations learned.
In any environment where the states do not change rapidly, with slow-moving features, LCR would show improvement.
Furthermore, the addition of LCR does not hamper performance in simpler environments, where learning representations with the main RL loss is sufficient.
However, the neighboring representations may not be linearly dependent on each other, and we can introduce non-linear locally constrained representations as well.
This will be a much softer constraint on the representations, as nonlinearity will make it much more flexible.
Another interesting future direction could be its extension to vision based tasks with the help of a decoder, where the state representations can be reconstructed with the help of its neighbors.