Inverse Rational Control with Partially Observable Continuous Nonlinear Dynamics

Abstract

A fundamental question in neuroscience is how the brain creates an internal model of the world to guide actions using sequences of ambiguous sensory information. This is naturally formulated as a reinforcement learning problem under partial observations, where an agent must estimate relevant latent variables in the world from its evidence, anticipate possible future states, and choose actions that optimize total expected reward. This problem can be solved by control theory, which allows us to find the optimal actions for a given system dynamics and objective function. However, animals often appear to behave suboptimally. Why? We hypothesize that animals have their own flawed internal model of the world, and choose actions with the highest expected subjective reward according to that flawed model. We describe this behavior as rational but not optimal. The problem of Inverse Rational Control (IRC) aims to identify which internal model would best explain an agent’s actions. Our contribution here generalizes past work on Inverse Rational Control which solved this problem for discrete control in partially observable Markov decision processes. Here we accommodate continuous nonlinear dynamics and continuous actions, and impute sensory observations corrupted by unknown noise that is private to the animal. We first build an optimal Bayesian agent that learns an optimal policy generalized over the entire model space of dynamics and subjective rewards using deep reinforcement learning. Crucially, this allows us to compute a likelihood over models for experimentally observable action trajectories acquired from a suboptimal agent. We then find the model parameters that maximize the likelihood using gradient ascent. Our method successfully recovers the true model of rational agents. This approach provides a foundation for interpreting the behavioral and neural dynamics of animal brains during complex tasks.

1. Introduction

Brains evolved to understand, interpret, and act upon the physical world. To thrive and reproduce in a harsh and dynamic natural environment, brains, therefore, evolved flexible, robust controllers. To be the controller, the fundamental function of the brain is to organize sensory data into an internal model of the outside world. The animals are never able to get complete information about the world. Instead, they only get partial and noisy observations of it. Thus, the brain should build its own internal model which necessarily includes uncertainties of the outside world, and base its decision upon that model [1]. However, we hypothesize that this internal model is not always correct, but the animals still behave rationally — meaning that animals act optimally according to their own internal model of the world, which may differ from the true world.

The goal of this paper is to identify the internal model of the agent by observing its actions. Unlike Inverse Reinforcement Learning (IRL) [2, 3, 4] which aims to learn only the reward function of target agent, or Inverse Optimal Control (IOC) [5, 6] to infer only unknown dynamics model, we use Inverse Rational Control (IRC) [7] to infer both. Since we consider neuroscience tasks which include naturalistic controls and complex physics of the world, we substantially extend past work [7] to include continuous spaces of state, action, and parameter with nonlinear dynamics. We parameterize nonlinear task dynamics and reward functions based on a physics model such that the family of tasks shares an overall structure but has different model parameters. In our framework, an experimentalist can observe state information of the environment and actions taken by the agent. On the other hand, the experimentalist cannot observe information about the agent’s internal model, such as its observations and beliefs. IRC infers the latent internal information of the agent using the data observable by the experimentalist.

The task is formulated as a Partially Observable Markov Decision Process (POMDP) [8, 9], a powerful framework for modeling agent behavior under uncertainty. In order to model an animal’s cognitive process whereby the decision-making is based on its own beliefs about the world, we reformulate the POMDP as a belief Markov Decision Process (belief MDP) [10, 11]. The agent builds its belief (i.e., its posterior distribution over world states) based on partial, noisy observations and its internal model, and the decision-making is based on its belief.

We construct a Bayesian agent to learn optimal policies and value functions over an entire parameterized family of models, which can be viewed as an optimized ensemble of agents each dedicated to one task. This then allows us to maximize the likelihood of the state-action trajectories generated by a target agent, by finding which parameters from the ensemble best explain the target agent’s data.

