Package: torchrl
TorchRLDocumentation|TensorDict|Features|Examples, tutorials and demos|Citation|Installation|Asking a question|ContributingTorchRLis an open-source Reinforcement Learning (RL) library for PyTorch.It provides pytorch andpython-first, low and high level abstractions for RL that are intended to beefficient,modular,documentedand properlytested.
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 forcost functions,returnsand data processing.TorchRL aims at (1) a high modularity and (2) good runtime performance. Read thefull paperfor a more curated description of the library.Documentation and knowledge baseThe TorchRL documentation can be foundhere.
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 outhere.We have some introductory videos for you to get to know the library better, check them out:TorchRL intro at PyTorch day 2022PyTorch 2.0 Q&A: TorchRLWriting simplified and portable RL codebase withTensorDictRL 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 throughTensorDict,
a convenient data structure(1)that can be used to streamline one's
RL codebase.
With this tool, one can write acomplete PPO training script in less than 100
lines of code!Codeimporttorchfromtensordict.nnimportTensorDictModulefromtensordict.nn.distributionsimportNormalParamExtractorfromtorchimportnnfromtorchrl.collectorsimportSyncDataCollectorfromtorchrl.data.replay_buffersimportTensorDictReplayBuffer,\LazyTensorStorage,SamplerWithoutReplacementfromtorchrl.envs.libs.gymimportGymEnvfromtorchrl.modulesimportProbabilisticActor,ValueOperator,TanhNormalfromtorchrl.objectivesimportClipPPOLossfromtorchrl.objectives.valueimportGAEenv=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)fordataincollector:# collect dataforepochinrange(10):adv_fn(data)# compute advantagebuffer.extend(data.view(-1))foriinrange(20):# consume datasample=buffer.sample(50)# mini-batchloss_vals=loss_fn(sample)loss_val=sum(valueforkey,valueinloss_vals.items()ifkey.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 theenvironment APIrelies on tensordict to carry data from one function to another during a rollout
execution:TensorDictmakes it easy to re-use pieces of code across environments, models and
algorithms.CodeFor 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 operationsUsing 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.CodeHere'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# stack and cattensordict=torch.stack(list_of_tensordicts,0)tensordict=torch.cat(list_of_tensordicts,0)# reshapetensordict=tensordict.view(-1)tensordict=tensordict.permute(0,2,1)tensordict=tensordict.unsqueeze(-1)tensordict=tensordict.squeeze(-1)# indexingtensordict=tensordict[:2]tensordict[:,2]=sub_tensordict# device and memory locationtensordict.cuda()tensordict.to("cuda:1")tensordict.share_memory_()TensorDict comes with a dedicatedtensordict.nnmodule that contains everything you might need to write your model with it.
And it isfunctorchandtorch.compilecompatible!Codetransformer_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"]TheTensorDictSequentialclass allows to branch sequences ofnn.Moduleinstances 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)asserttransformer.in_keys==["src","src_mask","tgt"]asserttransformer.out_keys==["memory","output"]TensorDictSequentialallows to isolate subgraphs by querying a set of desired input / output keys:transformer.select_subsequence(out_keys=["memory"])# returns the encodertransformer.select_subsequence(in_keys=["tgt","memory"])# returns the decoderCheckTensorDict tutorialsto
learn more!FeaturesA commoninterface for environmentswhich supports common libraries (OpenAI gym, deepmind control lab, etc.)(1)and state-less execution
(e.g. Model-based environments).
Thebatched environmentscontainers allow parallel execution(2).
A common PyTorch-first class oftensor-specification classis also provided.
TorchRL's environments API is simple but stringent and specific. Check thedocumentationandtutorialto learn more!Codeenv_make=lambda:GymEnv("Pendulum-v1",from_pixels=True)env_parallel=ParallelEnv(4,env_make)# creates 4 envs in paralleltensordict=env_parallel.rollout(max_steps=20,policy=None)# random rollout (no policy given)asserttensordict.shape==[4,20]# 4 envs, 20 steps rolloutenv_parallel.action_spec.is_in(tensordict["action"])# spec check returns Truemultiprocess and distributeddata 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):Codeenv_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,...)fori,tensordict_datainenumerate(collector):loss=loss_module(tensordict_data)loss.backward()optim.step()optim.zero_grad()collector.update_policy_weights_()Check ourdistributed collector examplesto
learn more about ultra-fast data collection with TorchRL.efficient(2)and generic(1)replay bufferswith modularized storage:Codestorage=LazyMemmapStorage(# memory-mapped (physical) storagecfg.buffer_size,scratch_dir="/tmp/")buffer=TensorDictPrioritizedReplayBuffer(alpha=0.7,beta=0.5,collate_fn=lambdax:x,pin_memory=device!=torch.device("cpu"),prefetch=10,# multi-threaded samplingstorage=storage)Replay buffers are also offered as wrappers around common datasets foroffline RL:Codefromtorchrl.data.replay_buffersimportSamplerWithoutReplacementfromtorchrl.data.datasets.d4rlimportD4RLExperienceReplaydata=D4RLExperienceReplay("maze2d-open-v0",split_trajs=True,batch_size=128,sampler=SamplerWithoutReplacement(drop_last=True),)forsampleindata:# or alternatively sample = data.sample()fun(sample)cross-libraryenvironment 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:Codeenv_make=lambda:GymEnv("Pendulum-v1",from_pixels=True)env_base=ParallelEnv(4,env_make,device="cuda:0")# creates 4 envs in parallelenv=TransformedEnv(env_base,Compose(ToTensorImage(),ObservationNorm(loc=0.5,scale=1.0)),# executes the transforms once and on device)tensordict=env.reset()asserttensordict.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())# inserts the NoopResetEnv transform at the index 0Nevertheless, transforms can access and execute operations on the parent environment:transform=env.transform[1]# gathers the second transform of the listparent_env=transform.parent# returns the base environment of the second transform, i.e. the base env + the first transformvarious tools for distributed learning (e.g.memory mapped tensors)(2);variousarchitecturesand models (e.g.actor-critic)(1):Code# create an nn.Modulecommon_module=ConvNet(bias_last_layer=True,depth=None,num_cells=[32,64,64],kernel_sizes=[8,4,3],strides=[4,2,1],)# Wrap it in a SafeModule, indicating what key to read in and where to# write out the outputcommon_module=SafeModule(common_module,in_keys=["pixels"],out_keys=["hidden"],)# Wrap the policy module in NormalParamsWrapper, such that the output# tensor is split in loc and scale, and scale is mapped onto a positive spacepolicy_module=SafeModule(NormalParamsWrapper(MLP(num_cells=[64,64],out_features=32,activation=nn.ELU)),in_keys=["hidden"],out_keys=["loc","scale"],)# Use a SafeProbabilisticTensorDictSequential to combine the SafeModule with a# SafeProbabilisticModule, indicating how to build the# torch.distribution.Distribution object and what to do with itpolicy_module=SafeProbabilisticTensorDictSequential(# stochastic policypolicy_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,)# Wrap the policy and value funciton in a common moduleactor_value=ActorValueOperator(common_module,policy_module,value_module)# standalone policy from thisstandalone_policy=actor_value.get_policy_operator()explorationwrappersandmodulesto easily swap between exploration and exploitation(1):Codepolicy_explore=EGreedyWrapper(policy)withset_exploration_type(ExplorationType.RANDOM):tensordict=policy_explore(tensordict)# will use eps-greedywithset_exploration_type(ExplorationType.MODE):tensordict=policy_explore(tensordict)# will not use eps-greedyA series of efficientloss modulesand highly vectorizedfunctional return and advantagecomputation.CodeLoss modulesfromtorchrl.objectivesimportDQNLossloss_module=DQNLoss(value_network=value_network,gamma=0.99)tensordict=replay_buffer.sample(batch_size)loss=loss_module(tensordict)Advantage computationfromtorchrl.objectives.value.functionalimportvec_td_lambda_return_estimateadvantage=vec_td_lambda_return_estimate(gamma,lmbda,next_state_value,reward,done,terminated)a generictrainer 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.variousrecipesto 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 ourcall for contributionsand ourcontributionpage.Examples, tutorials and demosA series ofexamplesare provided with an illustrative purpose:DQNDDPGIQLCQLTD3A2CPPOSACREDQDreamerDecision TransformersRLHFand many more to come!Check theexamples markdowndirectory for more details
about handling the various configuration settings.We also providetutorials and demosthat give a sense of
what the library can do.CitationIf 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}
}InstallationCreate a conda environment where the packages will be installed.conda create --name torch_rl python=3.9
conda activate torch_rlPyTorchDepending 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.
Seeherefor a detailed list of commands,
includingpip3or other special installation instructions.TorchrlYou can install thelatest stable releaseby usingpip3 install torchrlThis 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).Thenightly buildcan be installed viapip install torchrl-nightlywhich 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, callpip3 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 (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.)To install the library locally, start by cloning the repo:git clone https://github.com/pytorch/rlGo to the directory where you have cloned the torchrl repo and install it (after
installingninja)cd /path/to/torchrl/
pip install ninja -U
python setup.py develop(unfortunately,pip install -e .will not work).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 tryARCHFLAGS="-arch arm64" python setup.py developTo run a quick sanity check, leave that directory (e.g. by executingcd ~/)
and try to import the library.python -c "import torchrl"This should not return any warning or error.Optional dependenciesThe 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 wandbTroubleshootingIf aModuleNotFoundError: No module named ‘torchrl._torchrlerrors 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.One common reason might be that you are trying to import torchrl from within the
git repo location. The following code snippet should return an error if
torchrl has not been installed indevelopmode:cd ~/path/to/rl/repo
python -c 'from torchrl.envs.libs.gym import GymEnv'If this is the case, consider executing torchrl from another location.If you're not importing torchrl from within its repo location, it could be
caused by a problem during the local installation. Check the log after thepython setup.py develop. One common cause is a g++/C++ version discrepancy
and/or a problem with theninjalibrary.If the problem persists, feel free to open an issue on the topic in the repo,
we'll make our best to help!OnMacOs, we recommend installing XCode first.
With Apple Silicon M1 chips, make sure you are using the arm64-built python
(e.g.here).
Running the following lines of codewget https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py
python collect_env.pyshould displayOS: macOS *** (arm64)and notOS: macOS **** (x86_64)Versioning issues can cause error message of the typeundefined symboland such. For these, refer to theversioning issues documentfor a complete explanation and proposed workarounds.Asking a questionIf 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
thePyTorch forum.ContributingInternal collaborations to torchrl are welcome! Feel free to fork, submit issues and PRs.
You can checkout the detailed contribution guidehere.
As mentioned above, a list of open contributions can be found inhere.Contributors are recommended to installpre-commit hooks(usingpre-commit install). pre-commit will check for linting related issues when the code is committed locally. You can disable th check by appending-nto your commit command:git commit -m <commit message> -nDisclaimerThis 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.LicenseTorchRL is licensed under the MIT License. SeeLICENSEfor details.