Skip to main content

Optimal feedback control + interventions in JAX.

Project description

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 currently in active development. Expect some changes in the near future!

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.

Documentation

Documentation is available here.

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

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

feedbax-0.1.2.tar.gz (77.8 kB view details)

Uploaded Source

Built Distribution

feedbax-0.1.2-py3-none-any.whl (84.4 kB view details)

Uploaded Python 3

File details

Details for the file feedbax-0.1.2.tar.gz.

File metadata

  • Download URL: feedbax-0.1.2.tar.gz
  • Upload date:
  • Size: 77.8 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/5.0.0 CPython/3.11.0

File hashes

Hashes for feedbax-0.1.2.tar.gz
Algorithm Hash digest
SHA256 470a7dd6332e5f53ba3974c45d678d2a85c21f998f625d3243f71fae64d633cd
MD5 ad9be5a8e11a382dd4089e4ef69c11dc
BLAKE2b-256 b0daa0a502f599bc609966ce005dcc8abad88feb2fcfc528d661f766533b0175

See more details on using hashes here.

File details

Details for the file feedbax-0.1.2-py3-none-any.whl.

File metadata

  • Download URL: feedbax-0.1.2-py3-none-any.whl
  • Upload date:
  • Size: 84.4 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/5.0.0 CPython/3.11.0

File hashes

Hashes for feedbax-0.1.2-py3-none-any.whl
Algorithm Hash digest
SHA256 af785aadab5e86c20bec4176d859fc19d47587453cfe0c2fbcf46d137c6036a0
MD5 7d1eb72a34e86f2a50a96251d97ee8ad
BLAKE2b-256 19643eea7106ffa6cb359ebc528d682dc226bb35b3c33b10fbe40d805de835e1

See more details on using hashes here.

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page