Trust Guided Decision Transformer

AI summary

Overview

Research area: Offline reinforcement learning and return-conditioned sequence modeling (Decision Transformers), applied to D4RL benchmark suites.

Technical level: Advanced. The paper assumes familiarity with Decision Transformers, offline RL critics (IQL), behavior regularization, and split conformal prediction.

Scope: The paper diagnoses a failure mode in Decision Transformer evaluation rollouts and proposes a runtime context-selection rule that filters out unreliable context lengths before value-based action ranking.

What This Paper Is About

Decision Transformers (DTs) are trained on context windows taken from offline trajectories, but at evaluation time the context is generated by the model itself, one step at a time. The authors argue that this self-generated context can drift away from the offline distribution, so the policy ends up conditioning on histories that no training example resembles — a failure mode they call rollout context mismatch. The paper's goal is to detect this drift online using the model's own next-state prediction error and to use that signal to decide which context lengths are trustworthy before a critic ranks candidate actions.

Key Contributions

1. Identifying rollout context mismatch as a DT failure mode with a measurable precursor. The authors show that persistent rolling next-state prediction error precedes return degradation, and that the signal is observable from a frozen external probe — which establishes it as a property of the rollout rather than an artifact of any specific architecture.

2. Showing that training-side improvements alone are insufficient for reliability. Auxiliary next-state prediction, frozen IQL critic guidance, and a bounded behavior-regularized residual action head each improve return, but the resulting policies still produce long above-threshold runs of high prediction error during evaluation.

3. Introducing TGDT, a trust-filtered context-selection rule that conditions the critic's ranking on a held-out reliability threshold derived from offline data. On D4RL Maze2D, AntMaze, and MuJoCo locomotion, TGDT improves over the same trained model evaluated with full context, hard reset, and critic-only suffix selection.

4. Establishing that the contribution is the order of operations, not a new critic or a larger model. In a controlled ablation using the same trained model, frozen critic, and candidate suffix set, reordering the decision (filter by reliability, then rank by value) reduces the longest above-threshold run on Maze2D-medium by approximately 8 times and improves normalized return.

Main Findings

• Vanilla DT develops persistent prediction error during rollout. On D4RL Maze2D-medium, the median probe error rises above its held-out threshold during rollout and stays elevated for long stretches, with the violation fraction increasing over episode time. The failure is driven by persistent stretches, not isolated prediction spikes.

• The internal reliability head agrees with the external diagnostic probe. On Maze2D, the two signals identify similar persistent high-error regimes with Pearson correlation r in [0.82, 0.88] and threshold agreement of at least 0.86 (reported in Appendix A.2, using a reliability window K_e = 10).

• Training-side fixes do not eliminate the problem. On Maze2D-medium, evaluating DT, DT + SP, DT + Critic, and DT + Critic + SP under the "none" execution mode, adding critic guidance gives the largest gain in normalized score, while adding state prediction alone changes the reliability profile but does not solve the control problem. Critic-guided models still produce long above-threshold runs even after the score gap with TGDT is closed.

• Context reset alone is also insufficient. Hard reset reduces some mismatch but discards all context whenever the threshold fires, and falls short of TGDT on Maze2D-medium. TGDT keeps long contexts when they are reliable and shortens only when they are not.

• Critic-only suffix selection improves return but not reliability. Using the same trained DT + Critic + SP model and the same frozen critic, critic-only selection improves score over no selection, but its longest violation run remains long — it still selects actions from contexts the model has flagged as unreliable. TGDT reduces the longest violation run by roughly an order of magnitude over critic-only while obtaining the highest score.

• TGDT's gains are concentrated where context drift can accumulate. On Maze2D, average normalized score is 79.1 for TGDT versus 74.2 for the best sequence-modeling baseline (VDT). On AntMaze, TGDT averages 84.7 versus 76.7 for VDT, and roughly doubles VDT on antmaze-medium-diverse-v0 (30.0 to 60.0), although IQL still leads among value-based methods on that task at 70.0.

• On dense-reward MuJoCo locomotion the gain is smaller. TGDT leads on the Gym average (86.3 versus VDT at 84.1), with the largest task-level gaps on halfcheetah-medium-expert-v2 (96.8 ± 0.1 versus 93.9 ± 0.1) and walker2d-medium-expert-v2 (115.1 ± 1.2 versus 110.4 ± 0.9). VDT is higher on hopper-medium-replay-v2 (96.0 ± 1.9 versus 95.5 ± 2.2).

• Adroit and Kitchen follow the same pattern with smaller margins. Adroit average: TGDT 59.2 versus VDT 57.6, with VDT higher on pen-human-v1 (126.7 ± 4.3 versus 123.2 ± 5.2). Kitchen average: TGDT 73.2 versus VDT 71.0.

• AntMaze task-level results are mixed. TGDT is higher on antmaze-medium-diverse-v0 (60.0 ± 5.2 versus 30.0 ± 2.8) but slightly lower than VDT on antmaze-umaze-v0 (98.8 ± 2.3 versus 100.0 ± 5.5) and antmaze-umaze-diverse-v0 (95.2 ± 3.1 versus 100.0 ± 4.7).

