Library

Model-Based Transfer Learning for Contextual Reinforcement Learning

Deep reinforcement learning (RL) is a powerful approach to complex decision making. However, one issue that limits its practical application is its brittleness, sometimes failing to train in the presence of small changes in the environment. Motivated by the success of zero-shot transfer-where pre-trained models perform well on related tasks-we consider the problem of selecting a good set of training tasks to maximize generalization performance across a range of tasks. Given the high cost of training, it is critical to select training tasks strategically, but not well understood how to do so. We hence introduce Model-Based Transfer Learning (MBTL), which layers on top of existing RL methods to effectively solve contextual RL problems. MBTL models the generalization performance in two parts: 1) the performance set point, modeled using Gaussian processes, and 2) performance loss (generalization gap), modeled as a linear function of contextual similarity. MBTL combines these two pieces of information within a Bayesian optimization (BO) framework to strategically select training tasks. We show theoretically that the method exhibits sublinear regret in the number of training tasks and discuss conditions to further tighten regret bounds. We experimentally validate our methods using urban traffic and standard continuous control benchmarks. The experimental results suggest that MBTL can achieve up to 43x improved sample efficiency compared with canonical independent training and multi-task training. Further experiments demonstrate the efficacy of BO and the insensitivity to the underlying RL algorithm and hyperparameters. This work lays the foundations for investigating explicit modeling of generalization, thereby enabling principled yet effective methods for contextual RL.

cs.LGReinforcement Learning in RoboticsComputer Science
View graphSource

Paper

Similar papers

Structure Detection for Contextual Reinforcement LearningTianyue Zhou, Jung-Hoon Cho, Cathy Wu2026 · Proceedings of the AAAI Conference on Artificial Intelligence In Library
Structure Detection for Contextual Reinforcement LearningTianyue Zhou, Jung-Hoon Cho, Cathy Wu2026 · AAAI Conference on Artificial Intelligence · 0 citations In Library
Context-aware Dynamics Model for Generalization in Model-Based Reinforcement LearningKimin Lee, Younggyo Seo, Seunghyun Lee et al.2020 · arXiv (Cornell University) · 165 citations In Library
Contextualize Me -- The Case for Context in Reinforcement LearningCarolin Benjamins, Theresa Eimer, Frederik Schubert et al.2022 · arXiv (Cornell University) · 57 citations In Library
Zero-Shot Context Generalization in Reinforcement Learning from Few Training ContextsJames Chapman, Kedar Karhadkar, Guido Montúfar2025 · arXiv (Cornell University) · 0 citations In Library
Learning without Knowing: Unobserved Context in Continuous Transfer Reinforcement LearningChenyu Liu, Yan Zhang, Yi Shen et al.2021 · arXiv (Cornell University) · 6 citations In Library
Meta-Reinforcement Learning with Contextual Bias ReductionSiming Lan, Jiaming Guo, Rui Zhang et al.2026 In Library
A Survey on Model-based Reinforcement LearningFan-Ming Luo, Xu Tian, Hang Lai et al.2022 · arXiv (Cornell University) · 180 citations In Library
Context-Based Meta-Reinforcement Learning With Bayesian Nonparametric ModelsZhenshan Bing, Yuqi Yun, Kai Huang et al.2024 · IEEE Transactions on Pattern Analysis and Machine Intelligence · 16 citations In Library
Contextual Intelligence: The Next Leap for Reinforcement LearningAndré Biedenkapp2026 In Library

Peer review

Accept (poster)NeurIPS 2024 · via OpenReview
Reviewer ECSQ6/10 · confidence 3/52024-07-11

Summary

The paper proposes Model-Based Transfer Learning (MBTL) in Contextual RL, which solves multiple related tasks and enhances generalization across different tasks. MBTL strategically selects a set of source tasks to maximize overall performance and minimize training costs. The paper theoretically demonstrates that the method exhibits regret that is sublinear in the number of training tasks. MBTL achieves greater performance than other baselines (e.g., exhaustive training, multi-task training, and random selection) on urban traffic and control benchmarks.

