Paper: Uncertainty-Aware Decision Transformer for Stochastic Driving Environments (ICLR-2024)

Reviewer comment (Writing):
Where does the uncertainty in autonomous driving come from? My understanding of uncertainty is that with the same (s, a) pairs, the environment returns different rewards. Figure 1 tries to provide an example but a lot of information is missing. For example, what is the purpose of the traffic light? What does the x-axis label “return” mean for task 1 and task 2? Do the authors mean reward to go?

Author response:
Thank you for highlighting the problem, which is crucial for understanding our paper. Specifically, as discussed in the introduction, **uncertainty in the autonomous driving environment may arise from the hidden states of other vehicles and partial observability.** In other words, the transition $(s,a\rightarrow s')$ along with the reward $r$ is inpredictable (different possible outcome $s'$ and $r$) because the surrounding vehicles can behave stochastically and some information is not observable. 

Based on your suggestion, we have made revisions to Figure 1 and its corresponding description. Specifically, we remove the irrelevant traffic lights from the figure. Besides, we clarify that the term 'return' expresses exactly the same meaning as 'reward-to-go' in reinforcement learning, about the cumulative rewards in the episode. Regarding Figure 1(a), it represents a driving scenario composed of a sequence of driving actions, starting with going straight and followed by a right turn. In this scenario, the uncertainty in the environment is relatively low during the straight driving process, while the variance in the distribution of rewards increases during the right turn due to the uncertain behavior of other vehicles in the environment. Therefore, we propose to condition solely on the truncated return from Task I (going straight) rather than Task I & II to mitigate the influence of environmental stochasticity (less stochastic timesteps, lower return variance) for DT training, which still preserves rewards over sufficient timesteps for the optimization of current actions.

Reviewer comment (Experiments):
The motivating example of the uncertainty in the driving scenario seems to involve the tradeoff between reward (reaching point) and cost (collision). In my opinion, tuning the weights in reward functions or some other parameters could dramatically influence the “uncertainty” investigated in this paper. Basically, it will lead to either aggressive or conservative driving policy. However, cost or constraint is only discussed in the appendix. Therefore, I doubt that the proposed transformer model can correctly estimate the uncertainty.

Author response:
Thank you for pointing out the problem of our reward design. I believe it is necessary to clarify that **all the reinforcement learning algorithms we employ do not exhibit the so-called trade-offs. Instead, their primary objective is to purely seek strategies that maximize rewards.** 

As you said, we have discussed the reward and cost compositions of our dataset in Appendix D. **However, for the RL algorithms, the cost is integrated in the reward function as penalization for dangerous behaviors.** Otherwise, the RL agent may frequently conduct aggressive actions that may raise safety issues. So for the example driving scenario in Figure 1(a), the ego-vehicle will learn to yield to surrounding vehicles, thereby avoiding significant negative rewards resulting from collisions.

Furthermore, as you mentioned, adjusting the weights in the reward function will inevitably lead to variations in the estimation of uncertainties. However, **we do not consider this to be a problem because the essence of the effectiveness of RL lies in the design of a well-defined reward function.** While adjusting the weights in the reward function may harm uncertainty estimation, the more significant concern is that arbitrary adjustments can lead to suboptimal policy learning by the RL algorithm itself. For our approach, since we are identifying uncertainty in returns (rewards), when the reward function can serve as an effective conditioned target for DT policy learning, UNREST will also be capable of effective segmenting sequences for DTs.
  | Planner | CostWeight=0.1  | CostWeight=5.0 | CostWeight=10.0 | CostWeight=50.0  |
  | ------- | --------------------------- | --------------------------- | --------------------------- | --------------------------- |
  | DT  | $48.2\pm1.6$ ($40.4\pm2.4$) | $54.4\pm2.8$ ($47.4\pm1.7$) | $55.2\pm2.0$ ($47.6\pm1.2$) | $52.6\pm3.0$ ($46.5\pm2.3$) |
  | UNREST | $48.4\pm2.4$ ($41.3\pm1.5$) | $63.0\pm2.2$ ($61.9\pm3.0$) | $63.5\pm3.2$ ($62.9\pm4.0$) | $62.8\pm2.5$ ($61.5\pm1.9$) | 

The above table records our new testing driving scores on train and new (in brackets) scenarios w.r.t. different cost weights in reward functions. As can be seen from the table, both excessively high and excessively low weights can result in the inability of the DT to train effective policies. However, the introduction of uncertainty estimation in UNREST consistently brings improvements compared to DT across different weight settings.

