Tractable Joint Prediction and Planning over Discrete Behavior Modes for Urban Driving

There has been a substantial amount of work towards training trajectory forecasting models for urban driving [10, 1, 5, 3, 4, 2]. However, none of these papers performs evaluations on closed-loop planning tasks beyond doing log-replay on offline driving logs for short time horizons. Most closely related to the model we use are Scene Transformer [3] and AutoBots [4], which both use vector-based input representations and large transformer models. Our model can support many state-of-the-art forecasting approaches, and only requires minimal modification to most in order to facilitate planning. Note that although all of these approaches are trained to perform open-loop prediction, our approach demonstrates how they can be adapted to perform fully reactive closed-loop prediction and planning without needing to substantially modify the training procedure.

There is a substantial amount of literature on planning over learned forecasting models for driving. Several prior works plan over imitation models, but do not perform closed-loop planning, instead optimizing a fixed ego-trajectory [11, 12, 13, 14, 15]. Another common approach is to not model the influence of the ego-agent on other agents [16, 17, 18]. As we will demonstrate, these approaches are unable to perform proactive planning when forced into close interactions with dense traffic. Model-based reinforcement learning approaches such as [19, 20] can theoretically leverage learned models and perform closed-loop prediction and planning for autonomous driving. However, many of these approaches have been restricted to small state spaces and have not been shown to perform well on more difficult closed-loop driving environments. In particular, we argue that rich vector-based state representations such as the one used in this work are key to effective prediction and planning. [21] models joint interactions between agents using a game-theoretic framework and evaluates planning performance in simulation, but does not use online fully reactive closed-loop planning like our approach.

The most similar work to ours in the literature is CfO [6], which also does fully reactive closed-loop planning for driving in CARLA. The tasks they consider are significantly simpler than the ones considered in this paper – they only consider interactions with one other vehicle (we consider up to 100). Their approach also relies on an autoregressive normalizing flow model which, to our knowledge, has not yet been scaled to larger-scale datasets, whereas our model architecture is similar to other transformer-based forecasting models in the literature which have been deployed at scale.

Most competing approaches in CARLA [8] are based on imitation learning [22, 23, 24, 9, 25]. While most approaches focus on the visual imitation setting, our closest competitor is PlanT [26], which also does imitation learning, but operates in a similar setup to ours, where we assume the perception problem is solved. Imitation learning approaches rely on collecting optimal demonstrations, which can be challenging as we scale up to harder scenarios and more realistic data collection settings. We show that unlike imitation learning, our method is able to out-perform the demonstrator on our suite of challenging urban navigation tasks.

In this section, we present our model architecture used for trajectory prediction. Let $x_{t}^{a}$ denote the state of vehicle $a$ at timestep $t$ (we assume $x_{t}^{0}$ and $x_{t}^{1:A}$ are the vehicle states for the ego-vehicle and the rest of the vehicles respectively). We also assume that $t=0$ is the current timestep and $x_{\leq T}$ refers to future vehicle states $x_{1},...x_{T}$ (ignoring the current timestep). Our goal is to model the distribution over future vehicle states $P(x_{\leq T}|x_{0},c)$ where $c$ is some arbitrary conditioning information (omitted from now on for brevity). We can factorize this distribution autoregressively over time, as is common in conventional trajectory prediction approaches. In this work we only model marginal (as opposed to joint) distributions over agents, so we can additionally factorize our distribution as

We propose a simple model formulation which captures multimodality and is well-suited for downstream reactive planning. We represent the high-level multimodal behavior of the agents with categorical latent variables $z_{t}^{a}\in[1,...,K]$ for each agent:

Following [5], we implicitly characterize $P(x_{t+1}^{a}|x_{t},z_{t}^{a})$ using a deterministic mapping $x_{t+1}^{a}=f(x_{t},z_{t}^{a})$. Therefore, we only require our model to learn this 1-step prediction mapping $x_{t+1}^{a}=f(x_{t},z_{t}^{a})$, and estimate $P(z_{t}^{a}|x_{t})$. In practice, we find that training our model to predict $H$ steps open-loop into the future $x_{t+1:t+H}^{a}=f(x_{t},z_{t}^{a})$ is a useful auxiliary task that improves the training and generalization of the desired 1-step prediction mapping. Below we show how $f(x_{t},z_{t}^{a})$ and $P(z_{t}^{a}|x_{t})$ can be parameterized by a transformer model.

We design our architecture to predict each of the different modes for all agents independently, but all in one pass of the model. Thus, our model outputs a vector with shape $[K,A,H,4]$ of potential future states of all agents in the scene, where $K$ is the number of modes and $A$ is the number of agents. This corresponds to $K\times A$ trajectories of length $H$, where each waypoint consists of position, heading, and speed. The high-level structure of our model is depicted in Figure 1.

