Uncertainty-Weighted Decision Transformer (UWDT)

• The paper introduces a teacher-guided reweighting mechanism that uses predictive entropy to emphasize safety-critical action tokens during training.
• It addresses the imbalance between frequent low-risk transitions and rare high-risk decisions in dense, autonomous driving scenarios.
• Experimental results show improved reward, travel distance, exit rate, and a significant reduction in collision rates, especially in high-density traffic.

Uncertainty-Weighted Decision Transformer (UWDT) is a Decision Transformer variant for offline autonomous driving in dense roundabout scenarios that reweights action-prediction tokens by a teacher model’s predictive entropy. In the formulation reported in “Uncertainty-Weighted Decision Transformer for Navigation in Dense, Complex Driving Scenarios” (Zhang et al., 16 Sep 2025), a standard Decision Transformer is first trained as a teacher, then frozen, and its per-token action-distribution entropy is converted into normalized and clipped weights for a student model’s cross-entropy loss. The method is motivated by the imbalance between frequent low-risk transitions and relatively rare safety-critical decisions in offline driving data, especially in tactical merging and gap-acceptance settings.

1. Definition and problem scope

UWDT is defined in the context of autonomous navigation through a four-arm, two-lane roundabout in simulation, where the ego vehicle starts on the south inbound arm and must reach a fixed north exit while interacting with circulating, merging, and exiting traffic. The problem is explicitly framed as high-level decision-making rather than low-level control: the policy selects discrete maneuvers such as accelerate, decelerate, cruise, and lane changes, while a lower-level motion planner converts those maneuvers into continuous acceleration and steering commands (Zhang et al., 16 Sep 2025).

The central premise is that standard Decision Transformers, although well suited to offline reinforcement learning and long-horizon sequence modeling, optimize an action negative log-likelihood in which routine low-risk states dominate the loss. In dense roundabouts, however, the states that matter most for safety and success are tight merges, ambiguous gap-acceptance decisions, and unusual multi-vehicle interactions. UWDT addresses this imbalance without changing the DT architecture itself. Instead, it alters the training objective so that uncertain action tokens contribute more strongly to optimization.

The method therefore belongs to a specific class of uncertainty-aware offline sequence models: not architectures with explicit Bayesian heads or transition models, but teacher-guided loss reweighting schemes in which uncertainty enters through token-level importance weights. A plausible implication is that UWDT is best understood as a hard-example reweighting strategy embedded in a Decision Transformer training loop.

2. Task formalization, representations, and baseline DT

The driving problem is formulated as an MDP,

$$\mathcal{M} = (\mathcal{S}, \mathcal{A}, \mathcal{P}, r, \gamma).$$

The state is a bird’s-eye-view occupancy grid from highway-env of shape $4 \times 41 \times 50$. Its four channels are presence, longitudinal velocity $v_x$, lateral velocity $v_y$, and a binary on_road indicator. The grid is centered on the ego vehicle, covers $100\,\text{m} \times 82\,\text{m}$, is discretized at 2 m resolution, and all features are clipped to predefined ranges and linearly scaled to $[-1,1]$. No temporal stacking beyond the transformer’s sequential context is described (Zhang et al., 16 Sep 2025).

The action space is discrete:

$$\mathcal{A}=\{\text{llc},\ \text{rlc},\ \text{acc},\ \text{dec},\ \text{cruise}\}.$$

These correspond to left lane change, right lane change, accelerate, decelerate, and maintain speed. The reward at each step depends on a collision indicator $c_t$, a speed reward indicator $v_t$ that is 1 when speed lies in $[8,16]$ m/s, and a lane-change indicator $4 \times 41 \times 50$0. The raw reward is

$4 \times 41 \times 50$1

with

$4 \times 41 \times 50$2

It is then linearly rescaled to $4 \times 41 \times 50$3:

$4 \times 41 \times 50$4

The return-to-go is

$4 \times 41 \times 50$5

with $4 \times 41 \times 50$6 in the reported DT hyperparameters.

The baseline teacher is a standard Decision Transformer trained on an offline dataset

$4 \times 41 \times 50$7

Its causal context is written as

$4 \times 41 \times 50$8

and the teacher predicts a discrete action distribution

$4 \times 41 \times 50$9

The baseline objective is the action negative log-likelihood over valid tokens. At test time, the reported experiments use greedy rollout rather than stochastic sampling.

3. Teacher entropy and uncertainty-weighted loss

UWDT uses a three-stage teacher-student procedure: train a baseline DT teacher, freeze it, compute per-token predictive entropy, and train an architecturally identical student with a weighted action loss (Zhang et al., 16 Sep 2025).

For each valid token $v_x$0, the teacher produces a categorical distribution over the discrete action set. Uncertainty is measured by predictive entropy,

