Physics-enhanced reinforcement learning for real-time optimal control of dynamical systems

Unknown

Conceptio › Archive › arXiv CS

arXiv CSopen access

Physics-enhanced reinforcement learning for real-time optimal control of dynamical systems

Unknown · 2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

Open Source ↗Direct PDF ↓

machine learning, deep learning, neural networks

Physics-enhanced reinforcement learning for real-time optimal control of dynamical systems Matteo Tomasetto† , Nicolò Botteghi‡ , Gabriele Bruni‡ , and Andrea Manzoni‡

arXiv:2607.16177v1 [cs.LG] 17 Jul 2026

†

Department of Mechanical Engineering, Politecnico di Milano, Milano, Italy and ‡ MOX - Department of Mathematics, Politecnico di Milano, Milano, Italy

Reinforcement learning (RL) has recently emerged as a promising feedback control strategy for nonlinear and complex dynamical systems. However, RL algorithms are sample inefficient and require a large number of interaction with the environment to synthesize optimal control strategies. Consequently, applications of RL are typically limited to sparse sensors and actuators due to the curse of dimensionality entailed by the exploration-exploitation dilemma in high-dimensional spaces. In this work, we bridge RL and traditional optimal control for dynamical system with a novel Physics-EnhAnced Reinforcement Learning (PEARL) paradigm tailored to the control of highdimensional and parametric dynamical systems, exploiting the differentibility of their dynamics. Specifically, PEARL employs an actor-adjoint algorithm that leverages automatic differentiation to compute policy gradients over short horizons and adjoint-based sensitivities of future returns approximated via neural networks, significantly reducing the number of environment interactions, while mitigating long-term gradient instabilities. Through two challenging parametric navigation problems in unsteady flows, we show that PEARL (i) effectively exploits differentiable environments to outperform state-of-the-art RL algorithms, (ii) is sample efficient, thanks to the physics-guided policy learning, (iii) generalizes across multiple scenarios, which is crucial when dealing with parametric systems, and (iv) enables scaling RL to high-dimensional state and action spaces, without requiring low-dimensional state representations or multi-agent strategies.

INTRODUCTION

Reinforcement learning (RL) is widely gaining traction in solving complex decision-making tasks [13, 57, 94, 95]. Specifically, RL considers an agent in an environment learning optimal decision-making strategies to maximize the expected cumulative reward over time through trialand-error. The reward is a task-dependent scalar signal – possibly time-delayed, sparse, and noisy – evaluating the system response to different agent’s actions, which is enough for discriminating decisions and guiding agent learning [91]. In the model-free RL paradigm, an explicit model of the environment is not available to the agent. Therefore, optimal actuation strategies can only be learned from the experience – i.e., the data – obtained through multiple environment interactions. Although broadly applicable in practice, the model-free assumption comes at the cost of severe sample inefficiency, as many simulations – i.e., interactions with the environment – are often required to explore (potentially continuous and high-dimensional) action spaces. In recent years, deep reinforcement learning (DRL) has emerged as a scalable (feedback) control strategy, capable of coping with continuous states and actions thanks to a neural network-based agent approximation [3, 33, 58]. DRL has been employed in several applied sciences and engineering fields, such as games [56, 69, 70, 87–89, 100, 116], simulated and real-world robotics [42, 54, 59, 79, 99, 117, 118], autonomous driving [39], smart buildings [65], power systems [38], and large language models [4, 52]. However, neural network training may be data hungry, extremely unstable due to the online nature of the RL setting, and entails an intractable computational cost for complex and high-dimensional tasks.

To enhance data efficiency, model-based reinforcement learning (MBRL) approaches have been proposed [63, 94]. Specifically, instead of solely interacting with the (real) environment, the agent investigates and evaluates new decisions using a surrogate model, which is learned from the interaction data and emulates the underlying and unknown dynamics of the environment. Nevertheless, learning an accurate, robust, and computationally efficient proxy of the environment may be extremely challenging or even prohibitive, as in the case of complex systems with distributed actuation, biasing the agent training towards sub-optimal control strategies and degrading its performance when deployed on the real system [7]. Many strategies have been proposed to overcome the MBRL limitations, thus enhancing robustness, generalization, and interpretability, such as latent world models and planning [43–45], short model-based rollouts [49], sparse identification of nonlinear dynamics [20, 119], and unsupervised state representation learning [16]. RL has been recently applied to control dynamical systems governed by differential equations, with particular emphasis on fluid dynamics [9, 21, 24, 31, 81, 83, 101, 103, 113], and plasma confinement [27]. See [36, 82, 86, 105, 106] for an extensive overview of RL in scientific computing. However, balancing exploration and exploitation – i.e., trying new actions to gain knowledge about the environment and selecting the best-known actions to maximize immediate reward – in continuous and high-dimensional action spaces is one of the major challenges in DRL. While standard exploration strategies are effective for low-dimensional control problems, earning high rewards using randomly sampled actions is extremely unlikely when dealing with high-dimensional dynamical systems, thus leading to inaccurate gradient

AGENT ACTION

adjoint net

autodi

DIFFERENTIABLE ENVIRONMENTS single-agent leader-follower game in a double gyre ﬂow

REWARD

STATE where

mean-ﬁeld leader-follower game in a double gyre ﬂow

FIG. 1. Graphical summary of Physics-EnhAnced Reinforcement Learning (PEARL). Starting from the current state of the system yk and the scenario parameters µ, the agent’s policy determines the optimal action to apply to the environment, that is uk = π(yk , µ; θ). As a results, the agent collects the subsequent state yk+1 = F (yk , uk , µ) and the reward rk , enabling continuous control of the dynamics with the feedback loop. To train the policy, PEARL employs an actor-adjoint algorithm where the gradient of the loss function over the generic short-horizon k = k0 , . . . , k0 + h is computed via automatic differentiation, thus backpropagating the gradients through the dynamics, and the adjoint equation, whose terminal condition λ̂k0 +h is approximated with a neural network to account for long-term dependencies.

estimates, non-significant agent updates, and premature convergence to local minima. Moreover, determining which action contributed to success or failure, namely the so-called credit assignment problem, is often prohibitive in this context. To tackle the curse of dimensionality, multi-agent reinforcement learning (MARL) employs multiple decision-making agents interacting with the environment to achieve certain goals [2, 22, 23, 98]. Multi-agent systems find applications in a large variety of domains such as, e.g., robotics and swarm systems [48, 75], as well as cooperative and competitive environments [60, 97], and, more recently, multi-agent flow control [17, 77, 92, 102, 104]. However, agents’ interactions, communication and coordination, as well as partial or local observability, may compromise convergence and stability [46]. In contrast to model-free RL, where the environment is treated as a black box, traditional optimal control (OC) theory directly incorporates the governing dynamics of the environment – which is often explicitly available in this context, at least approximately, through physical laws – to discover optimal control strategies via physicsconstrained optimization problems [14, 53, 66]. Specifically, a set of Karush-Kuhn-Tucker optimality condi-

tions can be obtained through the Lagrange multipliers’ method or the so-called adjoint state method [14, 53, 66], which is closely related to the Pontryagin maximum principle [80]. Note that, in this context, the problem is typically formulated as the minimization of a taskspecific loss function, which is equivalent to maximizing the expected cumulative reward, up to a change of sign in the objective. However, the OC paradigm typically focuses on open-loop actuation in deterministic settings, thus preventing generalization to multiple scenarios and allowing only limited robustness to external disturbances. Conversely, physics-based closed-loop control strategies may be designed through the Hamilton-JacobiBellman equation [6, 32, 53] or model predictive control (MPC) [25]. However, these solutions become computationally intractable for (possibly parametric) moderate to large-scale dynamical systems. For instance, MPC performs multiple physics-constrained optimizations online, thus introducing latency and violating realtime requirements in challenging high-dimensional settings [64, 107]. Similarly to MBRL, in the specific case of low-dimensional actuation, surrogate models may be considered to mitigate complexity, while sacrificing fidelity and undermining control performance.

3 Within the scope of RL, physical priors have been mainly utilized to alleviate the curse of dimensionality, ensure safe and efficient explorations, design reward functions, learn low-dimensional state representations, and improve the accuracy of surrogate models in MBRL [5]. For instance, Liu and Wang [62] exploit governing equations and physical constraints to inform the model learning in MBRL, thus enhancing generalization and performance in the low-data limit. With the advent of differentiable simulators and surrogate models, automatic differentiation (AD) has been utilized to compute exact cumulative-reward gradients through the dynamics, thus effectively enabling the training of neural network-based open-loop controllers [15, 47, 50] and policies [28, 51, 72]. For instance, Mokbel et al. [71] compute the reward function when stabilizing chaotic events in differential fluid simulation via AD. Instead, analytic policy gradients have been exploited by Freeman et al. [34] and Wiedemann et al. [110], while they have been estimated by Eberhard et al. [29] to find optimal control actions in model-free open-loop RL. However, when employing AD over long rollouts – a procedure typically referred to as backpropagation through time (BPTT) – the recursive multiplication of derivatives arising from the chain rule can lead to vanishing or exploding gradients [11], compromising the optimization. To mitigate gradient instabilities over long horizons, optimization routines over smaller time windows are typically employed, at the cost of introducing a bias. For example, Liu and MacArt [61] consider multiple policy updates over short horizons for active flow control problems. Moreover, Xu et al. [114] propose the short-horizon actor critic (SHAC) algorithm for differentiable environments, which mitigates the short-horizon bias by approximating the expected cumulative rewards the agent would accumulate beyond the short horizon with a value function neural network, thus avoiding the learning of myopic policies. To overcome the limitations of the aforementioned approaches, in this work we bridge the gap between RL and OC by introducing a novel Physics-EnhAnced Reinforcement Learning (PEARL) algorithm, where we leverage the physical knowledge of the environment to guide the learning process in a RL context, thus enhancing sample efficiency, while enabling optimal and adaptive closed-loop control actions in real-time. In particular, we propose an actor-adjoint algorithm tailored to the control of differentiable environments. Similarly to SHAC, we exploit AD to compute physics-enhanced gradients over short horizons, thus preventing vanishing or exploding gradients. However, rather than approximating the long-term rewards accumulated beyond the short horizon, we employ a neural network to predict its gradient – i.e the so-called adjoint or value gradient –, thus accounting for long-term dependencies at the policy gradient level, speeding up optimization and significantly improving sample efficiency. Throughout our numerical experiments, we demonstrate that PEARL outperforms state-of-the-art model-free RL algorithms, such as prox-

