Distilling Reinforcement Learning Algorithms for In-Context Model-Based Planning
Authors: Jaehyeon Son
Soochan Lee
Gunhee Kim
Venue: ICLR 2025 Poster (2025)
Keywords: reinforcement learning
in-context learning
Primary area: reinforcement learning
TL;DR: This work is the first to propose model-based planning for in-context RL that imitates a source RL algorithm, leveraging Transformers to simultaneously learn environment dynamics and improve policy in-context.

Abstract: Recent studies have shown that Transformers can perform in-context reinforcement learning (RL) by imitating existing RL algorithms, enabling sample-efficient adaptation to unseen tasks without parameter updates. However, these models also inherit the suboptimal behaviors of the RL algorithms they imitate. This issue primarily arises due to the gradual update rule employed by those algorithms. Model-based planning offers a promising solution to this limitation by allowing the models to simulate potential outcomes before taking action, providing an additional mechanism to deviate from the suboptimal behavior. Rather than learning a separate dynamics model, we propose Distillation for In-Context Planning (DICP), an in-context model-based RL framework where Transformers simultaneously learn environment dynamics and improve policy in-context. We evaluate DICP across a range of discrete and continuous environments, including Darkroom variants and Meta-World. Our results show that DICP achieves state-of-the-art performance while requiring significantly fewer environment interactions than baselines, which include both model-free counterparts and existing meta-RL methods.

Published as a conference paper at ICLR 2025
DISTILLING REINFORCEMENT LEARNING ALGO- RITHMS FOR IN-CONTEXT MODEL-BASED PLANNING
Jaehyeon Son', Soochan Lee?, Gunhee Kim'
Scoul National University, ²LG AI Research jjh9876 (e07766snu..cm an..ac......... soochan 1eel Lgeeserrrachh gunhee snnne ac. kr
ABSTRACT
Recent studies have shown that Transformers can perform in-context reinforce- ment learning (RL) by imitating existing RL algorithms, enabling comple-efficiee adaptation unseen tasks without parameter updates. However, these models also inherit the suboptimal behaviors of the RL algorithms they imitate. This issue pri- marily arises due to gradual update rule employed by those algorithms. Model- based planning offers a promising solution to this limitation by allowing the mod- els to simulate potential outcomes before taking action, providing an additional mechanism to deviate from the suboptimal behavior. Rather than learning sepa- rate dynamics model, we propose Distillation for In-Context Planning (DICP), an in-context model- -ased RL framework where Transformers simultaneously learn environment dynamics and improve policy in-context. We evaluate DICP across a range of discrete and continuous environments, including Darkroom variants and Meta- World, Our results show that DICP achieves state-of-the-ar performance while requiring significantly fewer environment interactions than baselines, which include both model-free free counterparts and existing meta-RL methods. The code is available https://gghuhaccomm/ayyyooooooos
INTRODUCTION
Since the introduction of Transformers (Vaswani et al.) 2017), their versatility in handling diverse tasks has been widely recognized across various domains (Brown et al., 2020 Dosovitskiy et al.)
2021 Bubeck et al.] 2023). A key aspect of their success is in-context learning (Brown 20200) which enables models to acquire knowledge rapidly without explicit parameter updates through gradient descent. Recently, this capability has been explored in reinforcement learning (RL) (Chen jet al. 2021 Schulman et al. 2017 Lee et al. 2022, Reed et al.) 2022). where acquiring skills in sample-efficient manner is crucial. This line of research naturally extends to meta-RL, RL, which focuses on leveraging prior knowledge to quickly adapt to novel tasks.
In this context, Laskin al. (023) introduce Algorithm Distillation (AD), in-context RL ap- proach where Transformers sequentially model the entire learning histories of a specific RL algo- rithm across various tasks. The goal is for the models to replicate the exploration-exploitatio be- haviors of the source RL algorithm, enabling them to tackle novel tasks purely in-context. Beyond replication, Laskin et al. (2023) show that this approach can enhance sample efficiency by bypass- ing intermediate learning steps or distilling histories from multiple actors. This combination of an off-the-shelf RL algorithm with the in-context learning capability of Transformers has demonstrated
strong potential for improving the adaptability of meta-RI. (Lee et al. 2023a Liu & Abbeel 2023
Huang et 2024
Zisman al. 2024 Sinii et al. 2024).
However, prior in-context RL. approaches have a notable limitation: they tend to replicate the sub- optimal behaviors of the source algorithm. The source RL algorithm updates its policy gradually through gradient descent, sometimes deliberately preventing abrupt changes (Schulman al.) 2015 2017). As result, it may take multiple iterations to fully integrate newly discovered information, leading to repeated suboptimal actions during this process. Without mechanism to deviate from the source algorithm's 's behavior, existing in in-context RL methods (Laskin et al. 2023 Huang et al. 2024 Sinii al. 2024) inherit these inefficiencies.

---

