PaperlayerView exact revision on arXiv

Source-linked AI summary

Q-learning Decision Transformer: Leveraging Dynamic Programming for Conditional Sequence Modelling in Offline RL

Taku Yamagata, Ahmed Khalil, Raul Santos-Rodriguez

arXiv:2209.03993v4cs.LG

TL;DR

Decision Transformer can fail to stitch sub-optimal trajectories, while Q-learning approaches can have unstable learning behavior. QDT uses Q-learning estimates to relabel return-to-go values before training DT. It performs well across the evaluated environments, whereas DT and CQL each fail in some environment types.

  • Problem

    Decision Transformer’s stitching ability remains limited when offline datasets contain only sub-optimal trajectories.

    Source: §6 · p. 7

  • Method

    QDT uses Q-learning estimates to relabel return-to-go values in DT training data, then trains DT on the relabelled data.

    Source: §1 · p. 2

  • Results

    QDT performs well across all evaluated environments, whereas DT and CQL each fail in particular environment types.

    Source: §6 · p. 6

  • Takeaways & Limitations

    The authors conclude that QDT is a promising framework for improving offline reinforcement-learning algorithms.

    Source: §7 · p. 9

    Reinforcement-learning algorithms such as QDT risk application in controversial, high-impact areas of human life, including military uses.

    Source: §7 · p. 9

Abstract

from arXiv · show

Recent works have shown that tackling offline reinforcement learning (RL) with a conditional policy produces promising results. The Decision Transformer (DT) combines the conditional policy approach and a transformer architecture, showing competitive performance against several benchmarks. However, DT lacks stitching ability -- one of the critical abilities for offline RL to learn the optimal policy from sub-optimal trajectories. This issue becomes particularly significant when the offline dataset only contains sub-optimal trajectories. On the other hand, the conventional RL approaches based on Dynamic Programming (such as Q-learning) do not have the same limitation; however, they suffer from unstable learning behaviours, especially when they rely on function approximation in an off-policy learning setting. In this paper, we propose the Q-learning Decision Transformer (QDT) to address the shortcomings of DT by leveraging the benefits of Dynamic Programming (Q-learning). It utilises the Dynamic Programming results to relabel the return-to-go in the training data to then train the DT with the relabelled data. Our approach efficiently exploits the benefits of these two approaches and compensates for each other's shortcomings to achieve better performance. We empirically show these in both simple toy environments and the more complex D4RL benchmark, showing competitive performance gains.

1. Introduction

The introduction presents QDT as a way to combine DT’s sequence modelling with Q-learning’s ability to stitch suboptimal trajectories. It argues that relabelling training data with Q-learning estimates can address DT’s stitching failure while mitigating Q-learning’s difficulty with long-horizon and sparse-reward tasks.

1. Introduction

  • DT (Chen et al., 2021) lacks stitching: it may follow a high-return trajectory in the data rather than combine suboptimal trajectory segments into an optimal path.In the example, no training trajectory contains the entire optimal path, and DT follows trajectory 1 despite trajectory 2 providing the optimal first action.Source: §1 · p. 1, Figure 1 · p. 2
  • QDT performs consistently across the evaluated environments, while DT struggles to stitch trajectories and CQL fails in a sparse-reward scenario.The results are summarized for simple, maze2d, and delayed-reward settings.Source: Figure 2 · p. 2, §1 · p. 2
  • Q-learning can stitch trajectory segments by propagating future rewards backward, but long horizons and sparse rewards make this difficult with off-policy function approximation.The introduction notes that this backward propagation can struggle when values must travel many steps to the initial state.Source: §1 · pp. 1–2
  • QDT uses Q-learning estimates to relabel training-data returns-to-go, then trains DT to combine Q-learning’s value propagation with sequence modelling.The authors propose this to improve the offline dataset and compensate for the approaches’ respective weaknesses.Source: §1 · p. 2

2. Preliminaries

The preliminaries define offline RL as policy learning from a fixed dataset, describe Decision Transformers as return-conditioned sequence models, and introduce CQL as the Q-learning framework used in this work.

Offline Reinforcement Learning

  • Offline RL learns a policy solely from a static dataset collected by unknown behavior policies, unlike online or off-policy learning with environment access.The text notes offline RL can be valuable when environment interaction is expensive or high risk.Source: §2 · p. 3

Decision Transformers

  • Decision Transformers condition future-action predictions on target returns-to-go, past states, and actions, using a Transformer to model sequence context.The architecture uses stacked self-attention layers with residual connections.Source: §2 · p. 3