• TGDT is not model-based offline RL. The next-state predictor is never used to roll out trajectories, generate synthetic transitions, or optimize a policy through model rollouts. It serves only as a per-step reliability sensor, and the frozen critic only ranks candidate actions from trusted suffixes.

• The threshold is an empirical reliability scale, not a coverage guarantee. Because closed-loop rollouts are not exchangeable with teacher-forced validation trajectories, the authors make no formal coverage claim; three sources of non-exchangeability are named: temporal correlation through the rolling window, dependence of the rollout state distribution on the policy's own past actions, and a mismatch between the evaluation action distribution and the dataset's behavior policy.

Methodology in Plain English

Step 1 — Diagnose with an external probe. The authors train a small network that predicts the next state from a DT context plus the executed action. It is trained on the offline dataset under teacher forcing and then frozen. After each environment step, they measure the squared error between predicted and actual next state, averaged over a rolling window. This probe belongs to the rollout, not to TGDT, so it can be used to show that a plain vanilla DT also drifts.

Step 2 — Build the pieces into the DT. Three components are trained jointly: a frozen critic trained with implicit Q learning (used to rank candidate actions, never to search the continuous action space), a behavior-regularized action head that predicts a base action plus a bounded residual (keeping actions near the dataset so prediction error reflects context drift rather than action drift), and an internal reliability head that predicts the next state from the action token's representation, so the runtime rule does not need the external probe.

Step 3 — Calibrate a threshold from held-out data. The threshold is the (1 − α) quantile of rolling prediction errors collected on a held-out split under teacher forcing, using the same order statistic as split conformal prediction. With α = 0.05, this is the 95th percentile.

Step 4 — Decide at each step by filtering before ranking. The model is queried on several context-suffix lengths (default {1, 5, 10, 20}). Each suffix keeps its own rolling queue of recent prediction errors, all from steps strictly before the current one. Suffixes whose rolling score exceeds the threshold are rejected. Among the survivors, the frozen critic picks the highest-valued candidate action. If no suffix survives, the rule falls back to the shortest context (length 1). Because the score uses only past errors and the realized next state is used only after the action is committed, the rule is causal with no future leakage.

Step 5 — Compare against matched baselines. The key controlled comparison is critic-only suffix selection, which shares the same trained model, frozen critic, and candidate suffix set but drops the reliability filter.

Why This Matters

Impact on research. The paper reframes a class of Decision Transformer failures as an input-reliability problem rather than a policy-quality problem. It argues that the common question "which action is best given a context" should be preceded by "which contexts are worth ranking actions from," and it provides a controlled ablation, sharing model and critic, to support that claim. It also connects offline RL evaluation to the conformal prediction literature while explicitly declining to claim a coverage guarantee in the non-exchangeable closed-loop setting.

Real-world applications:

• Robot navigation and goal reaching — the Maze2D and AntMaze domains involve long-horizon, sparse-reward navigation where decisions late in an episode depend on earlier context.
• Continuous control for locomotion — the MuJoCo locomotion tasks represent dense-reward control, where the paper finds the benefit is smaller.
• Dexterous manipulation — the Adroit tasks cover robotic hand manipulation, where TGDT improves over the prior best on most tasks with smaller absolute margins.
• Multi-step task chains — the Kitchen tasks cover sequences of sub-tasks, matching the paper's claim that TGDT helps most when episode reward depends on decisions made after drift would have begun.

Industry relevance. The method adds only |L| extra forward passes per step (four with the default candidate set), which the authors describe as a small constant overhead for their default model size but one that grows with the candidate set and model size. That makes latency a practical design consideration for deployment. The paper's use of a frozen critic and a frozen offline dataset also fits the offline RL setting, where further environment interaction is unavailable.

Future Directions

• Restoring approximate coverage. The authors state that applying non-exchangeable conformal machinery is a natural extension, since the current threshold is not a formal coverage guarantee under the three named sources of non-exchangeability.

• Extending beyond state-based observations. Evaluation is restricted to state-based D4RL; whether the trust filter degrades gracefully with image-based observations and genuinely non-stationary environments is left as an open empirical question.

• Managing the per-step cost. With |L| = 4 the overhead is small at the default model size, but the cost grows with the candidate suffix set and model size, so latency-constrained settings may need a different operating point.

• Understanding per-suffix behavior in more depth. Appendix A.5 reports per-suffix selection frequencies showing that TGDT uses all four context lengths rather than collapsing to the shortest, which invites further study of how the trust filter redistributes context usage across tasks.

Target Audience

Researchers and practitioners working on offline reinforcement learning, Decision Transformers, and sequence-modeling policies, particularly those who evaluate policies over long rollouts and observe late-episode degradation. It is also relevant to readers interested in reliability monitoring, uncertainty calibration, and conformal prediction applied to sequential decision-making. Readers without background in offline RL, IQL critics, or transformers will find parts of the method section demanding, though the framing of the core idea — check whether the context is trustworthy before ranking actions — is stated plainly.

Authors’ abstract

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.