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

1 INTRODUCTION

• Standard REINFORCE misses how policy changes affect the population distribution, which also enters policies, dynamics, and rewards.The closest cited approach, MF-REINFORCE, addresses this dependence in finite state spaces by randomizing population-weight logits.Source: §1 · p. 1
• Transport REINFORCE randomizes the population law through transport maps, using population and random weights in finite spaces and Gaussian-mixture parameters in continuous spaces.Its explicit density score accounts for population dependence, while the transport couples original and perturbed laws to control their Wasserstein distance.Source: §1 · p. 2
• The method derives model-free policy-gradient estimators combining policy and population scores, with simulations at prescribed population arguments.It estimates population sensitivities using a probabilistic representation in finite spaces and centered policy differences in continuous spaces.Source: §1 · p. 2
• The paper proves objective and gradient consistency and derives bias and mean-square error bounds for the sample-based estimators.For objective perturbation, the stated orders are O(λ) in finite spaces and O(√λ) for a fixed Gaussian-mixture representation.Source: §1 · p. 2
• The experiments compare Transport REINFORCE with REINFORCE on seven discrete and continuous benchmarks and with MF-REINFORCE on finite-state benchmarks under matched simulator budgets.The suite also includes an explicitly solvable linear-dynamics setting with Gaussian initial conditions.Source: §1 · p. 2

2 BACKGROUND

• The stated assumptions support well-posedness of the finite-horizon control problem.The state and action spaces are assumed Polish, and the policy is selected from a compact admissible parameter set.Source: §2 · p. 2
• The control problem uses a parameterized randomized mean-field feedback policy over a finite horizon, with policy, dynamics, and rewards depending on the population distribution.The learner observes sampled transitions and rewards without analytic access to the model.Source: §2 · p. 2
• Standard REINFORCE captures the policy-score contribution while holding the population argument fixed, but the population flow itself also depends on the policy.The full gradient therefore includes a mean-field term collecting population-flow effects on rewards, policies, and state distributions.Source: §2 · p. 2

3 TRANSPORT MAP RANDOMIZATION

• The transport construction applies a measurable transformation of the population law, using the identity in finite spaces and an encoder–decoder map in continuous spaces.An auxiliary uniform variable lets the transport split atomic mass, which a deterministic state map alone cannot generally do.Source: §3 · p. 3

3.1 Finite-state space

• In finite state spaces, population randomization acts directly on simplex weights, and the transport retains the current state with probability 1 − λ or draws from randomized weights.The randomizer density must vanish on the simplex boundary for the score-function representation to avoid boundary terms; Gaussian randomization satisfies this condition, whereas a uniform simplex distribution does not.Source: §3.1 · pp. 3–4
• Finite-state perturbations are consistent as λ vanishes, and a model-free policy-gradient representation combines policy and population scores.The estimator has bias and mean-square error bounds; the scale choices λ⋆ ≍ B^−1/4 and η⋆ ≍ n^−1/4 balance terms in the stated MSE bound, with unknown constants.Source: §3.1 · pp. 4–5

3.2 Continuous-state space

• In continuous spaces, the population law is projected onto K-component Gaussian mixtures and the transport randomizes component weights, means, and covariances while preserving the mixture representation.The transport is the identity at λ = 0, and consistency results for the projected objective hold for fixed K.Source: §3.2 · pp. 5–6
• For continuous spaces, guarantees for the original objective also depend on mixture projection error εK, and a rate for εK requires additional approximation assumptions.The optimality loss has both a mixture-projection component controlled by K and a perturbation cost controlled by λ; gradient consistency and bias and MSE bounds are also established under stated assumptions.Source: §3.2 · pp. 6–7
• The continuous-state estimator’s MSE analysis requires control of fitted mixture-coordinate error, whose rate n^−1/2 needs an additional stability assumption.The scale choices λ⋆ ≍ B^−1/4 and η⋆ ≍ n^−1/6 balance scale-dependent terms under the stated conditions.Source: §3.2 · p. 8

