1 Introduction↩︎

Birds exhibit remarkable swarm intelligence to maintain compact formation in highly uncertain aerial environments; however, endowing UAV swarms with simultaneous capabilities remains a fundamental challenge [1]. In recent years, UAVs have demonstrated strong potential in applications such as low-altitude logistics, surveillance, inspection, and indoor operations. When deployed as swarms in 4D airspace, UAVs must operate under significant temporal uncertainty arising from dense interactions. To achieve compact formation maintenance, each UAV is required to arrive at its destination at the expected time point. But accurate 4D trajectory prediction for multiple UAVs becomes markedly more difficult, as modeling their joint spatiotemporal evolution requires reasoning over strongly coupled dynamics; the temporal dimension further exacerbates the difficulty of 4D trajectory prediction. Second, increased interaction density amplifies mutual collision risks, demanding predictive models that are compatible with coordinated control and formation maintenance. Third, additional sources of uncertainty, such as communication delay, downwash, and aerodynamic disturbances, introduce stochasticity that is difficult to model and propagate over time. Taken together, maintaining compact formation remains challenging for UAV swarms to operate safely and simultaneously in complex low-altitude environments.

Currently, most existing prediction approaches focus on 2D or 3D settings, either omitting control feedback or treating time as an auxiliary variable, as shown in Fig. 1 (a). Such formulations are inherently limited in their ability to capture long-horizon temporal dependencies, and coordinated group behaviors that are fundamental to UAV swarms. Moreover, although trajectory prediction for pedestrians and ground vehicles has been extensively studied, the advantages of these approaches do not readily translate into UAV swarms. These methods typically assume planar motion and weak inter-agent coupling, which are insufficient for aerial swarms, where formation maintenance, altitude variation, and aerodynamic coupling introduce additional complexity.

To address this problem, several model-based and data-driven approaches have been employed for UAV swarm 4D trajectory prediction; however, they still suffer from unstable performance [2], [3]. First, classical model-based methods rely heavily on faithful nominal dynamics and well-characterized uncertainty bounds, assumptions that are difficult to satisfy in real-world swarm flight due to nonlinear dynamics, aerodynamic disturbances, and communication delays. Second, many learning-based trajectory prediction methods adopt deterministic formulations or simplified uncertainty structures, limiting their ability to represent temporally correlated disturbances such as airflow interaction and downwash effects. Third, existing approaches often treat trajectory prediction and control as separate modules, making it difficult to propagate prediction uncertainty into downstream control decisions, especially under nonlinear dynamics, strong inter-agent coupling, and limited information exchange. These limitations highlight the necessity of integrating uncertainty-aware 4D trajectory prediction with closed-loop control in a unified and scalable manner.

In this study, we propose a unified framework for 4D UAV swarm trajectory prediction and uncertainty-aware control, which jointly addresses spatiotemporal forecasting and closed-loop control under uncertainty. By explicitly coupling prediction and control, the framework enables reliable multi-UAV coordination under nonlinear dynamics, strong inter-agent interactions, and limited information exchange. As shown in Fig. 1 (b), our framework decouples spatial axes for coarse prediction, refines temporally correlated residual dynamics via denoising, and ultimately embeds future-aware predictions into a DNMPC loop for closed-loop UAV swarm control. The key contributions are summarized as follows:

• Unified prediction-control framework: We propose a coarse-to-fine architecture that combines efficient axis-wise forecasting with diffusion-based residual modeling. This design effectively captures nonlinear swarm dynamics and temporally correlated disturbances without sacrificing real-time feasibility.

• Stability-Oriented DNMPC Integration: We theoretically analyze and empirically validate the integration of diffusion-based predictions into a DNMPC scheme. Compared with prediction-only forecasting modules that are typically evaluated without explicit control-oriented residual compensation, our framework refines coarse forecasts using diffusion-based residual dynamics and feeds the refined rollouts into a DNMPC loop.

• Dataset: We release a diverse multi-scenario dataset featuring synchronized 4D trajectories (3D position + time) with frame-level annotations of speed intention and target sectors, addressing the scarcity of realistic benchmarks for close-proximity swarm operations.

2 Related Work↩︎

