Title: Rebalancing Return Coverage for Conditional Sequence Modeling in Offline Reinforcement Learning
Venue: NeurIPS 2025
Accepted: yes
Keywords: Offline Reinforcement Learning

Conditional Sequence Modeling

Return Coverage Rebalancing
Abstract: Recent advancements in offline reinforcement learning (RL) have underscored the capabilities of conditional sequence modeling (CSM), a paradigm that models the action distribution conditioned on both historical trajectories and target returns associated with each state. However, due to the imbalanced return distribution caused by suboptimal datasets, CSM is grappling with a serious distributional shift problem when conditioning on high returns. While recent approaches attempt to empirically tackle this challenge through return rebalancing techniques such as weighted sampling and value-regularized supervision, the relationship between return rebalancing and the performance of CSM methods is not well understood. In this paper, we reveal that both expert-level and full-spectrum return-coverage critically influence the performance and sample efficiency of CSM policies. Building on this finding, we devise a simple yet effective return-coverage rebalancing mechanism that can be seamlessly integrated into common CSM frameworks, including the most widely used one, Decision Transformer (DT). The resulting CSM algorithm, referred to as Return-rebalanced Value-regularized Decision Transformer (RVDT), integrates both implicit and explicit return-coverage rebalancing mechanisms, and achieves state-of-the-art performance in the D4RL experiments.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nOne can roughly categorize offline RL algorithms as either \"reweighted imitation learning\" or \"Q learning constrained to in-distribution data\". The authors focus on the former, in a sub-area called Conditional Sequence Modeling (CSM). In CSM, we condition the trajectory on a desired Reward to Go (RTG) $g$, denoting the remaining reward in a trajectory from time $t$: $\\tau_t$\n\n$$ argmin_\\pi - \\log \\pi(a_t | s_t, g(\\tau_t), \\tau_{t-1}).  $$\n\nIt is difficult to predict $g$. Instead, one can learn what is effectively a value function $f_\\pi(s) = V_\\pi(s)$\n\n\n$$ argmin_\\pi - \\log \\pi(a_t | s_t, f_\\pi(s_t), \\tau_{t-1}).  $$\n\nThus, the problem reduced to finding the actions that lead to a specific return, where the (undiscounted) return is specified by $f$. By setting $f = V_*$, we can recover an optimal policy.\n\nFirst, the authors analyze the performance gap caused by non-optimal $f$, (i.e., $f \\neq V_*$). Then, they provide sample complexity bounds for CSM. Next, they provide the first algorithmic contribution of the paper, termed RDT.\n\nIn RDT, the authors select an expert subset of trajectories and learn a policy on this subset using imitation learning. More formally, they learn a policy that learns a return-conditioned policy on the full dataset, while also imitating the expert-subset policy. In other words, the authors can leverage the full dataset for training, while biasing the policy towards \"better\" trajectories. The authors utilize RDT to propose RVDT, the main contribution of the paper.\n\nRVDT takes the RDT objective and adds a term that represents the Q function. The full RVDT objective contains three terms -- two from RDT and a new Q value term:\n- The return-to-Go objective\n- The imitation learning objective, for the expert subset of trajectories\n- Maximization of $Q_\\pi$\n\nThe authors test their algorithm across standard offline RL benchmarks, demonstrating that their approach is state of the art. They perform ablations demonstrating that their method scales well with increasingly-expert data, as well as in low-data regimes.\n\nSTRENGTHS AND WEAKNESSES:\n### Strengths\n- The paper is very well written and I enjoyed reading it. It tackles many complex topics, yet the writing style made it easy to follow along.\n- Theorem 2 and 3 provide good theoretical as to why the authors' method should succeed\n- The RVDT objective seems fairly simple to implement, as it is just the combination of a few common objectives\n- RVDT performs very well compared with prior work\n- The authors provide sufficient ablations and analysis of their algorithm\n\n### Weaknesses\n- My biggest concern is that the author's method is very similar to QT. My understanding is that their objective just has one more term, corresponding to the imitation learning objective from an expert subset. Many experiments in Table 1 demonstrate that the performance of QT is only marginally less than RVDT.\n- Tuning hyperparameters $\\eta, \\alpha$ seems like it could be difficult. The automatic $\\eta$ normalization scheme still likely requires hyperparameter tuning based on the variance of Q values.\n\nQUESTIONS:\n- Contributions: Each contribution is a bit long and difficult to parse without reading the preliminary section. Could you simplify/summarize these a bit? \n- Section 3.2: Can you please add a sentence with a more intuitive explanation of the conditioning function $f$? The $f$ we aim to find is often the RTG of the optimal policy $\\pi_*$, correct? $f(s) = \\max_\\pi \\mathbb{E}_{\\pi}[g(s)]$?\n\nLIMITATIONS:\nI do not think the authors list any limitations. I would suggest they consider adding a few sentences about this. One limitation is that their method requires learning an extra imitation learning policy, which can be costly. Especially in scenarios where the performance gap betwee QT and RVDT is negligible.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 4, "ORIGINALITY": 2, "DATE": "2025-06-25"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper addresses the trade-off between return maximization and behavior coverage in conditional sequence modeling (CSM)-based offline RL methods such as Decision Transformer. They propose a test-time control framework that balances this trade-off by reranking sampled trajectories using a weighted combination of predicted return and a learned coverage score. Their method, which does not alter the training process, shows improved performance across several D4RL tasks.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\n1. The paper addresses a meaningful and underexplored problem in conditional sequence modeling (CSM) for offline RL—namely, the trade-off between return maximization and behavioral coverage—which is both practically relevant and conceptually important.\n2. The theoretical formulation is relatively complete and provides a useful lens to understand the inherent tension in CSM-based policies between imitating high-return behaviors and maintaining adequate state-action coverage.\n3. The experimental design is comprehensive, evaluating across multiple dimensions including data quality (expert vs mixed), data scale (low-data regime), and reward density (sparse vs dense).\n\nWeaknesses:\n\n1. The proposed rebalancing mechanism fundamentally relies on the availability and selection of high-quality, expert-level trajectories. In many real-world offline RL settings, such as medical or autonomous driving applications, “expert” behavior is often ambiguous.\n2. While the method claims to enhance expert-level coverage, in practice the KL regularization simply increases the weight of expert-like behaviors, but does not expand the policy’s support into previously unobserved state-action regions.\n\nQUESTIONS:\n1. The combined use of Q-value regularization and expert KL divergence may result in conflicting optimization signals, particularly when Q-value estimation is poor or inaccurate. How do you address or mitigate the potential instability and performance degradation that may arise from this issue?\n2. Besides using top-k return for expert trajectory selection, did the authors consider or evaluate any alternative definitions or methods for identifying expertise in the offline dataset?\n3. The bound in Theorem 2 shows a dependence on both expert-level and full-spectrum return coverage parameters (α_f and α_f*). Could you provide more intuition or empirical evidence about how trade-offs between these two types of coverage affect policy performance? Are there scenarios where improving one necessarily degrades the other?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper explore the CSM used in offline RL. The auther studies two types of return (expert-level and full spectrum) and provide some theoretical analysis on their impact. Additionally, auther provide a return converage rebalancing mechanism into DT framework.\n\nSTRENGTHS AND WEAKNESSES:\n>Strength\n\n1. the paper is clear\n2. auther did compare with various benchmarks and algorithms\n\n>Weaknesses\n\n1. The major concern is the assumptions in the theorem and the bound value $\\alpha_f$. In order for those assumption to hold, the $\\alpha$ maybe really small or even close to 0.  If this is the case, the bound for theorem 1 and lemma 1 become meaningless ($\\infty$).\n\n2. The method relies on training an expert policy via imitation learning on near-optimal trajectories, which is then used to regularize the main policy through a KL divergence term. The quality of the expert policy will affect the performance of RVDT. Additionally, the higher quality of expert policy will result in longer pretraining time and utilize larger dataset.\n\n3. For the evaluation, some of results have statistic value (std) some are not.\n\nQUESTIONS:\n1. Can auther discuss the validity of the return-coverage assumption? Those assumptions may hold in small or well-curated offline datasets, it becomes increasingly difficult to satisfy as the dataset grows in size and complexity. In such cases, $\\alpha$ may become vanishingly small, which would cause the theoretical performance bounds to become vacuous or uninformative. \n\n2. How the quality of the expert policy affects RVDT performance, How long it takes to train the expert policy, and the trade-off between expert policy quality, training time, and RVDT effectiveness?\n\n3. For Table 1 gym tasls, majority of RVDT performance are so close to the QT  if we consider QT's statistic value, it may hard to tell if RVDT has performance different with QT for the gym tasks. Same issue for Kitchen tasks and AntMaze tasks\n\n4. Why is  top 50% of trajectories to construct the expert dataset not more?\n\nI would be willing to raise my rating if the authors can address 1. the assumption bound issue 2. the inconsistance statistic result report, as they are my primary concern and the main reason for my initial assessment.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper investigates the distributional shift issue in conditional sequence modeling (CSM) for offline RL, particularly caused by imbalanced return distributions. The authors highlight the importance of covering both expert-level and full-spectrum returns, proposing a straightforward return-coverage rebalancing mechanism named Rebalancing Value-regularized Decision Transformer (RVDT). Extensive experiments on D4RL benchmarks validate that RVDT achieves state-of-the-art performance.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths:**\n\n1. The findings presented in this paper, particularly regarding the significant impact of expert-level coverage and full-spectrum conditional returns, are important and valuable to the research community.\n2. Based on the analysis provided above, the proposed RVDT method is both simple and effective.\n3. The experiments conducted are thorough and comprehensive.\n\n**Weaknesses:**\n\n1. Please correct me if I have overlooked anything, but an important ablation study seems to be missing. Specifically, an ablation on the two newly introduced terms in Equation (7)—the policy regularization term and the value-based guidance term—would be helpful. What would be the impact on performance if either of these two terms were removed?\n2. In the experimental section, the authors attribute RVDT's inferior performance compared to D-QL on the larger AntMaze tasks to D-QL's highly expressive policies. However, due to the lack of further experimental evidence, it would be helpful if the authors could clarify how D-QL performs relative to RVDT on larger Maze2D tasks, which are similarly large maze-like environments. Alternatively, could the authors provide additional qualitative or quantitative results to support this claim?\n3. What is the computational overhead introduced by RVDT? For instance, RVDT requires first training an expert policy \\$\\pi\\_e\\$ and then introduces two additional terms during offline RL training.\n4. Some related works also utilize a small portion of expert data to improve offline RL performance \\[1,2,3].\n\n\\[1] Kumar, Aviral, et al. \"Pre-training for robots: Offline RL enables learning new tasks from a handful of trials.\" *arXiv preprint arXiv:2210.05178* (2022).\n\n\\[2] Yang, Qisen, et al. \"Hundreds guide millions: Adaptive offline reinforcement learning with expert guidance.\" *IEEE Transactions on Neural Networks and Learning Systems* 35.11 (2023): 16288–16300.\n\n\\[3] Cheng, Peng, et al. \"Pushing the Limit of Sample-Efficient Offline Reinforcement Learning.\" *ICLR 2025 Workshop on World Models: Understanding, Modelling and Scaling.*\n\nQUESTIONS:\nSee the weaknesses part above.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 5, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}