[Paper Review] Q-value Regularized Transformer for Offline Reinforcement Learning

Shengchao Hu, Ziqing Fan+5 more|arXiv (Cornell University)|May 27, 2024

Elevator Systems and Control4 citations

TL;DR

This paper proposes QT, a Q-value regularized Transformer for offline reinforcement learning that combines conditional sequence modeling with dynamic programming via a dual-objective loss: behavior cloning for policy regularization and Q-value maximization for optimal action selection. Empirical results on D4RL benchmarks show QT outperforms both traditional offline RL methods and state-of-the-art CSM and DP approaches, achieving superior performance in long-horizon and sparse-reward settings by enabling robust trajectory stitching and improved return alignment.

ABSTRACT

Recent advancements in offline reinforcement learning (RL) have underscored the capabilities of Conditional Sequence Modeling (CSM), a paradigm that learns the action distribution based on history trajectory and target returns for each state. However, these methods often struggle with stitching together optimal trajectories from sub-optimal ones due to the inconsistency between the sampled returns within individual trajectories and the optimal returns across multiple trajectories. Fortunately, Dynamic Programming (DP) methods offer a solution by leveraging a value function to approximate optimal future returns for each state, while these techniques are prone to unstable learning behaviors, particularly in long-horizon and sparse-reward scenarios. Building upon these insights, we propose the Q-value regularized Transformer (QT), which combines the trajectory modeling ability of the Transformer with the predictability of optimal future returns from DP methods. QT learns an action-value function and integrates a term maximizing action-values into the training loss of CSM, which aims to seek optimal actions that align closely with the behavior policy. Empirical evaluations on D4RL benchmark datasets demonstrate the superiority of QT over traditional DP and CSM methods, highlighting the potential of QT to enhance the state-of-the-art in offline RL.

Motivation & Objective

• To address the limitation of Conditional Sequence Modeling (CSM) in offline RL, which fails to stitch optimal trajectories from sub-optimal ones due to inconsistency between sampled returns and optimal returns.
• To overcome the instability of Dynamic Programming (DP) methods in long-horizon and sparse-reward environments by integrating them with sequence modeling.
• To develop a unified framework that combines the trajectory modeling strength of Transformers with the optimal return prediction of Q-learning.
• To improve policy generalization and reduce distribution shift risk by aligning action sampling with both behavior policy distributions and learned Q-values.
• To achieve state-of-the-art performance in offline RL without requiring explicit environment interaction or additional behavior cloning components.

Proposed method

• The method employs a Transformer-based policy that predicts actions conditioned on a sequence of state-action-reward triplets and a return-to-go (RTG) token.
• It introduces a dual-objective training loss: a conditional behavior cloning term to match the behavior policy distribution and a Q-value regularization term to maximize expected returns.
• The Q-value regularization term is derived from a jointly trained Q-network that provides optimal future return estimates for each state-action pair.
• The model integrates Q-values directly into the action sampling process, biasing predictions toward high-reward actions while maintaining fidelity to the behavior policy.
• The framework avoids bootstrapping instability by relying on supervised learning on demonstrated trajectories, while still leveraging Bellman backup principles through the Q-value module.
• The approach enables effective trajectory synthesis by aligning expected returns of sampled actions with optimal returns across multiple trajectories.

Experimental results

Research questions

• RQ1 Can a Transformer-based policy with Q-value regularization improve trajectory stitching in offline RL compared to standard CSM methods?
• RQ2 How does integrating Q-value regularization into a sequence modeling framework affect performance in long-horizon and sparse-reward offline RL tasks?
• RQ3 To what extent can Q-value regularization mitigate distribution shift while maintaining fidelity to the behavior policy?
• RQ4 Does the joint optimization of behavior cloning and Q-value maximization lead to better sample efficiency and generalization than existing offline RL baselines?
• RQ5 Can the proposed method achieve state-of-the-art performance without requiring additional behavior cloning or model-based dynamics modeling?

Key findings

• QT achieves state-of-the-art performance across multiple D4RL benchmark environments, outperforming both conventional Dynamic Programming and Conditional Sequence Modeling methods.
• The method demonstrates superior trajectory stitching capability, effectively synthesizing high-return trajectories even from sub-optimal data sequences.
• The integration of Q-value regularization significantly improves return alignment, reducing the discrepancy between expected and optimal returns in action sampling.
• QT eliminates the need for explicit behavior cloning or conservative value function regularization by embedding policy regularization directly into the sequence modeling objective.
• Empirical evaluations show that QT maintains robust performance in long-horizon and sparse-reward settings, where traditional methods often fail.
• The ablation study confirms that both the behavior cloning and Q-value regularization components are essential for optimal performance, with each contributing uniquely to policy stability and return maximization.