Using Reinforcement Learning to Investigate Neural Dynamics During Motor Learning

Thank you for the review! Investigation of the properties of weight changes that reflect more or less brain-like learning activity is indeed interesting. We are excited to explore this in future work. We will expand the discussion to comment on this.

Thank you for your review, thoughtful questions, and suggestions. We reply to them in order.

how the 8 cost terms in [Berret et al. 2011] relate to the two proposed terms formulated as the divergence of policy needs more explanation and elucidation.

Connection to Berret et al., 2011: Berret et. al. 2011 find that of 8 cost terms, 3 terms make up a majority of the cost function:

$\theta$ corresponds to joint angles and $\tau$ to torque.

Our KL divergence terms simplify (see Appendix) to:

KL_zeroness

KL_smoothness

The KL_zeroness $\mu^2/2$ term is related to the squared angular acceleration cost by a linear model that is a function of the joint angles depending on the limb model. The smoothness term approximates the squared third derivative of position (jerk) and is thus related to the “angle jerk” term. Our model does not incorporate biomechanics; our primary goal is to analyze neural effects. We therefore directly apply this cost to the end effector’s acceleration. We will clarify this in the manuscript.

From Fig.2a,b,c, [without regularization] the model [... seems similar] to the monkey behavior. Therefore [the KL terms] need more justification

Additional justification of KL terms: To address your concern, we performed additional analyses (Appendix Figure 6) to show these KL terms increase similarity to the animal’s trajectories. First, the trajectories with neither constraints typically fail to converge to a reasonable trajectory. Appendix Figure 6a show unconstrained trajectories that curve in the direction of the curl field after significant learning (red curves). This clearly does not match the monkey’s curl field reaches which curve in the opposite direction. Second, the overcompensation of the curl field becomes more consistent, but excessive, with the zeroness constraint alone. As described in Fig 2f, it is optimal to “overcompensate” in a curl field to speed up the reach (Fig 2f, lines 310-312). Smoothness-only constrained models take advantage of this mechanism and exhibit obvious overcompensation beyond what is observed in biophysical reaches (Appendix Figure 6c). Finally, the smoothness constraint helps to reduce curvature to better match the curvature of the monkey’s reach (Appendix Figure 6ab). Together, these constraints produce more consistent and monkey-like learning trajectories. We have added these results and clearer connections to Berret et al., 2011 to the manuscript.

neural population trajectories in the PCA space [...] for motor learning and washout [were] not examined in the current modeling work. A comparison of the population activity for trained RNN and those from monkey experiments will provide more information and justification of the model.

Our initial submission included the neural population preparatory states, their stability over learning, and the structure of preparatory shifts over learning and washout (Figures 3a, b, e, g). To address your review for further comparisons, we added 2 figures of the PCA preparatory states across learning and washout (Appendix Figure 3 & 4) demonstrating stability in the force-predictive subspace, and one figure of the uniform shift along the “learning axis” and PC1 (Appendix Figure 5).

Could the model reproduce observation in the task variants including sequentially different or interfering curl field as in monkey experiments?

We agree that these are interesting experiments. For our manuscript, we chose to focus on the behavioral constraints and replication of key neural and behavioral phenomena during motor learning over learning, unlearning, and relearning of a curl field. We believe this paper forms a foundation to build upon to reproduce other experimental observations in later work.

The environment in the model is based on gym and the algorithm was an extension of PPO. But some details are not clear from the method. For example, the input dim is denoted either 9D (Line 145) or 10D (Line 183). It would be much better to provide the code for evaluation and reproducibility.

We apologize; the 10D was a typo. We will carefully correct any errors and clarify details in the Methods. We will post the code with the camera-ready submission.

Thank you for your comments and feedback. We fixed the typos you pointed out. We believe the new neural population activity figures will help clarify the relationship between the model and the monkey’s neural population activity. We hope these changes addressed your concerns.

