Trust Guided Decision Transformer

Chainesh Gautam; Raghuram Bharadwaj Diddigi; Chandramouli Kamanchi; Pankaj Dayama; Sumanta Mukherjee; Kameshwaran Sampath

cs.LGSep 25, 2026

Trust Guided Decision Transformer

Authors: Chainesh Gautam, Raghuram Bharadwaj Diddigi, Chandramouli Kamanchi, Pankaj Dayama, Sumanta Mukherjee, Kameshwaran Sampath

Organizations: International Institute of Information Technology Bangalore · IBM Research Bangalore

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.

Figures & tables

Dataset	Value-Based Methods					Conditional Sequence Modeling Methods
Gym Tasks	BEAR	BCQ	CQL	IQL	MoRel	BC	DT	StAR	GDT	CGDT	DC	VDT	TGDT
halfcheetah-medium-replay-v2	38.6	34.8	37.5	44.1	40.2	36.6	36.6	36.8	40.5	40.4	41.3	39.4 $\pm$ 2.0	45.1 $\pm$ 1.7
hopper-medium-replay-v2	33.7	31.1	95.0	92.1	93.6	18.1	82.7	29.2	85.3	93.4	94.2	96.0 $\pm$ 1.9	95.5 $\pm$ 2.2
walker2d-medium-replay-v2	19.2	13.7	77.2	73.7	49.8	32.3	79.4	39.8	77.5	78.1	76.6	82.3 $\pm$ 2.1	83.9 $\pm$ 2.7
halfcheetah-medium-v2	41.7	41.5	44.0	47.4	42.1	42.6	42.6	42.9	42.9	43.0	43.0	43.9 $\pm$ 0.7	45.9 $\pm$ 0.9
hopper-medium-v2	52.1	65.1	58.5	63.8	95.4	52.9	67.6	59.5	77.1	96.9	92.5	98.3 $\pm$ 0.1	99.2 $\pm$ 0.3

Table 1 : Offline D4RL performance. Scores are normalized D4RL scores. TGDT denotes our trust-filtered execution rule. See Section 4.1 for the evaluation protocol and baseline sourcing.

Panel of Figure 1 — Figure 1 : Rollout context mismatch and TGDT’s execution-time response. Left, vanilla DT shows persistent prediction error measured by a frozen diagnostic transition probe. Right, TGDT uses the reliability signal to adapt context length and reduce persistent high-error stretches.

Panel of Figure 2 — Figure 2 : Ablation study on Maze2D medium . Left, training-side components improve the model but do not eliminate persistent high-error stretches. Right, execution-time reliability filtering is necessary to reduce persistent mismatch. Both panels report normalized D4RL score and the longest consecutive violation run.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Hyperparameter	Value
Maximum context length $K$	20
Candidate suffix set $\mathcal{L}$	$\{1,5,10,20\}$
Reliability window $K_{e}$	10
Risk level $\alpha$	0.05
Held-out threshold percentile	95th percentile
Hard-reset cooldown	10 steps

Appendix

Table 2 : Default hyperparameters used for TGDT. The first block lists execution-time reliability parameters. The second block lists critic and behavior-regularization parameters. The third block lists transformer training parameters. Unless otherwise stated, all sensitivity experiments vary one hyperparameter while keeping the remaining values fixed at the defaults shown here.

Dataset	Pearson $r$	Spearman $\rho$	Threshold agreement	Probe-positive recall
maze2d-umaze-v1	0.88	0.81	0.86	0.85
maze2d-medium-v1	0.82	0.75	0.89	0.78

Appendix

Table 3 : Agreement between the external diagnostic probe and TGDT’s internal reliability head on Maze2D tasks. Both signals are evaluated on the same closed-loop rollout trajectories. Each signal is compared against its own held-out reliability threshold, computed with

\alpha=0.05

. Higher values indicate stronger agreement between the diagnostic signal used to motivate TGDT and the internal signal used by TGDT at runtime.

