Uncertainty-Aware Decision Transformer for Stochastic Driving Environments

Venue: robot-learning_org_CoRL_2024_Conference
Published: 2024-06-05 16:20:45

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.

Main paper: 
 introduction: 
 Safe and efficient motion planning has been recognized as a crucial component and the bottleneck in autonomous driving systems (Yurtsever et al., 2020) . Nowadays, learning-based planning algorithms like imitation learning (IL) (Bansal et al., 2018; Zeng et al., 2019) and reinforcement learning (RL) (Chen et al., 2019a; 2020) have gained prominence with the advent of intelligent simulators (Dosovitskiy et al., 2017; Sun et al., 2022b) and large-scale datasets (Caesar et al., 2021) . Building on these, offline RL (Diehl et al., 2021; Li et al., 2022a) 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 (Fujimoto et al., 2019; Kumar et al., 2020) .Nevertheless, the application of offline RL approaches still faces practical challenges. Specifically:(1) The driving task requires conducting long-horizon planning to avoid shortsighted erroneous decisions (Zhang et al., 2022) ; (2) The stochasticity of environmental objects during driving also demands real-time responses to their actions (Diehl et al., 2021; Villaflor et al., 2022) .The recent success of the Transformer architecture (Vaswani et al., 2017; Brown et al., 2020; Dosovitskiy et al., 2020) has inspired researchers to reformulate offline RL as a sequence modeling problem (Chen et al., 2021) , which naturally considers outcomes of multi-step decision-making and has demonstrated efficacy in long-horizon tasks. Typically, they train models to predict an action based on the current state and an outcome such as a desired future return (i.e. reward-to-go). However, existing works (Brandfonbrener et al., 2022; Paster et al., 2022; Yang et al., 2022) have pointed out that these decision transformers (DTs) 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. 1 (a), identical turning actions may yield entirely different outcomes w.r.t. the aggressive or cautious behavior of the surrounding blue vehicle.The key to overcoming the problem is to distinguish between outcomes of decisions and environment transitions, and train models to pursue goals that are not affected by environmental stochasticity. To our best knowledge, only limited works (Villaflor et al., 2022; Paster et al., 2022; Yang et al., 2022) attempt to solve the problem. Generally, they fit a state transition model, either for pessimistic planning (Villaflor et al., 2022) through sampling from VAEs (Zhang et al., 2018) , or to disentangle conditioning goals from environmental stochasticity (Paster et al., 2022; Yang et al., 2022 ) by adversarial training (Shafahi et al., 2019) . While adding complexity, these methods are applicable only when state transition functions can be learned adequately, which is often not the case for driving because of the uncertainty brought by complex interactions and partial observability (Sun et al., 2022a; Murphy, 2000) . Furthermore, driving trajectories encompass stochastic impact over an excessive number of timesteps. This dilutes the information related to current timestep decision-making in the outcome return and challenges VAE and adversarial training (Wu et al., 2017; Burgess et al., 2018) .In this paper, we take an initial step to customize DTs for stochastic driving environments without introducing transition or complex generative models. Specifically, our insight comes from the two driving tasks in Fig. 1 (a): When going straight, the cumulative rewards from Task I & II (termed as the global return) contains too much stochastic influence to provide effective supervision. Nonetheless, a viable strategy involves conditioning solely on the truncated return from Task I to mitigate the influence of environmental stochasticity (less stochastic timesteps, lower return variance), which still preserves rewards over sufficient timesteps for the optimization of current actions. Further experiments validate the point and we 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 2 (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 a sufficiently long horizon approximates global return optimization, as shown in Fig. 1(c) .The remaining problem is to specify the span of truncated returns. Specifically, our proposed UNcertainty-awaRE deciSion Transformer (UNREST) quantifies the impacts of uncertainties by the conditional mutual information between transitions and returns, which bypasses the complexity associated with transition and generative modeling. Sequences are then segmented into certain and uncertain parts accordingly. With the minimal impact of uncertainty in 'certain parts', we set the conditioning goal as the cumulative reward to the segmented position (with the number of timesteps), which can reflect true outcomes of decisions and 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 is adopted during inference for cautious planning. Highlights are:• We present UNREST, an uncertainty-aware sequential decision framework to apply offline RL in long-horizon and stochastic driving environments. The code will be public when published.• Recognizing the properties of driving environments, we propose a novel model-free environmental uncertainty measurement and segment sequences accordingly. Based on these, UNREST bypasses challenging generative training in previous works, by replacing global returns in DTs with less uncertain truncated returns (or dummy tokens) to learn from the true outcomes of agent actions.• We extensively evaluate UNREST on CARLA (Dosovitskiy et al., 2017) , where it consistently outperforms strong baselines (5.2% and 6.5% driving score improvement in seen and unseen scenarios). Additional experiments also prove the efficacy of our uncertainty estimation strategy. 
related works: 
 In this section, we review works about sequence modeling algorithms in RL and uncertainty estimation strategies as the foundation of our work and highlight the differences and contributions of UNREST. More related works on vehicle planning can be found in App. A.1.Offline RL as Sequence Modeling: Despite the potential to learn directly from offline data and the prospects of a higher performance ceiling (than IL) (Fujimoto et al., 2019) , the long-horizon and stochastic nature of driving environments still undermine the efficacy of current offline RL applications in autonomous driving tasks (Diehl et al., 2021; Shi et al., 2021; Li et al., 2022a) .Encouraged by the success of Transformer (Vaswani et al., 2017) in sequence modeling tasks (Brown et al., 2020; OpenAI, 2023) , a recent line of work (Chen et al., 2021; Janner et al., 2021; Furuta et al., 2022; Lee et al., 2022; Hu et al., 2023) adapts Transformer into RL by viewing offline RL as a sequence modeling problem. Typically, Decision Transformer (DT) (Chen et al., 2021) predicts actions by feeding target returns and history sequences, while Trajectory Transformer (TT) (Janner et al., 2021) further exploits the capacity of Transformer through jointly predicting states, actions, and rewards and planning by beam searching. The theoretical basis of these models is revealed in Generalized DT (Furuta et al., 2022) : they are trained conditioned on hindsight (e.g. future returns) to reach desired goals. DTs naturally consider outcomes of multi-step decision-making and have demonstrated efficacy in long-horizon tasks (Chen et al., 2021) .However, recent works (Brandfonbrener et al., 2022; Paster et al., 2022; Yang et al., 2022) have pointed out the fundamental problem with this training mechanism. Specifically, in stochastic environments like autonomous driving, certain outcomes can be achieved by accidental environment transitions, hence cannot provide effective action supervision. Notably, this is in fact a more general problem to do with all goal-conditioned behavior cloning and hindsight relabeling algorithms (Eysenbach et al., 2022; Štrupl et al., 2022; Paster et al., 2020) , but in this work we specifically focus on solutions to DTs. To tackle this issue, ESPER (Paster et al., 2022) adopts adversarial training (Shafahi et al., 2019) to learn returns disentangled from environmental stochasticity as a condition. Similarly, DoC (Yang et al., 2022) utilizes variational inference (Zhang et al., 2018) to learn a latent representation of the trajectory, which simultaneously minimizes the mutual information (so disentangled) with environment transitions to serve as the conditioning target. Besides, SPLT (Villaflor et al., 2022) leverages a conditional VAE (Zhang et al., 2018) to model the stochastic environment transitions. As the driving environment contains various interactions that make it difficult to model (Sun et al., 2022a) , and the long driving trajectory hinders generative training, our study proposes a novel uncertainty estimation strategy to customize DTs without transition or generative models. Besides, different from EDT (Wu et al., 2023) that dynamically adjusts history length to improve stitching ability, UNREST focuses on segmenting sequences and replacing conditions to address the overly optimistic problem.Uncertainty Estimation: Uncertainty estimation plays a pivotal role in reliable AI. One typical method for uncertainty estimation is probabilistic bayesian approximation, either in the form of dropout (Gal & Ghahramani, 2016) or conditional VAEs (Blundell et al., 2015; Dusenberry et al., 2020) , which computes the posterior distribution of model parameters. On the contrary, the deep deterministic methods propose to estimate uncertainty by exploiting the implicit feature density (Franchi et al., 2022) . Besides, deep ensembles (Lakshminarayanan et al., 2017a) train 𝐾 neural networks simultaneously, with each module being trained on a different subset of the data, and use the variance of the network outputs as a measure of uncertainty. In this work, we utilize the ensemble approach, which has been widely used in the literature for trajectory optimization (Vlastelica et al., 2021) , online and offline RL (Wu et al., 2021; An et al., 2021) , to jointly train 𝐾 variance networks (Kendall & Gal, 2017b) for estimating the uncertainty of returns. 
preliminary: 
 We introduce preliminary knowledge for UNREST in this section. is a collected interaction trajectory composed of transitions with horizon 𝑇.Following DTs (Chen et al., 2021) , we pose offline RL as a sequence modeling problem where we model the probability of the sequence token 𝑥 𝑡 conditioned on all tokens prior to it: 𝑝 𝜃 (𝑥 𝑡 |x <𝑡 ), where x <𝑡 denotes tokens from step 1 to (𝑡 -1), like the GPT (Brown et al., 2020) architecture. DTs consider a return-conditioned policy learning setting where the agent at step 𝑡 receives an environment state 𝑠 𝑡 , and chooses an action 𝑎 𝑡 conditioned on the future return 𝑅 𝑡 = 𝑇 𝑡 ′ =𝑡 𝑟 𝑡 ′ . This leads to the following trajectory representation:EQUATIONwith the objective to minimize the action prediction loss, i.e. maximize the action log-likelihood:EQUATIONThis supervised learning objective is the cause for DTs' limitations in stochastic environments, as it over-optimistically assumes actions in the sequence can reliably achieve the corresponding returns. At inference time, given a prescribed high target return, DTs generate actions autoregressively while receiving new states and rewards to update the history trajectory. 
approach: unrest: 
 In this section, we begin with an overview of UNREST's insight and architecture, followed by detailed explanations of the composition modules used in the approach.An overview of the proposed approach UNREST is illustrated in Fig. 2 . To address the overly optimistic issue, our key idea is to quantify the impact of environmental uncertainty, and learn to perform aggressively or cautiously at states with different levels of environmental impacts.To achieve this, we train 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 test time, an uncertainty predictor is involved to guide decision-making at different states. In the following, we introduce each module with details.Instead of conventional uncertainties that reflect variances of distributions (Wu et al., 2021; Kendall & Gal, 2017a) , in this paper we estimate the impact of environmental stochasticity, what we really care about for policy learning, as an indirect measure of environmental uncertainty. In particular, we propose to model the impact of the transition (𝑠 𝑡 -1 , 𝑎 𝑡 -1 → 𝑠 𝑡 ) on return 𝑅 𝑡 through their conditional mutual information (Seitzer et al., 2021; Duong & Nguyen, 2022) (3)Global-Return Emb.Return-Span Emb.h=1 h=80Uncertainty Evaluation. < 𝜖+ u !"# < 𝜖 u !$%$# < 𝜖 u !$%$& + 𝑅 !"# $ 𝑅 !$%$# ' s !"# 𝑎 !"# $ 𝑎 !"# $ 𝑎 ! $ 𝑎 !%& $ 𝑎 !%&%# s ! 𝑎 ! s !%& s !%&%# 𝑎 !%& 𝑎 !%&%# h=79 + 𝑅 !$%$& ' s !%&%' 𝑎 !%&%' $ 𝑎 !%&%' ≥ 𝜖 u ! < 𝜖 u !$% ••• Dum.Dum. Afterward, two models feed their respective outputs into variance networks (Kendall & Gal, 2017b) to predict Gaussian distributions N (•) of returns, which are optimized by maximizing log-likelihoods:EQUATIONPractically, the networks are implemented as ensembles, which together form Gaussian Mixture Models (GMM) (Mai et al., 2022) to better capture the return distribution as in App. A.2. Compared to challenging high-dimensional transition function learning and complex generative training, the return distributions can be effectively learned with fewer resources. Finally, the impact of the stochastic environmental transition (i.e. the conditional mutual information) is calculated through the Kullback-Leibler (KL) divergence (Joyce, 2011) between these two distributions, as a means to measure the environmental uncertainty at timestep 𝑡:EQUATIONwhere a larger divergence implies more influence on the return from the transition (𝑠 𝑡 -1 , 𝑎 𝑡 -1 → 𝑠 𝑡 ).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 to higher returns after training. Otherwise, in 'uncertain parts', the seemingly high return actions may easily cause safety issues due to environmental uncertainty, thus we simply ignore the stochastic returns, setting conditions as dummy tokens for behavior cloning.Segmentation strategy is therefore a crucial point to reinvent DTs. To distinguish between different levels of environmental stochasticity, we define an uncertainty threshold 𝜖. Next, we estimate uncertainties 𝑢 𝑡 by Eq. 5 for each timestep and record those larger than 𝜖 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 a 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 𝑐 -1 timesteps are all identified as certain. Since uncertain timesteps may occur intensively and intermittently over a particular duration of driving (e.g. at an intersection), the hyperparameter 𝑐 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:EQUATIONwhere the conditioning return is modified as 𝑅 ℎ 𝑡 = ℎ 𝑡 -1 𝑘=0 𝑟 𝑡+𝑘 , which only involves rewards in the next ℎ 𝑡 steps (called the return-span). In 'uncertain parts', ℎ is set as empty for meaningless conditions (i.e. the dummy tokens). In this way UNREST will learn to ignore the insignificant conditioned 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 to the next segmentation step, so that the conditioning return 𝑅 ℎ 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 Alignment Bound) Assuming that the rewards obtained are determined by transitions (𝑠, 𝑎 → 𝑠 ′ ) at each timestep and UNREST is perfectly trained to fit the expert demonstrations, then the discrepancy between target truncated returns and URNEST's rollout returns is bounded by a factor of environmental stochasticity and data coverage.It reveals that under natural assumptions in driving environments, UNREST can generalize to achieve high returns with bounded error at less stochastic states if expert demonstrations cover corresponding actions. The formal statement and proof of the proposition are left to App. B.Notably, due to sequence segmentation, the truncated returns no longer necessarily encompass all time steps from the present to the end of the sequence, which necessitates the return-span as a condition to provide information about the count of timesteps involved in returns. This enables UNREST to learn to get return 𝑅 ℎ 𝑡 over future ℎ 𝑡 timesteps. Otherwise, the model may be confused by the substantial differences in the magnitude of return conditions with varying timestep lengths (Hu et al., 2023) .Policy formulation: Unlike return transformers, DT takes the segmented sequence τ seg as input:EQUATIONHere we use similar notations as return transformers in Sec. 4.2 except that models are parameterized by 𝜃. In the above formulation, a return-span embedding 𝑓 ℎ 𝜃 (ℎ 𝑡 ) 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-return embedding 𝑓 𝑅 𝜃 (𝑅 𝑡 ) is optionally added to the output x𝑠 𝑡 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:EQUATIONThe learning objective is modified from Eq. 2 of DT, with detailed training process in App. C:EQUATIONDuring inference, we still differentiate between certain and uncertain states to account for the impacts of environmental stochasticity. To enable real-time prediction of environmental stochasticity, we introduce a lightweight uncertainty prediction model 𝑢 𝜓 (•). At each timestep, we query the predictor to obtain the uncertainty measure. If the current state transition is highly uncertain, we set the conditioning target to a dummy token to facilitate cautious planning, consistent with training. Conversely, at states with certain transitions, the planner acts aggressively to attain high target returns.The predictor can be implemented by various approaches, like a neural network or a heuristic estimator, whose results are summarized in Tab. 8 and 9. Practically, we choose the KD-Tree (Redmond & Heneghan, 2007) for its high computational efficiency and favorable estimation performance, with states as tree nodes and uncertainties estimated by return transformers as node values.Different from the conventional planning procedure of DTs, UNREST requires the specification of not only the target global return 𝑅 1 , but also the initial return-span ℎ 1 = 𝐻 and the target truncated return 𝑅 ℎ 1 . After segmentation, the effective planning horizon of the trained sequences is reduced to the return-span ℎ 𝑡 . Once ℎ 𝑡 reaches 1, we need to reset the target return and the return-span. Practically, we simply reset ℎ 𝑡 to a fixed return horizon 𝐻. For the truncated return, we train a return prediction model 𝑅 ℎ 𝜑 (•) similar to that defined in Eq. 4 and take the upper percentile 𝜂 of the predicted distribution as the new target return to attain. The hyperparameter 𝜂 can be tuned for a higher target return. We do not need to consider targets at 'uncertain states' since they are just set as dummy tokens. The complete planning process is summarized in Alg. 1. The offline dataset D is collected from the CARLA simulator (Dosovitskiy et al., 2017) with its built-in Autopilot. Specifically, we collect 30 hours of data from 4 towns (Town01, Town03, Town04, Town06) under 4 weather conditions (ClearNoon, WetNoon, HardRainNoon, ClearSunset), saving tuple (𝑠 𝑡 , 𝑎 𝑡 , 𝑟 𝑡 ) at each timestep. More details about the state, action compositions, and our reward definitions are left to App. D.We evaluate models at training and new driving scenarios and report metrics from the CARLA challenge (Dosovitskiy et al., 2017; team, 2020) to measure planners' driving performance: infraction score, route completion, success rate, and driving score. Besides, as done in (Hu et al., 2022) , we also report the normalized reward (the ratio of total return to the number of timesteps) to reflect driving performance at timestep level. Among them, the driving score is the most significant metric that accounts for various indicators like driving efficiency, safety, and comfort.Baselines: First, we choose two IL baselines: vanilla Behavior Cloning (BC) and Monotonic Advantage Re-Weighted Imitation Learning (MARWIL) (Wang et al., 2018) . Apart from IL baselines, we also include state-of-the-art offline RL baselines: Conservative Q-Learning (CQL) (Kumar et al., 2020) , Implict Q-Learning (IQL) (Kostrikov et al., 2021) . Constraints-Penalized Q-Learning (CPQ) (Xu et al., 2022) is chosen as a safe (cautious) offline RL baseline. Besides, we select two classic Transformer-based offline RL algorithms: Decision Transformer (DT) (Chen et al., 2021) and Trajectory Transformer (TT) (Janner et al., 2021) as baselines. Finally, we adopt three algorithms as rigorous baselines: Separated Latent Trajectory Transformer (SPLT) (Villaflor et al., 2022) , Environment-Stochasticity-Independent Representations (ESPER) (Paster et al., 2022) , and Dichotomy of Control (DoC) (Yang et al., 2022) . These algorithms fit state transition models and employ generative training to mitigate DTs' limitations in stochastic environments. 
driving performance: 
 Firstly, we implement all the models, then evaluate their performance at training (Town03) and new (Town05) scenarios (Q1), whose results are summarized in Tab. 1.Analyzing the results, we first notice that DT (Chen et al., 2021) performs worst among all sequence models, with a significant gap compared to the simple offline RL baseline IQL (Kostrikov et al., 2021) . BC 51.9 ± 1.9 (45.6 ± 5.2) 37.9 ± 4.7 (36.3 ± 8.6) 79.7 ± 6.0 (77.1 ± 7.8) 54.5 ± 1.8 (47.0 ± 4.8) 0.63 ± 0.02 (0.61 ± 0.04) MARWIL (Wang et al., 2018) 54.3 ± 2.0 (47.8 ± 3.4) 44.8 ± 3.2 (42.2 ± 2.2) 81.4 ± 3.2 (80.2 ± 3.8) 57.9 ± 2.0 (48.4 ± 4.4) 0.65 ± 0.02 (0.63 ± 0.01) CQL (Kumar et al., 2020) 55.0 ± 2.4 (50.7 ± 2.8) 48.6 ± 4.5 (42.8 ± 4.6) 80.5 ± 3.4 (78.7 ± 5.2) 62.4 ± 3.0 (57.7 ± 5.0) 0.65 ± 0.03 (0.62 ± 0.02) IQL (Kostrikov et al., 2021) 55.9 ± 3.3 (52.2 ± 3.3) 50.2 ± 6.2 (44.2 ± 2.2) 77.2 ± 4.6 (68.8 ± 4.2) 68.4 ± 2.4 (62.6 ± 6.2) 0.66 ± 0.03 (0.60 ± 0.01) CPQ (Xu et al., 2022) 54.7 ± 3.0 (53.8 ± 2.7) 46.6 ± 3.5 (40.3 ± 5.7) 79.4 ± 4.2 (78.0 ± 4.6) 66.9 ± 2.8 (64.4 ± 4.2) 0.63 ± 0.02 (0.64 ± 0.02)DT (Chen et al., 2021) 55.2 ± 2.0 (47.6 ± 1.2) 48.0 ± 4.7 (46.4 ± 0.3) 82.6 ± 1.0 (81.8 ± 4.9) 57.4 ± 1.3 (47.3 ± 3.4) 0.66 ± 0.01 (0.64 ± 0.02) TT (Janner et al., 2021) 58.3 ± 3.3 (54.8 ± 2.2) 45.9 ± 5.2 (52.0 ± 4.4) 78.4 ± 5.8 (77.0 ± 4.6) 63.6 ± 3.6 (56.0 ± 5.0) 0.74 ± 0.02 (0.64 ± 0.03) SPLT (Villaflor et al., 2022) 57.8 ± 4.9 (56.4 ± 6.1) 23.1 ± 8.1 (39.5 ± 9.5) 36.7 ± 8.7 (48.2 ± 9.7) 73.9 ± 1.4 (70.7 ± 8.3) 0.55 ± 0.03 (0.57 ± 0.05) ESPER (Paster et al., 2022) 54.8 ± 2.0 (51.2 ± 3.1) 48.5 ± 4.3 (53.3 ± 2.0) 77.3 ± 5.4 (84.7 ± 1.4) 56.2 ± 2.2 (46.9 ± 5.6) 0.64 ± 0.04 (0.63 ± 0.02) DoC (Yang et al., 2022) 56.9 ± 3.1 (54.4 ± 2.3) 49.9 ± 5. In new scenarios, it even shows a similar performance to BC. We attribute this to the uncertainty of global return it conditions on. When the conditioning target does not reflect the true outcomes of agent actions, DT learns to ignore the target return condition and makes decisions solely based on the current state like BC. Furthermore, ESPER (Paster et al., 2022) and DoC (Yang et al., 2022 ) also perform poorly in both scenarios, which may be a result of ineffective adversarial and VAE training of long-horizon and complex driving demonstrations.To tackle environmental 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 reward. TT instead learns a transition model regardless of environmental stochasticity and behaves aggressively. It attains the highest normalized reward at training scenarios since the metric prioritizes planners that rapidly move forward (but with low cumulative reward 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) like in Fig. 3(a) .Notably, UNREST achieves the highest driving score, route completion rate, success rate in both seen and unseen scenarios, and highest normalized reward in new scenarios without the need to learn transition or complex generative models. For the driving score, it surpasses the strongest baselines, TT (Janner et al., 2021) over 5% in training scenarios, and SPLT (Villaflor et al., 2022) over 6% in new scenarios (in absolute value), achieving a reasonable balance between aggressive and cautious behavior. For all the metrics, the results demonstrate that UNREST obtains significant improvement in terms of safety, comfort, and efficiency, and effectively increases the success rate of driving tasks. It demonstrates that the truncated returns successfully mitigate impacts of stochasticity and provide effective action supervision. In App. F.2, we test and find that UNREST only occupies slightly more resources than DT, and runs significantly faster while consuming less space than TT and SPLT, which additionally fit state transition models.The key components of UNREST include global return embedding, return-span embedding, ensemblebased uncertainty estimation, and the uncertainty-guided planning process. In this section, we conduct ablation experiments by separately removing these components to explore their impacts on the overall performance of UNREST (Q2). Results are shown in Tab. 2 and Tab. 3.The ablation results show pronounced differences, and it is apparent that the elimination of any part of the three components leads to a decline in UNREST's driving score in new scenarios. Among them, the global return embedding has the slightest impact, which suggests that the highly uncertain global return may not provide effective guidance, or that the truncated return is already sufficient for making reasonable decisions (temporal locality, Prop. 2). When the return-span embedding is removed, the absolute driving score drops by about 6%. This implies that the introduction of return-span embedding provides necessary information about the timesteps needed to achieve the target return. Removing the ensemble of return transformers (i.e. the GMM) induces a significant performance drop in both scenarios. It shows that a simple Gaussian distribution cannot well express the return distribution, resulting in poor uncertainty calibration capability (corresponding with results in Fig. 6 ). Finally, after canceling uncertainty estimation at test time, the driving and infraction scores of UNREST drop dramatically, which proves the importance of cautious planning.Finally, We verify the interpretability of UNREST's uncertainty estimation through visualizations (Q3). Typically, in Fig. 4 (a) we observe an initial increase in uncertainty as the ego-vehicle enters the lane to follow another vehicle, owing to the lack of knowledge about the other vehicle's behavior. This uncertainty gradually decreases back below the threshold when the vehicle stabilizes in the following state. Fig. 4 (b) shows a green light crossing scenario. While approaching the green light, the uncertainty about the light state causes the uncertainty to rise quickly. After the vehicle has moved away from the traffic light, the uncertainty immediately drops back below the threshold. 
conclusion: summary and limitations: 
 The paper presents UNREST, an uncertainty-aware decision transformer to apply offline RL in stochastic driving environments. Specifically, we propose a novel uncertainty measurement by computing the divergence of return prediction models, bypassing complex transition and generative training. Then based on properties we discover at driving, we segment sequences w.r.t. estimated uncertainty and adopt truncated returns as 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 this work is that the inference process of UNREST is somewhat complex with auxiliary return, uncertainty estimation models, and hyperparameters. One possible direction for improvement is to integrate return and uncertainty predictions into the model architecture, which we leave for future work. Although this work is evaluated in the CARLA simulator, we believe the proposed framework can surmount the sim-to-real gap and benefit practical autonomous driving. 
 
 Appendix: A.1 PLANNING FOR SELF-DRIVING Most autonomous driving systems (Montemerlo et al., 2008; Urmson et al., 2008; Ziegler et al., 2014) adopt a modular pipeline to break down the massive driving task into a set of submodules, with planning being one 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 is essentially a decision problem. Often engineers manually design rules (Ziegler et al., 2014; Chen et al., 2019b) for specific driving scenarios, which struggles to scale for complex tasks. In contrast, IL is the earliest (Pomerleau, 1988) and most widely (Bansal et al., 2018; Zeng et al., 2019; Hu et al., 2022; Bronstein et al., 2022; Codevilla et al., 2019; Zhang et al., 2021) used learning-based planning algorithm, which learns from offline data for either a cost function (Zeng et al., 2019) or an executable policy (Bansal et al., 2018; Hu et al., 2022; Bronstein et al., 2022; Codevilla et al., 2019) . However, they are strictly limited by expert quality (Ross et al., 2011) and will perform poorly when encountering out-of-distribution (OOD) states (Zhang et al., 2021; Ross et al., 2011; Prakash et al., 2020) . RL algorithms (Chen et al., 2019a; 2020) 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.Variance networks estimate uncertainty by loss attenuation (Kendall & Gal, 2017b) . Treating the output as a Gaussian distribution N (•), given input 𝑥 and its label 𝑦, the network outputs mean 𝜇(𝑥) and variance 𝜎 2 (𝑥) at its final layer. The network parameter 𝜙 is subsequently optimized by maximizing sample log-likelihood:EQUATIONAs the variance predicted by neural networks may be not well-calibrated or overestimated (Kuleshov et al., 2018) , in this paper we practically train an ensemble of 𝐾 variance networks with data drawn from different subsets of the dataset (implemented by masking) and estimate the variance according to (Mai et al., 2022) :𝜎 2 = 𝐾 𝑙=1 𝜎 2 𝑙 + 𝜇 2 𝑙 -𝜇 2, where 𝜇 = 𝐾 𝑙=1 𝜇 𝑘 /𝐾, (11) which has been proven to have strong estimation capability (Mai et al., 2022) .We next formally state our theorem, followed by detailed proofs of the proposition.Before stepping into the main theorem, we first introduce the problem setting that we 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 𝑔(𝑠 1 , 𝑎 1:ℎ ) to denote the cumulative rewards by rolling out the open loop sequence of actions 𝑎 1:ℎ under the deterministic dynamics 𝑃 : S × A → S and reward function 𝑟 : S × A → [0, 1] (so that the maximum return difference over future ℎ timesteps is ℎ). Moreover, 𝐽 ℎ (𝜋) represents the rollout rewards of policy 𝜋 over future ℎ timesteps. Next, we formally introduce the assumption, theorem and its corresponding proof, which is inspired by the reference (Brandfonbrener et al., 2022) :Theorem 1 (UNREST Alignment Bound) Consider an MDP (S, A, 𝑃, 𝑟, 𝛾), expert behavior 𝛽 : S × A → [0, 1] and conditioning function 𝑓 : S × N + → R. Assume the following:1. Return Coverage: 𝑝 𝛽 (𝑔 = 𝑓 (𝑠 1 , ℎ)|𝑠 1 ) ≥ 𝛼 𝑓 for the initial state 𝑠 1 and return-span ℎ.2. Near Determinism: 𝑝(𝑟 ≠ 𝑟 (𝑠, 𝑎) 𝑜𝑟 𝑠 ′ ≠ 𝑃(𝑠, 𝑎)|𝑠, 𝑎) ≤ 𝛿 at all state-action pairs (𝑠, 𝑎)for dynamics 𝑃 and reward function 𝑟.3. Consistency of 𝑓 : f(s,h)=f(s',h-1) for all states 𝑠.Then (URT shorted for UNREST):EQUATION)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 𝛼, and the horizon ℎ.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 1 (Determinism Equality) Assuming that the rewards obtained are uniquely determined by transitions (𝑠, 𝑎 → 𝑠 ′ ) at each timestep, then there exists 𝜖 > 0 and 𝛿 > 0, such that:EQUATION)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 its contrapositive proposition:If 𝑝(𝑟 ≠ 𝑟 (𝑠, 𝑎) 𝑜𝑟 𝑠 ′ ≠ 𝑃(𝑠, 𝑎)|𝑠, 𝑎) > 𝛿, then ∃𝜖, s.t. 𝐷 KL 𝑝(𝑅 𝑡 |τ <𝑡 ) || 𝑝(𝑅 𝑡 |𝑠 𝑡 , τ <𝑡 ) > 𝜖 .(14) First, since 𝑟 is uniquely (different transitions lead to distinct rewards) determined by (𝑠, 𝑎 → 𝑠 ′ ), we can induce that 𝑟 ≠ 𝑟 (𝑠, 𝑎) is equivalent with 𝑠 ′ ≠ 𝑃(𝑠, 𝑎) as the only difference can occur in 𝑠 ′ . Therefore, we need to prove:EQUATION)Considering the extreme case, set 𝛾 → 0, thus we have 𝑅 𝑡 → 𝑟 𝑡 . From the perspective of proof by contradiction, the KL-divergence cannot be smaller than any 𝜖 > 0, since 𝑝(𝑟 𝑡 |τ >𝑡 ) and 𝑝(𝑟 𝑡 |𝑠 𝑡 , τ >𝑡 ) must differ for that the probability 𝑟 𝑡 ≠ 𝑟 (𝑠 𝑡 -1 , 𝑎 𝑡 -1 ) is larger than 0 and can be determined by 𝑠 𝑡 . To this end, we have proved the contrapositive proposition of the forward direction, which is equivalent to the original proposition.Therefore, we can summarize from the above theorem and lemma that UNREST can generalize to achieve high returns at 'certain states' as long as the corresponding actions are covered by the expert.The next theorem proof is largely built on (Brandfonbrener et al., 2022) . Note that we omit the superscript 'URT' to simplify the equations. First, we expand the left term in Eq. 12:EQUATION)The last step follows by bounding the difference between 𝑔 1 and 𝑔(𝑠 1 , 𝑎 1:ℎ ) by the maximum return difference ℎ and a union of probability bound over ℎ timesteps based on the near determinism assumption:EQUATION)For the first term in Eq. 16, it can also be expressed and bounded by the maximum return difference:EQUATION)To bound this term, we need to further expand the distribution 𝑝 𝜋 𝑓 . First, with the assumption that UNREST is perfectly fitted to the expert dataset, we can deduce the following by the bayesian law:EQUATION)For simplification, we use 𝑠 𝑡 = 𝑃(𝑠 1 , 𝑎 1:𝑡 -1 ) to denote the state reached by following the deterministic dynamics defined by 𝑃 till timestep 𝑡. Next, based on the near determinism, consistency and coverage assumption, we can expand 𝑝 𝜋 𝑓 to get:EQUATION)Similarly, we can further expand 𝑝 𝜋 𝑓 (𝑎 2:ℎ | 𝑠 1 , s2 ) using the same rule:EQUATIONSubstituting this back to Eq. 20, we have:𝑝 𝜋 𝑓 (𝑎 1:ℎ | 𝑠 1 ) ≤ 𝛽 (𝑎 1 | 𝑠 1 ) 𝛽 (𝑎 2 | s2 ) 𝑝 𝛽 (𝑔 2 = 𝑓 ( s2 , ℎ -1) | s2 ) 𝑝 𝛽 (𝑔 1 = 𝑓 (𝑠 1 , ℎ) | 𝑠 1 ) • 𝑝 𝛽 (𝑔 2 = 𝑓 ( s2 , ℎ -1) | s2 , 𝑎 2 ) 𝑝 𝛽 (𝑔 2 = 𝑓 ( s2 , ℎ -1) | s2 ) 𝑝 𝜋 𝑓 (𝑎 3:ℎ | 𝑠 1 , s2 , s3 ) + 2𝛿 1 𝛼 𝑓 + 1 . . . recursively expand 𝑝 𝜋 𝑓 ≤ ℎ 𝑡=1 𝛽 (𝑎 𝑡 | s𝑡 ) 𝑝 𝛽 (𝑔 ℎ = 𝑓 ( sℎ , 1) | sℎ , 𝑎 ℎ ) 𝑝 𝛽 (𝑔 1 = 𝑓 (𝑠 1 , ℎ) | 𝑠 1 )+ ℎ𝛿 1EQUATIONThe last step is deduced by the trajectory determinism and the consistency of the conditioning function 𝑓 . Multiplying this back to Eq. 18 and noticing that the two indicator functions can never both be 1, we can yield that:EQUATIONThis can finally yield the bound in Eq. 12 by adding back to Eq. 16.To get a more clear understanding of our training pipeline, we present the training procedure of UNREST in Alg. 2.Update network parameters based on learning objective in Eq.9; Output: Trained return transformers 𝑅 𝜑 𝑎 (•), 𝑅 𝜑 𝑠 (•) and policy 𝜋 𝜃 (•).Our dataset is collected from the CARLA simulator, using its built-in Autopilot, which is a rulebased motion planner. Specifically, CARLA (Dosovitskiy et al., 2017) is an open-sourced simulator designed for autonomous driving research. It provides researchers with a realistic environment that faithfully simulates traffic dynamics, weather conditions, and high-fidelity 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 at a 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 (𝑠 𝑡 , 𝑎 𝑡 , 𝑟 𝑡 ), whose compositions will be introduced in detail in the next paragraphs.The symbol 𝑟 𝑡 ∈ R denotes the reward returned by the environment at the timestep 𝑡. Typically, our reward design is revised from Roach (Zhang et al., 2021) , encompassing factors such as speed, safety, and comfort.EQUATIONAmong them,𝑟 speed = 1 -|𝑣-𝑣 desired | 𝑣 maxrepresents the reward of approaching the target speed; 𝑟 position = -0.5Δ 𝑝 represents the reward obtained from driving in the correct position, where Δ 𝑝 is the lateral distance between the ego vehicle and the centerline of the target route; 𝑟 rotation = -Δ 𝑟 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; 𝑟 action incurs a penalty of -0.1 when the steering in current timestep differs from the previous step by more than 0.01, promoting smooth and comfortable driving; 𝑟 terminal 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) (Xu et al., 2022) algorithm also needs a constraint definition for training. Specifically, we adopt the same setting as in (Li et al., 2022b) , where each time of collision or traffic rule violation incurs a penalty of 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 (Villaflor et al., 2022) .Action: The action 𝑎 ∈ R 2 consists of two values, representing the target angle and target velocity that will be fed into two separate PID controllers (Johnson & Moradi, 2005) .The observed state 𝑠 ∈ R 37 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 if no leading vehicle is perceived in the horizon (iii) relative distance and velocity to the pedestrian ahead, which will be set to default maximum values if no pedestrian is perceived in the horizon (iv) relative velocities and distances to 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 to 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.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, all training processes are run on a NVIDIA A10 while inference is conducted on one NVIDIA 3090 GPU for fair comparison.Training Process: For all training processes, We employ an AdamW optimizer with a learning rate of 10 -4 , and a consistent batch size of 256. For the training of sequence models, a history length of 10 is carefully selected to be sampled each time, which corresponds to a real-world duration of 1 second. To determine the number of epochs for training, a dynamic value is assigned based on the dataset's size, defined by the following formula: 𝑛 epochs = int Details for segmentation: As discussed in Sec. 4.3, we train two return transformers for sequence segmentation. In practical implementation, to accurately quantify environmental uncertainty, we employ the ensemble method (Lakshminarayanan et al., 2017a) to train each return transformer. Using Eq. 11, we recalculate the variance and expectation of the returns. Then the sequences are segmented into certain and uncertain parts w.r.t. estimated uncertainties according to the principles we introduce in Sec. 4.3.In our experiments, we discover that 'uncertain states' tend to occur consecutively. To ensure meaningful segmentation, we enforce a minimum 'uncertain part' length of 𝑐 = 20 frames and select the uncertainty threshold 𝜖 = 3 according to Fig. 5(a) . Under these two hyperparameter settings, we find the segmented sequences often correspond to the completion of a significant driving task.Details for return-span embedding: We introduce the return-span embedding 𝑓 ℎ 𝜃 (ℎ 𝑡 ) 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 return-span embedding to the return embedding, as stated in the main text. The second variant considers treating 𝑓 ℎ 𝜃 (ℎ 𝑡 ) as a distinct 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).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: In the practical implementation, following DT (Chen et al., 2021) , 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.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 (team, 2020) . 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 at inference time to enable cautious planning. Typically, we implement more variants (results are shown in Tab. 8 and Tab. 9) other than KD-Tree (Redmond & Heneghan, 2007) for test-time uncertainty estimation:• Network prediction: 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) to each state in the training dataset. Subsequently, these labels are utilized to train an uncertainty network. • Heuristic: The heuristic method turns out to be a straightforward and intuitive alternative. Specifically, 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 is uncertain. • KD-tree: Ultimately, our chosen methodology employs a KD-tree for uncertainty estimation.Initially, we construct a KD-tree that stores both the states and the uncertainties computed by the trained return transformers. As test time, we sample the nearest neighbors of the current state from the tree and compute the mean value of neighbors as the current state's uncertainty.We select several well-established approaches to serve as our comparative baselines. Specifically, we include Behavior Cloning (BC), Implicit Q-Learning (IQL) (Kostrikov et al., 2021) , as well as prominent sequence models such as Decision Transformer (DT) (Chen et al., 2021) , Trajectory Transformer (TT) (Janner et al., 2021) , SPLT (Villaflor et al., 2022) , and ESPER (Paster et al., 2022) . To ensure a fair comparison, for all these baselines, we adopt the default hyperparameter settings from the repository https://github.com/avillaflor/SPLT-transformer.For Monotonic Advantage Re-Weighted Imitation Learning (MARWIL) (Wang et al., 2018) and Conservative Q-Learning (CQL) (Kumar et al., 2020) , we adopt the default hyperparameter from the Ray RLlib https://docs.ray.io/en/latest/rllib/index.html. Finally for the safe offline RL baseline Constraints-Penalized Q-learning (CPQ) (Xu et al., 2022) , we employ the default hyperparameters from the repository https://github.com/liuzuxin/OSRL.The full list of UNREST's hyperparameters can be found in Tab. 4, Tab. 5, and Tab. 6. In this section, we supplement more experimental results including the visualizations, complexity comparisons, sensitivity analysis of hyperparameters, and the results for more UNREST variants.Uncertainty visualizations: We present more visualizations of UNREST's uncertainty estimation in Fig. 5 to validate its interpretability. Initially, we visualize the uncertainty distribution in Fig. 5 (a).As we can see the majority of uncertainty associated with state transitions falls within the range of 0 to 10 -4 , aligning with intuitive expectations. In the context of autonomous driving environments with sparse vehicle presence and a 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 corner cases with large uncertainties.Then we seek uncertain transition states within the training dataset and visualize two typical scenarios. By associating the temporal steps with the uncertainty curves, it can be observed that as the ego-vehicle approaches the turning point, the uncertainty gradually increases since the ego-vehicle cannot forecast the traffic conditions in the lane after turning. After the turning behavior is almost completed, the uncertainty decreases below the threshold again. In the scenario of waiting at a red light, it can be observed that the uncertainty rapidly increases as the ego-vehicle approaches the red light due to the uncertainty in color changes of the traffic signal. However, after the ego-vehicle comes to a stop in front of the red light for approximately 15 frames, the uncertainty decreases sharply below the threshold. These visualization results effectively demonstrate the interpretability of our uncertainty measurement approach.Uncertainty Calibration: To provide a more intuitive impression of UNREST's effectiveness, we present in Fig. 6 the return distribution calibration results. Typically, the figure illustrates that using an ensemble of return transformers can significantly better calibrate 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 (Lakshminarayanan et al., 2017b) than the single return transformer variant. Notably, the uncertainty in the figure denotes the variance of the distribution, different from the so-called environmental uncertainty we used in the main text, that reflects the impact of environmental transitions (through conditional mutual information). The predicted return distribution has corresponding ground truth and thus can be directly calibrated like what we do in the figure. However, the impact of the environment has no ground truth value and can only be interpreted through visualizations like Fig. 4 and Fig. 5 . Impact of return horizon 𝐻: We investigate the influence of the return horizon length, denoted as 𝐻, on the performance of UNREST. The experimental results are depicted in Fig. 8 , where the return horizon varies from 50 to 1000. It is evident that a relatively modest value of 𝐻 (e.g., 𝐻=100) yields the most favorable performance, while both excessively small and excessively 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 to shortsightedness, thereby compromising performance.Impact of uncertainty threshold 𝜖: We next investigate the impact of the uncertainty threshold 𝜖 and present the results in Fig. 8 where we increase 𝜖 from 0.5 to 10. As is shown in the figure, a moderate value of 𝜖, namely 𝜖 = 3 contributes to the best performance, while both smaller and larger et al., 2022) . For ESPER, we report the results from their paper (Paster et al., 2022) . We report the mean and std for our method over 3 seeds with 10 trajectories for each seed. 'Stochastic' denotes customized stochastic environments according to (Yang et al., 2022) . The symbol '-' is used to indicate the omission of measurement results that are deemed unimportant. leading to the lowest driving score. Adopting a network to predict uncertainty at inference time yields better performance than the heuristic method, but it can often misclassify certain/uncertain states and still performs worse than the original model. Generally, the original UNREST model achieves the best driving scores in both training and new driving scenarios, which is the reason why we choose it as our ultimate planning framework. Finally, we also study a Reweighted BC variant, which additionally fit state value functions like that in MARWIL (Wang et al., 2018) and conduct 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, due to its significant increase in model complexity, we have refrained from incorporating this technique into the final model.For completeness, we evaluate our method on D4RL Mujoco tasks, though they are generally deterministic tasks. Specifically, we find our UNREST is generally competitive with SOTA baselines on these tasks, as shown in Tab. 10. The reason for his overall underperformance compared to DT lies in the unnecessarily forced sequence segmentation in such deterministic tasks, where it is more appropriate to retain global returns in order to extend the horizon length considered during planning. We also test UNREST in two customized stochastic D4RL environments from (Yang et al., 2022) , where UNREST achieves the highest return among the models. 


