Research map / Reinforcement learning / Visual and representation RL

In-Context Reinforcement Learning for Variable Action Spaces

ICML 2024 · Reinforcement learning › Visual and representation RL · approach: Inferring decision-making from behavior · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes Headless-AD, a transformer in-context RL agent that drops the output linear head and predicts action embeddings with random orthonormal action encodings and a contrastive loss, enabling generalization to unseen discrete action spaces of different size and order without retraining.

Problem

Algorithm Distillation and similar in-context RL transformers have a fixed output head tied to a predefined action set size and semantics, so new action spaces require data re-collection and retraining.

Approach

Remove the linear output head, encode each action with a random orthonormal embedding resampled every training step, feed the full action set as a prompt, predict the next action embedding, and train with InfoNCE contrastive loss so action choice is a similarity-based distribution over the action set.

Key results

• On Bernoulli bandits, Headless-AD nearly matches Thompson Sampling across reward distribution shifts and generalizes to 20-50 arms after training on up to 20 arms, outperforming AD retrained per arm count on larger sets.
• In Darkroom, Headless-AD maintains stable success rates across Train, Test, and All action sets (up to 125 actions, about 5x the training set) while AD degrades under permuted or substituted actions.
• Ablations show removing the action set prompt, replacing contrastive loss with MSE, or using non-orthogonal embeddings each reduces performance, most strongly on bandits with more arms.

Measured effects reported in the paper

Surprising findings

• Training AD on permuted action sets did not improve its performance on permuted action sets over standard AD training.

Limitations

• Number of actions is bounded by embedding dimension since only that many orthogonal embeddings exist; beyond that embeddings become linearly dependent.
• Inherits AD's fixed sequence length limitation, which may hinder long-episode environments.
• Evaluated only on toy environments (bandits and a small gridworld).

Experimental setup

Datasets: Bernoulli bandit trajectories from Thompson Sampling contextual bandit trajectories from LinUCB Darkroom trajectories from Q-learning

Benchmarks: Bernoulli bandit (4-50 arms) contextual bandit Darkroom gridworld with five action set variants visual Darkroom Darkroom alternative split

Metrics: regret success rate normalized return number of unique actions attempted

Baselines: Algorithm Distillation Thompson Sampling LinUCB Q-learning Random

Builds on: Algorithm Distillation Headless-LLM Wolpertinger InfoNCE Decision Pretrained Transformer

Closest papers by meaning

• XLand-100B: A Large-Scale Multi-Task Dataset for In-Context Reinforcement Learning (ICLR 2025 · similarity 0.88)
• Zero-Shot Context Generalization in Reinforcement Learning from Few Training Contexts (NeurIPS 2025 · similarity 0.88)
• Excluding the Irrelevant: Focusing Reinforcement Learning through Continuous Action Masking (NeurIPS 2024 · similarity 0.87)
• In-Context Reinforcement Learning From Suboptimal Historical Data (ICML 2025 · similarity 0.86)
• Zero-shot Model-based Reinforcement Learning using Large Language Models (ICLR 2025 · similarity 0.86)
• Vintix: Action Model via In-Context Reinforcement Learning (ICML 2025 · similarity 0.86)