Inverse Reinforcement Learning with Switching Rewards and History Dependency for Characterizing Animal Behaviors

5. Math: We appreciate the reviewer’s observation. All suggestions are valid, and we have addressed them accordingly.

6. Minor: We appreciate the reviewer’s observation. We have changed the color of the explore map to improve the visualization.

We 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 and articulates our contribution.

References: [1] Haarnoja et al. "Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor." Proceedings of the 35th International Conference on Machine Learning (2018). [2] Garg et al. "Iq-learn: Inverse soft-q learning for imitation." Advances in Neural Information Processing Systems, 34 (2021). [3] Wulfmeier et al. "Imitating language via scalable inverse reinforcement learning." NeurIPS 2024 (2024).

2. Complexity and scalability of the algorithm:
   2.1 Convergence guarantee: For a 50x50 gridworld, SWIRL has the convergence guarantee. We have added a convergence analysis discussion in Appendix A.2. The convergence of the SWIRL inference procedure can be established based on the convergence guarantee of the EM algorithm, MaxEnt IRL, and Soft-Q iteration.

   2.2 Complexity analysis: We have added a complexity analysis in Appendix A.3. The key bottleneck of SWIRL is in the RL inner loop. Current Soft-Q iteration has a time complexity of $O(ZIS^{L}A^2)$, where $Z$ is the number of hidden modes, $I$ is the number of Soft-Q iterations, $S$ is the number of states, $A$ is the number of actions, and $L$ is the history length. Although the transition model $P(s'|s, a)$ allows us to reduce the $S^L$ space by removing impossible $(s, s')$ pairs, we recognize that longer $L$ in SWIRL does require a considerable amount of computing resources.

   2.3 Empirical runtime: We have added an empirical runtime discussion in Appendix B.2. Empirically, SWIRL inference speed is not slow for typical animal behavior datasets. We implemented SWIRL efficiently with JAX, leveraging the features of vectorization and just-in-time compilation to improve our performance. Every SWIRL (S-2) experiment on all three datasets (simulated gridworld, labyrinth, and dopamine) can finish within 15-30 minutes on a V100 GPU. For longer $L$, an S-4 experiment on the labyrinth with 50 EM iterations takes 2-3 hours to finish on an L40S GPU. We switch to an L40S GPU for $L=4$ due to the V100 GPU's insufficient memory capacity.

   2.4 Integration with function approximators: Yes, it is possible. We have added a scalability and broader impact discussion in Appendix A.4. While scalability is not the main focus of this work, as typical behavior datasets in neuroscience have small or moderate state-action space and current SWIRL implementation can perform efficiently on them, we acknowledge the need for a more general and scalable implementation to address broader applications. In its current form, every step of the SWIRL inference procedure, except for the Soft-Q iteration, is compatible with large or continuous state-action spaces. However, the Soft-Q iteration is limited to discrete state-action spaces and can be slow with a large state-action space. Nevertheless, for applications requiring scalability and compatibility with general state-action spaces, alternative methods can be adapted to replace the Soft Q iteration in the RL inner loop. For instance, Soft Actor-Critic [1]. Furthermore, a promising future direction is to reformulate the standard MaxEnt IRL $r$-$\pi$ bi-level optimization problem in SWIRL as a single-level inverse Q-learning problem, based on the IRL approach known as IQ-Learn [2]. This method has been successfully adapted to large language models training, demonstrating great scalability [3].

3. Lack of transparency in choosing the number of hidden modes across experiments: We have included discussion on the number of hidden modes with additional experiment results in Appendix B.4. For the labyrinth dataset, in Appendix B.4.1, we tested hidden modes from 2 to 5 and found that the test log-likelihood (LL) plateaus at 4 hidden modes. However, after closely examining the hidden mode segmentation and reward map, we found that the 4 hidden modes result does not differ too much from 3 hidden modes: It mainly segments the explore mode into 2 explore modes with similar reward maps. As a result, we still show 3 hidden modes results in the main paper for simplicity. For the mouse spontaneous behavior experiment, in Appendix B.4.3, we show that we pick 5 hidden modes since the test LL plateaus at 5 hidden modes.

