Paper: Learning interpretable control inputs and dynamics underlying animal locomotion (ICLR-2024)

Reviewer comment (Experiments):
There is a lack of benchmarks comparisons to other techniques in the literature such as LFADs. Even though some suggest continuous inputs for their control scheme, could they be used to model the behavioral dynamics in Figure 3C for instance.

Author response:
That is a good point. We have now included a comparison to LFADS (where we used the standard LFADS implementation, with an autoregressive prior). In summary, we found that LFADS learns inputs throughout the whole time series that exhibit a strong correlation with the actual behavior (see **new Figure S6**). This is somewhat expected, as the autoregressive prior does not discourage the model from using inputs to fit the temporal structure in the data, that could alternatively be captured by the dynamics. To address the second part of your question, we found that such inputs allow us to accurately reconstruct the behavior, but lead to a representation in the LFADS model that is less predictive of the movement category than the iLQR-VAE MGU representation (see **updated Figure 2.C**).

Reviewer comment (Novelty):
I found the applications a bit limited. It wasn’t obvious from this analysis that there were fundamentally new types of experiments that were enabled by this approach.

Author response:
Thank you for highlighting the concern regarding limited application. To address this, we implemented our framework on a C. elegans dataset, capitalizing on a significant advantage — the presence of controlled external stimuli in certain trials. We exploited this well-timed laser heat shock as a ground truth control input. In **new Figure 4.D**, we show a salient control is inferred at the moment of the laser pulse, a feature not readily apparent in the posture space. This outcome is particularly encouraging as it suggests the potential for the model to be employed in conjunction with perturbation experiments in the future. For instance, it could be used alongside optogenetic stimulation applied to a population of neurons. In such experiments, the timing would be controlled by the experimentalist, while the impact on the dynamical system, though unknown, could be inferred by our model.

Reviewer comment (Experiments):
Overall, I think the lack of benchmarks combined with the lack of novel applications was a major sticking point for me and I view this manuscript as very borderline. I would be happy to reconsider the manuscript (which was very well written) if these were added or addressed. I see this as just 1 compelling figure short of a reasonable submission.

Author response:
Thank you for highlighting the concern regarding limited application. To address this, we implemented our framework on a C. elegans dataset, capitalizing on a significant advantage — the presence of controlled external stimuli in certain trials. We exploited this well-timed laser heat shock as a ground truth control input. In **new Figure 4.D**, we show a salient control is inferred at the moment of the laser pulse, a feature not readily apparent in the posture space. This outcome is particularly encouraging as it suggests the potential for the model to be employed in conjunction with perturbation experiments in the future. For instance, it could be used alongside optogenetic stimulation applied to a population of neurons. In such experiments, the timing would be controlled by the experimentalist, while the impact on the dynamical system, though unknown, could be inferred by our model.

Reviewer comment (Experiments):
Does the number of control inputs match the number of bout types people have reported in the literature?

Author response:
This is an interesting point you raised. To clarify, we found that rather than controlling distinct movements, the different control inputs were combined in different ratios to generate different categories of movements. We added a figure illustrating this (**Figure S3**).

Reviewer comment (Experiments):
There is no ground truth in behavior for the driving inputs and thus there is no way to validate that they are correct, however with an appropriate experiment (e.g. applying a perturbation or stimulus at a specific time) perhaps they could be useful.

Author response:
Thank you for highlighting the concern regarding limited application. To address this, we implemented our framework on a C. elegans dataset, capitalizing on a significant advantage — the presence of controlled external stimuli in certain trials. We exploited this well-timed laser heat shock as a ground truth control input. In **new Figure 4.D**, we show a salient control is inferred at the moment of the laser pulse, a feature not readily apparent in the posture space. This outcome is particularly encouraging as it suggests the potential for the model to be employed in conjunction with perturbation experiments in the future. For instance, it could be used alongside optogenetic stimulation applied to a population of neurons. In such experiments, the timing would be controlled by the experimentalist, while the impact on the dynamical system, though unknown, could be inferred by our model.

Reviewer comment (Evaluation):
Generally, more concretely linking any of the observations to real neural or behavioral variables would improve the manuscript's impact.

Author response:
Thank you for highlighting the concern regarding limited application. To address this, we implemented our framework on a C. elegans dataset, capitalizing on a significant advantage — the presence of controlled external stimuli in certain trials. We exploited this well-timed laser heat shock as a ground truth control input. In **new Figure 4.D**, we show a salient control is inferred at the moment of the laser pulse, a feature not readily apparent in the posture space. This outcome is particularly encouraging as it suggests the potential for the model to be employed in conjunction with perturbation experiments in the future. For instance, it could be used alongside optogenetic stimulation applied to a population of neurons. In such experiments, the timing would be controlled by the experimentalist, while the impact on the dynamical system, though unknown, could be inferred by our model.

