UNet Diffusion Trajectory Generator

• The paper presents a UNet-based diffusion model that generates high-fidelity state-action trajectories by inverting a progressive noise process.
• The methodology integrates DDPM and continuous-time SDE approaches with cross-attention conditioning to enhance offline RL, behavioral modeling, and mobility synthesis.
• Empirical results indicate significant improvements in metrics such as episode returns, Jensen–Shannon divergence, minADE, and minFDE, validating its robust performance.

A UNet-based diffusion trajectory generator is a class of neural generative model that synthesizes entire state-action-reward trajectories (or spatial or dynamical trajectories) by applying a UNet (or “U-shaped” encoder–decoder) architecture as the noise predictor (denoiser) within a diffusion probabilistic model, following either DDPM-style (Ho et al.) or more recent continuous-time score-based (EDM, SDE) methodologies. By learning to reverse a forward noise process that corrupts real trajectories, these architectures can produce new, high-fidelity samples under complex conditioning, with demonstrated advantages in offline reinforcement learning (RL), behavioral modeling, mobility synthesis, turbulence, and time-series smoothing.

1. Mathematical and Algorithmic Foundations

UNet-based diffusion trajectory generators operate by training a diffusion model—typically a Markovian chain or continuous SDE—over trajectory space. The forward process $q(x_t|x_{t-1})$ iteratively adds noise to a clean trajectory $x_0$, producing $x_T$ distributed approximately as white Gaussian noise. The reverse process $p_\theta(x_{t-1}|x_t)$, parameterized by a neural network (UNet), attempts to invert this chain by denoising, thus sampling realistic trajectories from noise.

In the archetypal DDPM parameterization, the forward diffusion is:

$$q(x_t|x_{t-1}) = \mathcal{N}\left(x_t; \sqrt{1-\beta_t} x_{t-1},\, \beta_t I\right), \quad \text{with } \beta_t > 0$$

with closed-form marginal: $$x_t = \sqrt{\bar \alpha_t} x_0 + \sqrt{1 - \bar \alpha_t} \epsilon, \quad \epsilon\sim \mathcal{N}(0, I)$$

The learned denoising network $\epsilon_\theta$ (the UNet) predicts the additive noise at each step. The loss is typically:

$$L(\theta) = \mathbb{E}_{t, x_0, \epsilon}\left[ \|\epsilon - \epsilon_\theta(\sqrt{\bar\alpha_t} x_0 + \sqrt{1 - \bar\alpha_t} \epsilon, t, c)\|^2 \right]$$

For conditional generation tasks, context $c$ (past trajectory, observed points, map context, attributes) is injected into the network via cross-attention, FiLM, or additional input channels (Zhu et al., 2023, Tamir et al., 2024, Qingze et al., 2024, Liu et al., 2024, Li et al., 25 Jul 2025, Zhu et al., 2024, Batool et al., 21 Jan 2026).

2. Representative UNet and Diffusion Architectures

Trajectory diffusion models diversify across application domains, but share core architectural conventions:

• 1D/2D UNet Backbone: Inputs are noisy trajectory tensors (typically shape $[B, d, n]$; $d$ = variables, $n$ = timesteps), processed via multi-level downsampling (using ResNet/CNN or self-attention blocks) and symmetric upsampling, with skip connections for multi-scale fusion (Zhu et al., 2023, Li et al., 25 Jul 2025, Batool et al., 21 Jan 2026).
• Diffusion Transformer (SDE-based UNet): In some RL scenarios, the backbone is a transformer-style UNet employing self-attention within blocks, skip connections, and context fusion via cross-attention (“diffusion transformer”) (Liu et al., 2024).
• GeoUNet and Spatial Variants: Incorporate specialized attention mechanisms (e.g., “Geo-Attention”) for integrating spatial constraints such as road networks (Zhu et al., 2024).
• Denoising Heads: The output is a tensor $\epsilon_\theta$ matching the input trajectory dimension, interpreted as the noise estimate at each timestep.

UNets are typically conditioned on timestep embeddings (sinusoidal or learned positional; injected via FiLM or additive bias), as well as conditioning signals (past observations, intent waypoints, maps, attributes), which may be processed via MLPs, CNNs, or attention modules before injection (Zhu et al., 2023, Liu et al., 2024, Tamir et al., 2024, Qingze et al., 2024).

3. Conditioning and Domain-Specific Extensions

A critical innovation across trajectory generation is tailoring the denoising UNet and sampling algorithms to admit conditioning information:

• Offline RL and Policy Learning: In Diffusion-Based Trajectory Branch Generation for Decision Transformer, conditioning is based on trajectory segments and a return-to-go scalar, with the UNet integrating these via cross-attention. Branches are filtered with a learned Value Function for reward-oriented bias ((Liu et al., 2024), Table below).

Application	Conditioning Mechanism	UNet Characterization
Offline RL (DT+BG)	Segment + Return-to-Go	Diffusion Transformer (SDE, cross-attn)
Mobility Synthesis	Road segment + trip attrs	GeoUNet (CNN+geo-attn)
Time-series Smoothing	Partial observations	UNet + cross-attn/FiLM
UAV Vision-Based Planning	Image + endpoint masks	UNet + ResNet img encoder