The main contributions of this paper are the following. First, our work is able to find both the reward function and internal dynamics model simultaneously in continuous nonlinear tasks. Note that continuous nonlinear dynamical systems are the most general form of tasks, so it is trivial to solve discrete and/or linear systems using the proposed approach. Second, we propose a novel approach to implement the Bayesian optimal control ensembles, including an idea of belief representation and belief updating method using estimators with constrained representational capacity (e.g., an extended Kalman filter). This allows us to build an algorithm that imitates the bounded rational cognitive process of the brain [12] and to perform belief-based decision-making. Lastly, we propose a novel approach to IRC combining Maximum Likelihood Estimation (MLE) and Monte Carlo Expectation-Maximization (MCEM). This method successfully infers the reward function and internal model parameters of the target agent by maximizing the likelihood of state-action trajectories under the assumption of rationality, while marginalizing over latent sensory observations. Importantly, this is possible because we trained ensembles of agents over entire parameter spaces using flexible function approximators. To the best of our knowledge, our work is the first study to infer both the reward and internal model of an unknown agent with partially observable continuous nonlinear dynamics.

2. Related Work

Inverse reinforcement learning (IRL).

The goal of IRL or imitation learning is to learn a reward function or a policy from expert demonstrations, and the goal of Inverse Optimal Control (IOC) is to infer an unknown dynamics model. Both approaches solve aspects of the general problem of inferring internal models of an observed agent. For example, some IRL works such as [13, 14, 15, 16] formulate the optimization problems to find features of reward or cost function that best explain the target agent’s state-action trajectories. Specifically, [13] finds reward features by solving a linear programming problem, and [15] uses a quadratic programming method to learn mappings from features to a cost function. In addition, [17] combines the principle of maximum entropy [18] to IRL so that the solution becomes as random as possible while still matching reward features the best. This guarantees avoiding the worst-case policy [19, 20]. Another stream of IRL is imitation learning [21, 22, 23, 24]. Typical IRL approaches use a two-step process: first learn the expert’s reward function first, and then train the policy with the learned reward. This could be slow, [21] directly extracts a policy from data. Across all of these methods, there is no a complete inverse solution that can learn how an agent models rewards, dynamics, and uncertainty in a partially observable task with continuous nonlinear dynamics and continuous controls.

Meta reinforcement learning (Meta RL).

The fundamental objective of Meta RL is to learn new skills or adapt to a new environment rapidly with a few training experiences. In order to efficiently adapt to the new tasks or environments, some Meta RL works try to infer tasks or meta parameters that govern the general task. For example, optimization-based meta RL works such as [25, 26, 27, 28] include a so-called ‘outer loop’ which optimizes the meta-parameters. In this sense, the meta RL is related to our goal since both work aim to infer the task parameters, although we use this parameterization to explain the actions of an agent. However, there are few studies to consider the partially observable setting of the agent. [29] includes both POMDP frameworks and meta-learning, but the partially observable information is the task information not the state information. [30] also considers a Bayesian approach with meta-learning, but it also uses Bayesian reasoning to infer the unseen tasks and learn quickly. Therefore, our paper differs from other Meta RL works in its task structure and goal. We allow partial observability about the world state, as occurs naturally in the animal’s decision-making process. More fundamentally, the goal of our work is not to find smarter agents, but rather to infer the internal model of an existing agent and to explain its behaviors.

Neuroscience and cognitive science

Neuroscientists aim to answer how the brain selects actions based on noisy sensory information and incomplete knowledge of the world. The hypothesis of the Bayesian brain [31] has been proposed to explain the brain’s functionalities with Bayesian inference and probabilistic representations of the animal’s internal state. Several studies propose mechanisms by which neurons could implement optimal behaviors despite pervasive uncertainty [11, 32, 33]. Despite the utility of having behavioral benchmarks based on optimality, animals often appear to behave suboptimally. Such suboptimality might come from the wrong internal model [7, 34] that is induced by a subjective prior belief of the animal [35, 36] and suboptimal inference [37]. The main goal of this paper is to infer the internal model of suboptimal agents using state-of-the-art deep reinforcement learning techniques and to provide a theoretical tool to interpret behavior and neural data obtained from ongoing neuroscience research.

