Original: Dynamic Inverse Reinforcement Learning for Characterizing Animal Behavior
Revised: Dynamic Inverse Reinforcement Learning for Characterizing Animal Behavior

Original: Abstract
Revised: Abstract

Original: Understanding decision-making is a core goal in both neuroscience and psychology, and computational models have often been helpful in the pursuit of this goal.
Revised: Understanding decision-making is a core objective in both neuroscience and psychology, and computational models have often been helpful in the pursuit of this goal.

Original: While many models have been developed for characterizing behavior in binary decision-making and bandit tasks, comparatively little work has focused on animal decision-making in more complex tasks, such as navigation through a maze.
Revised: While many models have been developed for characterizing behavior in binary decision-making and bandit tasks, comparatively little work has focused on animal decision-making in more complex tasks, such as navigation through a maze.

Original: Inverse reinforcement learning (IRL) is a promising approach for understanding such behavior, as it aims to infer the unknown reward function of an agent from its observed trajectories through state space.
Revised: Inverse reinforcement learning (IRL) is a promising approach for understanding such behavior, as it aims to infer the unknown reward function of an agent from its observed trajectories through state space.

Original: However, IRL has yet to be widely applied in neuroscience.
Revised: However, IRL has yet to be widely applied in neuroscience.

Original: One potential reason for this is that existing IRL frameworks assume that an agent’s reward function is fixed over time.
Revised: One potential reason for this is that existing IRL frameworks assume that an agent’s reward function is fixed over time.

Original: To address this shortcoming, we introduce dynamic inverse reinforcement learning (DIRL), a novel IRL framework that allows for time-varying intrinsic rewards.
Revised: To address this shortcoming, we introduce dynamic inverse reinforcement learning (DIRL), a novel IRL framework that allows for time-varying intrinsic rewards.

Original: Our method parametrizes the unknown reward function as a time-varying linear combination of spatial reward maps (which we refer to as “goal maps”).
Revised: Our method parametrizes the unknown reward function as a time-varying linear combination of spatial reward maps (which we refer to as “goal maps”).

Original: We develop an efficient inference method for recovering this dynamic reward function from behavioral data.
Revised: We develop an efficient inference method for recovering this dynamic reward function from behavioral data.

Original: We demonstrate DIRL in simulated experiments and then apply it to a dataset of mice exploring a labyrinth.
Revised: We demonstrate DIRL in simulated experiments and then apply it to a dataset of mice exploring a labyrinth.

Original: Our method returns interpretable reward functions for two separate cohorts of mice, and provides a novel characterization of exploratory behavior.
Revised: Our method returns interpretable reward functions for two separate cohorts of mice, and provides a novel characterization of exploratory behavior.

Original: We expect DIRL to have broad applicability in neuroscience, and to facilitate the design of biologically-inspired reward functions for training artificial agents.
Revised: We expect DIRL to have broad applicability in neuroscience, and to facilitate the design of biologically-inspired reward functions for training artificial agents.

Original:  Introduction
Revised:  Introduction

Original: Characterizing the decision-making behavior of humans and animals is a central goal in neuroscience and psychology [1, 2].
Revised: Characterizing the decision-making behavior of humans and animals is a central goal in neuroscience and psychology [1, 2].

Original: Decision-making tasks such as Two-Alternative Forced Choice (2AFC) and bandit problems have been widely studied [3–7], and previous work has developed a variety of models for behavior in these tasks [1, 7–9].
Revised: Decision-making tasks such as Two-Alternative Forced Choice (2AFC) and bandit problems have been widely studied [3–7], and previous work has developed a variety of models for behavior in these tasks [1, 7–9].

Original: The classic psychometric curve represents one such model [10], and more recent work has focused on models based on reinforcement learning [1, 8, 11, 12].
Revised: The classic psychometric curve represents one such model [10], and more recent work has focused on models based on reinforcement learning [1, 8, 11, 12].

Original: Such models allow us to understand and compare the decision-making strategies used by humans and animals, and can also provide a low-dimensional description of behavior that can be regressed against neural data [1, 13].
Revised: Such models allow us to understand and compare the decision-making strategies used by humans and animals, and can also provide a low-dimensional description of behavior that can be regressed against neural data [1, 13].

Original: Although a large literature has focused on models of decision-making in simple 2AFC and bandit tasks, comparatively few papers have sought to model behavior in larger, complex natural environments.
Revised: Although a large literature has focused on models of decision-making in simple 2AFC and bandit tasks, comparatively few papers have sought to model behavior in larger, complex natural environments.

Original: [Figure]
Revised: [Figure]

Original: In a recent study, Rosenberg et al.
Revised: In a recent study, Rosenberg et al.

Original: [15] introduced a novel experimental paradigm involving water-restricted mice navigating a 127-node labyrinth equipped with a water port at one terminal node (Fig. 1).
Revised: [15] introduced a novel experimental paradigm involving water-restricted mice navigating a 127-node labyrinth equipped with a water port at one terminal node (Fig. 1).

Original: At each node, the mouse could make up to 4 distinct decisions (‘stay’, ‘go left’, ‘go right’, or ‘reverse’).
Revised: At each node, the mouse could make up to 4 distinct decisions (‘stay’, ‘go left’, ‘go right’, or ‘reverse’).

Original: Navigation through a maze is a perfect example of complex yet natural decision-making behavior that remains poorly understood.
Revised: Navigation through a maze is a perfect example of complex yet natural decision-making behavior that remains poorly understood.

Original: Although reinforcement learning may seem like the natural framework for modeling such goal-driven behavior, the rewards experienced by these mice are not obvious to the experimenter.
Revised: Although reinforcement learning may seem like the natural framework for modeling such goal-driven behavior, the rewards experienced by these mice are not obvious to the experimenter.

Original: Indeed, as noted in Rosenberg et al.
Revised: Indeed, as noted in Rosenberg et al.

Original: [15], the observed trajectories indicate that mice are not only motivated by the extrinsic water reward, but also by intrinsic rewards such as their curiosity to explore the environment.
Revised: [15], the observed trajectories indicate that mice are not only motivated by the extrinsic water reward, but also by intrinsic rewards such as their curiosity to explore the environment.

Original: Inverse reinforcement learning (IRL) [17–19] addresses the problem of inferring the unknown reward function of an agent.
Revised: Inverse reinforcement learning (IRL) [17–20] addresses the problem of inferring the unknown reward function of an agent.

Original: Given access to the agent’s trajectories as it interacts with the environment,
Revised: Given access to the agent’s trajectories as it interacts with the environment,

Original: IRL identifies the states and actions that the agent finds rewarding.
Revised: IRL identifies the states and actions that the agent finds rewarding.

Original: While IRL has found many successes in robotics [20, 21] and healthcare contexts [22–24], existing IRL methods have not found broad application in neuroscience.
Revised: While IRL has found many successes in robotics [21, 22] and healthcare contexts [23–25], existing IRL methods have not found broad application in neuroscience.

Original: One plausible reason for this dearth is that IRL methods typically assume that the unknown reward function is fixed over the course of the agent’s trajectory.
Revised: One plausible reason for this dearth is that IRL methods typically assume that the unknown reward function is fixed over the course of the agent’s trajectory.

