Self-supervised In-context Operator Learning for Stochastic Mean-Field Control

Suyi Gao; Mo Zhou; Rongjie Lai

math.OCAug 18, 2026

Self-supervised In-context Operator Learning for Stochastic Mean-Field Control

Authors: Suyi Gao, Mo Zhou, Rongjie Lai

Organizations: Department of Mathematics, Purdue University, West Lafayette, USA · Department of Mathematics, University of California, Los Angeles, USA

Abstract

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 $\mathcal{O}(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.

Figures & tables

Panel of Figure 1 — Figure 1: Architecture of NFIST. Top: The solution operator is realized as a stack of affine coupling blocks. Middle: Each coupling block uses a transformer to produce the affine parameters $(\alpha,\beta)$ that define the coupling transformation. Bottom: The transformer conditions the flow on the task prompt. Query particles $(x,t)$ are processed by self-attention and cross-attention with a context $c$ constructed from either non-parametric particle-cloud prompts or parametric prompt tokens.

Figure 2: Single-well one-dimensional SOC problem, one solution operator netowrk (

\gamma=0.1

; random initial mean/variance and well location), on three tasks (strips). In both rows the solid curve is the operator output and the dashed curve the exact reference, drawn on top, so exact agreement appears as a dashed curve riding on a solid one.

Figure 3: One SOC solution operator evaluated zero-shot over the plane of well centers (

d=2

\gamma=0.1

). The SDE trajectories are simulated according to section 2.4 .

		Inference Time (s)		Relative error
dim	$\bm{\gamma}$	transport	$+$ score	trajectory	score
$1$	$0.01$	0.013	0.752	3.60e-4 $\pm$ 2.78e-4	1.11e-3 $\pm$ 7.99e-4
$1$	$0.1$	0.013	1.066	1.11e-3 $\pm$ 6.39e-4	3.29e-3 $\pm$ 1.74e-3
$1$	$1.0$	0.013	0.650	2.57e-3 $\pm$ 1.51e-3	6.84e-3 $\pm$ 3.40e-3
$2$	$0.01$	0.012	0.486	2.45e-3 $\pm$ 1.11e-3	9.67e-3 $\pm$ 4.13e-3
$2$	$0.1$	0.010	0.539	1.07e-3 $\pm$ 4.19e-4	4.83e-3 $\pm$ 1.50e-3

Table 1: Single-well SOC vs. the closed-form Gaussian reference: per-task inference time and the relative errors ( 33 ) and ( 34 ). The

d=1

family additionally randomizes the initial mean and variance.

Figure 4: Single-well SOC with the non-parametric particle-cloud prompt (

d=1

\gamma=0.1

): the operator is conditioned on raw

100

-point sample clouds

X_{0}

rather than on

(m_{0},\sigma_{0})

. Same three tasks and panel layout as the parametric Figure 2 .

Figure 5: One non-parametric SOC solution operator evaluated zero-shot on sampled tasks (

d=2

\gamma=0.1

). The SDE trajectories are simulated according to section 2.4 .

		Inference Time (s)		Relative error
dim	$\bm{\gamma}$	transport	$+$ score	trajectory	score
$1$	$0.01$	0.014	1.370	1.73e-2 $\pm$ 1.66e-2	4.87e-2 $\pm$ 4.59e-2
$1$	$0.1$	0.014	0.651	1.88e-2 $\pm$ 1.66e-2	4.93e-2 $\pm$ 4.00e-2
$1$	$1.0$	0.014	1.373	3.17e-2 $\pm$ 2.26e-2	8.25e-2 $\pm$ 5.72e-2
$2$	$0.01$	0.014	1.346	1.36e-2 $\pm$ 8.21e-3	4.72e-2 $\pm$ 2.42e-2
$2$	$0.1$	0.014	0.664	1.24e-2 $\pm$ 7.47e-3	4.00e-2 $\pm$ 2.02e-2

Table 2: Non-parametric (particle-cloud) single-well SOC vs. the exact Gaussian reference.

Figure 6: Multi-well SOC in the sharply multimodal

\gamma=0.5

regime of three tasks. Each strip samples the score along three high-density slices of the terminal law (as in the 2-D plot on the left).

		Inference Time (s)		Relative error
prompt type	dim	transport	$+$ score	trajectory	score
Parametric	$2$	0.012	1.331	2.16e-2 $\pm$ 4.32e-3	7.25e-2 $\pm$ 1.23e-2
Parametric	$5$	0.012	1.330	4.29e-2 $\pm$ 3.00e-3	1.10e-1 $\pm$ 7.47e-3
Parametric	$8$	0.012	1.364	4.44e-2 $\pm$ 2.51e-3	1.07e-1 $\pm$ 4.58e-3
Parametric	$11$	0.012	1.380	4.56e-2 $\pm$ 1.61e-3	1.02e-1 $\pm$ 2.78e-3
Parametric	$14$	0.012	1.311	4.48e-2 $\pm$ 1.97e-3	9.57e-2 $\pm$ 3.05e-3

Table 3: Multi-well SOC with wells spanning

\mathbb{R}^{d}

vs. the Monte-Carlo Hopf–Cole reference of B.2 : per-task inference time and the relative errors ( 33 ) and ( 34 ), the trajectory scored against the reference probability flow ( 39 ).

Figure 7: Schrödinger bridge, one-dimensional parametric solution operator (

\gamma=1.0

; random mean and variance at both endpoints).

Figure 8: One bridge solution operator (

d=2

\gamma=5.0

), solved zero-shot; each panel framed to its own task. The prescribed terminal law

P_{1}

(green shading, dashed

3\sigma

) is the bridge’s exact terminal marginal; trajectories (blue) carry each initial population (purple) onto it under fixed weights. The SDE trajectories are simulated according to section 2.4 .

		Inference Time (s)		Relative error
dim	$\bm{\gamma}$	transport	$+$ score	trajectory	score
$1$	$1.0$	0.014	0.714	3.27e-2 $\pm$ 2.47e-2	6.57e-2 $\pm$ 4.67e-2
$1$	$5.0$	0.013	1.362	3.13e-2 $\pm$ 1.83e-2	6.90e-2 $\pm$ 2.93e-2
$2$	$1.0$	0.012	1.379	2.12e-2 $\pm$ 1.16e-2	6.83e-2 $\pm$ 3.42e-2
$2$	$5.0$	0.012	1.316	2.16e-2 $\pm$ 1.04e-2	6.14e-2 $\pm$ 2.45e-2
$5$	$1.0$	0.012	1.322	2.28e-2 $\pm$ 7.57e-3	6.15e-2 $\pm$ 1.92e-2

Table 4: Schrödinger bridge vs. the closed-form Gaussian reference: per-task inference time and the relative errors ( 33 ) and ( 34 ), The

d=1

family randomizes mean and variance at both endpoints.

Figure 9: Schrödinger bridge with the non-parametric particle-cloud prompt (

d=1

\gamma=1.0

; endpoints prompted by

100

-point sample clouds).

Figure 10: One non-parametric Schrödinger bridge solution operator (

d=2

\gamma=5.0

), solved zero-shot from particle-cloud prompts; panels as in Figure 8 : prescribed terminal law

P_{1}

(green shading, dashed

3\sigma

), trajectories (blue), initial populations (purple). The SDE trajectories are simulated according to section 2.4 .

		Inference Time (s)		Relative error
dim	$\bm{\gamma}$	transport	$+$ score	trajectory	score
$1$	$1.0$	0.015	1.320	6.70e-2 $\pm$ 4.35e-2	1.53e-1 $\pm$ 7.97e-2
$1$	$5.0$	0.014	0.656	7.52e-2 $\pm$ 4.84e-2	1.70e-1 $\pm$ 1.01e-1
$2$	$1.0$	0.014	0.865	4.15e-2 $\pm$ 2.34e-2	1.31e-1 $\pm$ 5.48e-2
$2$	$5.0$	0.014	1.308	3.93e-2 $\pm$ 2.13e-2	9.66e-2 $\pm$ 3.90e-2
$5$	$1.0$	0.014	0.992	4.60e-2 $\pm$ 1.77e-2	1.24e-1 $\pm$ 3.76e-2

Table 5: Non-parametric (particle-cloud) Schrödinger bridge vs. the exact Gaussian reference,

N_{\rm pt}=100

prompt samples.

	Inference Time (s)		Relative error
Parameters	transport	$+$ score	trajectory	score
$a\in[0.033,\,0.300]$	0.012	0.901	6.36e-3 $\pm$ 2.15e-3	5.49e-2 $\pm$ 1.76e-2
$q\in[0.167,\,0.950]$	0.012	0.901	6.44e-3 $\pm$ 2.33e-3	5.56e-2 $\pm$ 1.76e-2
$\varepsilon\in[0.333,\,3.000]$	0.012	0.901	6.60e-3 $\pm$ 2.57e-3	5.56e-2 $\pm$ 1.85e-2
$c\in[0.333,\,3.000]$	0.012	0.901	1.29e-2 $\pm$ 8.19e-3	1.17e-1 $\pm$ 5.18e-2
$\sigma\in[0.167,\,1.500]$	0.012	0.901	6.76e-3 $\pm$ 3.23e-3	5.23e-2 $\pm$ 1.88e-2

Table 6: Systemic risk vs. the closed-form Riccati/Gaussian reference: per-task inference time and the relative errors ( 33 ) and ( 34 ). Each row sweeps one coordinate of the cost vector

\theta=(a,q,\varepsilon,c,\sigma)

log-uniformly over its training range, with the other four held at

\theta_{\rm base}=\{a=0.1,\,q=0.5,\,\varepsilon=1,\,c=1,\,\sigma=0.5\}

; All results are from the same trained solution operator network, so inference time is constant.

Figure 11: Systemic risk at the base cost

\theta_{\rm base}=\{a=0.1,\,q=0.5,\,\varepsilon=1,\,c=1,\,\sigma=0.5\}

, three tasks with initial mean

m_{0}\in\{-3.8,\,0.1,\,2.5\}

(top/middle/bottom strips).

Figure 12: The same operator prompted zero-shot with the low-volatility cost

\sigma=0.25

, on the tasks and panels of Figure 11 : in every task the density concentrates strongly and the score steepens. Contrast the high-volatility prompt of Figure 13 .

Figure 13: The same operator prompted with the high-volatility cost

\sigma=1.2

, on the tasks and panels of Figure 11 : with no retraining, the density stays markedly wider and the score gentler, exactly as the variance ODE prescribes.

Panel of Figure 14 — Figure 14: One path-planning solution operator, zero-shot on configurations: initial distribution $P_{0}$ (purple), target distribution $P_{T}$ (green, dashed $3\sigma$ ), obstacles (orange rings). Blue curves are Euler–Maruyama sample paths of the controlled diffusion ( 1 ) at the trained $\gamma=0.1$ ; the quadratic programming projection enforces the $3\sigma$ keep-out.

	Inference Time (s)		Peak memory	Parameters
dim	transport	score	(GB)	( $10^{6}$ )
$2$	0.0095	0.6302	0.06	4.55
$4$	0.0095	0.6177	0.06	4.56
$8$	0.0095	0.6192	0.06	4.56
$16$	0.0095	0.6184	0.06	4.58
$32$	0.0095	0.6197	0.06	4.61

Table 7: Per-task inference cost of the multi-well operator versus state dimension, on untrained (randomly initialized) models. “transport” is one batched rollout over all

n_{\rm time}

nodes and “score” is one analytic-score evaluation at each node. Everything but

d

is fixed:

B=1

N_{\rm pt}=100

n_{\rm time}=32

Figure 15: Per-task inference cost versus dimension for the multi-well operator, at

B=1

N_{\rm pt}=100

and

n_{\rm time}=32

. The grey guides are

\mathcal{O}(d)

\mathcal{O}(d^{2})

and

\mathcal{O}(d^{3})

, anchored at the transport measurement at

d=2048

(circled), where growth begins, so that each guide is compared with the data over the range in which an asymptotic rate is claimed.

Explore similar work

Randomized Transport Maps for Model-Free Policy-Gradient Mean-Field Control
Oct 8, 2026Adonis Jamal, Samy Mekkaoui, Yadh Hafsi +1Policy Gradient Methods Model-Free RL
École Polytechnique
Mesh-Free Convolution: Learned Spectral Attenuation and Transport
Oct 11, 2026Niloufar Zakariaei, Manav Manav, Eldad HaberPDE Operator Learning Spectral Methods
University of British Columbia · Computational Geoscience Inc.
Actor-Critic Learning for Extended Mean Field Control with Deterministic Policies
Jul 13, 2026Ziheng Cheng, Xin Guo, Huyên Pham +1Policy Gradient Continuous-Time RL
University of California, Berkeley · Ecole Polytechnique, CMAP · Imperial College London

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

Randomized Transport Maps for Model-Free Policy-Gradient Mean-Field Control

Adonis Jamal, Samy Mekkaoui, Yadh Hafsi, Huyên Pham

École Polytechnique

Policy Gradient Methods Model-Free RL

Open paper AskPDF arXiv

We develop a model-free policy gradient method for discrete-time mean-field control (MFC). In MFC, the policy affects the objective both through the controlled dynamics and through the population distribution. Standard REINFORCE estimators capture the first effect but not the second. We introduce Transport REINFORCE, a transport map-based approach that perturbs a suitable transformation of the population distribution to estimate this missing mean-field contribution. The method applies to both finite and continuous state spaces. In finite state spaces, we perturb the population distribution directly on the probability simplex through a convex combination of the current population weights and random weights. In continuous state spaces, we project the population distribution onto the manifold of Gaussian mixtures, and then randomize it via a transport map that ensures the perturbed law remains within this manifold. We prove consistency of the perturbed objective and gradient as the perturbation vanishes, and derive bias and mean-square error bounds for the resulting sample-based gradient estimator. Numerical experiments on several MFC benchmarks show that Transport REINFORCE improves over standard REINFORCE.

Abstract

Explore similar work

Randomized Transport Maps for Model-Free Policy-Gradient Mean-Field Control

Mesh-Free Convolution: Learned Spectral Attenuation and Transport

Actor-Critic Learning for Extended Mean Field Control with Deterministic Policies