$v_x$1

Larger entropy is interpreted as indicating that the teacher is less certain about which action is best in the current context. The paper associates such states with ambiguity, rarity, or safety-criticality.

Entropy is then converted into a raw weight through a monotone power map,

$v_x$2

The exponent $v_x$3 is selected to induce a desired dynamic range $v_x$4 between largest and smallest weights:

$v_x$5

Here $v_x$6 and $v_x$7 are obtained by deploying the trained teacher for 400 episodes in simulation and recording the minimum and maximum entropy over all timesteps. The raw weights are normalized per mini-batch,

$v_x$8

so that the batch-mean weight is 1, and then clipped:

$v_x$9

The reported hyperparameters are $v_y$0 and $v_y$1.

The student distribution is

$v_y$2

and its weighted loss is

$v_y$3

Weighting is therefore per-step and per-token, not per-sequence. It does not alter the transformer backbone, the state encoder, or the rollout policy structure; it alters only the scale with which individual action tokens contribute to the student’s training objective.

A recurring interpretive issue is the paper’s use of the term “epistemic.” The implemented quantity is single-model predictive entropy from a frozen teacher, which conflates epistemic uncertainty with action ambiguity or aleatoric structure. That distinction is explicitly left unresolved experimentally.

4. Architecture, dataset, and optimization protocol

The backbone shared by teacher and student combines a CNN state encoder with a causal transformer decoder (Zhang et al., 16 Sep 2025). Each occupancy grid is processed by three $v_y$4 convolutional layers with stride 2, channel progression $v_y$5, ReLU activations, batch normalization, and spatial dropout. The final feature map is flattened and projected to a 32-dimensional embedding,

$v_y$6

The transformer hyperparameters reported in Table I are a context length $v_y$7, embedding dimension $v_y$8, discount $v_y$9, $100\,\text{m} \times 82\,\text{m}$0 decoder layers, and $100\,\text{m} \times 82\,\text{m}$1 attention head. The text states that the teacher and student use the same architecture. Positional embeddings and modality embeddings are not clearly specified. A plausible implication is that standard DT-style positional handling is present, but this is not explicitly described.

Optimization uses AdamW with learning rate $100\,\text{m} \times 82\,\text{m}$2, batch size $100\,\text{m} \times 82\,\text{m}$3, $100\,\text{m} \times 82\,\text{m}$4 epochs, weight decay $100\,\text{m} \times 82\,\text{m}$5, warm-up ratio $100\,\text{m} \times 82\,\text{m}$6, and gradient norm clip $100\,\text{m} \times 82\,\text{m}$7. The paper also mentions “running-mean-square layer-norm” following Chen et al., but does not expand the implementation.

The offline dataset comprises 5,400 expert roundabout roll-outs generated by MCTS in highway-env, specifically roundabout-v0. The environment has inner and outer lane radii of $100\,\text{m} \times 82\,\text{m}$8 m and $100\,\text{m} \times 82\,\text{m}$9 m, simulation integration at 15 Hz, agent decision frequency at 2 Hz, and episodes lasting 11 s or terminating on collision. The ego vehicle begins 125 m before the south entrance at 8 m/s. Background vehicle speed is sampled from

$[-1,1]$0

with longitudinal position perturbations of standard deviation 1 m and perturbed IDM parameters. Traffic counts are sampled as

$[-1,1]$1

Several reproducibility details remain unspecified: the train/validation/test split of the offline dataset, exact MCTS configuration, dropout probabilities and padding details in the CNN, and sequence flattening mechanics for context length $[-1,1]$2. These omissions limit exact replication from the paper text alone.

5. Reported performance, behavioral findings, and interpretive cautions

The evaluation compares UWDT against CQL, SAC, BC Transformer, and a standard DT on accumulated reward, average speed, episode length, travel distance, reach exit rate, collision rate, time-to-exit, and halt duration (Zhang et al., 16 Sep 2025). The strongest direct comparison is against standard DT. UWDT improves reward from $[-1,1]$3 to $[-1,1]$4, average speed from $[-1,1]$5 to $[-1,1]$6, travel distance from $[-1,1]$7 to $[-1,1]$8, reach exit rate from $[-1,1]$9 to $$\mathcal{A}=\{\text{llc},\ \text{rlc},\ \text{acc},\ \text{dec},\ \text{cruise}\}.$$0, and reduces collision rate from $$\mathcal{A}=\{\text{llc},\ \text{rlc},\ \text{acc},\ \text{dec},\ \text{cruise}\}.$$1 to $$\mathcal{A}=\{\text{llc},\ \text{rlc},\ \text{acc},\ \text{dec},\ \text{cruise}\}.$$2. Time-to-exit improves from $$\mathcal{A}=\{\text{llc},\ \text{rlc},\ \text{acc},\ \text{dec},\ \text{cruise}\}.$$3 s to $$\mathcal{A}=\{\text{llc},\ \text{rlc},\ \text{acc},\ \text{dec},\ \text{cruise}\}.$$4 s, and UWDT reports a halt duration of $$\mathcal{A}=\{\text{llc},\ \text{rlc},\ \text{acc},\ \text{dec},\ \text{cruise}\}.$$5 s.

