Paper deep dive
Fast and Scalable Multi-Agent Distribution Matching via Partitioned Optimal Transport
Kooktae Lee, Ruchika Singh
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 86%
Last extracted: 10/1/2026, 7:07:47 AM
Summary
This paper proposes a scalable framework for terminal distribution matching in multi-agent systems using partitioned optimal transport (OT). To address the computational bottleneck of global discrete OT in large-scale systems, the authors partition agents and target samples into spatially corresponding blocks. By enforcing a mass-balance condition, the global problem decomposes into smaller local transport problems. The resulting block-restricted couplings provide an upper bound on the Wasserstein cost and generate barycentric targets for finite-horizon agent control. The framework establishes a cycle-to-cycle descent guarantee for the transport surrogate, validated through simulations for both linear and nonlinear agent dynamics.
Entities (8)
Relation Signals (6)
Kooktae Lee → affiliatedwith → Texas Tech University
confidence 99% · Kooktae Lee and Ruchika Singh are with the Department of Mechanical and Aerospace Engineering, Texas Tech University
Ruchika Singh → affiliatedwith → Texas Tech University
confidence 99% · Kooktae Lee and Ruchika Singh are with the Department of Mechanical and Aerospace Engineering, Texas Tech University
Partitioned Optimal Transport → providesupperboundon → Wasserstein Distance
confidence 96% · the resulting restricted coupling remains feasible for the global problem and provides an upper bound on the Wasserstein cost.
Partitioned Optimal Transport → reducescomputationalcostof → Global Discrete Optimal Transport
confidence 95% · While optimal transport provides a natural way... global discrete transport can become computationally expensive... We address this bottleneck by partitioning agents
Partitioned Optimal Transport → enables → Terminal Distribution Matching
confidence 94% · The proposed framework therefore enables scalable terminal distribution matching
Texas Tech University → receivedfundingfrom → NSF CAREER Grant
confidence 90% · This work was supported by NSF CAREER Grant CMMI-DCSD-2638508.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:This paper presents a scalable optimal-transport-based framework for terminal distribution matching in multi-agent systems. While optimal transport provides a natural way to measure distributional mismatch and assign agents to a desired spatial distribution, global discrete transport can become computationally expensive for large-scale systems. We address this bottleneck by partitioning agents and target samples into spatially corresponding blocks and solving smaller local transport problems. Under a mass-balance condition, the resulting restricted coupling remains feasible for the global problem and provides an upper bound on the Wasserstein cost. The local assignments generate target locations for finite-horizon agent control, applicable to both linear and nonlinear dynamics. By alternating local assignment and control, we establish a cycle-to-cycle descent guarantee for the resulting transport surrogate. The proposed framework therefore enables scalable terminal distribution matching while retaining a rigorous connection to the Wasserstein objective. The technical soundness of the proposed results is validated through simulations.
Tags
Links
- Source: https://arxiv.org/abs/2609.38960v1
- Canonical: https://arxiv.org/abs/2609.38960v1
Trouble viewing inline? Open PDF directly →
Full Text
45,678 characters extracted from source content.
Expand or collapse full text
Fast and Scalable Multi-Agent Distribution Matching via Partitioned Optimal Transport Kooktae Lee Ruchika Singh †thanks: *This work was supported by NSF CAREER Grant CMMI-DCSD-2638508.†thanks: Kooktae Lee and $ˆ1$Ruchika Singh are with the Department of Mechanical and Aerospace Engineering, Texas Tech University, Lubbock, TX 79409, USA, email: kooktae.lee@ttu.edu, ruchisin@ttu.edu. Abstract This paper presents a scalable optimal-transport-based framework for terminal distribution matching in multi-agent systems. While optimal transport provides a natural way to measure distributional mismatch and assign agents to a desired spatial distribution, global discrete transport can become computationally expensive for large-scale systems. We address this bottleneck by partitioning agents and target samples into spatially corresponding blocks and solving smaller local transport problems. Under a mass-balance condition, the resulting restricted coupling remains feasible for the global problem and provides an upper bound on the Wasserstein cost. The local assignments generate target locations for finite-horizon agent control, applicable to both linear and nonlinear dynamics. By alternating local assignment and control, we establish a cycle-to-cycle descent guarantee for the resulting transport surrogate. The proposed framework therefore enables scalable terminal distribution matching while retaining a rigorous connection to the Wasserstein objective. The technical soundness of the proposed results is validated through simulations. Index Terms: Multi-agent systems, optimal transport, Wasserstein metric, distribution matching, spatial coverage. I Introduction Coordinating a group of autonomous agents to achieve a desired spatial configuration is a fundamental problem in multi-agent systems. In many applications, the objective is not to drive each agent to a predetermined location, but to deploy agents across a region according to the spatial demand of the task. This type of spatial deployment arises in multi-robot coverage, surveillance, exploration, and environmental sampling, where coverage and sensing objectives can guide the spatial placement of multiple agents [1, 2, 3]. A number of approaches have therefore formulated multi-agent coordination in terms of spatial distributions, including density-driven optimal control (D2OC) for coverage under nonuniform reference densities [4, 5, 6]. These approaches primarily characterize desired spatial behavior through time-averaged agent trajectories. In this paper, we instead formulate the deployment objective as a terminal distribution matching problem, where the agent configuration at a prescribed terminal time is matched to a given target distribution. A closely related terminal-distribution formulation was developed by [7], where decentralized multi-agent distribution matching was achieved through an optimal-transport-based Wasserstein objective. Optimal transport (OT) [8] provides a natural framework for this problem by measuring the mismatch between the current and desired distributions [9, 10] and determining how agents should be associated with the target distribution. Thus, OT provides both a measure of distributional mismatch and an assignment mechanism connecting individual agents to the desired spatial distribution. Additional literature on OT-based distribution steering and density control can be found in [11, 12]. The computational cost of this assignment grows with the numbers of agents and target samples, since discrete OT considers associations between agents and target samples [10]. This becomes particularly challenging for large-scale multi-agent systems, motivating scalable and distributed OT methods [13]. Spatial structure has also been exploited in semi-discrete OT to obtain tractable transport partitions [14]. These approaches motivate exploiting spatial structure to reduce the computational burden of OT while retaining its role in assigning agents to a desired spatial distribution. This paper addresses the computational bottleneck of large-scale OT by exploiting spatial structure in the assignment problem. Our prior work [7] developed both a centralized sequential-assignment framework and a memory-based decentralized variant for terminal distribution matching. The present paper builds on the centralized formulation of [7], which is the natural choice for applications requiring structured, globally coordinated deployment. The decentralized variant of [7] instead targets settings where such global coordination is unavailable, and is not the focus here. Even in the centralized setting, however, the assignment step of [7] sequentially couples M agents with N target samples, and this coupling becomes computationally expensive as the number of target samples used to represent the desired distribution grows. To overcome this limitation, this work introduces a spatially partitioned computational architecture. The contributions of this paper are thus threefold: (i) a spatially partitioned OT formulation that strictly preserves global feasibility via block mass-balance constraints; (i) an integration of these localized assignments with the barycentric control and target-generation framework; and (i) a rigorous cycle-to-cycle descent guarantee for the transport surrogate, which maintains a direct upper bound on the Wasserstein distance. I Preliminaries and Problem Formulation This section introduces the multi-agent dynamics and terminal distribution matching problem, followed by the discrete optimal transport formulation used to measure distribution mismatch. We then summarize the sequential assignment–control framework of [7], which serves as the basis for the computational reduction developed in this work. I-A Terminal Distribution Matching Consider a swarm of M agents, where the state of agent i at discrete time k is xi(k)∈ℝdx_i(k) ^d. Each agent evolves according to either linear dynamics xi(k+1)=Axi(k)+Bui(k),x_i(k+1)=Ax_i(k)+Bu_i(k), (1) or control-affine nonlinear dynamics xi(k+1)=f(xi(k))+g(xi(k))ui(k).x_i(k+1)=f(x_i(k))+g(x_i(k))u_i(k). (2) Let N target samples yj∈ℝdy_j ^d represent the desired spatial distribution. The empirical agent distribution at time k and the target distribution, respectively, are defined by μk:=1M∑i=1Mδxi(k),ν:=1N∑j=1Nδyj, _k:= 1M _i=1^M _x_i(k), ν:= 1N _j=1^N _y_j, (3) where δz _z denotes the Dirac measure concentrated at z. Each agent carries a mass of 1/M1/M, while each target sample carries a mass of 1/N1/N. Given a terminal time T, the objective is to steer the terminal empirical distribution μT=1M∑i=1Mδxi(T) _T= 1M _i=1^M _x_i(T) toward the prescribed target distribution ν. I-B Discrete Optimal Transport The distribution mismatch between μ and ν is measured using the squared 2-Wasserstein distance. The set of transport couplings is Π(μ,ν)=π∈ℝ+M×N:∑j=1Nπij=1M,∑i=1Mπij=1N, (μ,ν)= \π _+^M× N: _j=1^N _ij= 1M,\; _i=1^M _ij= 1N \, (4) where the transport plan πij≥0 _ij≥ 0 denotes the amount of mass transported from agent i to target sample j. The prescribed row and column marginals preserve the total mass since they imply ∑i=1M∑j=1Nπij=1 _i=1^M _j=1^N _ij=1. The squared 2-Wasserstein distance between the empirical distributions is then 22(μ,ν)=minπ∈Π(μ,ν)∑i=1M∑j=1Nπij∥xi−yj∥2.W_2^2(μ,ν)= _π∈ (μ,ν) _i=1^M _j=1^N _ij\|x_i-y_j\|^2. (5) Note that for any feasible coupling π∈Π(μ,ν)π∈ (μ,ν), we have 22(μ,ν)≤∑i=1M∑j=1Nπij‖xi−yj‖2.W_2^2(μ,ν)≤ _i=1^M _j=1^N _ij\|x_i-y_j\|^2. (6) Thus, the transport cost of a feasible coupling provides an upper bound on the Wasserstein distance without requiring the coupling to be globally optimal. I-C Terminal Distribution Matching: Sequential Assignment and Control The terminal distribution matching framework of [7] addresses the distribution-level objective through an alternating assignment–control procedure. Rather than directly optimizing the terminal states of all agents jointly, which is computationally intractable due to high non-convexity and high dimensionality, target mass is assigned to the agents sequentially. Starting from the prescribed target sample weights, each agent receives transport mass according to its current assignment cost, and the assigned target samples determine a barycentric terminal reference bi=∑j=1Nπijyj∑j=1Nπij,b_i= _j=1^N _ijy_j _j=1^N _ij, where the sequential allocation of target samples ensures that assigned weights are appropriately decremented, preventing multiple agents from being assigned to the same target samples. The agents are then controlled toward these barycentric targets over a finite horizon. After the control horizon is completed, the target sample weights are reinitialized to their prescribed values and the sequential assignment–control procedure is repeated using the updated agent states. The resulting procedure converts the distribution-level matching problem into a sequence of agent-level finite-horizon control problems. However, the assignment step still sequentially couples M agents with N target samples, which can become computationally expensive as the problem size increases. This motivates the partitioned formulation developed in the next section, where the transport coupling is restricted to corresponding spatial blocks. Under a suitable mass-balance condition, the resulting problem decomposes into smaller local transport problems while preserving feasibility with respect to the global transport formulation. I Partitioned Optimal Transport I-A Spatial Partition and Mass Balance Suppose that the M agents and N target samples are partitioned into K corresponding spatial blocks. The blocks are chosen so that agents and target samples located in the same spatial region are grouped together. Let =⋃r=1K(r),=⋃r=1K(r),X= _r=1^KX^(r), = _r=1^KY^(r), (7) where the blocks are pairwise disjoint. Denote the numbers of agents and target samples in block r by Mr=|(r)|,Nr=|(r)|.M_r=|X^(r)|, N_r=|Y^(r)|. (8) We impose the following mass-balance condition on corresponding blocks: Assumption 1 (Block Mass Balance). For every block r, MrM=NrN=:qr. M_rM= N_rN=:q_r. (9) Thus, a block containing a fraction qrq_r of the agents is paired with a corresponding target block containing the same fraction of the target samples. Importantly, the corresponding blocks need not contain the same number of agents and target samples. Only their normalized empirical masses must agree. For example, if M=30M=30 and N=600N=600, a block containing 10 agents may be paired with a block containing 200 target samples, since both blocks contain one-third of their respective populations. The blocks may be spatially defined, for example, by clustering or a fixed spatial partition, but no particular partitioning procedure is required by the analysis. The key requirement is the mass-balance condition, which ensures that each agent block can be transported entirely to its corresponding target block. I-B Block-Restricted Coupling and Feasibility Define the block-restricted transport set Πblk(μ,ν)=π∈Π(μ,ν)|πij=0if xi∈(r),yj∈(m) with r≠m. _blk(μ,ν)= \π∈ (μ,ν)\; |\; aligned & _ij=0 x_i ^(r),\\ & y_j ^(m) with r≠ m aligned \. (10) The corresponding restricted OT problem is πblk⋆∈argminπ∈Πblk(μ,ν)∑i=1M∑j=1Nπij‖xi−yj‖2. _blk ∈ _π∈ _blk(μ,ν) _i=1^M _j=1^N _ij\|x_i-y_j\|^2. (11) The next result shows that this restriction costs nothing in terms of feasibility. Any block-diagonal coupling constructed under Assumption 1 is automatically a valid coupling for the unrestricted problem (5). Proposition 1 (Nonemptiness of the Coupling Set). Under Assumption 1, the set Πblk(μ,ν) _blk(μ,ν) defined in (10) is nonempty. Since Πblk(μ,ν)⊆Π(μ,ν) _blk(μ,ν) (μ,ν) by construction, every block-restricted coupling is therefore a feasible coupling for the original OT problem, and consequently provides an upper bound on the squared Wasserstein distance. Proof. The inclusion Πblk(μ,ν)⊆Π(μ,ν) _blk(μ,ν) (μ,ν) holds immediately from (10), which restricts Π(μ,ν) (μ,ν) by an additional sparsity condition. It remains to show Πblk(μ,ν)≠∅ _blk(μ,ν)≠ . Fix a block r. The total probability mass carried by the agent samples in this block is ∑xi∈(r)1M=MrM. _x_i ^(r) 1M= M_rM. Similarly, the total target mass in the corresponding block is ∑yj∈(r)1N=NrN. _y_j ^(r) 1N= N_rN. By Assumption 1, these quantities are equal. Therefore, a coupling between the normalized empirical measures supported on (r)X^(r) and (r)Y^(r) exists. After normalization by the common block mass, any standard finite discrete coupling can be constructed between the two finite uniform measures. Multiplying it by the common block mass gives a local coupling whose row sums are 1/M1/M and whose column sums are 1/N1/N. Perform this construction independently for every block and set all cross-block entries to zero. The resulting matrix π is nonnegative and satisfies ∑jπij=1M _j _ij= 1M for every agent i and ∑iπij=1N _i _ij= 1N for every target j, matching the marginal conditions in (4). Moreover, by construction πij=0 _ij=0 whenever xi∈(r)x_i ^(r) and yj∈(m)y_j ^(m) with r≠mr≠ m, which is exactly the block restriction in (10). Hence, π is a feasible global coupling satisfying the prescribed block restriction. In particular, such a coupling exists, thus Πblk(μ,ν) _blk(μ,ν) is nonempty. Finally, since a coupling satisfying the block restriction of (10) is, by definition, itself a coupling in Π(μ,ν) (μ,ν), and the Wasserstein cost is the minimum over Π(μ,ν) (μ,ν), any global coupling π satisfying the block restriction gives 22(μ,ν)≤∑i,jπij‖xi−yj‖2.W_2^2(μ,ν)≤ _i,j _ij\|x_i-y_j\|^2. -25.29494pt ∎ I-C Barycentric Targets Let π∈Π(μ,ν)π∈ (μ,ν) be the global transport plan obtained by assembling the local couplings from the corresponding blocks. Since its row marginal is uniform, ωi:=∑j=1Nπij=1M. _i:= _j=1^N _ij= 1M. The barycentric target associated with agent i is bi:=1ωi∑j=1Nπijyj=M∑j=1Nπijyj.b_i:= 1 _i _j=1^N _ijy_j=M _j=1^N _ijy_j. (12) Because πij≥0 _ij≥ 0 and ∑jπij=1/M _j _ij=1/M, the coefficients MπijM _ij form a probability vector. As a result, bi∈convyj:πij>0.b_i \y_j: _ij>0 \. (13) Every agent is therefore assigned a single reachable point rather than a distribution of targets. The barycenter also admits a natural variational characterization. Proposition 2. For every agent i, the point bib_i in (12) is the unique minimizer of Fi(z)=∑j=1Nπij‖z−yj‖2F_i(z)= _j=1^N _ij\|z-y_j\|^2 (14) over z∈ℝdz ^d. Proof. Differentiating gives ∇Fi(z)=2∑jπij(z−yj).∇ F_i(z)=2 _j _ij(z-y_j). The first-order optimality condition is therefore (∑jπij)z=∑jπijyj. ( _j _ij )z= _j _ijy_j. Since ∑jπij=1M>0, _j _ij= 1M>0, the unique stationary point is z=M∑jπijyj=bi.z=M _j _ijy_j=b_i. The Hessian is ∇2Fi(z)=2(∑jπij)I=2MI≻0,∇^2F_i(z)=2 ( _j _ij )I= 2MI 0, thus this stationary point is the unique global minimizer. ∎ I-D Barycentric Control Interpretation The control step requires a state-dependent objective that can be analyzed for a fixed transport assignment. Although the Wasserstein distance is the ultimate distribution-matching objective, its definition involves a minimization over the transport coupling, which generally depends on the agent states. Directly analyzing this nested optimization is therefore challenging for deriving agent-level control laws. To avoid the computational intractability, we freeze the transport coupling computed at the beginning of each control cycle. This design choice decouples the assignment from the ongoing motion, allowing us to treat the assigned targets as stationary references during the finite-horizon control phase. For any feasible coupling π∈Π(μ,ν)π∈ (μ,ν), its associated transport cost is ∑i=1M∑j=1Nπij‖xi−yj‖2. _i=1^M _j=1^N _ij\|x_i-y_j\|^2. Because π is itself a feasible coupling for the Wasserstein problem, this fixed-coupling cost provides an upper bound on the Wasserstein distance: 22(μ,ν)≤∑i=1M∑j=1Nπij‖xi−yj‖2.W_2^2(μ,ν)≤ _i=1^M _j=1^N _ij\|x_i-y_j\|^2. Thus, reducing the fixed-coupling transport cost reduces a certified upper bound on the Wasserstein objective. The connection between this fixed-coupling transport cost and barycentric target tracking has already been established in our previous work [7]. In particular, the weighted quadratic transport cost can be decomposed into a barycentric tracking term and a coupling-dependent residual. We restate this identity below in the notation used in this paper, as it provides the basis for the finite-horizon control analysis that follows. Defining the fixed-coupling transport surrogate cost Ψπ(x):=∑i=1M∑j=1Nπij‖xi−yj‖2, _π(x):= _i=1^M _j=1^N _ij\|x_i-y_j\|^2, (15) the following result holds. Proposition 3 (Barycentric Decomposition). For any feasible coupling π∈Π(μ,ν)π∈ (μ,ν), define the barycentric target of agent i as in (12). Then, the fixed-coupling transport cost satisfies Ψπ(x)=1M∑i=1M‖xi−bi‖2+Vπ, _π(x)= 1M _i=1^M\|x_i-b_i\|^2+V_π, (16) where Vπ=∑i=1M∑j=1Nπij‖yj−bi‖2≥0.V_π= _i=1^M _j=1^N _ij\|y_j-b_i\|^2≥ 0. (17) Moreover, VπV_π is independent of the agent states. The proposition follows directly from the weighted quadratic barycentric identity established in [7]. In particular, the cross term vanishes by the definition of bib_i, while the row marginal condition ∑jπij=1/M _j _ij=1/M yields the coefficient 1/M1/M in (16). The residual VπV_π represents the weighted dispersion of the target samples associated with each agent under the coupling. Consequently, driving the agents to their barycentric targets does not generally make the transport cost zero. If every row of π transports all of its mass to a single target, then Vπ=0V_π=0. In particular, for a one-to-one permutation coupling, exact barycentric reaching is equivalent to exact matching of the corresponding target samples. Because VπV_π is independent of the agent states for a fixed coupling, minimizing 1M∑i=1M‖xi−bi‖2 1M _i=1^M\|x_i-b_i\|^2 reduces Ψπ(x) _π(x) by exactly the same amount. Thus, the transport assignment determines independent barycentric targets, after which the control problem can be formulated as a collection of agent-level finite-horizon target-reaching problems. The novel aspect considered next is the use of the spatially partitioned coupling, which allows these transport assignments to be computed through smaller local problems. IV Finite-Horizon Barycentric Control Given the barycentric targets of Section I, this section addresses how each agent reaches, or contracts toward, its assigned target over a finite control horizon. Linear dynamics admit an exact closed-form controller while nonlinear dynamics are treated generically, through a contraction condition that any suitable controller or synthesis method may satisfy. As these results are standard finite-horizon control facts, we state them without proof and refer the interested reader to the full derivation in our related work [7]. IV-A Linear Dynamics Consider the LTI dynamics (1). Over a horizon of H steps, xi(k+H)=AHxi(k)+ΦHUi,x_i(k+H)=A^Hx_i(k)+ _HU_i, (18) where ΦH:=[AH−1BAH−2B⋯B] _H:= bmatrixA^H-1B&A^H-2B&·s&B bmatrix and Ui:=[ui(k)⊤ui(k+1)⊤⋯ui(k+H−1)⊤]⊤.U_i:= bmatrixu_i(k) &u_i(k+1) &·s&u_i(k+H-1) bmatrix . Assumption 2 (Finite-Horizon Reachability). The matrix ΦH _H has full row rank. This ensures reachability over horizon H, guaranteeing that arbitrary terminal states are achievable for feasible barycentric tracking. Lemma 1. (Minimum-Norm Exact Barycentric Control, Lemma 2 in [7]) Under Assumption 2, the minimum-ℓ2 _2-norm control sequence satisfying xi(k+H)=bix_i(k+H)=b_i is Ui⋆=ΦH⊤(ΦHΦH⊤)−1(bi−AHxi(k)).U_i = _H ( _H _H )^-1 (b_i-A^Hx_i(k) ). (19) Consequently, xi(k+H)=bix_i(k+H)=b_i for every agent i, implying that the barycentric target is reached exactly in H steps. IV-B Nonlinear Dynamics For the nonlinear dynamics (2), we characterize the control step through a terminal contraction condition with respect to the assigned barycentric target. Rather than specifying a particular nonlinear control synthesis method, we assume that the controller used during each control cycle satisfies the following condition. Assumption 3 (Barycentric Terminal Contraction). For some ρ∈[0,1)ρ∈[0,1), the controller satisfies ‖xi(k+H)−bi‖≤ρ‖xi(k)−bi‖\|x_i(k+H)-b_i\|≤ρ\|x_i(k)-b_i\| (20) for every agent i during the considered control cycle. This condition can, for example, be established through a finite-horizon feedback controller or verified for a nonlinear MPC implementation. Any control method satisfying (20) can be incorporated into the proposed assignment–control framework without changing the descent analysis in the following section. The exact LTI controller in Lemma 1 corresponds to the special case ρ=0ρ=0, yielding xi(k+H)=bix_i(k+H)=b_i at the end of the prescribed control horizon H. V Cycle-to-Cycle Descent and Wasserstein Guarantees Building on single-cycle results, this section establishes guarantees for the full alternating assignment-control iteration. Unlike standard receding-horizon control that updates actions at every step, the proposed framework applies the computed control sequence over the entire horizon H prior to re-solving the restricted optimal transport problem and updating barycentric targets. We demonstrate that the resulting surrogate cost Ψ decreases monotonically across cycles, bounds the true Wasserstein distance from above, and specializes to the exact LTI controller of Section IV. V-A Control Cycles and the Control-Step Descent Let the control cycles be indexed by ℓ=0,1,2,… =0,1,2,…, where each cycle spans a duration of H steps such that kℓ+1=kℓ+Hk_ +1=k_ +H. At the beginning of cycle ℓ , let πℓπ be an optimal global transport plan of the restricted OT problem at kℓk_ , i.e., πℓπ is a feasible global coupling whose nonzero entries satisfy the prescribed block restriction. Define the associated barycentric targets by biℓ=M∑j=1Nπijℓyj.b_i =M _j=1^N _ij y_j. During cycle ℓ , the coupling πℓπ and the barycentric targets biℓb_i are fixed. Define the cycle-level surrogate using (16) by Ψℓ:=Ψπℓ(x(kℓ)). _ := _π (x(k_ )). (21) To analyze reassignment, one additional feasibility condition is needed. The restricted coupling computed at the start of a cycle must remain a valid candidate coupling at the end of the cycle, so that the reoptimized coupling of the next cycle can be compared against it. Assumption 4 (Persistent Block Membership). The block membership of every agent remains unchanged over each control cycle. Consequently, the block-restricted sparsity pattern defining Πblk _blk does not change over the cycle, thus the coupling πℓπ computed at the beginning of the cycle remains feasible for the block-restricted OT problem defined at the new agent positions x(kℓ+1)x(k_ +1) at the end of the cycle. This assumption is naturally satisfied when partitions are defined by fixed agent assignments or when the control horizon is chosen such that agents do not migrate across block boundaries. Lemma 2 (Control-Step Descent). Under Assumption 3, the fixed-coupling surrogate satisfies Ψπℓ(x(kℓ+1))−Ψπℓ(x(kℓ)) _π (x(k_ +1))- _π (x(k_ )) (22) ≤−(1−ρ2)1M∑i=1M∥xi(kℓ)−biℓ∥2. ≤-(1-ρ^2) 1M _i=1^M\|x_i(k_ )-b_i \|^2. Proof. From (16), Ψπℓ(x(k))=1M∑i‖xi(k)−biℓ‖2+Vπℓ. _π (x(k))= 1M _i\|x_i(k)-b_i \|^2+V_π . Since the coupling is fixed throughout the cycle, the residual VπℓV_π is identical at kℓk_ and kℓ+1k_ +1. Assumption 3 gives, for every i, ‖xi(kℓ+1)−biℓ‖2≤ρ2‖xi(kℓ)−biℓ‖2.\|x_i(k_ +1)-b_i \|^2≤ρ^2\|x_i(k_ )-b_i \|^2. Summing over i and multiplying by 1/M1/M gives 1M∑i‖xi(kℓ+1)−biℓ‖2≤ρ2M∑i‖xi(kℓ)−biℓ‖2. 1M _i\|x_i(k_ +1)-b_i \|^2≤ ρ^2M _i\|x_i(k_ )-b_i \|^2. Subtracting the decomposition at kℓk_ from the decomposition at kℓ+1k_ +1 cancels the residual term and yields Ψπℓ(x(kℓ+1))−Ψπℓ(x(kℓ))≤ρ2−1M∑i‖xi(kℓ)−biℓ‖2, _π (x(k_ +1))- _π (x(k_ ))≤ ρ^2-1M _i\|x_i(k_ )-b_i \|^2, which is (22). ∎ V-B Cycle-to-Cycle Surrogate Descent Lemma 2 only accounts for the control step under a coupling that is fixed. The following theorem shows that reoptimizing the coupling at the end of the cycle can only help, so that the two effects combine into a single monotone descent guarantee for the full iteration. Theorem 1 (Cycle-to-Cycle Surrogate Descent). Suppose Assumptions 1, 3, and 4 hold. Let πℓπ be an optimal global transport plan satisfying the block restriction at the beginning of cycle ℓ , and let πℓ+1π +1 be an optimal global transport plan satisfying the block restriction computed at the beginning of cycle ℓ+1 +1. Then, the sequence Ψℓ≥0\ _ \_ ≥ 0 is monotonically nonincreasing, satisfying Ψℓ+1≤Ψℓ−(1−ρ2)1M∑i=1M‖xi(kℓ)−biℓ‖2. _ +1≤ _ -(1-ρ^2) 1M _i=1^M\|x_i(k_ )-b_i \|^2. (23) Proof. From Lemma 2, Ψπℓ(x(kℓ+1))≤Ψπℓ(x(kℓ))−(1−ρ2)1M∑i‖xi(kℓ)−biℓ‖2. _π (x(k_ +1))≤ _π (x(k_ ))-(1-ρ^2) 1M _i\|x_i(k_ )-b_i \|^2. (24) Consider the reassignment step at time kℓ+1k_ +1. By Assumption 4, the old global coupling remains feasible for the block-restricted transport problem at the new state. On the other hand, πℓ+1π +1 is defined as an optimizer over the same restricted feasible set. Therefore, ∑i,jπijℓ+1‖xi(kℓ+1)−yj‖2≤∑i,jπijℓ‖xi(kℓ+1)−yj‖2. _i,j _ij +1\|x_i(k_ +1)-y_j\|^2≤ _i,j _ij \|x_i(k_ +1)-y_j\|^2. By the definition in (21), Ψℓ+1≤Ψπℓ(x(kℓ+1)). _ +1≤ _π (x(k_ +1)). (25) Combining (24) and (25) yields Ψℓ+1 _ +1 ≤Ψπℓ(x(kℓ+1)) ≤ _π (x(k_ +1)) ≤Ψℓ−(1−ρ2)1M∑i‖xi(kℓ)−biℓ‖2, ≤ _ -(1-ρ^2) 1M _i\|x_i(k_ )-b_i \|^2, which proves (23) and completes the proof. ∎ Remark 1. Since 0≤ρ<10≤ρ<1, the coefficient 1−ρ21-ρ^2 is strictly positive. Consequently, if at least one agent is not positioned at its assigned barycentric target at the beginning of the cycle (xi(kℓ)≠biℓx_i(k_ )≠ b_i ), the decrement term is strictly positive, yielding strict descent (Ψℓ+1<Ψℓ _ +1< _ ). Remark 2. Theorem 1 is a statement about the complete assignment-control iteration. The first inequality (24) is due to barycentric control, while the second (25) follows from reoptimizing the restricted transport problem. Thus, unlike a fixed-coupling descent result, the theorem accounts for the reassignment step. The persistent block-membership condition is essential for this argument. If agents cross block boundaries between consecutive cycles, the previous coupling may no longer satisfy the new block sparsity constraint, and therefore it cannot in general be used as a feasible comparison point for the new restricted OT problem. V-C Wasserstein Upper Bound While the preceding analysis establishes the monotonic decrease of the transport surrogate Ψℓ _ across control cycles, its ultimate significance lies in how it bounds the true Wasserstein distance. Because every block-restricted transport plan constitutes a valid feasible point within the global transport set, the evaluated surrogate directly serves as an upper bound on the exact Wasserstein metric. Consequently, the cycle-to-cycle descent guarantees established for Ψℓ _ translate into certified performance bounds on the true distribution-matching objective without requiring the computationally expensive evaluation of the global optimal transport map at every step. Theorem 2 (Wasserstein Upper Bound). Let πℓπ be the global coupling used during cycle ℓ . Under Assumption 4, for every k∈[kℓ,kℓ+1]k∈[k_ ,k_ +1], 22(μk,ν)≤Ψπℓ(x(k)).W_2^2( _k,ν)≤ _π (x(k)). (26) In particular, 22(μkℓ+1,ν)≤Ψℓ+1≤Ψℓ.W_2^2( _k_ +1,ν)≤ _ +1≤ _ . (27) Proof. During cycle ℓ , the coupling πℓπ has row sums ∑jπijℓ=1M _j _ij = 1M and column sums ∑iπijℓ=1N. _i _ij = 1N. The agent dynamics change the locations xi(k)x_i(k) but do not change the mass associated with each agent. Therefore, the same matrix πℓπ has the prescribed marginal masses for the empirical distribution μk _k at every time in the cycle. By Assumption 4, the global coupling πℓπ continues to satisfy the required block restriction during the cycle. Consequently, it remains a feasible global transport coupling, yielding πℓ∈Π(μk,ν).π ∈ ( _k,ν). Using the definition of the Wasserstein distance, 22(μk,ν)=min∑i,jπ∈Π(μk,ν)πij‖xi(k)−yj‖2,W_2^2( _k,ν)= _π∈ ( _k,ν) _i,j _ij\|x_i(k)-y_j\|^2, and evaluating the objective at the feasible candidate πℓπ gives 22(μk,ν)≤∑i,jπijℓ‖xi(k)−yj‖2=Ψπℓ(x(k)).W_2^2( _k,ν)≤ _i,j _ij \|x_i(k)-y_j\|^2= _π (x(k)). This proves (26). At kℓ+1k_ +1, Theorem 1 gives Ψℓ+1≤Ψℓ. _ +1≤ _ . Since πℓ+1π +1 is itself a feasible global coupling, 22(μkℓ+1,ν)≤Ψℓ+1.W_2^2( _k_ +1,ν)≤ _ +1. Combining the inequalities proves (27). ∎ Remark 3. The preceding result does not imply that the true Wasserstein distance is itself monotonically decreasing at every cycle. The optimal global coupling can differ from the restricted coupling, and the upper-bound surrogate can decrease even when the Wasserstein optimum has a different cycle-to-cycle behavior. The rigorous guarantee is therefore monotonicity of the restricted surrogate together with an upper-bound certification of the Wasserstein distance. V-D The Linear Time-Invariant Case For the exact LTI controller of Lemma 1, the contraction factor is ρ=0ρ=0, and xi(kℓ+1)=biℓx_i(k_ +1)=b_i holds for every agent. Consequently, Assumption 3 holds with equality, and the control-step inequality of Lemma 2 reduces to an equality. Theorem 1 thus simplifies to Ψℓ+1≤Ψℓ−1M∑i=1M‖xi(kℓ)−biℓ‖2. _ +1≤ _ - 1M _i=1^M\|x_i(k_ )-b_i \|^2. (28) For the control step before reassignment, the inequality is in fact an equality: Ψπℓ(x(kℓ+1))−Ψπℓ(x(kℓ))=−1M∑i=1M∥xi(kℓ)−biℓ∥2. _π (x(k_ +1))- _π (x(k_ ))=- 1M _i=1^M\|x_i(k_ )-b_i \|^2. (29) The additional decrease in (28) stems from the subsequent restricted OT reassignment, which updates the coupling plan π at the new states to continually pull agents toward optimal target configurations across cycles. Thus, in the LTI case, the total cycle improvement consists of an exact barycentric control decrease followed by a further cost reduction from the reassignment, ensuring the monotonic decrease of the sequence Ψℓ\ _ \. VI Computational Complexity and Algorithm Summary VI-A Local Assignment and Computational Complexity A principal advantage of the proposed spatial partitioning architecture is its computational scalability. The block-restricted formulation can be implemented by solving one transport problem within each corresponding pair of spatial blocks. For each block r=1,…,Kr=1,…,K, the local transport problem is formulated as minπ(r)∑xi∈(r)∑yj∈(r)πij(r)∥xi−yj∥2, _π^(r) _x_i ^(r) _y_j ^(r) _ij^(r)\|x_i-y_j\|^2, (30) subject to ∑yj∈(r)πij(r)=1M,∑xi∈(r)πij(r)=1N,r=1,…,K. _y_j ^(r) _ij^(r)= 1M, 9.24994pt _x_i ^(r) _ij^(r)= 1N, 9.24994ptr=1,…,K. (31) At control cycle ℓ , the global restricted coupling πℓπ is obtained by assembling the K local optimal couplings π(r),ℓr=1K\π^(r), \_r=1^K along the block diagonal. Thus, the global restricted optimal transport problem is decoupled into K independent local transportation problems, which share no coupling variables and can be executed entirely in parallel at the beginning of each cycle. For balanced blocks of comparable size, the r-th local problem contains approximately M/KM/K agents and N/KN/K target samples. Thus, the proposed partitioning replaces one global transportation problem of size M×NM× N with K smaller problems of approximately (M/K)×(N/K)(M/K)×(N/K) size, substantially reducing the computational burden. In the special case where M=NM=N and each local problem reduces to a balanced one-to-one assignment, the dispersion term vanishes (Vπ=0V_π=0), and exact target-reaching by individual agents directly achieves exact empirical distribution matching. For this specific assignment regime, classical algorithms such as the Hungarian algorithm solve each local block of size M/KM/K with ((M/K)3)O((M/K)^3) worst-case complexity [15]. Consequently, solving the K local assignment problems yields a total computational complexity of K((MK)3)=(M3K2),KO\! ( ( MK )^\!3 )=O\! ( M^3K^2 ), compared with (M3)O(M^3) for the centralized assignment. This architectural decomposition thus achieves an ideal computational reduction of (K2)O(K^2) in the equal-cardinality regime. VI-B Algorithm Summary Building upon the localized sub-problem formulation via partitioning, the overall framework is implemented as an alternating assignment and control procedure. The spatial partition of the domain is constructed offline and fixed prior to execution, defining persistent agent-block memberships that satisfy Assumption 4. The complete computational pipeline is summarized in Algorithm 1. Algorithm 1 Partition-Restricted OT and Barycentric Control 1: Initial agent states xi(0)i=1M\x_i(0)\_i=1^M, target samples yjj=1N\y_j\_j=1^N, mass-balanced partition, horizon H, terminal cycle R 2: for ℓ=0,1,…,R−1 =0,1,…,R-1 do 3: Freeze current agent states xi(kℓ)i=1M\x_i(k_ )\_i=1^M at cycle ℓ . 4: parfor r=1,…,Kr=1,…,K do 5: Solve the local OT problem (30)–(31) using the current states in block (r)X^(r) to obtain π(r),ℓπ^(r), . 6: end parfor 7: Assemble the global coupling πℓ=blockdiag(π(1),ℓ,…,π(K),ℓ)π =blockdiag(π^(1), ,…,π^(K), ). 8: Compute the barycentric targets biℓ=M∑j=1Nπijℓyjb_i =M _j=1^N _ij y_j. 9: if linear dynamics then 10: Apply (19) for H steps. 11: else 12: Apply a controller satisfying (20). 13: end if 14: end for At the onset of each control cycle ℓ , the framework updates the transport coupling by solving independent local sub-problems, computes the resulting barycentric targets, and executes the finite-horizon control law. By Theorem 1, this iterative execution guarantees the monotonic nonincreasing behavior of the surrogate Ψℓ _ . VII Simulations This section evaluates the distribution-matching behavior, cycle-to-cycle surrogate descent, and computational benefit of the proposed partitioned transport formulation. The simulation uses M=30M=30 agents with the LTI dynamics (1), where A=[0.90.100.9] A= bmatrix0.9&0.1\\ 0&0.9 bmatrix, B=[00.1], B= bmatrix0\\ 0.1 bmatrix, and N=500N=500 target samples drawn from a mixture of four anisotropic 2-D Gaussians. The agents initially form a cluster outside the target support. Each control cycle uses H=50H=50 steps, and the assignment-control procedure is repeated for 20 cycles. For the partitioned case, the target samples are divided into K=5K=5 spatial blocks, and the corresponding agent blocks are chosen to satisfy the mass-balance condition in Assumption 1. Agent block membership is kept fixed across cycles, consistent with Assumption 4. Both centralized and local assignments are solved exactly as balanced transportation linear programs. (a) (b) Fig. 1: Multi-agent distribution matching trajectories: (a) centralized (non-partitioned) optimal transport and (b) partition-restricted transport with K=5K=5 blocks. Fig. 1 compares the centralized and partition-restricted solutions. In both cases, the agents disperse from the initial cluster and distribute across the four target modes. Proposition 1 guarantees that the mass-balanced restricted coupling is feasible for the original OT problem. Although the target samples are partitioned by K distinct blocks in the restricted OT method, its overall qualitative behavior for the terminal distribution matching shows a similar performance with respect to the original centralized OT method. Fig. 2: Cycle-to-cycle partitioned surrogate square root Ψℓ _ and the actual Wasserstein distances 2W_2 for partitioned and centralized methods over successive control cycles. Fig. 2 illustrates the square root of the restricted surrogate, Ψℓ _ in (21) together with the corresponding Wasserstein distances over successive control cycles. The partitioned surrogate square root Ψℓ _ decreases monotonically across all cycles, consistent with the cycle-to-cycle descent guaranteed by Theorem 1. Moreover, the computed Wasserstein distance at each cycle boundary satisfies 2(μkℓ+1,ν)≤Ψℓ+1W_2( _k_ +1,ν)≤ _ +1, as certified by Theorem 2. Thus, Fig. 2 numerically verifies both the monotone decrease of the partitioned surrogate and its upper-bound relationship with the true Wasserstein metric. While the centralized method achieves a slightly lower Wasserstein distance (better distribution matching) due to its unrestricted transport flexibility, the partitioned approach provides a favorable trade-off, securing substantial computational speedups with only a comparable sacrifice in matching performance. Fig. 3: Assignment solve times in logarithmic scale: (left) versus swarm size M, and (right) versus target resolution N. Fig. 3 illustrates the practical benefits of this restriction. At the largest tested scale (M=300M=300, N=1500N=1500), the centralized approach requires 5.48s5.48\,s, while K=5K=5 and K=10K=10 achieve solve times of 0.49s0.49\,s (11×11× speedup) and 0.21s0.21\,s (26×26× speedup), respectively. A comparable trend appears when scaling N at a fixed M=40M=40, where K=10K=10 operates roughly 22×22× faster than the centralized baseline, and in both sweeps the speedup grows with problem size rather than saturating. This confirms that partitioning into K blocks reduces each local problem to a 1/K1/K fraction of the agents and target samples, and the resulting reduction in aggregate solve time is substantial and improves further as K increases from 5 to 10. The empirical gains are therefore fully consistent with the computational motivation for the partitioned formulation. These simulation results demonstrate the three intended properties of the proposed framework: feasibility of the mass-balanced restriction, monotone descent of the restricted surrogate, and reduced assignment computation through spatial decomposition. VIII Conclusion This paper presented a scalable framework for optimal-transport-based distribution matching by decomposing a global assignment into K spatially restricted local subproblems. By enforcing a block mass-balance condition, the restricted coupling preserves feasibility for the original transport problem, certifying a rigorous Wasserstein upper bound. Combined with finite-horizon barycentric control, the formulation guarantees a monotonic cycle-to-cycle descent of the transport surrogate, which upper-bounds the true Wasserstein metric. Numerical evaluations demonstrated that this spatial restriction significantly enhances computational efficiency, reducing assignment solve times by an order of magnitude or more while maintaining matching performance comparable to global methods. Consequently, the proposed approach renders large-scale distribution matching tractable without compromising theoretical convergence or feasibility guarantees. References [1] J. Cortés, S. Martínez, T. Karatas, and F. Bullo (2004) Coverage control for mobile sensing networks. IEEE Transactions on Robotics and Automation 20 (2), p. 243–255. External Links: Document Cited by: §I. [2] D. Inoue, Y. Ito, and H. Yoshida (2021) Optimal transport-based coverage control for swarm robot systems: generalization of the voronoi tessellation-based method. IEEE Control Systems Letters 5 (4), p. 1483–1488. External Links: Document Cited by: §I. [3] S. Bandyopadhyay, S. Chung, and F. Y. Hadaegh (2014) Probabilistic swarm guidance using optimal transport. In 2014 IEEE Conference on Control Applications (CCA), p. 498–505. External Links: Document Cited by: §I. [4] S. Seo and K. Lee (2025) Density-driven optimal control for efficient and collaborative multiagent nonuniform coverage. IEEE Transactions on Systems, Man, and Cybernetics: Systems 55 (12), p. 9340–9354. External Links: Document Cited by: §I. [5] S. Seo and K. Lee (2026) Density-driven multidrone coordination for efficient farm coverage and management in smart agriculture. IEEE Transactions on Control Systems Technology 34 (2), p. 711–724. External Links: Document Cited by: §I. [6] K. Lee (2026) Density-driven optimal control: convergence guarantees for stochastic lti multi-agent systems. Automatica 193, p. 113231. External Links: Document Cited by: §I. [7] K. Lee (2026) Optimal transport-based decentralized multi-agent distribution matching. IEEE Transactions on Automatic Control 71 (8), p. 5478–5485. External Links: Document Cited by: §I, §I, §I-C, §I, §I-D, §I-D, §IV, Lemma 1. [8] C. Villani et al. (2009) Optimal transport: old and new. Vol. 338, Springer. Cited by: §I. [9] L. V. Kantorovich (2006) On the translocation of masses.. Journal of mathematical sciences 133 (4), p. 1381. Cited by: §I. [10] P. Gabriel and C. Marco (2019) Computational optimal transport with applications to data sciences. Foundations and trends® in machine learning 11 (5-6), p. 355–607. Cited by: §I, §I. [11] Y. Chen, T. T. Georgiou, and M. Pavon (2021) Optimal transport in systems and control. Annual Review of Control, Robotics, and Autonomous Systems 4, p. 89–113. External Links: Document Cited by: §I. [12] Y. Chen (2024) Density control of interacting agent systems. IEEE Transactions on Automatic Control 69 (1), p. 246–260. External Links: Document Cited by: §I. [13] V. Krishnan and S. Martínez (2025) Distributed online optimization for multi-agent optimal transport. Automatica. External Links: Document Cited by: §I. [14] V. Hartmann and D. Schuhmacher (2020) Semi-discrete optimal transport: a solution procedure for the unsquared euclidean distance case. Mathematical Methods of Operations Research 92 (1), p. 133–163. External Links: Document Cited by: §I. [15] H. W. Kuhn (1955) The hungarian method for the assignment problem. Naval research logistics quarterly 2 (1-2), p. 83–97. Cited by: §VI-A.