Published as a conference paper at ICLR 2022

R EINFORCEMENT L EARNING S TATE E STIMATION
FOR H IGH -D IMENSIONAL N ONLINEAR S YSTEMS
Saviz Mowlavi
Massachusetts Institute of Technology
Cambridge, MA 02139, USA
smowlavi@mit.edu

Mouhacine Benosman & Saleh Nabi
Mitsubishi Electric Research Laboratories
Cambridge, MA 02139, USA{benosman,nabi}@merl.com

A BSTRACT
High-dimensional nonlinear systems such as atmospheric or oceanic flows present
a computational challenge for data assimilation (DA) algorithms such as Kalmanfilters. A potential solution is to rely on a reduced-order model (ROM) of the dynamics. However, ROMs are prone to large errors, which negatively affects theaccuracy of the resulting forecast. Here, we introduce the reinforcement learning
reduced-order estimator (RL-ROE), a ROM-based data assimilation algorithm inwhich the correction term that takes in the measurement data is given by a nonlinear stochastic policy trained through reinforcement learning. The flexibility of
the nonlinear policy enables the RL-ROE to compensate for errors of the ROM,
while still taking advantage of the imperfect knowledge of the dynamics. We showthat the trained RL-ROE is able to outperform a Kalman filter designed using thesame ROM, and displays robust estimation performance with respect to different
reference trajectories and initial state estimates.

1

I NTRODUCTION

A key challenge in Earth sciences, particularly in meteorology and oceanography, is to predict the
spatio-temporal evolution of complex processes using physical knowledge as well as observational
data (Kalnay, 2003; Dueben & Bauer, 2018). A popular approach for solving such data assimilation
problems are state estimation algorithms, such as the Kalman filter and its numerous extensions,that use sparse sensor measurements to correct the state predicted by a forward dynamical model
every time new data becomes available (Zarchan & Musoff, 2015; Houtekamer & Zhang, 2016;
Carrassi et al., 2018). However, many processes arising in Earth sciences such as atmospheric or
oceanic flows are governed by partial differential equations (PDEs) which, when discretized, yieldhigh-dimensional and nonlinear dynamical models with thousands or more of state variables. Since
integrating these high-dimensional models with Kalman filter-based DA techniques is computationally expensive, dimensionality reduction approaches in which a reduced-order model (ROM) of the
dynamics is used instead of the full model have recently become popular (Ballabrera-Poy et al.,
2001; Cao et al., 2007; Fang et al., 2009; Ştefănescu et al., 2015).
A fundamental issue is that ROMs provide a simplified and imperfect description of the dynamics(Rowley & Dawson, 2017), which negatively affects the accuracy of the forecast. One potential
solution is to improve the accuracy of the ROM through the inclusion of additional closure terms(Ahmed et al., 2021). In this paper, we leave the ROM untouched and instead propose a new design paradigm for the DA component, resulting in a novel estimation algorithm which we call areinforcement-learning reduced-order estimator (RL-ROE). The RL-ROE is constructed from the
ROM in an analogous way to a Kalman filter, with the crucial difference that the linear filter gain
function is replaced by a nonlinear stochastic policy trained through reinforcement learning (RL).
The flexibility of the nonlinear policy enables the RL-ROE to compensate for errors of the ROM,
while still taking advantage of the imperfect knowledge of the dynamics. Our approach follows
recent trends to use machine learning to improve DA frameworks (Kashinath et al., 2021; Chattopadhyay et al., 2021; Farchi et al., 2021; Grooms, 2021; Hatfield et al., 2021). We describe how
we frame the problem as a stationary Markov decision process in order to enable RL training, which
is non-trivial since the RL-ROE must be able to estimate time-varying states. Finally, we show thatthe trained RL-ROE is able to outperform a Kalman filter designed using the same ROM, and dis1

Published as a conference paper at ICLR 2022

plays robust estimation performance with respect to different reference trajectories and initial state
estimates.

2

P ROBLEM FORMULATION

2.1

S ETUP

Consider the discrete-time nonlinear system given by
zk+1 = f (zk ),yk = Czk ,

(1a)
(1b)

