S4RL: Surprisingly Simple Self-Supervision for Ofﬂine Reinforcement Learning in Robotics
Abstract: Ofﬂine reinforcement learning proposes to learn policies from largecollected datasets without interacting with the physical environment.
These algorithms have made it possible to learn useful skills from data that can then bedeployed in the environment in real-world settings where interactions may be costlyor dangerous, such as autonomous driving or factories.
However, ofﬂine agentsare unable to access the environment to collect new data, and therefore are trainedon a static dataset.
In this paper, we study the effectiveness of performing dataaugmentations on the state space, and study 7 different augmentation schemes andhow they behave with existing ofﬂine RL algorithms.
We then combine the bestdata performing augmentation scheme with a state-of-the-art Q-learning technique,and improve the function approximation of the Q-networks by smoothening outthe learned state-action space.
We experimentally show that using this Surprisingly Simple Self-Supervision technique in RL (S4RL), we signiﬁcantly improve overthe current state-of-the-art algorithms on ofﬂine robot learning environments suchas MetaWorld [1] and RoboSuite [2, 3], and benchmark datasets such as D4RL[4].
Introduction
In reinforcement learning (RL), an agent is trained to interact with its environment and learn usefulskills that help with solving the given task.
However, interacting with the environment may be costly,or unsafe to do from scratch in real-world scenarios such as self-driving or industrial robotics.
Unlikedirectly interacting with the environment, it is much cheaper to collect and store data such that anagent can be trained on it and learn from the past experience.
In direct contrast to reinforcementlearning, ofﬂine RL (or batch RL) [5] allows a setup where a behavioral policy is used to collectand store the data such that a target policy can then be used to train on the data, without any furtherinteractions with the environments.
The behavioral policy can range from human demonstrations,random policy in the given environment, or near-optimal policies for the task.
However, learningfrom such demonstrations is a difﬁcult task since the data does not cover the full state-action space,and naive behavioral policies will not cover the state-action distribution for an optimal policy for atask.
Current ofﬂine model-free RL algorithms include learning a Q -function [6] where a parameterizedneural network is trained to learn the state-action values from data [7–12].
This family of algorithmssuffer from overestimating the true state-action values of data that is not in the same distributionas the ofﬂine dataset used.
Along with overestimation error, another source of error is with thefunction approximation of the neural networks that are typically used to parameterize the Q -functions.
This form of function approximation error may be of many forms, and recent pixel-based methodshave employed self-supervised learning and data augmentations to the image observations duringtraining, resulting in state-of-the-art results for normal reinforcement learning [13–15].
However,it is unclear to how perform such augmentations from proprioceptive information and the roleof augmentation in ofﬂine RL.
In this paper we investigate the role of data augmentations fromproprioceptive observations focused on ofﬂine RL.
Data augmentation also help perform better function approximation for Q -networks since it allows to smoothen out the state space by “visiting”the local regions and ensuring the learned Q -values are similar throughout, since small perturbationsin the given observation should not lead to drastically different Q -values.
Our main contribution is a Surprisingly Simple Self-Supervised ofﬂine RL (S4RL)
algorithm thatcombines studying data augmentations with a simple Q -learning method to signiﬁcantly improvethe performance of ofﬂine RL algorithm.
The proposed framework is general and can be added to anumber of off-the-shelf ofﬂine/batch RL algorithms for robot learning.
We present comparisons ourmodel to other comparable self-supervised learning strategies and state-of-the-art ofﬂine RL methodson the D4RL benchmark dataset [4], which consists of continuous control, navigation and roboticcontrol tasks collected from suboptimal and human agents.
We further experiment with two dexterousrobot manipulation environments: MetaWorld [1] and RoboSuite [2, 3].
Across all evaluations, weﬁnd that our proposed state-based data augmentation for Q -learning signiﬁcantly outperforms baseofﬂine RL algorithm as well as various competitive data augmentation baselines.
Related Work
Ofﬂine reinforcement learning: In ofﬂine RL, a static dataset of demonstrations is collected, and isused to train an agent for a given environment.
Popular actor-critic algorithms, such as Soft Actor
Critic (SAC) [16], tend to perform poorly on learning from ofﬂine datasets, because they are unableto generalize to out-of-distribution (OOD) data because of an overestimation bias: where the criticoverestimates the values of state-action pairs not previously encountered from the collected datawhich leads to a brittle policy.
The convergence properties of value functions have been proven in simple settings [17]; howeverlearning a policy using the learned value functions remains an active challenge in ofﬂine RL.
Simplealgorithms like ﬁtted Q -iterations have been able to learn state-action value functions, but do notwork well with sub-optimal data [18, 19].
To help solve the generalization and extrapolation problem,different ofﬂine RL algorithms have proposed to use constrained policy optimization between thetarget policy and the behaviour policy that was used to collect the data.
These constrains includeusing KL-divergence [10, 20], MMD [8] or f -divergence [21].
Some algorithms have also proposedto learn a state conditioned VAEs to minimize the extrapolation error [7, 8, 22].
Recent work by
Buckman et al.
[23] investigated using worst-case analysis to train a policy, denoted as pessimism .
Conservative Q -Learning (CQL) explicitly deals with the problem of overestimation for OOD stateaction distribution by pushing down the value of randomly sampled state-action distributions in thecritic objective [24].
Learning a dynamics model from ofﬂine data has been promising [25, 26].
Arecent review of ofﬂine RL outlines the state of the ﬁeld [5].
Recent work has also shown greatpromise using ofﬂine RL for real robotics tasks by learning policies that can directly be deployed onreal robots to perform complex tasks [27–29].
Similar to before, such work relies on collecting large scale robotic datasets [3, 30, 31], and learning policies that can then be safely deployed after a fewsteps of ﬁnetuning.
Representation learning in RL: Recent work in state and pixel-based RL suggests the need tolearn better representations from data usig better network architectures [32] or self-supervision [13].
Self-supervised representation learning has recently been applied to learn RL agents from pixeldata.
By using data augmentations and contrastive learning, Laskin et al.
[13] showed signiﬁcantimprovements on learning from pixel data.
The need for contrastive learning was simpliﬁed by RLwith Augmented Data (RAD) [15] and Data-Regularized Q -learning (DrQ) [14] as they providea simpler alternative that used data augmentations without a contrastive objective.
More recently,
Self-Predictive Representations (SPR) optionally uses data augmentations for a self-supervision[33].
The closest work to S4RL include DrQ [14] and RAD [15] in that we utilize their method for
Q -learning over augmentations, but in contrast we focus speciﬁcally on ofﬂine RL for robotic tasksand propose augmentations from states; we also benchmark different self-supervision techniquesfrom states.
Preliminaries
In reinforcement learning, an agent interacts with an environment to learn an optimal policy.
Thisis typically framed as a Markov Decision Process (MDP) which can be represented by a tuple ofthe form ( S , A , ρ , r , γ , ρ 0 ), where S is the state space, A is the action space, ρ ( s t +1 | s t , a t ) is thetransition function given the current state and action pair, r ( s t ) is the reward model given the state, γis the discount factor where γ ∈ [0 , 1) and ρ 0 is the initial state distribution of the MDP.
A policy, represented by π : S → A is trained to maximize the expected cumulative discountedreward in the MDP.
This can simply be formalized as: E (cid:104) (cid:80)
[Equation]
Furthermore, a state-action value function, Q ( s t , a t ) , is the value of performing a given action a tgiven the state s t .
The Q -function is trained by minimizing the Bellman Error over Q in a step calledpolicy evaluation
[Equation]
where Q i is the i -th step of policy evaluation.
The policy is then trained to maximize the state-actionvalues of performing an action a t given s t in a step called policy improvement:
[Equation]
where π i is the i -th step of policy improvement.
Unlike traditional reinforcement learning, in ofﬂine RL, the goal is not to learn an optimal policyfor the MDP, but rather to learn an optimal policy given the dataset.
A behaviour policy µ is used tocollect a static dataset D which is then used to train a target policy π .
The policy improvement cannow be stated as
[Equation]
where D is the static dataset.
Since the nature of the behaviour policy µ is unknown, and can becomposed of one (or more) sub-optimal policies, the RL task becomes challenging.
Since ofﬂine RLalgorithms tend to generalize poorly to OOD data, the nature of µ and the optimality of µ is important.
As shown Fu et al.
[4], typical ofﬂine RL methods perform poorly when a mixture of suboptimalpolicies or an untrained random policy is used to collect the dataset D .
In deep reinforcement learning,both the policy and the value functions are parameterized using neural networks and trained usinggradient descent.
Method
[Figure]
In this section we will ﬁrst consider a simple Q -learning technique based on DrQ [14], which helpscombine data augmentations and learning better value functions during training by encouraging localpoints around the state to have similar Q -values.
Then, we introduce 7 different data augmentationsstrategies and apply them on the otherwise static dataset D to improve function approximation bysmoothening the state-space for the Q -network.
4.1 Self-Supervision in Ofﬂine RL
We denote a data augmentation transformation as: T ( ˜ s t | s t ) where s t ∼ D , ˜ s t is the augmented state,and T is a stochastic transformation.
After producing one or multiple augmentations of a given state˜ s t ∼ T ( s t ) , we need a method to incorporate the augmentation to encourage smoothness and consistency.
For example, [13] uses contrastive learning to “pull” the representations of two augmentationsof the same state towards each other, while “pushing” two different state representations apart fromeach other.
We follow a similar scheme to DrQ [14], where we perform multiple augmentationsof a given state, and we simple average the state-action values and target values over the differentaugmetations of the state.
[Equation]
The main difference between our proposed objective and the normal Q -learning objective is themean over i different augmentations in the second term of the equation.
We augment the Bellmanerror in Equation 1 to simply be the mean error over the i different augmentations of the same state.
Intuitively, this will help improve the consistency of the Q -value within some perturbation ﬁeld ofthe current state s t , since the Bellman backups are taken over i -different views of the same state s t .
By simply averaging the state-action values and target values, we assume that the reward function islocally smooth to small perturbations in the state.
Utilizing multiple augmentations further allowsthe networks to learn different variations of the data; we explore the role of increasing the numberof different augmentations empirically in Section 5.1.
We use the trained Q -networks to performpolicy improvement using the objective in Equation 3, without augmentations.
The beneﬁts of theaugmentations are distilled to the policy since the value functions are used to train the policy.
Bycombining such local perturbations to the states with self-supervision, we are able to learn a morerobust policy that can generalize better to unseen data when deployed on robots.
An overview ofthe proposed method is available in Figure 1, where we disentangle the data augmentation andreinforcement learning steps to learn q -values and target values over i different augmentations of theoriginal state.
Without loss of generality, to show local perturbations, we draw an (cid:15) -ball
over thecurrent state s t , however different versions of augmentations are possible, as will be discussed.
4.2 Data Augmentations
In computer vision research, augmentation are commonly used as a way to obtain the same datapointfrom multiple viewpoints.
Transformations such as rotation, translations, color jitters, etc. arecommonly used to train neural networks.
Such transformations preserve the semantics of the imageafter the transformation since for example: an image of a cat rotated 15 ◦ remains an image of acat.
However, when working from only proprioceptive information of an agent (for example: thejoint angles and velocity information of an industrial robot), such transformations are semanticallymeaningless.
In ofﬂine RL, since an agent is unable to interact with the environment to collect new data samples,data augmentations serve as a simple technique that can allow the agent to do view different augmentations of the trajectories in the dataset.
Data augmentations from states, assumes that the output of asmall transformation to an input state results in a physically realizable state in the environment.
If theaugmented state is in the domain of the state space, ˜ s t ∈ S , then we consider it a valid transformation.
By performing valid augmentation stratgies, we are able to artiﬁcially increase the amount of dataavailable during training.
Even if a transformation is physically realizable in the environment, a data augmentation strategy thatis too aggressive may end up hurting the RL agent since the reward for the original state may notcoincide with the reward obtained from the augmented state.
That is r ( s t ) (cid:54) =
r ( ˜ s t ) .
Therefore, onekey assumption made when performing data augmentations on the states is that the reward model r ( s t )is smooth , that is r ( s t )
≈ r
( ˜ s t ) .
Therefore the choice of T ( ˜ s t | s t ) needs to be a local transformationto perturb the state without changing the semantics.
The choice of T ( ˜ s t | s t ) is more important inofﬂine RL than in traditional RL, since in traditional RL the agent may be able to self-correct for poorchoices of transformations by being able to visit those states and correct the value functions using
Bellman backups.
A comparison is presented in Figure 2, where we see that compared to normalofﬂine RL, adding augmentation allows for multiple paths in the policy evaluation step using Bellmanbackups.
By using data augmentation, we seek to exploit this smoothness property of the rewardmodel in the MDP, which enables us to artiﬁcially visit states that may be in the state space, but notin the dataset: s t , ˜ s t ∈ S but ˜ s t / ∈ D .
Following the constraints, we consider the following transformations T ( ˜ s t | s t ) :
zero-mean Gaussian noise to the state:
˜ s t ← s t + (cid:15) , where (cid:15) ∈ N (0 , σI ) and σ is an
important hyperparameter, which we will later discuss.
zero-mean Uniform noise to the state:
˜ s t ← s t + (cid:15) , where (cid:15) ∈ U ( − α, α ) and α is an
important hyperparameter.
random amplitude scaling as in RAD [15]: ˜ s t ← s t ∗ (cid:15) ,
where (cid:15) ∈ U ( α, β ) , which
preserves the signs of the states.
dimension dropout which simply zeros one random dimension in the state: ˜ s t ← s t ·
where 1 is a vector of 1s with one 0 randomly sampled from a Bernoulli distribution.
Thistransformation preserves all-but-one intrinsic values of the current state.
state-switch where we ﬂip the value of 2 randomly selected dimensions in the state.
This
naive transformation will likely break the “physical realizability” assumption, since it is
possible that the two randomly selected samples are semantically dissimilar properties of a
robot (such as joint angle and joint velocity).
To overcome this we hardcode the pairs of
dimensions that can be switched for each environment (details in Appendix F).
state mix-up where we use mixup [34] between the current and the next state: ˜ s t ←
λs t + (1 − λ ) s t +1 ,
where λ ∼ Beta ( α, α ) with α = 0 .
4 as in [34].
The next state s t +1remains unchanged.
adversarial state training which we take the gradient with respect to the value function to
obtain a new state.
This can be written as ˜ s t ← s t + (cid:15) ∇ s t J Q
π ( s t )) , where J Q is simplythe policy evaluation update as in Equation 2, and (cid:15) constraints the size of the gradient.
The
intuition behind this is to choose the direction within an (cid:15) -ball of s t where the Q -value
deviates the most.
This augmentation is intuitively the most useful in state spaces with high dimensionality where random noise may be less effective [35].
Experiments
[Figure]
In this section, we will ﬁrst describe the dataset (D4RL [4]), and some hyperparameters and information about the experimental setup.
Then we will use S4RL to investigate the effect of different dataaugmentation strategies on the OpenAI Gym subset of D4RL tasks to investigate which augmentationsare useful for learning.
Using the best data augmentation techniques found, we compare the algorithmto other self-supervision techniques that have been proposed for pixel-based RL, namely Contrastive
Unsupervised Reinforcement Learning (CURL) [13], and SAC+AutoEncoder (SAC+AE) wherewe feed the augmentations to the AutoEncoder to learn generalizable representations [36].
We add
S4RL framework to two state-of-the-art ofﬂine RL algorithms: Conservative Q -learning (CQL) [24]and Behaviour Regularized Actor Critic (BRAC) [10] to show the effectiveness over both baselineagents to show the generality of S4RL.
We add the CQL regularization to each self-supervisedbaseline considered, since the self-supervised baselines were considered for online RL and CQLadds a regularization term to help with overestimation.
We ﬁrst test each baseline on the full suiteof D4RL tasks [4], which includes environments such as a Maze environment where an ant agentmust learn locomotion and navigation skills to reach an environment goal which requires hierarchicalcontrol, Adroit domain which requires learning dextrous object manipulation for ﬁne-grained controlof a robotic hand [37], Kitchen environment which simulates a Franka Panda robot to performkitchen tasks.
We additionally experiment with more difﬁcult robotic environments: MetaWorld [1]which uses a simulated Sawyer robot and RoboSuite [2] which uses a Franka Panda robotic arm formanipulation.
We further benchmark each of the baselines on learning with limited data in Appendix
A, to investigate the effect of self-supervision when the agent has signiﬁcantly less data to trainfrom.
Finally, we experiment with different ways to regularize Q -networks typically used in neuralnetwork training such as Dropout and Weight Decay in Appendix D and include ablation experimentsinvestigating the number of augmentations and the role of hyperparameters in Appendix E.
Hyperparameters: For all experiments, we do not perform any hyperparamter tuning to the base
CQL agent; all agents are trained using the original hyperparameters.
Similar to DrQ, we use 2augmentations for obtaining the state-action values and the target values.
We use a σ and α value of3 × 10 − 4 for the zero-mean Gaussian and Uniform noise augmentation variant, respectively.
We usea value of 1 × 10 − 4 for (cid:15) for adversarial training; an ablation on σ and (cid:15) is available in Appendix E.
5.1 Results
Role of Data Augmentations We ﬁrst investigate different data augmentation schemes and itsperformance on the OpenAI Gym subset of the D4RL tasks.
We compare a base CQL agent with andwithout different forms of augmentation strategies, and tabulate the mean normalized performance,average normalized performance over all the tasks and relative rankings in Table 1.
We see that usingzero-mean Gaussian noise ( N ), zero-mean Uniform noise ( U ), state mix-up and adversarial statetraining consistently outperform the the baseline CQL agent, as well as different data augmentationvariants.
The average-ranking of S4RL-Adv and S4RL- N are 1.83 and 1.92 respectively, suggestingtheir effectiveness over a wide range of task and data distributions.
We also see that CQL+S4RLagent is able to learn useful policies given data collected from a random policy as evidenced by theperformance in “walker-random” where the base CQL agent is unable to learn any meaningful policy.
We also see that the S4RL agent is able to signiﬁcantly outperform the baseline CQL agent oncomplex data distributions such as “-medium-replay” where the data collected is from all the data
[Table] [Table]
collected while training a policy in the environment.
Therefore the data split consists of data thatranges from a completely untrained policy (random policy), to a “medium” trained policy.
It is also important to look at the different augmentation techniques that do not help policy learning.
Techniques that hurt the performance and perform worse than the baseline CQL agent include
Dimension-Dropout and State-Switch.
Both techniques are inspired by popular computer vision dataaugmentation algorithms, mainly MixMatch [38] and CutMix [39].
Since both techniques performelement-wise operations, it is likely that they omit important information about the state of the robotsuch as joint velocity.
Without such information, it is possible that the value function is unable toreason about the environment which results in poor value estimates.
[Figure]
Comparison to other self-supervision techniques Following insights from Table 1, we choosethree data augmentation approaches, namely S4RL- N , S4RL-MixUp and S4RL-Adv, to compareto other self-supervision techniques.
We present the results over all the different tasks in D4RL in
Table 2 using base CQL [24] and BRAC agents [10].
Despite its simplicity, we continuously observethat +S4RL- N and +S4RL-adv agents are able to signiﬁcantly outperform the baselines on almostall tasks, and is comparable to the best in the others for both ofﬂine RL algorithms.
We speciﬁcallynote that +S4RL-Adv signiﬁcantly outperforms all other baselines when the state-dimensionality ishigh, such as the dexterous manipulation environment of Adriot and on Franka robot.
In challengingenvironments that require hierarchical control such as the AntMaze environments, we see that
the
S4RL agent continues to be the best performing agent.
Interestingly, on certain environments namely“pen-cloned”, “pen-human”, “hammer-human” and “door-human”
the BRAC+S4RL agent is able tolearn some useful skills, whereas the base CQL agent performs poorly.
This further highlights theusefulness and generality of performing state-based augmentations for ofﬂine reinforcement learning.
Dexterous Manipulation Environments Finally, we perform more experiments in the roboticdomain using the MetaWorld [1] and RoboSuite environments on difﬁcult manipulation tasks.
For
MetaWorld, we train a SAC agent [16] for 1M steps and collect 1000 trajectories of 200 episodelength.
For RoboSuite, we collect data at 3 different instances during training to make the datadistribution more complex.
More details on the data collection is available in Appendix G.
The resultsare presented in Figure 4.
Similar to before we see signiﬁcant improvements over the baseline ineach of the 7 environments over both domains.
The proposed S4RL-Adv agent is the only agent thatis able to outperform the trained behaviour policy on 3 of the 5 MetaWorld tasks, while continuingto signiﬁcantly outperform the baselines by as much as 3 times, as in the case of “push-v1” and“pick-and-place-can”.
Full numerical results can be found in Appendix C
Conclusion
In this paper, we present S4RL: a Surprisingly Simple Self-Supervised ofﬂine RL method that usesdata augmentations to improve the function approximation for Q -learning algorithms in ofﬂine
RL.
S4RL offers simplicity and ease of implementation, and can be added to any ofﬂine agentthat requries Q -learning.
We ﬁrst compare and benchmark the effectiveness of 7 different dataaugmentation strategies from states in ofﬂine RL and then use the insights to compare againstdifferent self-supervised representation learning algorithms that have been proposed for pixel-basedonline RL.
We observe signiﬁcant performance gains over the benchmark D4RL dataset [4] and onmore dexterous robot manipulation environments.
We observe signiﬁcant performance gains whichmake ofﬂine policy learning increasingly more competitive and help robot learning from only pastdata.
Interesting future extension of this work can seek to use the suggested augmentation schemes tobuild better self-supervised learning algorithms.