Strengths

- The paper is well-written and organized, and includes a thorough illustrative diagrams and examples. - The problem formulation is well done, with clear mathematical representation. The analysis of Bayesian Optimization appears to be thorough. - The authors provide code for training and evaluating their proposed method. This significantly facilitates reproducing the experimental results and extending the introduced method.

Weaknesses

- The assumptions are too tight, particularly Assumption 3, which models the generalization gap using linear constraints. This approach is unsuitable for complex environments and therefore lacks generalizability. - The experimental environments (urban traffic and control benchmarks) are overly simplistic, consisting solely of vector environments with low-dimensional state spaces. The study lacks comparisons with complex tasks, such as those in the CARL benchmark, including games and a real-world application of RNA design. - The ablations on DRL algorithms (DQN, PPO, A2C) utilize outdated methods. Why not use more recent RL baselines?

Questions

Below, I have a few questions and feedback for the authors: - How does the computational time consumption compare between MBTL and other baselines (exhaustive training, multi-task training, and random selection)? - I am curious to see experimental results in complex environments, such as visual environments.

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

N/A

Authorsrebuttal2024-08-07

References

- [1] C. Benjamins et al., “Contextualize Me -- The Case for Context in Reinforcement Learning,” Transactions on Machine Learning Research, Jun. 2023. - [2] https://automl.github.io/CARL/main/source/environments/environment_families/rna.html - [3] Q. Li, Z. Peng, L. Feng, Q. Zhang, Z. Xue, and B. Zhou, “MetaDrive: Composing Diverse Driving Scenarios for Generalizable Reinforcement Learning,” IEEE Transactions on Pattern Analysis and Machine Intelligence, vol. 45, no. 3, pp. 3461–3475, Mar. 2023. - [4] V. Jayawardana, C. Tang, S. Li, D. Suo, and C. Wu, “The Impact of Task Underspecification in Evaluating Deep Reinforcement Learning,” in Advances in Neural Information Processing Systems, 2022.

Authorsrebuttal2024-08-14

Thank you again for your review. Please let us know if you have any further questions or comments. If you feel that your questions were sufficiently addressed, we would deeply appreciate it if you could consider raising the score.

Reviewer wJEK7/10 · confidence 4/52024-07-13

Summary

The paper proposes a new framework to estimate the expected generalization performance across different tasks where the differences have an explicit and model-based structure. To improve the expected generalization performance via training on selected tasks, the paper proposes naive and Bayesian optimization based method to effectively explore the task space to find a policy that have the optimal zero-shot performance given a selected task. The experiments demonstrate that the proposed MBTL method outperform other baseline method and reach the performance of the optimal transfer method assuming full knowledge.

Strengths

- By and large, the paper is written well. I especially appreciated the detailed discussion on SSTS problem formulation and its relations to robust RL training. - The idea of Bayesian Optimization to search training tasks using generalized performance estimation where tasks differences are “model-based” or have an explicit structure is novel and useful. - The section on analysis of BO and its comparison with ES and PS are also interesting, and highlights the sublinear regret theoretical results. - The theoretical results appear to be correct.

Weaknesses

- The empirical evaluation is by and large only on low dimensional systems. It would have been interesting to see how the method would scale with more challenging, high-dimensional tasks, such as those commonly found in vision control tasks in robotics. - It would have been interesting to see how well this method would work with policies/controllers not parametrized with neural networks (i.e. kernel machines). - Some additional commentary on how to use the GP to search for a trained policy given a selected task.

Questions

- Could you comment on the challenges of applying MBTL-GP on vision-control tasks?

Rating

7

Confidence

4

Soundness

4

Presentation

4

Contribution

3

Limitations

Yes, the author has discussed the limitations.

Reviewer 21Qs4/10 · confidence 4/52024-07-26

Summary

