Inverse Reinforcement Learning with Switching Rewards and History Dependency for Characterizing Animal Behaviors
Authors: Jingyang Ke
Feiyang Wu
Jiyi Wang
Jeffrey Markowitz
Anqi Wu
Venue: ICML 2025 poster (2025)
Keywords: neuroscience
decision-making
inverse reinforcement learning
naturalistic behavior
Primary area: applications->neuroscience_cognitive_science
TL;DR: We develop an novel inverse reinforcement learning framework that can model the history-dependent switching reward functions in complex animal behaviors

Abstract: Traditional approaches to studying decision-making in neuroscience focus on simplified behavioral tasks where animals perform repetitive, stereotyped actions to receive explicit rewards. While informative, these methods constrain our understanding of decision-making to short timescale behaviors driven by explicit goals. In natural environments, animals exhibit more complex, long-term behaviors driven by intrinsic motivations that are often unobservable. Recent works in time-varying inverse reinforcement learning (IRL) aim to capture shifting motivations in long-term, freely moving behaviors. However, a crucial challenge remains: animals make decisions based on their history, not just their current state. To address this, we introduce SWIRL (SWitching IRL), a novel framework that extends traditional IRL by incorporating time-varying, history-dependent reward functions. SWIRL models long behavioral sequences as transitions between short-term decision-making processes, each governed by a unique reward function. SWIRL incorporates biologically plausible history dependency to capture how past decisions and environmental contexts shape behavior, offering a more accurate description of animal decision-making. We apply SWIRL to simulated and real-world animal behavior datasets and show that it outperforms models lacking history dependency, both quantitatively and qualitatively. This work presents the first IRL model to incorporate history-dependent policies and rewards to advance our understanding of complex, naturalistic decision-making in animals.

Inverse Reinforcement Learning with Switching Rewards and History Dependency for Characterizing Animal Behaviors
Jingyang Ke
Feiyang Wu
Jiyi Wang' Jeffrey E. Markowitz 12
Angi Wu
Abstract
1. Introduction
Historically, decision making in neuroscience has been stud- ied using simplified assays where animals perform repet- itive, stereotyped actions (such licks, nose pokes, or lever presses) in response to sensory stimuli to obtain an explicit reward. While this approach has its advantages, it has limited our understanding of decision making to scenar- ios where animals are instructed to achieve an explicit goal over brief timescales, usually no more than tens of seconds, In contrast, in natural environments, animals exhibit much more complex behaviors that are not confined to structured, stereotyped trials. For example, a freely moving mouse may immediately rush toward the scent of food when hungry, but after eating, it might seek out quiet spot to rest for an ex- tended period. Thus, real-world animal behaviors form long sequences composed of multiple decision- makiinn processes. Each decision-making process involves a series of states and actions aimed at achieving goal, and such decision switching is unlikely to occur on very short timescales in simplified assays. Additionally, many of the goals animals pursue in natural settings are generated by intrinsic motiva- tions and thus unobservable. To truly understand animal's decision makiing in a naturalistic context, we need methods to uncover animals' intrinsic motivations during multiple decision-maln making processes.
Traditional approaches to studying decision- making in neuroscience focus on simplified be- havioral tasks where animals perform repeti- tive, stereotyped actions to receive explicit re- wards. While informative, these methods con- strain our understanding of decision- short timescale behaviors driven by explicit goals. In natural environments, animals exhibit more complex, long-term behaviors driven by intrin- sic motivations that are often unobservable. Re- cent works in time-varying inverse reinforce- ment learning (IRL) aim to capture shifting mo- tivations in long-term, freely moving behaviors. However, a crucial challenge remains: animals make decisions based on their history, not just their current state. To address this, we intro- duce SWIRL (SWitching IRI., a novel frame- work that extends traditional IRL by incorporating time-varying,
to
reward functions. SWIRL models long behavioral sequences as tran- sitions between short-term prcision-making pro- cesses, each governed by a unique reward func- tion. SWIRL incorporates biologically plausi- ble history dependency to capture how past deci- sions and environmental contexts shape behavior, offering a more accurate description of animal docision-making We apply SWIRL to simulated and world animal behavior datasets and show that it outperforms models lacking history de- pendency, both quantitatively and qualitatively. This work presents the first IRL model to incor- porate cestory-dependen policies and rewards to advance our understanding of complex, naturalis- tic decision- maainn in animals.
Inverse reinforcement learning (IRL), which infers agents' policies and intrinsic reward functions based on their inter- actions with the environment (Ng & Russell, 2000; Abbeel & Ng. 2004; Ziebart et al., 2008; 2010; Wu et al., 2024), has been shown to be effective in capturing animal decision- making intentions by learning reward functions from be- havioral trajectories (Sezener et al., 2014; Yamaguchi et al., 2018; Pinsler et al., 2018; Hirakawa et al., 2018). However, traditional IRL assumes a single static reward function over time, limiting its ability to account for shifts in intrinsic moti- vations. To address this limitation, recent IRL variants have aimed to uncover heterogeneous and time-varying reward functions (Babes- 2014; Nguyen al., 2015; Ashwood et al., 2022a; Zhu et al., 2024). Despite these advancements, significant challenge remains unaddressed: animals make decisions based on their history, not just their current state (Kennedy, 2022; Hattori
'Georgia Institute of Technology, Atlanta, GA, USA 'Emory University, Atlanta, GA, USA. Correspondence to: Anqi Wu <an- qiwu@ gatech.edu'>.
al., 2011; Surana & Srivastava,
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Inverse Reinforcement Learning with Switching Rewards and History Dependency for Characterizing Animal Behaviors
al., 2019). For example, in perceptual decision-making tasks, mice are found to make new decision based on reward, state and decision history (Ashwood et al., 2022b). Incor- porating historical context into the modeling could offer a more accurate representation of animal behavior.
that SWIRL outperforms alternative models when history dependency is not included, both quantitatively and qualita- tively. This underscores the necessity of incorporating this biologically plausible element when modeling long-ten gengttemm be- haviors. Additionally, for the first time, we will present the application of non-Markovian reward functions and state- action reward functions to model freely-moving animals, contrasting with previous works that only assume a single state-based reward.
To address the absence of history dependency in time- varying IRL models, we introduce a novel framework called SWitching IRL (SWIRL). Similar to Zhu et al. (2024), SWIRL models long recordings of animal behaviors sequence of short-term decision-making processes. Each decision-making process is treated as a Markov decision process (MDP) with a unique reward function that can be in- ferred using IRL. The segmentation of long recording into switching decision-making processes is unknown; therefore, each process is regarded as being associated with a hidden mode that must also be inferred. Most importantly, SWIRL incorporates biologically plausible history dependency, drawing on insights from animal behavior. The history dependency is added at two levels: the transitions between peeision-making processes (decision-level) and the actions taken to achieve a single goal within a decision- making process (action-level). Decision-level dependency is reflected in the transitions between decision-making processes over extended sequences of time bins, suggesting that an animal's 's current choice is shaped by its previous decisions and environmental feedback. Additionally, we posit that these transitions are influenced by the animal's location. For example, after a mouse drinks from a water port, if it stays nearby, is more likely to seek another goal. Conversely, if it is far from the port, indicating it has been away for some time, it may become thirsty again and return to search for water. For action-level history dependency, we will model the policy and reward functions as dependent on trajectory history within each cerision-mmainn process, using a non-Markovian decision framework. Such a dependency has been studied by existing reinforcement learning research, which often characterizes exploration with reward functions based on historical states and actions (Houthooft et al., 2016; Sharafeldin et al., 2024). Importantly, SWIRL is the first to incorporate history-dependent policies and rewards into IRI.
2. Related Work
IRL for animal behavior understanding. IRL has been widely used to infer animals' behavioral strategies and decision-making policies when the reward is unknown. For instance, Pinsler al. (2018) applies IRL to uncover the un- known reward functions of pigeons, explains and reproduces their flock behavior, and introduces a method to learn a leader-follower hierarchy. Similarly, Hirakawa et al. (2018) uses IRL to learn reward functions from animal trajectories, identifies environmental features preferred by shearwaters, and demonstrates differences in male and female migration route preferences based on the estimated rewards. In another study, Yamaguchi et al. (2018) applies IRL. C. elegans thermotactic behavior, revealing distinct behavioral strate- gies for fed and unfed worms. Additionally, Sezener et al. (2014) maps reward functions for rats freely moving in a square area, showing how these rewards changed before and after training. While these studies demonstrate the utility of IRL in uncovering behavioral strategies of freely moving an- imals, they share a key limitation: they all assume a single re- ward function governs all animal behaviors, which does not account for the complexities of long-term decision- -aaiing.
Heterogeneous and time-varying IRI. Recent works have extended traditional IRL, which assumes a constant reward, to models with time-varying or multiple reward functions driving behavioral trajectories. For example, Vaommn et al. (201 1) introduced Multi-intention IRL, which infers multiple reward functions across different trajectories but still assumes single reward function within each trajec- tory. On the other hand, the Dynamic IRL (DIRL) method (Ashwood et al., 2022a) models reward functions as a lin- ear combination of feature maps with cime-narying weights, addressing the issue of varying rewards within trajectory. However, DIRL requires trajectories be highly similar or clustered beforchand, significantly limiting its applicabil- ity. Moreover, it cannot capture switching decision-making processes over long-term periods where each process may vary in length. Additionally, BNP-IRL (Surana & Srivas- tava, 2014), locally consistent IRL (Nguyen et al., 2015) and inlti-intention inverse Q-learning (IQL) (Zhu al., 2024) all extended the multi-intention IRL. framework to allow for changing reward functions within trajectories, making
One key aspect we want to highlight in this paper is that the proposed SWIRL model has intriguing connections to traditional behavioral analysis methods in the animal neu- roscience literature. In Sec. 3.5, we will demonstrate that our SWIRL model offers a more generalized and princi- pled approach to characterize animal behaviors compared to existing autoregressive dynamics models (Wiltschko et al., 2015; Mazzucato, 2022; Stone, 2023; Weinreb et al., 2024).
In the Results section, we will apply our SWIRL a sim- ulated dataset as well two real-world animal behavior datasets. For both animal datasets, we will demonstrate
2

---

Inverse Reinforcement Learning with Switching Rewards and History Dependency for Characterizing Animal Behaviors
them the closest models our proposed SWIRL. How- ever, all models do not account for both decision-level and action-level history dependency, an important biologically plausible factor that SWIRL incorporates to achieve more accurate behavior modeling. In our experiments, we will use multi-intention IQL and locally consistent IRL as baseline models, as they are special cases of SWIRL.
A
#1
Hiddenn modes
AO
a1
a2
Action
30
91
States
Figure 1. SWIRL graphical model. Green dotted lines repre- sent transitions of the hidden modes depend on the previous state (decision-ley) dependency). Blue dotted lines represent that po- lices depend on past states (action-level dependency).
Dynamics-based behavior analysis in animal neuro- science. Traditional approaches to analyzing animal behav- ior in neuroscience often rely on autoregressive dynamics models. For instance, MoSeq and related works (Wiltschko et al., 2015; Weinreb et al., 2024) assume that animal be- havior consists of multiple segments modeled by a hidden markov model (HMM), with each segment evolving through an autoregressive process. Stone (2023) introduces a switch- ing linear dynamical system (SLDS), similar to ARHMM, but with an additional layer of continuous latent states be- tween the behavioral trajectories and the hidden states repre- senting behavioral segments. We argue that if each segment lasts only a few seconds, it represents meaningful action motifs, such as grooming and sniffing. However, if seg- ment is significantly longer and reflects a decision-making process, traditional dynamics-based models may not be suit- able for identifying these long-tern segments. However, these dynamics- basdd models are not entirely independent of SWIRL. We will demonstrate that SWIRL generalizes purely dynamics- based models by relying on a more princi- pled IRL framework to identify multiple decision-making processes. Our goal is to offer insights that bridge these traditional and new models for animal behavioral analysis.
from an expert agent x", comprising a set of state-action pairs ((sj,a))) The goal is estimate the policy and reward that generated these state-action pairs, often referred to expert demonstrations in IRL literature. We assume that we have collected N expert trajectories, denoted as D = {Ei, EEi,,,,...,,N {x}. Each trajectory con- sists of a sequence of state-action pairs, represented as En = [(s]] af), ($2, aj)....}, with T, time steps, which may vary across trajectories.
3.3. Switching Inverse Reinforcement learning
Our SWIRL model is built on the HM-MDP. Instead of explicitly knowing the reward for each mode, we will use IRL infer these rewards. Mathematically, each time step t. represent the agent's internal reward function Ts, with an additional dependency on the hidden mode z,- This means that the agent receives a reward rzz based on its cur- rent hidden mode z,, which indicates the decision-making state the animal is in (c.g., water seeking or home seeking), with rz, representing the corresponding intrinsic motiva- tion. Consequently, the optimal policy कe is determined by Tz,- However, SWIRL goes beyond merely embedding IRI. within HM-MDP: We also introduce two levels of history dependency into the model. The full graphical model is depicted in Fig. 1.
3. Methods
3.1. Hidden-Mode Markov Decision Process
A discounted Hidden-Mode Markov Decision Process (HM- MDP) is defined by the tuple M = S, P,,, 7)) Here, Z represents a finite set of hidden modes 2, S de- notes the finite state space, and A indicates the finite action space. r= represents the reward function under hidden mode The discount factor 7 is constrained to the interval [0, 1]. Starting from an initial state Su, the agent (animal) selects an action a based on its policy (behavioral strategy) t and subsequently receives a reward determined by Yz, E : {21, 22, ..,,, Am where m represents the total number of modes. The agent then transitions to the next state s' according the transition kernel P(s')s, (/|ss, a), while the agent's hidden mode also transitions to 2f based on the transition probability (=)))
The decision-level dependency is characterized by the idea that animals make new decisions based on their previous choices. The transitions between decision-making processes already account for this decision-level dependency since Pa (5 However, the hidden modes with such clas- sical transition are generated through an open-loop process: the mode 24+1 depends solely on the preceding mode 2., with 2++1 & being independent of the observation state 8t- Consequently, if a discrete switch should occur when the an- imal enters specific region of the state space, the classical transition will fail to capture this dependency. To address this, we extend the transition model to include the state 8, as condition, resulting in P= F2, st), which effectively captures the desired relationship between decisions and the animal's 's location.
3.2. Inverse Reinforcement Learning
Inverse Reinforcement Learning (IRL) addresses the scenario where we have gathered multiple trajectories
For action-level history dependency, we treat both the re-
3

