A Transformer-based Approach for Real-time Multi-agent Path Planning with Varying Initial/Final Conditions

Qian Zhang; Shanqing Liu; Jerome Darbon; George Karniadakis

math.OCOct 21, 2026

A Transformer-based Approach for Real-time Multi-agent Path Planning with Varying Initial/Final Conditions

Authors: Qian Zhang, Shanqing Liu, Jerome Darbon, George Karniadakis

Organizations: Division of Applied Mathematics, Brown University, Providence, RI 02912, USA.

Abstract

We address the problem of real-time multi-agent path planning under prescribed initial and final configurations, formulated as a high-dimensional optimal control problem subject to inter-agent collision avoidance and obstacle constraints. Classical solvers based on the Hamilton--Jacobi--Bellman equation or Pontryagin's maximum principle (PMP) become computationally intractable as the number of agents grows, due to the curse of dimensionality and the non-convexity introduced by interaction constraints. We propose OCFormer, a transformer-based neural network that parameterizes the inverse of a coordinate transformation mapping a simple, analytically tractable latent space, in which agents evolve independently under a decoupled Hamiltonian, to the physical space of the coupled, interacting system. The model is trained by minimizing a physics-informed loss comprising the Hamiltonian ODE residual from the PMP and a boundary matching term, without requiring pre-computed trajectory data. The self-attention mechanism endows OCFormer with permutation equivariance over agents, a structural symmetry that prior symplectic network approaches do not natively enforce. When trained on a distribution of boundary conditions, the amortized OCFormer produces near-optimal trajectories for unseen initial and final configurations in a single forward pass ( $\sim$ 220ms), enabling real-time inference without retraining. We further introduce a lightweight transfer learning framework that adapts a pretrained OCFormer to out-of-distribution perturbations in seconds via L-BFGS fine-tuning. Numerical experiments on two- and three-dimensional multi-agent navigation with up to 128 agents demonstrate the scalability and efficiency of the proposed approach.

Figures & tables

Figure 1. Flowchart of the OCFormer framework. For a pair of prescribed initial and terminal configurations, boundary data are first encoded into latent variables $(\boldsymbol{y},\boldsymbol{q})$ , which serve as the input to the transformer-based inverse coordinate map $\phi^{-1}$ . The network predicts the physical position and momentum trajectories $(\boldsymbol{x},\boldsymbol{p})$ over the time grid without solving an ODE at inference time. During training, the predicted trajectories are differentiated with respect to time and evaluated in the Hamiltonian system so that the model satisfies the physics-informed residual loss associated with Pontryagin’s maximum principle, while a boundary matching term enforces the initial and final constraints. The total loss is minimized over the network parameters, yielding a model that directly maps varying boundary conditions to feasible multi-agent trajectories in real time.

Figure 2. Coordinate-transformation view of the proposed trajectory representation. The forward map

\phi

sends the physical trajectory pair

(\boldsymbol{x},\boldsymbol{p})

to a simpler latent trajectory pair

(\boldsymbol{y},\boldsymbol{q})

associated with a decoupled Hamiltonian system, while the inverse map

\phi^{-1}

reconstructs the interacting multi-agent trajectories in physical space.

Figure 3. Architecture of OCFormer.

y,q

are first reshaped into

\mathbb{R}^{N\times d}

\mathcal{E}_{y},\mathcal{E}_{q}:\mathbb{R}^{d}\to\mathbb{R}^{D}

are tokenwise embedding layers, which are applied on the last axis of reshaped input.

z_{i},r_{i}\in\mathbb{R}^{N\times D}

. Partial transformer blocks

\mathcal{T}_{up}

and

\mathcal{T}_{low}

are depicted by the dashed rounded rectangle.

I_{D}

is identity mapping in

\mathbb{R}^{D}

\mathcal{T}_{i}

(

i=1,\cdots,n

) are the transformer blocks.

\mathcal{D}_{x},\mathcal{D}_{p}:\mathbb{R}^{D}\to\mathbb{R}^{d}

are tokenwise deembedding MLPs, whose output are flattened back into

\mathbb{R}^{Nd}

I_{Nd}

is the identity mapping in

\mathbb{R}^{Nd}

, serving as residual connection.

$M$	Layout	$r$	$J_{\text{LB}}$	OCFormer	Gap	SympOCNet	Feasible?
4	$1\times 4$	0.02	0.200	0.2003	0.2%	0.2001	Yes
8	$1\times 8$	0.02	0.400	0.4036	0.9%	0.4046	Yes
16	$1\times 16$	0.02	0.800	0.8307	3.8%	0.8345	Yes
32	$2\times 16$	0.02	1.156	1.2531	8.4%	1.4413	Yes
64	$2\times 32$	0.01	2.297	2.448	6.6%	—	—
128	$4\times 32$	0.01	4.704	5.1424	9.3%	13.3 1 1 1 SympOCNet did not converge for $M=128$ ; the best cost achieved was $13.3$ with constraint violation $-0.002$ , compared to OCFormer’s $5.14$ with constraint $+1.2\times 10^{-4}$ .	No

Table 1. Quantitative results for the 2D room symmetry setup, and direct comparison against SympOCNet (JAX implementation) on the identical circle benchmarks, using the same boundary data, parameters (

r

T

, room size,

\epsilon

l

\lambda

), and training budget (100k–500k Adam + L-BFGS on an NVIDIA RTX A6000).

J_{\text{LB}}=\sum_{i}\|\Delta x_{i}\|^{2}/(2T)

is the analytic lower bound (unconstrained optimum). Cost is evaluated via trapezoidal integration of

\frac{1}{2}\|v\|^{2}