2.1 Trajectory Prediction Methods↩︎

Trajectory prediction aims to forecast the future motions of agents based on their historical observations. In multi-agent settings, a central challenge lies in modeling inter-agent interactions and collective behaviors under uncertainty. To address the inherent diversity and stochasticity of future trajectories, a wide range of generative models has been proposed, including Generative Adversarial Networks (GANs) [4], Conditional Variational Autoencoders (CVAEs) [5], and diffusion-based models [6], [7]. Sequential models based on recurrent neural networks (RNNs) [8] and Transformers [9] further enable flexible handling of variable-length inputs. These approaches have demonstrated strong predictive performance on benchmark datasets, particularly in pedestrian and vehicle trajectory prediction.

Despite their success, most existing trajectory prediction methods are designed for offline forecasting or benchmarking purposes and primarily optimize prediction accuracy without explicitly considering deployment in closed-loop control systems. As a result, their adaptability often degrades under real-world conditions characterized by strong disturbances, sensing uncertainty, and complex agent interactions. In contrast, our work targets 4D multi-UAV swarm trajectory prediction, where temporal uncertainty, inter-agent coupling, and control-oriented robustness must be jointly addressed.

2.2 Diffusion Models for Motion Dynamics↩︎

Diffusion models have recently emerged as a powerful paradigm for motion generation across a variety of domains [10]. Existing diffusion-based approaches can be broadly categorized into two lines: holistic trajectory diffusion and end-to-end diffusion policy methods, which differ fundamentally in their modeling targets and system integration.

2.2.1 Holistic Trajectory Diffusion↩︎

Holistic trajectory diffusion methods aim to model the entire spatiotemporal trajectory as a single structured random variable, and directly learn its joint distribution through diffusion processes. Representative works such as SwarmDiff [11] and MotionDiffuser [12] employ diffusion models to generate full future trajectories by capturing high-order temporal correlations and cross-agent interactions in a unified manner.

However, the holistic formulation often treats trajectory generation as an open-loop process, without explicitly incorporating system dynamics or control feedback during generation.

2.2.2 End-to-End Diffusion Policy↩︎

In contrast, end-to-end diffusion policy methods use diffusion models to directly learn a mapping from observations to control actions, effectively treating diffusion as a stochastic policy representation. In [13], the action sequence is generated directly from sensory observations through a conditional diffusion process. Related works in motion planning further integrate diffusion models with reinforcement learning or cost-guided objectives [14], enabling the generation of action or trajectory samples that optimize task-specific rewards or constraints [15], [16]. Beyond robotics, diffusion-based policies have also demonstrated strong performance in human motion synthesis, where high-dimensional motion sequences are generated and edited under diverse conditioning signals, such as text descriptions [17] or partial state constraints [18]. While these end-to-end approaches offer seamless integration between perception and control [13], [19], they typically entangle dynamics modeling and decision-making into a single learned policy, which can limit robustness in safety-critical control scenarios.

Motivated by these gaps, we propose a unified framework that couples control-oriented 4D multi-UAV trajectory prediction with closed-loop tracking. Specifically, we first generate coarse axis-wise forecasts to reduce prediction complexity, then model temporally correlated inter-agent residual dynamics via diffusion, and finally integrate the refined predictions into a DNMPC scheme, shown in Fig. 2. The proposed method is detailed in the following section.

3 Method↩︎

We consider a swarm consisting of \(M\) UAVs, indexed by \(i\in\{1,\dots,M\}\). Let \(\mathbf{x}_{i,k}\) denote the true state of UAV \(i\) at discrete time step \(k\), and \(\bar{\mathbf{x}}_{i,k}\) the predicted state along the NMPC horizon. We denote reference trajectories by \(\mathbf{x}^{\mathrm{ref}}_{i,k}\). Position components are denoted by \(p\), with \(\bar{p}\) indicating predicted positions and \(\mathbf{p}^{obs}\) historical observations. Each UAV maintains formation by enforcing inter-UAV distance constraints with its neighbors. Let \(d^{\mathrm{ref}}_{ij}\) denote the desired relative distance between UAV \(i\) and UAV \(j\). At time step \(k\), the formation constraint is defined as: \[\mathcal{C}_{ij}(k) = \big\| \bar{p}_{i,k} - \bar{p}_{j,k} \big\|_2 - d^{\mathrm{ref}}_{ij}, \label{eq:formation95constraint}\tag{1}\] where \(\bar{p}_{i,k}\) is extracted from \(\bar{\mathbf{x}}_{i,k}\).