The authors have partially addressed my questions; however, two were only acknowledged without being fully addressed. The first pertains to reproducing neural population trajectories in the PCA space (Extended Fig. 9g in Sun et al., 2022, Nature), which serves as a valuable benchmark for comparing the trained RNN with neural population data. The second concerns comparing model predictions with task variants, such as sequentially different or interfering curl fields, as investigated in the monkey experiments in Sun et al., 2022, Nature. Both questions remain unresolved and require further clarification.

Thank you for your review! We will address the weaknesses first, followed by the questions.

The proposed model is pretty simple, almost a direct reuse of existing models. The mathematical introduction of the model is not very clear, mainly because many variables were not explained [...]

Thank you for this comment. We updated the methods section to clearly explain all variables. We will also clarify that, for the purposes of computational and systems neuroscience, we actually view model simplicity as a strength! It is desirable to have as simple a model that captures neural features as possible, since model complexities also complicate biological comparisons. The comparison did not discuss the impact of curriculum learning, which might not be used in the experiment with a monkey. We were empirically unable to train agents to perform this task without curriculum learning. We will state clearly that curriculum learning significantly facilitated training: A randomly initialized model struggled to learn delayed reaches without a curriculum if the acceptance window was too small, likely due to very sparse rewards.

We wish to point out that, generally, to train monkeys to perform a delayed reach task, a curriculum is used. Naive monkeys initially train to perform delayed reach tasks with larger targets, and the target is progressively shrunk until it can perform a delayed reach task with a smaller target. Although there are differences, such as that the monkey due to prior experience can make physical delayed reach tasks and this training is therefore to cognitively teach him how to perform this, it does also point to the fact that a monkey is not born with the ability to perform precise and timed reaches based on cues, but learns it through experience. We will mention this in the manuscript.

The paper mentioned "brain-like neural geometry", while the evidence is not solid. At least a method to compare the geometry should be suggested and treat proof of "brain-like neural geometry" as a future work.

With respect to comparing the geometry, we performed these comparisons (stated in Section 3.3): (1) preparatory state geometry, (2) stability of the force-predictive subspace, (3) uniformness of preparatory activity shifts after CF learning, and (4) orthogonality of learning-washout shifts.

The first 2 are qualitative, and the last 2 are quantitative, in that they involve either finding statistically insignificant difference in shift for target near to or far from the trained target (3), or finding that the orthogonal shifts are statistically significantly different (4). Certainly, this can be abstracted and quantified further in future works. We have added such a suggestion to the discussion!

Now we will address the questions. First, thank you for the style/typo suggestions. We have fixed the corresponding text in the updated manuscript.

What if the KL penalty terms were not used? A comparison with and without them can support the claim in Line 496 better.

Please see our additional analyses (Appendix Figure 6) to show these KL terms increase similarity to the animal’s trajectories. First, the trajectories with neither constraints typically fail to converge to a reasonable trajectory. Appendix Fig. 6a (top left) shows exemplar learned trajectories that curve in the direction of the curl field after significant learning (red curves). This clearly does not match the biological behavior: the monkey’s curl field reaches curve in the opposite direction. Second, the desirable overcompensation of the curl field becomes more consistent, but excessive, with the zeroness constraint alone. As described in Fig 2f, it is optimal to “overcompensate” in a curl field to speed up the reach. Smoothness-only constrained models take advantage of this mechanism and exhibit obvious overcompensation beyond what is observed in biophysical reaches (Appendix Figure 6c). Finally, the smoothness constraint helps to reduce curvature to better match the curvature of the monkey’s reach. Combining both constraints produces more consistent and monkey-like post-learning trajectories.

Is there any model of the monkey's arm in the implemented environment? Or the agent just conceptially moving a handle?

There is no model of the monkey’s arm in our model. We only model the end effector kinematics.

Line 233. No GPU is used? Is there any parallel computing with the CPU?

No GPU is used. Yes, we used vectorized environments, and ran 16 processes/environments in parallel.

Figure 3. How these dots are located? Define the learning axis and washout axis.

The learning and washout axis definitions are in the text at lines 403 and 419. We will reorganized some of the text to make this definition easier to find.