Reviewer comment (Evaluation):
In Figure 3 why does the classification accuracy decrease with time (observation time?)?

Author response:
In Figure 3, a distinct classifier is trained at each time point (excluding the dashed line that represents training from the complete time series of posture). Therefore as time advances toward the end of the movements, the classification gradually diminishes, approaching the chance level.

Reviewer comment (Presentation):
In Figure 4C and associated main text, the authors mention that the reduced model added eigenmodes with larger timescales, while it neglected the small timescale modes of the non-reduced system. While the second part of the previous sentence is expected, I am not sure how I should understand the first one. Can you provide some intuition. Also related to this, when I first looked at Figure 4C I thought that the large timescale values of the reduced and non-reduced model overlap, therefore the light grey ones are not visible. Can you probably make the circles of the reduced system non-opaque or non-filled to make the plot clearer?

Author response:
Thank you for this question. That is a good point, and indeed the behavior of the model reduction has interesting properties that are not all intuitive (which is another reason why we believe our example can be informative to other researchers in the area).

The reduced model learns modes that differ slightly from the top eigenmodes of the previous system, as it is a different linear system. In the original system, it is possible for multiple modes to combine with one another, leading to effectively longer timescales arising in the system. Since the reduced model looks for a minimal version of the dynamics, it might learn timescales that do not exist in the initial system (and that may be longer). Let us know if this clarifies things. We have, in addition, followed your advice and **updated Figure 3c** with open circles for the reduced model spectrum.

Reviewer comment (Reproducibility):
I think it would be helpful if the authors mention in the supplement how they fit the linear model of their framework, and provide a brief description of the iLQR-VAE framework that is a crucial component of the proposed approach.

Author response:
This is a good point, we have included this information in the form of supplementary sections. Briefly, here, we found the linear model by refitting a sparse linear RNN to the data.

Reviewer comment (Reproducibility):
Model Selection C.1. section is missing from the supplement.

Author response:
Thank you for spotting those, we have amended the manuscript to correct them!

Reviewer comment (Writing):
There is a typo in the subscripts in Eq. 5.

Author response:
Thank you for spotting those, we have amended the manuscript to correct them!

Reviewer comment (Novelty):
The paper utilizes iLQR-VAE for learning the parameters and inferring the control signals, but it is unclear what the difference is between the proposed model and the previous model. The differences seem minimal, making this a good application paper but one that does not introduce many novel elements from a modeling perspective.

Author response:
Thank you for your comment: indeed, our paper focuses on the application of iLQR-VAE to naturalistic animal behavior. We fine-tuned the model parameters to work on the datasets of interest to us, but we indeed did not change the existing method.

We would indeed describe our work as “*an application paper*”, more than proposing a new methodology. Note however that, **in the ICLR call for papers, “*applications in audio, speech, robotics, neuroscience,  biology, or any other field*” are considered and welcome. Therefore, we do not think that our focus on a neuroscience application should be ground for rejection.**

Reviewer comment (Novelty):
The model reduction technique for linear models, while not usually applied to behavioral modeling, is a very well-known concept and does not offer novelty from a methodological perspective.

Author response:
We agree that the model reduction (MOR) approach is certainly not new. However, we are not aware of other work applying it to behavioral modeling or neuroscience. We thus believe our paper would be very relevant to members of those communities, as a case study using MOR to extract insights from data. 

Additionally, applying model reduction to an input-driven dynamical system whose parameters are learned from data is to the best of our knowledge a novel approach. Asides from iLQR-VAE and LFADS, no methods have focused on learning input-driven dynamics from data.  We believe that the combination of such approaches with MOR opens up very exciting avenues for the field of systems identification, and is a valuable contribution.

Reviewer comment (Experiments):
There are interesting takeaways from a neuroscience perspective, but these would need to be validated more thoroughly and may be more suitable for a different venue.

Author response:
We agree that providing only the zebrafish as an example application somewhat limited the scope of the paper. We have addressed this by including a new application (see **new Figure 4** and **Figure S8**) to behavioral recordings from crawling C. elegans (worm). While the posture is still low-dimensional, the movement dynamics are more complex than the zebrafish because of their continuous nature, and the fact that the worm is changing its behavior in response to unknown environmental stimuli at variable time intervals.

We found that the model performed well on this dataset. It reconstructed the behavior accurately, using inputs that strongly correlated to ground truth aversive heat stimuli (**new Figure 4.D**).

Additionally, after performing the model reduction on the learned dynamical system, the top 3 modes of the reduced system matched the previously identified behavioral modes of the worm [`1`]. We hope that these results will convince you of the general applicability of our method, and the insights it can yield!

As we detail above, we believe that our new C. elegans results address this concern, both because of the ground truth input recovery, and the fact that the model reduction identifies modes that can be related to prior work.