3.1 Decoupled Trajectory Prediction↩︎

Predicting future UAV trajectories in four-dimensional space requires balancing model expressiveness and controllability. Existing keypoint-conditioned methods such as LBEBM [20] decode coupled 4D trajectories, which may entangle prediction errors across spatial dimensions. We therefore formulate trajectory forecasting as a dimension-factorized prediction problem.

We retain the encoder of LBEBM to extract interaction-aware representations from historical observations \(\mathbf{p}^{obs}_i\). Let \(\boldsymbol{\xi}_e\) denote the encoded latent feature. Axis-specific key points are predicted as: \[\mathbf{q}_{\kappa}^{\lambda} = \Phi_{\lambda}(\boldsymbol{\xi}_e), \quad \lambda \in \{x,y,z\},\] where \(\kappa\) indexes key points and \(\Phi_\lambda(\cdot)\) is an axis-dependent predictor.

The predicted key points are aggregated to form a compact structural descriptor: \[\boldsymbol{\xi}_k = \Psi\big([\mathbf{q}_{\kappa}^{x};\mathbf{q}_{\kappa}^{y};\mathbf{q}_{\kappa}^{z}]\big),\] where \(\Psi(\cdot)\) encodes multi-axis geometric cues. Trajectory decoding is then performed independently along each axis: \[\hat{\mathbf{T}}^{\lambda} = \Gamma_{\lambda}\big([\boldsymbol{\xi}_e,\boldsymbol{\xi}_k]\big), \quad \lambda \in \{x,y,z\}.\]

3.2 Residual Dynamics Refinement↩︎

To compensate for modeling errors and aerodynamic disturbances, we introduce a diffusion-based residual dynamics model. The translational dynamics of a quadrotor are given by: \[\dot{p}=v,\qquad m_b\dot{v}=m_bg+Rf_u+f_a,\] where \(m_b\) is the body mass, \(g\) is the gravity vector, and \(f_a\) aggregates unmodeled disturbances. Using a nominal mass \(\bar{m}_b\), the dynamics are rewritten as: \[\bar{m}_b\ddot{p}=u-\Delta(p,\dot{p},\ddot{p}),\] where \(\Delta\) collects all state-dependent uncertainties. \(\Delta(\cdot)\) is assumed to be an unknown but bounded and locally Lipschitz residual function capturing unmodeled dynamics and aerodynamic disturbances.

We model residual dynamics over a finite estimation horizon \(T\). Define the observation set \(\Omega=\{p,\dot{p},\ddot{p}\}\) and the joint sequence: \[\mathcal{S}= [\Omega_{1:T},\,u_{1:T},\,\Delta_{1:T}]^\top.\] For each UAV, \(p,\dot{p},\ddot{p} \in \mathbb{R}^3\), and the diffusion model is applied agent-wise with shared parameters across the swarm.

Residual estimation is formulated as sampling from \(p(\Delta_{1:T}\mid\Omega_{1:T},u_{1:T})\). The previous control input \(u_{t-1}\) is used as a proxy under receding-horizon operation.

3.3 DNMPC Control↩︎

The refined residual estimate \(\hat{\Delta}\) is incorporated into a DNMPC scheme. We use \(N\) to denote the NMPC prediction horizon, distinct from the residual estimation horizon \(T\). The residual horizon \(T\) determines the temporal extent of uncertainty estimation, while the NMPC horizon \(N\) governs control foresight, with \(T \leq N\) ensuring consistent disturbance compensation within receding-horizon optimization. Each UAV solves a local NMPC problem to track its reference trajectory while respecting formation constraints. The discrete-time prediction model is: \[\bar{\mathbf{x}}_{i,k+j+1} = f_d(\bar{\mathbf{x}}_{i,k+j},\bar{\mathbf{u}}_{i,k+j},\hat{\Delta}_{i,k+j},\delta t).\] Only the first control input is applied at each step, yielding a receding-horizon closed-loop control law.

