Model-Based Transfer Learning for Contextual Reinforcement Learning

Jung-Hoon Cho; Vindula Jayawardana; Sirui Li; Cathy Wu

Source-linked AI summary

Model-Based Transfer Learning for Contextual Reinforcement Learning

Jung-Hoon Cho, Vindula Jayawardana, Sirui Li, Cathy Wu

arXiv:2408.04498v4Aug 8, 2024cs.LG

TL;DR

Deep RL is brittle under task variation, making it important to select a small set of costly training tasks that generalizes across contexts. MBTL models source-task performance with Gaussian processes and the generalization gap with contextual similarity inside Bayesian optimization. Across traffic and continuous-control benchmarks, it achieves up to 43x improved sample efficiency and has sublinear regret in the number of training tasks.

Problem
Deep RL can be brittle under small environment changes, while selecting training tasks for generalization across contexts remains poorly understood and costly.
Method
MBTL combines Gaussian-process estimates of training performance with a linear contextual-similarity model of generalization gap in Bayesian optimization.
Results
MBTL achieves up to 43x improved sample efficiency and sublinear regret in the number of training tasks across reported benchmarks.
Takeaways & Limitations
Explicitly modeling generalization enables principled selection of an intermediate number of source tasks for contextual RL.
Takeaways & Limitations
MBTL is designed for single-dimensional context variation and relies on explicit similarity between context variables.

Abstract

from arXiv · show

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.

1 Introduction

Contextual RL aims to generalize across task variations, but deep RL can be brittle and existing training strategies trade off compute efficiency against negative transfer and instability. MBTL addresses this by selecting an intermediate set of source tasks while explicitly modeling generalization performance.

Deep RL can fail under small task variations, motivating contextual MDPs for parameterizing changes such as lanes, weather, or traffic flow.
Independent training is compute-intensive, whereas multi-task training is efficient but can suffer from limited model capacity and negative transfer.
Multi-policy training balances these extremes by training an intermediate number K of models, where N > K > 1.
MBTL strategically selects source tasks by explicitly modeling generalization performance and estimating the value of training each new source task.
The paper focuses on training reliability and within-domain generalization rather than model robustness or across-domain generalization.
MBTL combines a Gaussian-process acquisition function with theoretical regret analysis and empirical evaluation across urban traffic and continuous-control benchmarks.The reported experiments observe up to 43x improvements in sample efficiency.

2 Preliminaries and notation

The preliminaries define contextual MDPs, source-task performance, and zero-shot transfer across contexts. They frame generalization loss as the performance degradation incurred when a policy is transferred to a different task.

A contextual MDP is a collection of MDPs parameterized by a context variable that can affect dynamics, rewards, and initial-state distributions.
MBTL’s overview estimates training performance with Gaussian processes, combines it with transfer-gap information, and selects the next task through Bayesian optimization.
Source-task performance J(πx, x) is the expected return of a policy trained and evaluated on the same context using a specified RL algorithm.
Zero-shot transfer applies a policy trained on source context x to target context x′ without adaptation, often causing performance degradation called the generalization gap.
The generalization gap increases as the target context diverges from the source context, motivating strategic selection of source tasks.

3 Problem formulation

The sequential source task selection problem chooses distinct training contexts one at a time to maximize cumulative generalization performance across target contexts. Its state records current best performance, while training randomness is simplified through deterministic-transition assumptions.

Sequential source task selection chooses a distinct source task xk at each stage to maximize expected cumulative generalization performance across the contextual MDP.
The selection process uses feedback from observed performance on previously selected and trained tasks.
After training task xk, performance on each target context is updated to the better result from the new policy and earlier policies.
At each stage, the state contains the best known generalization performance for every target task, and the action is choosing a new task xk.
The formulation assumes deterministic transitions, so training a context always yields the same source performance and generalization gap; selection may terminate early under performance or budget conditions.

4 Model-Based Transfer Learning (MBTL)

MBTL uses Bayesian optimization to select source tasks by combining a Gaussian-process performance estimate with a similarity-based generalization-gap model. Its analysis establishes sublinear regret and shows that eliminating less promising regions can tighten the bound under stated assumptions.