Decision: 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.

Reviews: review: title: title: Review of paper 220
paper_summary: paper_summary: 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.
main_review: main_review: The paper is well written and in combination with the supplementary material is presented with a lot of interesting experiments and ablation studies.  Although there are no real world robot results, the simulation results presented seem convincing compared to the baselines presented. 

My first concern is about the dataset collected. The dataset uses an inbuilt CARLA autopilot to collect the expert trajectories - what steps were taken to make sure that there was enough environment stochasticity in the dataset? From what I understand the stochasticity here is the uncertainty over the ado agent behaviors (and not different towns or weather conditions). For example, in the turning example presented in figure 1, the uncertainty in making the right turn without collision comes from different behaviors that the ado agent exhibits (adversarial vs cooperative). Did the authors explicitly model different ado agent behaviors? was there any stochasticity added to the autopilot behavior? 

During inference, to estimate the uncertainty a KD-tree is used to mimic the uncertainty output from the return transformers. Presenting some details and evaluating the performance of the KD-tree like accuracy could be useful to understand the sensitivity of the planner to errors in uncertainty estimation.
questions: questions: 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? 
A 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? 

2) 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? 

3)  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?
limitations: limitations: 3
review_summary: review_summary: Weak accept
title: title: Uncertainty estimation to disentangle effect of stochastic environment transitions on reward
paper_summary: paper_summary: 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.
main_review: main_review: The paper is well written with clear explanations and concise notation to follow through. My questions are posted in the 'Questions For Rebuttal' section. 

Strengths
Explanations are clear without reliance on extensive jargon and equations. The Reinforcement learning formulation is standard.
The positioning of work with respect to past related work is commendable. Past work was explained clearly and how the authors' method differs from those was also succinctly described
The ablation experiments are extensive and demonstrate a clear advantage of the various modifications over the baseline DT method. 
The figures are easy to follow (except Fig 4) and do a good job at summarizing the method. 

Weaknesses
The appendix is too long.
The limitations are not addressed in the main paper, although it's present in the appendix.
I was unable to follow the example in Section 5.4 Uncertainty Visualization as the figure was unclear due to lack of legends and marginally inappropriate choice of colors.
questions: questions: 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.
2. Please move the limitations section in the main paper instead of being in the appendix.
3. I didn't understand how the 'lightweight uncertainty prediction model' in section 4.4 is trained, perhaps more details would be great.
4. How is the uncertainty threshold epsilon selected? 
5. How is the hyperparameter c (for frequent switching between certain and uncertain segments) selected?
6. Perhaps, section 5.4 Uncertainty Visualization could be expanded to show more instances of uncertainty on ego vehicle behavior.
7. How is the hyperparameter n (upper percentile of target return for planning) selected?
limitations: limitations: 3
review_summary: review_summary: It's a well written paper with very minor revisions needed for rebuttal.
title: title: good paper
paper_summary: paper_summary: 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.
main_review: main_review: ### Strengths:
- **Innovative Approach**: Introducing uncertainty estimation directly into decision transformers is a novel approach that addresses a critical challenge in RL.
- **Practical Relevance**: The focus on autonomous driving scenarios enhances the practical relevance and potential impact of the research.
- **Experimental Rigor**: Comprehensive experiments across diverse driving scenarios validate the effectiveness and robustness of UNREST.
- **Clear Presentation**: The paper is generally well-structured and clearly presents its methodology and findings, facilitating understanding and replication.