Reviewer comment (Novelty):
It seems that one important problem investigated in this paper is “Uncertainty Accumulation”, which is reasonable since the uncertainty increases as the trajectories accumulate more steps. However, why do we care about the accumulation of uncertainty in driving tasks? A normal driving algorithm just needs to make decisions based on current uncertainty and historical information. I guess one reason could be the DT framework uses reward-to-go as a condition. If this is the case, this paper may not be applicable to other frameworks.

Author response:
Thank you for pointing out the problem. Exactly as you said, the uncertainty accumulation is a normal characteristic within the context of autonomous driving environments. **However, our attention to this property is not solely due to the DT framework.** Typically, it's natural to consider uncertanity accumulation for DTs, since the conditioned reward-to-go includes the uncertainties at each time step from the current state to the end of the sequence. UNREST, however, addresses and reduces this conditioning uncertainty by segmenting the sequences, enabling more effective policy learning in stochastic environments.

**But for other RL algorithms, uncertainty accumulation is also a significant factor to take into consideration since the monte-carlo/trajectory statistics are fundamental factors in RL**, which utilize simulated trajectory return to update the value functions, or directly select a better action. In these simulated trajectories, the cumulated uncertainty can lead to large variance during backup/forward searching. **Our uncertainty estimation is applicable in these frameworks by accurately identifying uncertain timesteps and branching out to perform additional trajectory rollouts at those points.** This process helps improve the accuracy of the monte-carlo estimation results，which suggests the general use cases of our uncertainty estimation. For example, [1, 2] propose to learn transition models for monte-carlo simulation. Due to the significant bias associated with single-step rollouts, they have all opted for multi-step beam search to find the most effective and safe actions. By combining our uncertainty estimation method, we can effectively identify specific uncertain states that require branching and rollout. This can greatly helps reduce the number of simulations and improves planning efficiency.

Reviewer comment (Presentation):
A lot of important information is missing in figures, experiments, and conclusions.

Author response:
Thank you for pointing out our oversights in writing. We have made necessary modifications to the corresponding content in our updated version of paper (highlight in green and blue) and here we provide individual responses to all of your questions.

1. As discussed in R1, we have made revisions to Figure 1 and its corresponding description. Regarding Figure 1(a), it represents a driving scenario composed of a sequence of driving actions, starting with going straight (low uncertainty) and followed by a right turn (high uncertainty). We also remove the irrelevant traffic light from the figure and clarify the task and return distribution contents in the figure. For Figure 1(b) and Figure 1(c), we have modified their x-axis labels to provide greater precision, specifically as, 'Normalized Return (Ground Truth)' and 'Timestep', respectively. Typically, the x-axis in Figure (b) represents the ground truth return in original trajectories. More detailed explanations are expressed in the main context of the newly updated paper.
2. Yes, our training data is indeed collected from CARLA. And as we have discussed in R1, the uncertainty can arise from the stochastic behavior of surrounding vehicles and the partial observability. Two respective examples can be referred in Figure 4 in the paper. Typically, in Figure 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. In Figure 4(b), while approaching the green light, the partial observability about the light state change causes the uncertainty to rise quickly. After the vehicle has moved away from the traffic light, the uncertainty immediately drops back below the threshold.  
3. The meaning of $h$ is expressed as 'return-span token' in Figure 2, which is further explained below Equation (6). Specifically, the return-span $h_t$ denotes the timestep count involved in the truncated return $R^h_t$. As we discussed in Section 4.3, for states in certain parts, $h$ is set to the number of timesteps to next segmentation position to include the maximum duration that do not contain any uncertain steps; for states in uncertain parts, $h$ is set as empty for meaningless conditions, to conduct behavior cloning. So in Figure 2, we start with $h=1$, which is the final timestep in a certain part. Afterwards, the state is identified as uncertain and an uncertain part (orange background) begins with empty $h$. Finally, another certain part with 80 timesteps starts, with $h=80$, then sequentially decreases to $h=79, 78, ...$ until another uncertain timestep.
4. In Figure 3 the blud rectangles denote surrounding vehicles and the ego-vehicle is the white rectangle. Figure 3(a) illustrates a scenario where the presence of the surrounding vehicles in the adjacent lane prevents lane changing. In this scenario, the TT planner tends to make overly aggressive maneuvers, resulting in collisions with the vehicles in the adjacent lane. On the other hand, the UNREST planner adopts a cautious planning approach and takes timely actions to stop before any collision occurs. 

