[Paper Review] When Data Geometry Meets Deep Function: Generalizing Offline Reinforcement Learning

Jianxiong Li, Xianyuan Zhan+4 more|arXiv (Cornell University)|May 23, 2022

Reinforcement Learning in Robotics4 citations

TL;DR

This paper proposes DOGE, a novel offline reinforcement learning method that leverages dataset geometry through a state-conditioned distance function to improve generalization beyond the data distribution. By constraining policies based on geometric proximity to training data, DOGE enables safer extrapolation with deep Q functions, achieving state-of-the-art performance on D4RL benchmarks while reducing over-conservatism compared to prior methods.

ABSTRACT

In offline reinforcement learning (RL), one detrimental issue to policy learning is the error accumulation of deep Q function in out-of-distribution (OOD) areas. Unfortunately, existing offline RL methods are often over-conservative, inevitably hurting generalization performance outside data distribution. In our study, one interesting observation is that deep Q functions approximate well inside the convex hull of training data. Inspired by this, we propose a new method, DOGE (Distance-sensitive Offline RL with better GEneralization). DOGE marries dataset geometry with deep function approximators in offline RL, and enables exploitation in generalizable OOD areas rather than strictly constraining policy within data distribution. Specifically, DOGE trains a state-conditioned distance function that can be readily plugged into standard actor-critic methods as a policy constraint. Simple yet elegant, our algorithm enjoys better generalization compared to state-of-the-art methods on D4RL benchmarks. Theoretical analysis demonstrates the superiority of our approach to existing methods that are solely based on data distribution or support constraints.

Motivation & Objective

• Address the over-conservatism of existing offline RL methods that restrict policies within data distribution, harming generalization.
• Investigate whether deep Q functions can generalize effectively in out-of-distribution (OOD) regions when guided by dataset geometry.
• Develop a method that exploits the interpolation capability of deep networks while avoiding unreliable extrapolation in OOD areas.
• Design a plug-in constraint mechanism that integrates geometric inductive bias into standard actor-critic frameworks.
• Demonstrate improved generalization performance on D4RL benchmarks without requiring online interaction or complex architectural changes.

Proposed method

• Propose a state-conditioned distance function that measures the geometric proximity of a state-action pair to the training dataset's convex hull.
• Use the distance function as a learnable, differentiable constraint in the actor-critic framework to regularize policy updates.
• Introduce a hyperparameter $ G $, representing the upper quantile of distances in a mini-batch, to control the extent of OOD exploration.
• Formulate a constrained optimization objective where the policy is penalized for actions that are too far from the training data, based on learned distance.
• Train the distance function end-to-end with the main RL components using a Lagrangian relaxation scheme to balance constraint enforcement and policy optimization.
• Apply the method to standard offline RL algorithms (e.g., SAC) with minimal architectural modifications, enabling plug-and-play integration.

Experimental results

Research questions

• RQ1 Can deep Q functions generalize reliably in out-of-distribution regions when guided by geometric proximity to the training data?
• RQ2 Does incorporating dataset geometry into the policy constraint improve generalization in offline RL without sacrificing sample efficiency or safety?
• RQ3 How does the choice of distance threshold (via $ G $) affect the trade-off between conservatism and performance on OOD generalization?
• RQ4 Can a simple, geometry-aware constraint outperform existing methods that rely solely on data distribution or support constraints?
• RQ5 To what extent is the proposed method robust to hyperparameter variations in $ G $, $ eta $, and $ N $?

Key findings

• DOGE achieves state-of-the-art performance on the D4RL benchmark, outperforming SOTA methods such as SAC, TD3, and CQL across multiple Mujoco and AntMaze environments.
• The method significantly reduces over-conservatism: for example, on the hopper-m-r-v2 environment, DOGE achieves 76.2±17.7 with $ G=50 ext{th} $ quantile, outperforming more conservative baselines.
• Ablation studies show that $ G=50 ext{th} $ quantile provides a robust default setting, balancing exploration and safety across diverse tasks.
• Using $ G=30 ext{th} $ quantile leads to excessive constraint enforcement, causing $ heta o heta_{ ext{old}} $ and performance degradation due to suboptimality gap dominance.
• Using $ G=90 ext{th} $ or $ 100 ext{th} $ quantile increases risk of value overestimation and variance, leading to performance drops in high-variance tasks like hopper-m-r-v2 and walker2d-m-r-v2.
• The method is robust to changes in $ eta $ and $ N $, with minimal impact on performance, indicating strong hyperparameter stability.