Paper: L(M)V-IQL: Multiple Intention Inverse Reinforcement Learning for Animal Behavior Characterization (ICLR-2024)

Reviewer comment (Novelty):
The main weakness is the high similarity of this work to Ashwood et al. (NeurIPS 2022). Although the authors mention in the Appendix that the aforementioned approach “is limited to capturing continuous intra-episode variation of reward functions during navigation behavior, and difficult to adopt to other environments,” this claim requires further substantiation.

Author response:
*”The Authors mention in the Appendix that the aforementioned approach “is limited to capturing continuous intra-episode variation of reward functions during navigation behavior, and difficult to adopt to other environments” but, should that be true, that requires further substantiation.”*

Thanks for your question. The DIRL approach from Ashwood et al. (NeurIPS 2022) assumes the rewards vary continuously over time, i.e. the intention transition happens after each action execution. The limitation of their approach can be indicated from the Gridworld environment in our first experiment, where we consider each episode was demonstrated under one of two intentions, i.e. the intention is fixed within each episode, and the intention transition only happens between episodes. DIRL will not be feasible under such environments since the variation of reward functions is, however, discrete. In contrast, our method is viable in cases where the transition of intention occurs either after each action execution, i.e. continuous (as discussed in Remark 2 in Appendix B.2), or at the end of each episode, i.e. discrete (as theoretically demonstrated in Appendix B.2). We have elucidated this argument in the current revision to provide a clearer understanding of the adaptability and applicability of our approach in diverse settings.

Reviewer comment (Theory):
Page 6 under Figure 3: Why is it necessary to punish the types of reward irrelevant to the intentions? A more natural way seemingly would be to set them equal to zero, but the reviewer assumes this natural way has not worked for some reason.

Author response:
*”Page 6 under Figure 3: why is it necessary to punish the types of reward irrelevant to the intentions? A more natural way seemingly would be to set them equal to zero. I assume this natural way hasn’t worked out for some reason?”*

Thank you for your suggestion. We have conducted additional experiments on the Gridworld environment with the same resource distribution but removed the punishment on the intention irrelevant reward, i.e. replacing the -1 reward on the type of reward irrelevant to the intentions with 0. Related results are shown in Appendix C.2. Our algorithms still worked under this configuration. However, in this case there was an increased overlapping between some of the demonstrated trajectories under different intentions (Figure 6, Top), resulting in a slight decrease of the performance of LV-IQL, but LV-IAVI and LV-IQL still outperformed the single intention algorithm IAVI and IQL in trajectory clustering and recovering corresponding expert reward functions.

Reviewer comment (Evaluation):
Figure 4A: Why does the LL in the training curve drop? The reviewer finds it unlikely to be explained by overfitting as suggested in the text.

Author response:
*”Figure 4A: why does the LL in the training curve drop? That is unlikely to be explained by overfitting as suggested in the text.”*

Thanks for your question. As we consider each state $s_t$ in the state space $\mathcal{S}$ to be defined with a set of truncated history information to avoid explicitly describing a partially observable MDP formulation for the reversal-learning task, the size of the state space $\mathcal{S}$ will grow exponentially as the history truncation length $\ell_h$ increases. Thus the sampling on some states given the fixed animal demonstrations could be insufficient for the IQL to reconstruct a precise reward function via stochastic approximation. The abnormal drop of the training LL will be omitted given a larger animal demonstration dataset. We have updated this explanation in the current revision for better clarification.

Reviewer comment (Evaluation):
Figure 5C: Following up on the previous point, this figure suggests that the third intent is just spurious (not stable; immediately reverses to the first intent). The reviewer asks why it is considered important.

Author response:
*”Page 7 bottom line: “Although model performance continued to improve slightly with more latent states, we will focus […] on […] 3 states”. Wouldn’t it be easier to make this argument by using the Bayesian Information Criterion instead of the pure NLL to choose K? This way one can arrive at a principled number of intents that very well may turn out to be equal to 2.”*
*”Figure 5C. Following up on my previous point, this figure leaves me with the impression that the third intent is just spurious (not stable; immediately reverses to the first intent). Could you please comment on why you consider it important?”*

Thanks for your question. We have calculated the BIC for LMV-IQL across a varying number of latent states in the current revision (Figure 4C). It is shown that the BIC for the number of latent states $K = 2$ and $K = 3$ are similar, suggesting that both value are reasonable choices for $K$. The reason that the 3 latent states model was preferred was that the 3 latent states model provided a fine grained information about exploratory behavior, i.e. the disengaged, passive exploration (‘Win-stay’, performing random action selection after lose trails) and the active exploration (‘Exploration’, choosing the other side compared to the preceding trial regardless of win or lose) (Figure 4D). Figure 4H also shows that the mean posterior probability of the ‘Exploration’ latent state during the initial 5 trials of a new block showed a positive correlation with the average probability of the ‘Exploitation’ state during the final 5 trials of the preceding block, further supporting this point that the ‘Exploration’ latent state appears to involve a deliberate, exploration-oriented action selection only when mice are highly engaged and possess a good understanding of the environment. The ‘Exploration’ latent state gives the opportunity to characterize a type of learning strategy based on error. Since it has been shown that rodents (Kononowicz et al. PNAS 2022), as well as humans, can learn from their errors, we then think that maintaining the exploration latent state confers a clear advantage in terms of biological interpretability. Additionally, the instability of the ‘Exploration’ latent state as observed here could be due to our two-armed bandit reversal-learning task design. We considered a performance-determined reward switch here, i.e. the rewarded spout would be changed only when the subject reached an online performance of 75% correct in a 15-trials sliding average window and a minimum of 20-trials block. Thus the high fraction of ‘Exploration’ latents states would lead to unfinished sessions (each session was required to have 4 blocks), which was filtered out from our dataset.

Reviewer comment (Presentation):
Figure 3: The overlap of the crosses (x) and dots (.) is hard to follow. The reviewer requests using an alternative way to represent the data.

Author response:
*”Figure 3: the overlap of the crosses (x) and dots ( . ) is hard to follow. Could you please use an alternative way to represent this data?”*

Thanks, we have updated the data representation in the current revision. Instead of depicting a t-SNE visualization of expert trajectories (Figure 3 from the prior revision), the pertinent results are now depicted in Figure 2 (Top), where trajectories are directly plotted within the Gridworld environment as red lines, with the corresponding state-values displayed in the background.

Reviewer comment (Novelty):
Page 2: Regarding the claim “[IRL’s] adoption as mathematical behavior models in neuroscience research has been relatively limited” — the reviewer disagrees, noting that IRL seems to be an up-and-coming tool, with strong examples from Jon Pillow’s and Xaq Pitkow’s groups.