Figure 3 : Sensitivity to the reliability window $K_{e}$ and risk level $\alpha$ on Maze2D-medium and Maze2D-umaze. Each setting is evaluated over three independently trained seeds with 100 evaluation episodes per seed. Blue curves show normalized D4RL score, with shaded 95% confidence intervals across seeds. Orange dashed curves show the mean intervention rate per 1000 environment steps, where an intervention is a timestep with selected context length $L_{t}<K$ . For $K_{e}$ , intermediate windows perform best, while very short windows overreact and very long windows react late. For $\alpha$ , larger values correspond to lower held-out thresholds and more aggressive filtering. The default setting $K_{e}=10,\alpha=0.05$ lies in a stable region.

Panel of Figure 4 — Figure 4 : Training-side hyperparameter sensitivity on Maze2D tasks. Each subfigure contains four panels, varying one hyperparameter at a time while keeping the other three fixed at their default values: $\lambda_{Q}=0.01$ , $\lambda_{s}=1.0$ , $\delta_{\max}=0.05$ , and $\beta=0.05$ . Blue curves show normalized D4RL score, with shaded uncertainty across three independently trained seeds. Orange dashed curves show violation fraction, defined as the fraction of evaluation timesteps where the rolling reliability score exceeds the held-out threshold. Across both Maze2D tasks, moderate critic guidance, a trained reliability head, and bounded residual corrections improve return while keeping rollout reliability stable. Overly large critic weight or residual bound can increase action drift, raise violation fraction, and reduce performance.

Panel of Figure 5 — Figure 5 : Context-length usage on Maze2D tasks. Bars show the fraction of evaluation timesteps assigned to each candidate suffix length in $\mathcal{L}=\{1,5,10,20\}$ . TGDT uses multiple suffix lengths rather than collapsing to the shortest context, indicating that the method performs reliability-aware context selection rather than constant hard reset.

Mode	Forward passes / step	ms / step	Relative time	Est. time / 100 eps
Full context	1	$0.62\pm 0.03$	$1.0\times$	$37.2$ s
Value only, $\|\mathcal{L}\|=4$	4	$2.35\pm 0.08$	$3.8\times$	$2.35$ min
TGDT, $\|\mathcal{L}\|=4$	4	$2.48\pm 0.09$	$4.0\times$	$2.48$ min
TGDT dense, $\|\mathcal{L}\|=20$	20	$11.4\pm 0.4$	$18.4\times$	$11.4$ min

Appendix

Table 4 : Evaluation-time overhead on maze2d-medium-v1 . Estimated time is computed for 100 evaluation episodes with 600 steps per episode, giving 60k environment steps.

	Training ablations under none				Execution ablations on DT + Critic + SP
Dataset	DT	DT + SP	DT + Critic	DT + Critic + SP	Full context	Hard reset	Critic only	TGDT
MuJoCo medium-replay
halfcheetah-medium-replay-v2	35.2 $\pm$ 0.9 (31.4)	36.7 $\pm$ 0.8 (27.6)	40.1 $\pm$ 1.0 (29.8)	40.3 $\pm$ 0.9 (24.9)	39.8 $\pm$ 0.8 (25.1)	37.6 $\pm$ 1.2 (10.7)	40.8 $\pm$ 0.9 (23.6)	45.9 $\pm$ 1.1 (8.9)
hopper-medium-replay-v2	82.7 $\pm$ 1.2 (24.8)	83.1 $\pm$ 1.0 (22.1)	93.2 $\pm$ 0.9 (21.6)	93.8 $\pm$ 0.8 (18.4)	93.6 $\pm$ 0.9 (18.7)	89.5 $\pm$ 1.5 (7.9)	93.9 $\pm$ 0.8 (17.3)	94.8 $\pm$ 0.8 (6.8)
walker2d-medium-replay-v2	80.2 $\pm$ 1.4 (21.5)	79.1 $\pm$ 1.5 (20.2)	81.7 $\pm$ 1.2 (19.1)	82.1 $\pm$ 1.1 (16.8)	81.9 $\pm$ 1.2 (17.0)	83.1 $\pm$ 1.4 (7.5)	79.9 $\pm$ 1.6 (18.6)	84.4 $\pm$ 1.3 (6.9)
MuJoCo medium