Reviewer comment (Experiments):
The kinematics of the zebrafish are relatively simple, while the naturalistic behaviors of other animals (such as mice and monkeys) are more complex. It is unclear whether the model will still be able to capture these more complex dynamics, and one real-world dataset may not provide enough insight into the model's generalizability for this question.

Author response:
We agree that providing only the zebrafish as an example application somewhat limited the scope of the paper. We have addressed this by including a new application (see **new Figure 4** and **Figure S8**) to behavioral recordings from crawling C. elegans (worm). While the posture is still low-dimensional, the movement dynamics are more complex than the zebrafish because of their continuous nature, and the fact that the worm is changing its behavior in response to unknown environmental stimuli at variable time intervals.

We found that the model performed well on this dataset. It reconstructed the behavior accurately, using inputs that strongly correlated to ground truth aversive heat stimuli (**new Figure 4.D**).

Additionally, after performing the model reduction on the learned dynamical system, the top 3 modes of the reduced system matched the previously identified behavioral modes of the worm [`1`]. We hope that these results will convince you of the general applicability of our method, and the insights it can yield!

Reviewer comment (Experiments):
There is no comparison with other models presented.

Author response:
We agree that the lack of comparison with other methods was a weakness of the paper. We have now included a comparison with LFADS on the zebrafish dataset. In summary, we found that LFADS (which assumes an autoregressive prior over the inputs), can fit the behavioral time series but does so using continuous inputs which are strongly correlated to the posture (see **new Figure S6**). Moreover, LFADS leads to latent trajectories which are less predictive of the behavioral category (see **updated Figure 2.C**).

Reviewer comment (Presentation):
Section 4.1 contains mismatched figure labels: Figure 1C should be Figure 1B; Figure 1D should be Figure 1C.

Author response:
**Regarding Minor Changes**
Thank you for spotting the few typos and mislabelling of figures. We have made these amendments to the updated manuscript.

Reviewer comment (Writing):
In Figure 3 Caption, there is a missing space: ‘udriving’ should be ‘u driving’.

Author response:
**Regarding Minor Changes**
Thank you for spotting the few typos and mislabelling of figures. We have made these amendments to the updated manuscript.

Reviewer comment (Writing):
On Page 15, the phrase ‘cannot be reach via observations’ should be corrected to ‘cannot be reached via observations’.

Author response:
**Regarding Minor Changes**
Thank you for spotting the few typos and mislabelling of figures. We have made these amendments to the updated manuscript.

Reviewer comment (Writing):
On Page 16, the phrase ‘Therefore, we can the compute the’ should be corrected to ‘Therefore, we can compute the’.

Author response:
**Regarding Minor Changes**
Thank you for spotting the few typos and mislabelling of figures. We have made these amendments to the updated manuscript.

Reviewer comment (Experiments):
The paper could do with a proper model comparison, i.e. vs moseq.

Author response:
Thank you for your comment. We agree that the original version of the paper was lacking a comparison with other methods. We have now added a comparison to LFADS (a method similar to iLQR-VAE, but with a different recognition model, and which typically uses an autoregressive input prior).  We found that LFADS can reconstruct the behavior, but using inputs that are highly correlated with the posture (see **new Figure S.6**). Additionally, it gives rise to latent trajectories that are a lot less predictive of the behavior than the MGU model (see **updated Figure 2.C**). 

Note that Moseq is indeed a widely used method in the field of behavioral modeling, but its principal use is for segmenting movements into motifs. Thus, since the zebrafish behavior is by nature already segmented in bouts, we did not think that MoSEQ would yield additional insights. However, if you had a specific comparison in mind that you would like us to do we would be happy to hear that!

Reviewer comment (Writing):
The Balanced truncation bit needs to be better explained with some intuition.

Author response:
Thank you for that suggestion – we have tried our best to make the balanced truncation as intuitive as possible (see for instance the intuition given regarding the observability or controllability Gramians, or the intuition for the Hankel singular values). 

We have also moved the pendulum example to the Appendix, and included more details on how it illustrates an example case of the balanced truncation, hopefully providing further intuition for the behavior of the model.

However, if you have any suggestions on how to further improve the clarity of this section (e.g. specific places that you find confusing) we would be happy to further amend the paper!

Reviewer comment (Presentation):
Fig 1 A: what does blue and orange correspond to? Different sequences?

Author response:
That is a good point – blue and orange correspond to two different input / output channels of the pendulum. We have clarified this in the caption.

Reviewer comment (Presentation):
Some of the Figures seem to be mis-referenced, e.g., “As shown in Figure 3B, the MGU could reconstruct bouts using a sparser control compared to the LDS."

Author response:
Thank you very much for spotting those, we have amended the manuscript to fix them!