SINDy-RL for interpretable and efficient model-based reinforcement learning

Domain: physics
Published: 2025-11-28
Editorial outcome: minor_revision

Abstract
Deep reinforcement learning (DRL) has shown significant promise for uncovering sophisticated control policies that interact in complex environments, such as stabilizing a tokamak fusion reactor or minimizing the drag force on an object in a fluid flow. However, DRL requires an abundance of training examples and may become prohibitively expensive for many applications. In addition, the reliance on deep neural networks often results in an uninterpretable, black-box policy that may be too computationally expensive to use with certain embedded systems. Recent advances in sparse dictionary learning, such as the sparse identification of nonlinear dynamics (SINDy), have shown promise for creating efficient and interpretable data-driven models in the low-data regime. In this work, we introduce SINDy-RL, a unifying framework for combining SINDy and DRL to create efficient, interpretable, and trustworthy representations of the dynamics model, reward function, and control policy. We demonstrate the effectiveness of our approaches on benchmark control environments and flow control problems, including gust mitigation on a 3D NACA 0012 airfoil at Re = 1000. SINDy-RL achieves comparable performance to modern DRL algorithms using significantly fewer interactions in the environment and results in an interpretable control policy orders of magnitude smaller than a DRL policy.

Introduction
Much of the success of modern technology can be attributed to our ability to control dynamical systems: designing safe biomedical implants for homeostatic regulation, gimbaling rocket boosters for reusable launch vehicles, operating power plants and power grids, industrial manufacturing, among many other examples. Recently, advances in machine learning and optimization have rapidly accelerated our ability to tackle complicated data-driven tasks—particularly in the fields of computer vision1 and natural language processing2.

Reinforcement learning (RL) is at the intersection of both machine learning and optimal control, and the core ideas of RL date back to the infancy of both fields. An RL agent iteratively improves its control policy by interacting with an environment and receiving feedback about its performance on a task through a reward. Deep reinforcement learning (DRL) has shown particular promise for uncovering control policies in complex, high-dimensional spaces3–10. DRL has been used to achieve super-human performance in games11–15 and drone racing16, to control the plasma dynamics in a tokamak fusion reactor17, to discover novel drugs18, and for many applications in fluid mechanics19–29. However, these methods rely on neural networks and typically suffer from three major drawbacks: (1) they are infeasible to train for many applications because they require millions—or even billions15—of interactions with the environment; (2) they are challenging to deploy in resource-constrained environments (such as embedded devices and micro-robotic systems) due to the size of the networks and need for specialized software; and (3) they are “black-box” models that lack interpretability, making them untrustworthy to operate in safety-critical systems or high-consequence environments. In this work we improve the sample efficiency of reinforcement learning algorithms, even for high-dimensional problems, by leveraging sparse dictionary learning. Specifically, we build small, interpretable surrogate models for the environment dynamics, reward, and policy.

There has been significant research into reducing the amount of experience needed to train RL policies, such as offline RL30, experience replay methods31–33, transfer learning34, and meta-learning35,36. Training in a low-fidelity representation of the environment is perhaps the most common way to reduce the number of interactions in a full-order environment. However, there are many cases where an analytic reduced-order model does not exist and the dynamics must be learned from data to create a surrogate representation of the environment. Dyna-style model-based reinforcement learning (MBRL) algorithms iteratively switch between learning and improving a surrogate model of the environment and training model-free policies inside the surrogate environment by generating “imaginary” experience37. Deep MBRL algorithms have shown significant promise for reducing sample complexity on benchmark environments38 by simultaneously training neural network models of the environment. Although neural network models have recently become popular and are gaining wide adoption over traditional modeling methods, they are still overparameterized, data-inefficient, and uninterpretable.

In contrast, sparse dictionary learning provides an efficient and interpretable alternative to learn models from data, as in the sparse identification of nonlinear dynamics (SINDy)39. Sparse dictionary learning is a type of symbolic regression that learns a representation of a function as a sparse linear combination of pre-chosen candidate dictionary (or “library”) functions. A sparse, symbolic model lends itself naturally to interpretation and analysis—especially for physical systems where the dictionary terms have physical meaning. Importantly, SINDy has been extended to systems with control and used to design model predictive control (MPC) laws40,41. SINDy methods are incredibly efficient—both for model creation and deployment—making them promising for both online learning and resource-constrained control. In particular, SINDy has been used for online simultaneous dynamics discovery and optimal control42. Recent work used a single SINDy model to accelerate DRL and demonstrated its use on simple DRL benchmarks43. In this work, we generalize this DRL framework to include ensembles of dictionary models of both the dynamics and reward to accelerate learning in the low-data limit and quantify uncertainty.