where zk ∈ Rn and yk ∈ Rp are respectively the state and measurement at time k, f : Rn → Rn
is a time-invariant nonlinear map from current to next state, and C ∈ Rp×n is a linear map fromstate to measurement. In this study, we assume that the dynamics given in (1) are obtained from thenumerical discretization of a nonlinear partial differential equation (PDE), which typically requiresa large number n of state dimensions. Thus, the state zk will hereafter be referred to as the highdimensional state.
2.2

R EDUCED - ORDER MODEL

Because the high dimensionality of (1) makes integration with DA techniques computationallyexpensive, an alternative approach is to formulate a reduced-order model (ROM) of the dynamics (Rowley & Dawson, 2017). First, one chooses a suitable linearly independent set of modes{u1 , . . . , ur }, where ui ∈ Rn , defining an r-dimensional subspace of Rn in which most of the
dynamics is assumed to take place. Stacking these modes as columns of a matrix U ∈ Rn×r , one
can then express zk ' U xk , where the reduced-order state xk ∈ Rr represents the coordinates of
zk in the subspace. Finally, one finds a ROM for the dynamics of xk , which is vastly cheaper to
evolve than (1) when r  n.
There exist various ways to find an appropriate set of modes U and corresponding ROM for the
dynamics of xk (Taira et al., 2017). In this work, we employ the Dynamic Mode Decomposition(DMD), a purely data-driven algorithm that has found wide applications in fields ranging from fluid
dynamics to neuroscience (Schmid, 2010; Kutz et al., 2016) and is related to the linear inverse modeling (LIM) method from climate science (Tu et al., 2014). Starting with a collection of snapshots
Z = {z0 , . . . , zm } collected along a trajectory of (1a), the DMD seeks a best-fit linear model of the
dynamics in the form of a matrix A ∈ Rn×n such that zk+1 ' Azk , and computes the modes U as
the r leading principal component analysis (PCA) modes of Z. The transformation zk ' U xk andthe orthogonality of U then yield a linear discrete-time ROM of the form
xk+1 = Ar xk + wk ,yk = C r x k + v k ,

(2a)
(2b)

where Ar = U T AU ∈ Rr×r and Cr = CU ∈ Rp×r are the reduced-order state-transition and observation models, respectively. In order to account for the neglected PCA modes of Z as well as theunmodeled dynamics incurred by the linear approximation zk+1 ' Azk , we add (unknown) nonGaussian process noise wk and observation noise vk . Additional details regarding the calculationof Ar and U are provided in Appendix A.
2.3

R EDUCED - ORDER ESTIMATOR

This paper uses reinforcement learning (RL) to solve the following estimation problem: given a sequence of measurements {y0 , · · · , yk } from a reference trajectory {z0 , · · · , zk } of (1) and knowing
the ROM (2) defined by Ar , Cr and U , we want to estimate the high-dimensional state zk at current
time k. To this effect, we design a reduced-order estimator (ROE) of the formx̂k = Ar x̂k−1 + ak ,ak ∼ πθ ( · |yk , x̂k−1 ),

(3a)
(3b)

where x̂k is an estimate of the reduced-order state xk , and ak ∈ Rr is an action sampled from astochastic policy πθ which depends on the current measurement yk and the previous state estimate
2

Published as a conference paper at ICLR 2022x̂k−1 . The subscript θ denotes the set of parameters that defines the stochastic policy, whose goal
is to minimize the mean square error E[zk − ẑk ] over a range of reference trajectories and initial
reduced-order state estimates. Here, ẑk = U x̂k denotes the high-dimensional state estimate reconstructed from x̂k . A Kalman filter is a special case of such an estimator, for which the action in (3b)is given by
ak = Kk (yk − Cr Ar x̂k−1 ),(4)with Kk ∈ Rr×p the optimal Kalman gain. Although the Kalman filter is optimal when the statetransition and observation models are known exactly, its performance suffers in the presence of
unmodeled dynamics. In our case, such model errors are unavoidable due to the ROM (2) being
an inherent approximation of the high-dimensional dynamics (1), which motivates our adoption of
the more general form (3b). This form retains the dependence of ak on yk and x̂k−1 but is more
flexible thanks to the nonlinearity of the stochastic policy πθ , which we train with deep RL in an
offline stage. We call the estimator constructed and trained through this process an RL-trained ROE,
or RL-ROE for short.
2.4

S UMMARY OF THE PROPOSED METHODOLOGY

