1 Introduction↩︎

In real-world settings, robotic motion planning must produce trajectories that satisfy task- and user-specific objectives while remaining realistic and safe in complex environments. Classical model-based planners offer flexibility through explicit cost functions and constraints, and can be highly robust and performant when deployed in controlled settings [1]–[3]. However, purely model-based approaches struggle in unstructured and open-ended domains, where realistic behavior depends on rich priors learned from data. This has driven a shift toward learning-based approaches that leverage large-scale data to acquire such priors. However, many learning-based planners bake task structure into the model via conditioning [4], training-time reward shaping [5], or specialized inference procedures [6], which can make it nontrivial to accommodate new objectives or constraints.

In this paper, we propose a generative framework aimed at unifying these paradigms by separating trajectory realism from test-time control. A key ingredient enabling this separation is compression. Consider mapping high-dimensional trajectories to short sequences of low-dimensional or even discrete tokens: this simultaneously (i) reduces the dimensionality of the decision space the planner must search over and (ii) concentrates the burden of producing realistic motion in a decoder, i.e., a learned reconstruction model that expands tokens back into full trajectories. 1 previews our approach, which implements this principle by learning an environment-conditioned autoencoder that maps trajectories to a highly compressed, ordered latent representation; motion planning then reduces to tree search over latent tokens guided by a user-specified test-time objective. This yields a simple interface for incorporating domain knowledge or preferences while relying on the decoder to produce realistic, context-appropriate behavior.

From a motion-planning viewpoint, the environment-conditioned, ordered tokenization induces a context-dependent library of maneuvers. Unlike classical maneuver libraries with a small fixed set of hand-engineered primitives [7], [8], the library here is learned from data and automatically adapts to the scene through environment conditioning. Furthermore, the hierarchical ordering of this library enables efficient traversal of a combinatorially large set of maneuvers via coarse-to-fine refinement.

From the point of view of guided generation, our method chooses a different route than the iterative gradient-based guidance commonly used with diffusion and flow models, offering unique advantages. For example, partial token prefixes decode to trajectories in the original state space, so the objective can be applied directly without requiring additional estimators. Furthermore, latent token search naturally supports objectives that are not smooth or differentiable.

We evaluate our approach on the Waymo Open Motion Dataset [9] and nuPlan [10], using autonomous driving as a demanding testbed for guided generation with diverse objectives. Overall, our results indicate that a highly compressed, ordered latent representation can turn guided generation into a direct search problem. Throughout our experiments, we find that greedy search in the learned token space provides a simple yet effective mechanism for test-time control, supporting a variety of objectives, yielding strong performance on closed-loop planning, and achieving high realism metrics on guided multi-agent scenario synthesis, all without task-specific training.

2 Related Work↩︎

Test-time guidance for generative models. A large body of work studies how to steer generative models at test time. In diffusion models [11], guidance is often implemented by modifying the sampling dynamics using gradients of a conditioning signal, including classifier guidance [12], classifier-free guidance [13], and guidance using more general functions [14], [15]. These techniques have also been adopted for trajectory generation in robotics, e.g., Diffusion Policy [16], and for controllable motion prediction in multi-agent settings, e.g., MotionDiffuser [17]. Our work differs in the mechanism used to realize test-time guidance: instead of gradient-based steering during sampling, we perform decode-and-score search over ordered latent tokens, which naturally supports objectives that are non-differentiable and can be evaluated directly on decoded trajectories even in intermediate search iterations.

Tokenization and ordered representations. Tokenization via autoencoders is a standard technique for reducing the dimensionality of generative modeling problems. In image modeling, generators commonly operate in learned latent spaces of continuous VAEs [18] or discrete tokenizers such as VQ-VAE/VQGAN [19], [20] to improve efficiency and scalability [21]–[23]. Recently, nested dropout [24] has been used to learn ordered and variable-length image tokenizations, inducing a coarse-to-fine hierarchy in which shorter codes capture global structure and longer codes add detail [25]–[27]. We build on this perspective by learning a hierarchical trajectory representation. However, unlike image tokenizers, our trajectory tokenizer is environment-conditioned, which shifts part of the conditional modeling burden into the decoder.

Training-free generation in tokenizer latent spaces. Training-free test-time generation has been explored by optimizing latent codes of pretrained generators. Examples include optimizing GAN latents [28] with CLIP-based objectives [29], [30] and optimizing discrete tokenizer codes as in VQGAN-CLIP [31]. While such methods have historically been most effective for image editing, recent work suggests that extremely high compression tokenizers like TiTok [32] can make lightweight test-time latent manipulation surprisingly effective for generation [33]. Our work is inspired by the observation that extreme compression can shift generation complexity toward the decoder, but differs in the inference mechanism: we leverage ordering and discretization to perform tree search over tokens rather than using continuous gradient-based latent optimization. Separately, compact learned latents have also been used for control by planning over latent actions with learned dynamics using tree search [34], [35]. While algorithmically adjacent in their use of search over learned discrete representations, these methods target latent action planning rather than objective-guided trajectory generation, and are not directly comparable to our setting.

3 Conditional Trajectory Autoencoder↩︎

Trajectory prediction and planning problems in robotics must consider information such as sensor inputs or known maps. We therefore choose to represent a trajectory conditionally on given information about the environment and design a conditional autoencoder with the goal of learning a compact and expressive latent representation of the trajectory under a particular fixed environment (2).