Original: Yet in many real-world decision-making tasks, rewards can change over time.
Revised: Yet in many real-world decision-making tasks, rewards can change over time.

Original: For example: the goals of a mouse, and consequently its intrinsic rewards, can change with time depending on factors such as fatigue, satiation and curiosity.
Revised: For example: the goals of a mouse, and consequently its intrinsic rewards, can change with time depending on factors such as fatigue, satiation and curiosity.

Original: This motivated us to develop an IRL method for characterizing animal decision-making in complex environments such as the labyrinth, in which rewards may vary over time.
Revised: This motivated us to develop an IRL method for characterizing animal decision-making in complex environments such as the labyrinth, in which rewards may vary over time.

Original: Here we propose dynamic inverse reinforcement learning (DIRL), an IRL method that allows for time-varying reward functions.
Revised: Here we propose dynamic inverse reinforcement learning (DIRL), an IRL method that allows for time-varying reward functions.

Original: DIRL parametrizes the animal’s reward function as a time-varying linear combination of a small number of spatial reward maps, which we refer to as “goal maps”.
Revised: DIRL parametrizes the animal’s reward function as a time-varying linear combination of a small number of spatial reward maps, which we refer to as “goal maps”.

Original: A single goal map is a function specifying the amount of reward at each state in the environment.
Revised: A single goal map is a function specifying the amount of reward at each state in the environment.

Original: (Classic IRL methods thus seek to identify a single, fixed goal map.)
Revised: (Classic IRL methods thus seek to identify a single, fixed goal map.)

Original: DIRL posits the existence of multiple goal maps with time-varying weights, which allows the instantaneous reward function to vary in time.
Revised: DIRL posits the existence of multiple goal maps with time-varying weights, which allows the instantaneous reward function to vary in time.

Original: For example, a “thirst” goal map could have high levels of reward at the water port, while an “explore” reward map might have reward evenly dispersed throughout the environment.
Revised: For example, a “water” goal map could have high levels of reward at the water port, while an “explore” reward map might have reward evenly dispersed throughout the environment.

Original: Time-varying weights then modulate the extent to which each goal map is active at any timestep.
Revised: Time-varying weights then modulate the extent to which each goal map is active at any timestep.

Original: We introduce  a method for inferring both the goal maps and the time-varying weights from state-trajectory data, and demonstrate the application of our framework in simulation (gridworld and labyrinth environments), as well as on real mouse decision-making data from Rosenberg et al.
Revised: We introduce an inference method, which extends a fixed reward IRL framework by Ziebart et al. [20], to allow us to obtain both the goal maps and the time-varying weights from state-trajectory data. Finally, we demonstrate the application of our framework in simulation (gridworld and labyrinth environments), as well as on real mouse decision-making data from Rosenberg et al.

Original: Our method recovers interpretable reward functions for both cohorts of mice studied there, and reveals an “explore” goal map for one of the cohorts.
Revised: Our method recovers interpretable reward functions for both cohorts of mice studied there, and reveals an “explore” goal map for one of the cohorts.

Original: While exploration remains poorly understood in neuroscience [25], our method offers a powerful framework for characterizing exploratory, as well as exploitative behavior from an animal’s trajectories alone.
Revised: While exploration remains poorly understood in neuroscience [26], our method offers a powerful framework for characterizing exploratory, as well as exploitative behavior from an animal’s trajectories alone.

Original:  Related Work
Revised:  Related Work

Original: The neuroscientific literature describing models of decision-making is vast, including both normative [1, 8] and descriptive models [9, 10, 26–29].
Revised: The neuroscientific literature describing models of decision-making is vast, including both normative [1, 8] and descriptive models [9, 10, 27–30].

Original: While some of this work considers models of decisionmaking that vary over time, most such models have focused on 2AFC or bandit tasks, and do not scale easily to complex decision-making tasks such as navigation.
Revised: While some of this work considers models of decisionmaking that vary over time, most such models have focused on 2AFC or bandit tasks, and do not scale easily to complex decision-making tasks such as navigation.

Original: To the best of our knowledge, IRL has found relatively few applications in neuroscience.
Revised: To the best of our knowledge, IRL has found relatively few applications in neuroscience.

Original: One exception is Schultheis et al.
Revised: One exception is Schultheis et al.

Original: [30] which performed inverse optimal control to infer the cost function associated with sensorimotor behavior.
Revised: [31] which performed inverse optimal control to infer the cost function associated with sensorimotor behavior.

Original: Their targeted application—modeling the cost function optimized by a human performing a reaching task—is very different to ours, which results in different assumptions in their work.
Revised: Their targeted application—modeling the cost function optimized by a human performing a reaching task—is very different to ours, which results in different assumptions in their work.

Original: Another exception is Yamaguchi et al.
Revised: Another exception is Yamaguchi et al.

Original: [31] which described an IRL framework for identifying the thermotactic strategies used by C. elegans.
Revised: [32] which described an IRL framework for identifying the thermotactic strategies used by C. elegans.

Original: Relative to the work we present here, Yamaguchi et al.
Revised: Relative to the work we present here, Yamaguchi et al.

Original: [31] assumed a fixed reward function, and was restricted to the case of a linearly-solvable MDP.
Revised: [32] assumed a fixed reward function, and was restricted to the case of a linearly-solvable MDP.

Original: In Kwon et al.
Revised: In Kwon et al.

Original: [32], the authors considered the problem of “Inverse Rational Control” and inferred the parameters governing the evolution of an animal’s belief state and subjective reward function within a POMDP  framework.
Revised: [33], the authors considered the problem of “Inverse Rational Control” and inferred the parameters governing the evolution of an animal’s belief state and subjective reward function within a POMDP (partially observed MDP) framework.

Original: Finally, while not the main focus of their paper, Reddy [16] used the static IRL framework of Ziebart et al.
Revised: Finally, while not the main focus of their paper, Reddy [16] used the static IRL framework of Ziebart et al.

Original: [19] to infer the rewards optimized by mice navigating in the Rosenberg et al.
Revised: [19] to infer the rewards optimized by mice navigating in the Rosenberg et al.

Original: [15] task.
Revised: [15] task.

Original: The literature on IRL [17–19, 33–37] is extensive, but most studies have assumed that the true underlying reward function is fixed over time.
Revised: The literature on IRL [17–20, 34–37] is extensive, but most studies have assumed that the true underlying reward function is fixed over time.

Original: Of particular relevance to our work is the Maximum
Revised: Of particular relevance to our work is the Maximum

Original: Causal Entropy (MCE) framework of Ziebart et al.
Revised: Causal Entropy (MCE) framework of Ziebart et al.

Original: MCE infers a reward function under the assumption that an agent seeks to maximize both discounted future reward and the discounted future entropy of its policy.
Revised: MCE infers a reward function under the assumption that an agent seeks to maximize both discounted future reward and the discounted future entropy of its policy.