In summary, the methodology we propose consists of the following three steps. The first two arecarried out offline using full knowledge of {z0 , · · · , zk } from several trajectories of (1), and thethird is performed online using sole knowledge of measurements {y0 , · · · , yk } from a previously
unseen trajectory of (1).
1. Construction of an RL-ROE. A ROM of the form (2) is obtained by applying the DMD
to high-dimensional state snapshots zk from a single trajectory of (1). An RL-ROE (3) isthen designed based on this ROM.
2. Training of the RL-ROE. The policy πθ of the RL-ROE is trained using high-dimensional
snapshots zk from multiple reference trajectories of (1).
3. Deployment of the RL-ROE. Using measurements yk from a previously unseen reference
trajectory of (1), the trained RL-ROE returns an estimate ẑk = U x̂k for the corresponding(unobserved) high-dimensional state zk .
We note that previous studies (Morimoto & Doya, 2007; Hu et al., 2020) have already proposed
designing state estimators using policies trained through reinforcement learning. In appendix B, wedetail the novelties of our proposed methodology with respect to these studies.

3

F RAMING THE PROBLEM AS A STATIONARY MDPIn this section, we describe the offline training process for the policy πθ in the RL-ROE (3). In
order to train πθ with reinforcement learning, we need to formulate the problem as a stationary
Markov decision process (MDP). However, this is no trivial task given that the aim of the policy is
to minimize the error between the state estimate ẑk = U x̂k and a time-dependent reference state
zk . At first sight, such trajectory tracking problem requires a time-dependent reward function and,
therefore, a time-varying MDP.
To be able to use off-the-shelf RL algorithms, we introduce a trick to translate this time-varying
MDP to an equivalent stationary MDP based on an extended state. Indeed, we show hereafter thatthe problem can be framed as a stationary MDP by including zk into our definition of the MDP’sstate. Letting sk = (zk , x̂k−1 ) ∈ Rn+r denote an augmented state at time k, we can define an MDP
consisting of the tuple (S, A, P, R), where S = Rn+r is the augmented state space, A ⊂ Rr isthe action space, P(·|sk , ak ) is a transition probability, and R(sk , ak , sk+1 ) is a reward function.
At each time step k, the agent selects an action ak ∈ A according to the policy πθ defined in (3b),which can be expressed asak ∼ πθ ( · |ok ),(5)where ok = (yk , x̂k−1 ) = (Czk , x̂k−1 ) is a partial observation of the current state sk . The statesk+1 = (zk+1 , x̂k ) at the next time step is then obtained from equations (1a) and (3a) as
sk+1 = (f (zk ), Ar x̂k−1 + ak ),3(6)

Published as a conference paper at ICLR 2022

which defines the transition model sk+1 ∼ P(·|sk , ak ). Finally, the agent receives the reward
rk = R(sk , ak , sk+1 ) = −(zk − U x̂k )T Q(zk − U x̂k ) − aT
k Rak ,

(7)

where Q ∈ Rn×n and R ∈ Rr×r are positive semidefinite and positive definite matrices, respectively. The first term in the reward function R penalizes the difference between the high-dimensional
state estimate ẑk = U x̂k and the reference zk , which is only partially observed by the agent. The
second term favors smaller values for the action ak ; such regularization leads to more robust estimation performance in the presence of noise during online deployment of the RL-ROE, as we will see
later. Unless indicated otherwise, we will consider Q = I and R = I. Thanks to the incorporationof zk into sk , the reward function (7) has no explicit time dependence and the MDP is therefore
stationary.
Thanks to the stationarity of the MDP, one can use any deep RL algorithm to train the policy πθusing high-dimensional snapshots zk from multiple reference trajectories of (1). Here, we use theProximal Policy Optimization (PPO) algorithm (Schulman et al., 2017), which belongs to the classof policy gradient methods (Sutton et al., 2000). The setup of the RL training process, the parameterization of the policy πθ , and implementation details are discussed in Appendix C.
Remark. Since the policy (5) is conditioned on a partial observation ok of the state sk , the MDP
we have defined in this section is, in fact, a partially observable MDP (POMDP). In this case, it
is known that the optimal policy depends on a summary of the history of past observations andactions, hk = {o1 , a1 , . . . , ok }, rather than just the current observation ok (Kaelbling et al., 1998).
However, policies formulated based on an incomplete summary of hk are common in practice andstill achieve good results (Sutton & Barto, 2018). We therefore pursue this approach in the present
paper, and leave for future work testing the generalization of our policy input to a more completesummary of hk .

