SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Again, we compare the following fine-tuning variations: FT, FT-head, and FT-last+head. Additionally, we augment the PBT methods of the previous section with a prompt pool as in L2P, which enables learning of task-specific prompts. We refer to these methods as L2P + Prompt-tuning (L2P-PT), L2P + Prefix-tuning (L2P-PreT), L2P + P-tuning v2 (L2P-Pv2). We note that L2P-PT reflects the original version of L2P proposed by Wang et al. (2022c) which was originally applied in computer vision. Additionally, we compare L2M to two established methods from Continual RL (CRL), namely Elastic Weight Consolidation (EWC), and L2, both proposed by Kirkpatrick et al. (2017). EWC constrains task-specific weights to not change too much when adapting to a new task by estimating their individual importances. L2 applies a uniform constraint to all weights without considering their relevance. Moreover, we add another implementation of L2M, which is equipped with an oracle that provides information on what task is currently being observed in terms of a task index (L2M-Oracle). For L2M-oracle, the modulation pool contains as many modulators as there are tasks. At training time, the task index refers to the dataset the batches are sampled from. At inference time, the task index refers to the environment the DT is currently evaluated in. In contrast to Wolczyk et al. (2021), we do not use separate action heads per task for L2M to remain task agnostic. To illustrate an upper bound on performance, we provide scores for two multi-task baselines, which train on all tasks simultaneously, either from scratch, or after pre-training (Table 7 in Appendix H). We train each method for 100K steps per task and evaluate on all tasks every 50K steps. We provide further training details and hyperparameters in Appendix H.
In Figures 6(a) and 6(b), we show the performance scores of all methods on CW10 and DMC6, respectively. In addition, we report forgetting scores (as calculated by Wolczyk et al. (2021)), and rewards obtained at the end of training, in Appendix H. L2M outperforms all other approaches in both domains, attaining an average success rate of  $65\%$  across CW10 tasks and a normalized score of  $56\%$  across DMC6 tasks. Adding a task oracle to L2M increases the success rate to  $76\%$  and  $70\%$  on CW10 and DMC6, respectively, and closely matches the single-task performance of LoRA. L2P
Figure 7: Performance on the MT40 and DMC10 pre-training tasks after fine-tuning.
combined with different prompting approaches performs considerably worse than L2M. Interestingly, the established CRL method EWC does not mitigate forgetting sufficiently. Similar results for EWC were reported by Ben-Iwhiwhu et al. (2022) in an online RL setting. To the best of our knowledge, the results of our method, L2M, are the highest reported results on Continual-World v2 to date. Prior work reports average success rates of roughly  $40\%$  (Caccia et al., 2022; Ben-Iwhiwhu et al., 2022). Even though L2M outperforms other approaches, there is a considerable gap between L2M and L2M-oracle. We surmise this is due to conflation effects in the prompt pool and aim to investigate this in more detail in future work.
Performance on pre-training tasks. Finally, we evaluate performance on the pre-training tasks after fine-tuning, as shown in Figure 7. To remain task-agnostic, we train a set of 100 keys for the pre-training tasks, which we concatenate to the set of keys introduced during the fine-tuning stage. While FT, L2 and EWC experience a severe drop in performance, L2M and L2P-based approaches maintain a similar performance level as prior to fine-tuning. Thus, L2M preserves performance on the pre-training tasks, while it effectively adapts to new tasks.
To gain a better understanding of our method and its limitations, we conduct a number of additional ablations.
Modulation Targets. One important design choice in L2M is which weights of the pretrained model to modulate. In principle, LoRA can be applied to any weight matrix in the pretrained model. By default, we employ LoRA on the queries and values in the self-attention mechanism, and on the position-wise feedforward layer. In Appendix G.1, we conduct an ablation study in which we vary the modulation targets. Our analysis suggests that modulating the attention mechanism is not as important as modulating the position-wise feed-forward layer. However, modulating both, the attention mechanism, and the position-wise feed-forward layer yields the best results. Nevertheless, this performance gain comes at the cost of more trainable parameters.
Figure 8: Aggregation token ablation on CW10.
Query Representation. As specified in Section 2.2, we use embedded state tokens aggregated over a specified context window as input query to the prompt pool in L2M. This design choice is inspired by the observed task separation in the embedding layer after pre-training (see Figure 2). We investigate other choices of individual tokens and token combinations (state, action, reward, RTG, state-action, state-RTG) to represent the query (see Figure 8). Indeed, we observe that using the state embeddings
results in the best overall performance, as it aids task separation. In contrast, using RTG or action embeddings deteriorates performance. Furthermore, we also investigate the effect of using different layer representations and context windows to construct the query in Appendix H.1.
Alternative Modulators. We conduct an ablation study in which we compare L2M against L2M in combination with  $(\mathrm{IA})^3$  in Appendix H.2. Instead of the low-rank modulation matrices,  $(\mathrm{IA})^3$  employs elementwise multiplication with learnable modulation vectors. While performance decreases with L2M-(IA) $^3$ , it compares favourably in terms of parameter-efficiency. Depending on the task, this may be preferred.
Single-domain DT on Meta-World. To better understand the potential drawbacks of using a multi-domain model, we pre-train and fine-tune a single-domain model only on Meta-World (see Appendix H.3). Due to the common state and action spaces in Meta-World, we use a simpler, non-discretized, action space and training objective (MSE) for this experiment. This specialised single-domain model obtains considerably higher performance scores. However, these performance gains come at the cost of the loss of generality, as the specialised model can only handle the particular state/action space it was trained on. Thus, fine-tuning it to tasks with new state/action spaces is not possible. Nevertheless, this experiment highlights the drawbacks of our current mechanisms, such as action discretisation and autoregressive action-prediction, that enable multi-domain training.
Cross-domain FT. Finally, to better understand the potential upsides of using a multi-domain model, we pre-train a DT (with unified state space and action discretisation) on Meta-World only (MT40) and then fine-tune it on DMControl (DMC6). We observe that the fine-tuning performance on DMC6 (different domain) is considerably worse than for the MDDT (see Appendix H.4). In addition, we also fine-tune the pre-trained single-domain model on CW10 (same domain). Interestingly, the final performance on CW10 is also lower compared to the MDDT that was pre-trained on both domains. This experiment indicates, that multi-domain pre-training can, indeed, have a positive effect on the fine-tuning performance.
We provide additional details on our ablation studies in Appendix G and H.
Transformers in RL. Since the inception of Transformers (Vaswani et al., 2017) there has been a widespread adoption of the underlying architecture in various areas such as NLP (Devlin et al., 2019; Radford et al., 2019; Brown et al., 2020), computer vision (Dosovitskiy et al., 2021; He et al., 2022; Radford et al., 2021; Fürst et al., 2022; Ramesh et al., 2021; Rombach et al., 2022), speech recognition (Radford et al., 2022; Baevski et al., 2020) or video generation (Ho et al., 2022; Singer et al., 2022). More recently, Transformers have found their way into RL. Similar to the DT Chen et al. (2021), Trajectory Transformer (Janner et al., 2021) is based on the GPT architecture (Radford et al., 2018), but relies on dynamics modelling. Lee et al. (2022) extended DT to a multi-game setup to learn to play 46 Atari games. Meanwhile, a variety of DT-variants have been proposed (Zheng et al., 2022; Wang et al., 2022a; Shang et al., 2022; Meng et al., 2021). Siebenborn et al. (2022) replace the Transformer in DT with an LSTM (Hochreiter and Schmidhuber, 1997). PromptDT (Xu et al., 2022b) demonstrated that prompting a pre-trained DT model with expert trajectories can improve the agent's ability to generalize to new tasks. Jiang et al. (2022) presented a prompt-based Transformer for robot manipulation, that integrates multi-modal prompts via cross-attention. Xu et al. (2022a) propose to augment a DT with a hyper-network for parameter-efficient policy adaptation. A number of other works (which are largely orthogonal to our approach) instead aim to improve the pre-training stage, for example by predicting masked-out parts of the sequence (Carroll et al., 2022; Liu et al., 2022a; Sun et al., 2022). Reed et al. (2022) trained a Transformer that scaled to over 600 tasks. Most recently, Brohan et al. (2022) presented a scalable Transformer for real-world robotics manipulation. Laskin et al. (2022) and Liu and Abbeel (2023) make use of a multi-episodic context to elicit in-context learning capabilities. Other works use a Transformer backbone for history compression in online RL (Parisotto et al., 2020; Paischer et al., 2022, 2023). Li et al. (2023) and Yang et al. (2023) cover the landscape of Transformers in RL in more detail.
Continual and multi-task RL. A plethora of works have considered learning multiple tasks concurrently via RL (Tanaka and Yamamura, 2003; Rusu et al., 2016a; Borsa et al., 2016; Rajeswaran et al., 2017; El Bsat et al., 2017; Andreas et al., 2017; D'Eramo et al., 2020; Yu et al., 2020a; Igl et al., 2020; Sodhani et al., 2021). In CRL, however, tasks are considered to not be readily available at the same
time. Many continual learning methods were proposed for computer vision, but can also be applied to a CRL setting. Early works include regularization approaches, such as EWC (Kirkpatrick et al., 2017), Synaptic Intelligence (Zenke et al., 2017), Memory-aware Synapses (Aljundi et al., 2018), Gradient Episodic Memory (GEM, Lopez-Paz and Ranzato, 2017) and A-GEM (Chaudhry et al., 2019). Other approaches rely on adding new modules to existing architecture (Rusu et al., 2016b), iterative pruning (Mallya and Lazebnik, 2018), or improved exploration (Steinparz et al., 2022). A number of approaches have been tested on the Continual-World benchmark (Wolczyk et al., 2021), including 3RL (Caccia et al., 2022), ClonExSAC (Wolczyk et al., 2022) and Modulating masks (Ben-Iwhiuhu et al., 2022). L2P (Wang et al., 2022c) learns to prompt a frozen Vision Transformer (Dosovitskiy et al., 2021) and consistently outperforms prior methods on a number of continual learning benchmarks. Other follow-up works that rely on a prompting mechanism for CL have been proposed (Wang et al., 2022b; Smith et al., 2022; Razdaibiedina et al., 2023). Recent works provide a comprehensive overview of the field of CRL (Hadsell et al., 2020; Lesort et al., 2020; Khetarpal et al., 2022; Baker et al., 2023).
Parameter-efficient fine-tuning and Prompting. Large-language models (LLMs) are usually pretrained on vast amounts of data (Devlin et al., 2019; Radford et al., 2019; Brown et al., 2020). After pre-training, it is desirable to specialize or fine-tune the foundation model (Bommasani et al., 2021) to a down-stream task. A common way is fine-tuning all network weights or a fraction thereof. Fine-tuning the entire network is costly and suffers from catastrophic forgetting. Parameter-efficient fine-tuning methods and prompt-based tuning offer attractive alternatives. Houlsby et al. (2019) repurposed Adapter modules (Rebuffi et al., 2017) to interleave pretrained Transformer layers. Variations thereof have been proposed (Bapna et al., 2019; Pfeiffer et al., 2021, 2020; Karimi Mahabadi et al., 2021). Low Rank Adaptation injects trainable low-rank decomposition matrices into every layer of the model (Hu et al., 2022).  $(\mathrm{IA})^3$  modulates the inner activation flow of the Transformer layers by elementwise multiplication with learned modulation vectors (Liu et al., 2022b). Prompt tuning conditions the pre-trained model on new tasks by prepending learnable prompts to the embedded input sequence (Lester et al., 2021). Similarly, prefix-tuning adds learnable prefix vectors to the keys and values of each attention head input (Li and Liang, 2021). P-tuning v2, applies learnable prompts at each Transformer layer, instead of merely the input layer (Liu et al., 2021b). UniPELT combines different PEFT methods in a unified framework (Mao et al., 2022) Liu et al. (2021a) give a comprehensive overview of prompt-based learning for NLP.
Adapting agents to new tasks, while preserving performance on previously learned tasks, remains a major challenge towards more general RL agents. Consequently, we conduct a comprehensive evaluation of established fine-tuning, PEFT and PBT methods for Transformers in RL. We evaluate both how well new tasks are learned and how strongly performance deteriorates on the pre-training tasks. While full fine-tuning of a pre-trained model adapts well to new tasks, it suffers from catastrophic forgetting of pre-training tasks. Prompt-based tuning techniques preserve performance on pre-training tasks, but cannot compete on new tasks. We propose a novel method, L2M, which performs well in both dimensions. L2M enables efficient fine-tuning of pre-trained models via learnable modulation matrices, resulting in high performance both on new tasks and on pre-training tasks. Finally, we release a large dataset of trajectories for the Meta-World and DMControl benchmark suites to facilitate future research in the domain of offline RL.
For future work, we envision an investigation on the compositionality of learned modulation matrices. While currently a distinct set of modulators is selected for a given input sequence, we believe that by combining modulators across tasks, our method can improve in terms of performance and parameter-efficiency. Such an approach could improve forward-transfer, where new tasks benefit even more from previously learned tasks by combining existing skills. Furthermore, the context length of the models we consider is currently limited. Extending the context to leverage information from multiple episodes and tasks may improve the model's ability to generalize to new tasks or domains. Another promising avenue of research is to investigate more diverse domains and explore how different fine-tuning methods perform on more distinct fine-tuning domains.
The ELLIS Unit Linz, the LIT AI Lab, the Institute for Machine Learning, are supported by the Federal State Upper Austria. We thank the projects AI-MOTION (LIT-2018-6-YOU-212), DeepFlood (LIT-2019-8-YOU-213), Medical Cognitive Computing Center (MC3), INCONTROL-RL (FFG-881064), PRIMAL (FFG-873979), S3AI (FFG-872172), DL for GranularFlow (FFG-871302), EPILEPSIA (FFG-892171), AIRI FG 9-N (FWF-36284, FWF-36235), AI4GreenHeatingGrids(FFG-899943), INTEGRATE (FFG-892418), ELISE (H2020-ICT-2019-3 ID: 951847), Stars4Waters (HORIZONCL6-2021-CLIMATE-01-01). We thank Audi.JKU Deep Learning Center, TGW LOGISTICS GROUP GMBH, Silicon Austria Labs (SAL), University SAL Labs initiative, FILL Gesellschaft mbH, Anyline GmbH, Google, ZF Friedrichshafen AG, Robert Bosch GmbH, UCB Biopharma SRL, Merck Healthcare KGaA, Verbund AG, GLS (Univ. Waterloo) Software Competence Center Hagenberg GmbH, TÜV Austria, Frauscher Sensonic, TRUMPF and the NVIDIA Corporation. We acknowledge EuroHPC Joint Undertaking for awarding us access to Karolina at IT4Innovations (Czech Republic) and MeluXina at LuxProvide (Luxembourg).
Agarwal, R., Schwarzer, M., Castro, P. S., Courville, A. C., and Bellemare, M. (2021). Deep reinforcement learning at the edge of the statistical precipice. Advances in neural information processing systems, 34:29304-29320.
Aljundi, R., Babiloni, F., Elhoseiny, M., Rohrbach, M., and Tuytelaars, T. (2018). Memory aware synapses: Learning what (not) to forget. In Proceedings of the European Conference on Computer Vision (ECCV), pages 139-154.
Andreas, J., Klein, D., and Levine, S. (2017). Modular multitask reinforcement learning with policy sketches. In Precup, D. and Teh, Y. W., editors, Proceedings of the 34th International Conference on Machine Learning, ICML 2017, Sydney, NSW, Australia, 6-11 August 2017, volume 70 of Proceedings of Machine Learning Research, pages 166-175. PMLR.
Arjona-Medina, J. A., Gillhofer, M., Widrich, M., Unterthiner, T., Brandstetter, J., and Hochreiter, S. (2019). Rudder: Return decomposition for delayed rewards. Advances in Neural Information Processing Systems, 32.
Ba, J. L., Kiros, J. R., and Hinton, G. E. (2016). Layer normalization. arXiv preprint arXiv:1607.06450.
Baevski, A., Zhou, Y., Mohamed, A., and Auli, M. (2020). wav2vec 2.0: A framework for self-supervised learning of speech representations. Advances in Neural Information Processing Systems, 33:12449-12460.
Baker, M. M., New, A., Aguilar-Simon, M., Al-Halah, Z., Arnold, S. M., Ben-Iwhiwu, E., Brna, A. P., Brooks, E., Brown, R. C., Daniels, Z., et al. (2023). A domain-agnostic approach for characterization of lifelong learning systems. Neural Networks.
Bapna, A., Arivazhagan, N., and First, O. (2019). Simple, scalable adaptation for neural machine translation. arXiv preprint arXiv:1909.08478.
Bellemare, M. G., Candido, S., Castro, P. S., Gong, J., Machado, M. C., Moitra, S., Ponda, S. S., and Wang, Z. (2020). Autonomous navigation of stratospheric balloons using reinforcement learning. Nature, 588(7836):77-82.
Ben-Iwhiwhu, E., Nath, S., Pilly, P. K., Kolouri, S., and Soltoggio, A. (2022). Lifelong reinforcement learning with modulating masks. arXiv preprint arXiv:2212.11110.
Berner, C., Brockman, G., Chan, B., Cheung, V., Dkebiak, P., Dennison, C., Farhi, D., Fischer, Q., Hashme, S., Hesse, C., et al. (2019). Dota 2 with large scale deep reinforcement learning. arXiv preprint arXiv:1912.06680.
Bommasani, R., Hudson, D. A., Adeli, E., Altman, R., Arora, S., von Arx, S., Bernstein, M. S., Bohg, J., Bosselut, A., Brunskill, E., et al. (2021). On the opportunities and risks of foundation models. arXiv preprint arXiv:2108.07258.
Yu, T., Kumar, S., Gupta, A., Levine, S., Hausman, K., and Finn, C. (2020a). Gradient surgery for multi-task learning. In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H., editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual.
Yu, T., Quillen, D., He, Z., Julian, R., Hausman, K., Finn, C., and Levine, S. (2020b). Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on robot learning, pages 1094-1100. PMLR.
Zenke, F., Poole, B., and Ganguli, S. (2017). Continual learning through synaptic intelligence. In International Conference on Machine Learning, pages 3987-3995. PMLR.
Zheng, Q., Zhang, A., and Grover, A. (2022). Online decision transformer. In Chaudhuri, K., Jegelka, S., Song, L., Szepesvári, C., Niu, G., and Sabato, S., editors, International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, volume 162 of Proceedings of Machine Learning Research, pages 27042-27059. PMLR.
A Environments
A.1 Meta-World & Continual-World
Meta-World. For our experiments, we use the Meta-World benchmark (Yu et al., 2020b). Meta-World tasks consist of a Sawyer robotic arm simulated using the MuJoCo physics engine (Todorov et al., 2012). The observation space in Meta-World is a 39-dimensional vector. The action space is 4-dimensional, with all actions in range  $[-1,1]$ . For each task, the reward functions are scaled to have a maximum value of 10 and a minimum of 0. The exact reward-function definitions are provided in Yu et al. (2020b). Episodes last for 200 timesteps for each task. For all our experiments on Meta-World, we report success rates and mean rewards obtained.
Difference between Meta-World v1 and v2 environments. The Continual-World benchmark was published using v1 of Meta-World. However, in the meantime v2 has been released. While the v1 environments had a 9 dimensional observation space, the observation space is 39-dimensional in v2. Another major change in v2, was the introduction of dense and bounded reward functions.
Continual-World. The Continual-World benchmark was proposed by Wolczyk et al. (2021) and is built on top of Meta-World. Continual-World is a challenging CRL benchmark, consisting of 10 of the 50 tasks contained in Meta-World, denoted as CW10. We use the tasks contained in CW10 as fine-tuning set and select the remaining 40 tasks as pre-training set (MT40). The CW10 task sequence is:
hammer-v2, push-wall-v2, faucet-close-v2, push-back-v2, stick-pull-v2, stick-pull-v2, handle-press-side-v2, push-v2, shelf-place-v2, window-close-v2, and peg-unplug-side-v2.
A.2 DeepMind Control Suite (DMControl)
The DMControl Suite benchmark (Tassa et al., 2018) consists of 30 simulated continuous control tasks built on the MuJoCo physics engine. DMControl encompasses a wide variety of robot morphologies, ranging from a simple cartpole to a complex humanoid. Therefore, the observation and action spaces vary across environments. The action spaces vary between 1 (cartpole) and 21 continuous dimensions (humanoid), where each action dimension is bounded by  $[-1,1]$ . Agents can be trained based on continuous vector-based representations (state-based) or directly from pixels (pixel-based). In this work, we focus on the state-based representations, which vary between 3 (pendulum) and 67 continuous dimensions (humanoid). As described in Section 3, we select 16 of the 30 DMControl tasks and assign 10 to the pre-training set (DMC10) and 6 to the fine-tuning set (DMC6) (similar to Hafner et al. (2019)). We list the dimensionalities of their state and action spaces in Table 4.
The DMC10 tasks include:
cartpole-balance, finger-turn_easy, finger-turn_hard, fish-upright, hopper-stand, pendulum-swingup, point_mass-easy, reacher-hard, walker-run, walker-stand
Furthermore, the DMC6 tasks include:
ball_in_cup-catch, cartpole-swingup, cheetah-run, finger-spin, reacher-easy, walker-walk
For all environments in DMControl, episodes last for 1000 timesteps and the maximum achievable return is 1000. In addition to the mean reward obtained, we report data-normalized scores, as suggested by Fu et al. (2020). For each environment, the scores are normalized by the average scores achieved by the expert agent used to collect the dataset, and the performance attained by a random agent, normalized_score =  $\frac{\text{score-random_score}}{\text{expert_score-random_score}}$ . Therefore, the normalized scores are roughly between 0 and 1. The expert and random scores used for score normalization are available in the accompanying source code<sup>2</sup>.
Fine-tuning. In full fine-tuning, the entire pre-trained model is trained on the new task. We also try other common variations thereof, such as fine-tuning the action head, the last layer, or both.
2Available at: https://github.com/ml- jku/L2M
Adapters. Adapters are a parameter-efficient fine-tuning approach proposed by Houlsby et al. (2019). An adapter layer consists of a down-projection, a non-linearity and an up-projection along with a skip connection. Each Transformer block is interleaved with two Adapter modules. During training, only the Adapters (and optionally LayerNorms) are updated, while the remaining parameters remain frozen. This significantly reduces the number of trainable parameters, while preserving the ability to adapt to new tasks. Houlsby et al. (2019) adapt a BERT (Devlin et al., 2019) model to 26 different text classification tasks, while training only roughly  $4\%$  of the parameters and attaining performance close to full fine-tuning.