Notation. We denote the environment as \(\mathcal{E}\). In the typical autonomous driving setting, \(\mathcal{E}\) consists of static world features such as road edges, road lines, lane geometry, stop signs and traffic lights, as well as dynamic object history for each visible agent in the scene. The full trajectory of the agent of interest is denoted as \(\mathcal{T}\). The encoder \(\mathop{\mathtt{Enc}}\nolimits\) then produces a sequence of \(N\) \(D\)-dimensional latent tokens \(z := \{\mathbf{z}_i \in \mathbb{R}^D\}_{i=1}^N = \mathop{\mathtt{Enc}}\nolimits(\mathcal{T}, \mathcal{E})\) from the full trajectory \(\mathcal{T}\) of the agent of interest and the environment \(\mathcal{E}\). The decoder \(\mathop{\mathtt{Dec}}\nolimits\) attempts to reconstruct the trajectory \(\mathcal{T}\) from the latent tokens and the environment, producing a prediction \(\mathcal{T}_\text{pred} = \mathop{\mathtt{Dec}}\nolimits(\mathbf{z}, \mathcal{E})\). We represent trajectories as sequences of position samples uniformly spaced in time.

Prediction space and loss function. The decoder \(\mathop{\mathtt{Dec}}\nolimits\) predicts mean and variance parameters of a Gaussian distribution for each point along the trajectory, and is trained by minimizing the negative log-likelihood (NLL) of the ground truth reconstruction under the prediction. In practice, we make use of the \(\beta\)-NLL [36].

3.1 Compression via Adaptive Soft Quantization↩︎

To mitigate training challenges commonly associated with vector quantization while maintaining its desirable regularizing effects and enabling discrete search at test time, we use a form of soft quantization consisting of noise injection at the autoencoder bottleneck: \[\mathtt{corrupt}(\mathbf{z}) = \tanh(\mathbf{z}) + \boldsymbol{\epsilon}_t \label{eq:corrupt}\tag{1}\] with \(\boldsymbol{\epsilon}_t \sim \mathcal{N}(0, I\sigma_t^2)\). Note that a \(\tanh\) activation is applied pointwise to the input, effectively creating an amplitude-limited noisy channel.¹ The chosen noise level \(\sigma_t^2\) is picked adaptively during training to gradually ramp up from zero until a desired reconstruction accuracy is achieved. During training, we apply \(\mathtt{corrupt}\) to each token \(\mathbf{z}_i\) before feeding it to the decoder. At test time, we set \(\sigma_t = 0\). We find that this adaptive noise schedule outperforms choosing a fixed noise level (further details can be found in 7.4).

Hard quantization at test time. At test time, we leverage the decoder’s ability to predict starting from heavily noised input tokens by quantizing its input. For this purpose, we round each token elementwise to the nearest discrete level, out of \(N_l\) uniformly spaced levels. While reminiscent of finite scalar quantization (FSQ) [38], we apply hard quantization at test time and not during training.

3.2 Variable-Length Latent Encoding↩︎

To allow flexible reconstruction fidelity at test time and to facilitate structured exploration of the latent space, we choose to impose a causal ordering structure on the latent tokens \(\mathbf{z}_i\). Causality among the tokens \(\mathbf{z}_i\) is enforced via masked self-attention in the encoder and decoder networks. We additionally make use of nested dropout [24] to allow the decoder to operate on variable-length encodings. During training, nested dropout drops a random number of tail tokens from the encoding, forcing the decoder to reconstruct the trajectory from a truncated code.

3.3 Network Architecture↩︎

Our conditional autoencoder is implemented using three transformer models: an environment encoder which processes static world and dynamic object history information, and encoder and decoder models which attend to the environment information via cross-attention and process the latent tokens via causal self-attention. Our environment encoder follows Motion Transformer (MTR), making use of MTR’s local neighborhood attention [39]. At the bottleneck, we project transformer tokens to the desired low dimensionality \(D\) and during training apply nested dropout and noise corruption for soft quantization.

Multi-agent modeling. Joint tokenization of several agents’ trajectories offers potential for learning representations that exploit the correlations between trajectories present in multi-agent interactions. The environment encoder and cross-attention blocks are unchanged compared to the single-agent model, and process each agent’s trajectory individually in the agent-centric coordinate frame. To jointly model all agents, we use self-attention blocks with a hybrid attention mask allowing each agent trajectory token to attend to the other agent trajectories bidirectionally, while latent tokens can aggregate information from the set of all agents. As before, causal masking is imposed among latent tokens.

7.1 provides additional network architecture details.

4 Latent Token Search↩︎

Causal masking combined with nested dropout ensures that later tokens capture increasingly more fine-grained information about the encoded trajectory. Therefore, at test time, latent token sequences of any size may be used to obtain reconstructions of varying degrees of fidelity. For planning tasks, this property suggests a greedy latent space search strategy in which tokens can be picked one after another.