Published as a conference paper at ICLR 2025
BE11
Causal Transformer
Causal Transforner
P--1
To 9ce1
Pc+1
highest R
6c+2
0c42
highest + Reas
(a) Previous approaches
(b) Distillation for In-Context Planning (Ours)
Figure 1:
Comparison between previous approaches (Laskin et al. 2023 Lee et al. 2023a Huang
cct al. 2024) and our proposed Distillation for In-Context Planning (DICP). Unlike the previous approaches that directly predict actions without modeling dynamics, our approach leverages the in- context learned dynamics model for planning.
To address this limitation, we introduce Distillation for In-Context Planning (DICP), an in-context model-based RL framework where the dynamics model is learned in-context alongside policy im- provement. Unlike in-context policy improvement, in-context learning of the dynamics model does not inherit the suboptimal behaviors of the source algorithm, as modeling the environment's dynam- ics is independent of the inefficiency. By planning with this in-context learned dynamics model, our framework provides the agent with an additional mechanism to overcome the suboptimal behav- iors of the source algorithm. Furthermore, simulating potential outcomes before executing actions allows for more deliberate decision-making based on predicted future returns. To the best of our knowledge, ours is the first model-based approach on in-context RL using Transformers imitate the source algorithm.
To demonstrate the effectiveness of our framework, conduct experiments in both discrete and continuous environments, including Darkroom variants (Laskin et al. 2023) and Meta- World bench- mark suite (Yu al. (019)). The results indicate that our approach significantly surpasses both in- context model-free counterparts (e.g., AD (Laskin et al.) 2023). DPT (Lee et al. (023a)). and IDT (Huang et al. 2024) and existing meta-RL. methods (e.g. RI. (Duan et al.. 2016), MAML (Finn et al.) 2017). PEARL (Rakelly et al. 2019), MACAW (Mitchell al. 20211) FOCAL (Li et al., 2021, BOReL (Dorfman et al., 2021), MoSS (Wang et al., 2023b), and IDAQ (Wang et al. 2023a)). Notably, our method achieves ofate-of-the-at results the Meta-Weor World MLI and ML bench- marks, while requiring significantly fewer environmental interactions compared to the baselines.
RELATED WORK
RL as sequence modeling. With the advent of Transformers, which can learn from much larger datasets than what an agent can typically collect online, their application to offline RL (Levine
et al. 2020) has gained prominence (Chen et al.] 202T Janner et al.) 2021: Lee et al. 2022 Reed 2022p. Despite these advances, Laskin et al. (2023) point out key limitation: these ap- proaches struggle to improve their policy in-context through trial and error. The primary reason is that they are designed to imitate the dataset policy, which makes them unsuitable for performing in-context RL on novel tasks. To address this. [askin et al. (2023) propose Algorithm Distillation (AD), an in-context RL approach where a Transformer is trained to distill learning histories from a source RL algorithm across diverse tasks. Notably, AD becomes effective when the context length of Transformers exceeds the episode horizon. Instead of imitating the source algorithm's 's actions, -reetrinnd Transformer (DPT; Lee et al. (2023a)) is designed to predict optimal actions. In-context Decision Transformer (IDT; Huang et al. (2024)) implements hierarchical approach to in-context RI, where the Transformer predicts high-level decisions that guide sequence of actions,
2

---

Published as a conference paper at ICLR 2025
rather than predicting each action individually. In addition, Sinii et al. (2024) propose an in-context RL approach for variable action spaces, while Zisman et al. (2024) synthesize learning histories by gradually denoising policies rather than directly executing source algorithm.
Meta-RL. Deep meta-RL began with online approaches, where RNNs are employed to learn RL
algorithms that generalize across environments (Duan et al. 20166] Wang et al. 2016). In parallel, gradient-based approaches (Finn al. 2017 Nichol al. 2018) aim discover parameter initial- izations that rapidly adapt to new tasks. Recently, offtnne meta-RL has gained attention, leveraging
pre-collected meta-training datasets to address the meta-RL problem (Mitchell al. 2021] [Dorf man et al., 2021 Yuan & Lu 2022 Wang et al. 2023a b). Various methods have been explored for this problem, including gradient-bart based (Mitchell et al.
2021), Bayesian (Dorfman et al., 2021), and
contrastive learning approaches (Yuan & Lu 2022).
Furthermore, recent in-context RL approaches
(Laskin et al. 2023; Lee et al., 2023aa] Huang al. 2024 Sinii et al. 2024), including our own, fall
within the category of offline meta-RL.
Model-based meta-RL. In the prior research on model-based meta-RL, Zintgraf et al. (2020) 2021]: Dorfman et al. (202T) focus on learning to infer belief states about environment dynamics,
while Wang al. (2023b) aim to infer task representations. |Wang et al. (2023a) quantify uncertainty
using an ensemble of meta-training dynamics when facing novel meta-test tasks. Nagabandi et al.)
2219); Pinon al. (2022); Rimon al. (2024) learn to construct dynamics models for planning. Our approach aligns with the latter category but stands out by integrating both the policy and the dynamics model within the same sequence model. This contrasts with prior work that either relies
on separate modules for modeling dynamics (Nagabandi et al. (0199) Rimon et al. (024). or omits
a policy network while depending solely on the dynamics model (Pinon et al. (022). 2). We provide a
more comprehensive overview of related works in App.
PROBLEM FORMULATION
POMDP. We consider partially observable Markov decision process (POMDP): M AS O, r). At each time step t, agent interacts with the environment by selecting an action a, E A based on the current state 8t E S. The environment transitions to the next state 8141 S according to the transition model T($+11 St, az), and the agent receives reward ra E R, determined by the reward model R(r, $t, a.). However, the agent does not directly observe the true state 8,- Instead, it receives partial observation o, O of the state s,, The discount factor E [0, controls the relative importance of immediate versus future rewards. The tuple (T, R) is referred as the dynamics model or world model. The agent's 's history up time step is defined as h = (og, do, To. Ot, rz). The objective of RL is to find policy that maximizes the expected cumulative reward J ER.M
En y'r:] where T is the number of interactions. Throughout the paper, we use the terms environment and task interchangeably refer POMDP.
Meta-RL. We define an algorithm f O + A(A), where H is the space of histories. and ((A) represents the space of probability distributions over A. The objective of meta-RL is discover an algorithm f that maximizes the expected cumulative reward J over distribution of POMDPs p(M). During meta-training, the algorithm is optimized on set of tasks sampled from p(M), while during meta-test,
evaluated on another set of tasks sampled from p(M). For simplicity, we refer to meta-training and meta-test as training and test when the context is clear.
Following previous works (Laskin al.) 2023 Huang et al. 2024), we focus on scenarios where an offline dataset of learning histories 2 [hg (oh, do To ..., o'r, a4, 1/ is available, generated by source algorithm source for a set meta-training tasks {Mi ~ p(M)) The loss
function for AD (Laskin et al. (023) is defined as
(f)
log fo of K(-1);
where 9 represents a sequence model. Similarly, the loss functions for DPT (Lee et al. 2023a) and IDT (Huang et al. 2024) are respectively presented in Eq. 3 in App.
3

---

Published as a conference paper at ICLR 2025
Algorithm Meta- Trainnng Phase
Algorithm Distillation for In-Context Planning (DICP)
1: Required: task distribution p(M), source
algorithm fsouree. sequence model context
Input: Current observation Ot, learning his-
length &
tory 14-1, sequence model
2: D 2 0
2: Set beam size K, sample size L
3: for i 1 do
3: 4 0
4: Sample a training task M, ~ p(M).
4: for = 1..... do
5: Execute fsouree on M; and collect learning
5:
Oe ~ sol
0t, - 04,1411)
history 11 = (o): ag. ror.............. o7, ar. r+).
6: fr. 0111, Ru ~ gol 0z, ât h2-1)
6:
- D-DUK
7: h +- (h4-1, O2, (,,
7: end for
B B+B --BU (8,,,,00441) 9: end for
8: while not converge do
9:
Sample k-step segments from D.
10: while stopping criterion is not met do
10: Update for Eq. 2 with the sampled tra-
11: 3' < U
jectory segments.
12: for all (h, do
11: end while
13: for 1 = 1... do
14:
à, fo( θz ha-1)
Algorithm Meta-Test Phase
15:
is. Özt1, Rs go
Öss Ozs ht
Required: task distribution p(M), sequence
16:
h + (ha-1, 0g, f2)
model 0
17:
B'+ + B' (h δs+1)
2: for = 1... m do
18:
end for
3: Sample test task Mj p(M).
19:
end for
4: h <
20:
f- Top K elements of B' w.r.. Σe Pr. + R
5: fort ...
6: 0g f- DICP (or, 1,
7:
01 + Perform action a, MM.
21: end while
22: (he-1, Ot, âe, ft, ft,,,,,.,000 08-1) + Top element
8: h4 + (hr-1, oz, Tz)
9: end for
of B w.r.t. S-1
Rs
10: end for
23: Return: āz
DISTILLATION FOR IN-CONTEXT PLANNING
In this section, propose Distillation for In-Context Planning (DICP), where the agent utilizes an in-context learned dynamics model to plan actions. To learn the dynamics model in-context, we introduce meta-model go 7 O x A - A(R) A(O) A(R), which shares the same sequence model f with the algorithm fa- yields probability distributions over the spaces of reward rt, next observation 0t+1, and return-to-go R; = =t Σ and action a₃ H represents the time step at which the episode ends. We optimize A throughout the meta-training phase (Alg.) D with following loss:
based on history h2-1- current observation
n
T
= Lim Coy where LDyn = Σ
log (r oc+1, Relor, at f1-1).
(2)
Here, Lim corresponds to one of CAD. CDPT, or CIDT (Eq. balances the imitation loss Lim and the dynamics loss LDyn-
and À is hyperparameter that
After meta-training the sequence model is fixed and evaluated during the meta-test phase (Alg.
key distinction between our method and previous in-context RL approaches (Laskin et al. 2023
Lee et al.) 2023a Huang et al. 2024] lies in how actions are selected (1)).. Previous methods
generate actions directly from the sequence model: a, < ʃo{ ot, 14-1). These approaches may
replicate the inefficiencies arising from the gradual updates of gradicnt-base based RL algorithms, as they
imitate the source algorithm's 's behaviors.
In contrast, method employs additional DICP subroutine (Alg. 3 for action selection. As a planning strategy, we adopt Model Predictive Control (MPC) (Nagabandi et al. 2018 2019 Sa- mundsson al. 2018). At each time step, the agent simulates future outcomes, selects the path that maximizes predicted return, and executes the first action in that path. Specifically, upon receiving an observation ot, the agent samples J. candidate actions {á (
.......) áf), and predicts the consequences
4

---

Published as a conference paper at ICLR 2025
of the actions using the dynamics model, which are the next observations f o/411 and re- wards {fl..... }. The agent also predicts retur-to-go { ....,.... (2f1) } for each candidate action to evaluate each path. The process repeats for future time steps, generating new candidate actions ((41)) } and corresponding consequences, until either the prede-
[1 ..........11 )...)...., ..... (4 fined planning horizon is reached or the episode ends. This process constructs a planning tree, where each path is ranked based on the predicted return. The overall DICP algorithm is presented in Alg.
where superscripts are omitted for brevity.
After building the planning tree, the agent can apply suitable tree search algorithm that fits within computational and memory constraints. In our setup, we employ beam search as it is well-suited to Transformers, where decoding, sorting, and slicing operations for respective beams are paral- lelizable. For each planning step. the top K paths are selected based on the predicted return. As demonstrated in $55 when the environment provides well-structured dense rewards, a simple greedy search combined with the in-context learned dynamics model achieves competitive performance, although beam search remains applicable in such environments. Alg. greedy search repre- sented by skipping the while loop of LIO 21
The DICP algorithm provides a mechanism for the agent to deviate from the source algorithm's suboptimal behaviors and immediately pursue reward-maximizing actions. Since modeling the en- vironment's ' dynamics is independent of the source algorithm's 's gradual update rule, the agent can leverage the in-context learned dynamics model to enhance the decision-making process. When the agent identifies promising transitions that contain information for achieving higher returns, the dy- namics model assigns higher expected returns prospective state-action pairs that align with this information. When such pairs are sampled (L6 & L[4). the dynamics model assigns them higher returns 7 & L15). encouraging further exploration toward these promising areas (120 & 122
Architecture. We evaluate our framework using architectures from AD (Laskin et al. 2023). DPT (Lee al. 2023a), and IDT (Huang et al., 2024), incorporating modifications to include the addi- tional loss term Lyn shown in Eq. 2 These modifications are illustrated in Fig. 4 in App. A We refer to the resulting models as DICP-AD AD, DICP-DPT, and DICP-IDT. For DICP-AD and DICP- DPT, append the action ae to the input sequence and train the model to predict the reward rt. next observation 0++1, and return-to-gi R₂ conditioned on a, Additionally, following Lee et al. (023a)). compress each (o, r) tuple into single token before passing it through the Trans- former. For DICP-IDT, the model predicts the results of actions within its Decisions-to-CC module, which interacts with the environment through low-level actions. However, find that the original architecture struggles to propagate reward information, as the ceviewing-Decission module encodes only observations and actions. To address this, we modify the Making-Decissions Decisions module also en- code rewards. For completeness, we review the original contributions of Lee et al. (2023a) Huang al. (2024) in App. C
Distribution choice. For the distribution class of fe. action spaces and a Gaussian distribution with diagonal covariance for continuous ones. For ge. we apply the same approach, using categorical distribution for discrete state spaces and rewards, and a Gaussian for continuous ones. Thus, (I) and LDyn correspond to either cross-entropy loss or Gaussian negative likeliool,ood,, depending on the environment.
categorical distribution for discrete
EXPERIMENTS
5.1 ENVIRONMENTS
We evaluate our DICP framework across a diverse set of environments. For discrete environments, we use Darkroom, Dark Key-to-Door, and Darkroom- Darkroom-Permuted which are well-established bench-
marks for in-context RL studies (Laskin et al.) 2023 Lee et al. [203a Huang et al. 2024). For continuous ones, we test on the Meta- World benchmark suite (Yu et al. 2019).
Darkroom. Darkroom is 2D discrete environment where the agent should locate goal with limited observations. The state space consists of 9 x 9 grid, and the action space includes five actions: up, down, left, right, and stay. The agent earns reward of I whenever reaching the goal,
5

---

Published as a conference paper at ICLR 2025
with 0 reward otherwise. The agent's observation is restricted to its current position, which makes challenging to infer the goal's '' location. The agent always starts at the center of the grid, and the tasks distinguished by varying goal locations, resulting in 81 distinct tasks. The tasks are divided into disjoint training and test sets, with a 90:10 split. Each episode has a horizon of 20 steps.
Dark Key-to-Door. Dark Key-to-Door is variation of Darkroom, where the agent should find a key before reaching the goal. The agent receives reward of upon finding the key and an additional reward of 1 upon reaching the goal. Unlike Darkroom, the agent only earns a reward from the goal once per episode, with maximum reward of 2 per episode. This environment features 81 x 81 = 6561 distinct tasks, each defined by different key and goal locations. The train-test split ratio is 95:5. The episode horizon is extended to 50 steps.
Darkroom-Permuted is another variant of Darkroom, where the action space is randomly permuted. In this environment, the agent starts in fixed corner of the grid, with the goal located in the opposite corner. There are = 120 unique tasks, each corresponding to different permutation of the action space. Since the agent is tested on novel permutations, must explore the environment to figure out the effects of each action. The train-test split ratio is the same as in Darkroom. The episode horizon 50 steps.
Darkroom-Permnne
Worad. Meta-World Yu et al. 2019) is a robotic manipulation benchmark suite designed for meta-RL and multi-task RL. The suite includes 50 continuous control tasks and offers three distinct meta-RL benchmarks: ML1, ML10, and ML45. We focus MLI, which provides 50 pre-defined seeds for both training and test for each task. Each seed represents a different initialization of the object, goal, and the agent. Although the agent receives dense rewards tailored to each specific task, its performance is evaluated based on the average success rate across test seeds, which differs from the dense reward structure. The episode horizon is 500 steps. We also report results on MLI0 in App. B where the goal is to learn from 10 training tasks and generalize to 5 test tasks.
5.2 BASELINES
In-context RL methods. We compare our approach with recent in-context model-free methods, including AD (Laskin et al., 2023). DPT (Lee et al. (023a). and IDT (Huang al. (024)) We follow the evaluation protocols of previous works (Laskin et al. 2023: Huang et al. 2202)). For the source algorithms, data using 100 actors, accumulating a total of 100K to 1M environment steps, depending on the
PPO (Schulman al. 2017) generate learning histories. We collect
environment. The context length is set to four and ten episode horizons for discrete and continuous environments, respectively. We evaluate DPT only in Darkroom variants because they assume the availability of optimal actions, which are not applicable in general environments.
Other meta-RL. RL methods. For model-free meta-RI. methods, we include RI.2 (Duan et al. 2016 MAMI. (Finn et al.] (017)) PEARI. (Rakelly et al., 2019). MACAW (Mitchell et al. 202 1). and FOCAL Li al. 1021). l). For the model-based methods, compare against MuZero (Schritwiissrr ct al. 2019), BOReL (Dorfman et al. 2021). MoSS (Wang et al. 2023b), and IDAQ (Wang et al. 2023a). To ensure fair comparison,
report IDAQ results using SAC the source algorithm, rather than expert policies. The original BOReL. leverages an oracle reward function, so we report results without leading to failure on most tasks as discussed in Dorfman et al. (2021): Wang al.) (203a)). MuZero is not originally designed for meta-RL, but we include it due to its strong perfor- meta-training tasks followed by fine-tuning on meta-test tasks (Anand) et al. 2022). Due to the challenge in reproducing meta-RL approaches, we present the results re- ported in the original papers, following prior works (Anand et al. 2022 Wang et al. 2023b).
mance after pre-training
5.3 RESULTS
We begin by comparing our approach with model-free counterparts the Darkroom, Dark Key-to- Door, and Darkroom-Permutt environments. The first row of Fig. 2 displays the learning curves for these environments. Under unified configurations and implementations, the best algorithm among AD, DPT, and IDT varies according to environments. Nonetheless, our model-based ap- proach consistently outperforms the model-free counterparts. For planning, we use beam search

