In-Context Reinforcement Learning for Variable Action Spaces — part 1/4

Cleaned-text code points [0, 15861). Bibliography omitted here; source.md includes it.

# In-Context Reinforcement Learning for Variable Action Spaces

###### Abstract

Recently, it has been shown that transformers pre-trained on diverse datasets with multi-episode contexts can generalize to new reinforcement learning tasks in-context. A key limitation of previously proposed models is their reliance on a predefined action space size and structure. The introduction of a new action space often requires data re-collection and model re-training, which can be costly for some applications. In our work, we show that it is possible to mitigate this issue by proposing the Headless-AD model that, despite being trained only once, is capable of generalizing to discrete action spaces of variable size, semantic content and order. By experimenting with Bernoulli and contextual bandits, as well as a gridworld environment, we show that Headless-AD exhibits significant capability to generalize to action spaces it has never encountered, even outperforming specialized models trained for a specific set of actions on several environment configurations. Implementation is available at: https://github.com/corl-team/headless-ad .

###### Keywords:

## 1 Introduction

The transformer architecture, first introduced by Vaswani et al. (2017) , has been widely adopted in key areas of machine learning, including natural language processing ( Radford et al., 2018 ; Devlin et al., 2018 ) , computer vision ( Dosovitskiy et al., 2020 ) and sequential decision-making ( Chen et al., 2021 ) . One major feature of transformers is in-context learning (ICL), which makes it possible for them to adapt to new tasks after extensive pre-training ( Brown et al., 2020 ; Liu et al., 2023 ) . Recent developments, such as Algorithm Distillation (AD) by Laskin et al. (2022) and Decision Pretrained Transformer (DPT) by Lee et al. (2023) , have successfully employed transformer ICL abilities in sequential decision-making. These models are capable of predicting the next action based on a query state and history of environment interactions, which inform them about task objectives and environment dynamics. While effective at generalizing across various reward distributions ( Laskin et al., 2022 ; Lee et al., 2023 ) and transition functions ( Raparthy et al., 2023 ) , their adaptability to new action spaces remains unexplored and limited by architectural constraints.

Creating models that can adapt to new action spaces is essential for building the foundation of decision-making systems in order to enable large-scale pretraining across various environments and address real-world problems ( Jain et al., 2020 ; Chandak et al., 2020 ; London & Joachims, 2020 ; Jain et al., 2021 ) . With this in mind, our research focuses on variable discrete action spaces, with the notion of variability illustrated in the Figure 1 . In our study, we reveal the limitations of the Algorithm Distillation model from prior work, such as its diminished performance upon changes in action semantics as well as architectural constraints when handling varying action space sizes (see Figure 3 ).

Our solution, Headless-AD , is an architecture and training methodology tailored to effective generalization on new action spaces. We employ an approach similar to Wolterpinger ( Dulac-Arnold et al., 2015 ) and Headless-LLM ( Godey et al., 2023 ) by encoding actions with random embeddings ( Kirsch et al., 2023 ) and directly predicting these embeddings. This way, we remove the direct connection between the model output layer and the action space structure.

Through experiments using Bernoulli and contextual bandits, and a darkroom environment with changing action spaces, we demonstrate that Headless-AD is capable of matching the performance of the original data generation algorithm and scaling to action spaces up to 5x larger than those seen during training. We also observed that Headless-AD can even outperform AD when they are both trained for the same action space, especially when evaluated on larger action sets. To summarize, our contributions are as follows: • We show that AD struggles with generalization on novel action spaces ( Section 2 ).

• We extend AD with a modified model architecture and a training strategy, called Headless-AD , for it to acquire the ability to adapt to new discrete action spaces ( Section 3 ). We demonstrate the strong generalization capabilities of Headless-AD on Bernoulli and contextual bandits, and darkroom environments ( Section 4 ).

• We perform ablations on the loss and the prompt format to highlight the importance of Headless-AD’s design choices ( Section 5 ).

loh

## 2 Algorithm Distillation Struggles with Novel Action Spaces

Algorithm Distillation (AD) ( Laskin et al., 2022 ) is a transformer model trained to autoregressively predict the next action given the history of previous environment interactions and a current state. Formally, the history is defined as: h t = ( o 0 , a 0 , r 0 , … , o t , a t , r t ) , h_{t}=\left(o_{0},a_{0},r_{0},\dots,o_{t},a_{t},r_{t}\right), where o o are the observations, a a are the actions and r r are the rewards. The probabilities of each action are given by a model P θ ​ ( A = a t ( n ) | h t − 1 , o t ) P_{\theta}(A=a^{(n)}_{t}|h_{t-1},o_{t}) , where n n is the action index, t t is a timestamp and θ \theta are the model weights. AD is pretrained on data logged by a training agent, and the context size should be sufficiently large to span multiple episodes in order to capture policy improvement. This way, AD learns an improvement operator that increases performance entirely in-context when applied to novel tasks.