4

R ESULTS

We evaluate our proposed RL-ROE using simulations of the Burgers equation, a prototypical nonlinear hyperbolic PDE which takes the form
∂u
∂u
∂2u
+u
− ν 2 = 0,
∂t
∂x
∂x

(8)

where u(x, t) is the velocity at position x ∈ [0, L] and time t ∈ [0, T ], and the scalar parameter ν
acts like a viscosity. The boundary conditions are periodic and the initial condition u(x, 0) = u0 (x)will be specified later. We choose L = 1, T = 10, and ν = 0.01.
We discretize the Burgers equation (8) using a spectral method with n = 256 Fourier modes anda fifth-order Runge-Kutta time integration scheme, and we save snapshots every ∆t = 0.05 timeunits. The resulting discrete system takes the form (1) where the state vector zk ∈ Rn comprises
the values of u at the collocation points xi = iL/n and time t = k∆t, and the measurement vectoryk ∈ Rp comprises the values of u at p = 8 sensor locations x̄i = iL/p.
We then follow the procedure outlined in Section 2.4 to construct an RL-ROE, which we train
offline using PPO following the methodology presented in Section 3. The trained RL-ROE is finallydeployed online and compared against a Kalman filter under various reference trajectories as wellas initial state estimates. To ensure a proper comparison, the Kalman filter is constructed from thesame ROM on which the RL-ROE is based, and will therefore be referred to as KF-ROE. Moredetails regarding these various steps can be found in Appendix D.
The performances of the trained RL-ROE and KF-ROE are now compared. Figures 1(a,b,c) show
the L2 error of the RL-ROE and KF-ROE for specific reference trajectories of the Burgers equation,
initialized from (15a) using α = 0.5, 1, 2 where α is the amplitude of the initial condition (15a).
Importantly, we note that the ROM was computed from the trajectory initialized with α = 1, and istherefore very accurate (i.e., the model error is low) for this case. For each reference trajectory, we
consider 20 different initial state estimates sampled from (15b). The curves and shaded area reported
in Figures 1(a,b,c) indicate the mean and standard deviation of the error, defined at time step k by
|U x̂k − zk |, where x̂k is the reduced-order estimate given by the RL-ROE or KF-ROE, and zk
is the high-dimensional reference solution. Figures 1(d,e,f) show the trajectories of the reference,4

ROM
training
data

Published as a conference paper at ICLR 2022

(a)

(b)

(c)

(d)

(e)

(f)

Figure 1: Accuracy of the RL-ROE. (a,b,c) L2 error of the RL-ROE and KF-ROE with respect tospecific reference trajectories of the Burgers equation, initialized from (15a) using α = 0.5, 1, 2.
The RL-ROE and KF-ROE are evaluated using 20 different initial estimates sampled from (15b).
The curves and shaded area indicate the mean and standard deviation of the error, respectively.
(d,e,f) Phase space trajectories of the RL-ROE and KF-ROE predictions for 5 initial estimates, andthe reference solution.
RL-ROE and KF-ROE solutions in a two-dimensional slice of phase space spanned by the second
and third columns1 of U .
A few important observations emerge from Figure 1. First, the RL-ROE outperforms the KF-ROE
when the ROM suffers from large model errors, as is the case for α = 0.5 and α = 2. In the time
window t ≤ 2 during which most of the transient dynamics take place, the RL-ROE displays up
to an order of magnitude lower error than the KF-ROE. Second, when the ROM is very accurate,
as is the case for α = 1, the KF-ROE gives lower error for most of the time duration. Even then,however, Figure 1(e) shows that the RL-ROE converges faster to the reference trajectory. Finally,
the RL-ROE manages to keep the error at a low level in the time window t ∈ [5, 10], despite the factthat it was trained using trajectories that end at t = 5.

5