imal policy optimization [85] and twin-delayed deep deterministic policy gradient [35], and AD-based RL methods, such as BPTT and SHAC. A graphical summary of PEARL is available in Figure 1. The paper is organized as follows. Section II presents a unified overview of RL foundations and classical OC theory for dynamical systems, as well as RL algorithms for differentiable environments. Section III proposes a novel RL algorithm tailored to control differentiable environments. Section IV assesses the performance of the proposed control strategy in two challenging parametric navigation problems in unsteady flows. Eventually, Section V discusses results and future directions for advancing the proposed methodology toward large-scale and real-world problems.

II.

BRIDGING REINFORCEMENT LEARNING AND OPTIMAL CONTROL

FUNDAMENTALS OF REINFORCEMENT LEARNING

RL aims to design agents capable of making optimal decisions to maximize the cumulative reward. During training and deployment, agents interact with the underlying environment, which is typically modeled through Markov decision processes (MDPs) to account for stochasticity and uncertainties. Specifically, for a given time instant tk of the grid {t0 , . . . , tNt } discretizing the finite horizon [0, T ] with final time T > 0, the agent observes the state of the system yk ∈ RNy and takes action uk ∈ RNu . As a result, the system evolves according to the transition probability P(yk+1 |yk , uk ) : RNy × RNy × RNu → [0, 1]. Additionally, the agent receives the instantaneous scalar reward rk given by the problemspecific reward function R(yk , uk ) : RNy × RNu → R. Importantly, initial conditions y0 , environment dynamics P(yk+1 |yk , uk , µ), and reward functions R(yk , uk , µ) may depend on a set of parameters µ ∈ P ⊆ RNµ in the parameter space P, thus requiring different control strategies. Note that, for the sake of simplicity, we consider known and constant scenario parameters in this section, even though time-dependent parameters µ(t) can also be taken into account, as demonstrated in Section IV. The agent’s behavior is defined by the so-called policy function, which may be either stochastic – i.e., a distribution over the possible actions, conditioned on the current state – or deterministic. In the context of DRL, policies are modeled through neural networks with weights and biases θ ∈ RNθ . In this context, deterministic policies read as follows uk = π(yk , µ; θ)

for k = 0, . . . , Nt .

The agent’s goal is to find the optimal policy parameters maximizing the total expected cumulative reward G

4 across multiple initial states y0 and, if present, scenario parameters µ "N # t X k γ rk = E y∼ρ0 [Vπ(·;θ) (y, µ)], maxG(θ) = Eπ(·;θ) θ

µ∼ρµ

k=0

(1) where Eπ(·;θ) denotes the expected value given the policy π driving the system evolution, ρ0 and ρµ are, respectively, probability distributions accounting for different possible initial conditions and system parameters, and γ ∈ [0, 1] indicates the discount factor. Notably, the undiscounted setting, namely γ = 1, may be only considered for finite-horizon episodic tasks with guaranteed termination. Notice that, thanks to the law of total expectation, it is possible to formulate the optimization in terms of the so-called value function V : RNy × P → R: # "N t X Vπ (y, µ) = Eπ γ k rk |y0 = y, µ , (2) k=0

which measures how good the state y is in terms of the expected cumulative reward obtained when following the policy π starting from the state itself. RL methods can be classified in three main categories: (i) value based, (ii) policy based, and (iii) actor critic. Value-based approaches aim to estimate the value function V or the action-value function # "N t X γ k rk |y0 = y, u0 = u, µ , (3) Qπ (y, u, µ) = Eπ k=0

where Q : RNy × RNu × P → R is the quality function assessing the value of state-action pairs under the policy π driving decisions. After obtaining the value function, the optimal policy is simply derived by selecting the actions maximizing the value function, leading to the so-called greedy policy. Policy gradient methods, instead, maximize the expected cumulative reward directly through gradient ascent algorithms in the policy parameter space. To do so, the policy gradient theorem [96] is used to compute the gradient of the performance measure G with respect to the policy parameters θ, that is ∇θ G = E y∼ρπ [∇θ log π(u|y, µ; θ)Qπ (y, u, µ)] , u∼π µ∼ρµ

(4)

where u ∼ π(u|y, µ; θ) is the parametrized stochastic policy returning the action probability P(u|y, µ). Whenever a deterministic policy is employed, the gradient expression is given, instead, by the deterministic policy gradient theorem [90] h i ∇θ G = Ey∼ρπ ∇u Qπ (y, u, µ)|u=π(y,µ;θ) ∇θ π(y, µ; θ) . µ∼ρµ

(5) In practice, the gradients in Equations (4) and (5) are approximated through empirical means based on the trajectories sampled during training. For instance, the REINFORCE algorithm [111] exploits an unbiased policy

gradient approximation via Monte Carlo returns, at the cost of high variance, thus making learning unstable and data inefficient. Indeed, as shown in Appendix A, the rate of convergence of gradient ascent algorithms with inexact gradients is O √σN , where σ 2 is the variance of the gradient estimator and N is the numberof iterations, 2 resulting in a sample complexity of N = O σε2 to reach convergence up to a prescribed tolerance ε > 0. Different strategies have been proposed to reduce the variance of the gradient estimator, while introducing bias due to the bias-variance trade-off [40], such as control variates, advantage function approximation [84], and actorcritic strategies [41]. Actor-critic algorithms such as, e.g., proximal policy optimization (PPO) [85] and twindelayed deep deterministic policy gradient (TD3) [35] aim at jointly learning a policy – i.e., the actor – and a critic – i.e., the value function – with the goal of improving upon value- and policy-based methods. However, σ 2 generally scales with the amount of policy parameter Nθ [37, 74], as well as with the final time T and the problem dimensions Ny , Nu [1, 78, 93, 112], thus making model-free RL extremely inefficient, especially for high-dimensional problems over long horizons.

OPTIMAL CONTROL OF PARAMETRIC DYNAMICAL SYSTEMS

In this work, we take into account deterministic environments, such as those described by means of nonlinear, high-dimensional and parametric ordinary or partial differential equations (ODEs/PDEs) in the form ẏ = f (y, u, µ). Rather than relying on transition probabilities, the temporal evolution of the state y(t) ∈ RNy over t ∈ [0, T ] under control actions u(t) ∈ RNu is thus deterministic and governed by the vector field f : RNy × RNu × P → RNy modeling its time derivative ẏ. Note that the presented framework may be easily extended to stochastic or partially observable systems. Note also that, in the case of spatio-temporal systems, spatial discretization of the domain Ω through, e.g., finite element, finite volume or spectral methods [26, 55] is required to deal with finitedimensional variables. Moreover, similarly to MDPs, we discretize the dynamics over time, leading to yk+1 = F (yk , uk , µ),

(6)

where yk = y(tk ) and uk = u(tk ) for k = 0, . . . , Nt , while F : RNy × RNu × P → RNy is the discrete-time transition function. When dealing with deterministic environments with evolution described by Equation (6), the optimization problem in Equation (1) is analogous to the following constrained optimization [19, 86] – typically referred to

5 that is

as optimal control problem (OCP) [14, 53, 66] minJ(θ) = θ

N t −1 X

Z T minJ(θ) =

L(yk , π(yk , µ; θ), µ) + ϕ(yNt , µ),

k=0

s.t. ẏ = f (y, π(y, µ; θ), µ).

for k = 0, . . . , Nt − 1, (7) where the so-called cost or loss function J, with L : RNy × RNu × P → R the running cost and ϕ : RNy × P → R the final payoff, is negative of the performance measure G in deterministic and undiscounted settings. Notice that, while Equation (7) searches for optimal policy parameters across multiple scenarios and initial conditions, conventional OCPs directly optimize control actions in fixed configurations. When the environment dynamics is available, the gradient of the loss function in Equation (7) can be computed exactly through the adjoint state method [14, 53, 66] or, equivalently, backpropagation through time (BPTT) [108] and reverse-mode automatic differentiation (AD) [8, 67, 109], yielding s.t. yk+1 = F (yk , π(yk , µ; θ), µ)

∇θ J =

N t −1 X k=0

∂L ∂F + λ⊤ k+1 ∂uk ∂uk

∂π ∂θ

(8)

where λk (µ) ∈ RNy for k = 0, . . . , Nt are the closed-loop adjoint variables satisfying the (linear) adjoint equation running backward in time, that is λk =

∂F ∂F ∂π + ∂yk ∂uk ∂yk

⊤

λk+1 +

∂L ∂L ∂π + ∂yk ∂uk ∂yk

(10) As demonstrated in Appendix B, it is possible to compute the exact gradient of the loss function J with respect to the policy parameters θ through the adjoint state method, yielding Z T ∇θ J = 0

ADJOINT VARIABLE AS A SENSITIVITY MEASURE

This section delves into the role of the adjoint variable as a sensitivity measure, both in continuous and discrete-time, to enhance intuition and interpretability regarding physics-enhanced optimization, as well as to further bridge the gap between RL and traditional OC theory. For the sake of generality, we first consider the continuous-time extension of the OCP in Equation (7),

∂L ∂f + λ⊤ ∂u ∂u

∂π dt, ∂θ

(11)

where λ(t, µ) ∈ RNy is the closed-loop adjoint variable satisfying the (linear) adjoint ODE λ̇ = −

∂f ∂f ∂π + ∂y ∂u ∂y

⊤

λ−

∂L ∂L ∂π + ∂y ∂u ∂y

⊤ (12)

⊤

∂ϕ (y(T ), µ). As with terminal condition λ(T, µ) = ∂y demonstrated in Appendix D, the adjoint variable λ(t, µ) corresponds to the sensitivity of the value function to the state y, that is

λ(t, µ) = ∇y Vπ (y, t, µ)⊤ for t ∈ [0, T ],

(13)

where, in this context, the continuous-time value function Vπ (y, t, µ) under a generic fixed policy π reads as follows

⊤

(9) with terminal condition λNt = ∂y∂ϕN . See Appendix B for t the derivation of Equation (8) through the adjoint state method and Appendix C for the equivalence with respect to reverse-mode AD. Importantly, the computational cost of adjoint problems is independent of the number of policy parameters Nθ , in contrast to direct or forwardmode approaches. Moreover, in contrast to inexact gradient methods, first-order optimization with exact gradi 1 ents entails a convergence rate equal to O , yielding N a sample complexity equal to O 1ε [12, 37, 73], thus allowing us to cope with continuous and high-dimensional problems, as demonstrated in Section IV.