4. Spontaneous behavior experiment:
   4.1 We apologize for the confusion. The reward discussion was intended to illustrate how SWIRL can provide valuable insights into dopamine studies, as SWIRL is an application-to-neuroscience study. Following the reviewer’s suggestion, we have moved the reward discussion to Appendix C.2 and added a paragraph in the main paper emphasizing how this experiment demonstrates the value of SWIRL’s action-level history dependency, even in light of the lower performance of S-2.

   4.2 While the non-Markovian action-level history dependency introduced by SWIRL (S-2) does not demonstrate superior performance in this particular experiment, the findings showcase SWIRL’s unique contribution to neuroscience research. Specifically, SWIRL serves as a powerful tool for hypothesis testing in behavioral datasets, enabling researchers to validate or challenge hypotheses regarding decision-level dependency as well as non-Markovian action-level dependency. This versatility further confirms SWIRL’s great potential in advancing our understanding of complex behaviors.

We sincerely thank the reviewer NeTy for their detailed evaluation and for recognizing our work in advancing the application of IRL methods to model complex animal behaviors. We appreciate their constructive feedback and acknowledgment of SWIRL’s motivation and the improvements it offers in understanding animal decision-making processes through history-dependent policies and rewards.

We are sorry that the contributions of SWIRL may not have been fully recognized by the reviewer and apologize for not making SWIRL's contributions clearer in the paper. Below, we will respond to each concern raised in the Weaknesses and Questions and articulate the contributions of SWIRL:

1. Novelty
   1.1 Insufficient comparison with Nguyen et al.'s 2015 work: We apologize for the oversight. Nguyen et al. 2015 does have a very similar graphical model as SWIRL (S-1) and we have modified our paper to discuss that S-1 results can be used for Nguyen et al. 2015 as a baseline method. Despite this similarity, SWIRL still brings unique contributions to the field of computational neuroscience and the broader machine learning community. A key difference between SWIRL and Nguyen et al. 2015 is the action-level dependency we incorporated. In experiments, we show this action-level dependency is an important perspective in understanding animal behaviors. For instance, in the Labyrinth experiment, the hidden mode segments in Fig. 3.F illustrate that S-1 fails to learn meaningful hidden mode segments, while S-2 succeeds. Additional investigations into S-3 ($L=3$) and S-4 ($L=4$) in Appendix B.4.2 show that higher $L$ values improve test log-likelihoods (LL), suggesting that animals in the labyrinth task rely on longer history for decision-making. In the mouse spontaneous behavior experiment, while longer action-level dependency does not yield better results, the findings reinforce our claim that SWIRL serves as a versatile tool for hypothesis testing in behavioral datasets, enabling researchers to validate or challenge hypotheses regarding decision-level dependency as well as non-Markovian action-level dependency. This adaptability underscores SWIRL’s potential to advance our understanding of complex behaviors.

   1.2 We would like to further articulate SWIRL contributions here: As an application-to-neuroscience study submitted to the neuroscience track, SWIRL proposed a framework for IRL with switching reward functions that incorporated biologically feasible decision-level and action-level history dependencies, offering a valuable tool for analyzing animal behavior. SWIRL is the first model to incorporate action-level history dependency, and we have demonstrated its significance in understanding animal behavior, as discussed in the above response 1.1 regarding Nguyen et al.'s 2015 work. Additionally, we are the first to explore the connection between HM-MDPs and other popular dynamics models (e.g., ARHMMs, SLDS) used in animal behavior studies. By bridging these frameworks, we promote a more comprehensive understanding of behavior analysis tools and demonstrate how SWIRL can be adapted to experiments previously analyzed with these models. Moreover, our framework is easily extendable to large and continuous state and action spaces by replacing the Soft-Q iteration in the RL inner loop. This scalability positions SWIRL as not only a valuable tool for computational neuroscience but also a significant contribution to the broader machine learning community. Details on the scalability are described below in response 2.4.

We thank Reviewer NeTy for their valuable time and effort in reviewing our work and responding to our rebuttal. We appreciate the additional suggestions provided and have made corresponding clarifications regarding the concerns mentioned in their reply:

1. Experimental Details
   1.1 We have addressed Reviewer Qvd7's concerns regarding the experimental procedure in detail in our response to Qvd7 (https://openreview.net/forum?id=v7a4KET0Md&noteId=X5B2b73oi6).
   1.2 Number of random seeds: Following the suggestions of all three reviewers, we have conducted experiments for all three datasets using the same number of random seeds and updated the results in the Nov 27 revised version of the paper. The updated results do not impact the conclusion of those experiments.

2. Complexity and Scalability
   2.1 We acknowledge that solving a 50×50 gridworld environment (L=4) with the current implementation is not practical, although theoretically it can converge to at least a local optimum given sufficient time. As discussed in the scalability section of the appendix, larger environments are not the main focus of this study, as this work is centered on applications to neuroscience, particularly animal behavior studies. For scalability improvements, we propose exploring adaptations such as employing Soft Actor-Critic for the inner loops or IQ-Learn to convert IRL problem into a single-level inverse Q-learning problem as future work.
   2.2 We acknowledge the need of providing empirical results on larger gridworld environments to help readers better understand the practical limits of the current implementation. We will conduct additional experiments on larger gridworlds, such as 25×25, to explore and report these empirical limits in the future revision.

3. Minor Revisions
   3.1 We have removed the numbering of equations in Algorithm 1 as suggested by the reviewer.
   3.2 It should be $G(\theta, \theta^k)$. We apologize for the oversight and have corrected it.
   3.3 We have updated section 3.1 and 3.3 to ensure consistency in the $r$ function definition.

We thank again for the Reviewer NeTy's valuable feedback and hope above clarifications address the reviewer’s concerns.

We sincerely thank the reviewer Qvd7 for their detailed comments and analysis of our paper. We greatly appreciate the time and effort they dedicated to analyzing our work and providing constructive feedback. We are pleased that they found our paper to be interesting and well-written and recognized the potential of SWIRL.

Below, we provide detailed responses to each concern raised under Weaknesses and Questions:

1. POMDP: It is indeed possible to formulate our problem within a POMDP framework; however, such a formulation is nontrivial and requires further investigation. One key challenge is that representing our problem as a POMDP introduces significant complexity into the inference procedure, which could be unnecessary for our objectives. The main focus of SWIRL is to perform IRL on animal behavior data with multiple switching reward functions, incorporating both decision-level and action-level dependencies. For this purpose, the HM-MDP framework, which is simpler than POMDP, is well-suited to address our problem. The existing Inverse POMDP literature [1] 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.

2. Hierarchical RL: Hierarchical RL lacks a unified definition, as different Hierarchical RL frameworks introduce hierarchy in different ways. In SWIRL, HM-MDP provides a well-defined framework for our problem and can be viewed as one type of Hierarchical RL. Regarding H-AIRL [2], while it is not directly applicable to our setting, we acknowledge it has state-dependent option $z$ transition structure. However, its graphical model includes a future-option dependency, where the current action depends on the future option. This future-dependency differs from our history-dependency approach and is not biologically plausible in the context of our problem. As reviewer NeTy pointed out, another work [3] studies IRL with state-dependent hidden transitions, which is more aligned with our setup than H-AIRL. The graphical model of [3] can be viewed as S-1, making it a more relevant baseline for our framework. Consequently, we have cited [3] as a baseline work in our paper. Nonetheless, H-AIRL offers an insightful adversarial IRL perspective within the EM framework, which could inspire future development of SWIRL. We have cited [2] and included this discussion in the scalability and broader impact section of our paper.

3. Scalability: We have added a scalability and broader impact discussion in Appendix A.4. While scalability is not the main focus of this work, as typical behavior datasets in neuroscience have small or moderate state-action space and current SWIRL implementation can perform efficiently on them, we acknowledge the need for a more general and scalable implementation to address broader applications. In its current form, every step of the SWIRL inference procedure, except for the Soft-Q iteration, is compatible with large or continuous state-action spaces. However, the Soft-Q iteration is limited to discrete state-action spaces and can be slow with large state-action space. Nevertheless, for applications requiring scalability and compatibility with general state-action spaces, alternative methods can be adapted to replace the Soft Q iteration in the RL inner loop. For instance, Soft Actor-Critic [4]. Furthermore, a promising future direction is to reformulate the standard MaxEnt IRL $r$-$\pi$ bi-level optimization problem in SWIRL as a single-level inverse Q-learning problem, based on the IRL approach known as IQ-Learn [5]. This method has been successfully adapted to large language models training, demonstrating great scalability [6]. We also recognized the adversarial IRL perspective of H-AIRL and proposed it as a possible future direction for SWIRL.