C ONCLUSIONSIn this paper, we have introduced the reinforcement learning reduced-order estimator (RL-ROE), a
new methodology for forecasting the state of a high-dimensional nonlinear dynamical system using sparse observations. Our approach follows the recent trend of constructing a computationallyinexpensive reduced-order model (ROM) to approximate the dynamics of the system. The noveltyof our contribution lies in the design, based on this ROM, of a reduced-order estimator (ROE) in
which the feedback correction term is given by a nonlinear stochastic policy trained through reinforcement learning. To be able to use off-the-shelf RL algorithms, we introduce a trick to translate
this trajectory tracking problem, i.e., time-varying MDP, to an equivalent stationary MDP based on
an augmented state. We show using simulations of the Burgers equation that the trained RL-ROE
is able to outperform a Kalman filter designed using the same ROM and displays robust estimation
performance with respect to different reference trajectories and initial state estimates.
1
Since the columns of U approximate the PCA modes of the training snapshots Z DMD without centering, thefirst column will be dominated by the mean of the data. Thus, we display the trajectory coordinates associatedwith the second and third columns, which capture the largest amount of variance within the data.

5

Published as a conference paper at ICLR 2022

R EFERENCES
Shady E Ahmed, Suraj Pawar, Omer San, Adil Rasheed, Traian Iliescu, and Bernd R Noack. On
closures for reduced order models—a spectrum of first-principle to machine-learned avenues.
Physics of Fluids, 33(9):091301, 2021.
Joaquim Ballabrera-Poy, Antonio J Busalacchi, and Ragu Murtugudde. Application of a reducedorder kalman filter to initialize a coupled atmosphere–ocean model: Impact on the prediction ofel nino. Journal of climate, 14(8):1720–1737, 2001.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, andWojciech Zaremba. Openai gym. arXiv preprint arXiv:1606.01540, 2016.
Yanhua Cao, Jiang Zhu, I Michael Navon, and Zhendong Luo. A reduced-order approach to fourdimensional variational data assimilation using proper orthogonal decomposition. International
Journal for Numerical Methods in Fluids, 53(10):1571–1583, 2007.
Alberto Carrassi, Marc Bocquet, Laurent Bertino, and Geir Evensen. Data assimilation in the geosciences: An overview of methods, issues, and perspectives. Wiley Interdisciplinary Reviews:
Climate Change, 9(5):e535, 2018.
Ashesh Chattopadhyay, Mustafa Mustafa, Pedram Hassanzadeh, Eviatar Bach, and Karthik
Kashinath. Towards physically consistent data-driven weather forecasting: Integrating data assimilation with equivariance-preserving spatial transformers in a case study with era5. Geoscientific
Model Development Discussions, pp. 1–23, 2021.
Peter D Dueben and Peter Bauer. Challenges and design choices for global weather and climate
models based on machine learning. Geoscientific Model Development, 11(10):3999–4009, 2018.
F Fang, CC Pain, IM Navon, MD Piggott, GJ Gorman, PE Farrell, PA Allison, and AJH Goddard.
A pod reduced-order 4d-var adaptive mesh ocean modelling approach. International Journal for
Numerical Methods in Fluids, 60(7):709–732, 2009.
Alban Farchi, Patrick Laloyaux, Massimo Bonavita, and Marc Bocquet. Using machine learning tocorrect model error in data assimilation and forecast applications. Quarterly Journal of the Royal
Meteorological Society, 147(739):3067–3084, 2021.
Ian Grooms. Analog ensemble data assimilation and a method for constructing analogs with variational autoencoders. Quarterly Journal of the Royal Meteorological Society, 147(734):139–149,
2021.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy
maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning, pp. 1861–1870. PMLR, 2018.
Sam Hatfield, Matthew Chantry, Peter Dueben, Philippe Lopez, Alan Geer, and Tim Palmer. Building tangent-linear and adjoint models for data assimilation with neural networks. Journal ofAdvances in Modeling Earth Systems, 13(9):e2021MS002521, 2021.
Peter L Houtekamer and Fuqing Zhang. Review of the ensemble kalman filter for atmospheric data
assimilation. Monthly Weather Review, 144(12):4489–4532, 2016.
Liang Hu, Chengwei Wu, and Wei Pan. Lyapunov-based reinforcement learning state estimator.
arXiv preprint arXiv:2010.13529, 2020.
Leslie Pack Kaelbling, Michael L Littman, and Anthony R Cassandra. Planning and acting inpartially observable stochastic domains. Artificial intelligence, 101(1-2):99–134, 1998.
Eugenia Kalnay. Atmospheric modeling, data assimilation and predictability. Cambridge university
press, 2003.
K Kashinath, M Mustafa, A Albert, JL Wu, C Jiang, S Esmaeilzadeh, K Azizzadenesheli, R Wang,
A Chattopadhyay, A Singh, et al. Physics-informed machine learning: case studies for weather
and climate modelling. Philosophical Transactions of the Royal Society A, 379(2194):20200093,
2021.
6