Our contributions
In this work, we develop techniques at the intersection of sparse dictionary learning and DRL for creating trustworthy, interpretable, efficient, and generalizable models that operate in the low-data limit. Building on recent advances in ensemble dictionary learning44, we first introduce a Dyna-style MBRL algorithm that fits an ensemble of SINDy models to approximate an environment’s dynamics and uses modern model-free reinforcement learning to train agents in the surrogate environments. In systems where the reward is difficult to measure directly from the observed state (e.g. limited sensor information for flow control on an aircraft), we augment the Dyna-style algorithm by learning an ensemble of sparse dictionary models to form a surrogate reward function. Finally, after training a DRL policy, we use an ensemble of dictionary models to learn a lightweight, symbolic policy, which can be readily transferred to an embedded system. Figure 1 provides a schematic of the SINDy-RL framework. We evaluate our methods on benchmark environments for continuous control from mechanical systems using the dm_control45 and gymnasium suites46 as well as fluid systems from HydroGym47 and HydroGym-GPU48. We demonstrate that our methods can:

• Improves sample efficiency by orders of magnitude for training a control policy by leveraging surrogate experience in an E-SINDy model of the environment.
• Leverage the efficiency of the surrogate models to accelerate expensive hyperparameter tuning.
• Learn a surrogate reward when the reward is not directly measurable from observations.
• Reduce the complexity of a neural network policy by learning a sparse, symbolic surrogate policy, with comparable performance and smoother control.
• Quantify the uncertainty of models and provide insight into the quality of the learned models.

By leveraging sparse dictionary learning in combination with deep reinforcement learning, it can become feasible to rapidly train control policies in expensive, data-constrained environments while simultaneously obtaining interpretable representations of the dynamics, reward, and control policy.

Results
Reinforcement learning (RL) comprises a family of methods where an agent learns a policy, π, to perform a task through repeated interaction with an environment, E. Explicitly, at each environment state x_n, the agent samples an action u_n ~ π(x_n) and executes it in the environment, producing a new state x_{n+1} and reward r_n — an indication of how well the agent performed at that time. We define the value function to be the expected future return for taking actions from the policy, V^π(x) = E [sum_{k=0}^∞ γ^k r_k | x_0 = x] where 0 < γ ≤ 1 is the discount factor. RL methods seek a policy that maximizes this quantity, which can be a very challenging optimization problem, especially in the case of high-dimensional state-spaces, continuous action spaces, and nonlinear dynamics. Deep reinforcement learning (DRL) has made significant progress in addressing these problems by parameterizing functions as deep neural networks (DNNs), such as the policy π(x) ≈ π_φ(x), and training on collected experiences.

In contrast to neural networks, dictionary learning is an efficient form of symbolic regression that models a function y(x) = f(x) as a linear combination of d dictionary (or “library”) functions (such as polynomials, sines, cosines, etc.): Θ(x) = (θ_1(x), θ_2(x), … θ_d(x)). To learn a model from N data samples of inputs X = [x_1, x_2, ..., x_N]^T ∈ R^{N×m}, and their associated labels Y = [y_1, y_2, ..., y_N]^T ∈ R^{N×n}, we evaluate the dictionary at the data, X:

Θ(X) = [θ_1(X), θ_2(X), ..., θ_d(X)] ∈ R^{N×d}

to form the linear model Y = Θ(X)Ξ, where Ξ ∈ R^{d×n} are the coefficients to be fit. Sparse dictionary learning assumes that the desired function can be well approximated by a small subset of terms in the library, i.e. Ξ is a sparse matrix. For dynamics discovery, such as SINDy39,49 (where y = dx/dt), a sparse library is physically motivated by the observation that the governing equations for most physical systems have relatively few terms. To achieve this parsimony, a sparse optimization problem is formulated:

Ξ* = argmin_Ξ ||Y - Θ(X)Ξ||_F + R(Ξ) (1)

where ||⋅||_F is the Frobenius norm and R(Ξ) is a sparsity-promoting regularization.

Ensemble-SINDy (E-SINDy)44 introduced a way to select a model from an ensemble of SINDy models and is more robust to noise than SINDy, particularly in the low-data limit. E-SINDy can be generalized to arbitrary dictionary models by considering ensembles: Y^{(k)} = Θ^{(k)}(X^{(k)}) Ξ^{(k)}, for k = 1…N. It is important to note that there have been many proposed variants of SINDy to make the algorithm more robust50–55. These methods are generally compatible—if not synergetic—with E-SINDy; we therefore only present the simplest formulation in this work, though a practitioner may seek to amend our framework with a more specialized variant to best suit their purpose.

In this work, we introduce SINDy-RL: a unifying set of three techniques for applying sparse dictionary learning to DRL control tasks. As shown in Fig. 1, these techniques can be used either together or separately (depending on the context) to approximate the environment dynamics, reward function, and the final learned neural network policy using ensemble sparse dictionary learning. We frequently use the “hat” notation to indicate a surrogate approximation of a function, i.e. f̂(x) ≈ f(x). In the following, we introduce each of these new techniques and demonstrate their capability to quickly learn effective controllers for the set of five control problems depicted in Fig. 2. In particular, we demonstrate our ability to greatly reduce the amount of interaction needed with an expensive simulation environment while simultaneously extracting lightweight and interpretable representations of the dynamics, reward, and control policy.

