SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Michael Beukman<sup>1,2*</sup> Devon Jarvis<sup>1,3</sup>  
Richard Klein<sup>1</sup> Steven James<sup>1</sup> Benjamin Rosman<sup>1</sup>
<sup>1</sup>University of the Witwatersrand <sup>2</sup>University of Oxford <sup>3</sup>University College London
While reinforcement learning has achieved remarkable successes in several domains, its real-world application is limited due to many methods failing to generalise to unfamiliar conditions. In this work, we consider the problem of generalising to new transition dynamics, corresponding to cases in which the environment's response to the agent's actions differs. For example, the gravitational force exerted on a robot depends on its mass and changes the robot's mobility. Consequently, in such cases, it is necessary to condition an agent's actions on extrinsic state information and pertinent contextual information reflecting how the environment responds. While the need for context-sensitive policies has been established, the manner in which context is incorporated architecturally has received less attention. Thus, in this work, we present an investigation into how context information should be incorporated into behaviour learning to improve generalisation. To this end, we introduce a neural network architecture, the Decision Adapter, which generates the weights of an adapter module and conditions the behaviour of an agent on the context information. We show that the Decision Adapter is a useful generalisation of a previously proposed architecture and empirically demonstrate that it results in superior generalisation performance compared to previous approaches in several environments. Beyond this, the Decision Adapter is more robust to irrelevant distractor variables than several alternative methods.
Reinforcement learning (RL) is a powerful tool, and has displayed recent success in several settings [1-4]. However, despite its potential, RL faces a significant challenge: generalisation. Current agents and algorithms struggle to perform well beyond the narrow settings in which they were trained [5-7], which is a major limitation that hinders the practical application of RL. For RL algorithms to be effective in real-world scenarios, they must be capable of adapting to changes in the environment and performing well in settings that are different, but related, to those they trained on [7].
To illustrate this point, consider a legged robot trained to walk on a tiled floor. While the agent would perform well in these conditions, it will likely struggle in other reasonable settings, such as walking on asphalt or when carrying cargo. When walking on another surface, the friction characteristics would differ from what the robot trained on, which may cause it to slip. When carrying additional weight, the robot is heavier, with a potentially different center of mass. This may mean that the robot must exert more force to lift its legs or be more careful to prevent falling over. This limitation is far from ideal; we want our agents to excel in tasks with minor differences in dynamics (the effect of the agent's actions) without requiring additional training.
*Correspondence to mbeukman@robots.ox.ac.uk. Work done while at the University of the Witwatersrand, now at the University of Oxford.
Several methods have been proposed to address this problem, including training a single agent in various settings in the hope that it learns a generalisable and robust behaviour [3, 8-11]. However, this approach is not without its drawbacks, as it fails when variations are too large for a single behaviour to perform well in all settings [12, 13]. For example, using the same behaviour when walking unencumbered and under load may lead to suboptimal performance in one or both cases.
This observation has led to other methods that instead use contextual information to choose actions, allowing the agent to adapt its behaviour to the setting it is in. For instance, if the agent knows when it is walking unencumbered vs. carrying something heavy, it can perform slightly differently in each case [14]. The context differs conceptually from the state information, which generally includes data that the robot observes using its sensors, such as its joint angles, acceleration, and camera views. The crucial distinction between context and state is that context changes at a longer timescale compared to state; for instance, the floor friction and additional mass remain the same for a long time, whereas the robot's joint angles change after every action [15]. Currently, the predominant approach is to use context as part of the state information, ignoring the conceptual difference between these two aspects [16-18]. However, this may lead to the agent confounding the context and state, and thus exhibit worse generalisation [15]. This problem can be exacerbated in real-world settings, where identifying which variables affect the dynamics may be challenging [18]. This may lead to several irrelevant variables being added to the context, further expanding the state space and making learning more difficult.
To address this problem, we introduce an approach to incorporating context into RL that leads to improved generalisation performance. Our method separates the context and state, allowing the agent to decide how to process the state information based on the context, thus adapting its behaviour to the setting it is in. Our experimental results demonstrate the effectiveness of our approach in improving generalisation compared to other methods. We show that our approach outperforms (1) not incorporating context information at all; (2) simply concatenating context and state; and (3) competitive baselines. We also demonstrate that our approach is more robust to irrelevant distractor variables than the concatenation-based approach across multiple domains. We further provide a theoretical characterisation of problems where incorporating context is necessary and empirically demonstrate that a context-unaware method performs poorly. $^{1}$
Reinforcement learning problems are frequently modelled using a Markov Decision Process (MDP) [19, 20]. An MDP is defined by a tuple  $\langle S, \mathcal{A}, \mathcal{T}, \mathcal{R}, \gamma \rangle$ , where  $\mathcal{S}$  is the set of states,  $\mathcal{A}$  is the set of actions.  $\mathcal{T}: \mathcal{S} \times \mathcal{A} \times \mathcal{S} \to [0,1]$  is the transition function, where  $\mathcal{T}(s'|s,a)$  specifies the probability of ending up in a certain state  $s'$  after starting in another state  $s$  and performing a specific action  $a$ .  $\mathcal{R}: \mathcal{S} \times \mathcal{A} \times \mathcal{S} \to \mathbb{R}$  is the reward function, where  $R(s_t, a_t, s_{t+1}) = R_{t+1}$  specifies the reward obtained from executing action  $a_t$  in a state  $s_t$  and  $\gamma \in [0,1]$  is the environment discount factor, specifying how short-term and long-term rewards should be weighted. The goal in reinforcement learning is to find a policy  $\pi: \mathcal{S} \to \mathcal{A}$  that maximises the return  $G_t = \sum_{k=0}^{\infty} \gamma^k R_{t+k+1}$  [21].
We consider the Contextual Markov Decision Process (CMDP) [15] formalism, which is defined by a tuple  $\langle \mathcal{C},\mathcal{S},\mathcal{A},\mathcal{M}',\gamma \rangle$ , where  $\mathcal{C}$  is the context space,  $\mathcal{S}$  and  $\mathcal{A}$  are the state and action spaces respectively, and  $\mathcal{M}'$  is a function that maps a context  $c\in \mathcal{C}$  to an MDP  $\mathcal{M} = \langle S,\mathcal{A},\mathcal{T}^c,\mathcal{R}^c,\gamma \rangle$ . A CMDP thus defines a family of MDPs, that all share an action and state space, but the transition  $(\mathcal{T}^c)$  and reward  $(\mathcal{R}^c)$  functions differ depending on the context. Our goal is to train on a particular set of contexts, such that we can generalise well to another set of evaluation contexts. We focus solely on generalising over changing dynamics and therefore fix the reward function, i.e.,  $\mathcal{R}^c = \mathcal{R},\forall c\in \mathcal{C}$ .
Generalisation in RL is a widely studied problem, with one broad class of approaches focusing on robustness. Here, a single policy is trained to be robust to changes that may occur during testing, i.e., to perform well without failing catastrophically [3, 6, 8, 22-29]. Many of these approaches are successful in this goal and can generalise well when faced with small perturbations. However, using one policy for multiple environments (or contexts) may lead to this policy being conservative, as it cannot exploit the specifics of the current environment [12]. Furthermore, these policies are limited
<sup>1</sup>We publicly release code at https://github.com/Michael-Beukman/DecisionAdapter.
to performing the same action when faced with the same state, regardless of the setting. This may be suboptimal, especially when the gap between the different settings widens [13, 30].
This observation motivates the next class of methods, context-adaptive approaches. These techniques often use a version of the CMDP [15] formalism, and use the context to inform the agent's choice of action. This allows agents to exhibit different behaviours in different settings, which may be necessary to achieve optimality or generalisation when faced with large variations [13]. There are several ways to obtain this context, such as assuming the ground truth context [13], using supervised learning to approximate it during evaluation [14, 31], or learning some unsupervised representation of the problem and inferring a latent context from a sequence of environment observations [18, 32-34].
Despite the progress in inferring context, how best to incorporate context has received less attention. The most prevalent approach is to simply concatenate the context to the state, and use this directly as input to the policy [16-18, 35-38], which is trained using methods such as Soft Actor-Critic [39] or Proximal Policy Optimisation [40]. This ignores the conceptual difference between the state and context, and may lead to worse generalisation [15]. This approach may also be limited when there are many more context features compared to state features [41]. While concatenation is the most common, some other approaches have also been proposed. For instance, Biedenkapp et al. [41] learn separate representations for context and state, and then concatenate the representations near the end of the network. They also use another baseline that concatenates a static embedding of the context features to the state. Both of these techniques generalise better than the concatenation approach. Another method, FLAP [42], aims to improve generalisation by learning a shared representation across tasks, which is then processed by a task-specific adapter in the form of a linear layer. Given state  $s$ , the policy network outputs features  $\phi(s) \in \mathbb{R}^d$ , which are shared across tasks. The final action is chosen according to  $W_i \phi(s) + b_i$ , with a unique weight matrix  $W_i$  and bias vector  $b_i$  which are learned separately for each task. Concurrently, a supervised learning model is trained to map between transition tuples  $(s_t, a_t, r_{t+1}, s_{t+1})$  and these learned weights. At test time, this supervised model generates the weights, while the policy network remains fixed. FLAP generalises better to out-of-distribution tasks compared to other meta-learning approaches [36, 43, 44]. However, this approach requires a separate head for each task, which may scale poorly if we have many tasks or see each task only a few times.
One promising approach that has been used in other fields is Feature-wise Linear Modulation (FiLM) [45, 46]. In this method, features in one modality (e.g., natural language text) linearly modulate the neural-network features obtained from another modality (e.g., visual images). An example of this would be processing a visual scene, modulated by different natural language questions which would lead to a different final answer. In RL, Benjamins et al. [13] introduce cGate, which follows a similar procedure to FiLM; in particular, the neural network policy receives the state as input and outputs an action. Before the final layer, however, the context is first transformed by a separate network, and then used to modulate the state-based features. This approach showed promise, but restricted the context features' effect on the state-based features to be linear.
Another research area that has garnered more attention recently is learning general foundation models for control [48-50]. Inspired by recent work in fields such as natural language processing [51] and computer vision [52], these models aim to provide a general foundation that can be easily fine-tuned for particular tasks of interest. For instance, Gupta et al. [53] learns a general controller for several different robot morphologies. The robot's configuration is encoded and passed to a transformer [54], allowing one model to control a wide variety of morphologies. While these works directly train policies, Schubert et al. [55] instead train a dynamics model, and use a standard planning technique—model-predictive control [56-58]—to choose actions. They find that this approach achieves better zero-shot generalisation compared to directly learning a policy. On a different note, Sun et al. [59] pre-train a foundation model that predicts observations and actions. Then, during fine-tuning, they train a policy to perform a particular task using the pre-trained representations, which enables the use of either imitation learning or reinforcement learning in the downstream fine-tuning stage.
Finally, much work in meta-reinforcement learning [60] relates to the problem of generalisation. Meta-learning approaches generally aim to use multiple tasks during training to learn how to learn; this knowledge can then be used during testing to rapidly adapt to a new task [43, 60-67]. In this problem setting, Beck et al. [68] use a recurrent model to encode the current task, and generate the weights of the agent policy based on the task encoding. Sarafian et al. [69] take a similar approach, but
2Benjamins et al. [13] introduced the cGate method, and then published a revised version of the paper—not containing cGate—under the same name. We therefore cite both versions [13, 47].
instead generate the weights using the environment state, and processing the encoded context using the generated network. Another approach, MAML [44] aims to learn the weights of a neural network such that, for any new task, performing only a few gradient updates would lead to a high-performing task-specific network. While meta-learning can lead to highly adaptable agents, many methods require multiple episodes of experience in the target domain [36, 44], and may therefore be less well-suited to the zero-shot setting we consider [7].
4 Theoretical Intuitions
Now, while only a context-conditioned policy is guaranteed to be optimal on a general CMDP [47], in many cases simply training an unaware policy on a large number of diverse environments can lead to impressive empirical generalisation [5]. In this section, we characterise and unify these two observations. In particular, for a specific problem setting—where the context defines the target location an agent must travel to—we find that:
(i) For some context sets, an unaware policy will perform arbitrarily poorly on average, due to it being forced to take the same action in the same state, regardless of context.
(ii) However, for other context sets—where the different contexts are similar enough—an unaware policy can perform well on average, as it can simultaneously solve each task.
We defer the formal statement of this theorem and its proof to Appendix A, but here we briefly provide some intuition. Overall, our results rely on the fact that a context-unaware agent must perform the same action in the same state, regardless of context. Therefore, in a setting where we have different goals (indicated by the context), and making progress on one goal leads to making negative progress on another, a context-unaware agent cannot simultaneously perform well on both of these contexts. By contrast, if the goals are aligned in the sense that making progress on one leads to progress on another, a context-unaware agent can perform well on average. A context-aware agent, however, can always choose the correct action in either scenario.
(a) Non-overlapping
(b) Overlapping
Figure 1: An illustration of the two cases discussed above, with the agent's start state being the origin. The agent will receive a reward if it enters the circle corresponding to the current context and zero reward otherwise. In (a) there is no overlap between the goals, so an unaware policy (red) must visit each goal in sequence, whereas a context-aware policy (green) can go directly to the correct goal. In (b), the goals overlap, so an unaware policy can go directly to the joint intersection point of the goals, leading to only a slightly suboptimal value function compared to the optimal context-aware policy.
Fig. 1 visually illustrates these two cases and summarises the implications. In Fig. 1a, if there is no overlap between goals, an unaware policy is forced to visit each one in sequence. This is because it must perform the same action in the same state, regardless of the context, meaning that it cannot always travel to the correct goal. This policy will therefore have a low value function compared to the optimal context-aware policy. By contrast, in Fig. 1b, a single unaware policy can directly go to the intersection of all the goals, leading to only slightly less return than the optimal context-aware
policy. While the formal result in Appendix A merely shows the existence of some problems with these properties, we believe the concept generalises to more complex settings. In particular, if the problem exhibits some structure where the same policy can simultaneously make progress on each context, an unaware model can perform well. However, if making progress on one goal leads to the agent moving away from another, a context-unaware policy can at best perform arbitrarily worse than a context-aware one, and incorporating context is crucial. We empirically substantiate these findings in Section 7 and formally consider this example in Appendix A. Finally, see Appendix B for more details about how these theoretical results connect back to our empirical observations.
5 The Decision Adapter
An adapter is a (usually small) neural network with the same input and output dimensions. This network can be inserted between the layers of an existing primary network to change its behaviour in some way. These modules have seen much use in natural language processing (NLP) [70-74] and computer vision, usually for parameter-efficient adaptation [75-81]. In these settings, one adapter is usually trained per task, and the base model remains constant across tasks. Our method is inspired by the use of adapters in these fields, with a few core differences.
First, since our aim is zero-shot generalisation, having a separate adapter module for each task (corresponding to context in our case) provides no clear means for generalisation. Relatedly, since tasks are generally discrete in NLP, but we allow for continuous contexts, this approach would be problematic. Furthermore, the number of learnable parameters increases with the number of training contexts, presenting a computational limitation to the scalability. To circumvent these problems, we use a hypernetwork [82] to generate the weights of the adapter module based on the context. This allows the agent to modify its behaviour for each context by having a shared feature-extractor network which is modulated by the context-aware adapter. Moreover, since the hypernetwork is shared, experience gained in one context could transfer to another [72]. Second, instead of having separate pre-training and adaptation phases as is common in NLP [76], we do not alter the standard RL training procedure at all. We change only the architecture and train normally—allowing for the easy integration of our approach into standard RL libraries. Finally, our method is agnostic to the exact RL algorithm and can be applied to most off-the-shelf algorithms without much additional effort.
Algorithm 1 Decision Adapter—Changes to the standard forward pass in blue.
1: procedure ADAPTERFORWARD  $(s\in S,c\in \mathcal{C})$    
2:  $x_{1} = s$    
3: for  $i\in \{1,2,\dots ,n\}$  do   
4: if  $A_{i}\neq null$  then   
5:  $\theta_A^i = H_i(c)$  // Generate Weights   
6:  $x_{i}^{\prime} = A_{i}(x_{i}|\theta_{A}^{i})$  // Forward Pass   
7:  $x_{i} = x_{i} + x_{i}^{\prime}$  // Skip connection   
8: end if   
9:  $x_{i + 1} = L_i(x_i)$    
10: end for   
11: return  $x_{n + 1}$    
12: end procedure
Figure 2: An illustration of our network architecture, the Decision Adapter, alongside pseudocode for the forward pass. The bottom network is the hypernetwork, which generates the magenta weights for the top, primary, network. The light grey nodes in the top network are added by the adapter.
To formally describe our approach (illustrated in Fig. 2), suppose we have a state  $s \in S$ , context  $c \in \mathcal{C}$ , and a standard  $n$ -layer fully-connected neural network architecture that predicts our action  $a$ . Here  $x_{i+1} = L_i(x_i) = \sigma_i(W_i x_i + b_i)$ ,  $x_1 = s$  and  $x_{n+1} = a$ . We add an adapter module  $A_i$  between layers  $L_{i-1}$  and  $L_i$ . This adapter module is a multilayer neural network  $A_i: \mathbb{R}^{d_i} \to \mathbb{R}^{d_i}$  consisting of parameters  $\theta_A^i \in \mathbb{R}^{P_i}$ . These parameters are generated by the corresponding context-conditioned
${}^{3}{d}_{i}$  is the output dimension of layer  ${L}_{i - 1}$  and  ${P}_{i}$  is the number of parameters in the adapter.
hypernetwork  $H_{i}:\mathcal{C}\to \mathbb{R}^{P_{i}}$  (line 5 in Algorithm 1) and reshaped to form multiple weight matrices. Then, before we process  $x_{i}$  using layer  $L_{i}$ , we pass  $x_{i}$  through the adapter module  $A_{i}$  (line 6) defined by the generated weights and biases to obtain  $x_{i}^{\prime} = A_{i}(x_{i}|\theta_{A}^{i})$ . Furthermore, akin to the application of adapters in other fields [75, 76], we add a skip connection, by setting  $x_{i} = x_{i} + x_{i}^{\prime}$  (line 7). This updated  $x_{i}$  is then the input to layer  $L_{i}$  (line 9). Additionally, we could have multiple adapter modules in the primary network, with at most one module between any two consecutive layers (lines 3 and 4).
Finally, we note that the Decision Adapter model can learn to implement the same function as cGate [13], but is theoretically more powerful. While cGate uses a linear elementwise product between the context and state features, our model uses a hypernetwork to generate the weights of a nonlinear adapter module. Our hypernetwork is general enough to be able to recover this elementwise product, but is not constrained to do so. See Appendix C for a more formal treatment of this point.
6 Experimental Setup
6.1 Metrics
Our evaluation strategy is as follows: We take the trained model, and compute the average of the total episode reward over  $n = 5$  episodes on each evaluation context  $c$  to obtain a collection of contexts  $C$  and corresponding rewards  $R$ . We then calculate the Average Evaluation Reward (AER):  $AER = \frac{1}{c_{max} - c_{min}}\int_{c_{min}}^{c_{max}}R(c)dc$ , where  $c_{min}$  and  $c_{max}$  are the minimum and maximum context values in  $C$  respectively, and  $R(c)$  is the reward obtained on context  $c$ . This metric is high when the agent performs well across a wide range of contexts, and thus corresponds to generalisation performance. However, since we generally have  $\mathcal{C}_{train}\subset \mathcal{C}_{evaluation}$ , this metric also incorporates how well the agent performs on the training contexts. Despite this, it is still a useful metric for generalisation, as our training context set typically contains only a handful of contexts, with most of the evaluation contexts being unseen.
6.2Baselines