Appendix

Table 5: Internal ablation across all evaluated tasks. Each entry reports normalized D4RL score, with the selected-suffix violation rate in parentheses. The violation rate is the percentage of evaluation timesteps at which the rolling next-state prediction error of the executed context exceeds the held-out reliability threshold

\tau_{\alpha}

. The first four columns are training-side ablations evaluated with standard full-context execution, denoted none : vanilla DT, DT with a state-prediction head, DT with critic-guided residual action prediction, and DT with both critic guidance and state prediction. The last four columns are execution-time ablations using the same trained DT + Critic + State Prediction model. Full context always uses the maximum context length

K

and disables execution-time control. Hard reset uses the full context unless the previous rolling score exceeds

\tau_{\alpha}

, in which case it resets the next decision to the shortest context

L=1

. Critic only evaluates the same candidate suffix lengths

\mathcal{L}=\{1,5,10,K\}

as TGDT, but ranks all suffixes directly by

Q_{\phi}(s_{t},\hat{a}_{t}^{(L)})

without checking reliability. TGDT first filters candidate suffixes using

S_{t-1}^{(L)}\leq\tau_{\alpha}

, then applies critic ranking only among the trusted suffixes. This table separates gains from training-side components from gains due to trust-guided execution.

Explore similar work

Beyond Autoregressive RTG: Conditioning via Injection Outside Sequential Modeling in Decision Transformer
May 7, 2026Yongyi Wang, Hanyu Liu, Lingfeng Li +6Transformer Architectures Sequence Modeling
School of Computer Science Peking University Beijing, China 100871
Transformers Provably Implement In-Context Reinforcement Learning with Policy Improvement
May 7, 2026Haodong Liang, Lifeng LaiOffline Reinforcement Learning In-Context Learning
Department of Electrical and Computer Engineering, University of California, Davis
Statistical Convergence of Transformer Encoder-Accelerated Robust Reinforcement Learning
Sep 20, 2026Suman Banerjee, Hiroyasu TsukamotoOffline Reinforcement Learning Value Functions
Department of Aerospace Engineering, The Grainger College of Engineering, University of Illinois Urbana-Champaign, Urbana, IL, USA

May 7, 2026, cs.LGJ/K move · Enter open · S save

Beyond Autoregressive RTG: Conditioning via Injection Outside Sequential Modeling in Decision Transformer

Yongyi Wang, Hanyu Liu, Lingfeng Li, Bozhou Chen+5

School of Computer Science Peking University Beijing, China 100871

Transformer Architectures Sequence Modeling Offline Reinforcement Learning Autoregressive Model

Open paper AskPDF arXiv

Decision Transformer (DT) formulates offline reinforcement learning as autoregressive sequence modeling, achieving promising results by predicting actions from a sequence of Return-to-Go (RTG), state, and action tokens. However, RTG is a scalar that summarizes future rewards, containing far less information than typical state or action vectors, yet it consumes the same computational budget per token. Worse, the self-attention cost of Transformers grows quadratically with sequence length, so including RTG as a separate token adds unnecessary overhead. We propose SlimDT, which removes RTG from the autoregressive sequence. Instead, we inject RTG information into the state representations before the sequential modeling step, allowing the Transformer to process only a compact (state, action) sequence. This reduces the sequence length by one-third, directly improving inference efficiency. On the D4RL benchmark, SlimDT surpasses standard DT across various tasks and achieves performance comparable to existing state-of-the-art methods. Decoupling a sparse conditioning signal from an information-rich sequence thus yields both computational gains and higher task performance.

Abstract

Explore similar work

Beyond Autoregressive RTG: Conditioning via Injection Outside Sequential Modeling in Decision Transformer

Transformers Provably Implement In-Context Reinforcement Learning with Policy Improvement

Statistical Convergence of Transformer Encoder-Accelerated Robust Reinforcement Learning