For this reason, we test our approach by simulating an existing neuroscience task called ‘catching fireflies’ [38, 39], which is complex enough to require a sophisticated internal model, while being restrictive enough that animals can learn it and one can adequately constrain models of this behavior using feasible data volumes. Ultimately we will apply our approach to understand the internal models of behaving animals, where we do not know the ground truth. Before doing that, it is important to use simulated agents that allow us to validate the method when we do know the ground truth. Recently, a similar effort to build AI-relevant testbed for animal cognition and behavior is presented in [40, 41].

3. Bayesian Optimal Control Ensembles

Our method can be viewed as a search over an ensemble of agents, each optimally trained on different POMDP tasks, to find which of these agents best explain the experimentally observed behaviors of a target agent. The experimentalist is an external observer who has information about the world states and agent actions, but not about the agent’s internal model, noisy sensory observations, or beliefs.

3.1. Belief Markov Decision Process and optimal control

A POMDP is defined as a tuple $M=(\mathcal{S},\mathcal{A},\Omega ,R,T,O,\gamma )$ that includes states $s_t\in \mathcal{S}$, actions $a_t\in \mathcal{A}$, observations o_t ∈ Ω, reward functions R(s_t, a_t, s_t+1; θ), state transition probabilities T(s_t+1|s_t, a_t; θ), observation probabilities O(o_t|s_t; θ) at time t, and a temporal discount factor γ. Here, θ ∈ Θ denotes a vector of model parameters defining the rewards, state transitions and observations, and the state space $\mathcal{S}$ and action space $\mathcal{A}$ are considered to occupy a continuous space. Thus, θ parameterizes a POMDP family. A graphical model of a POMDP is presented in Figure 1.

Figure 1:

Graphical model of a POMDP. Solid circles denote observable variables to an experimentalist, and empty circles denote latent variables.

The state s_t is defined as the representation of the environment which can live in high dimensional spaces. It may be fully accessible by the experimentalist but not by the agent. The agent gets an observation o_t of the environment with state s_t, which is partial and noisy version of state s_t. Because of the partial observability, the dimension of o_t could be lower than the dimension of s_t. The observation process is modeled by the observation function O(o_t|s_t; θ). Note that any noise added from state to observation is the internal noise of the agent, i.e., the noise within the nervous system of the agent. Because of this noise, the observation is only known to the agent and the experimentalist can never access it directly.

Based on its observations and actions up to time t, a rational agent builds a posterior distribution B(s_t|o_1:t, a_1:t−1; θ) over the world state given the history of observations and actions, and it bases its actions upon that posterior. In practice this posterior is summarized by a belief b_t, defined as sufficient statistics for the posterior distribution over states, i.e., B(s_t|b_t) = B(s_t|o_1:t, a_1:t−1; θ). In principle, a belief b_t over a general continuous state could be infinite-dimensional, but we assume that the belief is continuous but finite-dimensional. Let B(s_t|b_t) be the probability that the environment is in the state s_t when the agent’s belief is b_t. By the Markov property, b_t is determined by b_t−1, a_t−1, o_t such that B(s_t|b_t) can be calculated as follows.

$$B\left(s_t\mid b_t\right)=B\left(s_t\mid b_{t-1},a_{t-1},o_t;\theta \right)$$ (1)

$$=\frac{1}{Z}O\left(o_t\mid s_t;\theta \right)\int d{s}_{t-1}T\left(s_t\mid s_{t-1},a_{t-1};\theta \right)B\left(s_{t-1}\mid b_{t-1}\right)$$ (2)