Original: We also use discounted future entropy to regularize the goal maps in DIRL, making our method an extension of MCE to the case of time-varying rewards.
Revised: We also use discounted future entropy to regularize the goal maps in DIRL, making our method an extension of MCE to the case of time-varying rewards.

Original: Finally, the IRL frameworks of Babes-Vroman et al.
Revised: Finally, the IRL frameworks of Babes-Vroman et al.

Original: [38], Choi and Kim [39] and Likmeta et al.
Revised: [38], Choi and Kim [39] and Likmeta et al.

Original: [40] pursued a related aim to ours, allowing for multiple agents with differing reward functions.
Revised: [40] pursued a related aim to ours, allowing for multiple agents with differing reward functions.

Original: However, these frameworks did not allow for single agents with time-varying rewards.
Revised: However, these frameworks did not allow for single agents with time-varying rewards.

Original: Surana and Srivastava [41] developed a Bayesian non-parametric method that assumed that an agent’s trajectory could be partitioned into distinct behavioral states, where each discrete state had a unique reward function; here, we instead consider the case where rewards vary continuously over time.
Revised: Surana and Srivastava [41] developed a Bayesian non-parametric method that assumed that an agent’s trajectory could be partitioned into distinct behavioral states, where each discrete state had a unique reward function; here, we instead consider the case where rewards vary continuously over time.

Original:  DIRL: Dynamic Inverse Reinforcement Learning 3.1 The Inverse Reinforcement Learning Problem
Revised:  DIRL: Dynamic Inverse Reinforcement Learning 3.1 The Inverse Reinforcement Learning Problem

Original: Let us consider a Markov Decision Process (MDP), M = {S , A , T , r, γ } , where S is the state space,
Revised: Let us consider a Markov Decision Process (MDP), M = {S , A , T , r, γ } , where S is the state space,

Original: A is the action space, T : S×S×A → [0 , 1] represents the probability of transitioning between states when a certain action is taken, r : S ×
Revised: A is the action space, T : S×S×A → [0 , 1] represents the probability of transitioning between states when a certain action is taken, r : S ×

Original: A → R is the reward function, specifying the reward obtained from taking action a ∈ A in state s ∈ S , and γ ∈ [0 , 1] is the discount factor.
Revised: A → R is the reward function, specifying the reward obtained from taking action a ∈ A in state s ∈ S , and γ ∈ [0 , 1] is the discount factor.

Original: Inverse reinforcement learning [17–19, 34, 38, 39] aims to infer the unknown reward function r ( a, s ) when given access to {S , A , T , γ } and N trajectories of agents navigating in this environment, D
Revised: Inverse reinforcement learning [17–19, 34, 38, 39] aims to infer the unknown reward function r ( s, a ) when given access to {S , A , T , γ } and N trajectories of agents navigating in this environment, D

Original: = { ζ 1 , ζ 2 , ..., ζ N } .
Revised: = { ζ 1 , ζ 2 , ..., ζ N } .

Original: Each trajectory is a sequence of state-action pairs, ζ i = { ( s 1 , a 1 ) , ( s 2 , a 2 ) , ... } .
Revised: Each trajectory is a sequence of state-action pairs, ζ i = { ( s 1 , a 1 ) , ( s 2 , a 2 ) , ... } .

Original: Typically, IRL frameworks assume that the reward function does not vary over the course of an agent’s trajectory.
Revised: Typically, IRL frameworks assume that the reward function does not vary over the course of an agent’s trajectory.

Original: As discussed earlier, this is a severe limitation when applying IRL in neuroscientific settings, where the subjective value of the rewards an animal receives may vary as a function of time, satiety, thirst, fatigue, curiosity, etc.
Revised: As discussed earlier, this is a severe limitation when applying IRL in neuroscientific settings, where the subjective value of the rewards an animal receives may vary as a function of time, satiety, thirst, fatigue, curiosity, etc.

Original: In the following, we introduce an extension of IRL that allows for time-varying reward functions.
Revised: In the following, we introduce an extension of IRL that allows for time-varying reward functions.

Original: 3.2 DIRL Generative Model: Goal Maps and Time-Varying Weights
Revised: 3.2 DIRL Generative Model: Goal Maps and Time-Varying Weights

Original: Here we extend the MDP defined above by assuming a time-varying reward function: r t ( a, s ) .
Revised: Here we extend the MDP defined above by assuming a time-varying reward function: r t ( a, s ) .

Original: Our goal is to develop methods for inferring this reward function from a set of observed trajectories
Revised: Our goal is to develop methods for inferring this reward function from a set of observed trajectories

Original: D = { ζ i } Ni =1 .
Revised: D = { ζ i } Ni =1 .

Original: Following previous work [17–19, 34, 38, 39], we assume that the reward function depends only on the state s , allowing us to write the dynamic reward function as r t ( s ) .
Revised: Following previous work [17–19, 34, 38, 39], we assume that the reward function depends only on the state s , allowing us to write the dynamic reward function as r t ( s ) .

Original: However, inferring the reward for every state at every timepoint still requires learning |S| × T parameters.
Revised: However, inferring the reward for every state at every timepoint still requires learning |S| × T parameters.

Original: We therefore make two assumptions to reduce the number of parameters and make inference tractable: (1) we model the reward function as low rank, parameterized by a small number of goal maps modulated by a set of time-varying weights; and (2) we impose a prior encouraging these weights to vary slowly in time.
Revised: We therefore make two assumptions to reduce the number of parameters and make inference tractable: [Equation] we model the reward function as low rank, parameterized by a small number of goal maps modulated by a set of time-varying weights; and (2) we impose a prior encouraging these weights to vary slowly in time.

Original: In our approach, the time-varying reward function has the following low-rank representation:
Revised: In our approach, the time-varying reward function has the following low-rank representation:

Original: [Equation]
Revised: [Equation]

Original: where u k ∈ R S represents the k ’th goal map, and α k,t ∈ R is the weight on this goal map at timestep t , and K is the rank of the representation.
Revised: where u k ∈ R S represents the k ’th goal map, and α k,t ∈ R is the weight on this goal map at timestep t , and K is the rank of the representation.

Original: Each goal map specifies a reward level for each state inthe environment, while weight α k,t specifies the contribution of goal map k to the animal’s reward function at t .
Revised: Each goal map specifies a reward level for each state in

Original: [Figure]
Revised: [Figure]

Original: [Equation]
Revised: [Equation]

Original: [Figure]
Revised: [Figure]

Revised: the environment, while weight α the contribution of goal map k to the animal’s reward

Original: To impose smoothness over time, we place a Gaussian random walk prior over weight trajectories:
Revised: To impose smoothness over time, we place a Gaussian random walk prior over weight trajectories:

Original: [Equation]
Revised: [Equation]

Original: where σ 2 k is a hyperparameter controlling the variance of the weight changes.
Revised: where σ 2 k is a hyperparameter controlling the variance of the weight changes.

Original: This prior reflects our belief that the factors influencing an animal’s subjective experience of reward (e.g., thirst, hunger, fatigue) vary slowly relative to the timescale of individual decisions.
Revised: This prior reflects our belief that the factors influencing an animal’s subjective experience of reward (e.g., thirst, hunger, fatigue) vary slowly relative to the timescale of individual decisions.

