Distilling Reinforcement Learning Algorithms for In-Context Model-Based Planning

Authors: Jaehyeon Son, Soochan Lee, Gunhee Kim
Primary area: reinforcement learning
Keywords: reinforcement learning, in-context learning

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.

Decision: accept (poster)
Reviewer scores: 6, 6, 8

Reviews and discussion:
<thread>
creation date:
1730051306549

thread_type:
official_review

<thread_object>
creator: reviewer_zbei

<content>
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.

strengths:
- Advances the performance of in-context RL, an exciting recent direction
- Provides a mechanism to overcome suboptimal behaviors inherited from the source RL algorithms
- Achieves state of the art on Meta-World, compared with a large variety of RL and in-context RL algorithms
- 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. 
- This seems like a well-done paper with a straightforward but impactful contribution.

weaknesses:
- Nothing major.

questions:
- How exactly are the actions sampled, and what is the sensitivity to the sampling approach? 

- What do you think would happen if the imitation loss and dynamics loss trained two separate models?

- 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?

- 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?

- Are there scenarios where having a world model might detract. For example, what happens if the world model is not accurate enough?

- What are possible explanations for why model-based performs worse on Pick-Out-Of-Hole?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730220776984

thread_type:
official_review

<thread_object>
creator: reviewer_xr4p

<content>
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.

strengths:
-Transformer simultaneously learns environment dynamics and improves policy in-context

-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.

-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.

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?

2) 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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730638293341

thread_type:
official_review

<thread_object>
creator: reviewer_jxdx

<content>
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.

strengths:
- The paper is clearly written, allowing readers to follow the main arguments.
- 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.

questions:
- Could the authors explain why improvements of DICP-AD over AD is not significant?
- Is it possible to evaluate on more challenging benchmarks like ML10?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732495601828

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Author Response

comment:
Dear Reviewers and AC,

We sincerely appreciate your invaluable feedback, which has been instrumental in improving our paper. In response, we have made the following updates:
- Added a computational analysis in Appendix A.
- Included an additional experiment on ML10 in Appendix B.
- Conducted an ablation study analyzing the impact of world model accuracy on final performance, detailed in Appendix C.

We hope these updates enhance the clarity and completeness of our work.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The authors present DICP, which builds upon decision transformers to enable in-context, model-based planning/RL.

The paper leverages model learning to accelerate and improve in-context adaptation to new tasks in RL. The method is general-purpose and the results show strong performance.

Reviewers highlighted several weaknesses. In particular, there were questions around whether all claims in the paper were justified by experimental evidence. There were also questions around the robustness of the learned model, although the authors added new results demonstrating this. 

All reviewers recommend acceptance, and the authors have largely addressed the major weaknesses mentioned.