Traffic density is partitioned into low density $$\mathcal{A}=\{\text{llc},\ \text{rlc},\ \text{acc},\ \text{dec},\ \text{cruise}\}.$$6, medium density $$\mathcal{A}=\{\text{llc},\ \text{rlc},\ \text{acc},\ \text{dec},\ \text{cruise}\}.$$7, and high density $$\mathcal{A}=\{\text{llc},\ \text{rlc},\ \text{acc},\ \text{dec},\ \text{cruise}\}.$$8, where $$\mathcal{A}=\{\text{llc},\ \text{rlc},\ \text{acc},\ \text{dec},\ \text{cruise}\}.$$9 is the number of interacting vehicles. The qualitative conclusion is density-dependent: at low density, UWDT, DT, and SAC are near-optimal; at medium density, BC and CQL deteriorate, SAC becomes more cautious, while UWDT and DT remain strong; at high density, UWDT dominates, DT remains good but shows slightly more collisions, and SAC often stops and waits too long. This is consistent with the paper’s motivating claim that uncertainty weighting matters most in dense, ambiguous interactions.

The paper also reports entropy statistics during testing. DT has min $c_t$0, max $c_t$1, and average $c_t$2, while UWDT has min $c_t$3, max $c_t$4, and average $c_t$5. The authors interpret the lower maximum entropy for UWDT as evidence of better calibrated uncertainty or more stable behavior. Strictly, these measurements are the deployed models’ action-distribution entropies during rollout, not a direct analysis of the teacher-generated training weights.

A notable limitation is the absence of controlled ablations beyond the DT-versus-UWDT comparison. There are no systematic studies of alternative weight mappings, different $c_t$6 or $c_t$7, teacher quality, entropy estimators, context length, or encoder variants. There is also no direct visualization of high-entropy states, no token-weight analysis, and no uncertainty heatmaps. Another interpretive caution concerns the paper’s conclusion, which states that UWDT conditions action selection on both desired return and an uncertainty measure and exploits epistemic uncertainty at inference time. The method description, however, primarily uses uncertainty during student training; the deployed student otherwise acts like a standard DT trained with weighted supervision.

6. Relation to adjacent uncertainty-aware and weighted DT methods

UWDT sits within a broader family of Decision Transformer modifications in which either uncertainty or weighting alters what the model learns from. Two particularly relevant comparisons are UNREST for stochastic driving and DT4Rec for retention-oriented recommendation.

Method	Signal	Primary mechanism
UWDT	Teacher predictive entropy	Per-token weighted action cross-entropy
UNREST	KL divergence between return predictors	Uncertainty-gated return conditioning and cautious planning
DT4Rec	Reward value	Weighted reward prompt and weighted contrastive loss

UNREST, introduced as an uncertainty-aware Decision Transformer for stochastic driving, is not a literal uncertainty-weighted loss method. Instead, it estimates environmental uncertainty by the KL divergence between two return distributions, segments trajectories into certain and uncertain parts, replaces global returns with truncated returns in certain regions, and drops return conditioning in uncertain regions via dummy tokens. It also uses uncertainty-guided planning at inference. In that sense, UNREST is better characterized as uncertainty-aware conditioning and policy switching than as token-level loss weighting (Li et al., 2023).

DT4Rec, by contrast, is a weighted Decision Transformer for recommendation rather than an uncertainty-aware one. Its weighting appears in two places: a reward prompt formed by a weighted aggregation of learnable meta embeddings, and a weighted contrastive loss in which lower-reward negatives receive larger penalties. The weights are deterministic functions of reward, not uncertainty, and the method contains no explicit uncertainty model (Zhao et al., 2023).

This comparison clarifies the terminological boundaries of UWDT. In the literal sense, UWDT is a loss-reweighted Decision Transformer driven by a teacher-derived uncertainty proxy. UNREST shows that uncertainty can instead alter return conditioning and planning, while DT4Rec shows that weighting can be important even without uncertainty estimation. Taken together, these lines of work suggest a spectrum of DT modifications: reward-weighted, uncertainty-weighted, and uncertainty-conditioned. A plausible implication is that future formulations may combine these ideas, for example by coupling token-level weighting with uncertainty-aware return selection or by replacing single-model entropy with better calibrated uncertainty estimators.