Since the number of agents and context objects in the scene is not fixed, we adopt a flexible vector-based representation of the scene. We consider two types of entities: vehicles $\mathcal{X}$ and scene context objects $\mathcal{C}$ which consist of road points, traffic lights, pedestrians, stop signs, and goal waypoints. Each object is represented by a feature vector that contains the relevant raw information. For vehicles, this is relative pose, bounding box, speed, and current effective speed limit. For road points this is relative pose, lane width, whether the point is in an intersection, and whether one can change lanes to the left and right. For traffic lights this is relative pose, affected bounding box, and light state. For pedestrians this is relative pose, bounding box, and speed. For stop signs this is relative pose and affected bounding box. For goal waypoints this is relative pose and lane width. We assume access to all agents and context objects within 50m of the ego-vehicle up to a hard cap for each class of object. Each object is encoded using a MLP (one for each object type) to ensure all object features are the same dimension $D$. Then the entire set of input features is processed by an encoder to produce encoded vehicle features: $\mathcal{X}^{\prime}=\textrm{Enc}(\mathcal{X},\mathcal{C})$.

The encoder consists of stacked cross-attention layers where each vehicle cross-attends to both the vehicle features $\mathcal{X}$ and the context features $\mathcal{C}$ concatenated together. In practice, we find that only allowing each vehicle to attend to nearby map features improves performance, so we only attend to the closest 50 map features. Additionally, we use relative positional embeddings to maintain translational invariance. Prior work [27] has demonstrated that translational invariance is a useful inductive bias when performing multi-agent trajectory prediction. When feature $u_{i}$ cross-attends to $u_{j}$, we compute query $Q_{ij}$, key $K_{ij}$, and value $V_{ij}$ as follows:

where $W^{q},W^{k},W^{v}$ are learned projection matrices and $p_{ij}$ is a learned relative positional embedding. All features have an associated spatial position $(x,y)$ as well as an orientation $\theta$. To compute $p_{ij}$, we compute the pose $(x_{j},y_{j},\theta_{j})$ of feature $u_{j}$ transformed to the frame of feature $u_{i}$ at pose $(x_{i},y_{i},\theta_{i})$ to get a relative pose $(x_{ij},y_{ij},\theta_{ij})$. We then construct a relative pose feature by concatenating $x_{ij}$ and $y_{ij}$, as well as $\sin(\theta_{ij})$ and $\cos(\theta_{ij})$. Finally, an MLP is used to encode this relative pose feature to obtain $p_{ij}$. Note that we do not include any absolute positional information in our features, so the network is only able to observe the relative poses of scene objects.

Following the encoder, our encoded vehicle features are of shape $[A,D]$. We expand these features to be of shape $[K,A,D]$ so that each unique combination of mode and actor corresponds to a single $D$-dimensional feature. Following [2, 3, 4], we introduce learnable anchor embeddings $M\in\mathbb{R}^{K\times D}$ for each of the $K$ modes, which are broadcasted and added to the encoded vehicle features in order to distinguish the different modes. We learn two separate sets of anchor embeddings: one for the ego-vehicle only, and one for all other vehicles.

For each combination of agent and mode we get the query features $Q_{k}^{a}=\mathcal{X}^{\prime a}+M_{k}$ with appropriate broadcasting. Then, we use the decoder to produce $K$ trajectory predictions for each agent: $Y_{k}=\textrm{Dec}(Q_{k}^{a},\mathcal{X}^{\prime})$. The decoder consists of two alternating operations: the query features $Q_{k}^{a}$ cross-attend to both the encoded vehicle features $\mathcal{X}^{\prime}$ and map features $\mathcal{C}$, and then the query features $Q_{k}^{a}$ self-attend along the $K$ modes dimension. In other words, for each agent, the different mode features self-attend to each other. Note that during self-attention, only features corresponding to the same agent can attend to each other. This ensures that each agent future is decoded independently.

Finally, the processed query features are passed to a MLP. Since each unique agent and mode has its own feature, each $D$-dimensional feature is mapped to a trajectory of length $H$, where each waypoint consists of positions, orientations, and speeds. For each agent, we predict relative to that agent’s coordinate frame, and also for each time step predict in-frame displacements rather than absolute positions, orientations, and speeds.

We train the model to predict the ground truth trajectories using a negative log-likelihood loss, where each prediction is parameterized by a Gaussian. Note that although we train our model to parameterize a Gaussian, we always take the mean when unrolling our model. In order to ensure our model can capture multimodal outcomes, we adopt a winner-takes-all objective [28] where only the closest of the $K$ predictions will backpropagate its loss. Since we are doing marginal prediction, our winner-takes-all objective takes the best prediction for each agent separately. Our model is also trained to output a logit for each mode to estimate the probability of that mode being used, and this is trained using a cross-entropy loss with a one-hot target for the closest of the $K$ predictions.

Our objective is to find the behaviors for the ego-vehicle that will maximize the discounted sum of rewards over $T$ timesteps, where $\gamma$ is a discount factor: $\sum_{t}^{T}\gamma^{t}R(x_{t},c)$. Because the ego-vehicle’s decision in the trajectory is determined by the chosen $z^{0}$, we can write our objective as

As illustrated in equation 2, we can sample $x_{\leq T}$ autoregressively using our transformer model. For each timestep, we just deterministically decode $x_{t+1}^{a}=f(x_{t},z_{t}^{a})$ for our latent sample $z_{t}^{a}$.