where Z = ∫ ds_t O(o_t|s_t; θ)∫ ds_t−1T(s_t|s_t−1, a_t−1; θ) B(s_t−1|b_t−1) is a normalizing constraint. In general this recursion is intractable, so we approximate it under tractable model assumptions, as we do in our application below. By replacing the state of the environment by the belief of the agent, the POMDP problem can be reformulated as a belief MDP problem, and the optimal policy can be found based on well-known MDP solvers [42, 43, 44, 45] applied to the fully observed belief state.

The optimal policy π*(a_t|b_t; θ) defines how the agent chooses an action $a_t^{*}$ that maximizes the temporally discounted total expected future reward, given the current belief b_t and internal model θ. This defines the Q-value Q(b_t, a_t; θ) as a belief-action value:

$$Q\left(b_t,a_t;\theta \right)=\int d{b}_{t+1}\overline{T}\left(b_{t+1}\mid b_t,a_t;\theta \right)\left(\overline{R}\left(b_t,a_t,b_{t+1};\theta \right)+\gamma \underset{a}{\text{max }}Q\left(b_{t+1},a;\theta \right)\right)$$ (3)

where $\overline{T}\left(b_{t+1}\mid b_t,a_t;\theta \right)$ is the belief transition probability and $\overline{R}\left(b_t,a_t,b_{t+1};\theta \right)$ is the reward as a function of belief, defined as follows.

$$\overline{T}\left(b_{t+1}\mid b_t,a_t;\theta \right)=\iiint d{s}_{t}d{s}_{t+1}d{o}_{t+1}B\left(s_t\mid b_t\right)T\left(s_{t+1}\mid s_t,a_t;\theta \right)O\left(o_{t+1}\mid s_{t+1};\theta \right)p\left(b_{t+1}\mid b_t,a_t,o_{t+1};\theta \right)$$ (4)

$$\overline{R}\left(b_t,a_t,b_{t+1};\theta \right)=\iint d{s}_{t}d{s}_{t+1}B\left(s_t\mid b_t\right)B\left(s_{t+1}\mid b_{t+1}\right)R\left(s_t,a_t,s_{t+1};\theta \right)$$

In (4), the belief update is expressed in a generalized form p(b_t+1|b_t, a_t, o_t+1; θ) that allows either deterministic optimal belief updates, or could even account for other constraints on the inference process, including stochasticity.

The optimal action from a belief state will be also defined by a deterministic policy ${\pi }^{*}\left(a_t\mid b_t;\theta \right)=\delta \left(a_t^{*}={\text{arg max}}_{a}Q\left(b_t,a;\theta \right)\right)$. In case of continuous belief and action spaces, it is hard both to compute an optimal Q-function and to maximize it. Thus, we will approximate both using neural networks.

3.2. Training Bayesian optimal control ensembles with partial noisy observations

To successfully design and train an ensemble of agents, we identify three major challenges and provide solutions.

First, how can we construct the optimal control ensembles that can solve a family of tasks? As discussed, the task can be parameterized by the model parameter θ ∈ Θ such that the family of tasks shares the model structure but has different model parameters. We use this model parameter as an external input to flexible function approximators (neural networks) to estimate values and policies (Critic and Actor). Thus, the agent can be trained over parameter spaces. As presented in Figure 2, Critic and Actor both take parameter vector θ as an input, and respectively calculate the Q-value and best action for the task with that θ.

Figure 2:

A block diagram of Algorithm 1.

Second, how should we represent and update the agent’s belief? For our concrete example application below, we use an extended Kalman filter [46] to provide a tractable Gaussian approximation for the belief state and its nonlinear dynamics. The resultant belief update is deterministic, p(b_t+1|b_t, a_t, o_t+1; θ) = δ (b_t+1 = f(b_t, a_t, o_t+1; θ)). Tests with more flexible particle filters showed that this approximation is reasonable in our target application. For other applications, different belief representations and dynamics may be more accurate [47], and in principle a family of agents could use representational learning [48].

