Skip to content

Distributed Reinforcement Learning accelerated by Lightning Fabric

License

Notifications You must be signed in to change notification settings

hanshuo-shuo/sheeprl_prey

Β 
Β 

Folders and files

NameName
Last commit message
Last commit date

Latest commit

Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 

Repository files navigation

⚑ SheepRL πŸ‘

What

An easy-to-use framework for reinforcement learning in PyTorch, accelerated with Lightning Fabric.
The algorithms sheeped by sheeprl out-of-the-box are:

Algorithm Coupled Decoupled Recurrent Vector obs Pixel obs Status
A2C βœ”οΈ βœ”οΈ βœ”οΈ βœ”οΈ βœ”οΈ 🚧
A3C βœ”οΈ ❌ ❌ βœ”οΈ βœ”οΈ 🚧
PPO βœ”οΈ βœ”οΈ ❌ βœ”οΈ βœ”οΈ βœ”οΈ
PPO Recurrent βœ”οΈ ❌ βœ”οΈ βœ”οΈ βœ”οΈ βœ”οΈ
SAC βœ”οΈ βœ”οΈ ❌ βœ”οΈ ❌ βœ”οΈ
SAC-AE βœ”οΈ ❌ ❌ βœ”οΈ βœ”οΈ βœ”οΈ
DroQ βœ”οΈ ❌ ❌ βœ”οΈ ❌ βœ”οΈ
Dreamer-V1 βœ”οΈ ❌ βœ”οΈ βœ”οΈ βœ”οΈ βœ”οΈ
Dreamer-V2 βœ”οΈ ❌ βœ”οΈ βœ”οΈ βœ”οΈ βœ”οΈ
Dreamer-V3 βœ”οΈ ❌ βœ”οΈ βœ”οΈ βœ”οΈ βœ”οΈ
Plan2Explore (Dreamer V1) βœ”οΈ ❌ βœ”οΈ βœ”οΈ βœ”οΈ βœ”οΈ
Plan2Explore (Dreamer V2) βœ”οΈ ❌ βœ”οΈ βœ”οΈ βœ”οΈ βœ”οΈ

and more are coming soon! Open a PR if you have any particular request πŸ‘

The actions supported by sheeprl agents are:

Algorithm Continuous Discrete Multi-Discrete
A2C βœ”οΈ βœ”οΈ βœ”οΈ
A3C βœ”οΈ βœ”οΈ βœ”οΈ
PPO βœ”οΈ βœ”οΈ βœ”οΈ
PPO Recurrent βœ”οΈ βœ”οΈ βœ”οΈ
SAC βœ”οΈ ❌ ❌
SAC-AE βœ”οΈ ❌ ❌
DroQ βœ”οΈ ❌ ❌
Dreamer-V1 βœ”οΈ βœ”οΈ βœ”οΈ
Dreamer-V2 βœ”οΈ βœ”οΈ βœ”οΈ
Dreamer-V3 βœ”οΈ βœ”οΈ βœ”οΈ
Plan2Explore (Dreamer V1) βœ”οΈ βœ”οΈ βœ”οΈ
Plan2Explore (Dreamer V2) βœ”οΈ βœ”οΈ βœ”οΈ

The environments supported by sheeprl are:

Algorithm Installation command More info Status
Classic Control pip install -e . βœ”οΈ
Box2D pip install -e . βœ”οΈ
Mujoco (Gymnasium) pip install -e . how_to/mujoco βœ”οΈ
Atari pip install -e .[atari] how_to/atari βœ”οΈ
DeepMind Control pip install -e .[dmc] how_to/dmc βœ”οΈ
MineRL pip install -e .[minerl] how_to/minerl βœ”οΈ
MineDojo pip install -e .[minedojo] how_to/minedojo βœ”οΈ
DIAMBRA pip install -e .[diambra] how_to/diambra βœ”οΈ
Crafter pip install -e .[crafter] https://github.com/danijar/crafter βœ”οΈ

Why

We want to provide a framework for RL algorithms that is at the same time simple and scalable thanks to Lightning Fabric.

Moreover, in many RL repositories, the RL algorithm is tightly coupled with the environment, making it harder to extend them beyond the gym interface. We want to provide a framework that allows to easily decouple the RL algorithm from the environment, so that it can be used with any environment.

How to use

Two options exist for using SheepRL. One can either clone the repo and install the local version, or one can pip install the framework using the GitHub clone URL. Instructions for both methods are shown below.

Cloning and installing a local version

First, clone the repo with:

git clone https://github.com/Eclectic-Sheep/sheeprl.git
cd sheeprl

From inside the newly create folder run

pip install .

Note

To install all the optional dependencies one can run pip install .[atari,mujoco,dev,test]

Installing the framework from the GitHub repo

If you haven't already done so, create an environment with your choice of venv or conda.

Note

The example will use Python standard's venv module, and assumes macOS or Linux.

# create a virtual environment
python3 -m venv .venv
# activate the environment
source .venv/bin/activate
# if you do not wish to install extras such as mujuco, atari do
pip install "sheeprl @ git+https://github.com/Eclectic-Sheep/sheeprl.git"
# or, to install with atari and mujuco environment support, do
pip install "sheeprl[atari,mujoco,dev]  @ git+https://github.com/Eclectic-Sheep/sheeprl.git"
# or, to install with minedojo environment support, do
pip install "sheeprl[minedojo,dev]  @ git+https://github.com/Eclectic-Sheep/sheeprl.git"
# or, to install with minerl environment support, do
pip install "sheeprl[minerl,dev]  @ git+https://github.com/Eclectic-Sheep/sheeprl.git"
# or, to install with diambra environment support, do
pip install "sheeprl[diambra,dev]  @ git+https://github.com/Eclectic-Sheep/sheeprl.git"
# or, to install all extras, do
pip install "sheeprl[atari,mujoco,miedojo,dev,test]  @ git+https://github.com/Eclectic-Sheep/sheeprl.git"
Installing on an M-series Mac

Note

if you are on an M-series Mac and encounter an error attributed box2dpy during install, you need to install SWIG using the instructions shown below.

It is recommended to use homebrew to install SWIG to support Gym.

# if needed install homebrew
/bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/HEAD/install.sh)"
# then, do
brew install swig
# then attempt to pip install with the prefered method, such as
pip install "sheeprl[atari,mujoco,dev,test] @ git+https://github.com/Eclectic-Sheep/sheeprl.git"
MineRL, MineDojo and DIAMBRA

Note

If you want to install the minedojo or minerl environment support, Java JDK 8 is required: you can install it by following the instructions at this link.

MineRL and MineDojo environments have conflicting requirements, so DO NOT install them together with the pip install -e .[minerl,minedojo] command, but instead install them individually with either the command pip install -e .[minerl] or pip install -e .[minedojo] before running an experiment with the MineRL or MineDojo environment, respectively.

Now you can use one of the already available algorithms, or create your own.

For example, to train a PPO agent on the CartPole environment with only vector-like observations, just run

python sheeprl.py exp=ppo env=gym env.id=CartPole-v1

You check all the available algorithms with

python sheeprl/available_agents.py

That's all it takes to train an agent with SheepRL! πŸŽ‰

Note

Before you start using the SheepRL framework, it is highly recommended that you read the following instructional documents:

  1. How to run experiments
  2. How to modify the default configs
  3. How to work with steps
  4. How to select observations

Moreover, there are other useful documents in the howto folder, which containes some guidance on how to properly use the framework.

πŸ“ˆ Check your results

Once you trained an agent, a new folder called logs will be created, containing the logs of the training. You can visualize them with TensorBoard:

tensorboard --logdir logs
tensorboard.mp4

πŸ€“ More about running an algorithm

What you run is the PPO algorithm with the default configuration. But you can also change the configuration by passing arguments to the script.

For example, in the default configuration, the number of parallel environments is 4. Let's try to change it to 8 by passing the --num_envs argument:

python sheeprl.py exp=ppo env=gym env.id=CartPole-v1 env.num_envs=8

All the available arguments, with their descriptions, are listed in the sheeprl/config directory. You can find more information about the hierarchy of configs here.

Running with Lightning Fabric

To run the algorithm with Lightning Fabric, you need to specify the Fabric parameters through the CLI. For example, to run the PPO algorithm with 4 parallel environments on 2 nodes, you can run:

python sheeprl.py fabric.accelerator=cpu fabric.strategy=ddp fabric.devices=2 exp=ppo env=gym env.id=CartPole-v1

You can check the available parameters for Lightning Fabric here.

πŸ“– Repository structure

The repository is structured as follows:

  • algos: contains the implementations of the algorithms. Each algorithm is in a separate folder, and (possibly) contains the following files:

    • <algorithm>.py: contains the implementation of the algorithm.
    • <algorithm>_decoupled.py: contains the implementation of the decoupled version of the algorithm, if present.
    • agent: optional, contains the implementation of the agent.
    • loss.py: contains the implementation of the loss functions of the algorithm.
    • utils.py: contains utility functions for the algorithm.
  • configs: contains the default configs of the algorithms.

  • data: contains the implementation of the data buffers.

  • envs: contains the implementation of the environment wrappers.

  • models: contains the implementation of the some standard models (building blocks), like the multi-layer perceptron (MLP) or a simple convolutional network (NatureCNN)

  • utils: contains utility functions for the framework.

Coupled vs Decoupled

In the coupled version of an algorithm, the agent interacts with the environment and executes the training loop.

In the decoupled version, a process is responsible only for interacting with the environment, and all the other processes are responsible for executing the training loop. The two processes communicate through distributed collectives, adopting the abstraction provided by Fabric's TorchCollective.

Coupled

The algorithm is implemented in the <algorithm>.py file.

There are 2 functions inside this script:

  • main(): initializes all the components of the algorithm, and executes the interactions with the environment. Once enough data is collected, the training loop is executed by calling the train() function.
  • train(): executes the training loop. It samples a batch of data from the buffer, compute the loss and updates the parameters of the agent.

Decoupled

The decoupled version of an algorithm is implemented in the <algorithm>_decoupled.py file.

There are 3 functions inside this script:

  • main(): initializes all the components of the algorithm, the collectives for the communication between the player and the trainers and calls the player() and trainer() functions.
  • player(): executes the interactions with the environment. It samples an action from the policy network, executes it in the environment, and stores the transition in the buffer. After a predefined number of iteractions with the environment, the player randomly splits the collected data in almost-equal chunks and send them separately to the trainers. It then waits for the trainers to finish the agent update.
  • trainer(): executes the training loop. It receives a chunk of data from the player, compute the loss and updates the parameters of the agent. After the agent has been updated, the first of the trainers sends back the updated agent weights to the player, which can interact again with the environment.

Algorithms implementation

You can check inside the folder of each algorithm the README.md file for the details about the implementation.

All algorithms are kept as simple as possible, in a CleanRL fashion. But to allow for more flexibility and also more clarity, we tried to abstract away anything that is not strictly related with the training loop of the algorithm.

For example, we decided to create a models folder with an already-made models that can be composed to create the model of the agent.

For each algorithm, losses are kept in a separate module, so that their implementation is clear and can be easily utilized also for the decoupled or the recurrent version of the algorithm.

πŸ—‚οΈ Buffer

For the buffer implementation, we choose to use a wrapper around a TensorDict.

TensorDict comes handy since we can easily add custom fields to the buffer as if we are working with dictionaries, but we can also easily perform operations on them as if we are working with tensors.

This flexibility makes it very simple to implement, with the classes ReplayBuffer, SequentialReplayBuffer, EpisodeBuffer, and AsyncReplayBuffer, all the buffers needed for on-policy and off-policy algorithms.

πŸ” Technical details

The tensor's shape in the TensorDict are (T, B, *), where T is the number of timesteps, B is the number of parallel environments, and * is the shape of the data.

For the ReplayBuffer to be used as a RolloutBuffer, the proper buffer_size must be specified. For example, for PPO, the buffer_size must be [T, B], where T is the number of timesteps and B is the number of parallel environments.

πŸ™‡ Contributing

The best way to contribute is by opening an issue to discuss a new feature or a bug, or by opening a PR to fix a bug or to add a new feature.

πŸ“­ Contacts

You can contact us for any further questions or discussions:

About

Distributed Reinforcement Learning accelerated by Lightning Fabric

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

 
 
 

Languages

  • Python 100.0%