Hope these clarifications could address your concerns and facilitate a more comprehensive understanding of the insights and contributions of this paper.

Reviewer comment (Presentation):
After checking Property 1 and Property 2, I still cannot fully understand Figure 1. For example, what does “normalized groud truth” mean (I assume “groud” is a typo of “ground”)? I highly suggest the authors provide a simplified and intuitive example in the first figure without adding any distracting information. The detailed explanation can be put in later sections.

Author response:
Thank you for pointing out our oversights in writing. We have made necessary modifications to the corresponding content in our updated version of paper (highlight in green and blue) and here we provide individual responses to all of your questions.

1. As discussed in R1, we have made revisions to Figure 1 and its corresponding description. Regarding Figure 1(a), it represents a driving scenario composed of a sequence of driving actions, starting with going straight (low uncertainty) and followed by a right turn (high uncertainty). We also remove the irrelevant traffic light from the figure and clarify the task and return distribution contents in the figure. For Figure 1(b) and Figure 1(c), we have modified their x-axis labels to provide greater precision, specifically as, 'Normalized Return (Ground Truth)' and 'Timestep', respectively. Typically, the x-axis in Figure (b) represents the ground truth return in original trajectories. More detailed explanations are expressed in the main context of the newly updated paper.
2. Yes, our training data is indeed collected from CARLA. And as we have discussed in R1, the uncertainty can arise from the stochastic behavior of surrounding vehicles and the partial observability. Two respective examples can be referred in Figure 4 in the paper. Typically, in Figure 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. In Figure 4(b), while approaching the green light, the partial observability about the light state change causes the uncertainty to rise quickly. After the vehicle has moved away from the traffic light, the uncertainty immediately drops back below the threshold.  
3. The meaning of $h$ is expressed as 'return-span token' in Figure 2, which is further explained below Equation (6). Specifically, the return-span $h_t$ denotes the timestep count involved in the truncated return $R^h_t$. As we discussed in Section 4.3, for states in certain parts, $h$ is set to the number of timesteps to next segmentation position to include the maximum duration that do not contain any uncertain steps; for states in uncertain parts, $h$ is set as empty for meaningless conditions, to conduct behavior cloning. So in Figure 2, we start with $h=1$, which is the final timestep in a certain part. Afterwards, the state is identified as uncertain and an uncertain part (orange background) begins with empty $h$. Finally, another certain part with 80 timesteps starts, with $h=80$, then sequentially decreases to $h=79, 78, ...$ until another uncertain timestep.
4. In Figure 3 the blud rectangles denote surrounding vehicles and the ego-vehicle is the white rectangle. Figure 3(a) illustrates a scenario where the presence of the surrounding vehicles in the adjacent lane prevents lane changing. In this scenario, the TT planner tends to make overly aggressive maneuvers, resulting in collisions with the vehicles in the adjacent lane. On the other hand, the UNREST planner adopts a cautious planning approach and takes timely actions to stop before any collision occurs. 

Hope these clarifications could address your concerns and facilitate a more comprehensive understanding of the insights and contributions of this paper.

Reviewer comment (Presentation):
Missing information in Figure 2. What does h represent? Why does h start from 1 but end with 80, 79?

Author response:
Thank you for pointing out our oversights in writing. We have made necessary modifications to the corresponding content in our updated version of paper (highlight in green and blue) and here we provide individual responses to all of your questions.

