Uncertainty-Aware Decision Transformer for Stochastic Driving Environments
Authors: Zenan Li
Fan Nie
Qiao Sun
Fang Da
Hang Zhao
Venue: CoRL 2024 (2024)
Keywords: Self-Driving
Decision Transformer
Uncertainty-Aware Planning
TL;DR: We present UNREST, an uncertainty-aware decision transformer to apply offline RL in stochastic driving environments.

Abstract: Offline Reinforcement Learning (RL) enables policy learning without active interactions, making it especially appealing for self-driving tasks. Recent successes of Transformers inspire casting offline RL as sequence modeling, which, however, fails in stochastic environments with incorrect assumptions that identical actions can consistently achieve the same goal. In this paper, we introduce an UNcertainty-awaRE deciSion Transformer (UNREST) for planning in stochastic driving environments without introducing additional transition or complex generative models. Specifically, UNREST estimates uncertainties by conditional mutual information between transitions and returns. Discovering 'uncertainty accumulation' and 'temporal locality' properties of driving environments, we replace the global returns in decision transformers with truncated returns less affected by environments to learn from actual outcomes of actions rather than environment transitions. We also dynamically evaluate uncertainty at inference for cautious planning. Extensive experiments demonstrate UNREST's superior performance in various driving scenarios and the power of our uncertainty estimation strategy.

Uncertainty-Awarc Decision Transformer for Stochastic Driving Environments
Zenan Lil.2, Fan Nie2.3, Qiao Sun', Fang Da', and Hang 7
TTinngua, SShanghai Qi Zhi, Stanford, QCraft, *Corresponding author
Emails: li-zn23 mails. tisillsssighha.eee edu edu.cn, en, niefan @ stanford. edu, hangzhao mail. tsingsua. edu,
Offline Reinforcement Learning (RL) enables policy learning without active interactions, making it especially appealing for self-driving tasks. Recent successes of Transformers inspire casting offline RI. as sequence modeling, which, however, fails in stochastic environments with incorrect assumptions that identical actions consistently achieve the same goal. In this paper, introduce an UNcertainty-awaHR deciSion Transformer (UNREST) for planning in stochastic driving environments without introducing additional transition or complex genera- tive models. Specifically, UNREST estimates uncertainties by conditional mutual information between transitions and returns. Discovering 'uncertainty accumula- tion' and 'temporal locality" properties of driving environments, we replace the global returns in decision transformers with truncated returns less affected by environments to learn from actual outcomes of actions rather than environment transitions. We also dynamically evaluate uncertainty at inference for cautious planning. Extensive experiments demonstrate UNREST's superior performance in various driving scenarios and the power of our uncertainty estimation strategy.
Abstract:
Keywords: Self-Driving, Decision Transformer, Uncertainty-Aware Planning.
Introduction
Safe and efficient motion planning has been recognized as a crucial component and the bottleneck in self-driving systems [1]. Nowadays, learning-basee planning algorithms like imitation learning (IL) [2] and reinforcement learning (RL) [3] have gained prominence with the advent of intelligent simulators [4] and large-scale datasets [5]. Building on these, offline RL [6, 7] becomes a promising framework for safety-critical driving tasks to learn policies from offline data while retaining the ability to leverage and improve over data of various quality [8, 9]. Nevertheless, the application of offline RL approaches still faces practical challenges: (1) The driving task requires conducting long-horizon planning to avoid shortsighted erroncous decisions [10]; (2) The stochasticity of environment objects during driving also demands real-time responses to their actions [6, 11].
Recent works propose to leverage the capability of Transformers 12, 13] by reformulating offline RL as a sequence modeling problem [14], which naturally considers outcomes of multi-step decision- making and has demonstrated efficacy long-term credit assignment [14, 15]. Typically, they train models to predict an action based on the current state and an outcome in hindsight such as desired future return (i.e., reward-to-go). However, existing works [16, 17, 18] have pointed out that these decision transformers (DTss) tend to be overly optimistic in stochastic environments because they incorrectly assume that actions, which successfully achieve a goal once, can consistently do so in subsequent attempts. This assumption is easily broken in stochastic environments, as the goal can be achieved by accidental environment transitions. In Fig. (()), identical turning actions may yield entirely different outcomes w.r.t. the aggressive or cautious behavior of the surrounding blue vehicle.
The key overcoming the problem is distinguishing between outcomes of decisions and environment transitions, and training models pursue goals not affected by environmental stochasticity. To the best of our knowledge, only limited works [11, 17, 18] attempt to solve the problem. Generally, they fit
8th Conference Robot Learning (CoRL. 2024), Munich, Germany.

---

12
202
50ovtam
200
(m)
Prodicted Retu
Roilout Retur
maps
162
Tax Figt
IA Soage tran
0.0
0.2
0.4
05 2.8
1.0
50
100
150 Timestep
200
253
300
Normalized Rotum (Ground Truth!
(a)
(b)
(c)
Figure Motivations of UNREST. (a): Example driving scenario where the variance of return increases when accounting for multiple tasks. (b): Calibration results of return distribution over
future 1,000 steps are more uncertain than 100 steps. (c): Rollout returns/distances of sequences maximizing the return of future 100, 500, and 1. 1,000 steps in the dataset are close to each other.
state transition model, either for sampling pessimistic actions [11] from VAES [19], or to disentangle conditioning goals from environmental stochasticity [17, 18] by adversarial training [20]. Besides adding complexity, these methods only applicable when transition functions and generative models can be learned adequately. This is often not the case for driving because of the uncertainty brought by complex interactions and partial observability [21, 22]. Furthermore, driving trajectories encompass stochasticity over an excessive number timesteps, challenging VAE/adversarial training [23], and diluting the information in hindsight regarding current step decision making..
In this paper, we initially customize DTs for stochastic driving environments without generative training. Specifically, our insight comes from Fig. (a)): When going straight, the cumulative rewards from Task & II (termed as the global return) contain too much stochastic influence to provide effective supervision. In contrast, a viable strategy involves conditioning solely on the truncated return from Task I to mitigate the impact of environmental stochasticity (less stochastic timesteps, lower return variance), which still preserves rewards over sufficient timesteps for action optimization. Experiments validate the point, and summarize the following properties of driving environments:
Property 1 (Uncertainty Accumulation) The impact of environmental stochasticity on the return distribution accumulates while considering more timesteps, as validated in Fig. 1(b).
Property (Temporal Locality) Driving can be divided into independent tasks, where we only need to focus on the current task without considering those much later. Hence, optimizing future returns over sufficiently long horizon approximates global return optimization, as shown in Fig. I(c).
The remaining problem is to specify the span of truncated returns. Specifically, our proposed UNcertainty-wwaRE deciSion Transformer (UNREST) quantifies the impacts of uncertainties by the conditional mutual information between transitions and returns, which bypasses the complexity associated with generative modeling. Sequences are then segmented into certain and uncertain parts accordingly. In 'certain parts' with minimal impact of uncertainty, we the conditioning goal as the cumulative rewards to the segmented position (with the number of timesteps), which reflects the actual outcomes of decisions and can be generalized to attain higher rewards. In contrast, in 'uncertain parts' where the environment is highly stochastic, we disregard the erroneous information from returns (with dummy tokens as conditions) and let UNREST imitate expert actions. Dynamic uncertainty evaluation integrated during inference for cautious planning. Key contributions are:
We present UNREST, an uncertainty-awar decision transformer to apply offline RI. in stochastic driving environments. Codes are public at /tap::///gttaan com htt///miiaalllnnieeeeee /o////CoRL222N2-322 Recognizing properties driving environments, we propose a model-free uncertainty measurement and segment sequences accordingly. Based on these. UNREST replaces global returns DTs with truncated returns less affected by environments learn from the actual outcomes of agent actions. We extensively evaluate UNREST CARLA [4], where it consistently outperforms strong baselines (5.2% and 6.5% absolute driving score improvement in seen and unseen scenarios). Additional experiments also prove the efficacy of our uncertainty estimation strategy.
2
Related Works
Offline RL Sequence Modeling: A recent line of works exploits Transformer's capability into RL by viewing offline RL sequence modeling problem. Typically, Decision Transformer

---

(DT) [14] predicts actions conditioned on target returns and history sequences. DTs naturally consider outcomes of multi-step decision-making and have demonstrated efficacy in long-term credit assignment [14]. Building on this, Trajectory Transformer (TT) [15] further exploits the capacity of Transformers through jointly predicting states, actions, and rewards and planning by beam searching. Generalized DT [24] reveals the theoretical basis of these models: they are trained conditioned on hindsight (e.g., future returns) to match the output trajectories with future information statistics.
However, in stochastic environments like autonomous driving, certain outcomes can be achieved by accidental environment transitions and, thus, cannot provide effective supervision for actions. Notably, this is in fact a more general problem to do with all goal-conditioned learning algorithms [25, 26].
To tackle this issue, ACT [27] and CGDT [28] approximate value functions as conditions without explicitly eliminating environmental stochasticity. ESPER [17] adopts adversarial training [20] to learn returns disentangled from environmental stochasticity condition. DoC [18] utilizes
variational inference [19] to learn latent trajectory representation, new condition that minimizes
the mutual information (so disentangled) with environment transitions. Besides, SPLT [11] leverages a conditional VAE [19] to model the stochastic environment transitions. As the driving environment contains various interactions that are difficult to model [21], our study proposes a novel uncertainty estimation strategy to customize DT's without transition or generative models. Besides, different from EDT [29] that dynamically adjusts history length during inference, we focus on segmenting sequences and replacing training conditions to address the overly optimistic problem.
Uncertainty Estimation: One typical method for uncertainty estimation is probabilistic bayesian approximation, either through dropout [30] or VAEs [31], which computes the posterior distribution of model parameters. Besides, the deep deterministic methods propose to estimate uncertainty by exploiting the implicit feature density [32]. In this work, we adopt the ensemble approach [33] widely used in the literature of RL [34], to jointly train K variance networks [35] for estimating the uncertainty of returns. Related works about vehicle planning are discussed App. A.I.
Preliminary
Keeping notations concise, we use subscripts or numbers for variables at specific timesteps. Greek letter subscripts for parameterized variables, and bold symbols for those spanning multiple timesteps.
3.1 Online and Offline RL
We consider learning in Markov decision process (MDP) [36] denoted by the tuple M AS p), where S and A the state space and the action space, respectively. Given states s, E S and action a A, P(s'Is,a) S x A x S - [0. 1] is the state transition function and r(s, a) : S x A - R defines the reward function. Besides, y E (0, 1] is the discount factor. The agent takes action a at state s according to its policy x<ass) : S x A - [0. 1J. At timestep I E [1,7], the accumulative discounted reward in the future, named return (i.e., reward-to-go), is R, = - y" rr.. The goal of online RL is to find policy π that maximizes the total expected return: J Ba,-n((Iss) (.ss. -.s.,,-((Ils
ax) by learning from the transitions (s, a, r') through interacting with the real environment. In contrast, Offline RL makes use of a static dataset with N trajectories D [TT) collected by certain behavior policy лい to learn policy that maximizes J, thereby avoiding safety issues during online interaction. Here {s), ar, Fr, ) is a collected interaction trajectory composed of transitions with horizon T. 3.2 Offline RL as Sequence Modeling
Following DT's [14], we pose offline RL as sequence modeling problem where we model the probability of the sequence token x, conditioned on all tokens prior to it: (x(x,(x<>) where x denotes tokens from step I to (1- 1). DTs learn policy under retum-condittoned setting where the agent at step / receives an environment state st, and chooses an action a, conditioned on the future return R, ET Fr. This leads the following trajectory representation:
(R 81, d1, R2, 82, .2...... R. R7.88.44))
with the objective minimize the action prediction loss, i.e. maximize the action log-likelihood:
LDT B Σ/= log pa s.)
This objective is the cause for DTs' limitations in stochastic environments, as it over-opiimittically assumes actions in the sequence can consistently achieve the corresponding returns. At inference
3