L(y(t), π(y(t), µ; θ), µ)dt + ϕ(y(T ), µ), 0

Z T Vπ (y, t, µ) =

L(y(τ ), π(y(τ ), µ; θ), µ)dτ +ϕ(y(T ), µ). t

(14) Significantly, the optimal value function V (y, t, µ) associated with the optimal policy is given by the HamiltonJacobi-Bellman equation [6, 32, 53] −

∂V = min[L(y, u, µ) + ∇y V f (y, u, µ))] ∂t u(t)

to be solved backward in time, with terminal condition V (y, T, µ) = ϕ(y(T ), µ). After time discretization, the OCP turns into Equation (7), with exact policy gradient in Equation (8) and adjoint variables λk (µ) = λ(tk , µ) for k = 0, . . . , Nt solving Equation (9). In this setting, as shown in Appendix D, the adjoint variables are equal to the sensitivity of the value function or the loss J with respect to the state yk , that is λk (µ) = ∇yk Vπ (yk , µ)⊤ = ∇yk J ⊤ for k = 0, . . . , Nt , (15) where the value function at the discrete level Vπ (yk , µ) for k = 0, . . . , Nt is defined as Vπ (yk , µ) =

N t −1 X κ=k

L(yκ , π(yκ , µ; θ), µ)+ϕ(yNt , µ). (16)

6 Notably, it is often convenient to write the value function in recursive form, thus obtaining the so-called Bellman equation [10] in deterministic and undiscounted settings Vπ (yk , µ) = L(yk , π(yk , µ; θ), µ) + Vπ (yk+1 , µ), which turns into the optimal Bellman equation [10] for the optimal value function V (yk , µ) associated with the optimal policy, that is

Note that, when considering the final payoff, the critic network prediction is added to ϕ instead of estimating the whole value function. To improve training stability and mitigate oscillations due to rapidly changing value estimates, SHAC considers two copies of the critic networks: the online network with parameters η, which is trained by minimizing the mean squared error Jψ (η) =

V (yk , µ) = min[L(yκ , π(yκ , µ; θ), µ) + Vπ (yk+1 , µ)]. π

SHORT-HORIZON POLICY OPTIMIZATION

When exploiting BPTT to learn optimal policy parameters over long rollouts (Nt ≫ 1), the recursive multiplication between derivatives can lead to vanishing or exploding gradients [11], severely compromising the optimization. To overcome this limitation, it is possible to truncate the loss computation, namely the cost/reward accumulation, over shorter horizons of h ≪ Nt steps, that is Jh (θ) =

k0X +h−1

L(yk , π(yk , µ; θ), µ) + ϕ(yk0 +h , µ),

k=k0

where k0 = nh is the starting point of the nth short horizon for n = 0, 1, . . . , Nh , and to optimize Jh in every short horizon. Specifically, one can consider gradient descent with descent direction k0X +h−1 ∂L ∂F ∂π ∇θ Jh = , + λ⊤ k+1 ∂uk ∂uk ∂θ

1 h

koX +h−1

V̄k − V̂ (yk , µ; η)

(17)

k=k0

where ∥·∥ denotes the Euclidean norm, the target values V̄k are computed via temporal-difference learning [94], and a target network with parameters η̄, which is employed to correct the loss function Jh and to compute the target values V̄k [70]. The target network is not update minimizing the mean square error in Equation (17) but, instead, softly updated using the online network parameters η. Despite providing an efficient and stable learning paradigm, the AD computation of the gradient ∇θ Jˆh relies on the value gradient ∇yk0 +h V̂ (yk0 +h , µ), which is derived by backpropagating through the critic network ψ. However, the value gradient (i.e., the adjoint) at the short-horizon terminal state yk0 +h can still be inaccurate even when the value itself is well approximated. The value gradient approximation can therefore be improved by leveraging the underlying physics via the adjoint equation in Equation (9), as proposed in the following section. III.

PHYSICS-ENHANCED REINFORCEMENT LEARNING

k=k0

where the adjoint variables are computed through the adjoint recursion introduced in Equation (9), starting from the terminal condition λk0 +h = ∂y∂ϕ . Importantly, k0 +h the adjoint variables here depend exclusively on the system dynamics within the short horizon, with no contribution from future information. Therefore, this strategy – commonly referred to as truncated BPTT – is capable of preventing gradient instabilities over long rollouts at the expense of introducing biased gradients and degraded long-term credit assignment. To mitigate the limitation of truncated BPTT, Xu et al. [114] propose the short-horizon actor critic (SHAC) algorithm for differentiable environments where a critic network ψ : RNy × P → R with weights and biases η̄ ∈ RNη corrects the value of the short-horizon terminal state yk0 +h in order to account for future dependencies, that is V̂ (yk0 +h , µ) = ϕ(yk0 +h , µ) + ψ(yk0 +h , µ; η̄), thus yielding the short-horizon loss Jˆh (θ) =

k0X +h−1 k=k0

L(yk , π(yk , µ; θ), µ) + V̂ (yk0 +h , µ).

In this section, we introduce the Physics-EnhAnced Reinforcement Learning (PEARL) algorithm tailored to control high-dimensional, distributed and parametric dynamical systems. To mitigate vanishing and exploding gradients over long rollouts, PEARL updates the policy parameters online after every short horizon of length H ≪ T , i.e., every h ≪ Nt time steps, while employing AD to compute the gradient descent directions that drive policy learning. Note that differentiable simulators are needed to backpropagate through the dynamical systems and compute derivatives. Similarly to SHAC, we take into account a neural network-based correction to mitigate the bias introduced with the short-horizon optimization. However, instead of approximating the value of the short-horizon terminal state V̂ (yk0 +h , µ), we directly estimate the adjoint – i.e., the value gradient – λ̂k0 +h that enters in the descent direction in order to account for long-term dependencies not observed yet, resulting in a value gradient learning strategy [30] or an actor-adjoint method. This paradigm shift is beneficial to achieve physics-informed value gradient estimates and, thus, more effective policy gradients. Indeed, rather than taking into account a temporal-difference (TD) scheme to compute target values and train the critic network, we

7 exploit the adjoint equation in Equation (9) to generate target adjoint variables and train the so-called adjoint network, better accounting for how state perturbations propagate through the underlying physical model.

ACTOR-ADJOINT METHOD

In this section, we detail the actor-adjoint method employed by PEARL to learn effective policies for steering dynamical systems, while drastically improving sample efficiency, robustness, and tackling the curse of dimensionality when compared to several state-of-the-are algorithms, such as PPO, TD3, BPTT, and SHAC. The policy parameters θ are updated online, while interacting with the environment, after every short horizon of h ≪ Nt time steps, towards the gradient descent direction b θ Jh = ∇

k0X +h−1 k=k0

⊤ ∂L ∂F + λ̂k+1 ∂uk ∂uk

∂π , ∂θ

(18)

which approximates the exact policy gradient in Equation (8) over short horizons to prevent instabilities. The adjoint approximations λ̂k for k = k0 + 1, . . . , k0 + h − 1 in Equation (18) are given in cascade by the adjoint recursion in Equation (9), with terminal condition λ̂k0 +h approximated through the so-called adjoint network φ : RNy × P → RNy having weights and biases ϑ̄ ∈ RNϑ , that is λ̂k0 +h =

∂ϕ + φ(yk0 +h , µ; ϑ̄). ∂yk0 +h

1 h

k0X +h−1

λ̄k − φ(yk , µ; ϑ)

(20) ∂ϕ where ∂y∂L = for the sake of compactness. Dif∂yk0 +h k0 +h ferently from SHAC, which considers TD schemes for the value function, Equation (20) allows us to provide physics-informed target value gradients when training the adjoint network φ. In particular, Equation (20) combines the current target adjoint variables and the corresponding adjoint net estimates in the right hand side through a parameter λ ∈ [0, 1], enabling credit assignment for longer temporal dependencies. Note that the differences between Equation (20) and the adjoint recursion in Equation (9) are due to the role of the adjoint net in Equation (19), which does not predict the entire adjoint variable but corrects the loss gradient available within the short horizon. Note also that, when considering a discount γ ∈ [0, 1), the target adjoint variables are rescaled by a factor γ k−k0 to remove the discounting contribution from the adjoint dynamics, thus facilitating the learning of the underlying value gradient by the adjoint network. The target adjoint network is then softly updated after each short horizon via the rule ϑ̄ ← αϑ̄ + (1 − α)ϑ with α ∈ [0, 1).

(19) IV.

We recall that, as discussed in Section II, the adjoint variables quantify the sensitivity of long-term losses (or rewards) to the current state. Therefore, when optimizing over short horizon, the adjoint variables are not available online due to their backward-in-time nature, and have to be approximated. Similarly to SHAC, to improve training stability and reduce oscillations caused by rapidly changing targets, we consider two copies of the adjoint network: an online network with parameters ϑ and a target one with parameters ϑ̄. Specifically, the former is updated through gradient descent after each short horizon, minimizing the mean squared error Jφ (ϑ) =

and more accurate long-term sensitivities. In particular, we consider the TD-λ scheme for k = k0 , . . . , k0 + h − 1: ⊤ ∂L ∂F ∂F ∂π λ̄k = + + λλ̄k+1 + ∂yk ∂uk ∂yk ∂yk+1 + (1 − λ)φ(yk+1 , µ; ϑ̄) ,

k=k0

, whereas the latter is employed to predict adjoint variables through Equation (19) and to compute target adjoint variables λ̄k for k = k0 , . . . , k0 +h−1. The target adjoint variables are computed through a TD scheme based upon the adjoint equation in Equation (9), thus providing sharper physics-enhanced gradients tied to the dynamics

NUMERICAL RESULTS

In this section, we assess the performance of PEARL in two parametric optimal navigation problems in unsteady flows, which are referred to as leader-follower game and mean-field leader follower game. Specifically, in Section IV A we aim at steering an agent – i.e., the follower – in a double gyre flow to track a target – i.e., the leader – passively transported by the underlying flow field. Note that the initial position of the leader and the follower may vary within the domain, therefore requiring a parametric policy capable of adapting to multiple scenarios. Moreover, Section IV B considers a high-dimensional variation of the leader-follower game, where we apply a distributed control action to steer a Gaussian density to follow a leader moving in the double gyre flow. In our numerical experiments, to deal with parametric control problems, a neural network architecture composed of two different branches to embed the state and the scenario parameters, which are modeled via feedforward neural networks with 2 layers and a problemspecific number of neurons. After encoding the state and the parameters, the corresponding embeddings are concatenated and further processed through a third 2layers neural network. For fairness of comparison, the

