Addressing Optimism Bias in Sequence Modeling for Reinforcement Learning

Two major works that have explored reformulating the offline RL problem as a single sequence modelling problem are Decision Transformer (Chen et al., 2021b) and Trajectory Transformer (Janner et al., 2021). In stochastic environments like autonomous driving, both approaches can act optimistically because they fail to account for the uncontrollable factors in the environment that can affect the ability to achieve a specific return.

Decision Transformer is a return-conditioned model-free method that learns a Transformer-based policy that takes in the historical states and actions and a target return, and outputs the action that is most likely to lead to a trajectory sequence that achieves the target return. In stochastic environments, the difficulties come from picking a suitable target return without strong prior knowledge of the testing domain. This is especially problematic because the distribution of possible returns is heavily dependent on the stochastic transitions of the environment. Thus, setting a large target return that is not always feasible can lead to overly aggressive and optimistic behavior, while setting any lower target return could lead to sub-optimal behavior in situations where the environment does unroll favorably.

Trajectory Transformer is a model-based method that trains a Transformer-based trajectory model that can hallucinate potential trajectories in the environment. In stochastic environments, the main concern is how you use this model to properly search for a suitable next action or trajectory-sequence given the uncontrollable stochastic transitions. Naively deploying NLP-style beam search as proposed in the original Trajectory Transformer paper (Janner et al., 2021) without accounting for these uncontrollable stochastic transitions will be biased to explore and pick trajectories where the environment just so happens to unroll favorably. This will similarly lead to optimistic behavior, which can be dangerous in safety-critical environments. To address this issue, a different searching procedure needs to be used that can reason about the states and actions separately. One possible avenue for doing this step-wise would be to explore the different variants of MCTS described in Ozair et al. (2021), but that would lead to an approach that scales exponentially with the search horizon, while our approach only scales linearly with the search horizon.

In order to illuminate the optimism bias in Decision Transformer and Trajectory Transformer, consider a simple discrete and stochastic MDP where there is only 5 states $\{s_{0},s_{11},s_{12},s_{21},s_{22}\}\in\mathcal{S}$ and 2 actions $\{a_{1},a_{2}\}\in\mathcal{A}$. The agent always starts in $s_{0}$. Then, if the agent takes $a_{1}$ it will stochastically transition from $s_{0}\rightarrow s_{11}$ or $s_{0}\rightarrow s_{12}$ with equal probability. Similarly, if the agent takes $a_{2}$ it will stochastically transition from $s_{0}\rightarrow s_{21}$ or $s_{0}\rightarrow s_{22}$ with equal probability. The reward at each state is: 0 for $s_{0}$, 10 for $s_{11}$, -10 for $s_{12}$, 6 for $s_{21}$, and 4 for $s_{22}$. The trajectory terminates upon reaching any state besides $s_{0}$. In this MDP, the expected return for taking $a_{1}$ is 0 and for taking $a_{2}$ is 5. Let’s assume that our dataset consists of samples from a uniform random policy. Now, the optimism bias in Trajectory Transformer comes from the modified beam search they use at test-time. They jointly unroll the next action and resulting state, reward, and return and then filter based on which trajectories have the highest estimated return. In this setting, the trajectory model should properly predict all possible transitions. Then, their beam search will chose to take action $a_{1}$ because it will lead to the highest return of 10 when it predicts the transition $s_{0}\rightarrow s_{11}$. However, picking $a_{1}$ is obviously suboptimal in expectation, even if $a_{1}$ could lead to the best possible return. The optimism bias for Decision Transformer comes from picking a fixed target return and their heuristic of setting it to be the highest in the dataset. If we condition on getting a return of 10 (the highest possible), then it will clearly pick action $a_{1}$ as that is the only action where the probability of achieving the desired outcome is nonzero. Note that this is an issue for any return-conditioned method, not just Decision Transformer. The inclination for both methods to select $a_{1}$ because there is a possibility it has the highest return, even though $a_{2}$ is a better action in expectation and in the worst-case scenario is the specific optimism bias we allude to in this paper.

We validate the optimism bias in these approaches in a toy multi-modal autonomous driving task, and further demonstrate how addressing these issues leads to our method achieving superior performance.