Original: The low-rank assumption reduces the number of parameters specifying the reward function from ( |S| × T ) to ( |S| + T ) K , a massive reduction provided K ≪ min( |S| , T ) .
Revised: The low-rank assumption reduces the number of parameters specifying the reward function from ( |S| × T ) to ( |S| + T ) K , a massive reduction provided K ≪ min( |S| , T ) .

Original: The smoothness assumption allows us to further reduce the effective number of parameters, as the weight trajectories α k,t become more correlated with decreasing variance σ 2 k [42].
Revised: The smoothness assumption allows us to further reduce the effective number of parameters, as the weight trajectories α k,t become more correlated with decreasing variance σ 2 k [42].

Original: Figure 2 illustrates the resulting generative model using a simplified 3 × 3 gridworld.
Revised: Figure 2 illustrates the resulting generative model using a simplified 3 × 3 gridworld.

Original: In this example, there are multiple goal maps, one with reward located only at the water port, a second with reward located only at the home state, and a third with reward distributed broadly across states, which is associated with exploratory behavior.
Revised: In this example, there are multiple goal maps, one with reward located only at the water port, a second with reward located only at the home state, and a third with reward distributed broadly across states, which is associated with exploratory behavior.

Original: Each map has a time-varying weight that determines its contribution to the animal’s total reward function at a given moment in time.
Revised: Each map has a time-varying weight that determines its contribution to the animal’s total reward function at a given moment in time.

Original: In this example, the “home” goal map dominates the reward function at the beginning of the session, while the “water” goal map dominates at the end of the session (e.g., reflecting the animal becoming thirsty).
Revised: In this example, the “home” goal map dominates the reward function at the beginning of the session, while the “water” goal map dominates at the end of the session (e.g., reflecting the animal becoming thirsty).

Original: To model decision-making behavior, we assume that animals seek to maximize the discounted expected future reward under a maximum entropy policy [43, 44], given by:
Revised: To model decision-making behavior, we assume that animals seek to maximize the discounted expected future reward under a maximum entropy policy [19, 20, 37, 43, 44], given by:

Original: [Equation]
Revised: [Equation]

Original: where Q t ( s, a ) is the soft Q-function for state s and action a at time t :
Revised: where Q t ( s, a ) is the soft Q-function for state s and action a at time t :

Original: [Equation]
Revised: [Equation]

Original: which arises from the reward function by performing soft value iteration [34] (see Supplementary
Revised: which arises from the reward function by performing soft value iteration [34] (see Supplementary

Original: Materials, SM, for further details).
Revised: Materials, SM, for further details).

Original: This is a common choice of policy in IRL frameworks as it is easily differentiable (unlike maximizing or “greedy” policies), and it has also been widely applied to data from both humans and animals [13, 27].
Revised: This is a common choice of policy in IRL frameworks as it is easily differentiable (unlike maximizing or “greedy” policies), and it has also been widely applied to data from both humans and animals [13, 28].

Original: Note that our formulation does not require a temperature parameter, as this is directly incorporated into the scale of the time-varying weights: larger (smaller) weights give rise to more (less) deterministic policies.
Revised: Note that our formulation does not require a temperature parameter, as this is directly incorporated into the scale of the time-varying weights: larger (smaller) weights give rise to more (less) deterministic policies.

Original: 3.3 DIRL inference procedure
Revised: 3.3 DIRL inference procedure

Original: During inference, our objective is to learn the time-varying weights { α k } Kk =1 , as well as the goal maps { u k } Kk =1 from the trajectories of an agent (animal).
Revised: During inference, our objective is to learn the time-varying weights { α k } Kk =1 , as well as the goal maps { u k } Kk =1 from the trajectories of an agent (animal).

Original: To do so, we alternately optimize the goal maps and time-varying weights to maximize the log-posterior of the observed trajectories in D under our model.
Revised: To do so, we alternately optimize the goal maps and time-varying weights to maximize the log-posterior of the observed trajectories in D under our model.

Original: Let u ∈ R S K be a vector of concatenated goal maps { u k } Kk =1 stacked vertically, and similarly let α ∈ R TK contain the concatenated time-varying weights { α k } Kk =1 .
Revised: Let u ∈ R S K be a vector of concatenated goal maps { u k } Kk =1 stacked vertically, and similarly let α ∈ R TK contain the concatenated time-varying weights { α k } Kk =1 .

Original: stacked vertically.
Revised: stacked vertically.

Original: We, first, initialize the parameters randomly, such that the elements of the goal maps are chosen from U (0 , 1) and the time-varying weights are Gaussian distributed.
Revised: We, first, initialize the parameters randomly, such that the elements of the goal maps are chosen from U (0 , 1) and the time-varying weights are Gaussian distributed.

Original: We then perform coordinate ascent to iteratively update the time-varying weights and the goal maps while holding the other set of parameters constant.
Revised: We then perform coordinate ascent to iteratively update the time-varying weights and the goal maps while holding the other set of parameters constant.

Original: Concretely, we obtain the goal map updates u by maximizing the following objective using gradient ascent:
Revised: Concretely, we obtain the goal map updates u by maximizing the following objective using gradient ascent:

Original: [Equation]
Revised: [Equation]

Original: where π t is the policy given by Eq.
Revised: where π t is the policy given by Eq.

Original: 3 and λ
Revised: 3 and λ

Original: || u
Revised: || u

Original: || 2 represents an L2 regularizer.
Revised: || 2 represents an L2 regularizer.

Original: We then update the time-varying weights α by first updating the reward function and policy (Eq. 3) with the new goal maps, r t ( s )
Revised: We then update the time-varying weights α by first updating the reward function and policy (Eq. 3) with the new goal maps, r t ( s )

Original: =
Revised: =

Original: (cid:80) k
Revised: (cid:80) k

Original: α t,k u ∗ k,s .
Revised: α t,k u ∗ k,s .

Original: We then use gradient ascent to perform the optimization:
Revised: We then use gradient ascent to perform the optimization:

Original: [Equation]
Revised: [Equation]

Original: The last two terms in this objective correspond to the negative log of the Gaussian prior on α (Eq. 2), where C − 1 = D ⊤ Σ − 1 D is the inverse prior covariance, with D a block diagonal matrix of K identical T × T first-order difference matrices (with 1s on the diagonal and -1s on the sub-diagonal), and Σ − 1 is a diagonal matrix with inverse noise-variances
Revised: The last two terms in this objective correspond to the negative log of the Gaussian prior on α (Eq. 2), where C − 1 = D ⊤ Σ − 1 D is the inverse prior covariance, with D a block diagonal matrix of K identical T × T first-order difference matrices (with 1s on the diagonal and -1s on the sub-diagonal), and Σ − 1 is a diagonal matrix with inverse noise-variances

Original: 1 /σ 2 k along the diagonal.
Revised: 1 /σ 2 k along the diagonal.

Original: We iteratively update the goal maps and time-varying weights using Eq.
Revised: We iteratively update the goal maps and time-varying weights using Eq.