Published as a conference paper at ICLR 2022

J Nathan Kutz, Steven L Brunton, Bingni W Brunton, and Joshua L Proctor. Dynamic mode decomposition: data-driven modeling of complex systems. SIAM, 2016.
Jun Morimoto and Kenji Doya. Reinforcement learning state estimator. Neural computation, 19(3):730–756, 2007.
Antonin Raffin, Ashley Hill, Maximilian Ernestus, Adam Gleave, Anssi Kanervisto, and Noah Dormann. Stable baselines3. https://github.com/DLR-RM/stable-baselines3, 2019.
Clarence W Rowley and Scott TM Dawson. Model reduction for flow analysis and control. Annual
Review of Fluid Mechanics, 49:387–417, 2017.
Peter J Schmid. Dynamic mode decomposition of numerical and experimental data. Journal of fluid
mechanics, 656:5–28, 2010.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy
optimization algorithms. arXiv preprint arXiv:1707.06347, 2017.
Răzvan Ştefănescu, Adrian Sandu, and Ionel Michael Navon. Pod/deim reduced-order strategies for
efficient four dimensional variational data assimilation. Journal of Computational Physics, 295:
569–595, 2015.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Richard S Sutton, David A McAllester, Satinder P Singh, and Yishay Mansour. Policy gradientmethods for reinforcement learning with function approximation. In Advances in neural information processing systems, pp. 1057–1063, 2000.
Kunihiko Taira, Steven L Brunton, Scott TM Dawson, Clarence W Rowley, Tim Colonius, Beverley J McKeon, Oliver T Schmidt, Stanislav Gordeyev, Vassilios Theofilis, and Lawrence S
Ukeiley. Modal analysis of fluid flows: An overview. Aiaa Journal, 55(12):4013–4041, 2017.
Jonathan H Tu, Clarence W Rowley, Dirk M Luchtenburg, Steven L Brunton, and J Nathan Kutz.
On dynamic mode decomposition: theory and applications. Journal of Computational Dynamics,
1(2):391–421, 2014.
Paul Zarchan and Howard Musoff. Fundamentals of Kalman filtering: a practical approach. Aiaa,
2015.

7

Published as a conference paper at ICLR 2022

A

DYNAMIC M ODE D ECOMPOSITIONIn this appendix, we describe the DMD algorithm (Schmid, 2010; Tu et al., 2014), which is a popular
data-driven method to extract spatial modes and low-dimensional dynamics from a dataset of highdimensional snapshots. Here, we use the DMD to construct a ROM of the form (2) given a snapshot
sequence Z = {z0 , . . . , zm } collected along a trajectory of (1a) and an observation model C.
Fundamentally, the DMD seeks a best-fit linear model of the dynamics in the form of a matrix
A ∈ Rn×n such that zk+1 ' Azk . Arranging the snapshots into two time-shifted matrices
X = {z0 , . . . , zm−1 },Y = {z1 , . . . , zm },

†

(9)

†

the best-fit linear model is given by A = Y X , where X is the pseudoinverse of X. The ROM
is then obtained by projecting the matrices A and C onto a basis U consisting of the r leadingleft singular vectors of X, which approximate the r leading PCA modes of Z. Using the truncatedsingular value decomposition
X = U ΣV T
(10)
n×r
r×r
where U , V ∈ R
and Σ ∈ R , the resulting reduced-order state-transition and observation
models are given by
Ar = U T AU = U T Y V Σ−1 ,Cr = CU .

(11a)
(11b)

Conveniently, the ROM matrices Ar and Cr can be calculated directly from the truncated SVD of
X, which avoids forming the large n × n matrix A.

B