---

Published as a conference paper at ICLR 2025
Darkroom
Dark Key-to-Door
Darkroom --moommeeee Permuted
Episode rewards
15
1.5
1.0
0.5
0
200
400
600
800
1000
0
500
1000
1500
2000
2500
500
1000
1500
2000
2500
Reach
Push
Pick-Place
1.0
0.8
0.9
0.6
0.6
0.8
0.4
0.4
0.7
0.2
0.2
0.6
0.0
2.5
5.0
7.5
10.0
10
20
30
50
100
150
200
1c3
lc3
1c3
Sweep
Beg-Inneer-Siie
Mean 50 Tasks
1.0
1.0
0.8
0.8
0.8
Success
0.6
0.6
0.2
0.4
0.4
0.0
0.2
0.2
50
100
150
20
Env steps
le3
Env steps
1e3
Env steps
AD
DPT IDT
DICP-AD (Ours)
DICP-DPT (Ous)
DICP-IDT (Ours)
Figure 2: Learning curves of in-context RL approaches during the meta-test phase discrete (1st row) and continuous (2nd and 3rd rows) environments. Our methods outperform model-free counter- parts in both sample efficiency and overall performance. Results averaged over and train-test splits for discrete and continuous benchmarks, respectively. We also report the mean success rate across all 50 tasks in Meta- World ML1. The final performance results for all ML1 benchmarks are presented in Table and Table g, Shaded areas represent 95% confidence intervals.
with a beam size 10, considering all 5 actions at each planning step. The ablation study over different beam sizes are discussed in $6.1
Furthermore, we evaluate our approach on Meta-World ML1 benchmarks against previous meta-Re RL methods, including both model-free and model-based methods. Table presents the mean success rates over the meta-test phase for Reach-v2, Push-v2, Pick-Place-v2, Sweep-v2, and Peg-Insert- Side-v2, following previous works (Yu al. 2019 Anand al. 2022 Wang et al.) 2023b a). Our method achieves state-of-the-ar performance on this benchmarks, surpassing all previous meta-RI. methods. A key highlight is that our approach reaches this performance with maximum of only 200K environment steps, whereas most other methods require at least 10M steps. To ensure a fair comparison with Wang et al. (2023a), whose meta-test is conducted with only 5K environment steps, we also report the results at 5K steps in Table Our approach consistently outperforms the baselines from Wang et al. (2223) at this step count as well. While in-context model-free baselines (Laskin) et al. 2023 Huang et al. 2024) are also relatively sample-efficient, their performance remains sub- optimal compared to our method. The learning curves are displayed in second and third rows of Fig. 2 For planning, employ greedy search with sample size of 10 each step. We also report additional results on all remaining MLI benchmarks Table and on ML10 in Table I0] respectively.