SINDy-RL: Learning surrogate dynamics
We propose a Dyna-style MBRL algorithm where we iteratively improve a surrogate dynamics model while learning a control policy. First, we collect offline data samples from the full-order environment by deploying a default policy (such as random input, an untrained neural network, Schroeder sweep56, etc.). We use the collected data to fit an ensemble of SINDy models to initialize a surrogate environment. Next, we iteratively train a policy for a fixed number of policy updates using the surrogate environment with a policy optimization algorithm, such as proximal policy optimization (PPO)57. By forcing the agent to only interact in the surrogate environment, we offload the majority of the expensive sample collection to the lightweight E-SINDy surrogate. Finally, we deploy the trained policy to the full-order environment, collect data for evaluation, and use the newly collected data to update the E-SINDy models. We repeat this process by iteratively updating and improving both the E-SINDy model and policy. In this work, we specify a fixed number of policy iterations before evaluating the agent, though a practitioner may consider an adaptive number of updates based on training metrics in the surrogate environment. A full summary of this process can be found in Algorithm 1 under Methods.

Either a continuous or discrete model can be fit; however, we learn discrete-time models where next-step updates are predicted explicitly, rather than integrating a continuous-time model forward in time, for ease of deployment in the environment. It is common for nonlinear dynamics models—especially learned models—to grow unbounded over long time horizons unless stability guarantees are enforced54. To accommodate this limitation, we bound the state space and reset the surrogate environment during training if a trajectory exits the bounding box.

SINDy-RL: Learning surrogate rewards
For Dyna-style MBRL, it is assumed that the reward function can be directly evaluated from observations of the environment. However, there are cases in which the reward function cannot be readily evaluated because the system may only be partially observable due to missing sensor information—as is the case in many fluids systems. There has been significant work on creating proxy rewards for controls tasks using “reward shaping”58 and learning an objective function through inverse reinforcement learning59. We propose learning a proxy reward, R̂(x_{k+1}, u_k), with supervised sparse dictionary learning when there are offline evaluations of the reward available: r_k = R(x_{k+1}, u_k). Our implementation uses sparse ensemble dictionary learning under the assumption that there is a sparse deterministic relationship between the reward function and observations from the environment. We incorporate this into Algorithm 1 by learning the reward function alongside the dynamics to create the surrogate environment.

SINDy-RL: Learning surrogate policies
Taking inspiration from behavior cloning algorithms in imitation learning60, we fit a sparse dictionary model approximation of the final neural network learned policy, π_φ. While imitation learning traditionally attempts to mimic the policy of an expert actor with a neural network student, we instead use the learned neural network as our expert and the dictionary model as our student, resulting in a lightweight, symbolic approximation. A dictionary model is not as expressive as a neural network; however, there has been previous work indicating that even limiting to purely linear policies can be a competitive alternative to deep policy networks61,62, showing that even complicated control tasks may have a simple controller. Likewise, there has been recent investigation into approximating neural networks with polynomials through Taylor expansion63, which has shown to provide sufficiently robust approximations. These reduced representations can be orders of magnitude smaller than a fully-connected neural network and they can be efficiently implemented and deployed to resource-constrained environments, such as embedded systems.

Many model-free algorithms assume a stochastic policy for optimally interacting with a stochastic environment and encouraging exploration; we proceed with a sparse ensemble fit of the expectation: π̂(x) ≈ E[π_φ(x)]. Because there is no temporal dependence on π_φ, we can assemble our data and label pairs (x, E[π_φ(x)]) by evaluating π_φ for any x.

Using SINDy to accelerate DRL training
By quickly developing a SINDy surrogate model using limited data, it is possible to offload the expensive interaction with a full-order model. In Fig. 2, we summarize how SINDy-RL reduces the number of interactions with our five evaluation environments compared to the popular model-free proximal policy optimization (PPO) algorithm57; in each case, SINDy-RL improves the sample efficiency by 10–100×. In Fig. 3, we also provide a more in-depth comparison our proposed SINDy-RL method from Algorithm 1 with a quadratic dynamics library to four different baseline experiments: (1) PPO, (2) Algorithm 1 with a linear SINDy library, (3) Algorithm 1 where SINDy models are replaced with neural network dynamics models (i.e. “Dyna-NN”), and (4) RLlib’s implementation of Model-Based Meta-Policy Optimization (MB-MPO)64. As shown in Fig. 3, SINDy-RL with a quadratic library learns a control policy with 100× fewer interactions in the full-order environment compared to model-free RL and greatly outperforms other model-based approaches. The linear models are incapable of approximating the global dynamics since there are multiple fixed points. In contrast, while neural network models should be more expressive, it appears that both the MB-MPO and Dyna-NN baselines fail to accurately capture the dynamics sufficient to achieve the task. For both comparisons, it is suspected that the dynamics may be captured with significantly larger data collections—as shown in previous work comparing SINDy with neural network models for MPC40. An investigation of the hyperparameters from our method can be found in Supplementary Section 4A.