### Weaknesses:
- **Complexity of Uncertainty Estimation**: The mathematical formulation of uncertainty estimation using conditional mutual information could be challenging for readers not familiar with information theory.
- **Limited Comparison**: While UNREST demonstrates superior performance, comparison with a wider range of state-of-the-art methods or alternative transformer adaptations could further strengthen the paper.
- **Generalization**: The extent to which UNREST can generalize to other domains beyond driving scenarios remains an open question.
questions: questions: I'm not an expert for  self-driving research. I try to give some suggestions:  

1. 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.  

2. 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.  

[1] Transformer in Transformer as Backbone for Deep Reinforcement Learning. Arxiv 2022.   
[2] PDiT: Interleaving Perception and Decision-making Transformers for Deep Rinforcement Learning. AAMAS 2024.   
[3] Sequential Asynchronous Action Coordination in Multi-Agent Systems: A Stackelberg Decision Transformer Approach. ICML 2024.   
[4] CoSLight: Co-optimizing Collaborator Selection and Decision-making to Enhance Traffic Signal Control. KDD 2024.   
[5] X-Light: Cross-City Traffic Signal Control Using Transformer on Transformer as Meta Multi-Agent Reinforcement Learner. IJCAI 2024.
limitations: limitations: 1
review_summary: review_summary: 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.