Conservative Q learning

  • CQL regularizes Bellman-error learning to lower values for out-of-distribution state-action pairs while maintaining values for pairs in the training data.Its resulting Q-function need not be a pointwise lower bound, but its action expectation under the learning policy lower-bounds that policy's value.Source: §2 · p. 3

3. Method

QDT uses Q-learning to relabel return-to-go values before training an otherwise unchanged Decision Transformer. Selective relabeling aims to let the DT combine optimal sub-trajectories from sub-optimal trajectories.

Method

  • QDT learns a value function with Q-learning, relabels dataset RTGs, then trains the DT without modifying the Q-learning or DT algorithms.The relabeling is intended to let the DT use optimal sub-trajectories drawn from sub-optimal trajectories, addressing its stitching limitation.Source: §3 · p. 3
  • Using CQL’s lower-bound value estimate, QDT replaces RTGs when the trajectory RTG is lower, then propagates replacements backward through rewards.The method avoids replacing every RTG because learned Q-functions may be inaccurate, particularly with long horizons and sparse rewards; it aims to replace values where the estimate is accurate or closer to truth.Source: §3 · p. 3
  • Working backward from a final RTG of zero, QDT compares the learned-policy state value with the current RTG and uses the larger value to update the preceding RTG.The state value is the learned Q-function averaged over actions under the learned policy; if it exceeds the current RTG, the preceding RTG becomes the preceding reward plus that value, otherwise it uses the current RTG.Source: §3 · pp. 3–4
  • For each DT context window, QDT regenerates RTGs backward from the final RTG so each input sequence preserves the relation between rewards and future returns.Relabeling can break the required recurrence between reward and RTG, so the window’s earlier RTGs are reconstructed by repeatedly adding each reward to the next RTG.Source: §3 · p. 4

Theoretical considerations of QDT

  • Under the stated optimal-training assumption, a reward-conditioning model outputs the optimal action when conditioned on the state and optimal value; QDT’s conservative estimates move behavior-policy RTGs closer to optimal values.The paper notes that behavior-policy RTGs are often below the corresponding optimal value, while CQL’s conservative value is a lower bound on the true value (Srivastava et al., 2019; Kumar et al., 2020).Source: §3 · p. 4

4. Related Work

The approach combines conditional sequence modelling with a learned pessimistic value function to relabel return-to-go values, and frames this data improvement as a data-centric tool for offline RL.

Related Work

  • Offline RL learns from static datasets and addresses distribution shift through policy constraints, conservative reward estimates, uncertainty penalties, weighted imitation, or conditional sequence modelling.Source: §4 · p. 4
  • Unlike weighted imitation learning, the method relabels return-to-go values using a pessimistic value function and trajectory traces, while retaining a conditional sequential policy.Weighted imitation instead uses value estimates to weight training data for non-conditional, non-sequential policies; applying relabelling to those algorithms is proposed as future work.Source: §4 · pp. 4–5
  • The method is presented as a data-centric offline RL approach that improves training-data quality while using existing models.Source: §4 · p. 5

5. Evaluation

The evaluation compares QDT with CQL and DT across sub-optimal-trajectory, stitching, delayed-reward, and dense-reward settings. QDT improves DT where stitching matters, while its performance is similar to DT on delayed and dense-reward control tasks.

Simple environment

  • On sub-optimal-only data, vanilla DT fails badly, while CQL performs well and QDT performs similarly to CQL, learning a near-optimal policy.The dataset contains 100 episodes from a uniformly random policy, with an episode achieving close to optimal reward removed.Source: §5 · p. 5, Table 1 · p. 5

Maze2D environments

  • Across Maze2D’s stitching tests, QDT clearly improves DT, especially where CQL works well, indicating that QDT brings stitching capability to DT.Maze2D includes open, umaze, medium, and large mazes, with normal or dense rewards.Source: §5 · pp. 5–6, Table 2 · p. 6

Open AI Gym MuJoCo environments with delayed (sparse) reward

  • With delayed sparse rewards, CQL struggles while DT performs well, and QDT performs similarly to DT by using CQL’s information where useful.QDT performs worse than DT on medium-replay-walker2d, where CQL overestimates the value function in most states.Source: §5 · p. 6

Open AI Gym MuJoCo environments

  • With dense rewards, CQL performs equally or better than DT and QDT, while QDT performs the same as DT.Results align with cited references except for medium-replayhopper, likely because of high variance from five random seeds.Source: §5 · p. 6

