Research map / Reinforcement learning / Offline RL
Reinformer: Max-Return Sequence Modeling for Offline RL
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Reinformer, a sequence model for offline RL that adds an expectile-regression return-prediction loss so that at inference it predicts the maximum in-distribution return-to-go and conditions action selection on it.
Problem
Return-conditioned sequence models (Decision Transformer and variants) maximize action likelihood but ignore the RL objective of maximizing return, so they set returns manually and struggle with trajectory stitching from sub-optimal data.
Approach
Add an expectile-regression returns-to-go loss to the sequence model so it predicts the maximum in-distribution return-to-go at each step (theorem: as m→1 the prediction approaches the dataset max), and condition the action on this predicted maximum rather than a manually chosen initial return.
Key results
- On D4RL Gym, Reinformer totals 709.46 vs 672.6 for DT and 687.3 for EDT (sequence models) across nine datasets.
- On Antmaze-v2 (four datasets), Reinformer totals 174.0 vs 142.6 for BC, 126.3 for DT, and 325.5 for IQL, so it beats sequence baselines but trails IQL.
- Reinformer reports +37.07% on Maze2d and +28.45% on Kitchen over the strongest sequence baseline; on Antmaze medium-play and medium-diverse, the max-return approach improves over the naive max variant by 70.03%.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add max-return sequence modeling to sequence modeling | Decision Transformer-style offline RL · Maze2d-v1 | normalized last score: +37.07% | Helped |
| Add max-return sequence modeling to sequence modeling | Decision Transformer-style offline RL · Kitchen-v0 | normalized last score: +28.45% | Helped |
| Add max-return sequence modeling instead of using a naive maximum-return initial condition (ablation) | Reinformer · Antmaze-v2 umaze and umaze-diverse | normalized best score: 19.43% | Helped |
| Add max-return sequence modeling instead of using a naive maximum-return initial condition (ablation) | Reinformer · Antmaze-v2 medium-play and medium-diverse | normalized best score: 70.03% | Helped |
| Increase expectile parameter m toward 1 (ablation) | Reinformer · D4RL datasets | normalized score | Helped |
| Set expectile parameter m=0.999 (ablation) | Reinformer · D4RL datasets | normalized score | Hurt |
| Use context length K=20 rather than a smaller context (ablation) | Reinformer · high-quality datasets such as Gym-medium-expert and Kitchen | normalized score | Helped |
| Use a smaller context length rather than K=20 (ablation) | Reinformer · low-quality datasets such as Gym-medium/medium-replay and Antmaze | normalized score | Helped |
| Change reward from r̂=100 × r to r̂=100 × r + 1 (ablation) | Reinformer · Antmaze-v2 umaze-diverse | normalized best score: NaN -> 65.8±4.1 | Helped |
| Change reward from r̂=100 × r to r̂=100 × r + 1 (ablation) | Reinformer · Antmaze-v2 medium-play | normalized best score: 11.4±3.78 -> 13.2±6.1 | Helped |
| Change reward from r̂=100 × r to r̂=100 × r + 1 (ablation) | Reinformer · Antmaze-v2 medium-diverse | normalized best score: 7.2±2.17 -> 10.6±6.9 | Helped |
| Change reward from r̂=100 × r to r̂=100 × r + 1 (ablation) | Reinformer · Antmaze-v2 umaze | normalized best score: 85±3.87 -> 84.4±2.7 | Hurt |
Surprising findings
- Existing sequence models such as DT and EDT underperform BC on Antmaze, indicating weak trajectory stitching.
- EDT, despite an expectile return loss, is highly sensitive to the initial return-to-go and still effectively relies on naive max-return selection.
Limitations
- Still falls short of IQL on Antmaze, with a considerable performance gap.
- Training is unstable on some datasets; Antmaze results report best rather than last score, and the method fails on antmaze-large.
- Many baseline results were reproduced by the authors, and hyperparameters (m, K, learning rate, training steps) were tuned per dataset.
Experimental setup
Datasets: D4RL Gym-v2 D4RL Maze2d-v1 D4RL Kitchen-v0 D4RL Antmaze-v2
Benchmarks: D4RL Gym D4RL Maze2d D4RL Kitchen D4RL Antmaze
Metrics: normalized score
Baselines: BC CQL IQL DT ODT EDT QDT
Builds on: Decision Transformer Online Decision Transformer Elastic Decision Transformer expectile regression upside-down reinforcement learning
Closest papers by meaning
- In-Dataset Trajectory Return Regularization for Offline Preference-based Reinforcement Learning (AAAI 2025 · similarity 0.89)
- A Tractable Inference Perspective of Offline RL (NeurIPS 2024 · similarity 0.89)
- Stitching Sub-trajectories with Conditional Diffusion Model for Goal-Conditioned Offline RL (AAAI 2024 · similarity 0.89)
- Offline RL with Observation Histories: Analyzing and Improving Sample Complexity (ICLR 2024 · similarity 0.88)
- Tackling Data Corruption in Offline Reinforcement Learning via Sequence Modeling (ICLR 2025 · similarity 0.88)
- DiffStitch: Boosting Offline Reinforcement Learning with Diffusion-based Trajectory Stitching (ICML 2024 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.