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

2026-05-14

AI Summary

This work addresses the challenge in offline-to-online reinforcement learning where static data mixing strategies often induce distributional shift, compromising both training stability and asymptotic performance. The authors propose ROAD, a novel framework that formulates data mixing as a bilevel optimization problem: the outer loop optimizes policy performance by dynamically adjusting the mixing strategy via a multi-armed bandit mechanism, while the inner loop performs Q-learning updates augmented with surrogate gradient estimation, offline prior preservation, and overestimation suppression techniques. ROAD introduces a plug-and-play adaptive replay mechanism that automatically adapts to diverse environments without manual hyperparameter tuning. Empirical results demonstrate that ROAD significantly outperforms existing replay strategies across multiple datasets, achieving superior final performance while maintaining training stability.

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.

Problem

• offline-to-online reinforcement learning
• distribution shift
• data mixing
• replay strategy
• stability-performance tradeoff

Innovation

• bi-level optimization
• adaptive data mixing
• offline-to-online reinforcement learning
• multi-armed bandit
• surrogate objective