The paper introduces a framework called Model-Based Transfer Learning (MBTL) for solving contextual reinforcement learning problems. By modelling the performance loss as a simple linear function of task context similarity, the authors leverages Bayesian optimization techniques and provides theoretical analysis showing that MBTL exhibits sublinear regret in the number of training tasks, and discusses conditions to further tighten the regret bounds. MBTL is validated on a range of simulated traffic management scenarios and standard control benchmarks, demonstrating superior performance compared to baselines like exhaustive training, multi-task training, and random task selection.

Strengths

1. In general, the paper is easy to follow and well-motivated. The figures are helpful for readers to quickly grasp the key concepts and problem settings. 2. Active and transfer learning for contextual RL is an interesting and practical problem, especially given that exhaustive/multi-task training on RL tasks can be computationally expensive. 3. The experiments are conducted on a relative broad range of benchmarks and context variations.

Weaknesses

I do think there are several major concerns regarding the current form of the paper: 1. The paper (especially the theoretical analysis) is based on too many assumptions: continuity of task space and performance function, linear genearalization gap and deterministic MDP transitions, which oversimplifies the problem and significantly limits the applicability range of the proposed method. The most concerning is Assumption 3 (Linear generalization gap), for which I have no clue why and when this can hold. Assuming Lipschitz continuity of the generalization gap seems much more reasonable to me, which makes assumption 3 an inequality instead of an (approximate) equality. 2. In line 96-98, the paper assumes that a vector representation/context feature $x$ of each task is given a priori, based on which the continuity assumptions 1-2 are made. Howeveer, as in line 333-335, context features are often not visible. Assuming prior knowledge of such information further limits the applicability range of the proposed method. 3. (2 continued) In fact, no matter whether context features are given a priori, there are many effective methods which perform transfer/meta-RL by learning a conditioned policy [1][2][3], which should be significantly better than the "multi-task" baseline adopted in the paper. However, there is currently no empirical evidence on how well the proposed method compared to these stronger baselines. 4. Even though Theorem 2 and Corollary 2.1 2.2 give quantitative regret bounds assuming the search space can be narrowed down at each training step, there is no theoretical guarantee on how the proposed method, especially MBTL-GP, can effectively realize such restricted search space. Hence the current analysis is not complete in terms of proving the effectiveness of proposed method. [1] Rakelly, Kate, et al. "Efficient off-policy meta-reinforcement learning via probabilistic context variables." International conference on machine learning. PMLR, 2019. [2] Shagun Sodhani, Amy Zhang, and Joelle Pineau. Multi-task reinforcement learning with context-based representations. In International Conference on Machine Learning, pages 9767–9779. PMLR, 2021. [3] Li, Lanqing, Rui Yang, and Dijun Luo. "Focal: Efficient fully-offline meta-reinforcement learning via distance metric learning and behavior regularization." ICLR 2021.

Questions

What's the foundamental difference between "exhaustive training" and "multi-task RL"? Based on the brief description in line 257-258, they seem similar to me. Also why is "multi-task RL" even worse than "random" most of the time (in Table 1, 2)?

Rating

4

Confidence

4

Soundness

2

Presentation

3

Contribution

2

Limitations

N/A

Reviewer 21Qs2024-08-12

