Optimal Feedback Control of the Fokker-Planck Equation for an SIRS Epidemic Model

Introduction

Mathematical modeling has long played an important role in understanding the transmission of infectious diseases and evaluating strategies for epidemic prevention and control. Beginning with the classical work of Kermack and McKendrick [16], compartmental epidemic models have provided a systematic framework for describing the evolution of susceptible, infected, and recovered populations. Among these, the susceptible–infected–recovered–susceptible (SIRS) model is particularly relevant to infectious diseases for which immunity is temporary or incomplete. By allowing recovered individuals to return to the susceptible class, the SIRS model accounts for reinfection and the possibility of recurrent epidemic outbreaks. Such dynamics motivate the development of intervention strategies that reduce transmission, limit the loss of immunity, and enhance recovery while accounting for the costs associated with implementing these measures.

Although deterministic epidemic models provide valuable insights into disease transmission, their predictions are based on prescribed parameter values and initial conditions. In practice, epidemic dynamics are influenced by substantial uncertainty arising from variations in contact patterns, transmission rates, recovery mechanisms, immunity duration, and environmental or behavioral factors. Stochastic epidemic models address these uncertainties by incorporating random perturbations into the governing dynamics. Several authors investigated stochastic epidemic dynamics using stochastic differential equations and established important results concerning positivity, extinction, persistence, stationary distributions and the effects of different types of noise [9, 13, 18, 23, 27]. See [1] for an overview of such models. Generally, stochastic differential equation models describe ensembles of possible epidemic trajectories rather than a single deterministic evolution. This distinction is particularly important when intervention strategies must account for the probability of severe outbreaks or other undesirable epidemiological outcomes.

Optimal control theory provides a mathematical framework for determining intervention strategies that balance epidemiological benefits against economic, logistical, and societal costs. In compartmental epidemic models, controls may represent nonpharmaceutical interventions that reduce transmission, vaccination or other measures that influence immunity, and treatment strategies that accelerate recovery [7, 20], and researchers have analyzed models incorporating confounding effects ranging from noncompliance with control measures [3, 21] to media coverage [8] to limited resources [14]. Classical optimal control formulations typically determine time-dependent interventions by minimizing a prescribed objective functional subject to the epidemic dynamics. In stochastic settings, optimal control can also be formulated through controlled diffusion processes and dynamic programming, as developed in the general theory of stochastic control [12]. However, time-dependent interventions do not explicitly adapt to the current epidemic state. Moreover, when the objective involves the likelihood of undesirable outcomes rather than only expected population levels, it is useful to consider the evolution of the entire probability distribution of epidemic states.

The Fokker–Planck (FP) equation offers a natural framework for this purpose. Associated with an underlying stochastic differential equation, the FP equation describes the temporal evolution of the probability density of the stochastic state variables. Consequently, an optimal control problem constrained by the FP equation allows interventions to influence the distribution of possible epidemic outcomes. This density-based perspective provides a direct way to incorporate objectives involving expected infection burden, probability mass in high-infection regions, and desirable terminal epidemic states. Furthermore, by allowing the controls to depend on both the epidemic state and time, the resulting formulation accommodates feedback intervention policies. Unlike purely time-dependent controls, these policies can assign different intervention intensities to different regions of the epidemic state space. While FP control has been applied biological and ecological models such as those for pedestrian motion [24] and biological imaging [2], it has not been broadly applied to epidemiological models with the exception of [22] where time-dependent controls were considered. Indeed, most work regarding FP control assumes time-dependent control or specially parameterized control ansatzes and, notably, does not consider control of the diffusive terms in the FP equation [5, 11].

Despite the modeling advantages, FP optimal control problems introduce mathematical and computational challenges. The probability density satisfies a time-dependent partial differential equation whose coefficients depend on the control variables. When the controls influence both the drift and diffusion coefficients, the resulting control-to-state mapping is highly nonlinear, and the diffusion operator may contain mixed spatial derivatives. Establishing existence of optimal controls and deriving necessary optimality conditions therefore requires an appropriate analytical framework for the controlled FP equation. The numerical approximation must also account for the nonnegativity and conservation properties of probability densities. The classical Chang–Cooper method [6] provides a useful structure-preserving discretization for drift–diffusion equations, although additional treatment is required for multidimensional systems with cross-diffusion.

Motivated by these considerations, we develop a FP optimal control framework for a stochastic SIRS epidemic model with three distinct intervention mechanisms. Stochastic perturbations are introduced into transmission, loss of immunity, and recovery, and the corresponding feedback controls modify both the deterministic transition rates and the stochastic fluctuations. The first control reduces transmission through nonpharmaceutical interventions, the second reduces loss of immunity, and the third increases the recovery rate. The resulting stochastic system induces a two-dimensional FP equation with a control-dependent diffusion matrix and nonzero cross-diffusion terms.

The principal contributions of this work are as follows. First, we formulate a distribution-based optimal control problem with space- and time-dependent feedback controls and a general objective functional incorporating intervention costs, running epidemic costs, and terminal costs. In particular, our problem includes controls in the diffusion terms of the FP equation. Second, we investigate the analytical properties of the associated control-to-state operator and develop existence and first-order optimality results using a forward FP equation, a backward adjoint equation, and a Hamiltonian minimization condition. Third, we construct a numerical framework combining Chang–Cooper discretization of the diagonal drift–diffusion terms, finite-difference approximation of the mixed diffusion term, and a sequential quadratic Hamiltonian (SQH) algorithm for computing feedback controls. We also investigate sufficient-decrease and convergence properties of the SQH iteration.