More precisely, consider evaluating the trajectory \(\mathcal{T}_\text{pred}\) according to some objective function \({f(\mathcal{T}_\text{pred}) \mapsto (C, \leq_C)}\) specified at test-time. Note that the output \(C\) can be any set admitting a total order \(\leq_C\) (e.g., \(f\) can be a scalar cost or a lexicographically ordered tuple, as in 5.3). We seek to find the trajectory that minimizes this objective by exploring the latent space of our decoder as follows. Starting with the first token \(\mathbf{z}_1\), we enumerate its possible quantized values to find the assignment minimizing the cost of the decoded trajectory \(f(\mathop{\mathtt{Dec}}\nolimits([\mathbf{z}_1], \mathcal{E}))\). Tokens are intentionally low-dimensional and heavily quantized, so enumeration is cheap (we score \(N_l^D\) candidates per step). We iteratively repeat this process for the next token while keeping the previous token assignments fixed. 1 illustrates this procedure, visualizing how successively selected tokens progressively refine a trajectory under a goal-reaching objective.

Variance prediction for out-of-distribution detection. Recall that our decoder is trained to predict the parameters of a Gaussian distribution for each point along the trajectory. We find that the predicted variance is correlated with the reconstruction error (see 7.5 for more details). Therefore, an effective strategy to detect and discard decoded trajectories of poor quality is to penalize high-variance predictions.

Optional continuous refinement. The decoder can also be evaluated on continuous latents, enabling optional gradient-based refinement of the tokens selected by discrete search, either for differentiable task objectives or auxiliary criteria such as predicted variance reduction. Discrete search remains essential, as continuous optimization from scratch performs substantially worse. See 7.7 for details.

5 Experiments↩︎

In our experiments, we answer the following three questions.

Reconstruction: does the learned encoding cover the space of trajectories sufficiently? In 5.1 we verify that this is the case by analyzing the autoencoder’s reconstruction accuracy. We also evaluate the ability of greedy token search to perform reconstruction, validating the effectiveness of our ordered latent representation.

Semantics: do encodings possess a high-level semantic interpretation? In 5.2 we qualitatively show that behavior can be transferred between different scenarios by simply copying tokens, indicating that certain encodings can be interpreted as mapping to certain high-level behavior.

Guidance: is token search adaptable to arbitrary objective functions, other than reconstruction? In [sec:exp-sa:guidance] [sec:exp-ma:guidance] we present results of behavior synthesis according to a variety of objectives, including a realistic route-following objective that enables the deployment of our framework as a self-driving planner in closed-loop.

In answering these questions, we consider two different datasets: the Waymo Open Motion Dataset (WOMD) [9] for open-loop experiments, and nuPlan [10] for closed-loop evaluation. We also consider two different domains: marginal single-agent as well as joint multi-agent trajectory modeling.

Finally, we ablate the main representation learning design choices in 5.5, finding that our proposed soft quantization and the use of a very low-dimensional bottleneck are key to the success of our approach.

5.1 Reconstruction↩︎

Hard quantization at test-time. Before attempting discrete tree search over quantized tokens, we analyze reconstruction accuracy under hard quantization on the WOMD validation set. As shown in ¿tbl:tab:recons?, we find that the conditional autoencoder’s drop in reconstruction accuracy is relatively modest even when heavily quantizing the tokens to \(N_l = 2\).

Greedy search. To check whether greedy token search is a suitable strategy for exploring the space of latent tokens, we consider a reconstruction objective that minimizes the average displacement error (ADE) of all position samples along the trajectory with respect to the ground truth. The greedy search with reconstruction objective forms a valid replacement for the learned encoder in our autoencoder. Indeed, ¿tbl:tab:recons? shows that greedy search significantly outperforms the learned encoder, demonstrating that greedy token selection is a valid approach thanks to the causal and noise-resilient structure of the autoencoder’s latent space.

Beam search. We find that using a less greedy search algorithm leads only to marginal improvements to reconstruction accuracy. This is demonstrated in 3 using beam search, which keeps a set of top candidates at each iteration, rather than just a single one.

Overall, we find the latent space of our decoder is sufficiently expressive using just three tokens, each heavily quantized to only three bits. This is evidenced by a reconstruction ADE that is significantly lower than state-of-the-art motion prediction models. For example, DriveGPT achieves a minimum ADE of 0.524 across 6 predictions [40], while our decoder is able to reconstruct the ground truth with an ADE of 0.386. We also find that our proposed representation can be efficiently traversed via greedy search, enabling our latent space to represent a large library of trajectories with an efficient hierarchical structure.

5.2 Semantics↩︎

We qualitatively show that latent token encodings learned by our conditional autoencoder carry high-level semantic information through token swapping. Concretely, consider decoding the latent token representation of a given trajectory \(\mathcal{T}_A\) in its environment \(\mathcal{E}_A\) under a different environment \(\mathcal{E}_B\): \[\mathcal{T}_{A \rightarrow B} = \mathop{\mathtt{Dec}}\nolimits(z_A, \mathcal{E}_B) \quad \text{with} \quad z_A = \mathop{\mathtt{Enc}}\nolimits(\mathcal{T}_A, \mathcal{E}_A). \label{eq:token-swap}\tag{2}\] [fig:token-copy-matrix:copy-matrix] shows the result of this simple manipulation for selected scenarios where the agent of interest is a vehicle at an intersection. We observe that decoding an encoding corresponding to a desired behavior in a different scenario can be used to transfer this behavior to a novel scenario.