The model output is a probability distribution across the action set, derived from a linear projection and a softmax function. As highlighted in Figure 3 , this structure causes a fundamental limitation in AD’s adaptability to new action spaces, as the output dimension is predetermined. To accommodate an action set of a different size than the one used during training, the model’s final layer must be redefined and the model retrained. Moreover, even with a constant action space size, the model’s efficacy diminishes if the action semantics are altered. The reason for this is the classifier nature of the model, which associates each dimension with a particular meaning of an action. Since augmenting the dataset with permuted action sets does not lead to improvement, it signifies that action set invariance should be enforced from a model design standpoint.

## 3 Headless-AD

To mitigate the action space limitations of AD, we propose Headless-AD, a new architecture that improves on AD by omitting the final linear layer and incorporating three key modifications. The Headless-AD architecture and data flow are visualized in Figure 2 .

Random Action Embeddings: To remove the dependence of the model on the pretrained action embeddings, we employed a dynamic mapping function g : 𝔸 → ℝ n g:~\mathbb{A}~\rightarrow~\mathbb{R}^{n} , which produces a unique random encoding for each action in a batch at the start of every training step. The mapping is shared across all batch instances and is consistent along the context sequence, i.e., actions with index i i in their respective action sets will all be mapped to the same embedding. During inference, we generated a single set of embeddings at the beginning and used them throughout the evaluation.

The core intuition behind employing random action embeddings is to eliminate any prior knowledge about the structure of the action space within our model. This approach stems from our observation that using learnable embeddings for actions becomes impractical when encountering new actions not seen during training. A new action would lack a pre-trained embedding, and assigning an arbitrary embedding could introduce an undesirable domain shift, as the model would not recognize it. Moreover, allowing the model to learn new embeddings on-the-fly would necessitate extra gradient steps, diverging from our goal of maintaining a zero-shot learning framework. The usage of random action embeddings ensures that the model does not depend on extracting any information from the embeddings themselves but rather relies on interpreting the context provided by historical interactions with the environment. Moreover, employing random embeddings enhances data variety, which has been demonstrated to boost in-context learning for RL agents ( Kirsch et al., 2023 ; Lu et al., 2023 ) .

We further refined this strategy by constraining the random embeddings to lie on a unit sphere and ensuring their orthogonality (see Appendix H ). The choice of a unit sphere normalizes the scale of the embeddings, while orthogonality allows the model to independently adjust the probability assigned to each action. This condition is crucial for preventing unintended probability mass allocation to multiple actions when the model’s prediction vector aligns with one embedding vector. A similar concept is explored by ( Elhage et al., 2022 ) in the context of feature interference.

Direct Prediction of Action Embeddings: The model output is modified to yield an action embedding a ^ t e ​ m ​ b \hat{a}^{emb}_{t} rather than a probability distribution over actions. This alteration makes the model independent of action set size and order, granting it permutation invariance. The InfoNCE Contrastive Loss ( Oord et al., 2018 ) , diverging from its usual role in representation learning ( Jaiswal et al., 2020 ) , serves as a regression objective to reinforce the similarity between the model prediction and the subsequent action in the data. All other actions are treated as negative samples. Thus, the objective is L = − 𝔼 [ log e f ⁡ ( a ^ t e ​ m ​ b , a t e ​ m ​ b ) / τ ∑ a ∈ A e f ⁡ ( a ^ t e ​ m ​ b , a e ​ m ​ b ) / τ ] , L=-\mathop{\mathbb{E}}\left[{\log{\frac{e^{f(\hat{a}^{emb}_{t},a^{emb}_{t})/\tau}}{\sum_{a\in A}{e^{f(\hat{a}^{emb}_{t},a^{emb})/\tau}}}}}\right], where a e ​ m ​ b = g ⁡ ( a ) a^{emb}=g(a) and τ \tau is a temperature parameter. We used dot-product as the similarity function f f .

Action Set Prompt: To address the model’s lack of awareness of the action space structure caused by the two previous changes, we prepend the input with a sequence of embeddings for all available actions.

The modified input format is thus represented as h t = ( a e ​ m ​ b , 0 , … , a e ​ m ​ b , N , o < t , a < t e ​ m ​ b , r < t , o t ) , h_{t}=\left(a^{emb,0},\dots,a^{emb,N},o_{<t},a^{emb}_{<t},r_{<t},o_{t}\right), where N N is the action set size. An illustrative code snippet with Headless-AD’s training procedure can be found in Appendix L . We suggest two methods to select actions during inference: (1) nearest neighbor selection: a = arg ⁡ max a ∈ A ⁡ f ⁡ ( a ^ e ​ m ​ b , a e ​ m ​ b ) a=\arg\max_{a\in A}f(\hat{a}^{emb},a^{emb}) and (2) probabilistic sampling based on the similarity to the predicted action embedding: P ⁡ ( a i ) = e f ⁡ ( a ^ e ​ m ​ b , a i e ​ m ​ b ) OPEN ∑ a ∈ A e f ⁡ ( a ^ e ​ m ​ b , a e ​ m ​ b CLOSE ) . P(a_{i})=\frac{e^{f(\hat{a}^{emb},a^{emb}_{i})}}{\sum_{a\in A}{e^{f(\hat{a}^{emb},a^{emb}}})}. We treat the specific choice of method as a hyperparameter.