The increased sample efficiency is particularly important for applications interacting with physical systems that require humans in-the-loop to monitor and reset environments or computationally expensive simulations—such as CFD solvers—that require significant time and resources to run. The benefits of increasing the sample efficiency can be seen using the fluid flow control environments. The 3D Airfoil environment (Fig. 4) uses the lattice Boltzmann method with over 72M cells, requiring 75GB of VRAM across four A100 GPUs. Even when using the GPUs, a single step in the full-order model takes on the order of 45s to update the environment. In comparison, the learned E-SINDy polynomials for the 3D Airfoil require less than 10kB of RAM on CPU, and a single step takes approximately 1-4 milliseconds using NumPy65—i.e., SINDy-RL experience collection is 10^4× faster than the full-order CFD models. The SINDy-RL agent was 14.47× more sample efficient for the 3D Airfoil environment; with only 25 dynamics updates (i.e. iterations of Step 2 from Algorithm 1), the median SINDy-RL agent reached the final median baseline agent performance. Physically, this was the difference between 14 hours of training using the SINDy-RL environment and 185 hours using the full-order environment, greatly reducing the total amount of time needed to learn capable policies. A comparison of clock times for different aspects of training the HydroGym environments can be found in Supplementary Table 4.

It is important to note the large variance in performance among the 20 trained SINDy-RL Cylinder agents depicted in the shaded region of Fig. 2. After interacting with the full-order Cylinder environment 13,000 times, the top 50% performing SINDy-RL agents were able to surpass the best baseline performance—some achieving 11% drag reduction, comparable with optimized solutions from previous literature66,67. However, the bottom 25% of agents performed poorly. Surrogate dynamics models can quickly diverge—especially early in training—which can provide misleading rewards. When these effects are further coupled with unfavorable policy network initializations, agent performance can severely degrade. The overall success of the majority of the agents indicates that methods like population-based training may be extremely effective at pruning these bad combinations of surrogate dynamics, rewards, and policies to quickly find exceptional policies using very few interactions in the full-order environment.

Accelerating hyperparameter tuning
The success of DRL to find effective control policies can often depend on the neural network initialization and choice of hyperparameters; otherwise training might lead to a suboptimal policy. To address this, it is common to use different random initializations for the network and hyperparameter tuning, and software packages have been created to facilitate these searches, such as Ray Tune68. While these tuning algorithms can be effective for uncovering sophisticated policies, they generally rely on parallelizing training and significant compute usage. Previous benchmarks69,70 have shown that the Swimmer-v4 environment is particularly challenging to train with PPO and it has been suggested71 that learned policies for Swimmer-v4 are especially sensitive to certain hyperparameters and tuning can achieve superior performance. We pursue this idea by comparing population-based training (PBT)72 to improve policies for agents interacting in the full-order Swimmer-v4 environment and a SINDy-RL environment with dictionary surrogates for the dynamics and reward. Both experiments used a population of 20 policies and periodically evaluated the policies using the rollouts from the full-order model. As shown in Fig. 2, the SINDy-RL training is able to achieve nearly 30% better maximal performance than the baseline using 50× fewer samples from the full-order environment. This indicates that SINDy-RL can provide a convenient way to accelerate hyperparameter tuning in more sophisticated environments. Additional training details for PBT can be found in Supplementary Section 4B.

Generalizing to unseen dynamics
It has been well-documented that while DNNs are excellent at performing on unseen data sampled from the same data distribution used for training, i.e. interpolation, they often fail at extrapolation to data beyond the convex hull of the training set. In Fig. 5, we investigate the ability of top-performing Pinball agents trained at Re = 100 for 20-second segments to extrapolate to unseen dynamics at Re = 150, 250, and 350 for a 100-second evaluation. It is well-known that the dynamics of the Pinball system undergo a bifurcation at around Re = 115 where the dynamics transition from being quasi-periodic to chaotic73; thus, the dynamics are fundamentally different at these Re values, which can be qualitatively seen—especially at Re = 350—in the initial conditions shown in Fig. 5(b). Despite these fundamental differences, the SINDy-RL policy is able to reasonably extrapolate. While performance does degrade compared to Re = 100, the agent is able to obtain a much better cumulative return of reward over the long trajectory compared to the baseline neural network counterpart.

Interpretable surrogate dynamics
Despite none of the evaluation environments having a polynomial representation of the dynamics, we find that surrogate models with polynomial dynamics often provide a sufficient representation of the environment to learn a control policy through repeated interaction. A thorough investigation of the learned dynamics from each environment can be found in Supplementary Section 5.

We find that for most systems, it is imperative to periodically refit the dynamics while the agent explores control strategies. For example, with the swing-up dynamics, the agent has no information about the goal state at the unstable equilibrium early in training. In Supplementary Fig. 13, we show that the first learned dynamics model provides a reasonable representation of the phase portrait, but the learned unstable fixed point is offset from the ground-truth position. Without refitting the dynamics, the agent would learn to drive the system to the wrong point and never stabilize the true system. However, by deploying the learned policy on the full-order system and gathering new data, the quality of the dynamics model improves and the agent is able to complete the task both using the surrogate and full-order dynamics. Finally, unlike a neural network representation of the dynamics, our approach provides a symbolic representation of the learned dynamics. In Supplementary Section 5, we provide the learned coefficients Ξ and show that the swing-up dynamics model is well-represented by an Eulerian integration scheme: x_{k+1} = x_k + Δt f(x_k).

