SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$(\mathbf{IA})^{3}$ : Infused Adapter by Inhibiting and Amplifying Inner Activations. Similar to LoRA,  $(\mathrm{IA})^{3}$  (Liu et al., 2022b) modulates the inner activation flow of the Transformer layers and can be applied to every component of the pre-trained model. Instead of low-rank modulation matrices as used in LoRA,  $(\mathrm{IA})^{3}$  performs elementwise multiplication with learnable modulation vectors. Thus,  $(\mathrm{IA})^{3}$  typically incurs less additional parameters compared to other methods.
Prompt-tuning. Lester et al. (2021) learn soft prompts to enable a frozen pre-trained GPT-3 (Brown et al., 2020) model to generalize to downstream tasks. In prompt-tuning, the learnable prompts are prepended to the input sequence, while the rest of the pre-trained model is kept frozen.
Prefix-tuning. Prefix-tuning (Li and Liang, 2021) is another PBT-approach. Similar to prompt-tuning, prefix-tuning prepends learnable vectors (i.e., soft-prompts) to the input. However, unlike in prompt-tuning, the learnable prompts are prepended to the keys and values of the attention head input. Therefore, prefix-tuning also incurs more parameters than prompt-tuning, but typically results in better performance.
P-tuning v2. Liu et al. (2021b) proposed P-tuning v2 as a successor of prefix-tuning. Instead of prepending continuous prompts only on the input layer, P-tuning v2 prepends them for every layer of the pre-trained model. This simple modification results in more parameters, but better performance. In particular, P-tuning v2 matches the performance of full fine-tuning on a number of NLP benchmarks and is effective across model sizes (330M to 10B parameters).
EWC and L2. EWC (Kirkpatrick et al., 2017) is an established technique for continual learning based on regularization. EWC helps to prevent forgetting of previous tasks when learning a new task by protecting parameters that are important for previously learned tasks. EWC uses the Fisher information matrix as a regularization term, which measures the sensitivity of each parameter with respect to each task and, thus, indicates which parameters require protection. L2 is used as a baseline in the original EWC publication (Kirkpatrick et al., 2017) and utilizes the L2 penalty to protect previously learned weights from being changed.
C Multi-Domain Decision Transformer (MDDT)
In this section, we describe our Multi-Domain Decision Transformer (MDDT) architecture in more detail. We pre-train a MDDT on the 50 collected pre-training datasets simultaneously (Section 3.1). Then we fine-tune the pre-trained model to 16 new tasks from two domains. To achieve this, we extend the DT architecture proposed by Chen et al. (2021) to handle inputs from multiple domains with varying state and/or action spaces.
Unified state-space. While all Meta-World tasks share a 39-dimensional state-space, the dimensionailities of the state-spaces in DMControl vary between 3 and 67 continuous values (See Appendix A). This variation reflects the diverse robot morphologies across DMControl tasks. To deal with these heterogeneous state spaces, we construct a unified state-space that spans all dimension of DMControl and Meta-World tasks. While the dimensionalities of state-spaces differ between tasks, they may share common observation types, such as orientation, current velocity, or arm positions of the robot. In total, DMControl encompasses 24 different observation types. Our constructed state-space unifies all of these observation-type attributes, amounting to 204 continuous dimensions in total. For example, the first 13 dimensions correspond to the current orientations of the robot and the last dimension corresponds to the current height. Dimensions that are unused by a certain task are zero-padded. Thus, the constructed state-space covers all environments considered in this work. Finally, as in the original DT publication (Chen et al., 2021), the states are embedded with a linear layer before serving as input to the MDDT.
We choose this scheme for handling varying state-spaces instead of discretisation followed by padding, as was used in prior work (Reed et al., 2022). Discretisation of states in combination with padding drastically increases the context length and, thus, the computational requirements. For example, when tokenising a trajectory of five timesteps in Meta-World, the context would expand to 195 tokens (39 dimensions * 5 timesteps) for states alone. In contrast, our scheme results in only 5 state tokens.
Action tokenisation & autoregressive prediction. To handle action spaces of varying dimension- alities, we tokenize each action dimension and autoregressively predict action tokens, similar to Reed et al. (2022) and Brohan et al. (2022). In the environments we consider, the action-spaces vary between 1 and 6 dimensions and all actions are bounded by [-1, 1]. We use a min-max tokenisation to discretise every continuous action dimensions into 64 bins and then pad to the maximum action dimension. Thus, we obtain 6 discrete action tokens from every action. During evaluation, we autoregressively predict the action tokens until the dimensionality of action space in the particular task is reached, and then execute the composed action. Note that autoregressive prediction of actions make the prediction task harder, compared to simply predicting the ground truth action, as is commonly done for continuous control tasks.
Context representation. Similar to Lee et al. (2022), the policy  $\pi_{\theta}$  with trainable parameters  $\theta$  predicts the next action  $a_{t}$  from a given context of the  $K$  most recent states, RTGs, actions and rewards. States and actions are pre-processed as described earlier, and then embedded using separate linear layers. Similarly, we embed the continuous scalars for RTGs and rewards using linear layers. Unlike Lee et al. (2022), we do not discretise rewards or RTGs. However, similar to Chen et al. (2021), we scale the rewards and RTGs such that all RTGs lie roughly within the range [0, 10]. To achieve this, we use separate reward scales of 200 and 100, for Meta-World and DMControl, respectively. Thus, each timestep is made up of 9 embedded tokens ( $1 \times$  state,  $1 \times$  RTG,  $6 \times$  action,  $1 \times$  reward). We use a context-length of 5 timesteps for all our experiments, and therefore, the final sequence length expands to 45 embedded tokens (5 timesteps * 9 tokens per timestep).
Positional encodings. Once the sequence is embedded, we add positional information to every embedded token. As Chen et al. (2021), we learn an embedding for each timestep and add this embedding to every embedded token (i.e., states/action/reward/rtg tokens) in a timestep. This provides the DT with vital information about the absolute position within the current episode. This is important since the choice for an action may differ at the beginning vs. at the end of an episode. We also explored other positional encodings in preliminary experiments, but found learned time-step based embeddings to work best.
Training & Evaluation. 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. During evaluation, we set the target return to the maximum observed return in the respective dataset, instead of utilizing the expert-action inference mechanism proposed by Lee et al. (2022). However, we also found that constant proxies per domain (e.g., 2000 and 1000) work well.
These simple yet effective modifications enable processing observations and actions that originate from different environments with varying state and/or action spaces at training and test time. However, a few notable limitations remain:
- The constructed state-space is specific to the environments we consider in our work and, thus, not directly transferable to other environments. Therefore, a new unified state-space may have to be constructed (same procedure), when using a different environment.
- The discretised actions are padded to the maximum action dimension. Therefore, an input sequence for an environment with a smaller than the maximum action dimension, may contain multiple redundant tokens. However, this design decision simplified our implementation considerably.
- During evaluation, we use autoregressive prediction of action dimensions. This makes the prediction task considerably harder, when compared to predicting all dimensions at once. Thus, this can result in lower overall performance (see Appendix 23).
D Dataset
To collect our dataset, we train expert agents via SAC (Haarnoja et al., 2018). We train a separate expert on each of the 66 considered tasks.
Table 2: Performance scores for data collection on MT40.
Task	|S|	|A|	Success Rate	Reward
assembly-v2	39	4	0.0	1206.9
basketball-v2	39	4	0.9	1375.95
bin-picking-v2	39	4	0.0	474.81
box-close-v2	39	4	0.0	759.15
button-press-topdown-v2	39	4	1.0	1299.24
button-press-topdown-wall-v2	39	4	1.0	1296.16
button-press-v2	39	4	1.0	1430.44
button-press-wall-v2	39	4	1.0	1508.16
coffee-button-v2	39	4	1.0	1499.17
coffee-pull-v2	39	4	1.0	1313.88
coffee-push-v2	39	4	0.6	508.14
dial-turn-v2	39	4	0.8	1674.29
disassemble-v2	39	4	1.0	1396.55
door-close-v2	39	4	1.0	1535.4
door-lock-v2	39	4	1.0	1712.65
door-open-v2	39	4	1.0	1544.32
door-unlock-v2	39	4	1.0	1733.64
drawing-close-v2	39	4	1.0	1845.92
drawing-open-v2	39	4	1.0	1710.65
f aquet-open-v2	39	4	0.9	1727.98
hand-insert-v2	39	4	1.0	1607.17
handle-press-v2	39	4	1.0	1854.79
handle-pull-side-v2	39	4	1.0	1613.72
handle-pull-v2	39	4	1.0	1581.75
lever-pull-v2	39	4	1.0	1449.05
peg-insert-side-v2	39	4	1.0	1545.19
pick-out-of-hole-v2	39	4	1.0	1435.64
pick-place-v2	39	4	0.0	6.59
pick-place-wall-v2	39	4	0.1	702.59
plate-slide-back-side-v2	39	4	1.0	1766.24
plate-slide-back-v2	39	4	1.0	1773.56
plate-slide-side-v2	39	4	1.0	1663.35
plate-slide-v2	39	4	1.0	1667.35
reach-v2	39	4	1.0	1858.99
reach-wall-v2	39	4	1.0	1831.14
soccer-v2	39	4	0.4	445.84
stick-push-v2	39	4	1.0	1470.71
sweepinto-v2	39	4	1.0	1761.69
sweep-v2	39	4	1.0	1458.35
window-open-v2	39	4	1.0	1537.59
Average	-	-	0.84 ± 0.34	1414.62 ± 439.39
Figure 9: Histogram for success rates over all tasks. The majority of all MT50 tasks can be solved successfully at the end of training by single-task SAC, while on a minority of tasks the agents fails to achieve successful completion.
Table 3: Performance scores for data collection on CW10.
Task	|S|	|A|	Success Rate	Reward
facet-close-v2	39	4	1.0	1768.87
hammer-v2	39	4	1.0	1632.21
handle-press-side-v2	39	4	1.0	1842.17
peg-unplug-side-v2	39	4	1.0	1373.45
push-back-v2	39	4	1.0	1373.32
push-v2	39	4	1.0	1672.88
push-wall-v2	39	4	1.0	1594.37
shelf-place-v2	39	4	1.0	1376.92
stick-pull-v2	39	4	1.0	1344.29
window-close-v2	39	4	1.0	1426.45
Average	39	4	1.0 ± 0.0	1540.49 ± 184.43
D.1 Meta-World
For each of the 50 Meta-World tasks, we train for 2M steps and record the entire replay buffer. Thus, the final Meta-World datasets contain 2M transitions (i.e., state-RTG-action-reward tuples). Each trajectory is 200 steps long, and each dataset consists of 10K trajectories. This amounts to 100M transitions in total for all 50 tasks, 80M and 20M for MT40 and CW10, respectively.
We use the same network architecture as Wolczyk et al. (2021), 4 linear layers with 256 neurons, LayerNorm (Ba et al., 2016) after the first layer, and LeakyReLU (Maas et al., 2013) activations for both, actor and critic. We use  $\alpha = 0.01$  for the LeakyReLU instead of  $\alpha = 0.2$  used by Wolczyk et al. (2021), since it leads to performance improvements.
After every 50 interaction steps, we perform 50 gradient steps using Adam (Kingma and Ba, 2015). We keep other parameters fixed at their default values in stable-baselines3 (Raffin et al., 2021). This includes the learning rate of  $3e^{-4}$ , batch size of 256, discount factor 0.99, and automatic entropy tuning. The target networks are synced after every update step with an update coefficient of 0.005. We evaluate the current policy after every 10K interaction steps. Each evaluation run consists of 10 individual evaluation episodes, and scores are averaged over all 10 episodes. The success rates provide a measure of trajectory quality (TQ, Schweighofer et al., 2022) for all individual datasets. Additional measures such as state-action coverage (SACo) may also be computed.
In Tables 2 and 3, we list the success rates and average rewards achieved by the expert agents on MT40 and CW10 tasks, respectively. In addition, we show the success rate distribution in Figure 9 and the learning curves on MT40 and CW10 in Figures 10 and 11, respectively. Across tasks, we observe notable differences in learning behaviour. The expert learns some tasks already after a few thousand interaction steps, while it requires more interaction steps for others. For some it does not
Figure 10: Learning curves for data-collection runs on all MT40 tasks with SAC. We train for 2M environment interaction steps on each task and record the entire replay buffer.
Figure 11: Learning curves for data-collection runs on all CW10 tasks with SAC. We train for 2M environment interaction steps on each task and record the entire replay buffer.
exhibit any learning progress at all (e.g. assembly-  $\nu 2$ ). Overall, the expert achieves average success rates of  $84\%$  and  $100\%$  on MT40 and CW10, respectively.
D.2 DMControl
For each of the 16 considered DMControl tasks, we train for 1M steps. Therefore, the final DMControl datasets contain 1M transitions per task, amounting to 16M transitions in total (10M for DMC10, 6M for DMC6). All trajectories are of the same length and contain 1000 timesteps.
We utilize the same network architecture as for Meta-World (Section D.1), using 1024 neurons per layer instead of 256. Furthermore, we perform one gradient step per interaction step and keep all other hyper-parameters the same as on Meta-World.
We list the average rewards achieved by the expert agents on DMC10 and DMC6 tasks in Table 4. The corresponding learning curves for all DMC10 tasks are available in Figure 12. On both splits, the agent performs close to the best possible return of 1000, achieving a mean reward of 788 on DMC10 and 840 on DMC6.
Table 4: Performance scores for data collection on (a) DMC10 and (b) DMC6.
(a) DMC10
Task	|S|	|A|	Reward
cartpole-balance	5	1	967.49
finger-turneasy	12	2	940.57
finger-turn-hard	12	2	940.15
fish-upright	21	5	787.87
hopper-stand	15	4	417.49
pendulum-swingup	3	1	682.85
point_mass-easy	4	2	842.03
reacher-hard	6	2	945.18
walker-run	24	6	402.36
walker-stand	24	6	957.61
Average	-	-	788 ± 219
(b) DMC6
Task	|S|	|A|	Reward
ball_in_cup-catch	8	2	971.1
cartpole-swingup	5	1	823.48
cheetah-run	17	6	411.76
finger-spin	9	2	961.85
reacher-easy	6	2	890.03
walker-walk	24	6	928.84
Average	-	-	840 ± 216
E Pre-training
We train our MDDT for a total of 1M update steps, with context length of 5 transitions (45 tokens). We use a learning rate of  $1e^{-4}$  and 4000 linear warm-up steps, followed by a cosine decay to  $1e^{-6}$ . Furthermore, we use gradient clip of 0.25, weight decay of 0.01, dropout of 0.2, a batch size of
Figure 12: Learning curves for data-collection runs on all DMC10 tasks with SAC. We train for 1M environment interaction steps on each task and record the entire replay buffer.
Figure 13: Learning curves for data-collection runs on all DMC6 tasks with SAC. We train for 1M environment interaction steps on each task and record the entire replay buffer.
1024 sequences and train using the AdamW optimizer (Loshchilov and Hutter, 2018). We base our implementation on the DT implementation in the transformers library (Wolf et al., 2020), and keep their default values for remaining parameters.
We use the same underlying GPT-2 like network architecture as Lee et al. (2022). However, we do not use their proposed expert-action inference mechanism. Instead, we set the target return to the maximum return in the respective dataset and use a reward scale of 200 and 100 for Meta-World and DMControl, respectively. Furthermore, we use timestep-based positional embeddings, as described in Appendix C.
In Table 5 and Figure 14, we show the aggregate scores over all tasks at the end of training, and learning curves across different model sizes, respectively. We vary the number of layers, the number of heads, and the embedding dimension. We also experimented with higher context lengths, but only found slight performance gains at higher computational cost. Due to the excessive computational training cost, we only train one model per size. Generally, there is a trend that more complex models
Table 5: Pre-training scores for multi-domain DT variants (40M-200M) trained on MT40 and DMC10. We vary the number of layers, the number of heads, and the embedding dimension. We report success rates for MT40, normalized scores for DMC10, and mean rewards across both domains.
Layers	Heads	Embedding Dim	Params	MT40	DMC10	Mean Reward
6	12	768	40M	0.6	0.84	983.92
8	12	768	58M	0.6	0.93	1005.6
8	16	1024	102M	0.61	1.21	1055.75
10	20	1280	200M	0.6	1.33	1030.1
(a) MT40
(b) DMC10
Figure 14: Learning curves for multi-domain pre-training runs on MT40 and DMC10 with different architectural choices of the Transformer network (L = # of layers, H = # of heads, E = embedding dimension). We train every model variant for 1M update steps with a batch size of 1024.
learn faster and attain higher performance on DMC10. However, for all our subsequent experiments, we used the model with 40M parameters, as it achieves a good trade-off between performance and training time. After pre-training, we evaluate this model across 3 environment seeds and attain an average success rate of  $60\%$  ( $\pm 0.5\%$ ) on MT40 and a normalized score of  $94\%$  ( $\pm 11\%$ ) on DMC10.
F Task Separation Analysis
To better understand what the DT learned, we visualize token embeddings using t-SNE (Van der Maaten and Hinton, 2008). First, we sample 256 state-RTG-action-reward sequences of length 5 (i.e., 45 tokens) for each task and record the hidden states of each (frozen) Transformer block. For each sequence, we then average the hidden states over the sequence, either across the entire sequence or individually per token-type (i.e., for states, RTGs, actions, rewards). We use the 40M pre-trained MDDT for our analysis. Thus, we end up with 256 768-dimensional vectors per task. We cluster all vectors using t-SNE projecting into a two-dimensional space with a perplexity of 30, 10K optimization steps, and the cosine distance.
In Figure 15, we show the t-SNE visualizations of all tokens combined, as well as for state, action, RTG, and reward tokens individually on 10 tasks in MT40 and DMC10. For embedded state tokens, we observe good cluster separation. Also, for all tokens combined and actions some clusters can be observed, but no task-specific separation. For rewards and RTG, we do not observe distinct clusters. This is expected, since rewards and RTGs do not vary substantially among tasks.
● assembly-v2	● button-press-topdown-v2	● coffee-button-v2	● finger_turn_easy	● point_mass_easy
● basketball-v2	● button-press-topdown-wall-v2	● coffeePull-v2	● finger_turn-hard	● reacher-hard
● bin-picking-v2	● button-press-v2	● cartpole_balance	● hopper_stand	● walker_run
● box-close-v2	● button-press-wall-v2	● fish_upright	● pendulum_swingup	● walker_stand
rreee eee
(b) States
(c) Return-to-go
(d) Reward
(e) Actions
(f) All
Figure 15: t-SNE clustering for state, action, RTG, and reward embeddings individually, as well as for all token embeddings combined for the first ten tasks in MT40 and DMC10.
Table 6: Aggregate scores for all single-task fine-tuning experiments on CW10 and DMC6 datasets.
Method	CW10	DMC6
Success Rate	Mean Reward	Norm. Score	Mean Reward
FT	0.83 ± 0.21	1449.15 ± 329.56	1.14 ± 0.09	832.46 ± 198.22
FT-head+last	0.7 ± 0.24	1232.9 ± 422.91	0.69 ± 0.37	547.76 ± 324.0
FT-head	0.03 ± 0.11	127.23 ± 205.68	0.08 ± 0.09	81.23 ± 76.37
Adapters	0.76 ± 0.23	1330.29 ± 346.19	0.85 ± 0.23	662.63 ± 255.08
LoRA	0.75 ± 0.24	1319.89 ± 365.68	0.79 ± 0.26	616.8 ± 276.89
(IA)3	0.64 ± 0.35	1131.7 ± 436.49	0.23 ± 0.26	195.95 ± 195.35
P-tuning v2	0.31 ± 0.39	556.97 ± 571.94	0.11 ± 0.11	102.33 ± 94.55
Prefix-tuning	0.12 ± 0.25	282.6 ± 468.18	0.12 ± 0.14	112.24 ± 117.92
Prompt-tuning	0.09 ± 0.23	178.25 ± 373.56	0.07 ± 0.08	76.27 ± 74.87
FT-MT-scratch	0.85 ± 0.07	1476.61 ± 395.13	0.92 ± 0.03	648.49 ± 354.7
FT-MT-pre-trained	0.92 ± 0.02	1491.76 ± 332.0	1.04 ± 0.17	764.38 ± 230.28
G Single-Task Fine-Tuning
We conduct a broad evaluation of different FT, PEFT and PBT methods to adapt the MDDT to each of the held-out fine-tuning tasks in a single-task setting (Section 3.2). In this setting, we present results for fine-tuning on all CW10 and DMC6 tasks individually.
We compare 9 different methods on CW10 and DMC6, as listed in Section 3.2. We train for 100K update steps on each task. Every 10K update steps, we evaluate the current model within the actual environment for 10 evaluation episodes and average performance across the evaluation episodes. The final performance scores are aggregated over all tasks in the task sequence. We show the aggregated task performances on CW10 and DMC6 in Table 6. In addition, we provide the respective learning curves in Figure 16. In Figure 17 we compare the fraction of trainable parameters vs. attained performance for all considered fine-tuning techniques on DMC6.
For FT variations, Adapters, LoRA and (IA) $^3$ , we use a learning rate of  $1e^{-4}$ . For all PBT methods, we increase the learning rate to  $1e^{-3}$  which resulted in better performance. For Adapters, we set a reduction factor of 16 for the down-projection. For LoRA, we set the rank to 8 by default, but conduct an ablation on this choice in Section G.1. For all prompt-tuning approaches, we use a prompt length of 25 and a dropout rate of 0.2. We found these values to work well in preliminary experiments. For LoRA, Adapters and (IA) $^3$  we do not use Dropout (Srivastava et al., 2014).
Moreover, we add the performance scores for two multi-task baselines in Table 6. While the first one is trained from scratch (FT-MT-scratch), the second is initialised from the pre-trained model (FT-MT-pre-trained). Both model variants are trained for 1M update steps on the multi-domain mix of CW10 and DMC6. The rationale for this comparison, is to investigate the effect of pre-training when learning new tasks. Indeed, we observe significantly higher fine-tuning scores on both domains for the pre-trained model.
G.1 What rank is required in LoRA?
To investigate the impact of the rank  $r$  in LoRA, we performed a hyper-parameter search varying  $r$  across the values  $\{1, 2, 4, 8, 16, 32\}$ . We present the results of our experiments on CW10 and DMC6 in Figure 18. Specifically, we plot the performance against the percentage of parameters trained. Note that the x-axis is on a log-scale. For all experiments presented in Sections 3.2 and 3.3, we used a rank of 8 for LoRA. As expected, increasing the rank can improve performance on both considered domains. However, this improvement comes at the expense of more trainable parameters.
G.2 What matrices should be modulated using LoRA?
Another important design choice when using LoRA is which of the weights of the pre-trained model to modulate. To investigate this, we conduct an ablation study in which we vary the modulation targets. In principle, LoRA can be applied to any weight matrix in the pre-trained model. In our
(b) CW10
(c) DMC6
Figure 16: Learning curves for single-task fine-tuning experiments on (a) CW10 and (b) DMC6. On every task, we train for 100K timesteps and consequently average the performances across all tasks in the data split.
Figure 17: Normalized score vs. fraction of parameters trained for various fine-tuning techniques on single-task experiments for DMC6.
ablation study, we specifically focus on modulating the keys, values, and queries in the self-attention mechanism, as well as the position-wise feedforward layer.