Our work is also related to the fields of risk-sensitive RL and conservatism for offline RL. Most prior works in risk-sensitive RL focus on learning a policy that is optimized not for expected return, but rather some risk-sensitive profile over the distribution of returns (Ma et al., 2021; Shen et al., 2014). Many approaches accomplish this by learning a distributional Q-function that estimates the full distribution of potential returns (Dabney et al., 2018) rather than just the expected return. Instead of learning a distributional or risk-sensitive Q-function, we use a learned Transformer world model to hallucinate possible future trajectories and pick a robust behavior that does well in the worst-case predicted future.

Prior works in offline RL have explored different methods of incorporating conservatism in both model-based (Yu et al., 2020; Kidambi et al., 2020) and model-free (Kumar et al., 2020; Kostrikov et al., 2021) RL methods. However, offline RL approaches generally leverage conservatism specifically to discourage the agent from visiting state-actions that are outside the training distribution where the learned models could fail to generalize (Levine et al., 2020). In contrast, in this works we explore using a conservative approach to address the difficulties of planning in stochastic environments in safety-critical domains, like autonomous driving.

Many recent works in trajectory prediction leverage attention-based models (Tang & Salakhutdinov, 2019; Choi et al., 2019) and in particular Transformer neural networks (Ngiam et al., 2021; Liu et al., 2021; Quintanar et al., 2021; age, ) to make accurate long-term predictions in complex traffic situations. These methods can learn to attend specifically to the relevant factors in both the target and surrounding vehicles’ recent trajectories in order to make predictions that are consistent with the surrounding traffic. Additionally, many trajectory prediction approaches have incorporated VAEs (Tang & Salakhutdinov, 2019; Quintanar et al., 2021; age, ) to facilitate covering the different possible modes of future traffic trajectories. However, these approaches have mostly been evaluated on different variations of prediction error on a held out test-set with little focus on the effectiveness of leveraging these models for long-term online planning, like for the CARLA (Dosovitskiy et al., 2017) NoCrash (Codevilla et al., 2019) or Leaderboard¹¹1https://leaderboard.carla.org/ benchmarks.

In this work, we also leverage Transformers and VAEs in order to make multi-modal trajectory predictions. However, we focus specially on learning models that can be effectively used for robust search during deployment on an ego vehicle.

Recently, there have been several works that have focused on learning behavioral policies for autonomous driving from offline datasets. Most of this work revolves around performing imitation learning on the behavioral policy of a privileged autopilot agent in simulation (Chen et al., 2020; Prakash et al., 2021; Chen et al., 2021a).

In our work, we also use a privileged autopilot to collect our offline dataset in simulation. Instead of just learning the behavioral policy, we use these data to learn both a multi-modal policy, and a multi-modal world model that allows us to do efficient and robust search at test-time. Later, we demonstrate how this search leads our method to outperform pure imitation learning.

In RL, we treat the environment as a Markov decision process (MDP)²²2While we do evaluate in environments that could be considered POMDPs, we use the MDP formalism in our paper to keep it consist with most prior work in offline RL. represented as $(\mathcal{S},\mathcal{A},T,\mathcal{R},\rho_{0},\gamma)$, where $\mathcal{S}$ denotes the state space, $\mathcal{A}$ denotes the action space, $T(s^{\prime}|s,a)$ represents the probabilistic transition dynamics, $\mathcal{R}$ is the reward function, $\rho_{0}$ is the initial state distribution, and $\gamma\in(0,1)$ is the discount factor.

For MDPs, a trajectory is a sequence of states and actions $\tau=\{s_{0},a_{0},s_{1},a_{1},\cdots,s_{T},a_{T}\}$. Each trajectory has corresponding rewards $r_{\tau}=\{r_{0},\cdots,r_{T}\}$. The discounted return for a specific timestep is

The goal in RL is to find a policy that maximizes the expected discounted return

In offline RL, we are a given a fixed dataset $\mathcal{D}_{\beta}$ consisting of trajectories collected by some behavioral policy $\pi_{\beta}$ in the environment. Without collecting additional data, we must learn a policy that will be effective immediately upon deployment.

