Papers
Topics
Authors
Recent
Search
2000 character limit reached

ROAD: Adaptive Data Mixing for Offline-to-Online Reinforcement Learning via Bi-Level Optimization

Published 14 May 2026 in cs.LG and cs.AI | (2605.14497v1)

Abstract: Offline-to-online reinforcement learning harnesses the stability of offline pretraining and the flexibility of online fine-tuning. A key challenge lies in the non-stationary distribution shift between offline datasets and the evolving online policy. Common approaches often rely on static mixing ratios or heuristic-based replay strategies, which lack adaptability to different environments and varying training dynamics, resulting in suboptimal tradeoff between stability and asymptotic performance. In this work, we propose Reinforcement Learning with Optimized Adaptive Data-mixing (ROAD), a dynamic plug-and-play framework that automates the data replay process. We identify a fundamental objective misalignment in existing approaches. To tackle this, we formulate the data selection problem as a bi-level optimization process, interpreting the data mixing strategy as a meta-decision governing the policy performance (outer-level) during online fine-tuning, while the conventional Q-learning updates operate at the inner level. To make it tractable, we propose a practical algorithm using a multi-armed bandit mechanism. This is guided by a surrogate objective approximating the bi-level gradient, which simultaneously maintains offline priors and prevents value overestimation. Our empirical results demonstrate that this approach consistently outperforms existing data replay methods across various datasets, eliminating the need for manual, context-specific adjustments while achieving superior stability and asymptotic performance.

Summary

  • The paper introduces a bi-level optimization formulation to automatically select data mixing ratios, effectively mitigating instability during the offline-to-online transition.
  • It proposes the ROAD algorithm, using a multi-armed bandit approach to dynamically balance offline and online replay for robust policy improvement.
  • Empirical results across various domains demonstrate that ROAD enhances asymptotic returns and training stability compared to static mixing methods.

Adaptive Data Mixing for Offline-to-Online RL via Bi-Level Optimization: A Technical Overview of ROAD


Motivation and Problem Statement

Offline-to-online (O2O) reinforcement learning is a scenario where an agent is first pretrained on a fixed dataset (offline phase) and then fine-tuned through direct online interaction (online phase). This paradigm is motivated by the need for stable and safe policy initialization while ultimately seeking higher asymptotic performance than what’s achievable from a static dataset alone. However, a fundamental challenge emerges: the non-stationary distribution shift between the offline data and the evolving online policy induces severe instability, including performance dips and "unlearning" during the transition to online fine-tuning.

Standard approaches typically employ static or heuristic-based mixing between offline and online replay buffers, lacking context sensitivity and adaptivity. These approaches suffer from an objective misalignment—minimizing Bellman error on the mixed dataset is not guaranteed to optimize the ultimate policy return due to the dependency between the mixing ratio and the loss landscape throughout training.

The core questions are: How should data from offline and online sources be mixed during fine-tuning? And, can this mixing process be made adaptive to environment, data quality, and training phase?


Bi-Level Optimization Formulation

The paper formulates the data-mixing problem as a functional bi-level optimization:

  • Inner Level: Standard Q-learning/Fitted Q Iteration applied to a mixture distribution over state-action pairs, parameterized by a data mixing strategy gg.
  • Outer Level: Aims to select the mixing strategy gg^* that maximizes expected policy performance, i.e., maxgGJ(π(f(g)))\max_{g \in \mathcal{G}} J(\pi(f^*(g))), where π\pi is the policy induced from the Q-function ff^* optimized at the inner level.

Direct gradient computation of the outer objective is generally intractable in realistic online RL settings due to high variance and the need for density estimation under fast-changing policies. The paper provides a theoretical treatment of this gradient, showing that the optimal data mixing strategy should function as a meta-decision variable aligning the geometry of the Q-learning loss landscape with policy performance maximization.


The ROAD Algorithm

To overcome intractability, the authors propose Reinforcement Learning with Optimized Adaptive Data-mixing (ROAD): a pragmatic, plug-and-play framework employing a multi-armed bandit (MAB) mechanism. ROAD discretizes mixing ratios, frames their selection as a bandit problem, and leverages a surrogate objective, RqR_q, guided by theoretical insights but grounded in in-training statistics. Specifically, ROAD:

  • Restricts the candidate mixing ratios to a finite set (e.g., Λ={0.1,0.2,0.3,0.4,0.5}\Lambda=\{0.1, 0.2, 0.3, 0.4, 0.5\}), representing the proportion of offline replay.
  • Uses a surrogate objective for each candidate at each period:

Rq=ΔoffκΔonR_q = \Delta_{\mathrm{off}} - \kappa \Delta_{\mathrm{on}}

Where Δoff\Delta_{\mathrm{off}} is a proxy for policy improvement in the offline region and Δon\Delta_{\mathrm{on}} penalizes overestimation or instability due to unreliable online estimates.

  • Selects mixing ratios adaptively via UCB over a sliding window, balancing exploration and exploitation in response to online training feedback.

ROAD’s online training scheme is depicted in (Figure 1): Figure 1

Figure 1: ROAD controls the offline data replay ratio using a bi-level optimization-derived surrogate, dynamically mixing offline and online samples during fine-tuning.


Theoretical Underpinnings