Furthermore, we observe that token swapping can result in highly repeatable behavior transfer. [fig:token-copy-matrix:swap-many] shows two examples: (i) we decode just three different encodings (corresponding to turn directions at an intersection) in many different intersection scenarios, and (ii) we decode a single encoding corresponding to harsh deceleration, again in many different environments. In both cases, decoding the reference encoding in novel environments reliably results in the desired behavior being transferred, suggesting a class of maneuvers may be characterized by a particular token sequence in an interpretable way.

5.3 Guidance↩︎

Token swapping offers a simple, single-step editing primitive, but it cannot in general optimize an arbitrary objective or incorporate complex preferences. In this section we therefore use guided search over tokens to produce trajectories that satisfy a desired behavior. We begin with a detailed analysis of open-loop maneuver editing performance according to two separate objectives (¿tbl:tab:maneuver-opt:left-turn? ¿tbl:tab:maneuver-opt:lane-change?) and conclude with closed-loop planning results on the closed-loop nuPlan benchmark (1). The experiments in this section make use of the variance-based out-of-distribution detection described at the end of 4.

Maneuver editing. We synthesize maneuvers according to the following objectives.

Left-turn objective — This objective directly maximizes the cumulative leftward heading change along the trajectory and penalizes excessive curvature. Since it does not take into account road geometry or obstacles at all and completely lacks realism or safety components, this objective is unsuitable for use in a traditional trajectory optimization framework.

Lane-change objective — We design a lexicographically ordered lane-change objective with a non-differentiable component. The objective first counts the integer number of lane changes as detected by changes in which lane center is closest. It then computes a lane-centering component which measures the deviation from the closest lane’s centerline. The minimization target is the tuple containing the (integer) count of excess or missing lane changes followed by the (real-valued) lane-centering residual.

¿tbl:tab:maneuver-opt:left-turn? presents the results of the left turn experiment, which demonstrates that token search with the simple left turn objective is able to achieve a high rate of success in cases where a left turn is indeed available, while remaining robust against invalid guidance provided by this objective in cases where a left turn is not available. For this evaluation, we filter the WOMD validation set to collect scenarios with the ego vehicle traversing an all-way stop intersection in the straight direction. We then subdivide these scenarios into (i) a subset where a legal left turn is possible and (ii) a subset in which a left turn is not possible. Success rate² is significantly improved with increasing search depth, demonstrating the importance of maintaining a large library of maneuvers as well as the effectiveness of greedy search in efficiently traversing it. Furthermore, note that the decoder automatically ensures that behavior is consistent with the given scenario as evidenced by near-zero rates of contact between the predicted trajectory and road edge geometry, and near-zero rates of left-turning in scenarios without an available left turn. Example scenarios showing the original and newly synthesized behavior can be found in 5.

¿tbl:tab:maneuver-opt:lane-change? presents the results of the lane-change experiment. Full-depth greedy search yields near-perfect compliance with the desired number of lane changes while minimizing lane center deviation, demonstrating that latent token search can effectively optimize non-differentiable, structured objectives. We select 208 WOMD validation scenarios with a ground-truth left lane change. In ¿tbl:tab:maneuver-opt:lane-change? we report average speed as a diagnostic signal and to confirm forward progress. The filtered scenarios often correspond to overtaking maneuvers and we observe that the generated speed profiles are consistent with slower driving when suppressing a lane change, while matching the ground-truth speed otherwise.

Throughout our experiments, we also include beam search with a beam size of four (Beam4 in ¿tbl:tab:maneuver-opt:left-turn? ¿tbl:tab:maneuver-opt:lane-change?). While Beam4 yields a small improvement in lane-centering in our lane-change experiments, the left-turn experiment shows that broader search can reduce validity, e.g., by increasing left-turn rates in scenes where a left turn is unavailable (¿tbl:tab:maneuver-opt:left-turn?). For these intentionally underspecified objectives, broader search can exploit loopholes in the objective and decoder imperfections. In contrast, greedy search is both more computationally efficient and a conservative default that biases solutions toward more robust, realistic behavior.

See 7.6 for details and experiments with an additional speed reduction objective.

Closed-loop planning. To evaluate latent token search in closed loop we turn to the nuPlan dataset and simulator [10]. As opposed to WOMD, this dataset includes target route information and proposes a closed-loop simulation challenge based on route following. For our nuPlan experiments, we train a trajectory autoencoder on the nuPlan dataset. Following standard practice, we apply ego history dropout during training [4], [5] and include an auxiliary object collision loss. However, in our case no route conditioning is included at training time. Instead, we implement a route-following planner by leveraging the test-time flexibility enabled by latent token search. In particular, we consider an objective that seeks to maximize progress along the route (measured as arclength traversed along a reference path extracted from nuPlan’s route specification) while penalizing excessive deviation from the route.

We find that latent token search with our simple route following objective performs on par with recent imitation learning planners despite being trained without any route conditioning. While recent RL-based approaches can achieve even better performance [5], we see RL and other architectural improvements as orthogonal to our approach, as they can be readily incorporated into our training procedure or decoder architecture in the future. Detailed results on nuPlan’s closed-loop benchmark with nonreactive agents, including the final aggregate score (CLS-NR) as well as collision rate with other agents or objects (S-CR) and route progress (S-PR) scores, are presented in 1.