Transformers (Vaswani et al., 2017) are a neural network architecture that use several stacks of self-attention blocks to process an arbitrary collection of inputs. By leveraging positional encodings and a causal-attention map, Transformers can be used as sequence models for autoregressive generation. In this work, we mostly use a similar GPT-based (Radford et al., 2018) architecture as DT (Chen et al., 2021b). In particular, we use the same linear layer + layer normalization to project the raw inputs into the embedding dimension. We use 4 self-attention blocks with 8 self-attention heads. While our high-level details are more inspired by DT, our code and GPT implementation are based on the publicly available TT (Janner et al., 2021) codebase³³3https://github.com/JannerM/trajectory-transformer.

Our method uses conditional variational autoencoders (CVAE) as generative policy and world models that we can use to generate realistic and multi-modal candidate trajectories for test-time search. Thus, in this section we give a general overview of variational autoencoders (VAE). The goal of a VAE model is to generate samples $\hat{d}$ that are within the distribution of a training set $\mathcal{D}=\{d_{i}\}_{i=1}^{N}$. Thus the model should be trained to maximize the likelihood of all training points $d_{i}$. VAEs accomplish this by sampling a latent variable $z$ from some prior $p(z)$ and training a decoder $p_{\theta}(d|z)$ to convert the latent variable into a sample. However, optimizing this model is often intractable and thus Kingma & Welling (2013) introduces an encoder $q_{\phi}(z|d)$ that can be used to instead optimize the evidence lower bound (ELBO) on the log-likelihood of the data

The expectation term represents the reconstructions loss, and the KL Divergence term acts as a regularizer that keeps the encoder output close to the prior. Thus with a trained VAE model, we can sample a $z$ from our prior $p(z)$ and pass it to our decoder $p_{\theta}(d|z)$ to generate a sample from the desired distribution.

When training VAEs with high-capacity models like Transformers or CNNs (Krizhevsky et al., 2012), the decoders often learn to ignore the high-entropy stochastic latent variables. Therefore, Higgins et al. (2016) proposed a modification to the standard ELBO loss and introduced a coefficient $\beta$ to the KL Divergence term. Lowering this hyperparameter $\beta$ below 1 reduces the regularizing effect on the latent variables and thus makes it easier for the models to incorporate the latent variables into their predictions. We find $\beta\in[0.0001,0.01]$ to be effective for our experiments.

In order for our models to be useful for search, they need to (1) be able to produce a good range of candidate behaviors for any given situation for the ego vehicle and (2) cover a majority of the different modes of potential responses from the agents in the environment. Towards this end we train separate Transformer-based VAEs for both the policy and world models. We make the specific design choice for the stochastic latent variables for both models to be discrete and consistent over the entire planning horizon. This allows us to tractably enumerate all possible candidate trajectories without exponential branching, which enables efficient search at test-time. The intuition is that the policy latent variables should correspond to different high-level intentions or policies for the ego-vehicle, like whether to tail another vehicle aggressively or keep your distance. Similarly, the world latent variables should correspond to different possible intentions of the observable vehicles, like whether a vehicle in the opposing lane will go straight or turn through the intersection. Additionally, the world latent variables should capture other events like lights changing or cars suddenly appearing in the sensing range after rounding a corner.

During training, we sample batches of $K$-length trajectories from our offline dataset. The states and actions of these trajectories are passed into the world and policy encoders in order to generate the corresponding $z^{w}$ and $z^{\pi}$s. Each $z^{\pi}$ and its corresponding trajectory are passed into the policy decoder, where it predicts all actions in the trajectory through the standard teacher-forcing procedure (Williams & Zipser, 1989). Each $z^{w}$ and its corresponding trajectory are passed into the world decoder, where it predicts all the next states, rewards, and discounted returns in the trajectory also through the teacher-forcing procedure. The policy decoder and encoder parameters are updated to minimize equation (1), and the world decoder and encoder parameters are updated to minimize equation (2). We train all our models with the Adam optimizer (Kingma & Ba, 2014) with a learning rate of $1\mathrm{e}{-4}$ and weight decay $0.1$. Additionally, we normalize all raw values by subtracting the mean and dividing by the standard deviation of the dataset.

