Documentation | Tensordict | Caractéristiques | Exemples, tutoriels et démos | Citation | Installation | Poser une question | Contributif
TORCHRL est une bibliothèque de renforcement open source (RL) pour Pytorch.
Lisez le papier complet pour une description plus organisée de la bibliothèque.
Vérifiez nos tutoriels de mise en place pour accélérer rapidement avec les fonctionnalités de base de la bibliothèque!
La documentation TORCHRL peut être trouvée ici. Il contient des tutoriels et la référence de l'API.
Torchrl fournit également une base de connaissances RL pour vous aider à déboguer votre code, ou simplement à apprendre les bases de RL. Vérifiez-le ici.
Nous avons des vidéos d'introduction pour que vous puissiez mieux connaître la bibliothèque, consultez-les:
Torchrl étant agnostique du domaine, vous pouvez l'utiliser sur de nombreux champs différents. Voici quelques exemples:
TensorDict
Les algorithmes RL sont très hétérogènes, et il peut être difficile de recycler une base de code entre les paramètres (par exemple de la ligne en ligne, de l'apprentissage basé sur l'État à l'apprentissage basé sur les pixels). Torchrl résout ce problème via TensorDict
, une structure de données pratique (1) qui peut être utilisée pour rationaliser la base de code RL. Avec cet outil, on peut écrire un script de formation PPO complet en moins de 100 lignes de 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 = { "low" : - 1.0 , "high" : 1.0 },
return_log_prob = True
)
buffer = TensorDictReplayBuffer (
storage = LazyTensorStorage ( 1000 ),
sampler = SamplerWithoutReplacement (),
batch_size = 50 ,
)
collector = SyncDataCollector (
env ,
actor ,
frames_per_batch = 1000 ,
total_frames = 1_000_000 ,
)
loss_fn = ClipPPOLoss ( actor , critic )
adv_fn = GAE ( value_network = critic , average_gae = True , gamma = 0.99 , lmbda = 0.95 )
optim = torch . optim . Adam ( loss_fn . parameters (), lr = 2e-4 )
for data in collector : # collect data
for epoch in range ( 10 ):
adv_fn ( data ) # compute advantage
buffer . extend ( data )
for sample in buffer : # consume data
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 } " )
Voici un exemple de la façon dont l'API de l'environnement s'appuie sur Tensordict pour transporter des données d'une fonction à une autre lors d'une exécution de déploiement:
TensorDict
facilite la réutilisation de code de code entre les environnements, les modèles et les algorithmes.
Par exemple, voici comment coder un déploiement dans 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
En utilisant cela, Torchrl résume les signatures d'entrée / sortie des modules, Env, les collectionneurs, les tampons de relecture et les pertes de la bibliothèque, permettant à toutes les primitives d'être facilement recyclées entre les paramètres.
Voici un autre exemple d'une boucle de formation hors politique dans Torchrl (en supposant qu'un collecteur de données, un tampon de relecture, une perte et un optimiseur ont été instanciés):
- 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()
Cette boucle de formation peut être réutilisée sur les algorithmes car il fait un nombre minimal d'hypothèses sur la structure des données.
Tensordict prend en charge plusieurs opérations de tenseur sur son appareil et sa forme (la forme du tendict, ou sa taille de lot, est la première dimension arbitraire arbitraire de tous ses tenseurs contenus):
# stack and cat
tensordict = torch . stack ( list_of_tensordicts , 0 )
tensordict = torch . cat ( list_of_tensordicts , 0 )
# reshape
tensordict = tensordict . view ( - 1 )
tensordict = tensordict . permute ( 0 , 2 , 1 )
tensordict = tensordict . unsqueeze ( - 1 )
tensordict = tensordict . squeeze ( - 1 )
# indexing
tensordict = tensordict [: 2 ]
tensordict [:, 2 ] = sub_tensordict
# device and memory location
tensordict . cuda ()
tensordict . to ( "cuda:1" )
tensordict . share_memory_ ()
Tensordict est livré avec un module tensordict.nn
dédié qui contient tout ce dont vous pourriez avoir besoin pour écrire votre modèle avec. Et c'est compatible functorch
et torch.compile
.
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"]
La classe TensorDictSequential
permet de ramifier les séquences des instances nn.Module
d'une manière très modulaire. Par exemple, voici une implémentation d'un transformateur à l'aide de l'encodeur et des blocs de décodeur:
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
permet d'isoler les sous-graphiques en interrogeant un ensemble de touches d'entrée / sortie souhaitées:
transformer . select_subsequence ( out_keys = [ "memory" ]) # returns the encoder
transformer . select_subsequence ( in_keys = [ "tgt" , "memory" ]) # returns the decoder
Vérifiez les tutoriels Tensordict pour en savoir plus!
Une interface commune pour les environnements qui prend en charge les bibliothèques communes (Openai Gym, DeepMind Control Lab, etc.) (1) et l'exécution sans état (par exemple, les environnements basés sur des modèles). Les conteneurs des environnements lots permettent l'exécution parallèle (2) . Une classe pytorch courante de classe de spécification du tenseur est également fournie. L'API des environnements de Torchrl est simple mais stricte et spécifique. Consultez la documentation et le tutoriel pour en savoir plus!
env_make = lambda : GymEnv ( "Pendulum-v1" , from_pixels = True )
env_parallel = ParallelEnv ( 4 , env_make ) # creates 4 envs in parallel
tensordict = env_parallel . rollout ( max_steps = 20 , policy = None ) # random rollout (no policy given)
assert tensordict . shape == [ 4 , 20 ] # 4 envs, 20 steps rollout
env_parallel . action_spec . is_in ( tensordict [ "action" ]) # spec check returns True
Multiprocess et collecteurs de données distribués (2) qui travaillent de manière synchrone ou asynchrone. Grâce à l'utilisation de Tensordict, les boucles d'entraînement de Torchrl sont rendues très similaires aux boucles de formation régulières dans l'apprentissage supervisé (bien que le "DatalOader" - lire le collecteur de données - soit modifié à la volée):
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_ ()
Consultez nos exemples de collection distribués pour en savoir plus sur la collecte de données ultra-rapides avec Torchrl.
Tampons de relecture efficaces (2) et génériques (1) avec stockage modularisé:
storage = LazyMemmapStorage ( # memory-mapped (physical) storage
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 , # multi-threaded sampling
storage = storage
)
Des tampons de relecture sont également proposés comme emballages autour des ensembles de données communs pour RL hors ligne :
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 : # or alternatively sample = data.sample()
fun ( sample )
Transformes d'environnement croisé (1) , exécutées sur l'appareil et de manière vectorisée (2) , qui traitent et préparent les données provenant des environnements à utiliser par l'agent:
env_make = lambda : GymEnv ( "Pendulum-v1" , from_pixels = True )
env_base = ParallelEnv ( 4 , env_make , device = "cuda:0" ) # creates 4 envs in parallel
env = TransformedEnv (
env_base ,
Compose (
ToTensorImage (),
ObservationNorm ( loc = 0.5 , scale = 1.0 )), # executes the transforms once and on device
)
tensordict = env . reset ()
assert tensordict . device == torch . device ( "cuda:0" )
D'autres transformations comprennent: la mise à l'échelle de récompense ( RewardScaling
), les opérations de forme (concaténation des tenseurs, le non-quête, etc.), la concaténation des opérations successives ( CatFrames
), le redimensionnement ( Resize
) et bien d'autres.
Contrairement à d'autres bibliothèques, les transformations sont empilées en tant que liste (et ne sont pas enveloppées les unes dans les autres), ce qui facilite les ajouter et les supprimer à volonté:
env . insert_transform ( 0 , NoopResetEnv ()) # inserts the NoopResetEnv transform at the index 0
Néanmoins, les transformations peuvent accéder et exécuter des opérations sur l'environnement parent:
transform = env . transform [ 1 ] # gathers the second transform of the list
parent_env = transform . parent # returns the base environment of the second transform, i.e. the base env + the first transform
Divers outils d'apprentissage distribué (par exemple les tenseurs mappés de mémoire) (2) ;
diverses architectures et modèles (par exemple, acteur-critique) (1) :
# create an nn.Module
common_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 output
common_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 space
policy_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 it
policy_module = SafeProbabilisticTensorDictSequential ( # stochastic policy
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 ,
)
# Wrap the policy and value funciton in a common module
actor_value = ActorValueOperator ( common_module , policy_module , value_module )
# standalone policy from this
standalone_policy = actor_value . get_policy_operator ()
Emballage et modules d'exploration pour échanger facilement entre l'exploration et l'exploitation (1) :
policy_explore = EGreedyWrapper ( policy )
with set_exploration_type ( ExplorationType . RANDOM ):
tensordict = policy_explore ( tensordict ) # will use eps-greedy
with set_exploration_type ( ExplorationType . DETERMINISTIC ):
tensordict = policy_explore ( tensordict ) # will not use eps-greedy
Une série de modules de perte efficaces et de rendement fonctionnel et de calcul d'avantage hautement vectoriel.
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 )
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 )
Une classe de formateurs génériques (1) qui exécute la boucle de formation susmentionnée. Grâce à un mécanisme d'accrochage, il prend également en charge toute opération de journalisation ou de transformation des données à tout moment.
Diverses recettes pour créer des modèles qui correspondent à l'environnement déployé.
Si vous pensez qu'une fonctionnalité est absente de la bibliothèque, veuillez soumettre un problème! Si vous souhaitez contribuer à de nouvelles fonctionnalités, consultez notre appel aux contributions et notre page de contribution.
Une série d'implémentations de pointe relève d'un objectif illustratif:
Algorithme | Compiler le support ** | API sans tenordict | Pertes modulaires | Continu et discret |
Dqn | 1,9x | + | N / A | + (via ActionDiscretizer Transform) |
Ddpg | 1.87x | + | + | - (continu uniquement) |
Iql | 3.22x | + | + | + |
Cql | 2.68x | + | + | + |
Td3 | 2.27x | + | + | - (continu uniquement) |
TD3 + BC | non testé | + | + | - (continu uniquement) |
A2C | 2.67x | + | - | + |
PPO | 2.42x | + | - | + |
SAC | 2.62x | + | - | + |
Redq | 2.28x | + | - | - (continu uniquement) |
Dreamer V1 | non testé | + | + (différentes classes) | - (continu uniquement) |
Transformateurs de décision | non testé | + | N / A | - (continu uniquement) |
Crossq | non testé | + | + | - (continu uniquement) |
Gail | non testé | + | N / A | + |
Impala | non testé | + | - | + |
IQL (MARL) | non testé | + | + | + |
DDPG (MARL) | non testé | + | + | - (continu uniquement) |
PPO (MARL) | non testé | + | - | + |
Qmix-vdn (marl) | non testé | + | N / A | + |
Sac (Marl) | non testé | + | - | + |
Rlhf | N / A | + | N / A | N / A |
** Le nombre indique une accélération attendue par rapport au mode avide lorsqu'il est exécuté sur CPU. Les nombres peuvent varier en fonction de l'architecture et de l'appareil.
Et bien d'autres à venir!
Des exemples de code affichant des extraits de code de jouet et des scripts de formation sont également disponibles
Vérifiez le répertoire des exemples pour plus de détails sur la gestion des différents paramètres de configuration.
Nous fournissons également des tutoriels et des démos qui donnent une idée de ce que la bibliothèque peut faire.
Si vous utilisez Torchrl, veuillez vous référer à cette entrée Bibtex pour citer ce travail:
@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}
}
Créez un environnement conda où les packages seront installés.
conda create --name torch_rl python=3.9
conda activate torch_rl
Pytorch
Selon l'utilisation de Functorch que vous souhaitez fabriquer, vous voudrez peut-être installer la dernière version (nocturne) Pytorch ou la dernière version stable de Pytorch. Voir ici pour une liste détaillée de commandes, y compris pip3
ou d'autres instructions d'installation spéciales.
Torchrl
Vous pouvez installer la dernière version stable en utilisant
pip3 install torchrl
Cela devrait fonctionner sur Linux, Windows 10 et OSX (Chips Intel ou Silicon). Sur certaines machines Windows (Windows 11), il faut installer la bibliothèque localement (voir ci-dessous).
La construction nocturne peut être installée via
pip3 install torchrl-nightly
que nous expédions actuellement uniquement pour les machines Linux et OSX (Intel). Surtout, les constructions nocturnes nécessitent également les constructions nocturnes de pytorch.
Pour installer des dépendances supplémentaires, appelez
pip3 install " torchrl[atari,dm_control,gym_continuous,rendering,tests,utils,marl,open_spiel,checkpointing] "
ou un sous-ensemble de ceux-ci.
On peut également souhaiter installer la bibliothèque localement. Trois raisons principales peuvent motiver ceci:
Pour installer la bibliothèque localement, commencez par cloner le repo:
git clone https://github.com/pytorch/rl
Et n'oubliez pas de consulter la branche ou la balise que vous souhaitez utiliser pour la construction:
git checkout v0.4.0
Accédez au répertoire où vous avez cloné le repo Torchrl et l'installez-le (après avoir installé ninja
)
cd /path/to/torchrl/
pip3 install ninja -U
python setup.py develop
On peut également construire les roues pour distribuer aux collègues en utilisant
python setup.py bdist_wheel
Vos roues y seront stockées ./dist/torchrl<name>.whl
et installable via
pip install torchrl < name > .whl
AVERTISSEMENT : Malheureusement, pip3 install -e .
ne fonctionne pas actuellement. Les contributions pour aider à résoudre ce problème sont les bienvenues!
Sur les machines M1, cela devrait être prêt à l'emploi avec la construction nocturne de pytorch. Si la génération de cet artefact dans MacOS M1 ne fonctionne pas correctement ou dans l'exécution le message (mach-o file, but is an incompatible architecture (have 'x86_64', need 'arm64e'))
apparaît, puis essayez
ARCHFLAGS="-arch arm64" python setup.py develop
Pour exécuter un contrôle de santé mentale rapide, laissez ce répertoire (par exemple en exécutant cd ~/
) et essayez d'importer la bibliothèque.
python -c "import torchrl"
Cela ne devrait renvoyer aucun avertissement ou erreur.
Dépendances facultatives
Les bibliothèques suivantes peuvent être installées en fonction de l'utilisation que l'on veut faire de 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
Dépannage
Si un ModuleNotFoundError: No module named 'torchrl._torchrl
ne se produit (ou un avertissement indiquant que les binaires C ++ n'ont pas pu être chargés), cela signifie que les extensions C ++ n'ont pas été installées ou non.
develop
: cd ~/path/to/rl/repo
python -c 'from torchrl.envs.libs.gym import GymEnv'
python setup.py develop
. Une cause courante est une différence de version G ++ / C ++ et / ou un problème avec la bibliothèque ninja
. wget https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py
python collect_env.py
OS: macOS *** (arm64)
OS: macOS **** (x86_64)
Les problèmes de version peuvent entraîner un message d'erreur du type undefined symbol
et autres. Pour ceux-ci, reportez-vous au document des problèmes de versioning pour une explication complète et des solutions de contournement proposées.
Si vous apercevez un bug dans la bibliothèque, veuillez soulever un problème dans ce dépôt.
Si vous avez une question plus générique concernant RL dans Pytorch, postez-la sur le forum Pytorch.
Les collaborations internes à Torchrl sont les bienvenues! N'hésitez pas à se nourrir, à soumettre des problèmes et des PR. Vous pouvez vérifier le guide de contribution détaillé ici. Comme mentionné ci-dessus, une liste de contributions ouvertes peut être trouvée ici.
Les contributeurs sont recommandés pour installer des crochets pré-engagés (en utilisant pre-commit install
). Pre-Commit vérifiera les problèmes liés à la liaison lorsque le code est engagé localement. Vous pouvez désactiver le chèque en ajoutant -n
à votre commande de validation: git commit -m <commit message> -n
Cette bibliothèque est publiée sous forme de fonctionnalité bêta pytorch. Les changements de rupture de BC devraient se produire, mais ils seront introduits avec une garantie de dépréciation après quelques cycles de libération.
Torchrl est autorisé sous la licence du MIT. Voir la licence pour plus de détails.