Closed-loop planner ablations. In ¿tbl:tab:nuplan-ablations?, we vary search depth, beam size and usage of a simple continuous latent token refinement step (details on this step can be found in 7.7). We find that a search depth of \(N = 3\) and a small beam size of two perform best, demonstrating the value of maintaining a large library of maneuvers combined with a greedy traversal strategy. Consistent with findings that compression is key to the decoder’s generative capabilities from the image domain [33], we find that reducing compression leads to reduced performance (see \(N=4\) experiment), and consistent with our prior experiments, we find that greedy search or a small beam size of two works best. We validate our route-following objective by disabling it, which leads to significantly reduced aggregate and route-progress scores. Adding a term to explicitly penalize road edge collisions can lead to a small improvement to the aggregate score. Finally, performing a small number of gradient-based optimization steps with a predictive variance minimization objective yields a further small improvement (note that this optional post-processing step does not require the main route following objective itself to be differentiable).

Computational efficiency. With the parameters of \(N = 3\), \(D = 3\) and \(N_\text{l} = 2\) used for most experiments in this section, greedy search requires just 24 evaluations of the decoder. On the NVIDIA RTX 6000 Ada GPU, this translates to about 115 trajectories per second, corresponding to about 2760 decoder calls per second. Note also that each call to the environment encoder is amortized across 24 decoder calls during search.

5.4 Guidance for Multi-Agent Trajectory Generation↩︎

We train a multi-agent trajectory autoencoder with \(N = 12\) tokens and token dimensionality \(D = 3\). We again impose binary quantization (\(N_l = 2\)) at test time.

Following MotionDiffuser’s evaluation protocol [17], we consider an objective which minimizes the final displacement error (FDE) defined as the Euclidean distance between each agent’s terminal position and the ground truth trajectory’s last point. We then validate the effectiveness of latent token search by evaluating realism, as measured by adherence of the generated trajectories to the ground truth trajectories quantified using the average displacement error (ADE) averaged over predicted agents, prediction horizons and object types, following the WOMD convention. Our baselines draw 64 samples for the evaluation of \(\operatorname{minADE}\) and \(\operatorname{meanADE}\) metrics. We therefore also generate 64 candidates by performing beam search with a beam size of 64, and then evaluate \(\operatorname{minADE}\), which returns the lowest ADE from among the 6 candidates with the best objective value. For the evaluation of \(\operatorname{meanADE}\) we generate a single sample using greedy search, since our method is deterministic. The min/mean FDE scores are defined analogously. Since terminal position guidance is an objective that can benefit from micro-adjustments to the final trajectory, we leverage our decoder’s ability to handle continuous tokens and optionally include gradient-based optimization of the latent tokens after running search (see 7.7).

As shown in 2, compared to controllable trajectory generation baselines from the literature, our latent token search achieves significantly better realism. The improvement in performance is particularly pronounced when considering mean metrics, highlighting that diffusion methods heavily rely on filtering of samples to achieve good performance, while our latent token search can reliably generate realistic samples without any filtering. As in the single-agent case, selecting the search depth can control the balance between realism and objective maximization, which we explore in 7.8.

Qualitative example. While in the previous experiment we impose guidance on every agent’s terminal position, in 6 we guide the terminal goal position for only a single agent in the scene. In this case, the joint trajectory decoder is able to synthesize joint trajectories in which the behavior of other agents, whose trajectories do not receive any direct guidance, is consistent.

5.5 Representation Learning Ablations↩︎

We conclude the experimental evaluation with ablations of the main representation-learning design choices underlying our conditional autoencoder and latent token search. All models in this section are trained for approximately \(30\%\) of the number of steps used for the main experiments. We evaluate each model on three tasks: (i) reconstruction using the learned encoder, both with and without test-time quantization; (ii) reconstruction using latent token search with a reconstruction objective; and (iii) guided generation using a goal-reaching objective that minimizes distance to the ground-truth terminal position.

2pt

Token dimensionality. We first vary the token dimensionality \(D\). As shown in 3, reconstruction performance improves with increasing token dimensionality, both for the learned encoder and for search with a reconstruction objective. However, the guided generation results in 3 show a trade-off between representation capacity and generation quality. While the goal-point error, measured by FDE, improves monotonically with increasing dimensionality, the full-trajectory ADE does not. This suggests that when dense supervision is available, as in reconstruction, higher-capacity tokens improve reconstruction quality. In contrast, when guidance is sparse, as in terminal goal reaching, it is useful to rely more strongly on the decoder by reducing the capacity of its input tokens. In our experiments, \(D=3\) provides the best trade-off between trajectory realism and controllability.

1pt

Architecture and training recipe. We next ablate the main architectural and training choices: causal masking, nested dropout, and the use of soft quantization during training rather than FSQ [38]. The results in 4 show that nested dropout is critical for inducing the desired coarse-to-fine hierarchy among latent tokens. Without nested dropout, learned-encoder reconstruction when given the full set of tokens remains strong, but search performance degrades substantially, indicating that the latent space is no longer organized in a way that can be efficiently explored by greedy token search. Causal masking has only a small effect on search performance, although removing it worsens reconstruction under quantization.

We also find that our proposed soft quantization during training, followed by hard quantization at test time, significantly outperforms applying FSQ during training across all metrics including both reconstruction and generation. Furthermore, FSQ appears to reduce the benefit of continuous refinement after search: unlike the soft-quantization setting, removing test-time quantization from the autoencoder does not substantially improve reconstruction.

6 Conclusion↩︎

