SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Throughout the experiments, we have used the following hyperparameters:
Table 2: Hyperparameters
Data set size (number of trajectories)	50
Number of datasets per evaluation	100
Number of time steps (T)	50 (reaching, navigation, pendulum, light-dark), 200 (cartpole)
Optimizer	L-BFGS-B (scipy wrapper from jaxopt)
Maximum entropy temperature	10-6(reaching, navigation), 10-3(pendulum, cartpole), 10-5(light-dark)
Optimizer restarts	50
J Tasks
J.1 Reaching task with biomechanical arm model
We implemented the non-linear biomechanical model for arm movements from Todorov and Li [8] and its partially observed version from Li and Todorov [43], which are described in more detail in the PhD thesis of Li [70]. The dynamics describe the movement of a two-link arm, which can be controlled by applying torques to the two joints. The task is to move the hand to a target location, as defined in the cost function
$$
J = \left\| \mathbf {e} _ {T} - \mathbf {e} ^ {*} \right\| ^ {2} + c _ {v} \left\| \dot {\mathbf {e}} _ {T} \right\| ^ {2} + c _ {a} \sum_ {t = 1} ^ {T - 1} \left\| \mathbf {u} _ {t} \right\| ^ {2},
$$
Figure S2: Maximum likelihood parameter estimates for the partially observable reaching task with different gradient computation methods
where  $\mathbf{e}_T$  is the position of the hand at the final time step in Cartesian coordinates,  $\mathbf{e}^{\star}$  is the position of the target, and  $\dot{\mathbf{e}}_T$  is the final velocity of the hand. That is, the task is to bring the hand to the target at the final time step  $(T = 50)$  and minimize the final velocity while making as little control as possible along the way. The parameters  $c_v$  and  $c_a$  trade off, how important the velocity and control parts of the cost function are relative to the main goal of reaching the target. Additionally, we parameterize the signal-dependent variability of the motor system with the parameter  $\sigma_m$ , which defines how strongly the variability is scaled by the control magnitude.
J.2 Navigation task
We implemented a simple navigation task in which an agent walks towards a target. The state  $\pmb{x}_t = [x_t, y_t, \theta_t, \omega_t]^T$  consists of the position in the horizontal plane, the heading angle, and the speed in the heading direction. The dynamics are
$$
f (\boldsymbol {x}, \boldsymbol {u}, \boldsymbol {v}) = \boldsymbol {x} + d t \left[ \begin{array}{c} \cos (\theta) \cdot \omega \\ \sin (\theta) \cdot \omega \\ u _ {1} + \sigma_ {m} u _ {1} v _ {1} \\ u _ {2} + \sigma_ {m} u _ {2} v _ {2} \end{array} \right],
$$
which means that the agent controls the velocity of the heading angle and the acceleration in heading direction. There is control-dependent noise, whose strength is determined by the parameter  $\sigma_{m}$ . The observation model is
$$
h (\boldsymbol {x}, \boldsymbol {w}) = \left[ \begin{array}{c} \sqrt {(x - k _ {1}) ^ {2} + (y - k _ {2}) ^ {2}} \\ \tan^ {- 1} (y - k _ {2}, x - k _ {1}) \\ \omega \end{array} \right] + \sigma_ {o} \boldsymbol {w},
$$
where  $\mathbf{k} = [k_1, k_2]^T$  is the position of the target and  $\sigma_o$  determines the magnitude of the observation noise. The cost function is
$$
J = \left(x _ {T} - k _ {1}\right) ^ {2} + \left(y _ {T} - k _ {2}\right) ^ {2} + c _ {v} \omega_ {T} + \sum_ {t = 1} ^ {T - 1} c _ {a} \boldsymbol {u} _ {t} ^ {T} \boldsymbol {u} _ {t},
$$
with parameters determining the cost of controls  $(c_{a})$  and the cost of the final velocity  $(c_{v})$ . The time horizon was  $T = 50$ .
J.3 Classic control tasks
We evaluate our method on two classic continuous control tasks. Specifically, we build upon the Pendulum and Cart Pole environments from the gym library [54]. Because these environments are not stochastic in their standard implementations and are typically used in an infinite-horizon reinforcement learning setting, we made the following changes to the environments. To account for the finite-horizon setting we are considering in this work, the state-dependent costs associated with the task goal are applied only to the final time step, while control-dependent costs are applied at every time step along the trajectory. To make the problems stochastic, we have added noise on the dynamics and the observation function.
J.3.1 Pendulum
The task is to make a pendulum, which is attached to a fixed point on one side, swing into an upright position, i.e. to reach an angle  $\theta = 0$ . The pendulum starts at the bottom  $(\theta = \pi)$  and can be controlled by applying a torque to the free end of the pendulum. We added control-dependent noise to the dynamics, where the parameter  $\sigma_{m}$  controls, how strongly a standard normal noise is scaled by the magnitude of the torque. In addition to this motor noise parameter, the task has two other free parameters: the cost of controls  $c_{a}$  and the cost of the final velocity  $c_{v}$ . In the partially observable version of the task, the agent receives the observation via the non-linear stochastic observation model  $\mathbf{y}_t = [\sin \theta_t, \cos \theta_t, \dot{\theta}_t]^T + 0.1 \mathbf{w}_t$ . The time horizon was  $T = 50$ .
J.3.2 Cart pole
In the Cart Pole task, a pole is attached to a cart that can be moved left or right by applying a continuous force. In our version of the task, the goal is to move the cart from the horizontal position  $x = 0$  to  $x = 1$  while balancing the pole in an upright position. Again, we add control-dependent noise to the force, parameterized the strength of the linear control-dependence  $\sigma_{m}$ . As above, the other two parameters are the control cost  $c_{a}$  and the cost of the final velocity  $c_{v}$ . In the partially observed version of the task, the agent receives a noisy observation with  $\mathbf{y}_t = \mathbf{x}_t + \boldsymbol{w}_t$ . The time horizon was  $T = 200$ .
J.4 Light-dark domain
We use a slightly modified version of the light-dark domain, which was originally introduced by Platt et al. [55]. We adapted the light-dark domain's motion model to include signal-dependent noise on the control
$$
\boldsymbol {x} _ {t + 1} = \boldsymbol {x} _ {t} + d t \boldsymbol {u} _ {t} + 0. 1 \boldsymbol {u} _ {t} \odot \boldsymbol {v} _ {t}
$$
The standard deviation of the perceptual uncertainty varies with the horizontal distance to the light source:
$$
\boldsymbol {y} _ {t} = \boldsymbol {x} _ {t} + \sigma | x _ {t, 1} - 5 | \boldsymbol {w} _ {t}.
$$
The cost function is composed of three terms. First, the agent should minimize the squared distance to the target  $\mathbf{p}$  at the final time step  $T$ . Second, the agent should minimize the squared control signals  $\boldsymbol{u}_t$ . And finally, the agent should minimize the horizontal distance to the light source:
$$
J = \underbrace {\left(\boldsymbol {x} _ {T} - \mathbf {p}\right) ^ {2}} _ {\text {f i n a l c o s t}} + \underbrace {\sum_ {t = 1} ^ {T - 1} \frac {1}{2} \boldsymbol {u} _ {t} ^ {2} + c \left(x _ {t , 1} - 5\right) ^ {2}} _ {\text {r u n n i n g c o s t}}.
$$
The time-horizon was set to  $T = 50$ .
We simulated the behavior of two agents in the light-dark domain in Fig. S3. First, the partially observable version of iLQG [43] takes the expected belief covariance into account for the computation of the policy. This agent moves towards the light source before approaching the target. Adding another cost term  $c$  that expresses a preference to be close to the light source does not change much about this behavior (top row). Second, the fully observable version of iLQG [8] computes the policy irrespective of the belief. To apply it in the partially observable setting, we simply combine it with an EKF for the state estimation. This agent does not move towards the light source by default and as a result, the belief uncertainty is higher. The agent only moves towards the light source before approaching the target if we add the additional cost term  $c > 0$ . One can also see in Fig. S3 that the agent's belief is more uncertain if it does not move to the light source before approaching the target. This results in a higher variability around the final position (inset plots in Fig. S3).
Figure S3: Trajectories and beliefs in the light-dark domain. Partially observable (top) vs. fully observable (bottom) iLQG in the light dark domain with different values for the cost parameter  $c$ , which expresses an inherent desire to be near the light source. The circles indicate the agent's belief covariance (2 standard deviations). One can also see that the variability around the final position is higher for agents that do not move towards the light source before approaching the target (inset plots).
K Additional results
The parameter estimates and true values are provided in Fig. S4, Fig. S5, Fig. S6.
K.1 Results for all tasks in the partially observable setting
Figure S4: Maximum likelihood parameter estimates for the Pendulum task (partially observable cases).
Figure S5: Maximum likelihood parameter estimates for the Cart Pole task (partially observable).
Figure S6: Maximum likelihood parameter estimates for the navigation task (partially observable).
K.2 Results for all tasks in the fully observable setting
Figure S7: Maximum likelihood parameter estimates for the reaching task (fully observable).
Figure S8: Maximum likelihood parameter estimates for the pendulum task (fully observable).
Figure S9: Maximum likelihood parameter estimates for the cart pole task (fully observable).
Figure S10: Maximum likelihood parameter estimates for the navigation task (fully observable).
K.3 Results for baseline with given control signals
Figure S11: Maximum likelihood parameter estimates for the baseline with given control signals (fully observable)
K.4 Quantitative results for the light-dark domain
As for the other tasks, we generated 100 sets of parameters sampled from uniform distributions in the following ranges
- horizontal target position  $p$ :  $(-1, 1)$
state-dependent cost parameter  $c$  ..  $(10^{-2},1)$
- perceptual uncertainty  $\sigma$ :  $(10^{-2}, 1)$
and generated a dataset of 50 trajectories using iLQG [43] for each set of parameters. For each set, we then performed inference using our method and the baseline.
While the baseline accurately infers the target position, our method shows more variability in its estimates of the target position. The estimates of the cost term  $c$  are comparable for both methods. However, the baseline completely fails to estimate the perceptual uncertainty, which our method manages to infer relatively accurately.
Figure S12: Maximum likelihood estimates for our method (top) and the baseline (bottom) in the light-dark domain.