Proposition 1: Practical tracking stability. Assume recursive feasibility of the DNMPC subproblems. Then the optimal value function \(V_{i,k}:=J_i^\star(\mathbf{x}_{i,k})\) is Lyapunov-like for the closed-loop error dynamics, i.e., there exist \(\alpha_i>0\) and \(\beta_i>0\), \[V_{i,k+1}-V_{i,k} \le -\alpha_i \|\mathbf{e}_{i,k}\|_2^2 + \beta_i\!\left(\bar{\delta}_\Delta^2+\bar{\delta}_n^2\right). \label{eq:lyap95decrease}\tag{2}\] Consequently, the closed-loop tracking error is input-to-state stable w.r.t. \((\tilde{\Delta}_{i,k},\,\bar{\mathbf{x}}_{j,k}-\mathbf{x}_{j,k})\) and is practically stable: \[\limsup_{k\to\infty}\|\mathbf{e}_{i,k}\|_2 \le c_i \sqrt{\bar{\delta}_\Delta^2+\bar{\delta}_n^2}, \label{eq:practical95bound}\tag{3}\] for some constant \(c_i>0\). In the ideal case \(\tilde{\Delta}_{i,k}\equiv 0\) and \(\bar{\delta}_n=0\), 2 implies asymptotic tracking of \(\mathbf{x}^{\mathrm{ref}}_{i,k}\).

4 4D Motion Trajectory Dataset↩︎

Our dataset comprises six scenarios featuring 4D trajectories collected from UAVs operating in real typical urban environments. The 4D trajectories, shown in Fig. 3, are recorded using positioning devices installed on the UAVs. Each scenario includes three UAVs following predefined formations. The three-UAV setting is chosen to reflect tightly coupled formations, while the distributed NMPC formulation does not impose an upper bound on swarm size.

The dataset contains a total of 7,900 frames, with scenario lengths ranging from 800 to 2,000 frames. As summarized in Table 1, the inter-UAV distance varies between \(1.3\,\mathrm{m}\) and \(5.8\,\mathrm{m}\) on average, reflecting realistic formation maintenance under environmental constraints. The recorded flight speed spans \(2.4\)–\(5.6\,\mathrm{m/s}\), while acceleration and turning intensity statistics reveal diverse maneuvering behaviors across scenarios. These characteristics jointly capture both smooth cruising and aggressive directional changes that frequently arise in practical low-altitude UAV operations.

6pt

5 Experiments↩︎

5.1 Implementation Details↩︎

We consider a short-horizon trajectory prediction setting, where the past 8 frames are used as observations to forecast the subsequent 12 frames. Rather than predicting all future states densely, the model outputs a set of representative future key points corresponding to the 3rd, 6th, 9th, and 12th prediction steps, which are subsequently interpolated and incorporated into the NMPC reference trajectory. This design reduces prediction redundancy while preserving sufficient temporal resolution for control.

Model training is performed for 100 epochs with a batch size of 70 using the Adam optimizer. The initial learning rate is set to \(5\times10^{-4}\) and decayed by a factor of 0.5 after the fifth epoch. Prediction accuracy is evaluated using Average Displacement Error (ADE) and Final Displacement Error (FDE), while closed-loop control performance is assessed using trajectory tracking error and RMSE.

The diffusion noise predictor \(\epsilon_{\theta}\) is implemented using a temporal U-Net architecture composed of three residual stages [14], each including two temporal convolution layers followed by group normalization [21] and Mish activation [22]. To encode diffusion timesteps and conditional inputs, we employ two separate two-layer MLPs with 64 hidden units, producing 16-dimensional embeddings. The noise predictor is optimized with Adam using a learning rate of \(2\times10^{-4}\) over 50,000 training steps with a batch size of 256 and \(K=20\) diffusion steps.