Thank you for the rebuttal, I appreciate the efforts for providing further experiments. After carefully reading through the rebuttal and general response, unfortunately, I feel some of my major concerns still remain. Most importantly, given the current form of the paper, I would expect more in-depth analysis and insight regarding the following subjects: 1. $\textbf{Theoretical justification of the linear generalization gap}$. This is claimed by the authors as a "key strength" of the paper. However, unlike a loss function with well-defined form, generalization gap is a complex function of model parameters, architectures, dataset sizes & distributions as well as context feature $\textit{without}$ a closed-form expression. Simply taking a linear assumption makes subsequent theoretical treatment easier, but without any convincing insight or theoretical justification, it would significantly limit the impact of the paper. My suggestion: First of all, I appreciate the additional result that the proposed algorithms perform reasonably well in the presence of a non-linear (e.g., quadratic) generalization gap function. However, to be more general, maybe consider using a separate nerwork to approximate this function, which in principle can be any function due to the universal approximation theorem, and then demonstrate the effectiveness of your methods. If you still want to hold on to the linear form to support your theoretical development, consider the Lipschitz constraint instead, which seems much more realistic. 2. $\textbf{Regarding MBTL-GP realizing restricted search space}$, the additional empirical evidence is nice. But since Theorem 2 and Corollary 2.1 2.2 are meant to theoretically ground the effectiveness of proposed methods, I think more rigorous proofs and insights instead of just empirical observations are still necessary to complete the whole argument. If you find it extremely hard or even impossible to bound/model the reduction rate in search space, at least state it clearly as an assumption, which of course, limits the impact of the paper. # Additional Concerns Given the authors explanation about "exhaustive training" and "multi-task RL", I realize that I previously misinterpreted the SSTS problem as a "continuous learning" + "active learning" setting, where you train a $\textbf{single}$ model on a sequence of tasks, in attempt to select the optimal next task to maximize generalization. However, now if I understand correctly, the proposed methods actually follow the "exhaustive training" paradigm, where for each new task, a new model is trained. This formulation of SSTS seems unconventional (or "novel" on the bright side), for which I have two major concerns: 1. According to Eqn 3, the task selection requires evaluation of all existing models (1-k) on every single target task $x'$, which may severely increase the total computational cost of SSTS. Suppose we have $N'$ target tasks, the proposed task selection requires $O(N'k)$ complexity to compute. If we want to select $N$ source tasks in total sequentially, it will end up with $O(N'N^2)$ complexity. Even though performing evaluation/inference is less computationally demanding than training, this additional cost can be non-neglectable especially when $N$ and $N'$ are large. This potentially makes the main motivation that "Our work has the potential to reduce the computational effort needed to solve complex real-world problems" much weaker. 2. Training a new model whenever encountering a new task is not scalable. In an idealized scenario, we would like to have a "universal model", like human brain, which can continuously learn by reusing prior knowledge to solve new tasks that are similar to tasks encountered before, and only make significant updates when the new task is completely beyond your current skill set (aka "out-of-distribution" in statistical words). This is the fundamental motivation of continual learning or life-long learning. The current setting of SSTS, if I understand correctly, seems less realistic or at least require significant justification for its practical impact.

Authorsrebuttal2024-08-12

We appreciate your valuable and insightful feedback. >Theoretical justification of the linear generalization gap We appreciate your concern regarding the linear assumption of the generalization gap. As discussed in our general response [GR1], the linear model was chosen for its simplicity and to streamline the algorithm design process, though various complex factors influence the generalization gap. We are currently exploring the possibility of extending the proposed methods using neural networks to approximate the generalization gap. But again, the key strength of our simple algorithm is that we don’t necessarily need pre-training of those parameters. >Regarding MBTL-GP realizing restricted search space We understand the importance of rigorous theoretical analysis and regret bounds. We agree that the paper would benefit from a more detailed analysis of how the MBTL-GP algorithm can systematically reduce the search space. While we offer examples of simpler algorithms like MBTL-GS to illustrate potential strategies, we face challenges in bounding the reduction rate for MBTL-GP. Therefore, we have provided empirical insights into the rate at which the search space for MBTL-GP can be reduced. In the revised paper, we clearly state this approach, along with formally defined assumptions and detailed explanations. >According to Eqn 3, the task selection requires evaluation of all existing models (1-k) on ... We appreciate your observation regarding the difference between our proposed method and conventional approaches such as multi-task training, where a single universal model is trained across multiple tasks, or independent (exhaustive) training, where separate models are trained for each task. As you mentioned, one of the most important motivations for our method is that evaluation is much cheaper than training RL models in terms of computational cost. We think that this paradigm of training multiple models and applying zero-shot generalization (or fine-tuning) has little been studied and is promising. The strength of our approach lies in its ability to achieve near-oracle performance with a significantly smaller number of trained models. For example, in our experiments on standard control benchmarks with 100 tasks, we achieved performance close to the oracle level with only 10-15 trained policies. While exhaustive evaluation is required to ensure the selection of the best model among all possible task-specific models, we believe that the computational cost of evaluation is relatively low compared to training the remaining RL policies. Moreover, once the generalized performance of the models is obtained, further evaluations are not necessary in subsequent steps. This reduces the evaluation complexity to $O(N')$ per step, resulting in a total complexity of $O(N N')$ when training $N$ tasks, as opposed to the $O(N' N^2)$ complexity you mentioned. Additionally, in practical scenarios, we typically consider cases where $k \ll N$, further reducing the complexity to $O(N' k)$, which is significantly smaller than $O(N N')$. This highlights the computational efficiency of our approach in scenarios with many tasks. >Training a new model whenever encountering a new task is not scalable. In an idealized... Thank you for your insightful comments. We recognize the importance of developing a "universal model" that can continuously learn and adapt to new tasks by reusing prior knowledge, as emphasized in recent research on continual and life-long learning. However, we believe that there are specific scenarios where task-specific training policies, rather than a universal model, are necessary and practical, especially when side information is available. For instance, in the design of traffic signal phases at 4-way intersections, it is crucial to train distinct reinforcement learning policies for different intersection configurations. Considering that there are approximately 16 million different intersections in the United States, training a universal model or separate 16 million models for all configurations would be highly expensive and almost impossible. However, by leveraging context information such as the number of lanes, lane lengths, expected traffic inflows, and speed limits, we can intelligently devise a training procedure that trains multiple models, while still achieving good generalization across the broader distribution. Empirically, we have observed that in such scenarios, our method proves to be computationally more efficient and performs effectively. This approach allows for the targeted training of models that are specialized but still generalize well to similar tasks, thus addressing practical constraints in real-world applications. We hope this conveys the core idea and motivation of our work. We also acknowledge the potential of context-free learning for scalable, universal models, as you mentioned, and consider it an important direction for future research.

Reviewer 21Qs2024-08-13

Thank you for the detailed response. My concern regarding the computational complexity is largely resolved. However, some of the major concerns remain, which I think the authors also agree upon. In summary, this is a novel paper which introduces the concept of "model-based transfer learning" based on "sequential source task selection", which I have not seen the exact same setting before, to my best knowledge. The authors propose to solve this problem by Bayesian optimization techniques, which are empirically proven effective, but the methodologies are not new. By assuming a linear generalization gap as well as bounded reduction of search space, the authors arrive at theoretical guarantees for sublinear regret. By the design, the proposed method can achieve better computational efficiency compared to conventional multi-task learning, in certain scenarios. I will give credit for all contributions above. However, as we discussed during the rebuttal, the paper still falls short in providing serveral key pieces of the story, such as the theoretical justification of the linear generalization gap (also needs to extend it to more generalized, realistic format as we discussed), theoretical guarantee for bounded search space, reliance on the existence of continuous context features, which significantly restrict the applicability and practical impact of the framework . Also for the problem setting of SSTS, the authors provided the example of "16 million different intersection configurations in the United States" to justify the need of training multiple models, which I find reasonable but $\textit{not convincing enough}$. A more promising and general approach to me to solve the same problem, is to leverage the power of pretrained large models. Specifically, one can use a pretrained context encoder model (e.g. large vision model for visual observation) to extract the context feature from the raw input, and use it to condition the downstream RL policy, instead of training a new model for each task, and carefully select the "optimal" next task, which is far less scalable. To conlcude, I believe there are novelties in this work and appreciate the authors' efforts in rebuttal, but remain conservative about its practical impact to the community. More importantly, I sincerely hope the authors to conduct further investigations to fix the issues we agreed upon to make the paper stronger, whether the paper is accepted. Given the reasoning above, I will keep my score but are open to further discussions.

Authorsrebuttal2024-08-14

Thank you for your thoughtful and comprehensive feedback. We appreciate your recognition of the novel aspects of our work. First of all, we see the reviewer’s concern on the generalization gap assumption. In situations where the true function is difficult to analyze, approximation methods are commonly used. Our MBTL algorithms approximate the generalization gap with a simple linear function of task context similarity. We show that even with the linear function, our MBTL framework works in various settings ranging from standard control tasks to complex real-world traffic applications. To further address your concern, we have also evaluated the MBTL-GP performance using non-linear approximations, including quadratic, cubic, $x^5$, and $x^{10}$ models along with the RMSE between the actual generalization gap. While higher-order approximation functions generally result in lower RMSE errors, we observed that the overall performance of MBTL on the SSTS problem does not consistently improve with more complex non-linear approximation. For example, the RMSE error of generalization gap with 10th-order polynomial approximation does not always perform best to solve CMDP tasks. This suggests that the simplicity and interpretability of the linear model can provide significant advantages without compromising effectiveness. | | Linear | Linear | Quadratic | Quadratic | Cubic | Cubic | $x^5$ | $x^5$ | $x^{10}$ | $x^{10}$ | | ----------------- | ----------- | ------ | ----------- | --------- | ----------- | ------ | ----------- | ------ | ----------- | ------ | | | Performance | RMSE | Performance | RMSE | Performance | RMSE | Performance | RMSE | Performance | RMSE | | Pendulum | 0.7555 | 0.1070 | 0.7423 | 0.0965 | **0.7615** | 0.0862 | 0.7494 | 0.0705 | 0.7558 | 0.0192 | | CartPole | 0.8896 | 0.2571 | 0.8102 | 0.1905 | **0.8941** | 0.1495 | 0.8926 | 0.1029 | 0.8761 | 0.0745 | | BipedalWalker | **0.9331** | 0.1422 | 0.9237 | 0.1201 | 0.9329 | 0.1016 | 0.9318 | 0.0771 | 0.9325 | 0.0246 | | HalfCheetah | 0.9165 | 0.0019 | 0.8426 | 0.0019 | 0.9260 | 0.0018 | 0.9271 | 0.0017 | **0.9295** | 0.0004 | | Traffic Signal | **0.8966** | 0.1780 | 0.8907 | 0.1606 | 0.8965 | 0.1409 | 0.8963 | 0.1106 | 0.8952 | 0.0676 | | Advisory Autonomy | 0.8177 | 0.0551 | 0.7782 | 0.0508 | 0.8214 | 0.0464 | **0.8245** | 0.0393 | 0.8245 | 0.0283 | | Eco-Driving | 0.5377 | 0.1071 | 0.4811 | 0.0950 | 0.5282 | 0.0832 | 0.5270 | 0.0660 | **0.5450** | 0.0434 | Additionally, we would like to highlight the potential for fine-tuning in our approach. When training a model at each step, it is possible to use the previous model or a model already trained in a closely related context as a starting point. This can significantly reduce the number of episodes required for training a new model, thereby improving efficiency and having the potential to be scalable. We are grateful for your constructive feedback and remain open to further discussions on these important points during the rebuttal period.

Reviewer wJEK2024-08-13

Thank you for your detailed response!

I wanted to thank author first for their detailed response. The response has addressed most of my questions and I would like to keep my original score to recommend accept this paper.

Authorsrebuttal2024-08-13

Thank you for your feedback and suggestions during the rebuttal period. It really helped us improve our work.

Reviewer ECSQ2024-08-14

Thank you for your detailed response to my review. The response has addressed most of my questions. However, I agree with **reviewer 21Qs** regarding the concern about the theoretical justification of the linear generalization gap. I am inclined to keep my original score.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC