A latent control model for realistic rodent motion

AI summary

Overview

Research area: Neuro-AI / embodied artificial intelligence — specifically physics-based control of biomechanically realistic virtual animals (a virtual rat) using imitation learning and reinforcement learning.

Technical level: Intermediate. The paper combines variational latent-variable models, distillation, and reinforcement learning; the core ideas are accessible, but the formal analysis (KL decompositions, information bottleneck) assumes comfort with probabilistic modeling.

Scope (1 sentence): The paper introduces SCAMPER, a four-stage training pipeline that gives a virtual rat both goal-directed control and diverse, naturalistic random exploration by combining a motion prior with a second "behavior prior."

What This Paper Is About

Researchers want embodied virtual animals whose movements actually look like real animals, so that artificial agents can be compared directly against biological ones. Existing systems either produce unnatural motion when trained with reinforcement learning, are locked into pure imitation, or — in the case of the PULSE framework that SCAMPER builds on — sample random motion that stays stuck in a single behavioral mode. SCAMPER's goal is a single controller that moves naturally both when pursuing a task goal and when exploring randomly, which is what sparse-reward neuroscience tasks demand.

Key Contributions

1. A dual-prior architecture for embodied rodent control. SCAMPER (Simulated Control via Augmented Motor Priors for Embodied Reinforcement learning) pairs a motion prior — used for residual goal-directed control — with a newly added behavior prior used for exploration, both conditioning on proprioception and both feeding a shared frozen decoder.

2. A formal account of why two priors are needed. Two propositions show that when the encoder and motion prior are trained jointly, the KL objective decomposes into a conditional mutual information term plus a marginal-posterior matching term, both weighted by the same hyperparameter λ, forcing a rate–distortion trade-off. Freezing the encoder removes the mutual-information term from the optimization, letting the behavior prior match the marginal posterior much more closely.

3. A learned gate between exploration and exploitation. A task policy outputs a continuous gating variable β_t ∈ [−1, 1], binarized by hysteresis thresholding into b_t ∈ {0, 1}, which selects at each timestep between a residually corrected motion-prior code and a stochastic sample from the behavior prior.

4. Empirical validation across locomotion, rearing, and sparse-reward foraging, with comparisons against MIMIC-MJX, a from-scratch RL model, and a motion-prior-only variant.

Main Findings

• The behavior prior is measurably more diverse. Across all sampling noise levels, the Vendi score (effective number of distinct behavioral modes) was significantly higher for the behavior prior than the motion prior. Motion-prior sampling scored below the training data for all σ < 1; the behavior prior matched training-data diversity at low noise and exceeded it at high noise. This was measured over 842 rollouts per noise level, each 5 s of simulated time.

• Diversity comes with travel and coverage. Motion-prior rollouts produced shorter net displacement than the training data at almost all noise levels, while behavior-prior rollouts roughly matched training-data distance at low noise and exceeded it at high noise. In spatial occupancy plots at σ = 0.22, motion-prior rollouts stayed largely inside the circular region occupied during training despite no physical barrier, whereas behavior-prior rollouts explored beyond it.

• Realism is preserved under stochastic sampling. Kernel Inception Distance (KID) was lower for the behavior prior than the motion prior at all noise levels, using a VAE trained on training-data joint angle trajectories as the inception model. The behavior prior also better matched the training data's rearing-height distribution, confirmed with Wasserstein distance on the rearing curves.

• The KL gap is quantified. The behavior prior reached a KL divergence of 5.88 nats from the encoder, versus 14.35 nats for the motion prior.

• Realistic gaits emerge at speeds absent from the data. On a dense-reward speed-tracking task, SCAMPER closely matched target speeds, while MIMIC-MJX failed to generalize beyond speeds in the training data and could not match targets above 0.5 m/s. The from-scratch model matched speeds but with unnatural motion and higher variance. SCAMPER produced walking, trotting, and galloping — even though the training data contained no sustained or high-speed locomotion. Gait analysis across 18 metrics showed SCAMPER closest to empirical rat data (e.g., "All" distance of 0.59 for walk, 1.19 for trot, 1.91 for gallop) versus MIMIC-MJX (1.28, 2.72, 2.74) and the from-scratch model (3.44, 11.01, 5.91).

• Rearing transfers with stability and precision. On a dense-reward rearing task with the target set at the maximum rearing height in the training data, only the SCAMPER policy solved the task consistently within the training budget. MIMIC-MJX rarely held the target long enough; the from-scratch model never reached it.

• Only the full dual-prior model solves sparse-reward foraging. In an arena foraging task with treats spawning outside a minimum radius and a Gaussian scent field, all models found treats at small spawn radii. Only full SCAMPER — able to switch to behavior-prior sampling — solved the task when the minimum spawn radius was large and extensive exploration was required. The motion-prior-only variant could not.

• Trade-off noted between the priors. The behavior prior is less stable under residual control than the motion prior, which is precisely what motivates using both rather than replacing one with the other.

Methodology in Plain English

SCAMPER trains in four stages, all using multi-layer perceptrons:

