TorchRL: A data-driven decision-making library for PyTorch
Authors: Albert Bou
Matteo Bettini
Sebastian Dittert
Vikash Kumar
Shagun Sodhani
Xiaomeng Yang
Gianni De Fabritiis
Vincent Moens
Venue: ICLR 2024 spotlight (2024)
Keywords: Reinforcement Learning
pytorch
control
robotics
Primary area: reinforcement learning
TL;DR: We present TorchRL, a new generalistic RL and control library for PyTorch. TorchRL offers a modular, lightweight, and agnostic tool for training reinforcement learning agents and other decision-making paradigms.

Abstract: PyTorch has ascended as a premier machine learning framework, yet it lacks a native and comprehensive library for decision and control tasks suitable for large development teams dealing with complex real-world data and environments. To address this issue, we propose TorchRL, a generalistic control library for PyTorch that provides well-integrated, yet standalone components. We introduce a new and flexible PyTorch primitive, the TensorDict, which facilitates streamlined algorithm development across the many branches of Reinforcement Learning (RL) and control. We provide a detailed description of the building blocks and an extensive overview of the library across domains and tasks. Finally, we experimentally demonstrate its reliability and flexibility, and show comparative benchmarks to demonstrate its computational efficiency. TorchRL fosters long-term support and is publicly available on GitHub for greater reproducibility and collaboration within the research community. The code is open-sourced on GitHub.

Published as a conference paper at ICLR 2024
TORCHRL: A DATA-DRIVEN DECISION-MAKING LIBRARY FOR PYTORCH
Albert Bou UPF, Acellera
Matteo Bettini University of Cambridge
Sebastian Dittert UPF
Vikash Kumar Meta Al
Shagun Sodhani Meta AI
Xiaomeng Yang Meta AI
Gianni De Fabritiis ICREA, UPF, Acellera
Vincent Moens Py Torch, Meta vmoens [at meta. com
ABSTRACT
PyTorch has ascended as premier machine learning framework, yet it lacks native and comprehensive library for decision and control tasks suitable for large development teams dealing with complex real-world data and environments. To address this issue, we propose TorchRL, a generalistic control library for PyTorch that provides well-integrated, yet standalone components. We introduce new and flexible Py Torch primitive, the TensorDict, which facilitates streamlined algorithm development across the many branches of Reinforcement Learning (RI. and control. We provide a detailed description of the building blocks and an extensive overview of the library across domains and tasks. Finally, we experimentally demonstrate its reliability and flexibility and show comparative benchmarks to demonstrate its computational efficiency. TorchRL fosters long-tern support and is publicly available on GitHub for greater reproducibility and collaboration within the research community. The code open-sourced on GitHub
INTRODUCTION
Originally. breakthroughs in AI were powered by custom code tailored for specific machines and backends. However, with the widespread adoption of Al accelerators like GPUs and TPUs, as well as user-friendly development frameworks such PyTorch (Paszke al., 2019), TensorFlow (Abadi) et al. 2015), and Jax (Bradbury et al.] (018). research and application has moved beyond this point. In this regard, the field of decision- making is more fragmented than other AI domains, such as computer vision or natural language processing, where few libraries have rapidly gained widespread recognition within their respective research communities. We attribute the slower progress toward standardization to the dynamic requirements of decision-making algorithms, which create trade-off between modularity and component integration.
The current solutions in this field lack the capability effectively support its wide range of applica- tions, which include gaming (Mnih et al.,
2013
Silver et al.) 2016 OpenAl et al. 2019b Espeholt
al. 2018). robotic control (Tobin al.
2017
OpenAl et al. (0191)) autonomous driving (Aradi
2022
finance (Ganesh et al. | (219)) bidding in online advertisement (Zhu & Roy 2202]) cooling
system control (Luo al.] (022), faster matrix multiplication algorithms (Fawzi al.
2022) or
chip design (Mirhoseini et al. (02T). Decision-making is also powerful tool to help train other
models through AutoML solutions (He et al. [2018] Zoph & Le. 2017 Zoph et al., 2018 Baker et 2017 Zhong et al. 2018) or incorporate human feedback in generative model fine-tuning (Christiano
al. 2017 von Werra et al., 20233] Nakano et al. (222). ) and is also used techniques such as
riml transfer (TTbii et al.]2201 Peng et al. 2018 Rudin et al., [022 OpenAl et al.) 2019a).
model-based RL (Brunnbauer et al. 2002 Wu et al. | Hansen et al. 2022a b) or offline data collection for batched RL and imitation learning (Kalashnikov et al. 2018 Shah et al. 2022). In
addition, decision- making algorithms and training pipelines must often meet challenging scaling requirements, such distributed inference (Mnih al. 2016:) Espeholt et al. 2018) distributed model execution (Nakano et al. 202T). They must also satisty imperative deployment demands when used on hardware like robots autonomous cars, which often have limited computational

---

Published as a conference paper at ICLR 2024
Environment
env (menn("Pendolun--0
observation
# Model: Actor and value
actiion
nlp_actor MLPInun_cclls-Gi degth-3,
in_ teaturesss,, out, featureat-)--1)
actor rensorD.cttotuleennnp.ettt
in_keys- fo__kess=[=obeen "oess=("oservaaionn] out out_keyss[""action]
nip_value MLPinun_cellsmüi, degthw2,
1n_ mo_teaturesws, but_features=-)
critic - Tonsortiictcquunttaallcooo
TensorDictModule le(tlodu (eiineeennlle,,
in_keyse ["observation", "action"], out keys=["staa action vateaaccvoouueee11_ value]]a
Deservat
coservation
cone
# Data CoLlector
co collector Synchataco
env, additiveeauussiiiiaaaprrr
observat
PELE
franes_per, batch=1000, total_ (ranestrame)))00000))
observattor
# Replay Buffer
buffer ConsorDictReplayBuffo
storageelayyTensorrSttto (198030))
Loss Module
loss fn D0PPLosssaccor critic)
optin torch.optin.Adam,
joss_fn.paranceers(/. 1r=2e-4)
* Train
led tesinnediat
for data in collector: # collect data
(affe # populate RB sample buffer. sarple(50)
on
mini-batch
observation reward
FHI
loss loss_fn(sanple) loss
loss loss("toss_actor) + "assl"loss_vaue"" loss. www.bkkkaardl step() optin. sriim.crr_ggrd11
Actor
Modale
Replay Butfer
Figure I: TorchRL overview. The left side showcases the key components of the library, demonstrating the data flow with TensorDict instances passing between modules. On the right side, a code snippet is provided as a toy example, illustrating the training of DDPG. The script provides users with full control over the algorithm's 's hyperparameters, offering a concise yet comprehensive solution. Still, replacing minimal number of components in the script enables seamless transition another similar algorithm, like SAC or REDQ.
resources and need to operate in real-time. Many libraries are either too high-level and difficult repurpose across use cases, or low-level for non-experts to use practically. Because of this intricate landscape, consensus adoption of frameworks within the research community remains limited (Brockman al. 2016). We believe that moving forward, the next generation of libraries for decision-making will need to be easy to use and rely on widespread frameworks while still being powerful, efficient, scalable, and broadly applicable, to meet the demands of this rapidly evolving and heterogeneous field. This is the problem that TorchRI addresses.
Training decision-makinn making models involves a sequentially driven algorithmic process, where the cooperation between framework components is vital for adaptability across different applications. For example, actors, collectors and replay buffers need to able to handle various modalities. Imposing fixed input/output format would restrict their reusability while eliminating all restrictions would create significant difficulties in designing adaptable scripts. Consequently, conventional decision- making libraries must choose between a well-integrated framework with high-level API or flexible framework with basic building blocks of limited scope. The former hinders quick coding and testing of new ideas, while the latter leads complex code and poor user experience. TorchRL effectively tackles these challenges with a powerful communication class that allows its components to operate independently while not limiting their scope and at the same time ensuring seamless integration with each other to flexibly create algorithms.
In this paper, we will detail our approach for truly general decision- making library, based primarily on module portability. efficient communication and code repurposing. We introduce TensorDict, a new Py PyTocc data carrier, which is central to achieving this goal. By adopting this new coding paradigm, we demonstrate how simple it is to reuse ML scripts and primitives across different tasks or research domains. Building upon this innovative technical foundation, we explore various independent components that rely on this abstraction. These components encompass range of advanced solutions, including distributed data collectors, cutting-eds models like Decision Transformers, and highly specialized classes such demonstrate that the current state the library encompasses diverse categories of decision making algorithms, caters wide range of applications, performs well on both small-scale machines and
environment API compatible with single and multi-agent problems. We
2

---

Published as a conference paper at ICLR 2024
distributed clusters, and strikes balance in terms of technical requirements that makes it accessible to practitioners with multiple levels of expertise.
Related work. Currently, there is vast array of libraries available for reinforcement learning. Some of these libraries are specifically designed for particular niches, such as Pyqlearning (Accel
Brain 2023) and Facebook/ELF (Tian al. 2017), only support small number of environments, like OpenSpiel(Laancoot et al. 2019) and reaver (Ring 2018), or lack the capability for parallel environments with distributed learning, such as MushroomRL (D' DrEamo et al. (021). Additionally, many previously popular libraries are no longer actively maintained, including. KerasRL (Plappert
2016), rlpyt (Stooke & Abbeel 2 ReAgent (Gauci et al. 2018). IntelLabs/coach (Caspi et al.
2017). Dopamine (Castro et al.
2018), ChainerRL (Fujita et al. 2021.. PyTorchRL (Bou et 20222.
2020). Therefore, for libraries aspiring to serve as comprehensive tools for decision-makinn it is crucial to offer robust long-term support while also being designed to ensure adaptability to the continuously evolving
TTAAeett (Guadarrama et al.. 2018), and autonomous- learning-
landscape of decision decision-making research.
Among currently popular solutions, there's s a strong trend toward high-level approaches that make easier to integrate different parts. This is particularly necessary for tools that have components that can't 't stand alone or rely on limited data carriers and require some extra code to make everything work smoothly together. Some well-know examples of these tools include Stable-baseling
(an Stable-
baselines3,, SB3) (Raffin et al. (021). ElegantRI. (Liu et al. (022)) Acme (Hoffman et al., 2020).
DI-engine (engine Contributors (021), garage (garage contributors] (019)), Tensorforce (Kuhnle et al. 2017), and RLlib (Liang et al. 2018). While high-level approaches can be welcoming to newcomers, they often obscure the inner workings of the code, which may not suit users who seek greater control over their scripts, and difficult code repurposing. TorchRL leverages its powerful data carrier to adopt highly modular design. Our primary goal is to provide well-tested building blocks that enable practitioners to create adaptable scripts. These building blocks can also be utilized by high-level tools to cater to users of all experience levels, avoiding the need to reimplement the same functionalities again. Finally, two additional popular libraries are Tianshou (Weng et al. 2022) and CleanRL (Huang et al. 2022). Tianshou shares similarities with TorchRL, offering a wide range of RL algorithms, a relatively versatile data carrier, and modular components. However, its primary focus is on building RL trainers to simplify problem-solving without extensive domain knowledge, which we believe can be complementary efficient single-file file implementations of RL. algorithms in Torch and JAX, enhancing code clarity. Yet, it sacrifices component reusability, making scaling challenging due to code duplication.
our approach. On the other hand, CleanRL provides
TORCHRL COMPONENTS
TorchRL is made up of separate, self-contained parts. Each of these parts, referred to as components, handles specific operation in the overall data workflow. While it's not always required, these components usually communicate using a dedicated tool designed to flexibly pass data back and forth, the TensorDict. The resulting atomicity and modularity of this approach facilitate versatile usage and a combination of individual components within most machine learning workflows. In this section, we provide an overview of the key components in the library, which can be used independently or as foundational elements for constructing decision -making algorithms, as depicted in Figure
Tensor Die Introducing seamless mode of communication among independent RL algorithmic components poses a set of challenges, particularly when considering the diversity of method signa- tures, inputs and outputs. shared, versatile communication standard is needed in such scenarios to ensure fluid interactions between modules without imposing constraints individual components. We address this problem by adopting a new data carrier for PyTorch named TensorDict, packaged ponent to be developed independently of the requirements of the classes it potentially communicates with. As a result, data collectors be designed without any knowledge of the policy structure, and replay buffers do not need any information about the data structure that is to be stored within them. Figure gives an overview of practical data flow using TensorDict communication tool.
the development of
separate open-source lightweight library. This library enables every com-
3

---

Published as a conference paper at ICLR 2024
A TensorDict is dictionary- liike object that stores tensors(-like) objects. It includes additional features that optimize its use with PyTorch. One of its notable features is the ability to handle the batch size (in contrast with the "feature" size), which means that it can be indexed not only by keys but also by shapes. The adopted convention is that a TensorDict hatch size represents the common leading dimensions of the tensors it contains. In the context of RI., the batch-size can include the minibatch or time dimension as well as the number of agents, tasks or processes. The shape-key duality is reflected in TensorDict's ability to handle both key-based and numerical indexnng. For both, this indexing can be performed along multiple dimensions (through nested tensordicts along the key-dimensions, and through the multiple tensor dimensions along the shape-dimension).
TensorDict provides a whole stack of extra functionality that naturally follows from these two basic features. These can roughly be split into shape-based and key-based methods. In the first bucket, one can find torch. Tensor-like features such as reshaping, (un)sequeezing, stacking, concatenating, masking or permuting. The latter contains methods that allow flatten or unflatten the tensordict structure (lo represent flat structure keys, Other utilities include methods to move a TensorDict from one device to another, point-to- point communication in distributed settings and many more. The class also offers set of efficient storage interfaces through memory- mapped tensors or HDF5 files, allowing to manipulate big data structures with little effort. More details about TensorDict functionality, well as benchmarks against comparable solutions, can be found in Appendix B
nested one or inversely), renaming or deleting
TensorDict enables the design of functions and classes with a generic signature, where we can safely restrict the input and output to be both TensorDict instances. This comes in handy in countless cases: straightforward example is the support that TorchRL provides for different simulator backends such as Gym (Brockman et al.] 2016), IsaacGym (Makoviychuk et al., 202T), DeepMind Control (Tunyasuvunakool et al., 2020). Brax (Freeman et al. (022)) and others. Each of these libraries has a dedicated step signature that is unified under the TorchRL EnvBase where the relative method expects a TensorDict instance as input and writes the data in it in standardized format. The input types supported are almost universal, which allows TorchRL's environments to receive complex data structures as input, thereby enabling a common API for stateful and stateless environment. Regarding the outputs, the info dictionary returned by Gym-like environment will is expanded in the output, making the whole step results readily available in a consistent way across libraries and environments. In practice, this means that in TorchRL, the same training script can be repurposed with minimal effort to train environments from various backends and domains.
TensorDict does not only provide clearer code but also comes with some performance improvements
that include better memory management, more efficient (a)synchronous point-to-point communication in distributed settings and device casting. Finally, the tensordi- ct library also comes with dedicated tensordict nn package that rethinks the way one designs PyTorch models. Various allow design of
primitives, such as TensorDicto cetsorriccMMuulle and TensorDict TensorriictSeqenttii complex PyTorch operations in an explicit and programmable way. TensorDict's nn primitives are fully compatible with torch. comp ee through a dedicated symbolic tracer available in the TensorDict package, making TorchRL itself compatible with the latest features of PyTorch 2.1.
Environment API, wrappers, and transforms. OpenAl Gym (Brockman et al. 2016) has become the most widely adopted environment interface in RL: its standardization of the environment API was a major leap forward that enabled research reproducibility. Its simplicity (only two functions are exposed during interaction) drove its success, but it presents several drawbacks such as fixed tuple signature for its step method or its heavy reliance on wrappers (see below).
The TorchRL environment interface aims to maintain the simplicity of Gym while addressing these drawbacks. As with Gym, only the reset () and step () methods are required to interact with these objects. Both of these methods rely on TensorDict to facilitate their integration, as discussed in the TensorDict section and presented FFiuure 1 TensorDicts enable carrying multidimensional like data, allowing batched/vectorized simulation (Freeman al., 2021 I Bettini et al. (022) and multidimensional input/output spaces. The native support for batched data shapes and nested data structures make TorchRL compatible by default with multi-agent and multi-task applications, where these features are key for a clear environment API.
Consequently, TorchRL is by means extension of Gym or any other simulation library, unlike other libraries that cover one specific simulator but not others. The generic environment API allows to easily support multitude of existing simulators: Gym and Gymnasium (Brockman al.) 2016) since
4

---

Published as a conference paper at ICLR 2024
v0.13, DMControl Tunyasuvvnakkoo et al. 2020). Habitat (Szot et al. 2021). 1). RoboHive (Rob] 2020
OpenML datasets \Vanschoren et al.] 2013). D4RL datasets (Fu et al. (2020) Brax (Freeman et al.
202I), Isaac (Makoviychuk et al. 20.1). Jumanji (Bonnet et al. and VMAS (Bettini et al. 2022)
some examples, but the list is constantly growing. Notably, the last four of these are vectorized and can pass gradients through the simulations, allowing compute reparameterized trajectories. We provide some more technical information about the environment API and its usage in Appendix C
TorchRL draws inspiration from other components of the PyTorch ecosystem (maintainers & con- tributors 2016: Paszke et al. 2019) which rely on the concept of transform sequences to modify module outputs. Significantly. TorchRL's Transforms, which should be seen as regular an. Module instances that can be applied wherever data transformation is required. Transforms are versatile and can be utilized in various components, including replay buffers, collectors, and even ported from environment to models, effectively connecting the training process with real-world applications. Examples include data transformations (e.g. resizing, cropping), target computation (reward-to-go) or even embedding through foundational models (Nair et al. 2022; Ma et al. (022). As we show in the Appendix, transforms offer a way to dynamically manipulate ate data at least as flexible as wrapping classes typically used in RL libraries.
Data collectors. TorchRL has dedicated classes for data collection that execute a policy in one or multiple batched environments and return batches of transition data in TensorDict format. Data collectors iteratively compute the actions to be executed, pass those to the environments (real or simulated), and can handle resetting when and where required. These collectors are designed for case of use, requiring only one or more environment constructors, a policy module and a target number of steps. However, developers can also specify the asynchronous or synchronous nature of data collection, the number of parallel environments, the resource allocation (e.g., GPU, number of workers) and the data postprocessing, giving them fine control and flexibility over the collection process. TorchRL offers distributed components that enable data collection at scale by coordinating multiple workers in a cluster. These components are compatible with various backends like gloo and multiple launchers and resource management solutions
or NCCL. through including submitit (Incubator (021) and Ray (Moritz et al., 2018). As scaling inevitably adds complexity, the distributed solutions in TorchRL are designed independent components that provide the same interface and data control as non-distributed components. This enables practitioners to work locally on projects with complete independence while also providing an easy way to scale up to distributed projects by just replacing non-distri non-distributed components with distributed counterparts, which will increase data collection throughput. Crucially, TorchRL distributed dependencies are optional and not required for non-distributed components.
Replay buffers and datasets. Replay buffers (RBs) crucial elements of RL that enable agents to learn from past experiences by storing, processing, and resampling data. However, creating flexible RB class that caters to various use cases without duplicating implementations can be challenging. To overcome this, TorchRL provides single RB implementation that offers complete composabilitty. Users can define distinct components for data storage, batch generation, pre-storage and post-sampling transforms, allowing independent customization of each aspect.
The RB class has a default constructor that creates buffer with a generic configuration, utilizing a versatile list storage and a sampler that generates batches uniformly. However, users can also specify the storage to be contiguous in either virtual or physical memory. This format provides faster performance and can handle data sizes up to terabytes. These features allow the creation of buffers that can be populated on the fly during training or used with static datasets for offline RL. TorchRL offers range of downloadable datasets for this purpose (such D4RL (Fu et al., 2020) or OpenML (Vanschoren et al., 2013)). Alternative samplers are also available, including a sampler without replacement that ensures all data is presented once before repeating, or C++ optimized prioritized sampler. RBs can store any Python object in their native form using the default constructor. Nevertheless, we encourage presenting the data in a TensorDict format, which streamlines the workflow (see Figure I) and benefits from TorchRL's optimized storage, sampling and transform techniques. TorchRL also integrates remote replay buffer component to gather data from distributed workers. This functionality is detailed in Appendix D)
Modules and models. In RL. the model architectures are distinct not only from those in other machine learning disciplines but also within RL itself. This creates two interconnected issues: Firstly, there is considerable architectural variability to contend with. Secondly, the nature of inputs and
5

---

Published as a conference paper at ICLR 2024
outputs for these networks can fluctuate based the specific environment and algorithm in use. Thus, both challenges necessitate a flexible and adaptable approach.
TorchRL responds the first problem by providing RL-dedicated neural network architectures, which are organically integrated into Py as native nn. .n.Mullles These are available at varying levels of abstraction, ranging from foundational building blocks like Multilayer Perceptron (MLP), Convolu- tional Neural Networks (CNN), Long short-term memory (LSTM) modules and Transformer, to com- or MorldModeeWreer
prehensive, high-level structures such as www.rrriiiiiitrrreeeoooooooeeeae but also exploration modules like NoisyLinear or Planners like CEMPlanne lanner.
To tackle the second challenge, TorchRI uses specialized TensorDi ctModule and TensorDict ure 6). TenscrDictMeculee wraps PyTorch modules, transforming them into tensordict- compatible objects for effortless integration into the TorchRL framework. Concurrently, concatenates TensorDictModules, functioning similarly to nn. Sequential. Importantly, these TensorDict modules maintain full compatibility with torch 2.0, torch. compiee and functorch. Vectorized maps (vmap) are also well-supported, enabling the execution of multiple value networks in a vectorized manner. All instances of TensorDietModus are functional and stateful, permitting parameters to be optionally inputted without additional transformations. This facilitates the design of meta-RLLen Alw.aaaggtttii algorithms within the TorchRL library.
primitives subclassed from the TensorDict library (see Fig TensorDiictsequetiiaa
Objectives and value estimators. In TorchRL, objective classes are stateful components that track trainable parameters from torch. models and handle loss computation, a erucial step for model optimization in any machine learning algorithm. These components share a common signature defined by two main methods: a constructor method that accepts the models and all relevant hyperparameters. and a forward computation method that takes a TensorDict of collected data samples as input and returns another TensorDict with one or more loss terms. Loss modules can accommodate any nn Module and process inputs that are not presented as TensorDict as it is shown in Appendix E
This design choice makes objective classes very versatile components, abstracting the implementation of a vast array of loss functions derived from various families of RL algorithms using a single template. These include on-policy algorithms like Advantage Actor-Critic (A2C) (Mnih al., 2016) or Proximal Policy Optimization (PPO) (Schulman et al. 2017), off-policy algorithms like Deep Q-Network (DQN) (Mnih et al.) 2013), distributional DQN (Bellemare al. (017)) and subsequent improvements of these (Hessel et al. 2017), Deep Deterministic Policy Gradient (DDPG) (Lillicrap)
ctt al., 2015). Twin Delayed Deep Deterministic Policy Gradient (TD3) (Fujimoto et al. 2018a), Soft
Actor-Criti (SAC) (Haarnoj et al.) 2018) Randomized Ensembled Double Q-learning (REDQ)
(Chen et al. 2021b), offline algorithms like Implicit Q-Learning (Kostrikov et al. 2021), Decision Transformer (Chen et al. 2021a)), and model-based algorithms like Dreamer (Hafner et al.) 2019).
Finally, TorchRL also provides a solution for state value estimations, a crucial step in the loss computation of many RL algorithms that can take various forms. As for the modules, the value esti- mators are presented both in a functional, explicit way as well as an encapsulated 'TensorDectModule class that facilitates their integration in TorchRL-based scripts. These replaceable objects, named ValueEst imat ors, can be called outside the objective class or dynamically during the objective forward call, providing a versatile solution accommodate different forms of loss computation.
ECOSYSTEM
Documentation, bnowledge-base, engineering. TorchRL has rich, near-complete documentation: each public class and function must have a proper set of docstrings. At the time of writing. the coverage is above 90%, and tested across Linux, MacOs and Windows platforms. All the optional dependencies (eg, environment backends or loggers) are being tested in dedicated workflows. Com- patibility with all versions of Gym is guaranteed starting from v0.13, including the latest transition to Gymnasium. If both libraries are present in the virtual environment, dedicated set of utilities can be used to control which backend to pick. Several runtime benchmarks have been put in place to ensure that modules keep a high high-efficiency standard. The ecosystem also includes datasets, issues- issues-tracking, tracking, model library and code examples. External forums, such as the Py Torch forum, are monitored for issues related to the library.

