Reinformer: Max-Return Sequence Modeling for Offline RL

TL;DR

The paper addresses the gap in offline RL where sequence-modeling approaches neglect the core objective of return maximization. It introduces Reinformer, a Reinforced Transformer that implements max-return sequence modeling by predicting the maximum in-distribution return via expectile regression and conditioning action selection on this predicted return during inference. Empirical results on D4RL show competitive performance with classical offline RL methods and clear advantages over prior sequence-modeling methods, particularly in trajectory stitching tasks such as Antmaze. The approach provides a principled mechanism to stitch suboptimal trajectories without relying on naive max-return initialization, offering practical benefits for learning from diverse offline data.

Abstract

As a data-driven paradigm, offline reinforcement learning (RL) has been formulated as sequence modeling that conditions on the hindsight information including returns, goal or future trajectory. Although promising, this supervised paradigm overlooks the core objective of RL that maximizes the return. This overlook directly leads to the lack of trajectory stitching capability that affects the sequence model learning from sub-optimal data. In this work, we introduce the concept of max-return sequence modeling which integrates the goal of maximizing returns into existing sequence models. We propose Reinforced Transformer (Reinformer), indicating the sequence model is reinforced by the RL objective. Reinformer additionally incorporates the objective of maximizing returns in the training phase, aiming to predict the maximum future return within the distribution. During inference, this in-distribution maximum return will guide the selection of optimal actions. Empirically, Reinformer is competitive with classical RL methods on the D4RL benchmark and outperforms state-of-the-art sequence model particularly in trajectory stitching ability. Code is public at https://github.com/Dragon-Zhuang/Reinformer.

Figures (10)

• Figure 1: The overview of ReinForced TransFormer (Reinformer) (a) Model Architecture: The returns-to-go is the second token of Reinformer inputs and the outputs contain returns and actions. (b) Train Loss: As a max-return sequence model, Reinformer not only maximizes the action likelihood but also maximizes returns by expectile regression. (c) Inference Pipeline: When inference, Reinformer first 1) gets state from the environment to predict the in-distribution maximum return. Then 2) predicted in-distribution max return is concatenated with state to predict the optimal action. Finally, 3) the environment executes the predicted action to return the next state.
• Figure 2: A maze example for trajectory stitching analysis.
• Figure 3: illustration of weight.
• Figure 4: The probability of improvement of Reinformer compared with other methods using rliable agarwal2021deep. The larger the probability is, the better our method performs.
• Figure 5: The comparison between different inference approaches on Antmaze dataset: Reinformer with predicted maximized returns versus Reinformer-n with naively maximized returns. Abbreviations: u $\rightarrow$umaze, ud $\rightarrow$umaze-diverse, mp $\rightarrow$medium-play, md $\rightarrow$medium-diverse.

Theorems & Definitions (2)

• Theorem 3.1
• proof