1. Expert imitator. A policy is trained with PPO to reproduce motion-capture trajectories that were registered onto a biomechanical rat model via inverse kinematics using STAC-MJX. The rat's observations are proprioceptive (joint angles, angular velocities, joint torques, body height, orientation, end-effector positions, kinematic sensors, touch sensors, previous action), and its goal is a set of future displacement targets; during imitation it sees the next five goal states.

2. Online distillation. The expert is distilled into a latent controller with three jointly trained pieces — an encoder that sees state and goal, a decoder that maps state and latent code to actions, and a proprioception-conditioned motion prior. Training rolls out the student in simulation while querying the frozen expert on visited states (DAgger), minimizing an action-matching loss, a KL term pulling the encoder toward the motion prior, and a temporal-smoothness term on latent means. This is the PULSE recipe, and the authors report it is more stable than training the latent controller from scratch with RL.

3. Behavior prior distillation. With everything else frozen, a second prior is trained to match the encoder's outputs by KL divergence — a pure marginal-posterior-fitting objective, free of the information-bottleneck pressure.

4. Transfer learning. A task policy, seeing task observations instead of goals, outputs a residual correction to the motion prior's mean plus the gate β_t. When the gate is "on," the system exploits a corrected motion-prior code; when "off," it samples a noisy code from the behavior prior. The policy learns when to do which. Everything downstream is trained with PPO on torque actions.

For evaluation they used the Vendi score for diversity, KID (with a VAE inception model trained on joint-angle trajectories) for realism, Wasserstein distance for rearing distributions, and standard gait metrics for comparing limb coordination against published rat data.

Why This Matters

Impact on research. SCAMPER closes a gap the authors identify explicitly: no prior embodied animal model could both mimic real motion and generate diverse samples suitable for exploration. Because most experimental neuroscience tasks offer only sparse extrinsic reward, a model that cannot explore cannot be trained on them — and a model that explores unnaturally would confound comparisons with neural recordings. The paper builds directly on the finding that a purely imitative virtual rat's internal representations matched rat motor cortex and dorsolateral striatum, extending that style of comparison to goal-directed tasks where model and animal train on the same task.

Real-world applications:

• Computational neuroscience: testing hypotheses about how animals balance exploration and exploitation, and comparing model learning trajectories against neural data.
• Behavioral phenotyping: generating synthetic but realistic movement distributions to benchmark pose-estimation and behavioral-analysis pipelines.
• Motor control and biomechanics: studying how gait transitions depend on body morphology and speed, using a body that can be altered and re-trained.
• Robotics and animation: the residual-control-plus-variational-prior recipe transfers to any physics-based character that needs both task performance and natural, sampleable idle behavior.

Industry relevance. Organizations building simulated agents, character animation, or legged-robot controllers face the same core problem — pretrained natural motion that does not collapse into a single mode when sampled, and that remains controllable downstream. The dual-prior and learned-gate design is a general pattern for that. The work was conducted with authors at Google, Mila, McGill, Harvard, and the Salk Institute, and one author is a Google Research Scientist.

Future Directions

• Scaling the behavior dataset. The training data was only 1 hour of a rodent freely behaving in an open arena, containing just walking, rearing, sitting still, and grooming. Richer datasets with more diverse behaviors, longer locomotion bouts at varied speeds, and more varied landscapes would broaden the reachable repertoire.

• Adding multimodal sensing. SCAMPER currently uses only proprioception plus a simple scent proxy — the authors note this effectively makes the rat blind and deaf, and call it an increasingly pressing limitation.

• Muscle-based actuation. The model uses torque actuators rather than muscles. The authors state SCAMPER is agnostic to actuator type, so it could be applied to muscle-driven bodies as they are developed.

• Better stability for the behavior prior under residual control. The behavior prior is less stable under residual control than the motion prior (reported in the appendix); improving this could allow finer-grained control during exploratory modes.

Target Audience

Neuro-AI researchers building embodied models of animal behavior; reinforcement learning researchers working on physics-based character control and motion priors; computational and systems neuroscientists who want models trainable on the sparse-reward tasks used in experiments; and robotics or graphics practitioners interested in pretrained natural motion that remains both controllable and sampleable. Readers without a background in variational inference will find the empirical sections accessible but the Section 3.3 analysis dense.

Authors’ abstract

The development of embodied models of animal behavior is a major goal of the neuro-AI community. Embodied models will enhance our ability to make direct comparisons between biological and artificial agents. To make such comparisons effective, though, we need embodied models that actually move like animals. Previous work has developed virtual animal bodies and techniques for training realistic motion through imitation. However, to date, there is no system for embodied animal modeling that produces natural movements both during trained goal-directed behavior, and during exploration driven by random generation. Here, we introduce Simulated Control via Augmented Motor Priors for Embodied Reinforcement learning (SCAMPER), a system for training an embodied virtual rat model, enabling natural movements during both goal-directed behavior and random exploration. SCAMPER uses a variational architecture with two distinct priors, a motion prior for control and a behavior prior for exploration. We show that combining both of these priors allows SCAMPER to generate realistic behavior during exploration and exploit the task structure effectively. SCAMPER can also learn from sparse rewards as a result of the diverse movements induced by the behavior prior, allowing the system to model realistic sparse-reward tasks from experimental neuroscience. Altogether, SCAMPER provides the neuro-AI community with an effective new system for training embodied models of rodents on realistic tasks.