Random Policy Enables In-Context Reinforcement Learning within Trust Horizons
Authors: Weiqin Chen
Santiago Paternain
Venue: Accepted by TMLR ()

Abstract: Pretrained foundation models (FMs) have exhibited extraordinary in-context learning performance, allowing zero-shot (or few-shot) generalization to new environments/tasks not encountered during the pretraining. In the case of reinforcement learning (RL), in-context RL (ICRL) emerges when pretraining FMs on decision-making problems in an autoregressive-supervised manner. Nevertheless, the current state-of-the-art ICRL algorithms, such as Algorithm Distillation,  Decision Pretrained Transformer and Decision Importance Transformer, impose stringent requirements on the pretraining dataset concerning the behavior (source) policies, context information, and action labels, etc. Notably, these algorithms either demand optimal policies or require varying degrees of well-trained behavior policies for all pretraining environments. This significantly hinders the application of ICRL to real-world scenarios, where acquiring optimal or well-trained policies for a substantial volume of real-world training environments can be prohibitively expensive or even intractable. To overcome this challenge, we introduce a novel approach, termed State-Action Distillation (SAD), that allows to generate an effective pretraining dataset guided solely by random policies. In particular, SAD selects query states and corresponding action labels by distilling the outstanding state-action pairs from the entire state and action spaces by using random policies within a trust horizon, and then inherits the classical autoregressive-supervised mechanism during the pretraining. To the best of our knowledge, this is the first work that enables effective ICRL under (e.g., uniform) random policies and random contexts. We also establish the quantitative analysis of the trustworthiness as well as the performance guarantees of our SAD approach. Moreover, our empirical results across multiple popular ICRL benchmark environments demonstrate that, on average, SAD outperforms the best baseline by 236.3% in the offline evaluation and by 135.2% in the online evaluation.

Published in Transactions on Machine Learning Research (04/2025)
Random Policy Enables In- In-Context Reinforcement Learning within Trust Horizons
Welqin Chen
chenw 8 edu
Deportment of Electricol, Computer, and Systems Engineering, Rensselaer Polytechmic Institute
Santiago Paternain
paters@rpi.cod edu
Deportment of Electrionl, Computer, and Systems Engineering. Rensselaer Polytechnic Institute
Reviewed on OpenReview: https: //tp:///pppnraee opwrrevew..........m net/ por/fo id- mA 95MM
Abstract
Pretrained foundation models (FMs) have exhibited extraordinary in-context learning per- formance, allowing zero-shot (or few-shot) generalization to new environments/taate not en- countered during the pretraining. In the case of reinforcement learning (RL), in-context RL (ICRL) emerges when pretraining FMs on pocision-makinn problems an autoregressive- supervised manner. Nevertheless, the current state-oftthe-ar. ICRL algorithms, such as Algorithm Distillation, Decision Pretrained Transformer and Decision Importance Trans- former, impose stringent requirements on the pretraining dataset concerning the behavior (source) policies, context information, and action labels, etc. Notably, these algorithms either demand optimal policies or require varying degrees of well-trained behavior policies for all pretraining environments. This significantly hinders the application of ICRL to real- world scenarios, where acquiring optimal or well-trained policics for substantial volume of real-world training environments can be prohibitively expensive even intractable. To over- come this challenge, we introduce novel approach, termed State-Action Distillation (SAD), that allows to generate an effective pretraining dataset guided solely by random policies. In particular, SAD selects query states and corresponding action labels by distilling the out- standing state-action pairs from the entire state and action spaces by using random policies within a trust horizon, and then inherits the classical anttorgreesive-uuprriiie mechanism during the pretraining. To the best of our knowledge, this the first work that enables fective ICRL under (e.g., uniform) random policies and random contexts. We also establish the quantitative analysis of the trustworthiness well the performance guarantees of our SAD approach. Moreover, environments demonstrate that, on average, SAD outperforms the best baseline by 236.3% in the offline evaluation and by 135.2% in the online evaluation.
empirical results across multiple popular ICRL benchmark
1 Introduction
Pretrained foundation models (FMs) have demonstrated promising performance across a wide variety of
domains in artificial intelligence including natural language processing (NLP) (Devlin] 2018
Radford] 2018
Radford al. 2019 Brown 2020), computer vision (CV) Yuuan al.
2021
Sammani et al.
2022
Ma et al.
2023
Chen et al.
2024b) and sequential decision-making (Chen et al.
202
Janner al.
2021
Xu et al.
2022b Yang et al.
2023 Light ct al. 2024a b). This success is attributed to FMs' impressive capability of
in-context learning (Dong el al. 2022
Liel el al.
2023 Wei et al. 2023 Wies et al. 2024) which refers to the
ability to infer and understand the new (unseen) tasks provided with the context information (or prompt) and
without model parameters updates. Recently, in-context reinforcement learning (ICRL) (Laskin et al. 2022
Grigsby et al.
2023 Lin
2023
Sinii et al.
2023 Zisman et al.) 2023 Lee al. 2024 Lu et al. 2024
Wang et. al.
2024
Dong al. 2024 has emerged when FMs are pretrained sequential decision-making

---

