Skip to content

Getting started¤

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!

What Feedbax can (and can't) do

Feedbax is designed for feedback control of biomechanical models by neural networks, to perform movement tasks in continuous spaces. However, it could also be used for other optimal control problems for which a suitable set of model, cost function, and task trials can be defined.

Feedbax is not designed for finding algebraic solutions to the classic formalisms of optimal control theory, such as the Linear-Quadratic-Gaussian (LQG) controller. Similarly, Feedbax does not enforce the separation principle between model components—that is, models don't need to include distinct components (e.g. neural network layers) for the controller and the observer (i.e. state estimator).

Feedbax is a JAX library¤

Feedbax uses JAX and Equinox.

Never used JAX before?

Please also check out MotorNet, a PyTorch library with many similarities to Feedbax.

Installation¤

pip install feedbax

Currently requires Python>=3.11.

For best performance, install JAX with GPU support.

Development¤

I've developed Feedbax over the last few months, 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!

Acknowledgments¤

  • 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