Ensemble Successor Representations for Task Generalization in Offline-to-Online Reinforcement Learning

TL;DR

The paper tackles reward generalization in offline-to-online RL where offline data come from a single task and online adaptation must reach new tasks with the same dynamics but different rewards. It introduces ESR-O2O, an ensemble-based framework that learns multiple successor representations $\psi_k$ and ensemble critics $Q_k$ from offline data, and uses online fine-tuning with fixed representations while optimizing a policy that maximizes the minimum ensemble value: $\pi(\cdot|s)=\arg\max_a \min_k Q_k(s,a)$. A theoretical bound on the online fine-tuning sub-optimality is presented: $|Q_j^{\pi_j^*}- Q_j^{\pi}| \leq \|w_j\|_{\infty}\|\psi^{\pi_j^*}-\hat{\psi}\|_1 + \|\hat{\psi}\|_{\infty}\|w_j-w^{\pi}\|_1$, highlighting the roles of representation and reward approximation errors. Empirically, ESR-O2O achieves superior generalization across diverse tasks (e.g., Quadruped, Walker, Reach, MetaWorld, AntMaze), especially under large reward gaps, while showing robustness to offline data coverage and providing meaningful ablations on ensemble components.

Abstract

In Reinforcement Learning (RL), training a policy from scratch with online experiences can be inefficient because of the difficulties in exploration. Recently, offline RL provides a promising solution by giving an initialized offline policy, which can be refined through online interactions. However, existing approaches primarily perform offline and online learning in the same task, without considering the task generalization problem in offline-to-online adaptation. In real-world applications, it is common that we only have an offline dataset from a specific task while aiming for fast online-adaptation for several tasks. To address this problem, our work builds upon the investigation of successor representations for task generalization in online RL and extends the framework to incorporate offline-to-online learning. We demonstrate that the conventional paradigm using successor features cannot effectively utilize offline data and improve the performance for the new task by online fine-tuning. To mitigate this, we introduce a novel methodology that leverages offline data to acquire an ensemble of successor representations and subsequently constructs ensemble Q functions. This approach enables robust representation learning from datasets with different coverage and facilitates fast adaption of Q functions towards new tasks during the online fine-tuning phase. Extensive empirical evaluations provide compelling evidence showcasing the superior performance of our method in generalizing to diverse or even unseen tasks.

Figures (16)

• Figure 1: Offline pre-training in a four-room maze navigation task. In the left picture, the blue block indicates the starting point, and the green one indicates the goal. Two middle pictures describe the values and policies learned by IQL from offline data, where the arrow represents the action to be taken at each position. The two right pictures show the values and policies learned by vanilla SR based methods. After offline training, both IQL and SR-based methods can learn correct values and the optimal policy.
• Figure 2: Online fine-tuning in a navigation task with a different goal. When the goal changes, the reward function changes accordingly. In this setting, IQL fails to learn correct values and policies in the new environment, while SR based methods re-learns of the values and leads to effective policies.
• Figure 3: Pre-training and Fine-tuning performance of vanilla SR based methods under different data distributions. The black dashed line indicates the transition from offline pre-training to online fine-tuning when the task changes. The shaded area represent the variance of performance across multiple experiments conducted with 5 random seeds. Representations learned from more diverse datasets, such as 'medium-replay' and 'replay' data, exhibit superior performance compared to representations learned from narrower datasets ('medium', 'expert').
• Figure 4: Frameworks of vanilla SR-based methods and ESR-O2O for offline-to-online learning. Red dashed boxes indicate the different parts of ESR-O2O from vanilla SR. While vanilla SR trains multiple representations from source environments, ESR-O2O extract ensemble representations using offline datasets from a single environment. Another difference lies in the construction of Q function, where vanilla SR considers linear composition, and ESR-O2O incorporates ensemble Q functions. Other parts are the same, including fine-tuning the value function with online interactions with $M_j$, which follows a greedy policy $\pi(\cdot|s)=\arg\max Q(s,a)$.
• Figure 5: Experimental environments, including Quadruped, Walker, Reach from UTDS utds2023, MetaWorld meta2020, and Antmaze from D4RL d4rl2020. These environments can be categorized into three classes based on the magnitude of the reward gap: small gap (Quadruped, Walker), big gap (Reach, MetaWorld), and no gap (Antmaze).

Theorems & Definitions (2)

• Proposition 1: Sub-optimality Gap.
• Lemma 1: psi_lemma2021w_lemma2011