SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Thomas Schmied<sup>1</sup>, Markus Hofmacher<sup>2</sup>, Fabian Paischer<sup>1</sup>, Razvan Pascanu<sup>3,4</sup>, Sepp Hochreiter<sup>1</sup>
<sup>1</sup> ELLIS Unit Linz and LIT AI Lab, Institute for Machine Learning,
$^{2}$  JKU LIT SAL eSPML Lab, Institute for Machine Learning, Johannes Kepler University, Linz, Austria
3 Google DeepMind,4 UCL
schmied@ml.jku.at
Reinforcement Learning (RL) has been successful in various domains like robotics, game playing, and simulation. While RL agents have shown impressive capabilities in their specific tasks, they insufficiently adapt to new tasks. In supervised learning, this adaptation problem is addressed by large-scale pre-training followed by fine-tuning to new down-stream tasks. Recently, pre-training on multiple tasks has been gaining traction in RL. However, fine-tuning a pre-trained model often suffers from catastrophic forgetting. That is, the performance on the pre-training tasks deteriorates when fine-tuning on new tasks. To investigate the catastrophic forgetting phenomenon, we first jointly pre-train a model on datasets from two benchmark suites, namely Meta-World and DMControl. Then, we evaluate and compare a variety of fine-tuning methods prevalent in natural language processing, both in terms of performance on new tasks, and how well performance on pre-training tasks is retained. Our study shows that with most fine-tuning approaches, the performance on pre-training tasks deteriorates significantly. Therefore, we propose a novel method, Learning-to-Modulate (L2M), that avoids the degradation of learned skills by modulating the information flow of the frozen pre-trained model via a learnable modulation pool. Our method achieves state-of-the-art performance on the Continual-World benchmark, while retaining performance on the pre-training tasks. Finally, to aid future research in this area, we release a dataset encompassing 50 Meta-World and 16 DMControl tasks.
Reinforcement Learning (RL) has been instrumental in training agents capable of achieving notable successes, both in simulation, and in the real-world (Silver et al., 2016; Vinyals et al., 2019; Berner et al., 2019; Arjona-Medina et al., 2019; Bellemare et al., 2020; Degrave et al., 2022). However, such agents are usually highly specialized and incapable of performing well outside of a narrowly-defined task. Furthermore, adapting a pre-trained agent to a new task by fine-tuning usually results in decreased performance on prior tasks. This effect is well-known in the literature as catastrophic forgetting (McCloskey and Cohen, 1989).
A common paradigm to learn multiple tasks concurrently is multi-task learning (Caruana, 1997). However, typically, not all tasks we want an agent to learn are available at training time. In this case, new tasks must be learned in a sequential manner. Learning a new task ideally exploits knowledge from previously learned tasks and does not adversely affect the performance on these prior tasks. Recent works have demonstrated that models based on the Transformer architecture (Vaswani et al., 2017) excel at learning multiple tasks concurrently from large offline datasets (Lee et al., 2022; Reed et al., 2022; Jiang et al., 2022; Brohan et al., 2022; Gupta et al., 2022; Shridhar et al., 2022). We want
Figure 1: General setup of the experiments. We pre-train a Multi-Domain Decision Transformer (MDDT) model on a large set of tasks from multiple domains, namely Meta-World and DMControl. Then we fine-tune this model to learn additional tasks from these domains.
to utilize this capability in RL, thereby obtaining agents that can quickly learn new skills and adapt to new tasks. However, it is still unclear how an agent can learn new tasks and skills without negatively impacting its performance on the tasks it has already learned.
Our aim is to find methods that enable an agent to efficiently learn new tasks without compromising its proficiency in previously acquired tasks. In this regard, we draw inspiration from fine-tuning (FT) techniques prevalent in supervised learning, such as parameter-efficient fine-tuning (PEFT, (Houlsby et al., 2019; Hu et al., 2022; Liu et al., 2022b)) and prompt-based tuning (PBT, (Lester et al., 2021; Li and Liang, 2021)). Both PBT, and PEFT, incorporate a small set of new parameters to adapt a pre-trained model to new tasks at low cost. Thus, they intrinsically avoid catastrophic forgetting. However, it is unclear how well these methods can be adopted for training RL agents from offline datasets. Therefore, we first conduct a comprehensive evaluation for fine-tuning a pretrained Decision Transformer (DT, Chen et al., 2021) on two established RL benchmarks, namely Meta-World (Yu et al., 2020b; Wolczyk et al., 2021) and DMControl (Tassa et al., 2018). We first pre-train a DT jointly on tasks from both domains, then we transfer the pre-trained model to new tasks using various FT, PEFT and PBT methods (Figure 1). We find that FT methods adjust well to new tasks, but performance on previous tasks generally decreases. PBT, in contrast, retains performance on previous tasks but does not adapt well to new tasks.
To combine the advantages of both approaches, we propose Learning-to-Modulate (L2M), which is based on two recent works from NLP and computer vision, namely Learning-to-prompt (L2P, Wang et al., 2022c) and Low Rank Adaptation (LoRA, Hu et al., 2022). L2M operates in a task-agnostic manner and efficiently adapts a pre-trained model to new tasks via a learnable modulation pool. The modulation pool consists of learnable keys that are associated with learnable modulation matrices. Based on the current context, L2M selects a set of modulators which are learned during fin-tuning on the new task. This assumes that tasks can be well discriminated, such that suitable modulators are selected. Indeed, we observe an emergent clustering of tasks in the embedding layer of our model after pre-training, as illustrated in Figure 2. In turn, L2M is capable of quickly adapting to new tasks, while retaining performance on prior tasks, and introduces only minimal additional parameters per task relative to the full model size.
Figure 2: t-SNE clustering of state embeddings for the first ten MT40 and DMControl tasks. Similar tasks are clustered together, while dissimilar tasks are apart (legend is available in Appendix F).
Figure 3: Illustration of L2M. (a) We construct a query  $\mathbf{q}_t$  by aggregating the embedded history of state-tokens. (b) The query  $\mathbf{q}_t$  is matched with learnable keys in a modulation pool, which map to learnable modulation matrices. We retrieve the modulation matrices with the highest similarity between  $\mathbf{q}_t$  and every key in the modulation pool. (c) Retrieved modulation matrices modify the pre-trained and frozen multi-domain DT.
To summarize, we make the following contributions:
- We conduct an extensive evaluation of fine-tuning, parameter-efficient fine-tuning, and prompting methods for Transformers in RL.
- We propose the novel L2M for efficient fine-tuning of a frozen pre-trained model by modulating the information flow via learnable modulation matrices.
- We release a dataset of trajectories for the Meta-World and DMControl benchmark suites.
We propose L2M, a method to adapt a pre-trained RL agent such that new tasks can be learned efficiently without compromising performance on pre-training tasks. To this end, we combine approaches from prompt-based tuning and parameter-efficient fine-tuning.
Reinforcement Learning. We consider a Markov decision process (MDP) represented by the tuple  $(\mathcal{S},\mathcal{A},\mathcal{P},\mathcal{R})$ .  $\mathcal{S}$  and  $\mathcal{A}$  denote state and action spaces, respectively, with states  $s_t\in S$  and actions  $a_{t}\in \mathcal{A}$  at timestep  $t$ . The transition function  $\mathcal{P}(s_{t + 1}\mid s_t,a_t)$  takes a state-action pair and yields a probability distribution over next states. The reward function  $\mathcal{R}(r_t\mid s_t,a_t)$  maps a given state-action pair to a scalar reward  $r_t$ . The objective is to learn a policy  $\pi (a_t\mid s_t)$  that predicts an action  $a_{t}$  given  $s_t$  that maximizes the reward  $r_t$ .
Decision Transformer. Decision Transformer (DT, Chen et al., 2021) rephrases RL as a sequence modelling problem. This assumes access to a dataset  $\mathcal{D} = \{\tau_i\mid 1\leq i\leq N\}$ , where  $\tau_{i} = (s_{0},a_{0},r_{0},\dots ,s_{T},a_{T},r_{T})$  represents a trajectory of length  $T$  consisting of state-action-reward triplets and  $N$  defines the number of trajectories in  $\mathcal{D}$ . Further, we augment each trajectory  $\tau_{i}$  with returns-to-go  $\hat{R}_t = \sum_{t' = t}^T r_{t'}$ . The policy  $\pi_{\theta}(a_t\mid s_{t - C:t},\hat{R}_{t - C:t},a_{t - C:t - 1},r_{t - C:t - 1})$ , with parameters  $\theta$  and context length  $C$ , is then trained in a supervised manner via upside-down-RL (Schmidhuber, 2019), minimising the cross-entropy  $\mathcal{L}_{CE}(\hat{a}_t,a_t)$  between predicted action  $\hat{a}_t$  and ground-truth action  $a_{t}$ . During inference, the policy produces a sequence of actions conditioned on a high return-to-go.
Low Rank Adaptation (LoRA). LoRA (Hu et al., 2022) is a parameter efficient method for fine-tuning pre-trained models. It performs a low-rank decomposition of a weight matrix that is utilized to modulate the information flow in the pre-trained model. Particularly, we assume a pretrained layer  $l$  with weight matrix  $\pmb{W}_l \in \mathbb{R}^{d_{\mathrm{in}} \times d_{\mathrm{out}}}$  that receives  $\pmb{x}_{l-1} \in \mathbb{R}^{1 \times d_{\mathrm{in}}}$  as input. LoRA adds learnable low rank modulation matrices  $\pmb{A}_l \in \mathbb{R}^{d_{\mathrm{in}} \times r}$  and  $\pmb{B}_l \in \mathbb{R}^{r \times d_{\mathrm{out}}}$ , where  $r \ll d_{\mathrm{in}}$ . The output  $\pmb{h}_l$  of layer  $l$  is then computed as:
$$
\boldsymbol {h} _ {l} = \boldsymbol {W} _ {l} \boldsymbol {x} _ {l - 1} + \boldsymbol {B} _ {l} \boldsymbol {A} _ {l} \boldsymbol {x} _ {l - 1}. \tag {1}
$$
During the fine-tuning stage, only the parameters of  $A_{l}$  and  $B_{l}$  are updated.
2.2 Learning-to-Modulate (L2M)
We propose L2M, a method that combines the advantages of PEFT and PBT approaches. In particular, we rely on LoRA (Hu et al., 2022), and refer to the trainable parameters it induces as modulators of the pre-trained model. Further, we draw inspiration from Learning-to-prompt (L2P, Wang et al., 2022c), and maintain a pool of modulators associated with learnable keys. For a given input sequence, we retrieve the best modulator weights from the pool, and leverage them to alter the behaviour of the pretrained model (Figure 3).
We define a modulation pool that contains a set of  $M$  learnable keys,  $\mathbf{K}_{\mathrm{pool}} = \{\mathbf{k}_i \mid 1 \leq i \leq M\}$ . All keys in  $\mathbf{K}_{\mathrm{pool}}$  are initialized uniformly from  $[-1, 1]$ . Each  $\mathbf{k}_i$  is associated with a set of modulation matrices  $\{\mathbf{A}_b^i, \mathbf{B}_b^i \mid 1 \leq b \leq B\}$ , for each layer block  $b$  of a DT with  $B$  layer blocks. As in Hu et al. (2022), we initialize  $\mathbf{A}$  according to a normal distribution around zero and  $\mathbf{B}$  with zeros. Thus, no modulation can occur at the beginning of training.
At time  $t$ , we generate a query vector  $\mathbf{q}_t \in \mathbb{R}^d$ , where  $d$  is the embedding dimension of the pre-trained model, in the following way. First, we construct a matrix  $\mathbf{S}_{t - C:t} \in \mathbb{R}^{C \times d}$  from the states of a trajectory  $\tau$  with the context length  $C$ , after they are processed via the embedding layer of the frozen pre-trained model. Then, we reduce the matrix to a query vector  $\mathbf{q}_t$  by an aggregation function  $g(\cdot)$ :
$$
\mathbf {q} _ {t} = g \left(\mathbf {S} _ {t - C: t}\right) \tag {2}
$$
where we use mean-pooling for  $g(\cdot)$ . We retrieve a set of modulation matrices  $\{\mathbf{A}_b^j, \mathbf{B}_b^j \mid 1 \leq b \leq B\}$  by the maximum similarity between each  $\mathbf{k}_i \in \mathbf{K}_{\mathrm{pool}}$  and the query  $\mathbf{q}_t$  at timestep  $t$ :
$$
j = \underset {\mathbf {k} _ {i} \in \mathbf {K} _ {\text {p o o l}}} {\operatorname {a r g m a x}} \sin \left(\mathbf {q} _ {t}, \mathbf {k} _ {i}\right) n \left(\mathbf {k} _ {i}\right) ^ {- 1} \tag {3}
$$
In our case,  $\mathrm{sim}(\cdot ,\cdot)$  corresponds to the cosine similarity and  $n(\mathbf{k}_i)$  represents the selection count for key  $\mathbf{k}_i$ . Every time a key  $\mathbf{k}_i$  is selected, we increase  $n(\mathbf{k}_i)$  by one. This way, we discourage that different queries always attend to the same key.
We use  $\{\mathbf{A}_b^j,\mathbf{B}_b^j\mid 1\leq b\leq B\}$  to modulate the pre-trained model according to Equation (1). More specifically, we apply LoRA on the queries and values in the self-attention mechanism, as well as on the activation in the feed-forward block of the pre-trained model. Only the modulators are learned via gradient descent, while the pre-trained model remains frozen. Following Wang et al. (2022c), the key is updated by maximizing the cosine similarity between  $\mathbf{q}_t$  and  $\mathbf{k}_j$  via an additional term in the end-to-end objective:
$$
\min  _ {\theta , \mathbf {k} _ {i}} \mathcal {L} _ {C E} \left(\hat {a} _ {t}, a _ {t}\right) - \lambda \operatorname {s i m} (\operatorname {s t o p g r a d} \left(\mathbf {q} _ {t}\right), \mathbf {k} _ {j}) \tag {4}
$$
where  $\lambda$  is a hyperparameter.
L2M unifies the benefits of LoRA and L2P, ensuring high-performance and few additional learnable parameters, while also avoiding forgetting on the pre-trained tasks.
2.3 Multi-Domain Decision Transformer (MDDT)
We extend the Decision Transformer architecture proposed by Chen et al. (2021) to handle inputs from multiple domains with varying state/action spaces. Since the dimensionality of states differ between Meta-World and DMControl we construct a unified state space, that comprises all dimension of DMControl and Meta-World environments. This state-space consists of 204 dimensions in total. Dimensions that are unused by a certain task are padded with zeros. Finally, the states are embedded with a linear layer before serving as input to the MDDT. Similar to Reed et al. (2022) and Brohan et al. (2022), we tokenise each action dimension and autoregressively predict action tokens. In the environments we consider, all actions are bounded by [-1, 1]. We use a min-max tokenisation, which comprises min-max normalisation and subsequent uniform discretisation into 64 bins. These changes enable processing observations and actions originating from different environments at training and test time. In line with prior work (Chen et al., 2021; Lee et al., 2022), we train the DT via return-conditioned upside-down RL using a cross-entropy loss to predict next actions autoregressively. We set the target return to the maximum observed return in the respective dataset, but also found that constant proxies per domain work well. We provide further implementation details on our MDDT as well as a discussion on limitations and implications of our design choices in Appendix C.
We consider two different benchmark suites that comprise 66 different tasks in total, namely Meta-World (Yu et al., 2020b; Wolczyk et al., 2021) and DMControl (Tassa et al., 2018). Meta-World consists of 50 diverse tasks for robotic manipulation, such as grasping, manipulating objects, opening/closing a window, pushing buttons, locking/unlocking a door, and throwing a basketball. We follow Wolczyk et al. (2021) and split them into 40 pre-training (MT40), and 10 fine-tuning tasks (CW10). We use Meta-World v2, instead of v1 used by Wolczyk et al. (2021), but will refer to it as Meta-World throughout this work. Further, we select 16 tasks from DMControl and assign 10 tasks to the pre-training set (DMC10) and 6 tasks to the fine-tuning set (DMC6). We elaborate on the different environments and on the data collection procedure in Appendix A, and Appendix D, respectively. Unlike Meta-World, which shares the same state and action spaces across tasks ( $|S| = 39$ ,  $|\mathcal{A}| = 6$ ), the state and action spaces of DMControl vary for each task (e.g., for hopper  $|S| = 15$ ,  $|\mathcal{A}| = 4$ , for cheetah  $|S| = 17$ ,  $|\mathcal{A}| = 6$ ), as listed in Appendix D.
We collect a dataset<sup>1</sup> by training task-specific agents on each task using an agent based on Soft Actor Critic (SAC) (Haarnoja et al., 2018). The average performance for SAC across the Meta-World and DMControl task splits is shown in Table 1. For each of the 50 tasks in Meta-World, we include 10K trajectories of length 200 in the dataset, amounting to 100M transitions in total. For DMControl, we include 1000 trajectories of length 1000 for each task, amounting to 16M transitions in total. The datasets contain the entire replay buffer of the task-specific agents, and consequently behaviours ranging from random to expert. We choose this collection scheme since prior work has illustrated the benefits of training agents on data that comprises mixed behaviour (Lee et al., 2022). We give further details on the training procedure, including hyperparameters and learning curves in Appendix D.
Table 1: Performance measures of task specific SAC-based agents used for data collection averaged over the different task splits. Mean and standard deviation are shown.
(a) Meta-World
Dataset	Success Rate	Mean Reward
MT40	0.84 ± 0.34	1414.62 ± 439.39
CW10	1.0 ± 0.0	1540.49 ± 184.43
(b) DMControl
Dataset	Mean Reward
DMC10	788.36 ± 219.11
DMC6	840.37 ± 216.63
Figure 1 illustrates our experimental design. First, we pre-train a MDDT on all pre-training datasets (MT40 and DMC10) simultaneously (see Section 3.1). Next, we conduct a broad evaluation of different FT, PEFT and PBT methods to adapt the pre-trained model to each of the held-out fine-tuning tasks (CW10 and DMC6). In Section 3.2, we show the results for fine-tuning on all CW10 and DMC6 tasks individually. Finally, we present results for training all methods in a continual RL setup, in which tasks are introduced sequentially (Section 3.3). All our models are trained offline on the generated datasets and evaluated online in the respective environments.
Unless mentioned otherwise, we report IQM and  $95\%$  bootstrapped confidence intervals, as proposed by Agarwal et al. (2021) (3 seeds per method and 50K bootstrap samples). We report success rates and data-normalized scores (normalized by mean score in dataset and random agent performance as suggested by Fu et al. (2020)) for Meta-World and DMControl, respectively.
We pre-train a MDDT on our datasets collected for MT40 and DMC10. This already results in strong performance on the pre-training tasks (60% for MT40, 94% for DMC10). We experiment with different model sizes (40M to 200M parameters), varying number of Transformer layers, number of heads per layer and embedding dimension. Generally, there is a trend that more complex models perform better. For all our subsequent experiments, we use a 40M parameter model, as it achieves a good trade-off between performance and training time. Performance scores, learning curves, implementation details and hyperparameters for our pre-training experiments are provided in Appendix E.
Source code and datasets are available at: https://github.com/mlj-kuL2M
After pre-training, we analyse the learned representations of the model via t-SNE (Van der Maaten and Hinton, 2008). Figure 2 illustrates emerging clusters of learned state embeddings for ten tasks of MT40 and DMC10. We observe that similar tasks cluster together (e.g., button-press-v2 and button-press-wall-v2), whereas dissimilar tasks are well separated (e.g., assembly-v2 and bin-picking-v2). We repeat this analysis for rewards, RTG and action token embeddings and explain the clustering procedure in Appendix F.
(a) CW10
(b) DMC6
Figure 4: IQM and  $95\%$  CIs for single-task fine-tuning on (a) CW10 and (b) DMC6. Models are pre-trained on MT40/DMC6 and are then optimized for each CW10/DMC6 task individually.
3.2 Single-Task Fine-Tuning
We evaluate the performance of various FT, PEFT and PBT strategies for fine-tuning on the held-out tasks. Overall, we compare a total of 9 methods: Full fine-tuning (FT), FT of action head (FT-head), FT of last Transformer layer and action head (FT-last+head), Adapters (Houlsby et al., 2019), LoRA (Hu et al., 2022), (IA) $^3$  (Liu et al., 2022b), Prompt-tuning (Lester et al., 2021), Prefix-tuning (Li and Liang, 2021), and P-tuning v2 (Liu et al., 2021b). We do not compare against meta-RL algorithms, as Mandi et al. (2022) showed that FT approaches perform on-par or better on several tasks. A detailed list of hyperparameters, training details and description for each method are provided in Appendix G. We fine-tune the pre-trained model for each task individually and aggregate the performance over all tasks within each domain.
In Figure 4, we show IQM over all tasks for all methods on CW10 and DMC6. Overall, FT attains the highest scores, since it can utilize the entire model capacity. Thus, it serves as an upper bound in terms of performance. Adapters achieve the second-highest scores on average, followed by LoRA,  $(\mathrm{IA})^3$  and FT-last+head. Notably, there is a large gap between PBT and PEFT methods, particularly on CW10. Finetuning the action head for each task separately, results in the worst performance. While the absolute scores differ between the two domains, the relative performance ranking among methods remains equivalent. On DMC6, we observe a larger performance gap between full FT and PEFT methods. While in Meta-World only the locomotion differs between tasks, in DMControl the entire robot morphology may change.
Figure 5: Success rate vs. fraction of parameters trained for various fine-tuning techniques on single-task experiments for CW10.
Therefore, more capacity of the model is required to sufficiently adapt to the new tasks.
We highlight the efficacy of the considered methods by comparing the fraction of trainable parameters against attained performance on CW10 in Figure 5 (see Appendix G for same comparison on DMC6). FT updates all parameters, while Adapters update two orders of magnitude less parameters  $(2.1\%)$
In contrast, LoRA requires less than half the amount of parameters of Adapters (1%), while reaching a similar level of performance. Notably,  $(\mathrm{IA})^3$ , trains approximately the same amount of parameters as Prompt-tuning (0.21%), but attains significantly higher performance. This result indicates that we can achieve decent performance with a small amount of trainable parameters if the pre-training and fine-tuning tasks are similar. Overall, LoRA compares favourably in both dimensions.
3.3 Continual Fine-Tuning
Ultimately, our goal is to adapt the pre-trained model to multiple novel tasks, while alleviating forgetting of tasks learned during pre-training. In this regard, we adapt the pre-trained model to all fine-tuning tasks in a sequential manner, as proposed by Wolczyk et al. (2021) for CW10. Moreover, we evaluate forgetting by measuring the performance on tasks acquired during pre-training after fine-tuning.
(b) DMC6
Figure 6: IQM and  $95\%$  CIs of success rates for CRL experiments on (a) CW10 and (b) DMC6. Models are pre-trained on MT40+DMC10 and are then trained on the tasks from CW10/DMC6 sequentially. On each task, we train for 100K steps and then move to the next task in the sequence.