Learning Bellman Complete Representations for Offline Policy Evaluation

TL;DR

This paper introduces BCRL, a method that learns Bellman complete representations for offline policy evaluation, improving accuracy and efficiency in complex, image-based continuous control tasks.

Contribution

The paper proposes a novel end-to-end algorithm that learns Bellman complete representations directly from data for offline RL, with theoretical guarantees and empirical validation.

Findings

BCRL outperforms previous representation learning methods like CURL and SPR in OPE tasks.

BCRL achieves competitive error with FQE and surpasses it in out-of-distribution evaluations.

Both Bellman completeness and coverage are essential for the method's success.

Abstract

We study representation learning for Offline Reinforcement Learning (RL), focusing on the important task of Offline Policy Evaluation (OPE). Recent work shows that, in contrast to supervised learning, realizability of the Q-function is not enough for learning it. Two sufficient conditions for sample-efficient OPE are Bellman completeness and coverage. Prior work often assumes that representations satisfying these conditions are given, with results being mostly theoretical in nature. In this work, we propose BCRL, which directly learns from data an approximately linear Bellman complete representation with good coverage. With this learned representation, we perform OPE using Least Square Policy Evaluation (LSPE) with linear functions in our learned representation. We present an end-to-end theoretical analysis, showing that our two-stage algorithm enjoys polynomial sample complexity provided some representation in the rich class considered is linear Bellman complete. Empirically, we extensively evaluate our algorithm on challenging, image-based continuous control tasks from the Deepmind Control Suite. We show our representation enables better OPE compared to previous representation learning methods developed for off-policy RL (e.g., CURL, SPR). BCRL achieve competitive OPE error with the state-of-the-art method Fitted Q-Evaluation (FQE), and beats FQE when evaluating beyond the initial state distribution. Our ablations show that both linear Bellman complete and coverage components of our method are crucial.