A key insight is that the optimal data mixing depends on the dynamics of the training process. Too much emphasis on offline data anchors early learning, mitigating overestimation and catastrophic forgetting. Later, more aggressive use of online samples is necessary for performance gains beyond dataset-imposed limitations. The surrogate objective gg^*0 is theoretically justified as follows:

  • gg^*1 reduces variance through an advantage-like baseline, normalizing for state-dependent value variations.
  • The penalty term gg^*2 is closely related to the signal-to-noise ratio of Q-estimates in areas not well-covered by offline data, hence suppresses harmful overestimation.
  • The algorithm's sliding-window bandit formulation enables responsiveness to temporal non-stationarity in environment or policy.

Empirical Results

Comprehensive benchmarking evaluates ROAD on domains including Antmaze, MuJoCo Locomotion, and FrankaKitchen, spanning diverse state-action space complexities and dataset qualities. The main numerical findings are:

  • ROAD consistently outperforms all static and heuristic data replay strategies across environments, both in terms of asymptotic policy return and training stability.
  • ROAD maintains high performance regardless of offline data quality, a regime where fixed mixing ratios or ad-hoc decrease schedules often perform suboptimally.
  • ROAD also generalizes robustly to different backbone algorithms (IQL, PEX, CQL, Cal-QL), always being at least as strong as the best static tuning.

Visualization of ROAD’s adaptivity is shown via mixing ratio heatmaps, revealing that in some scenarios ROAD converges to a near-optimal fixed ratio, while in others it dynamically shifts its strategy in response to data and policy evolution (Figure 2): Figure 2

Figure 2

Figure 2

Figure 2

Figure 2

Figure 2

Figure 2: ROAD adaptively selects different data mixing patterns, converging to fixed ratios when optimal or adapting to temporal training needs across environments.

The gg^*3 dynamics during fine-tuning (Figure 3) highlight stable growth in policy quality, with effective suppression of instability at critical early stages. Figure 3

Figure 3: The gg^*4 surrogate objective exhibits informative dynamics, indicating phases of improvement and stability regulation.


Ablation Studies

A series of ablations on the hyperparameter gg^*5, the window size gg^*6, and the exploration parameter gg^*7 demonstrate:

  • ROAD’s performance is robust to a broad range of hyperparameters, with only mild sensitivity when gg^*8 is pushed to extreme values, confirming the stability of the surrogate objective.
  • The algorithm remains effective across varying sliding-window sizes and exploration-exploitation tradeoffs, supporting its practical deployability.

Analysis of Data Quality and Task Regimes

A critical contribution is empirical evidence that no universal mixing strategy exists: optimal mixing is dataset-, environment-, and phase-dependent. ROAD's ability to automatically adapt mitigates the risk of performance collapse in regimes where offline data is of poor, diverse, or unknown quality. Further diagnostic experiments on simple MDPs show the nuanced impact of dataset bias and mixing on policy improvement and unlearning (Figure 4): Figure 4

Figure 4

Figure 4

Figure 4

Figure 4: Example illustrates the relationship between data quality, mixing patterns, and policy efficiency across different phases of offline-to-online RL.


Generalization and Implications

ROAD is readily composable with a variety of O2O RL algorithms and task domains, offering a unified adaptive replay mechanism that:

  • Maintains training stability under severe distribution shift
  • Accelerates asymptotic convergence, unlocking the full potential of policy fine-tuning
  • Reduces the need for manual, context-dependent engineering in data management

Its theoretical insights—formulating data replay as a bi-level meta-optimization and exploiting practical surrogates—invite avenues for further research, such as extending data selection granularity beyond buffer-level mixing or treating more expressive, possibly continuous, mixing strategies.


Limitations and Future Directions

ROAD's function space is restricted by its discretization of candidate mixing strategies and by the linearity assumption. For environments or tasks requiring highly non-linear, state- or transition-dependent mixing, richer function classes and continuous optimization techniques could be explored. Additionally, the reliance on trustworthy offline coverage cautions against deployment in settings where offline data are extremely sparse or highly biased.

Potential future work includes:

  • Sample-level adaptive replay strategies, leveraging density ratio estimation or importance weighting
  • Exploring analytical connections between ROAD’s surrogate and theoretical RL generalization bounds
  • Deployment on real-world robotic and safety-critical domains emphasizing both stability and performance

Conclusion

"ROAD: Adaptive Data Mixing for Offline-to-Online Reinforcement Learning via Bi-Level Optimization" introduces a theoretically motivated, empirically validated framework that automates the replay buffer mixing process in O2O RL, replacing static or heuristic choices with an adaptive, meta-optimized strategy. ROAD’s consistent improvements in stability, asymptotic return, and algorithmic flexibility position it as an effective and general plug-in for the next generation of data-driven RL agents.


Figure 5

Figure 5

Figure 5: ROAD’s normalized returns are superior to baseline methods when using Cal-QL as the offline-to-online backbone.

Figure 6

Figure 6

Figure 6: With CQL backbone, ROAD maintains its performance edge across tasks and data regimes.

Paper to Video (Beta)

No one has generated a video about this paper yet.

Whiteboard

No one has generated a whiteboard explanation for this paper yet.

Open Problems

We haven't generated a list of open problems mentioned in this paper yet.

Collections

Sign up for free to add this paper to one or more collections.

Tweets

Sign up for free to view the 2 tweets with 26 likes about this paper.