![Discord Shield](https://dcbadge.vercel.app/api/server/cZs26Qq3Dd)
TorchRL
Documentation | TensorDict |
Features | Examples, tutorials and demos | Citation | Installation |
Asking a question | Contributing
TorchRL is an open-source Reinforcement Learning (RL) library for PyTorch.
It provides pytorch and python-first, low and high level abstractions for RL that are intended to be efficient, modular, documented and properly tested.
The code is aimed at supporting research in RL. Most of it is written in python in a highly modular way, such that researchers can easily swap components, transform them or write new ones with little effort.
This repo attempts to align with the existing pytorch ecosystem libraries in that it has a dataset pillar (torchrl/envs), transforms, models, data utilities (e.g. collectors and containers), etc.
TorchRL aims at having as few dependencies as possible (python standard library, numpy and pytorch). Common environment libraries (e.g. OpenAI gym) are only optional.
On the low-level end, torchrl comes with a set of highly re-usable functionals for cost functions, returns and data processing.
TorchRL aims at (1) a high modularity and (2) good runtime performance. Read the full paper for a more curated description of the library.
Getting started
Check our Getting Started tutorials for quickly ramp up with the basic
features of the library!
Documentation and knowledge base
The TorchRL documentation can be found here.
It contains tutorials and the API reference.
TorchRL also provides a RL knowledge base to help you debug your code, or simply
learn the basics of RL. Check it out here.
We have some introductory videos for you to get to know the library better, check them out:
Writing simplified and portable RL codebase with TensorDict
RL algorithms are very heterogeneous, and it can be hard to recycle a codebase
across settings (e.g. from online to offline, from state-based to pixel-based
learning).
TorchRL solves this problem through TensorDict
,
a convenient data structure(1) that can be used to streamline one's
RL codebase.
With this tool, one can write a complete PPO training script in less than 100
lines of code!
Code
import torch
from tensordict.nn import TensorDictModule
from tensordict.nn.distributions import NormalParamExtractor
from torch import nn
from torchrl.collectors import SyncDataCollector
from torchrl.data.replay_buffers import TensorDictReplayBuffer, \
LazyTensorStorage, SamplerWithoutReplacement
from torchrl.envs.libs.gym import GymEnv
from torchrl.modules import ProbabilisticActor, ValueOperator, TanhNormal
from torchrl.objectives import ClipPPOLoss
from torchrl.objectives.value import GAE
env = GymEnv("Pendulum-v1")
model = TensorDictModule(
nn.Sequential(
nn.Linear(3, 128), nn.Tanh(),
nn.Linear(128, 128), nn.Tanh(),
nn.Linear(128, 128), nn.Tanh(),
nn.Linear(128, 2),
NormalParamExtractor()
),
in_keys=["observation"],
out_keys=["loc", "scale"]
)
critic = ValueOperator(
nn.Sequential(
nn.Linear(3, 128), nn.Tanh(),
nn.Linear(128, 128), nn.Tanh(),
nn.Linear(128, 128), nn.Tanh(),
nn.Linear(128, 1),
),
in_keys=["observation"],
)
actor = ProbabilisticActor(
model,
in_keys=["loc", "scale"],
distribution_class=TanhNormal,
distribution_kwargs={"min": -1.0, "max": 1.0},
return_log_prob=True
)
buffer = TensorDictReplayBuffer(
LazyTensorStorage(1000),
SamplerWithoutReplacement()
)
collector = SyncDataCollector(
env,
actor,
frames_per_batch=1000,
total_frames=1_000_000
)
loss_fn = ClipPPOLoss(actor, critic, gamma=0.99)
optim = torch.optim.Adam(loss_fn.parameters(), lr=2e-4)
adv_fn = GAE(value_network=critic, gamma=0.99, lmbda=0.95, average_gae=True)
for data in collector:
for epoch in range(10):
adv_fn(data)
buffer.extend(data.view(-1))
for i in range(20):
sample = buffer.sample(50)
loss_vals = loss_fn(sample)
loss_val = sum(
value for key, value in loss_vals.items() if
key.startswith("loss")
)
loss_val.backward()
optim.step()
optim.zero_grad()
print(f"avg reward: {data['next', 'reward'].mean().item(): 4.4f}")
Here is an example of how the environment API
relies on tensordict to carry data from one function to another during a rollout
execution:
![Alt Text](https://github.com/pytorch/rl/blob/main/docs/source/_static/img/rollout.gif)
TensorDict
makes it easy to re-use pieces of code across environments, models and
algorithms.
Code
For instance, here's how to code a rollout in TorchRL:
- obs, done = env.reset()
+ tensordict = env.reset()
policy = SafeModule(
model,
in_keys=["observation_pixels", "observation_vector"],
out_keys=["action"],
)
out = []
for i in range(n_steps):
- action, log_prob = policy(obs)
- next_obs, reward, done, info = env.step(action)
- out.append((obs, next_obs, action, log_prob, reward, done))
- obs = next_obs
+ tensordict = policy(tensordict)
+ tensordict = env.step(tensordict)
+ out.append(tensordict)
+ tensordict = step_mdp(tensordict) # renames next_observation_* keys to observation_*
- obs, next_obs, action, log_prob, reward, done = [torch.stack(vals, 0) for vals in zip(*out)]
+ out = torch.stack(out, 0) # TensorDict supports multiple tensor operations
Using this, TorchRL abstracts away the input / output signatures of the modules, env,
collectors, replay buffers and losses of the library, allowing all primitives
to be easily recycled across settings.
Code
Here's another example of an off-policy training loop in TorchRL (assuming
that a data collector, a replay buffer, a loss and an optimizer have been instantiated):
- for i, (obs, next_obs, action, hidden_state, reward, done) in enumerate(collector):
+ for i, tensordict in enumerate(collector):
- replay_buffer.add((obs, next_obs, action, log_prob, reward, done))
+ replay_buffer.add(tensordict)
for j in range(num_optim_steps):
- obs, next_obs, action, hidden_state, reward, done = replay_buffer.sample(batch_size)
- loss = loss_fn(obs, next_obs, action, hidden_state, reward, done)
+ tensordict = replay_buffer.sample(batch_size)
+ loss = loss_fn(tensordict)
loss.backward()
optim.step()
optim.zero_grad()
This training loop can be re-used across algorithms as it makes a minimal number of assumptions about the structure of the data.
TensorDict supports multiple tensor operations on its device and shape
(the shape of TensorDict, or its batch size, is the common arbitrary N first dimensions of all its contained tensors):
Code
tensordict = torch.stack(list_of_tensordicts, 0)
tensordict = torch.cat(list_of_tensordicts, 0)
tensordict = tensordict.view(-1)
tensordict = tensordict.permute(0, 2, 1)
tensordict = tensordict.unsqueeze(-1)
tensordict = tensordict.squeeze(-1)
tensordict = tensordict[:2]
tensordict[:, 2] = sub_tensordict
tensordict.cuda()
tensordict.to("cuda:1")
tensordict.share_memory_()
TensorDict comes with a dedicated tensordict.nn
module that contains everything you might need to write your model with it.
And it is functorch
and torch.compile
compatible!
Code
transformer_model = nn.Transformer(nhead=16, num_encoder_layers=12)
+ td_module = SafeModule(transformer_model, in_keys=["src", "tgt"], out_keys=["out"])
src = torch.rand((10, 32, 512))
tgt = torch.rand((20, 32, 512))
+ tensordict = TensorDict({"src": src, "tgt": tgt}, batch_size=[20, 32])
- out = transformer_model(src, tgt)
+ td_module(tensordict)
+ out = tensordict["out"]
The TensorDictSequential
class allows to branch sequences of nn.Module
instances in a highly modular way.
For instance, here is an implementation of a transformer using the encoder and decoder blocks:
encoder_module = TransformerEncoder(...)
encoder = TensorDictSequential(encoder_module, in_keys=["src", "src_mask"], out_keys=["memory"])
decoder_module = TransformerDecoder(...)
decoder = TensorDictModule(decoder_module, in_keys=["tgt", "memory"], out_keys=["output"])
transformer = TensorDictSequential(encoder, decoder)
assert transformer.in_keys == ["src", "src_mask", "tgt"]
assert transformer.out_keys == ["memory", "output"]
TensorDictSequential
allows to isolate subgraphs by querying a set of desired input / output keys:
transformer.select_subsequence(out_keys=["memory"])
transformer.select_subsequence(in_keys=["tgt", "memory"])
Check TensorDict tutorials to
learn more!
Features
-
A common interface for environments
which supports common libraries (OpenAI gym, deepmind control lab, etc.)(1) and state-less execution
(e.g. Model-based environments).
The batched environments containers allow parallel execution(2).
A common PyTorch-first class of tensor-specification class is also provided.
TorchRL's environments API is simple but stringent and specific. Check the
documentation
and tutorial to learn more!
Code
env_make = lambda: GymEnv("Pendulum-v1", from_pixels=True)
env_parallel = ParallelEnv(4, env_make)
tensordict = env_parallel.rollout(max_steps=20, policy=None)
assert tensordict.shape == [4, 20]
env_parallel.action_spec.is_in(tensordict["action"])
-
multiprocess and distributed data collectors(2)
that work synchronously or asynchronously.
Through the use of TensorDict, TorchRL's training loops are made very similar
to regular training loops in supervised
learning (although the "dataloader" -- read data collector -- is modified on-the-fly):
Code
env_make = lambda: GymEnv("Pendulum-v1", from_pixels=True)
collector = MultiaSyncDataCollector(
[env_make, env_make],
policy=policy,
devices=["cuda:0", "cuda:0"],
total_frames=10000,
frames_per_batch=50,
...
)
for i, tensordict_data in enumerate(collector):
loss = loss_module(tensordict_data)
loss.backward()
optim.step()
optim.zero_grad()
collector.update_policy_weights_()
Check our distributed collector examples to
learn more about ultra-fast data collection with TorchRL.
-
efficient(2) and generic(1) replay buffers with modularized storage:
Code
storage = LazyMemmapStorage(
cfg.buffer_size,
scratch_dir="/tmp/"
)
buffer = TensorDictPrioritizedReplayBuffer(
alpha=0.7,
beta=0.5,
collate_fn=lambda x: x,
pin_memory=device != torch.device("cpu"),
prefetch=10,
storage=storage
)
Replay buffers are also offered as wrappers around common datasets for offline RL:
Code
from torchrl.data.replay_buffers import SamplerWithoutReplacement
from torchrl.data.datasets.d4rl import D4RLExperienceReplay
data = D4RLExperienceReplay(
"maze2d-open-v0",
split_trajs=True,
batch_size=128,
sampler=SamplerWithoutReplacement(drop_last=True),
)
for sample in data:
fun(sample)
-
cross-library environment transforms(1),
executed on device and in a vectorized fashion(2),
which process and prepare the data coming out of the environments to be used by the agent:
Code
env_make = lambda: GymEnv("Pendulum-v1", from_pixels=True)
env_base = ParallelEnv(4, env_make, device="cuda:0")
env = TransformedEnv(
env_base,
Compose(
ToTensorImage(),
ObservationNorm(loc=0.5, scale=1.0)),
)
tensordict = env.reset()
assert tensordict.device == torch.device("cuda:0")
Other transforms include: reward scaling (RewardScaling
), shape operations (concatenation of tensors, unsqueezing etc.), concatenation of
successive operations (CatFrames
), resizing (Resize
) and many more.
Unlike other libraries, the transforms are stacked as a list (and not wrapped in each other), which makes it
easy to add and remove them at will:
env.insert_transform(0, NoopResetEnv())
Nevertheless, transforms can access and execute operations on the parent environment:
transform = env.transform[1]
parent_env = transform.parent
-
various tools for distributed learning (e.g. memory mapped tensors)(2);
-
various architectures and models (e.g. actor-critic)(1):
Code
common_module = ConvNet(
bias_last_layer=True,
depth=None,
num_cells=[32, 64, 64],
kernel_sizes=[8, 4, 3],
strides=[4, 2, 1],
)
common_module = SafeModule(
common_module,
in_keys=["pixels"],
out_keys=["hidden"],
)
policy_module = SafeModule(
NormalParamsWrapper(
MLP(num_cells=[64, 64], out_features=32, activation=nn.ELU)
),
in_keys=["hidden"],
out_keys=["loc", "scale"],
)
policy_module = SafeProbabilisticTensorDictSequential(
policy_module,
SafeProbabilisticModule(
in_keys=["loc", "scale"],
out_keys="action",
distribution_class=TanhNormal,
),
)
value_module = MLP(
num_cells=[64, 64],
out_features=1,
activation=nn.ELU,
)
actor_value = ActorValueOperator(common_module, policy_module, value_module)
standalone_policy = actor_value.get_policy_operator()
-
exploration wrappers and
modules to easily swap between exploration and exploitation(1):
Code
policy_explore = EGreedyWrapper(policy)
with set_exploration_type(ExplorationType.RANDOM):
tensordict = policy_explore(tensordict)
with set_exploration_type(ExplorationType.MODE):
tensordict = policy_explore(tensordict)
-
A series of efficient loss modules
and highly vectorized
functional return and advantage
computation.
Code
Loss modules
from torchrl.objectives import DQNLoss
loss_module = DQNLoss(value_network=value_network, gamma=0.99)
tensordict = replay_buffer.sample(batch_size)
loss = loss_module(tensordict)
Advantage computation
from torchrl.objectives.value.functional import vec_td_lambda_return_estimate
advantage = vec_td_lambda_return_estimate(gamma, lmbda, next_state_value, reward, done, terminated)
-
a generic trainer class(1) that
executes the aforementioned training loop. Through a hooking mechanism,
it also supports any logging or data transformation operation at any given
time.
-
various recipes to build models that
correspond to the environment being deployed.
If you feel a feature is missing from the library, please submit an issue!
If you would like to contribute to new features, check our call for contributions and our contribution page.
Examples, tutorials and demos
A series of examples are provided with an illustrative purpose:
and many more to come!
Check the examples directory for more details
about handling the various configuration settings.
We also provide tutorials and demos that give a sense of
what the library can do.
Citation
If you're using TorchRL, please refer to this BibTeX entry to cite this work:
@misc{bou2023torchrl,
title={TorchRL: A data-driven decision-making library for PyTorch},
author={Albert Bou and Matteo Bettini and Sebastian Dittert and Vikash Kumar and Shagun Sodhani and Xiaomeng Yang and Gianni De Fabritiis and Vincent Moens},
year={2023},
eprint={2306.00577},
archivePrefix={arXiv},
primaryClass={cs.LG}
}
Installation
Create a conda environment where the packages will be installed.
conda create --name torch_rl python=3.9
conda activate torch_rl
PyTorch
Depending on the use of functorch that you want to make, you may want to
install the latest (nightly) PyTorch release or the latest stable version of PyTorch.
See here for a detailed list of commands,
including pip3
or other special installation instructions.
Torchrl
You can install the latest stable release by using
pip3 install torchrl
This should work on linux, Windows 10 and OsX (Intel or Silicon chips).
On certain Windows machines (Windows 11), one should install the library locally (see below).
The nightly build can be installed via
pip3install torchrl-nightly
which we currently only ship for Linux and OsX (Intel) machines.
Importantly, the nightly builds require the nightly builds of PyTorch too.
To install extra dependencies, call
pip3 install "torchrl[atari,dm_control,gym_continuous,rendering,tests,utils,marl,checkpointing]"
or a subset of these.
One may also desire to install the library locally. Three main reasons can motivate this:
- the nightly/stable release isn't available for one's platform (eg, Windows 11, nightlies for Apple Silicon etc.);
- contributing to the code;
- install torchrl with a previous version of PyTorch (any version >= 2.0) (note that this should also be doable via a regular install followed
by a downgrade to a previous pytorch version -- but the C++ binaries will not be available so some feature will not work,
such as prioritized replay buffers and the like.)
To install the library locally, start by cloning the repo:
git clone https://github.com/pytorch/rl
and don't forget to check out the branch or tag you want to use for the build:
git checkout v0.4.0
Go to the directory where you have cloned the torchrl repo and install it (after
installing ninja
)
cd /path/to/torchrl/
pip3install ninja -U
python setup.py develop
One can also build the wheels to distribute to co-workers using
python setup.py bdist_wheel
Your wheels will be stored there ./dist/torchrl<name>.whl
and installable via
pip install torchrl<name>.whl
Warning: Unfortunately, pip3 install -e .
does not currently work. Contributions to help fix this are welcome!
On M1 machines, this should work out-of-the-box with the nightly build of PyTorch.
If the generation of this artifact in MacOs M1 doesn't work correctly or in the execution the message
(mach-o file, but is an incompatible architecture (have 'x86_64', need 'arm64e'))
appears, then try
ARCHFLAGS="-arch arm64" python setup.py develop
To run a quick sanity check, leave that directory (e.g. by executing cd ~/
)
and try to import the library.
python -c "import torchrl"
This should not return any warning or error.
Optional dependencies
The following libraries can be installed depending on the usage one wants to
make of torchrl:
# diverse
pip3 install tqdm tensorboard "hydra-core>=1.1" hydra-submitit-launcher
# rendering
pip3 install moviepy
# deepmind control suite
pip3 install dm_control
# gym, atari games
pip3 install "gym[atari]" "gym[accept-rom-license]" pygame
# tests
pip3 install pytest pyyaml pytest-instafail
# tensorboard
pip3 install tensorboard
# wandb
pip3 install wandb
Troubleshooting
If a ModuleNotFoundError: No module named ‘torchrl._torchrl
errors occurs (or
a warning indicating that the C++ binaries could not be loaded),
it means that the C++ extensions were not installed or not found.
Versioning issues can cause error message of the type undefined symbol
and such. For these, refer to the versioning issues document
for a complete explanation and proposed workarounds.
Asking a question
If you spot a bug in the library, please raise an issue in this repo.
If you have a more generic question regarding RL in PyTorch, post it on
the PyTorch forum.
Contributing
Internal collaborations to torchrl are welcome! Feel free to fork, submit issues and PRs.
You can checkout the detailed contribution guide here.
As mentioned above, a list of open contributions can be found in here.
Contributors are recommended to install pre-commit hooks (using pre-commit install
). pre-commit will check for linting related issues when the code is committed locally. You can disable th check by appending -n
to your commit command: git commit -m <commit message> -n
Disclaimer
This library is released as a PyTorch beta feature.
BC-breaking changes are likely to happen but they will be introduced with a deprecation
warranty after a few release cycles.
License
TorchRL is licensed under the MIT License. See LICENSE for details.