Author response:
*”Page 2: “[IRL’s] adoption as mathematical behavior models in neuroscience research has been relatively limited”. I had another impression – it seems to be an up-and-coming tool, as exemplified by some awesome works from Jon Pillow’s and Xaq Pitkow’s groups.”*
*”Page 2: “[our method presents a […] framework for characterizing the delicate balance between exploration and exploitation […] which constitutes a […] comparatively understudied aspect within the realm of neuroscience.” I’d say that, first, there’s a huge spillover from the machine learning field of intrinsic motivation (a.k.a. an internal reward for exploration); many of these works claim biological plausibility. There are some other nice works, e.g. Pisupati et al (eLife 2021) and references therein that directly address the issue. There’s also lots of work on Bayesian optimality that study the deviations from optimal exploitation to account for the environmental dynamics, e.g. Yu and Cohen (NeurIPS 2008).”*
*”Page 4: “In behavioral neuroscience, it is commonly considered that animals alternate between multiple intentions under the Markov property”. The entire reason why the Ashwood et al (Nat Neurosci 2022) paper cited there emerged is because that’s not how people used to characterize natural behaviors. This is reflected literally in the first sentence of the said paper. While this new work has gotten substantial traction in the field, I wouldn’t say that that new way to model data has completely wiped out the conventional approach.”*

Thanks for your suggestions. We have updated these statements.

Does this address your concerns?

Reviewer comment (Writing):
Page 4: The claim “In behavioral neuroscience, it is commonly considered that animals alternate between multiple intentions under the Markov property” is questionable. The reviewer points out that Ashwood et al. (Nat Neurosci 2022) emerged precisely because this was not the previous characterization of natural behaviors, and while the new approach has gained traction, it has not completely replaced the conventional one.

Author response:
*”Page 2: “[IRL’s] adoption as mathematical behavior models in neuroscience research has been relatively limited”. I had another impression – it seems to be an up-and-coming tool, as exemplified by some awesome works from Jon Pillow’s and Xaq Pitkow’s groups.”*
*”Page 2: “[our method presents a […] framework for characterizing the delicate balance between exploration and exploitation […] which constitutes a […] comparatively understudied aspect within the realm of neuroscience.” I’d say that, first, there’s a huge spillover from the machine learning field of intrinsic motivation (a.k.a. an internal reward for exploration); many of these works claim biological plausibility. There are some other nice works, e.g. Pisupati et al (eLife 2021) and references therein that directly address the issue. There’s also lots of work on Bayesian optimality that study the deviations from optimal exploitation to account for the environmental dynamics, e.g. Yu and Cohen (NeurIPS 2008).”*
*”Page 4: “In behavioral neuroscience, it is commonly considered that animals alternate between multiple intentions under the Markov property”. The entire reason why the Ashwood et al (Nat Neurosci 2022) paper cited there emerged is because that’s not how people used to characterize natural behaviors. This is reflected literally in the first sentence of the said paper. While this new work has gotten substantial traction in the field, I wouldn’t say that that new way to model data has completely wiped out the conventional approach.”*

Thanks for your suggestions. We have updated these statements.

Does this address your concerns?

Reviewer comment (Theory):
The fitted model is somewhat simplistic. Latent states are assumed to be multinomial or Markov, but the most plausible biological assumption would be that transitions between drives also interact with reward/satiety/recent history.

Author response:
Thank you for this great question. It is indeed a pertinent consideration to explore the utilization of a controlled latent state transition system, potentially incorporating external covariables, rather than the autonomous system currently under consideration in this paper. A promising avenue for future research could involve adopting a generalized linear model with inputs such as reward, satiety, recent history, etc., in lieu of the current generalized Bernoulli process or Markov process controlling the latent state transition dynamics. This adjustment would allow for the incorporation of the identification of potential external factors influencing intention transition dynamics. Furthermore, there is merit in exploring the feasibility of exerting control over the intention transition. Your suggestion is invaluable, and we will thoroughly explore these possibilities in our ongoing work.

Reviewer comment (Experiments):
There are only two experiments: one on simulated data (where it is) compared to IAVI and IQL but not the Ashwood et al. or other similar models that might be applicable. Similarly, the mouse behavior is quite limited in terms of the need for RL. Again, the Ashwood Nature Neuro paper or the Ebitz, Albarran, and Moore (2018) provide fairly flexible models that are likely to capture the data as well. Given the synthetic data, one would have expected a more challenging task here as a target for IRL.

Author response:
Thanks for your great suggestion. We have included an additional experiment on the application of LMV-IAVI to real mice navigation trajectories within a labyrinth from Rosenberg et al. (2021). This dataset was also considered as a benchmark for IRL algorithm comparison in Ashwood et al. (NeurIPS 2022), since it would be challenging to define a suitable reward function to apply forward models (such as the GLM-HMM model from Ashwood et al. (Nat. Neuroscience, 2022) and the framework from Ebitz et al. (2018)). In this experiment, we compared our approaches with the state-of-the-art Dynamic Inverse Reinforcement Learning (DIRL) framework presented by Ashwood et al. (NeurIPS 2022). Related results are shown in Section 5.2 (Page 6) and Appendix D.2 (Page 16). Our approaches outperform DIRL in predicting mice behavior (indicated by the log-likelihood on held-out trajectories, Figure 3B, and Figure 7A) and are able to provide interpretable reward functions (Figure 3D, Figure 7B).

Reviewer comment (Theory):
Is it possible to incorporate some recent reward history into the transition structure? Since the inference algorithm is EM, will any EM-compatible latent variable model work, in principle?

Author response:
Thank you for this great question. It is indeed a pertinent consideration to explore the utilization of a controlled latent state transition system, potentially incorporating external covariables, rather than the autonomous system currently under consideration in this paper. A promising avenue for future research could involve adopting a generalized linear model with inputs such as reward, satiety, recent history, etc., in lieu of the current generalized Bernoulli process or Markov process controlling the latent state transition dynamics. This adjustment would allow for the incorporation of the identification of potential external factors influencing intention transition dynamics. Furthermore, there is merit in exploring the feasibility of exerting control over the intention transition. Your suggestion is invaluable, and we will thoroughly explore these possibilities in our ongoing work.

Reviewer comment (Theory):
Where are the bottlenecks for the method in terms of inferential complexity? Is the limiting factor the IAVI regression (i.e., the size of the tabular problem) or the EM complexity?

Author response:
Thanks for your question. The computational complexity of IAVI is akin to that of standard value iteration, with the difference being an additive term for solving the SLE – a term that is often of minor importance given the usual predominance of the number of states over the number of actions. The primary contribution of Kalweit et al. was indeed the reduction in computational complexity relative to other IRL approaches, which was the reason for our selection of IAVI. The limiting factor is, therefore, rather identified as the complexity of the EM in the outer loop, although the computational demands of solving the MDP in the inner loop are of course inherently significant. Hence, the limiting part cannot be universally assigned without context; it depends on specific factors such as the size of the state space, the convergence rate of the EM, the convergence criteria of value iteration, and the complexity of computations within each E-step.

Reviewer comment (Presentation):
Some of the figures need additional details/components. E.g., Figure 1, 2 need color scalebars, Figure 3 would benefit from some explanation of the legend (where are the red and blue squares?), the colors on the state labels in Figure 4C are unnecessary and uncorrelated with the colors in the legend.

Author response:
*“Some of the figures need additional details/components. E.g., Figure 1, 2 need color scalebars, Figure 3 would benefit from some explanation of the legend (where are the red and blue squares?), the colors on the state labels in Figure 4C are unnecessary and uncorrelated with the colors in the legend, …”*

Thanks, we have updated the figures.

Reviewer comment (Reproducibility):
Definitions of the comparison methods were weak. For example, 'IAVI was further extended to the sampling-based model-free Inverse Q-learning (IQL) algorithm' with no citation or explanation of how the authors of this paper implemented those algorithms, is insufficient.

Author response:
*“Definitions of the comparison methods were weak. For example, 'IAVI was further extended to the sampling-based model-free Inverse Q-learning (IQL) algorithm' with no citation or explanation of how the authors of this paper implemented those algorithms, is insufficient.”*

Thanks for your suggestion. IQL was initially introduced alongside IAVI as a model-free extension in Kalweit et al. (2020), and we have faithfully implemented their approach in our current work without any modifications. In response to your feedback, we have diligently revised the relevant sections of our paper, providing more explicit citations for clarity.

Reviewer comment (Writing):
The primary metric, EVD, is cited but not defined.

Author response:
*“Similarly, the primary metric, EVD, is cited but not defined.”*

Thanks, we have updated the EVD definition in our paper for better clarity.

Reviewer comment (Experiments):
The authors only show an improvement over IAVI and IQL, and do not compare these other methods (including LV-IQL, which performed the same on the simulated dataset) in the experimental dataset.

Author response:
*“The authors only show an improvement over IAVI and IQL, and do not compare these other methods (including LV-IQL, which performed the same on the simulated dataset) in the experimental dataset.”*

Thanks for your great suggestions. We included an additional experiment on real mice navigation trajectories within a labyrinth from Rosenberg et al. (2021) as a benchmark to compare our algorithm with the state of the art — Dynamic Inverse Reinforcement Learning (DIRL) from Ashwood et al. (NeurIPS 2022). Related results are shown in Section 5.2 (Page 6) and Appendix D.2 (Page 16). Our approaches outperform DIRL in predicting mice behavior (indicated by the log-likelihood on held-out trajectories, Figure 3B, and Figure 7A) and are able to provide interpretable reward functions (Figure 3D, Figure 7B).

Reviewer comment (Experiments):
The authors motivate their method as extending beyond a single reward function, but then apply it only to the case of 2-3 rewards/intentions/states. Can this be extended easily to more than a small number?

Author response:
*“The authors motivate their method as extending beyond a single reward function, but then apply it only to the case of 2-3 rewards/intentions/states. Can this be extended easily to more than a small number?”*

Thank you for your insightful question. While our manuscript predominantly delves into a detailed exploration of the learned reward function and intentions from LMV-IQL with 2-3 latent states, we also conducted model fitting with a 4-latent-state LMV-IAVI on mice navigation trajectories and a 5-latent-state LMV-IQL on mice reversal-learning behavior. The primary practical constraint of EM based algorithms arises from computational time, particularly as the number of intentions increases—an inherent limitation shared with EM-MLIRL (Babes et al., 2011) and MI-Σ-GIRL (Likmeta et al., 2021). Notably, IAVI possesses the unique capability to ***analytically*** derive a matching reward function for observed behavior in ***closed-form***, thus omitting the computationally intensive inner loop at each iteration of EM. This analytical feature distinguishes our approaches, enabling their application on more expansive latent state spaces compared to existing algorithms.

Does this address your concerns?