8 same architecture is used for all the different algorithms. Moreover, as far as the hyperparameters are concerned, we employ a discount factor γ = 0.99, TD scheme with λ = 0.95, and target network update with α = 0.995 both for SHAC and PEARL.

LEADER-FOLLOWER GAME

This first control problem involves an optimal navigation problem where an agent – i.e., the follower – has to track a target particle – i.e., the leader. Both the leader and the follower are passively transported by an unsteady double gyre flow field, making the learning of the optimal policy extremely challenging due to the flow disturbances and the limited actuation capabilities. The dynamics of the agent position xF (t) in the rectangular domain Ω = (0, 2) × (0, 1) over the time interval [0, 100] is defined as   dxF (t) = v(xF (t), t) + u(t) (21) dt x (0) = x F

F,0

where v denotes the double gyre flow velocity defined as −πA sin πf (x1 , t) cos(πx2 ) v(x, t) = (22) (x1 ,t) πA cos πf (x1 , t) sin(πx2 ) ∂f∂x 1 for every x = (x1 , x2 ) ∈ Ω and every t ∈ [0, T ], with A = 0.1 being the intensity coefficient, f (x1 , t) = ϵ sin(ωt)x21 + (1 − 2ϵ sin(ωt))x1 , while the perturbation amplitude and the frequency of oscillation are set equal to ϵ = 0.25 and ω = π, respectively. Moreover, u(t) represents the control velocity steering the agent, whose components are bounded in the interval (−0.2, 0.2), while xF,0 represents the initial agent position uniformly sampled in the region [0.1, 1.9] × [0.1, 0.9]. The leader, instead, has no external actuation and it is solely transported by the gyre-field velocity, that is   dxL (t) = v(xL (t), t) (23) dt x (0) = x , L

L,0

where the initial position is uniformly sampled in the region [0.1, 1.9]×[0.1, 0.9] at the beginning of each rollout. Importantly, the variability of the initial conditions of the leader and the follower requires parametric policies, capable of adapting to different scenarios. In this context, the follower position is available within the agent state, while the leader position is regarded as time-dependent scenario parameter µ(t) = xL (t), thus resulting in timevarying inputs for the policy and the adjoint network. To solve the state equations for the leader and the follower in Equation (21) and Equation (23), we consider a uniform grid of temporal instances {tk }k=0,...,Nt with Nt = 1001 and ∆t = 0.1 seconds, and we employ forward Euler,

yielding the discrete-time equations xF,k+1 = xF,k + ∆t(v(xF,k , tk ) + uk ) xL,k+1 = xL,k + ∆tv(xL,k , tk ) where uk = u(tk ), xF,k = xF (tk ) and xL,k = xL (tk ) for every k = 0, . . . , Nt . The goal of the agent is to reach and track the leader moving in the flow over time. To this aim, we investigate two possible reward functions. First, we consider the dense reward function Rd Rd (xF,k , uk , µk ) = − ||xF,k − µk ||2 − β||uk ||2 , {z } | {z } | state term

(24)

action term

where µk = xL,k for k = 0, . . . , Nt , and β = 0.2 is a scalar coefficient balancing the contribution of state and action terms. Furthermore, we consider a sparse reward function Rs 2

−100||xF,k −µk || 2 Rs (xF,k , uk , µk ) = 100e k || , | {z } − β||u | {z } state term

action term

assigning a bonus equal to 100 at every time step where the follower is close to the leader. Note that, to preserve differentiability, we approximate the bonus via a scaled exponential function. The state of the agent is composed of the particle position xF,k and the flow velocity v(xF,k , tk ) for every k = 0, . . . , Nt . The neural networks into play – that are the policy, the critic and the adjoint network, depending on the RL method – are modeled with the double encoding described at the beginning of Section IV, with 64 neurons. Moreover, we consider learning rates equal to 0.0001, h = 16 steps in the short horizons, 1500 episodes of training and 10 episodes of evaluation. Figure 2 shows the training and evaluation rewards achieved by different agents over training and evaluation when employing the dense and the sparse reward functions. In both cases, PEARL is able to achieve the best rewards, outperforming not only the model-free baselines PPO and TD3, but also the gradient-based ones, namely BPTT, truncated BPTT, and SHAC. Due to the length of the control horizon, BPTT does not achieve good results due to exploding gradients which compromise optimization. Truncated BPTT improves the BPTT results, but it is not capable of retrieving robust optimal policies due to the myopic reward maximization over short horizons. On the other side, the model-free baselines, while achieving sufficient performance with dense rewards, are not able to learn acceptable policies when the rewards are sparse. Despite the similar working principles, PEARL outperforms SHAC, showing that learning the value gradients directly is more effective that learning the values, especially in the case of sparse rewards. Eventually, Figure 3 shows three examples of controlled solution in the leader-follow game when using PEARL. PEARL is capable of successfully and effectively controlling the follower to track the moving target in a double gyre flow, making their trajectories almost overlapping after few seconds of simulation.

FIG. 2. Leader-follower game. Training and evaluation rewards obtained by the different competing agents in the leaderfollower game with dense and sparse rewards.

MEAN-FIELD LEADER-FOLLOWER GAME

In this section we consider a high-dimensional variation of the leader-follower game in a double-gyre flow. The leader dynamics is the same as the one introduced in Section IV A, with initial position regarded as timedependent scenario parameters µ. The agent, instead, is a density defined over the whole rectangular domain Ω. Note that the mean-field description may be helpful to model pollutant concentration, as well as swarms at the macroscopic level. Specifically, starting from a Gaussian density y0 centered at (y0,1 , y0,2 ) with variance 0.05, the density dynamics is defined with the diffusion-advection PDE  dy   + ∇ · (−ν∇y + vy + uy) = 0 in Ω × (0, T ]  dt (−ν∇y + vy + uy) · n = 0 on ∂Ω × (0, T ]   y = y in Ω × {t = 0}, 0 (25) where ν = 0.001 is the diffusivity, n is the normal unit vector, and homogeneous Neumann boundary conditions are taken into account on the domain boundary ∂Ω. Similarly to the previous test case, v(x, t) represent the double gyre flow velocity, while u(x, t) is the distributed control action defined over the whole space-time domain. Regarding the time discretization, we consider a uniform grid within the horizon [0, T ], with final time T = 20 sec-

onds and time step ∆t = 0.2 seconds, that is Nt = 100. Moreover, we spatially discretize the problem with finite elements, ending up with Ny = 2145 degrees of freedom for the discrete states yk , and Nu = 4290 for the discrete control action uk for k = 0, . . . , Nt . The goal is to steer the density towards the moving target for different initial mean position of the density (y0,1 , y0,2 ) in the region [0.3, 1.7] × [0.3, 0.7] and different initial position of the leader in the range [0.1, 1.9] × [0.1, 0.9]. To do so, we aim at maximizing the reward function at the continuous level Z Z 2 R(yk , uk , µk ) = − 0.5 (yk − ȳk ) dx − 0.5 yk2 dΓ− Ω ∂Ω Z Z − 0.5β ||uk ||2 dx − βg ||∇uk ||2 dx, Ω

Ω

where β = βg = 0.1 balance state and control terms. The neural networks of the competing methods, i.e., PPO, TD3, PEARL and SHAC-MOD, are modeled with the double encoding described at the beginning of Section IV, with 1024 neurons per layer. To emphasize the importance of the proposed neural network architecture with double encoding, we test the original implementation of SHAC considering a single feed forward neural network with 4 layers of 1024 neurons. Moreover, we consider learning rates equal to 0.00001, h = 16 steps in the short horizons, 1000 episodes of training and 10 episodes of evaluation. Figure 4 shows the training and evaluation

Test 2

Test 3

100 seconds

10 seconds

5 seconds

0 seconds

Test 1

FIG. 3. Leader-follower game. Examples of different parametric controlled solutions in the leader-follower game using PEARL. The follower is indicated by the magenta circle and the leader by the yellow square. The first 4 rows show the state evolution over time, while the last shows the evolution of the leader and follower coordinates over time.

rewards achieved by different agents over training and evaluation when employing the dense reward function in Equation (IV B). PEARL and SHAC-MOD are able to achieve the best rewards over training, yet PEARL is the best-performing agent during evaluation. SHAC

relying on the original neural network architecture proposed in [115] is not able to achieve rewards comparable to PEARL. In this high-dimensional control problem, instead, the model-free RL baselines (PPO and TD3) are not capable of learning any meaningful policy, experienc-

11 ing exploding gradients due to the stochastic exploration policy compromising the numerical stability of the solver. In Figure 5 and 6 we show three examples of controlled solution in the mean-field leader-follow game when using PEARL. Despite the high-dimensionality of the problem, PEARL is capable of controlling the density to track the target moving in a double gyre flow. Eventually, in Figure 7 we show a comparison between the controls of PEARL and PPO at different stages of the training. It is worth noticing that after the first updates the control actions learn by PEARL are smooth and capable of preventing the explosion of the numerical solver, and, consequently, of the rewards, leading to the efficient learning of the optimal policy.

CONCLUSIONS

In this work, we bridge reinforcement learning and traditional optimal control theory for the closed-loop control of differentiable dynamical systems. After highlighting connections between the two frameworks in Section II, as well as pros and cons, we propose a novel PhysicsEnhAnced Reinforcement Learning (PEARL) algorithm tailored to control of high-dimensional, distributed and parametrized differentiable environments. Unlike classical model-free RL approaches, PEARL directly exploits the governing physics of the environment to compute policy gradients via automatic differentiation, speeding up the optimization, and tackling the curse of dimensionality. By combining short-horizon optimization, automatic differentiation, and a neural network-based approximation of the terminal adjoint of each short horizon to account for long-term dependencies, PEARL mitigates the vanishing or exploding gradients that plague backpropagation through time over long rollouts and prevent the learning of myopic contorl policies that affects truncated backpropagation through time, without sacrificing the closed-loop, real-time applicability that classical optimal control methods typically lack. Importantly, we exploit the adjoint equation to generate target data for adjoint network training, thus providing more informative and effective gradients with respect to value learning meth-

