Massively Parallel Imitation Learning of Mouse Forelimb Musculoskeletal Reaching Dynamics

TL;DR

This work presents MIMIC-MJX, a high-throughput neuromechanical platform that combines STAC-MJX pose-to-kinematics registration with Track-MJX imitation learning to reproduce mouse forelimb reaching in a musculoskeletal model. It demonstrates that physics-aware constraints on energy and control improve alignment between simulated muscle activations and in vivo EMG, and it leverages nonlinear dynamical forecasting to relate joint kinematics, muscle activity, and EMG. The results show accurate kinematic reproduction, scalable learning speed, and meaningful latent-space structure, suggesting a viable route for probing sensorimotor control under realistic biomechanics. The study also discusses limitations and future directions, including multiple targets and more physiologically grounded energetics, to enhance generalizability and interpretability of neuromuscular control strategies.

Abstract

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 behavior-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 a biomechanical model. We implement a imitation learning framework to perform a dexterous forelimb reaching task with a musculoskeletal model in a simulated physics environment. The mouse arm model is currently training at faster 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 energy and velocity lead to simulated musculoskeletal activity that better predict real EMG signals. This work provides evidence to suggest that energy and control constraints are critical to modeling musculoskeletal motor control.

Figures (4)

• Figure 1: A. Neural network encoder-decoder architecture for imitation learning to reproduce motion capture trajectories. B. Raw video data from mouse water reaching task with 3D multi-camera pose estimation is registered to the body model using STAC-MJX. C. Kinematics of the motion capture registered model and the imitation learning performance labeled as "track replay". D. Registration error and tracking error relative to original 3D pose data. E. Joint reward throughout training.
• Figure 2: A. Joint reward and high frequency activity by control cost parameter sweep, each point reflects the mean value across five independent random seeds, and the error bars denote the 95% confidence interval computed across those seeds. B. Mean absolute error between EMG and simulated muscle activity for joint reward only and physics aware models. C. Trial-by-Trial muscle activation comparison between EMG, joint reward only and physics-aware constraints. D. Average muscle activation over time for EMG, joint reward only and physics-aware constraints with standard error of the mean in shaded region. E-H. 3D PCA embeddings of mouse arm reaching trajectories, using representations from the intention space (E) and sequential decoder layers (F-H). The axes show % variance explained for each PC. The colormap visualizes the extent of the shoulder extension.
• Figure 3: A. Prediction performance (simplex $\rho$) using joint angles to predict simulated biceps and triceps activations. B. Prediction performance (simplex $\rho$) using reference joint angles and simulated actions to predict Real EMG signals. C. Predictions and observed values using joint angles to predict simulated biceps activity ($\rho$ = .802). D. Predictions and observed values using joint angles and simulated actions to predict biceps EMG ($\rho$ = .328). E. Predictions and observed values using joint angles to predict simulated triceps activity ($\rho$ = .789). F. Predictions and observed values using joint angles and simulated actions to predict triceps EMG ($\rho$ = .7).
• Figure 4: PyEDM parameter search. Embedding dimension and tau sweep with performance measured by simplex rho for the simulated muscle activity and EMG signals find an optimal tau of -1 and optimal embedding dimension of 3. Prediction horizon search for the simulated muscle activity and EMG signals show the presence of a nonlinearity.