The numerical experiments consider four epidemiological objectives that demonstrate the flexibility of the proposed framework. The first minimizes cumulative infection burden, while the second penalizes epidemic states in which the infected fraction exceeds a prescribed threshold. The third introduces a terminal reward for maintaining a sufficiently large susceptible population, and the fourth combines threshold-based infection regulation with a terminal susceptible-state objective. For each scenario, we compute the evolving probability density and the corresponding feedback control maps. Monte Carlo simulations of the underlying controlled and uncontrolled stochastic differential equations are then used to examine the effects of the optimized interventions on individual epidemic trajectories and their ensemble behavior.

The remainder of the paper is organized as follows. Section 2 introduces the controlled stochastic SIRS model, derives the associated FP equation, and formulates the feedback optimal control problem. Section 3 develops the analytical framework, including existence of an optimal control and first-order necessary optimality conditions. Section 4 presents the numerical approximation of the forward FP and backward adjoint equations together with the SQH algorithm and its convergence analysis. Section 5 presents numerical experiments under different epidemiological objectives and compares the resulting controlled and uncontrolled stochastic dynamics.

A controlled stochastic SIRS model and FP equation

We begin from the following controlled Susceptible-Infected-Recovered-Susceptible (SIRS) dynamics which model the spread of a generic disease in a generic region:

Here $b$ represents the constant introduction of new members into the population, and $\delta$ represents the natural death rate. A basic assumption of this model is that the population is homogeneous and well-mixed so that infection spreads via mass action with rate $\beta$ and infected individuals recover at rate $\gamma$. The parameter $\lambda$ represents loss of immunity whereby recovered individuals become susceptible again. We supply the policy-maker with three control variables $u := (u_1,u_2,u_3)$ by which they may inhibit disease spread:

• (a) $u_1 \in[0,u_1^{\max}]$ for some $u_1^{\max} \in[0,1]$ - this is reduction in infectivity due to governmentally enacted non-pharmaceutical intervention measures,

• (b) $u_2 \in[0,u_2^{\max}]$ for some $u_2^{\max} \in[0,1]$ - reduction in loss of immunity (perhaps due to vaccination or other treatment efforts),

• (c) $u_3 \in[0,u_3^{\max}]$ for some $u_3^{\max} \in[0,1]$ - increase in recovery rate due to treatment efforts.

Due to the highly volatile nature of epidemic disease spread, we also add stochasticity to the model, and specifically following [13, 23], we assume that there is some uncertainty in the rate parameters $\beta,\gamma,\lambda$. Thus we formally replace $\beta$ with $\beta+\sigma_\beta W$, and likewise with $\gamma,\lambda$. Finally, we assume that $b=\delta$ and that the initial conditions are normalized so that $S(0)+I(0)+R(0)=1$. Under these conditions, we see that the total population $N=S+I+R$ described by (1) (and indeed by the stochastic generalization) will be constant: $N(t)=1$. Thus we can write $R=1-S-I$, and drop this compartment. With all these assumptions, we introduce the controlled, stochastic SIRS model:

We write the dynamics in the general form in the form

where $x=(S,I)$, $u=(u_1,u_2,u_3)$, and $W$ is a 3 dimensional Brownian motion with independent components. In this notation, the drift and diffusion functions are

For notational convenience, we set $D(x,u)=\frac{1}{2}\Sigma(x,u)\Sigma(x,u)^{T}$. We record some important observations regarding the model.

• By the arguments in [13, 23], if $S_0,I_0,R_0>0$ and $S_0+I_0+R_0\le1$, then $[0,1]^3$ is a forward invariant set for the dynamics, meaning that we can take $\Omega=[0,1]^2$ and let $x_0=(S_0,I_0)\sim\rho_0$ where $\rho_0$ is supported in $\Omega$.

• Our admissible control actions take values in the rectangle $R_U=\prod_{i=1}^{3}[0,u_i^{\max}]$.

• All the components of the functions $F:\Omega\times R_U\to\mathbb{R}^2$ and $D:\Omega\times R_U\to\mathbb{R}^{2\times2}$ and their $x$-derivatives are continuous in $(x,u)$ and thus bounded on $\Omega\times R_U$.

The Fokker-Planck equation corresponding to (3) is

where the flux $E(x,u)$ is defined componentwise by

We now formulate an optimal control problem. We are interested in feedback control, so we consider $u_i:\Omega_T\to R_U$ for each $i=1,2,3$ where $\Omega_T=\Omega\times[0,T]$. Note that because the controls enter into (4) in the diffusion coefficients, we need very regular controls to guarantee nice properties of the solution. Some of this additional regularity will be required for the ensuing proof of existence of minimizers for the optimal control problem in theorem 3.5 (and the preceding lemmas), and yet more is required for the proof of the Pontryagin Minimum Principle in theorem 3.8. Accordingly, we take

for a fixed positive constant $M$. We consider a relatively generic cost functional

where $\ell:R_U\to[0,\infty)$ and $G,K$ are, at a minimum, bounded. For particular theorems, these functions will be required to satisfy further conditions as specified below.

Our optimization problem can then be stated abstractly as follows

Here the functions $\ell,G,K$ express the preferences of the policy-maker. The control cost $\ell(u)$ represents represents any political, economic, of logistical costs associated with implementing control measures. For example, this could encode the loss of gross domestic product (GDP) incurred when requiring businesses to close or the costs associated with manufacturing and distributing vaccines. The running cost $G(x,t)$ and final cost $K(x)$ implicitly define the desirable or undesirable epidemic outcomes. For example, recalling that $x=(S,I)$, a policy-maker may set $G(x,t)=I$, meaning that more cost is assessed when there is more probability density associated with states where the infectious population is larger.

Analysis of the FP optimal control problem

