• Stars
    star
    1,465
  • Rank 30,929 (Top 0.7 %)
  • Language
    Python
  • License
    Apache License 2.0
  • Created almost 5 years ago
  • Updated about 1 month ago

Reviews

There are no reviews yet. Be the first to send feedback to the community and the maintainers!

Repository Details

bsuite is a collection of carefully-designed experiments that investigate core capabilities of a reinforcement learning (RL) agent

Behaviour Suite for Reinforcement Learning (bsuite)

PyPI Python version PyPI version pytest

radar plot

Introduction

bsuite is a collection of carefully-designed experiments that investigate core capabilities of a reinforcement learning (RL) agent with two main objectives.

  1. To collect clear, informative and scalable problems that capture key issues in the design of efficient and general learning algorithms.
  2. To study agent behavior through their performance on these shared benchmarks.

This library automates evaluation and analysis of any agent on these benchmarks. It serves to facilitate reproducible, and accessible, research on the core issues in RL, and ultimately the design of superior learning algorithms.

Going forward, we hope to incorporate more excellent experiments from the research community, and commit to a periodic review of the experiments from a committee of prominent researchers.

For a more comprehensive overview, see the accompanying paper.

Technical overview

bsuite is a collection of experiments, defined in the experiments subdirectory. Each subdirectory corresponds to one experiment and contains:

  • A file defining an RL environment, which may be configurable to provide different levels of difficulty or different random seeds (for example).
  • A sequence of keyword arguments for this environment, defined in the SETTINGS variable found in the experiment's sweep.py file.
  • A file analysis.py defining plots used in the provided Jupyter notebook.

bsuite works by logging results from "within" each environment, when loading environment via a load_and_record* function. This means any experiment will automatically output data in the correct format for analysis using the notebook, without any constraints on the structure of agents or algorithms.

We collate all of the results and analysis in a pre-made jupyter notebook bit.ly/bsuite-colab.

Getting started

If you are new to bsuite you can get started in our colab tutorial. This Jupyter notebook is hosted with a free cloud server, so you can start coding right away without installing anything on your machine. After this, you can follow the instructions below to get bsuite running on your local machine.

Installation

We have tested bsuite on Python 3.6 & 3.7. To install the dependencies:

  1. Optional: We recommend using a Python virtual environment to manage your dependencies, so as not to clobber your system installation:

    python3 -m venv bsuite
    source bsuite/bin/activate
    pip install --upgrade pip setuptools
  2. Install bsuite directly from PyPI:

    pip install bsuite
  3. Optional: To also install dependencies for the baselines examples (excluding OpenAI and Dopamine examples), run:

    pip install bsuite[baselines]

Environments

Complete descriptions of each environment and their corresponding experiments are found in the analysis/results.ipynb Jupyter notebook.

These environments all have small observation sizes, allowing for reasonable performance with a small network on a CPU.

Loading an environment

Environments are specified by a bsuite_id string, for example "deep_sea/7". This string denotes the experiment and the (index of the) environment settings to use, as described in the technical overview section.

For a full description of each environment and its corresponding experiment settings, see the paper.

import bsuite

env = bsuite.load_from_id('catch/0')

The sequence of bsuite_ids required to run all experiments can be accessed programmatically via:

from bsuite import sweep

sweep.SWEEP

This module also contains bsuite_ids for each experiment individually via uppercase constants corresponding to the experiment name, for example:

sweep.DEEP_SEA
sweep.DISCOUNTING_CHAIN

In addition, sequences of bsuite_ids with the same tag can be loaded via:

from bsuite import sweep

sweep.TAGS

The TAGS variable groups bsuite environments together by their underlying tag, so all the basic tasks or scale tasks can be loaded with:

sweep.TAGS['basic']
sweep.TAGS['scale']

Loading an environment with logging included

We include two implementations of automatic logging, available via:

We also include a terminal logger in logging/terminal_logging.py, exposed via bsuite.load_and_record_to_terminal.

It is easy to write your own logging mechanism, if you need to save results to a different storage system. See the CSV implementation for the simplest reference.

Interacting with an environment

Our environments implement the Python interface defined in dm_env.

More specifically, all our environments accept a discrete, zero-based integer action (or equivalently, a scalar numpy array with shape ()).

To determine the number of actions for a specific environment, use

num_actions = env.action_spec().num_values

Each environment returns observations in the form of a numpy array.

We also expose a bsuite_num_episodes property for each environment in bsuite. This allows users to run exactly the number of episodes required for bsuite's analysis, which may vary between environments used in different experiments.