Original: 5 and Eq.
Revised: 5 and Eq.

Original: 6 until convergence.
Revised: 6 until convergence.

Original: (See Alg. 1 for pseudo-code).
Revised: (See Alg. 1 for pseudo-code).

Original: We consider the number of goal maps K , the discount factor γ , the strength of the goal map prior λ , as well as the noise variances associated with the time-varying weights, { σ k } Kk =1 , to be hyperparameters.
Revised: We consider the number of goal maps K , the discount factor γ , the strength of the goal map prior λ , as well as the noise variances associated with the time-varying weights, { σ k } Kk =1 , to be hyperparameters.

Original: To restrict the number of hyperparameters, we set σ k = σ ∀ k .
Revised: To restrict the number of hyperparameters, we set σ k = σ ∀ k .

Original: We then swept across a broad range of values for all hyperparameters (see SM for the full list of values considered) and selected the values that optimized the log-likelihood of a set of held-out trajectories.
Revised: We then swept across a broad range of values for all hyperparameters (see SM for the full list of values considered) and selected the values that optimized the log-likelihood of a set of held-out trajectories.

Original: Algorithm 1: DIRL Inference Procedure
Revised: Algorithm 1: DIRL Inference Procedure

Original: Let u =
Revised: Let u =

Original: [ u 1 , ... u k ] , α =
Revised: [ u 1 , ... u k ] , α =

Original: [ α 1 , ...
Revised: [ α 1 , ...

Original: α k ] ;
Revised: α k ] ;

Original: Initialize u 0 , α 0 ; for iter = 1 ...N iter do
Revised: Initialize u 0 , α 0 ; for iter = 1 ...N iter do

Original: Calculate rewards r t ( s )
Revised: Calculate rewards r t ( s )

Original: = (cid:80) k
Revised: = (cid:80) k

Original: α iter k,t u iter k,s ∀ s ∈ S , t ∈ { 1 ...T } ;
Revised: α iter k,t u iter k,s ∀ s ∈ S , t ∈ { 1 ...T } ;

Original: Get policy using soft value iteration : π t ( s, a ) =
Revised: Get policy using soft value iteration : π t ( s, a ) =

Original: e Q t ( s,a ) (cid:80) a ′ e Q t ( s,a ′ ) ∀
Revised: e Q t ( s,a ) (cid:80) a ′ e Q t ( s,a ′ ) ∀

Original: ( s, a, t ) ;
Revised: ( s, a, t ) ;

Original: Update u iter +
Revised: Update u iter +

Original: by maximizing the log-posterior of
Revised: by maximizing the log-posterior of

Original: Eq
Revised: Eq

Original: end
Revised: end

Original: Output u N α N iter , ;
Revised: Output u N α N iter , ;

Original:  Results 4.1 Application to a simulated gridworld environment
Revised:  Results 4.1 Application to a simulated gridworld environment

Original: B weights “home” “water”
Revised: B weights “home” “water”

Original: F
Revised: F

Original: A test LL true goal maps    recovered goal maps “water” “home” rewards end start example trajectories
Revised: A test LL true goal maps    recovered goal maps “water” “home” rewards end start example trajectories

Original: D
Revised: D

Original: C
Revised: C

Original: E true recovered time time #maps home water avg.
Revised: E true recovered time time #maps home water avg.

Original: We first demonstrate our method on simulated trajectories in a 5 × 5 gridworld environment, withactions per state (up, down, left, right, stay).
Revised: We first demonstrate our method on simulated trajectories in a 5 × 5 gridworld environment, withactions per state (up, down, left, right, stay).

Original: We generated two goal maps for this environment: a “home” map and a “water” map, which were rewarding only at the home state and the water state, respectively (Fig 3B).
Revised: We generated two goal maps for this environment: a “home” map and a “water” map, which were rewarding only at the home state and the water state, respectively (Fig 3B).

Original: Corresponding to these goal maps, we also generated time-varying weights (Fig.
Revised: Corresponding to these goal maps, we also generated time-varying weights (Fig.

Original: 3C, solid lines) for 50 timesteps with the random-walk prior of Eq.
Revised: 3C, solid lines) for 50 timesteps with the random-walk prior of Eq.

Original: 2 (for σ k = 2 − 3 . 5 ; this was chosen to provide adequate variation in the reward function during the time period considered).
Revised: 2 (for σ k = 2 − 3 . 5 ; this was chosen to provide adequate variation in the reward function during the time period considered).

Original: The weight for the water map started high but decreased over time, thus making the water state the most rewarding for the first ∼ 25 timesteps (Fig.
Revised: The weight for the water map started high but decreased over time, thus making the water state the most rewarding for the first ∼ 25 timesteps (Fig.

Original: 3E, blue solid line).
Revised: 3E, blue solid line).

Original: In contrast, the weight on the home map started small but increased so that the home state became the most rewarding state at the end of the 50 timesteps (Fig.
Revised: In contrast, the weight on the home map started small but increased so that the home state became the most rewarding state at the end of the 50 timesteps (Fig.

Original: 3E, red solid line).
Revised: 3E, red solid line).

Original: All of the other states in the environment had a constant reward of 0.
Revised: All of the other states in the environment had a constant reward of 0.

Original: In order to generate trajectories corresponding to this reward function (Fig. 3E), we used soft value iteration to learn the corresponding optimal time-varying policy (Eq. 3).
Revised: In order to generate trajectories corresponding to this reward function (Fig. 3E), we used soft value iteration to learn the corresponding optimal time-varying policy (Eq. 3).

Original: We then executed this policy in order to obtain 200 trajectories (a similar number to the number of trajectories we have for the real dataset discussed later), two examples of which are shown in Fig.
Revised: We then executed this policy in order to obtain 200 trajectories (a similar number to the number of trajectories we have for the real dataset discussed later), two examples of which are shown in Fig.

Original: 3A.
Revised: 3A.

Original: Next, we applied our IRL inference method so as to learn the goal maps and time-varying weights from 80% of the generated trajectories.
Revised: Next, we applied our IRL inference method so as to learn the goal maps and time-varying weights from 80% of the generated trajectories.

Original: Fig.
Revised: Fig.

Original: 3F shows the log-likelihood of the remaining 20% of trajectories as a function of the number of goal maps.
Revised: 3F shows the log-likelihood of the remaining 20% of trajectories as a function of the number of goal maps.

Original: The held-out test log-likelihood is equally high for 2 and 3 maps, so we focus on the 2-map solution in order to be able to compare with the generative parameters.
Revised: The held-out test log-likelihood is equally high for 2 and 3 maps, so we focus on the 2-map solution in order to be able to compare with the generative parameters.

Original: It is important to note that we can only recover rewards at each timestep up to an additive constant, as the policy remains unchanged upon the addition of a constant to the rewards.
Revised: It is important to note that we can only recover rewards at each timestep up to an additive constant, as the policy remains unchanged upon the addition of a constant to the rewards.

Original: Further, scaling of the goal maps accompanied by an inverse scaling of the time-varying weights also leaves the recovered rewards unchanged.
Revised: Further, scaling of the goal maps accompanied by an inverse scaling of the time-varying weights also leaves the recovered rewards unchanged.