We prove existence of a minimizer for (7) via analysis of the control-to-state map. Before this, we need a basic existence result.

The proof of this theorem is classical so we omit it, but comment on some facets. Existence and uniqueness of the solution follows as in [10 Ch. 7.1]. The uniform boundedness over choices of controls follows since $F(x,u)$ and $D(x,u)$ and their spatial derivatives are uniformly bounded for $(x,u) \in\Omega\times R_U$. Conservation of the integral follows by taking $\varphi\equiv1$ in (8), and nonnegativity follows by setting $\varphi=(\rho(\cdot,t))^- := -\min(\rho(\cdot,t),0)$ at each time $t$ and establishing the bound $\frac{d}{dt}\|\rho(\cdot,t)^-\|_{L^2(\Omega)}^2 \le C\|\rho(\cdot,t)^-\|_{L^2(\Omega)}^2$, whereupon the result follows by Grönwall’s inequality since $\rho_0^-=0$. Note that because $\rho_0$ is a probability density function, we will henceforth always assume that $\rho_0 \ge0$ and $\|\rho_0\|_{L^1(\Omega)}=1$, whereupon we have these same properties for $\rho(\cdot,t)$ for all times $t \in[0,T]$. We note that theorem 3.2 does not require much regularity for the controls $u$, so for this theorem the regularity in our control space $U_{\mathrm{ad}}$ is sufficient but certainly unnecessary.

From theorem 3.2, we have a well defined control-to-state operator

$u\in U_{\mathrm{ad}}\mapsto\rho\in W$, the weak solution of (4) with control map $u$. (9)

We prove existence of a solution of (7) through analysis of $S$. We proceed with two lemmas and then state existence of minimizers as a theorem. In all of the following work, we equip $U_{\mathrm{ad}}$ with the $L^2$-topology induced by the norm

Additionally, whenever a subsequence with some property is mentioned, we implicitly pass to this subsequence, but to avoid excessive notation, we do not rename the sequences.

Proof. Take a sequence of controls $u_k\in U_{\mathrm{ad}}$ such that $u_k\rightharpoonup u$ in $L^2(\Omega_T)$ and $\rho_k=S(u_k)\rightharpoonup\rho$ in $W$. We need to prove that $u\in U_{\mathrm{ad}}$ and $S(u)=\rho$.

By the Banach–Alaoglu theorem and compatibility of weak limits, along a subsequence, we can achieve $u_k\rightharpoonup^*u$ in $L^\infty(\Omega_T;R_U)$ and $(\nabla u_i)_k\rightharpoonup^*g_i$ in $L^\infty(\Omega_T;\mathbb{R}^2)$ for some $g_i\in L^\infty(\Omega_T;\mathbb{R}^2)$, $i=1,2,3$. Now for any $\varphi\in C_c^\infty(\Omega_T)$, $i=1,2,3$ and any first order partial derivative $\partial$, we see

By weak$^*$ convergence of $u_k$ to $u$ and of $(\nabla u_i)_k$ to $g_i$, this gives

implying that $u_i$ is weakly differentiable with weak partial $\partial_i(g_i)$, and

by weak$^*$ lower semicontinuity of the $L^\infty$ norm. Finally, weak $L^2$ convergence also preserves pointwise a.e. bounds, as follows by testing $u_k$ against appropriately designed indicator functions, and thus $u(x,t)\in R_U$ for almost every $(x,t)\in\Omega_T$. We conclude that $u\in U_{\mathrm{ad}}$.

It remains to prove that $S(u)=\rho$. By theorem 3.2, the sequence $\rho_k$ is bounded in $\mathcal{W}$. Since we have compact-continuous embeddings

the Aubin-Lions lemma [25 Ch. III, Prop. 1.3] yields a subsequence of $\rho_k$ which converges strongly to $\rho$ in $L^2(\Omega_T)$. Since $u_k(x,t)\in R_U$ for a.e. $(x,t)\in\Omega_T$ and $\|u_k\|_{W^{1,\infty}(\Omega_T)}\leq M$, the sequence $u_k$ is uniformly bounded and uniformly Lipschitz, whereupon the Arzela-Ascoli theorem gives $u_k\to u$ uniformly along a subsequence. Because $F$ is continuous and uniformly bounded, this allows us to conclude that $F(x,u_k)\to F(x,u)$ in $L^2(\Omega_T)$ via the dominated convergence theorem (and similarly for $D_{ij}$ and their spatial derivatives). Finally to prove that $S(u)=\rho$, we fix $\varphi\in H^1(\Omega)$. Considering (8), there are three terms to address. The pairing $(\rho_k)_t\varphi$ behaves nicely since $\rho_k\rightharpoonup\rho$ in $\mathcal{W}$. The term $\rho_kF(x,u_k)\cdot\nabla\varphi$ is controlled since both $\rho_k\to\rho$ and $F(x,u_k)\to F(x,u)$ strongly in $L^2(\Omega_T)$. And upon expanding, the terms

are controlled since $D_{ij}(x,u_k)$ is uniformly bounded and converges strongly to $D_{ij}(x,u)$ in $L^2(\Omega_T)$, and $\rho_k\to\rho$ weakly in $\mathcal{W}$. Thus we find that $\rho$ is a weak solution of (4) corresponding to $u$; that is, $S(u)=\rho$.

Lemma 3.4. The control-to-state operator $S$ is weak-to-strong continuous in the $L^2$-topology. That is, if $u_k\in U_{\mathrm{ad}}$ and $u_k\rightharpoonup u$ in $L^2(\Omega_T)$, then $S(u_k)\to S(u)$ in $L^2(\Omega_T)$.