1. As discussed in R1, we have made revisions to Figure 1 and its corresponding description. Regarding Figure 1(a), it represents a driving scenario composed of a sequence of driving actions, starting with going straight (low uncertainty) and followed by a right turn (high uncertainty). We also remove the irrelevant traffic light from the figure and clarify the task and return distribution contents in the figure. For Figure 1(b) and Figure 1(c), we have modified their x-axis labels to provide greater precision, specifically as, 'Normalized Return (Ground Truth)' and 'Timestep', respectively. Typically, the x-axis in Figure (b) represents the ground truth return in original trajectories. More detailed explanations are expressed in the main context of the newly updated paper.
2. Yes, our training data is indeed collected from CARLA. And as we have discussed in R1, the uncertainty can arise from the stochastic behavior of surrounding vehicles and the partial observability. Two respective examples can be referred in Figure 4 in the paper. Typically, in Figure 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. In Figure 4(b), while approaching the green light, the partial observability about the light state change causes the uncertainty to rise quickly. After the vehicle has moved away from the traffic light, the uncertainty immediately drops back below the threshold.  
3. The meaning of $h$ is expressed as 'return-span token' in Figure 2, which is further explained below Equation (6). Specifically, the return-span $h_t$ denotes the timestep count involved in the truncated return $R^h_t$. As we discussed in Section 4.3, for states in certain parts, $h$ is set to the number of timesteps to next segmentation position to include the maximum duration that do not contain any uncertain steps; for states in uncertain parts, $h$ is set as empty for meaningless conditions, to conduct behavior cloning. So in Figure 2, we start with $h=1$, which is the final timestep in a certain part. Afterwards, the state is identified as uncertain and an uncertain part (orange background) begins with empty $h$. Finally, another certain part with 80 timesteps starts, with $h=80$, then sequentially decreases to $h=79, 78, ...$ until another uncertain timestep.
4. In Figure 3 the blud rectangles denote surrounding vehicles and the ego-vehicle is the white rectangle. Figure 3(a) illustrates a scenario where the presence of the surrounding vehicles in the adjacent lane prevents lane changing. In this scenario, the TT planner tends to make overly aggressive maneuvers, resulting in collisions with the vehicles in the adjacent lane. On the other hand, the UNREST planner adopts a cautious planning approach and takes timely actions to stop before any collision occurs. 

Hope these clarifications could address your concerns and facilitate a more comprehensive understanding of the insights and contributions of this paper.

Reviewer comment (Presentation):
Information missing in Figure 3. What does the light blue rectangle mean? Which one is the ego vehicle? What does failure mean in Figure 3(a)? It seems that the lane change is not available as there is another car in the target lane. What’s the difference between Figure 3(a) and Figure 3(b)?

Author response:
Thank you for pointing out our oversights in writing. We have made necessary modifications to the corresponding content in our updated version of paper (highlight in green and blue) and here we provide individual responses to all of your questions.

1. As discussed in R1, we have made revisions to Figure 1 and its corresponding description. Regarding Figure 1(a), it represents a driving scenario composed of a sequence of driving actions, starting with going straight (low uncertainty) and followed by a right turn (high uncertainty). We also remove the irrelevant traffic light from the figure and clarify the task and return distribution contents in the figure. For Figure 1(b) and Figure 1(c), we have modified their x-axis labels to provide greater precision, specifically as, 'Normalized Return (Ground Truth)' and 'Timestep', respectively. Typically, the x-axis in Figure (b) represents the ground truth return in original trajectories. More detailed explanations are expressed in the main context of the newly updated paper.
2. Yes, our training data is indeed collected from CARLA. And as we have discussed in R1, the uncertainty can arise from the stochastic behavior of surrounding vehicles and the partial observability. Two respective examples can be referred in Figure 4 in the paper. Typically, in Figure 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. In Figure 4(b), while approaching the green light, the partial observability about the light state change causes the uncertainty to rise quickly. After the vehicle has moved away from the traffic light, the uncertainty immediately drops back below the threshold.  
3. The meaning of $h$ is expressed as 'return-span token' in Figure 2, which is further explained below Equation (6). Specifically, the return-span $h_t$ denotes the timestep count involved in the truncated return $R^h_t$. As we discussed in Section 4.3, for states in certain parts, $h$ is set to the number of timesteps to next segmentation position to include the maximum duration that do not contain any uncertain steps; for states in uncertain parts, $h$ is set as empty for meaningless conditions, to conduct behavior cloning. So in Figure 2, we start with $h=1$, which is the final timestep in a certain part. Afterwards, the state is identified as uncertain and an uncertain part (orange background) begins with empty $h$. Finally, another certain part with 80 timesteps starts, with $h=80$, then sequentially decreases to $h=79, 78, ...$ until another uncertain timestep.
4. In Figure 3 the blud rectangles denote surrounding vehicles and the ego-vehicle is the white rectangle. Figure 3(a) illustrates a scenario where the presence of the surrounding vehicles in the adjacent lane prevents lane changing. In this scenario, the TT planner tends to make overly aggressive maneuvers, resulting in collisions with the vehicles in the adjacent lane. On the other hand, the UNREST planner adopts a cautious planning approach and takes timely actions to stop before any collision occurs. 

Hope these clarifications could address your concerns and facilitate a more comprehensive understanding of the insights and contributions of this paper.

Reviewer comment (Writing):
The paper is not well written.

Author response:
Thank you for pointing out the issues in our writing. We have made necessary modifications to the corresponding content in our updated version of paper (highlight in green and blue). Moving forward, here we provide individual responses to all of your questions.