---

Inverse Reinforcement Learning with Switching Rewards and History Dependency for Characterizing Animal Behaviors
ward and policy under hidden mode as functions de- pendent the previous L states, specifically r= : 8 x A + R and u5 8 + A, where L E N and s denotes the cartesian product of L state spaces. To sim- plify the notation, we denote an element of S as s4, so that r=(s",,, a) := r= ,....., s", a) and T. (=(a|s")= := Tz (alss, 6², ;
Detailed SWIRL algorithm and EM derivation are provided in Appendix A. & A.2.
It is worth noting that this EM derivation seeks to maxi- mize the logelikelihood of policy Ts, (ay (T given the posterior of 2, (Eq. 3). This procedure can be incorporated into the Maximum Entropy IRL (MaxEnt IRL) (Ziebart et al., 2008: 2010) framework, since optimizing the reward by maximizing log corresponds to the Lagrangian dual of the standard MaxEnt IRL approach (Zeng et al., 2022). As a result, SWIRL can leverage a variety of IRL methods within the MaxEnt IRL framework: SWIRL can be inte- grated with model- methods to infer the policy used to optimize the reward, such as Soft-Q iteration (Haarnoja et al., 2017), or model-free methods, such as Soft-Q learn- ing (Haarnoja al., 2017) and Soft Actor-Critic (Haarnoja et al., 2018). Alternatively, we also convert the MaxEnt IRL. problem to a model-free inverse Q-learning problem, such as IQ-Learn (Garg et al., 2021). Although model-free methods are often more general and scalable, we find that the model- based Soft-Q iteration approach is more suitable for our behavioral neuroscience study setting. where envi- ronment sizes tend be small, data availability is limited, and the accuracy requirement for reward recovery is high. Accordingly, we use the Soft-Q iteration approach in our main experiments. Nevertheless, in Appendix C.1 we also demonstrate that SWIRL can be integrated with the model- free IQ-Learn to achieve solid performance, illustrating SWIRL's potential in broader applications.
, ) and pad both functions with dummy
variables if the current time step is less than L. While it is also possible to incorporate previous actions into the re- ward function, we focus on state-only rewards for simplicity, which is common practice in IRI. Nonetheless, this frame- work naturally extends from a single-state dependency to a history of states.
Furthermore, we can view the decision process as being non- Markovian, meaning that the current decision or action de- pends not only the current state but also on the history of previous states and actions. Noticeably, there are various ap- proaches to address non-Markovian decision processes, in- cluding state augmentation (Sutton, 1991), recurrent neural networks (Bakker, 2001; Hausknecht& & Stone, 2015), Neu- ral Turing Machines (Parisotto & Salakhutdinov, 2017) and so forth. In this paper, we adopt the most common approach- state augmentation; however, the framework can also be implemented using more advanced and scalable methods.
3.4. SWIRL Inference Procedure
The goal of inference is to learn the hidden modes and the model parameters 0- (Pz, 1',, T2. p(z)) given the collected trajectories 'D. Here, p(s1) and p(z1) repre- sent the probabilities of the initial state and hidden mode, respectively. The variables rN and n. denote the reward and policy associated with the hidden mode 2, while Pe is the transition matrix between hidden modes. We can maximize the likelihood of the demonstration trajectories D learn the optimal 0*, such that 6* = argmax, log P(D)0) (MLE). However, achieving this objective requires marginalizing over the hidden modes 2, which is intractable. To address (FM)
3.5. Connection to Dynamics- Behavior Methods
Traditional methods for analyzing animal behavior in neu- roscience often use autoregressive dynamics models, with the autoregressive hidden Markov model (ARHMM) being the most prevalent (Wiltschko al., 2015; Weinreb et al., 2024). ARHMMs assume that the animal behavior con- sists of multiple segments represented by hidden Markov model, where each segment evolves through an autoregres- sive process. Using the notation established earlier, we denote hidden modes as ž at time t, following the transi- tion 1(2(+))2)) At each time step t, the observation state 8, follows conditionally linear (or affine) dynamics, deter- mined by the discrete mode 2, This can be expressed as (1) S = As, St + Ut, where As, is the linear dynamics associ- ated with 2; and 2/ represents Gaussian noise. If 2, changes, the linear dynamics will also change accordingly. More
this issue, we employ the Exxpectatio-MMxxinizae algorithm, alternating between inferring the posterior distri- butions of the hidden modes and updating the parameters with the following auxiliary function:
G 0) = log + Σ plz1lk 0) log
log 18
(2) generally, we can represent the dynamics as p(s |st, z.). Consequently, the overall generative model for the ARHMM can be summarized as follows: (1) 2t p(ze|=t-)). 1). and (2) (3) St+1 ~ p(s++1 (st, 2)). Let's outline the generative model of SWIRL without history dependency:
Σ p(zreE 0og n2, Tz, (adl)brr))
(1) 2 p(=x ((=/ i), and (2) st+1 ~ Ea, p(s (a, The term (4) I T 2a) arises because the policy is derived from Tz. Consequently, the primary distinction between ARHMM
Zt
Σ
p( 2+1 k,, log P
1=1 Re.8e+1

---

Inverse Reinforcement Learning with Switching Rewards and History Dependency for Characterizing Animal Behaviors
and SWIRL lies in the dynamics used to generate 81+1-
4. Results
We can show that SWIRL is a more generalized ver- sion of ARHMM. In deterministic MDP, where p($t+1/st,ai) is a delta function and each action at uniquely determines St+1, 8t+1 directly implies at- Thus, En, 1(5c+1|st; de) (ar]s;, 2)) = T(ar|st, z+) =
Throughout the experiment section, we use the following terminology to denote our proposed algorithms and the base- line models we compare:
MaxEnt (Ziebart et al., 2008; 2010): Maximum Entropy IRL where the reward function is only a function of the current state and action. It is a singlc-mode IRL approach with a single reward function.
p($(+1|St, 2t). This effectively reduces SWIRL to
ARHMM. In the second real-world experiment (Sec. 4.3),
the MDP setup satisfies these assumptions. In such case, ARHMM can be seen as performing policy learn- ing through behavioral cloning without learning reward function, whereas SWIRL employs IRL. to learn the policy.
Multi-intention IQL (Zhu et al., 2024): learns time- varying reward functions based on HM-MDP. It is a SWIRL model with no history dependency.
Locally Consistent IRL (Nguyen et al., 2015): learns time vareina reward functions based on HM-MDP. It is SWIRL model with no action-level history dependency. ARHMM (Wiltschko et al., 2015): learns the segmen- tation of animal behaviors using autoregressive dynamics combined with a hidden Markov model.
Having established this connection, we can view SWIRL as more generalized version of ARHMM, as permits the MDP to be stochastic and allows multiple actions to result in the same preceding state. Additionally, explicitly mod- eling the policy introduces a reinforcement learning frame- work that better represents the decision-making processes of animals and reveals the underlying reward function. For SWIRL with history dependency, we can further connect it to the recurrent ARHMM (Linderman et al., 2016), which expands p(z++1| 2x) to p(ze+1 | 8t).
FARHMM (Linderman et al., 2016): recurrent ARHMM whose transition probability of the hidden modes also relies the state.
I-1, I-2: the baseline variant of our proposed SWIRL method which assumes the transition kemel Pa is indepen- dent of the state. The reward and policy depend either on the current state (in the case of I-l) or on both the current and previous states (in the case of 1-2). Note that I-1 represents the simplest version of SWIRL, which corresponds to Multi- intention IQL. Thus, we use I-1 to denote Multi-intention IQL. The model can incorporate an arbitrary history length L for the policy and reward; in this paper, we use L = 1 and L= 2.
Another advanced version of the ARHMM is the switching linear dynamical system (SLDS), which assumes that the state 8, is unobserved. Instead, the observed variable is a linear transformation of st. Thus, the complete genera- tive model for SLDS consists of: (1) 2 ~ p()) (2) 8141 ~ p(s++1(s;, 2t), and (3) yc+1 ~ p(8c+1/$(+1). This suggests that the representation 8t capturing the primary dynamics is, in fact, a latent representation of the observa- tion Vt- Building on this concept, we can extend SWIRL into a latent variable model, where St serves as the latent representation of the true observation state yt- This corre- sponds to the setup of Partial Observation Markov Decision Process (POMDP) in the literature. While SWIRL is not a POMDP method in the formal sense (c.g., modeling be- lief states or latent dynamics explicitly). it does address non-Markovian behavior by augmenting the state space with recent history (c.g., using $t, St-1, $t-2,..). This is a commonly used strategy to handle partial observability or Lon-Marrovianity--inne similar in spirit to how LSTMs or memory-based policies are used in POMDP settings. There- fore, SWIRL leverages temporal context in a manner similar to methods designed for POMDPs. This extension will link SWIRL to representation learning in reinforcement learning, which can be explored further in future work.
S-1, S-2: our proposed SWIRL method where Pa is state dependent. The suffix follows the same setup as above. S-1 corresponds Locally Consistent IRL.
SWIRL code can be found at ht / http:/ggttttbb.ec com/ BRA TTMI.GG/SWW RR.
4.1. Application to simulated gridworld environment
We begin by testing our method on simulated trajectories within 5 x 5 gridworld environment, where each state al- lows for five possible actions: up, down, left, right, and stay. The agent alternates between two reward maps: a home reward map and a water map (see Fig. 2A). Following the design of real animal experiments (Rosenberg et al., 2021), we assume that the water port provides water to the agent only once per visit. Therefore, under the water reward map, the agent receives a reward for (1) visiting the water state if it was not in the water state previously or (2) leaving the water state. The home reward map returns a reward at the home state. This leads to a non- Ma-Maaoovian reward func- tion that relies on both the current state and the previous state. We employed soft-Q iteration to determine the optimal policy for each reward function and generated 200 trajecto-
Thus, we argue that SWIRL offers a more generalized and principled approach to studying animal behavior compared to commonly used dynamics-basee models, as one can draw inspiration from the development of (latent) dynamics mod- els to enhance advanced IRL methods for analyzing animal decision- decision-making processes.
5

---

Inverse Reinforcement Learning with Switching Rewards and History Dependency for Characterizing Animal Behaviors
A
true
discovered
B
home reward
St1
R coreeaation
current
right
current
above
Test sagment
true
discovered
current
loft
water reward
current
below
current
current
00
MaxEnt
8/2
MaxEnt
S-1 $22
1-2
S:1
Figure 2. Simulation experiment on a 5 x gridworld. (A) Comparison between the true and discovered reward maps. The color scale represents reward values ranging from 0 to 1. The home reward is defined as r(*.). while the water reward depends on both the current and previous locations, r(St, 8t-1). To present the water reward, cach location is divided into five groups, as detailed in the table on the far right. For example, the polygon in the first row represents the reward value when is the current location and 8 is the location to the right. Light grey indicates an impossible transition where no reward exists. (B) Box plots illustrating the Pearson correlation between the true and recovered reward maps, test log-likelihood, and test segmentation accuracy over 10 runs. The x-axis represents the five different models. Outlier selection method is described in Appendix B.8.
ries based on the learned policy, using a history-dependent hideen-mode switching dynamic Pe (Z 12,, st). Accord- ingly, the agent is more likely to switch to the home map after visiting the water port and to switch to the water map after returning home. Each trajectory consists of 500 steps.
To assess the robustness of SWIRI., in Appendix B.7, we evaluated the performance of SWIRL (S-2) under increasing levels of random perturbations in this gridworld dataset. We found that SWIRL can tolerate moderate levels of noise or incomplete data (with 30% of the states and actions in train- ing data permuted), making suitable for real-world animal behavior datasets where such challenges are common.
We then used SWIRL to learn the reward functions and the transition dynamies between them, based 80% of the generated trajectories. As a baseline, we employed the Max- imum Entropy IRL (MaxEnt) method and tested four varia- tions of the SWIRL models (I-1, I-2, S-1, S-2), with I-1 rep- resenting multi-intention IQL. Fig. 2A displays compari- son between the true and discovered reward functions, while Fig. 2B presents boxplots showing the Pearson correlation between the true and recovered reward functions, along with the test log- likelihood (LL) and test segmentation accuracy (which measures the ability to predict the correct segments for home and water modes). The test performance was eval- uated using the remaining 20% of the trajectories. Notably, accurate reward recovery was only achieved with the S-2 model. All four SWIRL variations outperformed MaxEnt, indicating the presence of more than one hidden model. Both the state dependency of hidden-mode transitions (decision-level dependency) and the history dependency reward function (action-level dependency) contributed to further improvements test LL and segmentation accu- racy. Specifically, only the state-dependent models (S- S-2) could accurately and robustly recover test segments, while the independent models (I-1, I-2) exhibited lower accuracy with higher variance. This is attributed to the non- Markovian reward design, where the agent can only receive water once per visit. Notably, S-2, the full SWIRL model incorporating both decision-leved and action-level dependen- cies, demonstrated the best performance across all metrics.
4.2. Application to long, non-stereotyped mouse trajectories
We then applied SWIRL to the long, non-stereotyped trajec- tories of mice navigating a 127-node labyrinth environment with water restrictions (Rosenberg et al., 2021). In this ex- periment, a cohort of 10 water-deprived mice moved freely in the dark for 7 hours. A water reward was provided at an end node (Fig. 3A), but only once every 90 seconds at most. Similar the simulated experiment, the 90-second condi- tion forces the mice to leave the port after drinking water, leading to non-Markovian internal reward function. For our analysis, we segmented the raw node visit data into 238 trajectories, each comprising 500 time points. This data for- mat presents a considerably greater challenge compared to the same dataset processed with more handerafted methods in previous IRL applications (Ashwood et al., 2022a; Zhu et al., 2024), which were limited to clustered, stereotyped trajectories of only about 20 time points in length.
1.2.I. INFERRED INTERPRETABLEE HISTORY-DEPEEDENT REWARD MAPS
We applied SWIRL to 80% of the 238 mouse trajectories from the watereret labyrinth experiments. Accord- ing Rosenberg et al. (2021), mice quickly learned the labyrinth environment and began executing optimal paths
6

---

Inverse Reinforcement Learning with Switching Rewards and History Dependency for Characterizing Animal Behaviors
A
B
F:0.4
water
home
explore
water
110
115 57 oo
explore
n=0.7
MaxEni 1-1
1-2
S-1
8-2
117
58
11.
water
home
explore
end
H-1
500
1-2
S-1
S-2
bmc
500
water
home
explore
water port visit
home visi
start
Figure 3. Water-restrieted labyrinth experiment. (A) Setup for the labyrinth experiment. (B) Boxplot showing held-out test LL over 10 runs, with the x-axis representing the five different models. Outlier selection method is described in Appendix B.8. (C) History dependency inferred by SWIRI. (S-2), as reflected in reward maps of the water mode and the explore mode. (D) Inferred reward maps from SWIRL (S-2) under three hidden modes: water, home, and explore. To enhance visualization, the inferred reward r(s., 8(-1) was averaged over 8 to produce r(st). (E) Segments of a trajectory from held-out test data, predicted by four SWIRL models. The orange dot indicates when the mouse visits the water port, while the red cross denotes the mouse's visit to state 0 (home) that time. (F) Trajectories segmented into hidden modes based on SWIRL (S-2) predictions.
from the entrance to the water port within the first hour of the experiment. Therefore, we assume the mice acted opti- mally concerning the internal reward function guiding their behavior. We tested a number of hidden modes Z from 2 to in Appendix 1. and found Z = 3 to be the best fit. Fig. 3B displays the held-out test LL for MaxEnt and the SWIRL variations based on the remaining 20% of trajectories. The state dependency in hidden-mode switching dynamics and the history dependency in the reward function contributed to improved test performance. The final SWIRL model (S-2) successfully recovered a water reward map, a home reward map, and an explore reward map (Fig. 3D). For better visu- alization, we averaged the S-2-recovered history-dependent rewards across previous states and normalized the reward values to range of (0, 1). In the water reward map, mice received a high reward for visiting the water port. In the home reward map, there was a high reward for visiting state O at the center of the labyrinth, which also served as the entrance and exit. Mice occasionally went to state 0 to enter or leave the labyrinth and sometimes passed by on their way to other nodes. In the explore reward map, mice received high reward for exploring areas of the labyrinth other than state O.
substantial reward (0.7). This observation aligns with the water port design, as mice can only obtain water once every 90 seconds. Consequently, it makes sense that the mice would want to leave the water port after reaching it.
4.2.2. SWIRL INFERRED INTERPRETABLEE IISTORRY-DEPNNDEN IIDDEN-MODE SEGMENTS
We then visualized all mouse trajectories based the hidden-mode segments predicted by SWIRL (S-2) (Fig. 3F). In segments classified as water mode, mice start from var- ious locations in the labyrinth and move toward the water port. In segments identified as home mode, mice begin from distant nodes and head toward the center of the labyrinth (home). In segments categorized explore mode, mice start from junction nodes the water port and explore end nodes. This result demonstrates that SWIRL can identify sub- tub-trajectoriis of varying lengths from raw data spanning 500 time points, allowing us to visualize them together and reveal clustered behavioral strategies. This capability has not been achieved by previous studies on freely moving animal behavior over extended recording periods, and we conducted this analysis without prior knowledge of the lo- cations of the water port or home.
Importantly, SWIRL recovered an interpretable history- dependent reward map for the water port (Fig. 3C). In the water mode, mice always receive high reward for reaching the port, irrespective of where they were before. On the contrary, in the explore mode the reward depends their previous location: if the mouse was already at the water port, leaving it is valued more (1.0) than staying (0.4); if the mouse came from elsewhere, reaching the port still provides
We also provide a detailed visualization of the hidden-moot segments from an example trajectory in the held-out test data and compare the segmentation performance of the four SWIRL variations (Fig. 3E). In the S-2 segments, visits to the water port (indicated by orange dots) consistently occur near the end of water mode segment, while visits to state O (home) (indicated by red crosses) typically happen at the
7

---

Inverse Reinforcement Learning with Switching Rewards and History Dependency for Characterizing Animal Behaviors
conclusion of home mode segment. Notably, home mode segments that do not include visit to state 0 can still be valid, as these segments may end at state 1 or 2. In contrast, the 1-1, I-2, and S-1 segments exhibit instances of water segments that do not involve a visit to the water port, along with many home segments that lack clear interpretability. Overall, S-2 successfully identifies robust segments of rea- sonable length, avoiding the numerous rapid switches seen in the other variations. We attribute this to both the state dependency of hidden mode transitions and the history de- pendency in rewards. This suggests that mice are unlikely to make quick changes in their decisions; instead, they make choices based on their current location and take into account at least two locations while navigating the maze.
to compare the performance of SWIRL with ARHMM and its variant, rARHMM.
We applied SWIRL, rARHMM, ARHMM, and MaxEnt to 80% of the trajectories and assessed the held-out test LL on the remaining 20% (Fig. 4B). All four SWIRL models outperformed ARHMM and rARHMM on this dataset, indi- cating that learning rewards is more beneficial for behavior segmentation and explaining the behavior trajectories. In- terestingly, the history dependency in the reward function resulted in lower test LL. as S-1 and I-1 demonstrated higher test LL than S-2 and I-2. We believe this is attributable to the merging of consecutive identical syllables and the selec- tion of the top 9 syllables during the preprocessing phase for this dataset. As a result of these steps, the actual time interval between 9--1 and s, may vary significantly, lead- ing to poorly defined time concept that complicates the model's 's ability to capture the history dependency in the re- ward function. However, we can use SWIRL with different variations as hypothesii-testingg tool. The variation yield- ing the highest test LL may be regarded as more accurately reflecting the dynamics and structure of the data. Conse- quently, these results suggest that the behavior trajectories exhibit only Markovian dependency rather than long-term Mon-Markovian dependency. Since S-1 remains higher than I-1, we conclude that the state dependency in the hidden mode transition contributes to explaining the data.
Furthermore, in Appendix B.6.2, we found S-3 and S-4 ex- hibited even higher test LL than S-2, suggesting the presence of longer action-level history dependency (I. > 2) in this labyrinth dataset. This observation aligns with the partially observable nature of this 127-node labyrinth: The mouse may not know the whole environment, so it likely relies on a longer history of prior states to guide its decision-making However, hidden mode predictions of S-3 and S-4 are very close to that of S-2, suggesting that S-2 is already powerful enough for hidden mode prediction.
4.3. Application to mouse spontaneous behavior trajectories
The best SWIRL model (S-1) recovered reward maps and hidden-mode segments provide insights into the variability of dopamine impacts on animal spontaneous behavior: As illustrated in Fig. 4A, the reward maps exhibit some similar- ities along with meaningful differences. For certain reward maps, there is a relatively high correlation (e.g., 0.40 and 0.48) with dopamine fluctuations during the corresponding modes. This suggests that dopamine fluctuations can reflect a certain extent of reward during hidden modes I and 4.
We also employed SWIRL on a dataset in which mice wan- dered an empty arena without explicit rewards (Markowitz et al., 2023). In this experiment, mouse behaviors were recorded via depth camera video, and dopamine fluctuations in the dorsolateral striatum were monitored. The dataset includes behavior "syllables" inferred by MoSeq (Wiltschko et al., 2015), which indicate the type of behavior exhibited by the mice during specific time periods (e.g.. grooming, sniffing, etc.). Consequently, trajectories consist of behay- ioral syllables, with each time point representing syllable. We selected 1,59 trajectories, each comprising 300 time points, by retaining only the 9 most frequent syllables and merging consecutive identical syllables into a single time point. This method, also used in previous reinforcement learning studies on this dataset (Markowitz et al., 2023), sures that each syllable has sufficient data for learning and allows the model to concentrate on the transitions between different syllables.
Furthermore, the plot of hidden mode segments across all trajectories (Fig. 4C) shows how each spontaneous behav- ior trajectory is temporally segmented into distinct latent modes, each governed by one of the five learned reward functions (h1-h5). Early in the trajectories, we observe a higher prevalence of h1, h2, and h4, while h3 and h5 increasingly dominate toward the later stages. This tempo- ral pattern aligns well with the findings from prior work (Markowitz et al., 2023), which reported that mice exhibit greater activity at the beginning of spontaneous behavior sequences and progressively slow down and disengage over time. The reward maps provide mechanistic insight into these transitions: hl encourages selective transitions into pause while still supporting other state transitions, represent- ing a modulated rest phase; h2 discourages passivity, with low reward for transitions into pause and strong suppression of grooming; and h4 favors transitions into sniffing, reflect-
The MDP for this experiment comprises 9 states and 9 ac- tions, where the state represents the current syllable and the action signifies the next syllable. As mentioned in Sec. 3.5, the ARHMM can be viewed as variant of SWIRL that learns the policy through behavior cloning. In other words, the policy for this MDP aligns with the emission probability of the ARHMM. This setup offers an excellent opportunity
8

---

Inverse Reinforcement Learning with Switching Rewards and History Dependency for Characterizing Animal Behaviors
A
h1, DA corr=0.40
h2, DA corr=0.04
h3, DA corr=0.18 18
h4, DA corr=0.48 h5, DA cor:=0.27
walk
rear
state
groom
snitf?
pause
walk
real
groon
pause
walk
pauss
littr
Linus
action
-1.507
2.00
hidden mode
hidden made
12
hidden mode
Test
-1.500
ARHMM
hidden mode
hidden mode
-1.75
MaxEnt: -2.206
2
4
number modes
lime 300
Figure Mouse spontaneous behavior experiment. (A) SWIRL (S-1) inferred reward maps five hidden modes. denotes hidden mode 1, and For better visualization, reward values normalized to range of (0, 1). DA corr represents the Pearson correlation between the inferred reward map and the averaged dopamine fluctuation levels. (B) Held-out test I.I. each model across different number of hidden modes 10 runs. The shaded represents the total area that falls between standard deviation above and below the mean. (C) Inferred hidden-mode segments for all trajectories, with each row representing a trajectory.
ing exploratory drive. In contrast, h3 and h5 strongly reward transitions into pause and suppress departures from it, in- dicating globally suppressive states that enforce behavioral quiescence. The increasing expression of h3 and h5 over time reflects the mice's natural progression from active inactive states, supporting the interpretation that the learned rewards capture underlying behavioral structure consistent with known temporal dynamics.
combined with its connection to traditional dynamics-based behavior analysis methods in neuroscience, make SWIRL well-suited for behavioral neuroscience applications. Future extensions of SWIRL to partially observable Markov deci- sion processes (POMDPs) could further broaden its utility, offering enhanced insight into complex and nuanced animal behaviors.
Beyond its relevance to neuroscience study, SWIRL also benefits the broader machine learning community. It of- fers flexible framework with history-dependent switch- ing reward functions that can be integrated with various model-free IRI. algorithms. For example, in Appendix C.I, highlight SWIRL's potential be integrated with IQ- Learn, a state-of-the-arr scalable model-free IRL method that has been employed in training large language models (Wulfmeier et al., 2024). Future development of SWIRL that emphasizes scalability, such as incorporating RNNs or transformer-base policy model will particularly valuable for areas such robotics and large language models, where adaptable IRL strategies increasingly in demand.
While the non-Markovian action-level history dependency introduced by SWIRL does not demonstrate superior per- formance in this particular experiment, the findings show- case SWIRL's unique contribution to neuroscience research. Specifically, SWIRL serves powerful tool for hypothesis testing in behavioral datasets, enabling researchers vali- date or challenge hypotheses regarding decision-level depen- dency as well as non-Markovian action-level dependency. This versatility further confirms SWIRL's great potential in advancing our understanding of complex behaviors.
5. Discussion
We introduce SWIRL,, novel IRL. framework designed to capture history-dependent switching reward functions in complex animal behaviors. Our framework can infer interpretable switching reward functions from lengthy, non- stereotyped behavioral tasks, achieving reasonable hidden- mode segmentation-a feat that, the best of know]- edge, has not been accomplished previously. Those abilities,

---

Inverse Reinforcement Learning with Switching Rewards and History Dependency for Characterizing Animal Behaviors
Impact Statement
Garg, D., Chakraborty, S., Cundy, C., Song, J., and Ermon, S. Iq-learn: Inverse soft-q learning for imitation. In Ranzato, M., Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 4028-4039 Curran Associates, Inc., 2021.
SWIRL has the potential to advance neuroscience research by helping uncover connections between complex animal cerision-mmaing and corresponding brain activities. In broader context, SWIRL's capacity learn multiple history- dependent switching reward functions can also enable en- hanced policy learning and explainability in fields such as robotics and large language models. However, as with all methods that analyze complex behaviors, careful considera- tion of data privacy, fairness, and responsible deployment is imperative to ensure ethical use.
Haarnoja, T., Tang, H., Abbeel, P., and Levine, S. Rein- forcement learning with deep energy-based policies, In International conference on machine learning, pp. 1352- 1361. PMLR, 2017.
Haarnoja, T., Zhou, A., Abbeel, P., and Levine, S. Soft actor- critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In Dy. J. and Krause, A. (eds.), Proceedings of the 35th International Confer- ence on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 1861-1870. PMLR, 10-15 Jul 2018. URL.ttps:: /ttpps:///prrooee /pwoceddoeeeeeeeeooooooocccecece ngs mmlr. press/v80 http/haarnooaaaaaaaaaa html.
Acknowledgements
This work was supported by National Science Foundation NSF2336055 (AW), a Career Award at the Scientific Inter- face from the Burroughs Wellcome Fund (JEM), the David and Lucille Packard Foundation (JEM), and the Sloan Foun- dation (AW& & JEM).
References
Hattori, R., Danskin, B., Babic, Z., Mlynaryk, N., and Komiyama, T. Arca-specificity and plastic- ity of history-dependent value coding during learn ing. Cell, 17707: 2019. ISSN 0092-86774. doi: http:///ooooog//00000j...ee.... com/
Abbeel, P. and Ng, A. Y. Apprenticeship learning via inverse reinforcement learning. In Proceedings of the Twenty- First International Conference Machine Learning, ICML '04, pp. 1, New York, NY, USA, 2004. Asso- ciation for Computing Machinery. ISBN 1581 138385. doi: 10. 10.1145//115330.1155533 URL https: //oi.. 0rg/10.1145/1000
027. URL acienee/rrriidee///////222221213
1015430.
Hausknecht, M. and Stone, P. Deep recurrent q-learning for partially observable mdps. In AAAI 2015 Fall Symposium, 2015.
Ashwood, Z., Jha, A., and Pillow, J. W. Dynamic inverse reinforcement learning for characterizing animal behay- ior. In Oh, A. H., Agarwal, A., Belgrave, D., and Cho, K. (eds.), Advances in Neural Information Processing Systems, 2022a. URL https: / wwwp//wwweeeeeeeeeeeeeeeeeo n Forum??d-noonggu5Ki
Hirakawa, T., Yamashita, T., Tamaki, T., Fujiyoshi, H., Umezu, Y., Takeuchi, I., Matsumoto, S., and Yoda, K. Can ai predict animal movements? filling gaps in animal trajectories using inverse reinforcement learning, Eco- sphere, 910):e02447, 2018.
Ashwood, Z. C., Roy, N. A., Stone, I. R., Urai, A. E., Churchland, A. K., Pouget, A., Pillow, J. W., and Lab- oratory, T. I. B. Mice alternate between discrete strate- gies during perceptual decision-making. Nature Neuro- science, 25(2):201 212, Feb 2022b. ISSN 1546-1726. doi: 10011088/4411933-21--000000 doi.
Houthooft, R., Chen, X., Duan, Y., Schulman, J., De Turck, F., and Abbeel, P. Vime: Variational information max- imizing exploration. Advances in neural information processing systems, 29, 2016.
URL https:// 021000007---1-111111100100011011
Kennedy, A. The what, how, and why of natu- ralistic behavior. Current Opinion in Neuro- biology, 74:102549, 2022. ISSN 0959-4388. doi: URL https://www.ccienedd.cee httpa////.rtcceee///////92222220
Babes-Vrus Vroman, M., Marivate, V., Subramanian, K., and Littman, M. L. Apprenticeship learning about multiple intentions. In International Conference on Machine Learning, 2011. URL https://api. semanticscholes arg/CorppsID:993765
Linderman, S. W., Miller, A. C., Adams, R. P., Blei, D. M., Paninski, L., and Johnson, M. Recurrent switching lin- ear dynamical systems. arXiv preprint arXiv: 1610.08466, 2016.
Bakker, B. Reinforcement learning with Istm in non- markovian tasks with longterm dependencies. Memory, pp. 1-18, 2001.
10

---

Inverse Reinforcement Learning with Switching Rewards and History Dependency for Characterizing Animal Behaviors
Markowitz, J. E., Gillis, W. F., Jay, M., Wood, J., Harris, R. W., Cieszkowski, R.,, Scott, R., Brann, D., Koveal, D., Kula, T., Weinreb, C., Osman, M. A. M., Pinto, S. R., Uchida, N., Linderman, S. W., Sabatini, B. L., and Datta, S. R. Spontaneous behaviour is structured by reinforce- ment without explicit reward. Nature, 614(7946):10S- 117, January 2023.
Weinreb, C., Pearl, J. E., Lin, S., Osman, M. A. M., Zhang, L.., Annapragada, S., Conlin, E., Hoffmann, R., Makowska, S., Gillis, W. F., et al. Keypoint-moseq; pars- ing behavior by linking point tracking to pose dynamics. Nature Methods, 21) (1()))339-40) 1222999933 2024.
Wiltschko, A. B., Johnson, M. J., Jurilli, G., Peterson, R. E., Katon, J. M., Pashkovski, S. L., Abraira, V. E., Adams, R. P., and Datta, S. R. Mapping sub-second structure in mouse behavior. Neuron, 88(6): 121-1135, 2015.
Mazzucato, L. Neural mechanisms underlying the temporal organization of naturalistic animal behavior. Elife, 11:
Wu, F., Ke, J., and Wu, A. Inverse reinforcement learning with the average reward criterion. Advances in Neural Information Processing Systems, 36, 2024.
Ng. A. Y. and Russell, S. Algorithms for inverse rein- forcement learning. In Proceedings of the Seventeenth International Conference Machine Learning, ICML '00, pp. 663-670, San Francisco, CA, USA, 2000. Mor- gan Kaufmann Publishers Inc. ISBN 1558607072.
Wulfimeier, M., Bloesch, M., Vieillard, N., Ahuja, A., Born- schein, J., Huang, S., Sokolov, A., Barnes, M., Desjardins, G., Bewley, A., Bechtle, S. M. E., Springenberg, J. T., Momchev, N., Bachem, O., Geist, M., and Riedmiller, M. Imitating language via scalable inverse reinforcement learning. In The Thirty-eighth Annual Conference on Neu- ral Information Processing Systems, 2024. URL https: / / openrevieww.... net. dddmmmissddddddidee
Nguyen, Q P., Low, K. H., and Jaillet, P. Inverse reinforce- ment learning with locally consistent reward functions. In Neural Information Processing Systems, 2015.
Parisotto, E. and Salakhutdinov, R. Neural map: Structured memory for deep reinforcement learning. arXiv preprint arXiv:1702.088. 0000000806, 2017.
Yamaguchi, S., Naoki, H., Ikeda, M., Tsukada, Y., Nakano, S., Mori, I., and Ishii, S. Identification of animal behav- ioral strategies by inverse reinforcement learning. PLoS computational biology, 14(5)=e 2018.
Pinsler, R., Maag. M., Arenz, O., and Neumann, G. Inverse reinforcement learning of bird flocking behavior. In ICRA Swarms Workshop. 2018.
Rosenberg, M., Zhang, T., Perona, P., and Meister, M. Mice in a labyrinth show rapid learning, sudden insight, and efficient exploration. eLife, 10:e66175, jul 2021. ISSN 2050-084X. doi: 10. 10.7554/eL.ife.66177 //doi. wwo/oo/////o.o. 10. /004//eiffe www.e.e.......... 65175.
Zeng, S., Li, C., Garcia, A., and Hong, M. Maximum- likelihood inverse reinforcement learning with finite-time guarantees. In Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 10122 10135. Curran Associates, Inc., 2022.
URL https:
Sezener, C. E., Uchibe, E., and Doya, K. Obtaining reward functions of rats using inverse reinforcement learning. In Türkiye Autonomous Robots Conference, 2014.
Zhu, H., Crompe, B. D. L., Kalweit G., Schneider, A., Kalweit, M., Diester, I., and Boedecker, Multi- intention inverse q-learnng for interpretable behavior representation. Transactions Machine Learning Re- search, 2024. ISSN 2835-8856. URL https:// openrevien net wtr//orrmm//mm d=h
Sharafeldin, A., Imam, N., and Choi, H. Active sensing with predictive coding and uncertainty minimization. Patterns, 5(6), Jun 2024. ISSN 2666-3899, doi: 10. 10. (0.1016/j.pattee 20298100983. URL https: wrd.i..org/ii. 1016/ j.patter. 2 2024 100983.
Ziebart, B. D., Maas, A., Bagnell, J. A., and Dey, A. K. Maximum entropy inverse reinforcement learning. In Proceedings of the 23rd National Conference on Artificial Intelligence Volume 3, AAAI'08, pp. 1433-1438. AAAI Press, 2008. ISBN 9781577353683.
Stone, I. Latent Variable Models for Characterizing the Dynamic Structure Underlying Complex Behaviors. PhD thesis, Princeton University, 2023.
Surana, A. and Srivastava, K. Bayesian nonparametric inverse reinforcement learning for switched markov de- cision processes. In 2014 13th International Conference on Machine Learning and Applications, pp. 47-54, IEEE, 2014.
Ziebart, B. D., Bagnell, J. A., and Dey, A. K. Modeling in- teraction via the principle of maximum causal entropy. In Proceedings of the 27th International Conference on In- ternational Conference on Machine Learning, ICML' 10, pp- 1255-1262, Madison, WI, USA, 2010. Omnipress. ISBN 978 1755589977
Sutton, R. S. Dyna, an integrated architecture for learning, planning, and reacting. ACM Sigart Bulletin, 2(4):160~ 163, 1991.

---

Inverse Reinforcement Learning with Switching Rewards and History Dependency for Characterizing Animal Behaviors
A. Appendix A
A.I. The SWIRL algorithm
Following the EM update scheme, we derive the auxiliary function for the n-th trajectory during the E-step, where ........
Te Gn(0,0) log (8)1(( + Σ 1 log +
log P (sn,t+ 188,,, an,)
(5)
T
p(zn, log ann (an, (lm,
(6)
1=1
T
plan,e 2n,4t1 log Pz (2(2n,t41l2n, Sn.))
t=1 2m,11
Here are some remarks: (I) We incorporate state dependency into the hidden mode transition Pa. such that 2,+1 depends only on the previous hidden mode 24 but also on the current state st. This modification results longer segments of hidden modes with reduced fast-switching phenomena. (II) We can perform MaxEnt IRL based the policy term Tzn, (un, he ra) in Eq. 6, as discussed in Sec. 3.4. (III) If we have estimated the current policy T2, based on the current reward estimate Tz, we can apply any inference method to estimate the posterior probabilities p(zn,elEnn 0) and p(Zn,t, En,t+1 Pn,t+ilEn, (,,00) f). In this work, we use the standard nawarddbackwaad to note that (( (n,tt (n,z) represents the environment transition and is not involved in the optimization process.
passsag--assiigg algorithm. (IV) It is important
In the M-step, to maximize the auxiliary function G, we update all parameters by performing gradient descent on G. The inference procedure alternates between the E-step and M-step until convergence predetermined number of iterations. The algorithm summarized in Algorithm 1.
Algorithm The SWIRL Algorithm
Data: Expert demonstrations D = (21,,2,.... }
Result: Posterior probabilities of hidden modes z, rewards for each mode, and other parameters in 0 Initialize parameters 8
for & = 1,2,.... K do
E-step
For each hidden mode z and corresponding reward r.. infer the optimal soft Q-function Q using a RL method.
Compute the policy with the Boltzmann distribution:
exp1Q (s. a)/a}
Tz (als))
EreA exp (s.a')/a}
to obtain u. :r)s Vs S. is a predefined temperature parameter. For each trajectory En, forward-hackwart message passing to calculate p(=,, 0K) and p(=n,t, 2n,t+ 0%). Use the posteriors, 75 (a)sst; rf), and 0* to compute the auxiliary function G (Eqs. 5-7).
end
M-step
Update parameters 0 using gradient descent on the auxiliary function G with learning rate me:
gh-1 6r 7hh VeG 0*).
end
end
12

---

Inverse Reinforcement Learning with Switching Rewards and History Dependency for Characterizing Animal Behaviors
A.2. Detailed derivation of SWTRL objectives by EM algorithm
Here need learn 0 A (Ta, p(1)) Note that the total probability for a sequence {(S,, aes 2e) 2e) Jt= 1:77
T
pog 8, a) = pl pogg(()))(((())))))() (1r The expectation across all possible sequences is given by
t=2
N
Ellogp(z, 8,
log ((pp(nn.1)
Σ pf3n,1/51, a1) logp(zn.1)
N
log S,,01,
N
p(zn,e 2n,1-1 log do Pon..lpnnnnnn
n=1
E (l log Znn. la (ane
n=1 t=1
Now use Expectation-Maximizatioo to find 0. E step:
G(O, 0) = Σ p(2|S1NN 0) log p(2, EI:N,
N
Σ
II An=l
Kn 1: 0n)
p(z 1.
Tw
T
log (m,,)) log (og,((n,1)
(log Z.., (an, s f2))
(log P( an,))
+ log Pa P22,,nnt1/2nns (n,))) Sn.t))
N
log p(
+
p(2n,1 = =2 0) log p(zn.1)
p(2n,e 2|En, 0) log zn., (an, (0 rz)
p(zn,t = 2,,411 2'1kn: 0) P Pe(eent 8n.))
T--1
log
M step:
g++1 arg max G(0, G(0,0*).
For notational simplicity, only consider specific trajectory To compute C(0,0h) we need to estimate
(1:7,117 01:0; 0)
13

---

Inverse Reinforcement Learning with Switching Rewards and History Dependency for Characterizing Animal Behaviors
and
p(= 2, 2++1 =11TT 01:T, 0)- Thus we can use message passing algorithm, where we define the at for t = I.:
variables a and B. Forward variables
a1,= p(21 =2(0),
0t.= p(11: a1:11 24 = =20)
Σ
01:-1, 21 =10))pa = 2|81-1,
= 2 z')p(se 18-11 a0- 11-1) (arlst;
Σ at- - (a (o)
(a)lsf;rz).
Backward variables 2r for -1, =l....,T
ST,= 1,
Br, p(st+11TT...t.o a,
= -_ouuspssuuaaaaa:a:::.111111.11 2'120 s(41),
11 p(s1r, 01
Σ (4 21120 81)
Therefore,
plu = 2)
=p(z 2.E10)/r(E10)
1181,, = (001p((04112 04, = 2,0)/p(5 ()
Furthermore,
= =, ht 2 1194 8,, az)/p(E(0)
Xua. 2 (1:4111 de+1Tl2s. $t, as)
1 01-1, Ps /// (s (acls
p({]0)
And finally,
p(ε(0) Σ aT, Σ 21
14

---

Inverse Reinforcement Learning with Switching Rewards and History Dependency for Characterizing Animal Behaviors
B. Appendix B
B.1. Implementation details
We developed our SWIRL framework in JAX, leveraging its automatic differentiation and parallclization features to efficiently handle the computational demands of our algorithm. During each M-step update, we performed gradient descent optimize the model parameters. We employed the adam optimizer from Optax (5x5 Gridworld & Spontaneous Behavior) or the BFGS optimizer from jax. sci icipp py (Labbrinth) for learning the reward function r:z, using a learning rate of x 10-3 and momentum coefficients of (0.9, 0.999) (adan) or tolerance of 1 10 (BFGS), We backpropagated the gradient for the reward function from the policy term in loss function through 100 Soft-Q iterations. For optimizing the hidden mode transition Pa (Z++1 2t, 8t), we used the LBFGS optimizer from JAXopt: with the tolerance set to 1 x 10-1
B.2. Empirical runtime
Our SWIRL implementation leverages the advantages of JAX, including just-in-time (JIT) compilation and vectorization, to achieve high computational efficiency. For S-2 experiments across all three datasets (gridworld, labyrinth, and spontaneous behavior). SWIRL converges within 15-30 minutes on a V V100 GPU, which takes 50-100 EM iterations. For longer L, S-4 experiment on labyrinth with 50 EM iterations take 2-3 hours to finish on a L40S GPU.
B.3. Evaluation metrics
The test log-likelihood is computed from the forward variable a, where log log $1:T, a1:r) = log Σ. ar,,. Detail of the forward variable a computation is described Appendix A.2. For test hidden mode segmentation accuracy, we use the Viterbi algorithm to predict the hidden mode of each timepoint and compute the percentage of correct hidden mode prediction compared ground truth data. For reward correlation, compute the Pearson correlation between learnt reward and ground truth reward or dopamine fluctuation levels.
B.4. Environment setup
5x5 Gridworld:
State space: Discrete(25); Action space: Discrete(S). The 5 actions are 'up', 'left', 'down', 'right', "stay".
Labyrinth: State space: Discrete(1 (127); Action space: Discrete(4). The labyrinth has binary-tree structure and the 4 actions are 'move in left', 'move in right', 'move out left', 'move out right'.
Spontaneous Behavior: State space: Discrete(9); Action space: Discrete(9). Each state a MoSeq syllable that defines a behavior motif. Each action represents a next state that the agent transfers into.
B.5. Discussion on discount factor and temperature a
We set the discount factor = 0.95, a standard choice in RL and IRL literature. For the mouse spontaneous behavior dataset, we also tested smaller = 0.7, as previous literature Markowitz et al. (2023) suggested this value as optimal for the dataset. However, we observed that the results learned by SWIRL were highly similar for both discount factors, indicating that the choice of y had minimal impact on performance in this case.
We searched for the optimal temperature in {0.01, 0.1, 0.5, 1). For the labyrinth dataset, smaller values of led to better results for certain hidden modes. This observation aligns with the deterministic nature of behaviors in the labyrinth's 's tree-like structure. On the contrary, for the spontaneous behavior dataset, where animals exhibit more stochastic behavior patterns, we found higher values a were more appropriate.
B.6. Discussion on the number of hidden modes and history length
In this section, we discuss the selection and impact of the number of hidden modes Z = |Z| and action-level history length L.
B. Z IN LABYRINTH EXPERIMENT
We evaluated the test LL of SWIRL models on Z from 2 to 5 and found that the best model (S-2) plateaus beyond 4 (Fig. 5A). However, 7 = 4 result does not differ much from the Z = 3 result: Z 4 result mainly segments the explore
15

---

Inverse Reinforcement Learning with Switching Rewards and History Dependency for Characterizing Animal Behaviors
A
-070
B
3 hidden modes
-0.3
time
500
water
home
expllor
port visit
home
Teat
-0.8
-000
4 hidden modes
time
500
-1.00
4
water
home
explore1
explore2
water port visit
home
number of hidden modes
C
water
home
explore1
explore2
Figure 5. Water-restricted labyrinth experiment with different number of hidden modes Z. (A) Held-out test LL for each model across different number of hidden modes over 10 runs. The shaded area represents the total arca that falls between one standard deviation above and below the mean. Segments of trajectory from held-out test data, predicted by SWIRL (S-2) with Z and Z 4. The orange dot indicates when the mouse visits the water port, while the red cross denotes the mouse's 's visit to state 0 (home) at that time. (C) Inferred reward maps from SWIRL (S-2) with r( 8t-1) was averaged over - to produce r(s)).
water, home, and two explore maps. To enhance visualization, the inferred reward
mode of 7 = 3 into two explore modes with similar reward maps (Fig. SBC). As a result, still present 7 =3 3 as the primary result for simplicity.
B.6.2. L IN LABYRINTH EXPERIMENT
With Z = 3, we evaluated the test LL of SWIRL models on L from to 4 and found that the L = 4 (S-4) provides the best test LL (Fig. 6A). L = 3 and L = 4 provide similar hidden segments and reward maps (when averaged over ( to produce r(s:)) as = (Fig. 6BC). In the main paper, we present L = 2 (S-2) as the primary result as has effectively demonstrated the benefits of incorporating mon-Markovian action-level history dependency into SWIRL. However, we note that the test LL results in Fig. 6A suggest the presence of longer action-level history dependency (L 2) in this labyrinth dataset. This observation aligns with the partially observable nature of this 127-node labyrinth: The mouse may not know the whole environment, so it likely relies on longer history of prior states to guide its decision-making.
B.6.3. Z AND L IN MOUSE SPONTANEOUS BEHAVIOR EXPERIMENT
In the mouse spontaneous behavior experiment, we find that L = 1 (S-1) is the optimal choice, as L 1 (S-1) consistently provides higher test log-likelihood (L.L) compared to (S-2) (Fig. 4B). Additionally, we select = 5 for the number of hidden modes since the test LL plateaus Z 5 (Fig. 4B).
B.7. Robustness of SWIRL
To assess the robustness of SWIRL, we evaluated the performance of SWIRL (S-2) under increasing levels of random perturbations in the simulated gridworld dataset.
Specifically, we introduced random permutations to percentage of the states and actions in the training data, ranging from 0% to 50%. As expected, performance decreased as the level of permutation increased (Fig. 7). The model maintained high accuracy with less than 10% permutation. Between 10% and 30%, SWIRI demonstrated stable performance, achieving reasonable reward correlations and hidden mode segmentation accuracy despite the noise. Permutation beyond 30% led to very noisy data and it became hard for the model to maintain high performance.
These results suggest that SWIRI. can tolerate moderate levels of noise or incomplete data, making it suitable for real-world
16

---

Inverse Reinforcement Learning with Switching Rewards and History Dependency for Characterizing Animal Behaviors
A
B
S-1
-0.00
S-2
500
-0.70
Test
S-3
500
-072
0074
S-4
-0.76
time
500
S-1
S-2
S-3
S-4
water
home
explore
water port visit
home visit
C
water
home
explore
RO
Figure 6. Water-restricted labyrinth experiment with different action-level history length L. (A) Boxplot showing held-out LL over 10 runs, with the x-axis representing the four different models from L = to L = 4. Outlier selection method described in Appendix B.8. (B) Segments a trajectory from held-out data, predicted by the four SWIRL models. The orange dot indicates when the mouse visits the water port, while the red cross denotes the mouse's 's visit to state 0 (home) at that time. (C) Inferred reward maps from SWIRL (S-4): water, home, and explore. enhance visualization, inferred reward r(st, St-1; $4-2, St-3) was averaged over (11-1, St-3) to produce r(s.).
animal behavior datasets where such challenges common.
46
0.8
Test segment accuracy
0.8
correlation
0.6
Test L1
-150
0.5-
0.7-
0.4
0.3
9.6
03
-156
0.5
10 15
20
25
30
35
40
45 50
10 15 20
25
30
35 40
45
50
10
15
20
25
30
35
45 50
permutation %
permutation %
permulation %
Figure SWIRL (S-2) experiment correlation between the true and recovered reward maps, test log ligeliiklihodd, and test segmentation accuracy. The axis represents the percentage of states and actions permuted the training data. Outlier selection method described in Appendix B.8.
gridworld dataset with ten random permutations. Box plots illustrating Pearson
B.8. Outlier selection in box plot
All box plots in this paper are drawn by seaborn. www.rnt.ooooo...oooooomoo () with its default outlier selection method. Specifically, the upper quartile (Q3), lower quartile (Q1), and interquartile range (IQR) calculated. Values greater than 13+1.51QR less than Q1-1.51QR are considered as outliers.
17

---

Inverse Reinforcement Learning with Switching Rewards and History Dependency for Characterizing Animal Behaviors
C. Appendix C
C.1. Model-free SWIRL on a simulated gridworld environment
true
discovered
home reward
St
St-1
current
right
current
above
true
discovered
1
current
left
water reward
current
below
current
current
Figure 8. True and discovered reward maps in model-free SWIRL (S-2) simulation experiment. The color scale represents reward values ranging from 0 to 1. The home reward is defined as r(s:). while the water reward depends on both the current and previous locations, r(tt, 9--1)- To present the water reward, each location divided into five groups, detailed in the table on the far right. For example, the polygon in the first row represents the reward value when is the current location and 8t-1 is the location to the right. Light grey indicates an impossible transition where no reward exists.
To demonstrate that SWIRL also supports model-free IRI., we replaced the model-based soft-Q iteration IRL. approach used in main experiments with a model-free IRL method IQ-Learn (Garg et al., 2021). Garg et al. (2021) shows that we can convert the MaxEnt IRL objective L(TTrr) r) with a convex reward regularizer to an offline inverse Q-learning objective:
max J(Q) = (0
<V())) - Bic.o, (V ~V
with (s) = log E. exp (s, a). & is the concave function from convex conjugate. When using x >iddivergence, 6(2) = r g² The reward can be obtained as r(s, a) Q(s, a) "Bss E.e exp Q (s', a')). By marginalizing actions, one similarly obtain state-only reward r(s).
Since optimizing the reward by maximizing log corresponds to the Lagrangian dual of the standard MaxEnt IRL approach (Zeng et al., 2022), we can maximize the IQ-Learn objective (Q) instead of the log ग term in SWIRL EM auxiliary function (Bq. 3). Therefore, for a trajectory č, SWIRL can perform IRL with the following objective:
Σ p(z 0) (lo(Q ( -
max
with Vs, (s') log E. exp Q= (...aa)))))
We tested this IQ-Learn-based model-free SWIRL (S-2) simulated trajectories in the same 5x5 gridworld environment Sec. 4.1. The recovered reward maps achieved a Pearson correlation of 0.737 relative the true reward maps. Although the model-free SWIRL recovered reward maps (Fig. 8) appear noisier than those from the model-based approach in Sec. 4. (Fig, 2A), they still accurately capture the critical high-reward regions.
18


Reviews:
[{"review_id": "75ZNzYaDVR", "reviewer": "Reviewer_9ThW", "summary": "This paper presents an EM-based IRL algorithm SWIRL (SWItching IRL), for learning time-varying reward functions to model animal behavior.  The paper extends IRL by incorporating time-varying, history-dependent reward functions. \n\nA key contribution of this work is that it incorporates capturing the shifting motivations and history-dependent decision-making observed in animals by modelling long behavioural sequences as transitions between short-term decision-making processes, each governed by a unique reward function. \n\nThey do this by considering history dependency at both the decision level (transitions between decision-making processes depend on previous decisions and environmental context) and the action level (actions depend on the history of states within a decision-making process). Paper is well written, coherent, and technically sound.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer 9ThW for their detailed and thoughtful comments on our paper. We are pleased that the reviewer recognized the novelty of our work, the technical soundness, its potential impact on intelligent behavior research, and the overall paper presentation. Below, we address each of the questions and concerns:\n\n**1. Scalability and Generalization to Larger State Space**\n\n&nbsp;&nbsp;**1.1** Our SWIRL framework is compatible with larger state-action space. As discussed in Sec. 3.4, SWIRL inference procedure can be done in a scalable way (model-free IRL) instead of the model-based approach shown in the main body of the paper. It is worth noting that IRL in large and/or continuous state-action spaces is a very challenging problem. Even state-of-the-art standard IRL methods, when assuming a single static reward function, still cannot accurately recover the ground-truth reward. Despite limitations in recovering the true reward function, IRL methods can still produce policies that approximate the expert’s behavior. As a result, recent IRL literature [1][2] in such settings usually focuses on policy recovery performance rather than reward function.\n\n&nbsp;&nbsp;**1.2** In Appendix D.1, we present a detailed discussion of a scalable model-free variant of SWIRL and evaluate its performance on the same gridworld experiment as described in Section 4.1 of the main text. Despite being model-free, this scalable variant of SWIRL is still capable of achieving reasonable reward recovery. We expect that in environments with large state-action space, SWIRL will perform worse in reward recovery but still maintain reasonable policy recovery.\n\n&nbsp;&nbsp;**1.3** SWIRL has the potential to be further extended to continuous state-action space. We can change the DQN-style inverse-Q-learning approach in Appendix D.1 to SAC-style ([1] already shows SAC is compatible with inverse-Q-learning), building a SWIRL variant compatible with continuous state-action space. Further extension to RNN/transformer structure could improve SWIRL’s performance on data with longer history-dependency.\n\n**2. Selection of number of hidden modes**\n\nYes, it is a hyperparameter. We discussed the selection of the number of hidden modes in Appendix B.4.1. and Appendix B.4.3. It can be meta-learned based on the trend of test LL. In general, we recommend selecting the number of hidden modes at the point where the test log-likelihood curve plateaus.\n\n**3. Selection of L for labyrinth experiment**\n\nWe discussed longer L experiment results (S-3 & S-4) in Appendix B.4.2. Generally, L should be selected based on the trend of the test LL curve as well as the recovered hidden mode segments and reward maps. In this case since the hidden mode segments and reward maps remain similar with L >= 2, we choose to show S-2 as the main experiment.\n\n**4. Consideration of longer L**\n\n&nbsp;&nbsp;**4.1** For model-based SWIRL, the feasibility of longer L depends on the environment size. In the labyrinth environment (127 states, 4 actions), model-based SWIRL can perform reasonably fast (within 3 hours) on a L40S GPU with L up to 5. For a smaller environment such as the one in spontaneous behavior experiment (9 states, 9 actions), it is easy to test long L (e.g. L=10) with our model-based SWIRL. \n\n&nbsp;&nbsp;**4.2** For environments with longer L, we can use the scalable model-free SWIRL proposed in Appendix D.1, but at the cost of inaccurate reward recovery. On the labyrinth dataset, we find it is easy to test longer L (e.g. L=20) with the model-free SWIRL. However, mice behavior in this dataset does not show that long action-level history dependency and such long L actually leads to lower test LL. \n\nWe thank again for the reviewer’s effort and humbly request that the reviewer consider raising their score if the above reply adequately addresses their concerns.\n\n**Reference:**\n[1] Garg et al. \"Iq-learn: Inverse soft-q learning for imitation.\" Advances in Neural Information Processing Systems, 34 (2021). [2] Zeng et al. “Maximum-Likelihood Inverse Reinforcement Learning with Finite-Time Guarantees”. Advances in Neural Information Processing Systems, 35 (2022).", "claims_and_evidence": "The claims are well supported through multiple experiments, and various baselines. \n\nThe claims on originality and novelty is well supported ie. first IRL model to integrate both decision-level and action-level history dependency. Claims on principled method with empirical validation is also very well supported i.e. a mathematical formulation of SWIRL, including detailed explanations of how history dependency is incorporated at different levels, and a clear demonstration of improvements over baseline methods. In addition the paper is well written and presented. Claim on the relevance and importance of the work lie in the modelling of animal decision-making, which has the potential to further our understanding of intelligent behavior.", "methods_and_evaluation_criteria": "To test the efficacy of the proposed method, SWIRL is empirically tested on simulated data and real-world animal behaviour datasets. \n\nMultiple baselines have been compared such as latest concurrent work on Locally Consistent IRL (Nguyen et al., 2015),  ARHMM (Wiltschko et al., 2015), rARHMM (Linderman et al., 2016) I-1, I-2, S-1, S-2. The evaluation is competitive, and even though the approach compares similarly, it does consistently well and even favorably when compared to Nguyen et al., 2015 which is very positive.\n\nThe paper further provides connections between SWIRL and autoregressive dynamics models, arguing that SWIRL offers a more generalized and principled approach to characterizing animal behaviour. Results presented demonstrated that it outperforms existing models lacking history dependency, both quantitatively and qualitatively.", "theoretical_claims": "One of the major concerns of the method proposed is difficulty in scaling to larger state space? Would these results hold there? How might the authors propose to tackle this? Would it be possible to outline how the method generalizes to larger state space? \n\nIt seems that longer L in SWIRL does require a considerable amount of computing resources, how does the method work in this case? Have you considered environments where longer L was required? Was the method still competent compared to baselines? It would be useful to add discussion around this. Additionally, if you can point me to experiments that would be useful as well that consider longer L.", "experimental_designs_or_analyses": "Experiments are clear and results are reported with rigor. Complexity and Scalability of the method is unclear given the grid world experiments are not to scale with the dimensions considered. \n\nIt is not clear to use this method, how might one go about choosing the number of hidden modes? Was this a hyperparameter? Is there a way to meta learn this variable?", "relation_to_broader_scientific_literature": "The work is well positioned in the literature. Multiple baselines have been compared such as latest concurrent work on Locally Consistent IRL (Nguyen et al., 2015), ARHMM (Wiltschko et al., 2015), rARHMM (Linderman et al., 2016) I-1, I-2, S-1, S-2. The evaluation is competitive, and even though the approach compares similarly, it does consistently well and even favorably when compared to Nguyen et al., 2015 which is very positive.\n\nThe paper further provides connections between SWIRL and autoregressive dynamics models, arguing that SWIRL offers a more generalized and principled approach to characterizing animal behaviour. Results presented demonstrated that it outperforms existing models lacking history dependency, both quantitatively and qualitatively.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "Since a major contribution here is considering the Action-level history, for  experiments that only consider (S-2) even though it is better than  only current length of 1, it seems its too short a length to showcase the utility of history. Is this a result of the environments? How is this length determined or recommended by the method to be considered?", "other_comments_or_suggestions": "See above.", "questions_for_authors": "See above.", "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": "", "extra_scores": {}}, {"review_id": "IcEHSwDCDk", "reviewer": "Reviewer_FWjW", "summary": "In this work, the Authors extend the IRL framework, designed previously to consider multiple goal maps in real-world agents concurrently (Ashwood et al & Zhu et al), by explicitly incorporating history-dependent policies and rewards into the model. Using this new framework, the Authors model several standard datasets in the field (Rosenberg et al & Markowitz et al) where they show overall improvement in performance. The relation of the proposed method to other models, viewed as its edge cases, is discussed.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We greatly appreciate the time and effort reviewer FWjW dedicated to analyzing our work. We are sorry that the key contributions of SWIRL are likely misunderstood by the reviewer and apologize for not clarifying this fact more clearly in the paper. Below we will provide a thorough discussion on the relationship between Ashwood et al., Zhu et al. and SWIRL and clarify the novel contribution of SWIRL:\n\n**1. Comparison among Ashwood et al., Zhu et al. and SWIRL**\n\n&nbsp;&nbsp;**1.1** **Ashwood et al.** proposed the reward function as a dynamical combination of a number of reward maps. The inference method proposed in Ashwood et al. assumes the reward function changing dynamics is the **same for every trajectory**, which greatly limits its applicability. As discussed in Sec. 4.2 of our paper, Ashwood et al. used the **short, pre-clustered** and **stereotyped** labyrinth trajectories **trialized** and **sampled** from the original long trajectories. In their dataset, for water-restricted mice (with a water reward in labyrinth), each trajectory is **very short (20 steps)** and has **very similar reward changing dynamics and state visitation sequence** (always start from ‘home’ states e.g. 1,0,  go to water port, then go back to ‘home’). They then conduct a **separate** experiment on another dataset of water-unrestricted mice (with no water reward in labyrinth) to find ‘explore’ map, where the trajectories are also **very short, pre-clustered** and **stereotyped**. \n\n&nbsp;&nbsp;**1.2** In **Zhu et al.**, the agent is switching between a number of reward functions, and this model can handle trials with different reward switching dynamics. However, Zhu et al. only tested it on **the same labyrinth dataset as Ashwood et al.** and did not really show this ability. Furthermore, in Fig. 3e of our paper, I-1 (Zhu et al.) has a lot of unreasonable fast-switchings between hidden modes, which shows that the method of Zhu et al. **cannot** accurately recover the reward switching dynamics in the **long, non-stereotyped raw** behavior data. Also, Zhu et al. **cannot** recover the reward with **action-level history dependency** (e.g. Fig. 3c).\n\n&nbsp;&nbsp;**1.3** The introduction of **state-dependent decision mode transition** and **action-level history dependency** in **SWIRL** is critical and enables the effective reward learning on **long, non-stereotyped naturalistic trajectory** of the labyrinth dataset. Compared to the datasets used by Ashwood et al. and Zhu et al., in our experiment we simply segment the original data of water-restricted mice into trajectories of **500 time steps** each, **without any clustering**. Our trajectories are **much longer** and each trajectory maintains the original **different reward switching dynamics** from raw data. Our result shows that only models with both decision-level and action-level history-dependency can recover the switching rewards accurately.\n\n**2.** **Response to Questions For Authors:** SWIRL provides novel insights into real animal datasets\n\n&nbsp;&nbsp;**2.1** For labyrinth, as described above, SWIRL shows the switching between ‘home’, ‘water’ and ‘explore’ reward functions in the **long, naturalistic** water-restricted mice trajectories, which is a **new** result that has not been reported before. Rosenberg et al. does not have RL/IRL analysis and Zhu et al. & Ashwood et al. only show the switching between ‘home’ and ‘water’ in the **short, pre-clustered, handcrafted** version of water-restricted mice trajectories. Our result is clearly **much more naturalistic and reliable** than previous literature and supports the following important claim: Mice make decisions based on both current state and their history.\n\n&nbsp;&nbsp;**2.2** For spontaneous behavior, Markowitz et al. shows DLS dopamine activity correlates with reward through RL experiment. We verify the claim with IRL experiment and provide novel insights that there exist switching hidden decision modes with varying dopamine correlations. To the best of our knowledge, we are **the first** to perform switching reward analysis on this type of dataset, and our finding of switching hidden decision modes in those spontaneous behavior syllables is **new**.\n\n**3.** In addition, we are also the first to incorporate this switching reward idea with model-free inverse-Q-learning-based IRL method, enabling scalable application, which is an important contribution to the ML community. We discuss model-free SWIRL implementation in detail and exhibit its reasonable performance on the gridworld experiment (same env as Sec. 4.1) in Appendix D.1. \n\nWe thank again for the reviewer’s effort and humbly request that the reviewer consider raising their score if the above reply adequately addresses their concerns regarding the scope of the novelty.", "claims_and_evidence": "The main claims that (1) the consideration of history-dependent policies and rewards in this specific IRL framework is novel and (2) that the consideration of the history-dependency improves the model’s predictive power on the (Rosenberg et al) dataset are correct (I elaborate in the sections below). While history-dependency has indeed been considered in other prior work, including the models of neuroscience-relevant data (e.g. by “hacking” the definition of the Markovian state through defining it as a sequence of the past states), the addition of such historic dependency to the IRL framework is novel.", "methods_and_evaluation_criteria": "The methods and evaluation criteria here follow the well-established pattern used previously in prior works by Ashwood et al & Zhu et al. Specifically, the work here uses the datasets by Rosenberg et al & Markowitz et al, which have been shown to be adequate for this task and have been overall widely used in the field. As such, the methods and criteria here do make sense for the task at hand.", "theoretical_claims": "The methods used here make perfect sense for the task. The use of the E-M (Forward-Backward) algorithm has been previously argued to be the go-to tool for the POMDP-like problems, such as the problem here. I’ve scanned through the derivations in the Appendix; they looked overall correct.", "experimental_designs_or_analyses": "The experimental design here follows the standard practice in the field. Apart from the use of standard datasets, the work provides a reasonable set of baseline models and ablations.", "relation_to_broader_scientific_literature": "While the Authors correctly cite prior work (including works by Ashwood et al and Zhu et al) and correctly discuss the novelty of this new work in relation to prior work (the consideration of the history-dependency), it is specifically this part that is the most concerning for me. The addition of the history dependence is, arguably, a minor tweak upon Zhu et al’s work. Zhu et al’s work, in turn, was a minor (though a much bigger) tweak upon Ashwood et al’s work and, as such, to the best of my knowledge, hasn’t made it to one of these conferences. While adding the historical dependency is an important step in the right direction, this paper would’ve been much stronger if it had provided novel insights into the existing data, however, the qualitative conclusions for the Rosenberg et al & Markowitz et al data match those of the previous work. Thus, while the work appears completely correct to me, the scope of the novelty in the current format sadly precludes me from recommending the acceptance of this paper at this point.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "N/A", "other_comments_or_suggestions": "N/A", "questions_for_authors": "What are the novel, unique insights that can be obtained from the considered and/or other data using the proposed algorithm but not the previous algorithms?\n\n______________________\n\nPost-rebuttal:\n\nUpon the discussion with the Authors, I agree that 1) considering the history dependency is important and 2) that not having pre-clustered data has the potential to make a difference. Thus, I agree that the score can be raised and the work qualifies for a weak accept.", "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": "", "extra_scores": {}}, {"review_id": "EFWV6Dh7ay", "reviewer": "Reviewer_CRRd", "summary": "This paper addresses the limitation of traditional IRL, which assumes rewards depend only on the current state, making it insufficient for modeling long-term, history-dependent decision-making in animals. To capture this dependency, the paper introduces SWIRL, an IRL framework that models behavior as transitions between short-term decision-making processes, each governed by a distinct reward function. SWIRL infers these processes and their transitions using history-dependent reward modeling. The authors conducted experiments on simulated and real-world animal behavior datasets.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate the time and effort reviewer CRRd dedicated to analyzing our work. Below, we address each of the questions and concerns:\n\n**1. Hierarchical RL, RL+RNN and POMDPs:** We will refer to the works listed in the “Essential References Not Discussed” section using bracketed citations in this response.\n\n&nbsp;&nbsp;**1.1** [1][3][4][5][6][7]: We thank the reviewer for pointing out the RL methods with history dependency or multiple reward functions. However, it is worth noting that our method is solving the Inverse RL problem (given expert demonstrations, trying to recover the expert’s policy and the reward function) instead of the RL problem (given reward function, trying to find the optimal policy). IRL is a much harder problem and the goal is very different from RL. As a result, it is not feasible to apply those RL methods as baselines for our method.\n\n&nbsp;&nbsp;**1.2** [2]: it is worth noting that [2] infers reward functions from expert demonstrations while assuming the Observation function P(z|s, a) is already known—a condition that does not hold in our case. To the best of our knowledge, there is no existing Inverse POMDP method that can effectively address our problem. We recognize the need of discussing related POMDP literature and will add it in the next revision.\n\n**2. Scalability and longer history:** \nPlease refer to section 1 and section 4 of our reply to Reviewer 9ThW. We appreciate the idea of using RNNs or Transformers to learn sequential dependencies more effectively and will add it in the next revision as a future direction. \n\n**3. Hidden mode with infrequent occurrence:** \nIn cases where a hidden mode occurs infrequently, regularization or constraints can be incorporated to improve the reward recovery. For instance, if the reward is expected to be sparse, L1 regularization can be used; if large reward coefficients are undesirable, L2 regularization can help. When the occurrence of a mode is too low and prior knowledge is too limited, a common practice is to reduce the number of hidden modes, effectively merging the rare mode with a more frequent mode with similar reward function.\n\n**4. Section heading issue:** We apologize for the heading space issue and will correct it in our next revision.\n\n**5. Full name of IRL appears multiple times:** We will use the abbreviation consistently in our next revision. \n\n**6. ARHMM abbreviation issue:** ARHMM should be the abbreviation of the autoregressive hidden Markov model. We will address this issue in our next revision.\n\n**7. Response to Questions For Authors:** \n\nThank you for the thoughtful comment. Upon reflection, we agree that the original statement may have been imprecise. While SWIRL is not a POMDP method in the formal sense (e.g., modeling belief states or latent dynamics explicitly), it does address non-Markovian behavior by augmenting the state space with recent history (e.g., using {s_t, s_{t-1}, s_{t-2}, …}). This is a commonly used strategy to handle partial observability or non-Markovianity—similar in spirit to how LSTMs or memory-based policies are used in POMDP settings. So while SWIRL is not derived from a POMDP framework per se, it shares key characteristics with POMDP-appropriate methods in its use of temporal context. When we wrote that it “can be extended to POMDPs,” we were referring to the potential for a more principled approach—e.g., explicitly modeling latent states or belief updates—to further generalize the method. We will revise the discussion to clarify this distinction and avoid the implication that our current method is directly extendable in a formal POMDP sense.\n\nWe thank again for the reviewer’s effort and humbly request that the reviewer consider raising their score if the above reply adequately addresses their concerns.", "claims_and_evidence": "The paper makes well-motivated claims about the necessity of incorporating history-dependent rewards in IRL to model long-term animal decision-making. \nHowever, the paper does not provide a detailed analysis of how SWIRL scales to larger environments or higher-dimensional state spaces, which is crucial given its reliance on history-dependent modeling. Empirical results on computational efficiency would strengthen this claim.\nOverall, the paper provides substantial evidence to support its core claims but could improve in these areas to further solidify its contributions.", "methods_and_evaluation_criteria": "1. The proposed method increases the input state size by incorporating multiple past states, resulting in larger network inputs, higher computational complexity, and slower training and inference.\n2. Hierarchical RL decomposes complex decision-making into multiple levels of abstraction. A high-level policy selects sub-goals or temporally extended actions, while a low-level policy executes fine-grained actions to achieve them. In SWIRL, the hidden modes resemble the sub-task structure in HRL, as they dictate different behavioral phases. However, unlike HRL, where the hierarchy is explicitly defined, SWIRL infers latent behavioral modes from data. A direct comparison with HRL-based methods—particularly goal-conditioned RL or options-based methods—could clarify whether SWIRL provides advantages in learning structured behaviors without predefined hierarchical policies.\n3. POMDPs explicitly model uncertainty by maintaining a belief state over hidden variables. POMDP-based approaches have been widely used to model decision-making under partial observability, where the agent must infer unobserved environmental factors. SWIRL introduces hidden modes, which serve as an internal latent state governing behavior. Given this similarity, a comparison with POMDP-based RL methods could provide insights into whether belief-state tracking could enhance SWIRL’s performance.\n4. The experiments only evaluate small history lengths, which may not be sufficient for capturing long-term dependencies in sequential decision-making. Given that the paper aims to model history-dependent decision-making, it could include a discussion on the feasibility of using alternative approaches, such as RNNs or Transformers, to learn sequential dependencies more effectively.\n5. Since the reward network must learn multiple distinct reward functions corresponding to different hidden modes, there is a potential issue when some hidden modes appear infrequently. This could lead to underfitting in the reward networks for rarely occurring modes. Is there any discussion on how to mitigate this issue, such as through data augmentation, regularization, or constraints on reward function learning?", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "1. Since incorporating past states increases input size and computational complexity, it would be valuable to assess the efficiency of SWIRL in larger environments. A computational runtime analysis or an empirical study on training efficiency would provide useful insights.\n2. The paper primarily compares SWIRL against history-agnostic IRL models. While this highlights the benefit of incorporating history dependency, additional baselines, such as RNN-based RL methods [1, 2] and Transformer-based RL models [3, 4], could provide a stronger contextual comparison in handling long-term dependencies and sequential decision-making.\n3. The reward network must learn multiple distinct reward functions for different hidden modes. However, if certain modes appear infrequently, they may not receive sufficient training, leading to potential underfitting. A sensitivity analysis on the frequency of hidden modes and its impact on reward learning would help assess robustness.\n\n[1] Memory-based Control with Recurrent Neural Network.  \n[2] Deep Recurrent Q-Learning for Partially Observable MDPs.  \n[3] Decision Transformer: Reinforcement Learning via Sequence Modeling.  \n[4] Offline Reinforcement Learning as One Big Sequence Modeling Problem.", "relation_to_broader_scientific_literature": "1. The relationship between SWIRL and broader scientific literature, particularly Hierarchical RL and POMDPs, could be further clarified, as SWIRL’s hidden modes resemble the sub-task structures in HRL and the latent states in POMDPs, making a direct comparison with these methods valuable for distinguishing its contributions.\n2. This paper models history dependency by explicitly incorporating past states into the input, increasing the input size and computational cost. However, alternative sequence modeling methods, such as RNNs and Transformers, have been extensively studied for capturing long-term dependencies in reinforcement learning. Discussing whether such architectures could complement or improve SWIRL’s performance would provide valuable context and better position the work within the broader literature on history-aware decision-making.", "essential_references_not_discussed": "POMDPS:  \n[1] Learning Predictive State Representations.\n[2] Inverse Reinforcement Learning in Partially Observable Environments.\n\nHierarchical RL:  \n[3] Data-Efficient Hierarchical Reinforcement Learning.\n\nRL + RNN:  \n[4] Memory-based Control with Recurrent Neural Network.  \n[5] Deep Recurrent Q-Learning for Partially Observable MDPs.  \n\nRL + Transformers:  \n[6] Decision Transformer: Reinforcement Learning via Sequence Modeling.  \n[7] Offline Reinforcement Learning as One Big Sequence Modeling Problem.", "other_strengths_and_weaknesses": "See comments.", "other_comments_or_suggestions": "1. According to the submission guidelines: **Section headings should be numbered, flush left, and set in 11 pt bold type with the content words capitalized. Leave 0.25 inches of space before the heading and 0.15 inches after the heading. Similarly, subsection headings should be numbered, flush left, and set in 10 pt bold type with the content words capitalized. Leave 0.2 inches of space before the heading and 0.13 inches afterward.**  \nHowever, it seems that the author has modified these formatting rules.\n2. The full name of IRL appears multiple times throughout the paper, including in Section 1 (Introduction), Section 3.2, and Discussion. It is recommended to define it only once at the first occurrence and use the abbreviation consistently thereafter.\n3. In line 122, the term \"autoregressive process (ARHMM)\" appears, while in line 189, it is written as \"autoregressive hidden Markov model (ARHMM).\" Please clarify which is the correct full name for ARHMM to maintain consistency.", "questions_for_authors": "Although the discussion section mentions POMDPs, I believe the statement that \"SWIRL can be extended to POMDPs\" is not entirely accurate. POMDPs represent a parallel approach rather than a direct extension of SWIRL. It would be more appropriate to compare SWIRL directly with existing POMDP-based methods rather than framing it as a future extension. Could the authors clarify their perspective on this distinction?", "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": "", "extra_scores": {}}, {"review_id": "uYsGpfgnQ1", "reviewer": "Reviewer_5en4", "summary": "The paper introduces Switching Inverse RL (SWIRL), an inverse reinforcement learning framework for characterizing animal behaviors. In this problem setting, the goal is to infer reward functions and policies from animal behavior trajectories. To achieve this, SWIRL introduces two main design choices: time-varying rewards controlled by latent modes and biologically plausible history dependency. The time-varying rewards allow the model to capture switching behaviors. The history dependency occurs at two levels: the decision level (transitioning from one hidden mode to another conditioned on the previous state) and the action level (history-conditioned policies and rewards). The framework is optimized with an expectation-maximization procedure to recover the parameters for the hidden transition kernel, the policy, and the reward function. The paper compares SWIRL to prior IRL baselines and ablations of SWIRL on three sets of experiments: a 5x5 gridworld environment, a labyrinth dataset where water-deprived mice move freely based on varying internal goals, and a dataset where mice wandered an empty arena without explicit rewards. Across these experiments, they find SWIRL to outperform baselines in reward correlation, log-likelihood, and segment accuracy.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We greatly appreciate the time and effort reviewer 5en4 dedicated to analyzing our work and providing such constructive feedback. We are pleased that the reviewer recognized the novelty of our work, the technical soundness, its implications for both neuroscience and ML community and the experiment design & presentation. Below, we address each of the questions and concerns:\n\n**1. Scalability concern:** Please refer to section 1 of our reply to Reviewer 9ThW. \n\n**2. Evaluation metrics:** We will add the detailed description of evaluation metrics in our next revision. Test log-likelihood is computed from the forward variable, where $\\log p(s_{1:T}, a_{1:T}) = \\log \\sum_z \\alpha_{T, z}$, and the computation of forward variable $\\alpha$ is described in Appendix A.2. For test hidden mode segmentation accuracy, we use the Viterbi algorithm to predict the hidden mode of each timepoint and compute the percentage of correct hidden mode prediction compared to ground truth data. For reward correlation, we compute the Pearson correlation between learnt reward and ground truth reward.\n\n**3. Experiment setup:** We will add details of each environment in our next revision. For the state and action space:\n\n5x5 Gridworld: State space: Discrete(25); Action space: Discrete(5). The 5 actions are ‘up’, ‘left’, ‘down’, ‘right’, ‘stay’.\n\nLabyrinth: State space: Discrete(127); Action space: Discrete(4). The labyrinth has a binary-tree structure and the 4 actions are ‘move in left’, ‘move in right’, ‘move out from leaf node’, ‘move out from non-leaf node’. \n\nSpontaneous Behavior: State space: Discrete(9); Action space: Discrete(9). Each state is a MoSeq syllable that defines a behavior motif. The 9 actions are just the next state the agent transits into. \n\n**4. How is the reward function optimized?** The soft Q iteration is differentiable and we backpropagate the gradient for the reward function through the Q iteration. We utilized 100 iterations and found the soft Q iteration usually already converged before 100 iterations in our experiments. We will add this discussion in our next revision. \n\n**5. Why do you apply different optimizers to the transition kernel (L-BFGS) and the reward function (Adam)?** This is an empirical choice. A rational reason behind that could be that in the M-step, the transition kernel typically yields a smoother loss surface compared to the reward function. As a result, with second derivatives, second-order optimizers like L-BFGS can take advantage of the smooth curvature, leading to faster and more stable convergence. In contrast, the reward function often has a more complex loss landscape, where Adam tends to perform better due to their robustness and adaptive step sizes. We will add this discussion in our next revision. \n\n**6. Is there a reason why ARHMM and rARHMM are not included in experiments 1 and 2?** ARHMM/rARHMM models the emission probability as $P(s’|s, z)=\\sum_a P(s’|s, a)P(a|s, z)$, where the marginalization over the action space makes the log-likelihood incomparable to models like SWIRL in general case. Experiment 3 (spontaneous behavior) is the only experiment where action is effectively the same as the next state, making it possible to compare ARHMM/rARHMM with SWIRL by test LL.\n\nWe thank again for the reviewer’s effort and humbly request that the reviewer consider raising their score if the above reply adequately addresses their concerns.", "claims_and_evidence": "The central claim of this paper is that incorporating switching rewards and history dependency results in better inverse RL methods for characterizing animal behaviors. This claim is supported by three sets of experiments with tabular state and action spaces:\n- The first is a 5x5 gridworld environment where the underlying reward is accessible. Compared to baselines with constant rewards and less history dependence, SWIRL scores the highest correlation between predicted and ground truth rewards, log-likelihood of the held-out test data under the policy, and accuracy of mode prediction.\n- The second set of experiments is on a water-restricted labyrinth dataset. Since the ground truth rewards and behavior modes are not accessible, the paper only provides a quantitative comparison of the data log-likelihood under the policy, where SWIRL outperforms the baselines. They qualitatively visualize the mode switching in trajectories as well as the reward maps. Both indicate that the method learns a reasonable switching reward.\n- The third set of experiments is on a free-space wandering dataset. The held-out log-likelihoods and reward visualizations support their claim. They further plot the log-likelihood as a function of the number of hidden modes and show that SWIRL performance improves with more hidden modes, whereas removing the state dependence leads to degrading performance as the number of hidden modes increases.", "methods_and_evaluation_criteria": "The paper proposes an inverse RL framework for characterizing animal behaviors. Since animal behaviors are likely controlled by time-varying objectives (e.g. find water when thirsty, otherwise go home), the paper proposes to infer a switching reward conditioned on latent modes. Moreover, animal behaviors are influenced by the history of observations, and thus the paper introduces history dependence to the hidden transition kernel, the reward, and the policy. This framework is optimized by an expectation-maximization procedure, which iteratively performs RL on the current reward estimate, uses the policy to estimate an evidence lower bound on the data likelihood, and updates the parameters via gradient descent. The proposed method is biologically plausible and mathematically sound.\n\nThe evaluation criteria vary based on the problem setting. When the animal behavior data is synthetically generated by following a time-varying reward, then it is feasible to compare the correlation between predicted and ground truth rewards and the accuracy of mode switching. Otherwise, if one only has access to the dataset, then data likelihood under the policy is the only feasible evaluation criteria.", "theoretical_claims": "The paper did not make theoretical claims.", "experimental_designs_or_analyses": "As discussed in the \"Claims And Evidence\" section, the experimental design is sound and supports the claims made in this paper.", "relation_to_broader_scientific_literature": "This work has implications for both the neuroscience community and the machine learning community, offering a useful tool for understanding animal behaviors and a flexible framework for general inverse RL problems.", "essential_references_not_discussed": "The paper provides sufficient context for the reader to understand its contributions.", "other_strengths_and_weaknesses": "**Strengths**\n1. The paper introduces a novel inverse RL framework for characterizing animal behaviors. The proposed method models switching rewards using a hidden mode variable and incorporates history context into the hidden transition kernel, the policy, and the reward in a biologically plausible manner. \n2. The proposed method shows strong empirical results across three sets of experiments.\n3, The visualizations of learned rewards and mode transitions are particularly elucidating. \n\n**Weaknesses**\n1. The experiments in this paper feature tabular states and action spaces. The implications for continuous domains remain unclear.\n2. The proposed method is computationally heavy and requires manual tuning of the number of hidden modes.\n3. The method description is not quite clear. See questions below.", "other_comments_or_suggestions": "1. I suggest adding a more detailed description of evaluation metrics such as log-likelihood and segmentation accuracy. Currently, they are briefly introduced on lines 250-256 and lack clarity.\n2. I suggest adding a more detail description of the evaluation domains. For example, describe the state and action spaces for each environment used for evaluation.", "questions_for_authors": "- How is the reward function optimized? Looking at Algorithm 1 in the appendix and the auxiliary function (Eqn 5, 6, 7), it seems the gradient of the reward function is not directly passed through (because of the soft Q iteration). Do you backpropagate the gradient for the reward function through the Q iteration? If so, how many Q iterations do you use?\n- Why do you apply different optimizers to the transition kernel (L-BFGS) and the reward function (Adam)? Is it an empirical design choice?\n- Is there a reason why ARHMM and rARHMM are not included in experiments 1 and 2?\n\n## Updates After Rebuttal\nThe authors have adequately addressed my questions. I will maintain my score.", "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": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This work proposes SWIRL, an inverse reinforcement learning model with applications to the analysis of animal behavior. While previous work has modeled a single unobserved reward function underlying such behaviors, SWIRL posits both a latent state (e.g., hunger, thirst) with transitions and actions dependent on the previous sequence of states. (Thus, the method is not 1-Markov, as typically assumed.) The authors provide an efficient method for inference in this model in the tabular setting, with potential links to more complex problem setups, as well as a set of experiments on existing behavioral data.

Reviewers agreed that the experiments, collection of comparison models, and quality of the presentation were good. Some reviewers questioned the novelty of the present work relative to some comparison IRL methods, though experiments showed that the addition of a latent switching state improved the interpretability and accuracy of model fits to data.

Overall, a solid paper that contributes to an active area of research in behavioral neuroscience.
Decision: Accept (poster)