Original: Thus, to compare to the generative parameters, we perform a post hoc processing method to the recovered parameters to handle all such invariances (details in
Revised: Thus, to compare to the generative parameters, we perform a post hoc processing method to the recovered parameters to handle all such invariances (details in

Original: SM).
Revised: SM).

Original: Figures 3C and D show that our method allows us to accurately recover the generative goal maps and time-varying weights from the simulated trajectories.
Revised: Figures 3C and D show that our method allows us to accurately recover the generative goal maps and time-varying weights from the simulated trajectories.

Original: Finally, combining the goal maps and time-varying weights, we are able to accurately match the generative time-varying rewards for different states in the gridworld (Fig.
Revised: Finally, combining the goal maps and time-varying weights, we are able to accurately match the generative time-varying rewards for different states in the gridworld (Fig.

Original: 3E).
Revised: 3E).

Original: We also simulated trajectories in a 127-node labyrinth environment (akin to [15], Fig 1), and confirmed that we were able to recover goal maps and timevarying weights (see SM for details).
Revised: We also simulated trajectories in a 127-node labyrinth environment (akin to [15], Fig 1), and confirmed that we were able to recover goal maps and timevarying weights (see SM for details).

Original: We focus on the gridworld simulations here to demonstrate the versatility of our approach across environments.
Revised: We focus on the gridworld simulations here to demonstrate the versatility of our approach across environments.

Original: 4.2 Application of DIRL to real mouse trajectories
Revised: 4.2 Application of DIRL to real mouse trajectories

Original: Next, we applied our framework to the trajectories of real mice navigating (in the dark) in a 127-node labyrinth environment [15].
Revised: Next, we applied our framework to the trajectories of real mice navigating (in the dark) in a 127-node labyrinth environment [15].

Original: In this task, two cohorts of 10 mice moved freely through the labyrinth over the course of 7 hours.
Revised: In this task, two cohorts of 10 mice moved freely through the labyrinth over the course of 7 hours.

Original: The first cohort was water-restricted and the mice were provided with a water port at one of the terminal nodes (shown in gray in Fig. 1).
Revised: The first cohort was water-restricted and the mice were provided with a water port at one of the terminal nodes (shown in gray in Fig. 1).

Original: The second cohort was not water-restricted and did not have access to the water port.
Revised: The second cohort was not water-restricted and did not have access to the water port.

Original: We show an example trajectory for an animal in each cohort in Fig.
Revised: We show an example trajectory for an animal in each cohort in Fig.

Original: Over the course of the night, each animal completed over a hundred such trajectories, with some animals completing many more.
Revised: Over the course of the night, each animal completed over a hundred such trajectories, with some animals completing many more.

Original: [Figure]
Revised: [Figure]

Original: 4.2.1 Inferring interpretable reward functions from water-restricted mice
Revised: 4.2.1 Inferring interpretable reward functions from water-restricted mice

Original: We began our investigation by applying our method to a subset of trajectories for the water-restricted animals.
Revised: We began our investigation by applying our method to a subset of trajectories for the water-restricted animals.

Original: Here, we anticipated being able to identify the water port as highly rewarding (in contrast to the water-unrestricted animals).
Revised: Here, we anticipated being able to identify the water port as highly rewarding (in contrast to the water-unrestricted animals).

Original: Due to the high variability in trajectories across animals and over the course of the 7 hours, and to be able to obtain trajectories corresponding to similar goal maps and with a similar time course, we used an unsupervised clustering algorithm (based on dbscan [45] and using the Levenshtein distance metric; full details are in the SM) to identify similar trajectories.
Revised: Due to the high variability in trajectories across animals and over the course of the 7 hours, and to be able to obtain trajectories corresponding to similar goal maps and with a similar time course, we used an unsupervised clustering algorithm (based on dbscan [45] and using the Levenshtein distance metric; full details are in the SM) to identify similar trajectories.

Original: This procedure had the advantageous side effect of excluding the first 25 or so trajectories for each animal: our focus is not on characterizing learning.
Revised: This procedure had the advantageous side effect of excluding the first 25 or so trajectories for each animal: our focus is not on characterizing learning.

Original: Overall, we obtained 200 trajectories for this cohort.
Revised: Overall, we obtained 200 trajectories for this cohort.

Original: Each timestep in a trajectory was recorded at a one second interval.
Revised: Each timestep in a trajectory was recorded at a one second interval.

Original: We then fit the goal maps and time-varying weights to 160 of these trajectories, and held out the remaining 20% of trajectories as a validation set.
Revised: We then fit the goal maps and time-varying weights to 160 of these trajectories, and held out the remaining 20% of trajectories as a validation set.

Original: We found that validation log-likelihood began to level off at two maps (Fig. 4C), so we focus on the 2 map solution here.
Revised: We found that validation log-likelihood began to level off at two maps (Fig. 4C), so we focus on the 2 map solution here.

Original: The 3 map solution is shown in the SM (where the recovered water goal map is simply repeated).
Revised: The 3 map solution is shown in the SM (where the recovered water goal map is simply repeated).

Original: Overall, we found that we were, indeed, able to recover a “water” goal map (Fig 4A), with a large reward at the water state and small rewards elsewhere, as well as a “home” goal map with a large reward at the home state.
Revised: Overall, we found that we were, indeed, able to recover a “water” goal map (Fig 4A), with a large reward at the water state and small rewards elsewhere, as well as a “home” goal map with a large reward at the home state.

Original: The recovered time-varying varying weights (Fig 4B) for the water map were high at the beginning – the mouse likely entered the labyrinth due to being thirsty – and later on, as satiation and fatigue set in, tailed off.
Revised: The recovered time-varying varying weights (Fig 4B) for the water map were high at the beginning – the mouse likely entered the labyrinth due to being thirsty – and later on, as satiation and fatigue set in, tailed off.

Original: In contrast, the weights corresponding to the home map started low, but increased over time.
Revised: In contrast, the weights corresponding to the home map started low, but increased over time.

Original: The final reward function (Fig 4D) reflects the same dynamics: the water state was highly rewarding for the mouse at the start of its trajectory, while the home state became the more rewarding state after ∼ 6 seconds.
Revised: The final reward function (Fig 4D) reflects the same dynamics: the water state was highly rewarding for the mouse at the start of its trajectory, while the home state became the more rewarding state after ∼ 6 seconds.

Original: The other states in the labyrinth offered only a small intrinsic reward for the mouse throughout the course of its trajectory.
Revised: The other states in the labyrinth offered only a small intrinsic reward for the mouse throughout the course of its trajectory.

Original: 4.2.2 ‘Exploratory’ maps inferred from water-unrestricted mice
Revised: 4.2.2 ‘Exploratory’ maps inferred from water-unrestricted mice

Original: We then moved on to examining the trajectories of the water-unrestricted cohort.
Revised: We then moved on to examining the trajectories of the water-unrestricted cohort.

