Stable Offline Value Function Learning with Bisimulation-based Representations

2024-10-02

arXiv.org

AI Summary

This paper addresses the instability and divergence issues in value function learning within offline reinforcement learning. We propose KROPE, a bisimulation-based kernel representation method that constrains state-action representations to jointly preserve reward similarity and next-state dynamics similarity, thereby implicitly enforcing contraction of the Bellman operator in the representation space. We provide the first theoretical proof that state-action representations satisfying the bisimulation metric guarantee convergence of offline value function learning. Building on this insight, we design the KROPE algorithm to jointly optimize representation learning and value estimation in an end-to-end manner. Experiments demonstrate that KROPE significantly reduces Bellman error and policy evaluation bias, achieving superior stability and accuracy over leading offline RL methods across multiple benchmarks.

Technology Category

• Machine Learning: Online Learning & Bandits
• Search and Optimization: Sampling/Simulation-based Search
• Intelligent Robots: State Estimation

Application Category

• User Modeling, Personalization and Recommendation: User modeling and simulation for interactive and conversational systems
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

Abstract

In reinforcement learning, offline value function learning is the procedure of using an offline dataset to estimate the expected discounted return from each state when taking actions according to a fixed target policy. The stability of this procedure, i.e., whether it converges to its fixed-point, critically depends on the representations of the state-action pairs. Poorly learned representations can make value function learning unstable, or even divergent. Therefore, it is critical to stabilize value function learning by explicitly shaping the state-action representations. Recently, the class of bisimulation-based algorithms have shown promise in shaping representations for control. However, it is still unclear if this class of methods can stabilize value function learning. In this work, we investigate this question and answer it affirmatively. We introduce a bisimulation-based algorithm called kernel representations for offline policy evaluation (KROPE). KROPE uses a kernel to shape state-action representations such that state-action pairs that have similar immediate rewards and lead to similar next state-action pairs under the target policy also have similar representations. We show that KROPE: 1) learns stable representations and 2) leads to lower value error than baselines. Our analysis provides new theoretical insight into the stability properties of bisimulation-based methods and suggests that practitioners can use these methods for stable and accurate evaluation of offline reinforcement learning agents.

Problem

• Stable Offline Learning
• Reinforcement Learning
• Value Function

Innovation

• KROPE algorithm
• representation learning optimization
• offline reinforcement learning