SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We use Soft-Actor-Critic [39] for all methods, to isolate and fairly compare the network architectures. We aim to have an equal number of learnable parameters for each method and adjust the number of hidden nodes to achieve this. We use the following baselines, with more details in Appendix D:
Unaware This model simply ignores the contextual information and just uses the state. This approach allows us to evaluate how well a method that does not incorporate context performs.
Concat This method forms a new, augmented state space  $S' = S \times \mathcal{C}$ , and then gives the model the concatenation  $[s; c]$  of the state and context [16-18, 37]. This baseline allows us to compare our method to the current standard approach of using contextual information.
cGate cGate [13] learns separate state  $(\phi (s))$  and context  $(g(c))$  encoders, and predicts the action  $a = f(\phi (s)\odot g(c))$  , where  $f$  is the learned policy and  $\odot$  is the elementwise product.
FLAP FLAP [42] learns a shared state representation across tasks, which is then processed by a task-specific linear layer that is generated by conditioning on the context.
Our Adapter configuration is more fully described in Appendix D.3, and Appendix E contains ablation experiments comparing the performance of different hyperparameter settings.
6.3 Environments
6.3.1 ODE
This environment is described by an ordinary differential equation (ODE), parametrised by  $n$  variables making up the context. The dynamics equation is  $x_{t + 1} = x_t + \dot{x}_t dt$ , with  $\dot{x} = c_0 a + c_1 a^2 + c_2 a^3 + \ldots$ , truncated at some  $c_{n - 1}$ . The episode terminates after 200 timesteps, with the reward function:
4The AER can be computed over only the unseen evaluation contexts, but we include the training contexts to obtain a more general measure of performance. However, the difference between the two versions is minor, see Appendix F.4.
$$
R _ {t} = \left\{ \begin{array}{l l} 1 & \text {i f} | x | <   0. 0 5 \\ \frac {1}{4} & \text {i f} | x | <   0. 5 \end{array} \right.
$$
$$
\begin{array}{l l} \frac {1}{2} & \text {i f} | x | <   0. 1 \\ \frac {1}{2 0} & \text {i f} | x | <   2 \end{array}
$$
$$
\begin{array}{c c} \frac {1}{3} & \text {i f} | x | <   0. 2 \\ 0 & \text {o t h e r w i s e} \end{array}
$$
This incentivises the agent to give control  $a$  to keep the state close to  $x = 0$ . The context in this environment is  $c = [c_0, \dots, c_{n-1}] \in \mathbb{R}^n$ . The action space is two-dimensional, continuous and bounded, i.e.,  $a_0, a_1 \in [-1, 1]$ . The action is interpreted as a complex number,  $a = a_0 + a_1 i$ , to allow the dynamics equation to be solvable, even in normally unsolvable cases such as  $\dot{x} = a^2$ . Here, restricting  $a$  to  $\mathbb{R}$  would result in  $\dot{x} \geq 0$  and an unsolvable system for initial conditions  $x_0 > 0$ . We keep only the real part of the updated state and clip it to always fall between -20 and 20.
The ODE acts as a conceptually simple environment where we can arbitrarily scale both the number (i.e.,  $n$ ) and magnitude (i.e.,  $|c_i|$ ) of the contexts and precisely measure the effects. Furthermore, since many dynamical systems can be modelled using differential equations [83-85], the ODE can be considered a distilled version of these. Finally, context is necessary to perform well in this environment, making it a good benchmark. See Appendix D for more details.
6.3.2 CartPole
CartPole [86] is a task where an agent must control the movement of a cart to balance a pole vertically placed upon it. The observation space is a 4-dimensional real vector containing the cart's position  $x$  and velocity  $\dot{x}$ , as well as the pole's angle  $\theta$  and angular velocity  $\dot{\theta}$ . We use a continuous action space where  $a \in [-1,1]$  corresponds to the force applied to the cart (where negative values push the cart to the left and positive values to the right). The reward function is  $+1$  for each step that the pole is upright. An episode terminates when the pole is tilted too far off-center, the cart's position is outside the allowable bounds or the number of timesteps is greater than 500. We follow prior work and consider the variables Gravity, Cart Mass, Pole Length, Pole Mass and Force Magnitude collectively as the context [13], even if only a subset of variables change. Here, we change only the pole length and evaluate how susceptible each model is to distractor variables. In particular, during training, we also add  $k$  additional dimensions to the context, each with a constant value of 1. Then, during evaluation, we set these values to 0. This may occur during a real-world example, where accurately identifying or inferring the pertinent context variables may be challenging [18, 87], and lead to several irrelevant variables. Overall, this means that the context is represented as a  $5 + k$ -dimensional vector. This environment is useful as it (1) is a simple setting with models that can be trained quickly, and (2) still has desirable aspects, namely a way to perturb some underlying variables to change the dynamics of the environment. Furthermore, prior work has shown that even standard, unaware RL algorithms can generalise well in this environment [88] (in the absence of distractor context variables), which allows us to investigate the benefits of context and detriment of distractors.
6.3.3 Mujoco Ant
As a more complex and high-dimensional problem, we consider Ant from the Mujoco suite of environments [89, 90]. Here, the task is to control a 4-legged robot such that it walks. The observation space consists of the robot's joint angles and velocities, as well as contact forces and torques applied to each link, totalling 111 dimensions. The action space  $\mathcal{A} = [-1,1]^8$  represents the torque applied to each of the 8 joints. Furthermore, each episode is terminated after 1000 timesteps, with a positive reward for not having fallen over at each step and moving forward. The reward function also penalises large control forces. Here, the mass is the context variable, which is 5 by default. We train on the set  $\{5,35,75\}$  and evaluate on 200 evenly spaced points between 0.5 and 100.
7.1 Generalisation Performance
In this section, we consider the ODE environment, with a training context set of  $\{-5, -1, 1, 5\}$  and  $300k$  training timesteps. To measure generalisation, we use an evaluation context range consisting of 201 equally spaced points between  $-10$  and 10 (inclusive). In the left pane of Fig. 3, we plot the average performance across the entire evaluation range (as discussed in Section 6.1) at different points during training.
Firstly, we can see that the Unaware model fails to generalise well, since the optimal actions in the same state for two different contexts may be completely different. This result substantiates our theoretical analysis in Section 4, highlighting the importance of using context in this domain. Concat, cGate and FLAP perform reasonably well, but our Decision Adapter outperforms all methods and converges rapidly. In Appendix F, we plot separate subsets of the evaluation range: Train, Interpolation and Extrapolation, corresponding to the contexts in the training set, the contexts within the region  $[-5, 5]$  and the contexts outside the convex hull of the training contexts, containing  $c \in [-10, -5) \cup (5, 10]$ , respectively. There we find that most methods, with the exception of the Unaware model, perform well on the training contexts, and the Adapter outperforms all methods when extrapolating. These results also extend to the multidimensional context case, shown in Fig. 3 (right). We train on the context set given by  $\{1, 0, -1\}^2 \cup \{5, 0, -5\}^2$ , with the context  $(0, 0)$  being omitted due to it being unsolvable. The evaluation range consists of the cartesian product  $A^2$ , where  $A$  contains 21 equally-spaced points between  $-10$  and 10. The Adapter outperforms all methods, with the Concat model coming second. FLAP and the Unaware models struggle in this case.
Figure 3: The average reward over the evaluation range at various points during training for the ODE domain. Mean performance is shown, with standard deviation (across 16 seeds) shaded. The left pane shows the one-dimensional context results, while the right pane shows the two-dimensional results.
7.2 Robustness to Distractors
For this and the next section, we consider the Concat and cGate models as baselines. The reasons for this are that (1) the Concat model performed comparably to or better than the other baselines in the previous sections, making it a representative example of the context-aware approaches; and (2) the Concat model is currently a very prevalent way of incorporating context [16-18, 35-38]; and (3) cGate is the closest approach to our work in the current literature.
The results for CartPole are shown in Fig. 4. When there are no distractor variables, the Adapter, Concat and cGate models perform comparably. In this case, each architecture achieves near the maximum possible reward for the domain. Further, we see little effect after adding just one confounding context variable. However, as we add additional distractor variables, the Concat and cGate models' performances drop significantly, whereas the Adapter's performance remains relatively stable. Strikingly, the Adapter architecture trained with 100 distractor context variables is still able to perform well on this domain and significantly outperforms the Concat model with significantly fewer (just 20) distractor variables. This demonstrates that, given only useful context, concatenating state and context is a reasonable approach in some environments. However, this is a strong assumption in many practical cases where we may be uncertain about which context variables are necessary. In such cases, the consequences of ignoring the conceptual differences between state and context are catastrophic and it is necessary to use the Adapter architecture. We find a similar result in the ODE, which we show in Appendix F.5.
As an extension to the previous experiment, here we consider using non-fixed distractor variables, i.e., sampling them from Gaussian distribution with different means. The intuition behind this experiment is similar to before, but allows for small variations within training and testing. In particular, we use  $\sigma = 0.2$  and a mean of either 0 or 1. In this experiment, a new distractor context is sampled at the start of every episode. During training this is sampled from  $\mathcal{N}(1,0.2)$  and during evaluation we
Here, for a set  $S$ ,  $S^2 = S \times S = \{(a,b) | a, b \in S\}$  denotes the cartesian product of  $S$  with itself.
Figure 4: Showing the sensitivity of the (left) Concat, (middle) Adapter and (right) cGate models to distractor context variables in CartPole. The mean and standard deviation over 16 seeds are shown.
sample from  $\mathcal{N}(0,0.2)$ . In CartPole, the results in Fig. 5 show that the Adapter is still more robust to changing distractor variables compared to either Concat or cGate. See Appendix F.6 for more details, results on the ODE domain, and additional experiments when the mean of the Gaussian does not change between training and testing.
Figure 5: The sensitivity of Concat, Adapter and cGate to distractor context variables—sampled from a Gaussian distribution—in CartPole. We show the mean and standard deviation over 16 seeds.
7.3 High-Dimensional Control
We next consider a more complex and challenging continuous-control environment, that of the Mujoco Ant. We run the same experiment as in CartPole, where we add distractor variables, with no effect on the dynamics, and different values between training and testing. These results are shown in Fig. 6. Overall, we observe a similar result to that of CartPole—the Concat and cGate models are significantly less robust to having irrelevant distractor variables. By contrast, the Adapter consistently performs well regardless of how many distractor dimensions are added.
While our Adapter model performs well empirically, it does have some limitations, which we summarise here and expand on in Appendix G. First, we find that using an incorrect and noisy context (Appendix G.1) generally resulted in the Adapter model performing worse, particularly as we increased the level of noise. However, the Adapter model still performs comparably to the Concat model in this case. Second, we generally normalise the contexts before we pass them to the models, relative to the maximum context encountered during training. While our Adapter performs well when the context normalisation is incorrect by a factor of 2 or 3, its performance does suffer when the normalisation value is orders of magnitudes too small, leading to very large contexts being input into the model (see Appendix G.2 for further discussion). Third, when training on a context range that is too narrow (see Appendix G.3) or does not exhibit sufficient variation, our model is susceptible to overfitting, which
Figure 6: Showing the sensitivity of the (left) Concat, (middle) Adapter and (right) cGate models to distractor context variables in Ant. The mean and standard deviation over 16 seeds are shown.
leads to worse generalisation. Even if our model does not overfit, an overly narrow context range would also lead to poor generalisation. Thus, ensuring that agents are trained in sufficiently diverse environments is still important. Fourth, while we assumed access to the ground-truth context, it may not always be available in practice [31]. Therefore, integrating the Decision Adapter into existing context inference methods [18, 33] is a promising avenue for future work. A final limitation of our Adapter is its longer wall-clock training time compared to other methods (given the same training timesteps and a similar number of parameters). This is likely caused by the hypernetwork that we use. However, we did not utilise strategies such as caching the generated weights for a particular context (which would be possible since context remains constant for an episode), or more effectively batching the forward passes through the Adapter which could ameliorate this issue in practice.
Beyond addressing these limitations, there are several logical directions for future work. First, we only consider state and context features to be numerical vectors; therefore, extending our approach to different modalities—for instance, image observations and natural language context—would be a natural extension. In principle, our Adapter model could be directly applied in this setting without any fundamental modifications. One option would be to explore adding the hypernetwork directly into the modality-specific model (e.g., a convolutional neural network in the case of images), and the alternative would be to integrate the adapter module in the final (fully-connected) layers of the RL policy. Second, while we focused on the problem of zero-shot generalisation, our approach may also be suitable for efficient few-shot learning. In particular, adapters in NLP are often added to a large pre-trained model and then fine-tuned while freezing the primary model [76]. Using the Decision Adapter in this paradigm is a promising direction for future work, especially in cases where the difference between training and testing is so large that we cannot expect effective zero-shot generalisation. Combining our work with fine-tuning-based [91–93] or meta-learning [36, 60, 68, 69, 94–96] approaches, which generally focus more on this case, could likewise be promising.
In this work, we consider the problem of generalising to new transition dynamics. We first illustrate that for some classes of problems, a context-unaware policy cannot perform well over the entire problem space—necessitating the use of context. We next turned our attention to how context should best be incorporated, introducing the Decision Adapter—a method that generates the weights of an adapter module based on the context. When context is necessary, our Decision Adapter outperforms all of our baselines, including the prevalent Concat model. Next, when there are some irrelevant context variables that change between training and testing, the Concat and cGate models fail to generalise, whereas our Decision Adapter performs well. This result holds in several domains, including the more complex Ant. Furthermore, our Adapter is theoretically more powerful and empirically more effective than cGate, showing that our hypernetwork-based perspective is useful when dealing with context. Ultimately, we find that the Decision Adapter is a natural, and robust, approach to reliably incorporate context into RL, in a factored way. This appears to be of particular benefit when there are several irrelevant context variables, for instance, in real-world settings [97].
Computations were performed using High Performance Computing infrastructure provided by the Mathematical Sciences Support unit at the University of the Witwatersrand. M.B. is supported by the Rhodes Trust. D.J. is a Google PhD Fellow and Commonwealth Scholar. B.R. is a CIFAR Azrieli Global Scholar in the Learning in Machines & Brains program.
A Formalising the Theoretical Foundations
In this section, we more formally consider our theoretical results from Section 4. We first define our problem formulation in Assumption A.1, then state the theorem in Theorem A.5. Finally, we prove the theorem at the end of this section.
Assumption A.1. (Problem Formulation): Suppose we have a CMDP where  $S, \mathcal{A} \subseteq \mathbb{R}^d$  and  $\gamma \in (0, 1)$ . The context  $c \in \mathbb{R}^d$  defines a goal location, with a reward function  $R_c(s, a, s') = 1$  if  $||s' - c|| < \bar{\tau}$  and 0 otherwise. The episode terminates when the reward obtained is 1. Suppose further that  $\bar{\tau} = \tau D$ , with  $D \in \mathbb{R}^+$  being the distance the agent travels in a single step and  $\tau \in \mathbb{Z}^+$  being the number of steps it takes the agent to traverse a distance of  $\bar{\tau}$ . Here  $\tau$  and  $D$  are kept fixed across contexts.
Definition A.2. (Context Averaged Value Function): Denote  $\bar{V}_{\mathcal{C}}^{\pi}(s) = \mathbb{E}_{c\sim \mathcal{C}}[V_{c}^{\pi}(s)]$ , i.e., the expected value of the value function of policy  $\pi$  in state  $s$ , over all contexts in the set  $\mathcal{C}$ . Similarly  $\bar{V}_{\mathcal{C}}^{*}(s) = \mathbb{E}_{c\sim \mathcal{C}}[V_{c}^{*}(s)]$  is the average of the value functions of the optimal policies on each context.
Definition A.3. (Value Function Optimality Ratio): Let  $0 \leq \alpha_{\mathcal{C}}^{\pi}(s) = \frac{V_{\mathcal{C}}^{\pi}(s)}{V_{\mathcal{C}}^{*}(s)} \leq 1$  denote how close the average value function of policy  $\pi$  is to the optimal context-specific value function. We use the optimality ratio instead of the optimality gap [47] because the optimality gap depends heavily on the reward scale of the environment, whereas the optimality ratio depends only on the ratios.
Definition A.4. (Minimum Inter-context Distance): Let  $d_{\text{min}} = \min_{c_i \neq c_j \in \mathcal{C}} \left\lceil \frac{|c_i - c_j|}{D} \right\rceil$  be the minimum number of steps required to travel between any two context centers.
Theorem A.5. Then, in this problem space, we have the following:
(i) For some set of contexts  $\mathcal{C}_{far}$  with cardinality  $N$ , and some state  $s$ , all deterministic unaware policies  $\pi$  will have:
$$
\alpha_ {\mathcal {C} _ {f a r}} ^ {\pi} (s) = \frac {\bar {V} _ {\mathcal {C} _ {f a r}} ^ {\pi} (s)}{\bar {V} _ {\mathcal {C} _ {f a r}} ^ {*} (s)} \leq \frac {1}{N} \frac {1 - \gamma^ {N d _ {m i n}}}{1 - \gamma^ {d _ {m i n}}} \tag {1}
$$
$$
\lim  _ {d _ {\text {m i n}} \rightarrow \infty} \alpha_ {\mathcal {C} _ {\text {f a r}}} ^ {\pi} (s) \rightarrow \frac {1}{N} \tag {2}
$$
i.e., as the contexts within  $\mathcal{C}_{far}$  move further away from each other the agent's performance scales inversely with the number of contexts.
(ii) For some set,  $\mathcal{C}_{\text{close}}$  and all states  $s$ , at least one deterministic unaware policy  $\pi_g$  will have:
$$
\frac {\bar {V} _ {\mathcal {C} _ {\text {c l o s e}}} ^ {\pi_ {g}} (s)}{\bar {V} _ {\mathcal {C} _ {\text {c l o s e}}} ^ {*} (s)} > \gamma^ {\tau} \tag {3}
$$
regardless of the number of contexts in  $\mathcal{C}_{close}$ .
Proof. (i) Consider a set of contexts  $\mathcal{C}_{far}$  such that for all  $c_{i}, c_{j} \in \mathcal{C}_{far}, i \neq j, ||c_{i} - c_{j}|| > 4\bar{\tau}$ . Consider the state  $s = \frac{1}{N}\sum_{i=1}^{N}c_{i}$ , i.e., in the middle of all contexts. Now, suppose the optimal policy per context would require only  $\beta$  steps to reach the appropriate goal. An unaware policy that reaches every goal, by contrast, must travel to each context in sequence. If this policy visits  $c_{1}, c_{2}, \ldots, c_{N}$ , in order, then the value functions for each context would be:
-  $V_{c_1}^{\pi}(s) = \gamma^{\beta}\cdot 1 = V_{c_1}^{*}(s)$
V(s) = 1
$V_{c_3}^{\pi}(s) = \gamma^{\beta +\lceil \frac{||c_1 - c_2|}{D}\rceil +\lceil \frac{||c_2 - c_3|}{D}\rceil},$
and so on. Let  $d_{min}$  be the number of steps required to move between the two closest distinct contexts in  $\mathcal{C}_{far}$ . Thus,  $d_{min} \geq \frac{4\bar{\tau} - \bar{\tau} - \bar{\tau}}{D} = 2\tau$ . Then,  $V_{c_i}^{\pi}(s) \leq \gamma^{\beta + (i-1)d_{min}}$ . Now,
$$
\begin{array}{l} \bar {V} _ {\mathcal {C} _ {f a r}} ^ {\pi} (s) = \mathbb {E} _ {c \sim \mathcal {C} _ {f a r}} \left[ V _ {c} ^ {\pi} (s) \right] \\ = \frac {1}{N} \sum_ {i = 1} ^ {N} V _ {c _ {i}} ^ {\pi} (s) \\ \leq \frac {1}{N} \sum_ {i = 1} ^ {N} \gamma^ {\beta + (i - 1) d _ {m i n}} \\ = \gamma^ {\beta} \frac {1}{N} \sum_ {i = 1} ^ {N} \gamma^ {(i - 1) d _ {m i n}} \\ = \gamma^ {\beta} \frac {1}{N} \frac {1 - \gamma^ {N d _ {m i n}}}{1 - \gamma^ {d _ {m i n}}} \\ = \bar {V} _ {\mathcal {C} _ {f a r}} ^ {*} (s) \frac {1}{N} \frac {1 - \gamma^ {N d _ {m i n}}}{1 - \gamma^ {d _ {m i n}}} \\ \end{array}
$$
Hence,  $\alpha = \frac{\bar{V}_{\mathcal{C}_{far}}^{\pi}(s)}{\bar{V}_{\mathcal{C}_{far}}^{*}(s)} \leq \frac{1}{N} \frac{1 - \gamma^{Nd_{min}}}{1 - \gamma^{d_{min}}}$ . Finally, as  $d_{min} \to \infty$ ,  $\gamma^{d_{min}} \to 0$  and  $\gamma^{Nd_{min}} \to 0$ . Hence,  $\lim_{N \to \infty} \alpha = \frac{1}{N}$ .
(ii) Consider a set of contexts  $\mathcal{C}_{close}$  such that for all  $c_{i}, c_{j} \in \mathcal{C}_{close}, ||c_{i} - c_{j}|| < \bar{\tau}$ . Then, consider an unaware policy that always travels to the joint intersection point of these contexts. This intersection point is at most  $\bar{\tau}$  away from the closest context to the agent, therefore the agent will perform at most  $\bar{\tau}$  unnecessary steps compared to the optimal policy. Hence, its value function has a lower bound of  $V_{c_{i}}^{*}(s)\gamma^{\tau}$ , and we have
$$
V _ {c _ {i}} ^ {*} (s) \geq V _ {c _ {i}} ^ {\pi} (s) \geq V _ {c _ {i}} ^ {*} (s) \gamma^ {\tau}
$$
$$
\begin{array}{l} \bar {V} _ {\mathcal {C} _ {c l o s e}} ^ {\pi} (s) = \mathbb {E} _ {c \sim \mathcal {C} _ {c l o s e}} \left[ V _ {c} ^ {\pi} (s) \right] \\ \geq \mathbb {E} _ {c \sim \mathcal {C} _ {\text {c l o s e}}} \left[ V _ {c} ^ {*} (s) \gamma^ {T} \right] \\ = \gamma^ {\tau} \mathbb {E} _ {c \sim \mathcal {C} _ {\text {c l o s e}}} \left[ V _ {c} ^ {*} (s) \right] \\ = \gamma^ {\tau} \bar {V} _ {\mathcal {C} _ {\text {c l o s e}}} ^ {*} (s) \\ \end{array}
$$
And we have:  $\frac{\bar{V}_{\mathcal{C}_{close}}^{\pi}(s)}{\bar{V}_{\mathcal{C}_{close}}^{*}(s)}\geq \gamma^{\tau}$
B Linking Theory and Practice
The formulation in Assumption A.1 has consistent dynamics and context-dependent rewards (as the context defines the goal location, and the transition dynamics do not depend on the context). This is in contrast to the problem we aim to study, which has consistent rewards, but context-dependent dynamics; however, we can show that for the environment considered in Appendix A, these two are equivalent.