Reviewer comment (Writing):
A lot of notations are used before definition. For example, section 4.2 \tau^{ret}_{<t}.

Author response:
1. Regarding the issue of notation, we can ensure that all symbols used are clearly defined in the paper. Like the example $\boldsymbol{\tau_{<t}^{ret}}$ you raised in Section 4.2, we have defined in Section 3.2 that $\mathbf{x_{<t}}$ denotes tokens from step 1 to $(t-1)$, and in Section 4.2 we define $\boldsymbol{\tau^{\text{ret}}}=\{(s_t,a_t)\}_{t=1}^T$ to be all state-action pairs involved in the driving trajectory. Besides, we have clarified Task I, Task II, and Normalized Ground Truth in Figure 1. Specifically, Figure 1(a) is composed of two driving tasks: go straight then turn right, and the x-axis in Figure 1(b) denotes the ground truth return in original driving trajectories.

Reviewer comment (Writing):
Some "terms" and "notation" are used without definition or even description, which makes it very difficult to understand. For example, Figure 1(a) Task I and Task II, and (b) normalized ground Truth.

Author response:
1. Regarding the issue of notation, we can ensure that all symbols used are clearly defined in the paper. Like the example $\boldsymbol{\tau_{<t}^{ret}}$ you raised in Section 4.2, we have defined in Section 3.2 that $\mathbf{x_{<t}}$ denotes tokens from step 1 to $(t-1)$, and in Section 4.2 we define $\boldsymbol{\tau^{\text{ret}}}=\{(s_t,a_t)\}_{t=1}^T$ to be all state-action pairs involved in the driving trajectory. Besides, we have clarified Task I, Task II, and Normalized Ground Truth in Figure 1. Specifically, Figure 1(a) is composed of two driving tasks: go straight then turn right, and the x-axis in Figure 1(b) denotes the ground truth return in original driving trajectories.

Reviewer comment (Writing):
If DT is a fundamental base of the method, having a brief introduction is more friendly to new readers.

Author response:
3. We acknowledge your point that since DT is a crucial foundation of this paper and warrants a separate section to provide an overview of its fundamental concepts and contents. Specifically, we have written in Section 3.2 of the original paper a description of the training and testing details for DTs, which you may miss in the first time of reading.

Reviewer comment (Writing):
Does the "Transformer" term in section 4.2 mean two independent networks for x_{a_t} and x_{s_t}? Please give more concrete/clear description using notations flow of your method in general instead of long words description.

Author response:
4. We apologize for any misunderstanding caused by the informal expression used in our formula. Typically, we use two separate transformers to get $\tilde{x_{s_t}}$ and $\tilde{x_{a_t}}$. And to express the formula more concisely, we have used the notation $\mathcal T$ to replace the long word 'Transformer' in Equation (3) and (7).

Reviewer comment (Writing):
The exposition, particularly the first half of page 2, is hard to understand. For example, it is not explained what a highly representative environment model is and why it might not be learnable. Although this is related work, the reader is completely left in the dark about why the presented approach avoids this.