6. Discussion

The discussion presents QDT as more robust across environment types than DT and CQL, while noting that its performance depends on relabeling quality and experimental conditions.

Performance across environments

  • QDT is reported as the only algorithm performing well across all tested environments, while DT and CQL each fail in environments where the other succeeds.QDT improves on DT with sub-optimal datasets but generally does not match CQL; the discussion notes room for improvement when datasets lack near-optimal trajectories.Source: §6 · pp. 6–7

The role of the discount factor

  • Discounted value estimates can prevent QDT from using Q-learning information for relabeling, so the authors recommend accurate value estimation and task-appropriate discounts.A smaller γ aids Q-learning convergence but shortens the effective horizon; γ = 0.99 suits short-horizon tasks, while γ = 0.999 suits long-horizon tasks.Source: §6 · pp. 7–8
  • QDT with discounted RTG relabeling matches CQL on Maze2D umaze and medium at γ = 0.99 and large at γ = 0.999.This result requires different discount factors across tasks.Source: §6 · p. 7, Table 5 · p. 8

Conservative weight

  • 29.6 ± 15.5 to 46.9 ± 13.8: raising CQL’s conservative weight α from 5.0 to 100 improves QDT on delayed-reward medium-replay-waker2d, but remains worse than DT.The learned CQL values exceeded true values in many states, causing incorrect RTG relabeling; offline hyperparameter optimization remains an open challenge.Source: Table 3 · p. 7, §6 · p. 8

Alternative approaches

  • Hepburn & Montana (2022) use a learned transition model and value function for stitching; QDT is model-free and reuses Q-learning without a dedicated stitching mechanism.The discussion also identifies eligibility traces, multi-step temporal-difference methods, and Monte Carlo returns as possible approaches to improving Q-learning for delayed rewards.Source: §6 · p. 8

Stitching ability

  • In hopper-medium-replay-v2, DT does not exceed the dataset’s best trajectory, whereas QDT improves on DT and approaches CQL when the top 40–50% of trajectories are removed.CQL exceeds the best trajectory except at X = 0, when the dataset already contains its best score.Source: §6 · pp. 6–7, Figure 4 · p. 8

Alternative approaches

  • Eligibility traces, multi-step temporal-difference methods, and Monte Carlo returns have been studied off-policy but, to the authors’ knowledge, not in offline settings.Unlike DT and QDT, Q-learning assumes rewards satisfy the Markov property and depend on the current state and action.Source: §6 · p. 8

Reproducing results for benchmarking

  • Offline RL benchmark comparisons can conflict because reproductions vary in dataset versions and experimental details, and evaluations often use too few random seeds.The authors caution that high variance in complex environments warrants care when drawing conclusions from results.Source: §6 · p. 9

7. Conclusions

QDT combines Dynamic Programming with reward-conditioned sequence modeling to address weaknesses in offline RL. The results show benefits over existing offline algorithms, alongside environment-specific performance gaps and potential societal risks.

7. Conclusions

  • QDT outperforms DT on Maze2D but not CQL, and beats CQL on delayed-reward MuJoCo except walker2d, where it trails DT.The conclusion reports that these results are encouraging but that further work is needed to resolve the gaps across all environments.Source: §7 · p. 9
  • QDT brings Dynamic Programming benefits to reward-conditioned sequence modeling; the paper illustrates this framework using CQL and DT.The authors describe it as a novel framework for improving offline reinforcement learning algorithms.Source: §7 · p. 9
  • The authors identify further work as necessary to resolve QDT’s performance gaps across all environments.They note that QDT remains behind CQL on Maze2D and behind DT on walker2d with delayed rewards.Source: §7 · p. 9
  • QDT, like other autonomous-system improvements, could be applied in controversial, high-impact fields such as military applications.The paper identifies this as a potential negative societal impact inherited across work on autonomous systems.Source: §7 · p. 9

A. Simple environment example trajectory data and its computation … B.4. DT and QDT model details

The simple example shows how DT can choose a suboptimal action when its best return-to-go resembles a suboptimal trajectory, while Q-learning finds the shortest path. The evaluation describes the maze, offline dataset, and CQL, DT, and QDT model setups.