In this section, we describe how we can use our trained conditional policy and world model decoders in order to perform efficient and robust search at test-time.

We start with a toy autonomous driving problem that we designed to be very simple, but that still demonstrates the dangerous optimism bias in prior Transformer-based approaches.

In this toy problem, we have an ego vehicle trailing a leading vehicle with both travelling in the same direction on a $1$-D path. Both vehicles are represented using point-mass dynamics, but only the ego vehicle is controllable. Half of the time the leading vehicle will begin hard-braking at the last possible moment in order to stop just before the 70m mark before continuing. The other half of the time the leading vehicle will immediately speed up to the maximum speed and continue for the entire trajectory. The ego vehicle cannot infer beforehand whether the leading vehicle will brake or not, and thus this is a completely stochastic event from the perspective of the ego vehicle.

The observation space is the absolute position and velocity of both the ego and leading vehicle. The action space is just the acceleration for the ego vehicle clipped to $[-1,1]m/s^{2}$. The maximum velocity for both vehicles is $10m/s$ and the minimum velocity is $0m/s$, so the vehicles cannot travel backwards. The ego vehicle is rewarded for the distance traveled at each timestep, but is given a penalty of $-100$ if it crashes. The trajectory ends after $10s$ or if the ego vehicle crashes into the leading vehicle. The ego vehicle is initialized at $0m$ with a random velocity in $[7.5,10]m/s$. The leading vehicle is initialized randomly within $[10,20]m$ and with the same velocity as the ego vehicle.

For the offline dataset, we collected $\mathrm{\sim}{100000}$ steps with a distribution of different IDM (Treiber et al., 2000) controllers that demonstrate a wide range of aggressiveness, and includes some trajectories where the IDM controller is too aggressive and collides with the leading vehicle. We show our results in Table 1.

For DT, we find that conditioning on the maximum return in the dataset leads to crashes every time the leading vehicle brakes. If we condition on the mean return of the best controller used to collect the dataset, then we get the opposite behavior. The agent does not crash, but also does not take full advantage of the situations where the leading vehicle does not brake, and thus underperforms. While we were able to tune the conditional return in order to get reasonable results, we found the optimal value to be quite arbitrary. Thus, we believe this parameter will be very difficult to tune in more general and complex stochastic environments, where the possible returns would be very hard to estimate beforehand without significant prior knowledge of the specific testing scenario.

For TT, we find that the results depend heavily on the scope of the search used. When we use the default parameters from their codebase for the beam search, we get results very comparable to behavior cloning, which is quite suboptimal. When we reduce the low-probability filtering to allow for more aggressive search, we find that the method sometimes crashes into the leading vehicle because it picks the predicted trajectory where both vehicles will continue at max speed. Similar to DT, we expect tuning the search aggressiveness for TT to be difficult without significant prior knowledge of the intended testing scenario.

For our method, we find that our world VAE is able to predict both possible modes for the leading vehicle’s behavior. Additionally, the policy VAE seems to be able to predict a range of different trailing behaviors. Thus, our method is able to properly search for an effective and robust behavior and achieves results comparable to the best controller in the distribution used to collect the data.

Next, we evaluate our method on the CARLA (Dosovitskiy et al., 2017) NoCrash (Codevilla et al., 2019) benchmark. For these experiments, we run the $0.9.11$ version of CARLA at $5$fps. We assume access to a global route planner that can generate dense waypoints to our goal, an accurate localization system, and a perception system that can identify the state of any vehicle or traffic light directly in front of us within a limited sensing range. In our experiments, we obtain these ground truth quantities from the CARLA simulator, as commercial self-driving car efforts already have systems to provide these quantities. Thus, we leave implementing such systems as beyond the scope of this project.

The goal in the CARLA NoCrash benchmark is to navigate in a suburban town to a desired goal waypoint from a predetermined start waypoint. The benchmark takes place in the CARLA towns Town01 and Town02 and consists of 25 different routes in each town.