Author response:
Thank you for pointing out the issue with our unclear statement. We have modified the introduction in the newly updated version (highlighted in green and blue) and here we answer your problems respectively.
1. About the 'highly representative environment model', it usually means the state transition function $P(s'|s,a):\mathcal{S}\times\mathcal{A}\times\mathcal{S}\rightarrow[0,1]$ in reinforcement learning. Previous methods [2, 3, 4] attempt to fit the state transition model to deal with transition uncertainty. However, 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. Our method bypasses the challenging high-dimensional state transition function learning by computing uncertainty through conditional mutual information $u_t=D_{\text{KL}}\big(p_{\varphi_a}(R_{t}|\boldsymbol{\tau_{<t}^{\text{ret}}}),p_{\varphi_s}(R_{t}|\boldsymbol{\tau_{<t}^{\text{ret}}},s_{t})\big)$, where the return distributions can be effectively learned with fewer resources. 
2. Regarding Figure 1(a) and its corresponding explanation, it represents a driving scenario composed of two tasks (now denoted as Task I & II), starting with going straight (Task I) and followed by a right turn (Task II). In this scenario, the uncertainty in the environment is relatively low during the straight driving process, while the variance in the distribution of rewards increases substantially considering Task I & II as a whole due to the uncertain behavior of other vehicles in the environment during the right turn. Therefore, we propose to condition solely on the truncated return from Task I (going straight) rather than Task I & II to mitigate the influence of environmental stochasticity (less stochastic timesteps, lower return variance) for DT training, which still preserves rewards over sufficient timesteps for the optimization of current actions.

Hope these modifications in our presentation can help you better understand the main insights and contributions of our work.

Reviewer comment (Writing):
In Task I+II, the "+" is initially presented in a manner that does not make it clear that it means sequential composition (e.g., perform Task I and then Task II). This only makes sense once finishing the page and the segmentation, but the exposition could be improved significantly.

Author response:
Thank you for pointing out the issue with our unclear statement. We have modified the introduction in the newly updated version (highlighted in green and blue) and here we answer your problems respectively.
1. About the 'highly representative environment model', it usually means the state transition function $P(s'|s,a):\mathcal{S}\times\mathcal{A}\times\mathcal{S}\rightarrow[0,1]$ in reinforcement learning. Previous methods [2, 3, 4] attempt to fit the state transition model to deal with transition uncertainty. However, 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. Our method bypasses the challenging high-dimensional state transition function learning by computing uncertainty through conditional mutual information $u_t=D_{\text{KL}}\big(p_{\varphi_a}(R_{t}|\boldsymbol{\tau_{<t}^{\text{ret}}}),p_{\varphi_s}(R_{t}|\boldsymbol{\tau_{<t}^{\text{ret}}},s_{t})\big)$, where the return distributions can be effectively learned with fewer resources. 
2. Regarding Figure 1(a) and its corresponding explanation, it represents a driving scenario composed of two tasks (now denoted as Task I & II), starting with going straight (Task I) and followed by a right turn (Task II). In this scenario, the uncertainty in the environment is relatively low during the straight driving process, while the variance in the distribution of rewards increases substantially considering Task I & II as a whole due to the uncertain behavior of other vehicles in the environment during the right turn. Therefore, we propose to condition solely on the truncated return from Task I (going straight) rather than Task I & II to mitigate the influence of environmental stochasticity (less stochastic timesteps, lower return variance) for DT training, which still preserves rewards over sufficient timesteps for the optimization of current actions.

Hope these modifications in our presentation can help you better understand the main insights and contributions of our work.

Reviewer comment (Theory):
While there is ample empirical evidence for this approach's success, there is no clear theoretical framework to connect it to. There is a tentative connection to work on causal vs. non-causal entropy regularization and its similar effect on creating over-optimistic agents, but a deeper theoretical basis would be appreciated.

Author response:
We thank the reviewers for pointing out our oversight in the paper. In fact we have detailedly discuss UNREST's theoretical generalizability in our original paper, as you can find in Proposition 1. Due to the page limit, the proposition in the main context is explained in an informal manner, with the formal statement and proof left to Appendix B. Specifically, the proposition is inspired by [1] and 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 relevant content has been highlighted in the new version of our paper, where you can find further details for discussion. Besides, I find the causal/non-causal entropy regularization mentioned by you to be an exceptionally intriguing aspect. It is possible that we can explore this point further to explain how UNREST addresses the issue of overly-optimistic behavior of agents.

Reviewer comment (Experiments):
The experiments and setup appear to be for the fully observable context, except for the other driver's implicit latent states. However, in the automotive setting, a significant amount of uncertainty comes from the partial observability of the environment (e.g., the exact location of the crossing region in an intersection, the drivable area and its boundaries, the location of obscured pedestrians and bikes). It is unclear how the presented approach could be adapted to such settings.

Author response:
Thank you for highlighting this interesting question regarding our experimental design. As you aptly pointed out, the uncertainty in the driving environment can arise from both the latent states of surrounding vehicles and partial observability. In this response, I would like to elucidate the capability of UNREST in effectively handling the uncertainty introduced by partial observability.

**Theoretically,** we compute the uncertainty by the conditional mutual information $u_t=D_{\text{KL}}\big(p_{\varphi_a}(R_{t}|\boldsymbol{\tau_{<t}^{\text{ret}}},p_{\varphi_s}(R_{t}|\boldsymbol{\tau_{<t}^{\text{ret}}},s_{t})\big)$. So regarding the uncertainty introduced by partial observability, such as pedestrians who were obstructed and not visible in the previous step but are now visible, can result in a significant change in return distribution $p_{\varphi_s}(R_{t}|\boldsymbol{\tau_{<t}^{\text{ret}}},s_{t})$ from $p_{\varphi_a}(R_{t}|\boldsymbol{\tau_{<t}^{\text{ret}}})$ because of the newly captured information.

**From the experimental perspective,** we have discussed in Appendix D in the paper about our state space, which is similar to that in [2]. Specifically, the ego vehicle can only obtain distances of nearby vehicles within its limited field of view and perceive the boundaries of a limited portion of the target route, which is exactly consistent with your proposed partially observable setting. Building upon this, our experimental results further demonstrate the superior ability of UNREST compared to previous baselines in addressing the uncertainty brought by partial observability. Besides, we also visualize in Figure 4.(b) about the uncertainty curve when the ego-vehicle is passing the green traffic light. As you can see, while approaching the green light, the partial observability about the change of light state causes the uncertainty to rise quickly. But after the vehicle has moved away from the traffic light, the uncertainty immediately drops back below the threshold. This interpretable result also validate the effectiveness of our uncertainty measure.

Reviewer comment (Experiments):
Evaluated only on a single environment.

Author response:
Thank you for pointing out the problem of our experiments. The reason why we only study the driving environment is that, though as a general offline RL algorithm that can be applied to any RL environments, **UNREST is specially designed to address two major challenges in autonomous driving: long-horizon planning and environmental stochasticity.** To make the segmentation strategy proposed by UNREST effective, two essential factors are required：1）'uncertainty accumulation' property that ensures the truncated return is less affected by the environments than the global return, to provide effective supervison for policy learning; 2) 'temporal locality' property that makes optimizing future returns over a sufficiently long horizon similar to the optimization of global return. We have identified these two properties in the context of autonomous driving, but they may not necessarily apply in other environments. **We have also tested UNREST on the conventional D4RL benchmark (approximately deterministic environments, results in Appendix F.5)**, where it just attains similar results to vanilla DTs. Similarly, SPLT [3] is also only tested on the CARLA simulator. Therefore, we believe that the autonomous driving environment is the most suitable environment to test our proposed algorithm.
  | Environment | DT    | TT      | SPLT  | ESPER  | UNREST   |
  | ------- | --------------------------- | --------------------------- | --------------------------- | --------------------------- | ----------------------------- |
  | HalfCheetah-Stochastic  | $75.8\pm1.2$ | $83.7\pm2.8$ | $88.6\pm2.2$ | $77.4\pm5.6$ | $\mathbf{90.5\pm3.2}$ |
  | Hopper-Stochastic     | $89.4\pm3.3$ | $93.4\pm3.6$ | $94.8\pm2.7$ | $92.5\pm6.9$ | $\mathbf{97.4\pm4.4}$ | 

In the above table, **we test UNREST in the customized stochastic D4RL environments from DoC [4],** which add time-correlated Gaussian noise to actions. The results are also added in Appendix F.5, where UNREST achieves highest return among the models.

Reviewer comment (Novelty):
The idea of using truncated returns is not entirely novel though and has been used in the upside-down RL methods [1,2].

Author response:
We appreciate reviewer's attention to the novelty of our work. Here we need to clarify one significant point that the key novelty of UNREST lies in **the low-resource occupancy uncertainty estimation method, along with the proposed sequence segmentation and designs of conditions based on it.** 
- **Novelty 1:** Different from previous works [2, 3, 4] that attempt to fit high-dimensional state transition functions and introduce complex VAE or adversarial training, **we introduce for the first time the utilization of conditional mutual information to estimate uncertainty**. As discussed in the paper, the complex interactions, partial observability, and long trajectories challenges training of previous works, where our proposed mutual information can be effectively learned with fewer resources, as validated in experiments.
- **Novelty 2:** Back to the fundamental objective of this paper, we propose **a novel sequence segmentation and designs of conditions for customizing DTs in stochastic driving environments**. Specifically, we segment sequences into certain and uncertain parts according to our uncertainty measurements, and learn to behave aggressively and cautiously within them, respectively. 
- **Novelty 3:** We also **summarize two essential properties of driving environments**: 'temporal locality' and 'uncertainty accumulation' properties, based on which we **theoretically prove** that replacing conditioned returns in original DT with truncated returns that only consider returns to the next segmented positions can help UNREST generalize to attain high rewards. 

Therefore, while it is true that truncated return has been employed in previous works [1], in this paper based on our discovered properties, we introduce a novel specific segmentation method based on our proposed uncertainty estimation techiniques and provide corresponding theoretical justifications. This represents a distinct contribution that was absent in prior literature.

Hope this clarification will assist you in better understanding the key insights and novelty of this paper.

Reviewer comment (Experiments):
The algorithm assumes the data have to come from an expert if I understand correctly (especially for the logic to work in uncertain regions). The authors use expert data (from the autopilot) in the training process. Isn't this a drawback when compared to existing offline RL baselines which can also leverage suboptimal trajectories/data?

Author response:
Thank you for highlighting this crucial issue regarding the scope of applicability of UNREST. In fact, the decision to resort to direct behavior cloning (BC) in uncertain parts within UNREST can be attributed to the compromise made due to the unreliability of returns, which hinders the provision of effective supervision. We opted for this approach for two reasons: (1) Firstly, **BC is a lightweight and easily implementable algorithm** that simply requires setting the original conditioned return as a dummy token. (2) Secondly, **expert data collection is relatively easy in the context of autonomous driving environments**. However, it should be noted that the AutoPilot provided by CARLA has multiple fail cases and cannot be considered an optimal expert.

In order to propose a strategy that can be employed in uncertain regions when there exists suboptimal trajectories, we propose a simple yet effective UNREST variant that **additionally fit a state value function**. The remaining logic is consistent with MARWIL [5], where we utilize neural network estimation of advantages to reweight actions in the BC objective. In this way, UNREST can learn to identify and choose to imitate actions in optimal trajectories. The results are recorded in the following table. 
  | Planner | Driving Score $\uparrow$    | Success Rate$\uparrow$      | Route Completion$\uparrow$  | Infraction Score$\uparrow$  | Normalized Reward$\uparrow$   |
  | ------- | --------------------------- | --------------------------- | --------------------------- | --------------------------- | ----------------------------- |
  | Reweighted BC  | $64.2\pm2.4$ ($63.8\pm2.3$) | $57.5\pm3.3$ ($59.3\pm4.7$) | $85.5\pm2.7$ ($91.0\pm5.0$) | $71.4\pm2.5$ ($65.1\pm3.7$) | $0.68\pm0.03$ ($0.67\pm0.02$) |
  | Original Model     | $63.5\pm3.2$ ($62.9\pm4.0$) | $54.5\pm7.0$ ($57.5\pm5.4$) | $83.8\pm3.1$ ($90.0\pm6.0$) | $70.2\pm2.8$ ($62.9\pm3.8$) | $0.64\pm0.04$ ($0.65\pm0.03$) | 

Since the offline data collected by AutoPilot can be suboptimal, we find reweighted BC 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 still recommend the original UNREST as the base planner.

Reviewer comment (Experiments):
Would this method scale to other stochastic environments like the ones used in baseline algorithms [3] and [4]?

Author response:
Thank you for pointing out the problem of our experiments. The reason why we only study the driving environment is that, though as a general offline RL algorithm that can be applied to any RL environments, **UNREST is specially designed to address two major challenges in autonomous driving: long-horizon planning and environmental stochasticity.** To make the segmentation strategy proposed by UNREST effective, two essential factors are required：1）'uncertainty accumulation' property that ensures the truncated return is less affected by the environments than the global return, to provide effective supervison for policy learning; 2) 'temporal locality' property that makes optimizing future returns over a sufficiently long horizon similar to the optimization of global return. We have identified these two properties in the context of autonomous driving, but they may not necessarily apply in other environments. **We have also tested UNREST on the conventional D4RL benchmark (approximately deterministic environments, results in Appendix F.5)**, where it just attains similar results to vanilla DTs. Similarly, SPLT [3] is also only tested on the CARLA simulator. Therefore, we believe that the autonomous driving environment is the most suitable environment to test our proposed algorithm.
  | Environment | DT    | TT      | SPLT  | ESPER  | UNREST   |
  | ------- | --------------------------- | --------------------------- | --------------------------- | --------------------------- | ----------------------------- |
  | HalfCheetah-Stochastic  | $75.8\pm1.2$ | $83.7\pm2.8$ | $88.6\pm2.2$ | $77.4\pm5.6$ | $\mathbf{90.5\pm3.2}$ |
  | Hopper-Stochastic     | $89.4\pm3.3$ | $93.4\pm3.6$ | $94.8\pm2.7$ | $92.5\pm6.9$ | $\mathbf{97.4\pm4.4}$ | 

In the above table, **we test UNREST in the customized stochastic D4RL environments from DoC [4],** which add time-correlated Gaussian noise to actions. The results are also added in Appendix F.5, where UNREST achieves highest return among the models.