A. Simple environment example trajectory data and its computation

  • DT predicts up rather than the optimal right action because the best return-to-go is closer to the suboptimal training trajectory.With context length K = 1, evaluation supplies [s=S, R=-2], which is closer to trajectory 1’s [s=S, R=-3] than trajectory 2’s [s=S, R=-11]; the figure contrasts this failure with Q-learning finding the shortest path.Source: Figure 5 · p. 12, A. Simple environment example… · p. 12
  • The example computes return-to-go from rewards using Eq. 5 and then associates the resulting RTGs with the trajectories.The supplied passages do not give the equation itself.Source: A. Simple environment example… · p. 12

B.1. Environment

  • The 6-by-6 maze has eight actions, gives -10 per step and +100 at the goal, and yields +50 total reward on the optimal five-step path.Action indices are remapped by state, with each state’s mapping fixed across episodes.Source: §B.1 · pp. 12–13, Figure 6 · p. 13

B.2. Dataset

  • The offline RL dataset was prepared from 100 episodes of a uniformly random policy; one episode achieving positive total reward was removed, leaving 100 episodes and 4,454 time steps.The maximum return was -10.0, the minimum was -490 (episodes terminate at 50 time step), and the average was -415.5.Source: §B.2 · p. 13

B.3. CQL model details

  • The simple-environment CQL model uses Double Q-learning and an embedding lookup table for discrete states.Its detailed parameters are listed in Table 6, whose values are not included in the supplied passage.Source: §B.3 · p. 13, Table 6 · p. 13

B.4. DT and QDT model details

  • The simple-environment DT and QDT models are built on minGPT open-source code.Their detailed parameters are listed in Table 7, whose values are not included in the supplied passage.Source: §B.4 · p. 13, Table 7 · p. 14

B.5. Further evaluation results for simple environment … C.3. Evaluation Process

The supplementary material reports that DT and QDT show larger differences than CQL between best-training and end-of-training results, indicating overfitting and unstable learning. It also documents the evaluation implementations and training protocols for CQL, DT, and QDT.

B.5. Further evaluation results for simple environment

  • DT and QDT show larger differences than CQL between best-training and end-of-training results, indicating overfitting and unstable learning.Table 8 reports best-performing-model results, while Table 9 reports results at the end of training; no individual table values are provided in the passages.Source: §B.5 · p. 14, Table 8 · p. 14, Table 9 · p. 15

C.1. CQL model details

  • CQL evaluation uses d3rlpy scripts and the hyperparameters from Kumar et al. (2020) for MuJoCo Gym and Maze2D.The same d3rlpy script and settings are reused for Maze2D simulations.Source: §C.1 · p. 14

C.2. DT and QDT model details

  • DT simulations use code from the original paper authors for both MuJoCo Gym and Maze2D environments.QDT adds return-to-go relabelling to the DT script, with the relabelling code described in Algorithms 1 and 2.Source: §C.2 · p. 15

C.3. Evaluation Process

  • CQL is trained with five seeds for 500,000 updates at batch size 256 and evaluated with 10 end-of-training episode roll-outs.These settings are inherited from the d3rlpy offline RL library.Source: §C.3 · p. 15
  • DT is trained with five seeds for 100,000 updates at batch size 64 and evaluated with 100 end-of-training episode roll-outs.The settings are inherited from the DT authors’ source code (Chen et al., 2021).Source: §C.3 · p. 15
  • QDT uses five seeds, each with its own trained CQL model for dataset relabelling, and is evaluated with 100 end-of-training episode roll-outs.It trains for 100,000 updates on MuJoCo Gym and 150,000 on Maze2D at batch size 64, using the same roll-out count as DT.Source: §C.3 · p. 15

C.4. Hyper parameter search · D. Justification of replacing RTG with the learned value function · E. Further Discussions

The hyperparameter search retains the original MuJoCo conservative weights and selects α = 1.0 for Maze2D; the justification section argues that when the learned value exceeds return-to-go, its error is smaller. The supplied passages contain no material from Further Discussions.

C.4. Hyper parameter search

  • The search uses the original papers’ hyperparameters overall but tunes conservative weight α because its optimum may differ between CQL and QDT.The cited original papers are Kumar et al., 2020 and Chen et al., 2021.Source: §C.4 · p. 15
  • For MuJoCo, α = 10.0 on medium and α = 5.0 on medium-replay work well for QDT without significantly degrading CQL performance.These are the values used in the original CQL paper, so the authors retain them.Source: §C.4 · p. 15, Table 11 · p. 16, Table 10 · p. 16
  • For Maze2D, the search starts from α = 10.0, the value used for MuJoCo Gym environments in the CQL paper.The supplied passages do not state the Maze2D starting value for medium-replay datasets.Source: §C.4 · p. 15
  • For Maze2D, α = 1.0 performs best in the CQL sweep, and QDT improves over DT at that setting.The authors use α = 1.0 for the paper while noting that lower values may perform better and further optimization remains future work.Source: §C.4 · p. 16, Table 12 · p. 16