## 4 Experiments

As Headless-AD extends and improves on AD, we checked it in two different aspects. Firstly, it should maintain In-Context Learning abilities and thus generalize well to new tasks. Secondly, it should show high performance on action spaces different from the one seen during training. All of the following environments are designed specifically to check both of the above aspects. In our experiments, we used the TinyLLaMa ( Zhang et al., 2024 ) implementation of the transformer model and AdamW optimizer ( Loshchilov & Hutter, 2017 ) . All environment specific hyperparameters are listed in Appendix J .

### 4.1 Bernoulli Bandit

Motivation: This experiment checked Headless-AD’s abilities on a toy task, where the environment did not have a notion of state and returned binary rewards.

Setup 1: In our first experiment, we examined the model’s robustness to distributional shifts in rewards. We used a Bernoulli bandit where each arm is associated with a mean μ \mu and the reward after pulling an arm i i is generated by Bernoulli ​ ( μ ) \text{Bernoulli}(\mu) . The training dataset consisted of bandits with 4 − 20 4-20 arms, i.e., different action set sizes. Additionally, to evaluate the ICL capabilities of models, the reward distribution over the arms differed between train and test ( Laskin et al., 2022 ) . Specifically, the training data consisted of bandits where in 95 % 95\% of cases the odd-numbered arms were assigned a random μ ∈ [ 0.5 , 1 ] \mu\in[0.5,1] and even-numbered arms received μ ∈ [ 0 , 0.5 ] \mu\in[0,0.5] . For other 5 % 5\% of bandits, the ranges were swapped between odd and even arms. The test distribution also consisted of bandits with 4 − 20 4-20 arms, but the reward distribution was switched either to even arms or was uniform, i.e., all arms were assigned μ ∈ [ 0 , 1 ] \mu\in[0,1] . Learning histories were generated using Thompson Sampling algorithm for a total of 10,000 10,000 bandit instances with 300 300 steps each. The evaluation was performed on 100 100 bandits in each reward distribution, included 5 5 seeds, and the algorithms were rolled out for 300 300 steps.

Results and Discussion 1: As depicted in Figure 4 , the Headless-AD model demonstrates strong generalization capabilities by almost reaching the performance results set by the traditional Thompson Sampling algorithm under each test distribution.

Setup 2: In our second experiment, we evaluated the transferability of Headless-AD to new action set sizes. Training distribution remained the same as in the previous experiment. Each evaluation dataset consisted of a fixed amount of 20 20 , 25 25 , 30 30 , 40 40 and 50 50 arms and a uniform distribution of rewards. AD was trained on fixed-size bandits corresponding to each evaluation dataset. The training and evaluation reward distributions were the same as for Headless-AD.

Results and Discussion 2: Figure 5 shows that Headless-AD can effectively maintain its performance as the action space grows, without necessitating any retraining. Moreover, Headless-AD even outperforms a specially trained AD, especially on larger action sets. Note that Headless-AD’s performance curves resemble TS’s performance curves, signifying that it has indeed learned a policy improvement operator generic enough to apply to unseen action set sizes.

### 4.2 Contextual Bandit

Motivation: To validate the sustained performance of Headless-AD, we progressed to a more complex Multi-Armed Bandit (MAB) extension that integrated states and real-valued rewards.

Setup: Each time step presented a context with arm features modeled as two-dimensional vectors, and rewards were generated with a standard deviation σ = 1 \sigma=1 . The data were created using a LinUCB ( Li et al., 2010 ) algorithm trained over 300 300 steps, on bandits with 4 − 20 4-20 arms. We used 100 100 contextual bandits, 5 5 seeds and 300 300 evaluation steps.

Results and Discussion: As shown in Figure 6 , Headless-AD’s performance is on par with LinUCB across varied arm counts. While AD also reaches the performance of LinUCB on lower space sizes, it has problems converging on larger action space sizes, in addition to requiring a specific retraining. This serves to highlight the advantages of Headless-AD for use in even more complex environments than toy Bernoulli bandits.

### 4.3 Darkroom

Motivation: In this experiment, we delve into a more sophisticated Markov Decision Process (MDP) framework, constructing five distinct action spaces aimed at demonstrating the architectural constraints of AD. We then show how Headless-AD’s architecture is engineered to navigate the complexities presented by each of these diverse action sets.