4. Presentation of experiments:

   • 4.1 Label baseline as I-1: We articulated at the beginning of the results section that I-1 and S-1 can be viewed as baseline methods of Multi-intention IRL and locally consistent IRL (the [3] mentioned by reviewer NeTy). The notation I-1, S-1, I-2, and S-2 was chosen to highlight that SWIRL is a more general framework, with Multi-intention IRL and locally consistent IRL as specific cases within it. Researchers can leverage our SWIRL implementation to test various SWIRL models on animal behavior data, enabling the validation or rejection of hypotheses about history dependencies in animal behaviors. We apologize for not making this point clearer in the paper and will consider better ways to present I-1 and S-1 to avoid confusion in future revision.
   • 4.2 How often each experiment is run: We have included implementation details in Appendix B.1. Each experiment presented is the result of 10 runs initialized with different random seeds.
   • 4.3 How boxed plot created: We have included the relevant information in Appendix B.6 and referenced it in the captions of all box plots.
   • 4.4 Shaded areas in Figure 4B: The shaded area represents the total area that falls between one standard deviation above and below the mean. We have added this explanation to the captions of all plots with shaded areas to ensure clarity.
   • 4.5 MaxEnt baseline is missing in Figure 4: MaxEnt test LL is at the bottom right of the plot.

5. Minor points:

   • 5.1 "L453 refers to an appendix which seems to be missing?": We appreciate the reviewer’s observation. We have added the corresponding figure in Appendix C.1 and corrected this reference.
   • 5.2 "L202 $\xi$ is never defined in the paper?": It was defined in Section 3.2 and we apologize for not defining it explicitly in Algorithm 1. We have updated Algorithm 1 to include a clear definition of $\xi$.

Finally, we would like to highlight the key contributions of SWIRL in this work. As an application-to-neuroscience study submitted to the neuroscience track, it proposed the SWIRL framework for IRL with switching reward functions that incorporated biologically feasible decision-level and action-level history dependencies, offering a valuable tool for analyzing animal behavior.

SWIRL is the first model to incorporate action-level history dependency and we have demonstrated its significance in understanding animal behavior. For instance, in the Labyrinth experiment, the hidden mode segments in Fig. 3.F illustrate that S-1 fails to learn meaningful hidden mode segments, while S-2 succeeds. Additional investigations into S-3 ($L=3$) and S-4 ($L=4$) in Appendix B.4.2 show that higher $L$ values improve test log-likelihoods (LL), suggesting that animals in the labyrinth task rely on longer history for decision-making. In the spontaneous behavior experiment, while longer action-level dependency does not yield better results, the findings reinforce our claim that SWIRL serves as a versatile tool for hypothesis testing in behavioral datasets, enabling researchers to validate or challenge hypotheses regarding history dependencies. This adaptability underscores SWIRL’s potential to advance our understanding of complex behaviors.

Additionally, we are the first to explore the connection between HM-MDPs and other popular dynamics models (e.g., ARHMMs, SLDS) used in animal behavior studies. By bridging these frameworks, we promote a more comprehensive understanding of behavior analysis tools and demonstrate how SWIRL can be adapted to experiments previously analyzed with these models.

Moreover, our framework is easily extendable to large and continuous state and action spaces by replacing the Soft-Q iteration in the RL inner loop. This scalability positions SWIRL as not only a valuable tool for computational neuroscience but also a significant contribution to the broader machine learning community.

We 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 and articulates our contribution.

References: [1] Choi et al. "Inverse reinforcement learning in partially observable environments." JMLR (2011). [2] Chen et al. "Option-aware adversarial inverse reinforcement learning for robotic control." ICRA 2023 (2023). [3] Nguyen et al. "Inverse reinforcement learning with locally consistent reward functions." Advances in neural information processing systems, 28 (2015). [4] Haarnoja et al. "Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor." Proceedings of the 35th International Conference on Machine Learning (2018). [5] Garg et al. "Iq-learn: Inverse soft-q learning for imitation." Advances in Neural Information Processing Systems, 34 (2021). [6] Wulfmeier et al. "Imitating language via scalable inverse reinforcement learning." NeurIPS 2024 (2024).