• Offline, the algorithm estimates the population flow and its sensitivities at the current parameter θ.Source: §4 · p. 8
• Online, it holds those estimates fixed and uses B perturbed trajectories to estimate the gradient.Source: §4 · p. 8
• Both stages repeat after each policy update.Source: §4 · p. 8

5 NUMERICAL EXPERIMENTS

• The suite covers four finite-state and three continuous-state benchmarks, with five paired seeds and REINFORCE as a baseline.MF-REINFORCE is also compared on finite-state benchmarks; Transport-Proba is included on LQ control and portfolio.Source: §5 · p. 9
• Figure 1 pairs validation gaps over simulated transitions with the final policy’s population flow or terminal law for the first seed.Source: Figure 1 · p. 9
• After scale selection, Transport has the smallest mean optimality gap in every Table 1 row at matched simulator budgets.Source: Table 1 · p. 9, §5 · p. 9
• At matched simulator budgets, a Transport run takes 2 to 16 times REINFORCE’s wall-clock time because it also estimates sensitivities and population scores.Source: §5 · p. 10
• At the same scale, bimodal allocation’s gap falls from 0.1425 ± 0.1239 with K = 1 to 0.0020 ± 0.0009 with K = 2.The K = 1 projected objective is constant because the benchmark depends on terminal-law shape while its mean and variance are constant in θ.Source: §5 · p. 10
• Smaller λ reduces perturbation bias but raises population-score variance; the lowest gaps on four benchmarks occur at λ⋆/4 or λ⋆/2.The Figure 2 scale sweep plots final gap against λ/λ⋆, normalized by the smallest grid gap.Source: §5 · p. 10, Figure 2 · p. 10

Bias and mean-square error.

• Sampling largely follows n^-1, M^-1, and B^-1; small-scale perturbation bias is nearly linear in λ, and exact-flow sensitivity follows η^4.Table 2 compares predicted powers and fitted log-log slopes across six benchmarks, with the exact-flow η sweep on two continuous benchmarks.Source: Bias and mean-square error · p. 10, Table 2 · p. 10

• Transport REINFORCE randomizes the population argument through a transport map whose tractable density score recovers REINFORCE’s missing mean-field contribution.Source: §6 · p. 10
• The paper establishes objective and gradient consistency for finite-state and projected continuous-state problems, with bounds on estimator bias, mean-square error, and objective projection error.Source: §6 · p. 10
• The perturbation-scale bounds fix only the B^-1/4 order, and continuous-state guarantees concern J_K for a chosen mixture size K.The best scale constant is benchmark-dependent and selected after training; rates for ε_K remain future work.Source: §6 · pp. 10–11
• Potential extensions include non-exchangeable mean-field systems, graphon interactions, and cooperative multi-agent reinforcement learning.Source: §6 · p. 11

A.1 Canonical construction of the control problem

• The canonical construction uses measurable randomizations to generate actions and transitions, while perturbed population arguments are supplied separately from the law induced by the perturbed states.The perturbation randomizers are independent of the state and control process randomizers.Source: §A.1 · pp. 13–14
• The simulator returns rewards and independent transition draws, while terminal rewards are available at the final time; transition kernels and rewards are otherwise unknown.Policies are known and evaluable at prescribed population arguments, and unavailable population laws are estimated from particles.Source: §A.1 · p. 14

A.2 Well-posedness of the control problem

• The setup treats finite subsets of Rd and Rd as state spaces, with a nonempty closed action space, and addresses well-posedness separately for finite and continuous states.The finite-state argument is presented before the moment and continuity requirements for continuous states.Source: §A.2 · p. 14

A.2.1 Finite state space

• Under Assumption A.1, the original and perturbed finite-state processes are uniquely defined in law, and their objectives are bounded and continuous, attaining maxima when Θ is nonempty and compact.The objective bound is Tr + g.Source: §A.2.1 · p. 15