---

Published as a conference paper at ICLR 2024
Take-up. TorchRL has seen rapid growth since its initial open-sourcing At the time of writing, the library has more than 120 collaborators, is close to reaching 1500 stars on GitHub and has been adopted by many teams of researchers and practitioners. In Appendix J we express our gratitude to the contributors who have helped in the development and enhancement of TorchRI.
Applications. TorchRL is versatile library that prioritizes comprehensive coverage, striving to address a diverse array of decision-making scenarios. As demonstrated previously, its components exhibit remarkable adaptability, allowing it to fulfill its objective. Fine-tuned, simple examples of popular algorithms testify of the library's 's versatility. Illustrated through finely-tuned, straightforward examples of popular algorithms, the library's versatility becomes evident.
The scope of coverage extends from off off-policy model-free RL (eg, DQN (Mnih et al. 2013), Rain-
bow (Hessel et al. 2017). SAC (Haarnoja et al. 2018). DDPG (Lillicrap al. 2015). TD3 ( (Fujimoto
et al. 2018b) to on- on-policy model-free RL (PPO (Schulman et al. (017). A2C (Mnih et al. 2016)).
Offline RL, which has garnered significant attention, is also exemplified with algorithms such as
IQL (Tan 1993) and CQL (Kumar et al. 2020). Decision Transformer (Chen et al. 202 and
Online-Decision Transformer (Zheng et al. 2022). ) which serve as valuable resources to assist users in developing their own solutions. We ensure compatibility with model-based algorithms through state- of-the-art implementation of Dreamer (Hafner et al., 2019), which contrasts with most other libraries that typically only provide either model-free model- algorithms in isolation. The majority of TorchRL's components seamlessly integrate with Multi-Agent Reinforcement Learning (MARL) paradigms, and the library also offers specialized MARL components. The extent of our solution spectrum can be gauged through the rigorous benchmarking of six MARL algorithms in (Appendix H The library's primitives are also fit for training on distributed settings, as our implementation and results of IMPALA Espeholt et al. (2018) show.
Finally, RL from Human Feedback (RLHF) (Nakano et al. 2021) is also a feature of the library. Several modules and helper functions make it easy to interact with third party libraries such as Hugging-Fare Face transformers (Wolf et al., (020) or datasets (Lhoest al. 202T) fine-tune generative models with little effort. These examples showcase the wide applicability and usefulness of TorchRL and its underlying components. TorchRL, is currently used by some research groups on hardware.
4 RESULTS
In this section, we experimentally showcase some of our library's 's key features. We focus demon- strating code reliability, scalability, flexibility in supporting multiple decision-making paradigms, and component efficiency. In the repository, make available all scripts used in this section.
Online single-agent RL. The most common use case covered by decision making libraries in online RL. We conduct experiments using our components reproduce well-known results for A2C (Mnih eet al. 2016)), PPO (Schulman et al. 2017). DDPG (Liliicrap et al. 2015). TD3 (Fujimoto et al. 2018b), and SAC (Haarnoja al.] 2018) algorithms. These experiments are reported in Table I In each case, we closely follow the original implementations (including network architectures, hyperparameters and number of training steps) and obtain results that match those of their original papers.
Our training times, with implementations designed mainly to reproduce established benchmarks, including the architectures, are similar to those of other libraries when using a modern desktop GPU. For MuJoCo experiments, on-policy algorithms usually complete in about hour, while off off-co algorithms typically range from to hours. the case of Atari experiments, the training time is approximately 8 hours. Note that TorchRL offers additional options to accelerate training times, as can be seen for example in Table 5
which we did not use to maintain fidelity with the original works.
Distributed RL. TorchRL's distributed components enable the replication of scalable methods. As an illustration, we provide results for the IMPALA (Espeholt et al.) 2018)) algorithm, which utilizes distributed data collection and centralized learner. Our implementation adheres to the original paper's ' specifications, and we have validated it across several Atari environments, as shown also in More details as well as training plots are provided in suussecction G.1
Table I
Offline RL. TorchRL also includes support for offline methods. In Table 2 results are available for implicit Q-Learning (IQL.) (Kostrikov al.] 2021), the Decision Transformer (DTT (Chen al.) 2021a)
7

---

Published as a conference paper at ICLR 2024
Table Experimental training of multiple on-policy and off-policy algorithms. We run each training 5 times with different seeds and report the mean final reward and std.
HalfCheetahi
Hopper
Walker2D'
Ant'
Pong
Freeway!
Boxing'
Breakeut
A2C
836--964
4934-92
3814-109
54+-20
20.5740.65
30.58++144 1.40
88.76-+6.44
375.65-447.34
PPO
2770=821
1703+742
3336=9928
2469:46006
20.52+0..8
33.31+0.95
98.31+0093
335.71+-46.72
DDPG
10433=357
914=144
16834761
1169:1658
TD3
10285+837
2809+618
4409+256
5373-+165
SAC
11077::323
29631:644
4561:192
3467:1:54
IMPALA
20.5440.18
0.0+0.0
99.1940.54
525.57-+155477
Our implementations compute the Generalized Advantage Estimator (GAE) at every epoch.
MuJoCo environments v3, trained for IM steps.
Atari 2600 environment, trained for 40M frames (10M steps) for A2C and PPO and 200M for IMPALA.
Table 2:
Experimental training of offline algorithms. We run each training 5 times with different seeds and report the mean final reward and std.
HalfChectaht Hoppert
HalfChectah"
Hopper!
HalfChcetahl
Hopper!
DT 4916+33
1048+226
oDT 4968+:58
2830=96
IQL
4864+147
1418+131
D4RL offline RL (medium replay) datasets, trained for 50000 gradient updates.
and the Online Decision Transformer (oDT) (Zheng et al. 2022 on 2 MuJoCo environments, with all three implemented following their respective original papers. Plots are available subsection G.2
Multi-agent. To bootstrap the adoption of TorchRL for multi-agent use cases, we provide imple- mentations for several state-of-tt sfate-of-thh-ar MARL algorithms and benchmark them on three multi-robot
coordination scenarios in the VMAS (Bettini et al., (2022) simulator. The results, reported in Figure 2 show the correctness of TorchRL's implementations, matching the results from (Bettini et al. 2022). In the repository, in addition to making available the scripts used in this evaluation, examples and tutorials to illustrate how TorchRL can be scamlessly used in MARL contexts. Further details on MARL experiments and comparisons with other libraries are available in (Appendixx H
provide
Vectorized data collection. The philosophy of TorchRL and its core component, TensorDict, enables seamless compatibilityy with vectorized (batched) simulators (such IsaacGym (Makoviychuk et al.) 202T), Brax (Freeman et al. 202I). and VMAS (Bettini et al. 2022)). These vectorized simulators run parallel environment instances in batch, leveraging the SIMD paradigm of GPUs to greatly speed up execution. While this paradigm is becoming popular thanks to advances in hardware, not all RL libraries are ready to leverage its benefits. For instance, in Table 3 we report a comparison between TorchRL and RLlib (Liang et al. 2018) showing the collected frames per second as a function of the number of vectorized environments for the "simple spread" task in the VMAS simulator. The comparison shows that TorchRL is able to leverage vectorization to greatly increase collected frames. serves illustrate the usage of hardware accelerators in simulators, such as IsaacGym, which is also covered by the library's 's wrappers. The code used to generate the results is publicly available in the benchmark section of our repository.
Navigstion
Balance
Sampling
0.4
Moar maxdd
0.05
0.8
0.2
0.5
6.00
0.2
0.2
10
20
30
0
10
20
30
0
16
2)
30
Namber frames (M)
Nomber frans (M)
Number frames
MADDPG
IPPO
IQL
MAPPO
QMIX
VDN
Figure 2: MARL algorithm evaluation in three VMAS multi-robot control tasks. We report the mean and standard deviation reward over 6 random seeds. Each run consists of 500 iterations of 60, 000 steps each, with an episode length of 100 steps.
8

---

Published as a conference paper at ICLR 2024
Table Frames per second function of the number of vectorized environments. This evaluation shows that TorchRL. is able to leverage vectorized simulators to increase collected frames. The experiments were run on 100 steps of the "simple spread" MPE (Lowe al. 2017) scenario in the VMAS (Bettini et al. 2022 vectorized simulator.
of vectorized envs
3334
6667
10000
13333
16667
20000
23333
26666
30000
TorchRL fps
127
263848
398304
529872
444088
522537
538930
487261
384469
422003
RI.ibb fps
93
1885
1663
1675
1580
1523
1581
1447
1373
1336
Computational efficiency. To check that our solution does not come the cost of reduced throughput, we test two selected functionalities (advantage computation and data collection) against other popular libraries. These experiments were run on an AWS cluster with a single node (96 CPU cores and A100 GPUs per run). These results are reported in Table 4 and (Table 5| Overall, these results show that our solutions have a comparable not better throughput than others.
Table 4: Efficiency computing the Generalised Advantage Estimation (GAE) (Schulman et al. 2015). The data consisted of a batch of 1000 trajectories of 1000 time steps, with a "done" Irequency of 0,001, randomly spread across the data.
Using these implementations requires transforming back to tensors. Numpy and similar backends require moving data from and to GPU which can substantially impact performance. Those operations are unaccounted for in this table but can double the GAE runtime. On the contrary. TorchRL works on PyTorch tensors directly. A proper usage of Ray's 's GAE implementation would have needed split of the adjacent trajectories, which did not do to focus on the implementation efficiency.
Numpy arrays and then transforming the result
Library
Speed (ms)
Standalone
Meta-RL
Adjacent
Backend
Device
Tianshout
5.15
namba
CPU
SB31
14.1
numpy
CPU
RLLib (Ray)"
9.38
scipy
CPU
CleanRL
1.43
+
Jax
GPU
TorpheL-compiled
2.67
+
+
PyTorch
GPU
ToechRL-vec
1.33
+
PyTorch
GPU
Table 5: Data collection speed with common gym environments across common RL libraries. In each experiment, 32 workers were used. The scripts are available on TorchRL's benchmark folder.
Library
Breakout-v5
HalfChoetah- v4
Pendulum-vl
Tianshou
1212
4719
5823
RLLib
97
1868
1845
Gymnasium
9289
24388
25910
TorchRL (parallel env)
9092
21742
23363
TorchRL (syne collector)
10394
23527
24530
TorchRI. (asyne collector)
19401
32894
31584
CONCLUSION
TorchRL a modular decision-making library with composable, finllc-rsponnsibiity building pieces that rely on TensorDict for efficient class-to-class communication protocol. TorchRL's design principles are tailored to the dynamic field of decision-making, enabling the implementation of a wide range of algorithms. We believe TorchRL can be of interest researchers and developers alike. Its key advantage is that researchers can create or extend existing classes on their own and seamlessly integrate them with existing scripts without the need to modify the internal code of the library, allowing for quick prototyping while maintaining stable and reliable code base. We believe this feature will also be suited for developers looking apply an AI solution a specific problem. We demonstrate that TorchRI benefits from wide range of functionalities, is reliable and exhibits computational efficiency. The library is fully functional, and tested across various scenarios and applications.

---

Published as a conference paper at ICLR 2024
REFERENCES
Robohive - unified framework for robot learning.
robohi ve 2020. httts::::: http:///sse..ooooooooooooo Neoogle.ro wwt..///w//wooooo......oo
Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp, Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg, Dandelion Mané, Rajat Monga, Sherry Moore, Derek Murray, Chris Olah, Mike Schuster, Jonathon Shlens, Benoit Steiner, Ilya Sutskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda Viégas, Oriol Vinyals, Pete Warden, Martin Wattenberg, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. TensorFlow: Large-scale machine learning on heterogeneous systems, 2015. URL https: /www//wwwww.//.../ htwp:tensorrfow.coomm org/ from tensorflow.org.
Software available
Accel-Brain. Reinforcement Jeare learning code repository, bttp:p//////gggtttooooo/ccccaaaa
http:///rrrar.oooocceeee////aaaa LLearnnn 2023.
Szilárd Aradi. Survey of deep reinforcement learning for motion planning of autonomous vehicles.
IEEE Transactions on Intelligent Transportation Systems, 23(2):740-759, 2022. doi: 10.1
TITS.2020. 202..004655
Bowen Baker, Otkrist Gupta, Nikhil Naik, and Ramesh Raskar. Designing neural network architec- tures using reinforcement learning, 2017.
Marc G. Bellemare, Will Dabney, and Rémi Munos. A distributional persseetivee on reinforcement abs/1707. 018177777770877777
learning. CoRR, abs/1707.06887 2017. URL JRLtp::///axxi. /ttp///xwiiv../////////0011111
Daniel S Bernstein, Robert Givan, Neil Immerman, and Shlomo Zilberstein. The complexity of decentralized control of markov decision processes. Mathematics of operations research, (4)) 819-840, 2002.
Matteo Bettini, Ryan Kortvelesy, Jan Blumenkamp, and Amanda Prorok. Vmas: A vectorized multi- agent simulator for collective robot learning. The 16th International Symposium on Distributed Autonomous Robotic Systems, 2022.
Matteo Bettini, Ajay Shankar, and Amanda Prorok. Heterogeneous multi-robot reinforcement learning. In Proceedings of the 22nd International Conference on Autonomous Agents and Multiagent Systems, AAMAS '23. International Foundation for Autonomous Agents and Multiagent Systems, 2023.
Clément Bonnet, Donal Byrne, Victor Le, Laurence Midgley, Daniel Luo, Cemlyn Waters, Sasha Abramowitz, Edan Toledo. Cyprien Courtot, Matthew Morris, et al. Jumanji: Industry-driven hardwareaccellrrated environments, 2022. URL htps://githuu. comfnnstaleepuifuuanjj.
Albert Bou, Sebastian Dittert, and Gianni De Fabritii,, Efficient reinforcement learning experimenta- tion in pytorch, 2022. URL
fT_ fT_9HHp
James Bradbury, Roy Frostig. Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne and Qiao Zhang. JAX: composable transformations of Python+NumPP programs, 2018. URL
http:://ithub. woom/googee//a//oooo jax
Greg Brockman, Vicki Cheung. Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojeiech Zaremba. Openai gym. arXiv preprint arXiv:1606.01540 2016.
Axel Brunnbauer, Luigi Berducci, Andreas Brandstätter, Mathias Lechner, Ramin Hasani, Daniela Rus, and Radu Grosu. Latent imagination facilitates zero-shot transfer in autonomous racing, 2022.
Itai Caspi, Gal Leibovich, Gal Novik, and Shadi Endrawis, Reinforcement learning coach, December 2017. URL httl::ttt////oooo
Pablo Samuel Castro, Subhodeep Moitra. Carles Gelada, Saurabh Kumar, and Marc G. Bellemare. Dopamine: A Research Framework for Deep Reinforcement Learning. 2018. URL http:
7/a Tarxivv. //arxiv.orr//abs//82..000
10

---

Published as a conference paper at ICLR 2024
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learnin via sequence modeling. Advances in neural information processing systems, 34:15084-15097, 2021a.
Xinyue Chen, Che Wang, Zijian Zhou, and Keith Ross. Randomized ensembled double q-learning: Learning fast without a model. arXiv preprint arXiv:2101.05982, 2021b.
Paul F Christiano, Jan Leike, Tom Brown, Miljan Martic, Shane Legg, and Dario Amodei. Deep reinforcement learning from human preferences. Advances neural information processing systems, 30, 2017.
Christian Schroeder de Witt, Tarun Gupta, Denys Makoviichuk, Viktor Makoviychuk, Philip HS Torr, Mingfei Sun, and Shimon Whiteson. Is independent learning all you need the starcraft multi-agent challenge? arXiv preprint arXiv:2011.0953: 2020.
Carlo D'Eramo, Davide Tateo, Andrea Bonarini, Marcello Restelli, and Jan Peters. Mushroomrl: Simplifying reinforcement learning research. Jourmal of Machine Learning Research, 22(131): 1-5, 2021. URL /taltt:://jll: http://mmm...rrp/pppssss////////
056 html
DI engine Contributors. DI-engine: OpenDILab decision intelligence engine. com/opend lab/DI-eng 2021.
https:
Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Volodymir Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dunning, Shane Legg, and Koray Kavukcuoglu. Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures, 2018
Alhussein Fawzi, Matej Balog, Aja Huang, Thomas Hubert, Bernardino Komera-Paredess Mo- hammadamin Barckatain, Alexander Novikov, Francisco J R. Ruiz, Julian Schrittwieser, Grze- gorz Swirszez, David Silver, Demis Hassabis, and Pushmeet Kohli. Discovering faster ma- trix multiplication algorithms with reinforcement learning. Nature, 61007930):47-53, 2022. ISSN 1476-4687, doi: 10. 100.11/441566002-05577777 s4 641586-022-05172--
C. Daniel Freeman, Erik Frey, Anton Raichuk. Sertan Girgin, Igor Mordatch, and Olivier Bachem. Brax differentiable physics engine for large scale rigid body simulation, 2021. URL UR //gi thub. wtt.hh....o.aooo www..///moooooo.oooooooooo/..o. wow.o//bgl...o.ooooooooooooo....
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning, 2020.
Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor-
critic methods. In International conference on machine learning, pp. 1587-1596. PMLR, 2018a.
Scott Fujimoto. Herke Hoof, and David Meger. Addressing function approximation error in actor- critic methods. In International Conference on Machine Learning, pp. 1587-1596. PMLR, 2018b.
Yasuhiro Fujita, Prabhat Nagarajan, Toshiki Kataoka, and Takahiro Ishikawa. Chainerrl: A deep reinforcement learning library. Journal of Machine Learning Research, 22(77):1-14, 2021. 2021. URL
Sumitra Ganesh, Nelson Vadori, Mengda Xu, Hua Zheng, Prashant Reddy, and Manuela Veloso. Reinforcement learning for market making in multi-agent dealer market, 2019.
The garage contributors. Garage: A toolkit for reproducible reinforcement learning research. https
2019.
Jason Gauci, Edoardo Conti, Yitao Liang, Kittipat Virochsiri, Zhengxing Chen, Yuchen He, Zachary Kaden, Vivek Narayanan, and Xiaohui Ye. Horizon: Facebook's open source applied reinforcement learning platform. arXiv preprint arXiv: (1811100600 2018.

---

Published as a conference paper at ICLR 2024
Sergio Guadarrama, Anoop Korattikara, Oscar Ramirez, Pablo Castro, Ethan Holly, Sam Fishman, Ke Wang, Ekaterina Gonina, Neal Wu, Efi Kokiopoulou, Luciano Sbaiz, Jamie Smith, Gábor Bartók, Jesse Berent, Chris Harris, Vincent Vanhoucke, and Eugene Brevdo. TF-Agents: A library for reinforcement learning in tensorflow. /tttps://githbb..... http:itttwb.o.om..ooooem wwtp://so//o/te/eo/eeeeeeeeoeote ow wwwagentss 2018. htt.:htts::/// hnttps://github.... http://ttt..ccmeeeooooooo/////// 2019].
[Online; accessed 25-June-
T. Haarnoja, Aurick Zhou, Kristian Hartikainen, G. Tucker, Sehoon Ha, J. Tan, V. Kumar, H. Zhu, A. Gupta, P. Abbeel, and S. Levine. Soft actor-critic algorithms and applications. ArXiv, 2018.
Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint 1rX1v:1912.01003, 2019.
Nicklas Hansen, Yixin Lin, Hao Su, Xiaolong Wang, Vikash Kumar, and Aravind Rajeswaran. Modem: Accelerating visual model-based reinforcement learning with demonstrations, 2022a.
Nicklas Hansen, Xiaolong Wang, and Hao Su, Temporal difference learning for model predictive control, 2022b.
Yihui He, Ji tianshou, Zhijian Liu, Hanrui Wang, Li-Jia Li, and Song Han. Amc: Automl for model compression and acceleration on mobile devices. In Proceedings of the European Conference on Computer Vision (ECCV), September 2018.
Mattco Hessel, Joseph Modayil, Hado van Hasselt, Tom Schaul, Georg Ostrovski, Will Dabney, Dan Horgan, Bilal Piot, Mohammad Azar, and David Silver. Rainbow: Combining improvements in deep reinforcement learning, 2017.
Matthew W. Hoffman, Bobak Shahriari, John Aslanides, Gabriel Barth-Maron, Nikola Momchev, Daníla Sinopalnikov, Piotr Stańczyk, Sabela Ramos, Anton Raichuk, Damien Vincent, Léonard Hussenot, Robert Dadashi, Gabriel Dulac-Amold, Manu Orsini, Alexis Jacq, Johan Ferret, Nino Vieillard, Seyed Kamyar Seyed Ghasemipour, Sertan Girgin, Olivier Pietquin, Feryal Behbahani, Tamara Norman, Abbas Abdolmaleki, Albin Cassirer, Fan Yang, Kate Baumli, Sarah Henderson, Abe Friesen, Ruba Haroun, Alex Novikov, Sergio Gómez Colmenarejo, Serkan Cabi, Caglar Gulcehre, Tom Le Paine, Srivatsan Srinivasan, Andrew Cowie, Ziyu Wang. Bilal Piot, and Nando de Freitas. Acme: research framework for distributed reinforcement learning. arXiv preprint arXiv:2006.00997 2020. https://:
ab0p://.bbbbo2..oooo 00979
Shengyi Huang, Rousslan Fernand Julien Dossa, Chang Ye, Jeff Braga, Dipam Chakraborty, Kinal Mehta, and João G.M. Araújo. Cleanrl: High-quality single-file implementations of deep rein- forcement learning algorithms. Journal of Machine Learning Research, 23(274); 2022. URL
Facebook Incubator:
2021. [GitHub repositoy..
Dmitry Kalashnikov, Alex Irpan, Peter Pastor, Julian Ibarz, Alexander Herzog, Eric Jang, Deirdre Quillen, Ethan Holly, Mrinal Kalakrishnan, Vincent Vanhoucke, and Sergey Levine. Qt-opt: Scalable deep reinforcement learning for vision-based robotic manipulation, 2018.
Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learning arXiv preprint arXiv22110.06169.20 2021.
Alexander Kuhnle, Michael Schaarschmidt, and Kai Fricke. Tensorforce: tensorflow li-
brary for applied reinforcement learning. Web page, 2017. URL httpt:://gttuubbb com/
tsersorforrcettessotoooo
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning, 2020.
Karol Kurach, Anton Raichuk, Piotr Staúczyk, Michal Zajac, Olivier Bachem, Lasse Espeholt, Carlos Riquelme, Damien Vincent, Marein Michalski, Olivier Bousquet, et al. Google research football: A novel reinforcement learning environment. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pp. 4501-4510, 2020.
12

---

Published as a conference paper at ICLR 2024
Mare Lanctot, Edward Lockhart, Jean- Lean-Baptiste Lespiau, Vinicius Zambaldi, Satyaki Upadhyay, Julien Pérolat, Sriram Srinivasan, Finbarr Timbers, Karl Tuyls, Shayegan Omidshafiei, Daniel Hennes, Dustin Morrill, Paul Muller, Timo Ewalds, Ryan Faulkner, János Kramár, Bart De Vylder, Brennan Saeta, James Bradbury, David Ding, Sebastian Borgeaud, Matthew Lai, Julian Schrittweeerr, Thomas Anthony, Edward Hughes, Ivo Danihelka, and Jonah Ryan-Davis. Open- Spiel: A framework for reinforcement learning in games. CoRR,
2019. URL
http:/ /ttp///aavvvv/.a//////e...aaaa wwwpabs.///////oo.../t...t....et abs 1908., 09453
Quentin Lhoest, Albert Villanova del Moral, Yacine Jemnte,, Abhishek Thakur, Patrick von Platen, Suraj Patil, Julien Chaumond, Mariama Drame, Julien Plu, Lewis Tunstall, Joe Davison, Mario Saško, Gunjan Chhablani, Bhavitvya Malik, Simon Brandeis, Teven Le Scao, Victor Sanh, Canwen Xu, Nicolas Patry, Angelina McMillan Majorr, Phillipp Schmid, Sylvain Gugger, Clément Delangue, Théo Matussière, Lysandre Debut, Stas Bekman, Pierric Cistac, Thibault Goehringer, Victor Mustar, François Lagunas, Alexander Rush, and Thomas Wolf. Datasets: A community library for natural language processing. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pp. 175-184, Online and Punta Cana. Dominican Republic. November 2021. Association for Computational Linguistics. URL https
//aclantho www.hhhoooyy..ooo wtg.
21
Bric Liang, Richard Liaw, Robert Nishihara, Philipp Moritz, Roy Fox, Ken Goldberg, Joseph Gonzalez, Michael Jordan, and Ion Stoica. Rllib: Abstractions for distributed reinforcement learning. In International Conference on Machine Learning, pp. 3053-3062. PMLR, 2018
Timothy P Lillicrap, Jonathan J Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. arXiv preprint 1.Xiv::500.022.
2015.
Michael L Littmmn. Markov games as framework for multi-agent reinforcement learning. In Machine learning proceedings 1994, pp. 157-163. Elsevier, 1994.
Xiao- Yang Liu, Zechu Li, Ming Zhu, Zhaoran Wang, and Jiahao Zheng. ElegantRL: Massively parallel framework for cloud-native deep reinforcement learning. https://github. com/ AI4Fi nance- on/ElegantRI, 2021.
Ryan Lowe, Yi I Wu, Aviv Tamar, Jean Harb, OpenAl Pieter Abbeel, and Igor Mordatch. Multi-agent actor-critic for mixed processing systems, 30, 2017.
environments. Advances in neural information
Jerry Luo, Cosmin Paduraru, Octavian Voieu, Yuri Chervonyi, Scott Munns, Jerry Li, Crystal Qian, Praneet Dutta, Jared Quincy Davis, Ningjia Wu, Xingwei Yang, Chu-Ming Chang, Ted Li, Rob Rose, Mingyan Fan, Hootan Nakhost, Tingtiansho Liu, Brian Kirkman, Frank Altamura, Lee Ctianshoe, Patrick Tonker, Joel Gouker, Dave Uden, Warren Buddy Bryan, Jason Law, Deeni Fatiha, Neil Satra, Juliet Rothenberg, Mandeep Waraich, Molly Cartiansho, Satish Tallapaka, Sims Witherspoon, David Parish, Peter Dolan, Chenyu Zhao, and Daniel J. Mankowitz. Controlling commercial cooling systems using reinforcement learning. 2022.
Yecheng Jason Ma, Shagun Sodhani, Dinesh Jayaraman, Osbert Bastani, Vikash Kumar, and Amy Zhang. Vip: Towards universal visual reward and representation via value- cille-implicit pre-training. arXiv preprint arXiv:2210.000300 2022.
Torch Vision maintainers and contributors. Torchvision: Pytorch's computer vision library. https 2016.
Viktor Makoviychuk, Lukasz Wawrzyniak, Yunrong Guo, Michelle Lu, Kier Storey, Miles Macklin, David Hoeller, Nikita Rudin, Arthur Allshire, Ankur Handa, and Gavriel State. Isaac gym: High performance gpu-based physics simulation for robot learning, 2021.
Azalia Mirhoseini, Anna Goldie, Mustafa Yazgan, Joe Wenjie Jiang, Ebrahim Songhori, Shen Wang, Young-Joon Lee, Eric Johnson, Omkar Pathak, Azade Nazi, Jiwoo Pak, Andy Tong, Kavya Srinivasa, William Hang, Emre Tuncer, Quoe V Le, James Laudon, Richard Ho, Roger Carpenter, and Jeff Dean. A graph placement methodology for fast chip design. Nature, 594 (7862):207 2021. ISSN 1476- 7/doi . .rg 10. 1038 4 586 021 03544-
doi: 10. 10.1038/441586-005
-015---44-.. URL https
13

---

Published as a conference paper at ICLR 2024
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, loannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint arXiv:1312.5602, 2013.
Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. International conference on machine learning, pp. 1928-1937, 2016.
Philipp Moritz, Robert Nishihara, Stephanie Wang, Alexey Tumanov, Richard Liaw, Eric Liang, Melih Elibol, Zongheng Yang, William Paul, Michael Jordan, al. Ray: A distributed framework for emerging {AI} applications. In 13th { {USENIX] Symposium on Operating Systems Design and Implementatiin ((OSDI) 18), pp. 561-577, 2018.
Suraj Nair, Aravind Rajeswaran, Vikash Kumar, Chelsea Finn, and Abhinav Gupta. R3m: A universal visual representation for robot manipulation, 2022.
Rejichiro Nakano, Jacob Hilton, Suchir Balaji, Jeff Wu, Long Ouyang, Christina Kim, Christopher Hesse, Shantanu Jain, Vineet Kosaraju, William Saunders, Xu Jiang, Karl Cobbe, Tyna Eloundou, Gretchen Krueger, Kevin Button, Matthew Knight, Benjamin Chess, and John Schulman. Webgpt: Browser-assisted question with human feedback. CoRR, abs/211 2021. URL
attts://arriv..ot http:/ii.wrga/b//////1/2/ 09332
Chris Nota. The autonomous learning library. 2020.
nuttoomouus-eearring--iiii
OpenAI, Ilge Akkaya, Marein Andrychowicz, Maciek Chociej, Mateusz Litwin, Bob McGrew, Arthur Petron, Alex Paino, Matthias Plappert, Glenn Powell, Raphael Ribas, Jonas Schneider, Nikolas Tezak, Jerry Tworek, Peter Welinder, Lilian Weng, Qiming Yuan, Wojciech Zaremba, and Lei Zhang. Solving rubik's cube with robot hand, 2019a.
OpenAI, Christopher Berner, Greg Brockman, Brooke Chan, Vicki Cheung, Przemystaw Dębiak, Christy Dennison, David Farhi, Quirin Fischer, Shariq Hashme, Chris Hesse, Rafal Józefowicz, Scott Gray, Catherine Olsson, Jakub Pachocki, Michael Petrov, Henrique Pondé de Oliveira Pinto, Jonathan Raiman, Tim Salimans, Jeremy Schlatter, Jonas Schneider, Szymon Sidor, Ilya Sutskever, Jie Tang, Filip Wolski, and Susan Zhang. Dota 2 with Large Scale Deep Reinforcement Learning. 2019b. URL
06680
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein. Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. PyTorch: An Imperative Style, Righ-Performance Deep Learning Library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' Alché Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems 32, pp. 8024-8035. Curran Associates, Inc., 2019. URL http:::/papprss.euriiss....
9015 111----yttrcchaaae an- wtwprcc-aan-mmerattiiee style http::/ceeeeee..oooooooorroerrea pdf
Library.
Xue Bin Peng, Marein Andrychowiez, Wojciech Zaremba, and Pieter Abbeel. Sim-to-real transfer of robotic control with dynamics randomization. In 2018 IEEE International Conference on Robotics and Automation (ICRA). IEHE, may 2018 doi: 10.1 109//cra.2011.84652222 URL
https://dot .org/00. 10w.9111caa............ http:://ooo..or///0..0003//aaa.
Matthias Plappert. keras-rl. httpps://gittubbccmm/eeeee
2016.
Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dormann. Stable-baselines3: Reliable reinforcement learning implementations. Journal of Machine Learning Research, 22(268):1-8, 2021. URL http:://jmmi. pttpppppps///22.....
20 1364.34440 html
Tabish Rashid, Mikayel Samvelyan. Christian Schroeder De Witt, Gregory Farquhar, Jakob Foerster, and Shimon Whiteson. Monotonic value function factorisation for deep multi-agent reinforcement learning. The Journal of Machine Learning Research, 21(1):7234-7244 2020.
14

---

Published as a conference paper at ICLR 2024
Roman Ring. Reaver: Modular deep reinforcement learning framework. https://github.c Snoryy/reaver 2018.
Nikita Rudin, David Hoeller, Philipp Reist, and Marco Hutter. Learning walk in minutes using massively parallel deep reinforcement learning, 2022.
Mikayel Samvelyan, Tabish Rashid, Christian Schroeder de Witt, Gregory Farquhar, Nantas Nardelli, Tim GJ Rudner, Chia-Man Hung. Philip HS Torr, Jakob Foerster, and Shimon Whiteson. The the 18th International Conference on Autonomous Agents and MultiAgent Systems, pp. 2186-2188. 2019.
starcraft multi-agent challlen In Proceedings
John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. High-dimensional continuous control using generalized advantage estimation. arXiv preprint 150.::/50000244
2015.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: 1707.06777 2017.
Dhruv Shah, Arjun Bhorkar, Hrish Leen, llya Kostrikov, Nick Rhinehart, and Sergey Levine. Offtiane reinforcement learning for visual navigation, 2022.
David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George van den Driessche, Julian Schrittwesser, Joannis Antonogiou, Veda Panneershelvam, Marc Lanctot, Sander Dieleman, Dominik Grewe, John Nham, Nal Kalchbrenner, 1lya Sutskever, Timothy Lillicrap, Madeleine Leach, Koray Kavukcuoglu, Thore Graepel. and Demis Hassabis. Mastering the game of Go with deep neural networks and tree search. Nature. 529(7587):484 489, 2016. ISSN 1476-4687. doi: 10. 10.1038/nature 1694t
URL httpthttt//s/tt org 100 wtt.:/00000//atteeeeee11
Adam Stooke and Pieter Abbeel. ripyt: A research code base for deep reinforcement learning in pytorch. CoRR, abs/1909.01500, 2019. URL http:/:://aaziiv.o////a///////// http://gg/////990////// 01500
Peter Sunehag, Guy Lever, Audrunas Gruslys, Wojciech Marian Czarnecki, Vinicius Zambaldi, Max Jaderberg, Mare Lanctot, Nicolas Sonnerat, Joel Z Leibo, Karl Tuyls, et al. Value-deecomossitio networks for cooperative multi-agent learning based on team reward. Proceedings of the 17th International Conference on Autonomous Agents and MultiAgent Systems, pp. 2085-2087 2018.
Andrew Szot, Alex Clegg, Eric Undersander, Erik Wijmans, Yili Zhao, John Turner, Noah Maestre, Mustafa Mukadam, Devendra Chaplot, Oleksandr Maksymets, Aaron Gokaslan, Vladimir Vondrus, Sameer Dharur, Franziska Meier, Wojeiech Galuba, Angel Chang, Zsolt Kira, Vladlen Koltun, Jitendra Malik, Manolis Savva, and Dhruv Batra. Habitat 2.0: Training home assistants rearrange their habitat. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
Ming Tan. Multi-agent reinforcement learning: Independent cooperative agents. In Proceedings of the tenth international conference on machine learning, pp. 330-337, 1993.
Yuandong Tian, Qucheng Gong, Wenling Shang, Yuxin Wu, and Lawrence Zitnick. Elf: extensive, lightweight and flexible research platform for real-time strategy games. Advances in Neural Information Processing Systems (NIPS), 2017.
Josh Tobin, Rachel Fong, Alex Ray, Jonas Schneider, Wojeiech Zaremba, and Pieter Abbeel. Domain randomization for transferring deep neural networks from simulation to the real world, 2017.
Saran Tunyasuvunakool, Alistair Muldal, Yotam Doron, Siqi Liu, Steven Bohez, Josh Merel, Tom Erez, Timothy Lillicrap, Nicolas Heess, and Yuval Tassa. contool:::: Software and tasks for continuous control. Software Impacts, 6:100022, 2020. ISSN 2665 9638. doi: https://doi. www. www.scciiccc.... wttpett.oosscccciic.o
2020 100022. BTticcee/ji1/2265533200000
Joaquin Vanschoren, Jan N. van Rijn, Bernd Bischl, and Luis Torgo Openml: networked science in machine learning. SIGKDD Explorations, 15(2):49 60, 2013. doi: 10.1145/2641 190.2641198.
URLhttp::
acm. wtt....ro.....11155333366000 264111111196.2111991
15

---

Published as a conference paper at ICLR 2024
Leandro von Werra, Jonathan Tow, reciprocated, Shahbuland Matiana, Alex Havrilla, cat state, Louis Castricato, Alan, Duy Phung, Ayush Thakur, Alexey Bukhtiyarov, aaronrmm, Fabrizio Milo, Daniel, Daniel King, Dong Shin, Ethan Kim, Justin Wei, Manuel Romero, Nicky Pochinkov, Omar Sanseviero, Reshinth Adithyan, Sherman Siu, Thomas Simonini, Vladimir Blagojevic, Xu Song, Zack Witten, alexandremuzio, and crumb. CarperAl/trl:: v0.6.0: LLLaMa (Alpaca), Benchmark Util, T5 ILQL, Tests, March 2023, httttt::///oii.. wtg. 5281/zenodo. 7790115
Jiayi Weng, Huayu Chen, Dong Yan, Kaichao You, Alexis Dubureq, Minghao Zhang, Yi Su, Hang Su, and Jun Zhu. Tianshou: A highly modularized deep reinforcement learning library. Journal of Machine Learning Research, 23(267): 1-6, 2022. URLhttp:: /mml:::///m....//// org/ ptppers/v23.
21- 1127. htmm
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Transformers: State-of-the-ar art natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Pro- cessing: System Demonstrations, pp. 38-45, Online, October 2020. Association for Computational Linguistics. URL httph: www ac weeb httpb//rrgg/aoooooo/222222 wttp::ooooooo/222...eeeeeeeeeee emnlp-demos.
Philipp Wu, Alejandro Escontrela, Danijar Hafner, Ken Goldberg, and Pieter Abbeel. Daydreamer: World models for physical robot learning, 2022.
Chao Yu, Akash Velu, Eugene Vinitsky, Jiaxuan Gao, Yu Wang. Alexandre Bayen, and Yi Wu. The surprising effectiveness of ppo in cooperative multi-agent games. Advances in Neural Information Processing Systems, 35:2461 1-2424,, 2022.
Qinqing Zheng, Amy Zhang, and Aditya Grover. Online decision transformer, 2022.
Zhao Zhong, Junjie Yan, Wei Wu, Jing Shao, and Cheng-tianshen Liu. Practical block-wise neural network architecture generation. In Proceedings of the IEEE conference on computer vision and pultern recognition, pp. 2423-2432, 2018.
Zheqing Zhu and Benjamin Van Roy. Deep exploration for recommendation systems, 2021.
Barret Zoph and Quoc V. Le. Neural architecture search with reinforcement learning, 2017.
Barret Zoph, Vijay Vasudevan, Jonathon Shlens, and Quoc V. Le. Learning transferable architectures for scalable image recognition, 2018
16

---

Published as a conference paper at ICLR 2024
APPENDIX
A OPENSOURCE REPOSITORIES
TorchRL and TensorDict libraries opensource code can be found the following locations:
TorchRL: httppp/ss//ggthbb...oo///ooocc//
TecsorDict: httpps///ggghh.....///r//r/rrrrr
B TENSORDICT FUNCTIONALITY
TensorDiet's aim is abstract away functional details of Py'Torch workflows, enabling scripts be more easily reused across different tasks, architectures, or implementation details. While primarily tool for convenience, TensorDict can also provide notable computational benefits for specific operations.
B.I TENSORDICTBASE AND SUBCLASSES
The parent class of any TensorDict class TensorDi riterrsec...e...... This abstract class lists the common operations of all its subclasses and also implements a few of them. Indeed, some generic methods are implemented once for all subclasses, such as tensordict. apply, as the output will always be TensorDico instance. Others, such getitem have behavior that intrinsically linked to their specifics.
The following subclasses available:
TensorDics ct: this is the most common class and usually the only one users will need to interact with.
La TarysstackedFensorDictt This class is the result of calling torch. stack list of TensorDi ctasse subclass instances. The resulting object stores each instance independently and stacks the items only when queried through _getttem or similar. Heterogeneous TensorDict instances can be stacked together: in that case, get may fail to stack the tensors (if one missing from one tensordict instance if the shapes do not match). However, this class can still be used to carry data from object to object or worker to worker, even if the data heterogeneous. The original tensordicts can be recovered simply via indexing of the lazy stack: torch. (list_of_tde of list_of__dds
set eetitem_. get, set
0) [0] will return the first element
PersitenttTensoriie implements a TensorDict class with persistent storage. At the time of writing, only the H5 backend is implemented. This interface is currently being used to integrate massive datasets within TorchRL for offline RL and imitation learning.
CustammppTensorDier This abstract class is the parent of other classes that implement lazy operations. This is used to temporarily interact with a TensorDict on which a zero-copy reshape operation is be executed. Any in-place operations tensordict as well. For instance, tensordi et. permute (0, 1) set ("a", tensor) will set new key the parent tensordict as well.
these instances will affect the parent
Performance Through convenient and intuitive API, tensordict offers some out-of-the-box op- timizations that would otherwise be cumbersome achieve and, we believe, can have positive impact beyond RL usage. For instance, TensorDict makes it easy to store large datasets in contigu- ous physical memory through PyTorch interface with NumPy's memory-mapped arrays named MemmapTenne
available in the tensordict library. Unlike regular PyTorch datasets, TensorDict offers the possibility to index multiple items at a time and make them immediately available in memory in a contiguous manner: by reducing the I/O overhead associated with reading single files (which we preprocess and store in a mneoory-mapper array on disk), one can read, cast to GPU and preprocess multiple files all-at-once. Not only files, but we can also leverage this to apply transforms on batches of data instead of one item at time. These batched transforms device much faster to execute than their multiprocessed
amortize the time of reading independent
17

---

Published as a conference paper at ICLR 2024
Benchmarking of single Neeraioo over ImageNet, collection speed
8 colos
1096
Doraredoor
1516
tamsordes 3PJ ne extra wo oor
0661
GPJ,, worker translomm)
0051
d sirnucod tomorrcct tonsordicts GHA, nateers (joca trar slmmmy
15613
satttczdd berelomm
di siibbuee tinsorbiss 1 3-U vor (teesform by fran rodo)
0026
di shibuud lormardeess tronsform by trsii todo)
CaFJ, workars mrorstom titin iodni
1000
20300
4322
Roto (P4)
Figure 3: Dataloading speed with TensorDict.
@tensorclass
class MyClass:
image: Tensor Tensor
label: terch. Tensor
data = MyClass
image=torch.. zeros (100, 32, 32, 3, perpe-torch.uiateee label=torch randittt randint (1000, (100,) )), batch_size-
# MyData instance with floating point content data float data.apply (lambda x: x.float MyData instance on cuda data_cuda do t"
the first 10 elements of the original MyData instance
data_idx datal:
Figure 4: @tensorclass example.
counterparts. Figure shows per-image random transformations including cropping and flipping. The workflow used achieve this speedup is available in the repository documentation.
shows the collection speed over ImageNet with batches of 128 images and
tensorclass TensorDict also offers a @tensorciass decorator aimed at working as the 8dataclass decorator: it allows for the creation of specialized dataclasses with all the fea- tures of TensorDict: shape key (attribute) dimension, device, tensor operations and more. Stensorclass instances natively support any TensorDict operation. Figure shows a few exam- ples of @tensorclass usage.
B.2 TENSORDICT.NN. MODULES AND FUNCTIONAL API
The second pane the tensordict library is its module interface, on which TorchRL heavily relies. tensordict. aims to address
limitations of torch. nn. First, torch. offers a Module class which encourages users to create new components through inheritance. However, in many cases, this level of control in module construction is unnecessary and even hinder
18

---

Published as a conference paper at ICLR 2024
class MLF (nn. Module)
def _init_ (self);
self. 1ayer1 nn. (3, 4)
self. nn. (4,
self.activation torch.relu
def fsew, x)
y self. aellvactiittion (self. linearl (x))
return self. 1inear2 (y)
Figure 5: Programmable module design with TensorDDictModule
from tensordict.nn nn import TensorDictSequentiaal as Seq, TensorDictModule as Mod
module Seq
Mod (nn Lincar (3, 4) in out aidden1' hidden0'), out_keys= Mod Lnn. (4, 2), in_keys- mn_keys=("hiine ('bidden1'), out_keys sheys-[ output' )) }),
]),
Mod itorch.relu,, ('_dee==""
Figure Programmable module design with DendouDic
modularity. For example, consider the coding of simple multi-layered perceptron, depicted in figure 5, which is commonly encountered.
Such definition of dedicated class lacks flexibility. One solution commonly adopted to control the network hyperparameters (number of layers, the activation function, or layer width) is to pass them inputs the constructor and group the layers together within nn Sequent ialiilll instance. This solution works fine as long as the sequence is a simple chain of operations on a single tensor. As soon some operations need to be skipped, when modules output multiple tensors, nn Swquentiuintn.ee is not suited anymore and the operations must be explicitly written within the forward method, which the computational graph once and for all. The TensorDi etMouuullee class overcomes this limitation and provides fully programmable module design API, as shown in figure
Formulating module in this way brings all the flexibility we need address the shortcomings exposed above: it is easy to programmatically design new model architecture at runtime for any kind of model. Selecting subgraphs is easy and can be achieved regular indexing of the module or more advanced sub-graph selection. In the following example, call to welllleetssseeeeeeeeeeeeeee will output new operations depend the "hidden1 key:
ial where the modules are restricted only to those whose
sub_modutt module. sel wwweeeeetsessseueeeeeeeeeeeeeeee (inkeeee-eenneeeeee- (*) dddnn
To further back our claim that such model construction is beneficial both in terms of clarity and modularity, let us take the example writing residual connection with tensordict.nn:
block Seq
Mod (nn. (128, 128), in_keys ( ['x'), out ott_keye=("nntermddee
Mod forchhrelu,, in keys intermediatee ], out but_keys=[[ intermediate")) Mod Liner (128, 128), in keys= intermediate out_keys-1 intermediate
residual = Seq(
block, Mod (lambda xt2, in in_keys-('x', 'intermediates out_keys-['4 [(x[]
19

---

Published as a conference paper at ICLR 2024
For clarity, we have separated the backbone and the residual connection, but these could be grouped under the same TensorDict Sequent ial instance. Nevertheless, select subsequence will go through these various levels of nesting if ever called.
Functional usage Functional usage of these modules is a comerstone in TorchRL as it underpins efficient calls to the same module with different sets of parameters for off-policy algorithms where multiple critics are executed simultancously or in situations where target parameters are needed. tensord ct. provides make funct. ional function that transforms any torch. mod- ule a functional module that accepts an optional pars params return the parameters and buffers in TensorDict instance whose nested structure mimies the one of the originating module (unlike Module. state_dict () which returns them in a flat structure).
jarruddjrrgument. This function will
Using these functionalized modules is straightforward. For instance, can zero all the parameters of module and call it on some data:
params locke_funciionn (module) parame_zero params. slone zerc_ module (tensordict, params-pas aarams-paramsaaeroo
This allows us to group the calls multiple critics in vectorized fashion. the following example, stack the parameters of two critic networks and call the module with this stack using torch. vmap, the PyTorch vectorized-min map function:
make www.ffunniinnee (critic) vmap criticiiiice...... vmap (critic, (None, )
data vmap_critic( data, stack (Iparams_oritico, params_criticl), 0)
We refer to the torch.vmap documentation for further information on this feature. Similarly, calling the same module with trainable target parameters can done via functional call it.
Performance The read and write operations executed during the unfolding of SensorriettSequent
can potentially impact the execution speed of the underlying model. To address this concern, we offer a dedicated symbolic_tres operation graph to its essential elements. This function generates module that is fully compatible with corch. comp i le, resulting in faster execution. By combining these solutions, we can achieve module performance that on par, if not better, than their regular nn. Module counterparts.
function that simplifies the
Using TensorDict modules in tensordict-free code bases Because the library's goal is to be as min- imally restrictive as possible, we also provide the possibility of executing TensorDict-basee modules without explicitly requiring any interaction with TensorDict: the tensordict nn. di dispat ch. decorator allows interact with any module from the TensorDict ecosystem with pure tensors:
module TensorDictModuler lambda x+1, x-2,
in keys= <eys=|"x"1, "x" out keys= (""" "z"1)
y, module (x torch. randn (3))
VensorDictModuLa and associated classes also provide a select_ous keys method that allows to hide some specific keys of a graph to minimize the output. When used in conjunction with @di spatch, this enables a much clearer usage when multiple intermediate keys are present.
B.3 TENSORDICT VS OTHER TENSOR CARRIERS
TensorDict is of the many data carriers in RL and other machine learning fields. To avoid overfitting to a single comparison, we provide a generic comparison TensorDict in the RL ecosystem followed by comparison of TensorDict with PyTrees.
20

---

Published as a conference paper at ICLR 2024
B.3.1 TENSORDICT OTHER RL DATA CARRIERS
Examples of existing data carriers in RL are Tianshou's Batch, RLLib's TensorDict, Salina's Workspace. The goal of each of these classes somewhat similar, i.e. abstracting away the precise content of a module output to focus on the orchestration of the various components of the library.
First and foremost, TensorDict isn't focused on RL. specifically, while the others are. We did not restrict its content to be some derivation of a (observatinn ion, next_observation reward, done, action) tuple. It can store any numerical data, e.g., structures without reward for imitation learning, or with deeply nested data structures for Monte Carlo Tree Search algorithms. This allows TorchRL to cover much broader range of algorithms like Offline RL, imitation learning, optimal control and others, where other solutions may struggle given their content limitations.
Second, TensorDict behaves like a tensor: it supports shape manipulation, casting, copying, etc. All these operations are implemented in such way that in many places, tensors and TensorDicts are interchangeable. This allows us to build policies that return TensorDicts and not tensors for compound actors for instance. It also allows for certain implicit recursivity in TensorDict functionalities: for many operations, applying a transformation to all leaves in the TensorDict tree just boils down to calling that operation over all fields of the TensorDict, recursively.
B.3.2 TENSORDICT PYTREE
TensoDict also be compared to PyTree, concept shared by PyTorch and Jax that allows cast operations to any level of deeply nested data structures. Yet, TensorDict isn't 't based on PyTree. There are several things to take into consideration understand why:
Using PyTree implies decomposing the TensorDict, applying the function to its tensors (leaves), and recomposing it. This introduces some overhead, metadata need be saved, and the class reinstantiated which does not come for free. "TensorDict.apply can reduce that overhead by knowing advance what to do with the op that is applied it.
TensorDict operations are implemented in a way that makes the results more directly useful the user. Splitting (or unbinding) operations are a good example of this: using regular PyTree call over a nested dictionary would result in nested dictionaries of tuples of tensors. Using "TensorDiet.se will result in tuples of TensorDicts, which matches the 'torch.split" signature. This makes this operation more "intuitive" than the PyTree version and is more directly useful to the user. course, some Py Tree implementations offers tools to invert the structure order. Nevertheless, we believe having these features built-in to be useful.
TensorDicts carry important metadata, like "shape" (i.e., batch-size) and "device", whereas generic structures that are compatible with PyTree (such dict and lists) do not carry this information. In sense, it is less generic than a dictionary, but for its own good, Metadata allow us to leverage the concept of batch-size for instance, and separate it from the "feature-size" Among other advantages, this greatly facilitates vmap-ing ing over TensorDict, as we know what shape can be vmaped over advance. For example, TensorDict used TorchRL have by convention the "time" dimension positioned last. If we need to vmap over it, we can call vamp (op, (-1,) (data). This - 1 is generic (i.c., reusable across use cases, independently of the number dimensions). and does not depend on the batch size. This is anecdotal example would be harder to code using Py Trees.
We think that using PyTree requires more knowledge and practice than TensorDict. Sugar code like data. cuda () would be harder to code with Pree. Retrieving the resulting metadata (through data... , de) wouldn't 't be possible either, since PyTree does not carry metadata unless special class is written.
TensorDicts come with some handy features that would not be easy to come up with using PyTrees: they can be locked to avoid unwanted changes and they can be used as context managers for some operations (like functional calls over modules, or temporarily unlocking a tensordict using with data. unlock_ () ....).
Finally, note that DensorDict registered within PyTree in PyTorch, such that anything PyTree can do, TensorDiet can do too.
21

---

Published as a conference paper at ICLR 2024
B.4 TENSORDICT OVERHEAD MEASUREMENT
We measured TensorDict overhead for the lastest release at the time of writing the manuscript (v0.2.1) against other existing solutions.
B.4.1 TENSORDICT PYTREE
Dict has some similarities with the PyTree data-structure presented in Jax Bradbury et al. (001 or in PyTorch Paszke al. (2019), as it can dispatch operations over a of leaves. To measure the worst case overhead introduced by TensorDict (which handles meta-data such as batch-size and carrier device when executing these operations), we executed the same operations using Py PyTree and TensorDict over highly nested (100 levels) dictionary, where each leaf was scalar tensor. We consider this level of nesting to be above most of the usual use cases of TensorDict (in typical RI. setting, only couple of levels are needed for data representation and a tenth so for parameter representation with complex models). We tested the time takes to increase the value each leaf by 1.0 or to transfer all tensors from CPU to a GPU device. Additionally, required to split, chunk, and unbind a TensorDict along its batch-size compared to the equivalent Py Prere operation. We ran these operations 1000 times over 16 different runs and report the 95% confidence interval of the expected value.
also assessed the time
Results The value increment took 0.916 ms (0.91531-0.91628) with PyTree, compared with 1.0555 ms (1.05495 1. 1.56688 for TensorDict, which constitutes a 115% runtime increment. Casting to CUDA had similar impact, with PyTree running 1.181 ms (1.16242 1.19960) and TensorDict 1.302 ms (1.30065- 1.30242), a 10% runtime increment. For batch-size operations, the overhead was also worse for TensorDict: split took 1.015 ms (1.01436 1.01530) against 4. 4.254 ms (4.23771- 4,26700), chunk took 0.801 ms (0,80027 0.80084) vs 2. ms (2.14920- 2.16060) and unbind 1.418 ms (1.417299- 1.41778) vs 3.836 ms (3.834133.. 3.83615) for PyTTree and TensorDict in each case. Recall that those operations with PyTree returned nested dictionaries of tuples rather than tuples of dictionaries and that TensorDict had to re-compute the batch-size of each element of the batch: these differences explain the overhead observed. Nevertheless, will aim at addressing these issues in the next TensorDict release as there significant room for improvement. Our plan is detailed on the discussion page of the repository.
B.4.2 TENSORDICT FUNCTIONAL. CALLS VS FUNCTORCH
TensorDict offers a simple API to execute functional calls over a specific PyTorch module. Unlike torch. func. funct inoal call, TensorDict allows for functional calls over any method of the module, and its nested representation makes it possible to execute functional calls over sub- modules without effort. For the sake of completeness, we also compared how much functional call over a narrow (16 cells instead of 2048) but deep torch. nn. Trans fo Former r would be impacted by switch to TensorDict. The runtime for functorch was of 8.423 ms (8.38470- 8.46124) against 10.26071 ms (10.25635 10.26500) for TensorDict. This reduced speed should be put in perspective with the depth of the model and reduced runtime for the algebraic operations. Note that TensorDiet functional calls are currently in their first version and will be refactored soon to make them easier to use and more broadly applicable (eg, in distributed and graph settings). We have already developed prototype of this feature which will be included in the next release. After this refactoring, early experiments have already confirmed that the runtime will be at par with existing functional calls in functorch.
C ENVIRONMENT API
To bind input/output data to a specific domain, TorchRL uses TensorSpec, a custom version of the feature spaces found in other environment libraries. This class is specially tailored for PyTorch: its instances can be moved from device to device, reshaped, or stacked at will, as the tensors would be. TensorDict is naturally blended within TensorSped through a Compos i teSpec primitive which a metadata version of Dics The following example shows the equivalence between these two classes:
>>> from torchrl 1.ata import Subeunddddoont:immouseeeeeee
22

---

Published as a conference paper at ICLR 2024
Composite$pec
>>> spec Compositespec
oo (shape- (3, 4)),
shape (3,), ) device=' cpu'
>>> print (snec. TensorDict( fielas=(
obs: Tensor (shape=torch, (13, 41), device=cpu, dtype-toren. float32, is sa__saree-ralse)
batch_si sirch_shzeetrrhh.. Size ((31),
device-cpu, is (s_shared-False
TensorSpec comes with multiple dedicated methods run common checks and operations, such as spec i which checks if tensor belongs to the spec's domain, spec sopecpproipp project. which maps tensor onto
closest data point within spec's domain or spec. encode which creates tensor from NumPy array. Other dedicated operations such conversions from categorical encoding to one-hot are also supported.
Specs also support indexing and stacking, and for TensorDict, heterogeneous CompositeSpec instances can lazily be stacked together for better integration in multitask or multiagent frameworks.
Using these specs greatly empowers the library and is a key element of many of the performance optimizations in TorchRL. For instance, using the TensorSpecs allows us to preallocate buffers in shared memory executing the environment operations single time. For real-world deployment of TorchRL's solutions, these tools are essential as they allow us run algorithms on fake data defined by the environment specs, checking that there is no device or shape mismatch between model and environment, without requiring single data collection on hardware.
GPU maximize the throughput during parallel data collection, without
Designing environments in TorchRL To encode implement the internal __reeettt/___ () and __teepp () methods and provide the input and output domains through the specs for grounding the data. TorchRL provides set of safety checks grouped under the torchrl. envs werssuuill sutils....... www.cccheessceeeeeeeeeeeeeeeeeee
environment, developers only need to
which should be run before executing an environment for the first time. This optional validation step offloads many checks that would otherwise be executed at runtime to single initial function call, thereby reducing the footprint of these checks.
D REPLAY BUFFERS
TorchRL's replay buffers are fully composable: although they come with "batteries included", requiring a minimal effort be built, they also support many customizations such as storage type, sampling strategy or data transforms. We provide dedicated tutorials and documentation on their usage in the library. The main features can be listed as follows:
Various Storage types are proposed. We provide interfaces with regular lists, or contiguous physical or virtual memory storages through MemmapTennoo and torch. Tensor classes respectively.
At the time of writing, the Samplers include generic circular sampler, sampler without repetition and an efficient, C++ based prioritized sampler.
One can also pass Transforms which are notably identical those used conjunction with TorchRL's environments. This makes it easy to recycle data transform pipeline used during collection to one used offline to train from static dataset made of untransformed data. It also allows us to store data more efficiently: as example, consider the cost of saving images in uint8 format against saving transformed images in floating point numbers. Using the former reduces the memory footprint of the buffer, allowing it store more data and access it faster. Because the data is stored contiguously and the transforms are applied on-device. the reduction in memory footprint comes at little extra computational cost. One could also consider the usage of transforms
23

---

Published as a conference paper at ICLR 2024
from torchri, objectives import DONLoss from torchrl.date data import onellctDiscreeerrensotjpcc from torch import nn import torch
n_obs n_action =
astion_epec ouctooDissreteeeesssssecc
value_netword nn nn. Linei (n_obs, n_action) simple value model dgn_loss = DONLoss network,
define data
observation torch. randn (n_obs)
next observation torch. randn (n obs)
action acticn_speet rand)
next_reward torch. randn (1)
next_done torch. (1, dtype-torch.bool.
get loss value
loss_val don_loss((
observation-dsservatioon
rext_rewardmnex__ewwrry
pext_cone-next_ddone,
action-action)
Figure TensorDiet-free DQN Loss usage.
when working with Decision Transformers (Chen et al. 2021a), which are typically trained with different lookback window the collected data than the used during inference. Duplicating the same transform the environment and the buffer with a different set of parameters facilitates the implementation of these techniques.
Importantly, this modularity of the buffer class also makes it easy to design new buffer instances, which we expect have positive impact on researchers. For example, team of researchers developing a novel and more efficient sampling strategy can easily build new replay buffer instances using the proposed parent parent class while focusing on the improvement of a single of its pieces.
Distributed replay buffers: our buffer class supports Remote Procedural Control (RPC) calls through torch. di strriiii criibuued. rpc. In simple terms, this means that these buffers can be extended or accessed by distant nodes through a remote call to buffer. extend or bu feee.rr. sample. With the utilization of efficient shared ceemory-mapper storage when possible and multithreaded requests, the distributed replay buffer significantly enhances the transfer speed between workers, increasing by factor of three to ten when compared naive implementations.
USING LOSS MODULES WITHOUT TENSORDICT
In an effort free the library from its bounds to specific design choices, we make it possible to create certain loss modules without any interaction with TensorDict related classes. For example, figure 7 shows how a DQN loss function can be designed without recurring to a TensorDictModun instance.
Although the number of loss modules with this functionality currently limited, we are actively implementing that feature for larger set of objectives.
24

---

Published as a conference paper at ICLR 2024
F TRAINER
The modularity of TorchRL's components allows developers to write training scripts with explicit loops for data collection and neural network optimization the traditional structure adopted by most ML libraries. While this enables practitioners to keep tight control over the training process, might hinder TorchRL's adoption by first-time users looking for one-fits-all solutions. For this reason, we also provide a Trainer class that abstract this further complexity. By exposing simple train () method, trainers take care of running data collection and optimization while providing several hooks (i.c.. callbacks) different stages of the process. These hooks allow users to customize various aspects of data processing, logging, and other training operations. Trainers also encourage reproducibility by providing checkpointing interface that allows to abruptly interrupt and restart training at any given time while saving models and RB of any given size.
The trainer executes nested loop, consisting of outer loop responsible for data collection and an inner loop that utilizes this data or retrieves data from the replay buffer to train the model. The Trainer class does not constitute fundamental component; rather, it developed as simplified entry point to novice practitioners in the field.
In Figure 8 provide two code listing examples to train DDPG agent. The left approach utilizes the high-level Trainer class, On the other hand, the right listing explicitly defines the training loop, giving more control to the user over every step of the process.
G SINGLE-AGENT REINFORCEMENT LEARNING EXPERIMENTS
To test the correctness and effectiveness of TorchRI. framework components, we provide documented validations of online algorithms (PPO, A2C, IMPALA) and offline algorithms (DDPG, TD3 and SAC) for 4 MuJoCo environments (HalfCheetah-v3. Hopper-v3, Walker2d-v3 and Ant-v3) and Atari environments (Pong, Freeway, Boxing and Breakout). Training plots are shown in Figure 9
Figure 11 and Figure 12
G.1 ONLINE REINFORCEMENT LEARNING EXPERIMENTS
G.1.1 IMPLEMENTATION DETAILS
For A2C (Mnih al. 2016) and PPO (Schulman al.) 2017), we reproduce the results from the original work on MuJoCo and Atari environments, using the same hyperparameters and network architectures. We also reproduce Atari results for the IMPALA (Espeholt et al. 2018 distributed method. Ensuring a fair comparison among the off-policy algorithms DDPG (Lillicrap et al. 2015),, (D3(moomoot al. 2018b), and SAC(Haarojj et al. 2018), we uniformly apply the same archi- tecture, optimizer, learning rate, soft update parameter, and batch size as used in the official TD3 implementation. Explorationnspecifie parameters for DDPG and TD3 are taken directly from the according paper. To initiate the process, each algorithm begins with 10,000 steps of random actions, designed to adequately populate the replay buffer.
G.1.2 HYPERPARAMETERS
Tables 6 7 and & display all hyperparameters values and network architecture details required to reproduce online RL results.
G.2 OFFLINE REINFORCEMENT LEARNING EXPERIMENTS
G.2.1 IMPLEMENTATION DETAILS
To replicate the pre-training results of both the Decision Transformer Chen et al. (202 and the Online Decision Transformer Zheng al. (2022), we utilize the base GPT-2 transformer from Hugging Face Wolf et al. (2020). as presented in the official implementations. All parameters concerning architecture and training procedures
adopted directly from the specifications provided in the publications. Correspondingly, for IQL Kostrikov et al. (202I). we have selected the architecture and parameters mentioned in the paper to enable performance reproduction.
25

---

Published as a conference paper at ICLR 2024
Environment
Environment env GymEnv (' Pendulum-v1" & Model:
env GymEnv (' Penddluumvv
Model: Actor and value
Actor and value
mip_actor M
m mip_actor MLP
cem_cells=64, depth-3,
num _em_ceels=64, depth=3, in_features-3, cut_ cut_features-1
in un_feaaures-,, out put_features-1
actor TensorDictHodulett
actor TensorDictModule
mip_actor, in_keys= observationn
mip_acton
in_keys- (' observation)],
cut_keys- ['action']
out_keys- [' action'] ("attyy-("actnn))
ml p_value MIP
mip value M
num cells=64, depth=2, in in_ out_features-1
num cells=64, depth=2,
in features out_fcatures-1
critic TensorDicttequentiaal (
critic TensorDietsequuntiaal actor, TensorDictModule:
actor, TensorDictModulel
mlp_actor, keys
mlp_actor, in_keys 'observation', 'action',
'cbservation', 'action',
1, out keys_ ('state_action_value
1, out cut_keys ['state_action_value
Data Collector collector SyncDatacollecttor (
Data Collector collector SyneDataCollcctor
env, Moditiveaussian/rpppe actor
env, Addtiveaasssaaarzppprrr actor
)
) frames Baame__pr__atcc-11a1 total frames=10000000
Frames_per__a per___ee__ttttt00...
total frames=1000000,
Replay Buffer
Replay Buffer
buffer "ensorrticceeppayssfffeeer
buffer TensorriietRepplyyyafeerr
storage-laayyeensortttoagg max_size-100000,
storage-lagyTeenoosttooggg max_size-100000,
Loss Module loss_fr DDPGLosS actor, critite,
Loss Module loss_fn DDPGLOSS (
actor, critite,
ontin-torch.comm
optercchoppti...mm opt.m. Adam fass_fffarrmmmtterss 1r-2e-4,
i wwworcchopptm...mm Adam (
(),
1r-2e-4,
Trainer Trainer(
Training 1oop for data in collector: buffer.extent (data)
collectoor=collectoor
total__ramss-100000,
sample buffer. saefferssamlee (50)
frame_skip-1,
lcss loss_ (sample)
optin_stcos_ppr__aatee
lcss loss [' loss_acterr)
loss boss_module-loss_fi
loss loss value'
optimizer-optin,
loss.backward ()
optin.step
trainer.train'
optim.zero_gren
Figure 8: Trainer class example usage (left). Fully defined training loop (right).
26

---

Published as a conference paper at ICLR 2024
HarCheasah
Hopper
DGPS
35.00
DOPG
103
103
SAC
3000
SAC
8000
Roum
6000
2000
1500
4000
1000
2000
500
1.0
0.4
Enviromment Steps
conmenttt Stepe
Walce2dd
Ant
DDPS
DORG
TD3
TD3
SAC
4000
3000
Settim
5000
2000
1000
Envirorment Staps
Envicommnnt Steps
Figure 9: Online RL algorithms trained on MuJoCo environments. We report the mean and standard deviation reward over 5 random seeds. Each was trained for IM frames.
Table 6: Training parameters for single-agent on-policy algorithms on MuJoCo environments.
A2C
PPO
Discount (7)
0.99
Discount (7)
0.99
GAE A
0.95
GAE A
0.95
num envs
1
num envs
1
Horizon (T)
2048
Horizon (T)
64
Adam
3e 4
Adam
3e 4
Minibatch size
64
Minibatch size
64
Policy architecture
MLP
Policy architecture
MLP
Value net architecture
MLP
Value net architecture
MLP
Policy layers
[64, 64]
Policy layers
[64, 64]
Value net layers
(64, 64]
Value net layers
[64, 64]
Policy activation
Tanh
Policy activation
Tanh
Value net activation
Tanh
Value net activation
Tanh
Critic coef.
0.5
Critic coef.
0.5
Entropy coef.
0.0
Entropy coef.
0.0
Num. epochs
10
Clip E
0.2
27

---

Published as a conference paper at ICLR 2024
HalfCheetah-v3
16000
SAC
14000
12000
10000
Return
8000
6000
4000
2000
0.0
1.5
2.0
Environment Steps
Figure 10: To verify its paper performance we ran SAC on the HalfCheetah-v3 MuJoCo environment for 3M frames. We report the mean and standard deviation reward over 5 random seeds.
Pong
Freeway
30
25
20
Rubum
-10
PPO
PPO
A2C
A2C
4.0
Enntronmant Fromes Boxing
Framess
Breakout
1etum
50
PPO
Envirorment
2.5
Figure 11: Online RL algorithms with discrete action space trained on Atari environments. We report the mean and standard deviation reward over 5 random seeds. Algorithms were trained for 40M game frames (10 timesteps since we use frameskip 4).
28

---

Published as a conference paper at ICLR 2024
Pong
Boxing
20
80
Peturn
40
20
-20
IMPALA
IMPALA
Environment
Environment Frames
Breakout
800
Retum
800
IMPALA
1.00 Environment Frames
2,00
Figure 12: Distributed training results for IMPALA trained on Atari environments. We report the mean and standard deviation reward over 5 random seeds. IMPALA was trained for 200M game frames (50M timesteps since we use frameskip 4) for each environment, however, for Pong and Boxing environments rewards of only the first 50M frames are visualized IMPALA converged early. For Freeway, IMPALA is unable to get more that 0.0 reward, both in the original paper and in experiments.
Table 7: Training parameters for single-agent on-policy algorithms on Atari environments. For all cases, follow the environment data transforms from DeepMind (Mnih al. 2013
A2C
PPO
IMPALA
Discount (7)
0.99
Discount (v)
0.99
Discount (7)
0.99 (12 workers)
GAE X
0.95
GAE A
0.95
num envs
num envs
num envs
Horizon (T)
80
Horizon (T)
80
Horizon (T)
4096
RMSProp Ir
6e
Adam Ir
le 4
Adam Ir
2.5e
RMSProp alpha
0.99
Minibatch size
80
Minibatch size
1024
Minibatch size
32 x 80
Policy architecture
CNN
Policy architecture
CNN
Policy architecture
CNN
Value net architecture
CNN
Value net architecture
CNN
Value net architecture
CNN
Policy layers
[512]
Policy layers
[512]
Policy layers
[512]
Value net layers
[512]
Value net layers
[512]
Value net layers
[512]
Policy activation
ReLU
Policy activation
ReLU
Policy activation
ReLU
Value net activation
ReLU
Value net activation
ReLU
Value net activation
ReLU
Critic coef. Entropy coef.
0.5
Critic coef.
0.5
Critie coef.
0.5
0.01
Entropy coef.
0.01
Entropy coef.
0.01
Num. epochs
3
Clip
0.1
29

---

Published as a conference paper at ICLR 2024
Table 8: Training parameters for single-agent off off-policy algorithms.
SAC
TD3
DDPG
Discount (7)
0.99
Discount (c)
0.99
Discount (7)
0.99
Adam Ir (all nets)
3e
Adam Ir (all nets)
3e
Adam Ir (all nets)
3e 3e-
Batch size
256
Batch size
256
Batch size
256
Policy
ML.P
Policy net
ML.P
Policy net
ML.P
Q net
MLP
Q net
MLP
Q net
MLP
Policy layers
[256, 256]
Policy layers
[256, 256]
Policy layers
[256, 256]
Q net layers
[256, 256]
Q net layers
[256, 256]
Q net layers
[256, 256]
Policy activation
ReLU
Policy activation
ReLU
Policy activation
ReLU
Q net activation
ReLU
Q net activation
ReLU
Q net activation
ReLU
Target polyak
0.995
Target polyak
0.995
Target polyak
0.995
Buffer size
1000000
Buffer size
1000000
Buffer size
1000000
Init rand. frames
10000
Init rand. frames
10000
Init rand. frames
10000
Exploration noise
N0.00,0.11
Exploration noise
OU(0.0, 0.2)
Target noise
N(0.0, 0.2)
init
1.0
Noise clip
0.5
end
0.1
Policy delay
annealing steps
1000
0.15
3000
5000
4000
3000
2500
2000
1000
500
07
oDT
60T
10000
50030
20000
Gradlent
Figure 13: Offline RL algorithms pre-trained hopper-medium--2 (left) and waalichhetaa-mmdiiuee (right) D4RL datasets. We report the mean and standard deviation reward over random seeds. Each run consists of 50.000 network updates.
30

---

Published as a conference paper at ICLR 2024
Table 9: Training parameters for offline algorithms Decision Transformer (DT), Online Decision Transformer (oDT) and Implicit Q-Learning (IQL) for the hopper-medium--2 (Ho) and halfcheetah- medium-v2 (HC) D4RL datasets.
DT
oDT
IQL
Lamb Ir
le -4
Lamb Ir
le
Adam Ir
3e 3e-4
Batch size
64
Batch size
256
Batch size
256
Weight decay
5e
Weight decay
5e
Policy net
MLP
Scheduler
LamdaLR
Scheduler
L
Q net
MLP
Warmup steps
10000
Warmup steps
10000
Value net
ML.
Train context
20
Train context
20
Policy layers
[256, 256]
Eval context Ho
20
Eval context Ho
20
Q net layers
[256, 256]
Eval context HC
5
Eval context HC
5
Value net layers
[256, 256]
Policy net
GPT2
Policy net
GPT2
Policy activation
ReLU
Embd dim
128
Embd dim
512
Q net activation
ReLU
Hidden layer
Hidden layer
Value net activation
ReLU
Attn heads
Attn heads
4
Target polyak
0.995
Inner layer dim
512
Inner layer dim
2048
temperature (3)
3.0
Activation
ReLU
Activation
ReLU
expectile (r)
0.7
Resid. pdrop
0.1
Resid. pdrop
0.1
Discount (7)
0.99
Attn pdrop
0.1
Attn pdrop
0.1
Action head
[128]
Action head
[512]
Reward scaling
0.001
Reward scaling
0.001
Target return Ho
3600
Target return Ho
3600
Target return HC
6000
Target return HC
6000
init alpha
0.1
G.2.2 HYPERPARAMETERS
Table 9 displays all hyperparameter values and network architecture details required to reproduce our offline RL results.
H MULTI-AGENT REINFORCEMENT LEARNING EXPERIMENTS
Most of TorchRL's components are default-compatibl with Multi-Agent RL (MARL), and MARL- specifie components are also available. The requirement for nested data structures and multi- dimensional batch-size handling make TensorDict the ideal structure to support data representation in MARL contexts. TensorDicts allow to carry both per-agent data (e.g., reward in POMGs (Littman 1994)) and shared data (e.g.. global state, reward in Dec-POMDPs (Bernstein et al. 2202 by storing the tensors at the relevant nesting level, thus highlighting their semantic difference and optimizing storage. MARI also needs to cope with heterogeneous input/output domains, (ie, settings where the shape of the agent's is attributes differ). TorchRL. and TensorDict provide appropriate primitives to handle these cases with minimal disruption the code. As for single-agent solutions, MARL losses can easily be swapped and are semantically similar to their single-agent counterparts.
To showcase TorchRL's MARL capability, we implement six state-of-the-are algorithms and bench- mark them on three tasks in the VMAS simulator Bettini et al. (022).. Unlike existing MARL benchmarks (c.g., StarCraft (Sayvelyy et (2019), Google Research Football Kurach al. (2020). VMAS provides vectorized on-device simulation and set of scenarios focused on continuous and partially observable multi-robot cooperation tasks. TorchRI. and VMAS both use PyToorch backend, enabling performance gains when both sampling and training are run on-device. We implement MADDPG Lowe et al. (2017). IPPO de Witt et al. (2020), MAPPO Yu et al. (2022). IQL Tan [1993), VDN Sunnhaa al. (2018) 8), and QMIX Rashid et al. (2020) This selection of algorithms presents many of the different flavors discussed in this section. For uniform evaluation, we perform all training on-policy with the same hyperparameters and networks. For algorithms that require discrete actions, we use the default VMAS action discretization (which transforms 2D continuous actions into discrete directions). We evaluate all algorithms in three environments: (i) Navigation (Figure 14a).
31

---

Published as a conference paper at ICLR 2024
(a) Navigation.
(b) Balance.
(c) Sampling.
Figure 14: The three VMAS multi-robot control tasks used in the experiments.
where agents need to reach their target while avoiding collisions using LIDAR sensor, (ii) Balance (Figure 14b), where agents affected by gravity have to transport spherical package, positioned given goal the top, and (iii) Sampling (Figure 14c). where agents
randomly on top of line, need cooperatively sample probability density field. More details on the tasks, hyperparameters, and training scripts are available in the additional material. FFiuure shows the results. In all tasks, can observe the effectiveness of PPO-based methods with respect to Q-learning ones, this might be due to the suboptimality discrete actions in control tasks but also aligns with recent findings in the literature Yu al. ) de Witt al. (2020). In the sampling scenario, which requires more cooperation, we see how IQL fails due to the credit assignment problem, while QMIX and VDN are able to achieve better cooperation.
Overall, our experiments highlight how different MARL solutions can be compared on diverse tasks with minimal user effort. In fact, the scripts used for this comparison present minimal differences (mainly only in the loss and policy class installations), proving TorchRI. to flexible solution also MARL settings.
H.1 COMPARISON WITH RLLIB
To further assess the correctness of our implementations. we compare the IPPO algorithm (de Witl al. 2020) with the RLlib (Liang et al. 2018) library on the three VMAS tasks Figure 14 For the RLlib implementation, we use the code provided with (Bettini et al.) (023). The TorchRL code is publicly available in the repository examples. The comparison, shown Figure 15 demonstrates that, while TorchRI. and RLlib achieve similar returns, TorchRI. takes significantly lower time thanks its ability to leverage the vectorization of the VMAS simulator.
H.2 IMPLEMENTATION DETAILS
We now describe the implementation details for our experiments.
H.2.1 HYPERPARAMETERS
To uniform the training process across all algorithms, we train algorithms on-policy. The hyperpa- rameters used the same for all experiments and shown in Table 10 In general, the training scripts have the following structure; There is an outer loop performing sampling. At each iteration, Batch
Batch size frames are collected using # VMAS vectorized envs with Max episode steps. Mbuihath size optimization steps are then preformed for SDG Iterations using Adam optimizer. The training ends after # training iterations. Crities and actors (when used) are two-layer MLPs with 256 cells and tanh activation. Parameters are shared in all algorithms apart from MADDPG to follow the original paper implementation.
H.2.2 ENVIRONMENTS
The tasks considered scenarios taken from the VMAS Bettini et al. (2022) simulator. They all consider agents 2D continuous workspace. To move, agents take continuous 2D actions which represent control forces. Discrete action can be and will map to the five options: up, down, left,
32

---

Published as a conference paper at ICLR 2024
Navipution
Bolance
Sampling
Meen rounat
0.5
20
20
0.3
Namberr Frunnes (M)
Number framms (M)
Number of frumees (M)
RLJbb
Thechill.
(a) Return.
Navigation
Balance
Sampling
Ierrtin titar (e)
400
200
200-
1011-
10
Numbrr of Framus (M)
Number framss (M)
Nombr of framcs (M4)
TeechlLL
(b) Time.
Figure 15: Comparison of TorchRL. and RLlib on three VMAS tasks using the IPPO algorithm (de Witt al., 2020). The comparison shows that while TorchRL and RLlib achieve similar returns, TorchRL takes significantly less time thanks to its ability to leverage the vectorization of the VMAS simulator. We report the mean and standard deviation episodic return over 3 random seeds. Each run consists of 500 iterations of 60, 000 steps each, with an episode length of 100 steps.
Table 10: MARL training parameters.
Training
General
PPO
Batch size
60000
Discount y
0.9
Clip e
0.2
Minibatch size
4096
Max episode steps
100
GAE
0.9
SDG Iterations
45
NN Type
MLP
Entropy coeff
# VMAS vectorized envs
600
of layers
KL coeff
Learning rate
5e-5
Layer size
256
Max grad norm # training iterations
40
Activation
Tanh
500
# agents
3
right, and staying still. In the following, describe in detail the three environments used in the experiment:
Randomly spawned agents (circles with surrounding dots) need to navigate randomly spawned goals (smaller circles). Agents need to use LIDARs (dots around them) to avoid running into each other. For each agent, we compute the difference in the relative distance its goal over two consecutive timesteps. The mean of these values over all agents composes the shared reward, incentivizing agents to move towards their goals. Each agent observes its position, velocity, lidar readings, and relative position its goal.
Navigation (Figure 14a:
Balance (Fgurre 14b) Agents (blue circles) are spawned uniformly spaced out under line upon which lies spherical package (red circle). The team and the line are spawned at random position at the bottom of the environment. The environment has vertical gravity. The relative position of the package on the line is random. In the top half the environment, a goal (green circle) is spawned. The agents have to carry the package to the goal. Each agent receives the same reward which is proportional to the distance variation between the package and the goal over two consecutive timesteps. The team receives negative reward of -10 for making the package the line fall to the floor. The observations for each agent are: its position, velocity, relative position the package, relative position to the line, relative position between package and goal, package
33

---

Published as a conference paper at ICLR 2024
velocity, line velocity, line angular velocity, and line rotation modæ. The environment is done either when the package or the line falls or when the package touches the goal.
Sampling Figure 14c) Agents are spawned randomly in a workspace with an underlying Gaussian density function composed of three Gaussian modes. Agents need to collect samples by moving in this field. The field discretized to a grid (with agent-sized cells) and once an agent visits cell its sample is collected without replacement and given reward to the whole team. Agents can use lidar to sense each other in order to coordinate exploration. Apart from lidar, position, and velocity observations, each agent observes the values of samples in the 3x3 grid around it.
I COMPARISON OF DESIGN DECISIONS
As conclusion note, provide some more insight into the differences between TorchRL and other libraries UX and design choices.
In contrast with Stable-baselines (aselinesslaaiin et al. 2021 EfficientRL (Liu et al. 2021), TorchRL aims to provide researchers with the necessary tools to build the next generation of control algorithms, rather than providing precisely benchmarked algorithms. For this reason, TorchRL's code will usually be more verbose than SB3's because it gives users full control over the implementation details of their algorithm. For example, TorchRL does not make opinionated choices regarding architecture details or data collection setups. Nevertheless, the example repertory and tutorials are available to assist those who wish to get a sense of what configurations are typically deemed more suitable. The trainer class and multiple examples, tutorials, and rich documentation are available to help users get started with their specific problems.
On a similar note, we note that the Tianshou (Weng et al. 2022) entry point for most algorithms is usually Poli сy class that contains an actor, possibly a value network, an optimizer, and other components, In TorchRL, these items are kept separate to allow users orchestrate these components at will.
Another difference between TorchRL and other frameworks lies in the data carriers they use: relying on a common class to facilitate inter-object communication is not a new idea in the RL and control ecosystem. Nevertheless, we believe that TensorDict elegantly brings features that will drive its adop- tion by the ML community. In contrast, Tianshou's ' Bat ch has a narrower scope than TensorDict. Whereas the former is tailored for RI., the latter can be used across MI domains. Additionally, TensorDic ct has a larger set of functionalities and a dedicated tensordict.. nn package that blends within the PyTorch ecosystem as shown above.
Next, TorchRL is less an extension of a simulator than other libraries. For instance, Tianshou mainly supports Gym/Gymnasium environments (Brockman al. 2016) while TorchRL is oblivious to the simulation backend and works indifferently with DeepMind control, OpenAI Gym, or any other simulator.
Finally, TorchRL opts for a minimal set of core dependencies (PyTorch, NumPy, and tensordict) but a maximal coverage of optional external backends whether it is in terms of environments and simulators, distributed tools (Ray or submitit), or loggers. Restricting the core dependencies comes with multiple benefits, both terms of usability and efficiency: we believe that RLlib's choice of supporting multiple ML frameworks severely constrains code flexibility and increments the amount of code duplication. This has a significant impact on its Samp 1 eBat ch data carrier, which forced to be dictionary of NumPy arrays, thus leading multiple inefficient conversions if both sampling and training are performed on GPU.
J ACKNOWLEDGMENTS
We express our gratitude to the numerous contributors who have made TorchRL a reality through their dedicated efforts, whether by directly contributing to the codebase by providing valuable feedback and suggestions for its development. A complete list of these individuals can be found on the TorchRL GitHHub page. Furthermore, we extend our appreciation the PyTorch and Linux foundations for their support and confidence in this project, as they provide an ecosystem where the library can thrive (e.g., CI/CD, compute resources and technical support).
34


Reviews:
[{"review_id": "1eL41d1Kbj", "reviewer": "Reviewer_C5vY", "summary": "Authors develop a new unified library for RL training that introduces the TensorDict, which enables clean modular and composable implementations of all major RL algorithms and support for vectorized GPU environments such as IsaacGym and Brax. The TensorDict enables indexing by key like a regular python dict, but also indexing by shapes like a regular pytorch tensor, which allows it to inherit tensor-based operations like reshaping, concatenating, moving between devices, gradient storing, etc. TorchRL also includes TensorDictSequential as an nn.Sequential analogue operating on TensorDicts.\n\nTorchRL has implementations of nn.Modules including MLP, LSTMs, but also RL specific things like ActorCriticOperator.", "questions": "1. Will tensor-based operations only work on TensorDicts if the leading dimensions of every single tensor in the TensorDict shares the same leading dimension? Can these operations only be done on a subset of the items in the TensorDict?\n2. What are the next major features of TorchRL under development?\n3. How much resources will be dedicated to supporting TorchRL (github issues, new features) for the foreseeable future, compared to PyTorch as a whole?\n4. Are there computational efficiency results of TorchRL on accelerator-based simulators such as IsaacGym, Brax, and MJX?\n5. Typo in 3rd paragraph under Section 2, Subsection “Environment API, wrappers, and transformations” with the “Isaac” citation.\n6. In Section 2, subsection “Environment API, wrappers, and transformations,” how does the data transform involving “embedding through foundation models” “come at no cost in terms of functionality?” What does this sentence mean?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Writing contains detailed coverage of major existing RL frameworks and environments.\n- Substantial library implementation and development that is already having a large impact on the research community.\n- TorchRL includes a communication class.\n- Proposed TorchRL framework integrates well with existing pytorch code and all facets of RL are considered, including replay buffers, vectorized environments, and wall clock time.", "weaknesses": "### Overall\n(A1) Lacks a discussion of limitations. What are the current features that torchRL cannot handle/is bad at handling but should be improved upon in future releases?\n\n(A2) Details of TorchRL’s communication class could be more clear. What defines a communication class? What part of downstream RL algorithms does it accelerate?\n\n(A3) Prior work discussion should describe data structures similar to TensorDict that have been implemented in the past on accelerated ML libraries, such as perhaps pytrees in jax, which seem to share some similarities with TensorDict.\n\n### Computational efficiency results are a bit lacking.\n\n(B1) Paper did not measure the computational overhead/efficiency gain of using TensorDicts vs. regular python data structures containing pytorch objects. For instance, what is the efficiency of standard reshaping operations on TensorDicts of some set of tensors $X$, vs. the same reshaping operation on a regular python dict containing the same set of tensors $X$?\n\n(B2) How does an optimized TorchRL implementation compare efficiency-wise to the popularly-used pytorch implementations of different RL algorithms that are publicly available? Table 4-5 mostly compares with other RL frameworks that do not use pytorch. Without these results, it is hard to determine true efficiency effects of TensorDicts.\n\n(B3) A more in-depth computational efficiency comparison between TorchRL and Jax, on the same hardware, could help practitioners decide which library to use (although ease of use is a separate issue that cannot be easily quantifiable). There is a comparison to CleanRL (Jax) in Table 4, but not in Table 5.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "aEV2AWfMjv", "reviewer": "Reviewer_77b1", "summary": "The paper introduces TorchRL, a comprehensive decision-making library in the PyTorch ecosystem, designed to address fragmentization in the decision-making field by providing efficient, scalable, and flexible primitives suitable for complex real-world data and environments. The proposed library offers a modular design, combining components that maintain standalone functionality.  TorchRL introduces several core modular classes as well as TensorDict, a standalone library that can handle multiple tensors at the same time in the form of an improved dictionary class that supports batched tensor operations. An extensive experimental section demonstrates the correct reproduction of several famous single/multi-agent and offline RL as well as the library’s efficiency against strong competitors.", "questions": "1. Tables 3-5 show a performance comparison on a machine with 8 GPUs. Do the tested gaps with other libraries hold in single-GPU machines as well? In other words, could the gaps be in part due to the fact that TorchRL is using multiple GPUs?\n2. In the Appendix, Figure 14, it looks like TorchRL can outperform RLLib in speed but cannot reach the same rewards. Does it mean there may be some incorrect implementation?\n3. In Table 4, it is mentioned: \"Those data moves are unaccounted for in this table but can double the GAE runtime\" . Does this apply only to TorchRL?\n4. Could you quantify the overhead of using TensorDict a bit more? It would be useful to see (e.g., in Appendix E) how much the difference is between using TensorDict and not using it.\n5. You mentioned that some libraries make environments vectorizable and runnable directly on GPUs (i.e., with PyTorch or JAX). In my experience, CPU-GPU overhead is one of the major challenges in terms of efficiency. Are you planning to develop a “Gym-like” library of vectorized environments as well in the future, such as “pendulum” but with `_step` and `_reset` with tensors?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "The proposed library tries to tackle an important problem - the fragmentization of the RL field. While many libraries have been proposed over the years (in the form of the environment “Gyms” or full-stack RL solutions), TorchRL is, to the best of my knowledge, the most comprehensive and efficient with full support in the PyTorch ecosystem, which makes it exceptionally relevant given the prevalence of PyTorch implementations in the research community, as well as for practitioners. The experimental section touches on virtually all of the most important aspects for benchmarking, single-agent, multi-agent, and offline RL, as well as other benchmarks. The Github library has also received widespread support from several groups.\n\nFinally, I have personally used TorchRL and TensorDict in several projects and found them useful and efficient. I expect the library to have a great impact in streamlining the development of decision-making algorithms for practitioners and researchers alike.", "weaknesses": "The paper and its proposed library do not have any major weaknesses.\n\nThe main weakness - that does not influence my score since it is highly subjective, due to my experience while using it - is that the library’s implementation is at times hard to read - meaning that the code is at times a bit convoluted and not straightforward to understand. For instance, other platforms such as CleanRL are way easier, in my opinion (also due to their single-file implementations, which are orthogonal to the proposed library), in terms of how the whole RL process works from data collection to training, and as such perhaps more suitable for new researchers.\n\n---\n\nTypos (no influence to my score)\n\n- Page 4, `, Isaac makoviychuk2021isaac, Jumanji (Bonnet et al.),` missing `\\cite` and possibly the year for Jumanji\n- Page 22 . `implementations1`", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rKSP4TO0xN", "reviewer": "Reviewer_tnNy", "summary": "This paper introduces TorchRL, a PyTorch-based reinforcement learning/sequential decision-making library. The library supports many standard algorithms in model-free deep RL including the DQN family of algorithms, SAC, and PPO. Additionally, the library offers several additional features/algorithms that have emerged only in the past 2-3 years, including RLHF and decision transformers. There is also additional support for offline RL algorithms, multi-agent algorithms, and model-based RL algorithms such as Dreamer. \n\nCore to TorchRL is another contribution of the paper: the TensorDict (released as a second library), which is a flexible data carrier for PyTorch. It is a dictionary-style object that stores tensor-like objects.\n\nThe library is well-tested with good adoption by the community. Several standard algorithms from TorchRL are benchmarked in terms of performance and speed, demonstrating that TorchRL runs much faster than other RL libraries.", "questions": "- The introduction has strong statements indicating that other libraries are not being actively maintained or having fallen out of favor. Active maintenance is easier to see through version control, but how have the authors concluded that other libraries have fallen out of favor? My understanding is that many of these libraries are still being used.\n- For replay buffers, the paper says: \"To overcome this, TorchRL provides a single RB implementation that offers complete composability. Users can define distinct components for data storage, batch generation, pre-storage and post-sampling transforms, allowing independent customization of each aspect.\" Is there support for prioritized experience replay which is used in several model-free RL algorithms?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "- Comprehensivity: The library is extremely comprehensive. It spans different algorithm classes (e.g., offline RL, model-free Deep RL, multi-agent RL, etc.) whereas most algorithms focus on a single algorithm class (typically model-free RL). This makes this library more likely to be a single place where researchers and practitioners can go for a solution to their problem.\n- Modernity: This library includes multiple powerful tools introduced only in the past 2-3 years, including RLHF and decision transformers.", "weaknesses": "Presentation: I did not find the explanation of TensorDict, which seems integral to the paper to be very didactic. The text delves into the various functionalities of the TensorDict at a high-level, but I did not find it clear how this abstraction or these functionalities lends itself to ease development of RL algorithms. Figure 1 did not clarify this. I realize there is more context provided in the appendix, but the explanation in the main text needs to be improved. \n\nBenchmarking: For a software library, benchmarking is extremely important. It is well known that implementation details are critical for performance, making high quality algorithm implementations a core aim of deep RL libraries. All the experiments are run with a mere 3 seeds. Moreover, the Atari results are run on 40M timesteps as opposed to the usual 50M. Moreover, for some environments/algorithms, the performance seems slightly worse than these should be. For example, if we compare the SAC results from Table 1 to PFRL (https://github.com/pfnet/pfrl/tree/master/examples/mujoco/reproduction/soft_actor_critic), we see:\n- HalfCheetah: TorchRL performs at 11419 vs. around 15000 in PFRL and other papers\n- Ant: TorchRL scores about 3692 as opposed to the 5800+ seen in PFRL and other papers.\n\nSAC also slightly underperforms in the other two environments though the differences are not profound. Is there an explanation for this and for other environments? How are the hyperparameters found for these implementations? How do these implementations differ from other libraries?\n\nOf course, the results I am linking appear to be from the v2 of the environments, whereas the results here in the paper are from v3, which may explain the results. Do you know or have a citation for good results to be expected on the v3 mujoco environments? It is also possible that these differences are due to randomness, which is all the more reason to run more random seeds for benchmarking.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "q6zjfwhhCV", "reviewer": "Reviewer_gnLZ", "summary": "The paper introduces an new _modular_ reinforcement learning in PyTorch named TorchRL whose goal is provide a flexible, efficient and scalable library for rapid prototyping and research.\nThe key componnet here the TensorDict data structure enablibg easy efficient communication between different components like environments, buffers, models etc.\nIt includes many reference implementations including RL algorithms as well as distributed traning and other utilities.\nExperiments are performed to validate correctness and efficiency while showing competitive performance compared to other libraries.", "questions": "- While I understand it's not the focus, would be useful to also show how heteregoenous multi-agent problems can be setup in TorchRL's API while still potentially leveraging batching/vectorization capabilities?\n- Hyperparameters in RL are a big bane. Any plans to support automatic tuning, especially population based methods?", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Modular design: seems really well designed from the perspective of allowing new implementations from existing components vs the current standard practice of trying to fork an existing implementation where which component changed and mattered can be difficult to figure out. You either have single file implementations that don't really do distributed scaling well or you have weird nested inference to figure out how/where even is the actual algorithm implemented.\n- Modularity is demonstrated with reasonable tradeoffs for a very wide variety of RL methods and applications. \n- TensorDict as a concept is very generally useful and will hopefully see wider adoption in Torch ecosystem.\n- Distributed training and scaling building on `torch.distributed` is quite nice.  \n- Example implementations are short and clear while scaling well.", "weaknesses": "- Likely steeper learning curve and verbosity compared to higher level libraries. However, given the focus on research over applications, should not be that bad.\n- Limited to PyTorch, although DLPack has made things easier.\n- Smaller community and ecosystem. Would have been much better if this had been ready a couple of years ago that would have helped fix some of the fragmentation issues. If I were to guess, RL _algorithm_ research is somewhat waning (and has largely been not _that_ successful as a research endeavor because a lot of things that make RL work are considered engineering details).\n- Lots of tiny details matter when it comes to RL comparisons and more tooling and integration here would be useful. Reporting 3 average seed results is likely not the best way to report results when comparing RL methods.\n- Range of \"reasonableness\" is quite wide when it comes to RL implementations. Would be great if the paper also put the current best reported results for various environments to contextualize. That is also going to be helpful in convincing new algorithm implementations to start here rather than the best performing implementation in terms of reward scores etc.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (spotlight)