Controllable Sim Agents via Behavior Latents Juanwu Lu Purdue University [email protected]
Junyu Zhu University of Tokyo [email protected]
Ziran Wang Purdue University [email protected]
arXiv:2607.02496v1 [cs.RO] 2 Jul 2026
Abstract Realistic traffic simulation requires agents that imitate logged behavior and can also be steered along interpretable axes. Such controllability enables engineers to isolate variables, reproduce specific edge cases, and test autonomous systems without real-world risk. We introduce Controllable Neural Variational Agents (CNeVA), a controllable simulated-agent framework that learns to infer a peragent Gaussian behavior latent from per-channel discounted returns via a closedform conjugate variational update, conditioning a rectified-flow trajectory generator trained on a mixed channel-mask curriculum for classifier-free guidance. To tackle scarcity in reward signals, we propose soft eligibility gates that replace hard binary thresholds with smooth exponential decay, preserving the gradient signal for near-threshold agents. On the Waymo Open Motion Dataset, CNeVA attains competitive realism on the benchmark while exposing per-channel controllability that the higher-ranked imitation models lack. Speed- and acceleration-based steering produces monotone responses without stall-induced reward hacking. Safety controllability is monotone and substantial with the introduction of soft eligibility. We manage to achieve steerable map compliance under a context-residual return measure. Furthermore, our experiment demonstrates that steering metrics must be read alongside physical-plausibility guardrails to avoid reward-hacking confounds. Keywords Generative Models, Controllable Simulation, Autonomous Driving
1
Introduction
As autonomous vehicles (AVs) move toward large-scale deployment, ensuring their safety and reliability requires extensive validation. However, real-world testing is costly, while conventional traffic simulators remain limited in fidelity and scalability. Recent advances in generative models have made data-driven traffic simulation a promising alternative for AV testing, enabling simulated agents to both reproduce realistic behaviors and respond to AV actions. For adversarial evaluation, simulated agents are often expected to exhibit specific behavioral preferences as well. Although imitation learning can enable the generation of diverse and realistic traffic behaviors (Seff et al., 2023; Wu et al., 2024; Philion et al., 2024; Zhang et al., 2025a), it generally lacks fine-grained controllability over agent behavior and therefore struggles to adapt generated scenarios to user-specified objectives. Generative modeling has expanded the diversity of synthesized traffic scenarios. Existing methods using variational autoencoders (VAE) (Suo et al., 2021; Zhang et al., 2023; Liu et al., 2026) produce behavioral variation through latent-space resampling but cannot deliberately steer behavior toward specific driving preferences. Diffusion-based models (Jiang et al., 2023; Cao et al., 2024; Jiang et al., 2024; Tan et al., 2025a) improve controllability by injecting guidance during the generation process, yet achieving stylistic trajectories through guidance tuning is often nontrivial in practice, and excessive guidance can lead to unrealistic trajectories. Self-play reinforcement learning encourages desired behaviors through hand-crafted reward functions in closed-loop interaction (Cornelisse and Vinitsky, 2024; Cusumano-Towner et al., 2025; Chang et al., 2026). Despite their controllability, these methods rely on computationally expensive online training and typically require retraining whenever the reward specification changes. A practical, controllable simulator should provide flexi1
m m o<1
o<1
s1
...
s2 o1,n
o2,n
sT
s1 o1,n
oT,n
o2,n
sT
oT,n
λn
N
(a) Probabilistic Graphical Model for sim-agents.
...
s2
N
(b) Extended PGM with per-agent behavior latents.
Figure 1: Probabilistic graphical models for simulated agents. (a) The standard hidden Markov formulation from the Waymo Open Sim Agents (Montali et al., 2023). (b) Our extended model augments each agent with a per-agent latent behavior profile λn , which captures how the agent weights a shared set of reward channels. The latent profile is inferred from the observed trajectory through the reward-tilted preference factor and is later used as a controllable conditioning variable during generation. ble control of behavior while preserving the realism of generated trajectories (Rowe et al., 2024). Existing studies also lack comparable representations of preference profiles across agent populations, which are needed for compositional-style behaviors and for adapting to variations across scenarios. To address these limitations, we propose Controllable Neural Variational Agents (CNeVA) , a controllable traffic simulation framework that learns from offline log-replay data. In the generated scenarios, it encodes the driving behavior of each agent by a K-dimensional Gaussian latent λn over a predetermined reward basis, such as safety, map compliance, speed, and acceleration, whose posterior is identified in closed form from each demonstration’s per-channel discounted return Gn . The inferred posterior conditions a flow-based scene generator. We propose a mixed channel-mask curriculum that exposes classifier-free guidance to arbitrary subsets of channels, allowing a user to fix some steerable axes while leaving the rest unconstrained during generation. Our main contributions in this work are as follows: C1 We cast controllable behavior generation as variational inference over a Gaussian behavior latent with a closed-form conjugate posterior and provide a tilt-vs-regression explanatory framework that predicts which channels are identifiable from the return-shrinkage factor. C2 We introduce context-residual return labeling, soft eligibility gates, and contrastive conditioning as extensions to the reward and training pipeline, producing safety controllability substantially above the hard-eligibility ablation while retaining near-ground-truth speed. C3 We present a systematic accounting of the limits in our framework through the scope of drift-paired channel-steering matrices (CSMs) across three return measures, showing that map controllability is coordinate-specific and that CSM values must be interpreted jointly with physical-plausibility guardrails to avoid reward-hacking confounds.
2
Hierarchical Graphical Models for Simulated Agents
We model a multi-agent driving scenario as a Hidden-Markov process (Montali et al., 2023) with history o<1 (section 5.1), latent world state st , per-agent observations ot,n ∈ Rd , and time-invariant map m. Figure 1a shows the resulting PGM. A model that maximizes this HMM likelihood reproduces logged behavior but provides no explicit toggle for steering generation. We therefore extend the graph (fig. 1b) with a per-agent behavior profile λn ∈ RK , so that behavioral heterogeneity across drivers reduces to how each agent weights a shared set of reward signals, such as safety, map compliance, speed, and acceleration. With per-step preference factor ψt (ot,n , λn ) ≜ exp γ t−1 λ⊤ n r(ot,n ) ,
(1) 2
NFF
noitnettA-fleS laitapS
noitnettA -fleS laropmeT
redocnE yrotcejarT
Road Graph Embed
Training noitnettA-ssorC
ksaM lennahC dexiM
roiretsoP etagujnoC
CFM Decoder Block
redocnE yrotcejarT
Traffic Light Embed
Population Marginal Concat Add Unconditional Path Conditional Path
CFM Model
Realistic
CFM Model
UserControlled
Euler ODE Solver
Prior Sampling
Adversarial
Controllable Generation
Figure 2: CNeVA infers a per-agent behavior latent λn from standardized reward returns via a conjugate Gaussian posterior, with mixed channel masking for classifier-free CFM training. At generation time, conditional and unconditional velocity fields are combined via CFG and an Euler ODE solver to produce population-marginal, user-controlled, and prior-sampled rollouts. PT discount factor γ ∈ (0, 1], and the per-channel discounted return Gn ≜ t=1 γ t−1 r(ot,n ) ∈ RK , we apply belief propagation to collapse the trajectory product into a single inner product, ψ(τn , λn ) ≜
T Y
ψt (ot,n , λn ) = exp λ⊤ n Gn .
(2)
t=1
Marginalizing the latent state and conditioning on history yields the relaxed joint distribution p(τ , λ | m) ∝
N Y
p(λn ) ψ(τn , λn )
n=1
T Y
p(ot,n | o<t , m),
(3)
t=1
where o<t is a sufficient statistic for st . We defer the full HMM, the PGM, and the beliefpropagation derivation to section B. Our formulation approaches controllability by inference: How can we infer the latent preference that governs the behavior of each observed agent?
3
Controllable Neural Variational Agents
We propose to instantiate the factor-graph problem of section 2 with a neural network through two coupled objectives: variational inference of the per-trajectory posterior over λ (section 3.1) and a conditional generator that consumes λ as guidance (section 3.2). Sampling-time control of the trained model is described in section 3.3. Figure 2 presents our end-to-end architecture: the conjugate posterior of section 3.1 produces a per-agent λn from the standardized per-channel return, the mixed channel-mask curriculum of section 3.2 trains a single conditional flow-matching velocity field that covers both conditional and unconditional inputs, and an Euler integrator with classifierfree guidance at deployment supports population-marginal, user-controlled, and prior-sampled rollouts. 3.1
Variational Behavior Labeling
Controllable generation requires a conditional generative model p(τ | λ, m; θ) per agent. Because the latent λn is not observed in D, we treat each trajectory τn as evidence and infer a posterior label from the per-agent factor-graph sub-problem rooted at λn . Given the trajectory-level factor ψ(τn , λn ) and a population prior p(λn ), variational inference (Jordan, 1999) yields the per-agent approximation q ⋆ (λn ) ← arg min KL [q(λn ) ∥ p(λn | τn )] = arg max Eq λ⊤ n Gn − KL [q(λn ) ∥ p(λn )] . (4) q∈Q
q∈Q
3
This objective places mass on profiles that explain the observed returns while regularizing the inferred label toward the population prior. With a Gaussian prior, exponential-family conjugacy yields an analytic solution. Proposition 3.1 (Conjugate Gaussian posterior). Let p(λn ) = N (µ0 , Σ0 ). Under equation 1 to equation 4, the variational optimum coincides with the reward-tilted posterior and is the multivariate Gaussian q ⋆ (λn ) = N µ0 + Σ0 Gn , Σ0 . (5)
The posterior covariance is fixed at Σ0 because the reward-tilted factor is linear in λn , so the residual uncertainty represents population-level label noise rather than trajectory-specific epistemic uncertainty; we do not claim that λn recovers the agent’s underlying utility in the inverse-RL sense. Derivation is in section B. At training time we sample λn ∼ q ⋆ with the reparameterization trick: 1/2
λn = µ0 + Σ0 Gn + Σ0
ηn ,
ηn ∼ N (0, I).
(6)
Tilt-vs-regression interpretation. The conjugate posterior admits two equivalent readings. As a tilt, the exponential factor exp(λ⊤ n Gn ) shifts probability mass toward profiles that explain the observed returns; as Bayesian regression, λn on Gn is regularized toward µ0 by the prior precision −1 Σ−1 governs how 0 . Under the regression view, the return-shrinkage factor Σ0 (Σ0 + Σnoise ) strongly each channel’s return updates the posterior mean: channels with high signal-to-noise ratio (dense, trajectory-level penalties such as speed) experience strong shrinkage toward the observed return and are easy to steer, while channels with low signal-to-noise ratio (sparse, context-dependent events such as collisions) are prior-dominated and harder to control through λ alone. This asymmetry predicts the experimental controllability hierarchy of section 5.3. Per-channel standardization. The raw return Gn has very different per-channel scales: section D.1 reports empirical means |µk | ∈ [3.0, 50.1] and standard deviations σk ∈ [5.2, 33.6]. Using Gn verbatim in equation 5 allows one channel to dominate the conjugate update by an order of magnitude. We pre-condition with the empirical mean µG and diagonal scale diag(σG ) measured once on a fixed WOMD calibration split, and use the standardized return e n,k ≜ (Gn,k − µG,k ) / σG,k , G
k = 1, . . . , K,
(7)
in place of Gn inside equation 5 and equation 6. With Σ0 = I, this brings the per-channel posterior mean to an order-unity scale, so the four channels enter the conditional generator of section 3.2 on equal footing. 3.2
Conditional Behavior Generation
With the posterior q ⋆ (λn ) fixed, we train a posterior-labeled conditional generator against the reconstruction term of the variational lower bound, θ ∗ ← arg max Eτ ∼D Eλ∼q⋆ (λ|τ ) [log p(τ | m, λ; θ)] .
(8)
θ∈Θ
We optimize this objective via conditional flow matching (Lipman et al., 2023), a simulation-free surrogate that regresses a neural velocity field onto the rectified-flow target instead of evaluating the likelihood directly. We treat p(τ | m, λ; θ) as the implicit distribution induced by the trained velocity field. Let yn denote the future motion target of agent n as displacement features (yt,n = ot,n − ot−1,n ). For noise ϵ ∼ N (0, I) and flow time s ∈ (0, 1) (drawn from a logit-normal distribution in our implementation), the noisy latent on the linear interpolation path and its target velocity are xs = (1 − s)yn + sϵ,
vntarget ≜ ϵ − yn .
(9)
A neural velocity field vθ (xs , s, o<t,n , m, λn ) is trained to predict vntarget , transporting noisy trajectory samples back to realistic future displacements while λn selects among behavior modes compatible with the same scene history and map context. Architecturally, λn enters as an extra cross-attention token prepended to the per-agent scenario-context conditioning set. Inference uses classifier-free guidance (Ho and Salimans, 2022). The standard binary on/off recipe is too coarse for our K-dimensional latent, so we expose every channel subset during training 4
via a four-branch curriculum (i.e., null, single-channel, two-channel, and full, with proportions 0.2/0.4/0.2/0.2. See equation 25 in section D.4). We write bn ∈ {0, 1}K for the resulting mask with bn,k = 1 marking channel k as masked-out, and concatenate the value-zeroed label with the mask indicator, e n ≜ (1 − bn ) ⊙ λn ; bn ∈ R2K , λ (10) so the projection can distinguish “channel k has value 0” from “channel k is unobservable.” When bn = 1K , the λ token is overridden by a learned null embedding e∅ ∈ Rdh , recovering the unconditional path required for the CFG combination at inference. The full training loss is L(θ) =
N X
h Eq⋆ (λn ), bn , s, ϵ
i e n ) − v target 2 . vθ (xs , s, o<t,n , m, λ n 2
(11)
n=1
By design, the null branch, active when bn = 1K , learns the population path p(τ | m) on the training set. In contrast, the partially- and fully-conditional branches learn the steered path p(τ | m, λ). Architectural details are given in section D. 3.3
Inference and Controllable Generation
At deployment, the operator supplies a per-agent λop n in the standardized coordinates of equation 7; the trajectory is drawn from pθ (τn | m, λop n ) via the rectified-flow Euler sampler with classifier-free guidance (Ho and Salimans, 2022): op eθw (xs , λop v n ) ≜ (1+w) vθ (xs , s, o<t,n , m, λn ) − w vθ (xs , s, o<t,n , m, e∅ ),
(12)
with guidance scale w ≥ 0. The population-marginal, user-controlled, and prior-sampled operator regimes, the receding-horizon patch scheme, and the full sampler are detailed in section D.5. 3.4
Return Labeling and Eligibility
Context-residual returns. The raw return Gn,k conflates driving style with scenario difficulty. For example, a highway scene can yield low off-road penalties regardless of the individual driving style. Therefore, we residualize each per-channel return against a scene-context baseline, Gcr n,k = Gn,k − Ḡk (mn ), isolating the behavioral signal. We use this as the primary return measure for all CNeVA training and evaluation (section C.1). As a coordinate-specific alternative used only in the map conformity ablation in section F.5, we also define a lane-centerline return from the signed lateral offset to the nearest lane centerline (section C.2). Soft eligibility gates. The safety and map conformity channels use per-agent eligibility gates that decide which agents receive non-trivial labels. The hard-gated baseline labels only agents within fixed clearance, time-to-collision (TTC), or road-margin thresholds, introducing a discontinuity at the threshold and excluding the majority of safe, on-road agents from supervision. We replace these binary gates with a smooth exponential decay in the per-step pairwise risk, so that all valid agents receive labels. In contrast, agents far from hazards contribute negligibly. The design preserves gradient signal for near-threshold agents and fixes the safety-CSM erosion documented in section 5.4. We defer details about the decay formulas and scales (τc , τt , τm ) to section C.3. Contrastive conditioning. We further augment training with a paired steered-null forward pass over the same flow-time noise, encouraging the velocity field to learn the steered-versus-unsteered difference and improving robustness under closed-loop drift (section C.4).
4
Related Work
Data-Driven Traffic Simulation for AV Testing. Unlike open-loop trajectory forecasting (Nayakanti et al., 2023; Shi et al., 2024), scenario generation requires agents that react to AV actions, whether by modeling the joint trajectory distribution (Ngiam et al., 2022) or via a secondstage refinement (Zhang et al., 2022). Tokenized models that discretize actions (Seff et al., 2023; Philion et al., 2024; Wu et al., 2024; Zhang et al., 2025a; Chang et al., 2026; Zhou et al., 2024) or spaces (Hu et al., 2025; Zhou et al., 2024), and that address covariate shift via rollout during 5
training (Zhang et al., 2025a), improve the behavioral realism, but these autoregressive approaches can lose behavioral diversity to dominant motion patterns. Generative VAE (Suo et al., 2021; Zhang et al., 2023; Liu et al., 2026), diffusion (Jiang et al., 2023; Cao et al., 2024; Jiang et al., 2024; Tan et al., 2025a), and flow-matching (Lipman et al., 2023; Xing et al., 2025; Tan et al., 2025b) models capture multi-modal behavior well but expose little interpretability of driving style, and self-play RL often sacrifices realism or needs reference models (Cornelisse and Vinitsky, 2024; CusumanoTowner et al., 2025) for regularization (Chang et al., 2026). CNeVA instead infers a closed-form posterior behavioral latent over an explicit reward basis, which acts as a single guidance token for a rectified-flow generator, improving the realism and diversity of controllable generation. Controllable Agent Behavior Generation. Modeling driving style as a controllable parameter enables behavior-specific generation (Yin et al., 2021; Rempe et al., 2022; Chang et al., 2023; Liu et al., 2026). In generative models, controllability is typically injected through guidance (hard constraints (Jiang et al., 2024; Tan et al., 2025a), hand-crafted differentiable costs (Huang et al., 2026; Chang et al., 2024), STL rules (Zhong et al., 2023b), conditional priors (Pronovost et al., 2023), or LLM-prompted attractors (Zhong et al., 2023a)), but each typically requires redesigning the guidance mechanism for new behaviors. Driving style can also be read as explicit or implicit reward preferences (Rowe et al., 2024; Cusumano-Towner et al., 2025; Rowe et al., 2025), though perchannel rewards vary across scenarios. CNeVA instead learns the latent behavioral profile directly, keeping inference cheap and the editable interface compact while exposing per-channel controllability handles that map cleanly to interpretable driving-style axes.
5
Experiments
We evaluate CNeVA on the Waymo Open Motion Dataset (Ettinger et al., 2021) with four questions: RQ1 Does the unconditional path generate realistic behavior (section 5.2)? RQ2 Does the learned latent steer per-channel returns with physical plausibility (section 5.3)? RQ3 What do ablations without soft eligibility reveal about safety controllability erosion and reward-hacking confounds (section 5.4)? In addition, we investigate how sensitive map controllability is to the choice of return measure and defer our discussion and results to section F.5. 5.1
Experiment Setup
We evaluate on WOMD (Ettinger et al., 2021) under the WOSAC protocol (Montali et al., 2023): given 1.1 s of 10 Hz history, the model generates 8 s rollout for each scenario, with up to N = 128 agents, ot,n ∈ R9 (pose, heading, planar velocity, bounding box). We report open-loop fidelity as minADES and the WOSAC composite meta-metric. We evaluate controllability by the channel steering matrix (CSM) and the trajectory divergence D(ea , eb ) (section 5.3). We build the proposed CNeVA with dh = 512, six Transformer decoder layers with eight attention heads, and K = 4 reward channels. Each channel is a negative sum of penalties. Hence, a higher Gn,k denotes less of the penalized quantity: safety (collision count, proximity to nearest object, inverse time-to-collision), map (offroad indicator, road-edge distance, computed under the context-residual formulation of section 3.4), speed (linear and angular velocity magnitudes), accel (linear and angular acceleration magnitudes). We use standard normal behavior prior (µ0 , Σ0 ) = (0, I). Returns are standardized by the training dataset statistics (section D.1). The CFM head predicts a 6-d motion target, with bounding box length, width, and height agent-static and reused from history. During sampling, we apply the Euler integrator with 10 steps, and the receding-horizon patch length is ℓ = 16. The main model uses soft eligibility gates (section 3.4) with decay scales τc = 2.0 m, τt = 3.0 s, τm = 1.0 m, futurequery map conditioning with k = 16 nearest neighbors, and random mid-step anchor sampling. All main results are evaluated at 200, 000 training steps. 5.2
Benchmark
Table 1 positions CNeVA against the full WOSAC realism spectrum. On the official 2025 test server, our model attains a Realism meta-metric of 0.7145 with minADE = 1.80 m, using inference-time 6
Table 1: Waymo Open Sim-Agents Challenge Realism metrics on the test split of WOMD. We rank the baseline methods by realism on the 2025 challenge leaderboard. Oracle shows the results evaluated with the ground truth. Methods with full or partial controllability are marked in gold. Model
Kinematic (↑)
Interactive (↑)
Map-based (↑)
Realism (↑)
minADE (↓)
Constant Velocity TrafficBots V1.5 (Zhang et al., 2023) SceneDiffuser (Jiang et al., 2024) VBD (Huang et al., 2026) TrajTok (Zhang et al., 2025b) SMART-R1 (Pei et al., 2026) Oracle (Jiang et al., 2024)
0.2253 0.4304 0.4295 0.4169 0.4887 0.4940 0.5565
0.4327 0.7114 0.7681 0.7819 0.8116 0.8109 0.8576
0.4535 0.8360 0.7756 0.8137 0.9207 0.9194 0.9593
0.3985 0.6988 0.7030 0.7200 0.7852 0.7855 0.8330
7.5148 1.8825 1.7670 1.4743 1.3179 1.2990 0.0000
CNeVA (Ours)
0.4732
0.7482
0.8091
0.7145
1.8029
Table 2: Channel steering matrix (CSM) diagonals (∆Rk ) for CNeVA with soft eligibility gates, evaluated at 200K training steps. We report results with open-loop sampling and context-residual return at ρ = 1. The w = 0 column is latent-only (classifier-free guidance off); the w = 1.5 column is latent with classifier-free guidance. Open-loop Channel
Uncond.
Safety Map Speed Accel
+0.06±0.00 +0.06±0.00 −3.33±0.00 +4.77±0.00
Null ADE Offroad
CNeVA (w = 0) CNeVA (w = 1.5) +0.29±0.00 +0.24±0.00 +3.21±0.00 +4.19±0.00
+0.66±0.10 +0.61±0.14 +8.15±0.07 +8.76±0.07
1.113±0.011 32.5±0.2 %
guidance and best-of-128 selection on the frozen model. The benchmark result shows that CNeVA can achieve competitive behavioral realism. We notice that the top of the leaderboard is dominated primarily by non-controllable closed-loop and tokenized imitation models, and our remaining gap to them is concentrated in the collision and off-road components, consistent with open-loop error accumulation over the 8 s rollout. Crucially, the guidance and selection are post-hoc operations on the frozen model and leave the controllability results of section 5.3 unchanged, as they are measured on the same base model. Our CNeVA achieves this realism without closed-loop fine-tuning or autoregressive token-prediction objectives, while exposing a per-channel, reward-conditioned controllability interface that higher-realism imitation models lack. 5.3
Controllability
To quantify controllability, we compare steered CNeVA with the unconditional baseline using driftpaired channel-steering matrices (CSMs) on the WOMD validation split. For each channel k we feed λop n = ρek with the single-channel mask bn,j = I{j ̸= k} and compare against the null path under the same flow-time random seed (drift-paired evaluation). The CSM diagonal ∆Rk = Gsteered − Gbase k,k k measures per-channel response. All results use ρ = 1, w = 1.5, and context-residual returns unless stated otherwise. Per-channel response. Table 2 summarizes the evaluated four-channel CSM. Figure 7 in the appendix further contrasts these diagonals against the hard-eligibility ablation. The unconditional baseline produces near-zero diagonals on all channels, confirming that the null path does not steer by construction. CNeVA shows a clear hierarchy: the dense speed and acceleration channels respond strongest, followed by safety and map conformity, consistent with the tilt-vs-regression prediction in section 3.1 that dense kinematic penalties produce high signal-to-noise returns while sparse semantic penalties are prior-dominated. The w = 0 (latent-only) column isolates what the latent encodes without guidance. Specifically, the kinematic channels are already strongly steerable, while the sparse safety and map channels roughly double under w = 1.5. High-SNR channels are identifiable from the latent alone; sparse channels rely on classifier-free guidance. Results indicate that all four channels steer physically validly (section 5.4). 7
(a) Ground Truth
(b) Unconditional (λ = 0)
(c) Safety-Dominant
(d) Map-Dominant
(e) Speed-Dominant
(f) Accel-Dominant
Figure 3: CNeVA qualitative rollouts at ρ = 1, w = 1.5. Fidelity. CNeVA achieves a null-path minADE of 1.113 ± 0.011 m and an offroad rate of 32.5 ± 0.2% at 200K steps on the WOMD validation split and minADE = 1.80 m on the WOMD testing split. The off-road rate climbs over the rollout to roughly 1.3 to 1.9× the logged-data rate by 8 s (logged ≈ 17%), reflecting open-loop drift rather than a static gap. Under the identical local evaluation protocol, the hard-eligibility ablation reaches minADE = 1.112 ± 0.015 m (table 5), confirming that the soft eligibility gates do not degrade trajectory fidelity. Operating regime. All main results use the calibrated point ρ = 1, w = 1.5. At ρ = 5 all four channels stay monotonically positive with no sign-flip (∆Rspeed = +14.10, ∆Raccel = +11.73, ∆Rsafety = +0.85, ∆Rmap = +0.42; 5-seed means), but guardrails degrade (speed-steered stall +2.1 → +4.6 pp, retained speed 94.7% → 89.1%), so we report at the calibrated point. Qualitative results. Figure 3 shows rollouts under the four one-hot preferences against the groundtruth future and the null path. Safety and speed steering produce visible shifts matching the positive CSM diagonals. Acceleration yields smoother turning arcs, and the map panel is nearly indistinguishable from the null path (cf. section F.5). 5.4
Eligibility Gating and Reward Hacking
The CSM diagonals in table 2 measure return-space response but fail to verify whether the steered trajectories remain physically feasible. For instance, a model that maximizes the speed return by suppressing agent motion (stalling) achieves a large ∆Rspeed without producing useful behavior. We therefore complement the CSM with physical plausibility diagnostics: the stall fraction (proportion of agents with per-step displacement < 0.1 m), the mean speed of steered agents as a fraction of logged ground-truth speed, and the offroad rate at the 8 s horizon. We evaluate two ablation checkpoints alongside the main model to illustrate two failure modes that soft eligibility addresses. Reward hacking. Figure 4 (with table 5) reveals that the early-stage ablation’s speed CSM is inflated by reward hacking: it maximizes the speed return by suppressing motion (most steered agents stall) rather than by maintaining realistic velocities. CSM values must therefore be read jointly with the physical-plausibility guardrails, under which the speed response is physically valid. Safety erosion under hard eligibility. The hard-eligibility ablation adds future-query map conditioning and random mid-step anchoring while retaining hard binary gates. Its safety CSM drops to +0.21 ± 0.22, which is statistically indistinguishable from zero. Under hard gates, most agents 8
Speed channel response 50
reward hacking
80
40
30 physically honest
20
Percentage
Speed CSM ΔRspeed
Physical plausibility guardrails 100
60
Stall fraction Retained speed (v/vGT)
95%
76% 61%
65%
40
20
10
0
0 Early ablation (40K)
CNeVA (200K)
Early ablation (40K)
CNeVA (200K)
Figure 4: Reward-hacking contrast. Left: speed CSM is inflated 6× in the early checkpoint relative to the main model. Right: physical plausibility guardrails reveal that the early ablation achieves its CSM by stalling (76% stall, 61% of GT speed), while the main model retains 95% of GT speed. (clearance > 5 m, TTC > 6 s) receive no safety label, so the generator sees a near-zero safety signal. Replacing the thresholds with exponential decay (section 3.4) preserves the signal for nearthreshold agents and yields +0.66 ± 0.10. The two values are measured across different eligibility populations and are not directly comparable in absolute terms; the relevant observation is that soft eligibility yields a statistically significant positive response, whereas hard eligibility does not. Physical validity of safety steering. Under safety-channel steering, the main model’s stall fraction increases by only +0.9 percentage points relative to the unconditional baseline, and steered agents retain 97.9% of ground-truth speed. The speed reduction is consistent across all four channels, indicating that the context-residual labeling successfully decorrelates safety from speed. The safety CSM of +0.66 reflects genuine defensive-driving behavior (spacing, yielding) rather than a slowequals-safe confound.
6
Conclusion
We presented CNeVA, a controllable offline sim-agent framework that represents each traffic participant with a per-agent Gaussian behavior latent inferred in closed form from per-channel discounted returns. CNeVA conditions a mixed channel-mask flow-matching generator, supporting classifier-free guidance under null, partial, or fully specified operator preferences. Soft eligibility gates replace hard binary thresholds with smooth exponential decay, preserving gradient signal for near-threshold agents. Experiments on WOMD show that CNeVA reaches a WOSAC test-server realism of 0.71 (mid-spectrum; kinematic realism comparable to the leading imitation models, with collision/off-road and minADE limited by open-loop drift) while producing physically valid controllability: speed and acceleration respond monotonically without the stall-induced reward hacking observed in an early ablation, and safety steering is substantially above the hard-eligibility ablation while retaining near-ground-truth speed. Limitations. Map controllability is coordinate-specific: the context-residual return produces a positive response (+0.61), but the physical-offroad (−0.12) and lane-centerline (≈ 0) measures fail to produce controllable steering (section F.5). This sensitivity to the choice of reward definition exposes a structural limitation of the current framework. All four channels remain monotonically positive through ρ = 5; however, physical guardrails degrade at larger magnitudes. Therefore, the main results are reported at the calibrated operating point. Future work. Richer opportunity-aware reward decompositions that separate spatial from temporal map compliance, closed-loop training with contrastive objectives, and stronger latent-conditioning mechanisms beyond one-hot steering are natural next steps. 9
Acknowledgments The authors appreciate the Google TPU Research Cloud (TRC) for supporting access to TPUs.
References Yulong Cao, Boris Ivanovic, Chaowei Xiao, and Marco Pavone. Reinforcement learning with human feedback for realistic traffic simulation. In 2024 IEEE International Conference on Robotics and Automation (ICRA), pages 14428–14434, May 2024. doi: 10.1109/ICRA57147.2024.10610878. Wei-Jer Chang, Chen Tang, Chenran Li, Yeping Hu, Masayoshi Tomizuka, and Wei Zhan. Editing driver character: Socially-controllable behavior generation for interactive traffic simulation. IEEE Robotics and Automation Letters, 8(9):5432–5439, September 2023. doi: 10.1109/LRA.2023.3291897. Wei-Jer Chang, Francesco Pittaluga, Masayoshi Tomizuka, Wei Zhan, and Manmohan Chandraker. SAFE-SIM: Safety-critical closed-loop traffic simulation with diffusion-controllable adversaries. In European Conference on Computer Vision (ECCV), 2024. Wei-Jer Chang, Akshay Rangesh, Kevin Joseph, Matthew Strong, Masayoshi Tomizuka, Yihan Hu, and Wei Zhan. SPACeR: Self-play anchoring with centralized reference models. In International Conference on Learning Representations (ICLR), 2026. Daphne Cornelisse and Eugene Vinitsky. Human-compatible driving partners through dataregularized self-play reinforcement learning, 2024. Marco Cusumano-Towner, David Hafner, Alexander Hertzberg, Brody Huval, Aleksei Petrenko, Eugene Vinitsky, Erik Wijmans, Taylor Killian, Stuart Bowers, Ozan Sener, Philipp Kraehenbuehl, and Vladlen Koltun. Robust autonomy emerges from self-play. In Proceedings of the 42nd International Conference on Machine Learning, volume 267 of PMLR, pages 11710–11737, 2025. Scott Ettinger, Shuyang Cheng, Benjamin Caine, Chenxi Liu, Hang Zhao, Sabeek Pradhan, Yuning Chai, Benjamin Sapp, Charles R Qi, Yin Zhou, et al. Large scale interactive motion forecasting for autonomous driving: The Waymo open motion dataset. In IEEE/CVF International Conference on Computer Vision (ICCV), 2021. Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance, 2022. Yihan Hu, Siqi Chai, Zhening Yang, Jingyu Qian, Kun Li, Wenxin Shao, Haichao Zhang, Wei Xu, and Qiang Liu. Solving motion planning tasks with a scalable generative model. In Aleš Leonardis, Elisa Ricci, Stefan Roth, Olga Russakovsky, Torsten Sattler, and Gül Varol, editors, Computer Vision – ECCV 2024, pages 386–404, Cham, 2025. Springer Nature Switzerland. Zhiyu Huang, Zixu Zhang, Ameya Vaidya, Yuxiao Chen, Jaime Fernández Fisac, and Chen Lv. Versatile behavior diffusion for generalized traffic agent simulation. IEEE Transactions on Intelligent Transportation Systems, pages 1–17, 2026. Chiyu Max Jiang, Andre Cornman, Cheolho Park, Benjamin Sapp, Yin Zhou, and Dragomir Anguelov. MotionDiffuser: Controllable multi-agent motion prediction using diffusion. In IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2023. Chiyu Max Jiang, Yijing Bai, Andre Cornman, Christopher Davis, Xiukun Huang, Hong Jeon, Sakshum Kulshrestha, John Lambert, Shuangyu Li, Xuanyu Zhou, Carlos Fuertes, Chang Yuan, Mingxing Tan, Yin Zhou, and Dragomir Anguelov. Scenediffuser: Efficient and controllable driving simulation initialization and rollout. In Advances in Neural Information Processing Systems, volume 37, pages 55729–55760, 2024. Michael Irwin Jordan. Learning in graphical models. MIT press, 1999. Yaron Lipman, Ricky T Q Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. In International Conference on Learning Representations (ICLR), 2023. 10
Zihao Liu, Xing Liu, Yizhai Zhang, and Panfeng Huang. Stage: Style-controllable action generation for personalized autonomous driving. IEEE Robotics and Automation Letters, 11(2):2130–2137, February 2026. doi: 10.1109/LRA.2025.3640974. Nico Montali, John Lambert, Paul Mougin, Alex Kuefler, Nicholas Rhinehart, Michelle Li, Cole Gulino, Tristan Emrich, Zoey Yang, Shimon Whiteson, et al. The Waymo open sim agents challenge. In Advances in Neural Information Processing Systems (NeurIPS) Datasets and Benchmarks Track, 2023. Nigamaa Nayakanti, Rami Al-Rfou, Aurick Zhou, Kratarth Goel, Khaled S Refaat, and Benjamin Sapp. Wayformer: Motion forecasting via simple & efficient attention networks. In IEEE International Conference on Robotics and Automation (ICRA), 2023. Jiquan Ngiam, Benjamin Caine, Vijay Vasudevan, Zhengdong Zhang, Hao-Tien Lewis Chiang, Jeffrey Ling, Rebecca Roelofs, Alex Bewley, Chenxi Liu, Ashish Venugopal, David Weiss, Ben Sapp, Zhifeng Chen, and Jonathon Shlens. Scene transformer: A unified architecture for predicting multiple agent trajectories, 2022. Muleilan Pei, Shaoshuai Shi, and Shaojie Shen. Advancing multi-agent traffic simulation via r1style reinforcement fine-tuning, 2026. Jonah Philion, Xue Bin Peng, and Sanja Fidler. Trajeglish: Traffic modeling as next-token prediction. In International Conference on Learning Representations (ICLR), 2024. Ethan Pronovost, Meghana Reddy Ganesina, Noureldin Hendy, Zeyu Wang, Andres Morales, Kai Wang, and Nick Roy. Scenario diffusion: Controllable driving scenario generation with diffusion. In Advances in Neural Information Processing Systems (NeurIPS), 2023. Davis Rempe, Jonah Philion, Leonidas J. Guibas, Sanja Fidler, and Or Litany. Generating useful accident-prone driving scenarios via a learned traffic prior. In IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2022. Luke Rowe, Roger Girgis, Anthony Gosselin, Bruno Carrez, Florian Golemo, Felix Heide, Liam Paull, and Christopher Pal. CtRL-Sim: Reactive and controllable driving agents with offline reinforcement learning. In Conference on Robot Learning (CoRL), 2024. Luke Rowe, Roger Girgis, Anthony Gosselin, Liam Paull, Christopher Pal, and Felix Heide. Scenario dreamer: Vectorized latent diffusion for generating driving simulation environments. In IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2025. Ari Seff, Brian Cera, Dian Chen, Mason Ng, Aurick Zhou, Nigamaa Nayakanti, Khaled S Refaat, Rami Al-Rfou, and Benjamin Sapp. MotionLM: Multi-agent motion forecasting as language modeling. In IEEE/CVF International Conference on Computer Vision (ICCV), 2023. Shaoshuai Shi, Li Jiang, Dengxin Dai, and Bernt Schiele. Mtr++: Multi-agent motion prediction with symmetric scene modeling and guided intention querying. IEEE Transactions on Pattern Analysis and Machine Intelligence, 46(5):3955–3971, May 2024. doi: 10.1109/TPAMI.2024.3352811. Simon Suo, Sebastian Regalado, Sergio Casas, and Raquel Urtasun. TrafficSim: Learning to simulate realistic multi-agent behaviors. In IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2021. Shuhan Tan, John Lambert, Hong Jeon, Sakshum Kulshrestha, Yijing Bai, Jing Luo, Dragomir Anguelov, Mingxing Tan, and Chiyu Max Jiang. SceneDiffuser++: City-scale traffic simulation via a generative world model. In IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2025a. Tianyi Tan, Yinan Zheng, Ruiming Liang, Zexu Wang, Kexin ZHENG, Jinliang Zheng, Jianxiong Li, Xianyuan Zhan, and Jingjing Liu. Flow matching-based autonomous driving planning with advanced interactive behavior modeling. In Advances in Neural Information Processing Systems (NeurIPS), 2025b. 11
Wei Wu, Xiaoxin Feng, Ziyan Gao, and Yuheng Kan. Smart: Scalable multi-agent real-time motion generation via next-token prediction. In Advances in Neural Information Processing Systems, volume 37, pages 114048–114071, 2024. Zebin Xing, Xingyu Zhang, Yang Hu, Bo Jiang, Tong He, Qian Zhang, Xiaoxiao Long, and Wei Yin. GoalFlow: Goal-driven flow matching for multimodal trajectories generation in end-to-end autonomous driving. In IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2025. Zhao-Heng Yin, Lingfeng Sun, Liting Sun, Masayoshi Tomizuka, and Wei Zhan. Diverse critical interaction generation for planning and planner evaluation. In 2021 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pages 7036–7043, September 2021. doi: 10.1109/IROS51168.2021.9636266. Qichao Zhang, Yinfeng Gao, Yikang Zhang, Youtian Guo, Dawei Ding, Yunpeng Wang, Peng Sun, and Dongbin Zhao. Trajgen: Generating realistic and diverse trajectories with reactive and feasible agent behaviors for autonomous driving. IEEE Transactions on Intelligent Transportation Systems, 23(12):24474–24487, December 2022. ISSN 1558-0016. doi: 10.1109/TITS.2022.3202185. Zhejun Zhang, Alexander Liniger, Dengxin Dai, Fisher Yu, and Luc Van Gool. TrafficBots: Towards world models for autonomous driving simulation and motion prediction. In IEEE International Conference on Robotics and Automation (ICRA), 2023. Zhejun Zhang, Peter Karkus, Maximilian Igl, Wenhao Ding, Yuxiao Chen, Boris Ivanovic, and Marco Pavone. Closed-loop supervised fine-tuning of tokenized traffic models. In IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2025a. Zhiyuan Zhang, Xiaosong Jia, Guanyu Chen, Qifeng Li, and Junchi Yan. TrajTok: Technical report for 2025 Waymo open sim agents challenge. Technical report, Shanghai Jiao Tong University, 2025b. Ziyuan Zhong, Davis Rempe, Yuxiao Chen, Boris Ivanovic, Yulong Cao, Danfei Xu, Marco Pavone, and Baishakhi Ray. Language-guided traffic simulation via scene-level diffusion. In Conference on Robot Learning (CoRL), 2023a. Ziyuan Zhong, Davis Rempe, Danfei Xu, Yuxiao Chen, Sushant Veer, Tong Che, Baishakhi Ray, and Marco Pavone. Guided conditional diffusion for controllable traffic simulation. In 2023 IEEE International Conference on Robotics and Automation (ICRA), pages 3560–3566, 2023b. Zikang Zhou, Haibo Hu, Xinhong Chen, Jianping Wang, Nan Guan, Kui Wu, Yung-Hui Li, Yu-Kai Huang, and Chun Jason Xue. BehaviorGPT: Smart agent simulation for autonomous driving with next-patch prediction. In Advances in Neural Information Processing Systems (NeurIPS), 2024.
12
Supplementary Material A Notations
14
B Derivation of the Conjugate Gaussian Posterior
14
B.1 Completing the Square . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
14
B.2 Sum-Product Belief Propagation on the Simplified Factor Graph . . . . . . . . . .
14
C Extended Method Details
16
C.1 Context-Residual Return Labeling . . . . . . . . . . . . . . . . . . . . . . . . . .
17
C.2 Lane-Centerline Map Reward . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
17
C.3 Soft Eligibility Gates: Decay Formulas . . . . . . . . . . . . . . . . . . . . . . . .
17
C.4 Contrastive Conditioning . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
17
D Model Architecture and Training Details
18
D.1 Per-Channel Return Standardization . . . . . . . . . . . . . . . . . . . . . . . . .
18
D.2 Reward-Channel Construction . . . . . . . . . . . . . . . . . . . . . . . . . . . .
18
D.3 Lambda-Conditional Rectified-Flow Decoder . . . . . . . . . . . . . . . . . . . .
18
D.4 Mixed Channel-Mask CFG Curriculum . . . . . . . . . . . . . . . . . . . . . . .
19
D.5 Rectified-Flow Euler Sampler with Classifier-Free Guidance . . . . . . . . . . . .
20
E Pairwise Trajectory Divergence
21
F Supplementary Diagnostics
21
F.1
Interpretation and Limitations . . . . . . . . . . . . . . . . . . . . . . . . . . . .
21
F.2
Conditioning Ablation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
21
F.3
Multi-Seed CSM . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
22
F.4
Physical Plausibility Diagnostics . . . . . . . . . . . . . . . . . . . . . . . . . . .
22
F.5
Map Controllability Across Reward Measures . . . . . . . . . . . . . . . . . . . .
22
13
A
Notations
For reference, table 3 lists all notation used in this paper.
B
Derivation of the Conjugate Gaussian Posterior
This appendix derives the analytic posterior of Theorem 3.1 from two perspectives. Section B.1 gives the textbook completing-the-square argument used in the main text. Section B.2 re-derives the same posterior via sum-product belief propagation on the PGM in Figure 1b, which makes the conditional-independence assumptions and the role of the discount factor γ explicit. Variational identity. The objective in Eq. (4) is the KL projection onto the reward-tilted posterior: for any density q, the KL to the tilted posterior equals the KL to the Gaussian prior minus Eq [λTn Gn ], plus a log-normalizer independent of q. For a positive-definite Gaussian prior, the log-normalizer is finite and equal to µT0 Gn + 12 GTn Σ0 Gn . Thus, the maximizer in Eq. (4) is the same distribution recovered by completing the square below. Interpretation. Since the reward-tilted factor is linear in lambda, the conjugate update shifts the posterior mean but leaves the covariance fixed at the prior covariance. The residual uncertainty should therefore be read as population-level label noise, not trajectoryspecific epistemic uncertainty; CNeVA uses lambda as a behavior-profile coordinate induced by the chosen reward basis, not as an identified inverse-RL utility. B.1
Completing the Square
Combining the Gaussian prior p0 (λn ) = N (µ0 , Σ0 ) with the trajectory-level un-normalized factor ψ(τn , λn ) = exp(λ⊤ n Gn ) from equation 2, the un-normalized log-posterior writes ⊤ log p(λn | o1:T,n ) = − 21 λn − µ0 Σ−1 λn − µ0 + λ⊤ (13) n Gn + const. 0 Expanding the quadratic and collecting terms in λn yields a quadratic with precision matrix Σ−1 0 (unchanged because the likelihood is linear in λn ) and linear term Σ−1 µ + G . Completing the 0 n 0 square recovers a Gaussian with mean µ⋆n = Σ0 Σ−1 (14) 0 µ 0 + G n = µ 0 + Σ 0 Gn and covariance Σ0 , as claimed. B.2
Sum-Product Belief Propagation on the Simplified Factor Graph
The conjugate posterior is also recovered exactly by sum-product message passing on the treestructured per-agent factor subgraph induced by the relaxed joint equation 3; the simplified factor graph is shown in Figure 5. The per-step un-normalized factors ψt of equation 1 attach directly to the trajectory–latent pair; the geometric time-discount γ t−1 is folded into the factor itself, so no separate “discounted optimality” assumption is needed beyond the definition equation 1. Figure 5 makes the three factor families that remain after relaxation explicit: the autoregressive emission factors fo,t (ot,n | o<t,n , m), the un-normalized optimality factors {ψt }Tt=1 , and the behaviorlatent prior fλ (λn ) = p0 (λn ). Each filled square is one factor of the unnormalized joint equation 3, and each circle is a random variable (shaded if observed). Because the latent state s1:T and its transitions are absent, the per-agent optimality chains decouple immediately and sum-product is exact on each agent’s subgraph. Message from a single optimality factor. Because the observation ot,n is observed, the incoming message from the ot,n node to ψt is the Dirac indicator δ(ot,n −oobs t,n ). The outbound message that ψt sends to λn is the factor evaluated at the logged observation, t−1 ⊤ mψt →λn λn = ψt (oobs λn r(oobs (15) t,n , λn ) = exp γ t,n ) . The factor ψt is un-normalized on the λn –ot,n pair, so no additional summation step is introduced at the factor node. Product across time. The T outbound messages meet at the λn variable node, where sum-product multiplies them: T T Y X ⊤ ⊤ mψt →λn λn = exp λn γ t−1 r(oobs (16) t,n ) = exp λn Gn , t=1
t=1
14
Table 3: Table of Notations Notation
Description
Scalars d N T Tf K γ s
Dimension of motion state observations. Number of agents in a scenario. Total number of time steps in a scenario. Number of future time steps predicted per CFM rollout. Number of reward channels in the semantic decomposition. Discount factor for the per-channel return; γ ∈ (0, 1]. Rectified-flow time; s ∈ [0, 1].
Hidden-Markov Process and Sim-Agent Scenario st Latent world state at time step t (random variable). ot,n Observation of the n-th agent at time step t (random variable). m Static map context (random variable; observed). ot , ot,n Observation / observation of the n-th agent at time t (vector value). st Latent state at time t (vector value). m Static map context (vector value). τ Trajectory sequence [o1 , . . . , oT ]. S, O Latent state / observation space. H Hidden-Markov Process tuple over latent state S and observation space O. Training dataset of logged scenarios. D Reward and Behavior Latent K-channel per-step reward decomposition. r(ot,n ) Gn Per-channel discounted observed return of agent n. Per-agent channel mask, bn ∈ {0, 1}K ; bn,j = 1 marks channel j as unobservable. bn λn Per-agent behavior profile latent; λn ∈ RK . µ0 Mean of the Gaussian prior on λn . Σ0 Covariance of the Gaussian prior on λn . 2 σ0,k Per-channel prior variance (k-th diagonal entry of Σ0 ). q ⋆ (λn ) Analytic conjugate Gaussian posterior. µ⋆n Posterior mean µ0 + Σ0 Gn . τc Clearance decay scale for the safety soft eligibility gate (metres). τt Time-to-collision decay scale for the safety soft eligibility gate (seconds). τm Margin decay scale for the map soft eligibility gate (metres). Posterior Sampling ηn Standard-normal noise used in the reparameterized sample. λ-Conditional CFM Decoder vθ λ-conditional rectified-flow velocity field. xs Noisy rectified-flow latent at time s. yn Clean future displacement target for agent n, RTf ×d . ⋆ vn Rectified-flow velocity target ϵ − yn . ϵ Standard-normal noise sample for the rectified-flow forward process. dh Hidden dimension of the scene-decoder backbone. θ Parameters of the λ-conditional CFM decoder. ∆s Euler integrator step in flow time. Probability and Loss N (µ, Σ) Multivariate Gaussian distribution with mean µ and covariance Σ. U(0, 1) Uniform distribution on [0, 1]. Expectation operator. E DKL [ · ∥ · ] Kullback–Leibler divergence between two distributions. L(θ) Classifier-free CFM training loss.
15
m
o1,n
...
o2,n fo,2
fu,1
fo,3 fu,2
u1,n
u2,n
oT,n fo,T fu,T
uT,n
fλ
λn N
Figure 5: Simplified factor graph corresponding to the relaxed joint equation 3: the latent state s1:T and its transition kernel are marginalised away (section 2), leaving an autoregressive observation chain conditioned on the static map m. Filled squares are factors, circles are random variables (shaded if observed). The three factor families are: the autoregressive emission factors fo,t (ot,n | o<t,n , m) (modelled by the λ-conditional CFM), the time-discounted optimality factors fu,t (section B.2), and the behavior-latent prior fλ . Sum-product belief propagation on the per-agent subgraph rooted at λn recovers the analytic conjugate posterior of Theorem 3.1.
which reproduces equation 2 of the main text. The geometric discount enters the BP product exclusively through the per-step factor equation 1; if insteadPγ = 1 were chosen (time-homogeneous factor), the same algebra yields the undiscounted sum λ⊤ n t r(ot,n ). Combination with the prior factor. The variable node λn also receives a single inbound message from the behavior-prior factor fλ (λn ) = p0 (λn ) = N (µ0 , Σ0 ). The belief at λn is the product of all inbound messages: b(λn ) ∝ fλ (λn )
T Y
⊤ mψt →λn λn ∝ exp − 12 (λn − µ0 )⊤ Σ−1 0 (λn − µ0 ) + λn Gn . (17)
t=1
Completing the square (as in section B.1) gives b(λn ) = N µ0 + Σ0 Gn , Σ0 = q ⋆ (λn ),
(18)
which is the analytic posterior of Theorem 3.1. Because the subgraph rooted at λn , comprising fλ and the T factors {ψt }Tt=1 , is a tree, sum-product is exact and the belief b(λn ) coincides with the true marginal posterior. The remaining factors of the joint (initial-state prior, transition, observation) are independent of λn given the observed trajectory, so their contributions appear only in the normalization constant. Because the latent-state marginalization in equation 3 has already removed s1:T from the joint, the per-agent optimality chains in Figure 5 are independent across n, and sum-product is exact agentby-agent.
C
Extended Method Details
This appendix collects the return-labeling and eligibility details deferred from section 3: the contextresidual and lane-centerline return definitions, the soft-eligibility decay formulas, and the contrastive conditioning augmentation. 16
C.1
Context-Residual Return Labeling
The raw per-channel return Gn,k conflates agent behavior with scenario difficulty: a highway scene yields low offroad penalties regardless of driving style. We define the context-residual return as the per-channel discounted return residualized against the scene context, Gcr n,k ≜ Gn,k − Ḡk (mn ),
(19)
where Ḡk (mn ) is the mean return for agents sharing the same map context mn , estimated once on the calibration split. After standardization via equation 7, Gcr n,k isolates the behavioral signal from the structural component, producing posterior labels that are more informative for conditioning. The context-residual formulation is the primary return measure used for all CNeVA training and evaluation in section 5. C.2
Lane-Centerline Map Reward
As an alternative to the composite offroad-plus-road-edge penalty used in the context-residual return, we define a lane-centerline return based on the signed lateral offset to the nearest lane centerline: lc lane rt,n ≜ − d⊥ t,n / wt,n ,
(20)
where d⊥ t,n is the perpendicular distance from agent n at time t to the nearest lane centerline and lane wt,n is the corresponding lane half-width. This measure provides a purely geometric, coordinatespecific definition of map compliance. Below, we evaluate the map CSM diagonal under this measure alongside the context-residual and physical-offroad alternatives. C.3
Soft Eligibility Gates: Decay Formulas
Under the hard-gated baseline, safety labels are computed only for agents whose minimum clearance cn < 5 m or minimum time-to-collision ttcn < 6 s; map labels only for agents within a road-boundary margin or off-road. This binary gating introduces a discontinuity at the threshold and excludes the majority of agents from safety/map supervision. We replace the hard gates with smooth exponential decay applied to the per-step pairwise risk. For safety, each pairwise collision risk term is modulated by both clearance and TTC: max(c , 0) max(ttc , 0) t,ij t,ij safe safe r̃t,ij ≜ rt,ij · exp − · exp − , (21) τc τt with decay scales τc = 2.0 m and τt = 3.0 s. The product form ensures that both clearance and TTC must be large for the risk to vanish: a distant but fast-approaching agent (low TTC, high clearance) retains a high risk weight. For map compliance, the per-step risk is modulated by the signed margin mt,n to the nearest road boundary: max(m , 0) t,n map map r̃t,n ≜ rt,n · exp − , (22) τm with τm = 1.0 m. Under soft eligibility, all valid agents receive labels (the eligibility condition reduces to validity alone), but agents far from hazards contribute negligibly to the discounted return because the decay drives the per-step reward toward zero. Still, agents C.4
Contrastive Conditioning
The mixed channel-mask curriculum of section 3.2 exposes the unconditional and conditional paths during training, but does not explicitly contrast them within the same batch. We augment training with a contrastive conditioning step: for each agent n, the loss includes both the steered forward pass (with λn ∼ q ⋆ ) and the null forward pass (with e∅ ), computed over the same flow-time noise realization. This paired evaluation encourages the velocity field to learn the difference between steered and unsteered trajectories rather than their absolute positions, thereby improving the robustness of the conditioning signal during closed-loop rollout, where distribution shift would otherwise attenuate the guidance effect. The drift-paired CSM evaluation protocol of section 5.3 mirrors this training-time pairing at test time. 17
D
Model Architecture and Training Details
This appendix collects the parametric and operational details deferred from sections 3.1 to 3.3: the calibration of the prior covariance Σ0 , the λ-conditional rectified-flow decoder architecture, the classifier-free mixed-mask curriculum, and the deployment-time sampler. Training uses equation 11 on a v4-32 slice with batch 4 per device, AdamW (Loshchilov and Hutter, 2019) at peak learning rate 10−4 , 2000 warm-up steps, and weight decay 10−2 . D.1
Per-Channel Return Standardization
The raw per-channel discounted return Gn on WOMD demonstrations has substantially different scales across the four channels: the safety channel mean is around −37, map around −40, speed around −50, and accel only around −3; per-channel standard deviations span 5.2 (accel) to 33.6 (map). Feeding Gn verbatim into the conjugate update of equation 5 with Σ0 = I would let the speed and map channels dominate the per-agent λn by an order of magnitude, making the acceleration channel difficult for the conditional generator to use. We therefore standardize per channel using the empirical statistics measured on a fixed WOMD calibration split with γ = 0.99 and Tf = 80: µG ≈ −36.84, −39.81, −50.06, −3.03 , σG ≈ 19.40, 33.63, 12.93, 5.23 , (23) e n,k = (Gn,k − µG,k )/σG,k in channel order (safety, map, speed, accel). The standardized return G replaces Gn,k inside equation 5 and equation 6. With this preconditioning the prior covariance reduces to the only free hyperparameter; we use µ0 = 0 and Σ0 = I. The empirical histograms underlying equation 23 are shown in Figure 6; the bimodal shape of the map channel (with a heavy left tail from extended off-road episodes) motivated keeping the constituent offroad and dist-to-roadedge terms together rather than separating them into two reward channels.
Figure 6: Empirical histogram of the per-channel discounted return Gn,k over the fixed WOMD calibration split (N=35,774 simulated agents, γ = 0.99, Tf = 80). Per-channel mean and standard deviation are annotated; the dashed red line marks the mean. The four channels differ in scale by up to an order of magnitude (speed mean ≈ −50, accel mean ≈ −3), motivating the per-channel standardization of equation 23. The map channel is visibly bimodal: most agents stay on-map, while a tail of extended off-road episodes drags the mean down. D.2
Reward-Channel Construction
The reward basis used for posterior labeling is implemented as four sign-aligned penalty sums. The safety channel is collision plus distance-to-nearest-object plus time-to-collision shaping; the map channel is offroad plus distance-to-road-edge; the speed channel is linear plus angular speed; and the accel channel is linear plus angular acceleration. Each constituent is shaped into [−1, 0] before discounting, and a backward scan accumulates the per-channel discounted return Gn,k . Agents without a valid future horizon are masked out of the CFM loss, return-standardization statistics, and steering metrics; their reward tensors are zero-filled only after this masking step, so invalid agents do not contribute best-possible zero reward labels. Positive steering of channel k therefore means reducing the corresponding penalty mass. D.3
Lambda-Conditional Rectified-Flow Decoder
The velocity field vθ (xs , s, o<t,n , m, λn ) is a Transformer decoder that consumes (i) the noisy displacement latent xs , (ii) the flow time s embedded by Fourier features (sampled as sigmoid of 18
a standard normal in implementation), (iii) the cross-attention conditioning set, and produces the predicted velocity vn⋆ . The conditioning set condn ≜
DenseK→dh (λn ); embed(o<t,n ); embed(m)
(24)
folds the per-agent latent into an extra token alongside the road-graph and traffic-light embeddings, in the same way per-agent map context is consumed by BehaviorGPT (Zhou et al., 2024). Each decoder block stacks (a) a temporally causal self-attention over the future axis, (b) an unmasked spatial self-attention across agents, (c) a cross-attention to condn , and (d) a feed-forward sub-block, each preceded by an adaptive layer-norm whose scale and shift parameters are predicted from the flow-time embedding. For masked conditioning, the implementation concatenates the value-zeroed lambda vector with the binary channel mask before projection, giving a 2K-dimensional token input; Eq. (20) shows the unmasked special case.
D.4
Mixed Channel-Mask CFG Curriculum
Why not classifier guidance? A naive alternative would be to train an unconditional velocity vθun (xs , s, · ) together with a separate regressor pϕ (G | xs , s) on noised trajectories and steer sampling by ∇xs log pϕ , transposing the classifier guidance recipe of Dhariwal and Nichol (2021) from class labels to continuous returns. Two practical issues carry over. First, pϕ must be trained on the same diffused-trajectory distribution as the decoder, and its supervisory signal at large s, where xs is nearly isotropic Gaussian noise, is essentially uninformative. Meanwhile, stepping along ∇xs log pϕ in a continuous, high-dimensional λ-space resembles a gradient-based adversarial attack on pϕ (Ho and Salimans, 2022): high-confidence directions of the regressor are not the same as directions of high data density. Why a curriculum? The vanilla classifier-free guidance recipe (Ho and Salimans, 2022) flips between fully-conditional (λ = λn ) and fully-unconditional (λ = e∅ ) updates with a single Bernoulli mask. At inference, the operator-supplied one-hot probe λ = ρek , with only channel k informative and the remaining K −1 at value zero, is then strictly outside the training support. As a result, the network has only ever seen dense λ drawn from q ⋆ . We therefore expose every subset of channels during training by drawing un ∼ U(0, 1) and setting 1K 0 ≤ un < 0.2 (null path; learned null token wins) Mask1 (K) 0.2 ≤ un < 0.6 (exactly one channel kept) bn ∼ Mask2 (K) 0.6 ≤ un < 0.8 (exactly two channels kept) 0K 0.8 ≤ un ≤ 1.0 (all channels kept),
(25)
where bn,k = 1 marks channel k as masked-out and bn,k = 0 marks it as kept, and Maskj (K) is the uniform distribution over binary masks with exactly j kept entries. The one-hot probe at channel k now matches the single-channel training distribution exactly (mask everywhere except k). Mask indicator concat. For the projection to distinguish “λk = 0 (observed)” from “λk unobserved,” we concatenate the value-zeroed profile with the mask indicator before the R2K → Rdh Dense projection: e n ≜ (1 − bn ) ⊙ λn ; bn ∈ R2K . λ
(26)
When bn = 1K (the 20% null branch) we additionally override the projection output with the learned null embedding e∅ ∈ Rdh , recovering the unconditional path of Ho–Salimans. The conditional and unconditional velocity heads share all backbone parameters except e∅ . The training loop is summarized in Algorithm 1. 19
Algorithm 1 Training Require: Dataset D. Require: Reward channel count K; mask probabilities (p0 , p1 , p2 , pK ) = (0.20, 0.40, 0.20, 0.20). Require: Standardisation statistics µG , σG ; prior (µ0 , Σ0 ). 1: repeat 2: Sample a scenario (o1:T,1:N , m) ∼ D. 3: for all agent n = 1, . . . , N do e n,k = (Gn,k − µG,k )/σG,k for each k. 4: Compute G ▷ standardised return e n + Σ1/2 ηn . 5: Sample ηn ∼ N (0, I) and set λn = µ0 + Σ0 G 0 6: Sample un ∼ U(0, 1) and choose j ∈ {0, 1, 2, K} with probabilities (p0 , p1 , p2 , pK ). 7: Sample a uniformly random size-j subset Kn ⊆ {1, . . . , K} 8: for all k = 1, . . . , K do 9: Set bn,k = I{k ̸∈ Kn }. ▷ channel mask 10: end for e n = [(1 − bn ) ⊙ λn ; bn ] ∈ R2K ; if j = 0 flag agent n as null. 11: Form λ 12: Sample s as sigmoid of a standard normal and sample epsilon from a standard normal. 13: Form xs = (1 − s)yn + sϵ and the rectified-flow target vn⋆ = ϵ − yn . 14: end for e n by e∅ for every agent flagged null. 15: Inside vθ , replace the Dense projection of λ P 2 e 16: θ ← θ − α ∇θ n vθ (xs , s, o<t,n , m, λn ) − vn⋆ . 17: until converged
Score-form correspondence. The CFG combination equation 12 is the rectified-flow analogue of the noise-prediction form ϵe = (1 + w)ϵθ (z, c) − wϵθ (z, e∅ ) used by Ho and Salimans (2022): the rectified-flow velocity is an affine reparameterization of the score ∇xs log pθ (xs | λ) under the linear noise schedule of equation 9 (Lipman et al., 2023), so the linear extrapolation in velocity space inherits the implicit-classifier interpretation ∇xs log pi (λ | xs ) = ∇xs log p(xs | λ)−∇xs log p(xs ) of Ho and Salimans (2022); the resulting sample is approximately drawn from p(xs | λ) pi (λ | xs )w . D.5
Rectified-Flow Euler Sampler with Classifier-Free Guidance
Deployment recovers a trajectory from x1 ∼ N (0, I) by Euler-integrating the rectified-flow ODE in reverse time, with the bare conditional velocity replaced by the CFG combination of equation 12: eθw (xs , λop xs−∆s ← xs − ∆s v n ),
s ∈ {1, 1−∆s, . . . , ∆s},
(27)
eθw is evaluated by two forward passes of vθ with conditioning λop where v n and e∅ , respectively. Activating the learned null token for every agent recovers the unconditional population marginal; any other choice of (λop n , w) steers the trajectory along the requested preference direction without modifying the integrator or the noise schedule. The receding-horizon variant of section 3.3 simply re-evaluates equation 27 every ℓ output timesteps with a refreshed conditioning history. Algorithm 2 states the full sampler. Algorithm 2 One-Shot Sampling Require: Operator preference λop n . Require: Guidance scale w ≥ 0; integrator steps Tint . Require: History o<t,n and map m. 1: Sample x1 ∼ N (0, I); let ∆s = 1/Tint . 2: for s = 1, 1−∆s, . . . , ∆s do 3: v c ← vθ (xs , s, o<t,n , m, λop n ) 4: v u ← vθ (xs , s, o<t,n , m, e∅ ) e ← (1+w) v c − w v u 5: v e 6: xs−∆s ← xs − ∆s v 7: end for 8: return x0
▷ conditional pass ▷ unconditional pass ▷ CFG combination
20
The receding-horizon deployment regime of section 3.3 wraps algorithm 2 in a fixed-stride patch loop: every ℓ output timesteps, the model re-conditions on a history buffer that has been updated with its own previous predictions, so the CFG double forward is repeated within each patch and the conditioning is refreshed at every patch boundary. Algorithm 3 Receding-Horizon Patch Sampling Require: Operator preference λop n . Require: Guidance scale w ≥ 0 and integrator steps Tint Require: History o<1,n ; map m; horizon T ; patch length ℓ. 1: Initialise the history buffer ohist ← o<1,n . 2: for tstart = 1, 1+ℓ, 1+2ℓ, . . . , T −ℓ+1 do 3: Run algorithm 2 with conditioning ohist to obtain the next patch of ℓ displacement predictions ŷtstart :tstart +ℓ−1 . 4: Integrate displacements into per-step observations ôtstart :tstart +ℓ−1 via ôt = ôt−1 + ŷt . 5: Append ôtstart :tstart +ℓ−1 to ohist . 6: end for 7: return concatenated ô1:T
Setting ℓ = T recovers the single-pass open-loop sampler; setting ℓ = 1 corresponds to per-step replan. We use ℓ = 16 throughout, matching the training-time patch length used by the CFM loss in equation 11.
E
Pairwise Trajectory Divergence
Table 4 expands the pairwise trajectory-divergence summary from the main text. It confirms that speed and acceleration induce the largest geometric separation, while the safety-map pair remains near the noise floor, matching the weak safety/map CSM response.
F
Supplementary Diagnostics
This section provides additional multi-seed CSM diagnostics for the model in section 5, thereby expanding the aggregated results in table 2. These diagnostics illustrate the controllability hierarchy discussed in section 5. F.1
Interpretation and Limitations
These diagnostics support the limitation stated in the main text: CNeVA’s controllability depends strongly on the chosen reward basis and the identifiability of behavior labels under reconstruction training. Dense kinematic channels such as speed and acceleration yield consistent geometric responses. In contrast, safety and map compliance are sparse, context-dependent, and harder to capture with one-hot λ steering. The main model addresses the safety aspect of this limitation through soft eligibility gates (section 3.4), which yield a statistically significant safety CSM, whereas the hard-eligibility ablation does not (table 5). Map controllability remains a structural limitation of the current reward basis. F.2
Conditioning Ablation
As a minimal ablation on the model, we compare the steered CNeVA path against the unconditional path (λ = e∅ ) under identical flow-time noise; the CSM diagonals reported in table 2 are the result. Removing the conditioning entirely (i.e., zeroing λ and the mask indicator) recovers the unconditional baseline, confirming that the steering effect is attributable to the behavior latent and not to the noise realization. 21
Table 4: Pairwise trajectory divergence D(ea , eb ) at ρ = 1, w = 1.5 (m). Paired flow-time noise; single-channel-kept mask. Mean over pairs D̄ = 0.54 m. safety
map
speed
accel
–
0.10 –
0.63 0.62 –
0.49 0.47 0.96 –
safety map speed accel
Hard-elig. ablation
CNeVA (soft elig.)
Semantic channels
Kinematic channels 10
0.6
CSM diagonal ΔRk
CSM diagonal ΔRk
0.7
0.5 0.4 0.3 0.2 0.1
8
6
4
2
0.0 0 Safety
Map
Speed
Accel
Figure 7: CSM diagonal ∆Rk for CNeVA versus the hard-eligibility ablation at ρ = 1, w = 1.5 (openloop, context-residual return). Left: semantic channels (safety, map). Right: kinematic channels (speed, accel). Error bars show ±1 std over 5 seeds.
F.3
Multi-Seed CSM
Table 2 in the main text reports the CSM diagonals aggregated over 5 evaluation seeds (s = 42– 46). The per-seed standard deviations are small relative to the between-channel differences for the kinematic channels (speed std 0.07; accel std 0.07) and moderate for the semantic channels (safety std 0.10; map std 0.14). The unconditional baseline is near-zero across all seeds (safety: +0.06 ± 0.18; map: +0.06 ± 0.06), confirming that the null path has no directional steering bias. F.4
Physical Plausibility Diagnostics
Table 5 reports the physical-plausibility diagnostics for the main model and the two hard-eligibility ablations discussed in section 5.4: the stall fraction (per-step displacement < 0.1 m), the mean speed of steered agents as a fraction of logged ground-truth speed, and the offroad rate at 8 s under speed steering, alongside the speed and safety CSM diagonals. The early-stage ablation’s inflated speed CSM (+51.3) is achieved by stalling (75.9% of agents, 61% of GT speed), whereas the main model retains 94.7% of GT speed at a physically valid +8.15; the hard-eligibility ablation’s safety CSM collapses to +0.21 versus the main model’s +0.66. F.5
Map Controllability Across Reward Measures
This section evaluates the map CSM diagonal under two return definitions alternative to the contextresidual measure: the physical-offroad return (raw non-residualized offroad plus road-edge penalty) and the lane-centerline return (section C.2). Figure 8 shows a stark sensitivity. Under the contextresidual return the map diagonal is positive (∆Rmap = +0.61 ± 0.14); under the physical-offroad return it is slightly negative (−0.12); and under the lane-centerline return it is effectively zero (−0.002). The context-residual measure residualizes away the scenario-dependent baseline, so the remaining variation is behaviorally attributable; the coordinate-specific measures do not, and the model’s map-channel response fails to separate from the scenario effect. The result has two impli22
Table 5: Physical plausibility: CNeVA with soft eligibility versus two hard-eligibility ablations. Open-loop, ρ = 1, w = 1.5, mean ± std over 5 seeds. Soft and hard eligibility define different λ populations, so safety CSM values are not directly commensurable across rows. Speed steering Model
minADE↓
Stall%↓ v/vGT ↑
CNeVA
1.113 ± .011
65.1
94.7%
+8.15 ± 0.07 +0.66 ± 0.10
32.5
Ablations (hard eligibility) Early-stage (40K) 1.238 ± .009 Hard-elig. (200K) 1.112 ± .015
75.9 65.4
61% 94%
+51.3 ± 0.3 +9.5 ± 0.2
34.6 32.9
∆Rspeed
∆Rsafety
+2.18 ± 0.10 +0.21 ± 0.22
Offroad%↓
cations. First, map controllability in table 2 is genuine but measure-dependent, reflecting a shift relative to the scene-contextual expectation rather than absolute geometric compliance. Meanwhile, coordinate-level map control (e.g., lane-keeping in a specific lane) would require richer reward decompositions that separate spatial from temporal map compliance.
Map CSM diagonal ΔRmap
controllable
0.6
0.4
0.2
0.0
Context-residual return
Physical-offroad return
Lane-centerline return
Figure 8: Map CSM diagonal (∆Rmap ) across three return measures, evaluated on the 200K softeligibility model at ρ = 1, w = 1.5 (open-loop). Map steering is controllable only under the contextresidual definition (+0.61); the physical-offroad (−0.12) and lane-centerline (≈ 0) measures show no measurable physical-space response.
References Prafulla Dhariwal and Alex Nichol. Diffusion models beat GANs on image synthesis. In Advances in Neural Information Processing Systems (NeurIPS), 2021. Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance, 2022. Yaron Lipman, Ricky T Q Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. In International Conference on Learning Representations (ICLR), 2023. Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. International Conference on Learning Representations (ICLR), 2019. https://openreview.net/forum?id=Bkg6RiCqY7.
In URL
Zikang Zhou, Haibo Hu, Xinhong Chen, Jianping Wang, Nan Guan, Kui Wu, Yung-Hui Li, Yu-Kai Huang, and Chun Jason Xue. BehaviorGPT: Smart agent simulation for autonomous driving with next-patch prediction. In Advances in Neural Information Processing Systems (NeurIPS), 2024.
23