A.2.2 Continuous state space

• In continuous state space, Assumption A.3 imposes kernel and reward continuity plus quadratic-growth conditions; the resulting state laws vary continuously in W2 and the objective is finite, bounded, and continuous.The initial law has finite second moment and is independent of θ.Source: §A.2.2 · pp. 15–16

B.1 Common definitions

• The common definitions distinguish the state law being integrated from the population argument supplied to policies, transition kernels, and rewards.The policy-averaged transition kernel averages over the randomized action.Source: §B.1 · p. 17

B.2 Finite-state space

• In Proposition 3.1’s proof, the distance is zero at t = 0, and the argument invokes dTV(µ, ν) ≤ 1 for any µ, ν ∈ P(X).Source: §B.2.1 · p. 18

B.2.2 Perturbation consistency

• Under bounded rewards and total-variation Lipschitz conditions, Lemma B.2 establishes stability of the finite-state perturbed law.The supplied excerpt identifies the lemma and assumptions but does not provide the bound.Source: §B.2.2 · p. 18
• The perturbed objective satisfies |Jλ(θ) − J(θ)| ≤ C_T λ, where C_T = RLKT(T − 1) + LRT + 2gLKT + Lg.This bound is stated for the objective estimate in Theorem 3.2.Source: §B.2.2 · p. 19
• The proof establishes a gradient perturbation estimate by comparing perturbed and unperturbed gradients under the stated regularity conditions.The argument controls the difference between chain-rule gradient expressions and concludes with the gradient estimate.Source: §B.2.2 · p. 21
• As λ tends to zero, every accumulation point of maximizers of Jλ maximizes J.The proof uses compactness of Θ and continuity of J and Jλ.Source: §B.2.2 · p. 22

B.2.3 Density and score for the affine perturbation of the weights

• For λ in (0, 1), the affine perturbation of the population weights has a density in simplex coordinates.The law lemma applies to every z in the coordinate domain and identifies the density as hλ(· | z).Source: §B.2.3 · p. 22
• When β > N, the extended perturbation density is C1 across the simplex boundary and supported within the closed simplex.The density and its gradient vanish at the boundary under the stated tail condition.Source: §B.2.3 · pp. 22–23
• The simplex score satisfies a Stein identity, including E[H(Q)] = 0 for the random simplex weights Q.The identity applies to suitable continuous functions that are C1 in simplex coordinates with bounded gradient.Source: §B.2.3 · p. 23

B.2.4 Model-free policy-gradient representation

• The policy-gradient representation separates the policy-score contribution from a population-score contribution for the perturbed process.The policy term differentiates policy densities with the population flow fixed; the population term uses the simplex Stein identity.Source: §B.2.4 · p. 25

B.2.5 Bias and mean-square error

• The estimator analysis controls error through perturbation, population-flow, sensitivity, and online-sampling bounds; the online contribution is at most C(1 + λ^-2)/B.The population-flow and sensitivity estimates feed into the overall bias and mean-square error analysis.Source: §B.2.5 · pp. 27–30

B.3.1 Encoder, decoder, and numerical fitting

• Compact fitting constraints provide uniform moment bounds without requiring an inverse encoder-decoder identity; consistency requires η → 0 and a_n0/η → 0, while fixed-K results do not require approximation error to vanish.Mixture approximation error εK as K grows is a separate question.Source: §B.3.1 · pp. 30–31

B.3.2 Mixture coordinates and transport map

• The transport map preserves Gaussian-mixture form while perturbing component weights to (1 − λ)p_j + λq_j,t.The transported component retains a Gaussian conditional law.Source: §B.3.2 · p. 32

B.3.3 Perturbation estimate

• Under the stated moment bounds, the transport perturbation has Wasserstein displacement at most Cλ, with conditional mean-squared displacement bounded by Cλ².Source: §B.3.3 · p. 32

B.3.4 Objective perturbation consistency