Original: Relative to the water-restricted mice, the goal maps for these mice were not obvious a priori – there were no extrinsic reward locations for this cohort in the labyrinth.
Revised: Relative to the water-restricted mice, the goal maps for these mice were not obvious a priori – there were no extrinsic reward locations for this cohort in the labyrinth.

Original: We began by applying the same clustering algorithm as for the water-restricted animals (discussed in the SM) and identified 207 trajectories for our analysis.
Revised: We began by applying the same clustering algorithm as for the water-restricted animals (discussed in the SM) and identified 207 trajectories for our analysis.

Original: We fit the time-varying weights and goal maps to 80% of these trajectories, and held out the remaining 20% as a validation set.
Revised: We fit the time-varying weights and goal maps to 80% of these trajectories, and held out the remaining 20% as a validation set.

Original: As we show in Fig.
Revised: As we show in Fig.

Original: 5C, the validation log-likelihood is best for the 2 map solution.
Revised: 5C, the validation log-likelihood is best for the 2 map solution.

Original: While we recover (Fig 5A) a “home” map once again, a new map – that was not present for the water-restricted animals – also appears.
Revised: While we recover (Fig 5A) a “home” map once again, a new map – that was not present for the water-restricted animals – also appears.

Original: This map is rewarding at many states throughout the maze, but is very unrewarding at the home port, and is very rewarding at a state that is close to the home port (which we refer to as the “explore” state).
Revised: This map is rewarding at many states throughout the maze, but is very unrewarding at the home port, and is very rewarding at a state that is close to the home port (which we refer to as the “explore” state).

Original: Hence, when the animal places a positive weight on this map, the mouse wants to leave the home state and venture into the maze.
Revised: Hence, when the animal places a positive weight on this map, the mouse wants to leave the home state and venture into the maze.

Original: For these reasons, we refer to this map as the “explore” map.
Revised: For these reasons, we refer to this map as the “explore” map.

Original: The time-varying weights (Fig 5B) for the home map increase with time; this captures the tendency of these animals to go back to the home state towards the end of the trajectory.
Revised: The time-varying weights (Fig 5B) for the home map increase with time; this captures the tendency of these animals to go back to the home state towards the end of the trajectory.

Original: The weights corresponding to the explore map are positive at the beginning of the trajectory, capturing the exploratory behavior of these animals when they enter the labyrinth.
Revised: The weights corresponding to the explore map are positive at the beginning of the trajectory, capturing the exploratory behavior of these animals when they enter the labyrinth.

Original: Finally, the inferred rewards (Fig 5D) reveal that, as we would expect for this cohort, the time-varying reward for the water port is almost identical to the average reward function (across all other states in the labyrinth).
Revised: Finally, the inferred rewards (Fig 5D) reveal that, as we would expect for this cohort, the time-varying reward for the water port is almost identical to the average reward function (across all other states in the labyrinth).

Original: We can also see that the reward for the home state rises much more rapidly over time than the reward for other states.
Revised: We can also see that the reward for the home state rises much more rapidly over time than the reward for other states.

Original: 4.2.3 DIRL outperforms existing IRL approaches
Revised: 4.2.3 DIRL outperforms existing IRL approaches

Original: Finally, we compare the performance of our method with two popular IRL frameworks: the maximum entropy IRL framework of Ziebart et al.
Revised: Finally, we compare the performance of our method with two popular IRL frameworks: the maximum entropy IRL framework of Ziebart et al.

Original: [19] and the deep maximum entropy IRL framework of
Revised: [19] and the deep maximum entropy IRL framework of

Original: Wulfmeier et al.
Revised: Wulfmeier et al.

Original: In comparison to DIRL, these methods learn a static reward function.
Revised: In comparison to DIRL, these methods learn a static reward function.

Original: We use an open-source implementation of these frameworks [46] and infer the reward functions for each of these methods for the two cohorts of mice studied above.
Revised: We use an open-source implementation of these frameworks [46] and infer the reward functions for each of these methods for the two cohorts of mice studied above.

Original: Using the retrieved reward functions, we then obtain the corresponding optimal policy and use it to compute the log-likelihood of the validation set of trajectories.
Revised: Using the retrieved reward functions, we then obtain the corresponding optimal policy and use it to compute the log-likelihood of the validation set of trajectories.

Original: In Figures 4C and 5C, we demonstrate that our method dramatically outperforms these existing methods at explaining the held-out trajectories of both water-restricted and unrestricted mouse cohorts.
Revised: In Figures 4C and 5C, we demonstrate that our method dramatically outperforms these existing methods at explaining the held-out trajectories of both water-restricted and unrestricted mouse cohorts.

Original: In the SM, we also show that the trajectories generated by DIRL very well resemble the behavior of mice in the labyrinth, as compared to those generated from [19] and [34].
Revised: In the SM, we also show that the trajectories generated by DIRL very well resemble the behavior of mice in the labyrinth, as compared to those generated from [19] and [34].

Original:  Discussion
Revised:  Discussion

Original: In this work we develop DIRL, a novel inverse reinforcement learning framework for characterizing the behavior of animals during complex decision-making tasks.
Revised: In this work we develop DIRL, a novel inverse reinforcement learning framework for characterizing the behavior of animals during complex decision-making tasks.

Original: Our framework infers the timevarying reward functions of animals from their trajectories alone.
Revised: Our framework infers the timevarying reward functions of animals from their trajectories alone.

Original: We validated our framework on simulated data in a gridworld environment, and applied it to two cohorts of 10 mice navigating in a labyrinth [15].
Revised: We validated our framework on simulated data in a gridworld environment, and applied it to two cohorts of 10 mice navigating in a labyrinth [15].

Original: Our method provided distinct and interpretable reward functions for both cohorts: the water-restricted mice assigned a high reward to the water state upon entering the labyrinth, while the water-unrestricted mice were motivated to leave home and explore at the beginning of their trajectories.
Revised: Our method provided distinct and interpretable reward functions for both cohorts: the water-restricted mice assigned a high reward to the water state upon entering the labyrinth, while the water-unrestricted mice were motivated to leave home and explore at the beginning of their trajectories.

Original: As time passed and the mice became fatigued, the reward for the home state increased for animals in both cohorts.
Revised: As time passed and the mice became fatigued, the reward for the home state increased for animals in both cohorts.

Original: Our method dramatically outperfomed IRL methods with fixed reward
Revised: Our method dramatically outperfomed IRL methods with fixed reward

Original: [Figure]
Revised: [Figure]

Original: functions, as quantified by the log-likelihood of held-out trajectories, indicating a clear need for tailored IRL approaches for neuroscientific applications.
Revised: functions, as quantified by the log-likelihood of held-out trajectories, indicating a clear need for tailored IRL approaches for neuroscientific applications.

Original: Finally, our method is computationally efficient and infers the time-varying reward function from ∼ 4000 decisions in the 127-node labyrinth environment in 20 minutes on a laptop.
Revised: Finally, our method is computationally efficient and infers the time-varying reward function from ∼ 4000 decisions in the 127-node labyrinth environment in 20 minutes on a laptop.