We thank Reviewer Qvd7 for their valuable time in reviewing our work and responding to our rebuttal. Below, we provide further clarifications regarding our experimental approach and its validation:

1. Selection of Top Results Based on Training LL
Selecting results by training LL is a necessary practice to ensure reliable outcomes in the EM algorithm, as widely adopted in neuroscience research. We have discussed the rationale (EM can be sensitive to local optima) in the appendix. For instance, Keypoint-Moseq employs EM to infer ARHMMs for segmenting animal behaviors and uses a likelihood-based metric to remove outlier models [1].

2. Experiment Validation: Going Beyond Existing Literature
We would like to emphasize that our experimental validation goes ABOVE and BEYOND previous works using EM for IRL with multiple reward functions. Below, we detail how our baselines and the paper suggested by Reviewer Qvd7 (H-AIRL [4]) conduct their experiments:

• Multi-Intention IRL [2]
  The paper only mentions conducting 5-fold cross-validation without specifying the number of runs or random seed selections. Upon reviewing their code on GitHub, we found that for each fold, they perform 10 runs and select the best test LL as the test LL for that fold (top 1 from 10 runs based on test LL).

• Locally Consistent IRL [3]
  This work uses 3 random instances (likely equivalent to 3 random seeds) to perform experiments and plot error bars (3 seeds).

• H-AIRL [4]
  Similarly, H-AIRL experiments use 3 random seeds (3 seeds).

• SWIRL
  In contrast, we tested a larger number of random seeds (top 10 seeds from 22 seeds on Labyrinth) than these works to ensure reliability. For the spontaneous behavior dataset we only used 10 seeds and please see below response 3. for detailed explanation. Compared to previous work, we explicitly acknowledged this limitation of the EM algorithm (sensitive to local optima) in our analysis, which is a valuable contribution to the application of EM in IRL.

3. Adjusting number of seeds for each experiment We have outlined the reasoning behind the number of seeds used for each experiment in detail below. We acknowledge that using varying numbers of seeds seems random and may lead to confusion. To address this, we will conduct additional runs to ensure consistency across all three experiments by adopting the same strategy: selecting the top 10 training LL results from 20 random seeds. (UPDATE: We have updated the results in the Nov 27 revision.)

• Gridworld: We happened to conduct 1 more experiment when testing the swirl pipeline. As a result, we reported 11 random seeds. 10 random seeds instead should be totally fine. We didn't observe significantly lower training LL, which suggests that EM does not get stuck in bad local optima in this case.
• Labyrinth: On labyrinth experiment, we found that certain random seeds led to significantly lower training LL, which was in accordance with EM's limitation. To ensure we have enough meaningful results, we conducted more experiments with random seeds. Obviously we need more than 10 random seeds to report 10 meaningful results, but there is no specific reason for choosing the number of random seeds to be 22.
• Spontaneous Behavior: Spontaneous behavior dataset is a much smaller environment (only 9 states) than labyrinth and we didn't observe significantly lower training LL, which suggests that EM does not get stuck in bad local optima in this case. Therefore, we conclude that there is no need to do additional experiments and pick the top 10.

4. Mention in the Main Text
We acknowledge the importance of discussing these practices in the main text and will include them in a future revision due to time and page limits.

5. Outlier Selection in Seaborn Boxplots
We used the default outlier selection method in seaborn.boxplot. We would like to point out that outlier elimination in our plots does not affect their validity. For example, Reviewer Qvd7 noted that Figure 3E eliminates the best result for baseline I-1. However, including this outlier within the I-1 box would lead to only a minimal impact: I-1 would still outperform MaxEnt and remain inferior to I-2, S-1, and S-2. We will further clarify this point in future revisions.

We hope these clarifications address the reviewer’s concerns. Thanks again for the feedback.

References: [1] Weinreb et al. Keypoint-moseq: parsing behavior by linking point tracking to pose dynamics. Nature Methods (2024). [2] Zhu et al. Multi-intention inverse q-learning for interpretable behavior representation. Transactions on Machine Learning Research (2024). [3] Nguyen et al. Inverse reinforcement learning with locally consistent reward functions. Advances in neural information processing systems, 28 (2015). [4] Chen et al. "Option-aware adversarial inverse reinforcement learning for robotic control." ICRA 2023 (2023).