• Quadratic rewards in the linear–quadratic and portfolio benchmarks are not globally Wasserstein-Lipschitz, so those benchmarks fall outside the stated consistency, projection-error, and optimality-gap guarantees.Source: §B.3.4 · p. 33
• Under Assumption B.12, the proof establishes uniform-in-θ bounds for objective perturbation and projection error using time-dependent perturbation estimates.The projection-error bound follows from the triangle inequality and Theorem 3.8.Source: §B.3.4 · pp. 33–34

B.3.5 Objective gradient convergence

• Under Assumption B.13, objective gradients converge uniformly, and the objective error is additionally O(λ) under the stated stronger regularity assumptions.The objective-error bound is uniform in θ.Source: §B.3.5 · p. 36

B.3.6 Density and score for the randomized perturbation

• Lemma B.16 gives an explicit density for the randomized perturbation using inverse coordinates and a Jacobian, with the density zero outside its stated domain.For K = 1, the weight-density factor and free weight coordinates are omitted.Source: §B.3.6 · p. 37

B.3.7 Policy-gradient representation

• Under likelihood-ratio regularity assumptions, Theorem 3.12 represents the randomized objective’s policy gradient using trajectory returns and scores for the policy and perturbation law.The derivation uses dominated convergence; the return is held fixed when differentiating the trajectory likelihood.Source: §B.3.7 · pp. 38–39

B.3.8 Sensitivity estimation and statistical bounds

• For sensitivity estimation, the setup uses an M-particle base fit, n0-particle fits at θ ± ηeℓ, and conditionally independent online trajectories with fresh simulation randomness.The proof sets n0 = n/(2dθ).Source: §B.3.8 · p. 39
• The population-coordinate error is controlled directly by the fitting-accuracy assumption.This coordinate error is one of the offline fitting controls used in the statistical bounds.Source: §B.3.8 · pp. 41–42
• The sensitivity estimate uses centered finite differences; Taylor expansions cancel constant and quadratic terms, while fitting errors contribute to the remaining error.The proof decomposes the fitted-versus-exact sensitivity error into finite-difference approximation error and fitted-coordinate error.Source: §B.3.8 · p. 41
• The bias bound compares the conditional mean G = Γλ(bz, bD) with the true gradient using stability estimates and Jensen’s inequality.The fitted coordinates and sensitivity matrix are measurable with respect to the offline particles; Jensen’s inequality transfers the bound on their conditional mean to the estimator’s bias.Source: §B.3.8 · p. 42
• The mean-square-error proof decomposes the estimator error into online noise and conditional-mean error, with the cross term vanishing by conditional centering.The online estimator’s conditional variance is bounded using the B independent contributions and score-moment bounds.Source: §B.3.8 · p. 42
• The theorem establishes bias and mean-square-error bounds by combining offline fitting-error controls with the online estimator’s conditional variance.The online estimator averages B independent contributions conditional on the offline particles, and the theorem combines the resulting variance bound with the bias bound.Source: §B.3.8 · p. 42

C LINEAR DYNAMICS: A FULLY SOLVABLE FRAMEWORK

• Linear dynamics with Gaussian noise and Gaussian initial conditions keep the state law within the Gaussian manifold, characterized by its mean and variance.The framework is presented in one dimension for simplicity and is stated to extend to multiple dimensions.Source: C LINEAR DYNAMICS: A FULLY SOLVABLE… · pp. 42–43
• The LQ policy-gradient theorem gives a gradient representation whose reward-to-go form removes past rewards using conditional mean-zero score terms.The proof applies this argument to both policy scores and population scores.Source: C LINEAR DYNAMICS: A FULLY SOLVABLE… · p. 44
• Probabilistic approximations of moment sensitivities yield Transport-Proba, which estimates the mean-field score without the general continuous-state algorithm’s centered policy differences.Transport-Proba is compared with standard REINFORCE and the default Transport estimator in Appendix E.3.3.Source: C LINEAR DYNAMICS: A FULLY SOLVABLE… · pp. 44–45