Example run loop for a hypothetical agent with a step() method.

for _ in range(env.bsuite_num_episodes):
  timestep = env.reset()
  while not timestep.last():
    action = agent.step(timestep)
    timestep = env.step(action)
  agent.step(timestep)

Using bsuite in 'OpenAI Gym' format

To use bsuite with a codebase that uses the OpenAI Gym interface, use the GymFromDMEnv class in utils/gym_wrapper.py:

import bsuite
from bsuite.utils import gym_wrapper

env = bsuite.load_and_record_to_csv('catch/0', results_dir='/path/to/results')
gym_env = gym_wrapper.GymFromDMEnv(env)

Note that bsuite does not include Gym in its default dependencies, so you may need to pip install it separately.

Baseline agents

We include implementations of several common agents in the [baselines/] subdirectory, along with a minimal run-loop.

See the installation section for how to include the required dependencies at install time. These dependencies are not installed by default, since bsuite does not require users to use any specific machine learning library.

Running the entire suite of experiments

Each of the agents in the baselines folder contains a run script which serves as an example which can run against a single environment or against the entire suite of experiments, by passing the --bsuite_id=SWEEP flags; this will start a pool of processes with which to run as many experiments in parallel as the host machine allows. On a 12 core machine, this will complete overnight for most agents. Alternatively, it is possible to run on Google Compute Platform using run_on_gcp.sh, steps of which are outlined below.

Running experiments on Google Cloud Platform

run_on_gcp.sh does the following in order:

  1. Create an instance with specified specs (by default 64-core CPU optimized).
  2. git clones bsuite and installs it together with other dependencies.
  3. Runs the specified agent (currently limited to /baselines) on a specified environment.
  4. Copies the resulting SQLite file to /tmp/bsuite.db from the remote instance to you local machine.
  5. Shuts down the created instance.

In order to run the script, you first need to create a billing account. Then follow the instructions here to setup and initialize Cloud SDK. After completing gcloud init, you are ready to run bsuite on Google Cloud.

For this make run_on_gcp.sh executable and run it:

chmod +x run_on_gcp.sh
./run_on_gcp.sh

After the instance is created, the instance name will be printed. Then you can ssh into the instance by selecting Compute Engine -> Instances and clicking SSH. Note that this is not necessary, as the result will be copied to your local machine once it is ready. However, sshing might be convenient if you want to make local changes to agent and environments. In this case, after sshing, do

~/bsuite_env/bin/activate

to activate the virtual environment. Then you can run agents via

python ~/bsuite/bsuite/baselines/dqn/run.py --bsuite_id=SWEEP

for instance.

Analysis

bsuite comes with a ready-made analysis Jupyter notebook included in analysis/results.ipynb. This notebook loads and processes logged data, and produces the scores and plots for each experiment. We recommend using this notebook in conjunction with Colaboratory.

We provide an example of a such bsuite report here.

bsuite Report

You can use bsuite to generate an automated 1-page appendix, that summarizes the core capabilities of your RL algorithm. This appendix is compatible with most major ML conference formats. For example output run,

pdflatex bsuite/reports/neurips_2019/neurips_2019.tex

More examples of bsuite reports can be found in the reports/ subdirectory.

Citing

If you use bsuite in your work, please cite the accompanying paper:

@inproceedings{osband2020bsuite,
    title={Behaviour Suite for Reinforcement Learning},
    author={Osband, Ian and
            Doron, Yotam and
            Hessel, Matteo and
            Aslanides, John and
            Sezener, Eren and
            Saraiva, Andre and
            McKinney, Katrina and
            Lattimore, Tor and
            {Sz}epesv{\'a}ri, Csaba and
            Singh, Satinder and
            Van Roy, Benjamin and
            Sutton, Richard and
            Silver, David and
            van Hasselt, Hado},
    booktitle={International Conference on Learning Representations},
    year={2020},
    url={https://openreview.net/forum?id=rygf-kSYwH}
}

More Repositories

1

deepmind-research

This repository contains implementations and illustrative code to accompany DeepMind publications
Jupyter Notebook
12,817
star
2

alphafold

Open source code for AlphaFold.
Python
11,700
star
3

sonnet

TensorFlow-based neural network library
Python
9,691
star
4

pysc2

StarCraft II Learning Environment
Python
7,904
star
5

mujoco

Multi-Joint dynamics with Contact. A general purpose physics simulator.
Jupyter Notebook
7,202
star
6

lab

A customisable 3D platform for agent-based AI research
C
7,012
star
7

graph_nets

Build Graph Nets in Tensorflow
Python
5,325
star
8

graphcast

Python
4,242
star
9

learning-to-learn

Learning to Learn in TensorFlow
Python
4,063
star
10

open_spiel

OpenSpiel is a collection of environments and algorithms for research in general reinforcement learning and search/planning in games.
C++
4,019
star
11

alphageometry

Python
3,580
star
12

dm_control

Google DeepMind's software stack for physics-based simulation and Reinforcement Learning environments, using MuJoCo.
Python
3,473
star
13

acme

A library of reinforcement learning components and agents
Python
3,372
star
14

trfl

TensorFlow Reinforcement Learning
Python
3,139
star
15

dm-haiku

JAX-based neural network library
Python
2,697
star
16

alphatensor

Python
2,616
star
17

dnc

A TensorFlow implementation of the Differentiable Neural Computer.
Python
2,478
star
18

mctx

Monte Carlo tree search in JAX
Python
2,209
star
19

gemma

Open weights LLM from Google DeepMind.
Jupyter Notebook
2,061
star
20

code_contests

C++
2,010
star
21

kinetics-i3d

Convolutional neural network model for video classification trained on the Kinetics dataset.
Python
1,639
star
22

mathematics_dataset

This dataset code generates mathematical question and answer pairs, from a range of question types at roughly school-level difficulty.
Python
1,621
star
23

optax

Optax is a gradient processing and optimization library for JAX.
Python
1,492
star
24

penzai

A JAX research toolkit for building, editing, and visualizing neural networks.
Python
1,405
star
25

educational

Jupyter Notebook
1,382
star
26

jraph

A Graph Neural Network Library in Jax
Python
1,306
star
27

rc-data

Question answering dataset featured in "Teaching Machines to Read and Comprehend
Python
1,285
star
28

rlax

Python
1,185
star
29

tapnet

Tracking Any Point (TAP)
Python
1,033
star
30

scalable_agent

A TensorFlow implementation of Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures.
Python
972
star
31

neural-processes

This repository contains notebook implementations of the following Neural Process variants: Conditional Neural Processes (CNPs), Neural Processes (NPs), Attentive Neural Processes (ANPs).
Jupyter Notebook
966
star
32

android_env

RL research on Android devices.
Python
946
star
33

mujoco_menagerie

A collection of high-quality models for the MuJoCo physics engine, curated by Google DeepMind.
Jupyter Notebook
926
star
34

dramatron

Dramatron uses large language models to generate coherent scripts and screenplays.
Jupyter Notebook
904
star
35

tree

tree is a library for working with nested data structures
Python
891
star
36

xmanager

A platform for managing machine learning experiments
Python
794
star
37

mujoco_mpc

Real-time behaviour synthesis with MuJoCo, using Predictive Control
C++
771
star
38

materials_discovery

Python
770
star
39

chex

Python
716
star
40

reverb

Reverb is an efficient and easy-to-use data storage and transport system designed for machine learning research
C++
692
star
41

alphadev

Python
662
star
42

pycolab

A highly-customisable gridworld game engine with some batteries included. Make your own gridworld games to test reinforcement learning agents!
Python
654
star
43

ferminet

An implementation of the Fermionic Neural Network for ab-initio electronic structure calculations
Python
643
star
44

hanabi-learning-environment

hanabi_learning_environment is a research platform for Hanabi experiments.
Python
614
star
45

funsearch

Jupyter Notebook
611
star
46

ai-safety-gridworlds

This is a suite of reinforcement learning environments illustrating various safety properties of intelligent agents.
Python
577
star
47

dqn

Lua/Torch implementation of DQN (Nature, 2015)
Lua
546
star
48

ithaca

Restoring and attributing ancient texts using deep neural networks
Jupyter Notebook
540
star
49

meltingpot

A suite of test scenarios for multi-agent reinforcement learning.
Python
516
star
50

distrax

Python
509
star
51

recurrentgemma

Open weights language model from Google DeepMind, based on Griffin.
Python
505
star
52

surface-distance

Library to compute surface distance based performance metrics for segmentation tasks.
Python
493
star
53

tracr

Python
467
star
54

dsprites-dataset

Dataset to assess the disentanglement properties of unsupervised learning methods
Jupyter Notebook
463
star
55

alphamissense

Python
455
star
56

narrativeqa

This repository contains the NarrativeQA dataset. It includes the list of documents with Wikipedia summaries, links to full stories, and questions and answers.
Shell
432
star
57

lab2d

A customisable 2D platform for agent-based AI research
C++
415
star
58

open_x_embodiment

Jupyter Notebook
409
star
59

dqn_zoo

DQN Zoo is a collection of reference implementations of reinforcement learning agents developed at DeepMind based on the Deep Q-Network (DQN) agent.
Python
406
star
60

clrs

Python
376
star
61

spriteworld

Spriteworld: a flexible, configurable python-based reinforcement learning environment
Python
366
star
62

dm_pix

PIX is an image processing library in JAX, for JAX.
Python
363
star
63

concordia

A library for generative social simulation
Python
351
star
64

mathematics_conjectures

Jupyter Notebook
348
star
65

alphastar

Python
346
star
66

spiral

We provide a pre-trained model for unconditional 19-step generation of CelebA-HQ images
C++
327
star
67

dm_env

A Python interface for reinforcement learning environments
Python
326
star
68

dm_robotics

Libraries, tools and tasks created and used at DeepMind Robotics.
Python
315
star
69

uncertain_ground_truth

Dermatology ddx dataset, Jax implementations of Monte Carlo conformal prediction, plausibility regions and statistical annotation aggregation from our recent work on uncertain ground truth (TMLR'23 and ArXiv pre-print).
Python
315
star
70

long-form-factuality

Benchmarking long-form factuality in large language models. Original code for our paper "Long-form factuality in large language models".
Python
314
star
71

launchpad

Python
305
star
72

leo

Implementation of Meta-Learning with Latent Embedding Optimization
Python
302
star
73

streetlearn

A C++/Python implementation of the StreetLearn environment based on images from Street View, as well as a TensorFlow implementation of goal-driven navigation agents solving the task published in “Learning to Navigate in Cities Without a Map”, NeurIPS 2018
C++
279
star
74

gqn-datasets

Datasets used to train Generative Query Networks (GQNs) in the ‘Neural Scene Representation and Rendering’ paper.
Python
269
star
75

enn

Python
265
star
76

multi_object_datasets

Multi-object image datasets with ground-truth segmentation masks and generative factors.
Python
247
star
77

AQuA

A algebraic word problem dataset, with multiple choice questions annotated with rationales.
238
star
78

card2code

A code generation dataset for generating the code that implements Hearthstone and Magic The Gathering card effects.
236
star
79

grid-cells

Implementation of the supervised learning experiments in Vector-based navigation using grid-like representations in artificial agents, as published at https://www.nature.com/articles/s41586-018-0102-6
Python
236
star
80

arnheim

Jupyter Notebook
235
star
81

synjax

Python
233
star
82

torch-hdf5

Torch interface to HDF5 library
Lua
231
star
83

dm_memorytasks

A set of 13 diverse machine-learning tasks that require memory to solve.
Python
220
star
84

Temporal-3D-Pose-Kinetics

Exploiting temporal context for 3D human pose estimation in the wild: 3D poses for the Kinetics dataset
Python
214
star
85

opro

official code for "Large Language Models as Optimizers"
Python
199
star
86

dm_alchemy

DeepMind Alchemy task environment: a meta-reinforcement learning benchmark
Python
197
star
87

neural_testbed

Jupyter Notebook
187
star
88

kfac-jax

Second Order Optimization and Curvature Estimation with K-FAC in JAX.
Python
187
star
89

pg19

179
star
90

xquad

173
star
91

jmp

JMP is a Mixed Precision library for JAX.
Python
171
star
92

spectral_inference_networks

Implementation of Spectral Inference Networks, ICLR 2019
Python
165
star
93

abstract-reasoning-matrices

Progressive matrices dataset, as described in: Measuring abstract reasoning in neural networks (Barrett*, Hill*, Santoro*, Morcos, Lillicrap), ICML2018
162
star
94

neural_networks_chomsky_hierarchy

Neural Networks and the Chomsky Hierarchy
Python
162
star
95

xitari

This is the 0.4 release of the Arcade Learning Environment (ALE), a platform designed for AI research. ALE is based on Stella, an Atari 2600 VCS emulator.
C++
159
star
96

tensor_annotations

Annotating tensor shapes using Python types
Python
158
star
97

symplectic-gradient-adjustment

A colab that implements the Symplectic Gradient Adjustment optimizer from "The mechanics of n-player differentiable games"
Jupyter Notebook
150
star
98

mc_gradients

Jupyter Notebook
149
star
99

interval-bound-propagation

This repository contains a simple implementation of Interval Bound Propagation (IBP) using TensorFlow: https://arxiv.org/abs/1810.12715
Python
148
star
100

s6

C++
146
star