Efficient Cross-Episode Meta-RL

Authors: Gresa Shala, André Biedenkapp, Pierre Krack, Florian Walter, Josif Grabocka
Primary area: transfer learning, meta learning, and lifelong learning
Keywords: meta-reinforcement learning, transformers, in-context learning

Abstract:
We introduce Efficient Cross-Episodic Transformers (ECET), a new algorithm for online Meta-Reinforcement Learning that addresses the challenge of enabling reinforcement learning agents to perform effectively in previously unseen tasks. We demonstrate how past episodes serve as a rich source of in-context information, which our model effectively distills and applies to new contexts. Our learned algorithm is capable of outperforming the previous state-of-the-art and provides more efficient meta-training while significantly improving generalization capabilities. Experimental results, obtained across various simulated tasks of the MuJoCo, Meta-World and ManiSkill benchmarks, indicate a significant improvement in learning efficiency and adaptability compared to the state-of-the-art. Our approach enhances the agent's ability to generalize from limited data and paves the way for more robust and versatile AI systems.

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

Reviews and discussion:
<thread>
creation date:
1729887799993

thread_type:
official_review

<thread_object>
creator: reviewer_urva

<content>
summary:
This paper presents an online Meta-RL method. The authors suggested that decomposing the encoder into an intra-episode stage (where each episode is encoded separately) and cross-episode stage (where the vectors from the previous step are encoded into a single representation of the context) should improve performance both in terms of few-shot adaptation and in-distribution and out-of-distribution generalization.
The authors validated their claim on the Meta-World and ManiSkill benchmarks and compared against popular Meta-RL benchmarks.

strengths:
1. The paper is well written, especially the Method and Experiments sections. The formulation and explanation of the algorithm can be well understood
2. The results shown in the Experiments section look quite promising, including results over the widely used Meta-World benchmark 
3. The authors provided an open-sourced code to reproduce the results

weaknesses:
1) My main concern regarding this work is the novelty of the proposed method. The overall algorithm is very similar to previous lines of work, especially RL2 and TrMRL with the main difference from TrMRL is the decomposition of the transformer into two stages. For this change to have a profound enough impact on the field I would expect at least one of the following: 

      a) A theoretical/analytical analysis of why this specific architecture should generalize/perform better than others. The authors gave high-level explanations (e.g lines 204-209), but I found these not rigorous enough. For example, there is no sufficient motivation for hypothesis 1 (line 368), regular recurrent architectures should also be able to capture intra and cross episodic experiences, why should we expect ECET do it better? \
      b) A more thorough empirical study showing the benefits of the chosen design, including more baselines, benchmarks, and an ablation study (more on that below)  

2) The positioning of the paper concerning related work is not clear. For example, in lines 52-53/148-149 in the Introduction Section, the authors claim that previous methods used "single transitions or sequences of transitions from the same episode" to encode the context and state that incorporating cross-episodic information is one of the novelties of the proposed method. This claim is false as many previous works (e.g. VariBAD, TrMRL...) encoded transitions from multiple episodes into a single representation. 

3) The empirical study, which is shown in Section 5, is not sufficient to support the claims of the authors (mainly hypothesis 1 and 3):  \
    a) More benchmarks including ones from the original papers (e.g. VariBAD, RL2) are missing. It is not clear if the results will hold on benchmarks where sufficient hyperparameter tuning of the other baselines took place

    b) An ablation study is missing, especially over the disentangled transformer architecture

    c) Regarding hypothesis 3 - the authors did not compare to theoretical/empirical line of works in generalization in Meta-RL (e.g [1,2])

[1] Lee et al. Improving generalization in meta-rl with imaginary tasks from latent dynamics mixture \
[2] Rimon et al. - Meta Reinforcement Learning with Finite Training Tasks -- a Density Estimation Approach

questions:
1. In the title of the paper you suggested the method is "Efficient", but did not mention in what sense (not in the introduction at least). Do you mean in terms of time complexity (as discussed in lines 267-274)? If so, mention it in the introduction and add some results showing latency/FLOPs compared to "standard" (line 268) approaches.
2. In line 211 you mentioned you sample a sequence of $T$ transitions from each episode.\
    a) How do you sample the transitions? Uniform sampling? With replacement?\
    b) What is the chosen value for the different benchmarks? Did you test the effect of this value on the runtime/performance?\
    c) DId you also sample the transitions when running the other baselines? 
3. Why did you choose to add positional encoding to the IET?
4. What are the number of episodes, length of each episode, and $T$ in each benchmark?
5. Did you use the same hyperparameters for all benchmarks? If not, how sensitive is your method to different hyperparameters?
6. How did you choose the hyperparameters for the baselines? 
7. How many seeds did you use for the experiments in Section 5?  
8. In Figure 5 the results of TrMRL don't match the ones in the TrMRL paper. Do you have an explanation for this? Did you run this baseline as done in the original paper (connecting to questions 2c and 6)?  
9. How many learnable parameters does your architecture have compared to the other baselines (especially TrMRL)?  
10. Typo line 273: ihas -> has
</content>
</thread_object>