D.1 Sampling details

• Each update estimates population coordinates and sensitivities offline, then simulates perturbed trajectories online for the gradient estimator.Continuous-state sensitivities use centered differences at shifted parameters, while continuous-state population coordinates are fitted by constrained Expectation-Maximization on particles.Source: §D.1 · pp. 45–46

D.2 Computational Cost

• Finite-state fresh-batch updates use T(M + B) + nT(T + 1)/2 transitions, while shared-batch finite-state and continuous-state updates use T(M + n + B).The fresh-batch variant draws an auxiliary batch for each time; the shared-batch variant reuses trajectories.Source: §D.2 · p. 46
• The finite-state MSE theorem covers fresh auxiliary batches, whereas the experiments use a shared batch.The reported transition counts also exclude arithmetic costs of mixture fitting and score evaluation.Source: §D.2 · p. 46
• For fixed total n in continuous state space, increasing dθ decreases the particles per shifted system to n0 = n/(2dθ).Fresh finite-state auxiliary batches raise auxiliary simulation cost from nT to nT(T + 1)/2 transitions.Source: §D.2 · p. 46
• Beyond simulation, an update costs O(BTqdθ) for main scores, plus O(nT 2qdθ) for finite-state auxiliary recursion or O(Tqdθ) for centered differences.These arithmetic counts exclude mixture-fitting and score-evaluation costs.Source: §D.2 · p. 46

E.1 Protocol and settings

• Every benchmark gives methods the same simulator budget; five paired seeds and Adam are used, with frozen-policy validation every ten updates.Validation uses 10^5 interacting particles; the finite-state benchmarks resample initial laws during training, while continuous-state benchmarks keep them fixed.Source: §E.1 · pp. 46–47, Table 3 · p. 47

E.2.1 Two-state control

• Two-state control uses two states, stay/move actions, a five-step horizon, and a closed-form optimum that sends every initial law to (0.6, 0.4).Training samples the initial mass on state 1 from U([0.1, 0.9]); validation uses 0.8.Source: §E.2.1 · p. 48

E.2.2 Cybersecurity

• Cybersecurity models four defended/undefended and infected/susceptible computer states, with infection depending on the population and a three-step training horizon.Validation uses the uniform initial law and a 50-step horizon; no optimum is available.Source: §E.1 · p. 47, §E.2.2 · p. 48

E.2.3 Distribution planning

• Distribution planning moves agents on a ten-state cycle and penalizes distance from a target population distribution over a five-step horizon.The reference value is obtained by gradient ascent on the exact population recursion, but global optimality is not certified.Source: §E.1 · p. 47, §E.2.3 · p. 48

E.2.4 Targeted advertising

• Targeted advertising models customer conversion as population share plus an advertising effect, with a five-step horizon and validation starting at customer share 0.5.The reward discounts customer status and advertising cost; the terminal reward is zero.Source: §E.2.4 · p. 48

E.2.5 Linear–quadratic control

• Linear–quadratic control uses a 20-step mean-field linear system with a non-Gaussian initial population, and its reference policy follows from Riccati equations.Transport uses one Gaussian component.Source: §E.2.5 · p. 48

E.2.6 Mean–variance portfolio

• The portfolio benchmark uses Gaussian policies for wealth dynamics over ten steps, with an optimal feedback gain and an intercept found by concave quadratic optimization.Transport uses one Gaussian component.Source: §E.2.6 · pp. 48–49

E.2.7 Bimodal population allocation

• This benchmark tests Gaussian-mixture size: although the two-mode population’s mean and variance are policy-independent, a single Gaussian loses information needed to optimize, while K = 2 represents the terminal law exactly.The K = 1 projected objective is constant and provides no learning gradient; the unique maximizer is θ⋆ ≃ 0.868437, where J(θ⋆) = 0.Source: §E.2.7 · p. 49

