I build neural networks that learn solution operators: one trained model
that solves a whole family of control, transport and planning problems, trained on the problem's
own objective rather than on precomputed solutions.
In-context operator learning for stochastic mean-field control
With Mo Zhou and Rongjie Lai
Stochastic mean-field control coordinates large populations of interacting agents under
uncertainty. Existing numerical and deep-learning methods solve one problem instance at a time
and must be re-optimized whenever the task changes. We instead treat stochastic mean-field control
as an operator-learning problem: learn the map from a task to its solution, once.
The main obstacle is the diffusion term in the controlled Fokker–Planck equation, which rules
out deterministic transport maps. Our model, the Normalizing Flow Invertible Solution
Transformer (NFIST), combines the probability-flow ODE with an invertible
normalizing-flow-based transformer. This recasts the dynamics as a deterministic continuity
equation and gives the score in closed form, through the exact inverse and log-determinant of the
flow, at 𝒪(d) cost per particle.
Mesh-free and self-supervised: trained end-to-end on the stochastic control objective, with no
precomputed numerical solutions.
In-context learning: a task prompt, given as compact distribution parameters or as a raw
particle cloud, conditions the transport map, so one pretrained operator solves unseen tasks in a
single forward pass.
The operator-learning formulation is proven consistent with task-by-task optimization.
Zero-shot results on stochastic optimal control, Schrödinger bridges, systemic-risk control and
obstacle-avoiding path planning.
Neural operators for spacecraft-swarm trajectory planning
With Sidhdharth D. Sikka, Zehui Lu, Rongjie Lai and Shaoshuai Mou
Satellite constellations need orbital transfers that save fuel and avoid collisions. Classical
trajectory optimization scales poorly here, with both the number of spacecraft and the number of
obstacles, because every pair needs a safety constraint.
We train a permutation-equivariant neural operator that maps the initial states and targets of a
swarm, together with the obstacles, to collision-free, fuel-conserving trajectories. A batched
Gauss–Newton step then enforces the exact orbital dynamics and reduces fuel use further.
Self-supervised: trained without optimal trajectory labels.
Trained on 10 spacecraft, it generalizes zero-shot to swarms of 1,000 spacecraft and 11,000
obstacles.
Matches the accuracy of a per-agent optimal control solver while keeping collision
avoidance.
Combining diffusion models with classical optimization to recover point configurations from
incomplete or noisy data, a question in Euclidean distance geometry. This work is in progress.
Summers 2019 – 2021Undergraduate research, McDaniel College
Analysis on the 2-nacci fractal
The 2-nacci word fractal at the critical angle α = arccos(1/3): level 4 of its
iterated function system.
The Fibonacci word fractal is drawn from the Fibonacci words, each the concatenation of the
two before it (fn = fn−1fn−2), by a turning rule with a turning
angle α. I worked on a variation, the 2-nacci word fractal, built from the words
fn = fn−1fn−1fn−2. It is also the attractor of an
iterated function system of six similarities with ratio 1/c; at the critical angle
α = arccos(1/3), c = 3.
That system overlaps itself (it fails the open set condition), so with Benjamin Steinhurst I
built the fractal instead as a projective limit of graphs, and showed that the limit is
homeomorphic to the attractor. Its finitely ramified cell structure supports a Laplacian: we
assembled the Laplacian matrices of four graph approximations, computed their eigenvalues and
eigenvectors, and estimated the spectral dimension from the eigenvalue counting function.
I presented the projective-limit construction at MAA MathFest 2019, and the analysis at the
critical angle in a poster at the Joint Mathematics Meetings in 2022, held virtually.
Stochastic mean-field control (MFC) provides a fundamental framework for coordinating
large populations of interacting agents under uncertainty, with a wide range of
applications. Existing numerical and deep-learning methods solve one MFC problem instance
at a time and must be re-optimized whenever the task changes. In this work, we formulate
stochastic MFC as an operator-learning problem and develop, to the best of our knowledge,
the first mesh-free, self-supervised neural operator for stochastic MFC. The main
challenge is that the diffusion term in the controlled Fokker–Planck equation precludes
deterministic transport-map representations. We address this challenge by combining the
probability-flow ODE with an invertible normalizing-flow-based transformer, which recasts
the dynamics as a deterministic continuity equation and enables closed-form score
evaluation through the exact inverse and analytical log-determinant of the normalizing
flow, with 𝒪(d) cost per particle for networks of fixed
size. Through transformer-based in-context learning, task prompts, represented by compact
distribution parameters or raw particle clouds, condition the transport map, enabling a
single pretrained operator to solve unseen tasks in one forward pass. The resulting
Normalizing Flow Invertible Solution Transformer (NFIST) is trained end-to-end by
minimizing the stochastic control objective directly, requiring no precomputed numerical
solutions for training. We further prove the consistency of the proposed operator-learning
formulation with task-by-task optimization. Numerical experiments on stochastic optimal
control, Schrödinger bridge, systemic-risk control, and obstacle-avoiding path planning
demonstrate effective zero-shot generalization while substantially reducing the
computational cost of solving large families of stochastic MFC problems.
@misc{gao2026selfsupervised,
title = {Self-supervised In-context Operator Learning for
Stochastic Mean-Field Control},
author = {Gao, Suyi and Zhou, Mo and Lai, Rongjie},
year = {2026},
eprint = {2608.18282},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2608.18282}
}
Satellite constellations require orbital transfers that are both fuel efficient and
collision avoidant. Yet, the computational cost of optimization methods traditionally used
to plan their trajectories scales poorly with both the number of satellites as well as the
number of obstacles to avoid, due to the pairwise safety constraints. In this work, we
introduce a permutation-equivariant neural operator for trajectory planning of spacecraft
swarms. This neural operator maps distributions of spacecraft initial states, target
states, and obstacle initial states to trajectories which avoid collision and conserve
fuel. This neural operator output is then paired with a batched Gauss-Newton finish to
enforce exact orbital dynamics, and further reduce fuel use. The operator is
self-supervised, trained without optimal trajectory labels. When trained on ten
spacecraft, the proposed method generalized zero-shot to swarms of 1,000 spacecraft and
11,000 obstacles. The generated trajectories matched a per-agent optimal control solver's
accuracy while retaining collision avoidance. Operator learning grounded in physics may
offer a fast, scalable alternative to trajectory optimization in the increasingly crowded
orbits of the future.
@misc{sikka2026neural,
title = {Neural Operator Learning for Collision-Aware
Trajectory Planning of Spacecraft Swarms},
author = {Sikka, Sidhdharth D. and Gao, Suyi and Lu, Zehui
and Lai, Rongjie and Mou, Shaoshuai},
year = {2026},
eprint = {2608.00320},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2608.00320}
}
Talks & Posters
TalkSelf-supervised In-context Operator Learning for Stochastic Mean-Field ControlMinisymposium, SIAM Great Lakes Section Annual Meeting · Purdue UniversitySep 2026
PosterLaplacian on the 2-nacci Fractal at the Critical AngleJoint Mathematics Meetings (JMM 2022) · VirtualJan 2022
TalkA Projective Limit Construction of the 2-nacci FractalMAA MathFest 2019 · Cincinnati, OHAug 2019