Given our specific transformer architecture, we can perform each autoregressive step of this procedure with just one forward pass of the model. We generate all the multimodal predictions at once by leveraging the anchor embeddings and pick the mode prediction that corresponds to $x_{t+1}^{a}=f(x_{t},z_{t}^{a})$. Note that even though we trained these models to make multi-step open-loop predictions, we only need to take the first step of the prediction, which is $x_{t+1}^{a}$. Then, we can repeat this procedure up to the decided horizon $T$ in order to generate a sample of $x_{\leq T}$.

Now in order to perform planning, we need to evaluate Equation 3 for specific values of $z^{0}$. Since the ego latents $z^{0}=[z_{0}^{0},...z_{T-1}^{0}]$ can vary between timesteps, $z^{0}$ can take on $K^{T}$ possible values ($K$ possible ego-modes for $T$ timesteps), which is computationally expensive to fully enumerate.

Empirically, we find that keeping the latent modes consistent across time (i.e. $z_{0}=...=z_{T-1}$) for both the ego-vehicle and surrounding vehicles leads to good performance, and thus we use this scheme in our experiments to reduce complexity. We hypothesize that this works because the learned anchor embeddings encourage the model to learn locally consistent modes since the embeddings are shared globally, i.e., they are not state-conditioned. Additionally, we find that conditioning on these different modes leads to meaningfully diverse behaviors over time. Thus by keeping the latent modes consistent across time, we only need to compare $K$ ego latent modes to optimize Equation 3.

Specifically, we evaluate Equation 3 for any specific $z^{0}$ by generating $N$ samples of $x_{\leq T}$ with the previously described procedure. At the first time step, for each surrounding agent we sample $z_{0}^{a}\sim P(\cdot|x_{0})$ and for the ego-agent we set $z_{0}^{0}$ to the ego mode we are evaluating. Then, we continue to set $z_{t}=z_{0}$ at every subsequent timestep during autoregressive generation. We estimate $\mathbb{E}[\sum_{t}^{T}\gamma^{t}R(x_{t},c)]$ by taking the mean over these $N$ samples. We do this for all $K$ ego modes for $z^{0}$, and pick the mode with the maximum expected sum of rewards. We use MPC and execute the first step by using it as a target for a PID controller, and continue to replan at every step.

This approach enables us to perform fully reactive closed-loop planning. Note that we are planning over latent modes rather than open-loop trajectories, which means we can effectively model the causal influence of the ego on other vehicles and vice versa. Additionally, this allows us to be robust to different responses from other agents, since we can evaluate each candidate ego latent $z^{0}$ under a variable number of latent samples.

Our reward function is as follows:

where $\beta_{coll}$, $\beta_{lane}$, $\beta_{speed}$, and $\beta_{light}$ control the relative weight assigned to each reward term. The collision term $R_{coll}$ is -1 and leads to a termination if the ego collides with another object and 0 otherwise. The lane term $R_{lane}$ penalizes deviation from the route, where $\text{lateral}(x_{t},c)$ is the lateral distance of the ego from the closest route segment. The speed term $R_{speed}$ encourages the ego to drive at the speed limit. The light term $R_{light}$ penalizes the ego for having non-zero speed if the light is red; the penalty is scaled by the driving speed proportional to the speed limit.

For all experiments, we train our method with 4 encoder layers consisting of a cross-attention layer followed by a feedforward layer, and 4 decoder layers consisting of a cross-attention layer, self-attention layer, and a feedforward layer. Following [29], we use Pre-Layer Normalization. For the merging scenarios, we use an embedding size of 128 for a model with approximately 1.9 million parameters. For the CARLA Longest6 benchmark, we use an embedding size of 256 for a model with approximately 7.4 million parameters. We train our model with AdamW [30] [31] with an initial learning rate of $2e-4$ with cosine annealing to 0 as we train for 100 epochs. We train our multimodal model with $K=8$ latent modes and use $N=8$ samples to estimate expected rewards during planning.

For both experiment settings, we use a visualization range of 50m and observe the closest vehicles by distance up to a cap of 100 vehicles. For the merging scenarios, we set the coefficients of the cost function as $\beta_{coll}=20$, $\beta_{lane}=0.1$, $\beta_{speed}=1$, and $\beta_{light}=0$, as there are no traffic lights in these scenarios. For the CARLA Longest6 benchmark, we set the coefficients of the cost function as $\beta_{coll}=20$, $\beta_{lane}=1$, $\beta_{speed}=1$, and $\beta_{light}=4.$

We design a suite of challenging urban navigation scenarios in order to evaluate whether our approach can safely perform proactive maneuvers in dense traffic. We identify 10 distinct merging scenarios that each involve a high degree of interaction with other cars in the midst of dense urban traffic. These other cars are controlled by CARLA’s internal traffic manager, which uses an adaptive controller with randomly initialized minimum safety distances and target speeds. In each scenario, the goal is point-to-point navigation where the goals are specified a priori. Three of these scenarios are visualized in Figure 2.