---

61-1
Ge+c
Global Return Emb.
Decision Transformer
Return-Span
R
Rz-1
Dum.
Dum
Szec
Price
Peace
Prace
Pences
Percio
Breck
Return Transformers
Uncertainty Brahation.
Two return prediction transformers are trained for uncer- tainty estimation. The sequence then segmented into certain (no background) and uncertain (orange background) parts w.r.t. estimated uncertainties, with 'certain parts' conditioned returns to the next segmentation positions, and dummy tokens in 'uncertain parts'. Upper: The same architecture as DTs is used for action prediction, except that we add a return-span embedding to the truncated return embedding, and concatenate the global return embedding to the transformer output.
Figure 2:
Overview of UNREST. Lower:
time, given prescribed high target return, DTs generate actions autoregressively while receiving new states and rewards to update the history trajectory.
4 Approach: UNREST
4.1 Model Overview
An overview of the proposed approach UNREST is illustrated Fig. 2. To address the overly optimistic issue, our key idea is quantify the impact of environmental uncertainty, and learn to perform aggressively cautiously according different levels of environmental stochasticity.
To achieve this, we train two return transformers with different trajectory inputs to identify the impact of environmental uncertainty, obviating the need for complex transition and generative training. The expert sequences are then segmented into certain and uncertain parts w.r.t. estimated uncertainties, each with relabeled conditioning goals to facilitate the decision transformer to learn from outcomes of decisions rather than environment transitions. At time, an uncertainty predictor is involved for cautious decision-maker making at different states. In the following, we introduce each module with details. 4.2 Transformers for Uncertainty Estimation
Instead of conventional uncertainties that reflect variances of distributions [34], in this paper we estimate the impact of environmental stochasticity, what really matters for policy learning, as an indirect measure of environmental uncertainty. In particular, we propose to model the impact the transition (s,-1, a--1 > s,) on return R, through their conditional mutual information [37].
Specifically, two 'return transformers' are trained to approximate return distributions. and actions are embedded by linear layers (.) and f The obtained embeddings are then sequentially processed by two transformers Tg, (.) and Tp. (.) separately for return prediction:
Initially, states
Xx, fo(s), Na, fG
Ss, Tpp (.... Na, Sa, Ty. (a.
Denoting rot (( K=)
| two models feed their respective outputs into variance networks [35] to predict Gaussian distributions NV-) of returns, which are optimized by maximizing logeleiklihoos:::
Pia (R (Hou (Ka, D (Far1)), Lreum (va) ET-2 Σ log Pv. (R,
(4)
Pi, (Rr (Rr sx) = (Mp, (Z,,). T,,,, (Rs)),, Lretm (onnum(..) B--ll Σ log pv,, ( st)].
Practically, the networks implemented ensembles, which together form Gaussian Mixture Models (GMM) [38] to better capture the return distributions in App. A.2. Compared to challenging

---

high-diment sigh-dimensional transition function learning and complex generative training, the return distributions can be better learned with fewer resources. Finally, the impact of the environmental stochasticity, i.e. the conditional mutual information H (R,/TZZ) |T=)) - H(R, (R st) is estimated by the KL divergence between these two distributions, as a means to measure the environmental uncertainty at timestep r:
Ut D KL (p&, (R st), Pv. ( (R/l=2))
(5)
A larger divergence implies more information brought by &, for predicting R, condition on the history, i.c. more influence on the return from the stochastic transition (5-11. ar-1 - sy).
4.3 Transformer for Sequential Decision-Making In this section, we step to aid DT training in stochastic driving environments. As discussed in Sec. 4.1, we expect UNREST to segment sequences into certain and uncertain parts according to uncertainty estimations and learn to perform aggressively or cautiously within them, respectively. To achieve this, we propose to replace the conditioning global returns with truncated returns in 'certain parts', which are less affected by the environment due to uncertainty accumulation (Prop. 1). thus reliably helping the planner generalize higher returns after training. In contrast, in 'uncertain parts', the seemingly high return actions may cause safety issues due environmental uncertainty. Therefore, ignore the stochastic returns, setting conditions as dummy tokens for behavior cloning.
Segmentation strategy is thus a crucial point to reinvent DTs. To distinguish between different levels of environmental stochasticity, we define an uncertainty threshold Next, we estimate uncertainties u, by Eq. 5 for each timestep and record those larger than e as uncertain. The sequence is then divided into certain and uncertain parts according to these marked uncertain timesteps.
An intuitive example of our segmentation strategy is illustrated in the lower part of Fig. 2. Specifically, the 'certain part' begins at timestep with uncertainty below the threshold and persists until segmentation occurs at an uncertain timestep. Subsequently, the 'uncertain part' commences with the newly encountered uncertain timestep and encompasses subsequent ones until the final c- I timesteps are all identified as certain. Since uncertain timesteps may occur intensively and intermittently over a particular driving duration (e.g., at an intersection), the hyperparameter e is introduced to avoid frequent switching between certain and uncertain parts, and ensure a minimum length of segmented sequences. Finally, we represent the segmented sequence as follows:
- (h1, RT, S1,,1, h22 R/, S2, a2 a2,...,,yy Ri ST,S),
hc-1 where the conditioning return is modified as R = - TItR. which only involves rewards in the next h, steps (called the return-span). In 'uncertain parts', h is always set as zero (i.e., the dummy tokens). In this way, UNREST will learn to ignore the trivial conditioning targets and directly imitate expert actions instead of being misguided by the uncertain return information. In 'certain parts', , the return-span is set to the number of timesteps from the current to the next segmentation step, so the conditioning return R² incorporates the maximum duration that does not include any uncertain timesteps. With this segmentation and condition design, we derive the following proposition:
Proposition 1 (UNREST Aligmment Bound) Assuming that the rewards obtained are determined by transitions (s,a demonstrations, then the discrepancy between target truncated returns and URNEST's rollout returns is bounded by a factor of environmental stochasticity and data coverage.
s') at each timestep and UNREST is perfectly trained to fit the expert
It reveals that UNREST can generalize to achieve high returns with bounded error under natural assumptions in driving environments. The formal statement and proof are left to App. B.
Notably, due to sequence segmentation, the truncated returns no longer encompass all timesteps from the present to the end of the sequence, which necessitates the return-span as an additional condition to provide information about the count of timesteps involved in returns. This enables UNREST to learn to get return Rh over future h, timesteps. Otherwise, the model may be confused by the substantial differences in the magnitude of return conditions with varying timesteps [39].
Policy formulation:
Unlike return transformers, DT takes the segmented sequence - as input:
8x, ~ To(.,, (..,,R!! XR. Na, No,), where
(7)
RR =x(((R)) + fø (hr), Xs, (st), Xa, = (a,).
Xgh = f
5

---

Here we use similar notations as return transformers in Sec. 4.2 except that models are parameterized by 0. In the above formulation, a return-span embedding S (hy) is added to the return embedding, which bears the semantic meaning of how many timesteps are involved in the target return. Besides, to get the action distribution, the non-truncated global-retum embedding I (Re) is optionally added the output & to provide additional longer horizon guidance for planning. Using [ || +] to denote the concatenation of two vectors along the last dimension, the final predicted action distribution is:
sog Ro(a To(allT ................................ st) N (1 (R,))), or (R.))).
(8)
The learning objective is modified from Eq. 2 of DT, with detailed training process in App. C:
seg LUNREST (0) = B D Σ log zo (u (ee...) st)
4.4 Uncertainty-guit Planning During inference, we also differentiate be-
Algorithm 1: Uncertainty-guider planning
tween certain and uncertain states to ac-
Input: History т, return horizon H1, state St, policy πe. uncertainty model uy. return model R return percentile 17.
count for environmental stochasticity. To enable real-time uncertainty evaluation, we introduce a lightweight uncertainty
First, update target returns.
prediction model u We dynamically query the predictor at each timestep to ob- tain the uncertainty measure. If the current state transition is highly uncertain, we set the conditioning target to a dummy token
Update target global return R. < -R:-1 r--
if h,-1 == 1 or u (St-1) is True then Reset return span to return horizon h, H; Predict target return R R S1,
else
to facilitate cautious planning, consistent with training. Conversely, the planner acts aggressively at states with certain transi- tions to attain high target returns. Neural networks or heuristics can be used to im- plement the predictor, whose results are summarized in App. F.4. Practically, we 12
Update return span h, - h1-1 -
Update target return R/: - R/ Tr-15
Second, evaluate new state is uncertain or not. 10 if (w(s)) is True then
11
Reset conditioning target (h,, R²)) -
Finally, select action using trained policies. Sample and take action a, (a St);
choose the KD-Tree [40] for its compu- 13
tational efficiency and estimation perfor- 14 Update history r + T U (h, R., R. St, az).
mance, with states as tree nodes and un- Output: Action a, and history
certainties estimated by return transformers node values.
Different from the conventional planning procedure of DTs, UNREST requires the specification of not only the target global return R, but also the initial return-span h; = H and the target truncated return Rh. After segmentation, the effective planning horizon of the trained sequences is reduced to the return-span In- Once h, reaches 1, we need to reset the target return and the return-span. span. Practically, we simply reset h, to a fixed return horizon H. For the truncated return, we train return prediction model R (-) similar that defined in Eq. 4 and take the upper percentile 1) of the predicted distribution as the new target return to attain. The hyperparameter 1) can be tuned for higher target return. We do not need to consider targets at 'uncertain states' since they are just set as dummy tokens. The planning process is summarized in Alg. 1.
5 Experiments
In this section, we conduct extensive experiments to answer the following questions. Q1: How does UNREST perform in different driving scenarios? Q2: How do components of UNREST influence its overall performance? Q3: Does our uncertainty estimation possess interpretability?
5.1 Experiment Setup Datasets: The offline dataset is collected from CARLA [4] with its built-in Autopilot. Specifically, we collect 30 hours of data from 4 towns under 4 weather conditions, saving tuple (s,,a,,r) at each timestep. More details about the state, action, and reward definitions are left to App. D.
Metrics: We evaluate models at training and new driving scenarios and report metrics from the CARLA challenge [4, 41] to measure their driving performance: driving score (the most significant metric that accounts for various indicators like driving efficiency, safety, and comfort), infraction score, route completion, and success rate. Besides, we also report the normalized rewards (the ratio of total return to the number of timesteps) reflect driving performance timestep level [2].

---

Table 1: Driving performance on train (new) town and weather conditions in CARLA. Mean and standard deviation are computed over 3 seeds. All metrics are recorded in percentages (%) except the normalized rewards. The best results are bold and our method is colored in gray.
Planner
Driving Score
Success Rote|
Route Completion"
Infraction Score
Normalized Rewards
BC
$1.9+ (45.6 5.2)
37.9-4.7 (36.3 8.6)
6.
7.8)
54.5+ 1 (47.0 4.8)
0.63 0.02 (8:61 (D.04)
MARWIL [2]
24.3 (173344
41 (122222
8144
(
27.9 (18.4 4.4)
0.65 0.02 (0.63 0.01)
[9L
(5.722.(
80.5
(72.7
3..4
0.65 0.03 (0.62 = 0.02)
[QI-(43)
(6.2233
60.2-62(42++2)
2.2)
77.2+46
(6..8
(9.44
(62.6 6.2)
(0:60 0.01)
DT [14]
1.7
(7.024
0.66 0.01 (0. 0.02)
[15]
353433.3548+2.
45.95525500=44
784=581770=4.)
63.6 3.6 (56.0
0.74 0.02 (0 (0.64 0.03)
(39.5..33 (5.
47.
8.3.7
0.05)
ESPER (17)
43.5
2.0)
77.3
(847
56.2
2.22+2221 (46.9
0.64 0.04 (0.63 0.02)
ACT [27]
2.
4.1 3.3)
1.
61.8
2.5 (53.2
0.66 a 0.03 (0 68.65 0.05)
DoC[18)
36.9 (54.4
49.9
2...
60.3
(0.65 0.04)
UNREST
63.5 (62.9 4.0)
(57.5
83.8 3.1
(90.0 6.00
70.2
(62. 3.8)
0.64 0.04 (0.65 = 0.03)
Expent [4]
6.05
(6
22.5
0.72 0.01 (0.69 0.01)
Baselines: First, we choose two IL baselines: BC and MARWIL [42]. We also include state-of-
tae-art offline RL baselines: CQL [9] and IQL [43]. Besides, we compare two classic Transformer-
based offline RL algorithms: DT [14] and TT [15]. We also select ACT [27] as baseline, which
approximates value functions as the hindsight condition. Finally,
adopt three algorithms as
rigorous baselines: SPLT [11], ESPER [17], and DoC [18]. These algorithms fit state transition models and employ generative training to mitigate DT's' limitations in stochastic environments.
5.2 Driving Performance
We evaluate all the models' (baselines inherited from their public implementations) performance at training (Town03) and new (Town05) scenarios (Q1). The results are summarized in Tab. 1.
Analyzing the results, we first notice that DT performs worst among all sequence models, with a gap in driving score compared IQL. We attribute this to the uncertainty of global return it conditions on. ACT gains better performance through conditioning approximated value functions. However, it does not fundamentally decouple the effects of environmental stochasticity, resulting in suboptimal infraction scores. ESPER and DoC also perform poorly in both scenarios, which may result from ineffective generative training from long-horizon and complex driving demonstrations.
To tackle stochasticity, SPLT learns to predict the worst state transitions and achieves the best infraction score. However, its overly cautious planning process leads it to stand still in many scenarios like Fig. 3(c), resulting in an extremely poor route completion rate and normalized rewards. TT instead learns transition model regardless of environmental stochasticity and behaves aggressively. It attains the highest normalized rewards at training scenarios since the metric prioritizes planners that rapidly move forward (but with low cumulative rewards because of its short trajectory length caused by frequent infractions). In unseen driving scenarios, TT often misjudges the speed of preceding vehicles, resulting in collisions and lower normalized rewards (than ours) as in Fig. 3(a).
Notably, UNREST achieves the best driving score, route completion, and success rate in both scenarios, and the highest normalized rewards in new scenarios without the need to learn transition or generative models. Its driving score surpasses the strongest baselines, TT over 5% in training scenarios, and SPLT over 6% in new scenarios, achieving a reasonable balance between aggressive and cautious behavior. indicates that the truncated returns successfully mitigate the impacts of uncertainty and provide effective supervision. In App. F.2 we validate that UNREST runs significantly faster while consuming less memory than TT and SPLT, which additionally fit transition models. More results like uncertainty calibration, hyperparameters (e.g. the uncertainty threshold e), environments beyond driving (e.g. D4RL.), and new policies other than BC at uncertain states are studied in App. F. 5.3 Ablation Study
In this section, we conduct ablation experiments by removing key components (global return embed- ding, return-span embedding, ensemble-based uncertainty estimation, and the uncertainty-guiid guided planning process) to explore their impacts on UNREST (Q2). Results are shown in Tab. 2 and 3.
Eliminating any part of the four components harms the driving score in new scenarios. Among them, the giobal return embedding has the slightest impact, which suggests that the highly uncertain global return does not provide much new information over the truncated return due to temporal locality (Prop. 2). When the return-span embedding is removed, the absolute driving score drops by about 6%.
7

---

(a) Lane changing: Failing case TT
(b) Lane changing: UNREST's performance
'e--1 Ne !ay (c) Light crossing: Failing case SPLT
heset (d) Light crossing: UNREST's performance
herd hesk
Figure 3: UNREST performs well failing cases of TT and SPLT. White rectangles ego-vehicles.
Table 2: Ablation study results for UNREST on
Table 3: Ablation study results for UNREST on
train town and train weather conditions.
town and new weather conditions.
Planner
Stare
Roce
Co.
Infrecion Screff
Resankd
Planerr
Driving Samze
Readf
Screef
Noem. Rewank]
Winn global ench
64.5
5KR -
8-2-44
68.
0.6%
Win global onc
61.8
$5. 47.6=37
w 79.254
2.7
0.64 0.01
Wio iet spea eal.
57.0
23.3
23.6233
63.62.11
0.57 0.00
Wio 1tt spon call
$6.2=44
58.22.44
0.59 0.02
essenable
60.4-3 34
51.3263
802=44 00:531
0.02 0..0
Wio ensemable
57.8=22
488=41
83.4-3a
56.7=27
0.1 =004
gundeel plan
guided plan
model
28
Urovrsiney Volue
Timeetap
Fromo
Fraee
(a)
following:
Uncertainty curve and state changing process
Unorra my Thnutali
Unceratriy Vauue
Fame
(b) Light crossing: Uncertainty curve and state changing process
Figure 4: Visualizations of UNREST's uncertainty estimation results.
This implies that the introduction of return-span embedding provides necessary information about the timesteps needed achieve the target return. Removing the ensemble of return transformers induces significant performance drop in both scenarios, because simple Gaussian distribution cannot express the return distribution well. Finally, canceling test time uncertainty estimation causes substantial decline in driving and infraction scores, which proves the importance of cautious planning. 5.4 Uncertainty Visualization Finally, we verify the interpretability of our uncertainty estimation through visualizations (Q3). In Fig. 4(a), we observe an increase in uncertainty the ego-vehicle enters the lane, owing the lack of knowledge about the other vehicle's 's behavior. This uncertainty decreases below the threshold when the vehicle stabilizes in the following states. Fig. 4(b) shows green light crossing scenario. While approaching the light, the unpredictable light state causes the uncertainty value to rise quickly.
Conclusion and Limitations
The paper presents UNREST, an uncertainty-aware decision transformer to apply offline RL in stochastic driving environments. Specifically,, we propose novel uncertainty measurement by computing the divergence of return prediction models, bypassing complex transition and generative training. Then, based properties discover at driving, segment sequences w.r.t. estimated uncertainty and adopt truncated returns conditioning goals. This new condition helps UNREST learn policies that are less affected by stochasticity. Dynamic uncertainty estimation is also integrated at inference for cautious planning. Empirical results demonstrate UNREST's superior performance, lower resource occupation, and effective uncertainty estimation in various driving scenarios.
One limitation of UNREST is that its inference process is somewhat complex with auxiliary models and hyperparameters. A possible improvement direction is directly integrating return and uncertainty predictions into one single model. Besides, though UNREST is evaluated in the CARLA simulator, we believe it can surmount the sim-to-real gap and benefit practical autonomous driving.

---

Acknowledgments
We thank all the authors for their valuable contributions throughout the paper. This work started when Zenan Li interned the Shanghai Qi Zhi Institute. Thanks Shanghai Qi Zhi Institute and QCraft Inc. for their support in computation resources.
References
[1] E. Yurtsever, J. Lambert, A. Carballo, and K. Takeda. A survey of autonomous driving: Common practices and emerging technologies. IEEE access, 8:58443-58469 2020.
[2] A. Hu, G. Corrado, N. Griffiths, Z. Murez, C. Gurau, H. Yeo, A. Kendall, R. Cipolla, and
J. Shotton. Model-based imitation learning for urban driving. In Advances in Neural Information Processing Systems, 2022.
[3] J. Chen, B. Yuan, and M. Tomizuka. Model-free deep reinforcement learning for urban autonomous driving. In 2019 IEEE intelligent transportation systems conference (ITSC), pages 2765-2771. IEEE, 2019.
[4] A. Dosovitskiy, G. Ros, F. Codevilla, A. Lopez, and V. Koltun. Carla: An open urban driving simulator. In Conference on robot learning, pages 1-16. PMLR, 2017.
[5] H. Caesar, J. Kabzan, K. S. Tan, W. K. Fong, E. Wolff, A. Lang, L. Fletcher, O. Beijbom, and Omari. nuplan: closed-loop ml-based planning benchmark autonomous vehicles. arXiv preprint arXiv:2 1810, 2021.
[6] C. Diehl, T. Sievernich, M. Krüger, F. Hoffmann, and T. Bertran. Umbrella: Uncertainty- aware model-based offline reinforcement learning leveraging planning. arXiv preprint arXiv:2111.11097 2021.
[7] Li, C. Tang, M. Tomizuka, and W. Zhan. Hierarchical planning through goal-conditioned offline reinforcement learning. arXiv preprint arXiv:2205. 2022.
[8] S. Fujimoto, D. Meger, and D. Precup. Off-police deep reinforcement learning without explo- ration. In International conference machine learning, pages 2052-2062. PMLR, 2019.
[9] A. Kumar, A. Zhou, G. Tucker, and S. Levine. Conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems, 33:1 179-1191, 2020.
[10] C. Zhang, R. Guo, W. Zeng, Y. Xiong, B. Dai, Hu, M. Ren, and Urtasun. Rethinking closed-loop training for autonomous driving. European Conference on Computer Vision, Springer, 2022.
pages
[11] A. R. Villaflor, Z. Huang, S. Pande, J. M. Dolan, and J. Schneider. Addressing optimism bias in sequence modeling for reinforcement learning. In International Conference on Machine Learning, pages 22270-22283, PMLR, 2022.
[12] A. Vaswani, N. Shazeer, N. Parmar, Uszkoreit, L. Jones, A. N. Gomez, Kaiser, and I. Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[13] T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
[14] L. Chen, Lu, A. Rajeswaran, Lee, Grover, M. Laskin, Abbeel, A. Srinivas, and I. Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34:15084-1509, 2021.

---

[15] M. Janner, Q. Li, and S. Levine. Offline reinforcement learning as one big sequence modeling problem. Advances neural information processing systems, 34: 1243-122888 2021.
[16] D. Brandfoobrenerr, A. Bietti, J. Buckman, R. Laroche, and J. Bruna. When does return- conditioned supervised learning work for offline reinforcement learning? In Advances in Neural Information Processing Systems, 2022.
[17] K. Paster, S. Mcllraith, and J. Ba. You can't 't count on luck: Why decision transformers fail in stochastic environments. Advances in neural information processing systems, 2022.
[18] S. Yang, D. Schuurmans, P. Abbeel, and O. Nachum. Dichotomy of control: Separating what you can control from what you cannot. In The Eleventh International Conference on Learning Reprexentations, 2022.
[19] C. Zhang, J. Bütepage, H. Kjellström, and S. Mandt. Advances in variational inference. IEEE !ransactions on paltern analysis and machine intelligence, 41(8):2008-200 2026, 2018.
[20] A. Shafahi, M. Najibi, M. A. Ghiasi, Z. Xu, J. Dickerson, C. Studer, L. S. Davis, G. Taylor, and T. Goldstein. Adversarial training for free! Advances in Neural Information Processing Systems, 32, 2019.
[21] Q. Sun, X. Huang, J. Gu, B. C. Williams, and H. Zhao. M2i: From factored marginal trajectory prediction to interactive prediction. In Proceedings of the IEEE/CVV Conference on Computer Vision and Pattern Recognition, pages 6552, 2022.
[22] K. P. Murphy. A survey of pomdp solution techniques. environment, 2(10), 2000.
[23] Y. Wu, Y. Burda, R. Salakhutdinov. and R. Grosse. On the quantitative analysis of decoder-based generative models. In International Conference on Learning Representations, 2017.
[24] H. Furuta, Y. Matsuo, and S. S. Gu. Generalized decision transformer for offline hindsight information matching. In International Conference Learning Representations, 2022.
[25] B. Eysenbach, S. Udatha, R. Salakhutdinov. and S. Levine. Imitating past successes can be very suboptimal. Advances in Neural Information Processing Systems, 35:6047-6059, 2022.
[26] M. Strupl, F. Faccio, D. R. Ashley, J. Schmidhuber, and R. K. Srivastava. Upside-down reinforcement learning can diverge in stochastic environments with episodic resets. arXiv preprint arXiv:2205.06595, 2022.
[27] C.-X. Gao, C.
Wu, M. Cao, R. Kong, Z. Zhang, and Y. Yu. Act: Empowering decision transformer with dynamic programming via advantage conditioning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pages 12127-12135. 2024.
(28) Y. Wang, C. Yang, Y. Wen, Y. Liu, and Y. Qiao. Critic-guided decision transformer for offline reinforcement learning. In Proceedings of the AAAI Conference Artificial Intelligence. volume 38, pages 15706-15714, 2024.
[29] Y.H.. WU, X. Wang, and M. Hamaya. Elastic decision transformer. arXiv preprint 2023.
[30] Y. Gal and Z. Ghahramani. Dropout bayesian approximation: Representing model uncer- tainty in deep learning. In international conference machine learning, pages 1050-1059. PMLR, 2016.
[31] C. Blundell, J. Cornebise, K. Kavukcuoglu, and D. Wierstra. Weight uncertainty in neural network. In International conference machine learning, pages 1613-1622. PMLR, 2015.
[32] G. Franchi, X. Yu, A. Bursuc, E. Aldea, S. Dubuisson, and D. Filliat. Latent discriminant deter- ministic uncertainty. In European Conference Computer Vision, pages 243-260.. Springer, 2022.
10

---

[33] B. Lakshminarayanan, A. Pritzel, and C. Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. Advances in neural information processing systems, 30, 2017.
[34] Y. Wu, S. Zhai, N. Srivastava, J. M. Susskind, J. Zhang, R. Salakhutdinov, and H. Goh. Uncertainty weighted actor-critic for offline reinforcement learning. In International Conference on Machine Learning, pages 11319-11328, PMLR, 2021.
[35] A. Kendall and Y. Gal. What uncertainties do we need in bayesian deep learning for computer vision? Advances in neural information processing systems, 30, 2017.
[36] M. L.. Puterman. Markov decision processes: discrete stochastic dynamic programming. John Wiley & Sons, 2014.
[37] M. Seitzer, B. Schölkopf. and G. Martius. Causal influence detection for improving efficiency in reinforcement learning. Advances in Neural Information Processing Systems, 34:22905-222 229188, 2021.
[38] V. Mai, K. Mani, and L. Paull. Sample efficient deep reinforcement learning via uncertainty estimation. In International Conference on Learning Representations, 2022.
[39] K. Hu, R. C. Zheng, Y. Gao, and H. Xu. Decision transformer under random frame dropping. In The Eleventh International Conference Learning Representations, 2023.
[40] S. J. Redmond and C. Heneghan. A method for initialising the k-means clustering algorithm using kd-trees. Pattern recognition letters, 28(8):965-997))
2007. 2007.
[41] C. team. Carla autonomous driving leaderboard. 2020. URL https: carla. org/.
[42] Q. Wang, Xiong, L. Han, H. Liu, Zhang, et al. Exponentially weighted imitation learning for batched historical data. Advances in Neural Information Processing Systems, 31, 2018.
[43] Kostrikov, A. Nair, and S. Levine. Offline reinforcement learning with implicit q-learning. arXiv preprint arXiv:2110.06166 2021.
[44] Ziegler, P. Bender, M. Schreiber, H. Lategahn, T. Strauss, C. Stiller, Dang, U. Franke, N. Appenrodt, C. G. Keller, al. Making bertha drive-an an autonomous journey on a historic route. IEEE Intelligent transportation systems magazine, 6(2):8 20, 2014.
[45] W. Zeng, W. Luo, S. Suo, A. Sadat, B. Yang. Casas, and Urtasun. End-to-end interpretable neural motion planner. In Proceedings of the IEEE/CVF Conference on Computer Vision and Patter Recognition, pages
2019.
[46] Ross, G. Gordon, and Bagnell. reduction of imitation learning and structured prediction no-regret online learning. In Proceedings of the fourteenth international conference on artifi- cial intelligence and statistics, pages 627-635. JMLR Workshop and Conference Proceedings, 2011.
[47] V. Kuleshov, N. Fenner, and S. Ermon. Accurate uncertainties for deep learning using calibrated regression. In International conference machine learning, pages 2796-2804. PMLR, 2018.
[48] Z. Zhang, A. Liniger, D. Dai, F. Yu, and L. Van Gool. End-to-end urban driving by imitating reinforcement learning coach. In Proceedings the IEEE/CVF International Conference Computer Vision, pages 15222-15232, 2021.
[49] H. Xu, X. Zhan, and X. Zhu. Constraints penalized q-learning for safe offline reinforcement learning. In Proceedings of the AAAI Conference Artificial Intelligence, volume 36, pages 8753-8760, 2022.
[50] Q. Li, Z. Peng, and B. Zhou. Efficient learning of safe driving policy via human- copilot optimization. arXiv preprint arXhv:2202.103411 2022.

---

Supplementary Materials
A Additional Related Works
A.1 Planning for Self-Driving
Most autonomous driving systems [44] adopt modular pipeline break down the massive driving task into set of submodules, with planning being of the most fundamental components. The objective of motion planning is to efficiently drive the ego-vehicle to the destination while conforming to safety and comfort constraints, which essentially a decision problem. Often engineers manually design rules [44] for specific driving scenarios, which struggles to scale for complex tasks. In contrast, IL the earliest and most widely [2, 45] used learning- planning algorithm, which learns from offline data for either cost function [45] or an executable policy [2]. However, they are strictly limited by expert quality [46] and will perform poorly when encountering out-of-distribtiion (OOD) states [46]. RL algorithms [3] own better performance and generalizability with the downside of trial-and-error learning. To address these shortcomings, we employ offline RL as an alternative to train our planner.
A.2 Uncertainty Estimation
Variance networks estimate uncertainty by loss attenuation [35]. Treating the output as a Gaussian distribution N()., given input and its label y, the network outputs mean μ(x) and variance o² (x) its final layer. The network parameter subsequently optimized by maximizing sample log-
likelihood:
Lvamer (Φ) B(x,)) - log No(x|x)]
((ua(x) y22)
+n)
(10)
= (x.y))-
(x
As the variance predicted by neural networks may be not well-calibrated overestimated [47], in this paper we practically train ensemble of K variance networks with data drawn from different subsets of the dataset (implemented by masking) and estimate the variance according to [38]:
²= =K +M² n², where μ K Li= K.
(11)
which has been proven to have strong estimation capability [38].
B Proof for Proposition
We next formally state our theorem, followed by detailed proofs of the proposition.
B.1 Theorem Statement
Before stepping into the main theorem, first introduce the problem setting that undertake for the sake of concise proof. First, we assume the reward is within [0, 1] at each timestep for conveniently bounding the error. Besides, we use g(5), a1:1) to denote the cumulative rewards by rolling out the open loop sequence of actions a1:h under the deterministic dynamics P S x A - S and reward function r x A - [0. 1] (so that the maximum return difference over future timesteps h). Moreover, Jn (æt) represents the rollout rewards of policy π over future h timesteps. Next, we formally introduce the assumption, theorem and its corresponding proof, which is inspired by the reference [16]:
Theorem (UNREST Alignment Bound) Consider MDP (S.A, expert behavior B S x म + [0, 1] and conditioning function SxN+ R. Assume the following:
S-1

---

1. Return Coverage: Ps(g h)(ss) 2 a for the initial state st and return-span h.
2. Near Determinism: p(r # r(s, a) s' + a(s.a))s,a) s å at all state-action pairs (s, a) for dynamics P and reward function
3. Consistency of f: flss,h)=fsss'ee hh-) for) all states
Then (URT shorted for UNREST):
1 + +2)h² Ba [f(1, hi] ( s d( af
(12)
As shown by the theorem, the error between the specified target and UNREST's rollout return is bounded by the factor of environment determinism ö, data coverage a, and the horizon h.
Based on the theorem, we aim to demonstrate the generalization ability of UNREST at the identified 'certain states' by our uncertainty estimation stra. Specifically, we have the following lemma:
Lemma (Determinism Equality) Assuming that the rewards obtained uniquely determined by transitions (s, - s') at each timestep, then there exists e > 0 and 8 0, such that:
(Rell()) II p( [s,, Te)) e p(r r(s, a) or s' # P(s, a) 5. (13)
=
While the reverse direction is easy to see (the state transition probability is nearly deterministic so the transitions cannot provide additional information for return prediction), here we focus on proving the forward direction. Let us demonstrate contrapositive proposition:
If p(r # as, or s' # P(s, a) 6, then 3e, s.t. DKL [KL[p((R (R,RR p(, T<r)] (14)
First, since is uniquely (different transitions lead to distinct rewards) determined by (s, - s'), we can induce that r r(s, a) is equivalent with s' * P(s, a) the only difference can occur in s'. Therefore, need prove:
If (r # r(s, a)s, a) > 6, then 3e, s.t. ((R(/(()) || p(R Ter) >.
(15)
Considering the extreme case, set y - O, thus we have R, 7₁. From the perspective of proof by contradiction, the Kl.-divergence cannot be smaller than any must differ for that the probability # r(s;-1, a,a,-1) is larger than 0 and can be determined by s, To this end, we have proved the contrapositive proposition of the forward direction, which is equivalent to the original proposition.
since p(r, and p(rils,, T>r)
Therefore, we can summarize from the above theorem and lemma that UNREST can generalize to achieve high returns at 'certain states' long the corresponding actions are covered by the expert.
B.2 Proof for the Theorem
The next theorem proof is largely built on [16]. Note that we omit the superscript 'URT' to simplify the equations. First, we expand the left term in Eq. 12:
B,,h (s()), ) E,
Bxyls [[(1,, h) -i1]
E, Ban [[(), h) g(s,, an:)
+ Eyi Bana- [g(ss an.) 8i]
$ Ey Ba Arlss (s) h) 8()) + o1?.
(16)
The last step follows by bounding the difference between 81 and g()), by the maximum return difference h and union of probability bound over timesteps based on the near determinism assumption:
h susup U Jar-zyls (re ar(s. or Se4 # P(Se, a:)) s 812
(17)
S-2

---

For the first term in Eq. 16, it can also be expressed and bounded by the maximum return difference:
S Janh Pay Pny(aanlss) an)) # f(s1, hi]
By.hh Bana ayls [f(s,, s(s1, s Baihh
(18)
To bound this term, we need further expand the distribution Pxy. First, with the assumption that UNREST is perfectly fitted the expert dataset,
deduce the following by the bayesian law:
2 <a(alls))= B(ails)) (5(81 = f(11, ())s).)))
(19)
P5 = f(s), h)(ss)
For simplification, we use 3, = P(s1. (1:--1) to denote the state reached by following the deterministic dynamics defined by till timestep r. Next, based on the near determinism, consistency and coverage assumption, we can expand Pxy to get:
Pxy (atnh s1) x (a1 $1) [ p (s2 s1,a1) Pay (a2sh $2)
52
sxy (al
s1) Pxy (a2:k $1, 52) +
Pp (81 = f ($1, h) $1. a1)
=B s1)
Pp (g) = f (5), h) s1) Pxy (a2:h 51, 32) 6
6+ PB (81 - -r ) = (51, h) - $1, (1)
$ s1)
Pxy (azmh (1. +
PB (g1 f (si, s1)
8+ PB (g2 f (52, h - 1) 52)
=B(ai s1)
Paf (a2th 51) +
PB (81 = f (s1, h) ss
Pp (g2 f (5g, h
52)
+11
$ s1)
Pny (a2a I $1, 52) +
(20)
PB (g1 f (si. h) 51)
aa
Similarly, we can further expand Pay (#2:h 81, 52) using the same rule:
Pay (a2n
s1, = x x 52) p(s3 I 52, (2. (asch s1. 83)
s <ay(a2 (
= | 52) Pp(g2 f(52, 1) a2)
Pay (a3:h $1,32,33) + 6
(21)
pg = f(32, h 1)(32)
Substituting this back to Eq. 20, we have:
Pay (a1.h 1 s1)
PB (82 (32, h- 1) 52)
Ps (82 (52, h 1) 32, 42)
<a ( 52)
Pp (g2 = f (32, - 1) (22) Pay (a3:h 52,33,
PB (g1 = (s1, s1)
- -+1
+ 26
of
recursively expand Pzy
PB (gh = (5,, 1) 3k, an)
1
<
(a, 3,)
+ hō
+
Pp (g) f (s1,h) s1)
a
1=1
[ (5i, a1.)) = (51, hō PB (g) f (51, h) si)
[ B(a 5x) [=1
af +1
(22)
The last step is deduced by the trajectory determinism and the consistency of the conditioning function f. Multiplying this back to Eq. 18 and noticing that the two indicator functions can never both be 1, we can yield that:
Bi.h Bain-Ryin [f(s), h) - g(s), and)] < 128(-+1 1). af
(23)
This can finally yield the bound in Eq. 12 by adding back to Eq. 16.
S-3

---

Algorithm 2: UNREST training procedure
Input: Offlin dataset 2 = (T;) batch size B.
1 # Stage I: Train return transformers.
for each iteration do
Sample batch 8 = [T from D, where v = ((s,,,,))
Feed sampled sequences into Transformer and get embeddings x,,, Ku, by Eq. 3;
Feed Transformer outputs into variance networks predict return
Ryp (Re)Ter), Ry, Se));
Update network parameters based on the learning objectives in Eq. 4:
# Stage II: Segment sequence estimated uncertainty.
8 Compute uncertainty of each timestep M, + Dxi. (Ry, (T)), RR, (RelT<s s,)); seg
Segment sequence w.r.t. uncertainties 2 c [T; J |==
discussed in Sec. 4.3;
10 # Stage III: Train decision model.
for each iteration do Sample batch B (T, )
) from Do0e where ,/ [(h,, Rr, R/. St. az) )=) Feed sampled sequences into Transformer and get embeddings Sph, Rx, Exy Krt by Eq.
14
using global return then
Concatenate Transformer outputs with global return x, [Xs, f (R,)];
15
Feed & into variance network to predict action distribution Te (arlT 2eg H1, R R st); Update network parameters based on learning objective in Eq.9;
17
Output: Trained return transformers Rya (.), Rys (+) and policy xo
C UNREST Training Procedure
To get more clear understanding of our training pipeline, we present the training procedure of UNREST in Alg. 2.
D Dataset Information
Our dataset collected from the CARLA simulator, using its built-in Autopilot, which is rule-based motion planner. Specifically, CARLA [4] is open-sourced simulator designed for autonomous driving research. It provides researchers with realistic environment that faithfully simulates traffic dynamics, weather conditions, and high-fideliyy sensor data. The pre-built scenarios in CARLA offer a diverse set of driving scenes, while its high level of customizability and flexibility make it the preferred simulation platform for a majority of autonomous driving researchers.
We collect training data frequency of 10Hz, accumulating 30 hours of driving data from four distinct training towns (Town01, Town03, Town04, Town06) under four different weather conditions (ClearNoon, WetNoon, HardRainNoon, ClearSunset). At each timestep, we store the data tuple (St, (s,)),, whose compositions will be introduced in detail in the next paragraphs.
The symbol E R denotes the reward returned by the environment at the timestep r. Typically, our reward design revised from Roach [48], encompassing factors such as speed, safety, and comfort.
Reward:
= Espeed + possition
+ Frotation + Faction
Terminal
(24)
Among them, 'sspeed We-Weeee
represents the reward of approaching the target speed; "posiiion -0.5Ap represents the reward obtained from driving in the correct position. where Ap the lateral distance between the ego vehicle and the centerline of the target route; Frotation AA, represents the reward obtained from driving in the correct orientation, where is the angle difference between the ego vehicle and the centerline of the target route; Faction incurs penalty of -0.1 when the steering in current timestep differs from the previous step by more than 0.01, promoting smooth and comfortable
S-4

---

driving; Fierminal is 10 when reaching the destination, -10 when encountering a collision or violation of traffic rules, to encourage driving safely.
Safety constraints: In addition to the reward setting, the Constrained Penalized Q-learning (CPQ) [49] algorithm also needs constraint definition for training. Specifically. we adopt the same setting in [50], where each time of collision traffic rule violation incurs penalty of 1. 1.0. The constraint limit for CPQ is 1.0 (i.e. ensures no safety issue at each ride).
In order to avoid memory overflow in the sequence model, both the action space and state space have been specially designed. Specifically, we adopt a similar setting to SPLT [11].
Action: The action e R consists of two values, representing the target angle and target velocity that will be fed into two separate PID controllers.
State:
The observed state E R includes the following components: (i) ego vehicle velocity (ii) relative distance and velocity to the leading vehicle, which will be set to the default maximum value no leading vehicle is perceived in the horizon (iii) relative distance and velocity to the pedestrian ahead, which will be default maximum values if pedestrian is perceived in the horizon (iv) relative velocities and distances vehicles in four cardinal directions (front, rear, left, right), also set to default maximum values if no vehicles are perceived in the horizon (v) distance the traffic light ahead, set to default maximum value if no red light ahead (vi) distance to the stop sign ahead, set to default maximum value if no stop sign ahead (vii) relative angle difference w.r.t. the centerline of the target route (viii) distance to the centerline of the target route (ix) relative positions of the next 10 target waypoints.
As for stochasticity, introduce stochastic speed ranges in [20, 40]m/s and vehicle-specific random sizes of visibility areas in [10, m. This results in aggressive vehicles with higher speeds and smaller visibility ranges, as well as cautious vehicles with lower speeds and larger visibility ranges. These characteristics cannot be inferred from the observable state of the ego vehicle, leading to greater uncertainty during its interactions with the environment. Additionally, we also add Gaussian noise the control signals of the ego vehicle to incorporate stochasticity into its behavior. These special designs, integrating with the complex scenarios and interactions from our collected large dataset, ensure that our training data contains sufficient stochasticity.
E Implementation Details
In this section, we introduce the implementation details, including our training process, evaluation process, and hyperparameters. Specifically, UNREST and all the baselines are implemented with Python 3.7 and PyTorch 1.13.1. Besides, training processes are run on a NVIDIA A10 while inference conducted on one NVIDIA 3090 GPU for fair comparison.
E.1 Training Details
For all training processes, We employ AdamW optimizer with a learning rate of 10-4, and consistent batch size of 256. For the training of sequence models, history length of 10 is carefully selected to be sampled each time, which corresponds to a real-world duration of second. To determine number epochs for training, dynamic value is assigned based on
Training Process:
dataset's size, defined by the following formula: Repochs int len(dataset)
le6
nref). where frrr represents
reference epoch number that can be tuned.
Details for segmentation: As discussed in Sec. 4.3, we train two return transformers for sequence segmentation. practical implementation, to accurately quantify environmental uncertainty, we employ the ensemble method to train each return transformer. Using Eq. 11, we recalculate the variance and expectation of the returns. Then the sequences segmented into certain and uncertain parts w.r.t. estimated uncertainties according to the principles we introduce in Sec, 4.3.
S-5

---

In our experiments, we discover that 'uncertain states' tend to occur consecutively. To ensure meaningful segmentation, we enforce a minimum 'uncertain part' length of c = 20 frames and select the uncertainty threshold e = 3 according to Fig. 5(a). Under these two hyperparameter settings, we find the segmented sequences often correspond to the completion of significant driving task.
Details for return-span embedding: We introduce the return-span embedding f/(() to enhance the interpretability of truncated returns as conditions. Specifically, we explore two variants to incorporate this embedding. The first variant is to add the retum-span embedding to the return embedding, stated in the main text. The second variant considers treating f/ (h,) as distinet token input to the sequence model. However, we observe that the latter approach results in increased memory consumption without yielding improvements in performance (Tab. 8 and Tab. 9).
Details for global return embedding: The approach for obtaining the global return embedding differs from other embeddings due to its high uncertainty. Typically, instead of utilizing parameterized networks, we opt for a coarse uniform discretization approach, resulting in a corresponding one-hot vector (specifically, a 50-dimensional vector in our implementation). This strategy enables us to provide global guidance while mitigating the influence of its uncertainty on the training process.
Details for decision models:
the practical implementation, following DT [14], we directly predict the value of the action instead of its distribution. Specifically, we only retain the mean μ and assume unit variance. Consequently, the loss in Eq. 9 can be reformulated as the mean square loss, resembling the learning objective in DT.
E.2 Evaluation Details
Main setting: To comprehensively evaluate our models' planning capability, we select two distinct driving scenarios from CARLA: the training town under training weather conditions and the new town under new weather conditions. Typically, the training town is chosen as Town03, the most complex town among the training dataset, and the training weathers are 'WetNoon' and 'ClearSunset', The new town is Town05, the most complex town in the rest of the training set, and the new weathers are "SoftRainSunset" and 'WetSunset'. For each tested scenario, we carefully set up the vehicles and driving routes w.r.t. the specifications of CARLA Leaderboard [41]. Moreover, we conduct experiments with three different seeds (2022, 2023, 2024) to facilitate the calculation of mean values and variances.
Details for dynamic uncertainty estimation: As we have stated in the main text, we employ dynamic uncertainty estimation inference time to enable cautious planning. Typically, we imple- ment more variants (results are shown in Tab. 8 and Tab. 9) other than KD-Tree [40] for test-time uncertainty estimation:
The first approach entails training a neural network to forecast environmental uncertainty. Leveraging the uncertainty values computed by the aforementioned return transformers and employing the provided threshold, we assign binary class labels (certain and uncertain) each state in the training dataset. Subsequently, these labels are utilized to train an uncertainty network. Heuristic: The heuristic method turns out to be straightforward and intuitive alternative. Specif- ically, we observe that specific dimensions within received states have a strong correlation with uncertainties (e.g. distance to traffic light ahead). Thus, we directly compare the values of these particular dimensions to the preset uncertainty threshold to determine whether the state uncertain. KD-tree: Our chosen methodology employs a KD-tree for uncertainty estimation. Specifically, we use all state vectors to initialize the KD-tree. Subsequently, the tree will select a dimension each level and perform a binary split, ultimately placing all inputs in the leaf nodes (implemented by calling the Skleam package). At test time, query the nearest neighbors of the current state from the tree and compute the average uncertainty of its nearby 5 states as the estimated uncertainty.
Network prediction:
S-6

---

Table 4: Hyperparameters for training return transformers for Sequence Segmentation.
Hyperparameters
Value
Hyperparameters
Value
# Transformer layers
Transformer heads
Embedding dimension
128
Batch size B
256
Sampled Sequence length
10
Discount y
0.95
Learning rate
lc-4
Dropout
0.1
Optimizer Ensemble size # Reference epoch Reef
AdamW
Weight decay
False
Data mask probability
0.6
50
Transformer Activation
GELU
Table 5: Hyperparameters for training UNREST decision models.
Hyperparameters
Value
Hyperparameters
Value
Uncertainty threshold
3.0
Min 'uncertain part' length c
20
Transformer layers
Transformer heads
8
Embedding dimension
128
Batch size
256
Sampled Sequence length
10
Discount y
1.0
Optimizer
Adam
Weight decay
Fulse
Learning
lc-4
Dropout
0.1
Reference epoch Aiet
200
Transformer Activation
GELU
Global return dimension
50
Action Activation
Tanh
Table 6: Hyperparameters for UNREST's inference process.
Hyperparameters
Value
Hyperparameters
Value
Max history length
5
Uncertainty threshold
3.0
Return Horizon
100
Upper Percentile
0.7
Deterministic sample
True
# KD-Trree neighbor
5
E.3 Hyperparameters
We select several well-established approaches to serve as our comparative baselines. Specifically, we include Behavior Cloning (BC), Implicit Q-Learning (IQL) [43], as well as prominent sequence models such Decision Transformer (DT) [14]. Trajectory Transformer
(TT) [15], SPLT [11], and ESPER [17]. To ensure fair comparison, for all these baselines, adopt the default hyper- parameter settings from the repository https: wtwggtttubb..o wwt.l/co//vviiiiieee aa/
For Monotonic Advantage Re- re-Weighted Imitation Learning (MARWIL) [42] and Conservative Q-Leamnng (CQL) [9], we adopt the default hyperparameter from the Ray RLlib https:/ //cccs///ccc// /ttii///ddeiexeee...// htt///ndxx...e html. Finally for the safe offline RL baseline Constraints- Penalized Q-learning (CPQ) [49], we employ the default hyperparameters from the repository https://githut: cwm com/liuzuxin/OSn
The full list of UNREST's hyperparameters can be found in Tab. 4, Tab. 5, and Tab. 6.
F More Experimental Results
In this section, we supplement more experimental results, including the visualizations, complexity comparisons, sensitivity analysis of hyperparameters, the results for more UNREST variants, and results on environments other than driving (D4RL).
F1 Visualizations
Uncertainty visualizations: We present more visualizations UNREST's uncertainty estimation in Fig. to validate its interpretability, We first visualize the uncertainty distribution in Fig. 5(a). As see, the majority of uncertainty values associated with state transitions fall within the range of 0 to 10-4, aligning with intuitive expectations. In autonomous driving environments with sparse
S-7

---

Uncertainty Threshold
Uncerlaiity Threshald
Uncertaiity Threshold
Uncertainty Value
Uncertainty Value
Uncertainty
Timestep
Timestep
(a)
Uncertainty distribution
(b) Right turning: Uncertainty curve (c) Light crossing: Uncertainty curve
Frame
Frame
Freme 10
Frame 15
Frame
(d) Right turning: State changing process
Frame
10
Frame
Frame
(c) Light crossing: State changing process
Figure 5: More visualizations of UNREST's uncertainty estimation results. The black background indicates non-drivable areas, the dark gray areas represent drivable regions, the white rectangle denotes the ego vehicle, dark blue rectangles signify surrounding vehicles, the sky blue rectangles indicate pedestrians, the purple lines represent lane boundaries, and the red, green, and yellow markers indicate traffic lights and stop signs. Finally, the light gray path extending from the ego vehicle represents its intended route.
vehicle presence and limited number of traffic signals, the stochasticity of environment transitions tends to be relatively low. The uncertainty threshold is accordingly chosen as 3.0 to cover comer cases with large uncertainties.
Then, we seek uncertain transition states within the training dataset and visualize two typical scenarios. Fig. 5(b) and Fig. 5(d) correspond turning scenario, while Fig. 5(c) and Fig. 5(e) depict scenario of passing through red light. By associating the temporal steps with the uncertainty curves, it can be observed that the ego-vehicle approaches the turning point, the uncertainty gradually increases since the ego-vehicle cannot forecast the traffic conditions after turning. After the turning behavior is almost completed, the uncertainty decreases below the threshold again. In the scenario of waiting at red light, the uncertainty rapidly increases as the ego-vehicle approaches the red light due to the uncertain color changes of the traffic light. However, after the ego-vehicle decides to stop in front of the red light approximately the 15th timestep, the uncertainty decreases sharply below the threshold. These visualization results effectively demonstrate the interpretability of our uncertainty measurement approach.
Uncertainty Calibration: To provide more intuitive impression of UNREST's effectiveness, present the return distribution calibration results in Fig. 6. Notably, the uncertainty in the figure
S-8

---

20
Perfectly calibrated
20
Perfectly callbrated
Uncertainty
Uncertainty
Predicted Return
15
Prediction Mean Return
Predicted Return
Prediction Mean Return
10
10
RMSE: 0.851
RMSE: 1.851
-5
1.281
2.633
10 Groud Truth Return
15
20
10 Groud Truth Return
15
(a) Ensemble p ( Calibration
(b) Single p(R (())) Calibration
20
Perfectly calibrated
20
Perfectly calibrated
Uncertainty
Uncertainty
Predicted Return
15
Prediction Mean Return
Predicted Retur
15
Prediction Mean Return
10
RMSE: 1.170
RMSE: 2.794
NLL: 1.603
5
NLL: 4.8:2
10
15
20
10
15
20
Groud Truth Return
Groud Truth Return
(c) Ensemble p( st) Calibration
(d) Single ((R, (RR sy) Calibration
Figure 6: Calibration results return distribution using ensemble models are obviously better than that using single model. We the standard variance of the networks' predictions approximate indicator of uncertainty. The blue line signifies the ground truth, while the red dots denote the predicted mean returns. The areas shaded in orange represent the predicted mean coupled with their respective standard deviations.
denotes the variance of the distribution, different from the so-called environmental uncertainty we used in the main text that reflects the impact of environment transitions (through conditional mutual information). The impact of the environment has no ground truth value and can only be interpreted through visualizations like Fig. 5, while the predicted return distribution has corresponding ground truth and thus can be directly calibrated like what we do in Fig. 6. Typically, the figure illustrates that using an ensemble of return transformers can significantly better predict the return distribution than using a single model (closer to the ground truth, with tight uncertainty bands), where it achieves smaller (better) results on two widely adopted calibration metrics: RMSE and NLL than the single return transformer variant.
Sequence length distributions: We present different segmented sequence length distributions in Fig. 7. Analyzing the figures, it is evident that most uncertain sequences tend to possess lengths close to the state transition timestep (translate from uncertain certain states if the last c I steps all identified as certain), while certain sequences are almost uniformly distributed w.r.t. their lengths.
F.2 Complexity Comparisons
We provide empirical results about space/time usage as shown Tab. 7. Regarding GPU usage, UNREST demonstrates comparable levels to DT, while significantly less utilization compared to SPLT and TT. Regarding training and inference time, UNREST consumes slightly more time than DT, yet notably faster than TT and SPLT. These results signify that UNREST achieves superior performance while requiring relatively modest computational resources.
S-9

---

5,
Uncertaih Certain
Uncertain Centain
Uhcertai
Frequency
Frequency
2000
Frequency
Certain
5000
2000
6600
2000
1.00
10000
2000
12
10
50 60
59
10
6
19
<9
69 18
Truncated Sequence Lengtin
Truncated Sequence Length
Truncated Sequence
Figure 7: Segmented sequence length distributions with different threshold and state transition timestep c.
Table 7: Comparison results of training/inference time (10-2s) and training GPU memory (GB) between different sequence models. The training time calculated w.r.t. multiple iterations, while the inference time calculated w.r.t. multiple rollout steps. The GPU usage is fixed value for one training model.
Metric
BC
DT
TT
SPLT
UNREST
GPU usage
1.00
1.46
23.3
1.96
1.48
Training time
0.83 0.13
0.93 0.13
32.7 = 0.2
4.20 0.1
1.47 : 0.12
Inference time
0.15 = 0.01
0.16 0.02
25.6 + 0.7
0.89 = 0.02
0.24 0.03
Train condition
New conditoin
70
70
Driving Score
Driving Score
60
60
50
50
40
50
100
200
300
500
1000
0.5
1.0
2.0
3.0
5.0
10.0
Return Horizon H
Uncertainty Threshold E
Figure 8: The sensitivity analysis of hyper-parameters H/ and
F.3 Sensitivity Analysis
Impact of return horizon H: We investigate the influence of the return horizon length, denoted as H, on the performance of UNREST. The experimental results are depicted in Fig. 8, where the return horizon varies from to 1000. is evident that relatively modest value H (e.g., //=100) yields the best performance, while minimal and large values of exhibit noticeable performance deterioration. These findings align with our underlying assumption that shorter sequences can alleviate uncertainty, while excessively short sequences may give rise shortsightedness, thereby compromising performance.
Impact of uncertainty threshold e: We next investigate the impact of the uncertainty threshold and present the results in Fig. 8 where we increase e from 0.5 to 10. As shown in the figure, a moderate value e, namely e = gives the best performance, while smaller and larger values of exhibit obvious performance degradation. This behavior can be attributed the fact that the majority of uncertainty values fall within the range of 0 10-4. Consequently, excessively small values may incorrectly identify 'certain states' uncertain, while excessively large values may overlook important 'uncertain states'.
S-10

---

Table Driving performance of more UNREST variants on train town and train weather condition.
Planner
Driving Score
Success Rate?
Route Comp.
Infrac. Scoret
Norm. Rewards
Tokened ret-span emb.
62.8 3.6
56.2 4.8
81.1 4.7
66.5 2.5
0.62 0.03
Fixed horizon seg-
59.2 2.6
44.7 3.2
7 4.8
64.5 3.0
0.66 0.02
Separate BC model
61.0
5.3
82.8
64.6 2.4
0.68 0.04
Heuristic uncertainty
62. # 2.4
7.4
43.8
4.3
0.47 = 0.03
Predicted uncertainty
601.7 + 2.7
50.0
63.6
65.8 3.4
0.61 0,02
Reweighted BC
64.2 2.4
= 3.3
85.5 2.7
71.4 2.5
0.68 0.03
Original model
+ 3.2
54.5
83.8 3.1
70.2 2.8
0.04
Table Driving performance of more UNREST variants on new town and new weather conditions.
Planner
Driving Scoret
Success Rate?
Route Comp.1
Infrac. Scoret
Norm. Rewards
Tokened ret-span emb.
62.0 + 4.9
56.8 3.0
92.0
60.3 3.3
0.64 0.04
Fixed horizon seg.
58.5 3.6
46.7
80..3 2.8
64.7 2.8
0.64 0.02
Separate BC model
59.3 3.4
63.8
90.0
59.3
0.68 0.05
Heuristic uncertainty
2.6
4.3
6.2
7.0
0.52 0.02
Predicted uncertainty
5.4
5.5
8.70
5.1
0.64 0.04
Reweighted BC
63.8 2.3
4.7
9.0
65.1
0.67 0.02
Original model
62.9 +
$7.5
90.0 6.0
62.9 + 3.8
0.65 + 0.03
Table 10: Comparison to SOTA baselines on the standard D4RL Mujoco locomotion-v2 domain. For BC, MBOP CQL, DT, TT, IQL, and SPLT we use the results reported from the SPLT paper [11]. For ESPER, we report the results from their paper [17]. We report the mean and std for our method over seeds with 10 trajectories for each seed. 'Stochastic' denotes customized stochastic environments according to [18]. The symbol* is used to indicate the omission of mcasurement results that are deemed unimportant.
Datoset
Envirvement
BC
MROP
CQL.
DT
TT
IQ1.
SPL.T
ESPER
DoC
UNREST((URR)
Med- Export
HalfCheetah Hopper
85
Cap--Epert
79.6
55.1
105.4
107.6
1100:2.7
91.5
104.4126
19.95=13.91
91.313.3
93.5-5.4
Exprt
Walker24
36.6
70.2
108.8
108.1
101 916.8
109.6
108.641.1
106.87a1.26
107 3.2.1
105.7:4.4
Medium
HallChestah
43.1
44.6
44,0
42.6
46.9
47.4
44.310.7
42.51=0.08
43.2s0..
44.5=0.9
Medium
Hopper
50.57=3.43
Medium
Walker2d
77.3
41.0
72.5
74.0
79.0&2.8
78.3
77.940.3
69.8=1.2
72.212.2
70.6=3.9
Med- Med-Replay
HallChectah
4.3
42.3
45.5
36.6
11.912.5
44.2
42,740.3
35.922.0
38.7.0.7
39,400.5
Med-Ros Roplay
Hupper
27.6
12.4
95.0
82.7
91.5
94.7
75.0+23.8
20.2:16.1
$7 $7.3
65.7=21.9
Med- Replay
Walker2d
77.2
63.7.4.7
Averuge
47.7
47.8
77.6
74.7
78.9
76.9
72.9
64.2
70.4
73.0
Med- Stochastic
HalEChestah Hogper
75.821.2
83.7+2.8
88.6422.
77.4=5.6
84.4.4.1
90.5 a 3.2
Med- Stochastic
89,433.3
93.433.6
91.812.7
92.5=6.9
95,2
97.4
Med- Mecl-Sicc
Walker2d
91.762.8
95.743.4
100.2
E4 More UNREST Variants
Tab. and Tab. 9 record the detailed performance of more UNREST variants in both training and new scenarios.
As shown in the results, introducing return-span embedding as a separate token yields little im- provement but will increase memory usage. Fixed horizon segmentation means segmenting the sequences a fixed timestep 100. Although simple, it ignores environmental uncertainties and may lead to shortsighted problems. Therefore, it's reasonable that it performs notably worse than original model. The Separate BC model means separately training an offline RI model in 'certain parts' and a BC model in 'uncertain parts'. While achieving the highest normalized rewards, there still a significant gap in its driving score compared with the original model, which shows that the joint training of two models can benefit the planning process. Utilizing a heuristic method to estimate uncertainties at inference time obtains the best infraction score, which is apparent since is designed to promote cautious behaviors every possible uncertain scenario. However, this method is overly pessimistic and overlooks certain corner cases, ultimately leading to the lowest driving score. Adopting network to predict uncertainty at inference time yields better performance than the
$-11

---

heuristic method, but overall, due to the relatively low dimensionality of the inputs, the original model performs better than it with a faster inference speed. Finally, we also study a Reweighted BC variant, which additionally fits state value functions like that in MARWIL [42] and conducts advantage reweighted behavior cloning at uncertain parts. Since the offline data collected by AutoPilot can be suboptimal, we find reweighted behavior cloning can effectively down-weight those suboptimal actions and improve UNREST's driving scores at both training and new scenarios. However, since it significantly increases model complexity, we have refrained from incorporating this technique into the final model.
F,5 Results on D4RL
For completeness, we evaluate our method on D4RL Mujoco tasks, though they are mainly deter- ministic tasks. Specifically, our UNREST generally competitive with SOTA baselines on these tasks, as shown Tab, 10. The reason for its overall underperformance compared DT lies in the unnecessarily forced sequence segmentation in such deterministic tasks, where it is more appropriate to retain global returns in order extend the horizon length considered during planning. We also test UNREST in three customized stochastic D4RL environments from [18]. Compared to complex generative training (which also harms their performance in complex environments) of DoC [18] and ESPER [17], UNREST achieves the highest rewards with lightweight models.
S-12


Reviews:
[{"review_id": "WcHqlsNIkE", "reviewer": "Reviewer_hnBQ", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": 3, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "The paper aims to address a key challenge with decision transformer (DT), which is sensitive to stochasticity in the environment leading to poor performance.  The two key observations that authors use are 1) Uncertainty accumulation and 2) Temporal locality of policies, they suggest the use of truncated returns instead of global returns to account for overly optimistic actions from DT's in stochastic environments. The truncation is done using an estimate of uncertainty. Extensive evaluation is done in CARLA simulation to validate the results.", "summary_of_recommendation": "Weak accept", "technical_quality": 3, "clarity_of_presentation": 3, "potential_impact": 2, "robotics_focus": 3, "quality_of_limitations": 3, "questions_for_rebuttal": "1) One of the return transformers - p_phi_s (T_ret, s_t) - depends on the truncated return (T_ret) and the state at timestep t (s_t), but the truncated return also has the current state s_t as an input, what is the reason for having s_t twice? \nA broader question is, why cant the return transformer just be a probabilistic Q-function (under the expert policy) and then use the variance captured by it to identify certain vs uncertain states? \n\n2) How sensitive is UNREST to novel driving scenarios? i.e., scenarios that were not present in the dataset to learn the uncertainty estimators. For example, lets say there is a novel roundabout environment that wasn't present in the dataset. Is the uncertainty estimate less reliable? In uncertain regions, UNREST 'rests' on expert actions, but what if we don't have expert actions? \n\n3)  UNREST is evaluated on a few stochastic D4RL environments to demonstrate generalizability, what are examples of temporal locality in these environments/tasks analogous to driving scenarios?", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "renrBEL3F0", "reviewer": "Reviewer_1GUa", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": 4, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "The main idea of the paper is to avoid generative modelling or adversarial training for handling uncertainty in rewards due to stochastic transitions. This is done via estimating the effect of a transition on the truncated return via conditional mutual information between transitions and returns.", "summary_of_recommendation": "It's a well written paper with very minor revisions needed for rebuttal.", "technical_quality": 5, "clarity_of_presentation": 5, "potential_impact": 3, "robotics_focus": 3, "quality_of_limitations": 3, "questions_for_rebuttal": "1. The uncertainty estimation measures the discrepancy in estimating the reward from information before the transition (absence of s_t) and information of the transition (presence of s_t), is my understanding correct? perhaps, the notation in equation 3 was a bit confusing.\n2. Please move the limitations section in the main paper instead of being in the appendix.\n3. I didn't understand how the 'lightweight uncertainty prediction model' in section 4.4 is trained, perhaps more details would be great.\n4. How is the uncertainty threshold epsilon selected? \n5. How is the hyperparameter c (for frequent switching between certain and uncertain segments) selected?\n6. Perhaps, section 5.4 Uncertainty Visualization could be expanded to show more instances of uncertainty on ego vehicle behavior.\n7. How is the hyperparameter n (upper percentile of target return for planning) selected?", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MMZVbOzuOJ", "reviewer": "Reviewer_Jtpp", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": 3, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "The paper introduces UNREST, a novel approach for offline reinforcement learning (RL) in uncertain driving scenarios. Traditional RL models struggle in stochastic environments where identical actions may lead to varied outcomes. UNREST addresses this by integrating uncertainty estimation directly into decision-making. It measures uncertainty through conditional mutual information between transitions and returns, emphasizing actual outcomes over assumed transitions. This approach mitigates 'uncertainty accumulation' and 'temporal locality' issues typical in driving environments. Notably, UNREST replaces global returns with truncated versions less sensitive to environmental variability, enhancing robustness. Experimental results across diverse driving scenarios demonstrate UNREST's superior performance and highlight the efficacy of its uncertainty-aware planning strategy.", "summary_of_recommendation": "The recommendation for weak accept is grounded in several key factors. Firstly, the paper exhibits high quality evidenced by its innovative approach integrating uncertainty estimation into decision transformers, addressing a significant challenge in RL. The clarity of presentation, while generally effective, could benefit from clearer explanations in mathematical formulations. Originality shines through in adapting transformer models for stochastic environments without additional complex models, enhancing its significance in autonomous driving. The experimental rigor and practical relevance underscore its potential impact. Addressing complexities in uncertainty estimation and broadening comparative analyses would further strengthen its contribution to the field of autonomous systems and reinforcement learning.", "technical_quality": 3, "clarity_of_presentation": 3, "potential_impact": 4, "robotics_focus": 2, "quality_of_limitations": 1, "questions_for_rebuttal": "I'm not an expert for  self-driving research. I try to give some suggestions:  \n\n1. There are many methods improving Decision Transformer, e.g., TIT [1], PDiT [2], STEER [3] and their applications for Traffic Signal Control [4,5]. I think it will be better to compare these more powerful methods rather than just comparing with DT and TT.  \n\n2. In many general RL setting, the environment is also stochastic, could the proposed UNREST be used in Traffic Signal Control or other RL settings? This could improve the potential impact on other fields.  \n\n[1] Transformer in Transformer as Backbone for Deep Reinforcement Learning. Arxiv 2022.   \n[2] PDiT: Interleaving Perception and Decision-making Transformers for Deep Rinforcement Learning. AAMAS 2024.   \n[3] Sequential Asynchronous Action Coordination in Multi-Agent Systems: A Stackelberg Decision Transformer Approach. ICML 2024.   \n[4] CoSLight: Co-optimizing Collaborator Selection and Decision-making to Enhance Traffic Signal Control. KDD 2024.   \n[5] X-Light: Cross-City Traffic Signal Control Using Transformer on Transformer as Meta Multi-Agent Reinforcement Learner. IJCAI 2024.", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper introduces uncertainty into a decision transformer, to better equip the model to handle stochastic RL environments. Reviewers are all generally positive about this work.

Strengths:
* Well written, simple but good improvement over base DT, method seems like a good idea - a nice extension

Weaknesses:
* Only tested in driving scenarios (CARLA) - additional detail needed around the test conditions here and evidence for stochasticity in environment, unclear whether this would generalise.
* Some clarity issues in figures, requirements to refer to appendix for detail, limitations not discussed in main paper.

For the rebuttal, I recommend answering the reviewers technical questions, and addressing concerns around experimental tests and potential generality.

Post rebuttal:
The authors did a good job of addressing remaining concerns through their rebuttal and answering questions around test conditions and baselines. I think this paper has merit and is a valuable contribution.
Decision: Accept