Lastly, how should we train a rational model agent ensemble with continuous belief and action spaces? Here we use the model-free deep reinforcement learning algorithm called Deep Deterministic Policy Gradient (DDPG) [49]. This method is able to approximate the value function over continuous belief states, actions, and task parameters, all using one neural network (the Critic), and uses it to train a policy network (the Actor) which also receives inputs about the current belief and task parameters. Viable alternatives for continuous control in the deep reinforcement learning literature include [50, 51, 52, 53].

The training process for optimal control ensembles is summarized in Algorithm 1, and a block diagram is provided in Figure 2. The agent is trained on simulated experiences. Given the belief b_t and parameters θ, the Actor returns the best action a_t. As the agent performs the action a_t, it changes the world state to s_t+1 following the state transition probabilities T. The reward from the world R is given to the agent and fed back to the Critic to get a better estimation of the Q-value, which then improves the selection of the action in the Actor. From the new state s_t+1, the agent gets a partial and noisy observation o_t+1 with the observation probabilities O. Then, the Gaussian belief state is updated using the extended Kalman filter f, b_t+1 = f(b_t, a_t, o_t+1; θ). A new action a_t+1 is selected by the Actor, and these processes are iterated until the neural networks are fully trained. During this training, we sample new model parameters θ every episode so the agent can experience the entire space of tasks, and thus generalize better over that space.

4. Inverse Rational Control with Maximum Likelihood Estimation

Once an agent ensemble is fully trained over the entire parameter space, we can use this ensemble to find the internal model parameters of the best-fitting rational agent in that model family. We solve the continuous Inverse Rational Control problem by finding the parameters θ that have the highest likelihood for explaining an agent’s measured behavior.

4.1. Discrepancy between the true world and internal model

Recall that our core hypothesis is that animals have their own internal model of the world which may not be always correct, but they still behave rationally, choosing actions with the highest expected subjective reward according to their internal model. We must therefore distinguish between the two kinds of model parameters: the true ones ϕ which determine the world dynamics and are known to the experimentalist, and the agent’s internal model parameters θ which are latent for the experimentalist but governs all cognitive processes of the agent (Figure 3). The world parameters ϕ govern the world dynamics such as state transition probability T(s_t+1|s_t, a_t; ϕ) and reward function R(s_t, a_t, s_t+1; ϕ). On the other hand, the internal parameters θ govern the agent’s internal process such as the observation probability O(o_t|s_t; θ), the belief transition probability $\overline{T}\left(b_{t+1}\mid b_t,a_t;\theta \right)$, and the subjective reward as a function of belief $\overline{R}\left(b_t,a_t,b_{t+1};\theta \right)$, leading to a subjective belief update probability p(b_t+1|b_t, a_t, o_t+1; θ) and rational policy π(a_t|b_t; θ).

Figure 3:

An illustrative explanation of model discrepancy. The solid lines and circles are governed by the true world parameter ϕ which is known to the experimentalist. The dashed lines and empty circles are governed by internal model parameter θ which is latent to the experimentalist, and may differ from ϕ.

4.2. Inferring internal model parameter θ

To find the internal model parameters θ that maximize the log-likelihood of the experimentally observable data (s, a)_1:T, $\widehat{\theta }={\text{arg max}}_{\theta }\text{ln }p\left(s_{1:T},a_{1:T}\mid \varphi ,\theta \right)$ we use the Monte Carlo Expectation Maximization (MCEM) algorithm [54] to marginalize the complete data log-likelihood over latent observations o_1:T and beliefs b_1:T. This yields an iterative algorithm, which repeatedly maximizes

$${\widehat{\theta }}_{k+1}=\text{arg }\underset{\theta }{\text{max }}\int d{o}_{1:T}d{b}_{1:T}p\left(o_{1:T}{b}_{1:T}\mid s_{1:T},a_{1:T};{\theta }_k\right)\text{ln }p\left(s_{1:T},o_{1:T},b_{1:T},a_{1:T}\mid \varphi ,\theta \right)$$ (5)

