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

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

Setup: The Darkroom environment, inspired by Chevalier-Boisvert et al. (2018) and Jain et al. (2020) , consists of a N × N N\times N grid where the agent needs to reach a specific cell for a reward. In our experiment, the action space consisted of 3 3 -step sequences of 5 5 atomic actions: up , down , left , right , noop . As a result, the environment offered 5 3 5^{3} possible actions. The agent earned a reward of 1 1 if the trajectory induced by the action sequence passed through a goal cell, after which the episode finished. Otherwise, the reward was 0 0 . As an observation, the environment offered only the current coordinates of the agent, so the goal information could only be obtained from the agent’s memory. We divided the goals into disjoint sets used for training and testing in order to evaluate Headless-AD’s in-context learning abilities. Furthermore, the action set was randomly split into train and test sets, each including 50 50 and 75 75 actions respectively, to create five distinct spaces for assessing various generalization aspects. These action sets are visualized in Figure 1 .

Train Actions. Comprising the training split, this set represented the actions encountered during model training.

Test Actions. Comprising the test split, this set assessed Headless-AD’s generalization on novel and larger action sets.

All Actions. Combining both training and testing actions, this set contained 125 125 actions and was 2.5 2.5 times larger than the training set. Its aim was to challenge Headless-AD to effectively integrate seen and unseen actions.

Permuted Train Actions. Shuffled training set, meant to test the model’s adaptability to reordered action spaces that comprised of the same actions.

Sliced Test Actions. Tailored to match the training set size, this set contained a slice of the first 50 50 actions from the test set. Its aim was to check the models’ generalization abilities on unseen actions while maintaining the action set size.

In scenarios where the action space exceeded the size of the training set, we analyzed the performance of an AD model retrained from scratch. The data generation algorithm was Q-learning, executed over 200 200 episodes for each environment. Further details on model hyperparameters are available in Appendix J .

Results and Discussion: Figure 7 illustrates the generalization abilities of both AD and Headless-AD. First, note that both models maintain their performance when transitioning to novel tasks, thus fulfilling their purpose as ICL-RL models. However, this environment was mainly designed to challenge the models’ generalization abilities on novel action spaces. Following the scope of the action space novelty we set earlier, we studied the models’ ability to address changed action semantics and variable action set sizes.

While AD achieves high performance on the train action set, its limitations become evident when the action space is changed. The first limitation is AD’s action set size constraint. The linear layer at the end of its network fixes the size of the output dimension, something that cannot be modified once the model is trained. Increasing the amount of options, as was done in the Test Actions and All Actions sets, requires reinitializing the output dimension and thus retraining the model, which demands additional time and resources. In contrast, once trained, Headless-AD easily adapts to changes in the action set size without losing performance.

The second limitation is AD’s reliance on a stationary action space structure. Due to the classifier nature of AD’s network, it learns to associate each output dimension with a specific action meaning. When action semantics change either due to a permutation of seen actions, as in Permuted Train Actions , or due to a substitution of completely new actions, as in Sliced Test Actions , AD’s performance degrades. We checked whether this problem may be solved by biasing the training distribution to have a structure resembling the one during testing. We trained AD on permuted action sets while leaving the data and the model architecture the same. However, as one can see in Figure 3 , this training procedure resulted only in a slightly decreased performance of AD in the Permuted Train setting. An additional graph depicting the performance on each of the (action set, goal type) pairs can be found in Appendix C . Meanwhile, Permuted Train Actions does not pose a challenge for Headless-AD as, by design, it is invariant to specific action order. Most importantly, Headless-AD maintains its performance even when completely new actions are introduced, despite having never seen them in training. We attribute this feature to our success in making the model infer action meaning from the context.

## 5 Ablations

The hyperparameters tuned for each ablation can be found in Appendix J .

### 5.1 Action Set Prompt

We assessed the impact of omitting the action embedding enumeration from the model context.

Table 1 reveals that removing the action set prompt did not significantly alter the number of unique actions attempted by each model. This suggests that models without the prompt continue to sample a diverse range of actions, comparable to the original setup. However, the presence of the prompt did affect the performance. We hypothesize that, without the action set prompt, the model lacked explicit knowledge of the action space, which may have led to a more randomized sampling within the embedding space. Conversely, with the prompt in place, the model could directly target specific action embeddings.