Original: One exciting finding of our work is the discovery of an “explore” goal map for the water-unrestricted mouse cohort.
Revised: One exciting finding of our work is the discovery of an “explore” goal map for the water-unrestricted mouse cohort.

Original: In general, exploration in animals is not well understood [25].
Revised: In general, exploration in animals is not well understood [26].

Original: DIRL offers a novel unsupervised approach for characterizing exploration from behavior alone.
Revised: DIRL offers a novel unsupervised approach for characterizing exploration from behavior alone.

Original: Future work could build upon the framework we present here in order to provide normative explanations for the nature of the explore map.
Revised: Future work could build upon the framework we present here in order to provide normative explanations for the nature of the explore map.

Original: DIRL could also be used to explore the neural underpinnings of exploration, for example by correlating fluctuations in the inferred time-varying rewards with recorded neural activity.
Revised: DIRL could also be used to explore the neural underpinnings of exploration, for example by correlating fluctuations in the inferred time-varying rewards with recorded neural activity.

Original: Finally, while we developed DIRL with neuroscientific applications in mind, our framework is general and could be applied in other scientific fields.
Revised: Finally, while we developed DIRL with neuroscientific applications in mind, our framework is general and could be applied in other scientific fields.

Original: For example, DIRL could be applicable in healthcare settings, such as those considered in Hüyük et al.
Revised: For example, DIRL could be applicable in healthcare settings, such as those considered in Hüyük et al.

Original: In Hüyük et al.
Revised: In Hüyük et al.

Original: [47], the authors acknowledge that the policies used by healthcare providers to allocate organ donations have varied over time.
Revised: [47], the authors acknowledge that the policies used by healthcare providers to allocate organ donations have varied over time.

Original: By applying DIRL to such a dataset, it may be possible to recover the latent goals of policy-makers or medical practitioners when deciding how best to allocate organs to patients.
Revised: By applying DIRL to such a dataset, it may be possible to recover the latent goals of policy-makers or medical practitioners when deciding how best to allocate organs to patients.

Original: In addition to the above-mentioned scientific applications, we believe that DIRL also provides future directions for the development of RL algorithms.
Revised: In addition to the above-mentioned scientific applications, we believe that DIRL also provides future directions for the development of RL algorithms.

Original: Relative to animals, it is often challenging to get artificial agents to explore in sparse reward environments [48].
Revised: Relative to animals, it is often challenging to get artificial agents to explore in sparse reward environments [48].

Original: With access to the internal reward function that motivates mice, we hope that our framework may be useful for inspiring better reward functions for training artificial agents to navigate in novel environments (such as in [49] where artificial agents navigated in an analogous depth-6 binary tree environment).
Revised: With access to the internal reward function that motivates mice, we hope that our framework may be useful for inspiring better reward functions for training artificial agents to navigate in analogous environments (such as in [49] where artificial agents navigated in an analogous depth-6 binary tree environment).

Original: Another valuable future direction could involve developing alternative ways of formulating the problem of IRL with time-varying goals, such as casting our problem as a Hidden-Mode Hidden Markov Model [50] —a type of POMDP— and then performing IRL with POMDPs [51].
Revised: Another valuable future direction could involve developing alternative ways of formulating the problem of IRL with time-varying goals, such as casting our problem as a Hidden-Mode Hidden Markov Model [50] —a type of POMDP— and then performing IRL with POMDPs [51].

Original: We will close by discussing some limitations of our work.
Revised: We will close by discussing some limitations of our work.

Original: Firstly, our framework requires over a hundred trajectories to infer an animal’s time-varying reward function.
Revised: Firstly, our framework requires over a hundred trajectories to infer an animal’s time-varying reward function.

Original: While it may be possible to reduce this number (with, for example, a careful choice of prior), it is easy to conceive of failure modes where a single decision reveals nothing about the active goal map.
Revised: While it may be possible to reduce this number (with, for example, a careful choice of prior), it is easy to conceive of failure modes where a single decision reveals nothing about the active goal map.

Original: In practice, having access to multiple trajectories (or several decisions) can significantly reduce the uncertainty in the recovered time-varying weights and goal maps.
Revised: In practice, having access to multiple trajectories (or several decisions) can significantly reduce the uncertainty in the recovered time-varying weights and goal maps.

Original: Next, our approach may not scale well to high-dimensional state spaces, as we currently learn a separate reward for each state in each goal map.
Revised: Next, our approach may not scale well to high-dimensional state spaces, as we currently learn a separate reward for each state in each goal map.

Original: We don’t anticipate this being a problem in the neuroscientific applications that we discuss in this work (where the state-space is 127 dimensional) but an extension of our framework could involve learning a low-dimensional representation of each state in the environment via a deep network, and then learning time-varying rewards in this embedding space.
Revised: We don’t anticipate this being a problem in the neuroscientific applications that we discuss in this work (where the state-space is 127 dimensional) but an extension of our framework could involve learning a low-dimensional representation of each state in the environment via a deep network, and then learning time-varying rewards in this embedding space.

Original: This would allow us to scale our approach to higher dimensional state spaces, while also allowing for generalization across states.
Revised: This would allow us to scale our approach to higher dimensional state spaces, while also allowing for generalization across states.

Original: Finally, our inferred rewards rely on the animal’s policy being the Boltzmann policy of Eq.
Revised: Finally, our inferred rewards rely on the animal’s policy being the Boltzmann policy of Eq.

Original: However, assuming this form for the policy is not atypical in applications of reinforcement learning to human or animal decision-making data [13, 27].
Revised: However, assuming this form for the policy is not atypical in applications of reinforcement learning to human or animal decision-making data [13, 28].

Original: Overall, we believe that the advantages of using our method outweigh its limitations, and that we present a new, flexible framework for characterizing animal behavior in complex environments.
Revised: Overall, we believe that the advantages of using our method outweigh its limitations, and that we present a new, flexible framework for characterizing animal behavior in complex environments.

Original: Acknowledgments and Disclosure of Funding
Revised: Acknowledgments and Disclosure of Funding

Original: We thank Benjamin Cowley, Nathaniel Daw, Rachit Dubey and Karthik Narasimhan for useful comments and discussions at various points during this project.
Revised: We thank Benjamin Cowley, Nathaniel Daw, Rachit Dubey and Karthik Narasimhan for useful comments and discussions at various points during this project.

Original: We thank the anonymous NeurIPS reviewers for their insightful feedback and helpful suggestions for improvement.
Revised: We thank the anonymous NeurIPS reviewers for their insightful feedback and helpful suggestions for improvement.

Original: The mouse in Figure 1 was provided by BioRender.
Revised: The mouse in Figure 1 was provided by BioRender.

Original: This work was supported by grants from the Simons Collaboration on the Global Brain (SCGB AWD543027), the NIH BRAIN initiative (NS104899 and R01EB026946), and a U19 NIH-NINDS BRAIN Initiative Award (5U19NS104648).
Revised: This work was supported by grants from the Simons Collaboration on the Global Brain (SCGB AWD543027), the NIH BRAIN initiative (NS104899 and R01EB026946), and a U19 NIH-NINDS BRAIN Initiative Award (5U19NS104648).