SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Qiyang Li $^{\alpha \gamma *}$ , Jason Zhang $^{\alpha *}$ , Dibya Ghosh $^{\alpha}$ , Amy Zhang $^{\beta \gamma}$ , Sergey Levine $^{\alpha}$   
UC Berkeley $^{\alpha}$ , UT Austin $^{\beta}$ , Meta $^{\gamma}$
{qcli, jason.z, dibya.ghosh}@berkeley.edu
amy.zhang@austin.utexas.edu, svlevine@eecs.berkeley.edu
Learning to solve tasks from a sparse reward signal is a major challenge for standard reinforcement learning (RL) algorithms. However, in the real world, agents rarely need to solve sparse reward tasks entirely from scratch. More often, we might possess prior experience to draw on that provides considerable guidance about which actions and outcomes are possible in the world, which we can use to explore more effectively for new tasks. In this work, we study how prior data without reward labels may be used to guide and accelerate exploration for an agent solving a new sparse reward task. We propose a simple approach that learns a reward model from online experience, labels the unlabeled prior data with optimistic rewards, and then uses it concurrently alongside the online data for downstream policy and critic optimization. This general formula leads to rapid exploration in several challenging sparse-reward domains where tabula rasa exploration is insufficient, including the AntMaze domain, Adroit hand manipulation domain, and a visual simulated robotic manipulation domain. Our results highlight the ease of incorporating unlabeled prior data into existing online RL algorithms, and the (perhaps surprising) effectiveness of doing so.
Exploration, particularly in sparse reward environments, presents a major challenge for reinforcement learning, and standard exploration methods typically need to seek out all potentially novel states to cover all the places where high rewards may be located. Luckily, in many real-world RL settings, it is straightforward to obtain prior data that can help the agent understand how the world works. For example, if we are trying to find where we left our keys, we would not relearn how to navigate our environment, but rather might revisit locations that we recall from memory. If the data has reward annotations, pretraining with offline RL provides one solution to accelerate online finetuning. However, in many domains of interest, it is more likely for prior datasets to be task-agnostic, and not be labeled with the reward for the new task we hope to solve. Despite not having reward labels, such prior data can be useful for an agent exploring for a new task since it describes the
Figure 1: How to leverage unlabeled prior data for efficient online exploration? We label the prior data with an optimistic reward estimate and run RL on both the online and offline data. This allows more efficient exploration around the trajectories in the prior data, improving sample efficiency especially for hard exploration tasks.
environment dynamics and indicates regions of the state space that may be interesting to explore.
*Equal Contribution
How can this unlabeled data be utilized by a deep RL algorithm? Prior data informs the agent about potentially high-reward or high-information states, and behaviors to reach them. Initially, the agent may act to reach these states, leading to directed exploration in a more informative subspace rather than the complete state space. This should be an iterative process, since the agent can use the new online experience to refine what directions from the prior data it chooses to explore in. After the agent begins to receive reward signals, the prior data can help the agent to more reliably reach and exploit the reward signals, steering the exploration towards a more promising region.
Perhaps surprisingly, such informed exploration can be elicited in a very simple way: by labeling the prior data with optimistic rewards and adding it to the online experience of an off-policy RL algorithm (Figure 1). Specifically, we choose to impute the missing reward on the unlabeled prior data with an optimistic upper-confidence bound (UCB) of the reward estimates obtained through the agent's collected online experience. Early in training, these optimistic reward labels will be high for states from the prior data, and optimizing the RL objective with this data will elicit policies that attempt to reach regions of state present in the prior data. As states in the prior data are visited, the labeled reward for these states will regress towards their ground-truth values, leading to a policy that either explores other states from the prior data (when true reward is low) or returns to the region more consistently (when high). By training an RL policy with a UCB estimate of the reward, exploration is guided by both task rewards and prior data, focusing only on visiting states from the prior data that may have high reward. In practice, we find that this general recipe works with different off-policy algorithms [Ball et al., 2023, Kostrikov et al., 2021].
Our main contribution is a simple approach to leverage unlabeled prior data to improve online exploration and sample efficiency, particularly for sparse long-horizon tasks. Our empirical evaluations, conducted over domains with different observation modalities (e.g., states and images), such as simulated robot navigation and arm and hand manipulation, show that our simple optimistic reward labeling strategy can utilize the unlabeled prior data effectively, often as well as the prior best approach that has access to the same prior data with labeled rewards. In addition, we can leverage offline pre-training of task-agnostic value functions (e.g., ICVF [Ghosh et al., 2023]) from offline data to initialize our optimistic reward model to further boost the online sample efficiency in more complex settings, such as image-based environments. In contrast to other approaches in the same setting, our method is conceptually simpler and easier to be integrated into current off-policy online RL algorithms, and performs better as we will show in our experiments.
2 Problem Formulation
We formalize our problem in an infinite-horizon Markov decision process (MDP), defined as  $\mathcal{M} = (\mathcal{S},\mathcal{A},\mathcal{P},r,\gamma ,\rho)$ . It consists of a state space  $\mathcal{S}$ , an action space  $\mathcal{A}$ , transition dynamics  $\mathcal{P}:S\times \mathcal{A}\mapsto \mathcal{P}(S)$ , a reward function  $r:S\times \mathcal{A}\mapsto \mathbb{R}$ , a discount factor  $\gamma \in (0,1]$ , an initial state distribution of interest  $\rho$ . Algorithms interact in the MDP online to learn policies  $\pi (a|s):S\mapsto \mathcal{P}(\mathcal{A})$  that maximize the discounted return under the policy  $\eta (\pi) = \mathbb{E}_{s_{t + 1}\sim \mathcal{P}(\cdot |s_t,a_t),a_t\sim \pi (\cdot |s_t),s_0\sim \rho}[\sum_{t = 0}^{\infty}\gamma^t r(s_t,a_t)]$ . We presume access to a dataset of transitions with no reward labels,  $\mathcal{D} = \{(s_i,a_i,s_i')\}_i$  that was previously collected from  $\mathcal{M}$ , for example by a random policy, human demonstrator, or an agent solving a different task in the same environment. The core challenge in this problem setting is to explore in a sparse reward environment where you do not know where the reward is. While structurally similar to offline RL or online RL with prior data, where the agent receives a reward-labelled dataset, the reward-less setting poses a significant exploration challenge not present when reward labels are available. Typically offline RL algorithms must find the best way to use the prior data to exploit the reward signal. In our setting, algorithms must use the prior data to improve exploration to acquire a reward signal. Using this prior data requires bridging the misalignment between the offline prior dataset (no reward) and the online collected data (contains reward), since standard model-free RL algorithms cannot directly use experiential data that has no reward labels.
3 Exploration from Prior Data by Labeling Optimistic Reward (EXPLORE)
In this section, we detail a simple way to use the unlabeled prior data directly in a standard online RL algorithm. Our general approach will be to label the reward-free data with learned optimistic
rewards and include this generated experience into the buffer of an off-policy algorithm. When these rewards are correctly chosen, this will encourage the RL agent to explore along directions that were present in the prior data, focusing on states perceived to have high rewards or not yet seen in the agent's online experience.
Algorithm 1 EXPLORE
1: Input: prior unlabeled data  $\mathcal{D}_{\mathrm{offline}}$
2: Initialize the UCB estimate of the reward function:  $\mathrm{UCBR}(s,a)$
3: Online replay buffer  $\mathcal{D}\gets \emptyset$
4: Initialize off-policy RL algorithm with a policy  $\pi$ .
5: for each environment step do
6: Collect  $(s, a, s', r)$  using the policy  $\pi$  and add it to the replay buffer  $\mathcal{D}$ .
7: Update the UCB estimate  $\mathrm{UCBR}(s,a)$  using the new transition  $(s,a,s^{\prime},r)$
8: Relabel each transition  $(s, a, s')$  in  $\mathcal{D}_{\mathrm{offline}}$  with  $\hat{r} = \mathrm{UCBR}(s, a)$
9: Run off-policy RL update on both  $\mathcal{D}$  and the relabeled  $\mathcal{D}_{\mathrm{offline}}$  to improve the policy  $\pi$ .
10: end for
General algorithm. Since the prior data contains no reward labels, the agent can only acquire information about the reward function from the online experience it collects. Specifically, at any stage of training, the collected online experience  $\mathcal{D}$  informs some posterior or confidence set over the true reward function that it must optimize. To make use of the prior data, we label it with an upper confidence bound (UCB) estimate of the true reward function. Optimizing an RL algorithm with this relabeled prior experience results in a policy that is "optimistic in the face of uncertainty", choosing to guide its exploration towards prior data states either where it knows the reward to be high, and prior data states where it is highly uncertain about the reward. Unlike standard exploration with optimistic rewards for states seen online, the prior data may include states that are further away than those seen online, or may exclude states that are less important. Consequently, the RL agent seeks to visit regions in the prior data that are promising or has not visited before, accelerating online learning.
To explore with prior data, our algorithm maintains an uncertainty model of the reward, to be used for UCB estimates, and an off-policy agent that trains jointly on relabelled prior data and online experience. As new experience is added through online data collection, we use this data to update the agent's uncertainty model over rewards. When training the off-policy agent, we sample data from both the online replay buffer (already reward labeled) and from the prior data (no reward labels). The prior data is labelled with rewards from the UCB estimate of our uncertainty model, and then the online and prior data is jointly trained upon. Training with this data leads to a policy that acts optimistically with respect to the agent's uncertainty about the reward function, which we use to guide exploration and collect new transitions. Algorithm 1 presents a sketch of our method.
Notice that while this strategy is similar to exploration algorithms that use novelty-based reward bonuses [Bellemare et al., 2016, Tang et al., 2017, Pathak et al., 2017, Burda et al., 2018, Pathak et al., 2019, Gupta et al., 2022], since the optimism is being added to the prior data, the resulting exploration behaviors differ qualitatively. Adding novelty bonuses to online experience guides the agent to states it has already reached in the online phase, whereas adding optimism to the prior experience encourages the agent to learn to reach new states beyond the regions it has visited through online exploration. Intuitively, while standard optimism leads to exploration along the frontier of online data, optimism on prior data leads to exploration beyond the frontier of online data.
Practical implementation. In practice, we obtain UCB estimates of the rewards using a combination of a reward model and random-network-distillation (RND) [Burda et al., 2018], a novelty-based reward bonus technique. The former provides a reward estimate of the true reward value and the latter (RND) quantifies the uncertainty of the estimate, allowing us to get a UCB reward estimate.
For the reward model, we update a randomly initialized network  $r_{\theta}(s,a)$  simultaneously with the RL agent on the online replay buffer by minimizing the squared reward loss:
$$
\mathcal {L} (\theta) = \mathbb {E} _ {(s, a, r) \sim \mathcal {D}} \left[ (r _ {\theta} (s, a) - r) ^ {2} \right].
$$
RND randomly initializes two networks  $f_{\phi}(s,a), \bar{f}(s,a)$  that each output an  $L$ -dimensional feature vector for each state and action. We keep one network,  $\bar{f}$ , frozen and update the parameters of the
Figure 2: Visualizations of the agent's exploration behaviors using our optimistic reward labeling strategy on antmaze-medium. The dots shown (500 points each) are sampled uniformly from the online replay buffer (up to 60K environment steps), and colored by the training environment step at which they were added to the buffer. Both Online and Online + RND do not use prior data, and the latter uses RND reward bonus [Burda et al., 2018] on top of the received online reward to encourage online exploration. Naive learns a reward model and relabels the unlabeled prior data with its reward estimate with no optimism. Ours labels the prior data with an optimistic reward estimate.
Figure 3: Aggregated results for 6 AntMaze tasks, 3 Adroit tasks and 3 COG tasks. Ours is the only method that largely recovers the performance of vanilla RLPD (Oracle), which has access to true rewards (Ours does not). The baselines that do not use the prior unlabeled data (Online and Online + RND) perform poorly on both AntMaze and Adroit. Naively labeling the prior data without optimism (Naive) performs competitively on Adroit but poorly on AntMaze. Section 5.2 contains the description for the baselines we compare with.
other network  $\phi$  to minimize the mean-square-error between the predicted features and the frozen ones once on every new transition  $(s_{\mathrm{new}},a_{\mathrm{new}},s_{\mathrm{new}}^{\prime},r_{\mathrm{new}})$  encountered:
$$
\mathcal {L} (\phi) = \frac {1}{L} \| f _ {\phi} (s _ {\mathrm {n e w}}, a _ {\mathrm {n e w}}) - \bar {f} (s _ {\mathrm {n e w}}, a _ {\mathrm {n e w}}) \| _ {2} ^ {2}.
$$
During online exploration,  $\| f_{\phi}(s,a) - \bar{f} (s,a)\| _2^2$  forms a measure of the agent's uncertainty about  $(s,a)$ , leading us to the following approximate UCB estimate of the reward function
$$
\mathrm {U C B R} (s, a) \leftarrow r _ {\theta} (s, a) + \frac {1}{L} \| f _ {\phi} (s, a) - \bar {f} (s, a) \| _ {2} ^ {2}.
$$
For the RL agent, we use RLPD [Ball et al., 2023], which is specifically designed for efficient online RL and performs learning updates on both online and labeled offline data. It is worth noting that our method is compatible with potentially any other off-policy RL algorithms that can leverage labeled online and offline data. We experiment with an alternative based on implicit-Q learning in Appendix C.4. Although we only have access to unlabeled data in our setting, we constantly label the prior data with UCB reward estimates, enabling a direct application of RLPD updates.
4 Related Work
Exploration with reward bonuses. One standard approach for targeted exploration is to add exploration bonuses to the reward signal and train value functions and policies with respect to these optimistic rewards. Exploration bonuses seek to reward novelty, for example using an approximate density model [Bellemare et al., 2016, Tang et al., 2017], curiosity [Pathak et al., 2017, 2019], model error [Stadie et al., 2015, Houthooft et al., 2016, Achiam and Sastry, 2017, Lin and Jabri, 2023], or even prediction error to a randomly initialized function [Burda et al., 2018]. When optimistic rewards are calibrated correctly, the formal counterparts to these algorithms [Strehl and Littman, 2008, Dann et al., 2017] enjoy strong theoretical guarantees for minimizing regret [Azar et al., 2017, Jin et al., 2018]. Prior works typically augment the reward bonuses to the online replay buffer to encourage exploration of novel states, whereas in our work, we add reward bonuses to the offline data, encouraging the RL agent to traverse through trajectories in the prior data to explore more efficiently.
Using prior data without reward labels. Prior experience without reward annotations can be used for downstream RL in a variety of ways. One strategy is to use this data to pre-train state representations and extract feature backbones for downstream networks, whether using standard image-level objectives [Srinivas et al., 2020, Xiao et al., 2022, Radosavovic et al., 2022], or those that emphasize the temporal structure of the environment [Yang and Nachum, 2021, Ma et al., 2022, Farebrother et al., 2023, Ghosh et al., 2023]. These approaches are complementary to our method, and in our experiments, we find that optimistic reward labeling using pre-trained representations lead to more coherent exploration and faster learning. Prior data may also inform behavioral priors or skills [Ajay et al., 2020, Tirumala et al., 2020, Pertsch et al., 2021, Nasiriany et al., 2022]; once learned, these learned behaviors can provide a bootstrap for exploration [Uchendu et al., 2022, Li et al., 2023], or to transform the agent's action space to a higher level of abstraction [Pertsch et al., 2020, Singh et al., 2020a]. In addition, if the prior data consists of optimal demonstrations for the downstream task, there exists a rich literature for steering online RL with demonstrations [Schaal, 1996, Vecerik et al., 2017, Nair et al., 2017, Zhu et al., 2020]. Finally, for model-based RL approaches, one in theory may additionally train a dynamics model on prior data and use it to plan and reach novel states [Sekar et al., 2020, Mendonca et al., 2021] for improved exploration, but we are not aware of works that explicitly use models to explore from reward-free offline prior data.
Using prior data with reward labels. To train on unlabeled prior data with optimistic reward labels, we use RLPD [Ball et al., 2023], a sample-efficient online RL algorithm designed to train simultaneously on reward-labeled offline data and incoming online experience. RLPD is an offline-to-online algorithm [Xie et al., 2021, Lee et al., 2021, Agarwal et al., 2022, Zheng et al., 2023, Hao et al., 2023, Song et al., 2023] seeking to accelerate online RL using reward-labeled prior experience. While this setting shares many challenges with our unlabeled prior data, algorithms in this setting tend to focus on correcting errors in value function stemming from training on offline data [Nakamoto et al., 2023], and correctly balancing between using highly off-policy transitions from the prior data and new online experience [Lee et al., 2021, Luo et al., 2023]. Our optimistic reward labeling mechanism is largely orthogonal to these axes, since our approach is compatible with any online RL algorithm that can train on prior data.
Online unsupervised RL. Online unsupervised RL is an adjacent setting in which agents collect action-free data to accelerate finetuning on some downstream task [Laskin et al., 2021]. While both deal with reward-free learning, they crucially differ in motivation: our setting asks how existing (pre-collected) data can be used to inform downstream exploration for a known task, while unsupervised RL studies how to actively collect new datasets for a (yet to be specified) task [Yarats et al., 2022, Brandfonbrener et al., 2022]. Methods for unsupervised RL seek broad and undirected exploration strategies, for example through information gain [Sekar et al., 2020, Liu and Abbeel, 2021] or state entropy maximization [Lee et al., 2019, Jin et al., 2020, Wang et al., 2020] – instead of the task-relevant exploration that is advantageous in our setting. Successful mechanisms developed for unsupervised RL can inform useful strategies for our setting, since both seek to explore in directed ways, whether undirected as in unsupervised RL or guided by prior data, as in ours.
Offline meta reinforcement learning. Offline meta-reinforcement learning [Dorfman et al., 2021, Pong et al., 2021, Lin et al., 2022] also seeks to use prior datasets from a task distribution to accelerate learning in new tasks sampled from the same distribution. These approaches generally tend to focus on the single-shot or few-shot setting, with methods that attempt to (approximately) solve
the Bayesian exploration POMDP [Duan et al., 2016, Zintgraf et al., 2019, Rakelly et al., 2019]. While these methods are efficient when few-shot adaptation is possible, they struggle in scenarios requiring more complex exploration or when the underlying task distribution lacks structure [Ajay et al., 2022, Mandi et al., 2022]. Meta RL methods also rely on prior data containing meaningful, interesting tasks, while our optimistic relabeling strategy is agnostic to the prior task distribution.
Our experiments are aimed at evaluating how effectively our method can improve exploration by leveraging prior data. Through extensive experiments, we will establish that our method is able to leverage the unlabeled prior data to boost online sample efficiency consistently, often matching the sample efficiency of the previous best approach that uses labeled prior data (whereas we have no access to labeled rewards). More concretely, we will answer the following questions:
1. Can we leverage unlabeled prior data to accelerate online learning?
2. Can representation learning help obtain better reward labels?
3. Does optimistic reward labeling help online learning by improving online exploration?
4. How robust is our method in handling different offline data corruptions?
5.1 Experimental Setup
The environments that we evaluate our methods on are all challenging sparse reward tasks, including six D4RL AntMaze tasks [Fu et al., 2020], three sparse-reward Adroit hand manipulation tasks [Nair et al., 2021] following the setup in RLPD [Ball et al., 2023], and three image-based robotic manipulation tasks used by COG [Singh et al., 2020b]. See Appendix B for details about each task. For all three domains, we explicitly remove the reward information from the offline dataset so that the agent must explore to solve the task. We choose these environments as they cover a range of different characteristics of the prior data, where each can present a unique exploration and learning challenge. In Adroit, the prior data consists of a few expert trajectories on the target task and a much larger set of data from a behavior cloning policy. In AntMaze, the prior data consists of goal-directed trajectories with varying start and end locations. In COG, the prior data is mostly demonstration trajectories for sub-tasks of a larger multi-stage task. These tasks are interesting for the following reasons:
Sparse reward (All). Due to the sparse reward nature of the these tasks, the agent has no knowledge of the task, it must explore to solve the task successfully. For example, on AntMaze tasks, the goal is for an ant robot to navigate in a maze to reach a fixed goal from a fixed starting point. The prior data for this task consists of sub-optimal trajectories, using starting points and goals potentially different from the current task. Without knowledge of the goal location, the agent must explore each corner of the maze to identify the actual goal location before solving the task is possible.
High-dimensional image observation (COG). The three COG tasks use image observations, which poses a much harder exploration challenge compared to state observations because of high dimensionality. As we will show, even with such high-dimensional observation spaces, our algorithm is still able to efficiently explore and accelerate online learning more than the baselines. More importantly, our algorithm can incorporate pre-training representations that further accelerates learning.