Proof. Suppose $u_k\in U_{\mathrm{ad}}$ and $u_k\rightharpoonup u$ in $L^2(\Omega_T)$. Once again, $\rho_k:=S(u_k)$ is a bounded sequence in $\mathcal{W}$, and thus any subsequence is bounded in $\mathcal{W}$. Hence by the Aubin-Lions lemma [25 Ch. III, Prop. 1.3], we can pass to a further subsequence which converges strongly to some $\rho\in L^2(\Omega_T)$. By the closure guaranteed in lemma 3.3, we must have $\rho=S(u)$. Hence we have prove that every subsequence of $S(u_k)$ has a further subsequence which converges strongly to $S(u)$ in $L^2(\Omega_T)$. By the Urysohn subsequence principle [26 §2.1.17], we have that $S(u_k)\to S(u)$ in $L^2(\Omega_T)$.

We recall before stating our existence theorem that the functions $\ell,G,K$ are used in the definition of the cost functional $J(\rho,u)$ in (6).

Theorem 3.5. Assume that $\ell:R_U\to[0,\infty)$ is convex and lower-semicontinuous, $G\in L^\infty(\Omega_T)$, and $K\in L^\infty(\Omega)$. Then there exists a solution $u\in U_{\mathrm{ad}}$—and corresponding weak solution $\rho\in\mathcal{W}$ of (4)—of the minimization problem (7).

Proof. Define

Recall that $\rho(x,t)\geq0$ and $\|\rho(\cdot,t)\|_{L^1(\Omega)}=1$ for all $t\in[0,T]$. Then whenever $\ell,G,K$ satisfy the given assumptions,

Thus we can take a sequence $u_k\in U_{\mathrm{ad}}$ with corresponding $\rho_k:=S(u_k)\in W$ such that $J(\rho_k,u_k)\searrow m$. Since $u_k$ is, in particular, a bounded sequence in $L^2(\Omega_T)$, there exists a subsequence converging weakly to some $u\in L^2(\Omega_T)$ and by the weak closure in lemma 3.3, we have $u\in U_{\mathrm{ad}}$. By the weak-to-strong continuity in lemma 3.4, we have $\rho_k\to\rho:=S(u)$ in $L^2(\Omega_T)$. Considering the definition of the cost functional in (6), the conditions that $\ell$ is nonnegative, convex and lower-semicontinuous guarantee that the map $u \in U_{\mathrm{ad}} \mapsto\int_{\Omega_T} \ell(u)\,dx\,dt \in[0,\infty)$ is weakly lower-semicontinuous, and the integral terms involving the $\rho_k$ are preserved under limits simply by convergence of $\rho_k \to\rho$, since $G$, $K$ are bounded. Thus we have

proving that the pair $(\rho,u)$ is a solution of (7).

Given this, we would like to prove first order necessary conditions for optimality akin to the classical Pontryagin Minimum Principle. To do so, we need a notions of an adjoint solution and a Hamiltonian.

For functions $\rho(x,t)$, $q(x,t)$, $u(x,t)$, we define the formal Lagrangian by

Here the adjoint function $q(x,t)$ is playing the role of a Lagrange multiplier. Integrating by parts and expanding $J(\rho,u)$, we have

where we have enforced the zero flux boundary condition $E\cdot n=0$. The adjoint equation is derived by enforcing

for arbitrary smooth test functions $v$ satisfying $v(x,0)=0$. Computing formally, this yields

Solutions to (10) are considered in a weak sense analogous to (8) above. That is, $q$ is a weak solution of (10) if $q(\cdot,T)=-K(\cdot)$ and

for all $\psi\in H^1(\Omega)$ and a.e. $t\in[0,T]$. For fixed $u\in U_{\mathrm{ad}}$, existence and uniqueness follow directly as in [10 Ch. 7.1]. Regularity can then be upgraded to $W^{2,p}(\Omega_T)$ via standard parabolic regularity results [19 Thm. 7.35]. This is the main reason that we need so much additional regularity for our control space. With controls uniformly bounded $W^{1,\infty}(\Omega_T)$ and taking values in the compact rectangle $R_U$, the maps $(x,t)\mapsto F(x,u(x,t))$ and $(x,t)\mapsto D(x,u(x,t))$ are in $W^{1,\infty}(\Omega_T)$ which allows for the increased regularity.

Further, if $K \in W^{2,p}(\Omega)$, then $q \in W^{2,1}_{p}(\Omega_T)$ and $\|q\|_{W^{2,1}_{p}(\Omega_T)}$ is uniformly bounded over choices $u \in U_{\mathrm{ad}}$.

We henceforth assume that $K \in W^{2,\infty}(\Omega)$ so that we have uniformly bounded $\|q\|_{W^{2,1}_{p}(\Omega_T)}$ for any finite $p$.

We then formally define the Hamiltonian

Note that $H$ depends on $(x,t)$ implicitly through $\rho,q,u$; we have suppressed this dependence to reduce the notational burden.

We prove that the optimal control map can be characterized as a pointwise minimizer of this Hamiltonian. Before doing so, we prove a lemma relating the cost functional and the Hamiltonian.

Proof. We compute

where in the middle equality we have used (11) with $\psi=\rho-\hat{\rho}$, and in the last we have integrated by parts in $t$ and canceled boundary terms. But now with $q$ playing the role of the test function $\varphi$ in (8), we see that all terms with $\rho$ drop out of (13) since $\rho$ is the weak solution of (4) corresponding to $u$. Thus we arrive at

where in the last equality we have replaced $\int_{0}^{T}\int_{\Omega}\hat{\rho}_{t}q\,dxdt$ using that $\hat{\rho}$ is a weak solution of (4). This completes the proof.