We propose a guided generative model in which test-time objectives are optimized by decode-and-score search over an ordered, highly compressed latent code. By making partial token prefixes directly decodable and scorable in trajectory space, this interface supports a large class of test-time objectives without task-specific training, while relying on a learned decoder to enforce realism. We deploy our method in autonomous driving, but hope the underlying idea can find broad applications in controllable generation.

Limitations and future work. Our experiments focus on autonomous driving, where trajectories are relatively low-dimensional and strong contextual priors are available from maps and agent histories. An important direction for future work is therefore to extend compressed, ordered trajectory representations to other robotics domains, such as manipulation, which introduce higher-dimensional state and action spaces or contact-rich dynamics. Furthermore, our experiments show success of greedy search for a variety of driving-related test-time objectives. However, future work could include evaluation on tasks with even richer and longer-horizon structure, and explore additional classes of latent search and optimization algorithms.

Impact Statement↩︎

This work supports test-time controllable trajectory generation for robotics. In safety-critical domains (e.g., autonomous driving), responsible deployment requires extensive validation and system-level safeguards beyond the scope of this paper. The methods presented in this work may find additional applications beyond robotics. As with other work whose goal it is to advance the field of machine learning, controllable generation in other domains may have many potential societal consequences, none of which we feel must be specifically highlighted here.

7 Appendix↩︎

7.1 Conditional Autoencoder Model↩︎

A graphical overview of the design of our trajectory conditional autoencoder is shown in 7.

Trajectory and environment tokenization. We use a similar strategy to MTR [39]. Agent trajectories are augmented with additional metadata (object type, timestep index, heading angle, and relative position between consecutive samples) and encoded with a PointNet architecture [45]. Road edges and road lines are split into 20-point segments and also encoded using a small PointNet-style network. In contrast to the original MTR, we encode each feature in the local coordinate frame placed at its initial pose (for agent trajectories) or center of mass (for road geometry).

Transformer models. The environment encoder, taking as input the tokenized static world and agent history representations, uses local self attention layers to improve efficiency and incorporate a spatial locality bias. We again closely follow MTR [39] in this regard.

Our trajectory encoder and decoder make use of blocks of causal self-attention followed by cross-attention. In the case of the encoder, the queries are initialized from the following tokens: (a) the token encoding of the full input trajectory, (b) learnable positional encodings for the latent tokens and (c) one or more additional register tokens [46] (not shown in the diagram in 7). Even though the number of key/value tokens produced by the environment encoder may be large, we do not use any form of local attention or geometric aggregation, since cross-attention with respect to our small number of query tokens is sufficiently fast.

The latent tokens are then processed by the bottleneck consisting of projection to the desired token dimensionality and noise injection (during training) or hard quantization (at test time). After projecting back to the transformer feature dimensionality, the decoder processes the latent tokens with the same design as the encoder, noting that we again include one or more register tokens with learnable positional embedding. Finally, we regress the final trajectory from one of the register tokens using a small MLP.

In both our trajectory encoder and decoder’s self-attention layers, causal masking is enforced between the latent tokens, while register and input trajectory tokens may attend to all tokens.

7.2 Multi-Agent Model Details↩︎

As illustrated in 8, the multi-agent version of our conditional trajectory autoencoder consists of cross-attention layers to aggregate environment information (identical to the single-agent case) interleaved with self-attention layers that can aggregate information from all trajectories among each other and into the causally masked latent tokens. For invariance to global coordinate transformations, we encode the environment in each agent’s local coordinate frame. For invariance to different ordering of the input trajectories, we do not add any positional encoding to the trajectory embeddings fed to the encoder, relying on the transformer’s natural permutation equivariance. To also preserve equivariance in the decoder, we again do not add a traditional positional encoding to the input trajectory tokens. Instead, we compute a relative pose embedding for each agent \(a\) by first computing the pose of all other modeled agents in \(a\)’s frame and then processing those poses with a permutation equivariant PointNet-style encoder.

7.3 Model Hyperparameters↩︎

The following hyperparameters are shared across the environment encoder model, the trajectory encoder, and the trajectory decoder:

• Layers: 6

• Heads: 8

• Transformer feature dimensionality: 256

• Feed-forward MLP width: 2048

We use these parameters for both the single- and multi-agent experiments.

Token dropout schedule. During training, we perform token dropout with \(50\%\) probability. If dropout is enabled, we choose the number of tokens to keep according to an exponential schedule with probabilities proportional to \((1/2)^{N_\text{drop}}\), where \(N_{\text{drop}}\) denotes the number of tokens to drop.

Other training details. We train our models using the AdamW optimizer [47] with a batch size of 64, learning rate of \(10^{-5}\) and weight decay of \(0.02\). For adaptive soft quantization, we set \(\gamma = 0.9995\) and \(\Delta\sigma = 0.01\) (see 7.4). Our single-agent WOMD model is trained on around 70M samples, and we train our nuPlan and multi-agent WOMD models for around 300M samples.

7.4 Adaptive Noise Injection↩︎

