Return Augmented Decision Transformer for Off-Dynamics Reinforcement Learning

TL;DR

The paper tackles offline off-dynamics reinforcement learning by leveraging data from a source domain with different dynamics to improve policy learning in a data-scarce target domain. It introduces RADT, a return-augmented Decision Transformer framework, with two concrete implementations: RADT-DARA (dynamics-aware reward augmentation) and RADT-MV (mean-variance return matching). The authors provide suboptimality guarantees under standard data-coverage assumptions and demonstrate, via D4RL MuJoCo experiments, that RADT variants often outperform baselines and DT alone, especially when target data is limited. The work offers a principled way to fuse source and target data for return-conditioned policies, enabling more reliable deployment in shifted dynamics settings.

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 the decision transformer (DT), which can predict actions conditioned on desired return guidance and complete trajectory history. Previous works tackle 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 Decision Transformer (RADT) method, 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 RADT achieves the same level of suboptimality as would be obtained without a dynamics shift. We introduce two practical implementations RADT-DARA and RADT-MV respectively. Extensive experiments conducted on D4RL datasets reveal that our methods generally outperform dynamic programming based methods in off-dynamics RL scenarios.

Figures (11)

• Figure 1: Rank scores for all baseline algorithms and our methods across the Random, Medium, Medium-R, and Medium-E datasets under BodyMass and JointNoise shift settings in the Walker2D, Hopper, and HalfCheetah environments. Ranks were assigned within each dataset, with the highest-performing algorithm receiving rank 1, followed by rank 2, and so on. In cases of tied scores, algorithms were assigned the same rank, and subsequent ranks were adjusted accordingly. Lower rank scores indicate better overall performance.
• Figure 2: Performance of RADT methods under varying BodyMass shift settings in the Walker2D Medium and Hopper Medium environments. "B-x" denotes that the body mass in the simulator is set to x. The target body mass is 2.94 in the Walker2D environment and 5 in the Hopper environment.
• Figure 3: Performance of RADT methods across varying JointNoise shift settings in the Walker2D and Hopper Medium environments. "N-x" denotes the addition of random noise in the range (-x, +x) to the action.
• Figure 4: Performance of RADT methods across Kinematic and Morphology Shift settings in the Medium Walker2D and Medium Hopper environments. For more details about these new shift settings, please refer to \ref{['sec:experiment_setting']}.
• Figure 5: Performance of RADT-MV under varying target return settings in the Walker2D and Hopper Medium Replay environments with the BodyMass Shift setting. The x-axis represents the target return values, the y-axis represents the normalized return. The dash line represents the line $y = x$.

Theorems & Definitions (9)

• Remark 4.1
• Theorem 4.5
• Remark 4.6
• Corollary A.1
• Theorem A.2
• Remark A.3
• Lemma B.1: Corollary 1 of brandfonbrener2022does
• Lemma B.2: Lemma 1 of brandfonbrener2022does
• Theorem B.3