It is important to note a limitation of using dictionary dynamics; for large observation spaces, the size of a polynomial dictionary increases combinatorially. Therefore for large spaces, dimensionality reduction is critical. We demonstrate the viability of this approach for both the Pinball and 3D Airfoil environments by linearly projecting the 70- and 318-dimensional observations onto the dominant SVD modes of the system and found this to be sufficient for purposes of training. For more complicated dynamics, autoencoders have been found to be useful for finding nonlinear projections74.

Interpretable reward functions
In the setting where exact rewards are not analytically expressible from observations (e.g., all of our evaluation environments except swing-up), we periodically fit the surrogate dictionary reward r̂ = R̂(x, u) alongside and independent of the SINDy dynamics. For Swimmer-v4, the rewards in the full-order environment are given by the agent’s body-centered horizontal velocity. However, this quantity is not provided by the observation space. It has been well-documented that—even with access to the exact rewards—solving the Swimmer-v4 task is a considerable challenge because of the placement of the velocity sensors; researchers have even created their own modified versions of the environment when performing benchmarks38. Our method identified that a sparse reward of r̂ ≈ v_x—the horizontal velocity of the leading segment (not the body)—was a suitable surrogate and remained stable throughout the entirety of the population-based training.

In contrast to the Swimmer-v4 environment, the Cylinder, Pinball, and 3D Airfoil environments are governed by nonlinear PDEs, where the rewards are scalar measurements evolving with the dynamics on the domain and are not analytically expressible in terms of the provided observations. The lack of available information makes modeling the environment and the reward a very challenging task. Just like the dynamics, Fig. 6(i) demonstrates the necessity of periodically updating the surrogate models. At the beginning of training, the learned reward is actually anti-correlated with the full-order reward; however, the learned reward becomes well-correlated after subsequent updates. Despite the inability to learn the exact reward in the fluid environments, the surrogate reward provides a sufficient enough learning signal for training a comparable control policy when having access to the exact rewards. We also highlight the interpretability of our method in Fig. 6(ii); the learned reward function is a quadratic polynomial of the observations, and we can analyze the explicit influence of the control. The bowl-shaped function (given by the quadratic) gradually shifts in response to the increasingly positive control. Furthermore, it is clear that the drag is minimized (i.e., reward is maximized) for large values of |u|, indicating that an optimal control strategy would apply maximal control input to stay in the bowl’s minimum. For the Pinball and 3D Airfoil environments, the sparse sensors in the flow are not sufficient to describe the net forces acting on surfaces; indeed, the discovered reward functions are dominated primarily by terms coupling the actuation (local information) to the sensors. The learned rewards and a detailed investigation for each environment can be found in Supplementary Section 6.

Interpretable surrogate policies
While DNN policies can find solutions to complicated control problems, they are typically over-parameterized, black-box models that lack interpretability. There have been several approaches to combine symbolic regression and DNNs to improve interpretability55,75–79, but they tend to focus on discovering the dynamics rather than a symbolic form for a controller. In this work, we discover lightweight models of the control policies with sparse dictionary learning using behavior cloning; for each environment, we leverage the E-SINDy dynamics and the neural network policy obtained from Algorithm 1 to collect data for fitting the dictionary model.

Figure 5 compares the performance of the SINDy-RL neural network policy to the sparse dictionary policy distilled from it in the Pinball environment. As shown in Fig. 5e, the dictionary approach is able to consistently stabilize the wake across various values of Re compared to its neural network counterpart. This is shown quantitatively in Fig. 5f where the dictionary policy consistently receives higher cumulative returns during long evaluations of the policy. The final dictionary control policy is represented as a quadratic polynomial library containing 66 elements, nearly two orders of magnitude fewer parameters than the corresponding neural network policy with over 5000 parameters. This is also considered a relatively small neural network model; for deeper or wider networks, one can reasonably expect even larger gains. In Supplementary Section 7B, we provide further evidence that a dictionary policy can produce smoother and more consistent control inputs—outperforming the neural network when it has overfit to a particular trajectory in the swing-up environment.

However, the Cylinder environment provides an example where this method may struggle; the surrogate policy only reduces drag by about 3.7% compared to the neural network policy’s 8.7% reduction. The original neural network agent learned a bang-bang control policy; this is very difficult to approximate with a smooth polynomial due to the bounded derivatives. One of the key challenges that supervised imitation learning faces (e.g., behavior cloning) is the issue of compounding errors during policy deployment. Because our policy approximation is a type of behavior cloning, our method inherits this challenge; the approximate policy slowly drifts away from the state-action pairs it was trained on and ultimately performs suboptimally. A detailed investigation of the learned dictionary policies for each environment (including a comparison of different data sampling strategies used for the behavior cloning) can be found in Supplementary Section 7.