---

Published as a conference paper at ICLR 2025
Table Meta-test success rates on Meta- ML1. The success rates are reported as percentages, averaged over 3 train-test splits with 50 test seeds. The results marked with are taken from Yu et al. (2019), while the ones with are from Wang et al. (2023a). The results for MoSS and MuZero are from Wang et al. (2023b) and Anand et al. (2022) respectively. Our approach achieve state-of-the- art performance with much fewer environment steps. Results at 5K steps are presented in Table We omit '-22' from the task names.
Method
Reach
Push
Pick-Place
Sweep
Peg- Peg-Insert-Sice -ide
Max Steps
RL
100
96
98
300M
MAML
100
94
80
300M
PEARL
68
44
28
300M
MACAW
4
0
5K
FOCAL
38
10
5K
MuZero
100
100
100
10M
MoSS
86
100
100
40M
BoREL
0
0
5K
IDAQ
59
30
5K
AD
86
88
96
67
73
200K
IDT
100
94
95
96
76
200K
DICP-AD (Ours)
84
87
97
71
83
200K
DICP-IDT (Ours)
100
100
100
100
87
200K
ABLATION STUDY
We perform an ablation study to evaluate key design choices in our approach. We examine the effect of model-based planning at different scales ($6.1), the effect of context length the accuracy of the in-context learned dynamics model ($6.2). and the effect of different source algorithms ($6.3)
6.1 EFFECT OF MODEL-BASED PLANNING AT DIFFERENT SCALES
We evaluate our approach using search algorithms at varying scales, as shown in the first row of Experiments are conducted in two environments: Darkroom (discrete) and Reach-v2 (con- tinuous). We use DICP-AD for Darkroom and DICP-IDT for Reach-v2, the best-perforning models for each environment. The results demonstrate that incorporating model-based planning signifi- cantly improves performance compared baselines that train the dynamics model but do not use it for planning. While increasing the beam size sample size improves performance, the gains plateau after a size of 10. When the dynamics model is not used for planning, the performance is comparable that of model-free baselines.
Fig. 3
6.2 EFFECT OF CONTEXT LENGTHS ON ACCURACY THE DYNAMICS MODEL
We investigate the effect of context length on the accuracy of the in-context learned dynamics model. Our hypothesis is that extending the context length enables the model to capture more information about the environment, improving prediction accuracy. To verify this, evaluate the test loss by forwarding the learning histories of the source algorithm on meta-test tasks into the meta-trained DICP-IDT on Reach-v2. The test loss is measured as the mean squared error between the predicted and actual rewards and observations. The models are trained with a context length of 1000, and the loss is calculated by averaging across every 10 positions within the context. The results in the second row of Fig. support our hypothesis: longer context lengths yield more accurate predictions for both observations and rewards. Given that the effectiveness of model-based planning heavily depends on the dynamics model's bias (Janner et al.) 2019 Hiraoka et al. 2021). our framework benefits from longer context lengths. Furthermore, in-context model-free RL methods have also been shown to improve with extended context lengths (Laskin al. 2023), our combined framework gains substantial advantages from using longer context lengths.
8

---

