SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Boyuan Chen*
boyuanc@mit.edu
Chuning Zhu*
Pulkit Agrawal
pulkitag@mit.edu
Kaiqing Zhang†
Abhishek Gupta†
Model-free reinforcement learning algorithms have exhibited great potential in solving single-task sequential decision-making problems with high-dimensional observations and long horizons, but are known to be hard to generalize across tasks. Model-based RL, on the other hand, learns task-agnostic models of the world that naturally enables transfer across different reward functions, but struggles to scale to complex environments due to the compounding error. To get the best of both worlds, we propose a self-supervised reinforcement learning method that enables the transfer of behaviors across tasks with different rewards, while circumventing the challenges of model-based RL. In particular, we show self-supervised pretraining of model-free reinforcement learning with a number of random features as rewards allows implicit modeling of long-horizon environment dynamics. Then, planning techniques like model-predictive control using these implicit models enable fast adaptation to problems with new reward functions. Our method is self-supervised in that it can be trained on offline datasets without reward labels, but can then be quickly deployed on new tasks. We validate that our proposed method enables transfer across tasks on a variety of manipulation and locomotion domains in simulation, opening the door to generalist decision-making agents.
As in most machine learning problems, the ultimate goal of building deployable sequential decision-making agents via reinforcement learning (RL) is its broad generalization across tasks in the real world. While reinforcement learning algorithms have been shown to successfully synthesize complex behavior in single-task sequential decision-making problems [31, 36, 45], their performance as generalist agents across tasks has been less convincing [9]. In this work, we focus on the problem of learning generalist agents that are able to transfer across problems where the environment dynamics are shared, but the reward function is changing. This problem setting is reflective of scenarios
that may be encountered in real-world settings such as robotics. For instance, in tabletop robotic manipulation, different tasks like pulling an object, pushing an object, picking it up, and pushing to different locations, all share the same transition dynamics, but involve different reward functions. We hence ask the question - Can we reuse information across tasks in a way that scales to high dimensional, longer horizon problems?
Figure 1: RaMP is a self-supervised reinforcement learning algorithm that pre-training on exploration trajectories without (or with unknown) reward. It then quickly adapts to new tasks through interaction with the environment.
A natural possibility to tackle this problem is direct policy search [54, 44]. Typical policy search algorithms can achieve good performance on a single task. However, the policy is optimized for a particular reward and may be highly suboptimal in new scenarios. Other model-free RL algorithms like actor-critic methods [23, 15, 13] or Q-learning [51, 31] may exacerbate this issue, with learned Q-functions entangling dynamics, rewards, and policies. For new scenarios, an ideal algorithm should be able to disentangle and retain the elements of shared dynamics, allowing for easy substitution of rewards without having to repeat significant computations. While multi-task policy search techniques like goal-conditioned RL address this issue to some extent, they are restricted to particular tasks like goal-reaching and usually require prior knowledge about the tasks one may encounter in testing.
Model-based RL arises as a natural fit for disentangling dynamics and rewards [33, 10, 17, 19, 20]. These algorithms directly learn a model of transition dynamics and leverage the learned model to plan control actions [35, 17, 10, 8]. In particular, the models can be used to re-plan behaviors for new rewards. However, these learned dynamics models are brittle and suffer from compounding error [25, 3] due to their autoregressive nature. States predicted by the model are iteratively used as inputs for subsequent predictions, leading to compounding approximation errors during planning [25]. This makes it challenging to apply them to problems with long horizons and high dimensions.
In this work, we thus ask - Can we build RL algorithms that disentangle dynamics, rewards, and policies for transfer across problems, while retaining the ability to solve problems with high dimensional observations and long horizons? To this end, we propose a self-supervised RL algorithm that can leverage collected data to implicitly model transition dynamics without ever having to generate future states, and then use this to quickly transfer to a variety of different new tasks with varying reward functions that may be encountered at test time.
Our method relies on the following key insight: to prevent tying specific reward functions and policies to the RL algorithm, one can instead model the long-term evolution of randomly chosen basis functions of all possible reward functions [6]. Such basis functions can be generated using random features [39-41]. The accumulation of such basis functions over time, which we refer to as a Q-basis, is task-agnostic [48], and implicitly models the transition dynamics of the problem. We show that with a large enough number of random features and model-free RL, the value function for any reward function that may be encountered in testing could be estimated, via a simple linear combination of these Q-basis functions. This is important since it allows us to obtain the benefits of transfer, without explicitly learning a dynamics model that may suffer from compounding error. Moreover, with randomly chosen basis functions that are agnostic of downstream tasks, this is not restricted to just a certain distribution of tasks, as in most multi-task or goal-conditioned RL [49, 2]. Finally, we note that the random features (as well as the corresponding Q-basis functions) can be generated using datasets that are not necessarily labeled with reward signals. Indeed, these random-feature-based pseudo rewards can be viewed as self-supervised signals in training, making the approaches based on them naturally compatible with large (task-agnostic) offline datasets.
Based on the aforementioned insights, we propose a new algorithm Random Features for Model-Free Planning (RaMP) that allows us to leverage unlabelled offline datasets to learn reward-agnostic Q-bases. These can be used to estimate the test-time value functions for new reward functions using linear regression, enabling quick adaptation to new tasks under the same shared transition dynamics. We show the efficacy of this method on a number of tasks for robotic manipulation and locomotion in simulation, and highlight how RaMP provides a more general paradigm than typical generalizations of model-based or model-free RL.
Model-based RL is naturally suited for transfer across tasks with shared dynamics, by explicitly learning a model of the transition dynamics, which is disentangled with the reward functions for each task [17, 35, 10, 16, 8, 47, 27]. These models are typically learned via supervised learning on one-step transitions and then used to extract control actions via planning [28, 42] or trajectory optimization [46, 59, 38]. The key challenge in scaling lies in the fact that they sequentially feed model predictions back into the model for sampling [50, 3, 25]. This can often lead to compounding errors [25, 3, 55], which grow with the horizon length unfavorably. In contrast, our work does not require autoregressive sampling, making it easier to scale to longer horizons and higher dimensions.
On the other hand, model-free RL avoids the challenge of compounding error by directly modeling either policies or Q-values without autoregressive generation [54, 43, 44, 30, 31, 15]. However, these approaches entangle rewards, dynamics, and policies, making them hard to transfer. While attempts have been made to build model-free methods that generalize across rewards, such as goal-conditioned value functions [21, 1, 37, 14] or multi-task policies [22, 18], they only apply to restricted classes of reward functions and particular training distributions. Our work aims to obtain the best of both worlds (model-based and model-free RL), learning some representation of dynamics independent of rewards and policies, while using a model-free algorithm for learning.
Our notion of long-term dynamics is connected to the notion of state-action occupancy measure [34, 58], often used for off-policy evaluation and importance sampling methods in RL. These methods often try to directly estimate either densities or density ratios [20, 34, 58]. Our work simply learns the long-term accumulation of random features, without requiring any notion of normalized densities. Perhaps most closely related work to ours is the framework of successor features, that considers transfer from a fixed set of source tasks to new target tasks [5, 26, 57, 11]. Like our work, the successor features framework leverages the linearity of rewards to disentangle dynamics from rewards using model-free RL. However, transfer using successor features is dependent on choosing (or learning) the right featurization and forces an implicit dependence on the policy. Our work leverages random features and multi-step Q-functions to alleviate the transfer performance of successor features.
Model. We consider the standard Markov decision process (MDP) as characterized by a tuple  $\mathcal{M} = (\mathcal{S},\mathcal{A},\mathcal{T},R,\gamma ,\mu)$ , with state space  $\mathcal{S}$ , action space  $\mathcal{A}$ , transition dynamics  $\mathcal{T}:S\times \mathcal{A}\to \Delta (\mathcal{S})$ , reward function  $R:\mathcal{S}\times \mathcal{A}\rightarrow \Delta ([-R_{\max},R_{\max}])$ , discount factor  $\gamma \in [0,1)$ , and initial state distribution  $\mu \in \Delta (\mathcal{S})$ . The goal is to learn a policy  $\pi :S\to \Delta (\mathcal{A})$ , such that it maximizes the expected discounted accumulated rewards, i.e., solves  $\max_{\pi}\mathbb{E}_{\pi}\left[\sum_{h = 1}^{\infty}\gamma^{h - 1}r_h\right]$  with  $r_h\coloneqq r(s_h,a_h)\sim R_{s_h,a_h} = \operatorname *{Pr}(\cdot |s_h,a_h)$ . Hereafter, we will refer to an MDP and a task interchangeably. Note that in our problem settings, different MDPs will always share transition dynamics  $\mathcal{T}$ , but will have varying reward functions  $R$ .
Q-function Estimation. Given an MDP  $\mathcal{M}$ , one can define the state-action  $Q$ -value function under any policy  $\pi$  as  $Q_{\pi}(s,a) := \mathbb{E}_{a_h \sim \pi(\cdot \mid s_h)} \left[ \sum_{h=1}^{\infty} \gamma^{h-1} r_h \mid s_1 = s, a_1 = a \right]$  which denotes
the expected accumulated reward under policy  $\pi$ , when starting from state-action pair  $(s, a)$ . By definition, this  $Q$ -function is inherently tied to the particular reward function  $R$  and the policy  $\pi$ , making it challenging to transfer for a new reward or policy. Similarly, one can also define the multi-
$$
step\left(\tau -step\right)Q - function\quad Q_{\pi}\left(s,\widetilde{a}_{1},\widetilde{a}_{2},\dots ,\widetilde{a}_{\tau}\right) = \mathbb{E}_{\substack{a_{\tau +h}\sim \pi (\cdot \mid s_{\tau +h})\\ s_{h + 1}\sim \mathcal{T}(\cdot \mid s_{h},a_{h})}}\Big[\sum_{h = 1}^{\infty}\gamma^{h - 1}r_{h}\big|s_{1} = s,a_{1} =
$$
$$
\left. \widetilde {a} _ {1}, a _ {2} = \widetilde {a} _ {2}, \dots , a _ {\tau} = \widetilde {a} _ {\tau} \right].
$$
One can estimate the  $\tau$ -step  $Q_{\pi}$  by Monte-Carlo sampling of the trajectories under  $\pi$ , i.e., by solving
$$
\min  _ {\widehat {Q} \in \mathcal {Q}} \frac {1}{N} \sum_ {j = 1} ^ {N} \left\| \widehat {Q} \left(s, \widetilde {a} _ {1} ^ {j}, \widetilde {a} _ {2} ^ {j}, \dots , \widetilde {a} _ {\tau} ^ {j}\right) - \frac {1}{M} \sum_ {m = 1} ^ {M} \sum_ {h = 1} ^ {\infty} \gamma^ {h - 1} r _ {h} ^ {m, j} \right\| _ {2} ^ {2}, \tag {2.1}
$$
where  $\mathcal{Q}$  is some function class for  $Q$ -value estimation, which in practice is some parametric function class, e.g., neural networks;  $r_h^{m,j} \sim R_{s_h^{m,j},a_h^{m,j}}$  and  $(s_h^{m,j},a_h^{m,j})$  come from  $MN$  trajectories that
are generated by  $N$  action sequences  $\{(\widetilde{a}_1^j,\widetilde{a}_2^j,\dots ,\widetilde{a}_\tau^j)\}_{j = 1}^N$  and  $M$  trajectories following policy  $\pi$  after each action sequence. A large body of work considers finding this  $Q$ -function using dynamic programming, but for the sake of simplicity, this work will only consider Monte-Carlo estimation.
In practice, the infinite-horizon estimator in (2.1) can be hard to obtain. We hence use a finite-horizon approximation of  $Q_{\pi}$  (of length  $H$ ), denoted by  $Q_{\pi}^{H}$ , in learning. We will treat this finite-horizon formula in the main paper for practical simplicity. We also provide a method to deal with the infinite-horizon case directly by bootstrapping the value function estimate. We defer the details to Appendix B.1. Note that if one chooses  $H = \tau$ , then the  $\tau$ -step  $Q$ -function defined above becomes
$$
Q _ {\pi} ^ {H} (s, \widetilde {a} _ {1}, \widetilde {a} _ {2}, \dots , \widetilde {a} _ {H}) := \mathbb {E} _ {s _ {h + 1} \sim \mathcal {T} (\cdot | s _ {h}, a _ {h})} \Big [ \sum_ {h = 1} ^ {H} \gamma^ {h - 1} r _ {h}   |   s _ {1} = s, a _ {1} = \widetilde {a} _ {1}, \dots , a _ {H} = \widetilde {a} _ {H} \Big ].
$$
Note that in this case, the  $Q$ -function is irrelevant of the policy  $\pi$ , denoted by  $Q^H$ , and is just the expected accumulated reward executing the action sequence  $(\widetilde{a}_1, \widetilde{a}_2, \dots, \widetilde{a}_H)$  starting from the state  $s$ . This Q-function can be used to score how "good" a sequence of actions will be, which in turn can be used for planning.
Problem setup. We illustrate our problem setup in Figure 1. Consider a transfer learning scenario, where we assume access to an offline dataset consisting of several episodes  $\mathcal{D} = \{(s_h^m, a_h^m, s_{h+1}^m)\}_{h \in [H], m \in [M]}$ . Here  $H$  is the length of the trajectories, which is large enough, e.g., of order  $\mathcal{O}(1/(1 - \gamma))$  to approximate the infinite-horizon setting;  $M$  is the total number of trajectories. This dataset assumes that all transitions are collected under the same transition dynamics  $\mathcal{T}$ , but otherwise does not require the labels, i.e., rewards, and may even come from different behavior policies. This is reminiscent of the offline RL problem statement, but offline RL datasets are typically labeled with rewards and are task-specific. In contrast, in this work, the pre-collected dataset is used to quickly learn policies for any downstream reward, rather than for one specific reward function. The goal is to make the best use of the dataset  $\mathcal{D}$ , and generalize the learned experience to improve the performance on a new task  $\mathcal{M}$ , with the same transition dynamics  $\mathcal{T}$  but an arbitrary reward function  $R$ . Note that unlike some related work [5, 4], we make no assumption on the reward functions of the MDPs that generate  $\mathcal{D}$ . The goal of the learning problem is to pre-train on the offline dataset such that we can enable fast (even zero-shot) adaptation to arbitrary new reward functions encountered at test time.
3 RaMP: Learning Implicit Models for Cross-Reward Transfer with Self-Supervised Model-Free RL
We now introduce our algorithm, Random Features for Model-Free Planning (RaMP), to solve the problem described in Section 2 – learning a model of long-term dynamics that enables transfer to tasks labeled with arbitrary new rewards, while making use of the advantages of model-free RL. Clearly, the problem statement we wish to solve requires us to (1) be able to solve tasks with long horizons and high-dimensional observations and (2) alleviate policy and reward dependence to be able to transfer computation across tasks. The success of model-free RL in circumventing compounding error from autoregressive generation raises the natural question: Is there a model-free approach that can mitigate the challenges of compounding error and can transfer across tasks painlessly? We answer this in the affirmative in the following section.
3.1 Key Idea: Implicit Model Learning with Random Feature Cumulants and Model-Free RL
The key insight we advocate is that we can avoid task dependence if we directly model the long-term accumulation of many random functions of states and actions (treating them as the rewards), instead of modeling the long-term accumulation of one specific reward as in typical Q-value estimation. Since these random functions of the state and actions are task-agnostic and uninformed of any specific reward function, they simply capture information about the transition dynamics of the environment. However, they do so without actually requiring autoregressive generative modeling, as is commonplace in model-based RL. Doing so can effectively disentangle transition dynamics from reward, and potentially allow for transfer across tasks, while still being able to make use of model-free RL methods. Each long-term accumulation of random features is referred to as an element of a "random" Q-basis, and can be learned with simple modifications to typical model-free RL algorithms. The key is to replace the Q-estimation of a single task-specific reward function with estimating the Q-functions of a set of task-agnostic random functions of the state and actions as a random Q-basis.
At training time, the offline dataset  $\mathcal{D}$  can be used to learn a set of "random" Q-basis functions for different random functions of the state and action. This effectively forms an "implicit model", as it carries information about how the dynamics propagate, without being tied to any particular reward function. At test time, given a new reward function, we can recombine Q-basis functions to effectively approximate the true reward-specific Q-function under any policy. This inferred Q-function can then be used for planning for the new task. As we will show in Section 3.4, this recombination can actually be done by a linear combination of Q-basis functions for a sufficiently rich class of random features, reducing the problem of test-time adaptation for new rewards to a simple linear regression problem. We detail the two phases of our approach in the following subsections.
Figure 2: RaMP: Depiction of our proposed method for transferring behavior across tasks by leveraging model-free learning of random features. At training time, Q-basis functions are trained on accumulated random features. At test time, adaptation is performed by solving linear regression and recombining basis functions, followed by online planning with MPC.
3.2 Offline Training: Self-Supervised Learning Random Q-functions
Without any prior knowledge about the downstream test-time rewards, the best that an agent can do is to model the evolution of the state (i.e., model system dynamics). The key insight we make is that the evolution of state can be implicitly represented by simply modeling the long-term accumulation of random features of state to obtain a set of Q-basis functions. Such a generation of random features is fully self-supervised. These functions can then be combined to infer the task-specific Q-function.
Given the lack of assumptions about the downstream task, the random features being modeled must be expressive and universal in their coverage. As suggested in [39-41], random features can be powerful in that most nonlinear functions can be represented as linear combinations of random features effectively. As we will show in Section 3.4, modeling the long-term accumulation of random features allows us express the value function for any reward as a linear combination of these accumulations of random features. In particular, in our work, we assume that the random features are represented as the output of  $K$  neural networks  $\phi(\cdot, \cdot; \theta_k): S \times \mathcal{A} \to \mathbb{R}$  with weights  $\theta_k \in \mathbb{R}^d$  and  $k \in [K]$ , where  $\theta_k$  are randomly sampled i.i.d. from some distribution  $p$ . Sampling  $K$  such weights  $\theta_k$  with  $k \in [K]$  yields a vector of scalar functions  $[\phi(\cdot, \cdot; \theta_k)]_{k \in [K]} \in \mathbb{R}^K$  for any  $(s, a)$ , which can be used as random features. To model the long-term accumulation of each of these random features, we note that they can be treated as reward functions (as can any function of state and action [48]) in model-free RL, and standard model-free policy evaluation to learn Q-functions can be reused to learn a set of  $K$  Q-basis functions, with each of them corresponding to the Q-value of a random feature.
We note that done naively, this definition of a Q-basis function is tied to a particular policy  $\pi$  that generates the trajectory. However, to transfer behavior one needs to predict the accumulated random features under new sequences of actions, as policy search for the new task is likely to involve evaluating a policy that is not the same as the training policy  $\pi$ . To allow the modeling of accumulated features that is independent of particular policies, we propose to learn multi-step Q-basis functions for each of the random features (as discussed in Section 2), which is explicitly dependent on an input sequence of actions  $(\widetilde{a}_1,\widetilde{a}_2,\dots ,\widetilde{a}_{\tau})$ . This can be used to search for optimal actions in new tasks.
To actually learn these Q-basis functions (one for each random feature), we use Monte-Carlo methods for simplicity. We generate a new dataset  $\mathcal{D}_{\phi}$  from  $\mathcal{D}$ , with  $\mathcal{D}_{\phi} = \{( (s_{1}^{m},a_{1:H}^{m}), \sum_{h\in [H]}\gamma^{h - 1}\phi (s_{h}^{m},a_{h}^{m};\theta_{k}))\}_{m\in [M],k\in [K]}$ . Here we use
$\sum_{h \in [H]} \gamma^{h-1} \phi(s_h^m, a_h^m; \theta_k)$  as the accumulated random features for action sequences  $\{a_1, \dots, a_H\}$  taken from state  $s_1$ . We then use  $K$  function approximators representing each of the  $K$  Q-basis functions, e.g., neural networks  $\psi(\cdot, \cdot; \nu_k): S \times \mathcal{A}^H \to \mathbb{R}$  for  $k \in [K]$ , to fit the accumulated random features. Specifically, we minimize the following loss  $\min_{\{\nu_k\}} \frac{1}{M} \sum_{m \in [M], k \in [K]} \left( \psi(s_1^m, a_{1:H}^m; \nu_k) - \sum_{h \in [H]} \gamma^{h-1} \phi(s_h^m, a_h^m; \theta_k) \right)^2$ .
This objective is a Monte Carlo estimate, but we can also leverage dynamic programming to do so in an infinite-horizon way, as we show in Appendix B.1.
3.3 Online Planning: Inferring Q-functions with Linear Regression and Planning
Our second phase concerns online planning to transfer to new tasks with arbitrary rewards using our learned Q-basis functions. For a sufficiently expressive set of random non-linear features, any non-linear reward function can be expressed as a linear combination of these features. As we outline below, this combined with linearity of expectation, allows us to express the task-specific Q-function for any particular reward as a linear combination of learned Q-bases of random functions. Therefore, we can obtain the test-time Q-function by solving a simple linear regression problem from random features to instantaneous rewards, and use the same inferred coefficients to recombine Q-bases to obtain task-specific Q-functions. This inferred Q-function can then be used to plan an optimal sequence of actions. We describe each of these steps below.
3.3.1 Reward Function Fitting with Randomized Features
We first learn how to express the reward function for the new task as a linear combination of the random features. This can be done by solving a linear regression problem to find the coefficient vector  $w = [w_{1}, \dots, w_{K}]^{\top}$  that approximates the new task's reward function as a linear combination of the (non-linear) random features. Specifically, we minimize the following loss