Source-linked AI summary
Trust Guided Decision Transformer
Chainesh Gautam, Raghuram Bharadwaj Diddigi, Chandramouli Kamanchi, Pankaj Dayama, Sumanta Mukherjee, Kameshwaran Sampath
TL;DR
Long-rollout Decision Transformer performance can degrade as self-generated context drifts from offline training data. TGDT filters suffixes using next-state prediction error before critic ranking, and improves over full-context, reset-based, and critic-only context selection across the evaluated D4RL domains.
Problem
During long evaluation rollouts, Decision Transformer context can drift from offline training histories, leaving actions poorly constrained by the training objective.
Source: §1 · p. 1Method
TGDT checks suffix reliability using rolling next-state prediction error against a held-out offline threshold, then uses a frozen critic to rank actions from surviving suffixes.
Source: §3 · p. 3, §3.3 · p. 5Results
TGDT improves over full context, hard reset, and critic-only suffix selection on D4RL Maze2D, AntMaze, and MuJoCo locomotion.
Source: §1 · p. 2Takeaways & Limitations
Filtering unreliable contexts before value ranking reduces persistent high-error runs and improves normalized return over critic-only suffix selection on Maze2D-medium.
Source: §1 · p. 2The held-out teacher-forced threshold is not a formal coverage guarantee for closed-loop evaluation because the scores are non-exchangeable.
Source: §3.3 · p. 5, §4.7 · p. 9
Abstract
from arXiv · showhide
Decision Transformer performance degrades on long rollouts because the conditioning context drifts out of the training distribution. We show that this drift is visible through the model's own next state prediction error, which rises during rollout and stays elevated, giving a direct signal of when context has become unreliable. We introduce Trust Guided Decision Transformer (TGDT), which selects context before applying value guidance. At each step, TGDT evaluates several recent context suffixes using rolling next state prediction error, calibrated against held out offline data via split conformal prediction. It keeps only suffixes whose error stays within the calibrated threshold, then uses a frozen critic to choose the highest value action among the trusted suffixes. This reverses the order used by value only elastic selection, where the critic may choose an action generated from a context the model itself has flagged as unreliable. Experiments on D4RL navigation and locomotion tasks show that state prediction, critic guidance, and hard context reset each solve only part of the problem. TGDT reduces persistent high error runs and improves return over vanilla Decision Transformer, reset based context control, and value only context selection.
1 Introduction
Long-rollout degradation is traced to rollout context mismatch: self-generated histories drift from offline training contexts, leaving actions poorly constrained. TGDT filters contexts for reliability before using a critic to rank their actions.
1 Introduction
- Long evaluation rollouts can drift from offline trajectory contexts, leaving DT actions on histories the training objective does not constrain well.The paper calls this failure mode rollout context mismatch.Source: §1 · p. 1
- Value-guided and behavior-regularized fixes primarily improve action quality within a context, while critic value does not establish that the action-generating context is reliable.TGDT therefore rejects unreliable suffixes before critic ranking.Source: §1 · pp. 1–2
- On Maze2D-medium, trust filtering reduces the longest above-threshold error run by approximately 8× and improves normalized return over critic-only suffix selection.The controlled ablation uses the same trained model, frozen critic, and candidate suffix set; the difference is filtering before ranking.Source: §1 · p. 2
- Persistent rolling next-state prediction error precedes return degradation and is detectable with a frozen external probe, including for vanilla DT.The probe signal is a property of the rollout rather than an artifact of a particular architecture.Source: §1 · pp. 1–2
- Auxiliary prediction, frozen IQL critic guidance, and a behavior-regularized residual action head each improve return, but still leave long runs of high prediction error.The prediction head is used as a reliability sensor, not for model rollouts or policy optimization through simulated trajectories.Source: §1 · p. 2
- TGDT improves over full context, hard reset, and critic-only suffix selection on D4RL Maze2D, AntMaze, and MuJoCo locomotion.It uses the same trained model and ranks actions only after rejecting suffixes whose recent prediction error exceeds a held-out offline threshold.Source: §1 · p. 2
2 Related work
Related work adds value guidance, behavior regularization, or adaptive context length to sequence-modeling policies. TGDT instead uses prediction error to assess context reliability before ranking actions, and uses its dynamics predictor only as a reliability sensor.
2 Related work
- Decision Transformer and value-guided sequence models use return-conditioned sequence modeling, with variants adding Q learning, critic alignment, advantage conditioning, or value-guided training.These approaches address limitations of pure sequence imitation, including poorly supported requested returns and trajectory stitching; Trajectory Transformer uses its learned model for planning [Janner et al., 2021].Source: §2 · pp. 2–3
- Behavior-regularization methods constrain actions toward offline support, an ingredient TGDT uses to keep prediction error interpretable as context mismatch rather than arbitrary action drift.TGDT does not claim novelty in these ingredients.Source: §2 · p. 3
- EDT adapts history length at test time, while TGDT uses recent next-state prediction error to filter suffixes before critic ranking.This makes context selection a reliability decision before it becomes a value-ranking decision.Source: §2 · p. 3
- TGDT calibrates its held-out threshold using the split-conformal order statistic but makes no online coverage guarantee because closed-loop rollouts are not exchangeable with teacher-forced validation trajectories.The threshold is used as a reliability signal, evaluated by its effect on persistent prediction error and return.Source: §2 · p. 3
- Unlike model-based offline RL methods that use dynamics models for rollouts or policy optimization, TGDT uses its next-state predictor only to assess context reliability.Its policy remains a return-conditioned sequence model, and its critic ranks actions proposed by trusted suffixes.Source: §2 · p. 3
3 Method
TGDT turns persistent next-state prediction error into a context-reliability check, then lets a frozen critic rank actions only from trusted suffixes. Its method combines an offline-trained predictor, behavior-regularized action proposals, an internally predicted next state, and a held-out error threshold.
3.1 Detecting rollout context mismatch with a transition probe
- An external transition probe detects persistent rollout context mismatch through elevated next-state prediction error, while TGDT uses its own reliability head at runtime.The probe and internal head identify similar persistent high-error regimes on Maze2D, with Pearson correlation r ∈[0.82, 0.88] and threshold agreement at least 0.86.Source: §3.1 · p. 4
3.2 Training a value guided DT with a reliability head
- TGDT jointly trains a frozen offline critic, a behavior-regularized action head with bounded residuals, and an internal next-state prediction head.The critic ranks candidate actions at evaluation, while behavior regularization keeps prediction error informative about context drift rather than action drift.Source: §3.2 · pp. 4–5
3.3 Held out reliability threshold
- A held-out offline split sets the threshold from rolling prediction errors, using their 95th percentile when α = 0.05.Because the closed-loop scores are non-exchangeable, the threshold is an empirical reliability scale, not a formal coverage guarantee.Source: §3.3 · p. 5
3.4 Trust filtered context selection
- TGDT rejects suffixes whose recent rolling error exceeds the held-out threshold before the critic ranks actions, falling back to the shortest suffix if none remain trusted.Each suffix score uses errors from earlier steps, so the action is committed before the observed next state updates the score.Source: §3.4 · pp. 5–6
4 Experiments
The experiments test whether Decision Transformer rollout mismatch is measurable, whether training-side improvements address it, and whether execution-time reliability filtering helps. TGDT improves benchmark performance most where long rollouts allow context mismatch to accumulate, while its runtime filtering reduces persistent high-error runs.
4 Experiments
- The experiments ask whether rollout reliability failures are measurable, whether training-side changes resolve them, and whether filtering beats critic-only suffix selection.The critic-only comparison holds the trained model, critic, and candidate suffix set fixed.Source: §4 · p. 6
4.1 Setup
- The experiments compare five D4RL domains and value-based and sequence-modeling baselines, using three seeds and 100 evaluation episodes per task and method.Execution modes include full context, hard reset, critic-only suffix selection, and TGDT.Source: §4.1 · pp. 6–7
4.3 Rollout context mismatch has a measurable precursor
- Vanilla DT’s median probe error rises above the held-out threshold during Maze2D-medium rollouts and remains elevated in persistent runs.The violation fraction increases over episode time, indicating mismatch is not merely isolated prediction spikes.Source: §4.3 · p. 7
4.2 Main D4RL results
- On AntMaze medium-diverse, VDT scores 30.0 →60.0 while IQL leads value-based methods at 70.0; TGDT gains most on long-horizon tasks.Table 1 reports normalized D4RL scores across benchmark domains; the AntMaze result is consistent with sparse-reward settings requiring sub-trajectory stitching.Source: §4.2 · p. 7, Table 1 · p. 8
4.4 Training-side fixes are not sufficient
- Training-side state prediction and critic guidance improve aspects of performance but do not eliminate long above-threshold error runs.The critic gives the largest normalized-score gain; state prediction alone changes reliability but does not solve the control problem.Source: §4.4 · pp. 7–8, Figure 2 · p. 9
4.5 Trust-filtered selection reduces persistent mismatch
- TGDT achieves the highest Maze2D-medium score and cuts the longest above-threshold run by roughly an order of magnitude versus critic-only selection.The modes share the same trained model, frozen critic, and candidate suffixes; TGDT filters unreliable suffixes before critic ranking.Source: §4.5 · p. 8, Figure 2 · p. 9
4.6 Mechanism of TGDT
- TGDT shortens context when rolling error exceeds the threshold and returns to longer contexts when scores fall below it.This selective adaptation outperforms always resetting to a one-step context.Source: §4.6 · p. 9, Figure 1 · p. 9
4.7 Limitations
- Held-out thresholds are calibrated on teacher-forced offline scores, so closed-loop evaluation violates the exchangeability assumption and no formal coverage guarantee is claimed.Evaluation is limited to state-based D4RL; image-based and non-stationary settings remain open empirical questions, and extra forward-pass cost grows with suffix count and model size.Source: §4.7 · p. 9
A.1 Default hyperparameters
The default settings specify execution-time reliability filtering and training-side action regularization, with reported ablations evaluated across multiple seeds and episodes.
Execution-time parameters
- The candidate suffix set determines which context lengths TGDT evaluates, while the reliability window averages recent prediction errors before judging a suffix.The risk level sets the held-out threshold; α = 0.05 corresponds to the 95th percentile.Source: §A.1 · p. 12
Training-side parameters
- The critic weight and state-prediction weight scale their respective training objectives, while the residual bound and penalty constrain action corrections.Together, these settings keep actions close to the data while permitting local value-guided improvement.Source: §A.1 · p. 12
Evaluation protocol
- Internal ablations and sensitivity analyses use three independent seeds and 100 evaluation episodes per seed.Reported values are means and standard errors across seeds.Source: §A.1 · p. 12
A.2 Agreement between the diagnostic probe and the internal reliability head
The appendix compares TGDT’s internal reliability head with an external diagnostic probe on the same closed-loop Maze2D trajectories. Their rolling error magnitudes and reliability decisions agree substantially, supporting the internal head’s use at runtime.
Signals and evaluation
- The internal head is TGDT’s runtime reliability signal, whereas the external probe diagnoses rollout mismatch but is not used by TGDT at evaluation.Both signals are evaluated on the same closed-loop trajectories and compared with their own held-out thresholds.Source: §A.2 · p. 12, Table 3 · p. 13
Agreement metrics
- Probe-positive recall focuses on whether the internal head also flags timesteps identified as unreliable by the external probe.This avoids relying only on overall threshold agreement, which can be high when most timesteps are below threshold.Source: §A.2 · p. 13
Agreement results
- The internal head and external probe have Pearson correlation 0.82, Spearman correlation 0.75, threshold agreement 0.89, and probe-positive recall 0.78.The measures capture agreement in error magnitude, ranking, binary decisions, and probe-flagged timesteps, respectively.Source: §A.2 · p. 13, Table 3 · p. 13
A.3 Sensitivity to reliability window and risk level
Sensitivity tests vary the reliability window and risk level while holding other settings fixed. Intermediate windows and moderate risk levels provide stable performance without reacting too readily to isolated errors or allowing mismatch to persist.
Reliability window
- Intermediate windows, especially Ke = 5 and Ke = 10, achieve the best scores on both Maze2D tasks.Short windows react to isolated spikes and intervene frequently; long windows smooth errors but may detect accumulating mismatch too late.Source: §A.3 · p. 14
Risk level
- Performance is stable across α ∈ {0.01, 0.05, 0.10} and drops at α = 0.20 on both Maze2D tasks.A low threshold filters too aggressively and can reject useful contexts, while a high threshold may allow unreliable contexts to persist.Source: §A.3 · p. 14
Default choice
- The default Ke = 10 and α = 0.05 lie in a stable intermediate region rather than a narrow optimum.This setting reacts to persistent mismatch without overreacting to isolated errors.Source: §A.3 · p. 14
A.4 Sensitivity to training-side regularization
Training-side regularization balances critic guidance, state-prediction learning, and bounded residual corrections against rollout reliability. Sensitivity results also show that reliability-window settings matter: intermediate windows perform best, while extreme windows react too early or too late.
Training-side hyperparameter sensitivity
- The sweeps vary one hyperparameter at a time on maze2d-medium-v1 and maze2d-umaze-v1, holding the others at λQ = 0.01, λs = 1.0, δmax = 0.05, and β = 0.05.Each setting uses three independent seeds and 100 evaluation episodes per seed.Source: §A.4 · p. 14
- Moderate critic guidance improves return, with the best performance around λQ = 0.01; larger values reduce performance as critic pressure moves residual actions away from the near-data region.The λQ = 0 setting, which removes critic guidance, gives the lowest score.Source: §A.4 · p. 14
- When λs = 0, both return and reliability degrade; λs ∈{0.5, 1.0} improves performance, while very large values can slightly reduce return.Over-weighting prediction can trade off against action learning.Source: §A.4 · pp. 14–15
Reliability-window sensitivity
- Intermediate reliability windows perform best: very short windows overreact, while very long windows react late.The default setting Ke = 10, α = 0.05 lies in a stable region.Source: Figure 3 · p. 15
Training-side hyperparameter sensitivity
- A residual bound that is too small leaves little room to improve over behavior cloning, while an overly large bound increases violation fraction and reduces return.The default δmax = 0.05 lies near the best region, balancing local action improvement with reliability.Source: §A.4 · p. 15
- Removing the residual penalty allows larger corrections and increases violation fraction; overly large penalties suppress useful corrections and lower return.Moderate penalties are stable on both Maze2D tasks.Source: §A.4 · p. 15
A.5 Context-length usage
TGDT selects context length online and uses a mixture of short, medium, and full contexts rather than always resetting to the shortest suffix. It shortens context when longer histories are unreliable and preserves history when the reliability score remains below threshold.
A.5 Context-length usage
- The context-length diagnostic aggregates selected candidate suffixes L = {1, 5, 10, 20} across three seeds and 100 episodes per seed on each Maze2D task.Bars report the fraction of evaluation timesteps assigned to each suffix length.Source: §A.5 · p. 15, Figure 5 · p. 16
- TGDT uses a mixture of short, medium, and full contexts instead of collapsing to the shortest suffix.The diagnostic supports interpreting TGDT as elastic reliability-aware context selection rather than a reset heuristic.Source: §A.5 · pp. 15–16, Figure 5 · p. 16
- TGDT shortens context when long histories become unreliable and preserves longer context when the reliability score remains below the held-out threshold.This contrasts with hard reset overuse of L = 1 and value-only selection overuse of long contexts despite high recent prediction error.Source: §A.5 · p. 16
A.6 Evaluation-time runtime overhead
TGDT's extra evaluation cost comes from querying the transformer on multiple candidate suffixes, while its reliability filter adds only scalar operations. Against value-only selection with the same suffix set, runtime is nearly the same; enumerating every suffix is substantially slower.
A.6 Evaluation-time runtime overhead
- Runtime is measured on maze2d-medium-v1 using an NVIDIA RTX A4500 20GB GPU, batch size one, and 60k environment steps.The estimate covers 100 evaluation episodes with 600 steps per episode.Source: §A.6 · p. 16, Table 4 · p. 17
- TGDT and value-only selection have nearly the same runtime when they use the same candidate suffix set.Both query the transformer once per candidate suffix and query the critic for each candidate action; TGDT adds only scalar reliability comparisons and queue updates.Source: §A.6 · pp. 16–17
- The dense variant is substantially slower because runtime scales linearly with the number of candidate suffixes.The default multiscale set {1, 5, 10, 20} avoids enumerating every suffix length.Source: §A.6 · pp. 16–17
A.7 Internal ablation across different tasks
Across the execution-time ablation, TGDT improves over critic-only selection on both Maze2D variants, isolating the contribution of reliability filtering. The reliability-window sweep also finds a stable range of settings, with overly short or long windows and aggressive filtering reducing return.
A.7 Internal ablation across different tasks
- TGDT improves over critic-only selection by 5 to 7 normalized points on both Maze2D variants, the largest jump in the execution-time ablation.The two modes share the same trained model and frozen critic; TGDT filters suffixes by reliability before critic ranking.Source: §A.7 · p. 17, Table 5 · p. 18, §A.8 · p. 18
A.8 Sensitivity to reliability window and risk level
- Settings Ke ∈{5, 10, 20} and α ∈{0.01, 0.05, 0.10} are stable within the seed standard deviation of the best result.Very short windows react to isolated spikes, very long windows react late, and α = 0.20 rejects more suffixes than calibration justifies, reducing return; the selected default is Ke = 10, α = 0.05.Source: §A.8 · p. 17