During training, we apply a simple elementwise corruption \[\tanh(\mathbf{z}) + \boldsymbol{\epsilon}_t \quad\text{where}\quad \boldsymbol{\epsilon}_t \sim \mathcal{N}(0, I\sigma_t^2),\] to each token \(\mathbf{z}\) in order to support test-time quantization. We find that choosing \(\sigma_t\) to adaptively ramp up significantly improves convergence during training. Concretely, we dynamically choose \(\sigma_t\) based on the average displacement error (ADE) averaged across every prediction in the current training minibatch \(\mathop{\mathrm{ADE}}_t\): \[\begin{align} \sigma_t &= \max(0, \gamma \sigma_{t-1} + (1 - \gamma) \hat{\sigma}_t) \\ \text{with}\quad \hat{\sigma}_t &= \begin{cases} \sigma_{t-1} + \Delta\sigma & \mathop{\mathrm{ADE}}_t \leq \mathop{\mathrm{ADE}}_\text{target} \\ \sigma_{t-1} - \Delta\sigma & \mathop{\mathrm{ADE}}_t > \mathop{\mathrm{ADE}}_\text{target} \end{cases}. \end{align}\] Here, \(\mathop{\mathrm{ADE}}_\text{target}\) denotes the target training ADE and is a fixed hyperparameter. Likewise, the decay factor \(\gamma \in [0, 1)\) and the noise increment \(\Delta\sigma > 0\) can be used to tune the responsiveness of the adaptive schedule to changes in \(\mathop{\mathrm{ADE}}_t\).

A unique advantage of this approach is that it yields a decoder that can optionally accept hard-discretized latents at test-time, while retaining the ability to use continuous latents (we exploit this property in 7.7). Therefore, our diagnostic experiment in 9 does not directly compare against methods using hard quantization such as VQ-VAE [19] or FSQ [38].

7.5 Predictive Uncertainty Calibration↩︎

We empirically validate the quality of our model’s predictive variance by performing a risk-coverage analysis using  1500 randomly selected WOMD validation scenarios. Our results indicate that the model’s variance output is predictive of the reconstruction error and motivates our variance-based token rejection. This holds when not using any test-time quantization as well as when aggressively quantizing to \(N_l = 2\), as shown in 5. FDE refers to the mean absolute deviation of the final sample of the trajectory.

7.6 Details on Planning Objective Functions↩︎

Our framework supports general objectives of the form \[f(\mathcal{T}_{\text{pred}}, z) \mapsto (C, \leq_C),\] where \(\mathcal{T}_{\text{pred}} = \mathop{\mathtt{Dec}}\nolimits(z, \mathcal{E})\) refers to the decoded trajectory, and \(z\) is the tokenized representation. In the main text we have considered the left-turn objective with \(C = \mathbb{R}\) as well as a lane-change objective with \(C = \mathbb{N} \times \mathbb{R}\) under a lexicographic order that first prioritizes the correct number of lane changes and then minimizes a continuous lane centerline deviation.

In our experiments, we consider objectives that penalize high predictive variance, taking the following form. Letting \(n\) denote the length of the token sequence \(z\) (starting at \(n = 1\) at the beginning of search and increasing to \(n = N\) once the maximum depth is reached), \[f_g(\mathcal{T}_{\text{pred}}, z) := g(\mathcal{T}_{\text{pred}}^{(\mu_{\text{xy}})}) + \lambda \mathbb{1}[\mathcal{T}_{\text{pred}}^{(\sigma_{\text{xy}})} > \sigma_{\text{max}}(n)].\] Here, \(\mathcal{T}_{\text{pred}}^{(\mu_{\text{xy}})}\) refers to the decoder’s predicted mean of the trajectory. We use \(\mathcal{T}_{\text{pred}}^{(\sigma_{\text{xy}})}\) to denote the magnitude of the predicted covariance for the final sample of the trajectory, which we check against the threshold \(\sigma_{\text{max}}(n)\) to impose a heavy penalty \(\lambda \gg g(\mathcal{T}_{\text{pred}}^{(\mu_{\text{xy}})})\).

We find that an uncertainty threshold which depends on the token sequence length is beneficial, as the distribution of the predicted variance depends strongly on the length of the latent token encoding (see 10). This leaves us to freely choose the cost function \(g\) based on the desired target application.

Left turn maneuver optimization. The heading \(\theta[i]\) along each segment \(i\) of the polyline defined by the current candidate trajectory \(\mathcal{T}_{\text{pred}}^{(\mu_{\text{xy}})}\) is first computed using finite differences. We similarly compute the curvature \(\kappa[i]\). We can then compute the total cumulative heading change in the counterclockwise direction as \(\operatorname{CCW} = \sum_i \max(0, \theta[i + 1] - \theta[i])\). This allows us to define the very straightforward cost function \[g_{\text{left}}(\mathcal{T}_{\text{pred}}^{(\mu_{\text{xy}})}) := -\min\{{\operatorname{CCW}}, \theta_{\text{min}}\} + \lambda_\kappa \mathbb{1}\{\max_i \kappa[i] > \kappa_{\text{max}}\},\] which encourages turns with a leftward heading change of at least \(\theta_{\text{min}}\) (which we set to \(\frac{\pi}{4}\)) while rejecting excessive curvature maneuvers (we set \(\kappa_\text{max} = 0.35\) very high so that this constraint is rarely active, but \(\kappa_\text{max}\) could be tuned to match actual vehicle kinematics). In addition to ¿tbl:tab:maneuver-opt:left-turn?, we present visualizations of the left turn maneuver optimization in 11 across all scenarios, including those in which a left turn is not available. We highlight that the fact that token search is not always able to achieve the desired heading change is a desirable property: as shown in 11 (c), our examples include cases where left turns are illegal or impossible.