For control performance evaluation, we benchmark the proposed framework against a physics-driven DNMPC baseline, which has demonstrated strong stability and coordination capabilities in multi-agent control scenarios [23]. To assess the benefit of incorporating diffusion-based residual modeling, we further compare it with representative data-driven dynamics learning methods, including Gaussian Processes (GP) models [24] and deep neural network (DNN) predictors [25], both of which have been widely adopted for quadrotor dynamics approximation. An autoregressive sequence modeling baseline based on GPT4 [26] is also included. The GPT-based autoregressive model serves as a strong sequence modeling baseline to assess whether generic temporal predictors can approximate residual dynamics. For a fair tracking comparison, all learning-based models are trained on the same offline dataset to predict residual dynamics \(\Delta\), and are integrated into the identical control law in (8). Each method is evaluated over 10 independent trials.

5.2 Comparisons to State-of-the-art Methods↩︎

We compare our method with representative baselines spanning classical filtering, deterministic sequence models, and generative trajectory predictors, as detailed in Table 2. Kalman Filter (KF) and Unscented Kalman Filter (UKF) serve as classical model-based baselines for state estimation [27]. Temporal Convolutional Networks (TCN) and MSRL represent deterministic learning-based predictors [28], [29]. LBEBM and NPSN capture energy-based and probabilistic sequence modeling paradigms [20], [30]. GAN- and CVAE-based methods are included as canonical generative baselines for multimodal trajectory prediction [31], [32], while MTRTraj represents recent large-capacity transformer-based trajectory models [33].

We modify all compared methods to accept 4D observed trajectories as inputs and predict 4D future trajectories. Additionally, we adjust their data augmentation strategies, such as rotation, flipping, and reversing, to their 4D counterparts for consistency. The experimental results indicate that our method outperforms all evaluated approaches in 4D trajectory prediction. Specifically, our method achieves a 16.7% improvement in the ADE metric, reducing it from 0.18 to 0.15, and a 18.2% enhancement in the FDE metric, lowering it from 0.33 to 0.27, compared to the second-best method, MTRTraj. While Unscented Kalman Filter (UKF) shows an ADE metric close to ours, our FDE metric surpasses it by 48%, decreasing from 0.52 to 0.27.

Table 2 also compares prediction accuracy and computational efficiency across representative baselines. Classical filters (KF/UKF) achieve low latency but suffer from large prediction errors, while lightweight neural models such as TCN and MSRL offer moderate accuracy gains at increased cost. Our diffusion-based prediction improves accuracy but incurs substantial runtime overhead due to iterative sampling. Our method adopting DDIM acceleration [34] attains the best accuracy-efficiency trade-off, achieving the lowest ADE/FDE with sub-30 ms inference latency at about 34 FPS. These results indicate that reusing multistep diffusion predictions effectively reduces sampling overhead while preserving real-time feasibility for control deployment.

3.5pt

Thus, the diffusion-based predictor provides NMPC with more reliable and temporally consistent state rollouts, reducing abrupt prediction errors that may otherwise lead to infeasible or aggressive control actions. This effect becomes particularly pronounced under tight formation constraints and high-speed maneuvers, where small modeling errors can be significantly amplified by inter-UAV interactions. As a result, DM-NMPC achieves improved tracking accuracy and robustness by implicitly regularizing the NMPC optimization landscape, leading to smoother control inputs and enhanced stability in multi-UAV swarm.

5.3 Robustness and Adaptation Analysis↩︎

Table 3 compares learning-assisted NMPC controllers across four motion scenarios. Overall, DM-NMPC achieves the lowest prediction and tracking errors, indicating improved compensation for unmodeled dynamics and inter-UAV interactions. The prediction error, measured in Newtons, quantifies the accuracy of estimating interaction-induced and environment-dependent disturbance forces between neighboring UAVs. Accurate force prediction is critical for DNMPC, as these estimates are directly used in the feedback compensation term of the control law. Compared with GP-NMPC and DNN-NMPC, our DM-NMPC reduces trajectory tracking RMSE by 15-22% in Urban and Factory scenarios and by 10-15% in Lab and Rural environments.

3pt

Fig. 4 compares inter-UAV distance convergence under Gaussian (left) and heavy-tailed (right) disturbances. Under Gaussian noise, most learning-based compensation strategies achieve stable convergence, with GP-based methods benefiting from smoothness assumptions. However, under heavy-tailed Cauchy disturbances, GP- and DNN-based compensators exhibit degraded performance, characterized by persistent oscillations or steady-state offsets. In contrast, diffusion-based compensation consistently maintains bounded inter-UAV distance with reduced oscillations across both regimes, demonstrating superior robustness to non-Gaussian and extreme disturbances in multi-UAV formation control.