R ELATED WORKIn this appendix, we compare our work with two previous studies that have already proposed designing state estimators using policies trained through reinforcement learning. Morimoto & Doya(2007) introduced an estimator of the form x̂k = f (x̂k−1 ) + L(x̂k−1 )(yk−1 − C x̂k−1 ), wheref (·) is the state-transition model of the system, and the state-dependent filter gain matrix L(x̂k−1 )is defined using Gaussian basis functions whose parameters are learned through a variant of vanillapolicy gradient. Their reward function, however, was calculated using the measurement error instead of the state estimate error, potentially limiting the performance of the trained estimator. Hu
et al. (2020) proposed an estimator of the form x̂k = f (x̂k−1 ) + L(xk − x̂k )(yk − Cf (x̂k−1 )),where L(xk − x̂k ) is approximated by neural networks trained with a modified Soft-Actor Critic
algorithm (Haarnoja et al., 2018). Although they derived convergence properties for the estimate
error, the dependence of the filter gain L(xk − x̂k ) on the reference state xk limits its practicalapplication. A major difference between these past studies and our work is that they do not construct a ROM of the dynamics and only consider low-dimensional systems with four state variables
at most, in comparison with the hundred or more state dimensions that our RL-ROE can handle.
Therefore, RL-ROE represents the first application of reinforcement learning to state estimation for
high-dimensional systems, which makes it applicable to systems governed by PDEs such as fluid
flows.

CS ETUP OF THE RL TRAINING PROCESSIn this appendix, we describe the setup of the RL training process as well as implementation details.
The goal of the training is to find the optimal policy parametersθ ∗ = arg max E [R(τ )],
τ ∼πθ

θ

(12)

PK
where the expectation is over trajectories τ = (s1 , a1 , s2 , a2 , . . . ), and R(τ ) = k=1 rk is thefinite-horizon undiscounted return, with the integer K denoting the length of each training trajectory.
Contrary to conventional RL notation, trajectories here start at time k = 1. Indeed, the environmentis initialized at time k = 0 according to the distributions
z0 ∼ pz0 (·),
x̂0 ∼ px̂0 (·),
8

(13a)
(13b)

Published as a conference paper at ICLR 2022

from which the augmented state s1 = (z1 , x̂0 ) = (f (z0 ), x̂0 ) follows immediately. Thus, s1
constitutes the start of the trajectory of agent-environment interactions. This sequence of operations
mirrors that of a Kalman filter: the calculation of the current state estimate uses the previous state
estimate as well as the current measurement, and therefore begins from the second time step, whichis k = 1 in our case.
To find the optimal policy parameters θ ∗ , we employ the Proximal Policy Optimization (PPO) algorithm (Schulman et al., 2017), which belongs to the class of policy gradient methods (Sutton et al.,2000). PPO alternates between sampling data by computing a set of trajectories {τ1 , τ2 , τ3 , . . . } using the most recent version of the policy, and updating the policy parameters θ in a way that increases
the probability of actions that led to higher rewards during the sampling phase. The policy πθ encodes a diagonal Gaussian distribution described by a neural network that maps from observation
to mean action, µθ0 (ok ), together with a vector of standard deviations σ, so that θ = {θ 0 , σ}. We
utilize the Stable Baselines3 (SB3) implementation of PPO (Raffin et al., 2019) and define our MDP
as a custom environment in OpenAI Gym (Brockman et al., 2016). Besides the discount factor γ,
all results to follow are obtained with the default PPO hyperparameters in SB3, which demonstrates
the robustness of our approach with respect to the RL hyperparameters.

D

P ROTOCOL FOR THE RESULTS ON THE B URGERS EQUATION