Published as a conference paper at ICLR 2025
Darkroom
Reach
Episode rewards
15
0.8
Success rate
10
0.6
K-15
L-15
K-10
0.4
L-10
0.2
No phaningg
planning
0.0
200
400
600
800
1000
1000
2000
3000
4000
5000
Env steps
Env steps
Test Loss for Reward
Test Loss for Observation
1e-3
1o-4
1.8
MSE loss
2.0
MSB loss
1.7
1.5
1.6
1.5
1.0
200
800
200
400
1000
Context length
Context length
Figure 3: First row: The effect of model-based planning at different scales. We present learning curves with varying beam sizes K and sample sizes L. For the case labeled "No planning." the dynamics model is not utilized for planning, while the meta-model is still trained. The dashed vertical line marks the time step when planning begins, coinciding with the point where the context is fully filled. Second row: The effect of context lengths on the accuracy of the in-context learned dynamics model. Results are averaged over 3 train-test splits.
Table 2: Meta-test success rates of in-context RL approaches using SAC as the source algorithm. Our approach outperforms original model-free IDT, The results are averaged over 3 train-test splits.
Method
Reach
Push
Pick-Place
Sweep
Peg-Inseen Peg-Insert-Side
IDT
85
65
17
4
1
DICP-IDT (Ours)
85
69
20
3
3
6.3 EFFECT OF DIFFERENT SOURCE ALGORITHMS
Our framework is agnostic to the choice of the source algorithm, are previous model-free ap- proaches (Laskin et al. 2023 Lee et al. 2023a Huang et al., 2024). To empirically verify this, train IDT and DICP-IDT with SAC (Haarnoja al. (018) the source algorithm. As shown in Table 2 our method continues to outperform model-frea free counterpart. These results demonstrate that our approach robust to the choice of source algorithm, outperforming model-free baselines regardless of which source algorithm is used. It is worth noting that since our SAC configuration is specifically tuned for these tasks, the overall performance may be underestimated. The hyperpa- rameters for SAC are detailed in App. A
7 CONCLUSION
We introduced an in-context model-based RL framework that leverages Transformers to not only learn dynamics models but also improve policies in-context. By incorporating model-hased plan- ning, our approach effectively addresses the limitations of previous in-context model-free RI. meth- ods, which often replicate the suboptimal behavior of the source algorithms. Our framework demon- strated superior performance across various discrete and continuous environments, establishing in- context RL as one of the most effective meta-RL approaches.

---

Published as a conference paper at ICLR 2025
One limitation of our approach is the additional computational cost incurred during action selection compared to model-free methods. However, this trade-off aligns with the broader trend of leverag- ing increased inference-tine computation to maximize the reasoning capabilities of Transformers (Brown et al. 2020. Wei et al. 2022). Future work could explore adaptive planning strategies that dynamically adjust the planning scale based context. Another promising direction is in- corporating expert demonstrations to accelerate learning. Additionally, investigating offline dataset construction strategies enable our method adapt to changing dynamics would be a meaningful research direction (Lee et al. 2023b 2024 Son et al. (0244.. Finally, exploring more advanced or efficient sequence models could further enhance in-context RL.
REPRODUCIBILITY STATEMENT
We committed ensuring the full reproducibility of our research. Our open-sourced code ables easy replication of all experiments, and all datasets used in our experiments can be generated using our code. Detailed configurations are provided in both the code and App. A We hope this resource serves as valuable reference for the research community, particularly for newcomers, by providing unified implementation of in-context RL methods.
ACKNOWLEDGMENTS
We thank Jaekyeom Kim, Yeda Song, Jinuk Kim, Fatemeh Pesaran zadeh, Wonkwang Lee, Dong- wook Lee, and the anonymous reviewers for their thoughtful feedback. This work was partly sup- ported by Samsung Advanced Institute of Technology, Institute of Information & Communications Technology Planning & Evaluation (ITP) grant funded by the Korea government (MSIT) (No. RS- 2022-11220156, Fundamental research on continual meta-learning for quality enhancement of ca- sual videos and their 3D) metaverse transformation; No. RS-2019-11191082 SW StarLab; No. RS- 2021-11211333, Artificial Intelligence Graduate School Program (Seoul National University)), and Korea Radio Promotion Association (Development of Intelligent Docent Service for Information- Disadvantaged Groups). Gunhee Kim is the corresponding author.
REFERENCES
Ankesh Anand, Jacob C. Walker, Yazhe Li, Eszter Vértes, Julian Schrittwieser, Sherjil Ozair, Theo- phane Weber, and Jessica B. Hamrick. Procedural generalization by planning with sulf-supervised world models. In International Conference on Learning Representations, 2022.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In Neural Informa- tion Processing Systems, 2020.
Sébastien Bubeck, Varun Chandrasekaran, Ronen Eldan, John A. Gehrke, Eric Horvitz, Ece Ka- mar, Peter Lee, Yin Tat Lee, Yuan-Fang Li, Scott M. Lundberg, Harsha Nori, Hamid Palangi, Marco Tulio Ribeiro, and Yi Zhang. Sparks of artificial general intelligence: Early experiments with gpt-4. ArXiv preprint, 2023.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. In Neural Information Processing Systems, 2021.
Ignasi Clavera, Jonas Rothfuss, John Schulman, Yasuhiro Fujita, Tamim Asfour, and P. Abbeel. Madel-based reinforcement learning via meta-policy optimization. In Conference on Robot Learning, 2018
Ron Dorfman, Idan Shenfeld, and Aviv Tamar. Offline meta reinforcement learning identifiability challenges and effective data collection strategies. In Neural Information Processing Systems, 2021.
10

---

Published as a conference paper at ICLR 2025
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenbor,, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021.
Yan Duan, John Schulman, Xi Chen, Peter L. Bartlett, Ilya Sutskever, and P. Abbeel. RL Fast Reinforcement Learning via Slow Reinforcement Learning. ArXiv preprint, 2016.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In International Conference on Machine Learning, 2017.
David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In Neural Information Processing Systems, 201 8.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International Confer- ence on Machine Learning, 2018.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy P. Lillicrap. Mastering diverse domains through world models. ArXiv preprint, 2023.
Takuya Hiraoka, Takahisa Imagawa, Voot Tangkaratt, Takayuki Osa, Takashi Onishi, and Yoshi- masa Tsuruoka. Meta-model-based meta-policy optimization. In Asian Conference on Machine Learning. 2021.
Sili Huang. Jifeng Hu, Hechang Chen, Lichao Sun, and Bo Yang. In-context decision transformer: Reinforcement learning via hierarchical chain-of-thought. In International Conference on Ma- chine Learning, 2024.
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model- based policy optimization. In Neural Information Processing Systems, 2019.
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. In Neural Information Processing Systems, 2021.
Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy H. Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, Afroz Mohiuddin, Ryan Sepassi, George Tucker, and Henryk Michalewski. Model based reinforcement learning for atari. In International Conference on Learning Representations, 2020.
Michael Laskin, Luyu Wang, Junhyuk Oh, Emilio Parisotto, Stephen Spencer, Richie Steigerwald, DJ Strouse, Steven Stenberg Hansen, Angelos Filos, Ethan A. Brooks, Maxime Gazeau, Hi- manshu Sahni, Satinder Singh, and Volodymyr Mnih. In-context reinforcement learning with algorithm distillation. In International Conference on Learning Representattonss 2023.
Jonathan Lee, Annie Xie, Aldo Pacchiano, Yash Chandak, Chelsea Finn, Ofir Nachum, and Emma Brunskill. Supervised pretraining can learn in-context reinforcement learning. In Neural Infor- mation Processing Systems, 2023a.
Kuang-Huei Lee, Ofir Nachum, Mengjiao Yang, L. Y. Lee, Daniel Freeman, Winnie Xu, Sergio Guadarrama, Ian S. Fischer, Eric Jang. Henryk Michalewski, and Igor Mordatch. Multi-game decision transformers. In Neural Information Processing Systems, 2022.
Soochan Lee, Jachyeon Son, and Gunhee Kim. Recasting continual learning as sequence modeling. In Neural information Processing Systems, 2023b.
Soochan Lee, Hyeonseong Jeon, Jachyeon Son, and Gunhee Kim. Learning to continually learn with the bayesian principle. In International Conference on Machine Learning, 2024.
Sergey Levine, Aviral Kumar, G. Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives open problems. ArXiv preprint, 2020.
Lanqing Li, Rui Yang, and Dijun Luo. FOCAL: efficient fully reiffrreinoreeementt learning via distance metric learning and behavior regularization. In International Conference on Learning Representations, 2021.
11

---

Published as a conference paper at ICLR 2025
Zichuan Lin, Garrett Thomas, Guangwen Yang, and Tengyu Ma. Model-based adversarial meta- reinforcement learning. In Neural Information Processing Systems, 2020.
Hao Liu and P. Abbeel. Emergent agentic transformer from chain of hindsight experience. In International Conference Machine Learning, 2023.
Eric Mitchell, Rafael Rafailov, Xue Bin Peng. Sergey Levine, and Chelsea Finn. Offline meta- reinforcement learning with advantage weighting. In International Conference on Machine Learning. 2021.
Anusha Nagabandi, Gregory Kahn, Ronald S. Fearing, and Sergey Levine. Neural network dynamics for model-based deep reinforcement learning with model-free fine-tuning. IEEE International Conference on Robotics and Automation, 2018.
Anusha Nagabandi, Ignasi Clavera, Simin Liu, Ronald S. Fearing, Picter Abbeel, Sergey Levine, and Chelsea Finn. Learning adapt in dynamic, real-world environments through meta- reinforcement learning. In International Conference on Learning Representations, 2019.
Alex Nichol, Joshua Achiam, and John Schulman. On first-order meta-learning algorithms. ArXiv preprint, 2018.
Brieuc Pinon, reinforcement learning: Transformers and tree search. European Symposium on Artificial Neural Networks, Computational Intelligence and Machine Learning, 2022.
Delvenne, and Raphaël M. Jungers. model-bases approach to meta-
Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dor- mann. Stable-baselines:: Reliable reinforcement learning implementations. Journal of Machine Learning Research, 2021.
Kate Rakelly, Aurick Zhou, Chelsea Finn, Sergey Levine, and Deirdre Quillen. Efficient off- meta-reinforrcement learning via probabilistic context variables. In International Conference on Machine Learning. 2019.
Scott Reed, Konrad ZoIna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Barth-Maron, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, Tom Eccles, Jake Bruce, Ali Razavi, Ashley D. Edwards, Nicolas Manfred Otto Heess, Yutian Chen, Raia Hadsell, Oriol Vinyals, Mahyar Bordbar, and Nando de Freitas. A generalist agent. Trans- actions on Machine Learning Research, 2022.
Zohar Rimon, Aviv Tamar, and Gilad Adler. Meta reinforcement learning with finite training tasks density estimation approach. In Neural Information Processing Systems, 2022.
Zohar Rimon, Tom Jurgenson, Orr Krupnik, Gilad Adler, and Aviv Tamar. Mamba: effective world model approach for meta-reiiforrcemnnt learning. International Conference on Learning Representations, 2024.
Steindór Sæmundsson, Katja Hofmann, and Mare Peter Deisenroth. Meta reinforcement learning with latent variable gaussian processes. In Conference on Uncertainty in Artificial, 2018.
Julian Schrittwieser, loannis Antonoglou, Thomas Hubert, Karen Simonyan, L. Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, Timothy P. Lillicrap, and David Silver. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 2019.
Julian Schrittwieser, Thomas Hubert, Amol Mandhane, Mohammadamin Barekatain, Ioannis Antonoglou, and David Silver. Online and offline reinforcement learning by planning with learned model. In Neural Information Processing Systems, 2021.
John Schulman, Sergey Levine, Pieter Abbeel, Michael I. Jordan, and Philipp Moritz. Trust region policy optimization. In International Conference on Machine Learning, 2015.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. ArXiv preprint, 2017.
12

---

Published as a conference paper at ICLR 2025
Viacheslav Sinii, Alexander Nikulin, Vladislav Kurenkov, Ilya Zisman, and Sergey Kolesnikov. In-context reinforcement learning for variable action spaces. In International Conference on Machine Learning. 2024.
Jachyeon Son, Soochan Lee, and Gunhee Kim. When meta-learning meets online and continual learning: A survey. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2024.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. In Neural Information Processing Systems, 2017.
Jane X. Wang, Zeb Kurth-Nelson, Hubert Soyer, Joel Z. Leibo, Dhruva Tirumala, Rémi Munos, Charles Blundell, Dharshan Kumaran, and Matthew M. Botvinick. Learning to reinforcement learn. ArXiv preprint, 2016.
Jianhao Wang, Jin Zhang, Haozhe Jiang, Junyu Zhang, Liwei Wang, and Chongjic Zhang. Offline meta reinforcement learning with in-distribution online adaptation. In International Conference on Machine Learning. 2023a.
Mingyang Wang. Zhenshan Bing, Xiangtong Yao, Shuai Wang, Hang Su, Chenguang Yang, Kai Huang, and Alois Knoll. Meta-reinforremmnnt learning based on self-supervised task representa- tion learning. In AAAI Conference on Artificial Intelligence, 2023b.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Ed Huai hsin Chi, F. Xia, Quoe Le, and Denny Zhou. Chain of thought prompting elicits reasoning in large language models. In Neural Information Processing Systems, 2022.
Weirui Ye, Shaohuai Liu, Thanard Kurutach, Pieter Abbeel, and Yang Gao. Mastering atari games with limited data, In Neural Information Processing Systems, 2021.
Tianhe Yu, Deirdre Quillen, Zhampeng He, Ryan C. Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-word A benchmark and evaluation for multi-tast and meta reinforcement learning. Conference on Robot Learning. 2019.
Haoqi Yuan and Zongqing Lu. Robust task representations for offline meta-reinforcement learning via contrastive learning. In International Conference on Machine Learning, 2022.
Peiyuan Zhang, Guangtao Zeng. Tianduo Wang, and Wei Lu. Tinyllama: An open-source small language model. ArXiv preprint, 2024.
Luisa M. Zintgraf, Kyriacos Shiarlis, Maximilian Igl, Sebastian Schulze, Yarin Gal, Katja Hofmann, and Shimon Whiteson. Varibad: A very good method for bayes-adaptive deep RL via meta- learning. In International Conference Learning Representations, 2020.
Luisa M. Zintgraf, Leo Feng, Cong Lu, Maximilian Igl, Kristian Hartikainen, Katja Hofmann, and Shimon Whiteson. Exploration in approximate hyper-state space for meta reinforcement learning. In International Conference on Machine Learning. 2021.
Ilya Zisman, Vladislav Kurenkov, Alexander Nikulin, Viacheslav Sinii, and Sergey Kolesnikov. Emergence of in-context reinforcement learning from noise distillation. In International Confer- ence on Machine Learning, 2024.
13

---

Published as a conference paper at ICLR 2025
A IMPLEMENTATION DETAILS
A.1 SOURCE ALGORITHMS
For the source algorithms (Schulman et al. 2017 Haarnoja et al. 2018), we use the implementation of Stable Baselines 3 (Rathin et al. (002]. The hyperparamter settings are detailed in Table B Unless otherwise specified, most hyperparameters are set to default values.
Table 3: Hyperparamters the source PPO algorithm.
Hyperparmater
Darkroom
Dark Key-to-Door
Darkroom-Permuiee
Meta-Workd
learning_rate
3e-4
3e-4
3e-4
3e-4
nusteps
20
.00
50
100
batch size
100
epochs
20
10
20
20
0.99
0.99
0.99
0.99
total timestops
100K
100K
100K
1M
Table 4: Hyperparamters for the source SAC algorithm.
Hyperparmater
Meta-Was World
learning raae
3e-4
lcarning_starts
100
hatch_size
128
train freq
10
gradient steps
1
buffer size
IM
0.99
total timesteps
IM
A.2 TRANSFORMERS
We implement Transformers using the open-source TinyLlama (Zhang al. 2024 The hyperpa- rameters are provided Table along with additional parameters specific to IDT.
Table 5: Hyperparamters for Transformers. All three Transformer modules in IDT share the same of hyperparameters.
Hyperparmater
AD/DPT (disc.)
IDT (disc.)
AD (cont.)
IDT (cont.)
llaye n_head n_emhed intermediate size dropout
4
4
4
4
4
4
32
32
64
32
128
128
128
0.1
0.1
0.1
0.1
attention dropout optimizer
0.1
0.1
0.1
0.1
AdamW cosine decay lc-2
Adam/ cosine decay lc-3
AdamıW cosine decay lc-2
Adam
scheduler
cosine decay 1c-3
dearningraaee (at start) 31
0.9
0.9
0.9
0.9
32
0.99
0.99
0.99
0.99
weight_decar decay
0.01
0.01
0.01
0.01
nLactions per high-level decision dim_z.
10
14

---

Published as a conference paper at ICLR 2025
To Oze1.
1..0 R
artc-1
Causal Transformer
Decssons Go
O Pink
Op-1 1-
0:46-1
0:46-1
Λ
t9c
(a) DICP-AD
Making Decisions
A
Ozos. R
R
Retc
Causal Transformer
Opec-1 dstc-1
Reviewing Decisions
d--k
dr-1
dg
(b) DICP-DPT
(c) DICP-IDT
Figure 4: Model-based adaptation in-context RL methods for DICP. The notations are defined as follows: 2: high-level decisions, R: return-to-go, k: the number of transitions within the context, e: the number of low-level actions guided by single high-level decision, and It': relabeled return-to- go, following (Huang et al. 2024). Newly introduced components highlighted in blue.
Table 6: Hyperparamters for planning.
Hyperparmater
Darkroom
Dark Key-to-Door
Darkroom Darkroom-Permuted
Meta- Woold
planning horizon
16
16
beam size (K)
10
10
sample size (L)
5
10
A.3 ENVIRONMENTS
In the ML1 benchmarks, we observe that using shorter horizons for the source algorithm accelerates the learning of in-context RL methods. Specifically, terminate the source algorithm after 100 steps, as opposed to the original 500-step horizon provided by Meta-World. To minimize train-test also meta-test the in-context RL methods using 100-step horizons. Even under these less favorable conditions, in-context RL methods still outperform the baselines, as shown in
discrepancies,
Fig. 2 and Table
A.4 PLANNING
We find that predicting only the immediate reward, rather than both the immediate reward and the return-to-ga go, is sufficient to outperform baselines. Summation of predicted immodiate rewards up to the current step can room variants, the binary nature of rewards often leads to different beam paths achieving the same
myopic proxy for the total return of planning path. In Dark-
Table 7: FLOPs per action selection.
Method
Darkroom
Dark Key-to-Door
Barkroom-Permutee
Meta-World
AD
6M
20M
20M
709M
DPT
6M
20M
20M
709M
IDT
8M
8M
8M
3M
DICP-AD
2G
18G
18G
8G
DICP-DPT
2G
18G
I8G
8G
DICP-IDT
147M
147M
147M
15M
15

---

Published as a conference paper at ICLR 2025
cumulative reward during planning. When such ties arise, we break them by following the model's preference during action generation. This approach allows the agent to behave like model-free counterpart when the goal key is distant, while benefiting from model-based planning when the dynamics model can confidently predict rewards over shorter horizons. In Meta-World, the dense, human-designed reward structure further improves the effectiveness of this myopie return estima- tion. Since the rewards are designed to guide the agent toward the goal by greedily following reward increases, ranking planning paths based immediate rewards often sufficient to improve in-context RL approaches. Hyperparamters for planning are detailed in Table
Table 8: Results on all benchmarks of Meta- World ML1. We omit the rows for the 5 tasks reported in Table while they are still included the mean calculation.
Task
AD
IDT
DICP-AD (Ours)
DICP-IDT (Ours)
Assembly
0
2
28
4
Basketball
50
54
48
66
Bin-Picking
0
Box-Close
34
44
38
Button-
Cirtoo-Prees-Topeee Pross-Too
98
100
100
100
Button-Prec
-rrtro---es--Toooow------ Wall
98
100
100
100
Button-Press
100
100
100
100
Button-Press- Wall
100
100
Coffee-Button
100
100
100
100
Coffee-Pull
88
78
42
90
Coffec-Push
98
100
92
100
Dial-Turn
76
30
94
Disassemble
Door-Close
100
100
100
100
Door-Lock
98
100
88
98
Door-Open
78
14
90
80
Door-Unloct
72
74
82
Hand-Insert
84
70
86
68
Drawer-Close
100
100
100
100
Drawer-Open
26
46
62
90
Faucet-Open Open
100
100
96
98
Faucet-Close
100
86
100
Hammer
8
22
52
10
Handle-Press-Se
100
100
94
100
Handle-Press
100
100
100
100
Handle-Pull-Sid
100
82
100
Handle-Pull
56
66
74
78
Lever-Pull
96
96
100
98
Pick-Place-Wai
6
66
6
70
Pick-Out-Of-Hole
36
98
2
86
Plate-Slide
92
96
94
Hate-Slide-Siie
100
100
100
100
Site-Slide-Back
92
98
88
86
Plate-Slide-Bacd Slade-lleBB--iiieee
94
100
100
96
Peg- Unplug-Side
86
100
38
98
Soccer
94
96
92
100
Stick-Push
2
0
0
2
Stick-Pull
2
14
0
50
Push-Wall
86
94
92
100
Push-Back
0
Reach-Waid
100
98
94
100
Shelf-Place
68
0
62
Sweep-Into
86
86
88
88
Window-Open
100
100
100
100
Window-Close
100
100
88
98
Mean
69
While model-based planning requires additional computation, the cost is negligible. Specifically, method does not increase the number of training parameters compared model-free counter-
16

---

Published as a conference paper at ICLR 2025
parts, and the primary difference lies in the increased number of Transformer inferences per action selection. As shown Table] 7 the maximum computation per action selection is approximately 18 GFLOPs in experiment. Given that modern GPUs process hundreds of TFLOPs per sec- ond, the computational expense is minimal in practice while the performance gains are substantial, making the trade-off highly favorable in framework. The difference becomes less signifi- cant when using architectures like IDT, which are specifically designed to handle longer sequences efficiently.
B ADDITIONAL EXPERIMENTS
We conducted experiments across all 50 environments of Meta-World MLI, using PPO (Schulman ctt al.) (0177) as the source algorithm with the same hyperparameters as those used in Table Each method is evaluated with 20K environment steps during the meta-test phase. As shown in Table g our method outperforms model-free counterparts. The mean learning curve is displayed in Fig. 2
is noteworthy that since configuration for the source algorithm is not tuned for every task, the results may not fully reflect the full potential of in-context RI. methods, including ours.
Table 9: Meta-test success rates with 5K environment steps on benchmarks of ML1. The results marked with are taken from Wang et al. (2023a). Our approach consistently outperforms all the baselines at 5K steps as well.
Method
Sweep
Peg-Inserts Inserrttt-SS
Steps
MACAWt FOCAL
38
10
5K
BoREL IDAQ
0
0
5K
59
30
5K
IDT
5K
DICP-IDT (Ours)
87
45
5K
Furthermore, we also conducted experiments on ML10 from Meta- Meta-World, a benchmark designed to evaluate the generalization ability of meta-RL methods using ten predefined training tasks and five test tasks. The results, presented Table TO demonstrate that our approach surpasses the model-free counterpart and achieves state-of-the ofte-offthheart performance on this benchmark with significantly fewer environment steps than the baselines. Notably, our method does not rely on expert demonstrations or task descriptions for the test tasks.
Table 10: Meta-test success rates on Meta- ML10. The success rates are reported as percent- ages, averaged over different seeds. The results marked with are taken from Yu et al. (2019).
Method
Success Rate
Steps
PEARL*
13.0
350M
MAML
31.6
350M
RI.2*
35.8
350M
36.7
500K
DICP-IDT (Ours)
46.9
500K
C EXTENDED RELATED WORK
DPT. Decision-Prette Pretriin- Transformer (DPT; Lee et al. (023aa) is an in-context RL approach that meta-trains (or pre-trains) Transformers to predict optimal actions based on contextual trajecto- ries. Consequently, DPT requires access to the optimal actions during training. The authors suggest several sources for gathering the meta-traie trti-training dataset, including (i) randomly sampled trajectories directly from the environment, (ii) trajectories generated by a source RL algorithm, and (iii) trajec- tories of expert policy. In this work, concentrate the second source minimize reliance
17

---

Published as a conference paper at ICLR 2025
direct access environment dynamics and expert policies. While the authors also discuss deploying DPT when offline datasets of the meta-test tasks are available, our focus is on scenarios where only offline datasets for meta-train tasks available for fair comparison. In this scheme, the loss function of DPT' defined as
LDPT (0)
ΣΣ log fo og h-1)
where oj denotes the optimal action the true state 8t.
IDT. In-Context Decision Transformer (IDT; Huang et al. 2024) addresses the high computa- tional costs of previous in-context RL methods, particularly when handling long-horizon tasks. IDT restructures the secision-making process to predict high-level decisions instead of individual ac- tions. These high-level decisions then guide multiple action steps in a separate module, dividing the Transformer used in the earlier approach (Laskin et al.] (023) into modules with shorter sequences. The architecture of IDT consists of three modules: (i) the Making Deeiinn-eesssssss module predicts high- level decisions, (ii) the Decisions-to-Go module decodes these high-level decisions into low-level actions, and (iii) the Reviewing-Decisteee high-level decisions. The loss function for IDT is structured similarly to that of AD, while the loss can be divided across the three modules:
module encodes the resulting low-level actions back into
(0)
log fo ht-1)
a
T
Σ
[log fo(a) z4, of) log fo of 2-1) + log fo( h//1)]
(4)
where 2, high-level decision and 2 an encoded context by the Reviewing- Decissins mod- ule.
TTO. TTO (Janner et al. 202T is an offline RL approach that imitates the policy of the offline dataset, while simultancously learning dynamics model to facilitate planning to improve upon that policy. Our approach similar to this in that we also use model-based planning to improve upon offline datasets. However, framework focuses improving distilled algorithm rather than specific policy. This distinction makes our problem setting more general and challenging compared single-task RL like TTO, since our agent should learn the dynamics model primarily in-context when encountering novel tasks.
Model-based RL. Model-based RL has gained significant attention due to its ability improve sample efficiency by utilizing dynamics model to generate simulated experiences (Ha & Schmid- huber 2018 Kaiser al. 2020; Janner et al., 2019 Schrittwieser al. 2019 et al. [202]] Hafner
et al. 2023). Some approaches further incorporate planning with search algorithms, such MCTS (Schrittwieser et al., 2019: 2021 Ye et al., 202T) or beam search (Janner et al.) [202)) to explore the planning space effectively. However, key challenge in model-based RL is managing the bal- ance between sample efficiency and model bias, as errors in the learned dynamics model can lead to suboptimal decision-making (Janner al. 2019)
Model-based Meta-RL. Previous approaches have explored model-based approach to meta-RI RL (Sæmundsson al. 2018 Lin et al., 2020 Hiraoka al. 2021 Rimon et al.) (022). Clavera al. (018) aim to overcome the limitations of model-based RL (Janner et al. 2019) by incorporating meta-model. Nagabandi et al. (2019) improve upon model-free meta-RI. methods (Duan et al.,
2016 Finn et | (017) 7 to adapt shifting dynamics. Anand et al. (2022) focus zero-shot general-
ization, adapting MuZero (Schrittwieser et al., (019)) to meta-RL. More recently, Rimon et al. (024) introduce a model-based meta-RL approach that combines meta-RL method (Zintgraf et al. with a model-based RL method (Hafner et al. 2023
Axcetly,,Rmoon
2020
18


Reviews:
[{"review_id": "9udnfw4VwJ", "reviewer": "Reviewer_jXdx", "summary": "This paper proposes a model-based in-context reinforcement learning method called Distillation for In-Context Planning (DICP). With a dynamics model for planning, it provides the ability to deviate from the source algorithm's behavior. The authors show that DICP achieves better performance  on Darkroom and Meta-World benchmarks.", "questions": "- Could the authors explain why improvements of DICP-AD over AD is not significant?\n- Is it possible to evaluate on more challenging benchmarks like ML10?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is clearly written, allowing readers to follow the main arguments.\n- It provides a comprehensive experiments and ablations to demonstrate the effectiveness of DICP compared with the model-free counterparts.", "weaknesses": "I think the experimental results of this paper do not strongly support the main motivation. The authors claim that: “Model-free in-context reinforcement learning methods are trained to mimic the source algorithm, they also reproduce its suboptimal behaviors. Model-based planning offers a promising solution to this limitation by allowing the agents to simulate potential outcomes before taking action, providing an additional mechanism to deviate from the source algorithm’s behavior.” However, in the experiments section, DICP does not show significant performance advantages over its model-free counterparts. For example, as shown in Appendix B, the success rate of DICP-AD compared to AD only improves from 68% to 69%, and DICP-IDT compared to IDT only improves from 75% to 80%. Therefore, I believe model-based planning does not significantly enhance the policy beyond the source behavior.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "3CwWnNm3GA", "reviewer": "Reviewer_XR4p", "summary": "The paper proposes a novel method called Distillation for In-Context Model-Based Planning (DICP) to improve the efficiency and effectiveness of in-context reinforcement learning. DICP leverages a learned dynamics model to predict the consequences of actions and uses this information to plan more effectively.", "questions": "1) Can you comment on how much the trade-off between performance vs computation is required in your approach and other comparisons? Do the gains outweigh computational expense?\n\n2) What happens when the world model is incorrect? How is the performance affected? What steps are taken to ensure model-based planning can be accurate?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "-Transformer simultaneously learns environment dynamics and improves policy in-context\n\n-Avoid sub-optimal behavior of the source algorithm", "weaknesses": "-One of the flaws in model based planning is that the model might not be perfect. Errors of the world model might lead to sub-optimal policies as well which haven’t been discussed in the paper.\n\n-Some analysis on how many times sub-optimal behvior from source algorithm was discovered and your approach was able to learn the optimal policy would be important to ensure extra computation of model based planning is worth it.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "UlrjhrVfG4", "reviewer": "Reviewer_zBei", "summary": "This paper extends the use of decision transformers for in-context meta-task learning to incorporate model-based planning. The main innovation here is to have the transformer output predicted state values (r, o, R) in addition to the next action, and to use this state-transition model to select better actions. This can be applied to multiple different transformer-based agents and yields improvements both in terms of sample efficiency and overall score.", "questions": "- How exactly are the actions sampled, and what is the sensitivity to the sampling approach? \n\n- What do you think would happen if the imitation loss and dynamics loss trained two separate models?\n\n- I would be interested in more discussion of how this method might apply to online learning. For example, how might it interact with intrinsically-rewarded exploration to improve the world model? How much does this method depend on the quality of the offline dataset? How effectively would this approach adapt to an environment where the dynamics change?\n\n- Would MCTS perhaps have benefits relative to beam search? Is there a way to not build the whole planning tree as an initial step, or what is the advantage to doing so?\n\n- Are there scenarios where having a world model might detract. For example, what happens if the world model is not accurate enough?\n\n- What are possible explanations for why model-based performs worse on Pick-Out-Of-Hole?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- Advances the performance of in-context RL, an exciting recent direction\n- Provides a mechanism to overcome suboptimal behaviors inherited from the source RL algorithms\n- Achieves state of the art on Meta-World, compared with a large variety of RL and in-context RL algorithms\n- This is an important new innovation that makes sense, and as far as I can tell (though I do not have full knowledge of the literature) is the first demonstration of incorporating model-based planning into transformer-based in-context RL. \n- This seems like a well-done paper with a straightforward but impactful contribution.", "weaknesses": "- Nothing major.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)