; Gap is relative to

J_{\text{LB}}

. SympOCNet was not evaluated at

M=64

. For every configuration, OCFormer’s fine-grid constraint minimum

\min_{i\neq j,t}\|x_{i}-x_{j}\|-2r

is positive, in the narrow range

[1.2,1.5]\times 10^{-4}

, confirming collision-free trajectories throughout.

Figure 4. Predicted trajectories at

t=0,T/3,2T/3,T

for

M=128

drones, the largest and most demanding configuration in Table 1 . Drones shown as colored circles with paths as gray curves.

Panel of Figure 5 — Figure 5. Predicted multi-agent trajectories for the random-perturbation experiment with $n=32$ drones at times $t=0,\frac{1}{3},\frac{2}{3},1$ from left to right. Starting from the perturbed initial condition, the transferred model generates collision-free paths that progressively reorganize the swarm toward the prescribed terminal arrangement. The intermediate snapshots show that the learned corrections adapt the reference solution while preserving the overall interaction structure encoded by the original OCFormer.

$M$	Test cost (mean $\pm$ std)	Gap vs. $J_{\text{LB}}$	Feasibility	Inference (ms)
8	$0.4053\pm 2.6\!\times\!10^{-6}$	$1.3\%$	$100\%$	$218.6\pm 1.3$
16	$0.8399\pm 7.6\!\times\!10^{-6}$	$5.0\%$	$75\%$	$220.9\pm 1.2$

Table 2. Zero-shot amortized inference results. Trained on 100 rotated boundary pairs, tested on 20 held-out rotations. Inference time is the wall-clock time for a single forward pass (JIT-compiled, averaged over 100 runs).

Panel of Figure 6 — Figure 6. Predicted trajectories for the 3D room experiment with $n=100$ drones at times $t=0,\frac{1}{3},\frac{2}{3},1$ from left to right. The snapshots illustrate how the learned model coordinates a large collection of agents in a confined three-dimensional domain, moving them from their initial configuration to the terminal one while maintaining collision-free motion throughout the room. The intermediate frames highlight the scalability of OCFormer to high-dimensional multi-agent path-planning problems.

Panel of Figure 7 — Figure 7. Predicted trajectories for the 3D ball-obstacle experiment at times $t=0,\frac{1}{3},\frac{2}{3},1$ from left to right. The figure shows the agents moving from their initial to terminal positions while simultaneously avoiding pairwise collisions and steering around the spherical obstacle. The intermediate configurations make clear that the learned dynamics capture both inter-agent avoidance and obstacle-avoidance behavior in a genuinely three-dimensional setting.

Decentralized Multi-Agent Motion Planning

See all 20 papers

Permutation Robustness Is Not Enough: Action Collapse in Multi-Agent Transformer Policies
Oct 2, 2026Amit Thakur, Mukesh SinghalMulti-Agent Systems Multi-Agent Reinforcement Learning
University of California, Merced
PRIMAL3: Pathfinding via Reinforcement and Imitation Multi-Agent Learning - Leveraging LaCAM3
Aug 5, 2026Chengyang He, Tanishq Duhan, Gadiel Sznaier Camps +6Multi-Robot Systems Multi-Agent Path Finding
Multi-Agent Robotic Motion Lab (MARMot), National University of Singapore · Multi-robot Systems Lab (MSL), Stanford University · Robotics and Machine Intelligence Lab (ROMI), Hong Kong Polytechnic University
NaviFormer: A Deep Reinforcement Learning Transformer-like Model to Holistically Solve the Navigation Problem
Apr 18, 2026Daniel Fuertes, Andrea Cavallaro, Carlos R. del-Blanco +2Reinforcement Learning Robot Navigation
Grupo de Tratamiento de Im´agenes, Information Processing and Telecommunications Center, ETSI Telecomunicaci´on, Universidad Polit´ecnica de Madrid, 28040 Madrid, Spain · Idiap Research Institute, Martigny, Switzerland

Oct 2, 2026J/K move · Enter open · S save

Permutation Robustness Is Not Enough: Action Collapse in Multi-Agent Transformer Policies

Amit Thakur, Mukesh Singhal

University of California, Merced

Multi-Agent Systems Multi-Agent Reinforcement Learning Cooperative MARL Transformer-Based RL Multi-Agent Coordination

Open paper AskPDF arXiv

Transformer policies are attractive for multi-agent robot learning because self-attention can model interactions among agents. However, multi-agent teams are unordered, while transformers typically process agents as ordered token sequences. We study how this mismatch affects cooperative navigation policies under agent-order permutations. Our results show that low permutation error alone can be misleading: policies may appear robust simply because all agents choose the same action. We therefore evaluate policies using both permutation-consistency metrics and action-collapse diagnostics, including action diversity, same-action fraction, and maximum action frequency. A PPO-ID baseline yields non-collapsed behavior but remains order-sensitive, while strong equivariance regularization can still induce homogeneous behavior. A weak equivariance penalty improves the robustness while preserving more diverse actions for teams with

N=3

agents, whereas teams with

N=4

agents require substantially smaller regularization weights. These findings suggest that multi-agent transformer policies should be evaluated not only by return and permutation robustness, but also by whether they maintain non-collapsed, differentiated multi-agent behavior.

Abstract

Decentralized Multi-Agent Motion Planning

Explore similar work

Permutation Robustness Is Not Enough: Action Collapse in Multi-Agent Transformer Policies

PRIMAL3: Pathfinding via Reinforcement and Imitation Multi-Agent Learning - Leveraging LaCAM3

NaviFormer: A Deep Reinforcement Learning Transformer-like Model to Holistically Solve the Navigation Problem