7. Biological plausibility: We have added additional biological evidences about history dependencies in the introduction section and provided more discussion on non-Markovian labyrinth results in Appendix B.4.2.

8. Minor: We thank the reviewer for pointing it out. We have fixed it according to the reviewer’s suggestion.

9. Discussion suggestion: Yes, SWIRL could indeed be used to evaluate whether models of intrinsic reward, such as expected free energy or empowerment, effectively capture animal behavior. For instance, SWIRL can be applied to infer reward functions directly from animal behavior data. Separately, intrinsic rewards can be computed using various models, including expected free energy or empowerment. By comparing the reward functions inferred by SWIRL with those computed by intrinsic reward models, we can assess the validity of these models. An appropriate intrinsic reward model should yield reward functions that align with those inferred by SWIRL.

We 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.

References: [1] Haarnoja et al. "Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor." Proceedings of the 35th International Conference on Machine Learning (2018). [2] Garg et al. "Iq-learn: Inverse soft-q learning for imitation." Advances in Neural Information Processing Systems, 34 (2021). [3] Wulfmeier et al. "Imitating language via scalable inverse reinforcement learning." NeurIPS 2024 (2024).

3. Scalability. We have added a scalability and broader impact discussion in Appendix A.4. While scalability is not the main focus of this work, as typical behavior datasets in neuroscience have small or moderate state-action space and current SWIRL implementation can perform efficiently on them, we acknowledge the need for a more general and scalable implementation to address broader applications. In its current form, every step of the SWIRL inference procedure, except for the Soft-Q iteration, is compatible with large or continuous state-action spaces. However, the Soft-Q iteration is limited to discrete state-action spaces and can be slow with large state-action space. Nevertheless, for applications requiring scalability and compatibility with general state-action spaces, alternative methods can be adapted to replace the Soft Q iteration in the RL inner loop. For instance, Soft Actor-Critic [1]. Furthermore, a promising future direction is to reformulate the standard MaxEnt IRL $r$-$\pi$ bi-level optimization problem in SWIRL as a single-level inverse Q-learning problem, based on the IRL approach known as IQ-Learn [2]. This method has been successfully adapted to large language models training, demonstrating great scalability [3].

4. Convergence analysis: We have added a convergence analysis discussion in Appendix A.3. The convergence of the SWIRL inference procedure can be established based on the convergence guarantee of the EM algorithm, MaxEnt IRL and Soft-Q iteration.

5. Limitations:

   5.1 SWIRL may face challenges in modeling complex animal behaviors where decision-making and hidden mode transitions are influenced by external environmental factors, internal alertness and physiological states. For instance, an animal foraging during daylight may exhibit different behaviors compared to nighttime, even when guided by the same internal food reward map. Another example involves the presence of predators. In an environment where predator locations are unknown, an animal's optimal foraging policy may not be heading directly to the food source. Instead, it might need to make decisions based on its alertness to the predators' presence. Additionally, physiological factors can further complicate behavior. For example, an animal intending to move right under a specific reward map might fail to execute the intended action due to poor physiological conditions, such as fatigue after prolonged activity or illness, resulting in its body to move straight instead. In such scenarios, SWIRL may struggle to accurately infer the true internal reward maps driving the animal's decisions.

   5.2 As we discussed in the scalability section, the Soft-Q iteration in SWIRL is limited to discrete state-action spaces. However, we can replace the Soft-Q iteration by other methods such as Soft Actor-Critic and make SWIRL compatible with continuous state-action space.

6. Robustness to noisy data: We discuss additional experiment results in appendix B.5. In this gridworld experiment, we introduced random permutations to a 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. The model maintained high accuracy with less than 10% permutation. Between 10% and 30%, SWIRL still 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 SWIRL can tolerate moderate levels of noise or incomplete data, making it suitable for real-world animal behavior datasets where such challenges are common.

We sincerely thank the reviewer trZk for their detailed and thoughtful comments on our paper. We greatly appreciate the time and effort they dedicated to analyzing our work and providing such constructive feedback. We are particularly encouraged by the reviewer’s recognition of SWIRL’s novelty and “potential to advance our understanding in neuroscience and behavioral science.”