4 Model-Based Transfer Learning (MBTL): MBTL combines Gaussian-process estimates of source-task performance with a linear contextual-similarity model of generalization loss inside Bayesian optimization.The acquisition function sequentially selects training tasks using the estimated performance set point and generalization gap.
4 Model-Based Transfer Learning (MBTL): The generalization-gap model assumes performance loss grows linearly with source-target context distance, simplifying analysis while relying on smooth task variation.Relaxing this assumption could improve efficiency but would increase modeling complexity.
4.3 Regret analysis: MBTL’s cumulative regret is bounded by √KC1βKγK with probability at least 1 −δ, establishing sublinear regret under the specified UCB scaling.The result implies that average performance approaches the optimum as the number of selection iterations increases.
4.3 Regret analysis: Restricting the search space to promising regions yields tighter or equivalent regret bounds than searching the full task space.The bound tightens when the reduced search space is smaller, provided task elimination is justified by prior selections and modeled transfer performance.
4.3 Regret analysis: If the search space shrinks geometrically as |Xk| ≤2−⌊log2 k⌋|X|, MBTL obtains a corresponding specialized regret bound.The paper links this reduction to progressively eliminating partitions of the task space when selected source policies sufficiently address remaining targets.

5 Experiments and analysis

Experiments evaluate MBTL on urban traffic and continuous-control CMDPs, comparing strategic source-task selection with canonical and heuristic multi-policy baselines. MBTL consistently reduces training requirements while remaining effective across algorithms and acquisition functions.

Traffic benchmark experiments: Up to 30x improved sample efficiency lets MBTL match or exceed independent and multi-task baselines on traffic CMDPs.MBTL rapidly approaches the Oracle within approximately 10 transfer steps, while independent training often outperforms multi-task training.
Traffic benchmark experiments: Multi-policy baselines generally outperform independent and multi-task training, while MBTL often surpasses heuristic random and greedy source-task selection.These comparisons evaluate both the value of training an intermediate number of models and the value of Bayesian-optimization-based selection.
Continuous control benchmark experiments: Up to 43x fewer samples achieve performance comparable to independent or multi-task baselines on continuous-control CMDPs.MBTL generally shows superior performance across context variations, with exceptions, and often closely trails the Oracle with few trained policies.
Sensitivity analysis: MBTL remains effective with DQN, PPO, and A2C as underlying deep-RL algorithms.The sensitivity analysis tests these algorithms on Cartpole with varying pole lengths.
Sensitivity analysis: MBTL is not particularly sensitive to the acquisition function’s optimism representation, indicating weak dependence on hyperparameters.Expected improvement emphasizes marginal gains, whereas UCB balances exploration and exploitation using mean and variance.

6 Related work

Related work frames contextual reinforcement learning through CMDPs, multi-task training, zero-shot transfer, and source-task selection. Existing approaches address shared structure and transfer, but face challenges including unobserved contexts, negative transfer, and selecting suitable sources.

Contextual Reinforcement Learning: CMDPs formalize contextual-RL generalization by incorporating context-dependent dynamics, rewards, and initial-state distributions.
Multi-task training: Multi-task methods exploit shared structure through policy decomposition and distillation, but unobserved contexts create partial observability and negative transfer challenges.
Zero-shot transfer and policy reuse: Zero-shot transfer applies policies directly to unseen environments without additional training and has been combined with approximations, hypernetworks, and metadata.
Source task selection: Prior source-task selection work emphasizes choosing relevant or structurally aligned tasks to avoid degradation and improve transfer.

7 Conclusion

The paper presents MBTL as an intermediate-model strategy for CMDPs that explicitly models generalization and selects source tasks through Bayesian optimization. It reports strong sample efficiency and sublinear regret, while limiting the current method to one-dimensional context variation with explicit similarity.

Conclusion: MBTL trains an intermediate number of models using a Gaussian process for training performance and an explicit model of the generalization gap.It layers on existing RL methods rather than using independent or single universal-policy training.
Conclusion: MBTL achieves up to 43x improved sample efficiency on standard and real-world benchmarks and sublinear regret in the number of training tasks.
Conclusion: The current method is designed for single-dimensional context variation and relies on explicit context-variable similarity.Future directions include high-dimensional contexts, formalized task similarity, and new real-world CMDP benchmarks.

A.1 Notation

This notation section introduces the paper’s notation and illustrates how observed and predicted generalized performance, together with marginal improvement, support source-task selection.

A.1 Notation: Table 3 defines the notation used in the problem formulation.
A.1 Notation: Figure 10 contrasts observed generalized performance J with predicted generalized performance ˆJ after zero-shot transfer from x1 to x′.
A.1 Notation: Figure 11 illustrates selecting x2 by maximizing estimated marginal improvement, ˆV(x; π1) − V(x1).The estimated value corresponds to the red area under the predicted curve, while V(x1) is the area under J(π1, x′).

A.2 Model-Based Transfer Learning (MBTL) Algorithm

MBTL selects training tasks strategically for contextual reinforcement learning by estimating performance and using an acquisition function to guide Bayesian optimization.

Algorithm: MBTL initializes performance estimates and an empty policy set across the task context set before sequential task selection.The procedure begins with J and V set to zero, an empty π, and iteration counter k = 1.
Algorithm: At each iteration, MBTL uses a Gaussian-process estimate of training performance to calculate marginal generalized performance and an acquisition function.
Algorithm: The algorithm selects the next task by maximizing the acquisition function, trains a policy on that task, and adds the policy to the collection.
Algorithm: After exhausting the training budget, MBTL performs zero-shot transfer and calculates generalization performance.

A.3 Theoretical analysis

The theoretical analysis develops regret bounds for MBTL by adapting existing results to iterations whose search space is restricted.

Theoretical analysis: The section provides detailed proofs for the regret bounds stated in Theorem 2, Corollary 2.1, and Corollary 2.2.
Theoretical analysis: The analysis adapts key results from to settings where the search space is restricted at each iteration.
Theoretical analysis: The proof framework therefore addresses regret analysis when task selection operates over iteration-specific restricted search spaces.

A.3.1 Proof of Theorem 2

The proof establishes a high-probability cumulative-regret bound for MBTL using lemmas that handle restricted search spaces and Gaussian-process estimation.

Proof of Theorem 2: Theorem 2 bounds cumulative regret RK under a confidence parameter δ′ and scaling factor βk defined from the task-set size and iteration index.
Proof of Theorem 2: The proof extends results from through lemmas designed for restricted search spaces Xk ⊆ X at each iteration.
Proof of Theorem 2: Instantaneous regret is bounded using the size of the restricted search space and the Gaussian-process confidence parameter β.
Proof of Theorem 2: With probability at least 1 −δ′, the cumulative bound follows by combining Gaussian-process deviation control, restricted-space regret summation, and the Cauchy–Schwarz inequality.

A.3.2 Proof of Corollary 2.1

Across traffic, advisory-autonomy, Cartpole, and control benchmarks, MBTL generally adapts well to context variation and often outperforms standard baselines. Transfer performance improves with selected source tasks, although the strength of transfer depends on the task and variation.

Traffic signal control: In traffic control, MBTL improves with additional transfer steps and performs especially well for inflow and road-length variations.Performance is relatively stable for speed-limit variation, while MBTL shows the clearest upward trends for inflow and road length.
Scope of results: MBTL’s advantage is not uniform: transfer quality varies by domain and parameter, with some variations showing weak or unstable adaptation.The reported results note task-specific differences in effectiveness, including stronger and weaker transfer patterns across physical and environmental changes.
Baseline comparison: Oracle Transfer remains strongest in several comparisons, while independent and multi-task training show larger drops, instability, or variance under context changes.These patterns appear in traffic, advisory autonomy, Cartpole, and BipedalWalker, where universal or separately trained policies can be less stable across variations.
Benchmark results: MBTL generally achieves the strongest normalized generalized performance across traffic, advisory-autonomy, Cartpole, and several continuous-control variations.The reported comparisons identify MBTL as consistently superior or generally outperforming other strategies across these benchmark families.
Transferability patterns: Transferability decreases when source and target contexts diverge, including larger cart-mass differences in Cartpole and changing traffic conditions.The heatmaps characterize how task-parameter differences affect source-to-target transfer effectiveness.