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

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:

1. Scalability concern: Please refer to section 1 of our reply to Reviewer 9ThW.

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.

3. Experiment setup: We will add details of each environment in our next revision. For the state and action space:

5x5 Gridworld: State space: Discrete(25); Action space: Discrete(5). The 5 actions are ‘up’, ‘left’, ‘down’, ‘right’, ‘stay’.

Labyrinth: 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’.

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

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.

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.

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.

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.

We sincerely appreciate the time and effort reviewer CRRd dedicated to analyzing our work. Below, we address each of the questions and concerns:

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.

  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.

  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.

2. Scalability and longer history: Please 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.

3. Hidden mode with infrequent occurrence: In 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.

4. Section heading issue: We apologize for the heading space issue and will correct it in our next revision.

5. Full name of IRL appears multiple times: We will use the abbreviation consistently in our next revision.

6. ARHMM abbreviation issue: ARHMM should be the abbreviation of the autoregressive hidden Markov model. We will address this issue in our next revision.

7. Response to Questions For Authors:

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

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.

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:

1. Comparison among Ashwood et al., Zhu et al. and SWIRL

  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.

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

  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.

2. Response to Questions For Authors: SWIRL provides novel insights into real animal datasets

  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.

  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.

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.

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 regarding the scope of the novelty.

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:

1. Scalability and Generalization to Larger State Space

  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.

  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.

  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.

2. Selection of number of hidden modes

Yes, 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.

3. Selection of L for labyrinth experiment

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

4. Consideration of longer L

  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.

  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.

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.

Reference: [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).