Uncertainty quantification
Treating a set of dictionary coefficients, Ξ, as random variables, an ensemble acts as an empirical approximation for the distribution of likely Ξ values. From this, one can derive an efficient framework for analytically approximating a model’s pointwise variance. Explicitly, for Ξ = [ξ_1, … ξ_n], and letting Cov(ξ_i) denote the sample covariance matrix of ξ_i computed from the ensemble, the uncertainty of y given x is:

Tr(Var_Ξ [y(x)|x]) = sum_{i=1}^n Θ(x) Cov(ξ_i) Θ(x)^T (2)

In Fig. 7, we visualize the evolution of the uncertainty landscape for the swing-up dynamics when training the dynamics and policy with SINDy-RL. At the beginning of training, the dynamics are confined to a narrow band of low-variance states where data has been previously collected. During training, the variance landscape changes in response to new, on-policy data. By the time the agent has learned to solve the swing-up task, the learned dynamics have become more certain and the agent takes trajectories over regions that have less variance.

In this work, we only use variance to inspect the learned dictionary functions, although there are further opportunities to exploit the variance. There is a trade-off in sample efficiency between refining the learned dynamics and improving the policy early in training. With a poor representation of the surrogate dynamics, we risk overfitting to the surrogate and discovering policies that do not generalize well to the real environment. Analogous to curiosity-driven learning80, the estimated uncertainty of the dynamics and reward models may strategically guide the exploration of the environment and rapidly improve them—making it less necessary to query the environment later. Likewise, we may encourage “risk-averse” agents by penalizing regions with high uncertainty—encouraging agents to take trajectories where the dynamics are trusted and agree with the full-order model.

Methods
Learningsurrogatedynamics
Algorithm1givesanoverviewoftheDyna-styleMBRLalgorithm.EachdynamicsmodelintheensembleisfitusingthediscreteformofSINDywithcontrol(SINDy-C)40andSTRidge49.HyperparametersusedtolearnthedynamicsforeachenvironmentcanbefoundinSupplementarySection5.

Algorithm1.(Dyna-styleSINDy-RL)
Input:E ▹full-orderenvironment
N_off,N_collect ▹#ofoff-andon-linepolicysteps
n_batch ▹#ofpolicyiters/SINDyupdate
Θ ▹dictionaryfunctions
A ▹policyoptimizationalgorithm
π_0 ▹defaultpolicy
Step1:InitializeSurrogateEnvironment
D_off =CollectData(E,π_0,N_off)
D=InitializeDatastore(D_off)
Ξ=ESINDy(D,Θ)
Ê =Surrogate(Ξ)
Step2:ModelImprovement
π=InitializePolicy()
whilenotdone:do
π=A(Ê,π,n_batch)
D_on =CollectData(E,π,N_collect)
D=UpdateStore(D,D_on)
Ξ=ESINDy(D,Θ)
Ê =Surrogate(Ξ)
endwhile
Output:Optimizedpolicy,π;SINDyenvironment,Ê.

Learningsurrogatepolicies
Tobuildourdataset,wesamplepointsfromtrajectoriesoftheagentinteractingintheenvironment,ensuringthatthestatesaredirectlyrelevanttothecontrolstask.Toavoidthecostofcollectingmoredatafromthefull-orderenvironment,E,weproposesamplingnewtrajectoriesfromalearnedensembleofdynamicsmodels.Explicitly,wesampleN_τ trajectoriesbypropagatingtheE-SINDymodel

τ_i = { (x_k, u_k), x_{k+1}, r_k }^T_{k=1, x_{k+1} = f̂(x_k, u_k)}

andusethecollecteddata(x,u)fromeachτ tofitthedynamicsmodel.DuringtheDRLtraining,theneuralnetworkpolicymayoverfittospecificregionsofthespaceandbiasourdatacollection;thus,wedrawinspirationfromtubeMPC88,whereanMPCcontrollerattemptstostaywithinsomeboundedregionofanominaltrajectory.Insteadofonlysamplingpointsfromthesurrogatetrajectories,wecansampleplausiblepointsinaneighborhoodofthetrajectoriestocreateamoreaccurateapproximationwhilestillavoidingregionsofthespacethatcannolongerbetrusted.