[1] A. Agarwal, S. M. Kakade, J. D. Lee, and G. Mahajan. On the theory of policy gradient methods: Optimality, approximation, and distribution shift, 2020. [2] S. V. Albrecht, F. Christianos, and L. Schäfer. Multiagent reinforcement learning: Foundations and modern approaches. MIT Press, 2024. [3] K. Arulkumaran, M. P. Deisenroth, M. Brundage, and A. A. Bharath. Deep reinforcement learning: A brief survey. IEEE Signal Processing Magazine, 34(6):26–38, 2017. [4] Y. Bai, A. Jones, K. Ndousse, A. Askell, A. Chen, N. DasSarma, D. Drain, S. Fort, D. Ganguli, T. Henighan, N. Joseph, S. Kadavath, J. Kernion,

ods. Through two challenging parametric navigation problems in unsteady double-gyre flows – a leader-follower game and its mean-field, high-dimensional counterpart – we demonstrated that PEARL outperforms state-ofthe-art model-free algorithms, such as PPO and TD3, as well as AD-based alternatives, such as BPTT, truncated BPTT and SHAC. Notably, PEARL achieves outstanding performance without resorting to low-dimensional state representations or multi-agent decompositions, directly handling high-dimensional states and distributed actions, and generalizes across multiple scenario parameters within a single policy. These results indicate that embedding adjoint-based sensitivities directly into the policy gradient offers a principled and effective route to sample-efficient, physics-consistent policy learning. As future work, we will apply PEARL to larger-scale and more complex control problems, for example, arising in active flow control. In addition, we will extend our algorithm to deal with partially observable dynamics, in the direction of real-world settings. Moreover, modelbased extensions may be developed with differentiable surrogate models emulating the real environment dynamics in order to further enhance sample efficiency.

CODE AVAILABILITY

The code can be found in our repository https://github.com/MatteoTomasetto/pearl. ACKNOWLEDGEMENTS

NB and AM acknowledge the Project “Reduced Order Modeling and Deep Learning for the real-time approximation of PDEs (DREAM)” (Starting Grant No. FIS00003154), funded by the Italian Science Fund (FIS) - Ministero dell’Università e della Ricerca. AM also acknowledges the project “Dipartimento di Eccellenza” 2023-2027 funded by MUR.

T. Conerly, S. El-Showk, N. Elhage, Z. Hatfield-Dodds, D. Hernandez, T. Hume, S. Johnston, S. Kravec, L. Lovitt, N. Nanda, C. Olsson, D. Amodei, T. Brown, J. Clark, S. McCandlish, C. Olah, B. Mann, and J. Kaplan. Training a helpful and harmless assistant with reinforcement learning from human feedback, 2022. [5] C. Banerjee, K. Nguyen, C. Fookes, and M. Raissi. A survey on physics informed reinforcement learning: Review and open problems. arXiv preprint arXiv:2309.01909, 2023. [6] M. Bardi, I. C. Dolcetta, et al. Optimal control and viscosity solutions of Hamilton-Jacobi-Bellman equations, volume 12. Springer, 1997.

FIG. 4. Mean-field leader-follower game. Training and evaluation rewards obtained by the different competing agents in the leader-follower game with dense and sparse rewards.

[7] B. Barkley and D. Fridovich-Keil. Stealing that free lunch: Exposing the limits of dyna-style reinforcement learning. In A. Singh, M. Fazel, D. Hsu, S. LacosteJulien, F. Berkenkamp, T. Maharaj, K. Wagstaff, and J. Zhu, editors, Proceedings of the 42nd International Conference on Machine Learning, volume 267 of Proceedings of Machine Learning Research, pages 2978– 3002. PMLR, 2025. [8] A. G. Baydin, B. A. Pearlmutter, A. A. Radul, and J. M. Siskind. Automatic differentiation in machine learning: a survey. Journal of machine learning research, 18(153):1–43, 2018. [9] G. Beintema, A. Corbetta, L. Biferale, and F. Toschi. Controlling rayleigh–benard convection via reinforcement learning. Journal of Turbulence, 21(9-10):585–605, 2020. [10] R. Bellman. Dynamic Programming. Princeton University Press, Princeton, NJ, USA, 1 edition, 1957. [11] Y. Bengio, P. Simard, and P. Frasconi. Learning longterm dependencies with gradient descent is difficult. IEEE Transactions on Neural Networks, 5(2):157–166, 1994. [12] D. Bertsekas. Nonlinear Programming. Athena scientific optimization and computation series. Athena Scientific, 2016. [13] D. Bertsekas. Reinforcement learning and optimal control, volume 1. Athena Scientific, 2019. [14] D. P. Bertsekas. Dynamic programming and optimal control. Athena Scientific, Belmont, MA, 1995. [15] L. Böttcher, N. Antulov-Fantulin, and T. Asikis. Ai pontryagin or how artificial neural networks learn to control dynamical systems. Nature Communications, 13(1), 2022. [16] N. Botteghi, M. Poel, and C. Brune. Unsupervised representation learning in deep reinforcement learning: A review. IEEE Control Systems, 45(2):26–68, 2025. [17] N. Botteghi, M. Tomasetto, U. Fasel, F. Braghin, and A. Manzoni. Hypemarl: Multi-agent reinforcement learning for high-dimensional, parametric, and distributed systems, 2025. [18] J. Bradbury, R. Frostig, P. Hawkins, M. J. Johnson, Y. Katariya, C. Leary, D. Maclaurin, G. Necula, A. Paszke, J. VanderPlas, S. Wanderman-Milne, and Q. Zhang. JAX: composable transformations of Python+NumPy programs, 2018.

[19] S. L. Brunton and J. N. Kutz. Data-Driven Science and Engineering: Machine Learning, Dynamical Systems, and Control. Cambridge University Press, 2022. [20] S. L. Brunton, N. Zolman, J. N. Kutz, and U. Fasel. Machine learning for sparse nonlinear modeling and control. Annual Review of Control, Robotics, and Autonomous Systems, 8(Volume 8, 2025):127–152, 2025. [21] M. A. Bucci, O. Semeraro, A. Allauzen, G. Wisniewski, L. Cordier, and L. Mathelin. Control of chaotic systems by deep reinforcement learning. Proceedings of the Royal Society A, 475(2231):20190351, 2019. [22] L. Buşoniu, R. Babuška, and B. De Schutter. A comprehensive survey of multiagent reinforcement learning. IEEE Transactions on Systems, Man, and Cybernetics, Part C (Applications and Reviews), 38(2):156–172, 2008. [23] L. Buşoniu, R. Babuška, and B. De Schutter. Multiagent reinforcement learning: An overview. Innovations in multi-agent systems and applications-1, pages 183– 221, 2010. [24] M. Buzzicotti, L. Biferale, F. Bonaccorso, P. Clark di Leoni, and K. Gustavsson. Optimal control of point-topoint navigation in turbulent time dependent flows using reinforcement learning. 2020. [25] E. Camacho and C. Bordons. Model Predictive Control. Springer London, 2004. [26] R. Courant and D. Hilbert. Methods of mathematical physics: partial differential equations. John Wiley & Sons, 2008. [27] J. Degrave, F. Felici, J. Buchli, M. Neunert, B. Tracey, F. Carpanese, T. Ewalds, R. Hafner, A. Abdolmaleki, D. de Las Casas, et al. Magnetic control of tokamak plasmas through deep reinforcement learning. Nature, 602(7897):414–419, 2022. [28] M. P. Deisenroth and C. E. Rasmussen. Pilco: a modelbased and data-efficient approach to policy search. In Proceedings of the 28th International Conference on International Conference on Machine Learning, ICML’11, page 465–472, Madison, WI, USA, 2011. Omnipress. [29] O. Eberhard, C. Vernade, and M. Muehlebach. A pontryagin perspective on reinforcement learning. In N. Ozay, L. Balzano, D. Panagou, and A. Abate, editors, Proceedings of the 7th Annual Learning for Dynamics & Control Conference, volume 283 of Proceedings of Machine Learning Research, pages 233–244.

Test 2

Test 3

20 seconds

10 seconds

5 seconds

2 seconds

1 second

0 seconds

Test 1

0.0

0.5

1.0

1.5

2.0

2.5

3.0

FIG. 5. Mean-field leader-follower game. Examples of different parametric controlled solutions in the mean-field leader-follower game using PEARL. The leader is denoted by the yellow square. The different rows show the state evolution over time.

PMLR, 2025. [30] M. Fairbank. Reinforcement learning by value gradients, 2008. [31] D. Fan, L. Yang, Z. Wang, M. S. Triantafyllou, and G. E. Karniadakis. Reinforcement learning for bluff body active flow control in experiments and simulations. Proceedings of the National Academy of Sciences,

117(42):26091–26098, 2020. [32] W. H. Fleming and H. M. Soner. Controlled Markov processes and viscosity solutions. Stochastic Modelling and Applied Probability. Springer New York, NY, 2006. [33] V. François-Lavet, P. Henderson, R. Islam, M. G. Bellemare, J. Pineau, et al. An introduction to deep reinforcement learning. Foundations and Trends® in Ma-

Test 2

Test 3

20 seconds

10 seconds

5 seconds

2 seconds

1 second

0 seconds

Test 1

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

FIG. 6. Mean-field leader-follower game. Examples of different control actions in the mean-field leader-follower game using PEARL. The leader is denoted by the yellow square. The different rows show the action evolution over time.

chine Learning, 11(3-4):219–354, 2018. [34] C. D. Freeman, E. Frey, A. Raichuk, S. Girgin, I. Mordatch, and O. Bachem. Brax – a differentiable physics engine for large scale rigid body simulation, 2021. [35] S. Fujimoto, H. Hoof, and D. Meger. Addressing function approximation error in actor-critic methods. In International conference on machine learning, pages

1587–1596. PMLR, 2018. [36] P. Garnier, J. Viquerat, J. Rabault, A. Larcher, A. Kuhnle, and E. Hachem. A review on deep reinforcement learning for fluid mechanics. Computers & Fluids, 225:104973, 2021. [37] S. Ghadimi and G. Lan. Stochastic first- and zerothorder methods for nonconvex stochastic programming.

PEARL Control – 10 seconds

PEARL State – 10 seconds

Evaluation

Episode 10

Episode 5

Episode 3

Episode 2

Episode 1

PPO Control – 10 seconds

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

0.0

0.5

1.0

1.5

2.0

2.5

3.0

FIG. 7. Mean-field leader-follower game. Comparison between the controls of PEARL and PPO at several stages of the training. It is worth noticing that after the first updates the control actions learn by PEARL are smooth and capable of preventing the explosion of the numerical solver, and, consequently, of the rewards.

SIAM Journal on Optimization, 23(4):2341–2368, 2013. [38] M. Glavic. (deep) reinforcement learning for electric power system control and related problems: A short review and perspectives. Annual Reviews in Control, 48:22–35, 2019.

[39] S. Govinda, B. Brik, and S. Harous. A survey on deep reinforcement learning applications in autonomous systems: Applications, open challenges, and future directions. IEEE Transactions on Intelligent Transportation Systems, 26(7):11088–11113, 2025.

16 [40] E. Greensmith, P. Bartlett, and J. Baxter. Variance reduction techniques for gradient estimates in reinforcement learning. In T. Dietterich, S. Becker, and Z. Ghahramani, editors, Advances in Neural Information Processing Systems, volume 14. MIT Press, 2001. [41] I. Grondman, L. Busoniu, G. A. D. Lopes, and R. Babuska. A survey of actor-critic reinforcement learning: Standard and natural policy gradients. Trans. Sys. Man Cyber Part C, 42(6):1291–1307, 2012. [42] S. Gu, E. Holly, T. Lillicrap, and S. Levine. Deep reinforcement learning for robotic manipulation with asynchronous off-policy updates. In 2017 IEEE international conference on robotics and automation (ICRA), pages 3389–3396. IEEE, 2017. [43] D. Hafner, T. Lillicrap, J. Ba, and M. Norouzi. Dream to control: Learning behaviors by latent imagination, 2020. [44] D. Hafner, T. Lillicrap, I. Fischer, R. Villegas, D. Ha, H. Lee, and J. Davidson. Learning latent dynamics for planning from pixels, 2019. [45] D. Hafner, J. Pasukonis, J. Ba, and T. Lillicrap. Mastering diverse control tasks through world models. Nature, 640(8059):647–653, 2025. [46] P. Hernandez-Leal, B. Kartal, and M. E. Taylor. A survey and critique of multiagent deep reinforcement learning. Autonomous Agents and Multi-Agent Systems, 33(6):750–797, 2019. [47] P. Holl, N. Thuerey, and V. Koltun. Learning to control pdes with differentiable physics. In International Conference on Learning Representations, 2020. [48] M. Hüttenrauch, A. Šošić, and G. Neumann. Deep reinforcement learning for swarm systems. Journal of Machine Learning Research, 20(54):1–31, 2019. [49] M. Janner, J. Fu, M. Zhang, and S. Levine. When to trust your model: Model-based policy optimization, 2021. [50] W. Jin, Z. Wang, Z. Yang, and S. Mou. Pontryagin differentiable programming: An end-to-end learning and control framework. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 7979–7992. Curran Associates, Inc., 2020. [51] P. Karnakov, L. Amoudruz, and P. Koumoutsakos. Optimal navigation in microfluidics via the optimization of a discrete loss. Phys. Rev. Lett., 134:044001, 2025. [52] T. Kaufmann, P. Weng, V. Bengs, and E. Hüllermeier. A survey of reinforcement learning from human feedback, 2025. [53] D. E. Kirk. Optimal control theory: an introduction. Courier Corporation, 2004. [54] J. Kober, J. A. Bagnell, and J. Peters. Reinforcement learning in robotics: A survey. The International Journal of Robotics Research, 32(11):1238–1274, 2013. [55] J. N. Kutz. Data-driven modeling & scientific computation: methods for complex systems & big data. Oxford University Press, 2013. [56] G. Lample and D. S. Chaplot. Playing fps games with deep reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 31, 2017. [57] F. L. Lewis and D. Liu. Reinforcement learning and approximate dynamic programming for feedback control. John Wiley & Sons, 2013. [58] Y. Li. Deep reinforcement learning: An overview. arXiv preprint arXiv:1701.07274, 2017.

[59] L. Lin. Reinforcement learning for robots using neural networks. Carnegie Mellon University, 1992. [60] M. L. Littman. Markov games as a framework for multiagent reinforcement learning. In Machine learning proceedings 1994, pages 157–163. Elsevier, 1994. [61] X. Liu and J. F. MacArt. Adjoint-based machine learning for active flow control. Phys. Rev. Fluids, 9:013901, 2024. [62] X.-Y. Liu and J.-X. Wang. Physics-informed dyna-style model-based deep reinforcement learning for dynamic control. Proceedings of the Royal Society A: Mathematical, Physical and Engineering Sciences, 477(2255), 2021. [63] F.-M. Luo, T. Xu, H. Lai, X.-H. Chen, W. Zhang, and Y. Yu. A survey on model-based reinforcement learning. Science China Information Sciences, 67(2), 2024. [64] M. Löhning, M. Reble, J. Hasenauer, S. Yu, and F. Allgöwer. Model predictive control using reduced order models: Guaranteed stability for constrained linear systems. Journal of Process Control, 24(11):1647–1659, 2014. [65] A. Manjavacas, A. Campoy-Nieves, J. Jiménez-Raboso, M. Molina-Solana, and J. Gómez-Romero. An experimental evaluation of deep reinforcement learning algorithms for hvac control. Artificial Intelligence Review, 57(7), 2024. [66] A. Manzoni, A. Quarteroni, and S. Salsa. Optimal control of partial differential equations. Springer, 2021. [67] C. C. Margossian. A review of automatic differentiation and its efficient implementation. Wiley interdisciplinary reviews: data mining and knowledge discovery, 9(4):e1305, 2019. [68] S. Mitusch, S. Funke, and J. Dokken. dolfin-adjoint 2018.1: automated adjoints for fenics and firedrake. Journal of Open Source Software, 4(38):1292, 2019. [69] V. Mnih, K. Kavukcuoglu, D. Silver, A. Graves, I. Antonoglou, D. Wierstra, and M. Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint arXiv:1312.5602, 2013. [70] V. Mnih, K. Kavukcuoglu, D. Silver, A. A. Rusu, J. Veness, M. G. Bellemare, A. Graves, M. Riedmiller, A. K. Fidjeland, G. Ostrovski, et al. Human-level control through deep reinforcement learning. Nature, 518(7540):529, 2015. [71] S. Mokbel, C. Lagemann, E. Lagemann, and S. L. Brunton. Controlling chaotic energy events in fluids with reinforcement learning. In Proceedings of the 16th ACM International Conference on Future and Sustainable Energy Systems, E-Energy ’25, page 954–958. ACM, 2025. [72] M. A. Z. Mora, M. Peychev, S. Ha, M. Vechev, and S. Coros. Pods: Policy optimization via differentiable simulation. In M. Meila and T. Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 7805–7817. PMLR, 2021. [73] Y. Nesterov. Introductory Lectures on Convex Optimization: A Basic Course, volume 87 of Applied Optimization. Springer Science & Business Media, New York, 2013. [74] Y. Nesterov and V. Spokoiny. Random gradient-free minimization of convex functions. Foundations of Computational Mathematics, 17(2):527–566, 2015. [75] J. Orr and A. Dutta. Multi-agent deep reinforcement learning for multi-robot applications: A survey. Sensors, 23(7):3625, 2023.

17 [76] A. Paszke, S. Gross, S. Chintala, G. Chanan, E. Yang, Z. DeVito, Z. Lin, A. Desmaison, L. Antiga, and A. Lerer. Automatic differentiation in pytorch. 2017. [77] S. Peitz, J. Stenner, V. Chidananda, O. Wallscheid, S. L. Brunton, and K. Taira. Distributed control of partial differential equations using convolutional reinforcement learning. Physica D: Nonlinear Phenomena, 461:134096, 2024. [78] J. Peters and S. Schaal. Reinforcement learning of motor skills with policy gradients. Neural Networks, 21(4):682–697, 2008. [79] A. S. Polydoros and L. Nalpantidis. Survey of modelbased reinforcement learning: Applications on robotics. Journal of Intelligent & Robotic Systems, 86(2):153– 173, 2017. [80] L. Pontryagin, V. G. Boltyanskii, R. V. Gamkrelidze, and E. F. Mishchenko. The mathematical theory of optimal processes. Wiley, NY, 1962. [81] J. Rabault, M. Kuchta, A. Jensen, and U. Réglade. Artificial neural networks trained through deep reinforcement learning discover control strategies for active flow control. Journal of Fluid Mechanics, 865:281–302, 2019. [82] J. Rabault, F. Ren, W. Zhang, H. Tang, and H. Xu. Deep reinforcement learning in fluid mechanics: A promising method for both active flow control and shape optimization. Journal of Hydrodynamics, 32:234–246, 2020. [83] F. Ren, C. Wang, and H. Tang. Bluff body uses deep-reinforcement-learning trained active flow control to achieve hydrodynamic stealth. Physics of Fluids, 33(9):093602, 2021. [84] J. Schulman, P. Moritz, S. Levine, M. Jordan, and P. Abbeel. High-dimensional continuous control using generalized advantage estimation, 2018. [85] J. Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017. [86] O. Semeraro. Reinforcement Learning for Fluid Mechanics: an overview on Fundamentals from a Control Perspective. In Machine Learning for Fluid Dynamics. 2025. [87] K. Shao, Z. Tang, Y. Zhu, N. Li, and D. Zhao. A survey of deep reinforcement learning in video games. arXiv preprint arXiv:1912.10944, 2019. [88] D. Silver, A. Huang, C. J. Maddison, A. Guez, L. Sifre, G. van den Driessche, J. Schrittwieser, I. Antonoglou, V. Panneershelvam, M. Lanctot, S. Dieleman, D. Grewe, J. Nham, N. Kalchbrenner, I. Sutskever, T. Lillicrap, M. Leach, K. Kavukcuoglu, T. Graepel, and D. Hassabis. Mastering the game of go with deep neural networks and tree search. Nature, 529(7587):484–489, 2016. [89] D. Silver, T. Hubert, J. Schrittwieser, I. Antonoglou, M. Lai, A. Guez, M. Lanctot, L. Sifre, D. Kumaran, T. Graepel, T. Lillicrap, K. Simonyan, and D. Hassabis. Mastering chess and shogi by self-play with a general reinforcement learning algorithm, 2017. [90] D. Silver, G. Lever, N. Heess, T. Degris, D. Wierstra, and M. Riedmiller. Deterministic policy gradient algorithms. In International conference on machine learning, pages 387–395. Pmlr, 2014. [91] D. Silver, S. Singh, D. Precup, and R. S. Sutton. Reward is enough. Artificial Intelligence, 299:103535, 2021. [92] P. Suárez, F. Alcántara-Ávila, A. Miró, J. Rabault,