Proof. The proof utilizes the classical needle variation technique. Supposing that $(\rho,u)$ is optimal, we fix a point $(x_0,t_0)\in\Omega_T$ and $w\in R_U$. Define the set

and the control map

In what follows, $w$ is the fixed constant value in $R_U$, while $u=u(x,t)$ depends on $(x,t)$. Assuming $\rho_k$ is the state map corresponding to $u_k$ and $q_k$ is the adjoint map corresponding to $(\rho_k,u_k)$, lemma 3.7 tells us that

where $I_1,I_2$ are the indented integrals. Now the map

$(x,t)\mapsto\left[H(\rho(x,t),q(x,t),w)-H(\rho(x,t),q(x,t),u(x,t))\right]$

is in $L^1(\Omega_T)$, so by the Lebesgue differentiation theorem

as long as $(x_0,t_0)\in\Omega_T$ is a Lebesgue point for this function. Thus, taking the limit as $k\to\infty$ in (16), we arrive at (15) so long as $I_1,I_2\to0$. We see

where $E$ is the Fokker-Planck flux defined in (5). Now $q_k\to q$ strongly in $L^2(0,T;H^1(\Omega))$ by a standard parabolic stability estimate [10 Ch. 7, Thm. 2]. However, we recall that for sufficiently regular $K$, $\|q\|_{W_p^{2,1}(\Omega_T)}$ is uniformly bounded over choices of controls for any $p>1$. Taking $p$ large enough, we have that $W_p^{2,1}(\Omega_T)$ is compactly embedded in $C^{1,1/2}(\overline{\Omega_T})$ [15 Thm. 3.14], [17 Ch. II, Lem. 3.3]. This implies that $\nabla q_k\to\nabla q$ uniformly along a subsequence. Assuming $(x_0,t_0)$ is a Lebesgue point of $|E(x,w)|$, the term in the parentheses remains bounded as $k\to\infty$ and thus $I_1\to0$ by uniform convergence of $\nabla q_k$. The integral $I_2$ is treated identically. Because such points $(x_0,t_0)$ comprise almost all of $\Omega_T$, passing to the limit in (16) shows that

as desired.

Given this, the first order necessary optimality conditions for this system are as follows.

Numerical approximation

In this section, we describe the numerical methods used to solve the FP optimal control problem (7). The numerical procedure consists of finite-difference approximations of the forward FP equation (4) and the backward adjoint equation (10), coupled through the sequential quadratic Hamiltonian (SQH) method. The latter iteratively updates the space- and time-dependent feedback control $u=(u_1,u_2,u_3)$ by pointwise minimization of a quadratically augmented Hamiltonian.

Numerical approximation of the forward FP equation

Let

denote a uniform discretization of $\Omega_T=\Omega\times(0,T)$, and let $f_{i,j}^n$ denote the numerical approximation of the probability density $\rho(x_i,y_j,t_n)$. Since $\rho$ represents a probability density, the numerical treatment of the forward FP equation should, as far as possible, retain the fundamental structural properties of the continuous problem. In particular, the numerical solution should remain nonnegative, conserve total probability under the zero-flux boundary condition, and accurately balance drift and diffusion without introducing spurious oscillations. For this reason, we use a two-dimensional Chang–Cooper discretization for the drift and diagonal diffusion terms.

We write the controlled FP equation as

where

The diagonal contributions in (18) can be written in conservative flux form as

and

This motivates the modified drift coefficients

In the implementation, these quantities are evaluated at the corresponding cell faces using finite differences of the diagonal diffusion coefficients. At an $x_1$-face $(x_{i+1/2},y_j)$, define

and analogously at an $x_2$-face,

The Chang–Cooper weights are determined by

with the limiting value

The quantity $w$ measures the relative strength of drift and diffusion across a grid cell. Thus, unlike a standard centered discretization, the Chang–Cooper weighting automatically biases the numerical flux in a manner determined by the local drift–diffusion balance. For example, the discrete flux through the face $(x_{i+1/2},y_j)$ may be written in the form

where

Similarly,

The diagonal drift–diffusion contribution at an interior grid point is therefore represented by the conservative flux difference

A principal advantage of the Chang–Cooper construction is that the drift and diffusion terms are not discretized independently. Instead, the weighting in (19) is chosen from their local ratio. For one-dimensional FP equations, this construction is well known to preserve nonnegativity under the appropriate time-step restriction and to reproduce the correct discrete equilibrium associated with zero probability flux. In the present two-dimensional implementation, the same construction is applied direction-by-direction to the diagonal drift–diffusion contributions. This significantly reduces the artificial numerical diffusion and spurious oscillations that may arise from a straightforward centered approximation.

The diffusion matrix in the present stochastic SIRS model is not diagonal, so the FP equation also contains the mixed term

This contribution is treated separately. Setting

we use the centered approximation

Consequently, if $\mathcal{L}^{n}_{\mathrm{CC}}$ denotes the Chang–Cooper discretization of the diagonal drift–diffusion terms and $\mathcal{L}^{n}_{\mathrm{mix}}$ denotes the approximation (20), the full explicit update can be written schematically as

Thus the diagonal part is advanced using Chang–Cooper fluxes, while the cross-diffusion generated by the common transmission noise is included through a centered mixed-derivative correction. This is the same decomposition used in the numerical implementation.

The time discretization in (21) is explicit, and therefore the time step must be chosen sufficiently small relative to the spatial mesh and the magnitudes of the drift and diffusion coefficients. Under the corresponding stability restriction, the Chang–Cooper component of the scheme retains the positivity-preserving character of the underlying discretization. Because the mixed derivative is added separately by a centered explicit approximation, positivity of the complete two-dimensional update is not automatic for an arbitrary time step. We therefore use a sufficiently fine temporal grid and monitor the minimum value of the computed density throughout the calculation.

