Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Offline Meta-Reinforcement Learning with Flow-Based Task Inference and Adaptive Correction of Feature Overgeneralization

Min Wang, Xin Li, Mingzhong Wang, Hasnaa Bennis

TL;DR

FLORA tackles two core challenges in offline meta-reinforcement learning: inferring complex, multimodal task distributions and mitigating extrapolation errors caused by feature overgeneralization. It introduces flow-based task inference (FTI) to capture rich task structures and adaptive correction of overgeneralization (ACO) to regulate decomposed Q-values through epistemic uncertainty and return-based feedback, all within a decoupled dynamics-reward representation $Q^{\pi}(s,a,\mathbf{z})=\psi^{\pi}(s,a,\mathbf{z})^\mathsf{T}W(\mathbf{z})$. The method uses a chain of invertible flows for task representation, double-feature learners for uncertainty estimation, a distributional TD target, and a KL-regularized policy objective to guarantee improvement; these components collectively yield rapid adaptation and superior meta-policy performance. Empirically, FLORA outperforms strong baselines on challenging Meta-World and MuJoCo tasks, with ablations confirming the critical roles of FTI and ACO in handling OOD actions and maintaining stable learning.

Abstract

Offline meta-reinforcement learning (OMRL) combines the strengths of learning from diverse datasets in offline RL with the adaptability to new tasks of meta-RL, promising safe and efficient knowledge acquisition by RL agents. However, OMRL still suffers extrapolation errors due to out-of-distribution (OOD) actions, compromised by broad task distributions and Markov Decision Process (MDP) ambiguity in meta-RL setups. Existing research indicates that the generalization of the $Q$ network affects the extrapolation error in offline RL. This paper investigates this relationship by decomposing the $Q$ value into feature and weight components, observing that while decomposition enhances adaptability and convergence in the case of high-quality data, it often leads to policy degeneration or collapse in complex tasks. We observe that decomposed $Q$ values introduce a large estimation bias when the feature encounters OOD samples, a phenomenon we term ''feature overgeneralization''. To address this issue, we propose FLORA, which identifies OOD samples by modeling feature distributions and estimating their uncertainties. FLORA integrates a return feedback mechanism to adaptively adjust feature components. Furthermore, to learn precise task representations, FLORA explicitly models the complex task distribution using a chain of invertible transformations. We theoretically and empirically demonstrate that FLORA achieves rapid adaptation and meta-policy improvement compared to baselines across various environments.

Offline Meta-Reinforcement Learning with Flow-Based Task Inference and Adaptive Correction of Feature Overgeneralization

TL;DR

FLORA tackles two core challenges in offline meta-reinforcement learning: inferring complex, multimodal task distributions and mitigating extrapolation errors caused by feature overgeneralization. It introduces flow-based task inference (FTI) to capture rich task structures and adaptive correction of overgeneralization (ACO) to regulate decomposed Q-values through epistemic uncertainty and return-based feedback, all within a decoupled dynamics-reward representation . The method uses a chain of invertible flows for task representation, double-feature learners for uncertainty estimation, a distributional TD target, and a KL-regularized policy objective to guarantee improvement; these components collectively yield rapid adaptation and superior meta-policy performance. Empirically, FLORA outperforms strong baselines on challenging Meta-World and MuJoCo tasks, with ablations confirming the critical roles of FTI and ACO in handling OOD actions and maintaining stable learning.

Abstract

Offline meta-reinforcement learning (OMRL) combines the strengths of learning from diverse datasets in offline RL with the adaptability to new tasks of meta-RL, promising safe and efficient knowledge acquisition by RL agents. However, OMRL still suffers extrapolation errors due to out-of-distribution (OOD) actions, compromised by broad task distributions and Markov Decision Process (MDP) ambiguity in meta-RL setups. Existing research indicates that the generalization of the network affects the extrapolation error in offline RL. This paper investigates this relationship by decomposing the value into feature and weight components, observing that while decomposition enhances adaptability and convergence in the case of high-quality data, it often leads to policy degeneration or collapse in complex tasks. We observe that decomposed values introduce a large estimation bias when the feature encounters OOD samples, a phenomenon we term ''feature overgeneralization''. To address this issue, we propose FLORA, which identifies OOD samples by modeling feature distributions and estimating their uncertainties. FLORA integrates a return feedback mechanism to adaptively adjust feature components. Furthermore, to learn precise task representations, FLORA explicitly models the complex task distribution using a chain of invertible transformations. We theoretically and empirically demonstrate that FLORA achieves rapid adaptation and meta-policy improvement compared to baselines across various environments.
Paper Structure (32 sections, 4 theorems, 39 equations, 9 figures, 5 tables, 1 algorithm)

This paper contains 32 sections, 4 theorems, 39 equations, 9 figures, 5 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Offline Meta-Reinforcement Learning
  4. Successor Features
  5. Problem Formulation
  6. Methodology
  7. Overgeneralization Problem
  8. Adaptive Correction of Overgeneralization (ACO)
  9. Flow-based Task Inference (FTI)
  10. Experiment
  11. Experimental Settings
  12. Setups
  13. Offline Data Collection
  14. Baselines
  15. Performance Comparison
  16. ...and 17 more sections

Key Result

Theorem 1

(Policy Superiority Guarantee) Consider a meta-task $\mathcal{M}_i$ with an optimal policy $\pi^{*}$ whose action-value is $Q_{i}^{\pi^{*}}$. Let $Q_{i}^{\pi_j^{*}}$ be the action-value of an optimal policy of $\mathcal{M}_j$ when performed on $\mathcal{M}_i$. Given the set $\{\hat{Q}_{i}^{\pi_1^{*}

Figures (9)

  • Figure 1: Motivating example: In environments with high-quality datasets or narrow task distributions (e.g. Drawer-Close), the decomposed $Q$ value accelerates the adaptability of the training policy $\pi_\varphi$, and the converged optimal policy outperforms the behavior policy $\pi_\beta$. However, in environments with low-quality datasets or broad task distributions (e.g. Door-Close), the decomposed $Q$ value exacerbates the overestimation issue, ultimately leading to divergence of $Q$ value and failure of $\pi_\varphi$.
  • Figure 2: Testing average performance on 8 ML1 environments over 6 random seeds.
  • Figure 3: Testing average performance of FLORA and baselines on MuJoCo over 6 random seeds.
  • Figure 4: Ablation study.
  • Figure 5: Visualization of sparse 2D navigation.
  • ...and 4 more figures

Theorems & Definitions (7)

  • Theorem 1
  • Lemma 1
  • proof
  • Lemma 2
  • proof
  • Theorem 2
  • proof