The absence of the prompt was more detrimental in the Bernoulli Bandit environment, where each decision has a direct impact on the final outcome due to the environment’s single-episode structure. However, in a multi-episode environment such as the Darkroom, the early lack of action space information becomes less impactful over time, as the model eventually encounters all actions through the context. This divergence highlights the prompt’s critical role in enabling the model to make informed choices in environments where each option has immediate and lasting consequences. Note that the impact of action set prompt ablation on the performance got more pronounced as the number of arms grows.

### 5.2 Contrastive Loss

In principle, the prediction of action embedding can be facilitated by various ways. Here, we considered the performance of a model that was asked to directly copy the corresponding embedding from its context to the output. To achieve it, we used mean squared error (MSE) loss instead of contrastive loss.

In this case, probabilistic interpretation became less meaningful, which is why the nearest neighbor of the predicted vector was chosen as an action index, without sampling.

Table 1 illustrates a significant decline in the variety of actions attempted by the model under this configuration within the bandit environment, showing that the model concentrated only on a fraction of the action set. Though the final regrets for this model were far from random, we point out that this result is because the model skipped exploration phase and went directly to an exploitation of suboptimal actions. That allowed it not to waste time on even more suboptimal ones via exploration. Appendix G shows the in-context curves of both models, providing more insight into MSE’s effect on model quality.

Conversely, on the Darkroom environment, the loss substitution led to a marginal increase in the amount of attempted actions. However, this did not lead to an increased performance. In fact, the performance of this model on the Darkroom environment was similar to a random behavior.

In summary, although the employed neural network architecture was capable of learning an improvement operator and demonstrating ICL capabilities, the implemented loss function played a crucial role in the success of this learning process. We emphasize the importance of our design choice to use contrastive loss by showing that a more naive approach of directly copying the action embeddings, as incentivized by MSE loss, resulted in underperformance.

### 5.3 Orthonormal Action Embeddings

In this ablation study, we underscored the significance of employing orthonormal vectors for action embeddings by contrasting them with vectors derived from a standard normal distribution. Our preference for the former stemmed from their ability to simplify the approximation of action probabilities because of their property of independence. Unlike embeddings from a standard normal distribution, orthonormal vectors ensure that assigning a probability weight to one vector does not unintentionally influence another (we illustrate it in Appendix K ). This concept echoes the principle of superposition , observed when models incorporate more features than available dimensions, leading to feature interference ( Elhage et al., 2022 ) .

Table 1 illustrates the consequences of using linearly dependent action embeddings, manifesting as diminished performance across Bernoulli bandits and Darkroom scenarios. Notably, these results bear a striking resemblance to ”Prompt Ablation”, with both indicating a slight performance dip in Darkroom, a more noticeable decline in Bernoulli Bandits, and a general downtrend as the number of bandit arms grows. This parallel underscores a shared objective between these design choices: refining the model’s capacity for precise action selection.

## 6 Related Work

Here, we discuss previous research in adapting RL to environments with variable action spaces. For an extended Literature Review, see Appendix B .

Recent research by Chandak et al. (2020) ; Ye et al. (2023) has focused on scenarios where the amount of available actions grows during evaluation by introducing new actions to an existing set. However, their model requires fine-tuning when new actions are introduced, while our model does not require parameter updates. Additionally, our research explores a broader spectrum of dynamic action spaces, encompassing the addition, removal, and substitution of actions. Lastly, we work in a setting where the action set remains constant throughout model evaluation.

Kirsch et al. (2022) present the concept of SymLA, a methodology that ensures resilience to changes in input and output sizes and permutations. This is achieved by integrating symmetries into a neural network model by representing each neuron with uniformly structured RNNs and data flow with message-passing connections. However, we suggest a simpler approach by extending the well-known transformer architecture. Additionally, we demonstrate the high performance of our model on more complex environments.

Further developments by Jain et al. (2020) include a specific module designed to generate action representations informed by observed trajectories after action application. Additionally, Jain et al. (2021) delve into variable action spaces, placing a specific emphasis on modeling the interconnections among actions to improve the quality of action representations. Our work, however, adopts a more implicit approach to inferring action-related information.