$$\approx \text{arg }\underset{\theta }{\text{max }}\frac{1}{L}\sum _{l=1}^L\text{ln}p\left(s_{1:T},o_{1:T}^{(l)},b_{1:T}^{(l)},a_{1:T}\mid \varphi ,\theta \right)$$ (6)

where the sum is over samples (o^(l), b^(l))_1:T drawn from a posterior distribution p(o_1:T db_1:T|s_1:T, a_1:T; θ_k) determined by parameters θ_k from previous iterations.

The log-likelihood of the complete data (including the l-th samples of observations and beliefs based on parameter θ_k) can be decomposed using the Markov property into

$$\text{ln }p\left(s_{1:T},o_{1:T}^{(l)},b_{1:T}^{(l)},a_{1:T}\mid \varphi ,\theta \right)$$ (7)

$$=\text{ln }p\left(s_0,o_0^{(l)},b_0^{(l)},a_0\right)+\sum _{t=1}^T\left(\text{ln}T\left(s_t\mid s_{t-1},a_{t-1};\varphi \right)+\text{ln}O\left(o_t^{(l)}\mid s_t;\theta \right)+\text{ln}p\left(b_t^{(l)}\mid b_{t-1}^{(l)},a_{t-1},o_t^{(l)};\theta \right)+\text{ln}\pi \left(a_t\mid b_t^{(l)};\theta \right)\right).$$ (8)

Note that the only terms depending on the agent’s parameters θ are the latent observations probabilities, belief dynamics, and policy. So when we optimize over θ, all other terms vanish. Moreover, since we use deterministic belief updates, the belief update term in (8) is also independent of θ when evaluated on sampled beliefs. The only terms that survive are

$$\widehat{\theta }=\text{arg }\underset{\theta }{\text{max }}\sum _{l=1}^L\sum _{t=1}^T\left(\text{ln }O\left(o_t^{(l)}\mid s_t;\theta \right)+\text{ln }\pi \left(a_t\mid b_t^{(l)};\theta \right)\right).$$ (9)

To optimize (9), we use gradient ascent over parameter space, $\theta \leftarrow \theta +\alpha {\nabla }_{\theta }\mathcal{L}$ with learning rate α. This is explained in Algorithm 2.

5. Demonstration task: ‘Catching fireflies’

To verify the proposed method, we carefully select a relevant task. Our application focus on continuous world states, actions, and beliefs makes standard RL testbeds (e.g. Nintendo, MuJoCo) more difficult. Common tasks like gridworld or tiger do not exhibit continuous properties and remain excessively small toy problems. Standard continuous control tasks do not use partially observability; tasks that do would likely generate beliefs that would be substantially harder to interpret. Additionally, there is a ready application to existing neuroscience experiments based on ‘catching fireflies’ in virtual reality [55, 39], which is complex enough to be interesting to animals, requires a continuous representation of uncertainty and continuous control, and yet remains tractable enough that we can assess the fidelity of recovered beliefs.

In our task, an agent must navigate through a virtual environment to reach a transiently visible target, called the ‘firefly’ (Figure 4A). At the beginning of each trial, a firefly blinks briefly at a random location on the ground plane. The agent is able to control its forward and angular velocities to freely navigate the world. If the agent stops sufficiently close to the invisible target, it receives a reward. As the agent moves, a sparse ground plane texture generates an optic flow pattern, a vector field of local image motion. This allows the agent to estimate its speed up to some perceptual uncertainty. However, there is no direct access to information about its current location because the ground plane texture is transient and does not provide spatial landmarks. Thus, the agent must integrate optic flow to estimate its current position relative to the firefly target, as well as its uncertainty about that position.

Figure 4:

A. An illustration of the ‘catching fireflies’ task from the agent’s point of view. To reach the transiently visible firefly target, an agent must navigate by optic flow over a dynamically textured ground plane. The agent is rewarded if it stops close enough to the target location. B. Converging trajectory of IRC estimates of the agent’s parameters θ. We use gradient ascent to find θ that maximizes approximated log likelihood $\mathcal{L}(\theta )$ in Algorithm 2. C. Successful recovery of individual agent parameters. The black line is the identity, meaning that true values and estimated values are equal. Across all parameter spaces, the proposed approach accurately recovers the agent’s internal model parameters given limited data.

We demonstrate the efficacy of our approach using a simulated agent for which ground truth is known. Thus, we verify our method by showing the successful recovery of the internal model parameters since we know the ground truth. Note that there are no comparisons to alternative methods because no other algorithms exist that solve the IRC problem in continuous state and action spaces. Figure 4B shows a two-dimensional contour plot of the approximate log-likelihood of observable data $\mathcal{L}(\theta )$. Recall that the model parameters θ are high dimensional, so here we plot only two dimensions of θ. The red line shows an example trajectory of parameters θ as IRC Algorithm 2 converges. Our approach estimates θ that maximizes the log-likelihood of the observable data $\mathcal{L}(\theta )$. Figure 4C shows that the estimated parameters recovered by our algorithm closely match the agent’s true parameters.

6. Conclusion

This paper introduces a novel framework to infer the internal model of agents from their behaviors. We infer not only the subjective reward function of the agent, but we also simultaneously infer the task dynamics that the agent assumes. To accomplish this, we first train Bayesian optimal control ensembles that generalize over the space of task parameters. Since the target agent is only exposed to partial information about the world state, the agent chooses the best action based on its belief about the world and its assumptions about the task. By using this optimally trained agent ensemble, our approach to Inverse Rational Control with continuous state and action spaces can infer the internal model parameters that best explain the collected behavioral data. With a simulated agent where we know the ground truth, we confirm that our approach successfully recovers the internal models. This success encourages us to apply this method to real behavioral data as well as to new tasks and applications.

Broader Impact

We have implemented IRC for neuroscience applications, but the core principles have value in other fields as well. We can view IRC as a form of Theory of Mind, whereby one agent (a neuroscientist) creates a model of another agent’s mind (for a behaving animal). Theory of Mind is a prominent component of human social interactions, and imputing rational motivations to actions provides a useful description of how people think [56, 57, 58]. Using IRC methods to provide a better understanding of people’s motivations could yield important insights for understanding and improving social and political interactions, as well as raising possible ethical concerns if used for manipulation. The design of agents interacting with humans would also benefit from being able to attribute rational strategies to others. For example, recent work uses a related approach to impute purpose to a neural network [16]. One important practical example is self-driving cars, which currently struggle to handle the perceived unpredictability of humans. While humans do indeed behave unpredictably, some of this may stem from ignorance of the rational computation that drives their actions. The IRC provides a framework for interpreting agents, and serves as a valuable tool for greater understanding of unifying principles of control.

Appendix. A Derivation of Monte Carlo Expectation Maximization (MCEM)

A derivation from (5) to (6) is based on MCEM. We here provide more details of the MCEM.

Let x be the observable data, z be the latent variable and θ be the parameters that govern the process. The goal is to find θ that maximizes the log likelihood of the observable data.

$$\theta =\text{arg }\underset{\theta }{\text{max }}\text{ln }p(x\mid \theta )$$

The log likelihood of the observable data can be reformulated as follows.

$$\begin{gathered} \text{ln }p(x\mid \theta )=\int dzq(z)\text{ ln }p(x\mid \theta ) \\ =\int dzq(z)[\text{ln }p(x,z\mid \theta )-\text{ln }p(z\mid x,\theta )] \\ =\int dzq(z)[\text{ln }p(x,z\mid \theta )-\text{ln }q(z)+\text{ln }q(z)-\text{ln }p(z\mid x,\theta )] \\ =\int dzq(z)\left[\text{ln}\frac{p(x,z\mid \theta )}{q(z)}-\text{ln}\frac{p(z\mid x,\theta )}{q(z)}\right] \\ =\int dzq(z)\text{ln}\frac{p(x,z\mid \theta )}{q(z)}-\int dzq(z)\text{ln}\frac{p(z\mid x,\theta )}{q(z)} \end{gathered}$$ (10)