An equally important property is conservation of total probability. For the continuous FP equation,

and hence the zero-flux boundary condition implies

The discrete Chang–Cooper terms are written as differences of neighboring numerical fluxes, so the interior contributions telescope when summed over all grid cells. Thus conservation reduces to the treatment of the boundary fluxes.

Accordingly, after every forward Euler update, the zero-flux boundary condition from (4) is imposed directly. Because the diffusion matrix contains the off-diagonal term $D_{12}$, the normal flux cannot be enforced using only the diagonal diffusion coefficients. For example, on a vertical boundary the condition involves

whereas on a horizontal boundary,

The boundary values of the density are therefore reconstructed so that the complete normal probability flux, including the cross-diffusion contribution, vanishes. The four corner values are obtained by averaging the neighboring boundary values. With this treatment, the spatial discretization retains the conservative character of the FP equation.

Numerical approximation of the adjoint equation

The adjoint equation (10) is solved backward in time from the terminal condition

Unlike the forward equation, which is written in conservative FP form, the adjoint equation is discretized directly in its backward Kolmogorov form,

The face-centered drift coefficients used by the forward Chang–Cooper solver are first averaged to the spatial grid points. For interior grid points, the first derivatives are approximated by centered differences,

while the diagonal second derivatives are approximated by the usual three-point formulas. The mixed derivative is approximated by

Consequently, if $\mathcal{L}_h$ denotes the resulting spatial finite-difference operator, the backward time update takes the form

The adjoint boundary condition

is also enforced directly at every backward time step. Thus, on the vertical boundaries the discretization accounts for the coupling between $D_{11}q_{x_1}$ and $D_{12}q_{x_2}$, while on the horizontal boundaries it accounts for the coupling between $D_{12}q_{x_1}$ and $D_{22}q_{x_2}$. The corner values are again obtained by averaging adjacent boundary values.

A sequential quadratic Hamiltonian method

We employ the sequential quadratic Hamiltonian (SQH) method to numerically resolve the optimal feedback maps. A full exposition of the SQH method with worked examples can be found in [4].

The SQH method is an iterative procedure based on successive quadratic perturbations of the Hamiltonian (12). For a penalization parameter $\varepsilon> 0$ and a current control approximation $\tilde{u}$, we define the augmented Hamiltonian

The quadratic term restricts the magnitude of the change in the control between consecutive SQH iterations. In particular, increasing $\varepsilon$ forces the new control to remain closer to the previous iterate, whereas decreasing $\varepsilon$ permits a larger update.

Starting from an initial control $u^{0}$, the corresponding forward FP solution $\rho^{0}$ and adjoint solution $q^{0}$ are computed. At iteration $k$, the new feedback control is obtained pointwise in $(x,t)$ by solving

This pointwise update is particularly convenient for the present problem because the three control components can be obtained explicitly and subsequently projected onto their admissible intervals.

For the mixed $L^{1}/L^{2}$ control cost used here, the pointwise minimization leads to explicit candidate values for each control. Writing

the updates implemented numerically are

and

where $\Pi_{[a,b]}$ denotes projection onto the interval $[a,b]$. After updating the controls, the drift and diffusion coefficients are recomputed and the forward FP equation is solved again.

The magnitude of the control update is measured by

The iteration terminates when $\tau_{k}<\kappa$ or when a prescribed maximum number of iterations has been reached.

To determine whether the candidate control is accepted, we employ a sufficient-decrease criterion. Specifically, the candidate is accepted when

If this condition fails, the candidate control and state are rejected, the previous iterate is restored, and the quadratic penalization parameter is increased according to

If the sufficient-decrease condition is satisfied, the candidate control is accepted, the corresponding adjoint equation is solved, and the penalization parameter is decreased according to

Thus, the SQH procedure adaptively balances the size of the control update against the decrease obtained in the objective functional.

Analysis of the SQH method

We next establish some properties of the SQH iteration. The first result gives a sufficient-decrease property analogous to that obtained for time-dependent controls, with the control norm now taken over the full space-time domain $\Omega_T$.

Proof. By definition of the SQH update, for almost every $(x,t)\in\Omega_T$,

In particular, choosing $w=u^k$ gives

Using the definition of the augmented Hamiltonian, we therefore obtain

for almost every $(x,t)\in\Omega_T$. Integrating over $\Omega_T$ and applying the Hamiltonian identity derived in § 3 gives

which proves (26).

As a consequence of Theorem 4.1, we obtain a basic convergence property of the accepted SQH iterates.

Proof. Statement (a) follows immediately from (26), since the objective functional is bounded below. For statement (b), using $\varepsilon_k\geq\eta$ gives

Summing from $k=0$ to $K$ yields the telescoping estimate

Since the objective functional is bounded below, the right-hand side remains bounded as $K \to\infty$. Hence

which implies

In particular, for any prescribed tolerance $\kappa> 0$, the stopping criterion

is eventually reached along the accepted SQH iterates. Moreover, the continuity properties of the control-to-state map established in § 3 imply that the corresponding changes in the FP states also tend to zero.

We finally give a stronger convergence statement for the mixed $L^2$–$L^1$ control cost used in the numerical experiments. Define the SQH update map

by

where $\rho_u = S(u)$ and $q_u$ denotes the corresponding adjoint state.

Proof. Under assumption (28), $\mathcal{T}_{\varepsilon}$ is a contraction on $(L^2(\Omega_T))^3$. Hence, by the Banach fixed-point theorem, there exists a unique $u^*$ such that

and

At the fixed point,

Since the quadratic SQH penalty vanishes at $w=u^*$,