Evaluationenvironments
WeevaluateeachourSINDy-RLtechniquesonfiveenvironmentsdepictedinFig.2:(1)dm_controlswing-up45balancesapoleonacartintheunstableuprightpositionstartingfromthestabledownpositionatrest,(2)gymnasiumSwimmer-v446controls3-segmentrobottotravelasfaraspossibleinthehorizontaldirection(alongthex-axis)forafixedtime,(3)HydroGymCylinder47reducesthedragforce,C_D,ofarotatingcylinderinanunsteadyfluidflowatRe=100withmeasurementsoftheliftforce,C_L,(4)HydroGymPinballreducesthenetdragforceC_{D1}+C_{D2}+C_{D3}onthesystemofcylinderstostabilizethequasi-periodicflowatRe=100,and(5)HydroGym-GPU3DAirfoil48mitigatestheeffectofalargeincominggustinanunsteadyflowatRe=1000byminimizing|C_L-C_L^ref|+0.25|C_D-C_D^ref|,i.e.deviationsofC_L andC_D fromreferencevalues.ForPinballand3DAirfoil,stateinformationisobtainedfromameshofvelocityprobesintheflow,forminga70-and318-dimensionalmeasurementspace,respectively.Thisdataisprojectedontotheleadingsingularvaluedecompositionmodes(SVD)andusesthecoefficients,a_i=v_i^Tx,toforma10-and2-dimensionalobservationspacefortherespectivesystems.DetailsforeachenvironmentcanbefoundinSupplementarySection3.WhereasTheswing-uprewardisanalyticallyexpressibleintermsoftheobservedvariables,allotherrewardsarecomputedfromenvironmentinformationthatcannotretrievedanalyticallyfromtheobservation.Wetreattherewardsasfunctionsoftheobservationsandactionsbyapproximatingthemwithanensembleofsparsemodels.

Benchmarking
PPOisusedasthepolicyoptimizationalgorithm,A,forallDyna-styleexperimentswithidenticalhyperparametersbetweencomparisons.Foreveryalgorithmandexperiment,multipleindependentinstantiationsareruntoprovideadistributionofperformanceacrossrandomseeds.Twentyinstantiationswereusedforallenvironments,except3DAirfoilwherewearelimitedtofourinstantiationsduetothecomputationaldemandoftheenvironment.Specificdetailsaboutthetraining,hyperparameters,andmorecanbefoundinSupplementarySection4.

