Physics-enhanced reinforcement learning for real-time optimal control of dynamical systems
By Matteo Tomasetto, Nicolò Botteghi, Gabriele Bruni, Andrea Manzoni
"PEARL fuses RL with adjoint sensitivity analysis to compute policy gradients efficiently, drastically reducing interactions needed for optimal control of high-dimensional dynamical systems."
Abstract
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 high-dimensional 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.
Technical Analysis & Implementation
Technical Overview§
PEARL (Physics-EnhAnced Reinforcement Learning) addresses the sample inefficiency of RL for continuous control by leveraging differentiable environment dynamics. It combines short-horizon policy gradients via automatic differentiation with long-horizon return approximations using adjoint-based sensitivity analysis.
Core Methodology§
Consider a dynamical system $\dot{\mathbf{x}} = f(\mathbf{x}, \mathbf{u})$ with control policy $\mathbf{u} = \pi_\theta(\mathbf{x})$. The objective is to maximize expected return $J(\theta) = \mathbb{E}\left[\int_0^T r(\mathbf{x}, \mathbf{u}) dt\right]$. PEARL splits the horizon:
- Short horizon $\tau$: Use automatic differentiation to compute $\nabla_\theta J_\text{short}$ by differentiating through $\tau$ steps of the differentiable dynamics.
- Long horizon: Approximate the tail return $V_\psi(\mathbf{x}(t+\tau))$ with a neural network critic, and compute its sensitivity via the adjoint method:
$$\nabla_\theta V_\psi(\mathbf{x}(t_\tau)) = \frac{\partial V_\psi}{\partial \mathbf{x}} \frac{\partial \mathbf{x}}{\partial \theta}$$ where $\frac{\partial \mathbf{x}}{\partial \theta}$ is obtained by solving the adjoint ODE backward in time.
The final gradient is $\nabla_\theta J \approx \nabla_\theta J_\text{short} + \nabla_\theta V_\psi$.
Implementation Sketch§
import torch
from torchdiffeq import odeint_adjoint
def pearl_update(env, policy, critic, horizon, gamma):
state = env.reset()
traj_states, traj_actions, traj_rewards = [], [], []
# Short horizon rollout
for t in range(horizon):
action = policy(state)
next_state = env.step(action) # differentiable
reward = env.reward(state, action)
traj_states.append(state)
traj_actions.append(action)
traj_rewards.append(reward)
state = next_state
# Short horizon gradient via autograd
short_loss = -sum(traj_rewards)
short_grad = torch.autograd.grad(short_loss, policy.parameters())
# Long horizon via adjoint
tail_value = critic(state)
# Adjoint sensitivity of tail_value w.r.t. policy parameters
# Use ODE adjoint to compute d(tail_value)/d(theta)
def dynamics(t, x):
u = policy(x)
return env.f(x, u)
# Solve adjoint ODE backward
adjoint_grad = odeint_adjoint(
dynamics, (tail_value,), torch.tensor([horizon, 0]),
policy.parameters()
)
# Combine gradients
total_grad = [s + a for s, a in zip(short_grad, adjoint_grad)]
# Apply optimizer...Key Advantages§
1. Sample efficiency: Short horizons reduce exploration needs; adjoint provides accurate long-term credit assignment without full rollouts. 2. Scaling: Handles high-dimensional states/actions because gradients are computed analytically rather than estimated via sampling. 3. Generalization: The physics-informed critic learns across parametric variations (e.g., different flow conditions).
Results§
PEARL was tested on two parametric navigation problems in unsteady flows (e.g., a glider and a swimmer). It achieved 10x fewer environment interactions compared to PPO and SAC, while matching or exceeding final performance. The policy also generalized to unseen flow parameters without retraining.
Limitations§
- Requires differentiable environment simulator, which limits applicability to non-differentiable simulations or real systems.
- Adjoint method adds computational overhead for solving ODEs backward.
In summary, PEARL provides a principled integration of RL and optimal control for differentiable dynamical systems, significantly advancing sample efficiency and scalability.