We also highlight that the ultimate heading change achieved by our token search is not simply clustered around the threshold \(\theta_{\text{min}}\), indicating that token search finds solutions that align with the correct road geometry rather than blindly optimizing the objective which, if used on its own to optimize a less robust trajectory representation, would be far too simple to produce correct behavior (11 (a)).

Speed reduction. We select \(\sim\)​800 test scenarios in which a vehicle is traveling at an initial speed of around 9 m/s and maintains a similar average speed throughout its full trajectory. In this case, our objective is to slow down to a lower final speed of 5 m/s maintained for the last three seconds of the trajectory.

Our choice of \(g\) is again very simple: \[g_{\text{slowdown}}(\mathcal{T}_{\text{pred}}^{(\mu_{\text{xy}})}) := \max_{i \in I} \max\{0, v_i - v_\text{max}\},\] with \(v_i\) denoting the magnitude of the velocity along the \(i\)th segment of the trajectory, again computed using finite differences, \(v_{\text{max}}\) denoting the maximum speed constraint value (5 m/s in our experiment) and \(I\) denoting the range of timesteps to apply the constraint over (in our experiment we apply it from \(\require{physics} t = \qty{6}{s}\) to the end of the predicted trajectory at \(\require{physics} t = \qty{9}{s}\)).

r0.5

As shown in 6, token search does not enjoy a 100% success rate. However, we argue that this behavior is desirable, allowing even very naively specified objective functions to produce reasonable real-world behavior thanks to the decoder serving as a form of “learned guard rails.” Indeed, we show in 12 (c) that our framework makes a best effort to increase deceleration while maintaining smooth and safe longitudinal jerk and acceleration values.

7.7 Gradient-Based Token Optimization↩︎

In both our closed-loop nuPlan evaluation as well as the multi-agent guided generation experiments we make use of optional gradient-based latent token optimization to boost performance. This simple post-processing step uses the solution found by search as a high quality seed that can be iteratively optimized further according to the same or a different objective. We implement this token optimization via gradient descent, directly backpropagating gradients from the computed objective value to the latent tokens.

In practice, we use Adam with a learning rate of 0.1, optimizing for 5 iterations in the nuPlan experiment (1 ¿tbl:tab:nuplan-ablations?) and for 40 iterations in the multi-agent generation experiment from 2.

Predictive variance minimization. For the nuPlan experiment (1 ¿tbl:tab:nuplan-ablations?), we choose an objective used during gradient-based optimization that differs from the search objective. While search maximizes route-progress, the continuous optimization’s objective is to minimize the terminal predicted variance of the solution found by search. This choice leads to a noticeable boost in nuPlan’s aggregate closed-loop score without imposing any additional restrictions on the base objective used in search (e.g., it may still be non-differentiable, or use lexicographic ordering of multiple components).

Final displacement error minimization. For the multi-agent guided generation experiment from 5.4, we use gradient-based optimization to further reduce the same final displacement error already minimized by search. While this optimization does not significantly improve realism metrics, it significantly helps to accurately enforce the prescribed final position. We conclude that gradient-based token optimization is a useful post-processing step in cases where objectives are to be tightly enforced.

Importance of discrete search prior to optimization. We find that performing discrete search first is essential in order for the optimization to produce high quality results. To demonstrate this, we compare optimization as refinement of search solutions to standalone optimization with randomly initialized tokens. For this experiment, we use the same multi-agent goal reaching objective as used in 2. As illustrated in 7, search with no post-optimization beats all of our optimization-only results on realism (as measured by ADE) while being vastly more efficient: searching to a depth of 8 requires only 8 batched sequential calls to the decoder, taking 25 ms per scenario on our hardware (Nvidia RTX 6000 Pro). Meanwhile, the optimization-only method is only able to achieve comparable FDE at significantly worse ADE when selecting the best output among several initializations and running for an extended number of iterations, a configuration which requires 415 ms per scenario (over \(16 \times\) slower). Search followed by optimization as a post-processing step achieves the best realism and objective satisfaction.

7.8 Search Depth Selection for Multi-Agent Guided Generation↩︎

To explore the dependence of both realism and constraint satisfaction metrics on the search depth, we repeat the experiments from 5.4 for a range of possible maximum depth values. For this experiment, we evaluate a random subset of WOMD’s interactive validation split containing 1024 scenarios and do not enable gradient-based token optimization.

13 presents results of our sweep. In all cases, we find that increasing the search depth leads to improved FDE, as expected due to the coarse-to-fine structure of our latent space. However, when using greedy search to generate a single sample per scenario, we observe that ADE (our proxy for realism) stops significantly improving after a depth of just four out of the maximum of \(N=12\) tokens (13 (a)). When combining this effect with the generation of multiple candidates via beam search, we find that it results in a reduction of realism (13 (b)). This can be explained by considering the decreasing diversity of candidates available at the end of search at increasing search depths, i.e., at a high search depth, we expect candidates to cluster tightly around an FDE-minimizing solution, but this comes at the expense of reducing coverage and therefore increasing the likelihood of excluding the most realistic solution.

Based on this observation, rather than running search for the full depth of \(N=12\) tokens, our experiments in 2 use a depth of \(N=8\) for generation using search only, and \(N=7\) for generation using search and gradient-based optimization.

References↩︎