Peer review rounds
[{"round": 1, "reviewer_comments": {"reviewer_1": "Now the draft looks better, only one thing which I think it is better to clarify it for the readers is this part:\n``In this work, we simply set a fixed number of dynamics updates, but in practice one would either stop training once the\nperformance has plateaued or reached some task-specific criteria\ndetermined by the practitioner. We have clarified this in the draft.''", "reviewer_2": "The new draft of the manuscript is great! I really like the inclusion of the fluidic pinball with 70 dimensions! It drives the\npurpose of the paper very well in a more real-world environment.\nI had three minor comments on: including clock time information, longer captions for each figure to make it informative, and\ntitles for missing figures. All of these have been addressed. The longer captions on the figures were very helpful.\nAll my comments and concerns have been addressed!", "reviewer_3": "I believe that the argument if it works in atoy problem it is useful for solving complex problems is not valid any more in 2025.\nI thank the authors for considering my comments. However, the core of my arguments has not been answered. The authors\nmake major claims about their method when it is only performing at toy problems.\nToday RL has been used to model complex 3D flows (cylinders, channels, boundary layers)\nso I fail to understand the merit of a methodology that is only shown to work in simple 2D flows and toy environments.\nThe authors present no evidence that the methodology is potent enough to solve complex, non-linear problems.\nThe authors show again flows at Re around 100. The 2D flow past a single cylinder at Re up to 500 have the same\ndynamics as Re=100 - one can verify this by projecting the dynamics in a two-three dimensional latent space. And the\n\"pinball\" flow with the arrangement of three cylinders does not add anything in terms of complexity or non-linearity as the\ndynamics amount to a linear superposition of linear dynamics of a single cylinder. Finally the additional paper that is brought\nforward by Wolf et al. also shows simple dynamics including for a example a lid driven cavity flow in the laminar regime.\nI regret that I maintain my objections and cannot recommend publication of this paper."}, "author_response": "Response to referee comments on “SINDy-RL: Interpretable and Efficient\nModel-Based Reinforcement Learning”\nReferee 1\nNow the draft looks better, only one thing which I think it is better to clarify it for the\nreaders is this part:\n“In this work, we simply set a fixed number of dynamics updates, but in practice one\nwould either stop training once the performance has plateaued or reached some task-specific\ncriteria determined by the practitioner. We have clarified this in the draft.”\nAuthor Response: Thank you again for your careful reading of this paper, and we\nare pleased to hear that you believe our last update was a significant improvement. We\nappreciate your latest comment, and we have updated the latest draft to have a clearer\ndescription of the method.\nThe most significant change to the latest draft comprises a new environment to test our\nmethod (gust mitigation for a 3D airfoil). We include these results while making minimal\nchanges to the main text to maintain the clarity from our previous draft.\nThe new draft of the manuscript is great! I really like the inclusion of the fluidic pinball with\n70 dimensions! It drives the purpose of the paper very well in a more real-world environment.\nI had three minor comments on: including clock time information, longer captions for\neach figure to make it informative, and titles for missing figures. All of these have been\naddressed. The longer captions on the figures were very helpful.\nAll my comments and concerns have been addressed!\nAuthor Response: Thankyouagainforyourexcellentfeedbackduringourfirstrevision.\nWebelievethatitmadeourpreviousdraftsignificantlystronger, andwearepleasedtoreceive\nyour latest enthusiastic response! In our most recent draft, we continue to demonstrate the\neffectiveness of our method on a new aerodynamics problem (gust mitigation of a 3D airfoil\nin an unsteady flow). Whereas the previous result from the Pinball was on a 70-dimensional\nobservation space, the new environment has a 318-dimensional space that we project down\nintoalowerdimensionalspace beforetrainingouragent; further demonstrating thescalability\nof our method. As a 3D fluids environment, it was also significantly more expensive to run.\nIn addition to including the new clock times in our SI table, we also discuss the resources\nneeded to run this example. To efficiently run the environment, we needed to distribute the\nsimulation across four NVIDIA A100 GPUs on an HPC node, requiring at least 75GB of\nVRAM, whereas simulating the surrogate model is nearly negligible and only requires access\nto CPUs. This setup provides conditions closer to what modern practitioners face when\nusing DRL for aerodynamics and significantly strengthens our claims.\nI thank the authors for considering my comments. However, the core of my arguments has\nnot been answered. The authors make major claims about their method when it is only\nperforming at toy problems. Today RL has been used to model complex 3D flows (cylinders,\nchannels, boundary layers)\nso I fail to understand the merit of a methodology that is only shown to work in simple 2D flows and toy environments.\nThe authors present no evidence that the methodology is potent enough to solve complex, non-linear problems.\nThe authors show again flows at Re around 100. The 2D flow past a single cylinder at Re up to 500 have the same\ndynamics as Re=100 - one can verify this by projecting the dynamics in a two-three dimensional latent space. And the\n\"pinball\" flow with the arrangement of three cylinders does not add anything in terms of complexity or non-linearity as the\ndynamics amount to a linear superposition of linear dynamics of a single cylinder. Finally the additional paper that is brought\nforward by Wolf et al. also shows simple dynamics including for a example a lid driven cavity flow in the laminar regime.\nI regret that I maintain my objections and cannot recommend publication of this paper.\nAuthor Response: Thank you again for the careful reading of our paper. We sincerely\nappreciate your constructive feedback, and we recognize the need for demonstrating modern\nmethods on challenging engineering problems. We have taken your feedback seriously, and\nwe have included an example of controlling the fully unsteady flow around a 3D\nNACA 0012 airfoil at Re = 1000 for the purposes of gust mitigation.\nIn response to your feedback from our first draft, we chose the airfoil geometry to be the\nstandard NACA 0012 because of its prevalence in the fluids community for a wide-range of\nengineering applications. The geometry also induces symmetry-breaking in the flow, and we\nbelieve it to be a more realistic representative of bluff-body dynamics compared to a cylinder.\nThe flow is simulated at M = 0.2 with a 20◦ angle-of-attack, well-within the post-stall\nregime (typically 14−16◦ for NACA 0012), with gust factor G = 2.0. Under these conditions,\nthe flow exhibits complex, unsteady dynamics—exhibiting 3D instabilities that are not\npossible in 2D flows. In addition to creating an efficient implementation of the environment,\nwe have thoroughly validated the CFD solver (and have included a comparison to the\nliterature in the supporting materials) and have verified the robustness of the numerics to\nstochastic control from an agent actuating the jets.\nThe scale of the environment is also much closer to what engineers might practically\nface when seeking to use DRL for flow control with CFD, requiring 72 million cells with the\nlattice Boltzmann method. To efficiently run the environment, we distributed the simulation\nacross four NVIDIA A100 GPUs on an HPC node, requiring at least 75GB of VRAM,\nwhereas simulating the surrogate model is nearly negligible and only requires access to CPUs.\nConsistent with our previous results, our method utilized far fewer evaluations with the\nfull-environment to train an effective 3D airfoil agent—taking a total of 14 hours to train\ncompared to 185 hours for the baseline DRL method.\nWhile this significantly accelerates the time it takes to discover sophisticated controllers\nfor fluid environments, our success also demonstrates the practicality of training agents with\n3\nsurrogate models using limited computational resources. The SINDy model can easily be\nevaluated on a single-CPU on a laptop in order to mitigate the real financial costs of running\nfull 3D CFD on HPC systems—effectively lowering the barrier for research groups who can\nonly afford a limited amount of time on HPC clusters.\nFinally, we further demonstrate the scalability of our approach to high-dimensional\nmeasurement spaces by obtaining observations from 318 sensor measurements from the flow\n(compared to 70 measurements in the largest 2D example) and projecting into a lower-\ndimensional space before training. Despite the complex physics, this method is able to\ncapture the dominant behavior of the system needed to learn a controller.\nWe thank you again for encouraging us to deploy our method on more demanding\napplications. We believe the addition of the 3D airfoil results significantly strengthens our\nclaims and the relevancy of the paper to the aerospace and fluids community.\n4"}, {"round": 2, "reviewer_comments": {"reviewer_2": "All my concerns have been addressed. I appreciate the paper including another higher dimensional domain in the 3D Airfoil,\nincluding its accompanying analysis. I believe this makes the paper stronger."}, "author_response": ""}]

Decision letter