We will now respond to each concern raised in the Weaknesses and Questions.

1. Methodology: We apologize for not making this clearer in the paper. We have added implementation details and hyperparameters selection discussion to the Appendix B. Specifically, for questions the reviewer raised:

   1.1 Number of hidden modes: We have included additional experiment results in Appendix B.4. For the labyrinth dataset, in Appendix B.4.1, we test the hidden modes from 2 to 5 and find that the test log-likelihood (LL) plateaus at 4 hidden modes. However, after taking a close look at the hidden segmentation and reward map, we find the 4 hidden modes result does not differ too much from 3 hidden modes: It mainly segments the explore mode into 2 explore modes with similar reward maps. As a result, we still show 3 hidden modes results in the main paper for simplicity. For the mouse spontaneous behavior experiment, in Appendix B.4.3, we show that we pick 5 hidden modes since the test LL plateaus at 5 hidden modes.

   1.2 Action-level history $L$: In the main paper, we present results for $L=2$ (S-2) as it effectively demonstrates the benefits of incorporating non-Markovian action-level history dependency into SWIRL. S-2 provides meaningful hidden mode segments with interpretable reward maps, whereas $L=1$ (S-1) fails to capture these dynamics adequately. To further strengthen this claim, we included additional experimental results in Appendix B.4.2, where we evaluated history lengths ranging from $L=1$ to $L=4$ for the labyrinth dataset. Our findings indicate that $L=4$ (S-4) achieves the highest test log-likelihood (LL), suggesting that animal behavior in this dataset exhibits a longer non-Markovian dependency. This observation aligns with the partially observable nature of this 127-node labyrinth: The mouse may not know the whole environment, so it tends to rely on longer state history to inform its decision-making. Due to the mouse's limited knowledge of the entire environment, it likely relies on a longer history of prior states to guide its decision-making. For the mouse spontaneous behavior experiment, in Appendix B.4.3, we show that $L=1$ (S-1) is the optimal choice, as $L=2$ (S-2) consistently provides lower test log-likelihood (LL) compared to $L=1$ (S-1).

   1.3 Hyperparameter selection: We have added the discussion of the selection of discount factor $\gamma$ and temperature $\alpha$ in Appendix B.3.

2. Computational complexity and empirical runtime:

   2.1 Complexity analysis: We have added a complexity analysis in Appendix A.3. The key bottleneck is in the RL inner loop. Current Soft-Q iteration has a time complexity of $O(ZIS^{L}A^2)$, where $Z$ is number of hidden modes, $I$ is number of Soft-Q iterations, $S$ is number of states, $A$ is number of actions and $L$ is the history length. Although as we know the transition model $P(s'|s, a)$, the $S^L$ space can be reduced by removing impossible $(s, s')$ pairs, we recognize that longer $L$ in SWIRL does require a considerable amount of computing resources.

   2.2 Empirical runtime: We have added an empirical runtime discussion in B.2. Empirically, SWIRL speed is not slow for typical animal behavior datasets. We implemented SWIRL efficiently with JAX, leveraging the features of vectorization and just-in-time compilation to improve our performance. Every SWIRL (S-2) experiment on all 3 datasets (simulated gridworld, labyrinth and dopamine) can finish within 15-30 minutes on a V100 GPU. For longer $L$, a S-4 experiment on labyrinth with 50 EM iterations take 2-3 hours to finish on a L40S GPU. We switch to a L40S GPU for $L=4$ due to the V100 GPU's insufficient memory capacity.

We thank Reviewer trZk for their valuable time and effort in reviewing our work and responding to our rebuttal. We have made the following clarifications regarding the concerns mentioned in their reply:

1. Following the suggestions of all three reviewers, we have conducted experiments for all three datasets using the same number of random seeds and updated the results in the Nov 27 revised version of the paper. The updated results do not impact the conclusion of those experiments.

2–3. We acknowledge the need to discuss those limits and relationships more in the main text. Given the time and page limit, we kept most of the updated content in the appendix but will definitely integrate some into the main text in future revisions to enhance clarity and accessibility.

We thank again for the Reviewer trZk's valuable feedback and hope above clarifications address the reviewer’s concerns.