</thread>
<thread>
creation date:
1730241915325

thread_type:
official_review

<thread_object>
creator: reviewer_khht

<content>
summary:
This paper proposes a new algorithm; Efficient Cross-Episodic Transformers (ECET) for online meta-reinforcement learning. Their method tackles the problem of enabling reinforcement learning agents to perform effectively in previously unseen tasks. In order to achieve that, they take past episodes and use them as in-context information, considering different task variations and single transitions by not only extracting information from within sequences of transitions but also incorporating information across sequences sampled from different episodes; thus, integrating insights from prior meta-RL research. Their main contribution is improving the generalization capabilities and providing more efficient meta training, outperforming previous state-of-the-art methods. They perform an empirical study in the Meta-World and Maniskill Benchmarks to conclude that their extended algorithm captures more general task features and outperforms the state-of-the-art in online meta-RL algorithms in few-shot adaptation to parametric variations of tasks and out-of-distribution tasks. 
Overall, this paper could be a significant algorithmic contribution, with the caveat of some clarifications on the theory and experiments. Given these clarifications in an author's response, I would be willing to increase the score.

strengths:
For originality, the paper presents a novel method for learning meta-features through efficient cross-episode meta-RL, using a novel transformer architecture that enhances the adaptability of RL agents across diverse sets of tasks. For significance, their approach leverages intra-episode experiences and cross-episode experiences, providing a more comprehensive learning process and integrating ideas from previous meta-RL research. From a perspective of clarity and quality, their learned algorithm is capable of outperforming the previous state-of-the-art and providing more efficient meta-training while significantly improving generalization capabilities. In general, the proposed method enhances the agent’s ability to generalize from limited data supported with sufficient empirical experiments and paves the way for more robust and versatile AI systems.

The paper has some good points to highlight and summarize:
- Improving generalization capabilities and providing efficient meta-training. 
- Outperforming previous state-of-the-art methods by capturing complex, multilevel patterns in the data.
- Integrating multiple ideas from previous meta-RL research.

weaknesses:
Based on my experience, while I truly appreciate the authors’ effort, I believe that their paper lacks novelty and can be considered as an improvement but not as a real contribution for several reasons. First of all, there is no central contribution beyond improving generalizing capabilities by introducing a new architecture and outperforming previous SOTA–which was supported in the empirical results in figures 4, 5, and 6— still, their method’s generalization to tasks that are completely different from the meta-training set remains limited as their approach is tailored to exploiting similarities as shown in hypothesis 3: ECET outperforms the state-of-the-art online meta-RL in out-of-distribution (OOD) tasks. This leads to Section 1, Related Work, it is still unclear how their method differs from other baselines in solving the same problem, specifically the TrMRL method Melo (2022), which is the closest to the authors’ proposed algorithm. In other words, the authors’ addition and value are still vague compared to previous research tackling the same research question. Having read Melo (2022), I see that their cross-episode experience equals the episodic memory system proposed in the TrMRL method. More discussion on these topics would be helpful. 


Minor comments:
1. Page 6: Line 4, ‘ihas’ should be corrected.

questions:
For the theory, there are a few steps that need clarification and further explanations on novelty. For novelty, it needs to be clarified if their method is being stated as a novel result. It looks like the method has already been shown in "Transformers are Meta-Reinforcement Learners.” There is a statement that “TrMRL is the closest baseline to our proposed method (Melo, 2022). It uses a transformer in the inner loop of the RL2 algorithm. It takes a sequence of the 5 most recent transitions as input to the transformer. We use the implementation provided by the authors.” Is your algorithm 2 somehow an extension? Is your theorem 3 completely new?

In Appendix A.2.3 Procedural Components, you mentioned that “SampleTransitions: We sample the transition sequences of length T from each episode in H. For past episodes, we randomly” You randomly did what?

There is also one concept in the algorithm that I cannot verify. How did you apply the intra-episode and cross-episode simultaneously? More about this would be illuminating

For the experiments, the following should be addressed:
- The central contribution is enhancing the generalizing capabilities. The authors concluded that their method is capable of adjusting to completely unseen tasks that share the same state and action spaces as the training tasks. Despite outperforming state-of-the-art methods in online meta-RL by efficiently compressing cross-episodic knowledge, their method’s generalization to tasks that are completely different from the meta-training set remains limited as it is tailored to exploiting similarities. I would like to see some discussions on this and I think it would be beneficial to demonstrate that empirically.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730556692840

