Return Augmented Decision Transformer for Off-Dynamics Reinforcement Learning

TL;DR

This paper introduces the Return Augmented (REAG) method to improve Decision Transformer-based offline reinforcement learning in scenarios with dynamics shift, demonstrating consistent performance gains on benchmark datasets.

Contribution

The paper proposes the REAG approach to align return distributions across source and target domains in RCSL, with theoretical guarantees and practical implementations.

Findings

REAG improves DT-based RL performance in off-dynamics settings.

Theoretical analysis shows policy suboptimality is maintained with REAG.

Experiments on D4RL datasets validate the effectiveness of REAG methods.

Abstract

We study offline off-dynamics reinforcement learning (RL) to utilize data from an easily accessible source domain to enhance policy learning in a target domain with limited data. Our approach centers on return-conditioned supervised learning (RCSL), particularly focusing on Decision Transformer (DT) type frameworks, which can predict actions conditioned on desired return guidance and complete trajectory history. Previous works address the dynamics shift problem by augmenting the reward in the trajectory from the source domain to match the optimal trajectory in the target domain. However, this strategy can not be directly applicable in RCSL owing to (1) the unique form of the RCSL policy class, which explicitly depends on the return, and (2) the absence of a straightforward representation of the optimal trajectory distribution. We propose the Return Augmented (REAG) method for DT type frameworks, where we augment the return in the source domain by aligning its distribution with that in the target domain. We provide the theoretical analysis demonstrating that the RCSL policy learned from REAG achieves the same level of suboptimality as would be obtained without a dynamics shift. We introduce two practical implementations REAG$_\text{Dara}^{*}$ and REAG$_\text{MV}^{*}$ respectively. Thorough experiments on D4RL datasets and various DT-type baselines demonstrate that our methods consistently enhance the performance of DT type frameworks in off-dynamics RL.