Dynamic Inverse Reinforcement Learning for Characterizing Animal Behavior
Authors: Zoe Ashwood
Aditi Jha
Jonathan W. Pillow
Venue: NeurIPS 2022 Accept (2022)
Keywords: Neuroscience
decision-making
inverse reinforcement learning

Abstract: 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. 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. 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. However, IRL has yet to be widely applied in neuroscience. One potential reason for this is that existing IRL frameworks assume that an agent's reward function is fixed over time. To address this shortcoming, we introduce dynamic inverse reinforcement learning (DIRL), a novel IRL framework that allows for time-varying intrinsic rewards. Our method parametrizes the unknown reward function as a time-varying linear combination of spatial reward maps (which we refer to as "goal maps"). We develop an efficient inference method for recovering this dynamic reward function from behavioral data.  We demonstrate DIRL in simulated experiments and then apply it to a dataset of mice exploring a labyrinth. Our method returns interpretable reward functions for two separate cohorts of mice, and provides a novel characterization of exploratory behavior. We expect DIRL to have broad applicability in neuroscience, and to facilitate the design of biologically-inspired reward functions for training artificial agents. 

Dynamic Inverse Reinforcement Learning for Characterizing Animal Behavior
Zoe C. Ashwood 1.2,
Aditi 1.3,
1.3,*
Jonathan W. Pillow!
Princeton Neuroscience Institute, Princeton University
Dept. of Computer Science, Princeton University
Dept. of Electrical and Computer Engineering, Princeton University
{zashwood, aditijha, pillow)@princeton.cc
Abstract
Understanding decision decision-making is a core objective in both neuroscience and psy- chology, and computational models have often been helpful in the pursuit of this goal. While many models have been developed for characterizing behavior binary decision-making and bandit tasks, comparatively little work has focused on animal decision-mane mariing in more complex tasks, such as navigation through a maze. Inverse reinforcement learning (IRL) is promising approach for under- standing such behavior, as it aims to infer the unknown reward function of an agent from its observed trajectories through state space. However, IRL has yet be widely applied in neuroscience. One potential reason for this is that existing IRL frameworks assume that agent's reward function is fixed over time. To address this shortcoming, we introduce dynamic inverse reinforcement learning (DIRL), a novel IRL framework that allows for time-varying intrinsic rewards. Our method parametrizes the unknown reward function as time-varying linear combination of spatial reward maps (which we refer as "goal maps"). We develop an efficient inference method for recovering this dynamic reward function from behavioral data. We demonstrate DIRL simulated experiments and then apply it to dataset mice exploring labyrinth. Our method returns interpretable reward functions for two separate cohorts mice, and provides novel characterization of exploratory behavior. We expect DIRL to have broad applicability in neuroscience, and to facilitate the design of biologically-inspire reward functions for training artificial agents.
Introduction
Characterizing the decision-making behavior of humans and animals is central goal in neuroscience and psychology 2 Decision-making tasks such as Two-Alternative Forced Choice (2AFC) and bandit problems have been widely studied [7 and previous work has developed variety of models for behavior in these tasks 9 The classic psychometric curve represents one such model [TO), and more recent work has focused models based reinforcement learning [I] [1, Such models allow us to understand and compare the decision makiing strategies used by humans and animals, and can also provide low-dimensional description of behavior that can be regressed against neural data 1
Although large literature has focused models of decision-making in simple 2AFC and bandit tasks, comparatively few papers have sought to model behavior in larger, complex natural environments.
'These authors contributed equally to this work.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

Figure 1: Example mouse trajectories in the labyrinth task of Rosenberg et al. [15). Water- restricted and unrestricted mice moved freely through a 127 node maze environment for 7 hours. Over the course of the night, each mouse com- pleted over 100 such trajectories, which began at the home port (indicated by an image of house). For the water-restricted mice, a water port existed the terminal node in the environment that is shaded grey (next to the water drop image).
labyrinth task: example trajectories water-restricted
water-unrestricted
Start
end
In recent study, Rosenberg et al. [15] introduced a novel experimental paradigm involving water-restricted mice navigating 127-node labyrinth equipped with a water port at one terminal node (Fig. D. At each node, the mouse could make up to distinct decisions ('stay', 'go left', 'go right', or 'reverse"). Navigation through a maze is a perfect example of complex yet natural decision-make behavior that remains poorly understood. Although reinforcement learning may seem like the natural framework for modeling such goal-driven behavior, the rewards experienced by these mice not obvious to the experimenter. Indeed, as noted in Rosenberg et al. [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.
14
Inverse reinforcement learning (IRL) (7-20 addresses the problem of inferring the unknown reward function of an agent. Given access the agent's 's trajectories as it interacts with the environment, IRL. identifies the states and actions that the agent finds rewarding. While IRL. has found many successes in roboties [21 (2 and healthcare contexts 25) existing IRL. methods have not found broad application in neuroscience. 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. Yet in many real-world decision-making tasks, rewards can change over time. 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. This motivated to develop an IRL method for characterizing animal decision -aaiin in complex environments such as the labyrinth, in which rewards may vary over time.
Here we propose dynamic inverse reinforcement learning (DIRL), an IRL method that allows for time-varying reward functions. DIRL parametrizes the animal's 's reward function as a time-varying linear combination of a small number of spatial reward maps, which we refer to as "goal maps". A single goal map is a function specifying the amount of reward at each state in the environment. (Classic IRL methods thus seek to identify a single, fixed goal map.) DIRL posits the existence of multiple goal maps with time-varying weights, which allows the instantaneous reward function to vary in time. 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. varying weights then modulate the extent to which each goal map is active at any timestep. We introduce an inference method, which extends fixed reward IRL framework by Ziebart et al. (20), 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 interpretable reward functions for both cohorts of mice studied there, and reveals an "explore" goal map for one of the cohorts. While exploration remains poorly understood in neuroscience [26], our method offers a powerful framework for characterizing exploratory, well exploitative behavior from an animal's 's trajectories alone.
data from Rosenberg et al. [1]. Our method recovers
2 Related Work
The neuroscientific literature describing models of decision- decision-making is vast, including both normative and descriptive models 9
10 27 330 While some of this work considers models of decision-
making 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. To the best of our knowledge, IRL has found relatively few applications in neuroscience. One exception is Schultheis et al. [3]

---

which performed inverse optimal control to infer the cost function associated with sensorimotor behavior. Their targeted poppiiation-modeeee
the cost function optimized by a human performing a reaching task-is very different to ours, which results in different assumptions in their work. Another exception Yamaguchi et 32] which descrihed an IRI. framework for identifying the thermotactic strategies used by C. elegans. Relative to the work we present here, Yamaguchi et al. [32) assumed a fixed reward function, and was restricted to the case of a linearly-solvable MDP. In Kwon et al.
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. Finally, while not the main focus of their paper, Reddy [16] used the static IRL framework of Ziebart et al. [9] to infer the rewards optimized by mice navigating in the Rosenberg et al. [15] task.
The literature on IRL. 20
is extensive, but most studies have assumed that the true underlying reward function is fixed over time. Of particular relevance to our work is the Maximum Causal Entropy (MCE) framework of Ziebart et al. (20). 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. 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.
Finally, the IRL frameworks of Babes-Vroman et al. 38 Choi and Kim [39 and Likmeta et al. [40
pursued a related aim to ours, allowing for multiple agents with differing reward functions. However, these frameworks did not allow for single agents with time-varying rewards. 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 unique reward function; here, we instead consider the case where rewards vary continuously over time.
3 DIRL: Dynamic Inverse Reinforcement Learning
3.1 The Inverse Reinforcement Learning Problem
Let us consider a Markov Decision Process (MDP), M = AS where S is the state space, A is the action space, T: S S x A - [0, 1] represents the probability of transitioning between states when a certain action is taken, r : S x A + IR is the reward function, specifying the reward obtained from taking action a E A in state 8 S, and y [0, 1] is the discount factor. Inverse reinforcement learning 17 34 38 39 aims to infer the unknown reward function r(s, a) when given access to AS,A, (S,,,T,7) } and N trajectories of agents navigating in this environment, D {Ci, S2 ...; C,): Each trajectory is a sequence of state-action pairs, C = {(s,, a1), (52, (2).... Typically, IRL frameworks assume that the reward function does not vary over the course of agent's 's trajectory. 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 function of time, satiety. thirst, fatigue, curiosity, etc. In the following, we introduce an extension of IRL that allows for time-varying reward functions.
3.2 DIRL Generative Model: Goal Maps and Time- Varrin Weights
Here we extend the MDP defined above by assuming time-varying reward function: (a s). Our goal is to develop methods for inferring this reward function from a set of observed trajectories D J Following previous work 11-11
38 39), we assume that the reward function depends only on the state 8, allowing us to write the dynamic reward function as (s). However, inferring the reward for every state at every timepoint still requires learning [S] x T parameters. 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 small number of goal maps modulated by a set of time-vare varying weights; and (2) we impose a prior encouraging these weights to vary slowly in time.
In our approach, the time- varying reward function has the following low-rank representation:
(s()
ag. Uk,
where Un E R represents the k'th goal map, and ORR E RR the weight on this goal map at timestep t, and K is the rank of the representation. Each goal map specifies reward level for each state in
3

---

DIRL
goal maps
weights
water
rewards
home
home
water
avg.
other states
time
explore
time
Figure 2: DIRL model schematic. We parametrize the time-varying reward function as linear combination of a small number of spatial goal maps, where the weights on each map can vary over time. A goal map is map of the environment, indicating the reward available in each state.
the environment, while weight ai specifies the contribution of goal map & to the animal's 's reward function at 6.
To impose smoothness over time, we place a Gaussian random walk prior over weight trajectories:
agec 0 + Ebs
E/ (N,00
where of is a hyperparameter controlling the variance of the weight changes. This prior reflects our belief that the factors influencing an animal's subjective experience of reward (e.g.. thirst, hunger, fatigue) vary slowly relative the timescale of individual decisions. The low-rank assumption reduces the number of parameters specifying the reward function from ([S) x T) to (IS) + T) K, massive reduction provided K < mi, T). The smoothness assumption allows us further reduce the effective number of parameters, as the weight trajectories ark. become more correlated with decreasing variance of 442.
Figure 2 illlusrrates the resulting generative model using a simplified 3 x 3 gridworld. In this example, there are multiple goal maps, one with reward located only at the water port, second with reward located only at the home state, and third with reward distributed broadly across states, which associated with exploratory behavior. Each map has a time-varying weight that determines its contribution the animal's 's total reward function at given moment time. In this example, the "home" goal map dominates the reward function at the beginning of the session, while the "water" goal map dominates the end of the session (e.g., reflecting the animal becoming thirsty).
To model decision- maniing behavior, we assume that animals seek maximize the discounted expected future reward under maximum entropy policy [19 20
44), given by:
(a.a)
(s, a)
S A,E (1,...T).
(3)
Erees (s.ar))
where Q'(s, a) is the soft Q-function for state s and action at time
Q'(s, a) r(( + PP(s s,)) log Σ Exsper
(4)
which arises from the reward function by performing soft value iteration (see Supplementary Materials, SM, for further details). This is 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 [3] 28 Note that our formulation does not require a temperature parameter, as this directly incorporated into the scale of the time-varying weights: larger (smaller) weights give rise to more (less) deterministic policies.

---

3.3 DIRL inference procedure
During inference, our objective to learn the time- timevnrying weights {ax) as well the goal maps {uz Ju from the trajectories of an agent (animal). To do so, we alternately optimize the goal maps and viry-ng weights to maximize the log-posterior of the observed trajectories in 'D under our model. Let RSS be a vector of concatenated goal maps {u=) stacked vertically, and similarly let a E RTK contain the concatenated time-varying weights [= stacked vertically. We, first, initialize the parameters randomly, such that the elements of the goal maps are chosen from (0, 1) and the time-varying weights are Gaussian distributed. We then perform coordinate ascent to iteratively update the time-varying weights and the goal maps while holding the other set of parameters constant.
Concretely, we obtain the goal map updates u by maximizing the following objective using gradient ascent:
u arg Max
Σ
log (st, az) - A/|u||2
(5)
where n. is the policy given by Eq. 3 and 2||u||2 represents L.2 regularizer.
We then update the tire-narying weights a by first updating the reward function and policy (Fq. with the new goal maps, (s(s)= En a, We then gradient ascent to perform the optimiza- tion:
a arg
Σ Σ
logme((s a.) 2 log -C -
(6)
i=1 (Se.Re)NG,
The last terms in this objective correspond the negative log of the Gaussian prior (Eq. 2
where C-1 = D E the inverse prior covariance, with D a block diagonal matrix of identical T' first-order difference matrices (with Is on the diagonal and -Is the sub-diagonal), and Σ- is a diagonal matrix with inverse noise-variances 1/0 along the diagonal.
We iteratively update the goal maps and time-varying weights using Eq. Fland Eq. convergence. (See Alg. for pseudo-code). We consider the number goal maps K. the discount factor y the strength of the goal map prior X, well as the noise variances associated with the time-varying weights, {o to be hyperparameters. To restrict the number of hyperparameters, we set a a vk. We then swept across broad range of values for all hyperparameters (see SM for the full list of values considered) and selected the values that optimized the log- liglliiel a set of held-out trajectories.
Algorithm 1: DIRL Inference Procedure
Input 1: MDP state and action spaces, transition matrix: (S,
Input 2: N trajectories, = (
Input 3: Hyperparameters: no. of goal maps K, noise variances {σk), discount factor 7,
strength of goal map prior, A;
Output: Parameters governing the rewards (u,, as Ik=1: where E R$, ar E IRT;
Let = [ui,.. up]] a [a],...a]] Initialize u a
for Naer do Calculate rewards re(s) Ex
a iee iter Vs E S, (1...T):
('(()))
Get policy using soft value iteration a)) Σn, (st() ((s,a,t):
E.,
Update u ier+1 by maximizing the log-posterior trajectories (Eq. 50:
Update rewards re(s) = Ee o,, "R,
iter (eertl and learn new policy गे (s, a);
Update a ier+1 by maximizing the log-posterior, with noise variances { (Eq. 60.
end
Output u Nurr Nos,
5

---

Results
4.1 Application to a simulated gridworld environment
A
example
B
true goal maps
C
weights
trajectories
end
"home"
"water"
"home"
"water"
2
true
recovered
start
0
time
50
50
D
recovered goal maps
E
rewards
F
test LL
2
water
home
-2.1-
avg.
-2.15
50
3 #maps
2
4
time
Figure 3: Simulations a 5 x 5 gridworld. (A) Example expert trajectories when the time-varying reward function is obtained using the generative goal maps shown in B and the time-varying weights shown in C. (B) Generative goal maps: the first map has a high reward at the "home state" (upper left), the other has high reward at the the "water state" (bottom center). (C) Time varying weights for the home and water goal maps: solid lines show the generative parameters, while dotted lines show the recovered parameters along with a 95% confidence interval. Error bars are computed via the inverse Hessian of the log-posterior of Eq. 6 at the MAP estimate of the weights.
(D) Recovered goal maps. (E) Rewards for the home and water states are shown in red and blue respectively. The average reward for the remaining states is shown in green. Solid lines show the generative rewards, while the dotted lines show the inferred rewards. (F) Held-out test set performance as function of the number of goal maps. Higher values are better; units are bits per decision.
We first demonstrate our method on simulated trajectories in 5 gridworld environment, with 5 actions per state (up, down, left, right, stay). We generated two goal maps for this environment: a "home" map and "water" map, which were rewarding only at the home state and the water state, respectively (Fig 3). Corresponding to these goal maps, we also generated time-varying weights (Fig. C solid lines) for 50 timesteps with the saddom-wak prior of Eq. 2 (for σ 2 2-3.5, this was chosen to provide adequate variation in the reward function during the time period considered). 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. E blue solid line). In contrast, the weight on the home map started small but increased so that the home state became the most rewarding state the end of the 50 timesteps (Fig. BE, red solid line). All of the other states in the environment had a constant reward of In order to generate trajectories corresponding to this reward function (Fig. (3.), we used soft value iteration to learn the corresponding optimal time-varying policy (Eq. (3). We then executed this policy in order obtain 200 trajectories (a similar number the number of trajectories we have for the real dataset discussed later), two examples of which are shown in Fig. 3
Next, applied IRL. inference method from 80% of the generated trajectories. Fig. 3 shows the log-likelihood of the remaining 20% of trajectories as function of the number of goal maps. The held-out test log- 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. It is important to note that can only recover rewards at each timestep up to an additive constant, as the policy remains unchanged upon the addition of constant to the rewards. Further, scaling of the goal maps accompanied by an inverse scaling of the time-varying weights also leaves the recovered rewards unchanged. 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 SM). Figures C and D show that our method allows us to accurately recover the generative goal maps and tire-ig weights from the simulated trajectories. Finally, combining the goal maps and vireinarying weights, we are able to accurately match the generative time- varying rewards for
to learn the goal maps and time-varying weights

---

different states in the gridworld (Fig. EE). We also simulated trajectories in 127-node labyrinth environment (akin [15]. Fig D and confirmed that we were able recover goal maps and time- varying weights (see SM for details). We focus on the gridworld simulations here to demonstrate the versatility of our approach across environments.
4.2 Application of DIRL to real mouse trajectories
Next, we applied our framework the trajectories of real mice navigating (in the dark) in 127-node labyrinth environment [15]. In this task, two cohorts of 10 mice moved freely through the labyrinth over the course of hours. 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, D. The second cohort was not water-restricted and did not have access the water port. We show example trajectory for an animal in each cohort in Fig. [
Over the course of the night, each animal completed over hundred such trajectories, with some animals completing many more.
water-restricted animals
A
recovered goal maps
test
"water"
"home"
Lrandom
max ent deep max ent
1 map
DIRL
4 maps
-2
-1.5
-1
-0.5
B
D
rewards
25
25
weights
water
home
avg. other states
time 20
0
20
time
Figure 4: Inferred varying rewards for water-restricted mice. (A) Inferred goal maps for the water-restricted mice: a "water" map and "home" map. Dots indicate each of the 127 nodes in the labyrinth environment. (B) Recovered time-varying weights for the same animals. Each timestep corresponds to second. (C) Model comparison: comparison of DIRL held-out trajectories to random policy, as well as the Maximum Entropy IRL framework of [19] and the Deep Maximum Entropy IRL framework of [34]. Higher test log- likelihodo is better; test log-like likelihood has units of bits/decision. (D) Inferred time-varying rewards for the home state, the water state, and the average reward for the remaining states.
4.2.1 Inferring interpretable reward fuunctions from water-restricted mice
We began our investigation by applying our method to a subset trajectories for the water-restricted animals. Here, we anticipated being able to identify the water port highly rewarding (in contrast to the water-unrestricted animals). Due to the high variability trajectories across animals and over the course of the 7 hours, and 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 metrie; full details are in the SM) to identify similar trajectories. This procedure had the advantageous side effect of excluding the first 25 or so trajectories for each animal: focus is not characterizing learning. Overall, obtained 200 trajectories for this cohort. Each timestep in trajectory was recorded at a one second interval.
We then fit the goal maps and time-varying weights to 160 of these trajectories, and held out the remaining 20% of trajectories as validation set. We found that validation log-likelihood began to level off at two maps (Fig. 4C), we focus the map solution here. The map solution is

---

shown in the SM (where the recovered water goal map is simply repeated), Overall, we found that we were, indeed, able to recover "water" goal map (Fi with 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. The recovered time-varying varying weights (Big(B3) 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. In contrast, the weights corresponding to the home map started low, but increased over time. The final reward function (Fig #D) 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. The other states in the labyrinth offered only small intrinsic reward for the mouse throughout the course of its trajectory.
4.2.2 'Exploratory' maps inferred from water-unrestricted mice
We then moved on to examining the trajectories of the water-unrestricted cohort. 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. 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.
We fit the time- varying weights and goal maps 80% of these trajectories, and held out the remaining 20% as a validation set. As we show in Fig. SC. the validation log-likelihood is best for the 2 map solution. While we recover (A) a "home" map once again, new map that for the water-restricted animals- also appears. This map is rewarding at many states throughout the maze, but is very unrewarding at the home port, and is very rewarding at state that is close to the home port (which we refer to as the "explore" state). Hence, when the animal places a positive weight on this map, the mouse wants to leave the home state and venture into the maze. For these reasons, we refer to this map as the "explore" map. The vire-ig 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. 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. Finally, the inferred rewards (Fig 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). We can also see that the reward for the home state rises much more rapidly over time than the reward for other states.
not present
capturingg
4.2.3 DIRL outperforms existing IRL. approaches
Finally, we compare the performance of our method with two popular IRL frameworks: the maximum
entropy IRL framework of Ziebart al. [19] and the deep maximum entropy IRI. framework of Wulfmeier et al. [34] In comparison to DIRI., these methods learn a static reward function. 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. Using the retrieved reward functions, we then obtain the corresponding optimal policy and use to compute the log- log-likelioood of the validation set of trajectories. In Figures 4C and we demonstrate that our method dramatically outperforms these existing methods at explaining the held-out trajectories of both vater-restricted and unrestricted mouse cohorts. 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 and
5 Discussion
In this work we develop DIRL, a novel inverse reinforcement learning framework for characterizing the behavior of animals during complex decision-makee makinn tasks. Our framework infers the time- varying reward functions of animals from their trajectories alone. 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]. 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. As time passed and the mice became fatigued, the reward for the home state increased for animals in both cohorts. Our method dramatically outperfomed IRL methods with fixed reward
8

---

water-unrestricted animals
A
recovered goal maps
test LL
"explore"
"home"
random
max ent
deep max ent
map
DIRL
4 maps
-2
-1.5
-0.5
B
0
D
rewards
5
weights
water
home
weights
0
avg. other
explore state
states
-5
-5
10
0 time 10
time
Figure 5: Inferred time-varying rewards for water-unrestricted mice. (A) Inferred goal maps for this cohort: an "explore" map and a "home" map are returned. Dots indicate each the 127 nodes in the labyrinth environment. (B) Recovered time-varying weights for the same animals. (C) Model comparison: comparison of DIRL on held-out trajectories random policy, as well as the Maximum Entropy IRL framework of [19] and the Deep Maximum Entropy IRL framework of [34]. Higher test log-lokelliee lieeliiiihood is better; test log-likelihood has units of bits/decision. (D) Inferred time-varying rewards for the home state, the water port, the 'explore' state (the highly rewarding state in the 'explore' goal map that closest to home) and the average reward for the remaining states. Here the average reward line is almost top of the line corresponding to the water port.
functions, quantified by the log-like regliikli of held-out trajectories, indicating clear need for tailored IRL approaches for neuroscientific applications. 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 laptop.
One exciting finding of our work is the discovery of an "explore" goal map for the water- water-unnestricted mouse cohort. general, exploration in animals is not well understood (26). DIRL offers a novel unsupervised approach for characterizing exploration from behavior alone. Future work could build upon the framework we present here in order to provide normative explanations for the nature of the explore map. DIRL could also be used to explore the neural underpinnings of exploration, for example by correlating fluctuations in the inferred lime-varying rewards with recorded neural activity. Finally, while we developed DIRL with neuroscientific applications in mind, our framework is general and could be applied in other scientific fields. For example, DIRL could applicable in healthcare settings, such those considered in Hüyük al. [47]. In Hüyük al. [47, the authors acknowledge that the policies used by healthcare providers to allocate organ donations have varied over time. By applying DIRL to such a dataset, it may be possible to recover the latent goals of policy-makers makers or medical practitioners when deciding how best to allocate organs to patients.
In addition to the above-mentioned scientific applications, we believe that DIRT. also provides future directions for the development of RL algorithms. Relative to animals, it often challenging to get artificial agents to explore in sparse reward environments [48). 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 navigate in analogous environments (such as in [49] where artificial agents navigated in an analogous 6 binary tree environment). 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-Moo Hidden Markov Model 50 type of POMDP-.. and then performing IRL. with POMDPs (51)
We will close by discussing some limitations of our work. Firstly, our framework requires over hundred trajectories to infer an animal's varying reward function. While may be possible to

---

reduce this number (with, for example, careful choice of prior), it is easy to conceive of failure modes where a single decision reveals nothing about the active goal map. In practice, having access to multiple trajectories (or several decisions) can significantly reduce the uncertainty in the recovered time-varying weights and goal maps. 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. We don't t anticipate this being 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 low-dimensional representation of each state in the environment via a deep network, and then learning time-varying rewards in this embedding space. This would allow to scale approach to higher dimensional state spaces, while also allowing for generalization across states. Finally, our inferred rewards rely on the animal's 's policy being the Boltzmann policy of Eq. However, assuming this form for the policy is not atypical in applications of reinforcement learning to human or animal decision-making data [3 28 Overall, we believe that the advantages of using method outweigh its limitations, and that we present a new, flexible framework for characterizing animal behavior in complex environments.
Acknowledgments and Disclosure of Funding
We thank Benjamin Cowley, Nathaniel Daw, Rachit Dubey and Karthik Narasimhan for useful comments and discussions at various points during this project. We thank the anonymous NeurIPS reviewers for their insightful feedback and helpful suggestions for improvement. The mouse in Figure was provided by BioRender. This work was supported by grants from the Simons Collaboration on the Global Brain (SCGB AWD543027), the NIH BRAIN initiative (NS104899 and R01EB026946), and U19 NIH-NINDS BRAIN Initiative Award (5U19NS104648)
References
[1] Yael Niv. Reinforcement learning in the brain. Journal of Mathematical Psychology, 53 (3):139-154, June 2009. ISSN 0022-2496.
doi: 10. 10.10 2008. 12.005. URL httpa
/www.ac www.sscienceirrct..omm com/ wttpssciencceeeeeee
[2] Robert C Wilson and Anne GE Collins. Ten simple rules for the computational modeling of behavioral data, eLife, 8:44954, November 2019. ISSN 2050-084X. doi: 10.7554/eL.it 49547. URL httpps:///oi...... org 10. 7554/ eLife. 49547, Publisher: eLife Sciences Publications, Ltd.
[3] K. H. Britten, M. N. Shadlen, W. T. Newsome, and J. Movshon. Responses of neurons in macaque MT to stochastic motion signals. Visual Neuroscience, 10(6):157-1166, December 1993. ISSN 0952-5238. doi: 10. 10.1017/5005252380000222
[4] Jerome R. Busemeyer and Julie C. Stout. A contribution of cognitive decision models to clinical assessment: decomposing performance on the Bechara gambling task. Psychological assessment, 14(3):253, 2002. Publisher: American Psychological Association.
[5] Eldad Yechiam, Jerome R. Busemeyer, Julie C. Stout, and Antoine Bechara. Using cognitive models to map relations between neuropsychological disorders and human decision-making deficits. Psychological science, 16(12):9739978 2005. Publisher: SAGE Publications Sage CA: Los Angeles, CA.
[6] Joshua Gold and Michael N. Shadlen. The neural basis of decision making. Annual Review of Neuroscience, 30:535- 574, 2007. ISSN 0147-006X. doi: 10.1 btt.aaanreev.nuuro29.9666 113038.
[7] Mattco Carandini and Anne K. Churchland. Probing perceptual decisions in rodents. Nature Neuroscience, 16(7):824-831 July 2013, ISSN 1546-1726. doi: 10.11038/m..4111 URL htttp::/www.natiie www.erticceeeeeeeeeeeeeeeeeeeeee nn. 3410 Group.
Number: 7 Publisher: Nature Publishing
[8] Peter Dayan and Nathaniel D. Daw. Decision theory, reinforcement learning, and the brain. Cognitive, Affective, Behavioral Neuroscience, 8(4):499 453, December 2008 ISSN 1531- 135X. doi: 10.3758/CABN.8A. 8. 4.429. URL https://://o hwt.d/doi.o.ogg//oooooo 10. 3758/35583333 CABN. 8. 8.4. 4.9
10

---

[9] Laura Busse, Asli Ayaz, Neel T. Dhruv, Steffen Katzner, Aman B. Saleem, Marieke L. Schölvinck, Andrew D. Zaharia, and Matteo Carandini. The detection of visual con- trast in the behaving mouse. The Journal of Neuroscience: The Official Journal of the Society for Neuroscience, 1131111133--11361, August 2011. ISSN 1529-2401. doi:
10.1 SANNNEEUUROSC..69-111....0
[10] Felix A. Wichmann and N. Jeremy Hill. The psychometric function: I. Fitting, sampling. and goodness of fit. Perception & Psychophysics, 63(8):1293-1313, November 2001. ISSN 1532-5962. doi: 10.3758/BF03194544
https://://// http:///al...rr/....3///////1110
[11] Nathaniel D. Daw. Trial-by-trial data analysis using computational models. Decision making, affect, and learning: Attention and performance XXIII, 23(1), 2011. Publisher: Oxford University Press Oxford.
[12] Armin Lak, Emily Hueske, Junya Hirokawa, Paul Masset, Torben Ott, Anne E Urai, Tobias H Donner, Matteo Carandini, Susumu Tonegawa, Naoshige Uchida, and Adam Kepecs. Reinforce- ment biases subsequent perceptual decisions when confidence is low, widespread behavioral phenomenon. eLife, 9:e49834, April 2020. ISSN 2050-084X. doi: 10.7554/eLife.44 49834, URL https://doi. org/ 7554/eLife 49834
Publisher: eLife Sciences Publications, Ltd.
[13] Charles Findling, Vasilisa Skvortsova, Rémi Dromnelle, Stefano Palminteri, and Valentin Wyart. Computational noise in reward-guided learning drives behavioral variability in volatile environments. Nature neuroscience, 22(12):2066-200) 2019. Publisher: Nature Publishing Group.
[14] Dean Mobbs, Pete C. Trimmer, Daniel T. Blumstein, and Peter Dayan. Foraging for foundations in decision neuroscience: insights from ethology. Nature Reviews Neuroscience, 19(7):419 427, July 2018. ISSN 1471-0048. doi: 10. 1038//415188-011 10118-8811100------111 10100-1--7 URL TRLhttpps: // .atuur
com wwt.aaatiiilsseeee setiicceees/4158--1 s41533-113333 018-
Number: 7 Publisher: Nature Publishing Group.
[15] Matthew Rosenberg, Tony Zhang, Pietro Perona, and Markus Meister. Mice a labyrinth show rapid learning, sudden insight, and efficient exploration. eLife, 10:e661 175, July 2021. ISSN 2050-084X. doi: 7 URL htttps://ooi.... org/ 7554/ eLife....f.e..eee Publisher: eLife Sciences Publications, Ltd.
eLi...6175... 66175
[16] Gautam Reddy. Reinforcement waves as a mechanism for discontinuous learning. Technical report, bioRxiv, May 2022. URL BtLL:htpss://ww/ http:://wwwwbbooii..../. biorxiv.orgg Tttpttennt/0.1111///11 05. 06.06.4 490910v11
Section: New Results Type: article.
[17] Andrew Y. Ng and Stuart J. Russell. Algorithms for inverse reinforcement learning. In Icml,
volume 1, page 2, 2000.
[18] Pieter Abbeel and Andrew Y. Ng. Apprenticeship learning via inverse reinforcement learning. In Proceedings of the twenty-first international conference on Machine learning, page 1, 2004.
[19] Brian D. Ziebart, Andrew Maas, J. Andrew Bagnell, and Anind K. Dey. Maximum entropy inverse reinforcement learning. In Proc. AAAI, pages 1433-1438, 2008.
[20] Brian D. Ziebart, J. Andrew Bagnell, and Anind K. Dey. Modeling Interaction via the Prin- ciple of Maximum Causal Entropy. 2010. URL BRWMT5Wu-S
[21] Pieter Abbeel, Adam Coates, Morgan Quigley, and Andrew Ng. An Application of Reinforce- ment Learning to Aerobatic Helicopter Flight. Advances in Neural Information Processing Systems, volume 19. MIT Press, 2006. URL http:ttt/////ceeeiiiiiii .ew.eeddddinnnnnnnnn..eee..ee 20000neee/8888901111187777788880 Abstaaat....tttttttttttt...e.
neurips cwturips.ccppppp..... cc/paper////pppp
[22] Adam Coates, Pieter Abbeel, and Andrew Y. Ng. Learning for control from multiple demonstrations. Proceedings of the 25th international conference on Machine learning ICML '08, pages 144-151, Helsinki, Finland, 2008. ACM Press. ISBN 978-1-60558-205. doi: 14.1145/13901 1390175. URI.http: /ttttp://portall acm. haem..ogg//iieeeo org/citation. wtt.taatioo..fmdiii...... 1390156. 1390175

---

[23] Niranjani Prasad, Barbara Engelhardt, and Finale Doshi Velez. Defining admissible rewards for high-confidence policy evaluation batch reinforcement learning. In Proceedings of the ACM Conference Health, Inference, and Learning, pages 1-9, 2020.
[24] Alex Chan and Mihaela van der Schaar. Scalable bayesian inverse reinforcement learning. arXiv preprint arXiv:2102.06483 2021.
[25] Aishwarya Mandyam, Didong Li, Diana Cai, Andrew Jones, and Barbara Engelhardt. Effi- cient Bayesian Inverse Reinforcement Learning via Conditional Kernel Density Estimation. November 2021.
[26] Robert N Hughes. Intrinsic exploration in animals: motives and measurement. Be- havioural Processes, 41(3):213-22) December 1997. ISSN 0376-6357. doi: 10.1016/
$0076-663577970000555--1 URI. http:ppp:///...ccccceeeeeeeeeete wwwp::/wsieenediirre.... com/ htt.:/rrrrttcocc/cccnce/eeeeeeee http://cciicce/ariclee/.// p11/S0376635797000555
[27] Florian Kattner, Aaron Cochrane, and C. Shawn Green. Trial-dependent psychometric functions accounting for perceptual learning in 2-AFC discrimination tasks. Journal of Vision, 17(11):3, September 2017. ISSN 1534-7362. doi: 10.1167/17.11.3.
[28] Amir Dezfouli, Kristi Griffiths, Fabio Ramos, Peter Dayan, and Bernard W. Balleine. Models that learn how humans learn: The case of decision-making and its disorders. PLOS Computa- tional Biology, 15(6)ze 1006903, June 2019. ISSN 1553-7358. doi: 10.1371/journal,
1006903. URL
plos Publisher: Public Library of Science.
journal pcbi.cpb............ 1006903
[29] Nicholas A. Roy, Hyun Bak, Athena Akrami, Carlos D. Brody, and Jonathan W. Pillow. Extracting the dynamics of behavior in sensory decision-making experiments. Neuron, 109 (4):597-610.c6, February 2021. ISSN 0896-6273, doi: htwp://6.nnuroo..2220220 12.004. URL
https://wwv httpptttt////aaadaaaaaerrerrooo httpp//rweeeeeeeeeeeeeeeeee////e
[30] Zoe Ashwood, Nicholas Roy, Iris R. Stone, Anne Urai, Anne K. Churchland, Alexandre Pouget, and Jonathan W. Pillow. Mice alternate between discrete strategies during percep- tual prccisionmmainng Nature Neuroscience, 25(2):201-22) February 2022. ISSN 1546-
1726. doi: 941593-021-01007-
URL https: //www. nature. com/ articlest Number: Publisher: Nature Publishing Group.
[31] Matthias Schultheis, Dominik Straub, and Constantin A. Rothkopf. Inverse Optimal Control Adapted the Noise Characteristics of the Human Sensorimotor System. October 2021. doi: 10.48550/arXiv.211 1111101111000 11100. URL https: /ttp:///axivvvv.a http://aaxvvvoooraaa.a aataaa/s......ae..aeaaaeaaaaaaea 2110. 1111110.1110000 11130v1
[32] Shoichiro Yamaguchi, Honda Naoki, Muneki Ikeda, Yuki Tsukada, Shunji Nakano, Ikue Mori, and Shin Ishii Identification of animal behavioral strategies by inverse reinforcement learning. PLOS Computational Biology, 14(5)=e1006122 May 2018 ISSN 1553-7358. doi: 10.1371/ journal. pwb.. 1006122....... URL https: htttps://jjurnalll/ journals... plos org http://oooooooo/rrrrreeo htmmpiiio//rticleiiet
10. 10.1371/journal. pcbi. 1006122 Publisher: Public Library of Science.
[33] Minhae Kwon, Saurabh Daptardar, Paul R Schrater, and Xaq Pitkow. Inverse Ra- tional Control with Partially Observable Continuous Nonlinear Dynamics. In Ad- vances in Neural Information Processing Systems, volume 33, pages 7898-7909. Cur- ran Associates, Inc., 2020. URL httpps///rrr//eeeinngg wttp:iinnnniis.cccpeeeeeeeeeeeee http:ccccpapee/////aaaeee 550100977cccc333228877444483 wwt.///tbbttttttttaa
[34] Markus Wulfmeier, Peter Ondruska, and Ingmar Posner. Maximum entropy deep inverse reinforcement learning. arXiv preprint arXiv: TeXXv:/5577048800 2015.
[35] Chelsea Finn, Sergey Levine, and Pieter Abbeel. Guided cost learning: Deep inverse optimal control via policy optimization. In International conference on machine learning, pages 49-58. PMLR, 2016.
[36] Yunzhu Li, Jiaming Song, and Stefano Ermon. Infogail: Interpretable imitation learning from visual demonstrations. arXiv preprint arXiv:1703.0883. 2017.
12

---

[37] Adam Gleave and Sam Toyer. A Primer on Maximum Causal Entropy Inverse Rein- forcement Learning, March 2022. URL http://arxiv wwwpriii...rggaae www.aa/a.a/aeeaaea 11409 Number: arXiv:2203.1 1409 arXiv:2203.1 11409 [es].
[38] Monica Babes- Vroman, Vukosi Marivate, Kaushik Subramanian, and Michael L. Littman. Apprenticeship learning about multiple intentions. ICML, 2011.
[39] Jaedeug Choi and Kee-Eung Kim. Nonparametric bayesian inverse reinforcement learning for multiple reward functions. In NIPS, 2012.
[40] Amarildo Likmeta, Alberto Maria Metelli, Giorgia Ramponi, Andrea Tirinzoni, Matteo Giu- liani, and Marcello Restelli. Dealing with multiple experts and non-stationarity in inverse reinforcement learning: application to real-life problems. Machine Learning, 1 110(9): 2541-2576, September 2021. ISSN 1573-0565. doi: 10. 10.10077/6199944420--99999994449 URL
fettp:://do..rrg111
059399----
[41] Amit Surana and Kunal Srivastava. Bayesian nonparametric inverse reinforcement learning for switched markov decision processes. In 2014 13th International Conference on Machine Learning and Applications, pages 47-54, 2014, doi: 10.1109/CMLAA.2144.1.MMM
[42] James S. Hodges and Daniel J. Sargent. Counting degrees of freedom in hierarchical and other Sechly-prrmmttrrseed models. Biometrika. 88(2):367-379, 06 2001. ISSN 0006-3444. doi:
10. wt..00//biomet/8.22..o 2.3667 URL https:::: org/ 1093/ 10.000//1000 http9///iomet///// bwtpbt/t..tttt....//////
[43] Tuomas Haarnoja, Haoran Tang, Pieter Abbeel, and Sergey Levine. Reinforcement learning with deep energy-based policies. 2017.
[44] Divyansh Garg, Shuvam Chakraborty, Chris Cundy, Jiaming Song, and Stefano Ermon. IQ- Learn: Inverse soft-Q Learning for Imitation. arXiv:2106. 12142 [cs], December 2021. URL arXiv: 2106.12142.
httpr//ab//210... 01122216611244 12142
[45] Martin Ester, Hans-Peter Kriegel, and Xiaowei Xu, A Density- -ased Algorithm for Discovering Clusters in Large Spatial Databases with Noise. page 6, 1996.
[46]
Yiren Lu. IRL-Imitation, 2017. URL
[47]
Alihan Hüyük, Daniel Jarrett, and Mihaela van der Schaar, Inverse Contextual Bandits: Learning How Behavior Evolves over Time. In Proceedings of the 39th International Conference Machine Learning, pages ISSN: 2640-3498.
PMLR, June 2022. UtLhttpps:// / //pppce//p/cccdiggss mlr
press/v162/huyukk20 html
[48) Deepak Pathak, Pulkit Agrawal, Alexei A. Efros, and Trevor Darrell. Curiosity-driven ploration by self-supervised prediction. In Doina Precup and Yee Whye Teh, editors, Pro- ceedings of the 34th International Conference on Machine Learning, volume 70 of Pro- ceedings of Machine Learning Research, pages 2778-2787. PMLR, 06-11 Aug 2017. URL httts::/proceediie .mr
mlr prwss//70/pathakkee html
[49] Runzhe Yang, Xingyuan Sun, and Karthik Narasimhan. Generalized Algorithm for Multi- Objective Reinforcement Learning and Policy Adaptation. arXiv:/908, 08342 [cs], November 2019. URL /ttt:://taxx/ org/ abs/
abs/1908. 1908. 08342
arXiv: 1908.08342.
[50] Samuel M. Choi, Dit- Yan Yeung, and Nevin Zhang. Hidden-Mo Mode Markov Decision Processes for Nonstationary Sequential Decision Making. In Ron Sun and C. Lee Giles, editors, Sequence Learning: Paradigms, Algorithmss, and Applications, Lecture Notes in Computer Science, pages 264-287. Springer, Berlin, Heidelberg, 2001. ISBN 978-3-540-4455544 doi: 10.100773-540-44555 X_, URL httppp///bb..org//0000///------0
[51] Jaedeug Choi and Kee-Eung Kim. Inverse Reinforcement Learning in Partially Observable Environments. Journal of Machine Learning Research, 12:691-730, March 201
13

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] please refer to the Discussion
section.
(c) Did you discuss any potential negative societal impacts of your work? [N/A] We do not foresee any negative societal impacts of our work.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes] We confirm that we have read through the ethics review guidelines and that our paper conforms to them.
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [N/A]
(b) Did you include complete proofs of all theoretical results? [N/A]
3. If you ran experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material code as a supplemental material. The real data studied in this work is publicly available at hattps://githen
URL)? Yes]] We provide the
com/ www.//smeestte///oooete htw.sstteer/Rsoebbrree /tt////oseenbeeeeeee 2021 Repository
(b) Did you specify all the training details (c.g.. data splits, hyperparameters, how they were chosen)? [Yes] See our Methods section, as well as the SM.
(c) Did you report error bars (e.g.. with respect to the random seed after running experi- ments multiple times)? [Yes] See Figures 2, and 4 for examples of this.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, cloud provider)? [Yes] See the Discussion section, well as the SM.
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) your work uses existing assets, did you cite the creators? [Yes] We cite the curators of the real-world mouse dataset many times.
(b) Did you mention the license of the assets? [Yes] In the SM, we make clear that this dataset is publicly available and comes with an MIT open source license.
(c) Did you include any new assets either in the supplemental material or as URL? [Yes] We provide our code as a supplemental material.
(d) Did you discuss whether and how consent was obtained from people whose data you're using/curating? [N/A]
(c) Did you discuss whether the data you using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given participants and screenshots, if applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent participant compensation? [N/A]
14


Reviews:
[{"review_id": "Q_bxawdzWO", "reviewer": "Reviewer_Rzrz", "summary": "The paper proposed a new method, Dynamic IRL, and applied it to model mouse behaviour in a spatial navigation task to recover the goals of the animals. Unlike the previous IRL methods, the Dynamic IRL allows the reward to change over time slowly. To make the inference trackable, the authors model the reward function as a linear combination of weights and a few \"goal maps\". The authors demonstrated and validated the model on a simulated gridworld task, and apply it to a real mouse dataset to successfully recover the goals of the animals in two different conditions. ", "questions": "\"Building computational models of decision-making is a core objective in both neuroscience and psychology\" - this is a substantial statement and I believe it is not 100% true. It might be true for some computational neuroscientists/psychologists, but the goal of building models is not to build models, it is for understanding the biological system or psychological system. I would suggest authors make this softer, especially when this is the first sentence of the abstract and introduction.\n\nFrom the paper, it seems that the author separate the data into 80% training data, and 20% testing data (where test LL was computed). When tuning the hyperparameters, the authors did it on a validation set. Did the author use a separate validation set from the 80% training set, not using the 20% testing data? It would be great if the author can make this very clear in the writing. \n\nFollowing the training data split, I'm wondering if the non-distinguishable 2-map and 3-map can be due to a lack of statistical power. For example, in the simulation, the authors can try using longer sequences or more sequences, just to test this hypothesis (even though longer or more sequences are not available in the real dataset). Or instead of an 80/20 split, k-fold cross-validation might provide more statistical power. \n\nIf this non-distinguishable is intrinsic not from model recovery noise, I would be happy to see (brief) discussion on future work on addressing this, maybe somehow penalize for model complexity in the model?\n\nThe data selection part can seem a bit arbitrary. It makes sense to remove the first 25 or so trajectories, but I'm wondering what else (if any pattern) is removed. I'm also wondering if the results change dramatically if this data selection is not done (maybe only removing the first 25), i.e. is this step necessary when apply the dynamic IRL model? If so, I would expect the model can capture some noise to some extent, e.g. the reason to allow reward function to vary is to capture fatigue. I would like to see a (brief) discussion on future work on a more principal way to capture noise. Maybe a noise map, or allowing the reward function to not just change slowly across time, but also can undergo change point?\n\nIt is exciting to see the model is already applied to a real animal dataset. However, the conclusion/implication we can draw from the modelling results seems to only confirm what we can observe from the behavioural data alone, and mostly serve to validate the model's ability to recover something meaningful. I can imagine the recovered weights can be used to correlate with some neural recordings. The authors touched on this a bit by saying the method can provide a normative account, but since this is mainly a methodology paper, it would be great to add a sentence or two, detailing what kind of scientific findings your model can provide in addition to only analysing behavioural data. \n\n\n\n\n\n", "limitations": "Another limitation might be, that in lots of behavioural experiments, the state might be hidden/unobservable. I think this is definitely out of the scoop of this paper, that the authors do not need to include/address. But if authors have comments on this, it can be very interesting to read.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The paper is well-written and easy to follow. The figures are also high quality.\n- The contribution is clear, allowing the reward to vary across time in the IRL model.\n- The authors performed experiments on both synthetic data and real animal datasets, and compared to an existing IRL model. \n\nWeaknesses:\n(I only list them here, and I will explain more and put my questions and suggestions in the next section)\n- The number of goal maps seems to be non-distinguishable to some extent, e.g. in the simulation, the 2-map and 3-map fit the 2-map data equally good.\n- When applying it to real data, it requires a procedure to cluster the data first.\n- The modelling results seem to mainly confirm the hypothesis from the behaviour and descriptive. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "99LzOB0Z9-z", "reviewer": "Reviewer_PZrL", "summary": "This paper describes an application of inverse reinforcement learning to maze exploration by mice. It proposes a novel IRL method to model changes in the reward function over time, by representing the reward as a linear combination of multiple goals, such that the importance of each goal can change over time. This model jointly infers the possible goals, and the dynamic reward function from mouse trajectory. The model can be applied to free exploration, where mice are not given an explicit goal on each trial. In this case it still recovers reasonable goal maps, such as a \"home reward\" and a random-looking exploration reward map. The model assumes that mice have multiple goals and choose between the set of actions associated with these goals using a Boltzman policy, but makes no other assumptions. ", "questions": "1. Why is there a log in equation (4)? I am not familiar with a formulation of a Q function with a log.\n\n2. Line 141: \"coordinate ascent\", line 144: \"gradient ascent\" -- this should have probably been \"coordinate descent\"?\n\n", "limitations": "I believe this work has minimal negative societal impact -- it is extremely unlikely that IRL can be used to infer human goals. The main limitation of IRL as acknowledged by the authors, is the need to see many instances of a behaviour. This will work well in contexts where the goal set is relatively small and constant through the observation period. However, IRL can easily miss fine-grained variation in goals, such as when animals see exploration as play,  trying out multiple ways to achieve a goal. This will show up as random actions when averaged over trials.", "rating": 10, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: This is very interesting work -- among the most interesting NIPS submissions I have seen over the years, and a great application of IRL to a real problem. I strongly recommend this work to be presented at the conference. Although the explore map tells us little in practice, it is very interesting in theory and takes an important step toward understanding exploration. This is a short review because I think the paper is really great.\n\nWeaknesses: the exploration map leaves some open questions  -- \ncan this map emerge from a random walk? \nfrom self-avoiding random walk? \nare mice intentionally seeking out intersections, or other structural elements that allow them to encode the spatial map?\nThis is a minor weakness, since it is impossible to address all of these within 9 pages.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "3T_FmqYeW_U", "reviewer": "Reviewer_9v8k", "summary": "This paper introduces an inverse reinforcement learning framework that infers the time-varying reward function of the animal. The reward is designed as a linear combination of rewards maps, which can be weighted differently. The proposed method is confirmed using simulated data and real mouse trajectories. ", "questions": "Overall, the reviewer enjoys reading this paper. It is excellent work, and the NeurIPS community will like it. However, there are a few clarification questions as follows.   \n- How is the artificial mouse trained in Section 4.1? The reviewer assumes that the authors train the artificial mouse to generate the trajectories. Is any specific RL algorithm used to train the agent? \n- All the results included in the paper are based on two maps scenario. The inferred weights in Figure 4B and Figure 5B are only two types. Three maps scenario is presented in the Supplementary material. However, Is it possible to extend this method to multiple maps? Maybe more than 5? If then, this work can be useful in the ML and neuroscience communities. \n-  Some essential pieces of literature are missing. There are recent NeurIPS papers that study RL or IRL in decision-making problems in neuroscience. These can be discussed in related works.\n \"A neurally plausible model learns successor representations in partially observable environments.\" NeurIPS 2019, \n \"Inverse rational control with partially observable continuous nonlinear dynamics.\" NeurIPS 2020.\n\n\n", "limitations": "The authors describe the limitation in the last section. ", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength\n-      Interesting story with important motivation. The inverse RL approach to time-varying reward is novel. \n-      Well-written paper. Easy to follow. All figures are easy to read and well-organized. \n-      Supplementary material is helpful and well-organized.\n\nWeakness\n-      Little impact on the ML community (in terms of ML algorithm). Nevertheless, it is a valuable work to the neuroscience community. \n-      Labels are missing in some figures, which makes it hard to interpret the graph.\n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "3E4cx6Rtwcx", "reviewer": "Reviewer_yX2q", "summary": "# Summary\nThis paper presents an inverse reinforcement learning algorithm designed specifically for navigation tasks involving non-stationary reward functions.\nThe approach is motivated by applications in neuroscience, for example understanding the (time-varying) goal-seeking or exploratory behavior of animals.\nThe algorithm builds on maximum entropy IRL, it assumes access to the transition model of the environment (in performing soft value iteration), and is designed for discrete state spaces (in how it defines features).\nThe rewards are defined as a linear function of learned goal-map features and learned weights associated with those features.\nThe non-stationarity of the environment rewards is handled through learning feature weights at each timestep.\nThe approach is evaluated  in a simulated domain matching the assumptions of the algorithm as well as on real-world data collected from (water restricted and unrestricted) mice navigating a maze.\n\n# Contribution\n- I believe the contribution is a method (based on maximum causal entropy IRL) for inferring (time-varying) reward functions of animals performing navigation tasks, specifically in settings with discrete states, deterministic transitions, and where time-alignment of trajectories is practical\n", "questions": "The section above includes a number of questions. In addition to those:\n1. What are the convergence properties of the algorithm (specifically the dual ascent of the features / weights)? If you run it multiple times, how much do the learned features differ?\n\nAnd some minor questions:\n1. Section 3.1 introduces MDPs, but it’s not clear to me whether the paper deals with finite horizon MDPs (as would be indicated by the use of a timestep-dependent reward function and Q-function) or infinite horizon MDPs (as would be indicated by the use of a discount factor)\n2. In equation 4, should the q-function on the rhs have time index t + 1?\n", "limitations": "Thank you for including a section discussing the limitations of the approach\n", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "# Strengths\n1. The paper is clear and well-written\n2. The paper is motivated by an important problem (better understanding human/animal behavior)\n3. The paper focuses on less common aspects and applications of IRL. For example, the application of IRL in a neuroscience context is interesting and less common than applications in other areas (e.g., robotics or autonomous driving), and applying IRL in a context where learning rewards is itself a primary goal is interesting (for example, the differences in the learned goal maps between the two mice cohorts).\n4. The approach does seem applicable to understanding animal behavior in situations similar to those considered in the paper\n5. The supplementary material provides extensive details on the algorithm and experiments (including well-documented code allowing for reproducing experiments), and execution of the experiments seems sound\n\n\n# Weaknesses\n## Originality\n- The main algorithmic novelty seems to be learning time-varying rewards; however, it seems plausible to me that this non-stationary aspect could instead be addressed by treating the problem as partially observable (where the timestep or the history of the agent are unobserved). \n- In the context of _RL_ non-stationary MDPs can be formulated as _RL_ within a specific type of partially observable MDP [4]. Does that connection hold for _IRL_? For example, can you simply specify the timestep as a feature (or specify a feature capturing historical information) and thereby remove the need for time-varying rewards? Can you instead use IRL methods designed for partially-observed environments [5]? What's the benefit of the time-varying reward formulation over these alternatives?\n- Separately, what’s the relationship with existing work that discusses IRL in the context of non-stationary reward functions such as [3]?\n\n## Relationship with prior work\n- The paper would be improved by clarifying the relationship between the proposed algorithm and existing work (citations to relevant work are provided in section 3.2, but describing the relationship specifically would help)\n- The algorithm seems to be a special case of maximum causal entropy IRL [1]. Both algorithms alternate between the two steps of (i) computing a policy on the current set of rewards via soft VI and (ii) updating rewards to maximize the likelihood of the observed data under the learned policy. This approach is discussed in [1], and a more recent discussion that makes the connection clear is given in [2] (section 3.2). In particular, the proposed algorithm assumes a deterministic transition function, and therefore employs simplified likelihood calculations in eqs (5) and (6) as discussed in section 3.3 of [2].\n- The algorithm learns parameters of both the features and their weights, which existing work does as well (e.g., reference 30 in the paper), though with a particular structure and using coordinate ascent.\n- Additionally, it would help to distinguish from IRL work in understanding human behavior. For example, what’s the relationship with IRL methods used for predicting pedestrian behavior in navigation tasks, for example [6,7]?\n\n## Significance\n- The approach doesn't seem to scale to large / continuous state/action spaces (due to the use of soft value iteration and defining features per-state). Section 5 discusses an extension that would learn embeddings, but it’s not clear to me how this would work.\n- The approach doesn't seem applicable in cases where the transition model is stochastic and/or unknown (due to using the deterministic transition model form of the likelihood in eq (5) and (6), and due to using soft VI). \n- The structure of the features seem specific to discrete state-space navigation tasks\n- Time-varying rewards seem difficult to apply in many situations. For example, the clustering and alignment procedure used in the paper (that seems to require each trajectory to be aligned and of the same length) would be difficult to apply to trajectories of animals/humans in the real world (where natural start and end points are likely unavailable).\n- Despite the above limitations, it does seem the approach is applicable to understanding animal behavior in situations similar to those considered in the paper. I think this is significant, but it's not clear to me how significant. I think the paper would benefit from further arguing why this particular setting is important.\n\n## Experimental quality / analysis\n- Why do MaxEntIRL and DeepMaxEntIRL perform poorly? Section 4.2.3 claims they perform poorly because they learn a static reward function, but is it possible the issue can instead be framed as one of partial observability?\n- For example, if these baselines were provided with features capturing the timestep or history, how would they perform? To capture the history a map of normalized state-visitation counts could be provided as a feature.\n- If addressing partial observability (of the timestep or state-history) addresses the issue, why should we prefer learning time-varying rewards?\n- Another way to potentially answer this question would be to train a (recurrent) policy via behavior cloning, and then compute likelihood on the validation set.\n- Providing an additional experiment on a different real-world dataset would improve the paper by illustrating the generality of the approach\n\n\n[1] Modeling Interaction via the Principle of Maximum Causal Entropy\n\n[2] A Primer on Maximum Causal Entropy Inverse Reinforcement Learning\n\n[3] Dealing with multiple experts and non-stationarity in inverse reinforcement learning: an application to real-life problems\n\n[4] Hidden-Mode Markov Decision Processes for Nonstationary Sequential Decision Making\n\n[5] Inverse Reinforcement Learning in Partially Observable Environments\n\n[6] Intent-aware long-term prediction of pedestrian motion\n\n[7] Learning How Pedestrians Navigate: A Deep Inverse Reinforcement Learning Approach", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept