S4RL: Surprisingly Simple Self-Supervision for Offline Reinforcement Learning

TL;DR

Offline RL suffers from overfitting and poor generalization when learning from fixed datasets. The authors propose S4RL, a surprisingly simple self-supervised framework that applies state-space augmentations and averages Q-value estimates over augmented views to regularize learning. Empirically, S4RL yields significant improvements over state-of-the-art offline RL methods on D4RL benchmarks and dexterous robotic tasks (MetaWorld, RoboSuite), including scenarios with limited data. The approach is easy to implement and can be integrated with existing offline RL algorithms to improve robustness and performance in real-world robotic applications.

Abstract

Offline reinforcement learning proposes to learn policies from large collected datasets without interacting with the physical environment. These algorithms have made it possible to learn useful skills from data that can then be deployed in the environment in real-world settings where interactions may be costly or dangerous, such as autonomous driving or factories. However, current algorithms overfit to the dataset they are trained on and exhibit poor out-of-distribution generalization to the environment when deployed. In this paper, we study the effectiveness of performing data augmentations on the state space, and study 7 different augmentation schemes and how they behave with existing offline RL algorithms. We then combine the best data performing augmentation scheme with a state-of-the-art Q-learning technique, and improve the function approximation of the Q-networks by smoothening out the learned state-action space. We experimentally show that using this Surprisingly Simple Self-Supervision technique in RL (S4RL), we significantly improve over the current state-of-the-art algorithms on offline robot learning environments such as MetaWorld [1] and RoboSuite [2,3], and benchmark datasets such as D4RL [4].

Figures (5)

• Figure 1: Overview of the proposed augmentation scheme. Here $T_1(s_t)$ and $T_2(s_t)$ represent different augmentation schemes. We simply perform $k-$ data augmentations to the observations $s_t$, and use that the augmented states to calculate the q-values and target values for better function approximation of the Q-network by ensuring that local perturbations in the state space have similar $Q$-value estimates.
• Figure 2: A diagram illustrating the difference between a Bellman backup for offline RL versus offline RL with augmentations. Using data augmentations, we are able smoothen out the local state-space around the observations in the dataset $s_t$.
• Figure 3: Illustration of the robotic environments used in the experiments. The environments range from robotic arms for difficult manipulation tasks such as opening a door or placing a can into a bin or dexterous manipulation of a pen to a desired configuration using a robot hand.
• Figure 4: Results on Dexterous Manipulation Environments from MetaWorld metaworld and RoboSuite robosuitemandlekar2018roboturk. We report the fraction (as %) of goals reached in each of the environments during evaluation. Notably with a fixed amount of data, S4RL augmentation improves policy success rate over base CQL by approx. 20% across different tasks.
• Figure 5: Learning from limited data. We train an S4RL agent and the baselines on a small $5\%, 10\%,$ and $25\%$ of the different D4RL data splits. We see that S4RL is consistently the best performing algorithm when learning from limited data.