$$=\mathcal{L}(q,\theta )+KL(q\parallel p)$$ (11)

Since KL divergence is always non-negative value, $\mathcal{L}(q,\theta )$ is the lower bound of ln p(x|θ). The complete data log likelihood ln p(x, z|θ) is easier to handle than the observed data log likelihood ln p(x|θ). Thus, instead of maximizing ln p(x|θ), we aim to maximize its lower bound $\mathcal{L}(q,\theta )=\int dzq(z)\text{ln}\frac{p(x,z\mid \theta )}{q(z)}$.

A.1 E-step

As KL(q‖p) gets smaller, we have a tighter lower bound. If KL(q‖p) = 0, $\text{ln }p(x\mid \theta )=\mathcal{L}(q,\theta )$. KL(q‖p) = 0 is satisfied only if q = p. Thus q(z) = p(z|x, θ) from (10). This is the E-step of the EM algorithm [59]. Note that in this step, q(z) is a function only of z, which means both x and θ are used as given variables. Thus, we denote θ_old as a fixed parameter that is used to specify q(z). Once q(z) = p(z|x, θ_old) is used in $\mathcal{L}(q,\theta )$ of (11), ln p(x|θ) can be expressed as follows.

$$\begin{gathered} \text{ln }p(x\mid \theta )=\mathcal{L}(q,\theta ) \\ =\int dzp\left(z\mid x,{\theta }_{\text{old}}\right)\text{ln}\frac{p(x,z\mid \theta )}{p\left(z\mid x,{\theta }_{\text{old}}\right)} \\ =\int dzp\left(z\mid x,{\theta }_{\text{old}}\right)\text{ln }p(x,z\mid \theta )-\int dzp\left(z\mid x,{\theta }_{\text{old}}\right)\text{ln }p\left(z\mid x,{\theta }_{\text{old}}\right) \\ =\int dzp\left(z\mid x,{\theta }_{\text{old}}\right)\text{ln }p(x,z\mid \theta )+H\left(z\mid x,{\theta }_{\text{old}}\right) \\ =\mathcal{Q}\left(\theta ,{\theta }_{\text{old}}\right)+H\left(z\mid x,{\theta }_{\text{old}}\right) \end{gathered}$$ (12)

A.2 M-step

Next, we want to find θ that maximizes ln p(x|θ). This is the M-step of the EM algorithm. Since H(z|x, θ_old) is a constant (i.e., not a function of θ),

$$\begin{gathered} \theta =\text{arg }\underset{\theta }{\text{max }}\text{ln }p(x\mid \theta ) \\ =\text{arg }\underset{\theta }{\text{max }}\mathcal{Q}\left(\theta ,{\theta }_{\text{old}}\right) \\ =\text{arg }\underset{\theta }{\text{max }}\int dzp\left(z\mid x,{\theta }_{\text{old}}\right)\text{ln }p(x,z\mid \theta ). \end{gathered}$$ (13)

If p(z|x, θ_old) is hard to get analytically, (13) can be approximated by the Monte Carlo approach. The resultant optimization is called the MCEM algorithm.

$$\begin{gathered} \theta =\text{arg }\underset{\theta }{\text{max }}\int dzp\left(z\mid x,{\theta }_{\text{old}}\right)\text{ln }p(x,z\mid \theta ) \\ \approx \text{arg }\underset{\theta }{\text{max }}\frac{1}{L}\sum _{l=1}^L\text{ln}p\left(x,z^{(l)}\mid \theta \right) \end{gathered}$$ (14)

where z^l is l-th particle for the latent variable z.