Published in Transactions on Machine Learning Research (04/2025)
problems. Whereas FMs use texts the context/prompe in NLP, ICRL treats the state-action-reward tuples as the contextual information for decision-making.
However, the current state-of-the-art (SOTA) ICRL algorithms impose strict requirements on the pretraining datasets. More specifically, Algorithm Distillation (AD) (Laskin et al. 2022) requires the context to contain the complete learning history (from the initial policy to the final-trained policy) of the source (or behavior) RT. algorithm for all pretraining environments. In addition, AD requires environments to have short episodes, allowing the context to capture cross-episodic information. This enables AD to learn the improvement operator of the source RL algorithm. Conversely, Decision Pretrained Transformer (DPT) (Lee et al. 2024 partially relaxes the requirement on the context, permitting it to be gathered by random policies and without needing to adhere to the transition dynamics. Nevertheless, DPT necessitates the optimal policy to label an optimal action for any randomly sampled query state across all pretraining environments. To explore the feasibility of ICRL in the absence of optimal policies, Decision Importance Transformer (DIT) (Dong et al. 2024) proposes to leverage the observed state-action pairs in the context data as query states and corresponding action labels. Each state-action pair within the context is assigned a weight. in the training process. This weight is proportional to the return-to-go of the pair, Thus, DIT prioritizes the training high-return pairs. Despite not demanding optimal policies, DIT still requires a substantial context dataset to comprehensively cover all state-action pairs from the state and action spaces. Furthermore, DIT still mandates that more than 30% of the context data comes from well-trained policies to ensure the coverage of good action labels, and the context should originate from a complete episode.
Notably, acquiring either optimal policies or well-trained policies across multitude of pretraining environ- ments in real-world scenarios can be prohibitively expensive even intractable. On the other hand, the transition data available in real-world problems collected as the context may not originate from a complete episode. These stringent requirements on the pretraining dataset of the SOTA ICRL algorithms severely limit their practical applications to the real world, especially for those where the transition data exhibits high variance to train an effective policy. This work thus aims to relax these requirements by considering an untrained random policy. Although the random policy cannot solve all decision--aainng problems, it is worth noting that the random policy and the optimal policy select the same optimal actions in certain problems, e.g. the grid world navigation problem (Laskin et al., 2022 Lee et al.] 2024 Dong et al., 2024 where the reward is received only upon reaching the unique goal. In these problems, the optimal action induced by both policies corresponds to navigating to the goal as quickly as possible. Consequently, this paper centers on the ICRL that operates without the need for optimal (or any degree of well-trained) policies or episodic context, placing its emphasis on the scenarios under (c.g., uniform) random policies and random contexts only. Our applicable domains in this work focus on the multi-armed bandits and the grid world navigation problems with a single sparse reward received only upon achieving the unique goal.
1.1 Main Contributions
The main contributions of this work are summarized as follows.
We propose a novel approach termed State-Action Distillation (SAD) to generate the pretraining dataset of ICRL under random policics. Notably, SAD distills the outstanding state-action pairs over the entire state and action spaces for the query states and corresponding action labels (refer to Figure by executing all possible actions under the random policies within a trust. horizon.
To the best of our knowledge, SAD stands the first method that enables effective ICRL under (e.g,, uniform) random policies and random contexts.
We establish the quantitative analysis of the trustworthiness as well as the performance guarantees of our SAD approach. We substantiate the efficacy of SAD by empirical results on several popular ICRL benchmark environments. On average, SAD significantly outperforms all existing SOTA ICRL algorithms. More concretely, SAD surpasses the best bascline by 236.3% in the offine evaluation and by 135.2% in the online evaluation.
2

---

Published in Transactions on Machine Learning Research (04/2025)
Envs
Context
Trust Horizon
iii)
RS
8p
Random
FM
Action Label
argmax
R
fi)
State
Query State
Space
Space
(v)
Figure 1: Schematic of the State-Action Distillation approach: i) Collecting the context by using the random policy to interact with pretraining environments. ii) Sampling a query state randomly from the state space. iii) Starting from the query state and any action in action space, running trust horizons under the random policy, and distilling the action label by the action that yiclds the maximal return. iv) Pretraining foundation models in a supervised mechanism, which predicts the action label given the context and query state.
2
Related Work
2.1 Offline Reinforcement Learning
In contrast the unlimited interactions with the environment online RL, offline RL seeks to learn optimal
policies from pre-collected and static dataset (Fujimoto al., 2019 Levine al. 2020 Kumar et al. 2020:
Kostrikov et al. 2021 Chen al., 2024a). One of the critical challenges in offline
RL is with bootstrapping
from out-of-distributin (OOD) actions (Levine et al. 2020 Kumar et al. 2020 Xu et al. 2022a Liu et al.
2024h due to the mismatch between the behavior policies and the learned policies. To address this issue, the current SOTA offline RL algorithms propose to update pessimistically by cither adding a regularization or underestimating the Q-value OOD actions.
2.2
Sntregressive-Supervisee Decision Making
In addition to the traditional offline RL methods, autoreerresiive-suervviss mechanisms based on the trans- former architecture [Vaswani 2017) have been successfully applied to offline decision making domains by their powerful capability sequential modeling. The pioncering work in the autoregrestive-suerriess deci- sion making the Decision Transformer (DT) (Chen et al., 2021). DT autoregressively models the sequence of actions from the historical offline data conditioned on the sequence of returns in the history. During the inference, the trained model be queried based on pre-defined target returns, allowing it to generate actions aligned with the target returns. The subsequent works such Multi-Game Decision Transformer (MGDT) [Lee et al.) 2022) and Gato (Reed et al. 2022) have exhibited the success of the autoregressive- supervised mechanisms in learning multi-lask policies by fine-tuning or leveraging expert demonstrations in the downstream tasks.
2.3 In-Context Reinforcement Learning
However, both traditional offline RL and autorreressive--uperrvidd decision making mechanisms suffer from the poor zcro-shot generalization and in-context lcarning capabilities to new environments, neither can improve the policy, with a fixed trained model, in context by trial and error. In-contexi. reinforcement learning (ICRL) aims to pretrain a transformer- based FM, such as GPT2 (Radford et al. 2019), across a wide range of pretraining environments. During the evaluation (or inference), the protrained model can directly infer the unseen environment and learn in-context without the need for updating model parameters. The SOTA ICRL algorithms including AD Laskin et. al. 2022), DPT [Tee et al. 2024) and DIT (Dong al.

---

Published in Transactions on Machine Learning Research (04/2025)
2024h have demonstrated the potential of the ICRL framework. Nevertheless, each of these methods imposes distinct yet strict. requirements on the pretraining dataset e.g., requiring well-trained (or even optimal) behavior policies, the context to be episodic and/or substantial, which significantly restrict their practicality in real-world applications. Accordingly, mastering and executing ICRL under (o.g., uniform) random policies and random contexts remains a crucial direction and a critical challenge.
3 In-Context Reinforcement Learning
This section introduces the background of ICRL mechanisms and three SOTA ICRL algorithms. We start by presenting the preliminaries of ICRL.
3.1 Preliminaries
RL problems are generally formulated as Markov Decision Processes (MDPs) Sutton 2018). An MDP can be reprosented by a tuple (S, A, R, P.P). where S and A denote finite state and action spaces, R: S x A + R denotes the reward function that evaluates the quality of the action, P: S x A x S - [o, 1] denotes the transition probability that describes the dynamics of the system, and p: S + [0, 1] denotes the initial state distribution.
A policy defines a mapping from states to probability distributions over actions, providing a strategy that guides the agent in the decision making. The agent interacts with the environment following the policy x and the transition dynamics of the system, and then generates an episode of the transition data (So, 40 To: ...). The performance measure (x) is defined by the expected discounted cumulative reward under the policy x
80
J(x) = E.
Σ 1=0
(1)
The goal of RL to find an optimal policy x* that maximizes (Tx). It is crucial to recognize that T often varios across different MDPs (environments). Thus, the optimal policy for standard RL must be re-lcarned each time new environment is encountered. Under this circumstance, ICRL proposes to pretrain a FM on a wide variety of pretraining environments, and then deploy it in the unseen test environments without updating parameters in the pretrained model, i.e., zero-shot generalization (Sohn et al. 2018
Mazoure et al.
2022
Zisselman et al. 2023 Kirk et al. 2023
3.2
Supervised Pretraining Mechanism
In this subsection, we introduce the methodology behind ICRL a supervised pretraining mechanism. Con- sider two distributions over environments Terain and Trest for pretraining and test (evaluation), respectively. Each environment, along with its corresponding MDP 7, can be regarded an instance drawn from the environment distributions, where cach environment may exhibit distinct reward functions and transition dynamics. Given an environment т, a context/ prommpt C = {si, 0i, ri, s(]ve(x) refers to a sample from a pretraining context dataset Derai | T), i.e., C ~ Derain | r), which are collected through interactions between a behavior policy and the environment 7 (see e.g., Algorithm
[]. Notably, Dieain (- T) contains the contextual information regarding the environment т. We next consider a query state distribution D and a label policy that maps the query state to the distribution of the action label, i.c., mi : S - Da; (A), where (.)()) denotes the probability distribution and at denotes the rrediction/outpa goal of the FM. Then, the joint distribution over the environment T, contexl. C, query state Sqz and action label as is given by
(T, Sq 01) = Trrain(T) Dirain (Clr) DT, m](ay(sa)-
(2)
ICRL follows supervised pretraining mechanism. More concretely, FM with parameter 0 (denoted by Fo C S + A(ai)) is pretrained to predict the action label a given the context C and query state sy To do so, the current literature (Laskin et al. 2022 Lee et al. 2024 Dong et al. 2024 often considers the following objective function
0t arg min Ep Braa [Fo( C, ())
(3)
4

---

Published in Transactions on Machine Learning Research (04/2025)
where 1(:-)) represents the loss function, for instance, negative log-likelihood (NLL) for discrete-action prob- lems and mean square error (MSE) for continuous-action scenarios. We note that while the current. SOTA ICRL algorithms (AD, DPT, and DIT) adhere to a common objective function equation 3 they differ signif- icantly in constructing the context, query state, and action label. Furthermore, it is important to highlight that selecting appropriate action labels in existing ICRL algorithms can be prohibitively expensive. We discuss this in more detail for cach of the ICRL basclines that follow.
Algorithm Distillation. Instead of learning an optimal policy for a specific environment, AD proposes to learn a RL, algorithm itself across a wide range of environments. This is, the FM in AD is pretrained to imitate the source (or behavior) algorithm over the pretraining environment distribution Traain In general, AD demands a well-trained source algorithm with its complete learning history (from the initial policy the final-traiedd policy). Additionally, AD is restricted to the environments with short episode length, as the context must capture cross-cpisode information of the source algorithm, while standard transformers often have limited context length. In termss of the objective function equation 3 AD lakes the state al the time step as the query state, ag from the source algorithm as the action label, and the episodic history data (So, du, To. 1. -1 to be the context.
Decision Pretrained Transformer. Instead of being stringent on the context and the environment itself, DPT handles the context and query state in more general manner. Specifically, in terms of the objective function equation 3 DPT can consider random collection of transitions as the context C, random query state 8q drawn from Dg, and an optimal action label corresponding to 8q- Despite less requirements on the context, DPT necessitates access to optimal policies for the optimal action labels in all pretraining environments, which may not be available in real-world applications.
Decision Importance Transformer. DIT' proposes to learn ICRL without optimal action labels, follow- ing the same supervised protraining mechanism as DPT. To that end, DIT chooses to consider every possible stale and action in the context as the query state and correspouding action label. It is important to point out that DIT requires (partially) complete episode to form the context, enabling the computation of a return-to-go for each state-action pair. By mapping the return-to-go to a weight that reflects the quality of each state-action pair, DIT can pretrain the FM using the DPT structure, augmented by the weight assigned to each query state and action label. In other words, DIT prioritizes the training on high-return pairs. Notably, DIT still mandates that more than 30% of the context data comes from well-trained policies to ensure the coverage of good action labels.
To summarize, it is worth highlighting that all these SOTA ICRL algorithms necessitate varying degrees of well-trained, or even optimal policies during the pretraining phaso. However, obtaining such policies for real-world applications is often prohibitive, it demands extensive training across a vast number of real- world environments. This challenge becomes even more pronounced in the domains where the transition data exhibits high variance to train an effective policy. Thus, executing ICRL under (c.g., uniform) random policies and random contexts is crucial for enabling the practical application of ICRL in the real-world.
4 State-Action Distillation
In this scction, we propose the State-Action Distillation (SAD), an approach for generating the pretraining dataset for ICRL under random policies and random contexts (see Figure I
As indicated in equation the pretraining data consists of the context, the query state and the action label. We start by introducing the generation of the context under random policies in SAD (refer to Algorithm J. It is important to highlight that the context is collected through interactions with the environment under any given random policy (e.g., uniform random policy), and notably, the context not necessarily originates from a complete episode. These benefits make SAD potentially well-suited for ICRL's real-world applications with random transition data only. Having collected the random context, we are now in the stage of collecting query states and corresponding action labels for the pretraining of FM under the random policy.
5

---

Published in Transactions on Machine Learning Research (04/2025)
Algorithm Collecting Contexts under Random Policy
1: Require: Random policy и, context horizon length T, state space S, environment empty context
C C=0
2: for in [T] do
3: Sample a state ~~ S and an action a T(-ss) Collect sr) by executing action in the environment 7 Add (s,a, to C
end for
7:
Return C
We proceed by recalling that DIT prioritizes the training high-return pairs from the context data collected. To that end, DIT assigns weight w to the loss function during the pretraining phase that is proportional to the return-to-go ie., w(st, at) x Ee 7 'r. Nonetheless, we acknowledge that DIT may not explore to train on good state-action pairs under the random policy for two reasons: (t) DIT solely considers to train on the state-action pairs that are observed in the context, which now contains limited and, more importantly, random transition data. (ti) Even for the state-action pairs in the collected context, the return-to-go does not necessarily prioritize the optimal pair but rather promotes the pair with high immediate reward, as the discount factor applies starting from the current time step with horizon of (T + 1) only, instead of T. This issue becomes even more critical in the problems with sparse rewards.
Under this circumstance, our SAD approach advocates for distilling the outstanding query states and action labels by searching across the entire state and action spaces under the random policy. Before proceeding, we recall the definition of the optimal action in the problems of multi-armed bandit (MAB) and MDP. For any query state 8q, the optimal action for 8y corresponds the action that maximizes the optimal Q-function
MMAB
e arg max E [r((q) a)} aEA QMAn
GMDP A arg max E Σ Er'rols 8qx ao
(4)
aEA
(x
where 8q in the MAB problem refers to the singleton state of bandits, and T² denotes the optimal policy the MDP. However, both AMAB and MDD are intractable to obtain our problem of interest for reasons: (i) computing the expectation in the Q-function demands to sample infinite episodes; (ii) one can only have access to the random policy, instead of n² Therefore, we instead consider (i) stochastic approximation that uses the average the unbiased estimate of the expectation due to the law of large numbers; (ii) maximizing the episodic return under the random policy.
4.1 Trustworthiness of the Random Policy
Subsequently, the crucial question arises: when can we trust the random policy? We claim: The random policy is probabilistically trustworthy for the MAB and MDP problems within a trust horizon. Wo formalize this claim for MAB and MDP this subsection, which relies on the following assumptions.
Assumption 1. The absolute value of the reward r(s, a) is bounded by positive constant B for all state- action pairs in the MAB and MDP, i.e., |r(s,a)) < (s ES A.
Note that Assumption is common in the literature (Azar et al. 2017 Wei et al. 2020 Zhang et al. 2021 In particular, in the case of finite state-action spaces, is always possible to design the reward avoid the possibility of being unbounded.
Assumption 2. Given random policy n, assume that
arg max QMDP((s, = arg max QMDP((S; a), Vsg S. dEA
(5)
aEA
It worth highlighting first that Assumption g may not hold universally for all MDP problems. In fact, it is unlikely that random policy and the optimal policy select the same optimal actions for all problems. Nev- ertheless, acknowledge that Assumption does hold in the MDP problems like the grid world navigation

---

Published in Transactions on Machine Learning Research (04/2025)
where the reward is received only upon achieving the unique goal (Laskin al. 2022 Lee al 2024 Dong
et al. 2024), which fall into the applicable domains of this work. We visualize in Figure grid world navigation problem as an example. In particular, the action derived from maximizing the return under the random policy becomes equivalent to that guided by maximizing the return under the optimal policy, as the maximal return induced by both policies corresponds to navigating to the goal as quickly as possible. We formalize this Proposition 1 and provide a thcoretical proof and empirical validation (refer Appendix A.2).. Although Assumption 2 holds for the grid world navigation problems considered in this work, it is crucial to point out that Assumption 2 may not hold for all grid world navigation problems. To this end, we further discuss the validity of Assumption for the grid world navigation with two sparse rewards (see Appendix A. and demonstrate that the validity of Assumption depends on problem-specific factors such as the rewards and the discount factor 7-
Se
$2
S3
S4
a1
Figure 2: A single-dimensional grid world MDP comprising five states {so, $2, 83, 94), where 5o represents the goal state (golden star). The environment offers two possible actions: ao (go lcft), and aj (go right). Any transitions that would result in (left right) boundary crossing will be confined to the current position. The reward structure is sparse, with value of received solely upon reaching the unique goal state Sp and value of 0 otherwise. We consider infinite time horizon with discounter factor 7.
Having introduced the necessary assumptions, we are in the stage of investigating the trustworthiness of the random policy in both MAB and MDP problems. To proceed, we rely on the following two definitions.
Definition (MAB). Denote by Sq and a' the singleton state and optimal arm in the MAB problem. Consider random policy л. Suppose that each orm has been selected N, times (Na E N+) under The trastworthiness of the random policy in the the MAB problems represents the probability of selecting the optimal arm under the random policy, i.e.,
N.-
1
Σ a')
max (E*)((a)
Σ
(6)
N
N.
The corresponding trust horizon denotes the minimall number of selections over all arms, i.e., N A minagg N. that guarantees such trustworthiness.
Definition | iimlies by the law of large numbers that N + serves trust horizon, ensuring the selection of the optimal arm with probability one. In practice, achieving large trustworthiness requires a sufficiently large trust horizon. The intuition is that a large enough horizon approximates the infinite step problem. The same intuition holds for the MDP problems. We formalize this concept in the next definition, which builds upon the Q-function of the finite-horizon MDP
T (sg: a) EE
Σ (rrss Sy. a0 = Va E A,
(7)
Nep
N
1
QUN (sp: a) a)= Nop
ΣΣ
(1'r(5t,
ar)lo = ao a,) Va E A,
(8)
t=
where *T.N (sg, a) is an unbiased estimate of Qii a) and Nep denotes the number of episodes. Definition 2 (MDP). Denote by a* the optimal action given a query state Sq. Consider a random policy स as well as its Q-function UMDE. The trustworthiness of the random policy in the the MDP problems represents the probability of selecting the optimal action for any query state in the state space, i.c.,
(MDp(sq, max @MD a), Vs, S. aEAl(a")
(9)

---

Published in Transactions on Machine Learning Research (04/2025)
The corresponding trust horizon denotes the minimal horizon length of the MDP such that
2 max MDDp o), Vsg S. oEA((a))
(10)
In both Definitions 1 and 2 the trustworthiness refers to the probability of selecting the optimal action. While the trust horizon in both cases approxinates the infinite-horizon problem, they show subtle differences in the relationship to the trustworthiness. In the MAB problem, which consists of a single episode, the trustworthiness depends on the trust horizon only. However, the MDP, computing QNiPP (s,, a) generally intractable, thus requiring the use of an unbiased estimate (as defined in equation B) which depends on the number of episodes Nep- Consequently, the trustworthiness in the MDP setting is influenced by both the trust. horizon and the number of episodes. Next, formally quantify the relationship between the trustworthiness and the trust horizon for the MAB and MDP problems.
Theorem 1 (MAB). Let Assumption [ hold. The random policy is at least (1 8) -)-trustwyrrhy in Defimilion 7 when the trust. horizon N salisfies
8B2
1 + v 8
N 2
log
(11)
B(r(s,, a²)l
max Blr(sq, aEAi(a')
Proof. See Appendix
A.1
Theorem implies that the trust horizon N quantifies the trustworthiness of the decision making under the random policy æ for MAB problems. Indeed, a larger implies a higher probability (smaller 8) that the average reward of the optimal arm under exceeds that of the next-best arm, therefore, making a more reliable decision. We substantiate this claim by empirical evidence (depicted in (igure[aa). In the practical implementation, we simply execute the MAB under the random policy π until every action in the action space A selected at least. N limes. Subsequently, select the action with the maximal average reward the action label. The detailed procedure for collecting such action labels in MAB is outlined in Algorithm
Algorithm Collecting Query States and Action Labels under Random Policy (MAB)
1: Require: Random policy v, singleton query stale 8q, action space A, enviromment т, emply average reward list Ly, trust horizon Exccute the MAB in under the random policy until every action A selected at least N times for in [A] do Record the average reward associated with the action in the history, and add it to I.p
end for
6: Obtain = A(arg max(L,))
7: Return (8,, a1)
Theorem 2 (MDP). Let Assumptions 1 and 2 hold. Define x min (QxDp((s,aa)) max o,ES
(MDr(s(,))) oEA)(e")
Consider the trust horizon > log, (M(1 (2(BB)) The random policy is at least (1 8) -ttssttoorrhy in Definition g when the number of episodes Nep satisfies
211
1 -
Nap (m (1 77) (2B) N+1)
log
(12)
5
Gi
Proof. See Appendix
A.4
Theorem implics that the trust horizon N and the number of episodes Nep quantify the trustworthiness of the decision making under the random policy for MDP problems. Notice that G2 in equation 12 is monotonically decreasing with respect to the trust horizon N when N > log, (xs1 -2)((2B)
(see

---

Published in Transactions on Machine Learning Research (04/2025)
Lemma in Appendix A.5 Thus, with fixed number of episodes, larger N corresponds to higher probability (smaller 8) that the average reward of the optimal action under и exceeds that of the next-best action, indicating a more reliable decision. This aligns with the intuition that larger N corresponds to a closer approximation of the infinite-horizon MDP, where the random policy π selects the same optimal action as that of the optimal policy (by Assumption 2 However, it is worth noting that Theorem considers therefore
the worst-case guarantee of selecting the optimal actions over all states the state space. still hints to the fact that one should select the action randomly select
could be conservative. That being said, Theorem 2 that yields the largest average reward in practice. In the practical implementation, query state from the state space and execute an episode of N steps for each action in the action space. When the maximal return across all actions is no less than a pre-designed return threshold R, we choose the action that maximizes (M a) across the entire action space A. Otherwise, we randomly sample another query state and repeat the process of evaluating each action in a horizon N, The implementation details are summarized in Algorithm
Algorithm 3 Collecting Query States and Action Labels under Random Policy (MDP)
Require: Random policy म, state space S, action space A, environment T, trust horizon N, return threshold R
Set max_return -
3: while max_return < R do
4: Sample query state 8q S
5: Empty a return list L,
6: for a in [A] do
7: Initialize the state and action as 80 = 8.g ag =
Run an episode of N steps in 7 under the random policy Add the discounted episodic return to L,
10: end for
11: max_ return max(Lr)
12: end while
13: Obtain a = A(arg mL)(L,))
14: Return (sy. a1)
Notably, this work focuses on the grid world navigation problems where the reward is received solely upon reaching the unique goal. In this sparse reward MDP, it is worth highlighting that the action that maximizes the discounted return is essentially the same as that reaches the goal using fewest steps. This is because that the discounted return Errrr, has a term 7' where is the time step. Thus, fewer steps (smaller t) corresponds to larger discounted return (larger 7²). Therefore, for any query state Sy, we prioritize actions that can achieve the goal within N steps, with the actions consuming fewer steps being preferred. If no action can accomplish the goal within N steps, we sample another query state until qualified action is identified. This action is then designated as the action label associated with the query state Sq. Details of this implementation are outlined in Algorithm 5 (see Appendix B
Having introduced the processes for collecting the context, query state, and action label under the random policy, we can now generate the pretraining dataset by integrating the aforementioned procodures (refer Algorithm 4 Given the pretraining dataset, the model pretraining procedure as well the offline and online deployment for SAD are summarized in Algorithm 6 (see Appendix B).
4.2 Performance Guarantees
This subsection establishes the performance guarantees of SAD, offering deeper insights into its efficacy. Corollary 1.
Let hypotheses of Theorems 7 and Ø hold. Denote by I the length of the trajectory. For ony environment 7 and history data H, SAD and the well-specified posterior sampling follow the same trajectory distribution with probability (1 - -8,
Pro (trajectory T, H) Pps (Irajectory r, H), Vtrajectory.
(13)

---

Published in Transactions on Machine Learning Research (04/2025)
Algorithm 4 State-Action Distillation (SAD) under Random Policy
Require: Empty pretraining dataset with size (D), pretraining environment, distribution Terain ran- dom policy π, context horizon length T, state space S, action space A, trust horizon N
for in [[D]] do
3: Sample enviromment
Tusin
4: Collect the context C under the environment and the random policy through Algorithm 1 Collect the query state and the action label or under the environment T and the random policy π through Algorithm 2 for MAB (Algorithm 3 for MDP)
5:
Add (C, 0t) to the pretraining dataset D end for Return D
Proof. See Appendix A.6
Notice that DPT takes the same distribution as that from a well-specified postcrior sampling Loo et al. 2024), which is widely recognized as a provably sample-effccint RL algorithm (Osband et al. 2013). However, crucial highlight that our SAD approach focuses only on the pretraining dataset generation, while maintaining the same pretraining process as that of DPT. Building upon Theorems and implying that SAD probabilistically selects the optimal action (required by DPT), Corollary subsaaniaaes that our SAD approach follows the same trajectory distribution as that of the posterior sampling with probability (1 11-6y 8)).
Having established the corollary above, we next investigate the regret bound of SAD in the finite MDP setting (see details in Appendix A.7) Consider the online cumulative regret of SAD over K episodes in the environment 7 Regret,, Ek=o k=0 V+(m;) - V* where TT(- st) = Fo Ck-1, Then, the regret
bound of SAD is formally stated as follows.
Corollary 2. Let hypotheses of Theorems and 2 hold. Given environment t and a conslant B' > 0, suppose
that sup, (eai(T)/Tmusslll S B'. In the finite MDP with horizon T, it holds with probability (1 - 6) KT that
TTe [Regret, (Fo)) S O(8" (
(14)
Proof. See Appendix A.8
Analogous to Corollary Corollary 2 implics that SAD achieves the same regret bound as that of DPT
(equation 14
with probability (1 -8) 8)*T as SAD and DPT undergo the identical pretraining process.
5
Experiments
In this section, we substantiate the efficacy of our proposed SAD method on five ICRL benchmark problems: Gaussian Bandits, Bernoulli Bandits, Darkroom, Darkroom-Larye Miniworld, which are commonly consid- ered in the ICRL literature (Laskin et al. 2022 Lee al.) 2024 Dong al. 2024). All these problems are challenging to solve in-context, as the test environments differ from the pretraining environments, while the parameters of the FM remain frozen during the test.
5.1 Environmental Setup
Gaussian Bandits. We investigate a five-armed bandit problem in which the state space S consists solely of singleton state Sq- With each arm (action) pulled, the agent receives reward. The goal is to identify the optimal arm that can maximize the cumulative reward. We consider the reward function for each arm following Gaussian distribution with mean Ha and variance o², i.e., R(-Is;,a) a) = NiPa,, 02). Each arm possesses means Pa drawn from a uniform distribution 0 1] and all arms share the same variance a = 0.3. We consider the pretraining and test data to have distinet Gaussian distributions with different means.
Bernoulli Bandits. We adopt the same setup as in Geussion Bandits, with the exception that the reward function does not follow Gaussian distribution. Instead, we model the reward function using Bernoulli
10

---

Published in Transactions on Machine Learning Research (04/2025)
distribution. Specifically, the mean of each Ma is drawn from Beta distribution Beta(1, 1), and the reward function follows a Bernoulli distribution with probability of success Ma- To validate the capability of SAD tackling OOD scenarios, we consider the test data drawn from the Bernoulli distribution while the pretraining data drawn from the Gaussian distribution as in the Gaussian bandits.
Darkroom. Darkroom (Laskin et al. 2022 Zintgraf et al.) 2019) is a vro-dimensional navigation task with discrete state and action spaces. The room consists of 7 7 grids ((S) 49), with an unknown goal randomly placed at any of these grids. The agent can select 5 actions: go up, go down, go left, go right, or stay. The horizon length for Darkroom is 49, meaning the agent must reach the goal within 49 moves. The challenge of this task arises from its sparse reward structure, i.., the agent receives a reward of 1 solely upon reaching the unique goal, and 0 otherwisc. Given 7 x 7 == 49 available goals, we utilize 39 of these goals (~ 80%) for pretraining and reserve the remaining 10 (~ 20%) (unseen during pretraining) for test.
Darkroom-Large. We adopt the same setup as in Darkroom, yet with expanded state space of 10 x 10 and a longer horizon T 100. Consequently, the agent must explore the environment more extensively due to the sparse reward setting, making this task more challenging than Darkroom. We still consider 80% of the 100 available goals for pretraining and the remaining unseen 20% goals for test.
Mintworld. Miniworld is a three- diree-dimeliional pixel-based navigation task. The agent is situated a room with four differently colored boxes, one of which is the target (unknown to the agent). The agent must navigate to the target box using 25 25 image observations and by selecting from 4 available actions: turn left, turn right, move forward, or stay. Similar to Darkroom, the agent receives reward of 1 only upon approaching the unique target box, and otherwise. The high-dimeusional pixel inputs clearly render Miniworld a much more challenging task than Darkroom and Darkroom-Large.
5.2 Numerical Results
We consider four SOTA ICRL algorithms basclines in this work: AD, DIT, DPT, and DPT Since all these methods are FM-based, employ the same transformer architecture (cansal GPT2 model (Radford]
et al. 2019h) and hyperparameters (number of attention layers, number of attention heads, embedding dimensions, etc) across all experiments to ensure fair comparison. The main hyperparameters employed
in this work are summarized in Tables 2 (refer to Appendix C.I.
In all experiments, we employ a uniform random policy to collect context, query states, and action labels, as indicated in Algorithms 3 Then, we pretrain the FM and deploy it with two options: online and offline. In online deployment, the pretrained model collects its own conlext by iteratively interacting with the test environments. In offline deployment, the pretrained model directly uses a sampled context from offline context dataset, which is pre-collected using the random policy to interact with the test environments. The details of online/offlte /oflline deployments are deferred to Algorithm in Appendix
Offilie Evaluation
Online Evaluaaion
Offline Evaluation
Online Evaluation
0.35
Cumu atuue Regitt
160
Cummul otive Rogret
175
150
0.20
0.30
0.25
D77
A0
100
120
80
SAD (ous)
Suboglimaiity
0.3
100
125
75
SAD
BAT
loura)
0.15
0.1
0.10
0.05
0.0
5XB
(om
0-
100
200
300
400
500
100
200
300
400
500
100
200
300
400
500
100
200
300
400
500
Context Horizon
Context Horizon
Context Horizon
Context Horizon
(a)
(b)
(c)
(d)
Figure 3: Offline and online evaluations of ICRL algorithms trained under a uniform random policy: AD, DPT, DIT, DPT* and SAD (ours). Each algorithm contains four independent runs with mean and standard deviation. Gaussian Bandits: (a) and (b), Bernoulli Bandits: (c) and (d).
11

---

Published in Transactions on Machine Learning Research (04/2025)
Bandits. We adhere to offiine and online evaluation metrics for Bandits established in (Lee et al. 2024 In the offline evaluation, we utilize the suboplimality over different context horizon, defined by Pa Ma: where Pa and Me represent the mean rewards over 200 test environments of the optipal arm and the selected arm, respectively. In online evaluation, we employ cumulative regret, defined by Σi - pa.), where a denotes the selected arm at time L. Figures B(a) and (()) remonstrate that our SAD approach significantly outperforms three SOTA basclines under uniform random policy, by achieving much lower suboptimality and cumulative regret. More specifically, let define the performance improvement of SAD over baselines in the offline evaluation by (suboptimallityasellee / performance improvement in the online evaluation is to simply replace the suboptimalily by curvulative regret. Then, SAD surpasses the best baseline, DIT, by achieving 354.0% performance improvements in the offline evaluation and 273.9% in the online evaluation (refer to the first row of Tables 3 and 4 in Appendix C.2 To evaluate the dito---distribution performance of SAD compared to other baselines, we test the models pretrained by all methods on the Gaussion Bandits and assess their performance on the Bernoulli Bondits with no further fine-tuning. Figures (()) and 3jd) illustrate that SAD still achieves lower suboplimabity and cumulative regret than all other baselines, demonstrating more robust performance in handling out- of-distribution scenarios. More specifically, SAD surpasses the best baseline, DIT, by 289.5% in the offline evaluation and 313.9% in the online evaluation (refer to the second row of Tables 3 and 4 in Appendix C.2
suboptimalityssad) suboptimalitysaa
Likewise, the
In the environments of Darkroom and Miniworld, we use return as the evaluation metrie. Moreover, we define the performance improvement of our SAD approach over the baseline methods both offline and online evaluations by (ReturnsaD Returnbaseltee) Returnbaseline.
Offline Evaluation
Online Evaluation
Offliie Evaluation
Online Evaluation
30
60
25
40
50
Retur
20
30
Delum
60
DP
A0
15
30
40
20
20
10
540 (sun)
10
AD
DPT
DIT
SAD
DPT
0
0
5
10
15 20
25
30
35
40
0
AD
DPT
DIT
SAD
DPT
0
5
10
15
20 Episode
25
30
35
40
Episode
(a)
(b)
(c)
(d)
Figure 4: Oflline and online evaluations of ICRL algorithuns trained under a uniform random policy: AD, DPT, DIT, DPT*, and SAD (ours). . Each algorithm contains four independent runs with mean and standard deviation. DarkRoom: (a) and (b). DarkRoom-Large: (c) and (d).
Darkrooms. Figure demonstrates that our SAD approach significantly outperforms three SOTA basc- lines in the Darkroom and Largroom-Large under uniform random policy, by achieving much higher return. In the Darkroom, SAD surpasses the best baseline, DIT, by 149.3% in the offline evaluation and 41.7% in the online evaluation (refer to the third row of Tables 3 and 4 in Appendix C.2).
(.2e) Likewise, in the Darkroom- Large, SAD outperforms the best baseline, DIT, by 266.8% in the offline evaluation and 24.7% in the online evaluation (refer to the fourth row of Tables 3 and 4 in Appendix
C.2
Miniworld. Figure 5 demonstrates that our SAD approach outperforms the three SOTA baselines in the Minsworld under uniform random policy, by achieving a higher return. More specifically, SAD surpasses the best bascline, DIT, by 122.1% the offline evaluation and 21.7% in the online cvaluation (refer to the fifth row of Tables and in Appendix C.2
Tables 3 and also imply that SAD significantly outperforms all baselines average across the five ICRL benchmark cnvironments. In the offline evaluation, SAD exceeds the best bascline DIT by 236.3% on average, the second-best. DPT by 2015.9%, and the third-best AD by 2075.2%. In the online evaluation, SAD surpasses DIT by 135.2%, DPT by 3093.8%, and AD by 3208.8% on average. In addition to comparing SAD with the three SOTA ICRL algorithms under the uniform random policy, we also include the empirical performance of the DPT with optimal action labels (DPT*) as the oracle upper bound of SAD. We observe that SAD demonstrates performance comparable to DPT in tasks such Gaussian Bandils, Bernoulli
12

---

Published in Transactions on Machine Learning Research (04/2025)
Offline Evaluation
Online Evaluation
30
33
25
Return
20
23
20
15
13
10
DT
540 teurs
DPT
DIT
SAO
15
25
35
Episade
(a)
(b)
Figure 5: Offline and online evaluations of TCRL algorithms trained under a uniform random policy: AD, DPT, DIT, DPT*, and SAD (ours). Each algorithm contains four independent runs with mean and standard deviation. Environment: Miniworld.
Bandits, DarkRoom, and DarkRoom--arge, Although Miniworld introduces challenges due to its pixel-based inputs and complex environments, SAD under the random policy still achieves approximately 50% of the performance of DPT Overall, SAD is within 18.6% of the performance of DPT in the offline evaluation across five ICRL tasks, and within 12.3% in the online evaluation (sce details in Tables and
5.3 Ablation Studies
Trust Horizon. Theorems I limplies that the uniform random policy is probabilistically trustworthy within horizon N, with monotonically increasing probability of selecting the optimal action with N in the MAB problem. We substantiate this observation from the theorem by empirical evidence, as presented in Fig- ure a). Furthermore, we conduct empirical investigations into the influence of the trust horizon N on the performance of the MAB problem, which considers the environments of Gaussian Bondits. As expected, a larger N in the MAB problem leads to a better performance (see Figures (a) and 7b)).
MAB
MDP
100
100
(85)
Actuncy 190)
2
60
5
4 3
4
4
40
20
20
400
(a)
(b)
(c)
(d)
Figure 6: (a) and (b): The accuracy (probability) of selecting the optimal action in the MAB and MDP problems with varying trust horizon N. (c) and (d): The minimal number steps required for a query state reach the goal (the golden star) the upper-right corner and in the middle.
We then shift to the MDP problem with the environment of Darkroom, Notice that in our practical algorithm for the sparse-reward MDP like Darkroom (Algorithm 5), we only utilize the state-action pairs that can reach the goal within trust horizon. Therefore, we solely record the probability /accuracy of selecting the optimal actions those states, presented in Figure b).. It shows that the accuracy monotonically decreases with respect to the trust horizon N, which, at first glance, may lead to monotonically decreasing performance well. Nonetheless, we acknowledge that this is not the case. In particular, large trust horizon N in the MDP leads to the low accuracy of selecting the optimal action, whereas a small N may induce the partially short-sighted training of FM, as Algorithm [ solely trains on the states at most N steps from the goal, instead of all states (refer to Figures (c) and d)). Our numerical results Figures ( and ((d) validate this with the fact N performing best, and provide the empirical evidence that either an excessively large small trust horizon can lead to suboptimality.
13

---

Published in Transactions on Machine Learning Research (04/2025)
Offline Evaluation
Online Evaluation
Offline Evalluation
Online Evaluation
Suboptimaltt
0.25
0.20
K-46
Cumulaauue Ragmt
50
40-
60
Vei
Return
25
20
Retum
40-
30
0.15
30
0.10
20
10
0.05
10
km11
100
200
400
500
100
200
300
500
N-1
N-3
N-7
N-
N-15
15
25 30 35
Context Horizon
Context Horizon
Episode
(b)
Figure 7: Offline and online evaluation of SAD with varying trust horizon N for MAB: (a) and (b), and for MDP: (c) and (d), Each N contains four independent runs with mean and standard deviation.
Transformer Hyperparameters. We aim to validate the robustness of our proposed SAD approach with respect to the hyperparameters in the transformer block. Concretely, we focus on the number of attention heads and attention layers, as they have large impacts on the model size of the transformer. As depicted in Figure g our empirical results Darkroom demonstrate a robust performance across varying numbers of attention heads and attention layers.
Offline Evaluation
Online Evaluation
Offline Evaluation
Online Evaluation
30
30
25
40
25
40
30
Retur
30
15
20
15
20
10
h
headt
10
4 Hywr How
hvade
hhooh
15
25
35
5
15 Episode
25
30 35
Number of Heads
Episode
Number of Layers
(b)
Figure 8: Offline and online evaluations of SAD with different transformer hyperparameters: the number of attention heads ((a) and (b)); the number of attention layers ((c) and (d)). Each hyperparameter contains four independent runs with mean and standard deviation.
6 Conclusion
In this work, we propose State-Action Distillation (SAD), a novel approach for generating the pretraining dataset for ICRL, which designed overcome the limitations imposed by the existing ICRL algorithms like AD, DPT, and DIT in terms of relying on well-trained or even optimal policies to collect the pretraining dataset. SAD leverages solely random policies to construct the pretraining data, significantly promoting the practical application of ICRL in real-world scenarios. We also provide the quantitative analysis of the trustworthiness as well as the performance guarantees of SAD. Moreover, our empérical results multiple popular ICRL benchmark environments demonstrate significant. improvements over the existing baselines terms of both performance and robustness. Nevertheless, we note that SAD currently limited to the discrete action space. Extending SAD to handle the continuous action space well as more complex environments presents promising direction for the future research.
14

---

Published in Transactions on Machine Learning Research (04/2025)
References
Mohammad Gheshlaghi Azar, lan Osband, and Rémi Munos. Minimax regret bounds for reinforcement learning. In International conference on machine learning, pp. 263-272. PMLR, 2017.
Tom Brown. Language models are few-shot learners. arXiv preprint or arXiv:2005.14100 2020.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Ar- avind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling, Advances in neural informalion processing systems, 34:15084- 15097, 2021.
Weiqin Chen, Sandipan Mishra, and Santiago Paternain. Domain adaptation for offline reinforcement learn- ing with limited samples. arXiv preprint Xrv:::2408.122 12136, 2024a.
Zhe Chen, Jiannan Wu, Wenhai Wang, Weijie Su, Guo Chen, Sen Xing, Muyan Zhong, Qinglong Zhang, Xizhou Zhu, Lewei Lu, et al. Internvl: Scaling up vision foundation models and aligning for generic linual-linguistic tasks. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 24185 24198, 2024b.
Jacob Devlin. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint wxXi:11800008800 2018.
Juncheng Dong, Moyang Guo, Ethan X Fang, Zhuoran Yang, and Vahid Tarokh. In-context reinforcement learning without optimal action labels. In ICML 2024 Workshop In-Contert Learning. 2024.
Qingxiu Dong, Lei Li, Damai Dai, Ce Zheng, Zhiyong Wu, Baobao Chang, Xu Sun, Jingjing Xu, and Zhifang Sui. A survey on in-conlext learning. arXie preprint wxXiv2301100233, 2022.
Scott Fujimoto, Edoardo Conti, Mohammad Ghavanzadeh, and Joelle Pineau. Benchmarking batch deep reinforcement learning algorithms. arXiv preprint ar Xix::1910.00 708, 2019.
Jake Grigsby, Linxi Fan, and Yuke Zhu. Amago: Scalable in-context reinforcement learning for adaptive agents. arXiv preprint arXiv:2310. 099791 2023.
Wassily Hoeffding. Probability inequalities for sums of bounded random variables. The collected works of Wassily Hoeffding, pp. 409 426, 1994.
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems, 34:1273-11 111127-1286, 2021.
Robert Kirk, Amy Zhang, Edward Grefenstette, and Tim Rocktäschel. A survey of zero-shot generalisation in deep reinforcement learning. Journal of Artificial Intelligence Research, 76:201-2644 2023.
Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learning. 2021.
arXiv preprint ar orXii:2110.000
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforce- ment learning. Advances Neural Information Processing Systems, 33:1179-1191, 2020.
Michael Laskin, Luyu Wang, Junhyuk Oh, Emilio Parisotto, Stephen Spencer, Richie Steigerwald, DJ Strouse, Steven Hansen, Angelos Filos, Ethan Brooks, et al. In-context reinforcement learning with algorithm distillation. arXiv preprint arXiv:2210. 14215, 2022.
Jonathan Lee, Annie Xie, Aldo Pacchiano, Yash Chandak, Chelsea Finn, Ofir Nachum, and Emma Brun- skill. Supervised pretraining can learn in-context reinforcement learning. Advances in Neural Information Processing Systems, 36, 2024.
Kuang-Huci Loe, Ofir Nachum, Mengjiao Sherry Yang, Lisa Lee, Daniel Freeman, Sergio Guadarrama, Ian Fischer, Winnie Xu, Eric Jang, Henryk Michalewski, al. Multi-gama decision transformers. Advances in Neural Information Processing Systems, 15:27921-27988 2022.
15

---

Published in Transactions on Machine Learning Research (04/2025)
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offine reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprinl arXiv:2005006443, 2020.
Yingcong Li, Muhammed Emrullah Iidiz, Dimitris Papailioponlos, and Samet Oymak. Transformers as algorithms: Generalization and stability in in-context learning. In International Conference on Machine Learning, pp. 19565- 19594. PMLR, 2023.
Jonathan Light, Min Cai, Weiqin Chen, Guanzhi Wang, Xiusi Chen, Wei Cheng, Yisong Yue, and Ziniu Hu. Strategist: Learning strategic skills by llms via bi-level tree search. arXiv preprint arXiv:2408.00 10635, 2024a.
Jonathan Light, Sixue Xing, Yuanzhe Liu, Weiqin Chen, Min Cai, Xiusi Chen, Guanzhi Wang, Wei Cheng, Yisong Yue, and Ziniu Hu. Pianist: Learning partially observable world models with Ilms for multi-agent decision making. arXiv preprint orXtv:2411.1990 2024b.
Licong Lin, Yu Bai, and Song Mei. Transformers as decision makers: Provable in-context reinforcement learning via supervised pretraining- arXiv preprint orXim:230.00550 2023.
Jinxin Liu, Ziqi Zhang, Zhenyu Wei, Zifeng Zhuang, Yachen Kang, Sibo Gai, and Donglin Wang. Beyond ood state actions: Supported cross-domain offline reinforcement learning. In Proceedings of the AAAI Conference Artificial Intelligence, volume 38, pp. 13945-13953, 2024.
Chris Lu, Yannick Schroecker, Albert Gu, Emilio Parisotio, Jakob Foerster, Satinder Singh, and Feryal Behbahani. Structured state space models for in-context reinforcement learning. Advances in Neural Information Processing Systems, 36, 2024.
Zixian Ma, Jerry Hong, Mustafa Omer Gul, Mona Gandhi, Irena Gao, and Ranjay Krishna. Crepe: Can vision-language foundation models reason compositionally? In Proceedings of the IEEE/CVP Conference on Computer Vision and Pattern Recognition, pp.
2023.
Bogdan Mazoure, Nya Kostrikov, Ofir Nachum, and Jonathan Tompson. Imprroving zero-shot generalization offline reinforcement learning using generalized similarity functions. Advances in Neural Information Processing Systems, 35:25088 25101, 2022.
Ian Osband, Daniel Russo, and Benjamin Van Roy. (morc) efficient reinforcement learning via postcrior sampling. Advances in Newral Information Processing Systems, 26, 2013.
Alec Radford. Improving language understanding by generative pre-training. OpenAl blog, 2018
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAT 7 blog, 188):9, 2019.
Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Barth- Maron, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, et al. A generalist agent. arXiv preprint XrXi:2205.001199 2022.
Fawaz Sammani, Tanmoy Mukherjee, and Nikos Deligiannis. Nix-ppt: A model for natural language planations in vision and vision-language tasks. In proceedings of the IEEF/CVF conference computer vision and pattern recognition, pp. 8322 8332, 2022.
Viacheslav Sinii, Alexander Nikulin, Vladislav Kurenkov, Ilya Zisman, and Sergey Kolesnikov. In-context reinforcement learning for variable action spaces. arXio preprint orXiw:2812.1392 2023.
Sungryull Sohn, Junkyuk Oh, and Honglak Lee. Hierarchical reinforcement learning for zero-shot general- ization with subtask dependencies. Advances in neural information processing systems, 31, 2018
Richard S Sutton. Reinforcement learning: An introduction. A Bradford Book, 2018
A Vaswani. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
16

---

Published in Transactions on Machine Learning Research (04/2025)
Jiuqi Wang, Ethan Blaser, Hadi Daneshmand, and Shangtong Zhang. Transformers learn temporal difference methods for in-context reinforcement learning, arXiv preprint. arXiv:2405. 13861, 2024.
Chen- Yu Wei, Mehdi Jafarnia Jahromi, Haipeng Luo, Hiteshi Sharma, and Rahul Jain. Model-free reinforce- ment. learning in infinite-horizon average-reward markov decision processes. In International conference machine learning, pp. 10170-10180. PMLR, 2020.
Jerry Wei, Jason Wei, Yi Tay, Dustin Tran, Albert Webson, Yifeng Lu, Xinyun Chen, Hanxiao Liu, Da Huang. Denny Zhou, et al. Larger language models do in-context learning differently. arXiv preprint arXiv:2303.08846 2023.
Noam Wies, Yoav Levine, and Amnon Shashua. The learnability of in-contexi learning. Advances in Newral Information Processing Systems, 36, 2024.
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. arXiv preprint arXim:2111.0200 2021.
Haoran Xu, Li Jiang, Li Jianxiong, and Xianyuan Zhan. A policy-guided imitation approach for offline reinforcement learning. Advances in Neural Information Processing Systems, 35:4085-4009, 2022a.
Mengdi Xu, Yikang Shen, Shun Zhang, Yuchen Lu, Ding Zhao, Joshua Tenenbaum, and Chuang Gan. Prompting decision transformer for few-shot policy generalization. In international conference on machine PMLR, 2022b.
learning, pp.
Sherry Yang, Ofir Nachum, Yilun Du, Jason Wei, Pieter Abbeel, and Dale Schuurmans. Foundation models for decision making: Problems, methods, and opportunities. orXiv preprint arXiv:2303.044 2023.
Lu Yuan, Dongdong Chen, Yi-Ling Chen, Noel Codella, Xiyang Dai, Jianfeng Gao, Houdong Hu, Xuedong Huang, Boxin Li, Chunyuan Li, et al. Florence: A new foundation model for computer vision. arXiv preprint 1rXic:2111.11490 2021.
Zihan Zhang, Simon Du, and Xiangyang Ji. Near optimal reward-free reinforcement learning. In International Conference Machine Learning, pp. 12402-12412. PMLR, 2021.
Luisa Zintgraf, Kyriacos Shiarlis, Maximilian Igl, Sebastian Schulze, Yarin Gal, Katja Hofmann, and Shimon Whiteson. Varibad: very good method for bayes-adaptive deep rl via meta-learning. arXiv preprint
arXiv:1910008346 2019.
Tly Zisman, Vladislav Kurenkov, Alexander Nikulin, Viacheslav Sinii, and Sergey Kolesnikov. Emergence in-context reinforcement learning from noise distillation. orXiv preprint arXiv:2312.122 2023.
Ev Zisselman, Itai Lavie, Daniel Soudry, and Aviv Tamar. Explore to generalize in zero-shot rl. Advances in Neural Information Processing Systems, 36:63174-6319,, 2023.
17

---

Published in Transactions on Machine Learning Research (04/2025)
A Omitted Proofs
A.1 Proof of Theorem 1
Thcorem (MAB). Let Assumption hold. The random policy is at least (1 8) -trustworrty as in
Definition when the trust horizon N salisfies
8B
1 + v1
N2
log Blr(sg, a)]
(15)
6
B[r(sg, a")] aEA)(a*)
max
Proof. For any action E A {a'), consider two positive constants
E1 (E]r((s a"]] Blr(s,, a)]),
(16)
= (1 a) (E(r(( a c")] - Blr(s a)l),
(17)
where E [0, 1].
Consider the following two inequalities
N
No rsg: a²) 2 E[r((y: a")l c1
(18)
1
Σ r(g. a) S E(r(sq a)l + cz.
(19)
N
We note that the two inequalities above are the sufficient but not necessary conditions for N.. En r(sg a²) N. Exir(se a) to hold. Therefore, we simply obtain that
P
(5gg a²)
1
(s)
No
N₂*
No
P
(8gg. a²) 2 B(r(sg. a²)]
61, Na
Σ
r(sg. a) S B[r(sy, a)] + 62
(20)
N
1=1
N₂* a²) B(r(s,, a")] -
= a²)
P
r() E(r(s,, a)l +
(21)
N
where the last equation follows from the fact that each arm is independent other arms. bound the two probabilities in the previous expression using Hoeffding's inequality (Hoeffding
We then lower
1994)
Since Assumption
1 implics that r(,)) E [-B, B), Hoeffding's inequality yields
r r(5g a) Elr(sq a)] S 5e2 2 exp (-TB- 2N 2 (- (3
P
(22)
i.e.
1 N₂ N
Na
P
Σ
r ((g.. a) - Eir(s a)] s c2 2 exp
(23)
2B²
1=
Likewise, we have
1
N,*
(-13 2N
P
r(8 a²) Elr(s a" a²)] ] 2 -1 21 eep
-e-exp
(- (-B))
(24)
N
1=1
18

---

Published in Transactions on Machine Learning Research (04/2025)
i.e.,
N P
1
Σ
( a²) - Elr(sg a d')] 2 -E1
) 21 21lexp exp
(-
(-2
.Na.ci
(25)
212
Since N A minled Nay it then follows from the monotonicity of the exponential function that
1
P
(8gg
a-)
r(sq; a)
N
No
j=1
Na
No
exp
exp
(26)
2B22
2
Nci
1
.Ne
exp
exp
(27)
2B2
2B2
Therefore, it holds for any E
[0, 1] that
P
( a²)
Na Σ a))
Na² (E[r((gq a" } E [[((sq) a)))
)
(1-exp exp
N a(1-a)² a)² (B[TBBBrss a²)] - Bir a)))
1 exp
2B²
2
(28)
Notice that the maximal value of the previous equation with respect to a reaches at 0.5. Then it holds that
N
No r(sg a)
N (Blr(s,, a")] Bir(sg a1))² ))
2
P
r(gg a')2 2
2 1 exp
va A AA\oo {a'}.
N
Na
8B
(29)
Since the previous inequality holds for any a A {a*}, let define
1
arg max aEA)(a'} N₂
r(gg: a).
(30)
Then it also holds that
N.
1 No
P
N Σ
r(sy: a²) a²)
2
max
Σ
r(sg. a)
aEAifa') N₂ 1=1
i=1
N (E(r(s a d1)] Eir(s, a))) 2
- exp
(31)
8B2
N (B(r(sg a a²)] --maxxe (a-} Elr(sn a)))²² 2
1 exp
exp
(32)
8B²
where the last inequality follows from the monotonicity.
To make the previous probability greater than 1 - 5, we require
( 1- exp (- [ (B[r(s,, - max B(r(s,,))]
))
21 -8. 8.
(33)
8B²
19

---

Published in Transactions on Machine Learning Research (04/2025)
Hence, require the trust horizon N - min Na to satisfy
8B2
1+ 1+
N > [R(r(s a*)] - Mo (8 [))
log
(34)
This completes the proof.
A.2 Validity of Assumption 2 in the grid world MDP with a single sparse reward
For the sake of simplicity and without loss of generality, we consider a singleddimensional grid world MDP with the understanding that Assumption holds for the two-dimensional grid world MDP well, which is considered in our numerical experiments. The environmental details of the single-dimensional grid world MDP can be found in Figure 2 where the reward is received only upon reaching the unique goal. To proceed, we rely on the lemma below, and Assumption 2 is then validated by Proposition [ that follows.
Lemma 1. Consider the MDP of single-dimensionae grid world with S = { $1 $2, 83, 84. and A = {ao, a as depicted in Figure Consider the random policy in Assumption 2 It holds that
VMDD(>0)> (8)) 2 VMDP($1) 2 (MD)((>) 2 VMDP(>s) 2 VMDD((1))
(35)
Proof. For simplicity, consider to be uniform random policy, i.e. P(ao 8) P(ai 18) =05, As E S. Recall the Bellman expectation equation
ViDP (s)
Σ a(a (s) (r(s,a) + +BBB (MDP(()))
(36)
By combining the previous Bellman expectation equation with Figure 2 yields
VāDP (so) (r(so, 0n) + VDD r(so, 01) + + MDD (s1)),
VTDP(8 (81) 2 (r(s), 00) + (MPP(()) rrsr + VDD (82)),
VTDP (52) (r(s2, op) + (s) r(), + MDP (53))
(37)
(83) 3 (r(s3, 00) + +>XDDr(52) ((33, + V (84)),
VT (s1) = (r(s1, 0o) + +V/DD r(ss. + VMDP (81)).
Substituting all rewards from Figure 9 (left) into the previous equations yields
VMDP (5p) (1 + 7 MMPP (s)) +
VMDr (1+>V/Dr(* + V (82)),
V (s2) (>VMDDP(11) >VVDD (sz))
(38)
(M3) (TV)Dp((5)) + > (54))
Va (s4) (>VMD)()s)) + <VVDD
Given 7 E (0, 1), the last equation of equation 38
implics that
VMPP =
(MD)(ss) 2 1a (sL)
(39)
Then, the fourth equation of equation 38 yiclds
2
V (82)
(82)
V (MD((53) (MDP(sa)
(40)
7
2
V (MD)
VNDP
(41)
7
ViDD((s))
(42)
20

---

Published in Transactions on Machine Learning Research (04/2025)
Likewise, the third equation of equation 38
can be rewritten
2
(1) MDD((32) - VMDP
(43)
2 V MDP VNDP((S)
(44)
7
2 (M)D
(45)
Combining the previous inequality with the first two equations of equation 38 directly yields
(80) 2 (s1)
(46)
This completes the proof.
Proposition 1. Consider the MDP of single-dimensional grid world with S {so, 81, 82, S4) and A {uo,ai depicled in Figure 2 Consider a uniform random policy n. It holds that
arg max @ump(o, a) arg max Qumplss E
(47)
Proof. By the definition of @uDP and x² we oblain
@UDr(s, a) = (M, V(s,a) E x A.
(48)
Since the objective of the agent in Figure is reach the goal state quickly as possible, and stay still, x (s) is given by
n (*)) do: Vs E S.
(49)
Moreover, we have
arg max (MDp(s,a) arg max @MDP(( a) =s) ag, Vs S.
(50)
aEA
acA
We then turn to consider the learning of Q-function under the random policy т. For simplicity, we consider æ to a uniform random policy, i..,, P(ao 8) = P(a1 {) = 0.5, Vs e S.
We next prove that (s, an) 2 (s.r(s,a1). Vs S. We start with the state sp. Notice that
VMDD(ss)) @MDP a1) = r(so, 01) + + VDDP (81) = - (
QMDD( 0o) = r(so, 0n) + > VMDr
(51)
(52)
Lemma iimplies that Vxx MMD((s)) V (s1) The previous equations then directly indicate
QaDr(so, an) 2 (8)
(53)
For the state $1, we have
QTDP (s1, ag) = r(81, an) + 7 VVDP (so) + (MDD Qun (s1. (11. r($1, 01) + VMDP (s2) 7 Virr
(54)
(55)
Lemma implics that VinD 2 (MD 2 (MDr(s2). Then it holds that
@uDr(s1, ao) 2 ())
(56)
For the state 82, we have
@ump ao) r(52 an) V MPP (81) 7 WWDD Qhp a1) r(52, a1) V VMDP (s3)
(57)
(58)
21

---

Published in Transactions on Machine Learning Research (04/2025)
Employing Lemma 1
directly yields
(xmp((2, an) 2 (M2 a1).
(59)
For the state 83, we have
( ao) = r(53, 0o) + y MDP = (s2) QiP (33, 01) = r(83, 01) + y (sL) = 7 ViDDp(s))
(60)
(61)
Employing Lemma Last but not least, for the state 84, we have
directly yields
(MDp(s3, au) 2 (53 41).
(62)
(MD ao) = r(51, ao) + >V VV (s)) 7s)D
(63)
(upp(ss, 01) = a() + + (SL = >VLD
(64)
Employing Lemma
1
directly yields
(MDP(S4, ao) 2 as)
(65)
Hence we obtain
arg max (MDP(s,a) = do, Vs S. agA
(66)
and then
arg max aEA
QUNP =0 arg max QMmP a), Vs S.
(67)
aEA
This completes the proof.
Empirical validation of Assumption 2 in the grid world MDP. addition to the theoretical proof above, we also provide an empirical validation of Assumption 2 in the grid world MDP in Figure 2
We start by considering the Bellman expectation equation
(s, a) r(s, a) + ylE, (Q ((s E SxA.
(68)
Let us define the temporal-difference (TD) error as follows
a(s,a) |r(s, a) + E.r [QkDp('' a')] (UDp(s, ((s,a) E
(69)
With the understanding that the TD error E(s, a) of (s zero, we iteratively learn and update Q @MDr(s a) by minimizing the TD error. To that end, we consider ? = 0.99 and a convergence threshold EQ 10-1 (can be arbitrarily small). We initialize QuDr(s, a) to be full of zeros as in Figure (middle-left). Subsequently, we consistently update QNDP a) under the uniform policy æ, until convergence as follows
E(s,a) S EQ, (
(70)
Empirically, we obscrve that the Q-table (MDD(s, a) converges after the 1216th iteration; demonstrated in Figure 10 As depicted in Figure (middle-right), the convergent QuDp(s a) implies that the optimal action for any state under the uniform policy a is ag (see the golden stars), ie.,,
arg max @MDP((s a) an, Vs E S. aEA
(71)
22

---

Published in Transactions on Machine Learning Research (04/2025)
Likewise, follow the same process above using the Bellman optimality equation
QMDP a) = r((,,a) + max (MDP(S', ('))) V(s, a) E S A.
(72)
e'EA
The convergent MD a) Figure (right) demonstrates that the optimal action for any state under the optimal policy is ag as well, i.e,
arg max (MDP(s, ag. as E S.
(73)
This validates Assumption
empirically.
Reward
0
QDPP
Q
1.0
0,100
50
1.00
0.00
50
0.00
0.075
22.32
20.75
22.0
50
99.5
0.8
0.050
99.0
$1
0.00
0.00
5x
St
100.00
98.01
0.6
98.5
0.00
0.00
States
0.00
0.00
10.000
States
20.75 *
18.85
20.5
States
99.00
97.03
20.0
97.5
0.00
0.00
0.00
0.00
19.60
16.48
98.01
96.06
-0.050
0.2
0.00
0.00
0.00
0.00
18.85
18.48
19.0
Sg
97.03
96.06
96.5
0.0
-0.100
18.5
a1
Actions
Actions
Actions
Actions
Figure 9: Left: reward table. Middle-Left: initial Q-table. Riglle convergent Q-table under uniform
policy v. Right: convergent optimal Q-table.
1.0
TD
=
10
0.6
2 0.44
0.2
0.0
200
lteration
Figure 10: The learning curve of the maximal TD error of @M a) over the entire statc-action spaces: (a.))x) (r(s,a) + 7 [MmDp(s', (s',"'s',a)) ad)] (x, a)).
A.3 Discussion of Assumption
2 in the grid world MDP with two sparse rewards
consider the same problem setting as in Appendix A.2 except that assign the rewards R₁ and R, to the states so and 84, respectively. In the following three cases, we demonstrate that Assumption holds for some (but not all) navigation problems that have two sparse rewards. The validity of Assumption depends on the problem-specific factors such as the rewards and the discount factor y.
As depicted in Figure 11
(a) Consider R = 1, R. = 0.1, = 0.99. The action that maximizes the Q table for both random and optimal policies is the same, i.e., ag arg max, QM a) arg max, Qirpp(s a) for all states as shown Figure 12
(b) Consider R = =, 1, R, 0.9, = 0.8. In this case, the random policy selects as at the states s0, $1, 82 and a at the states $3, 84 (sce QMPP in Figure However, notice that the optimal policy selects exactly the same optimal actions as that of the random policy, i.c., ag for states so, 81, 82 and aj for states 83, 84 (see QMPP in Figure 13 thus Assumption 2 still holding.
23

---

Published in Transactions on Machine Learning Research (04/2025)
So
51
S2
S3
S4 *
ao
a1
Figure 11: single-dimensional grid world MDP comprising five states (80, 81; 82, 83, 84). The environment offers two possible actions: ag that corresponds moving left, and ai that corresponds Crossing the boundaries is strictly prohibited. Any transitions that would result in boundary crossing will be confined to the current position. The reward structure is sparse, with value of R2 received upon reaching 8p, value of Rs received upon reaching and a value of O otherwise. We consider an infinite time horizon with a discounter factor 7.
moving right,
Q MMPP
MPP
24.0
24.17
22.64
100.00
99.00
23.5
24.17
21.56
23.0
S1
100.00
98.01
99.0
98.5
States
States
22.64
20.92
22.5
99.00
97.03
22.0
97.5
21.56
20.71
98.01
96.16
21.5
*
97.0
20.92
20.71
21.0
97.03
96.16
01 Actions
dg Actions
Figure 12: When R₁ 1, R, 0.1, 0.99. convergent optimal Q-table.
Left: convergent Q-table under uniform policy Right:
(c) Consider R = 1, R. = 0.9. 0.97. The convergent Q-tables in Figure 14 indicate that under the random policy, the optimal action for states 83 and 84 is 0j, whereas under the optimal policy, it is ay- This discrepancy indicates that Assumption does not hold.
24

---

Published in Transactions on Machine Learning Research (04/2025)
Q
MOP
5g
2.72
1.58
$6
5,00
4.00
4.8
4.6
$1
2.72
1.23
$1
5,00
3.20
States
1.58
1.49
4,00
3.60
1.23
2.49
3.20
4.50
1.6
54
1.1
4.50
*
*
31
3c
a2
Actions
Actions
Figure 13: When Rí 1, R, = 0.9, = 0.8. Left: convergent Q-table under uniform policy n. Right: convergent optimal Q-table.
MDP
QaPp
13.60
12.38
1344
S0
33.33
32.33
33.0
13.2
5
11.92
S1
33.33
31.36
13.0
States
12.38
12.20
12.8
32.33
30.42
12.6
31.5
11.92
13.24
12.4
31.36
30.41
31.0
12.2
13.24
Sc
*
30.5
@1 Actions
ag
a1 Actions
Figure 14: When Rí = 1, R. 0.9, 7. = 0.97. Left: convergent Q-table under uniform policy Right:
convergent optimal Q-table.
25

---

Published in Transactions on Machine Learning Research (04/2025)
A.4 Proof of Theorem 2
Theorem 2 (MDP). Let Assumptions 1 and 2 hold. Define min (Q a²) max ThrD((s a)).
e
(fa)
Consider the trust horizon N > log, (M (2BB) 1. The random policy is at least (1 - 6) -trustworrthy as in Definition when the number of episodes Nep satisfies
2(1 - 2 Nep 2 (n (1(1-)//22)) (2 - N+1)²
1 v
log
(74)
8
G₁
Proof.
For any action E A {a'), let us define
N
QMpr a) = E
Σ
(s ar) s 00.00 00
=
+E
>rr(s,, az) |00 -g..
(75)
1=0
N+1
QMDP(S =
rr(s ar) 550 - 55: 00
a²
+E
rr(( az) 80 Sq ao =a"
(76)
@MDP (MpDp()))")) (8g:0²))
6..
Then Vsg E S
Msp a²) a-) QMPP a) QUDD (8g. a²) Ea. (s a) - (a)
(77)
(78)
QMDP (8g. a²) QNIP (80, a) + Ea Ea-
(79)
@MDP(80 a²) QMIP a) +
7' (B(r(s), s0 =..
89.00 ao a] - E
(r(5,, ax)
8g. 00 = a²))
1=
(80)
WW
PMDp(sy: a²) - QMp a) - N+1 N 2B
(81)
1 1-7
@um a²) aEAlfa')
max
QUDD((8. a) N
2B
(82)
1 1-7
N+1 2B
min Fi
aEAlfa*) (NDD(()))
max
(83)
1-7
N+1 2B
(84)
1 1-
(c) 0.
(85)
where (a) follows from Assumption and the properties geometry series, (b) follows from the definition of N, (c) is due to > log, (x(1 (2B())) 1. Therefore, let us consider two positive constants as follows
61 = a (QNiN) a²) QuiDe (M) a)).
(86)
E2 = (1 - a) (MND a²) )
Q (x))))
(87)
where a E [0, 1]. Consider the following two inequalities
Q (8g. a 2 QMDY (Ag: a²) 61,
(88)
Q (sg, a) < QuNN (3q. a) E2.
(89)
26

---

Published in Transactions on Machine Learning Research (04/2025)
We acknowledge that equation 88 and equation 89 are the sufficient but not necessary conditions for @NDr (8g: a²) 2 QMDr a) to hold. Thus,
MIDD a²) (M))) a))
(90)
2 P (QNID (sg: a²) QNDP (sg: a²) - E1. Q (s a) S 1MDDP(sp a) + +
(91)
P (QMDP(Ss a²) 2 QMD a²) a²) -1) P (MMDP(ssg a) & Q (s a) + +22
(92)
where the last equation follows from the fact that each action is independent other actions. Assumption implies that Ee (7'r(8e, ar) = 84, 4n = a, π) E |- B B
X J, Va e A. We there-
fore lower bound the two probabilities in the previous expression using Hoeffding's inequality (Hooefdingg) 1994h
P (OMiDe(sy.a' a") QNDD 1M a²) -a) 2 exp
Nepeř
(93)
1 1-y I-7
2B
P ((M a) QDDP a) S 5 exp
C2p 2
(94)
1-,
2B²
I-7
Then, it holds for any E [0, 1] that
@ (Sg. a²) 2 @MDD (M a)
(95)
Ea
Nop
2 exp
1xepp
exp
(96)
2B²
1-yNi112
2
2B²
1-9811
2
1-7
1-7
Q.i. (Sq. a) )
22 (MMD a*)
M
exp
(97)
1-9811) 2
2B2
I-
(1 - ax)²
(QMD a²)
( a)
1 exp
(98)
N- 2
2Ba
Notice that the maximal value of the previous expression with respect reaches at 0.5. Then it holds that
Q...V (Sys a) 12) [1 I-
Nop (QNid a²)
(QNi a²) (²) (s.. 2 1 1 exp
exexp
Va E A {a').
8B2
(99)
Since the previous incquality holds for any a E A {a*}, let us define
= arg max (sp a). dEAl(a*)
(100)
27

---

Published in Transactions on Machine Learning Research (04/2025)
Then it also holds that
P
QTNY MDP (sg. a²)²
max Q aEAlfa))
(NIDP((g:0)
Nop (ONIDp a²) QMPP (sg: â)
1 exe exp
(101)
(- 8B²
1-7
2
Nop (QDD a²) maxxe CMDP a)
1 1 exe
(102)
8 (1 2
where the last inequality follows from the monotonicity and the fact that ( a²) - 1NDDp(s a) N O, A {.
Let us define
a = arg max Q QMiD a). aEAl(a*)
(103)
Then Vig E S we obtain
Q F.N (sg a²) - max
QMDP (Sn, a)
(104)
aEA)((a)
ONDP (Sq. a²) Q.N QMDP(s5. (s.. a)
(105)
(b)
QUID (80, a²) - E.- (M ū) Ea)
(106)
QMDP(8 (8g. a²) - -MMDP (8,, a) + Ea- Ea-
(107)
2
(s a²) QMDP a) +
Σ
y¹ (B(t 80 = Sq.
=g. ag = al - 1E [r(2,, at)
300 = Sg. a0 = a²))
N=
(108)
@UDP (s @MDD N+1 2B
(109)
1 1-7
(4)
(8p: a) N 2B
QMDP a²)
max (EA\(a-)
(110)
1 7-7
2B
min 5,
(ADD(()) a*)
max Qkar((s..a) aEA((a*)
N+1
(111)
1-7
N+1 2B
(112)
1-7
(
0.
(113)
where (a) follows from the definition of à, (b) follows from the definition of QNDP (c) follows from Assump-
tion 1 and the properties of geometry series, (d) follows from the fact that may not be the maximizer of Q MPDPl (zq: -), (e) follows from the definition of (f) is due to N > log, (M(1 (2B((B)) 1.
Consoquently, equation 102 is monotonically increasing with QT (s a²) max, QXD a). Sub-
stituting the previous inequality into equation 102 and by the monotonicity yields
28 Nop -
N P
(x) max d'd'(a')
ON (M)
> 1- exep
1-7
(114)
8B2
1_g8+1 I-7
28

---

Published in Transactions on Machine Learning Research (04/2025)
To make the previous probability greater than 1 - 8, we require
2
Nep K - 2B
1-7
exp
1-g0t11)2 2 1-y
2 1 - 8.
(115)
8B2
Thus,
8B2 1-1
MI1
1 v 6
Nep log
(116)
8
y N+1 +1 2B T-7
2(1 -1 21
1+ v1 -
(x (1- 2 log
(117)
This completes the proof.
A.5 Technical Lemma
Lemma 2. Given N >> log, (s(1 - (7)((BB) - 1, Gi in equation 12 is monolonically decreasing with respect to the trust horizon N.
Proof. We proceed by defining Y = - - Since N > log, (m(1-7)/(23)) - 1, we can obtain
Y K k(1 (1-17)
E ( (NN+,,111
(118)
2B
Let Z N + and then we can rewrite G1 as
2 2
G =
where Y E (72 1].
(119)
(Y (Y-22)²
The chain rule implies that
0G1
0G 0z
(120)
0N
20
ON
0G
(121)
07
1 90 7
- log (Y - -2²)
log
(122)
Y -
(Y -x-22
1 1- 2
7 3% "oogni 111-1(111 Y1
(123)
アー
(Y 72)3
<0,
(124)
where the last incquality follows from E (0, 1] and Y E (,,,1).
This completes the proof.
A.6 Proof of Corollary I
Let hypotheses of Theorems ! and hold. Denote by the length of the trajectory. For any environment and history data H, SAD and the well well-specified posterior sampling follow the same trajectory
Corollary 1
29

---

Published in Transactions on Machine Learning Research (04/2025)
distribution with probability (1 11 8)
Pry (trajectory | H) = Pps (trajectry T, H), Vtrajectory.
(125)
Proof.
To proceed, we rely on the following assumption.
Assumption 3. Denote by Fø the pretrained FM. sg), assume Parain Sg) = Fola I C, Sg) for all
a A.
Note that Assumption 3 is common assumption in the in-context learning literature (Xie et al. 2021
Lee et al., 2024), assuming that the pretrained FM fits the pretraining distribution exactly provided with sufficient coverage and data, where the SAD fits with a sufficiently large trust horizon N.
With Assumption 3 established, Theorem 1 of (Lee et al. 2024) implies that equation 125 holds when the optimal action is selected at each step. addition, Theorems and @ indicate that the FM trained by SAD selects the optimal action label with probability 5 at each step. Consequently, equation 125 holds for SAD with probability (1 (1-68 8y.
This completes the proof.
A.7 Finite MDP setting from Osband et al. (2013)
Let consider the finite MDP setting as in OOsand et al. 2013), where [[r(s;,a)) E [0. 1]. Denote by S, A, the state space, action space, and time horizon. Consider the uniform random policy for sampling the context C and query state 8q- Denote by (T) and (Train(T) the test and pretraining distribution over the environment t, respectively. Consider the online cumulative regret of SAD over episodes in the environment 7
K Σ - k=
Regrot, (Fo)
- V+(T)),
(126)
where Th(- (t) = Fe(. St)
A.8 Proof of Corollary 2
2 Let hypotheses of Theorems 1 and 2 hold. Given the environment t and a constant > 0, suppose that sup, ( S B'. In the finite MDP setting ahove, holds with probability (1 (1-5)KT that
Corollary
ET (Regret, (Fo)) < /
(127)
Proof. Thcorems and 2 imply that the FM trained by SAD selects the optimal action label with probability 1-6 5 at each step, while the finite MDP setting above comprises K K steps. Therefore, it holds with probability (1 - KT that the trained FM Fo is equivalent to the posterior sampling established in Corollary [ Then, it follows directly from Corollary 6.2 of (Lee et al. 2024] that with probability (1 (1-8) 8) it holds that
E T [Regret, < 0
(128)
where the notation Ô omits the polylogarithmic dependence. Subsequently, by using the bounded likelihood ratio between the test and pretraining distributions yiclds
/
ET RRegret, (Fo))=
(129)
B' / Regget,
(130)
= B'R.... [Regret, (Fo))
(131)
(132)
This completes the proof.
30

---

Published in Transactions on Machine Learning Research (04/2025)
B Implementation Details
We provide below the pseudo-codes that are omitted in the main body of the paper.
Algorithm 5 Collecting Query States and Action Labels under Random Policy (Sparse-Rewarl Reward MDP)
1: Require: Random policy n, stale space S, action space A, environment т, trust horizon N N + 3: while min_sten > N do
Set
4: Sample query state Sq ~ Empty a step list l.
for a in [A] do Initialize the state and action as sp 84, ao = Run episode of N steps in under the random policy n, and terminate the episode early upon recciving a reward
9:
Add consumed steps to Ls (add "N + 1 no reward is received)
10: end for
11: min_steg = min(L.)
12: end while
13: Obtain a A(arg min(Ls))
14: Return (sg, ar)
Algorithm Pretraining and Deployment SAD (Inspired by (Lee et al. 2024))
1: Require: Pretraining dataset D, initial model parameters 8, test environment distribution Trest, number of episodes Ne
// Model pretraining
while converged do Sample (C, Sq. a)) from the pretraining dataset D and predict actions by the model Fa(-IC,, sq) for all
E [[C]]
5: Compute the loss in equation with respect to the action label a and backpropagate to update 0.
6: end while
7: // Offline deployment
Sample unseen environments Trest Sample context C ~ Tuss( 7)
10: Deploy Fg in by selecting a; E arg mag Fo(a C, at time step t
11: // Online deployment
12: Sample unseen environments T ~ Trest and initialize empty context C = 0
13: for in [Ng] do
14: Deploy Fo by sampling a, ~ Fo(: C, 8t) at time step t
15: Add (so, on ro....)
16: end for
C Experimental Details
C.1 Hyperparameters
The main hyperparameters employed in this work summarized in Tables 2
C.2 Additional Results
We provide in this subsection additional experimental results that omitted in the main body of paper.
31

---

Published in Transactions on Machine Learning Research (04/2025)
Table The main hyperparameters of each algorithm
Hyperparameters
AD
DPT
DIT
SAD (ours)
Causal transformer
GPT2
GPT2
GPT2
GPT2
Number of attention heads
3
3
3
3
Number of atlention layers
3
3
3
3
Embedding size
32
32
32
32
DIT Weight
N/A
N/A
500
N/A
Learning rate
0.001
0.001
0.001
0.001
Dropout
0.1
0.1
Table 2:
The main hyperparameters of each environment
Hyperparameters
Gaussian Bandits
Bernoulli Bandits
Darkroom
Darkroom-Large
Miniworld
Action dimension
5
5
Pixel-based
X
X
x
X
/
Trust Horizon
320
320
10
# of epochs
100
100
100
100
200
Context horizon
500
500
49
100
50
Training/Test ratio
0.8/0.2
0.8/0.2
0.8/0.2
0.8/0.2
0.8/0.2
# of enviromments
100000
100000
24010
100000
40000
Figure 15 presents the dataset generation time of our SAD approach, compared with other SOTA baselines (AD, DPT, DIT), where the MAB and MDP problems consider the environments of Gaussian Bandits and Derkroom respectively. SAD (avg) represents the average of consumed time over different trust horizons (consistent with those Figure 6). On average, SAD requires the most significant amount of time in the MDP problem. While the MAB problem, SAD ranks as the second most time-consuming method, with its duration surpassing all except DIT. Notice that the additional computational time of SAD aligns with the prevailing trend of leveraging increased computation to fully harness the advanced reasoning capabilities of FMs
Brown 2020
Figure 16 compares our SAD approach with the SOTA basclines (AD, DPT, DIT) under the same time budget (38 seconds, which is consumed by DIT). Despite with degraded performance and increased vairance compared to before, our SAD method still outperforms all baselines, especially in the offline deployment. This implics that our method has a more robust performance than the baselines when tackling the random contexts as the offline deployment uses pre-collected random contexts.
Table demonstrates that our method with N 7 uses 21,445,147 transition data only and consumes 82 seconds only in the dataset generation. In contrast, obtaining well-trained policies for all pretraining environments in the bascline methods (AD) uses 2,401,000,000 transition data and consumes 3452 seconds in the dataset generation. Furthermore, the baseline methods need to learn and maintain policy and value networks for all pretraining environments (additional costs), which not the case our method since the only thing we do collecting at random. It true that the bascline methods take about 20-40 scconds (see Figure 15p if they just employ the random policy (no well-trained policies), however, accompanied by catastrophic performance depicted in Figure 4
32

---

Published in Transactions on Machine Learning Research (04/2025)
Table 3: Performance improvements of SAD compared to baseline algorithms in the offline evaluation.
Environment
SAD AD
SAD vs DPT
SAD vs DIT
SAD DPT
Gaussian Bandits
647.0%
508.9%
354.0%
-5.2%
Bernoulli Bandits
553.4%
426.8%
289.5%
-18.7%
Darkroom
2162.2%
2069.6%
149.3%
-1.3%
Darkroom-Large
6325.5%
6389.9%
266.8%
-4.99
Miniworld
687.7%
684.2%
122.1%
-2..9%
Average
2075.2%
2015.9%
236.3%
-186%
Table 4: Performance improvements of SAD compared to baseline algorithms in the online evaluation.
Environment.
SAD vs AD
SAD vs DPT
SAD vs DIT SAD vs DPT*
Gaussion Bandils
933.6%
942.5%
273.9%
-0.4%
Bernoulli Bandits
846.8%
830.9%
313.9%
-0.2%
Darkroom
3053.9%
2893.8%
41.7%
-3.4%
Darkroom-Lary
10626.9%
10221.8%
24.7%
-0.1%
Miniworld
582.9%
580.1%
21.7%
-57.3%
Average
3208.8%
3093.8%
135.2%
-2.33%
MAB
MDP
2500
120
2000
100
Time (6)
80
1500
Time
60
1000
40
500
0
(avg)
(avg)
Figure 15: The dataset generation time consumed by SAD averaged over varying trust horizons (as in
Figure 60, compared with AD, DPT, and DIT.
Offline Evaluation
Online Evaluation
25
40
35
30
Return
25
15
20
10
15
10
5
DIT
SAD
0
AD
DIT
SAD
0
10
15 20
25
30 35 40
Episode
Figure 16:
Offline and online evaluations of ICRL algorithms under the same time budget (38 seconds).
Table 5: The amount of transition data and computation time of SAD (N 7) compared to AD with well-trained policies.
Algorithms
Amount of Transition Data
Computation Time (s)
AD SAD (ours)
401, 000, 000
3452
21, 445, 147
82
33


Reviews:
[{"review_id": "jiiUEqjO4j", "reviewer": "Reviewer_SRJi", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: \n   - Interesting and useful problem\n   - Theoretical guarantees (while using somehow irrealistic assumptions)\n\nWeakness: \n   - Presentation/Formalization of the problem difficult to follow. The problem is not well presented, which complicates understanding of what follows. Finally, I had to refer to the formalization from the DIT paper to understand the problem setting. In particular, what is C is not very clear and how we use the methods at test time is not discussed. Also, I felt very confusing to avoid distinguishing the behavior policy (that does not need to be optimal in concurrent approaches), from the action labeller (for which we require good accuracy). The introduction of $\\Delta$ in the formalization, and the notion of action label, are also particularly unclear. \n   - Theoretical assumption 2 is far from being realistic. I appreciated that authors studied its validity for a very simple minigrid problem, but even for a straight minigrid with two states with rewards (the highest the most far from start), I feel it does not hold yet, right ? The random policy will indicate going in the direction of the low but close reward... It would be useful that authors study conditions of validity to complement the sudy. \n  - Theoretical study indicates that with an horizon sufficiently long, the estimation of Q from a random policy is accurate. But it really looks counter-intuitive, or at least additionnal assumptions are missing to me, especially when authors set Nep=1 in their algorithms. Please clarify (I suppose that Nep is recovered by considering a high number of weak labels during training, which explains the accuracy loss observed in fig 5b, but I feel that  authors should give something more to connect everything) \n - Finally, the approach comes down to replace the need for having a well trained labeller by a huge dataset collected at random. Isn't this at least as constraining ?  I feel this is very costly. I cannot find any place where authors give the size of the collected dataset. Authors give only a quick result at the end of appendix (additional results) regarding computing ressources, while it is core for analysis the interest of the approach. The average time given for MDP looks significantly greater as concurrent approaches. And it is given only for the sparse problem, for which authors use expert knowledge to start the random policy close to the goal (from my point of view, this trick looks cheating, as getting this knowledge is close to requiring the optimal). At least, authors should consider to give performance metrics for a given time budget, including everything (dataset collection + training), common for every approach. With a infinite budget, it looks obvious that we can reach better estimations for small problems.\n- Experimental details are missing. For instance, how is set the context for the offline experiment ?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper \"Random Policy Enables In-Context Reinforcement Learning within Trust Horizons\" focuses on the pretraining of fundation decision models, that aim at solving a variety of tasks in various environnements, given a prompt specifying the context as input. Typically, the context prompt contains a set of pre-collected transitions (state-action-reward) from the environement to be solved. The training task is to learn to select the best action given a current state and a context. Concurrent approaches from the litterature usually leverage optimal - or well trained - policies to label actions corresponding to each query state, that the process attempts to imitate from the context. Authors claim that they propose the first ICRL approach based on fully random policies for collecting training data.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Improve the formalization of the problem, the presentation of the concurrent approaches and of the experimental methodology\n- Further study conditions of validity of assumption 2 \n- I would suggest to permute algo3 and 5 (algo3 in appendix and algo5 in main), since algo5 is more easy to connect with the rest of the discussion. algo 3 coming after theorem 2 is very confusing. This makes feel that pratical algorithms are full disconnected from the theoretical analysis. \n- Further analysis to connect theorem 2 with algo 5\n - Algo 3 looks extremly costly,  and can loop infinitely if start is a bit far from the goal. A random policy cannot reach any far state in a reasonnable time. Please discuss (or remove that algo which looks not convincing) \n- The approach requires that we use environements that allow to reset anywhere. Please discuss. \n- Additionnal results obtained with a fixed time budget common for all approaches looks mandatory. \n- Sentence :  \"It shows that the accuracy monotically decreases with N, which ideally should lead to... \". Why ideally ? I would rather say fortunately... And why this is not the case ?  \n- From my point of view theorem 1 is already well known and given in many bandit papers. Also, for the MAB problem, it looks natural that it works, while using UCB policies could look better to collect good data. I would suggest to focus on MDP, which is clearly more interesting (but for which the approach looks rather limitative). \n  \n\n\n\nMinor : \n - in (6), r is given depending on \\pi. why ? It is not the case elsewhere...", "broader_impact_concerns": ".", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "ps6XKD6a8J", "reviewer": "Reviewer_4kRM", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n1. The paper is conceptually sound.  \n2. It also addresses an important aspect of ICRL by relaxing the context and policy requirements.  \n3. The results appear promising in the presented domains.\n\n\nWeaknesses:\n1. The paper seems artificially constrained to 12 pages, resulting in a non-linear and confusing presentation. For example, the related work—as well as some referenced algorithms and results—has been relegated to the appendix. Moreover, the motivating example for the main assumption (Assumption 2) appears only later in the paper.\n2. Assumption 2 is very strong. While I understand it holding in the example provided, and in certain navigation tasks and gridworlds, introducing it six pages into the paper is problematic. Because it is critical to the mathematical framework, it should be discussed in greater detail in the main text, and the scope of the applicable domains should be outlined in the introduction. I am curious what the authors think about this, and if I am misunderstanding something.\n3. The definition of ‘trust horizon’ is confusing due to its different interpretations in the contexts of MABs and MDPs.\n4. The setup for sparse rewards is not clearly defined. It is implicitly assumed to be the case in which a reward is received only upon reaching the goal. However, what about cases where a reward is obtained only twice in a trajectory? Would that not also constitute a sparse reward setting? And which action-labeling algorithm should practitioners use in such cases?\n5. The paper does not discuss Corollaries 1 and 2.\n\n6. The motivating example involving healthcare applications in the introduction appears misplaced. **How is a uniform policy a reasonable choice in healthcare?** Moreover, how does the healthcare example relate to Assumption 2, and can it satisfy that assumption? The authors should either revise the motivating example or ensure consistency by justifying Assumption 2 as well as the use of a uniform random policy in healthcare applications.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "1. The paper presents a method to train foundation models (FMs) for in-context reinforcement learning (ICRL) using random contexts and random query states.  \n2. In addition to constructing the context and query state, the paper presents a method to obtain action labels. It proposes three distinct procedures:\n   1. For the multi-armed bandit (MAB) case, choose the action that yields the maximum average reward after each action has been taken N times.\n   2. For MDPs with sparse rewards, select the action that requires the fewest steps to reach a goal. Restart the search if no action can achieve the goal in fewer than N steps.\n   3. For MDPs with dense rewards, sample trajectories of length N—referred to as the ‘trust horizon’.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Please reorganize the paper so that the critical components are included in the main text rather than being relegated to the appendix (see the comments above for specific examples).\n2. Please introduce and discuss the RL environments in the main paper. Additionally, explain why these environments were chosen (i.e., what makes them particularly interesting and challenging).\n3. Please provide a more detailed discussion of Assumption 2 in the main text, addressing the concerns noted above.\n4. Please include a discussion of Corollaries 1 and 2. How do they compare with results in the existing literature? The discussion should help the reader contextualize whether these results are truly useful.\n5. You justified not using AD by noting that it only allows short episode lengths. However, why can’t a longer context length be used? For the domains under consideration, how long would the context need to be, and why would that be prohibitive?\n6. I am curious why samples beyond the N collected for computing the average reward (in Algorithm 2, step 2) are discarded. Why would extra samples be problematic? Moreover, if exactly N samples per action are needed, why is there a need to sample from a random policy at all? Wouldn’t a simple for-loop suffice?\n7. For Theorem 1, the notation is inconsistent: you sum over the index $i \\in {1, \\ldots, N}$, yet there is no corresponding index in the reward function. Additionally, the reward function is denoted as $r(s,a)$ in one instance and as $r(s,a \\mid \\pi)$ in another. Please ensure that the notation is consistent.\n8. Please provide explicit definitions for ‘trust horizon’ and ‘trustworthiness’ to ensure clarity for the reader.", "broader_impact_concerns": "I do not have any concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "l9QVyH1vCZ", "reviewer": "Reviewer_T9U3", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths:**\n1. This work addresses the important issue of previous ICRL methods requiring an optimal policy by using a random policy to obtain action labels.\n2. It provides theoretical support for performance guarantees and evaluates the proposed method through comprehensive experiments, achieving remarkable improvements over baselines.\n\n**Weaknesses:**\n1. To obtain action labels, SAD needs to explore the entire action space within the N trust horizon. This process appears computationally expensive and may lead to unsafe situations in certain environments.\n2. This work aims to improve ICRL using foundational models. However, the experiments are conducted only with GPT-2, a relatively small-scale model. The performance of SAD with larger models, such as 7B or larger, remains unexplored.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This work focuses on in-context reinforcement learning (ICRL), where previous methods require a well-trained or optimal policy during the pretraining phase. It proposes a novel approach, State-Action Distillation (SAD), which uses a random policy to explore good state-action pairs via a trust horizon. The effectiveness of SAD is demonstrated from both theoretical and empirical perspectives.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Figure 1 depicts the overall framework of the proposed SAD. However, it does not clearly demonstrate how SAD works, especially the concept of the Trust Horizon, which is the core of SAD but is not adequately explained in the current figure. A clearer illustration describing the algorithm flow is encouraged.", "broader_impact_concerns": "No ethical concerns are noted.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The recommendation is based on the reviewers' comments, the action editor's evaluation, and the authors’ response. 

This paper proposes State-Action Distillation (SAD) for in-context reinforcement learning (ICRL), which can generate an effective pretraining dataset guided solely by random policies. SAD shows notable performance improvements in ICRL benchmarks. All reviewers find the studied setting novel and the results provide new insights. The authors’ rebuttal has successfully addressed the major concerns of reviewers. Therefore, I recommend acceptance of this submission. I also expect the authors to include the new results and suggested changes during the rebuttal phase to the final version.
Decision: Accept (poster)