SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Multi-stage exploration (COG). The robotic manipulation tasks in COG have two stages, meaning that the robot needs to perform some sub-task (e.g., opening a drawer) before it is able to complete the full task (e.g., picking out a ball from the drawer). See Appendix B, Figure 9 for visualizations of images from the COG dataset showing the two stage structure.
The prior data consists only of trajectories performing individual stages by themselves, and there is no prior data of complete task executions. The agent must learn to explore past the first stage of the task without any reward information, before it can receive the reward for completing the full task. Therefore, this multi-stage experiment requires the agent to be able to effectively make use of incomplete trajectories to aid exploration. This setting is closer to the availability of the data one may encounter in the real world. For example, there may be a lot of data of robots performing tasks that are only partially related to the current task of interest. This data could also be unlabeled, as the rewards for other tasks may not match the current task. Despite being incomplete, we would still like to be able to utilize this data.
5.2 Comparisons
While most prior works do not study the setting where the prior data is unlabeled, focusing on either the standard fully labeled data setting or utilizing an online unsupervised learning phase, we can adapt some of these prior works as comparisons in our setting (see Appendix A for details):
Online: A data-efficient off-policy RL agent that does not make use of the prior data at all.
Online + RND: An augmentation of online RL with RND as a novelty bonus. This baseline uses only online data, augmenting the online batch with RND bonus [Burda et al., 2018]. To be clear, of the comparisons listed here in bold, this is the only one that uses an online RND bonus.
Naïve Reward Labeling: This comparison labels the prior data reward with an unbiased reward estimate; this is implemented using our method but omitting the RND novelty score.
Naïve + BC: This baseline additionally uses a behavioral cloning loss to follow the behaviors as seen in prior data, inspired by similar regularization in offline RL [Vecerik et al., 2017].
MinR: This is an adaptation of UDS [Yu et al., 2022] to the online fine-tuning setting. The original UDS method uses unlabeled prior data to improve offline RL on a smaller reward-labeled dataset, using the minimum reward of the task to relabel the unlabeled data. MinR uses the same labeling strategy, but with RLPD for online RL.
$\mathbf{BC} + \mathbf{JSRL}$ : This baseline is an adaptation of JSRL [Uchendu et al., 2022]. The original JSRL method uses offline RL algorithms to pre-train a guide policy using a fully labeled prior dataset. Then, at each online episode, it uses the guide policy to roll out the trajectory up to a random number of steps, then switches to an exploration policy such that the initial state distribution of the exploration policy is shaped by the state visitation of the guide-policy, inducing faster learning of the exploration policy. Since in our setting the prior data has no labels, we use behavior cloning (BC) pre-training to initialize the guide-policy instead of using offline RL algorithms.
Oracle: This is an oracle baseline that assumes access to ground truth reward labels, using the same base off-policy RL algorithm [Ball et al., 2023] with true reward labels on prior data.
We also include a behavior prior baseline (inspired by the behavior prior learning line of work [Ajay et al., 2020, Tirumala et al., 2020, Pertsch et al., 2021, Nasiriany et al., 2022, Singh et al., 2020a]) comparison in Appendix C.5 on AntMaze.
5.3 Does optimistic labeling of prior data accelerate online learning?
Figure 3 shows the aggregated performance of our approach on the two state-based domains and one image-based domain. On AntMaze and COG domains, our optimistic reward labeling is able to outperform naive reward labeling significantly, highlighting its effectiveness in leveraging the prior data to accelerate online learning. Interestingly, in the two state-based domains, we show that without access to prior rewards, our method can nearly match the performance of the oracle baseline with prior rewards. This suggests that optimistic reward labeling can allow RL to utilize unlabeled prior data almost as effectively as labeled prior data on these three domains. In particular, on the sparse Adroit relocate task, using optimistic rewards can even outpace the oracle baseline (Appendix C.2, Figure 13). In the COG domain, there is a larger performance gap between our method and the oracle. We hypothesize that this is due to the high-dimensional image observation space, causing the exploration problem to be more challenging. On the Adroit domain, we find that periodic resetting of the learned reward function is necessary to avoid overfitting to the collected online experience – with the resetting, we find that even naive reward relabeling results in exploratory behavior, diminishing the gap when performing explicit reward uncertainty estimation with RND.
5.4 Can representation learning help obtain better reward labels?
We now investigate whether we can improve exploration by basing our UCB reward estimates on top of pre-trained representations acquired from the prior data. We chose to use representations from ICVF [Ghosh et al., 2023], a method that trains feature representations by pre-training general successor value functions on the offline data. Value-aware pre-training is a natural choice for downstream reward uncertainty quantification, because the pre-training encourages representations to obey the spatiality of the environment – that states close to one another have similar representa
Environment Steps  $(\times 10^{6})$
Figure 4: Accelerating exploration with pre-trained representations on three visual-input COG tasks. Ours + ICVF uses optimistic reward labeling with a pre-trained representation initialization; Ours uses optimistic reward labeling without the pre-trained representation initialization. The same applies for Naive and Naive + ICVF. Overall, initializing the reward model and the RND network using pre-trained representations greatly increases how quickly the model learns.
Figure 5: The effect of pre-trained representations on the relabeled reward. We show the normalized relabeled rewards (by standardizing it to have zero mean and unit variance) of an optimal prior trajectory from the closed drawer task in the COG environment. Using pre-trained representations to learn the optimistic rewards leads to smoother rewards over the course of an optimal trajectory. See more examples in Appendix D.
tions, while states that are difficult to reach from one another are far away. We pre-train an ICVF model on the prior dataset and extract the visual encoder  $\xi_{\mathrm{ICVF}}(s)$ , which encodes an image into a low-dimensional feature vector. These pre-trained encoder parameters are used to initialize the encoder for both the reward model  $r_{\theta}(\xi_{\theta}(s))$  and the RND network  $f_{\phi}(\xi_{\phi}(s))$ . Note that these encoder parameters are not shared and not frozen, meaning that once training begins, the encoder parameters for the reward model will be different from the encoder parameters for the RND model, and also diverge from the pre-trained initialization. As shown in Figure 4, ICVF consistently improves the online learning efficiency across all image-input tasks on top of optimistic reward labeling. It is worth noting that the encoder initialization from pre-training alone is not sufficient for the learning agent to succeed online (Naive + ICVF still fails almost everywhere). On most tasks, without optimistic reward labeling, the agent simply makes zero progress, further highlighting the effectiveness of optimistic labeling in aiding online exploration and learning. To understand how the pre-trained representation initialization influences online exploration, we analyze its effect on the UCB reward estimate early in the training. Figure 5 shows the UCB reward estimates on an expert trajectory of the Grasp from Closed Printer task when the reward model and the RND model are initialized with the pre-trained representations (in blue) and without the pre-trained representations (in orange). When our algorithm is initialized with the ICVF representations, the UCB reward estimates are increasing along the trajectory time step, labeling the images in later part of the trajectory with high reward, forming a natural curriculum for the agent to explore towards the trajectory end. When our algorithm is initialized randomly without the pre-trained representations, such a clear monotonic trend disappears. We hypothesize that this monotonic trend of the reward labeling may account for the online sample efficiency improvement with a better shaped reward.
Figure 6: State coverage for 3 AntMaze tasks, estimated by counting the number of square regions that the ant has visited, normalized by the total number of square regions that can be visited. Ours optimistically labels the prior data, whereas Naïve naively labels the prior data without optimism. Optimistic labeling enables the RL agent to achieve a significantly higher coverage. Online RND also explores well in the maze, but with a much slower rate compared to our approach, highlighting the importance of leveraging unlabeled prior data for efficient online exploration. Full results in Appendix C.1.
Figure 7: The normalized return on antmaze-large-diverse-v2 under different offline data corruptions (see visualization of the offline data for each kind in Appendix B, Figure 8).
5.5 Does optimistic reward labeling help online learning by improving online exploration?
We have shown that optimistic reward labeling is effective in accelerating online learning, does it actually work because of better exploration? To answer this question, we evaluate the effect of our method on the state-coverage of the agent in the AntMaze domain. In this domain, the agent is required to reach a particular goal in the maze but has no prior knowledge of the goal, and thus required to explore different parts of the maze in the online phase in order to locate the goal. The optimal behavior that we expect the agent to perform is to quickly try out every possible part of the maze until the goal is located, and then focus entirely on that goal to achieve  $100\%$  success rate.
To more quantitatively measure the quality of our exploration behavior, we divide up the 2-D maze into a grid of square regions, and count the number of regions that the Ant agent has visited (using the online replay buffer). This provides a measure for the amount of state coverage that the agent has achieved in the environment, and the agent that explores better should achieve higher state coverage. Figure 6 shows how optimistic reward labeling affects the state coverage, and Figure 2 shows an example visualization of the state-visitation for each method on antmaze-medium. Across all six AntMaze tasks, our approach achieves significantly higher state-coverage compared to naively label without optimism both in the early stage of training and asymptotically compared to the baseline that does not leverage the unlabeled prior data. In particular, we find that naively labeling reward with no optimism is doing poorly on larger mazes, highlighting the importance of optimistic reward labeling for effective online exploration.
5.6 How robust is our method in handling different offline data corruptions?
To further test the capability of our method in leveraging prior data, we corrupt one of the D4RL offline datasets, antmaze-large-diverse-v2, and test how our method performs under the corruption. We experimented with three different corruptions (and reported results in Figure 7).
Orthogonal transitions: all the transitions that move in the up/right (towards the goal) are removed from the offline data. In this setting, stitching the transitions in the offline data would not lead to a trajectory that goes from the start to the goal. This is a good test for the robustness of the algorithm for exploration because the agent can not rely on the actions in the prior data and must explore to discover actions that move the ant in the correct direction.
Insufficient coverage: all the transitions around the goals are removed from the offline data. This requires the algorithm to actively explore so that it discovers the desired goal that it needs to reach.
1% data: A random 1% of the transitions are kept in the offline data. This tests the capability of our method in handling the limited data regime.
In addition to the main comparisons above, we consider two modifications specifically for these experiments to address the more challenging tasks.
Ours + Online RND: This is an extension of our method, with an RND bonus added to the online data in addition to the prior data. Similar to how we can extend our method with representation learning, having exploration bonuses on the online data does not conflict with exploration bonuses added to the offline data.
Naïve + Online RND: This is an extension of Naïve Reward Labeling. An RND bonus added to the online data, and the prior data is labeled with the unbiased reward estimate as before.
Directly using our method in the Insufficient Coverage setting yields zero success, but this is to be expected as our algorithm does not incentivize exploration beyond the offline data. We find that this can be patched by simply combining our method with an additional RND bonus on online data, allowing our method to learn even when data has poor coverage and no data near the goal is unavailable. Similarly, we have found that combining our method with online RND can also boost the performance in the Orthogonal Transitions, where no transitions move in the direction of the goal. For the  $1\%$  Data setting, our method can still largely match the oracle performance without the need for an online RND bonus.
It is also worth noting that in the Orthogonal Transitions setting, our method achieves good performance whereas the oracle fails to succeed. The failures of the oracle are not unexpected because the prior data do not contain any single trajectories, or even trajectories stitched together from different transitions, that can reach the goal. However, since our method succeeds, this indicates that exploration bonuses applied only to the prior data can still be utilized, despite there being no forward connection through the prior data. These additional results, while initially surprising even to us (especially the "orthogonal" setting, which seems very difficult), strongly suggest that our method can lead to significant improvement even when the data is not very good - certainly such data would be woefully inadequate for regular offline RL or naïve policy initialization.
6 Discussion
We showed how we can effectively leverage unlabeled prior data to improve online exploration by running a standard off-policy RL algorithm on the data relabeled with UCB reward estimates. In practice, the UCB estimates can be approximated by combining the prediction of a reward model and an RND network, both can be trained online with little additional computational cost. We demonstrated the surprising effectiveness of this simple approach on a diverse set of domains. Our instantiation of the optimistic reward labeling idea presents a number of avenues for future research. First, on the relocate Adroit task, we found that naively fitting the reward model on the online replay buffer without any regularization leads to poor performance (possibly due to catastrophic overfitting). While we have found a temporary workaround by periodically re-initializing the reward model, such a solution may seem ad-hoc and could be disruptive to the RL learning due to the sudden change in the learning objective (which is shaped by the reward model). Next, the UCB estimates of the state-action can rapidly change, especially in the beginning of the training, which may cause learning instability in the RL algorithm. There is no mechanism in our current algorithm to tackle such a potential issue. Nevertheless, our work indicates that there exist simple ways of incorporating prior data for online RL agents even when no reward labels exist. Scaling these mechanisms to more complex prior datasets is an exciting direction towards the promise of RL agents that can leverage general prior data.
Acknowledgement. This work was partially done while QL was a visiting student researcher at FAIR, Meta. We would like to thank Seohong Park for providing his implementation of goal-conditioned IQL (GC-IQL) and pre-trained GC-IQL checkpoints (which were used in producing Figure 17). We would like to thank Fangchen Liu for the acronym of the method. We would also like to thank Laura Smith, Toru Lin, Oleg Rybkin, Kuba Grudzien, Aviral Kumar, and Joey Hong for discussion on the method and feedback on the early draft of the paper. We would also like to thank the members of the RAIL lab for insightful discussions on the paper. This research was partially supported by the Office of Naval Research under N00014-21-1-2838 and N00014-22-1-2773, and ARO W911NF-21-1-0097.
Joshua Achiam and S. Shankar Sastry. Surprise-based intrinsic motivation for deep reinforcement learning. *ArXiv*, abs/1703.01732, 2017.
Abien Fred Agarap. Deep learning using rectified linear units (ReLU). arXiv preprint arXiv:1803.08375, 2018.
Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C. Courville, and Marc G. Bellemare. Reincarnating reinforcement learning: Reusing prior computation to accelerate progress. In Neural Information Processing Systems, 2022.
Anurag Ajay, Aviral Kumar, Pulkit Agrawal, Sergey Levine, and Ofir Nachum. OPAL: Offline primitive discovery for accelerating offline reinforcement learning. arXiv preprint arXiv:2010.13611, 2020.
Anurag Ajay, Abhishek Gupta, Dibya Ghosh, Sergey Levine, and Pulkit Agrawal. Distributionally adaptive meta reinforcement learning. ArXiv, abs/2210.03104, 2022.
Mohammad Gheshlaghi Azar, Ian Osband, and Rémi Munos. Minimax regret bounds for reinforcement learning. In International Conference on Machine Learning, pages 263-272. PMLR, 2017.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.06450, 2016.
Philip J Ball, Laura Smith, Ilya Kostrikov, and Sergey Levine. Efficient online reinforcement learning with offline data. arXiv preprint arXiv:2302.02948, 2023.
Marc Bellemare, Sriram Srinivasan, Georg Ostrovski, Tom Schaul, David Saxton, and Remi Munos. Unifying count-based exploration and intrinsic motivation. In Advances in Neural Information Processing Systems, pages 1471-1479, 2016.
David Brandfonbrener, Stephen Tu, Avi Singh, Stefan Welker, Chad Boodoo, N. Matni, and Jacob Varley. Visual backtracking teleoperation: A data collection protocol for offline image-based reinforcement learning. ArXiv, abs/2210.02343, 2022.
Yuri Burda, Harrison Edwards, Amos Storkey, and Oleg Klimov. Exploration by random network distillation. arXiv preprint arXiv:1810.12894, 2018.
Christoph Dann, Tor Lattimore, and Emma Brunskill. Unifying PAC and regret: Uniform PAC bounds for episodic reinforcement learning. In NIPS, 2017.
Ron Dorfman, Idan Shenfeld, and Aviv Tamar. Offline meta reinforcement learning - identifiability challenges and effective data collection strategies. In Neural Information Processing Systems, 2021.
Yan Duan, John Schulman, Xi Chen, Peter L. Bartlett, Ilya Sutskever, and P. Abbeel. RL<sup>2</sup>: Fast reinforcement learning via slow reinforcement learning. ArXiv, abs/1611.02779, 2016.
Jesse Farebrother, Joshua Greaves, Rishabh Agarwal, Charline Le Lan, Ross Goroshin, Pablo Samuel Castro, and Marc G Bellemare. Proto-value networks: Scaling representation learning with auxiliary tasks. arXiv preprint arXiv:2304.12567, 2023.
Bradly C Stadie, Sergey Levine, and Pieter Abbeel. Incentivizing exploration in reinforcement learning with deep predictive models. arXiv preprint arXiv:1507.00814, 2015.
Alexander L. Strehl and Michael L. Littman. An analysis of model-based interval estimation for markov decision processes. J. Comput. Syst. Sci., 74:1309-1331, 2008.
Haoran Tang, Rein Houthooft, Davis Foote, Adam Stooke, OpenAI Xi Chen, Yan Duan, John Schulman, Filip DeTurck, and Pieter Abbeel. # Exploration: A study of count-based exploration for deep reinforcement learning. In Advances in neural information processing systems, pages 2753-2762, 2017.
Dhruva Tirumala, Alexandre Galashov, Hyeonwoo Noh, Leonard Hasenclever, Razvan Pascanu, Jonathan Schwarz, Guillaume Desjardins, Wojciech M. Czarnecki, Arun Ahuja, Yee Whye Teh, and Nicolas Manfred Otto Heess. Behavior priors for efficient reinforcement learning. *ArXiv*, abs/2010.14274, 2020.
Ikechukwu Uchendu, Ted Xiao, Yao Lu, Banghua Zhu, Mengyuan Yan, Josephine Simon, Matthew Bennice, Chuyuan Fu, Cong Ma, Jiantao Jiao, et al. Jump-start reinforcement learning. arXiv preprint arXiv:2204.02372, 2022.
Mel Vecerik, Todd Hester, Jonathan Scholz, Fumin Wang, Olivier Pietquin, Bilal Piot, Nicolas Heess, Thomas Rothörl, Thomas Lampe, and Martin Riedmiller. Leveraging demonstrations for deep reinforcement learning on robotics problems with sparse rewards. arXiv preprint arXiv:1707.08817, 2017.
Ruosong Wang, Simon Shaolei Du, Lin F. Yang, and Ruslan Salakhutdinov. On reward-free reinforcement learning with linear function approximation. *ArXiv*, abs/2006.11274, 2020.
Tete Xiao, Ilija Radosavovic, Trevor Darrell, and Jitendra Malik. Masked visual pre-training for motor control. arXiv preprint arXiv:2203.06173, 2022.
Tengyang Xie, Nan Jiang, Huan Wang, Caiming Xiong, and Yu Bai. Policy finetuning: Bridging sample-efficient offline and online reinforcement learning. In Neural Information Processing Systems, 2021.
Mengjiao Yang and Ofir Nachum. Representation matters: offline pretraining for sequential decision making. In International Conference on Machine Learning, pages 11784-11794. PMLR, 2021.
Denis Yarats, David Brandfonbrener, Hao Liu, Michael Laskin, P. Abbeel, Alessandro Lazaric, and Lerrel Pinto. Don't change the algorithm, change the data: Exploratory data for offline reinforcement learning. ArXiv, abs/2201.13425, 2022.
Tianhe Yu, Aviral Kumar, Yevgen Chebotar, Karol Hausman, Chelsea Finn, and Sergey Levine. How to leverage unlabeled data in offline reinforcement learning, 2022.
Han Zheng, Xufang Luo, Pengfei Wei, Xuan Song, Dongsheng Li, and Jing Jiang. Adaptive policy learning for offline-to-online reinforcement learning. *ArXiv*, abs/2303.07693, 2023.
Zhuangdi Zhu, Kaixiang Lin, Bo Dai, and Jiayu Zhou. Learning sparse rewarded tasks from suboptimal demonstrations. ArXiv, abs/2004.00530, 2020.
Luisa M. Zintgraf, Kyriacos Shiarlis, Maximilian Igl, Sebastian Schulze, Yarin Gal, Katja Hofmann, and Shimon Whiteson. Varibad: A very good method for bayes-adaptive deep rl via meta-learning. ArXiv, abs/1910.08348, 2019.
A Experiment Setup Details
Our codebase is based on the official RLPD codebase https://github.com/ikostrikov/rlpd with minor modifications. All configurations in our experiments use 10 seeds (20 seeds for all COG experiments) and we compute the standard error to plot the confidence interval in all our plots.
A.1 Architecture Details
RLPD extends a standard off-policy actor-critic algorithm which trains an actor network and a critic network simultaneously. Our extension of the RLPD requires training three additional networks: a reward prediction network, a termination prediction network and an RND network. For state-based domains, each of the actor, critic, reward and RND network is multilayer perceptron (MLP) with ReLU [Agarap, 2018] activation. For the COG (image-based) domain, each network uses a pixel encoder followed by a MLP.