E.3.1 Learning curves and population dynamics

• In distribution planning, Transport learns a terminal law closer to the target than REINFORCE, which remains near its initial law; MF-REINFORCE overconcentrates mass near the center.The reward charges for movement while the benefit depends on the resulting population distribution.Source: §E.3.1 · pp. 49–50
• In the portfolio, REINFORCE increases mean wealth despite its running square penalty, while population correction substantially reduces that excess.Figure 3 separately plots cybersecurity and advertising validation objectives over simulated transitions and first-seed infected mass and customer share.Source: §E.3.1 · p. 50, Figure 3 · p. 50
• Cybersecurity gives similar final objectives for Transport and REINFORCE, while both population-corrected methods improve over REINFORCE on advertising.Neither corrected method recovers advertising until customer share reaches 0.9; MF-REINFORCE has the higher mean objective with substantial seed variation.Source: §E.3.1 · p. 50, Figure 3 · p. 50

E.3.2 Perturbation scale and mixture size

• On advertising, the largest tested perturbation scale yields perturbed optimization error below 10^-5 but leaves an original gap of 0.05.Further optimization of that same perturbed objective retains the loss; larger perturbations can change which policy the objective favors.Source: §E.3.2 · pp. 50–51, Table 6 · p. 51
• For bimodal allocation, K = 1 makes the fitted objective constant across scales, whereas K = 2 retains policy dependence and optimizes the perturbed objective to 10^-4.The gap increases with λ as the learned parameter moves away from θ⋆.Source: §E.3.2 · p. 50, Table 5 · p. 51

E.3.3 Sensitivity estimator and runtime

• Centered differences give a smaller final gap and shorter runtime than Transport-Proba at every tested scale on linear–quadratic control; the portfolio comparison is closer relative to seed variation.Transport-Proba uses likelihood-ratio moment sensitivities; both estimators use K = 1, which retains relevant population moments in these benchmarks.Source: §E.3.3 · p. 52
• Transport adds arithmetic overhead beyond simulator transitions because each update estimates sensitivities and evaluates population scores.The overhead is particularly large on distribution planning and bimodal allocation, so matched simulator budgets do not imply the same gain per second.Source: §E.3.3 · p. 52, Table 7 · p. 53

E.3.4 Auxiliary scale and reuse of trajectories

• At small η, the factor (1 − η)/η amplifies noise in earlier sensitivities, while advertising sensitivity error is minimized at an intermediate scale.Shared and fresh auxiliary batches have comparable errors in the sweep; reuse cuts auxiliary transitions from nT(T + 1)/2 to nT but makes successive sensitivity estimates dependent.Source: §E.3.4 · p. 52

E.3.5 Perturbation of the law and gradient

• At fixed policies, finite-state total variation scales proportionally with λ, single-Gaussian W2 scales approximately linearly, and gradient perturbation is approximately linear on evaluated benchmarks.The two-component mixture changes weights and redistributes mass between modes; its law-distance behavior differs across the tested scales.Source: §E.3.5 · pp. 52, 54, Figure 4 · p. 54

E.3.6 Bias and mean-square error

• On advertising, sensitivity error decreases only from 0.47 to 0.38 as n rises from 130 to 260, with fitted slope −0.51.The result is consistent with systematic error remaining in the perturbed sensitivity at fixed η.Source: §E.3.6 · p. 54
• At benchmark budgets, continuous-state fitting error dominates truncation error at every tested η, and total sensitivity error decreases as η^-2 up to 0.95.With exact flows, centered differences have squared error of order η^4; fitted endpoint errors divided by 2η add a term that grows as η decreases.Source: §E.3.6 · p. 54, Figure 6 · p. 55
• Table 9 reports log-log slopes for all benchmarks, varying one quantity at a time while holding others at benchmark allocations; for continuous-state λ sweeps, B is fixed at 4096.The λ slopes are fitted on the three smallest scales at which the bias is resolved.Source: Table 9 · p. 55