We adopted the following steps when obtaining the results shown in Section 4:
1. Construction of an RL-ROE. Starting from the pulse-shaped initial condition1
u0 (x) =,(14)
cosh(20(x − L/2))
we calculate one solution trajectory of (1) for t ∈ [0, T /2] = [0, 5], and we denote as
DMD} the resulting solution snapshots at times k = 0, . . . , m =Z DMD = {z0DMD , . . . , zm
T /2∆t. The DMD is then applied to these snapshots, yielding a ROM of the form (2)defined by matrices Ar , Cr and U . The ROM governs the evolution of a reduced-orderstate xk ' U T zk ∈ Rr . We pick r = 15 for the dimensionality of the reduced-ordersubspace (which corresponds in this case to 99.99% of the energy of the snapshots Z DMD
being included in the modes U ), giving the ROM a significant computational advantage
compared with the high-dimensional system (1) of size n = 256. An RL-ROE (3) is thendesigned based on this ROM.
2. Training of the RL-ROE. We train the stochastic policy πθ of the RL-ROE (3) using PPO,
as described in Section 3. In order for the resulting estimator to perform well under various
reference trajectories and initial estimates, we initialize each trajectory of the MDP during
the offline training process withz0 = αz0DMD ,(15a)x̂0 = U T z0DMD + β,(15b)where α ∼ U(0.5, 2) ∈ R and β ∼ N (0, 0.1I) ∈ Rr , which defines the distributions givenin (13). In other words, the reference trajectories are initialized as a randomly scaled up
or scaled down version of the pulse defined in (14), and the reduced-order state estimate
is initialized as a reduced-order projection of that pulse, polluted with additive Gaussian
noise. During training, we limit each trajectory to the same time window t ∈ [0, T /2] thatwas used in constructing the ROM – that is, we pick K = T /2∆t in the finite-horizon
return. We end the training when the return no longer increases on average. The training
hyperparameters and learning curves are presented in Appendix E.
3. Deployment and evaluation of the RL-ROE. We evaluate the trained RL-ROE against
a time-dependent Kalman filter constructed from the same ROM, which we refer to as
KF-ROE. The KF-ROE is given by equations (3a) and (4), with the calculation of the timevarying Kalman gain detailed in Appendix F. The RL-ROE and KF-ROE are comparedonline based on three specific reference trajectories initialized from (15a) using α = 0.5,
1, and 2. For each reference trajectory, we consider 20 different initial state estimates sampled from (15b), and feed measurements yk to the RL-ROE and KF-ROE. Their tracking
performance of the reference state zk is then evaluated and compared over the full time
window t ∈ [0, T ].
9

Published as a conference paper at ICLR 2022

E

T RAINING HYPERPARAMETERS AND LEARNING CURVES

The stochastic policy πθ is trained with PPO using the default hyperparameters from Stable Baselines3,
except for the discountfactor γ which we choose
as 0.75. The mean(d)output of the stochastic
(a)
(b)
(c)
policy and the value function are approximated by two neural networks, each containing two hidden
layers with 64 neurons and tanh activation functions. The training process alternates between sampling data for 20 trajectories (of length 100 timesteps each) and updating the policy. Each policy
update consists of multiple gradients steps through the most recent data using 10 epochs, a minibatch
size of 64 and a learning rate of 0.0003. The policy is trained for a total of one million timesteps,
corresponding to 10000 trajectories. Figure 2 reports the learning curves. During training, the policy
is tested (with stochasticity switched off) after each update using 10 separate test trajectories, andis saved if it outperforms the previous best policy. Finally, the RL-ROE is defined using the latest
saved policy upon ending of the training process, and the stochasticity of the policy is switched off
during subsequent evaluation of the RL-ROE.

Figure 2: Learning curves for the stochastic policy. The line and shaded area show the mean andstandard deviation of the results over 10 runs, each one smoothed with a moving average of size 100
episodes.

F

K ALMAN FILTER

The time-dependent Kalman filter that we use as a benchmark in this paper, KF-ROE, is based on thesame ROM (2) as the RL-ROE, with identical matrices Ar , Cr and U . Similarly to the RL-ROE,
the reduced-order estimate x̂k is given by equation (3a), from which the high-dimensional estimate
is reconstructed as ẑk = U x̂k . However, the KF-ROE differs from the RL-ROE in its definition of
the action ak in (3a), which is instead given by the linear feedback term (4). The calculation of the
optimal Kalman gain Kk in (4) requires the following operations at each time step:Pk− = Ar Pk−1 AT
r + Qk ,(16)

Sk = Cr Pk− CrT + Rk ,Kk = Pk− CrT Sk−1 ,Pk = (I − Kk Cr )Pk− ,

(17)
(18)
(19)

where Pk− and Pk are respectively the a priori and a posteriori estimate covariance matrices, Sk is
the innovation covariance, and Qk and Rk are respectively the covariance matrices of the process
noise wk and observation noise vk in the ROM (2). Since these noise covariance matrices are
unknown, we choose Qk = Rk = I for all k after verifying empirically that these values yield
the best possible results. At time step k = 0, the a posteriori estimate covariance is initializedas P0 = cov(U T z0 − x̂0 ), which can be calculated from the initial reference and estimated state
distributions (15).

10