• Spatio-Temporal Context: For map-related tasks (e.g., TrajDiffuse, ControlTraj), UNets are conditioned on semantic or distance-transform maps and road-segment embeddings to ensure environment-compliant generation (Qingze et al., 2024, Zhu et al., 2024).
• Guidance and Filtering: Several implementations employ guidance (TVF, classifier-free, or map-based) to steer sampling toward desired or high-reward regions, optionally with acceptance/rejection filtering (Liu et al., 2024, Qingze et al., 2024, Zhu et al., 2023).
• Observation and Smoothing Constraints: Time-indexed observations are encoded and used as cross-attentive context for trajectory smoothing under state-space models (Tamir et al., 2024).

4. Training Protocols and Inference Procedures

Training consists of minimizing the denoising loss (L₂ or score matching) over a large set of real trajectories, with noised inputs generated according to the prescribed diffusion schedule. Best practices include large batch sizes (e.g., 1k+), learning rate scheduling, and utilization of time/context embeddings at each block (Zhu et al., 2023, Li et al., 25 Jul 2025, Liu et al., 2024).

Inference (trajectory sampling) proceeds in the following high-level sequence:

1. Initialization: Start from Gaussian noise or prescribed prior (standard normal or terminal state) as input trajectory.
2. Reverse Sampling: For each diffusion step (DDPM: stochastic; DDIM: deterministic), update the trajectory by applying the UNet denoiser, possibly integrating guidance or cross-attention conditioning, and optionally accelerating with reduced-step schedules ((Li et al., 25 Jul 2025) provides precise DDIM formulas).
3. Post-Processing: Some workflows apply “hard inpainting” (fixing endpoints), map-based guidance, or acceptance filtering (e.g., reward continuity, environment constraints) in between steps (Liu et al., 2024, Qingze et al., 2024, Batool et al., 21 Jan 2026).

Accelerated inference via DDIM and step-skipping is widely adopted; empirical evidence suggests UNet-based models retain fidelity even with ∼3% of the original number of steps (Li et al., 25 Jul 2025). Exponential moving averages of weights further stabilize predictions.

5. Empirical Results and Performance Assessment

UNet-based diffusion trajectory generators consistently demonstrate state-of-the-art or highly competitive performance across domains:

• Offline RL (DT+BG): Branch Generation (BG) coupled with a Decision Transformer outperforms all prior sequence modeling baselines on D4RL. Ablations show disabling the TVF guidance substantially reduces total episode return from 300.0 to 249.4 on Maze2d; disabling branch filtering further drops it to 228.2. On sparse-reward AntMaze, BG+DT achieves goal-reaching where standard DT fails (Liu et al., 2024).
• Urban Mobility/Traffic: DiffTraj and ControlTraj’s UNet variants yield lower Jensen–Shannon divergence (JSD) across point-density, trip, and length distributions compared to LSTM- and GAN-based alternatives. On real datasets, classifier-free guidance can balance diversity and fidelity (Zhu et al., 2023, Zhu et al., 2024).
• Trajectory Prediction with Constraints: TrajDiffuse achieves near-perfect compliance (ECFL ≈ 99.1–99.6%) with environmental constraints. Performance in minADE and minFDE is within 0.05–0.09 m on PFSD, matching or slightly exceeding previous bests (Qingze et al., 2024).
• Turbulence and Physical Processes: U-Net backbones recover both coarse and fine-scale statistics of Lagrangian turbulence, matching structure functions and intermittency flatness across scales. Transformer backbones slightly underestimate small-scale features (Li et al., 25 Jul 2025).
• Human-Aware UAV Planning: HumanDiffusion’s UNet planner produces trajectories with pixel-space MSE = 0.02 and 80% real-world mission success rate. Ablations show that removing image conditioning or skip connections degrades MSE and reduces success (Batool et al., 21 Jan 2026).

6. Variants, Robustness, and Future Directions

The UNet-based diffusion paradigm exhibits architectural robustness—swap-in transformer-style blocks or geo-attention as needs dictate, with only minor degradation at extreme scales (Li et al., 25 Jul 2025). Sample diversity and complexity are governed by conditioning mechanisms (context fusion, map encoding), reward guidance, and variance schedules.

Open challenges identified include:

• Support for variable-length or continuous-time trajectories instead of fixed-length interpolation (Zhu et al., 2023).
• Integration of learned or adaptive variance schedules for higher sample quality.
• Formal incorporation of privacy guarantees in synthetic data applications.
• Direct inclusion of spatio-temporal graph modules or richer attention for multi-agent and structured domains (Zhu et al., 2024, Lee et al., 29 Sep 2025).
• Further improving efficiency for inference in large-scale or time-critical environments, especially by distillation or learned simulators.

This synthesis draws upon and integrates results reported in (Liu et al., 2024, Zhu et al., 2023, Qingze et al., 2024, Zhu et al., 2024, Li et al., 25 Jul 2025, Batool et al., 21 Jan 2026), and (Tamir et al., 2024).