In a similar work, Lu et al. (2023) ; Kirsch et al. (2023) employ random projections for action encoding, a technique we also utilize. This method facilitates training across multiple domains with actions of varying sizes. Nonetheless, their focus is predominantly on continuous spaces, while our focus is on discrete action spaces.

To the best of our knowledge, Headless-AD is the first study to explore variable discrete action spaces for in-context reinforcement learning.

## 7 Conclusion

In our work, we have introduced a new architecture that extends Algorithm Distillation (AD) for environments with variable action spaces, achieving invariance to their structure and size. Our approach consists of discarding the last linear layer, granting invariance to the action space structure, and making the model infer action semantics from the context, preparing it for the introduction of novel actions. We demonstrated Headless-AD’s capability to generalize across new action spaces on a set of environments. We also observed its performance gains over vanilla AD, especially on larger action spaces. We hypothesize that this is due to the augmentation of the dataset by random embeddings, which was shown to improve the generalization abilities of agents ( Kirsch et al., 2023 ) . Headless-AD marks the progress toward versatile foundational models in RL, ones capable of operating across an expanded range of environments. We hope that Headless-AD inspires further development of models that can adapt to any action space beyond discrete ones.

Limitations. Headless-AD shares AD’s limitation of fixed sequence lengths, which may limit its effectiveness in environments with long episodes. A unique constraint of Headless-AD is its limit on the number of actions, dictated by the dimensionality of the action embeddings – only as many actions as there are dimensions can be orthogonally represented. Going over this limit causes embeddings to become linearly dependent, unintentionally distributing probability across multiple actions. While action spaces in RL environments typically do not become excessively large, and increasing the dimension of embeddings could mitigate this issue, it remains a point for consideration.

Future Work. In our study, we demonstrated the ability of our algorithm to generalize to new action spaces, as shown by its performance on elementary tasks. To extend and validate these findings, future research should focus on more complex environments. This will offer a deeper insight into the algorithm’s versatility and robustness in diverse and complex settings.

Additionally, to ensure the broader applicability and adaptability of our approach, it is essential to examine its compatibility and performance with various models beyond Algorithm Distillation (AD) ( Laskin et al., 2022 ) , such as the Decision Pretrained Transformer (DPT) ( Lee et al., 2023 ) . This exploration will provide a more comprehensive understanding of the algorithm’s strengths and limitations, potentially leading to further improvements and a wider scope of applications in different contexts.

## Impact Statement

Strong safeguards are necessary to prevent unauthorized users from manipulating the model, such as adding harmful action embeddings that could lead to negative outcomes.

Another concern is how the model handles out-of-distribution data. If the model encounters a new action that is significantly different from the training actions, it may take a while to understand its effects. Since our model learns by trying out actions, there is a risk it might perform harmful actions before learning they are inappropriate.

Any application of Headless-AD in real-life scenarios should be aware of these potential risks.

## Appendix A Background

Partially Observable Markov Decision Process. A Markov Decision Process (MDP) is defined by a tuple ( S , A , P , R ) (S,A,P,R) , where s t ∈ S s_{t}\in S denotes a state, a t ∈ A a_{t}\in A an action, p ⁡ ( s ′ | s t = s , a t = a ) p(s^{\prime}|s_{t}=s,a_{t}=a) the transition probability from state s t s_{t} to s t + 1 s_{t+1} after taking action a t a_{t} , and R ⁡ ( s , a ) R(s,a) the reward for action a a in state s s ( Sutton & Barto, 2018 ) . An agent π \pi observes the state s t s_{t} , selects action a t ∼ π ( ⋅ | s t ) a_{t}\sim\pi(\cdot|s_{t}) , and receives the subsequent state s t + 1 ∼ P ( ⋅ | s t , a t ) s_{t+1}\sim P(\cdot|s_{t},a_{t}) and reward R ⁡ ( s t , a t ) R(s_{t},a_{t}) . In POMDPs, the agent receives an observation o t o_{t} instead of the full state s s , which contains partial information about the MDP’s real state. In the context of our work, o t o_{t} may lack goal information, requiring inference from the agent’s memory.