In-Context Reinforcement Learning for Variable Action Spaces — part 3/4
Cleaned-text code points [31777, 47461). Bibliography omitted here; source.md includes it.
Multi-Armed Bandits. A Bernoulli multi-armed bandit (MAB) environment consists of N N arms a i ∈ A a_{i}\in A , each associated with a mean μ i \mu_{i} ( Sutton & Barto, 2018 ) . Pulling an arm a i a_{i} yields a reward r i ∼ Bernoulli ( μ i ) r_{i}\sim\text{Bernoulli}(\mu_{i}) . The agent’s objective is to identify the arm with the highest μ i \mu_{i} . Performance is measured using regret, calculated as ∑ t ( μ a ∗ − μ a ^ t ) \sum_{t}(\mu_{a^{*}}-\mu_{\hat{a}_{t}}) . Unlike MDPs, MABs lack states.
In a Contextual MAB, each arm a a has a feature vector x a ∈ R d x_{a}\in R^{d} ( Sutton & Barto, 2018 ) . At each step t t , the agent observes a context state s t s_{t} . Selecting action a i a_{i} results in a reward from a normal distribution with mean μ t = ⟨ s t , a i ⟩ \mu_{t}=\langle s_{t},a_{i}\rangle and standard deviation σ \sigma . Here, unlike in MDPs, the actions influence immediate rewards but not future states.
In-Context Learning. In-Context Learning describes the capability of a model to infer its task from the context it is given. For instance, the GPT-3 model ( Brown et al., 2020 ) can be prompted in natural language to perform a variety of functions such as text classification, summarization, and translation, despite not being explicitly trained for these specific tasks. One of the possible prompts, which is used in our paper in some form, is a list of example pairs ( x i , y i ) (x_{i},y_{i}) ending with a query x q x_{q} for which the model is expected to generate a corresponding prediction y ^ q \hat{y}_{q} .
Contrastive Learning. Contrastive learning focuses on creating representations where “similar” examples are close together in the feature space, while “dissimilar” ones are far apart ( Weinberger & Saul, 2009 ; Schroff et al., 2015 ) . This concept may be encapsulated in a triplet loss formula, where the similarity between an anchor and a positive example is maximized, and that between an anchor and a negative example is minimized, expressed as L = s i m ( x , x + ) − s i m ( x , x − ) L=sim(x,x^{+})-sim(x,x^{-}) . Oord et al. (2018) has developed a variant of contrastive loss called InfoNCE. This loss was lately widely adopted for representation learning ( Jaiswal et al., 2020 ) .
## Appendix B Related Work
### B.1 Transformers in Reinforcement Learning
According to the survey by Li et al. (2023) , Transformers ( Vaswani et al., 2017 ) are increasingly utilized in reinforcement learning (RL) for various tasks, including representation learning of individual observations and their histories, as well as model learning, as seen in Dreamer by Hafner et al. (2019) . In our research, we focus on the application of Transformers in sequential decision-making and in developing generalist agents. Same as AD ( Laskin et al., 2022 ) , Headless-AD also utilizes transformers as the base model for our approach.
The incorporation of transformers as a sequence modeling tool in RL began with the Decision Transformer (DT) by ( Hu et al., 2022 ) , which is trained autoregressively on offline datasets of state, action, and return-to-go tuples. Unlike conventional RL, which focuses on return maximization, DT generates appropriate actions during inference by conditioning on specified return-to-go values. The Trajectory Transformer by Janner et al. (2021) is an alternative approach using beam search to bias trajectory samples based on future cumulative rewards. Building on this, the Multi-Game Decision Transformer (MGDT) by Lee et al. (2022) improves upon DT with enhanced transfer learning capabilities for new games, eliminating the need for manual return specification. Similarly, the Switch Trajectory Transformer by ( Lin et al., 2022 ) expands on the Trajectory Transformer to facilitate multi-task training. However, when transitioning to novel tasks, these approaches require fine-tuning the model.
### B.2 Offline Meta-RL
Traditional RL agents, tailored to specific environments, struggle with novel tasks. In contrast, by training across a diverse array of tasks, Meta-RL equips agents with adaptable exploration strategies and an understanding of common environmental patterns ( Duan et al., 2016 ; Wang et al., 2016 ; Rajeswaran et al., 2017 ; Zhang et al., 2018 ; Team et al., 2021 ; Nikulin et al., 2023 ) . Offline Meta-RL, a subset of this approach, trains agents solely on pre-existing datasets, without direct environmental interaction. A critical challenge here is MDP ambiguity, where task-specific policies misinterpret data due to dataset biases. Dorfman et al. (2021) propose a data collection method to mitigate this issue, and suggest an approach that treats Meta-RL as a Bayesian RL problem for optimal exploration in new tasks. Li et al. (2020) introduce FOCAL, a framework separating task identification from control. However, their method relies on a strict mapping assumption that fails in certain scenarios, such as those with sparse rewards. Li et al. (2021) enhance FOCAL with an attention mechanism and improved metric learning, showing greater robustness in scenarios with sparse rewards and domain shifts. However, most Offline Meta-RL methods rely on explicit task modeling, which can introduce limiting biases. Alternatively, In-Context Learning implicitly infers tasks from environmental interactions, offering a potentially more flexible approach.
### B.3 In-Context Learning in RL
In-Context Learning (ICL) is an ability of a pretrained model to adapt and perform a new task given a context with examples { x i , f ( x i ) } n \{x_{i},f(x_{i})\}^{n} , where f f is a function that gives ground truth targets ( Brown et al., 2020 ; Wang et al., 2023 ) . Previous work on ICL ( Mirchandani et al., 2023 ) showed that Large Language Models operate as General Pattern Machines and are able to complete, transform and improve token sequences, even when the sequences consist of randomly sampled tokens. This ability supports our design choice to randomly encode the actions of agents. The research on Transformer Circuits ( Elhage et al., 2021 ) gives evidence of Transformers’ ability to copy tokens, either literally or on a more abstract level, explaining their ICL capabilities. This is particularly useful for our model, which explicitly incentivizes an abstract copying of the action embeddings.
Efforts to blend In-Context Learning (ICL) abilities of transformers with reinforcement learning (RL) are gaining traction, promising to create adaptable RL agents for real-world scenarios with varying conditions. Attempts to transfer ICL features to RL include Prompt-Based Decision Transformers ( Xu et al., 2022 ) , which leverage task-specific demonstration datasets as prompts, exhibiting strong few-shot learning without weight updates. Laskin et al. (2022) introduced Algorithm Distillation (AD) that trains a policy improvement operator using data from agent-environment interactions of learning RL algorithm, predicting the next action autoregressively. A key strategy here is across-episode training for capturing policy improvement. Lee et al. (2023) developed Decision Pretrained Transformer (DPT) using supervised training with unordered environmental interactions as context to predict optimal actions. Under certain conditions, DPT can be proved to implement posterior sampling, resulting in near-optimal exploration. However, DPT’s reliance on an optimal policy during training, not always available, is a limitation, similar to AD’s constraint on action space structure, restricting their use as fully generalist agents. Lin et al. (2023) provided theoretical analysis of AD and DPT, offering guarantees about learning from base algorithms. Wang et al. (2024) also provided a theoretical analysis showing that transformers may implement several RL algorithm in-context. Zisman et al. (2023) propose a method for distilling the improvement operator from a demonstrator whose actions are initially noisy, with the noise level decreasing progressively throughout the data collection phase. This approach simplifies the data generation compared to AD as it does not require logging the training process of RL models and relaxes DPT’s limitation because the demonstrator may be suboptimal. Kirsch et al. (2023) demonstrated that data augmentation, through random projection of observations and actions, enhances task distribution and generalization on new domains, boosting ICL capabilities. Raparthy et al. (2023) study the effect of model size and dataset properties on the success of in-context learning. Other research, acknowledging Transformers’ limitations when it comes to long sequences, explores alternative models such as S4 ( Gu et al., 2021 ) . Lu et al. (2023) adapted S4 for RL, showing that it is capable of surpassing LSTM in performance and Transformers in runtime, indicating potential in RL ICL applications.
In contrast to previous works that tackle changing reward distributions or environment dynamics as changing goals, we expand the application of ICLRL to changing action spaces. Our use of random embeddings, motivated by the preparation of the model to unseen actions, may also improve ICL abilities by data augmentation, as suggested by Kirsch et al. (2023) .
### B.4 Discarding the Linear Layer
The Headless LLM approach, introduced by Godey et al. (2023) , utilizes Contrastive Learning to eliminate the language head from the model architecture. Rather than generating a probability distribution over tokens, it directly predicts token embeddings. This strategy aims to enhance training and inference speeds by discarding a substantial linear layer. While in the RL context, where action spaces are typically small, this might not significantly impact runtime, it does offer a key advantage, as the model is no longer dependent on the number of actions and can therefore handle variable-length action spaces. Unlike Headless LLM, our approach does not rely on contrastive loss for learning action representations. Instead, we use fixed action embeddings, and the model is tasked to output the embedding of the next action. During inference, an action is chosen from a categorical distribution, where the logits are the similarities between the predicted embedding and available actions. Thus, the role of contrastive loss is to enhance the likelihood of selecting the correct action while diminishing the probabilities of other actions.
Wolpertinger by Dulac-Arnold et al. (2015) also employs a similar approach of removing the linear head to improve the train and inference speeds. The authors suggest associating each action with an embedding and performing the training in a continuous action space. A specific action index is chosen as a nearest neighbor of the predicted embedding, and its effect is treated as environment dynamics. Most importantly, this nearest neighbor selection process can be optimized using an approximate algorithm. This leads to logarithmic time complexity, which is particularly advantageous in environments with large action spaces. Thus, we believe that Headless-AD is also capable of performing well on large action spaces, benefiting from runtime gains of approximate nearest neighbor lookup. However, Wolpertinger’s usage of fixed action embeddings prevents introduction of new actions, highlighting the significance of Headless-AD’s usage of randomized action embeddings.
## Appendix C Algorithm Distillation on Permuted Train Sets
## Appendix D Across-Environment Generalization
We evaluated Headless-AD’s ability to generalize skills learned in one domain to another. We trained Headless-AD on Contextual Bandits and subsequently assessed its performance on Bernoulli Bandits. To align with the Contextual Bandit setting, we introduced a random vector to serve as the state in the Bernoulli Bandit environment, effectively transforming it into a single-context Contextual Bandit scenario with Bernoulli-distributed rewards.
Figure 9 presents the results of these cross-domain experiments. Headless-AD not only adapted to the new domain but also maintained a satisfactory performance level, highlighting the model’s capacity for cross-environment application.
## Appendix E Visual Darkroom
We adapted the Darkroom environment to produce high-dimensional visual observations to demonstrate Headless-AD’s capability to generalize to new action spaces in more complex settings. In this modified version, the grid is visually represented, with the agent’s position indicated in red (see Figure 10 ). This change introduces a more complex observation space. For the model to handle these high-dimensional visual observations, we integrated a convolutional network that converts the visual input into an embedding of size ’token_dim’.
Both Headless-AD and AD were configured with the identical hyperparameters previously applied in the Darkroom experiments. The results at Figure 11 show that Headless-AD maintains its superior performance in this more complex observational setting, whereas AD experiences a notable drop in performance. This discrepancy likely stems from our decision to retain the original hyperparameters without tuning them for the new environment. Based on our experience, AD’s performance is particularly sensitive to hyperparameter settings.
## Appendix F Darkroom. Alternative Split
A random split of actions in Darkroom environment may lead to a data leakage as multiple action sequences result in the same endpoint. To address this, we conducted an experiment with distinct, non-overlapping action sets for training and testing. We trained Headless-AD with actions that cause movements of 0 0 , 1 1 , and 3 3 cells, and tested on movements of 2 2 cells, introducing scenarios not encountered during training.
Figure 12 shows that although Headless-AD experienced a slight decrease in performance on the test action set, it still remained significantly above the random level and outreaching the AD-level, underscoring its capability to generalize beyond the training action space. This experiment intends to make the capabilities and limitations of Headless-AD clearer.
loh loh
## Appendix G MSE-Headless-AD In-Context Curves on Bernoulli Bandit
## Appendix H Sampling of Orthonormal Vectors
To sample the orthonormal vectors used as action embeddings, we use the torch.nn.init.orthogonal function from PyTorch ( Paszke et al., 2019 ) that utilizes a specific algorithm from Saxe et al. (2013) .
## Appendix I Algorithms’ Training Times
## Appendix J Model Hyperparameters
In this section, we detail the hyperparameters for our models, each tuned for specific environments and design configurations. The tuning process utilized Bayesian optimization via the wandb sweep tool ( Biewald, 2020 ) . The optimization objective was chosen to maximize both the performance score achieved by the model and the efficiency in the number of actions utilized. The objective function was structured as follows:
n a / N + final_normalized_return , n_{a}/N+\text{final\_normalized\_return},
where n a n_{a} represents the total number of actions attempted by the model during evaluation, and N N signifies the amount of possible actions within an environment. Moreover, we took the return from the final episode in evaluation and normalized it to the [ 0 , 1 ] [0,1] range, with the lower bound 0 0 corresponding to the performance of a random agent, and the upper bound 1 1 denoting the efficiency of our data generation algorithm. Darkroom environment already has the returns in the range [ 0 , 1 ] [0,1] , so we did not perform normalization for this environment.
## Appendix K Linear Dependence of Different Types of Action Embeddings