Feedbax is a JAX library for optimal feedback control with neural networks.
Feedbax makes it easy to:
- train a neural network to control a simulated limb (biomechanical model) to perform movement tasks;
- intervene on existing models and tasks—for example, to:
- add force fields that disturb a limb;
- alter the activity of a single unit in a neural network;
- perturb the sensory feedback received by a network;
- add any kind of noise to any part of a model's state;
- schedule an intervention to occur on only a subset of task trials or time steps;
- specify which parts of the model are trainable, and which states are available as sensory feedback;
- train multiple replicates of a model at once;
- swap out components of models, and write new components.
Feedbax is in active development. Expect some changes in the near future. In particular, I intend to replace the staged model approach with a simpler lazy-executing DAG framework.
Feedbax uses JAX and Equinox.
Please also check out:
- MotorNet, a PyTorch library with many similarities to Feedbax.
- Collimator, a more mature JAX library for composing and optimizing dynamical systems.
pip install feedbax
Currently requires Python>=3.11.
For best performance, install JAX with GPU support.
Documentation is available here.
I started to develop Feedbax while learning JAX. My short-term objective has been to support my own use cases—graduate research in the neuroscience of motor control—but I've also tried to design something reusable and general.
I've added GitHub issues to document some of my choices and uncertainties. For an overview of major issues in different categories, check out this GitHub conversation. Refer also to this page of the docs, for an informal overview of how Feedbax objects relate to each other.
There are many features, especially pre-built models and tasks, that could still be implemented. Some of the models and tasks that are implemented have yet to be fully optimized. So far I've focused more on the overall structure, than on coverage of all the common use cases I can imagine. If there's a particular model, task, or feature you'd like Feedbax to support, let us know, or contribute some code!
- Thanks to my PhD supervisor Gunnar Blohm and to the rest of our lab, as well as to Dominik Endres and Stephen H. Scott for discussions that have directly influenced this project
- Special thanks to Patrick Kidger, whose JAX libraries and their documentation often serve as examples to me