B. Font, O. Lehmkuhl, and R. Vinuesa. Active flow control for drag reduction through multi-agent reinforcement learning on a turbulent cylinder at r e d= 3900. Flow, Turbulence and Combustion, pages 1–25, 2025. [93] H. J. Suh, M. Simchowitz, K. Zhang, and R. Tedrake. Do differentiable simulators give better policy gradients? In K. Chaudhuri, S. Jegelka, L. Song, C. Szepesvari, G. Niu, and S. Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 20668–20696. PMLR, 2022. [94] R. S. Sutton and A. G. Barto. Reinforcement learning: An introduction, volume 1. MIT press Cambridge, 1998. [95] R. S. Sutton, A. G. Barto, and R. J. Williams. Reinforcement learning is direct adaptive optimal control. IEEE control systems magazine, 12(2):19–22, 1992. [96] R. S. Sutton, D. McAllester, S. Singh, and Y. Mansour. Policy gradient methods for reinforcement learning with function approximation. Advances in neural information processing systems, 12, 1999. [97] A. Tampuu, T. Matiisen, D. Kodelja, I. Kuzovkin, K. Korjus, J. Aru, J. Aru, and R. Vicente. Multiagent cooperation and competition with deep reinforcement learning. PloS one, 12(4):e0172395, 2017. [98] M. Tan. Multi-agent reinforcement learning: Independent vs. cooperative agents. In Proceedings of the tenth international conference on machine learning, pages 330–337, 1993. [99] C. Tang, B. Abbatematteo, J. Hu, R. Chandra, R. Martı́n-Martı́n, and P. Stone. Deep reinforcement learning for robotics: A survey of real-world successes. Annual Review of Control, Robotics, and Autonomous Systems, 8(Volume 8, 2025):153–188, 2025. [100] H. Van Hasselt, A. Guez, and D. Silver. Deep reinforcement learning with double q-learning. In Proceedings of the AAAI conference on artificial intelligence, volume 30, 2016. [101] P. Varela, P. Suárez, F. Alcántara-Ávila, A. Miró, J. Rabault, B. Font, L. M. Garcı́a-Cuevas, O. Lehmkuhl, and R. Vinuesa. Deep reinforcement learning for flow control exploits different physics for increasing reynolds number regimes. Actuators, 11(12), 2022. [102] J. Vasanth, J. Rabault, F. Alcántara-Ávila, M. Mortensen, and R. Vinuesa. Multi-agent reinforcement learning for the control of three-dimensional rayleigh-bénard convection. arXiv:2407.21565, 2024. [103] S. Verma, G. Novati, and P. Koumoutsakos. Efficient collective swimming by harnessing vortices through deep reinforcement learning. Proceedings of the National Academy of Sciences of the United States of America, 115(23):5849–5854, 2018. [104] C. Vignon, J. Rabault, J. Vasanth, F. Alcántara-Ávila, M. Mortensen, and R. Vinuesa. Effective control of two-dimensional Rayleigh–Bénard convection: Invariant multi-agent reinforcement learning is all you need. Physics of Fluids, 35(6), 2023. [105] C. Vignon, J. Rabault, and R. Vinuesa. Recent advances in applying deep reinforcement learning for flow control: Perspectives and future directions. Physics of Fluids, 35(3), 2023. [106] R. Vinuesa, O. Lehmkuhl, A. Lozano-Durán, and J. Rabault. Flow control in wings and discovery of novel approaches via deep reinforcement learning. Flu-

18 ids, 7(2), 2022. [107] E. Weinan, J. Han, and J. Long. Empowering optimal control with machine learning: A perspective from model predictive control. IFAC-PapersOnLine, 55(30):121–126, 2022. 25th International Symposium on Mathematical Theory of Networks and Systems MTNS 2022. [108] P. J. Werbos. Backpropagation through time: what it does and how to do it. Proceedings of the IEEE, 78(10):1550–1560, 2002. [109] P. J. Werbos. Backwards differentiation in ad and neural nets: Past links and new opportunities. In M. Bücker, G. Corliss, U. Naumann, P. Hovland, and B. Norris, editors, Automatic Differentiation: Applications, Theory, and Implementations, pages 15–34, Berlin, Heidelberg, 2006. Springer Berlin Heidelberg. [110] N. Wiedemann, V. Wüest, A. Loquercio, M. Müller, D. Floreano, and D. Scaramuzza. Training efficient controllers via analytic policy gradient, 2023. [111] R. J. Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine learning, 8(3-4):229–256, 1992. [112] C. Wu, A. Rajeswaran, Y. Duan, V. Kumar, A. M. Bayen, S. Kakade, I. Mordatch, and P. Abbeel. Variance reduction for policy gradient with action-dependent factorized baselines, 2018. [113] C. Xia, J. Zhang, E. C. Kerrigan, and G. Rigas. Active flow control for bluff body drag reduction using rein-

forcement learning with partial measurements. Journal of Fluid Mechanics, 981:A17, 2024. [114] J. Xu, V. Makoviychuk, Y. Narang, F. Ramos, W. Matusik, A. Garg, and M. Macklin. Accelerated policy learning with parallel differentiable simulation, 2022. [115] J. Xu, V. Makoviychuk, Y. Narang, F. Ramos, W. Matusik, A. Garg, and M. Macklin. Accelerated policy learning with parallel differentiable simulation. arXiv preprint arXiv:2204.07137, 2022. [116] D. Ye, Z. Liu, M. Sun, B. Shi, P. Zhao, H. Wu, H. Yu, S. Yang, X. Wu, Q. Guo, et al. Mastering complex control in moba games with deep reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 6672–6679, 2020. [117] F. Zhang, J. Leitner, M. Milford, B. Upcroft, and P. Corke. Towards vision-based deep reinforcement learning for robotic motion control. In Australasian Conference on Robotics and Automation 2015. Australian Robotics and Automation Association (ARAA), 2015. [118] W. Zhao, J. P. Queralta, and T. Westerlund. Sim-toreal transfer in deep reinforcement learning for robotics: a survey. In 2020 IEEE symposium series on computational intelligence (SSCI), pages 737–744. IEEE, 2020. [119] N. Zolman, C. Lagemann, U. Fasel, J. N. Kutz, and S. L. Brunton. Sindy-rl for interpretable and efficient modelbased reinforcement learning. Nature Communications, 16(1), 2025.

Appendix A: Sample complexity of inexact gradient algorithms

In this section, we derive the sample complexity of the gradient ascent algorithm when the ascent direction is approximated through a stochastic estimator. Let us consider the gradient ascent algorithm in the policy parameter space θ (n+1) = θ (n) + ηĝn for n = 1, . . . , N

(A1)

where η > 0 is the learning rate, N > 0 is the number of iterations, while ĝn ∈ RNθ is an unbiased estimator of the policy gradient ∇θ G(θ (n) ) for every iteration n = 1, . . . , N , i.e. ĝn = ∇θ G(θ (n) ) + ξ n with E[ξ n ] = 0 ∈ RNθ . Moreover, let us assume that the gradient estimate has bounded variance, that is E[∥ξ n ∥2 ] ≤ σ 2 for every iteration n = 1, . . . , N , where ∥·∥ denotes the Euclidean norm. Notice that, while throughout this work we consider gradients as row vectors, in this section they are regarded as column vectors to simplify the notation. Consistently with [73, 74], we assume that G is L-smooth, that is ∥∇θ G(θ 1 ) − ∇θ G(θ 2 )∥ ≤ L∥θ 1 − θ 2 ∥ for all θ 1 , θ 2 ∈ RNθ with constant L > 0. Notice that the negative of G is also L-smooth. The L-smoothness of the negative of G is equivalent to [73, 74] G(θ (n+1) ) ≥ G(θ (n) ) + ∇θ G(θ (n) )⊤ (θ (n+1) − θ (n) ) −

L (n+1) ∥θ − θ (n) ∥2 . 2

Exploiting the inexact gradient ascent update in Equation (A1), we obtain G(θ (n+1) ) ≥ G(θ (n) ) + η∇θ G(θ (n) )⊤ ĝn −

Lη 2 ∥ĝn ∥2 . 2

Taking the expectation of Equation (A2), we get G(θ

(n+1)

) ≥ G(θ

(n)

Lη Lη 2 2 )+η 1− ∥∇θ G(θ (n) )∥2 − σ 2 2

(A2)

19 It is now possible to average over the iterations n = 1, . . . , N , exploiting the telescopic sum and obtaining N Lη X G(θ (N +1) ) − G(θ (1) ) Lη 2 N 2 ≥η 1− σ , ∥∇θ G(θ (n) )∥2 − N 2 n=1 2 that is, assuming η < L2 and since G(θ ∗ ) ≥ G(θ (N +1) ) due to the optimality of θ ∗ , N X

∥∇θ G(θ (n) )∥2 ≤

n=1

Lη 2 N 2(G(θ ∗ ) − G(θ (1) )) + σ2 η(2 − Lη)N η(2 − Lη)

The learning rate η can be selected in order to minimize the right-hand side of the above inequality and achieve a sharper upper bound, even though the same result can be proven also with different step size rules [73, 74]. Doing the math, one can obtain N 1 X σ min ∥∇θ G(θ (n) )∥2 ≤ , ∥∇θ G(θ (n) )∥2 = O √ n=1,...,N N n=1 N 2 thus guaranteeing the convergence of the algorithm with rate O √σN . In other words, N = O σε2 iterations are required to achieve

min ∥∇θ G(θ (n) )∥2 ≤ ε, with ε > 0 being the prescribed tolerance.

n=1,...,N

Appendix B: Adjoint state method