D. Justification of replacing RTG with the learned value function

  • When the learned lower bound ˆV (st) exceeds return-to-go Rt, the error in ˆV (st) is smaller than the error in Rt.The section defines V ∗(st) as the optimal state value and ˆV (st) as its learned lower bound.Source: D. Justification of replacing RTG with… · p. 16
  • Because ˆV (st) is a lower bound on V ∗(st), both sides of the derivation are non-negative, permitting the absolute-value step.This is the stated intermediate step in the argument comparing value-function error with return-to-go error.Source: D. Justification of replacing RTG with… · p. 16

E.1. Why CQL outperforms DT/QDT on Maze2D, but fails on MuJoCo Gym delayed reward? … E.5. Consistency relabelling ablation experiment

The appendix explains QDT’s mixed performance through dataset and environment characteristics, reports no clear QDT advantage in additional MuJoCo tests, and finds limited ablation gains from consistency relabelling despite a rationale for retaining it.

E.1. Why CQL outperforms DT/QDT on Maze2D, but fails on MuJoCo Gym delayed reward?

  • Maze2D is simpler and has shorter episodes than MuJoCo, whose tasks have higher action and state dimensions and longer episodes.Sparse Maze2D medium and large also lose performance relative to their dense counterparts, consistent with their longer episodes than open and umaze.Source: §E.1 · p. 17, Table 13 · p. 17
  • Dataset quality helps explain the gap: DT performs well when optimal trajectories are present, while QDT can help DT on suboptimal-only Maze2D by using CQL information.Maze2D is described as containing suboptimal trajectories except open and open-dense; MuJoCo Gym has some optimal trajectories except halfcheetah.Source: §E.2 · p. 17
  • Halfcheetah medium and medium-replay remain exceptions: QDT matches DT despite lacking optimal trajectories because CQL itself struggles and offers little help.The open and open-dense Maze2D datasets are another exception to the suboptimal-data pattern, and DT performs well on them.Source: §E.2 · p. 17

E.3. Why QDT performs close to DT, not CQL in Fig. 4 (Gym hopper)?

  • A possible explanation for QDT remaining near DT is that a few critical states drive returns; Kumar et al. (2022) theoretically favor dynamic programming over imitation learning in such environments.QDT still uses DT as its agent algorithm, so environment properties that hinder DT can also hinder QDT.Source: §E.3 · p. 17

E.4. Extra results for removing top X%

  • Removing top-X% trajectories in halfcheetah-medium-replay and walker2d-medium-replay showed no clear QDT benefit over DT.Figure 7 reports evaluation scores for CQL, DT, and QDT on these datasets, with one-standard-deviation shading and dataset maximum scores as references.Source: §E.4 · p. 18, Figure 7 · p. 18

E.5. Consistency relabelling ablation experiment

  • Consistency relabelling showed a nonsignificant average benefit in the Simple Environment and no clear benefit in the more complex environments.The ablation used ten random seeds for Simple Environment and three for the other environments.Source: §E.5 · pp. 18–19, Table 15 · p. 19
  • The authors recommend retaining Algorithm 2 for now because it may prevent inconsistent, non-realistic training samples, even though its changes are minor relative to original RTG variations.The proposed explanation for the limited complex-environment gains is that relabelling changes are small compared with the original return-to-go variations.Source: §E.5 · p. 19

E.6. Aggregated evaluation results

Across maze2d, MuJoCo Gym delayed reward, and MuJoCo Gym, QDT improves on DT in maze2d, CQL fails in delayed reward, and DT and QDT perform similarly in MuJoCo Gym; QDT does not fail in any group.

E.6. Aggregated evaluation results

  • Aggregated scores cover maze2d, MuJoCo Gym delayed reward, and MuJoCo Gym, using median, IQM, and mean with 95% stratified bootstrap confidence intervals.The confidence-interval method is attributed to Agarwal et al. (2021).Source: §E.6 · p. 19
  • QDT improves DT in maze2d with help from CQL, while CQL fails in MuJoCo Gym delayed reward and DT and QDT perform similarly in MuJoCo Gym.The authors also report that QDT never failed in any of the three environment groups.Source: Figure 8 · p. 19, §E.6 · p. 19

Related papers

Loading 2209.03993v4…