For our observation we use a low-dimensional vector representation consisting of: (1) the relative heading error to the next target waypoint, (2) the distance from the center of the target lane, (3) the ego vehicle speed, (4) the relative distance to the leading vehicle or the max sensing range if there is no leading vehicle in range, (5) the speed of the leading vehicle or the max speed if there is no leading vehicle in range, (6) the distance to the upcoming red light or the max sensing range if there is no red light in range. There are 2 actions: the steering and the target velocity for a PID controller. The ego car is rewarded for traveling faster and receives a small penalty for deviating from the target lane and a large penalty for crashing or incurring a traffic infraction. We terminate the trajectory whenever the car crashes, incurs an infraction, times out, or reaches the goal.

We collect our offline dataset with autopilot agents with a distribution of different levels of aggressiveness. This aggressiveness corresponds to the parameters of a time-to-collision-based controller that is used to adjust the vehicle’s speed in response to any leading vehicles. Additionally, the autopilots use a PID controller for steering, and always immediately brake if they are too close to a red light. We collect $\mathrm{\sim}{300000}$ steps with these autopilots in the Town01 routes in the dense traffic setting. The autopilots are imperfect and have an average success rate of $95.45\%$ in Town01.

We evaluate all methods by training on this Town01 dataset and then running in the unseen Town02 routes with the dense traffic setting. We depict our results in Table 2. We find that TT and DT conditioned on the max return have a lower success rate than BC, which we suspect is due to the optimism bias we previously described leading to unnecessary collisions and infractions. When we tune the target return for DT it can outperform BC in terms of average success rate and speed. However, similar to the toy problem, we found the tuned return to be quite arbitrary. Without online evaluation or prior knowledge of the testing domain, it would be quite difficult to estimate the best target return, especially considering that the testing scenarios are different from the training scenarios. Our SPLT method is also able to achieve a higher average success rate and speed compared to BC. We suspect that our positive results are due to our method’s planning procedure which avoids the optimism bias of TT’s naive beam search.

IQL is a non-recurrent model-free offline RL method. Instead of using model-based planning or a return-conditioned policy, it leverages a conservative Q-learning approach to perform advantage weighted regression(Peng et al., 2019) in order to learn a policy that improves over the data-collection policy. IQL’s superior performance compared to all the offline Transformer-based approaches demonstrates the benefits of leveraging a conservatively trained Q-function. We believe the contributions of our work related to addressing the optimism bias in prior Transformer-based offline RL methods to be mostly orthogonal to IQL. Thus, we leave potential extensions of incorporating an IQL style Q-function into the planning or policy learning procedure as future work.

Next, we evaluate our method on a modified version of the CARLA Leaderboard⁶⁶6https://leaderboard.carla.org/ benchmark. For these experiments, we run the $0.9.10.1$ version of CARLA at $10$fps. The only additional assumption we make is that our perception system can identify the state of any vehicles or pedestrians directly surrounding us in all directions within a limited sensing range.

The CARLA Leaderboard benchmark is much more comprehensive and requires the agent to perform more involved maneuvers like lane-changing in urban and highway situations. The major difference from our NoCrash setup is that we introduce 8 additional variables to the observation, corresponding to the distance and speed of surrounding vehicles in each of the 4 diagonal directions. We collect $\mathrm{\sim}{1.2}$ million time steps using the autopilot from the Transfuser (Prakash et al., 2021) codebase⁷⁷7https://github.com/autonomousvision/transfuser in the CARLA Challenge 2021 training routes for our offline dataset. We evaluate all methods on the officially released devtest routes. We depict our results in Table 3.

Besides DT with a tuned conditional-return, our method achieves the best overall driving score among all methods. Driving score is a comprehensive indicator for driving quality that accounts for route completion, collisions, and traffic infractions. The increased complexity in the Leaderboard scenarios leads to the world being less predictable and cooperative from the ego-vehicle’s perspective. Thus, our SPLT Transformer method, which disentangles the world dynamics and the agent decision-making process, is better equipped to handle this stochastic and uncooperative environment.

We believe that TT in particular underperforms relative to our SPLT method because its naive beam search does not plan appropriately for the range of possibly uncooperative multi-modal outcomes. On the other hand, our method is even able to achieve results on par with the privileged autopilot we used to collect the dataset.