Therefore the fixed-point condition gives

for almost every $(x,t) \in\Omega_T$, which is precisely the pointwise Hamiltonian minimum condition.

Numerical results

In this section, we present numerical results for the FP optimal control framework computed using the SQH algorithm. We set $b = \delta= 0.01$, so that the total population in the underlying SIRS model approaches unity. The state variables are the susceptible and infected fractions, $x = (S,I)$, and the FP equation is solved on the computational domain $\Omega= [0,1] \times[0,1]$. We discretize each spatial direction using 50 grid points and take the final time to be $T = 2$, with 10001 temporal grid points.

We take the transmission, recovery, and loss-of-immunity parameters to be $\beta= 3$, $\gamma= 1$, and $\lambda= 0.5$, respectively. The three feedback controls $u_1(S,I,t)$, $u_2(S,I,t)$, and $u_3(S,I,t)$ act by reducing transmission, reducing loss of immunity, and increasing recovery, respectively. Their admissible ranges are

Thus, unlike time-dependent open-loop interventions, the controls obtained from the FP formulation are feedback policies whose values depend on both the current epidemic state $(S,I)$ and time.

The initial probability density is taken to be a two-dimensional Gaussian centered at $(S_0,I_0) = (0.99,0.01)$. Specifically, the numerical implementation uses

restricted to the computational domain and normalized numerically to have unit mass. This represents uncertainty in the initial epidemic state around a population consisting predominantly of susceptible individuals with a small initially infected fraction.

Stochasticity is introduced through uncertainty in the transmission, loss-of-immunity, and recovery mechanisms. In the controlled stochastic system, the corresponding diffusion amplitudes are modulated by the same feedback controls that modify the drift, so that the transmission noise is scaled by $(1-u_1)$, the loss-of-immunity noise by $(1-u_2)$, and the recovery noise by $(1+u_3)$. In the simulations presented here, the baseline noise amplitudes are taken to be $\sigma_\beta= \sigma_\lambda= \sigma_\gamma= 0.5$.

The objective functional combines the state-dependent epidemic cost with a mixed $L^1/L^2$ penalty on the three feedback controls. In particular, the control contribution has the form

where we take $\beta_1 = 0.02,\beta_2 = 0.05$. The $L^1$ contribution penalizes the overall amount of intervention and can promote sparse or concentrated control actions, while the quadratic term regularizes the controls and discourages excessively large interventions.

Scenario 1. In the first scenario, we consider the running cost $G(S,I,t) = I$, $K(S,I) = 0$. Thus, infection is penalized continuously throughout the time interval, with larger values of the infected fraction incurring a proportionally larger state cost. No terminal-state cost is imposed. This provides a baseline test of the feedback-control framework in which the objective is to reduce the cumulative infection burden while balancing this reduction against the cost of the three interventions. The resulting controlled FP density and optimal feedback controls are shown in Figure 1. The controlled density evolves toward states with smaller values of $I$, while the feedback maps show how the three interventions are deployed over the $(S,I)$ state space and time. The control $u_1$ reduces transmission and hence suppresses the transition from susceptible to infected individuals, $u_2$ reduces the loss of immunity and therefore limits the return of recovered individuals to the susceptible class, and $u_3$ increases the recovery rate and promotes the transition from infected to recovered individuals. The distinct spatial and temporal structures of these controls illustrate how the optimizer combines prevention of new infections, retention of immunity, and enhanced recovery to reduce the cumulative infection burden.

The corresponding controlled and uncontrolled stochastic simulations are shown in Figure 2. In the uncontrolled system, the susceptible fraction decreases substantially while the infected fraction undergoes a pronounced increase, whereas the optimized feedback policy considerably limits both the loss of susceptible individuals and the growth of the infected population. The ensemble means exhibit the same qualitative behavior as the individual realizations. The mean controls evaluated along the controlled stochastic trajectories show how the state-dependent feedback policy is implemented dynamically: $u_1$ provides the principal

mechanism for suppressing new infections, while $u_2$ acts by retaining individuals in the recovered class and $u_3$ promotes removal from the infected class through enhanced recovery. Thus, the stochastic simulations confirm that the feedback controls reduce the infection burden relative to the uncontrolled dynamics while adapting the intervention to the evolving epidemic state.

Scenario 2. In this scenario, we consider running cost function $G(S,I,t)=100\,1_{\{I\geq0.08\}}(I-0.08)$, and zero terminal cost. This imposes no state penalty while the infected fraction remains below $0.08$, but strongly penalizes excursions above this threshold. This is meant to represent the goal of “flattening the curve” so as to not overwhelm the healthcare system. The controlled FP density and optimal feedback controls are shown in Figure 3. In contrast to Scenario 1, the optimizer is not rewarded for reducing $I$ throughout the entire state space; instead, control is valuable primarily when it reduces the probability of entering or remaining in the penalized region $I>0.08$. This is reflected in the state-dependent structure of the feedback maps. The control $u_1$ reduces transmission and therefore limits the flow from the susceptible to the infected class, $u_2$ reduces loss of immunity and retains individuals in the recovered class, and $u_3$ increases recovery and promotes removal from the infected class. Together, these mechanisms reshape the probability density so as to limit its movement into states with infection levels above the prescribed threshold, while avoiding unnecessary intervention in regions where the infection-dependent running cost vanishes.