thread_type:
official_review

<thread_object>
creator: reviewer_iqfe

<content>
summary:
This paper proposes a method for learning features for meta-reinforcement learning in latent space by training a two-level hierarchy of transformer networks. The lower layer is trained to represent trajectories experienced while doing a specific task (i.e. experiencing a single MDP) while the upper layer is trained on all tasks. This method is shown to have significantly better performance than several other meta-RL methods on a variety of meta-RL benchmarks, with the primary hypothesis for this being that the proposed architecture allows the networks to learn more general features that are informative across a variety of tasks. This hypothesis is supported by some experiments showing that embedded representations learned by the proposed method tend to cluster together by task more than representations learned by other methods.

strengths:
Paper is overall clearly written and explained.

Related work positions the paper well.

Experiments are well designed.

I am a big fan of explicit hypothesis statements, and very much appreciate this paper for being clear about what it is investigating.

The idea is relatively simple and seems like a clear opportunity for more research.

The performance is impressive and the results are quite thorough.

weaknesses:
There are few weaknesses in my opinion. Instead, I will give some minor feedback.

- “By considering both intra-episode experiences (information within a single episode) and cross-episode experiences (information across multiple episodes), our method can learn richer meta-features for the environment.” The opposite of intra is inter, not cross. The title already contains 'cross', but consider using this term when explaining (I believe it already occurs once towards then end).

- Consider citing Bhatia et al. (RL^3: Boosting Meta Reinforcement Learning via RL inside RL^2) which also uses transformers and learns state-augmentation features for meta-RL, albeit not through inter-episode data sharing.

- Figure 2 bleeds into the margins.

- It could be clearer which baselines (or other meta-RL methods) this method is compatible with and which are mutually exclusive.

questions:
“Hypothesis 1: ECET captures more general task features through its ability to capture intra- as well as cross-episodic experiences.” 

I recognize claims about feature generality are difficult to substantiate, and that it is somewhat of a catch-22 to try to deal with these topics experimentally. I appreciate the authors for making an attempt. I am interested in their thoughts regarding essentially more rigorous versions of such experiments.

- Do you have a definition for the generality of a task feature? How does one measure this? 

- You’ve shown that embedding vectors tend to cluster for some tasks. Do you think there are any possible confounders beyond generality of features?

- Moreover, if one were to take a more systematic and focused approach to quantifying the difference between tasks so that embeddings generated from those tasks may be compared to the embedding distances --- similar to what is presented in figure 13 but with some notion of how similar two columns/rows ought to be --- do you suspect they would encounter anything surprising?

Is there anything that makes such an experiment difficult to execute other than initially deciding how to measure difference between tasks? I agree these results seem promising w.r.t. this hypothesis, but there are definitely a few more logical steps to nail down before this can be determined without doubt (which would be an important result in it own right, of course).
</content>
</thread_object>

</thread>
<thread>
creation date:
1730649472286

thread_type:
official_review

<thread_object>
creator: reviewer_ixfq

<content>
summary:
The paper proposes new online Meta-RL algorithm that is able to combine intra- and cross-episodic information to better make decisions under uncertainty. Each intra-episode is encoded into a latent variable by a transformer model, which are then processed altogether by another transformer and a context-specific latent is formed. Authors claim SOTA results on Maniskill and challenging Meta-World ML45 benchmarks, surpassing previous SOTA methods.

strengths:
- The paper reports state of the art results on challenging tasks, including challenging environment Meta-World ML45.
- The method is explained clearly and visually explained.

weaknesses:
[W1] Authors compare their method with TrMLR, which uses RNN in the outer-loop, while authors of the proposed method use transformer. I believe that the difference in the capacity can explain the performance distinction between two methods. To my mind, it is early to claim SOTA results without trying to match architectures first.

[W2] A strong baseline is missing from the paper: AMAGO [1]. The [1] authors do not report the results on ML45, stating it is left for the future work. However, since the main contribution of the paper is SOTA results, I believe AMAGO should be included in the baselines. Also, mind that AMAGO was introduced in ICLR”24.

[W3] Authors do not state whether the baselines (especially TrMLR) were tuned to get the best performance out of them. Again, as the paper claims SOTA results, it is vital to have this information.

questions:
1. I would appreciate if authors include an additional baseline method on ML45, AMAGO [1], to ensure the result of the proposed method is indeed SOTA.
2. To justify SOTA performance, I would also like to see the experiments on TrMLR with transformer or mamba architecture as an outer-loop memory model.
3. I would also appreciate it if authors provide how many hyperparameters were swept for the baseline TrMLR method and which method for HPO was used.

Although my initial recommendation is to reject the paper, I would be happy to revise my score as the authors address my questions. 

