Sarah Dean PRO
asst prof in CS at Cornell
Recovering and Steering Belief States in Partially Observed Systems
Sarah Dean, Cornell University
CS Colloquium at RPI, September 2026
"Next token prediction"
Auto-regressive architectures trained with next element prediction loss have had great success in many domains
$$ (o_0, o_1, o_2, o_3, o_4, o_5 ..., o_t) \rightarrow o_{t+1} $$
Eyjolfsdottir & Branson, Agent-Centric Animal Pose Forecasting
Haffner et al., Mastering Diverse Domains through World Models
Radford et al., Improving Language Understanding by Generative Pre-Training
natural language
robotics
biology
Many researchers are interested in intermediate network activations, whether for use in downstream tasks or for scientific understanding
Inside the black box...
Levy et al., Simulation Distillation: Pretraining World Models
in Simulation for Rapid Real-World Adaptation
Eyjolfsdottir & Branson, Agent-Centric Animal Pose Forecasting
robotics
biology
Outline:
1. Recovering Latent States
2. Control from Partial Observation
filter
controller
\(y\)
\(u\)
system
\(x\)
Partially observed stochastic dynamical system
$$ x_{t+1} = f(x_t,u_t,w_t),\quad y_t = g(x_t, v_t) $$
- Inputs \(u_t\) and outputs \(y_t\) are observed
- States \(x_t\), process noise \(w_t\), measurement
noise \(v_t\) are unobserved
Setting: stochastic dynamics
ex: UAV
\(x\)
\(f\)
\(g\)
\(w\)
\(v\)
Irreducible uncertainty in state is described by the posterior distribution
$$P(x_t | y_0, u_0, y_1, u_1,...,y_t) $$
\(u\)
\(y\)
$$x_0\sim \mathcal D_0, w_t\sim \mathcal D_w, v_t\sim \mathcal D_v$$
Computing this posterior is called filtering $$P(x_t|y_0,...,y_t) \propto P(x_{t-1}|y_0,...,y_{t-1}) P(x_t|x_{t-1} , u_t, y_{t})$$ but for general distributions it can be hard
Setting with no models, only observations
$$ (y_0, u_0, y_1, u_1, \dots, y_T, u_T) $$
- Idea: train a deep auto-regressive model on this sequence, predicting \(y_t\) from \(y_{0:t-1},u_{0:t-1}\)
- Question: will model activations recover something about the latent state?
Unknown dynamics/noise models
ex: UAV
\(x\)
\(?\)
\(?\)
\(w\)
\(v\)
For any similarity transform, \(\tilde x = S(x)\) is an equivalent state space representation with identical input/output behavior
$$ \tilde f = S \circ f \circ S^{-1}, \qquad \tilde g = g \circ S^{-1} $$
\(u\)
\(y\)
future
\(h_t \approx x_t??\)
past
\(\exists S ~\text{s.t.}~ S(h_t)\approx \arg\max_{x_t} P(x_t|\text{past})??\)
posterior?
\(h_t \approx \arg\max_{x_t} P(x_t|\text{past})??\)
posterior, up to similarity
transform?
Linear-Gaussian Systems
ex: UAV
Kinematic equations:
\(f_t = m a_t\)
\(v_{t+1} = v_t + a_t\)
\(p_{t+1} = p_t + v_t\)
\( \begin{bmatrix}p_{t+1}\\ v_{t+1} \end{bmatrix} = \begin{bmatrix} 1 & 1\\ 0 & 1\end{bmatrix} \begin{bmatrix}p_t\\ v_t\end{bmatrix} + \begin{bmatrix}0\\ 1/m\end{bmatrix}f_t\)
\((u_t+w_t)\)
\(y_t = \begin{bmatrix} 1 & 0 \end{bmatrix} x_t + v_t\)
\(=: x_t\)
Consider linear dynamics and measurement functions: $$ x_{t+1} = A x_t + B u_t + w_t, \qquad y_t = C x_t + v_t $$
Guassian initial state and noise: $$x_0\sim \mathcal N(0,\Sigma_0), \quad w_t\sim\mathcal N(0,\Sigma_w),\quad v_t\sim \mathcal N(0,\Sigma_v)$$
\(x\)
\(A, B\)
\(C\)
\(w\)
\(v\)
\(u\)
\(y\)
Kalman Filter
For linear-Gaussian systems, efficient posterior computation via Kalman filter:
- \(\hat x_{t+1} = A\hat x_t + Bu_t - L_t\big(y_t-C\hat x_t\big)\)
- \(\Sigma_{t+1} = (A+ L_tC)\Sigma_tA^\top + \Sigma_w\)
- \(L_t = -A\Sigma_tC^\top(C\Sigma_tC^\top+\Sigma_v)^{-1}\)
We have that \(P(x_t \mid y_{0:t-1}, u_{0:t-1})\) = \(\mathcal N(\hat x_{t}, \Sigma_{t})\)
Quickly converges to steady-state Kalman filter (SS-KF) where \(\Sigma_t=\Sigma\) and \(L_t=L\)
does not depend on inputs or outputs
mean/mode
Auto-regressive model
Idea: train an auto-regressive model on \((y_0, u_0, y_1, u_1, \dots, y_T, u_T)\)
\(L\) past inputs/outputs
\(H\) future outputs
two layer linear network
Question: will model activations recover the Kalman filter state estimate, up to similarity transform?
squared loss \( \displaystyle \min \frac{1}{2T} \sum_{t=1}^T \Vert{}y_{t:t+H-1} - G_2 G_1 {z}_t\Vert{}_{2}^2\)
\(z_t:=\)
Main result
SS-KF on \(\{u_k, y_k\}_{k=0}^{t-1}\)
Assumptions: Linear-Gaussian system with
- Non-explosive, i.e. \(\rho(A) = \max |\lambda_i(A)| \leq 1\)
- Observable, i.e. \(\begin{bmatrix}C^\top & (CA)^\top & \dots &(CA^{d_x-1})^\top\end{bmatrix}^\top\) is full rank
- Persistent excitation, \(\Sigma_w, \Sigma_v\succ 0\) and training \(\{u_t\}_{t=0}^{T-1} \sim \mathcal N(0,I)\)
similarity transform
model dimension
Proposition (Informal). The optimization landscape is benign.
Yahya Sattar
ICML '26
Theorem (Informal). Let \(d_h=d_x\) and \(H \geq d_x\) and \(L \gtrsim \log(T) / (1 - \rho_\mathrm{KF})\).
- Fix \(\delta,\epsilon>0\). Assume \(T \gtrsim \max\left\{ L \left(d_{z} + \log\left(\frac{T}{\delta}\right)\right), \frac{BH}{\epsilon} \left(d_x (H d_{y} + d_{z}) + \log \left(\frac{T}{\delta}\right)\right) \right\}.\)
- Then, for any \(B\)-bounded \(\{u_k, y_k\}_{k=0}^{t-1}\), with probability at least \(1 - \delta\), $$\min_{S\in\text{GL}_{d_x}(\mathbb R)}\Vert{}Sh_t-\hat{x}_{t|t-1}\Vert{}_{2}^2 \lesssim \epsilon$$
filter stability
Proof outline
\(\|S h_t - \hat x_t \|^2 = \|S \hat G_1 z_t - G^\star_1 z_t + G^\star_1 z_t - \hat x_{t|t-1} \|^2\)
- \( \leq 2\|S \hat G_1 - G_1^\star \|^2\|z_t\|^2 +2\|G^\star_1 z_t - \hat x_{t|t-1} \|^2\)
- \(\leq \frac{10\|z_t\|^2}{\sigma_{\min}(G_2^\star G_2^\star{}^\top, G_1^\star G_1^\star{}^\top)}\| \hat G_2 \hat G_1 - G_2^\star G_1^\star\|^2\ + 2\|G^\star_1 z_t - \hat x_{t|t-1} \|^2\)
Lemma: Procrustes distance bound holds if \(\| \hat G_2 \hat G_1 - G_2^\star G_1^\star\|^2\) small
- Recall that for the steady state Kalman filter, $$\hat x_{t+1|t} = \underbrace{(A-ALC)}_{\bar A} \hat x_{t|t-1} + B u_{t} + \underbrace{AL}_{F} y_{t} ,\qquad y_t = C \hat x_{t|t-1} + C(x_t-\hat x_{t|t-1}) + v_t$$
- As a consequence, $$\hat{x}_{t|t-1} = \underbrace{\begin{bmatrix}F & \dots & \bar{A}^{L-1}F & B & \dots & \bar{A}^{L-1}B \end{bmatrix}}_{G_1^\star} {z}_t + \bar{A}^L \hat{x}_{t-L},\qquad y_{t:t+H-1} =\underbrace{ \begin{bmatrix} C \\ CA \\ \vdots \\ CA^{H-1} \end{bmatrix}}_{G_2^\star }\hat{x}_{t|t-1} + \xi_t $$
Proof outline
$$\hat{x}_{t|t-1} = {G_1^\star} {z}_t + \bar{A}^L \hat{x}_{t-L},\qquad y_{t:t+H-1} ={G_2^\star }\hat{x}_{t|t-1} + \xi_t $$
\(\|S h_t - \hat x_t \|^2 = \|S \hat G_1 z_t - G^\star_1 z_t + G^\star_1 z_t - \hat x_{t|t-1} \|^2\)
- \( \leq 2\|S \hat G_1 - G_1^\star \|^2\|z_t\|^2 +2\| \bar{A}^L \hat{x}_{t-L} \|^2\)
- \(\leq \frac{10\|z_t\|^2}{\sigma_{\min}(G_2^\star G_2^\star{}^\top, G_1^\star G_1^\star{}^\top)}\| \hat G_2 \hat G_1 - G_2^\star G_1^\star\|^2\ + 2\| \bar{A}^L \hat{x}_{t-L}\|^2\)
Lemma: Procrustes distance bound holds if \(\| \hat G_2 \hat G_1 - G_2^\star G_1^\star\|^2\) small - \(\leq\frac{ 10 \|z_t\|^2}{\sigma_{\min}(G_2^\star G_2^\star{}^\top, G_1^\star G_1^\star{}^\top)}\frac{1}{\lambda_{\min}(\sum_{t=1}^T z_tz_t^\top)} \sum_{t=1}^T \| (\hat G_2 \hat G_1 - G_2^\star G_1^\star)z_t\|^2 + 2\| \bar{A}^L \hat{x}_{t-L} \|^2\)
properties of optimal KF
truncation error, decays geometrically in \(L\)
Theorem: in-sample prediction error $$\lesssim H\left(d_x (H d_{{y}} + d_{{z}}) + \log \left(\frac{T}{\delta}\right)\right)$$
Theorem: persistent excitation $$\sum_{t=1}^T z_t z_t^\top \succsim T I$$
Proof outline
$$\hat{x}_{t|t-1} = {G_1^\star} {z}_t + \bar{A}^L \hat{x}_{t-L},\qquad y_{t:t+H-1} ={G_2^\star }\hat{x}_{t|t-1} + \xi_t $$
\(\|S h_t - \hat x_t \|^2 = \|S \hat G_1 z_t - G^\star_1 z_t + G^\star_1 z_t - \hat x_{t|t-1} \|^2\)
- \( \leq 2\|S \hat G_1 - G_1^\star \|^2\|z_t\|^2 +2\| \bar{A}^L \hat{x}_{t-L} \|^2\)
- \(\leq \frac{10\|z_t\|^2}{\sigma_{\min}(G_2^\star G_2^\star{}^\top, G_1^\star G_1^\star{}^\top)}\| \hat G_2 \hat G_1 - G_2^\star G_1^\star\|^2\ + 2\| \bar{A}^L \hat{x}_{t-L}\|^2\)
- \(\leq\frac{ 10 \|z_t\|^2}{\sigma_{\min}(G_2^\star G_2^\star{}^\top, G_1^\star G_1^\star{}^\top)}\frac{1}{\lambda_{\min}(\sum_{t=1}^T z_tz_t^\top)} \sum_{t=1}^T \| (\hat G_2 \hat G_1 - G_2^\star G_1^\star)z_t\|^2 + 2\| \bar{A}^L \hat{x}_{t-L} \|^2\)
- \(\lesssim \frac{H}{T} \left(d_x (H d_{{y}} + d_{{z}}) + \log \left(\frac{T}{\delta}\right)\right) \Vert{}{z}_t\Vert{}_{2}^2 \) as long as $$L \gtrsim \log(T) / (1 - \rho),\qquad T \gtrsim L \left(d_{{z}} + \log(T/\delta)\right)\:.$$
Main result
Theorem (Informal). Let \(d_h=d_x\) and \(H \geq d_x\) and \(L \gtrsim \log(T) / (1 - \rho_\mathrm{KF})\).
- Fix \(\delta,\epsilon>0\). Assume \(T \gtrsim \max\left\{ L \left(d_{z} + \log\left(\frac{T}{\delta}\right)\right), \frac{BH}{\epsilon} \left(d_x (H d_{y} + d_{z}) + \log \left(\frac{T}{\delta}\right)\right) \right\}.\)
- Then, for any \(B\) bounded \(\{u_k, y_k\}_{k=0}^{t-1}\), with probability at least \(1 - \delta\),
$$\min_{S\in\text{GL}_{d_x}(\mathbb R)}\Vert{}Sh_t-\hat{x}_{t|t-1}\Vert{}_{2}^2 \lesssim \epsilon$$
1. Sequences have low rank auto-regressive structure
\(y_{t:t+H-1} \approx G_2^\star G_1^\star {z}_t\)
2. Small prediction errors imply small parameter errors (excitation)
\(\hat G_2\hat G_1 z_t \approx G_2^\star G_1^\star z_t \) for all \(t\)
\(\implies\)
\(\hat G_2\hat G_1 \approx G_2^\star G_1^\star \)
Three ingredients:
\(z_t\)
$$\min \frac{1}{2T} \sum_{t=1}^T \Vert{}y_{t:t+H-1} - G_2 G_1 {z}_t\Vert{}_{2}^2$$
3. Small end-to-end errors imply small first layer errors
\(\hat G_2\hat G_1 \approx G_2^\star G_1^\star \)
\(\implies\)
\(S\hat G_1 \approx G_1^\star \)
- Opportunity for new provable learning and control algorithms based on (deep) end-to-end training
- KF mean \(\hat x_t\) is sufficient for optimal quadratic control
- Hidden state (posterior) recovery is limited by data
- cannot recover latent quantities absent from data, whether due to lack of observability or excitation
- Open question: expand theory beyond linear-Gaussian?
- Key challenge: tractability of optimal filter (nonlinear)
- Idea: study best filter within function class, e.g. class of linear filters
- Open question: what part of posterior do we expect to recover in general? (\(\hat x_t\) vs. \(\Sigma_t\)?)
Implications and open questions
Outline:
1. Recovering Latent States
2. Control from Partial Observation
filter
controller
\(y\)
\(u\)
system
\(x\)
Observer effects occur when there is coupling between actuation and observation
- examples: electronic circuits, quantum wave collapse, human psychology, robotics ....
Observer Effects
- Linear state update with \(A\in\mathbb R^{n\times n}\), \(B\in\mathbb R^{n\times p}\) $$x_{t+1} = Ax_t + Bu_t + w_t $$
- Bilinear measurements with \(C_i\in\mathbb R^{m\times n}\) $$y_t = \underbrace{\Big(C_0 + \sum_{i=1}^p u_t[i] C_i \Big)}_{C(u_t)}x_t + v_t$$
- Gaussian process noise, measurement noise, and initial state $$\{w_t\} \sim \mathcal N(0,\Sigma_w),\quad \{v_t\} \sim \mathcal N(0, \Sigma_v),\quad x_0 \sim \mathcal N(0,\Sigma_0)$$
Bilinear Observation
Belief state via the Kalman Filter: \(x_t|\mathcal I_t \sim \mathcal N(\hat x_t,\Sigma_t)\)
\(\hat x_{t+1} = A\hat x_t + Bu_t - L_t\big(y_t-C(u_t)\hat x_t\big)\)
\(\Sigma_{t+1} = (A+ L_tC(u_t))\Sigma_tA^\top + \Sigma_w\)
\(L_t = -A\Sigma_tC(u_t)^\top(C(u_t)\Sigma_tC(u_t)^\top+\Sigma_v)^{-1}\)
Problem Setting: Optimal Control
$$\min_{u_t=\pi_t(\mathcal I_t)} \mathbb E\left[x_T^\top Q x_T+ \sum_{t=1}^{T-1} x_t^\top Q x_t + u_t^\top R u_t \right]\\ \text{s.t.} \quad x_{t+1} = Ax_t + Bu_t + w_t \\\qquad\qquad\qquad y_t =\Big(C_0 + \sum_{i=1}^p u_t[i] C_i \Big)x_t + v_t$$
Small departure from classic LQG control
Sunmook Choi
Yahya Sattar
Yassir Jedra
Maryam Fazel
Leo Maynard-Zhang
$$\mathcal I_t = \{u_0,...,u_{t-1}, y_0, ..., y_{t-1}\}$$ information set for decision-making
quadratic costs (\(Q,R\succ 0\))
"Certainty Equivalent" Separation Principle
- Separation principle (SP): independently design
estimation & control- Optimal for partially observed
LQ control (\(C_i=0\) for \(i>0\))
- Optimal for partially observed
- The SP policy has two components:
- State estimation \(\hat x_t = \mathbb E[x_t|\mathcal I_t]\) (Kalman filter)
- State dependent policy \(u_t = K^\star_t \hat x_t\) (LQR \(K^\star\))
Kalman
Filter
State Feedback
\(y\)
\(\hat x\)
\(u\)
Pitfalls of Separation
- Theorem: For \(T\geq 2\), the optimal policy is not affine in the estimated state.
- As a consequence, the SP policy is not optimal
- In addition, we construct examples where
- The SP policy locally maximizes the cost
- The SP policy is essentially an open-loop policy
Examples where SP is bad
$$\begin{align*} x_{t+1} &= \begin{bmatrix} 1 & 0.3 \\ 0 & 1\end{bmatrix} x_t + \begin{bmatrix}0.3 \\ 0 \end{bmatrix} u_t + w_t \\ y_t &= u_t\begin{bmatrix} 1 & 0\end{bmatrix} x_t + v_t \end{align*}$$
with \(Q=I\) and \(R=1000\)
$$\begin{align*} x_{t+1} &= x_t + \begin{bmatrix}0 & 1 \end{bmatrix} u_t + w_t \\ y_t &= u_t^\top \begin{bmatrix} 1\\ 0\end{bmatrix} x_t + v_t \end{align*}$$
with \(Q=\frac{1}{2}\) and \(R=I\)
\(\implies\)
infinite horizon \(K_\star = \begin{bmatrix} 0 \\ \frac{1}{2}\end{bmatrix}\)
\(y_t = \hat x_t \begin{bmatrix} 0 \\ \frac{1}{2}\end{bmatrix}^\top \begin{bmatrix} 1 \\ 0\end{bmatrix} x_t + v_t\)
\(0\) only noise is observed!
- Select a fixed "probing" action whenever all three hold:
- no probing action has been taken recently
- \(|C(K_t^\star \hat x_t)|\) is small
- \(\bar\Sigma_{OL} - \Sigma_t\) is small
- Otherwise, follow SP \(K_t \hat x_{t}\)
- Theoretically and empirically improved performance
\(\propto 1/\sigma_z^2\)
\(\propto 1/\sigma_z^2\)
Heuristic Policy
open loop steady-state covariance
Leo Maynard-Zhang
Approximately Optimal Control
Sunmook Choi
Yahya Sattar
$$\begin{align*}\min_{u_t=\pi_t(\hat x_t,\Sigma_t)}~~& \mathbb E\left[ \sum_{t=1}^{T-1} \hat x_t^\top Q \hat x_t + \mathrm{tr}(\Sigma_t) + u_t^\top R u_t \right]\\ \text{s.t.}~~& \hat x_{t+1} = A\hat x_t + Bu_t + L_t(C(u_t)\hat x_t - y_t)\\ &\Sigma_{t+1} = (A+ L_tC(u_t))\Sigma_tA^\top + \Sigma_w \end{align*}$$
$$\mathbb E\left[x_t^\top Q x_t \mid \mathcal I_t \right]= \hat x_t\top Q\hat x_t + \mathrm{tr}(Q\Sigma_t)$$
$$\begin{align*} \min_{u_t=\pi_t(\mathcal I_t)} ~~&\mathbb E\left[\sum_{t=1}^{T-1} x_t^\top Q x_t + u_t^\top R u_t \right]\\ \text{s.t.}~~& x_{t+1} = Ax_t + Bu_t + w_t \\& y_t =C(u_t)x_t + v_t\end{align*}$$
\(L_t = -A\Sigma_tC(u_t)^\top(C(u_t)\Sigma_tC(u_t)^\top+\Sigma_v)^{-1}\)
Andrew Lowitt
Beixi Du
Daniel Cao
Rewrite OCP in terms of the belief space from KF
This is now a fully observed, nonlinear, stochastic optimal control problem
innovations \(\sim \mathcal N(0, C(u_t)\Sigma_t C(u_t)^\top + \Sigma_v)\)
Approximately Optimal Control
$$\begin{align*}\pi(\hat x_t,\Sigma_t)=\arg \min_{u_k}~~& \mathbb E\left[ \sum_{k=1}^{H} \bar x_k^\top Q \bar x_k + \mathrm{tr}(\bar\Sigma_k) + u_k^\top R u_k \right]\\ \text{s.t.}~~& \bar x_{k+1} = A\bar x_k + Bu_k\\ &\bar\Sigma_{k+1} = (A+ L_kC(u_k))\Sigma_kA^\top + \Sigma_w \\ &\bar x_0 = \hat x_t ,~~ \bar \Sigma_0=\Sigma_t\end{align*}$$
\(L_t = -A\Sigma_tC(u_t)^\top(C(u_t)\Sigma_tC(u_t)^\top+\Sigma_v)^{-1}\)
Rewrite OCP in terms of the belief space from KF
Expected dynamics, open loop inputs, finite horizon \(\rightarrow\) MPC
"Solve" MPC problem with autograd and L-BFGS
- Foundations of a simple model for studying partially observed control
- Open question: do "world models" learn belief states, or only the mean/MLE?
- For purely linear dynamics, uncertainty (e.g. \(\Sigma_t\)) is not naturally in hidden activations
- Open question: do reinforcement learning algorithms automatically recover uncertainty-aware behavior?
- e.g. caution under uncertainty, probing to gather information
Implications and open questions
- Two-Layer Linear Auto-Regressive Models Estimate Latent States at ICML26 (arxiv:2606.12691) with Yahya Sattar, Sunmook Choi, Leo Maynard-Zhang, Yassir Jedra, Maryam Fazel
- Sub-optimality of the Separation Principle for Quadratic Control from Bilinear Observations at CDC25 (arxiv:2504.11555) with Yahya Sattar, Sunmook Choi, Yassir Jedra, Maryam Fazel
- Dual Control of Linear Systems from Bilinear Observations with Belief Space Model Predictive Control (arXiv:2604.24663) with Daniel Cao, Beixi Du, Andrew Lowitt, Sunmook Choi, Yahya Sattar
- Performance Gaps Between Open-Loop, Certainty-Equivalent, and
Optimal Control Under Action-Dependent Observations with Leo Maynard-Zhang, Sunmook Choi, Yahya Sattar, Yassir Jedra, Maryam Fazel
Thanks! Questions?
Sunmook Choi
Yahya Sattar
Yassir Jedra
Maryam Fazel
Leo Maynard-Zhang
Andrew Lowitt
Beixi Du
Daniel Cao
- Setting: generate sequence with arbitrary Hidden Markov Model \((A,B)\). Provide LLM with sequence only $$(o_1,o_2,...,o_t)$$
- Observation: as LLMs observe longer sequences, their predictions converge towards that of the true posterior
"Real" Models, Synthetic Data
Yijia Dai
Dai et al., Pre-trained Large Language Models Learn Hidden Markov Models In-context
- Setting: generate sequences with fixed stochastic Linear Dynamical System \((A,B,C,\Sigma_w,\Sigma_v)\).
- Train two-layer linear network
- Observation: for sufficiently large training data, the hidden activations align with the Kalman filter state estimate
Simple Models, Synthetic Data
Sattar et al., Two-Layer Linear Auto-Regressive Models Estimate Latent States
\(x\)
\(A, B\)
\(C\)
\(w\)
\(v\)
\(u\)
\(y\)
Sunmook Choi
Leo Maynard-Zhang
Recovering and Steering Belief States in Partially Observed Systems
By Sarah Dean