The corresponding controlled and uncontrolled stochastic simulations are shown in Figure 4. The uncontrolled trajectories exhibit a pronounced increase in the infected fraction, with trajectories entering the region $I>0.08$ that is strongly penalized by the objective, together with a substantial decrease in the susceptible population. Under the optimized feedback policy, the growth of the infected population is markedly suppressed and the stochastic trajectories are largely kept away from the high-infection region, while the decline in the susceptible population is also reduced. The mean controls along the stochastic trajectories illustrate the mechanism by which this threshold regulation is achieved: $u_1$ suppresses the generation of new infections, $u_2$ reduces the return of recovered individuals to the susceptible class, and $u_3$ enhances the removal of infected individuals through recovery. The resulting behavior highlights the principal feature of this scenario: rather than minimizing infection uniformly, the feedback strategy allocates control effort specifically to reduce the probability and extent of excursions above the prescribed infection threshold.

Scenario 3. In this scenario, we consider $G(S,I,t) = 0$, $K(S,I) = -10\max(S-\bar{S},0)$, so that there is no running state cost and the objective instead provides a terminal reward for states in which the susceptible fraction exceeds the prescribed level $\bar{S}$. The resulting controlled FP solution and optimal feedback controls are shown in Figure 5. Since the objective acts only through the terminal susceptible population, the feedback controls are selected according to their ability to steer the probability density toward states satisfying $S > \bar{S}$ at the final time. The control $u_1$ contributes directly to this objective by reducing transmission and hence limiting the loss of susceptible individuals to the infected class. The control $u_2$ has a different effect: it reduces the return of recovered individuals to the susceptible class and therefore does not directly promote a larger value of $S$, while $u_3$ increases recovery from the infected to the recovered class and can influence the subsequent evolution of both $I$ and $S$. The resulting feedback maps therefore reflect the different contributions of these three mechanisms to the terminal susceptible-state objective rather than a direct penalty on the infected population during the time interval.

The corresponding controlled and uncontrolled stochastic simulations are shown in Figure 6. In the uncontrolled system, transmission produces a substantial decline in the susceptible population, whereas the optimized feedback policy limits this loss and maintains the controlled trajectories at larger values of $S$. This behavior is consistent with the terminal reward, which favors realizations that reach the final time with $S > \bar{S}$. Although the objective contains no direct running penalty on $I$, the controlled infected trajectories are also affected because the interventions alter the transitions among the susceptible, infected, and recovered populations. In particular, $u_1$ suppresses the $S \to I$ transition, $u_2$ suppresses the $R \to S$ transition, and $u_3$ enhances the $I \to R$ transition. The mean controls along the stochastic trajectories consequently illustrate a policy designed primarily to preserve the susceptible population at the terminal time, rather than to minimize the infected fraction throughout the entire time horizon.

Scenario 4. In this scenario, we consider $G(S,I,t)=\mathbf{1}_{\{I\geq0.08\}}(I-0.08)$, $K(S,I)=-\max(S-\bar{S},0)$, combining a threshold-based running penalty on infection with a terminal reward for maintaining the susceptible fraction above the prescribed level $\bar{S}$. The controlled FP density and optimal feedback controls are shown in Figure 7. In this case, the feedback policy must balance two state objectives: limiting the probability of occupying states with $I > 0.08$ during the evolution and steering the terminal density toward states with $S > \bar{S}$. The control $u_1$ is particularly relevant to both objectives because reducing transmission suppresses the $S \to I$ transition, simultaneously limiting the growth of infection and preserving susceptible individuals. The control $u_2$ reduces the loss of immunity and therefore suppresses the $R \to S$ transition, while $u_3$ increases recovery and enhances the $I \to R$ transition. The different state- and time-dependent structures of the three feedback maps reflect how the optimizer combines these mechanisms to balance suppression of

infection above the threshold with the terminal susceptible-population objective. The corresponding controlled and uncontrolled stochastic simulations are shown in Figure 8. The uncontrolled trajectories exhibit a substantial decrease in the susceptible fraction together with an increase in infection that carries the population into the penalized region $I > 0.08$. Under the optimized feedback policy, the growth of the infected population is reduced while a larger susceptible fraction is maintained, reflecting the simultaneous influence of the running and terminal objectives. The mean feedback controls along the stochastic trajectories show how the three intervention mechanisms contribute to this behavior: $u_1$ reduces the generation of new infections and the associated loss of susceptible individuals, $u_2$ retains recovered individuals in the immune class by reducing the $R \to S$ transition, and $u_3$ accelerates recovery through the $I \to R$ transition. Thus, Scenario 4 illustrates how the feedback framework balances short-term regulation of excursions above the infection threshold with the terminal objective of maintaining a sufficiently large susceptible population.

Conclusion

In this work, we developed a FP optimal control framework for a stochastic SIRS epidemic model with three intervention mechanisms targeting disease transmission, loss of immunity, and recovery. By incorporating space- and time-dependent feedback controls into both the drift and diffusion coefficients, the proposed approach enables interventions to influence the probability distribution of epidemic outcomes rather than individual trajectories alone. We investigated the existence of optimal controls and derived first-order necessary optimality conditions involving a forward Fokker–Planck equation, a backward adjoint equation, and a Hamiltonian minimization condition. For numerical computation, we combined a Chang–Cooper discretization with a finite-difference treatment of cross-diffusion and a SQH algorithm, together with an analysis of its convergence properties. Numerical experiments involving four epidemiological objectives demonstrated how the resulting feedback policies adapt to different goals, including reducing cumulative infection, limiting excursions above an infection threshold, and maintaining a desired susceptible population at the terminal time. Monte Carlo simulations further illustrated the effectiveness of the optimized controls relative to uncontrolled stochastic dynamics. Future work will focus on extending the framework to more complex epidemic models, strengthening the analysis and numerical treatment of degenerate diffusion, and incorporating epidemiological data for model calibration and validation.

Acknowledgements

S. Roy was supported by the US National Science Foundation grant nos. 2212938, 2309491, 2230790.