Where are Figures S1 to S4?

We have fixed the supplementary figures to start at S1 instead of S5. We have renamed them as Appendix Figures.

Thank you for your review! We address the weaknesses and question below.

There are very few baseline comparisons - which parts of the framework (loss / architecture / training) are necessary to see their main results? The paper makes reference to Richards et al., 2019, but does not adequately address this. The networks trained with [RL] are not compared with networks that are trained with supervised learning, which is more common in this field.

Regarding the loss and architecture, we will clearly state the RNN (rather than any feedforward architecture) is necessary because we are concerned about neural dynamics. We will also provide further detail that we were only successful training an LSTM (and were unable to train a vanilla RNN). For the reward, we directly mimicked the monkey’s juice reward for a successful target reach in order to maximize ecological validity. We recognize it is possible to perform reward shaping, but we view this as introducing unnecessary complexity for our goal of analyzing the effects of prior experience on neural activity. We will state that future work may explore the interesting question of how modifying rewards and even more complex RNN architectures may affect results.

Regarding training, we are unaware of any other baseline for this motor learning task. Although supervised learning is appropriate for other tasks studying an already learned behavior, this task is unique in that it studies neural population activity during learning dynamics. In other studies in this field, we typically study the post-learning RNNs, but here, the difference in pre- and post-learning states are vital. Supervised learning would introduce many undesired complexities. For example, to match monkey learning, since the output is acceleration, we would need to define the entire target acceleration profile under a curl field. This is not realistic as the solution is known to the agent a priori through gradients that point to optimal behavior, rather than learning it through exposure (where we aim to study the emergent learning dynamics). It would also preclude analysis of degree of trajectory curvature under a curl field, since we would have already defined the final behavior supervision. This makes the supervised baseline unnatural for our task. We will further emphasize this in the manuscript.

Finally, as it relates to Richards et al., we emphasize that their manuscript also stated “it will also be crucial to the development of a deep learning framework for systems neuroscience to focus on how a given animal might solve tasks in its appropriate Brain Set.” Our results serve as a piece of evidence that, with the loss, architecture, and learning rules fixed, the addition of prior experience (expansion of the Brain (task) Set to more than just movements performed or learned within an experiment, to those possibly seen during prior experience), can induce more brain-like learning activity.

unclear which [acceleration profiles] are more akin to the experiments since only a qualitative comparison is performed (and it seems like 'neither' is the most similar [...]

To clarify the importance of the KL terms, we performed additional analyses (Appendix Figure 6) to show these KL terms increase similarity to the animal’s trajectories during learning of a curl field. First, the KL terms induce a better match to the monkey’s overcompensation behavior. As described in Fig 2f, it is optimal to “overcompensate” in a curl field to speed up the reach. Non-regularized models do not converge reliably. Smoothness-constraint models overly exploit this mechanism and exhibit overcompensation beyond biophysical reaches, while models constrained by both converge reliably while utilizing a level of overcompensation comparable to the monkey’s. We will provide quantitative statistics for this result in the camera-ready submission. difference between learning and relearning are not quantified adequately in Figure 2g. Is there a significant difference? We added S.E.M bars to Fig 2g. Relearning converges faster than learning at epoch 40 with (p<0.05).

All the comparisons [...] are qualitative; there is a lack of quantitative comparisons throughout the paper.

In addition to the previous reply, we also included statistics for the uniform shifts (See Fig. 3g).

The effect of the curl field seems to be quite small in the networks in Figure 3b, as compared to the monkey in Figure 3c.

The main neural effect of learning the curl field is apparent when examining preparatory state shifts outside of the PCA/TDR plane. This shift can be seen and compared to the shifts in monkey neural activity in the new Appendix Figure 5.

While the authors examine the organization of preparatory states using TDR, previous [...] studies have also analyzed this using PCA [...] Does this structure only emerge using TDR or also PCA?

The structure of preparatory states is organized similarly in the PCA space, see our new Appendix Figures 3, 4, and 5.