SlowGuess/NeurIPS_2023_OCR huggingface dataset:
[60] Tadahiro Taniguchi, Shingo Murata, Masahiro Suzuki, Dimitri Ognibene, Pablo Lanillos, Emre Ugur, Lorenzo Jamone, Tomoaki Nakamura, Alejandra Ciria, Bruno Lara, et al. World models and predictive coding for cognitive and developmental robotics: frontiers and challenges. Advanced Robotics, pages 1-27, 2023.
[61] John W Krakauer, Asif A Ghazanfar, Alex Gomez-Marin, Malcolm A MacIver, and David Poeppel. Neuroscience needs behavior: correcting a reductionist bias. Neuron, 93(3):480-490, 2017.
[62] Paul Cisek and Alexandre Pastor-Bernier. On the challenges and mechanisms of embodied decisions. Philosophical Transactions of the Royal Society B: Biological Sciences, 369(1655):20130479, 2014.
[63] Cory T Miller, David Gire, Kim Hoke, Alexander C Huk, Darcy Kelley, David A Leopold, Matthew C Smear, Frederic Theunissen, Michael Yartsev, and Christopher M Niell. Natural behavior is the language of the brain. Current Biology, 32(10):R482-R493, 2022.
[64] Christopher M Bishop. Pattern recognition and machine learning. Springer, 2006.
[65] Roy Frostig, Matthew James Johnson, and Chris Leary. Compiling machine learning programs via high-level tracing. Systems for Machine Learning, 4(9), 2018.
[66] Pauli Virtanen, Ralf Gommers, Travis E Oliphant, Matt Haberland, Tyler Reddy, David Cournaeau, Evgeni Burovski, Pearu Peterson, Warren Weckesser, Jonathan Bright, et al. Scipy 1.0: fundamental algorithms for scientific computing in python. Nature methods, 17(3):261-272, 2020.
[67] Mathieu Blondel, Quentin Berthet, Marco Cuturi, Roy Frostig, Stephan Hoyer, Felipe Llinares-López, Fabian Pedregosa, and Jean-Philippe Vert. Efficient and modular implicit differentiation. arXiv preprint arXiv:2105.15183, 2021.
[68] Steven George Krantz and Harold R Parks. The implicit function theorem: history, theory, and applications. Springer Science & Business Media, 2002.
[69] Mathieu Blondel, Quentin Berthet, Marco Cuturi, Roy Frostig, Stephan Hoyer, Felipe Llinares-López, Fabian Pedregosa, and Jean-Philippe Vert. Efficient and modular implicit differentiation. In Advances in Neural Information Processing Systems, volume 35, pages 5230–5242, 2022.
[70] Weiwei Li. Optimal control for biological movement systems. PhD thesis, UC San Diego, 2006.
Table 1: Notation
xt∈Rn	state at time t
ut∈Ru	control at time t
yt∈Rm	observation at time t
T∈N	number of time steps
f: Rn×Ru×Rv→Rn	system dynamics function
h: Rn×Rw→Rm	observation function
vt∈Rv	standard multivariate normal dynamics noise
wt∈Rw	standard multivariate normal observation noise
ct: Rn×Ru→R	cost function at intermediate time steps
cT: Rn→R	cost function at final time step
bt∈Rb	summary statistics of agent's belief, p(xt|y1:t-1)
βt: Rb×Ru×Rm→Rb	belief dynamics
πt: Rb×Rj→Ru	policy of the agent
θ∈Rp	model parameters
g: Rn×Rb×Rv×Rw×Rj→Rn×Rb	joint dynamics of states and beliefs
B MCE IRL Baseline
Maximum causal entropy inverse reinforcement learning (MCE IRL) [47] can be formulated as maximizing the likelihood
$$
p \left(\boldsymbol {x} _ {1: T} \mid \boldsymbol {\theta}\right) = p \left(\boldsymbol {x} _ {0} \mid \boldsymbol {\theta}\right) \prod_ {t = 0} ^ {T - 1} p \left(\boldsymbol {x} _ {t + 1} \mid \boldsymbol {x} _ {t}, \boldsymbol {u} _ {t}, \boldsymbol {\theta}\right) \Pi_ {t} ^ {\boldsymbol {\theta}} \left(\boldsymbol {u} _ {t} \mid \boldsymbol {x} _ {t}\right), \tag {10}
$$
$$
\Pi_ {t} ^ {\boldsymbol {\theta}} (\boldsymbol {u} _ {t} \mid \boldsymbol {x} _ {t}) = \exp (Q _ {t} ^ {\boldsymbol {\theta}} (\boldsymbol {x} _ {t}, \boldsymbol {u} _ {t}) - V _ {t} ^ {\boldsymbol {\theta}} (\boldsymbol {x} _ {t})).
$$
Here,  $Q_{t}^{\theta}$  is the soft Q-function at time  $t$ , given by
$$
Q _ {t} ^ {\boldsymbol {\theta}} (\boldsymbol {x} _ {t}, \boldsymbol {u} _ {t}) = - c _ {t} (\boldsymbol {x} _ {t}, \boldsymbol {u} _ {t}) - \mathbb {E} \left[ V _ {t + 1} ^ {\boldsymbol {\theta}} (\boldsymbol {x} _ {t + 1}) \right]
$$
and  $V_{t}^{\theta}$  the normalization,
$$
V _ {t} ^ {\pmb {\theta}} (\pmb {x} _ {t}) = \log \int_ {\pmb {u} _ {t}} \exp (Q _ {t} ^ {\pmb {\theta}} (\pmb {x} _ {t}, \pmb {u} _ {t})) \mathrm {d} \pmb {u} _ {t}.
$$
For arbitrary systems, computing the soft Q-function exactly is infeasible, therefore common methods apply approximations such as linearization [18], importance sampling [19], or deep function approximation [23]. For the case of linear dynamics and quadratic reward, the optimal policy is given by a Gaussian distribution  $\Pi_t(\boldsymbol{u}_t \mid \boldsymbol{x}_t) = \mathcal{N}(\boldsymbol{u}_t; L_t\boldsymbol{x}_t, -H_t^{-1})$ , where  $L_t$  is the controller gain and  $H_t$  a matrix resulting from the computation of the LQG controller [48]. As the tasks we consider can be well solved by linearizing the dynamics locally, we choose an approximation by linearization and use the optimal MCE controller for the linearized dynamics to compute the likelihood function for the parameter set  $\theta$ .
To apply this baseline to the setting where control signals are missing, we use the estimates of the controls as we determine in our proposed method for the data-based linearization (Section 3.3 and Appendix G).
We can compute the approximate likelihood function by performing the following steps:
1. Estimate the missing control signals using Eq. (11)
2. Linearize the system as described in Section 3.3
3. Compute the MCE policy for the linearized system, i.e., compute  $Q_{t}^{\theta}$  and  $V_{t}^{\theta}$ .
4. Compute the likelihood (in log space) using Eq. (10)
To maximize the (log) likelihood efficiently, one needs to compute the gradient of the likelihood function, which is straightforwardly achieved by backpropagating the gradient using automatic differentiation in step 4.
Note that there is no explicit model of partial observability and we use point estimates for the control signals, therefore the likelihood in Eq. (10) decomposes in independent factors given states and controls. In contrast, in our approach, when incorporating partial observability and missing controls, one has to compute approximate likelihood contributions within a forward pass.
C Algorithm to compute an approximate likelihood for IOC
Algorithm 1 Approximate likelihood computation
Output: Approximate likelihood of parameters  $p(\pmb{x}_{1:T} \mid \pmb{\theta})$   
Input: Parameters  $\pmb{\theta}$ , Data  $\pmb{x}_{1:T}$ , Model  $f, h$   
1: Determine the policy  $\pi$  using iLQG  
2: Determine the belief dynamics  $\beta$  using the EKF  
3: for  $t$  in  $\{1, \dots, T-1\}$  do  
4: Compute  $p(\pmb{x}_{t+1}, \pmb{b}_{t+1} \mid \pmb{x}_{1:t})$  using Eq. (4)  
5: Update  $p(\pmb{b}_{t+1} \mid \pmb{x}_{1:t+1})$  using Eq. (5)  
6: Obtain  $p(\pmb{x}_{t+1} \mid \pmb{x}_{1:t})$  using Eq. (3)  
7: end for
D Inverse optimal control derivations
We start with defining the joint dynamics of states and estimates, which returns the next state and estimate as a function of the previous state and estimate and the noises,
$$
\left[ \begin{array}{c} \boldsymbol {x} _ {t + 1} \\ \mathbf {b} _ {t + 1} \end{array} \right] = g (\boldsymbol {x} _ {t}, \mathbf {b} _ {t}, \boldsymbol {v} _ {t}, \boldsymbol {w} _ {t}, \boldsymbol {\xi} _ {t}).
$$
To do so, we insert the policy and observation function into the system and belief dynamics, giving
$$
\boldsymbol {x} _ {t + 1} = f (\boldsymbol {x} _ {t}, \pi_ {t} (\mathbf {b} _ {t}, \boldsymbol {\xi} _ {t}), \boldsymbol {v} _ {t}),
$$
$$
\mathbf {b} _ {t + 1} = \beta_ {t} \left(\mathbf {b} _ {t}, \pi_ {t} \left(\mathbf {b} _ {t}, \boldsymbol {\xi} _ {t}\right), h \left(\boldsymbol {x} _ {t}, \boldsymbol {w} _ {t}\right)\right).
$$
The individual factors of the likelihood function can be determined as
$$
p (\boldsymbol {x} _ {t + 1} \mid \boldsymbol {x} _ {1: t}) = \int p (\boldsymbol {x} _ {t + 1}, \mathbf {b} _ {t + 1} \mid \boldsymbol {x} _ {1: t}) \mathrm {d} \mathbf {b} _ {t + 1},
$$
where  $p(\mathbf{x}_{t + 1},\mathbf{b}_{t + 1}\mid \mathbf{x}_{1:t})$  is given by marginalizing over the current belief  $\mathbf{b}_t$  as
$$
\begin{array}{l} p \left(\boldsymbol {x} _ {t + 1}, \mathbf {b} _ {t + 1} \mid \boldsymbol {x} _ {1: t}\right) = \int p \left(\boldsymbol {x} _ {t + 1}, \mathbf {b} _ {t + 1}, \mathbf {b} _ {t} \mid \boldsymbol {x} _ {1: t}\right) d \mathbf {b} _ {t} \\ = \int p \left(\mathbf {x} _ {t + 1}, \mathbf {b} _ {t + 1} \mid \mathbf {x} _ {1: t}, \mathbf {b} _ {t}\right) p \left(\mathbf {b} _ {t} \mid \mathbf {x} _ {1: t}\right) d \mathbf {b} _ {t} \\ = \int p (\boldsymbol {x} _ {t + 1}, \mathbf {b} _ {t + 1} \mid \boldsymbol {x} _ {t}, \mathbf {b} _ {t}) p (\mathbf {b} _ {t} \mid \boldsymbol {x} _ {1: t}) d \mathbf {b} _ {t}. \\ \end{array}
$$
To derive a tractable approximation of the distribution  $p(\mathbf{x}_{t + 1},\mathbf{b}_{t + 1}\mid \mathbf{x}_{1:t})$ , we model the initial belief of the agent's state estimate  $p(\mathbf{b}_1)$  as a Gaussian and linearize the joint dynamics function, leading to a Gaussian approximation of the desired quantity, i.e.,  $p(\mathbf{x}_{t + 1},\mathbf{b}_{t + 1}\mid \mathbf{x}_{1:t})\approx \mathcal{N}(\boldsymbol {\mu}_t,\boldsymbol {\Sigma}_t)$ .
First, we apply a first-order Taylor expansion of the joint dynamics  $g$  around the observed state  $\boldsymbol{x}_t$  and the mean of the belief  $\mu_t^{(b)}$  and the noises:
$$
g (\pmb {x} _ {t}, \mathbf {b} _ {t}, \pmb {v} _ {t}, \pmb {w} _ {t}, \pmb {\xi} _ {t}) \approx g (\pmb {x} _ {t}, \mu_ {t} ^ {(b)}, 0, 0, 0) + J _ {\mathbf {b}} (\mathbf {b} _ {t} - \mu_ {t} ^ {(b)}) + J _ {\pmb {v}} \pmb {v} _ {t} + J _ {\pmb {w}} \pmb {w} _ {t} + J _ {\pmb {\xi}} \pmb {\xi} _ {t},
$$
where  $J_{\bullet}$  denotes the Jacobian of  $g$  w.r.t.  $\bullet$ , evaluated at  $(x_{t},\mu_{t}^{(b)},0,0,0)$ .
To derive an explicit representation of the Jacobians, we insert the filtering and control law obtained by the Kalman filter and maximum entropy iLQG controller and find
$$
\begin{array}{l} g (\boldsymbol {x} _ {t}, \mathbf {b} _ {t}, \boldsymbol {v} _ {t}, \boldsymbol {w} _ {t}, \boldsymbol {\xi} _ {t}) = \left[ \begin{array}{c} f (\boldsymbol {x} _ {t}, \pi_ {t} (\mathbf {b} _ {t}, \boldsymbol {\xi} _ {t}), \boldsymbol {v} _ {t}) \\ \beta_ {t} (\mathbf {b} _ {t}, \pi_ {t} (\mathbf {b} _ {t}, \boldsymbol {\xi} _ {t}), h (\boldsymbol {x} _ {t}, \boldsymbol {w} _ {t})) \end{array} \right] \\ = \left[ \begin{array}{c} f (\boldsymbol {x} _ {t}, \boldsymbol {u} _ {t}, \boldsymbol {v} _ {t}) \\ \beta_ {t} (\boldsymbol {b} _ {t}, \boldsymbol {u} _ {t}, \boldsymbol {y} _ {t}) \end{array} \right], \\ \end{array}
$$
$$
\begin{array}{l} \boldsymbol {y} _ {t} = h (\boldsymbol {x} _ {t}, \boldsymbol {w} _ {t}), \\ \boldsymbol {u} _ {t} = \pi_ {t} (\mathbf {b} _ {t}, \boldsymbol {\xi} _ {t}) = L _ {t} (\mathbf {b} _ {t} - \bar {\boldsymbol {x}} _ {1: T}) + \mathbf {m} _ {t} + \bar {\boldsymbol {u}} _ {1: T} - \tilde {L} _ {t} \boldsymbol {\xi} _ {t}, \\ \beta_ {t} \left(\mathbf {b} _ {t}, \boldsymbol {u} _ {t}, \boldsymbol {y} _ {t}\right) = f \left(\mathbf {b} _ {t}, \boldsymbol {u} _ {t}, 0\right) + K _ {t} \left(\boldsymbol {y} _ {t} - h \left(\mathbf {b} _ {t}, 0\right)\right), \\ \end{array}
$$
leading to the equations
$$
\begin{array}{l} \mathbb {J} _ {\mathbf {b}} = \left[ \begin{array}{c} \nabla_ {\boldsymbol {u}} f (\boldsymbol {x} _ {t}, \boldsymbol {u} _ {t}, 0) \nabla_ {\mathbf {b}} \boldsymbol {u} _ {t} \\ \nabla_ {\mathbf {b}} \beta_ {t} (\mathbf {b} _ {t}, \boldsymbol {u} _ {t}, \boldsymbol {y} _ {t}) + \nabla_ {\boldsymbol {u}} \beta_ {t} (\mathbf {b} _ {t}, \boldsymbol {u} _ {t}, \boldsymbol {y} _ {t}) \nabla_ {\mathbf {b}} \boldsymbol {u} _ {t} \end{array} \right] \\ = \left[ \begin{array}{c} \nabla_ {\pmb {u}} f (\pmb {x} _ {t}, \pmb {u} _ {t}, 0) L _ {t} \\ \nabla_ {\pmb {x}} f (\pmb {b} _ {t}, \pmb {u} _ {t}, 0) - K _ {t} \nabla_ {\pmb {x}} h (\pmb {b} _ {t}, 0) + \nabla_ {\pmb {u}} f (\pmb {b} _ {t}, \pmb {u} _ {t}, 0) L _ {t} \end{array} \right], \\ \end{array}
$$
$$
\mathbb {J} _ {\boldsymbol {v}} = \left[ \begin{array}{c} \nabla_ {\boldsymbol {v}} f (\boldsymbol {x} _ {t}, \boldsymbol {u} _ {t}, 0) \\ \nabla_ {\boldsymbol {v}} f (\boldsymbol {b} _ {t}, \boldsymbol {u} _ {t}, 0) \end{array} \right],
$$
$$
\mathbb {J} _ {\boldsymbol {w}} = \left[ \begin{array}{c} 0 \\ \nabla_ {h} \beta_ {t} (\boldsymbol {b} _ {t}, \boldsymbol {u} _ {t}, \boldsymbol {y} _ {t}) \nabla_ {\boldsymbol {w}} h (\boldsymbol {b} _ {t}, 0) \end{array} \right] = \left[ \begin{array}{c} 0 \\ - K _ {t} \nabla_ {\boldsymbol {w}} h (\boldsymbol {b} _ {t}, 0) \end{array} \right],
$$
$$
\mathbb {J} _ {\boldsymbol {\xi}} = \left[ \begin{array}{c} \nabla_ {\boldsymbol {u}} f (\boldsymbol {x} _ {t}, \boldsymbol {u} _ {t}, 0) \nabla_ {\boldsymbol {\xi}} \boldsymbol {u} _ {t} \\ \nabla_ {\boldsymbol {u}} f (\mathbf {b} _ {t}, \boldsymbol {u} _ {t}, 0) \nabla_ {\boldsymbol {\xi}} \boldsymbol {u} _ {t} \end{array} \right] = \left[ \begin{array}{c} \nabla_ {\boldsymbol {u}} f (\boldsymbol {x} _ {t}, \boldsymbol {u} _ {t}, 0) L _ {t} \\ \nabla_ {\boldsymbol {u}} f (\mathbf {b} _ {t}, \boldsymbol {u} _ {t}, 0) L _ {t} \end{array} \right].
$$
Propagating the Gaussian belief over the agent's state estimate,  $p(\mathbf{b}_t \mid x_{1:t})$  through the linearized dynamics model (Eq. (4)) can be done by applying standard identities for linear transformations of Gaussian random variables [64] and gives
$$
p \left(\boldsymbol {x} _ {t + 1}, \mathbf {b} _ {t + 1} \mid \boldsymbol {x} _ {1: t}\right) \approx \mathcal {N} \left(\mu_ {t}, \Sigma_ {t}\right),
$$
with  $\mu_t = g(\pmb{x}_t, \mu_t^{(b)}, 0, 0, 0)$  and  $\Sigma_t = \mathbb{J}_{\mathbf{b}} \Sigma_t^{(b)} \mathbb{J}_{\mathbf{b}}^T + \mathbb{J}_{\mathbf{v}} \mathbb{J}_{\mathbf{v}}^T + \mathbb{J}_{\mathbf{w}} \mathbb{J}_{\mathbf{w}}^T + \mathbb{J}_{\boldsymbol{\xi}} \mathbb{J}_{\boldsymbol{\xi}}^T$ . Marginalization over  $\mathbf{b}_{t+1}$  gives the desired likelihood factor, while conditioning on  $\pmb{x}_{t+1}$  gives the belief statistic for the following time step (see Section 3.1).
E Special case: full observability
If the state  $\boldsymbol{x}_t$  is fully observable to the agent, the problem is simplified significantly. The control problem from the agent's perspective (Fig. S1, left) can be solved by applying iLQG [8] to the observed states directly (Section 2.2)
$$
\boldsymbol {u} _ {t} = \pi_ {t} (\boldsymbol {x} _ {t}, \boldsymbol {\xi} _ {t}),
$$
This also simplifies the IOC problem from the researcher's perspective because evaluating the approximate likelihood becomes straightforward as we do not need to marginalize over the agent's belief.
Recall that the likelihood can be decomposed as
$$
p \left(\boldsymbol {x} _ {1: T} \mid \theta\right) = p \left(\boldsymbol {x} _ {0} \mid \theta\right) \prod_ {t = 0} ^ {T - 1} p \left(\boldsymbol {x} _ {t + 1} \mid \boldsymbol {x} _ {1: t}, \theta\right)
$$
and the dynamics are given as
Figure S1: Left: decision network from the agent's perspective [following the notation used in 9]. The agent directly observes the state  $\boldsymbol{x}_t$ . At each time step, they perform a control  $\boldsymbol{u}_t$  and incur a cost  $c_t$ . Right: probabilistic graphical model from the researcher's perspective, who observes a trajectory  $\boldsymbol{x}_{1:T}$  from an agent. The control signals  $\boldsymbol{u}_t$  are not directly observed.
$$
\boldsymbol {x} _ {t + 1} = f (\boldsymbol {x} _ {t}, \boldsymbol {u} _ {t}, \boldsymbol {v} _ {t}) = f (\boldsymbol {x} _ {t}, \pi (\boldsymbol {x} _ {t}, \boldsymbol {\xi} _ {t}), \boldsymbol {v} _ {t}).
$$
We can approximate the likelihood contribution at each time step as
$$
p \left(\boldsymbol {x} _ {t + 1} \mid \boldsymbol {x} _ {1: t}, \theta\right) \approx \mathcal {N} \left(\mu_ {t}, \Sigma_ {t}\right),
$$
$$
\mu_ {t} = f \left(\boldsymbol {x} _ {t}, \pi \left(\boldsymbol {x} _ {t}, 0\right), 0\right)
$$
$$
\Sigma_ {t} = \mathbb {J} _ {\boldsymbol {v}} \mathbb {J} _ {\boldsymbol {v}} ^ {T} + \mathbb {J} _ {\boldsymbol {\xi}} \mathbb {J} _ {\boldsymbol {\xi}} ^ {T},
$$
where  $\mathbb{J}_{\bullet}$  denotes the Jacobian of  $f$  w.r.t.  $\bullet$ , evaluated at  $(x_{t},\pi (x_{t},0),0)$ :
$$
\mathbb {J} _ {\boldsymbol {v}} = \nabla_ {\boldsymbol {v}} f (\boldsymbol {x} _ {t}, \boldsymbol {u} _ {t}, 0),
$$
$$
\mathbb {J} _ {\boldsymbol {\xi}} = \nabla_ {\boldsymbol {u}} f (\boldsymbol {x} _ {t}, \boldsymbol {u} _ {t}, 0) \nabla_ {\boldsymbol {\xi}} \boldsymbol {u} _ {t} = \nabla_ {\boldsymbol {u}} f (\boldsymbol {x} _ {t}, \boldsymbol {u} _ {t}, 0) L _ {t}.
$$
F Implementation
We provide a flexible framework for defining non-linear stochastic dynamical systems of the form introduced in Section 2.1 by implementing the dynamics, observation function, and cost function. The optimal estimation control methods based on iterative linearization are implemented using the automatic differentiation library jax [65], so that Jacobians and Hessians for the linearization of the dynamics and quadratization of the cost do not have to be specified manually. Furthermore, the Jacobians needed for linear-Gaussian approximation in the IOC likelihood (Section 3.2) are also computed using automatic differentiation. Gradient-based maximization of the log likelihood was performed using the L-BFGS-B scipy [66] from the jaxopt library [67]. The implementation is provided in the supplementary material and will be made public upon publication.
For computations, we used an Intel Xeon Platinum 9242 Processor, using 1 core per run. The mean run times for computing maximum likelihood estimates (partially observable setting) were as follows:
- Reaching: 140 s
- Navigation: 64 s
Pendulum: 31 s
CartPole: 115 s
G Estimating controls for linearization
For every evaluation of the likelihood function as described in Section 3.1, we would naively need to solve the forward control problem once by running the iLQG algorithm given the current parameter values. This is computationally inefficient due to the iterative procedure of iLQG, which requires multiple forward and backward passes. Instead of performing the iterative procedure, which computes the locally optimal control law  $\{L_{1:T-1}, \mathbf{m}_{1:T-1}\}$  around a nominal trajectory  $\{\bar{x}_{1:T}, \bar{u}_{1_T-1}\}$ , then computes the new nominal trajectory given the current control law and so on, we realize that in the IOC setting, we already have observed the actual trajectory  $x_{1:T}$  performed by the agent. To make solving the forward control problem more efficient, we only compute the locally optimal control law once around this trajectory. This would require the control signals  $u_{1:T}$  to be given, but they are unobserved in our problem setting. To obtain a proxy for the controls for the purposes of linearization, we solve for the controls given the trajectory  $x_{1:T}$  in the system dynamics equation  $x_{t+1} = f(x_t, u_t, 0)$  using the Gauss-Newton method for non-linear least squares as implemented in jaxopt, i.e.,
$$
\hat {\boldsymbol {u}} _ {t} = \underset {\boldsymbol {u}} {\arg \max } \left(\boldsymbol {x} _ {t + 1} - f \left(\boldsymbol {x} _ {t}, \boldsymbol {u}, 0\right)\right) ^ {2}. \tag {11}
$$
H Efficient gradient computation
In the case where the full state is not observable, one cannot linearize around the given trajectory as described in Section 3.3. The forward optimal control problem is commonly solved by starting with a randomly initialized nominal trajectory and iterating between computing the locally optimal control law and linearization until convergence. To compute gradients, a naive approach would be to differentiate through the loop of the forward procedure. In our experiments, we found, however, that this approach is numerically unstable and does not yield correct gradients (see Fig. S2). Here, we provide two approaches to compute gradients in this case.
H.1 Differentiating through the last iteration
In this approach, one first solves the forward problem as usual to determine the linearization. One then executes one additional backward pass for the optimal controller, which one can differentiate using automatic differentiation.
H.2 Applying the implicit function theorem
Determining the optimal linearization and controller is essentially a fixed point problem. To differentiate through the optimal linearization and controller, we apply the implicit function theorem [68, 69].
We assume the forward equation to determine the linearization is given by
$$
\boldsymbol {x} _ {t + 1} = \hat {f} (\boldsymbol {x} _ {t}, \boldsymbol {u} _ {t}, \boldsymbol {\theta}),
$$
$$
\boldsymbol {u} _ {t} = \hat {\pi} _ {t} (\boldsymbol {x} _ {t}, \boldsymbol {s} _ {t}, \boldsymbol {\theta})
$$
where  $\theta$  are the parameters we want to differentiate with respect to and  $s_t$  is a vector determining the optimal controller, such that it can be determined via a backward pass with
$$
\boldsymbol {s} _ {t - 1} = \gamma (\boldsymbol {x} _ {t}, \boldsymbol {u} _ {t}, \boldsymbol {s} _ {t}, \boldsymbol {\theta}).
$$
For the LQG controller,  $s_t$  would consist of the (vectorized) matrix characterizing the optimal value function.
The fixed point problem of jointly determining the optimal linearization and controller can be formulated as determining  $X \coloneqq [u_{1}, x_{2}, u_{2}, \ldots, u_{T-1}, x_{T}, s_{1}, \ldots, s_{T}]$  such that
$$
\begin{array}{l} X = \eta (X, \boldsymbol {\theta}) := [ \hat {\pi} _ {1} (\boldsymbol {x} _ {1}, \boldsymbol {s} _ {1}, \boldsymbol {\theta}), \hat {f} (\boldsymbol {x} _ {1}, \boldsymbol {u} _ {1}, \boldsymbol {\theta}), \hat {\pi} _ {2} (\boldsymbol {x} _ {2}, \boldsymbol {s} _ {2}, \boldsymbol {\theta}), \dots , \hat {f} (\boldsymbol {x} _ {T - 1}, \boldsymbol {u} _ {T - 1}, \boldsymbol {\theta}), \\ \left. \gamma \left(\boldsymbol {x} _ {2}, \boldsymbol {u} _ {2}, \boldsymbol {s} _ {2}, \boldsymbol {\theta}\right), \dots , \gamma \left(\boldsymbol {x} _ {T}, \boldsymbol {u} _ {T}, \boldsymbol {s} _ {T}, \boldsymbol {\theta}\right) \right]. \\ \end{array}
$$
The optimal solution of the fixed point problem satisfies
$$
\kappa (X ^ {*} (\boldsymbol {\theta}), \boldsymbol {\theta}) := X ^ {*} (\boldsymbol {\theta}) - \eta (X ^ {*} (\boldsymbol {\theta}), \boldsymbol {\theta}) = 0,
$$
where  $X^{*}(\pmb{\theta})$  is a function yielding the optimal solution depending on  $\pmb{\theta}$ , i.e.,  $X^{*}: \mathbb{R}^{p} \to \mathbb{R}^{(n + u + s)t}$  where  $s$  is the number of elements of  $\pmb{s}_t$ .
The implicit function theorem [68] gives, that for  $(X_0,\pmb{\theta}_0)$  satisfying  $\kappa (X_0,\pmb {\theta}_0) = 0$  with the requirement that  $\kappa$  is continuously differentiable and the Jacobian  $\partial_{\pmb{\theta}}\kappa (X_0,\pmb {\theta}_0)$  is a square invertible matrix, then there exists a function  $X^{*}(\cdot)$  on a neighborhood of  $\pmb{\theta}_0$  such that  $X^{*}(\pmb{\theta}_{0}) = X_{0}$ . Furthermore, for all  $\pmb{\theta}$  in this neighborhood, we have that  $\kappa (X^{*}(\pmb {\theta}),\pmb {\theta}) = 0$  and  $\partial X^{*}(\pmb {\theta})$  exists.
By applying the chain rule, the Jacobian  $\partial X^{*}(\pmb{\theta})$  needs to satisfy
$$
\begin{array}{l} 0 = \partial_ {1} \kappa (X ^ {*} (\boldsymbol {\theta}), \boldsymbol {\theta}) \partial X ^ {*} (\boldsymbol {\theta}) + \partial_ {2} \kappa (X ^ {*} (\boldsymbol {\theta}), \boldsymbol {\theta}) \\ = (I - \partial_ {1} \eta (X ^ {*} (\boldsymbol {\theta}), \boldsymbol {\theta})) \partial X ^ {*} (\boldsymbol {\theta}) - \partial_ {2} \eta (X ^ {*} (\boldsymbol {\theta}), \boldsymbol {\theta}). \\ \end{array}
$$
$\partial_1 \eta(X^*(\pmb{\theta}), \pmb{\theta}))$  is given by a sparse matrix, as each element of  $\eta(X, \pmb{\theta})$  only depends on few elements of  $X$ . The desired gradient  $\partial X^*(\pmb{\theta})$  (or vector-Jacobian products with it) can then be computed using linear system solvers [69].
H.3 Results with different gradient computation methods
We ran the MLE for the partially observable reaching task with different gradient computation methods. The results are shown in Fig. S2. While unrolling the loop led to numerical problems, so the estimates are far off the true values, all other methods led to quite similar results. The mean run times were as follows:
Unrolled loop: 141 s
- Differentiating through the last iteration: 103 s
- Implicit differentiation: 243 s
Data-based linearization: 129 s
I Hyperparameters