SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Dominik Straub*
Centre for Cognitive Science Technische Universität Darmstadt dominik.straub@tu-darmstadt.de
Heinz Koeppl
Centre for Cognitive Science Technische Universität Darmstadt heinz.koepl@tu-darmstadt.de
Matthias Schultheis *
Centre for Cognitive Science Technische Universität Darmstadt matthias.schultheis@tu-darmstadt.de
Constantin A. Rothkopf
Centre for Cognitive Science Technische Universitat Darmstadt Constantin.rothkopf@tu-darmstadt.de
Inverse optimal control can be used to characterize behavior in sequential decision-making tasks. Most existing work, however, is limited to fully observable or linear systems, or requires the action signals to be known. Here, we introduce a probabilistic approach to inverse optimal control for partially observable stochastic non-linear systems with unobserved action signals, which unifies previous approaches to inverse optimal control with maximum causal entropy formulations. Using an explicit model of the noise characteristics of the sensory and motor systems of the agent in conjunction with local linearization techniques, we derive an approximate likelihood function for the model parameters, which can be computed within a single forward pass. We present quantitative evaluations on stochastic and partially observable versions of two classic control tasks and two human behavioral tasks. Importantly, we show that our method can disentangle perceptual factors and behavioral costs despite the fact that epistemic and pragmatic actions are intertwined in sequential decision-making under uncertainty, such as in active sensing and active learning. The proposed method has broad applicability, ranging from imitation learning to sensorimotor neuroscience.
Inverse optimal control (IOC) is the problem of inferring an agent's cost function and other properties of their internal model from behavior. While IOC has been a fundamental task in artificial intelligence and machine learning, particularly reinforcement learning (RL) and robotics, it has widespread applicability in several scientific fields including behavioral economics, psychology, and neuroscience. For example, in cognitive science and sensorimotor neuroscience, optimal control models have explained key properties of behavior, such as speed-accuracy trade-offs [1] or the minimum intervention principle [2]. But, while researchers usually build an optimal control model and compare its predictions to behavior, certain parameters of the agent's internal processes are typically unknown. For example, an agent might have uncertainty about their perception or experience intrinsic costs of behavior. These parameters are different between individuals and inferring them from observed behavior can help to understand internal tradeoffs between behavioral goals, perceptual and cognitive
processes, and predict behavior under novel conditions. Applying IOC in these domains poses several challenges that make most previous methods not viable.
First, most IOC methods assume the agent's control signals to be known. This assumption, while convenient in simulations or robotics, where the control signals may be easily quantified, does not hold in many other real-world applications. In transfer learning or behavioral experiments, the control signals are internal quantities of an animal or human, e.g., neural activity or muscle activations, and are therefore not straightforwardly observable. Thus, we consider the scenario where a researcher has observations of the system's state only, i.e., measurements of behavior.
Second, most IOC methods model the variability of the agent using a stochastic policy in a maximum causal entropy formulation [MCE; 3]. Behavioral variability in biological systems, however, is known to arise from multiple distinct sources [4]. There is noise in the sensory system, which makes the state of the world partially observable, and in the motor system. In sensorimotor neuroscience, the uncertainty in the sensory and motor systems can be characterized quantitatively by formulating accurate models, which are helpful to understand behavioral variability [5].
Third, many IOC methods are based on matching feature expectations of the cost function between the model and data [3], and are thus not easily adapted to infer other model parameters. In a behavioral experiment, however, researchers are often interested in inferring the noise characteristics of the sensorimotor system or other properties of the agent's internal model besides the cost function [6].
Fourth, while the theory of linear-quadratic-Gaussian (LQG) control [7] is suited to deal with the issues above, many real-world systems are not captured by the LQG assumptions. First, the dynamics may be non-linear, e.g., in robotics and motor control when controlling joint angles in a kinematic chain. Second, the variability of the system may not be captured by normal distributions, e.g., in sensorimotor control, where the standard deviation of sensory and control signals scales with their means. While iterative methods for solving the optimal control problem such as iterative variants of LQG [8] exist, here we consider the corresponding inverse problem.
To address these issues, we adopt a probabilistic perspective. We distinguish between the control problem faced by the agent and the IOC problem faced by the researcher. From the agent's perspective, the problem consists of acting in a partially observable Markov decision process (POMDP; Fig. 1 A). We consider the setting of continuous states and controls, stochastic nonlinear dynamics, partial observations, and finite horizon. For this setting, there are efficient approximately optimal solutions to the estimation and control problem (see Section 2). The researcher, on the other hand, is interested in inferring properties of the agent's model and cost function. The IOC problem from their perspective can be formulated using a probabilistic graphical model (Fig. 1 B), in which the state of the system is observed, while variables internal to the agent are latent.
Figure 1: A Decision network from the agent's perspective [notation from 9]. At each time step, the agent receives a noisy observation  $\mathbf{y}_t$  of the state  $\mathbf{x}_t$ , performs a control  $\mathbf{u}_t$ , and incurs a cost  $c_t$ . B Probabilistic graphical model from the researcher's perspective, who observes a trajectory  $\mathbf{x}_{1:T}$  of an agent. Quantities internal to the agent, i.e. observations  $\mathbf{y}_t$ , beliefs  $\mathbf{b}_t$  and control signals  $\mathbf{u}_t$ , are not directly observed.
Here, we unify MCE models, which are agnostic regarding the probabilistic structure causing the observed stochasticity of the agent's policy, with IOC methods that involve an explicit observation model. First, we define the IOC problem where we allow for both: We employ an explicit observation model, but also allow the agent to have additional stochasticity through an MCE policy. Second, we provide a solution to the IOC problem in this setting by approximate filtering of the agent's state estimate via local linearization, which allows marginalizing over these latent variables and deriving an approximate likelihood function for observed trajectories given parameters (Section 3). By maximizing the approximate likelihood function, an estimate of the optimal parameters can be determined. Third, we evaluate our proposed method on two classic control tasks, pendulum and cart pole, and on two human behavioral tasks, navigation and a manual reaching (Section 4.2 - Section 4.3). Fourth, we show that our approach allows disentangling the influences of perceptual uncertainty and behavioral costs on information-seeking behavior in the light-dark domain (Section 4.4).
Inferring costs or utilities from behavior has been of interest for a long time in several scientific fields, such as behavioral economics, psychology, and neuroscience [10-12]. More specific to the problem formulation adopted here, estimating objective functions in the field of control was first investigated by Kalman [13] in the context of deterministic linear systems with quadratic costs. More recent formulations were developed first for discrete state and control spaces under the term inverse reinforcement learning [IRL; 14, 15], including formulations allowing for stochasticity in action selection [16]. In this line, the maximum entropy [ME; 17] and MCE formulation [3] gave rise to many new methods, e.g., for non-linear continuous systems via linearization [18] or importance sampling [19] for fully observable deterministic systems.
IOC methods for stochastic systems have been developed in the setting of affine control dynamics [20, 21]. Arbitrary non-linear stochastic dynamics in the infinite horizon setting have been approached using model-free deep MCE IRL [22, 23]. The latter approaches, however, do not yield interpretable representations, as the cost function is represented by a neural network. Further, past methods based on MCE are limited to estimating cost functions and cannot be used to infer other latent quantities, such as noises or subjective beliefs. The partially observable setting for IOC has previously been addressed for discrete state-action spaces [24] and continuous states with discrete actions [25]. Schmitt et al. [26] addressed systems with linear dynamics and continuous controls for a linear switching observation model. Other work has considered partial observability from the researcher's perspective, e.g., through occlusions [27, 28]. There are some IOC methods which are applicable to partially observable stochastic systems: In our previous work [29] we regarded LQG systems, while the work of Chen and Ziebart [30] can be used to estimate cost functions that depend on the state only. Non-linear dynamics in the infinite-horizon setting and the joint estimation of model parameters have been approached by Kwon et al. [31] by training a policy network as a function of the whole parameter space. This work, however, also assumes the control signals to be given and a stationary, deterministic policy.
Applications of IOC methods range from human locomotion [32] over spatial navigation [33], table tennis [34], to attention switching [35], and target tracking [36]. Other work has been aimed at inferring other properties of control tasks, e.g., the dynamics model [6], learning rules [37], or discount functions [38]. Several subfields of robotics including imitation and apprenticeship learning [39] as well as transfer learning [40] have also employed IOC.
Before we introduce our probabilistic approach to inverse optimal control, we give an overview of the control and filtering problems faced by the agent and algorithms that can be used to solve it. For a summary of our notation in this paper, see Appendix A.
2.1 Partially observable Markov decision processes (POMDPs)
We consider a special case of POMDPs [41, 42], a discrete-time stochastic non-linear dynamical system (Fig. 1 A) with states  $\boldsymbol{x}_t \in \mathbb{R}^n$  following the dynamics equation  $\boldsymbol{x}_{t+1} = f(\boldsymbol{x}_t, \boldsymbol{u}_t, \boldsymbol{v}_t)$ , where  $f$  is the dynamics function,  $\boldsymbol{u}_t \in \mathbb{R}^u$  are the controls and  $\boldsymbol{v}_t \sim \mathcal{N}(0, I)$  is  $v$ -dimensional Gaussian noise. We assume that the agent has only partial observations  $\boldsymbol{y}_t \in \mathbb{R}^m$ , following  $\boldsymbol{y}_t = h(\boldsymbol{x}_t, \boldsymbol{w}_t)$ , with  $h$  the stochastic observation function and  $\boldsymbol{w}_t \sim \mathcal{N}(0, I)$ $w$ -dimensional Gaussian noise. While  $\boldsymbol{v}_t$  and  $\boldsymbol{w}_t$  are standard normal random variables, the system can incorporate general control- and state-dependent noises through non-linear transformations within the dynamics function  $f$  and observation function  $h$ . The agent's goal is to minimize the expected cost over a time horizon  $T \in \mathbb{N}$ , defined by
$$
J = \mathbb {E} \left[ c _ {T} (\pmb {x} _ {T}) + \sum_ {t = 1} ^ {T - 1} c _ {t} (\pmb {x} _ {t}, \pmb {u} _ {t}) \right],
$$
consisting of a final state cost  $c_{T}(\pmb{x}_{T})$  and a cost at each time step  $c_{t}(\pmb{x}_{t},\pmb{u}_{t})$ .
2.2 Iterative linear-quadratic Gaussian (iLQG)
The control problem from Section 2.1 can be solved approximately using iLQG [8, 43]. This method iteratively linearizes the dynamics and quadratizes the costs around a nominal trajectory,  $\{\bar{x}_i,\bar{u}_i\}_{i = 1,\dots,T}$ , with  $\bar{\boldsymbol{x}}_i\in \mathbb{R}^n$ ,  $\bar{u}_i\in \mathbb{R}^u$ , and computes the optimal linear control law,  $\pmb{u}_t = \pi_t(\pmb {x}_t) = L_t(\pmb {x}_t - \bar{\pmb{x}}_t) + \mathbf{m}_t + \bar{\mathbf{u}}_{1:T}$  for the approximated system. The quantities  $L_{t}$  and  $\mathbf{m}_t$  are the control gain and offset, respectively, and determined through a backward pass for the current reference trajectory. In the following iteration, the determined optimal control law is used to generate a new reference trajectory and the process is repeated until the controller converges.
2.3 Maximum causal entropy (MCE) reinforcement learning
The goal of MCE RL is to minimize the expected cost as in Section 2.2, while maximizing the conditional entropy of the stochastic policy  $\Pi_t(\boldsymbol{u}_t \mid \boldsymbol{x}_t)$ , i.e., to minimize  $\mathbb{E}[J(\boldsymbol{x}_{1:T}, \boldsymbol{u}_{1:T}) - \sum_{t=1}^{T-1} \mathcal{H}(\Pi_t(\boldsymbol{u}_t \mid \boldsymbol{x}_t))]$ . This formulation has been used to treat RL as probabilistic inference [44-46] and model the stochasticity of the agent in IRL [17, 3]. The objective of IRL is to maximize the likelihood of states and controls  $\{\boldsymbol{x}_t, \boldsymbol{u}_t\}_{t=1,\dots,N}$ , induced by the maximum entropy policy. The resulting optimal policy is given by the distribution  $\Pi_t(\boldsymbol{u}_t \mid \boldsymbol{x}_t) = \exp(Q_t(\boldsymbol{x}_t, \boldsymbol{u}_t) - V_t(\boldsymbol{x}_t))$ , where  $Q_t$  is the soft Q-function and  $V_t$  the normalization [47]. For general dynamics and reward functions, it is not feasible to compute the soft Q-function exactly. Approximate solutions have been derived using linearization [18], importance sampling [19], or deep function approximation [23]. For linear dynamics and quadratic costs, the optimal policy is 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$  and  $H_t$  result from the optimal LQG controller [48]. More detailed formulas are provided in Appendix B.
2.4 Extended Kalman filter (EKF)
Given the system defined in Section 2.1, the optimal filtering problem is to compute a belief distribution of the current state given past observations, i.e.,  $p(\mathbf{x}_t \mid \mathbf{y}_{1:t-1})$ . For linear-Gaussian systems, the solution is given in closed form and known as the Kalman filter [49]. In case of non-linear systems as in Section 2.1, a Gaussian approximation to the optimal belief can be computed using the extended Kalman filter (EKF) via  $\mathbf{b}_{t+1} = f(\mathbf{b}_t, \mathbf{u}_t, 0) + K_t(\mathbf{y}_t - h(\mathbf{b}_t, 0))$ , where  $\mathbf{b}_t \in \mathbb{R}^n$  denotes the mean of the Gaussian belief  $p(\mathbf{x}_t \mid \mathbf{y}_1, \dots, \mathbf{y}_{t-1})$ . The matrix  $K_t$  denotes the Kalman gain for time  $t$  and is computed by applying the Kalman filter to the system locally-linearized around the nominal trajectory obtained by the approximate optimal control law of iLQG (Section 2.2).
3 Probabilistic inverse optimal control
We consider an agent acting in a partially observable Markov decision process (POMDP) as introduced in Section 2.1. We assume that the agent acts at time  $t$  based on their belief  $\mathbf{b}_t$  about the state of the system  $\boldsymbol{x}_t$ , which evolves according to  $\mathbf{b}_{t + 1} = \beta_t(\mathbf{b}_t,\boldsymbol {u}_t,\boldsymbol {y}_t)$ . While the belief of the agent is defined commonly as a distribution over the true state, here we model  $\mathbf{b}_t$  as a finite-dimensional summary statistics of the distribution, i.e.,  $\mathbf{b}_t\in \mathbb{R}^b$ . The function  $\beta_{t}:\mathbb{R}^{b}\times \mathbb{R}^{u}\times \mathbb{R}^{m}\to \mathbb{R}^{b}$  is called belief dynamics. We further assume that the agent follows a time-dependent policy  $\pi_t:\mathbb{R}^b\times \mathbb{R}^j\to \mathbb{R}^u$ , i.e.,  $\boldsymbol {u}_t = \pi_t(\boldsymbol {b}_t,\boldsymbol {\xi}_t)$ , which can be stochastic with  $\pmb{\xi}_t\sim \mathcal{N}(0,I)$ .
In the inverse optimal control problem, the goal is to estimate parameters  $\pmb{\theta} \in \mathbb{R}^p$  of the agent's optimal control problem given the model and trajectory data. These parameters can include properties of the agent's cost function, the sensory and control systems of the agent, or the system's dynamics. More precisely, we assume that we are given a parametric form of the system dynamics, belief dynamics, cost function, and initial belief of the agent, which all might depend on the unknown parameters that we aim to infer. Further we assume a set of state trajectories. Importantly, we do not assume knowledge of the agent's belief. We follow a probabilistic approach to inverse optimal control, i.e., we consider the likelihood function
$$
p \left(\boldsymbol {x} _ {1: T} \mid \boldsymbol {\theta}\right) = p \left(\boldsymbol {x} _ {1} \mid \boldsymbol {\theta}\right) \prod_ {t = 1} ^ {T - 1} p \left(\boldsymbol {x} _ {t + 1} \mid \boldsymbol {x} _ {1: t}, \boldsymbol {\theta}\right), \tag {1}
$$
describing the probability of the observed trajectory data  $\boldsymbol{x}_{1:T} \coloneqq \{\boldsymbol{x}_1, \dots, \boldsymbol{x}_T\}$  given the parameters. For a set of trajectories, we assume them to be independent given the parameters, so that the likelihood factorizes into single trajectory likelihoods of the form in Eq. (1). In this equation, generally, each state  $\boldsymbol{x}_{t+1}$  depends on all previous states  $\boldsymbol{x}_1, \dots, \boldsymbol{x}_t$ , because the agent's internal noisy observations and control signals are not accessible to the researcher (Fig. 1 B). Therefore, the Markov property does not hold from the researcher's perspective, rendering computation of the likelihood function intractable. To deal with this problem, we employ two key insights: First, the joint dynamical system of the states and the agent's belief is Markovian [50]. Second, by keeping track of the distribution over the agent's belief, i.e., by performing belief tracking [29], we can iteratively compute the individual factors of the likelihood function in Eq. (1). In our IOC method, the goal is to maximize the likelihood w.r.t. the parameters  $\theta$ . To do so, we use gradient-based optimization with automatic differentiation to differentiate through the likelihood for computing the optimal parameters (Algorithm 1). An implementation of our algorithm is publicly available<sup>2</sup>.
We first introduce a general formulation of the IOC likelihood involving marginalization over the agent's internal beliefs in Section 3.1. Then, we show how to make the computations tractable by local linearization in Section 3.2. In Section 3.3, we provide details for suitable linearization points, which enables us to evaluate the approximate likelihood within a single forward pass.
3.1 Likelihood formulation
We start by defining a joint dynamical system of states and beliefs [50], in which each depends only on the state and belief at the previous time step and the noises. For that, we insert the policy into the dynamics and the policy and observation function into the belief dynamics, yielding the equation
$$
\left[ \begin{array}{c} \boldsymbol {x} _ {t + 1} \\ \mathbf {b} _ {t + 1} \end{array} \right] = \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] =: g (\boldsymbol {x} _ {t}, \mathbf {b} _ {t}, \boldsymbol {v} _ {t}, \boldsymbol {w} _ {t}, \boldsymbol {\xi} _ {t}). \tag {2}
$$
For given values of  $\mathbf{x}_t$  and  $\mathbf{b}_t$ , this equation defines the distribution  $p(\mathbf{x}_{t+1}, \mathbf{b}_{t+1} \mid \mathbf{x}_t, \mathbf{b}_t)$ , as  $\mathbf{v}_t, \mathbf{w}_t, \boldsymbol{\xi}_t$  are independent of  $\mathbf{x}_{t+1}$  and  $\mathbf{b}_{t+1}$ . Importantly, with this formulation, the control signals are only implicitly regarded through the policy, as we assume them to be latent for the researcher. In Section 3.2 we will introduce an approximation via linearization, which leads to a closed-form expression for  $p(\mathbf{x}_{t+1}, \mathbf{b}_{t+1} \mid \mathbf{x}_t, \mathbf{b}_t)$ .
One can use this Markovian joint dynamical system to compute the likelihood factors for each time step [29]. To this end, we first rewrite the individual likelihood terms  $p(\pmb{x}_{t+1} \mid \pmb{x}_{1:t})$  of Eq. (1) by marginalizing over the agent's belief at each time step, i.e.,
$$
p \left(\boldsymbol {x} _ {t + 1} \mid \boldsymbol {x} _ {1: t}\right) = \int p \left(\boldsymbol {x} _ {t + 1}, \mathbf {b} _ {t + 1} \mid \boldsymbol {x} _ {1: t}\right) \mathrm {d} \mathbf {b} _ {t + 1}. \tag {3}
$$
As the belief is an internal quantity of the agent and thus not observable to the researcher, we keep track of its distribution,  $p(\mathbf{b}_t \mid \mathbf{x}_{1:t})$ . For this, we rewrite
$$
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} \mid \boldsymbol {x} _ {t}, \mathbf {b} _ {t}\right) p \left(\mathbf {b} _ {t} \mid \boldsymbol {x} _ {1: t}\right) d \mathbf {b} _ {t}, \tag {4}
$$
where we have exploited the fact that the joint dynamical system of states and beliefs is Markovian. The distribution  $p(\mathbf{b}_t \mid \mathbf{x}_{1:t})$  acts as a summary of the past states and can be computed by conditioning on the current state, i.e.,
$$
p \left(\mathbf {b} _ {t} \mid \boldsymbol {x} _ {1: t}\right) = \frac {p \left(\boldsymbol {x} _ {t} , \mathbf {b} _ {t} \mid \boldsymbol {x} _ {1 : t - 1}\right)}{p \left(\boldsymbol {x} _ {t} \mid \boldsymbol {x} _ {1 : t - 1}\right)}. \tag {5}
$$
After determining  $p(\mathbf{b}_t \mid \mathbf{x}_{1:t})$ , we can propagate it through the joint dynamical system to arrive at the distribution  $p(\mathbf{x}_{t+1}, \mathbf{b}_{t+1} \mid \mathbf{x}_{1:t})$ . To obtain the belief distribution of the following time step,  $p(\mathbf{b}_{t+1} \mid \mathbf{x}_{1:t+1})$ , we condition on the observed state  $\mathbf{x}_{t+1}$ . To obtain the likelihood contribution, on the other hand, we marginalize out  $\mathbf{b}_{t+1}$ . To summarize, starting with an initialization  $p(\mathbf{b}_0)$ , we can compute the individual terms  $p(\mathbf{x}_{t+1} | \mathbf{x}_{1:t})$  of the likelihood by executing Algorithm 1 (Appendix C).
$^{2}$ https://github.com/RothkopfLab/nioc-neurips
3.2 Tractable likelihood via local linearization
While the marginalization and propagating operations in the previous section can be done in closed form for linear-Gaussian systems, this is no longer feasible for non-linear systems. Therefore, we follow the approach of local linearization used in iLQG (Section 2.2) and the EKF (Section 2.4). For the belief statistics, we consider the mean of the agent's belief, i.e.,  $\mathbf{b}_t = \mathbb{E}[\pmb{x}_t\mid \pmb{y}_1,\dots ,\pmb{y}_{t - 1}]$  and initialize the distribution for the first time step as a Gaussian,  $p(\mathbf{b}_1) = \mathcal{N}(\boldsymbol{\mu}_1^{(b)},\Sigma_1^{(b)})$ . We then approximate the distribution  $p(\pmb{x}_{t + 1},\mathbf{b}_{t + 1}\mid \pmb{x}_t,\mathbf{b}_t)$  as a Gaussian by applying a first-order Taylor expansion of  $g$ .