Scaling Algorithm Distillation for Continuous Control with Mamba
Extracted source text · arXiv 2506.13892
\documentclass{article}
% if you need to pass options to natbib, use, e.g.:
% \PassOptionsToPackage{numbers, compress}{natbib}
% before loading neurips_2024
% ready for submission
%\usepackage{neurips_2024}
% to compile a preprint version, e.g., for submission to arXiv, add add the
% [preprint] option:
\usepackage[preprint]{neurips_2024}
% to compile a camera-ready version, add the [final] option, e.g.:
% \usepackage[final]{neurips_2024}
% to avoid loading the natbib package, add option nonatbib:
%\usepackage[nonatbib]{neurips_2024}
\usepackage[utf8]{inputenc} % allow utf-8 input
\usepackage[T1]{fontenc} % use 8-bit T1 fonts
\usepackage{hyperref} % hyperlinks
\usepackage{url} % simple URL typesetting
\usepackage{booktabs} % professional-quality tables
\usepackage{amsfonts} % blackboard math symbols
\usepackage{nicefrac} % compact symbols for 1/2, etc.
\usepackage{microtype} % microtypography
\usepackage{xcolor} % colors
\usepackage{graphicx}
\usepackage{subcaption}
\usepackage{amsmath}
\title{Scaling Algorithm Distillation for Continuous Control with Mamba}
% The \author macro works with any number of authors. There are two commands
% used to separate the names and addresses of multiple authors: \And and \AND.
%
% Using \And between authors leaves it to LaTeX to determine where to break the
% lines. Using \AND forces a line break at that point. So, if LaTeX puts 3 of 4
% authors names on the first line, and the last on the second line, try using
% \AND instead of \And before the third author name.
\author{
Samuel Beaussant\\
Akkodis Research\\
France\\
\texttt{samuel.beaussant@akkodis.com} \\
\And
Mehdi Mounsif\\
Akkodis Research\\
France\\
\texttt{mehdi.mounsif@akkodis.com} \\
}
\begin{document}
\maketitle
\begin{abstract}
Algorithm Distillation (AD) was recently proposed as a new approach to perform In-Context Reinforcement Learning (ICRL) by modeling across-episodic training histories autoregressively with a causal transformer model. However, due to practical limitations induced by the attention mechanism, experiments were bottlenecked by the transformer's quadratic complexity and limited to simple discrete environments with short time horizons. In this work, we propose leveraging the recently proposed Selective Structured State Space Sequence (S6) models, which achieved state-of-the-art (SOTA) performance on long-range sequence modeling while scaling linearly in sequence length. Through four complex and continuous Meta Reinforcement Learning environments, we demonstrate the overall superiority of Mamba, a model built with S6 layers, over a transformer model for AD. Additionally, we show that scaling AD to very long contexts can improve ICRL performance and make it competitive even with a SOTA online meta RL baseline.
\end{abstract}
% rajouter que DT ne scale pas bien (exemple ant-dir ou il est pire que MACAW)
\section{Introduction}
Since its inception, the Transformer architecture \cite{Attention} has become the de facto model for many sequence modeling tasks across a vast range of domains. This widespread adoption can be partly explained by its impressive performance and versatility. However, its most striking feature is arguably the ability to perform "in-context" learning (ICL) after self-supervised pre-training on enough data. This departure from traditional in-weights learning has ignited significant interest as a novel form of meta-learning and made it a very good backbone for many foundation models \cite{DinoV2, SAM, vuong2023open}. In particular, Algorithm Distillation (AD) \cite{AD} leverages the in-context abilities of transformers to adapt and self-improve in new tasks from a few environment interactions. It has demonstrated very promising results from purely offline data in a meta reinforcement learning setting, but experiments have been limited to small and simple environments such as gridworlds. Hence, it is not clear if these appealing results could be further extended to more challenging and realistic environments with continuous state-action spaces.
One major difficulty arises from the fact that the context accessible to the meta-learner needs to be cross-episodic to enable long-range credit assignment and policy improvement. However, transformer models have been shown to perform poorly on benchmarks requiring modeling long-range dependencies \cite{LRA, LostInMiddle}. Moreover, by design, the attention mechanism powering its effectiveness scales quadratically with respect to the sequence length. As such, inference could be prohibitively slow for Reinforcement Learning (RL) tasks or even real-time Computer Vision tasks \cite{VisionMamba}.
Recently, State Space Models (SSMs) \citep{S4, Mamba} have outperformed transformer-based models in long-range sequence modeling tasks \citep{LRA}. In particular, the Mamba architecture based on the S6 layer \citep{Mamba} achieved impressive results in natural language processing (NLP) as well as DNA sequence and audio waveform modeling. Moreover, a recent work \citep{MambaICL} demonstrated that for simple supervised tasks, a Mamba model can exhibit similar ICL abilities as transformers. State Space Models also enjoy superior inference speeds and an efficient hardware implementation, making them potentially useful in many real-time control settings.
In this work, we study the synergy of Algorithm Distillation with Mamba for long-horizon tasks. We aim to investigate the impact of architecture inductive bias in this setting by comparing Mamba with a memory-efficient causal transformer \citep{FlashAttention}. Through four continuous control tasks, we show that, given similar model sizes, Mamba achieves better asymptotic performance in all settings considered. Furthermore, it seems that more complex tasks require longer contexts, which highlights the importance of efficiently modeling long sequences for optimal results when scaling AD. Finally, we compare the asymptotic performance of AD with two other SOTA meta-RL baselines and found that it is competitive with online meta-RL while training purely offline. It also outperforms prior SOTA offline meta-RL when compared in similar settings.
\section{Background}
\subsection{Meta Reinforcement Learning}
We consider the problem of \textbf{Meta Reinforcement Learning} and assume a distribution of tasks $p_{\mathcal{T}}$. Tasks are generally assumed to share a some fundamental similarities. Each task can be represented by a Markov Decision Process (MDP), defined as 5-tuple $\mathcal{T} = <S,A,P,R,\rho_0>$ where $S$ is the set of valid (continuous) states, $A$ is the set of valid (continuous) actions, $P(s_{t+1}|s_t, a_t)$ is the transition probability distribution describing the MDP dynamic, $R: S \times A \times S \mapsto \mathbb{R}$ is the reward function and $\rho_0$ is the initial state distribution. Usually in this framework, an agent (or decision-maker) is presented with a task ${\mathcal{T}}$ drawn from $p_{\mathcal{T}}$ during meta-training and must find a policy that maximizes its expected reward $G_t = \mathbb{E} \left[ \sum_{t}^T \gamma^t r_t\right]$. If meta-training succeeds, the meta-learner has learned the underlying structure and can find quickly a high-performing policy for a (potentially unseen) task $\mathcal{T}' \sim p_{\mathcal{T}}$. Ideally, adaptation to a new task (i.e meta-testing) should be possible from a small amount of data (referred to as few-shot learning). \textbf{In-Context Reinforcement Learning} is a family of methods dealing with Meta-RL. In this framework, adaptation to a new task happens entirely in the context of a sequence model, without any update to the model parameters. The term "in-context learning" was first coined in prior work on sequence models \citep{GPT3} to describe the meta-learning abilities acquired by large language models (LLMs) pre-trained on vast amount of data. It is opposed to "in-weights learning", i.e gradient-based learning with parameters updates. Depending on the sequence model used for adaptation, the task and the context length required, ICL performances may vary.
\subsection{Mamba}
In this work, we leverage the Mamba architecture \citep{Mamba} based on the structured state-space models class \citep{S4}. This family of sequence models is related to RNNs, CNNs and classical state-space models. They map 1-D sequence $x_k \mapsto y_k \in \mathrm{R}$ using a latent state variable $h_k \in \mathrm{R}^N$ through the following relations:
\begin{equation}
h_k = A h_{k-1} + B x_k
\end{equation}
\begin{equation}
y_k = C h_k
\end{equation}
where $A \in \mathrm{R}^{N\times N}$, $B \in \mathrm{R}^{N\times1}$ and $C \in \mathrm{R}^{1\times N}$ are learnable parameters. When dealing with multi-dimensional data, each input channels is modeled by its own 1-D SSM. Importantly, the state matrix $A$ is structured i.e initialized to a diagonal matrix following the HiPPO theory \citep{Hippo} such that $A_{n,n} = -(n+1)$. Mamba is a fully recurrent model which enjoys fast inference, linear scaling in sequence length and an efficient hardware-aware implementation on GPU in order to be computationally efficient during training. Hence, it appears better suited than transformers in many reinforcement learning settings.
\section{Algorithm distillation}
Broadly speaking, Algorithm Distillation tackles the meta-learning problem by learning to imitate a Reinforcement Learning algorithm. Over the course of training, an agent will learn to explore, plan, and correlate rewards to actions (also known as credit assignment). All of these complex behaviors are reflected in the learning data distribution of an agent. Modeling these histories using a sequence model thus entails learning to imitate these behaviors, which can later be repurposed during inference to in-context learn new tasks. This is similar to the few-shot learning abilities demonstrated by Large Language Models (LLMs) \citep{GPT3}. But instead of manually designed prompts or demonstrations, the model collects its own experience by interacting with the environment to bootstrap its own Reinforcement Learning in-context. More formally, we call learning histories (or learning trajectories) the sequence of experiences collected by running a Reinforcement Learning algorithm $\mathcal{P}^{\textit{source}}$ on multiple individual tasks ${\mathcal{M}}^{N}_{n=1}$ producing dataset $D$:
\begin{equation}
D := \{(o_0^{(n)}, a_0^{(n)}, r_0^{(n)}, ..., o_{T-1}^{(n)}, a_{T-1}^{(n)}, r_{T-1}^{(n)}, o_T^{(n)}, a_T^{(n)}, r_T^{(n)}) \sim \mathcal{P}^{\textit{source}}_{{\mathcal{M}}^{n}}\}^N_{n=1}
\end{equation}
Given a dataset $D$, a source algorithm can be distilled into a model by autoregressively predicting actions conditioned on a cross-episodic context. The original AD paper only considered discrete actions and thus trained their sequence model to map learning histories to probabilities over actions with a negative log likelihood (NLL) loss:
\begin{equation}
L(\theta) := - \sum_{n=1}^N \sum_{t=1}^T \log P_\theta \left(A=a_t^{(n)} | h_{t-1}^{(n)}, o_t^{(n)}\right)
\end{equation}
with
\begin{equation}
h_t := (o_0, a_0, r_0, ..., o_{t-1}, a_{t-1}, r_{t-1}, o_t, a_t, r_t) = (o_{\leq t}, a_{\leq t}, r_{\leq t})
\end{equation}
a cross-episodic learning history. In this work, we consider only continuous-control problems. As such, we perform AD by optimizing the following loss:
\begin{equation}
L(\theta) := \sum_{n=1}^N \sum_{t=1}^T \left(a_t^{(n)} - f_\theta(h_{t-1}^{(n)}, o_t^{(n)})\right)^2
\end{equation}
The simple imitation-learning objective optimized during meta-training departs from traditional meta-RL methods that usually try to maximize the expected return $G_t$. The rationale justifying this behavior cloning loss is that correct action prediction given a learning history requires the model to approximate the policy improvement operator of the source algorithm. After the pre-training phase (i.e., the RL algorithm distillation), the sequence model can populate its own context by interacting with a new and unseen task and sequentially predicting actions. By imitating the source algorithm, it will effectively explore to gather diverse data at the beginning, assign credit to actions, and then slowly exploit successful policies.
\newline \newline
However, in practice, learning histories can be composed of millions of steps, which would require a significant amount of compute to model, particularly with causal transformers, which are quadratic in the sequence length. Rather than training on full sequences, we can keep one in every $k$ training episodes to effectively subsample the learning trajectories. The downsampling parameter $k$ is a hyper-parameter that needs to be tuned. Too high and the in-context learning becomes unstable as the learning histories are sparse. Too low significantly increases the compute burden and could even be detrimental to the sample efficiency of in-context learning. Additionally, one may also choose to randomly sample across-episodic subsequences $\bar{h_t} := (o_j, a_j, r_j, ..., o_{j+c}, a_{j+c}, r_{j+c})$ of length $c < T$ from $D$ rather than training with entire learning trajectories. However, as we show in section \ref{results}, this can lead to sub-optimal in-context learning or even prevent it.
\section{Related Works}
\label{sec:related_works}
In this section, we review works and methodologies broadly related to AD and Meta-RL. To better contextualize our own work, we divide the meta-RL domain into two categories, namely \textit{Context-based} Meta-RL and \textit{Gradient-based} Meta-RL. We also review some relevant work from the Offline RL field.
\textbf{Context-based Meta Reinforcement Learning: }
Context-based (i.e., gradient-free) adaptation was first proposed and explored simultaneously in two concurrent works, $\text{RL}^2$ \citep{RL2} and \citep{wang2016learning}. Both approaches frame meta-learning as a Reinforcement Learning problem that can be solved using reinforcement learning. Using a memory-based neural network (such as an LSTM \citep{LSTM}), the meta-agent can discover the latent but shared structure underlying multiple related tasks to adapt quickly to unseen goals. Given the advantage of SSMs for sequence modeling compared to alternative models, a recent work \citep{metaS5} investigated the choice of S5 \citep{S5} as a drop-in replacement for online in-context meta-learning ($\text{RL}^2$) and found that it compares favorably to LSTMs while being faster than transformers. Additionally, the authors of \cite{PEARL} depart from the popular memory-based approach and propose using a probabilistic context variable encoded in a VAE's \cite{VAE} latent space to perform posterior sampling for adaptation. However, these works only consider the online Meta-RL setting, which is notoriously more expensive to run compared to the offline setting, as the meta-learner needs to interact with the environment at meta-training time. In our particular setting, AD requires a curated dataset of learning trajectories generated by source RL agents. However, meta-training (or pre-training) is still performed offline.
\textbf{Gradient-based Meta Reinforcement Learning: }
Gradient-based methods, unlike in-context approaches, achieve adaptation by fine-tuning the weights of the meta-learner at inference time (this is also sometimes referred to as \textit{in-weights} adaptation \citep{AD}). A number of works have explored this paradigm. Most notably, MAML \cite{MAML} learns a good parameter initialization point such that downstream adaptation requires as few gradient steps as possible. Many successor meta-learning works build on this framework. For instance, MACAW \citep{MACAW} adapts the popular offline RL algorithm Advantage-Weighted Regression \citep{AWR} to the offline Meta-RL setting by using MAML as their meta-optimizer. For online approaches, meta Q-Learning \citep{MQL} uses a simple context history encoded in a GRU \citep{GRU} along with an off-policy RL algorithm and was shown to outperform PEARL and ($\text{RL}^2$). Nevertheless, gradient-based and particularly MAML-flavored algorithms incur significant computational costs despite their effectiveness.
\textbf{Sequence Modeling for Offline Reinforcement Learning: }
Multiple works have recast Offline Reinforcement Learning as a sequence modeling problem with transformers. Most notably, Decision Transformer \citep{DT} demonstrated that transformers can learn high return policies from offline data. Concurrently, the Trajectory Transformer framework \citep{TT} was proposed to predict the dynamics of the environment for model-based planning. These Transformer-based approaches have demonstrated comparable or superior performance in standard Offline RL benchmarks. But none of these previous approaches demonstrated the emergence of in-context RL. Prompt Decision Transformer (PDT) \citep{PDT} also tackles the in-context offline Meta-RL setting. It investigates the ability of Transformer-based architecture to leverage medium to expert prompt demonstrations for few-shot learning. Their approach achieves new SOTA results in standard benchmarks. However, contrary to AD, Prompt DT requires at least medium-level demonstrations to adapt to a new task, defined as the middle 500 trajectories in the full replay buffer generated by the RL agents. If we consider a task that requires an agent to interact with a real physical system (say a robot), any extra learning on the physical system may requires human supervision, manual resetting and extensive care in how the robot operates to not damage itself or his surroundings. As such, reaching a "medium-level" performance in a test task with this setting to prompt a transformer could be prohibitively expensive or unpractical. To the best of our knowledge, only \citep{metaS5} considered SSMs for RL. This is likely due to the fact that current RL benchmarks are made up of environments with relatively short time horizons, hence not requiring long-range credit assignment. In contrast, AD natively requires modeling long-range dependencies and could benefit from SSMs.
\section{Experiments}
\label{sec:xp}
In this section, we investigate the in-context Reinforcement Learning performance of AD in multiple settings. When adapting to a test task, there is no fine-tuning nor gradient descent involved, contrary to the considered baselines. We pre-train our models with varying context lengths to study how it affects the performance of the downstream in-context learned policies. Following this, we perform a fair comparison between Mamba-based architectures and transformer-based models to assess the potential benefits of State-Space Models (SSMs). Both models have the same number of parameters and are trained using the same dataset. Finally, we explore AD as a competitive meta-RL approach and evaluate it against two prior SOTA methods, MQL \citep{MQL} and MACAW \citep{MACAW}. All experiments were run with a consumer-grade Nvidia RTX 3060 graphics card.
\subsection{Tasks and environments}
\label{sec:tasks}
We use slightly modified versions of four continuous-control RL benchmark environments, adapted to our meta-learning setting. We focus our study on robotic manipulation and locomotion. Just like the original AD paper, these environments support multiple tasks that cannot be solved in a zero-shot manner as observations are not sufficient to infer the task. These prerequisites ensure that the agent has to adapt to an unseen task by trial-and-error. However, unlike previous work, we consider tasks with a longer time horizon which makes across-episodic history modeling more challenging. This addresses one of the limitations acknowledged by the authors of AD. The four environments are briefly described below:
\textbf{Reacher-Goal:} A 2-link robot arm must reach a 2D target in an unknown location inside a circle. As such, the agent needs to explore and perform credit assignment to successfully solve the task. Each episode is 50 timesteps long.
\textbf{Pusher-Goal:} A 7 DoF human-like robot arm has to push a target cylinder to a target location using its end-effector. Both the target cylinder location and the target location are unknown and must be in-context learned. The episodes last 100 timesteps.
\textbf{Half-Cheetah-Vel:} A multi-jointed 2D cheetah-like robot must run at a desired but not observed speed. Once again, exploration and credit assignment are mandatory to learn a successful policy. Episodes are 200 steps long.
\textbf{Ant-Dir:} A high-dimensional locomotion task where a 3D robot has to coordinate its four legs to run as fast as possible in the goal directions (randomly sampled between 0 and 360 degrees). The goal is concealed from the agent and must be inferred through interactions with the environment. Episodes are 200 steps long.
Both Half-Cheetah-Vel and Ant-Dir have already been explored in previous works such as \citep{MAML, MACAW, MQL, PEARL} and are popular benchmark tasks for meta-RL. In this work, we introduce the Pusher-Goal environment to provide a hard exploration problem and challenge the in-context learning abilities of AD. All tasks are run using the Mujoco simulator \citep{Mujoco} with OpenAI Gym \citep{gym}. %Additional details regarding our environments setup can be found in appendix \ref{appendix_task}.
\subsection{Pre-training}
Algorithm Distillation is a two-step process involving first data generation followed by autoregressive model training.
\textbf{Data generation:}
For each environment, a learning trajectory is generated by running a SOTA Reinforcement Learning algorithm on a uniformly sampled task. More specifically, we used Proximal Policy Optimization (PPO) \cite{PPO}, Soft Actor Critic \citep{SAC}, and DroQ \citep{DroQ} depending on which was more appropriate for a particular MDP, both in terms of sample efficiency and training stability. For instance, PPO was very stable (i.e increases the mean return monotonically) but converges slower than off-policy algorithms such as SAC and DroQ. As such the length of the learning trajectories also increases which incur additional cost during pre-training. In contrast, both SAC and DroQ are very sample-efficient but unstable in some cases, with random and sharp performance volatility. This unstable behavior might also emerged during ICRL (as AD imitate the RL algorithm) and decrease performance. Hence the choice to chose different algorithms depending on the environment to find a trade-off between stability and sample-efficiency. Characterizing the impact of data distribution is an interesting and important question which let for future work. Exact hyper-parameters, including the downsampling parameter $k$ and dataset-related technical details, can be found in Appendix \ref{appendix_data}. After data collection, the learning trajectories of the collected tasks are downsampled using the same environment-specific factor $k$. At training time, we sample a batch containing multiple (potentially sub)sequences of different tasks and autoregressively predict the continuous actions.
\textbf{Model training:}
We compare Mamba with a similarly sized Decision Transformer trained on the same data and with the same context length. Exact parameter counts, hyper-parameters, and architecture can be found in Appendix \ref{appendix_models}. Both models output continuous actions conditioned on a sequence of tokens. Unlike the original AD work, we do not consider that each token represents either a state, an action, or a reward. Instead, we concatenate one-step transitions into a single token such that $c^i = (s_{t}^i, a_{t}^i, r_{t}^i, s_{t+1}^i)$. In other words, transition tuples are seen as a single token. In our preliminary experiments, both approaches were competitive, but the former entails using a context $3 \times$ longer, which exceeds our compute budget for all but the Reacher-Goal environment. To improve robustness, we add a small amount of isotropic Gaussian noise to each token of the input sequences.
\subsection{Results}
\label{results}
We analyze AD's adaptation efficacy and focus our discussion on the sample-efficiency and asymptotic performance (i.e., mean return) metrics. We also report the asymptotic performance of both baselines on the same benchmark. For all of our evaluations, we aggregate the performance over 3 pre-training random seeds and for 10 unseen test tasks each.
\begin{figure*}[!ht]
\centering
\begin{subfigure}{0.48\textwidth}
\centering
\includegraphics[width=1\textwidth]{images/reacher_goal.png}
\label{reacher}
\end{subfigure}
\hspace{0.3cm}
\begin{subfigure}{0.48\textwidth}
\centering
\includegraphics[width=1\textwidth]{images/pusher_goal.png}
\label{pusher}
\end{subfigure}
\par\bigskip
\begin{subfigure}{0.48\textwidth}
\centering
\includegraphics[width=1\textwidth]{images/half_cheetah_vel.png}
\end{subfigure}
\hspace{0.3cm}
\begin{subfigure}{0.48\textwidth}
\centering
\includegraphics[width=1\textwidth]{images/ant_dir.png}
\end{subfigure}
\caption{Main results: each learning curve is obtained by pre-training and testing the models with the best performing context length on validation tasks. We evaluate AD with Mamba and a Decision Transformer architecture against MACAW and MQL asymptotic performance.}
\label{inference_xp_perf}
\end{figure*}
\newline\newline
\textbf{What is the best suited architecture for AD ?} Figure~\ref{inference_xp_perf} reports the in-context RL performance for both a Mamba-based architecture and a similarly sized Decision Transformer. More specifically, we plot the best performing context length for each model. As shown, the Mamba variant outperforms the Transformer variant on all the environments. Differences in performance range from a small edge in Reacher-Goal and Half-Cheetah-Vel to a more significant improvement in Ant-Dir and Pusher-Goal. Interestingly, using a Decision Transformer with AD leads to the worst results amongst all baselines in Ant-dir, which is arguably the most complex environment. We believe that this could be evidence for Transformer-based AD not scaling well above a certain task complexity. These results are consistent with multiple other works that found that SSMs are better suited for modeling long-range dependencies, as required for credit assignment. Moreover, from a practical standpoint, we found Mamba to be easier to train, requiring less hyper-parameter tuning and being overall more stable during training.
\begin{figure*}[!ht]
\centering
\begin{subfigure}{0.48\textwidth}
\centering
\includegraphics[width=1\textwidth]{images/reacher_goal_ctx_len.png}
\end{subfigure}
\hspace{0.3cm}
\begin{subfigure}{0.48\textwidth}
\centering
\includegraphics[width=1\textwidth]{images/pusher_goal_ctx_len.png}
\end{subfigure}
\par\bigskip
\begin{subfigure}{0.48\textwidth}
\centering
\includegraphics[width=1\textwidth]{images/half_cheetah_vel_ctx_len.png}
\end{subfigure}
\hspace{0.3cm}
\begin{subfigure}{0.48\textwidth}
\centering
\includegraphics[width=0.99\textwidth]{images/ant_dir_ctx_len.png}
\end{subfigure}
\caption{Impact of context size on ICRL performance. Each learning curve is obtained by pre-training and testing a Mamba model with a particular context length on validation tasks. Except for the simpler Reacher-goal environment, context size seems correlated with task complexity.}
\label{inference_xp_ctx_length}
\end{figure*}
\textbf{How does context length impact ICL performance ?} In-context RL was shown to emerge with AD when the context size was large enough, spanning across episodes, to effectively capture improvements in the training data. Prior work concluded that a 2-4 episodes long (200 steps) context was enough to learn a near-optimal in-context RL algorithm but did not study longer contexts. Besides, the experiments were conducted on relatively simple discrete MDPs. Therefore, we wonder how these observations may evolve when dealing with more challenging environments. We re-investigate this matter by training several AD variants with different context lengths. To go a step further, we study how very long context sizes, ranging from a few episodes to entire (sub-sampled) learning trajectories, affect downstream in-context RL. Results for the Mamba-based architecture are plotted in Figure~\ref{inference_xp_ctx_length}.
Nearly all the context sizes experimented with enable in-context RL, except for the Half-Cheetah-Vel task where only AD with a context containing the (down-sampled) full history seems to effectively learn. For all other smaller context sizes, the AD variants underperform and seem to plateau after trying to zero-shot the task. Even then, performance correlates positively with context length. Although the best performing AD variants in Ant-Dir have the longest context, shorter context lengths also exhibit in-context RL but to a lesser extent. Surprisingly, in Pusher-Goal, shorter context sizes do very well in terms of sample efficiency and asymptotic performance, nearly zero-shotting the new tasks. Still, AD with full context performs better as it seems to explore more extensively. On the other hand, for the simpler Reacher-Goal environment, a shorter context size results in more efficient in-context RL. We hypothesize that due to the relative simplicity of the task, long-range credit assignment may not be required. Overall, the optimal context length seems related to the task complexity, which could be defined in terms of the state-action space dimensionality and the task time-horizon.
\textbf{How does AD compares to other Meta-RL approaches ?} We perform a comparative study of AD adaptation efficiency against MACAW and MQL, two off-policy SOTA meta-RL methods. Rigorously, both are not directly comparable to AD due to incompatible algorithmic designs. For instance, MQL learns online and gets to interact with the environment at meta-training time, unlike MACAW and AD, which both learn from offline data. Nonetheless, our aim is instead to roughly situate the AD approach in relation to other SOTA methods when it comes to overall adaptation efficiency. As such, we analyze their respective asymptotic performance in the most similar configurations possible. As such, MACAW is trained on data coming from the same dataset as AD, and MQL is trained online on the same number of tasks. Results are reported in Figure~\ref{inference_xp_perf}. For AD, we plot the variants with the best performing context length. As a reference, we also plot on the same figure the asymptotic performance of the source RL algorithm. This baseline provides an 'oracle' for asymptotic performance. Both MQL and AD are close to optimal performance in Reacher-Goal. But the gap widens in all other environments except for MQL in Half-Cheetah-Vel. Surprisingly, AD is competitive with MQL even though it was trained fully offline. It even outperforms it on 2 out of 4 tasks but struggles on Half-Cheetah-Vel. Regarding MACAW, AD systematically achieves higher mean returns on all environments. It is worth mentioning that MACAW was trained with significantly less data than in its original paper, which explains its relative under-performance. Indeed, in \cite{MACAW}, MACAW was meta-trained with 45 tasks and evaluated on 5. We found this experimental setting to be very unbalanced and not challenging enough to reflect the extrapolation and adaptation capabilities of AD and other approaches. As such, we limited the amount of meta-training data. For instance, we used 30 tasks during training for our Ant-dir experiments and tested on 30 unseen direction. Similarly, we used only 15 velocities with HalfCheetah-vel.
These results highlight that, while being very simple, AD is a competitive Meta-RL approach when used in conjunction with a powerful sequence model such as Mamba.
\section{Conclusion}
\label{sec:conclusion}
In this work, we addressed arguably one of the main limitations of AD as proposed in the original paper and extended the range of possible applications. We showed that Mamba, a recently proposed SSM, can efficiently handle multi-episodic context and exhibit ICRL when powering AD in complex environments. Experimentally, AD Mamba demonstrated clear advantages over our baseline Decision Transformers, achieving better performance while also being more scalable. Furthermore, experiments showcased that modeling longer contexts tends to improve ICRL on more complex tasks, highlighting the benefit that powerful long-horizon sequential models can bring to AD. This relates to recent findings on many-shot ICL for NLP tasks \cite{ManyShots}. Finally, we showed that the synergy between Mamba and AD enables it to outperform previous offline SOTA approaches and to compete with MQL.
\textbf{Limitations and Future Work:} Despite appealing results from only offline data, some limitations remain. This study considered widely used Meta-RL environments to benchmark AD. However, these environments are still manageable for a transformer-based AD (even if it barely learn in Ant-dir). Evaluating AD on environments with potentially very long horizon and/or which requires extensive exploration is worth while and an interesting extension to this work. Gathering enough learning trajectories to distill a reinforcement learning algorithm could be impractical. Even though the initial cost is amortized by producing a sample-efficient in-context RL algorithm, this procedure may not be suited for applications where online learning is prohibitively expensive. Future work could investigate how AD may leverage existing demonstration datasets that perhaps do not contain enough behavioral diversity (e.g., exploration-exploitation trade-offs) for ICRL. Additionally, future work may explore how the data and the sampling scheme used for pre-training, shape the ICLR abilities. Indeed, a contemporary study \cite{SkillIt} showed that skill or behavior acquisition during pre-training can be impacted by how the data is sampled. Similarly, it could be interesting to research how to mix and sample beginner/medium/expert trajectories for optimal ICRL with a data-driven approach. On the same note, incorporating prompting strategies during inference to improve the sample efficiency of ICRL is also an exciting direction of research. Finally, recent work \citep{RL_diffusion} unveiled the strong performance of diffusion models for behavioral cloning. As such, replacing the simple autoregressive-based distillation with a diffusion-based approach may improve the performance and generalization ability of AD.
\newpage
\bibliography{neurips_2024.bib} % .bib
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\appendix
\newpage
\section{Dataset Generation}
\label{appendix_data}
\begin{table}[!ht]
\centering
\begin{minipage}[t]{0.32\textwidth}
\caption{PPO hyper-parameters}
\label{table:ppo}
\centering
\begin{tabular}{ll}
\toprule
Parameter & Value \\
\midrule
num. epochs & 10 \\
steps per env. & 500 \\
num. envs & 5 \\
entropy coeff. & 5e-3 \\
learning rate & 1e-4 \\
batch size & 64 \\
gamma & 0.99 \\
gae lambda & 0.95 \\
epsilon (clip) & 0.2 \\
value coeff & 0.5 \\
\bottomrule
\end{tabular}
\end{minipage}%
\hfill
\begin{minipage}[t]{0.32\textwidth}
\caption{SAC hyper-parameters}
\label{table:sac}
\centering
\begin{tabular}{ll}
\toprule
Parameter & Value \\
\midrule
buffer size & 1e6 \\
num warm-up steps & 400 \\
num. envs & 5 \\
learning rate & 1e-4 \\
batch size & 256 \\
gamma & 0.99 \\
gae lambda & 0.95 \\
epsilon (clip) & 0.2 \\
value coeff & 0.5 \\
\bottomrule
\end{tabular}
\end{minipage}%
\hfill
\begin{minipage}[t]{0.32\textwidth}
\caption{DroQ hyper-parameters}
\label{table:droq}
\centering
\begin{tabular}{ll}
\toprule
Parameter & Value \\
\midrule
buffer size & 1e6 \\
num. warm-up steps & 400 \\
num. gradient steps & 3 \\
policy update freq. & 3 \\
num. envs & 5 \\
learning rate & 1e-4 \\
batch size & 256 \\
gamma & 0.99 \\
gae lambda & 0.95 \\
epsilon (clip) & 0.2 \\
value coeff & 0.5 \\
\bottomrule
\end{tabular}
\end{minipage}
\end{table}
\begin{table}[!ht]
\caption{Environments hyper-parameters}
\centering
\begin{tabular}{lcccc}
\toprule
Settings & Reacher-goal & Pusher-goal & HalfCheetah-vel & Ant-goal \\
\midrule
RL algorithm & SAC & PPO & SAC & DroQ \\
num. train tasks & 70 & 35 & 15 & 30 \\
Train steps per task & 1e5 & 5e5 & 4e5 & 5e5 \\
\bottomrule
\end{tabular}
\end{table}
\newpage
\section{Models pre-training}
\label{appendix_models}
\begin{table}[!ht]
\caption{Common hyper-parameters}
\centering
\begin{tabular}{lcccc}
\toprule
Settings & Reacher-goal & Pusher-goal & HalfCheetah-vel & Ant-goal \\
\midrule
Optimizer & \multicolumn{4}{c}{Adam} \\
Epochs & \multicolumn{4}{c}{450} \\
Learning rate Scheduler & \multicolumn{4}{c}{cosine with warmup} \\
Downsampling rate k & 4 & 10 & 8 & 10 \\
weight decay & 1e-4 & 2e-4 & 5e-4 & 5e-4\\
\bottomrule
\end{tabular}
\end{table}
\begin{table}[!ht]
\caption{Mamba hyper-parameters}
\centering
\begin{tabular}{lcccc}
\toprule
Settings & Reacher-goal & Pusher-goal & HalfCheetah-vel & Ant-goal \\
\midrule
Embedding dim. & \multicolumn{4}{c}{32} \\
Dim. model & 384 & 512 & 512 & 512 \\
num. layers & 6 & 8 & 8 & 8\\
num. params. & 6M & 12M & 12M & 12M \\
\bottomrule
\end{tabular}
\end{table}
\begin{table}[!ht]
\caption{Decision Transformer hyper-parameters}
\centering
\begin{tabular}{lcccc}
\toprule
Settings & Reacher-goal & Pusher-goal & HalfCheetah-vel & Ant-goal \\
\midrule
num. heads & \multicolumn{4}{c}{4} \\
pos. encoding & \multicolumn{4}{c}{Absolute} \\
feed fwd. dim. & \multicolumn{4}{c}{2048} \\
Embedding dim. & 42 & 64 & 64 & 128 \\
num. layers & 8 & 9 & 9 & 6\\
num. params. & 6M & 12M & 12M & 14M \\
\bottomrule
\end{tabular}
\end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% \newpage
% \section*{NeurIPS Paper Checklist}
% \begin{enumerate}
% \item {\bf Claims}
% \item[] Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
% \item[] Answer: \answerYes{} % Replace by , \answerNo{}, or \answerNA{}.
% \item[] Justification: The abstract and introduction claims that Mamba can be used to scale Algorithm Distillation to more complex Meta RL environments and make it competitive with SOTA baselines. We believe our results support well our claims.
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the abstract and introduction do not include the claims made in the paper.
% \item The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
% \item The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
% \item It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
% \end{itemize}
% \item {\bf Limitations}
% \item[] Question: Does the paper discuss the limitations of the work performed by the authors?
% \item[] Answer: \answerYes{}{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification: Limitations are discussed in the Final section along future work?
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
% \item The authors are encouraged to create a separate "Limitations" section in their paper.
% \item The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
% \item The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
% \item The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
% \item The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
% \item If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
% \item While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
% \end{itemize}
% \item {\bf Theory Assumptions and Proofs}
% \item[] Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
% \item[] Answer: \answerNA{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification: This paper only exposes experimental results.
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the paper does not include theoretical results.
% \item All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced.
% \item All assumptions should be clearly stated or referenced in the statement of any theorems.
% \item The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
% \item Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
% \item Theorems and Lemmas that the proof relies upon should be properly referenced.
% \end{itemize}
% \item {\bf Experimental Result Reproducibility}
% \item[] Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
% \item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification: All hyper-parameters are disclosed in the appendices. We also detailed our experimental setup. The code will be made public upon publication.
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the paper does not include experiments.
% \item If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
% \item If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
% \item Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
% \item While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
% \begin{enumerate}
% \item If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
% \item If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
% \item If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
% \item We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
% \end{enumerate}
% \end{itemize}
% \item {\bf Open access to data and code}
% \item[] Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?
% \item[] Answer: \answerNo{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification: The code with all relevant scripts for reproductibility will be made public upon publication. Right now it is still very much a draft and not self contained.
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that paper does not include experiments requiring code.
% \item Please see the NeurIPS code and data submission guidelines (\url{https://nips.cc/public/guides/CodeSubmissionPolicy}) for more details.
% \item While we encourage the release of code and data, we understand that this might not be possible, so “No” is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
% \item The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (\url{https://nips.cc/public/guides/CodeSubmissionPolicy}) for more details.
% \item The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
% \item The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.
% \item At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
% \item Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
% \end{itemize}
% \item {\bf Experimental Setting/Details}
% \item[] Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
% \item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification: All the important training and test details are either included in the main text or the appendices.
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the paper does not include experiments.
% \item The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.
% \item The full details can be provided either with the code, in appendix, or as supplemental material.
% \end{itemize}
% \item {\bf Experiment Statistical Significance}
% \item[] Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
% \item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification: Results for the experiments are reported for 3 training seeds with 10 test seeds each.
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the paper does not include experiments.
% \item The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
% \item The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
% \item The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
% \item The assumptions made should be given (e.g., Normally distributed errors).
% \item It should be clear whether the error bar is the standard deviation or the standard error of the mean.
% \item It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96\% CI, if the hypothesis of Normality of errors is not verified.
% \item For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
% \item If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
% \end{itemize}
% \item {\bf Experiments Compute Resources}
% \item[] Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
% \item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification: As mentioned in the experiments section, we only used a Nvidia RTX 3060 GPU.
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the paper does not include experiments.
% \item The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
% \item The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
% \item The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
% \end{itemize}
% \item {\bf Code Of Ethics}
% \item[] Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics \url{https://neurips.cc/public/EthicsGuidelines}?
% \item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification: It does not infringe any of the NeurIPS ethics guidelines.
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
% \item If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
% \item The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).
% \end{itemize}
% \item {\bf Broader Impacts}
% \item[] Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
% \item[] Answer: \answerNA{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification: We believe our work does not have potential positive societal impacts or negative societal impacts.
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that there is no societal impact of the work performed.
% \item If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
% \item Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
% \item The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
% \item The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
% \item If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
% \end{itemize}
% \item {\bf Safeguards}
% \item[] Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
% \item[] Answer: \answerNA{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification: We believe our work does not have a high risk for misuse
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the paper poses no such risks.
% \item Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
% \item Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
% \item We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.
% \end{itemize}
% \item {\bf Licenses for existing assets}
% \item[] Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
% \item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification: We cited the authors of Mamba and FlashAttention and recognize that our work is an extension of Algorithm Distillation, a previously published paper. All libraries used are also credited and MIT licensed.
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the paper does not use existing assets.
% \item The authors should cite the original paper that produced the code package or dataset.
% \item The authors should state which version of the asset is used and, if possible, include a URL.
% \item The name of the license (e.g., CC-BY 4.0) should be included for each asset.
% \item For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
% \item If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, \url{paperswithcode.com/datasets} has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.
% \item For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
% \item If this information is not available online, the authors are encouraged to reach out to the asset's creators.
% \end{itemize}
% \item {\bf New Assets}
% \item[] Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
% \item[] Answer: \answerNA{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification: We did not introduced new assets.
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the paper does not release new assets.
% \item Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
% \item The paper should discuss whether and how consent was obtained from people whose asset is used.
% \item At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
% \end{itemize}
% \item {\bf Crowdsourcing and Research with Human Subjects}
% \item[] Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
% \item[] Answer: \answerNA{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification: The paper does not involve crowdsourcing nor research with human subjects.
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
% \item Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.
% \item According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
% \end{itemize}
% \item {\bf Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects}
% \item[] Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
% \item[] Answer: \answerNA{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification: The paper does not involve crowdsourcing nor research with human subjects.
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
% \item Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
% \item We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
% \item For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
% \end{itemize}
% \end{enumerate}
\end{document}
\begin{thebibliography}{37}
\providecommand{\natexlab}[1]{#1}
\providecommand{\url}[1]{\texttt{#1}}
\expandafter\ifx\csname urlstyle\endcsname\relax
\providecommand{\doi}[1]{doi: #1}\else
\providecommand{\doi}{doi: \begingroup \urlstyle{rm}\Url}\fi
\bibitem[Vaswani et~al.(2017)Vaswani, Shazeer, Parmar, Uszkoreit, Jones, Gomez,
Kaiser, and Polosukhin]{Attention}
A.~Vaswani, N.~Shazeer, N.~Parmar, J.~Uszkoreit, L.~Jones, A.~N. Gomez,
{\L}.~Kaiser, and I.~Polosukhin.
\newblock Attention is all you need.
\newblock \emph{Advances in neural information processing systems}, 30, 2017.
\bibitem[Oquab et~al.(2024)Oquab, Darcet, Moutakanni, Vo, Szafraniec, Khalidov,
Fernandez, HAZIZA, Massa, El-Nouby, Assran, Ballas, Galuba, Howes, Huang, Li,
Misra, Rabbat, Sharma, Synnaeve, Xu, Jegou, Mairal, Labatut, Joulin, and
Bojanowski]{DinoV2}
M.~Oquab, T.~Darcet, T.~Moutakanni, H.~V. Vo, M.~Szafraniec, V.~Khalidov,
P.~Fernandez, D.~HAZIZA, F.~Massa, A.~El-Nouby, M.~Assran, N.~Ballas,
W.~Galuba, R.~Howes, P.-Y. Huang, S.-W. Li, I.~Misra, M.~Rabbat, V.~Sharma,
G.~Synnaeve, H.~Xu, H.~Jegou, J.~Mairal, P.~Labatut, A.~Joulin, and
P.~Bojanowski.
\newblock {DINO}v2: Learning robust visual features without supervision.
\newblock \emph{Transactions on Machine Learning Research}, 2024.
\newblock ISSN 2835-8856.
\bibitem[Kirillov et~al.(2023)Kirillov, Mintun, Ravi, Mao, Rolland, Gustafson,
Xiao, Whitehead, Berg, Lo, et~al.]{SAM}
A.~Kirillov, E.~Mintun, N.~Ravi, H.~Mao, C.~Rolland, L.~Gustafson, T.~Xiao,
S.~Whitehead, A.~C. Berg, W.-Y. Lo, et~al.
\newblock Segment anything.
\newblock In \emph{Proceedings of the IEEE/CVF International Conference on
Computer Vision}, pages 4015--4026, 2023.
\bibitem[Vuong et~al.(2023)Vuong, Levine, Walke, Pertsch,
et~al.]{vuong2023open}
Q.~Vuong, S.~Levine, H.~R. Walke, K.~Pertsch, et~al.
\newblock Open x-embodiment: Robotic learning datasets and {RT}-x models.
\newblock In \emph{Towards Generalist Robots: Learning Paradigms for Scalable
Skill Acquisition @ CoRL2023}, 2023.
\bibitem[Laskin et~al.(2023)Laskin, Wang, Oh, Parisotto, Spencer, Steigerwald,
Strouse, Hansen, Filos, Brooks, maxime gazeau, Sahni, Singh, and Mnih]{AD}
M.~Laskin, L.~Wang, J.~Oh, E.~Parisotto, S.~Spencer, R.~Steigerwald,
D.~Strouse, S.~S. Hansen, A.~Filos, E.~Brooks, maxime gazeau, H.~Sahni,
S.~Singh, and V.~Mnih.
\newblock In-context reinforcement learning with algorithm distillation.
\newblock In \emph{The Eleventh International Conference on Learning
Representations}, 2023.
\bibitem[Tay et~al.(2020)Tay, Dehghani, Abnar, Shen, Bahri, Pham, Rao, Yang,
Ruder, and Metzler]{LRA}
Y.~Tay, M.~Dehghani, S.~Abnar, Y.~Shen, D.~Bahri, P.~Pham, J.~Rao, L.~Yang,
S.~Ruder, and D.~Metzler.
\newblock Long range arena: A benchmark for efficient transformers.
\newblock In \emph{International Conference on Learning Representations}, 2020.
\bibitem[Liu et~al.(2024)Liu, Lin, Hewitt, Paranjape, Bevilacqua, Petroni, and
Liang]{LostInMiddle}
N.~F. Liu, K.~Lin, J.~Hewitt, A.~Paranjape, M.~Bevilacqua, F.~Petroni, and
P.~Liang.
\newblock Lost in the middle: How language models use long contexts.
\newblock \emph{Transactions of the Association for Computational Linguistics},
12:\penalty0 157--173, 2024.
\bibitem[Zhu et~al.(2024)Zhu, Liao, Zhang, Wang, Liu, and Wang]{VisionMamba}
L.~Zhu, B.~Liao, Q.~Zhang, X.~Wang, W.~Liu, and X.~Wang.
\newblock Vision mamba: Efficient visual representation learning with
bidirectional state space model.
\newblock \emph{arXiv preprint arXiv:2401.09417}, 2024.
\bibitem[Gu et~al.(2022)Gu, Goel, and Re]{S4}
A.~Gu, K.~Goel, and C.~Re.
\newblock Efficiently modeling long sequences with structured state spaces.
\newblock In \emph{International Conference on Learning Representations}, 2022.
\bibitem[Gu and Dao(2023)]{Mamba}
A.~Gu and T.~Dao.
\newblock Mamba: Linear-time sequence modeling with selective state spaces.
\newblock \emph{arXiv preprint arXiv:2312.00752}, 2023.
\bibitem[Grazzi et~al.(2024)Grazzi, Siems, Schrodi, Brox, and Hutter]{MambaICL}
R.~Grazzi, J.~Siems, S.~Schrodi, T.~Brox, and F.~Hutter.
\newblock Is mamba capable of in-context learning?
\newblock \emph{arXiv preprint arXiv:2402.03170}, 2024.
\bibitem[Dao et~al.(2022)Dao, Fu, Ermon, Rudra, and R{\'e}]{FlashAttention}
T.~Dao, D.~Fu, S.~Ermon, A.~Rudra, and C.~R{\'e}.
\newblock Flashattention: Fast and memory-efficient exact attention with
io-awareness.
\newblock \emph{Advances in Neural Information Processing Systems},
35:\penalty0 16344--16359, 2022.
\bibitem[Brown et~al.(2020)Brown, Mann, Ryder, Subbiah, Kaplan, Dhariwal,
Neelakantan, Shyam, Sastry, Askell, et~al.]{GPT3}
T.~Brown, B.~Mann, N.~Ryder, M.~Subbiah, J.~D. Kaplan, P.~Dhariwal,
A.~Neelakantan, P.~Shyam, G.~Sastry, A.~Askell, et~al.
\newblock Language models are few-shot learners.
\newblock \emph{Advances in neural information processing systems},
33:\penalty0 1877--1901, 2020.
\bibitem[Gu et~al.(2020)Gu, Dao, Ermon, Rudra, and R{\'e}]{Hippo}
A.~Gu, T.~Dao, S.~Ermon, A.~Rudra, and C.~R{\'e}.
\newblock Hippo: Recurrent memory with optimal polynomial projections.
\newblock \emph{Advances in neural information processing systems},
33:\penalty0 1474--1487, 2020.
\bibitem[Duan et~al.(2016)Duan, Schulman, Chen, Bartlett, Sutskever, and
Abbeel]{RL2}
Y.~Duan, J.~Schulman, X.~Chen, P.~L. Bartlett, I.~Sutskever, and P.~Abbeel.
\newblock Rl 2: Fast reinforcement learning via slow reinforcement learning.
\newblock \emph{arXiv preprint arXiv:1611.02779}, 2016.
\bibitem[Wang et~al.(2016)Wang, Kurth-Nelson, Tirumala, Soyer, Leibo, Munos,
Blundell, Kumaran, and Botvinick]{wang2016learning}
J.~X. Wang, Z.~Kurth-Nelson, D.~Tirumala, H.~Soyer, J.~Z. Leibo, R.~Munos,
C.~Blundell, D.~Kumaran, and M.~Botvinick.
\newblock Learning to reinforcement learn.
\newblock \emph{arXiv preprint arXiv:1611.05763}, 2016.
\bibitem[Hochreiter and Schmidhuber(1997)]{LSTM}
S.~Hochreiter and J.~Schmidhuber.
\newblock Long short-term memory.
\newblock \emph{Neural computation}, 9\penalty0 (8):\penalty0 1735--1780, 1997.
\bibitem[Lu et~al.(2024)Lu, Schroecker, Gu, Parisotto, Foerster, Singh, and
Behbahani]{metaS5}
C.~Lu, Y.~Schroecker, A.~Gu, E.~Parisotto, J.~Foerster, S.~Singh, and
F.~Behbahani.
\newblock Structured state space models for in-context reinforcement learning.
\newblock \emph{Advances in Neural Information Processing Systems}, 36, 2024.
\bibitem[Smith et~al.(2023)Smith, Warrington, and Linderman]{S5}
J.~T. Smith, A.~Warrington, and S.~Linderman.
\newblock Simplified state space layers for sequence modeling.
\newblock In \emph{The Eleventh International Conference on Learning
Representations}, 2023.
\bibitem[Rakelly et~al.(2019)Rakelly, Zhou, Finn, Levine, and Quillen]{PEARL}
K.~Rakelly, A.~Zhou, C.~Finn, S.~Levine, and D.~Quillen.
\newblock Efficient off-policy meta-reinforcement learning via probabilistic
context variables.
\newblock In \emph{International conference on machine learning}, pages
5331--5340. PMLR, 2019.
\bibitem[Kingma and Welling(2014)]{VAE}
D.~P. Kingma and M.~Welling.
\newblock Auto-encoding variational bayes.
\newblock In \emph{2nd International Conference Learning Representation
(ICLR)}, 2014.
\bibitem[Finn et~al.(2017)Finn, Abbeel, and Levine]{MAML}
C.~Finn, P.~Abbeel, and S.~Levine.
\newblock Model-agnostic meta-learning for fast adaptation of deep networks.
\newblock In \emph{International conference on machine learning}, pages
1126--1135. PMLR, 2017.
\bibitem[Mitchell et~al.(2021)Mitchell, Rafailov, Peng, Levine, and
Finn]{MACAW}
E.~Mitchell, R.~Rafailov, X.~B. Peng, S.~Levine, and C.~Finn.
\newblock Offline meta-reinforcement learning with advantage weighting.
\newblock In \emph{International Conference on Machine Learning}, pages
7780--7791. PMLR, 2021.
\bibitem[Peng et~al.(2019)Peng, Kumar, Zhang, and Levine]{AWR}
X.~B. Peng, A.~Kumar, G.~Zhang, and S.~Levine.
\newblock Advantage-weighted regression: Simple and scalable off-policy
reinforcement learning.
\newblock \emph{arXiv preprint arXiv:1910.00177}, 2019.
\bibitem[Fakoor et~al.(2019)Fakoor, Chaudhari, Soatto, and Smola]{MQL}
R.~Fakoor, P.~Chaudhari, S.~Soatto, and A.~J. Smola.
\newblock Meta-q-learning.
\newblock In \emph{International Conference on Learning Representations}, 2019.
\bibitem[Cho et~al.(2014)Cho, van Merrienboer, Gulcehre, Bahdanau, Bougares,
Schwenk, and Bengio]{GRU}
K.~Cho, B.~van Merrienboer, C.~Gulcehre, D.~Bahdanau, F.~Bougares, H.~Schwenk,
and Y.~Bengio.
\newblock Learning phrase representations using rnn encoder--decoder for
statistical machine translation.
\newblock In \emph{Proceedings of the 2014 Conference on Empirical Methods in
Natural Language Processing (EMNLP)}, page 1724. Association for
Computational Linguistics, 2014.
\bibitem[Chen et~al.(2021)Chen, Lu, Rajeswaran, Lee, Grover, Laskin, Abbeel,
Srinivas, and Mordatch]{DT}
L.~Chen, K.~Lu, A.~Rajeswaran, K.~Lee, A.~Grover, M.~Laskin, P.~Abbeel,
A.~Srinivas, and I.~Mordatch.
\newblock Decision transformer: Reinforcement learning via sequence modeling.
\newblock \emph{Advances in neural information processing systems},
34:\penalty0 15084--15097, 2021.
\bibitem[Janner et~al.(2021)Janner, Li, and Levine]{TT}
M.~Janner, Q.~Li, and S.~Levine.
\newblock Offline reinforcement learning as one big sequence modeling problem.
\newblock \emph{Advances in neural information processing systems},
34:\penalty0 1273--1286, 2021.
\bibitem[Xu et~al.(2022)Xu, Shen, Zhang, Lu, Zhao, Tenenbaum, and Gan]{PDT}
M.~Xu, Y.~Shen, S.~Zhang, Y.~Lu, D.~Zhao, J.~Tenenbaum, and C.~Gan.
\newblock Prompting decision transformer for few-shot policy generalization.
\newblock In \emph{international conference on machine learning}, pages
24631--24645. PMLR, 2022.
\bibitem[Todorov et~al.(2012)Todorov, Erez, and Tassa]{Mujoco}
E.~Todorov, T.~Erez, and Y.~Tassa.
\newblock Mujoco: A physics engine for model-based control.
\newblock In \emph{2012 IEEE/RSJ international conference on intelligent robots
and systems}, pages 5026--5033. IEEE, 2012.
\bibitem[Brockman et~al.(2016)Brockman, Cheung, Pettersson, Schneider,
Schulman, Tang, and Zaremba]{gym}
G.~Brockman, V.~Cheung, L.~Pettersson, J.~Schneider, J.~Schulman, J.~Tang, and
W.~Zaremba.
\newblock Openai gym.
\newblock \emph{arXiv preprint arXiv:1606.01540}, 2016.
\bibitem[Schulman et~al.(2017)Schulman, Wolski, Dhariwal, Radford, and
Klimov]{PPO}
J.~Schulman, F.~Wolski, P.~Dhariwal, A.~Radford, and O.~Klimov.
\newblock Proximal policy optimization algorithms.
\newblock In \emph{arXiv preprint arXiv:1707.06347}, 2017.
\bibitem[Haarnoja et~al.(2018)Haarnoja, Zhou, Abbeel, and Levine]{SAC}
T.~Haarnoja, A.~Zhou, P.~Abbeel, and S.~Levine.
\newblock Soft actor-critic: Off-policy maximum entropy deep reinforcement
learning with a stochastic actor.
\newblock In \emph{International conference on machine learning}, pages
1861--1870. PMLR, 2018.
\bibitem[Hiraoka et~al.(2022)Hiraoka, Imagawa, Hashimoto, Onishi, and
Tsuruoka]{DroQ}
T.~Hiraoka, T.~Imagawa, T.~Hashimoto, T.~Onishi, and Y.~Tsuruoka.
\newblock Dropout q-functions for doubly efficient reinforcement learning.
\newblock In \emph{International Conference on Learning Representations}, 2022.
\bibitem[Agarwal et~al.(2024)Agarwal, Singh, Zhang, Bohnet, Chan, Anand, Abbas,
Nova, Co-Reyes, Chu, et~al.]{ManyShots}
R.~Agarwal, A.~Singh, L.~M. Zhang, B.~Bohnet, S.~Chan, A.~Anand, Z.~Abbas,
A.~Nova, J.~D. Co-Reyes, E.~Chu, et~al.
\newblock Many-shot in-context learning.
\newblock \emph{arXiv preprint arXiv:2404.11018}, 2024.
\bibitem[Chen et~al.(2023)Chen, Roberts, Bhatia, WANG, Zhang, Sala, and
Re]{SkillIt}
M.~F. Chen, N.~Roberts, K.~Bhatia, J.~WANG, C.~Zhang, F.~Sala, and C.~Re.
\newblock Skill-it! a data-driven skills framework for understanding and
training language models.
\newblock In \emph{Thirty-seventh Conference on Neural Information Processing
Systems}, 2023.
\bibitem[Chi et~al.(2023)Chi, Feng, Du, Xu, Cousineau, Burchfiel, and
Song]{RL_diffusion}
C.~Chi, S.~Feng, Y.~Du, Z.~Xu, E.~Cousineau, B.~Burchfiel, and S.~Song.
\newblock Diffusion policy: Visuomotor policy learning via action diffusion.
\newblock \emph{Robotics: Science and Systems 2023 (RSS)}, 2023.
\end{thebibliography}