NeuroGym is a curated collection of neuroscience tasks with a common interface. The goal is to facilitate the training of neural network models on neuroscience tasks.
| Badges | |
|---|---|
| fairness |  |
| package |  |
| docs |    |
| tests |       |
| running on |  |
| license |  |
NeuroGym is a curated collection of neuroscience tasks with a common interface. The goal is to facilitate the training of neural network models on neuroscience tasks.
NeuroGym inherits from the machine learning toolkit Gymnasium, a maintained fork of OpenAI’s Gym library. It allows a wide range of well established machine learning algorithms to be easily trained on behavioral paradigms relevant for the neuroscience community. NeuroGym also incorporates several properties and functions (e.g. continuous-time and trial-based tasks) that are important for neuroscience applications. The toolkit also includes various modifier functions that allow easy configuration of new tasks.
Please see our extended project documentation for additional details.
Create and activate a virtual environment to install the current package, e.g. using conda (please refer to their site for questions about creating the environment):
conda activate # ensures you are in the base environment
conda create -n neurogym python=3.11 -y
conda activate neurogym
Install the latest stable release of neurogym using pip:
pip install neurogym
NeuroGym includes optional reinforcement learning (RL) features via Stable-Baselines3. To install these, choose one of the two options below depending on your hardware setup:
pip install neurogym[rl]
NOTE for Linux/WSL users: If you do not have access to a CUDA-capable NVIDIA GPU (which is the case for most users), above line will install up to 1.5GB of unnecessary GPU libraries. To avoid excessive overhead, we recommend first isntalling the CPU-only version of PyTorch:
pip install torch --index-url https://download.pytorch.org/whl/cpu
pip install neurogym[rl]
To contribute to NeuroGym or run it from source with live code updates:
git clone https://github.com/neurogym/neurogym.git
cd neurogym
pip install -e .
This installs the package in editable mode, so changes in source files are reflected without reinstalling.
To include both RL and development tools (e.g., for testing, linting, documentation):
pip install -e .[rl,dev]
NOTE: psycohopy installation is currently not working
If you need psychopy for your project, additionally run
pip install psychopy
Currently implemented tasks can be found here.
Wrappers (see their docs) are short scripts that allow introducing modifications the original tasks. For instance, the Random Dots Motion task can be transformed into a reaction time task by passing it through the reaction_time wrapper. Alternatively, the combine wrapper allows training an agent in two different tasks simultaneously.
🧪 Beta Feature — The configuration system is optional and currently under development. You can still instantiate environments, agents, and wrappers with direct parameters.
It is only used in a small portion of the codebase and is not required for typical usage.
See the demo.ipynb notebook for the only current example of this system in action.
NeuroGym includes a flexible configuration mechanism using Pydantic Settings, allowing configuration via TOML files, Python objects, or plain dictionaries.
Using a TOML file can be especially useful for sharing experiment configurations in a portable way (e.g., sending config.toml to a colleague), reliably saving and loading experiment setups, and easily switching between multiple configurations for the same environment by changing just one line of code. While the system isn't at that stage yet, these are intended future capabilities.
Create a config.toml file (see template) and load it:
from neurogym import Config
config = Config('path/to/config.toml')
You can then pass this config to any component that supports it:
from neurogym.wrappers import monitor
env = gym.make('GoNogo-v0')
env = monitor.Monitor(env, config=config)
Or directly pass the path:
env = monitor.Monitor(env, config='path/to/config.toml')
from neurogym import Config
config = Config(
local_dir="logs/",
env={"name": "GoNogo-v0"},
monitor={"name": "MyMonitor"}
)
from neurogym import Config
config_dict = {
"env": {"name": "GoNogo-v0"},
"monitor": {
"name": "MyMonitor",
"plot": {"trigger": "step", "value": 500, "create": True}
},
"local_dir": "./outputs"
}
config = Config.model_validate(config_dict)
NeuroGym is compatible with most packages that use gymnasium. In this example jupyter notebook we show how to train a neural network with RL algorithms using the Stable-Baselines3 toolbox.
Creating custom new tasks should be easy. You can contribute tasks using the regular gymnasium format. If your task has a trial/period structure, this template provides the basic structure that we recommend a task to have:
from gymnasium import spaces
import neurogym as ngym
class YourTask(ngym.PeriodEnv):
metadata = {}
def __init__(self, dt=100, timing=None, extra_input_param=None):
super().__init__(dt=dt)
def new_trial(self, **kwargs):
"""
new_trial() is called when a trial ends to generate the next trial.
Here you have to set:
The trial periods: fixation, stimulus...
Optionally, you can set:
The ground truth: the correct answer for the created trial.
"""
def _step(self, action):
"""
_step receives an action and returns:
a new observation, obs
reward associated with the action, reward
a boolean variable indicating whether the experiment has terminated, terminated
See more at https://gymnasium.farama.org/tutorials/gymnasium_basics/handling_time_limits/#termination
a boolean variable indicating whether the experiment has been truncated, truncated
See more at https://gymnasium.farama.org/tutorials/gymnasium_basics/handling_time_limits/#truncation
a dictionary with extra information:
ground truth correct response, info['gt']
boolean indicating the end of the trial, info['new_trial']
"""
return obs, reward, terminated, truncated, {'new_trial': new_trial, 'gt': gt}
For the authors of the package, please refer to the zenodo DOI at the top of the page.
Training Artificial Neural Networks to Uncover Behavioral Strategies in neuroscience
ANNUBeS is a deep learning framework meant to generate synthetic data and train on them neural networks aimed at developing and evaluating animals' training protocols in neuroscience.