In this section, we derive the loss gradient formulas in Equations (11) and (8), both in continuous-time and after time discretization. To compute the exact gradient of the loss function J with respect to the policy parameters θ, one may exploit a Hamiltonian formulation, in accordance with the Pontryagin maximum principle [80], a Lagrangian multiplier approach or, as considered in the following, the adjoint state method. Let us consider the OCP in Equation (10), where the control is determined through the policy u = π(y, µ; θ). Importantly, the adjoint state method is traditionally applied to directly determine the open-loop optimal control actions, i.e. to compute the gradient ∇u J, while we optimize with respect to the policy parameters θ, thus focusing on closed-loop control strategies. The loss gradient ∇θ J may be initially computed via the chain rule, yielding Z T ∂ϕ dy ∂L ∂L ∂π dy ∂L ∂π + + dt + (y(T )) (T ), ∇θ J = ∂y ∂u ∂y dθ ∂u ∂θ ∂y dθ 0 ∂ϕ dµ ∂L where ∂J ∂θ = ∂θ = ∂θ = 0 and dθ = 0. Note that we omit the dependence on µ for the sake of compactness. Note also that, in order to have a well-defined expression, we assume that all the components – i.e. the policy π and the loss function terms L and ϕ are differentiable. The rationale behind the adjoint state method is to introduce an auxiliary time-dependent variable, that is the so-called adjoint or co-state variable λ(t, µ) ∈ RNy , in order to get rid of the Jacobian dy dθ , which is prohibitive to compute. Note that, in general, the semi-discrete adjoint variable may have a different dimension than Ny . However, we take into account the same spatial discretization for state and adjoint variables. Let us compute the so-called tangent equation by differentiating the state equation with respect to the policy parameters  dẏ ∂f ∂f ∂π dy ∂f ∂π   = + + dθ ∂y ∂u ∂y dθ ∂u ∂θ (B1)   dy (0) = dy0 = 0 dθ dθ dµ where ∂f ∂θ = 0 and dθ = 0. It is now possible to add Equation (B1) multiplied by the negative of the adjoint variable to the time integral of the loss gradient formula above Z T ∂L ∂L ∂π dy ∂L ∂π dẏ ∂f ∂f ∂π dy ∂f ∂π ∂ϕ dy ∇θ J = + + − λ⊤ − + − dt + (y(T )) (T ). ∂y ∂u ∂y dθ ∂u ∂θ dθ ∂y ∂u ∂y dθ ∂u ∂θ ∂y dθ 0

After integrating by parts the term λ⊤ ddθẏ , we have Z T ⊤ ∂f ∂f ∂π ∂L ∂L ∂π dy ∂L ∂π ∂ϕ dy ⊤ ⊤ ∂f ∂π ⊤ ∇θ J = λ̇ + λ + + + + +λ dt + (y(T )) − λ (T ) (T ). ∂y ∂u ∂y ∂y ∂u ∂y dθ ∂u ∂θ ∂u ∂θ ∂y dθ 0

FIG. 8. One-step computational graph of a discrete-time optimal control problem. Starting from the state yk and the policy parameters θ, it is possible to retrieve the control action uk = π(yk ; θ). Moreover, the subsequent state yk+1 is given by the discrete-time transition function, that is yk+1 = F (yk , uk ). Finally, it is possible to compute the loss function, which depends on the states yk , yk+1 and the policy parameters θ through the control action uk . After the forward pass (solid magenta arrows), it is possible to compute the sensitivities of the loss J to the policy parameters θ in reverse-mode (dashed teal arrows).

Selecting the adjoint variable as solution of Equation (12), we get rid of the term dy dθ , and we end up with the loss gradient formula in Equation (11). Notice that, in general, one has to consider the conjugate transpose of the matrix and the vector in the adjoint equation. The same arguments can be applied to the discrete-time optimal control problem in Equation (7) with uk = π(yk , µ; θ) for k = 0, . . . , Nt . Specifically, via chain rule, one can obtain ∇θ J =

N t −1 X k=0

∂L ∂L ∂π + ∂yk ∂uk ∂yk

∂L ∂π ∂ϕ dyNt dyk + + . dθ ∂uk ∂θ ∂yNt dθ

The tangent equation is now given by  ∂F dy ∂F ∂π dyk ∂F ∂π   k+1 = + + dθ ∂yk ∂uk ∂yk dθ ∂uk ∂θ   dy0 = 0 dθ for k = 0, . . . , Nt − 1. We can now add the tangent equation multiplied by the negative of the adjoint variables λk ∈ RNy for k = 0, . . . , Nt in the loss gradient formula above, yielding ∇θ J =

N t −1 X k=0

∂L ∂L ∂π + ∂yk ∂uk ∂yk

dyk ∂L ∂π + − λ⊤ k+1 dθ ∂uk ∂θ

dyk+1 − dθ

∂F ∂F ∂π + ∂yk ∂uk ∂yk

dyk ∂F ∂π − dθ ∂uk ∂θ

∂ϕ dyNt . + ∂yNt dθ

Rearranging the terms in the sum and exploiting the homogeneous initial condition of the tangent equation above, this is equivalent to ∇θ J =

N t −1 X k=0

⊤ −λ⊤ k + λk+1

∂F ∂π ∂F + ∂yk ∂uk ∂yk

∂L ∂L ∂π + + ∂yk ∂uk ∂yk

dyk ∂L ∂π ∂F ∂π ∂ϕ dyNt ⊤ ⊤ + + λk+1 + − λ Nt . dθ ∂uk ∂θ ∂uk ∂θ ∂yNt dθ

Selecting the adjoint variables as solutions of Equation (??), it is possible to retrieve the loss gradient formula in Equation (8).

Appendix C: Equivalence between adjoint state method and reverse-mode automatic differentiation

While the adjoint method provides the mathematical formulas for sensitivities, automatic differentiation provides the software infrastructure to evaluate them efficiently[8, 67, 109]. Automatic differentiation breaks computer programs into a sequence of simple differentiable operations, and systematically applies the chain rule to retrieve the exact derivatives up to floating-point precision. Specifically, it considers a computational graph where nodes represent variables or intermediate values, while edges represent differentiable operations. Figure 8 shows the one-step computational graph with reference to the discrete-time optimal control problem in Equation 7, assuming no dependence on scenario parameters µ without loss of generality.

21 Differently from forward-mode approaches where derivatives are computed alongside the forward computation, reverse-mode automatic differentiation first stores the variable values through a forward pass, and then propagates gradients backward. This second approach is more efficient when dealing with computational graphs with many inputs and a single output, as typically happens in deep learning, at the cost of a higher memory usage. With reference to the computational graph in Figure 8, the gradient ∇θ J can be computed by summing the sensitivities of the paths going from θ to the loss J, that are θ → uk → yk+1 → J and θ → uk → J. In particular, exploiting the chain rule over those paths, one can retrieve ∇θ J =

dJ ∂J ∂F ∂π ∂J ∂π = + , dθ ∂yk+1 ∂uk ∂θ ∂uk ∂θ

Similarly, the gradient ∇yk J can be computed backpropagating the derivatives in the graph, that is ∇yk J =

dJ ∂J ∂F ∂π ∂J ∂π ∂J ∂F ∂J = + + + . dyk ∂yk+1 ∂uk ∂yk ∂uk ∂yk ∂yk+1 ∂yk ∂yk

∂J ∂L ∂J ∂L Notice that, since J depends on the past through the running cost L, we have ∂u = ∂u and ∂y = ∂y . Therefore, k k k k dJ if we now define the adjoint variables as the sensitivities λk = dy k ⊤

⊤

= ∇yk J ⊤ and

⊤

λk+1 = dydJ = ∇yk+1 J ⊤ = ∂y∂J , we recover the loss gradient in Equation (8) and the adjoint equation in k+1 k+1 Equation (9), thus proving the equivalence between the adjoint state method and reverse-mode automatic differentiation, as well as the equivalence in Equation (15). The same procedure can be repeated when taking into account multi-step computational graphs, achieving the same result. In our setting, reverse-mode automatic differentiation allows us to easily retrieve the gradient of the loss function with respect to the policy parameters. Importantly, differentiable environments are required to propagate the derivatives through the dynamics. To this aim, one may consider numerical physical solvers defined in, e.g., PyTorch [76], JAX [18] or FEniCS-adjoint [68]. Appendix D: Equivalence between adjoint and value-gradient

In this section, we show the equivalence between the closed-loop adjoint variable λ(t, µ), which satisfies Equation (12), and the value gradient ∇y Vπ (y, t, µ), where the (differentiable and smooth enough, for the sake of argument) value function Vπ (y, t, µ) is given by Equation (14). Let us compute the total time derivative of the value function using the chain rule ∂Vπ ∂Vπ ∂Vπ ∂Vπ dVπ = V˙π = + ẏ = + f (y, π(y; θ), µ), dt ∂t ∂y ∂t ∂y dVπ π where ∂V ∂y = dy = ∇y Vπ and

dVπ d = dt dt

"Z T

# L(y(τ ), π(y(τ ); θ), µ)dτ + ϕ(y(T ), µ) = −L(y(t), π(y(t); θ), µ)

thanks to the fundamental theorem of calculus. Combining these equations, we obtain the following PDE −

∂Vπ = L(y, π(y; θ), µ) + ∇y Vπ f (y, π(y; θ), µ), ∂t

(D1)

which corresponds to the Hamilton-Jacobi-Bellman equation equation whenever the optimal policy is taken into account. Differentiating Equation (D1) with respect to the state, we get ∂Vπ ∂L ∂L ∂π ∂f ∂f ∂π d ∂Vπ = −∇y = + + ∇2y Vπ f + ∇y Vπ + , (D2) − dy ∂t ∂t ∂y ∂u ∂y ∂y ∂u ∂y where ∇2y Vπ is the Hessian matrix of the value function. Notice that, since d ∇y Vπ = ∇y V˙π = ∇y dt

∂Vπ ∂Vπ + f ∂t ∂y

= ∇y

∂Vπ + ∇2y Vπ f, ∂t

22 we can rewrite Equation (D2) as follows d ∇y Vπ = −∇y Vπ dt

∂f ∂f ∂π + ∂y ∂u ∂y

−

∂L ∂L ∂π + ∂y ∂u ∂y

that is equivalent to Equation (12). Due to the existence and uniqueness of the solutions of these equations, we conclude that λ = ∇y Vπ⊤ . Notably, similar statement can be proven when taking into account the open-loop adjoint RT variable [80] and the value function Vu (y, t, µ) = t L(y(τ ), u(τ ), µ)dτ + ϕ(y(T ), µ) with no explicit dependence on the parametric policy [19, 32, 86]. The equivalence between the adjoint variables λk (µ) = λ(tk , µ) and the sensitivity of the value function in Equation (16) to the state yk holds also after time discretization. Indeed, differentiating the Bellman equation for a generic policy π with respect to the state yk , one can directly retrieve Equation (9) requiring that λk (µ) = ∇yk Vπ (yk , µ)⊤ . Moreover, as shown in Appendix C, the discrete adjoint variables are also equal to the loss gradient λk = ∇yk J ⊤ . Indeed the loss terms at t0 , . . . , tk−1 vanish when differentiating with respect to ↷k due to time causality.

Record · ID 381723 · SHA-256 e099030b36a3efa1
Retrieved via Conceptio — every document is proof-bundled with source, license, and retrieval metadata.