Massively Parallel Imitation Learning of Mouse Forelimb Musculoskeletal Reaching Dynamics

The brain has evolved to effectively control the body, and in order to understand the relationship we need to model the sensorimotor transformations underlying embodied control. As part of a coordinated effort, we are developing a general-purpose platform for data-driven simulation modeling high fidelity behavioral dynamics, biomechanics, and neural circuit architectures underlying embodied control. We present a pipeline for taking kinematics data from the neuroscience lab and creating a pipeline for recapitulating those natural movements in physics simulation. We implement an imitation learning framework to simulate a dexterous forelimb reaching task with a musculoskeletal model in the Mujoco physics environment. The imitation learning model is currently training at more than 1 million training steps per second due to GPU acceleration with JAX and Mujoco-MJX. We present results that indicate that adding naturalistic constraints on control magnitude lead to simulated muscle activity that better predicts real EMG signals. This work provides evidence to suggest that control constraints are critical to modeling biological movement control.