Under different reference spacing and speed settings, Table 4 reports the formation constraint RMSE of various NMPC variants. When the speed is increased from 8 m/s to 10 m/s, the RMSE of all methods increases. The baseline methods exhibit a more pronounced increase in error under the combination of high speed and large spacing, reflecting their limited ability to compensate for unmodeled disturbances or residual dynamics. In contrast, DM-NMPC achieves the lowest RMSE across all conditions and maintains a relatively stable error level even in the more challenging setting (10 m/s, 0.8 m), indicating its capacity to more effectively absorb prediction errors and disturbances during the closed-loop control process, thereby enhancing the robustness and consistency of the formation constraints.

4pt

5.4 Prediction Efficiency and Runtime Performance↩︎

As illustrated in Fig. 5, increasing the prediction horizon initially improves tracking accuracy because the controller can anticipate future motion more effectively. However, when the horizon becomes too long, performance deteriorates as modelling errors accumulate and the system is more likely to violate real-time assumptions. In parallel, the solver time per step rises almost linearly with the horizon length and exceeds the 50 ms real-time limit when \(N \ge 32\). Overall, a moderate horizon of \(N\), from 8 to 16, offers the most balanced compromise between tracking quality and computational feasibility for DNMPC in UAV swarms. This decline at large \(N\) is mainly due to estimating \(\Delta\) from increasingly outdated state observations and control inputs.

As shown in Fig. 6, GP and DNN tend to concentrate in the inner region or exhibit distorted coverage, while GPT produces scattered interior mass. In contrast, the diffusion-based model aligns well with the high-density manifold, particularly along the outer ring, indicating improved distributional fidelity.

5.5 Failure Case Analysis↩︎

We further analyze a representative failure case in Fig. 7. Although our method accurately captures short-term motion trends, prediction errors accumulate in highly curved segments where inter-UAV interactions intensify. This suggests that the current residual dynamics model, trained under fixed formation assumptions, may be insufficient to fully capture abrupt interaction changes or rare coordination patterns. We leave adaptive interaction modeling and uncertainty-aware horizon adjustment as future work.

5.6 Generalization Discussion↩︎

To evaluate cross-dataset generalization, we additionally conduct trajectory prediction experiments on the public CityNav dataset [35]. Since CityNav does not provide synchronized multi-UAV trajectories or control annotations, this evaluation focuses on open-loop trajectory prediction rather than closed-loop control. As shown in Table 5, despite the domain gap between datasets, our method achieves consistently lower prediction errors, suggesting enhanced robustness to scenario variations.

6pt

6 Conclusion and Future Work↩︎

In this paper, we present a unified framework for 4D UAV swarm trajectory prediction and distributed control. By decoupling the trajectory prediction task and employing a diffusion model solely for residual dynamics refinement, our approach balances uncertainty modeling with the real-time demands of onboard control. Experimental results confirm that our method outperforms existing transformer and generative baselines. The distributed DNMPC formulation does not impose a fixed upper bound on the swarm size, and the residual model shares parameters across UAVs. However, the current dataset and empirical validation are limited to three-UAV formations, scaling to larger swarms remains future work.

7 Acknowledgement↩︎

This work was partially supported by the National Natural Science Foundation of China Grant 62373315 and U24A20252, and National Key Research and Development Program of China Grant 2024YFB4707603.

References↩︎

---

1. \(^\ast\) Corresponding author.↩︎

2. \(^{1}\) Tianshun Li is with The Hong Kong University of Science and Technology (Guangzhou), China tli449@connect.hkust-gz.edu.cn↩︎

3. \(^{2}\)Hongliang Lu is with the Southern University of Science and Technology, China and MoSense Technologies, China honglianglu36@gmail.com↩︎

4. \(^{3}\)Haoang Li and Xinhu Zheng are with The Hong Kong University of Science and Technology (Guangzhou), China {haoangli,xinhuzheng}@hkust-gz.edu.cn↩︎