ePoster

A GPU-Accelerated Deep Reinforcement Learning Pipeline for Simulating Animal Behavior

Charles Zhang, Elliott Abe, Jason Foat, Bing Brunton, Talmo Pereira, Bence Olveczky, Emil Warnberg
COSYNE 2025(2025)
Montreal, Canada

Conference

COSYNE 2025

Montreal, Canada

Resources

Authors & Affiliations

Charles Zhang, Elliott Abe, Jason Foat, Bing Brunton, Talmo Pereira, Bence Olveczky, Emil Warnberg

Abstract

To study how motor circuits control bodies in robust, adaptive, and flexible ways, we often break down the complex control system into individual components to be studied in isolation. However, it is increasingly recognized that we need to consider biomechanically realistic models when interpreting the neural computations underlying movement and motor learning. Training an artificial neural network (ANN) to control a biomechanical model and have it reproduce the behavior of real animals allows the probing of neural control of movement in a fully transparent and configurable system. To make this approach accessible to the broader neuroscience community, we present a two-step pipeline that leverages advances in GPU-accelerated physics simulation : 1) stac-mjx, a fast, animal-agnostic Python package for performing skeletal model registration and keypoint retargeting, converting arbitrary 3D keypoint data into joint angles for any corresponding biomechanical model. 2) tracking-mjx, an RL framework with extensible environments, training scripts, and logging for GPU-accelerated training of animal behavior imitation. Our training pipeline is many times faster than previous implementations that require massively distributed training on thousands of CPUs. It supports the addition of new tasks, neural network architectures, and model organisms to facilitate rapid iteration of hypothesis generation and testing.

Unique ID: cosyne-25/gpu-accelerated-deep-reinforcement-28dfa0e4