Controllable Sim Agents with Behavior Latents

Overview§

CNeVA is a controllable simulated agent framework for traffic simulation. It learns per-agent Gaussian behavior latents from per-channel discounted returns using a closed-form conjugate variational update, and conditions a rectified-flow trajectory generator via a mixed channel-mask curriculum for classifier-free guidance. Soft eligibility gates replace hard thresholds to preserve gradients for near-threshold agents.

Core Methodology§

Behavior Latent Inference§

Each agent $i$ has a Gaussian behavior latent $z_i \sim \mathcal{N}(\mu_i, \sigma_i^2 I)$. Given per-channel discounted returns $R_i^c$ (e.g., speed, acceleration), the posterior is updated in closed form using conjugate priors: $$p(z_i | R_i) \propto p(R_i | z_i) p(z_i)$$ where $p(R_i | z_i) = \mathcal{N}(z_i, \tau^{-1} I)$ and $p(z_i) = \mathcal{N}(0, I)$. The posterior is $\mathcal{N}(\mu_i', \sigma_i'^2 I)$ with $\mu_i' = \frac{\tau}{1+\tau} R_i$ and $\sigma_i'^2 = (1+\tau)^{-1}$. This is a conjugate variational update.

Rectified-Flow Trajectory Generator§

A rectified-flow model generates future trajectories conditioned on the latent $z_i$ and past context. The flow transforms noise $x_0$ to data $x_1$ via ODE: $$dx_t = v_{\theta}(x_t, t, c) dt$$ where $c = \text{concat}(\text{encoder}(z_i), \text{context})$. Training uses classifier-free guidance by randomly masking channels during training.

Soft Eligibility Gates§

To handle sparse reward signals, soft eligibility gates use exponential decay: $$g_c = 1 - \exp(-\alpha \cdot R_c)$$ where $R_c$ is the per-channel return and $\alpha$ is a hyperparameter. This replaces hard thresholds with smooth gradients.

Implementation Details§

Architecture§

• Encoder: MLP with LayerNorm for latent inference.
• Trajectory Generator: Rectified-flow with transformer backbone. Conditioning via cross-attention.
• Training: Mixed curriculum with random channel masking (prob 0.15).

Code Snippet§

import torch
import torch.nn as nn

class CNeVA(nn.Module):
    def __init__(self, latent_dim=32, hidden_dim=256):
        super().__init__()
        self.latent_encoder = nn.Sequential(
            nn.Linear(6, hidden_dim),  # returns: speed, accel, etc.
            nn.ReLU(),
            nn.Linear(hidden_dim, latent_dim * 2)  # mu, logvar
        )
        self.flow_net = RectifiedFlowTransformer(latent_dim, hidden_dim)
        
    def infer_latent(self, returns):
        # returns: [batch, channels]
        params = self.latent_encoder(returns)
        mu, logvar = params.chunk(2, dim=-1)
        # Conjugate update (simplified): mu' = tau/(1+tau)*returns_mean
        tau = 1.0
        mu_prime = tau / (1 + tau) * returns.mean(dim=-1, keepdim=True)
        logvar_prime = -torch.log(1 + tau)
        return mu_prime, logvar_prime
    
    def sample_latent(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std

    def forward(self, past, returns):
        mu, logvar = self.infer_latent(returns)
        z = self.sample_latent(mu, logvar)
        traj = self.flow_net(past, z)
        return traj, mu, logvar

Training Objectives§

• Reconstruction loss: $\| \hat{x}_1 - x_1 \|^2$ for flow.
• KL divergence: $\mathcal{D}_{KL}(\mathcal{N}(\mu', \sigma'^2) \| \mathcal{N}(0, I))$.
• Soft eligibility regularization: $\sum_c \mathcal{L}_{BCE}(g_c, \mathbb{1}_{R_c > 0})$.

Results§

On Waymo Open Motion Dataset, CNeVA achieves competitive realism (ADE 1.32, FDE 2.89) while enabling per-channel controllability. Speed steering shows monotonic response (corr=0.97), safety steering reduces collisions by 40% without reward hacking.

Conclusion§

CNeVA provides a principled framework for controllable agents via behavior latents and rectified flows, with soft eligibility addressing sparse rewards.