[1] Grigsby, Jake, Linxi Fan, and Yuke Zhu. "Amago: Scalable in-context reinforcement learning for adaptive agents.”

*After rebuttal update:* I changed my score to 6 since the questions were mostly addressed, see the comments for details.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732357149955

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Response to all reviewers

comment:
Dear Reviewers,

We sincerely thank you for your constructive feedback and thoughtful comments. Your suggestions have been invaluable in improving the clarity, rigor, and scope of our manuscript. We have carefully considered each of your points and have made significant updates to address your concerns comprehensively. Below, we summarize the key changes and additions made to the revised manuscript, now available on OpenReview.

---

**On Including AMAGO as a Baseline**

We appreciate the suggestion to include **AMAGO** as a baseline. In response:  
- We have evaluated **AMAGO** on both the **ML10** and **ML45** benchmarks. The results are presented in **Figure 6** and other figures summarizing ML10 and ML45 outcomes.  
- **Key Findings:**  
  - Transformer-based methods, including ECET and AMAGO, outperform non-transformer approaches.  
  - ECET consistently achieves superior performance compared to AMAGO, underscoring the effectiveness of our disentangled transformer architecture in capturing task representations efficiently.  

By adding AMAGO as a baseline, we strengthen the scope and rigor of our comparisons, ensuring our results are well-positioned within the state-of-the-art.

---

**On Additional Benchmarks**

To address concerns about the breadth of our evaluation:  
- We added locomotion tasks such as **HalfCheetahVel**, **HalfCheetahDir**, and **AntDir** to further evaluate the generalization and meta-training capabilities of our approach.  
- **Figure 4** compares generalization performance at test time across these tasks, while **Figure 12** illustrates the meta-training curves.  
- Our results show that ECET, along with TrMRL and VariBAD, demonstrates robust performance across these tasks, further validating the versatility of our approach.

---

**On the Positioning of Our Approach**

We appreciate your feedback regarding the positioning of our method and its relation to prior work. In response, we have revised the **Introduction** and **Related Work** sections to clarify our contributions and distinguish our approach from existing methods such as VariBAD, TrMRL, and AMAGO. Key updates include:  

1. **Acknowledgment of Prior Work:**  
   - We detail the contributions of prior works like VariBAD, TrMRL, and AMAGO in encoding transitions across episodes. This provides a more accurate context for understanding the advancements offered by ECET.

2. **Specificity of Our Contribution:**  
   - **Compared to VariBAD and RNN-based Methods:**  
     Our transformer-based architecture effectively handles long-term task dynamics across multiple episodes, avoiding limitations like vanishing gradients and memory constraints that affect RNN-based approaches.  
   - **Compared to TrMRL and AMAGO:**  
     ECET achieves a better balance of computational efficiency and performance by reducing runtime complexity while maintaining the ability to process inter-episode information.  

3. **Practical Relevance:**  
   - Our method addresses key challenges in Meta-RL, such as improving generalization and reducing computational overhead.  
   - Empirical results demonstrate faster convergence during meta-training and fewer timesteps required for generalization to out-of-distribution tasks.  
   - Comparisons with TrMRL and AMAGO consistently show ECET’s superiority in both efficiency in terms of timesteps and final performance.

---

**Closing Remarks**

We hope that these updates address the concerns raised and clarify the contributions and strengths of our work. The revised manuscript reflects a broader evaluation, a clearer positioning of our method, and a more transparent experimental setup.  

We are grateful for the opportunity to further refine our study based on your feedback and look forward to your comments. We hope these revisions meet your expectations and provide a comprehensive basis for evaluating our contributions.  

Sincerely,  
The Authors
</content>
</thread_object>

</thread>
<thread>
creation date:
1775718989335

thread_type:
public_comment

<thread_object>
creator: ~jingyang_you1

<content>
comment:
Would it be possible for the authors to release the source code for ECET? I cannot seem to replicate the mujoco results. For example, in halfcheetahdir, VariBAD reaches ~1500 using only 4e6 steps, while in the JMLR version of VariBAD, ~1500 can only be obtained with ~4e7 steps, which differs by a factor of 10. 

The reason I can come up with is ECET reports the #meta-steps for mujoco environments, instead of the total number of steps. With T=5 and #rollouts=2, the factor 10 is explainable.  

The metaworld experiments seem to have reported the total steps.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This work proposes an algorithm for learning features for meta-reinforcement learning in latent space. This is achieved by training a two-level hierarchy of transformer networks.  The lower layer represents experienced trajectories while the upper one is trained on all tasks. This method achieves a great good performance in meta-RL benchmarks. It is also demonstrated experimentally that the performance gains are due to the improved representation learning algorithm. Although we had a dissenting reviewer that raised concerns about its novelty and relationship with related work. These were outweighed by the support of other reviewers that highlighted the competitiveness of the empirical results.