Paper: TorchRL: A data-driven decision-making library for PyTorch (ICLR-2024)

Reviewer comment (Reproducibility):
Limited to PyTorch, although DLPack has made things easier.

Author response:
On the library being limited to PyTorch, our focus is to maintain PyTorch as the core of RL research; hence, we currently do not plan to extend support to other backends like Jax or DLPack.

Reviewer comment (Evaluation):
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.

Author response:
With respect to your concerns regarding the nuances in RL comparisons, we have expanded our experiments to include 2 additional seeds. By providing results based on a total of  5 seeds, we aim to offer a more thorough and reliable evaluation. At the moment of writing this response, we are running the experiments and will update the paper by the end of the discussion period. Early results from these new seeds support our initial results. Regarding tooling and integration, our goal is to provide an extensive toolbox to easily and reliably experiment with how changes in a configuration can improve or deteriorate the results of a specific algorithm.

Reviewer comment (Experiments):
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.

Author response:
On providing experiments that replicate the most popular results vs the best performing implementations available:  This is a very valid point. We struggled with the decision of whether to target the best outcomes or strictly replicate the most recognized results in the literature (such as the original implementations of each algorithm). For the paper, we opted for the latter as we believed it would be a more convincing approach to demonstrate the correctness of our implementations, given that they are better known and tested. However, we intend to include state-of-the-art examples in the repository alongside the paper implementation, recognizing the potential interest of many users. We hope our library achieves this difficult balance between showcasing what we already know to work and making it possible to code the “next big thing”.

Reviewer comment (Reproducibility):
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?

Author response:
Addressing the question of heterogeneous multi-agent problems, we believe this is also a great point and it has been the focus of intensive development over this summer. The core functionality introduced was the grouping mechanism, which allows MARL users to decide how to group agents’ data in tensordicts. Data belonging to agents within the same group will be stacked, while data from agents in different groups will be kept as separate entries. This way users can decide whether to benefit from vectorization in the agent dimension, by grouping agents together, or whether to keep data separate to allow for heterogeneity. In any case, all groups can be vectorized over an environment dimension, meaning that heterogeneity choices will not impact environment vectorization benefits. While we have already produced examples and tutorials for homogeneous agents (all stacked in the same group), it is in our priorities to produce tutorials and examples for heterogeneous agents. It is also important to notice that this paradigm is already used in the BenchMARL library, which uses the TorchRL MARL grouping API to allow the training of homogeneous and heterogeneous tasks from VMAS and PettingZoo.

Reviewer comment (Experiments):
Hyperparameters in RL are a big bane. Any plans to support automatic tuning, especially population based methods?

Author response:
Finally, on the topic of hyperparameter tuning in RL, we acknowledge its critical role. While automatic tuning, particularly through population-based methods, is not currently a feature of our library, we recognize its potential and are considering it for future updates. In the meantime, we support robust logging tools like Weights & Biases for manual hyperparameter optimization.

Reviewer comment (Presentation):
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.

Author response:
Regarding the presentation of TensorDict, we understand the need for a clearer introduction and contextualization of TensorDict in the main text. To address this, we've revised the explanation, starting by introducing the challenges in coding decision-making algorithms and then transitioning to explain how TensorDict addresses these issues. This approach should provide the reader with a better context before diving into TensorDict's features.

Reviewer comment (Experiments):
Benchmarking: 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:  
- HalfCheetah: TorchRL performs at 11419 vs. around 15000 in PFRL and other papers  
- Ant: TorchRL scores about 3692 as opposed to the 5800+ seen in PFRL and other papers.  

SAC 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?

Author response:
On the benchmarking front, we recognize the critical importance of robust benchmarking for software libraries. As per your suggestion, we have expanded our experiment runs to include a total of 5 seeds for each experiment. This should provide a more reliable assessment of performances and enhance the robustness of our results. At the moment of preparing this response, we are in the process of running the experiments and plan to update the paper by the end of the discussion period. Initial results from these new seeds align with our original findings.

Concerning the Atari results, our choice of 40 million steps aligns with the methodology in the "Proximal Policy Optimization Algorithms" paper by Schulman et al., which we used as a reference point to replicate results for A2C and PPO and ensure the correctness of our implementations. This approach is consistent with some studies in the field, although we acknowledge that 50 million steps are also common. We want to clarify that the choice is not arbitrary.

In response to the performance comparisons with other environments and algorithms, we have run additional tests and can confirm that our implementations and hyperparameter choices are in line with established practices. We looked specifically at the SAC results mentioned for HalfCheetah in PFRL. Here the final performance values are for 3 million environment steps. However, we limited our simulations to 1 million steps similar to the TD3 paper, to optimize time due to constraints. This is also consistent with the main convergence pattern of off-policy algorithms in MuJoCo environments. When comparing the plots of PFRL and TorchRL for HalfCheetah, up to 1 million steps they lay in the same range of performance. Further, as mentioned, our results are obtained using the v3 of MuJoCo environments. We chose this version to match the version of environments at the time of the paper release. To directly compare the results of our implementation we additionally ran the HalfCheetah-v2 for 5 seeds and 1M steps. The results of v3 and v2 for HalfCheetah for 1 million environment steps do not show a significant difference in performance.  As asked we can point to other resources of algorithm benchmarks trained on the MuJoCo v3 like Open AI spinning up https://spinningup.openai.com/en/latest/spinningup/bench.html. Important to note, that developers used slightly different hyperparameter settings, like a higher learning rate for example.

Reviewer comment (Experiments):
Of 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.

Author response:
Further, as mentioned, our results are obtained using the v3 of MuJoCo environments. We chose this version to match the version of environments at the time of the paper release. To directly compare the results of our implementation we additionally ran the HalfCheetah-v2 for 5 seeds and 1M steps. The results of v3 and v2 for HalfCheetah for 1 million environment steps do not show a significant difference in performance.  As asked we can point to other resources of algorithm benchmarks trained on the MuJoCo v3 like Open AI spinning up https://spinningup.openai.com/en/latest/spinningup/bench.html. Important to note, that developers used slightly different hyperparameter settings, like a higher learning rate for example.

Reviewer comment (Novelty):
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.

Author response:
Regarding the maintenance and popularity of other libraries, we have revised our statements to focus solely on the aspect of active maintenance, removing any assertions about their popularity or favorability among users.

Reviewer comment (Reproducibility):
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?

Author response:
Lastly, for replay buffers, TorchRL not only supports prioritized experience replay but our implementation is backed by an efficient C++ implementation of Sum and MinTree that allows for a fast, multithreaded retrieval of priorities for millions of items. The option to use prioritized experience replay or not can be directly selected in the configuration of each off-policy algorithm. Regarding the phrase "single RB implementation," we intend to convey that TorchRL utilizes a unified Replay Buffer class, which can be configured with a variety of storage options (such as physical memory and lists), sampling methods (including circular and prioritized sampling), along with different writers and transformation capabilities. This flexible architecture allows us to offer numerous Replay Buffer implementations.

Reviewer comment (Reproducibility):
The main weakness — which does not influence the reviewer’s score since it is highly subjective and based on their experience using it — is that the library’s implementation is at times hard to read, meaning that the code can be a bit convoluted and not straightforward to understand. For instance, other platforms such as CleanRL are, in the reviewer’s opinion, way easier to follow (also due to their single-file implementations, which are orthogonal to the proposed library) in terms of understanding the RL process from data collection to training, and thus perhaps more suitable for new researchers.

Author response:
Firstly, regarding the readability of our library's implementation, we acknowledge that it can be challenging at times. We are committed to enhancing the code's clarity and documentation and are completely open to suggestions from the community, but striking a balance between a modular, flexible design and simplicity is a delicate task. We would like to emphasize that: 
(1) we do not seek the same "single file, all-inclusive" implementation as CleanRL, and instead focus on flexibility, modularity, breadth and reusability.
(2) Some users find a "trainer" class with a single entry point more readable than self-contained scripts. Others prefer single scripts. A third category likes scripts with a separate “utils” script that contains the constructors and some extra features. It is very hard (if not impossible) to satisfy them all with a single solution.
In any case, we will make sure to focus more on the readability of our examples in the future and improve the library's documentation at the same time.

Reviewer comment (Writing):
Typo on Page 4: `, Isaac makoviychuk2021isaac, Jumanji (Bonnet et al.),` is missing `\cite` and possibly the year for Jumanji.

Author response:
Finally, thank you for pointing out the typos; we have rectified them in the revised version of the paper.

Reviewer comment (Writing):
Typo on Page 22: `implementations1`.

Author response:
Finally, thank you for pointing out the typos; we have rectified them in the revised version of the paper.

Reviewer comment (Experiments):
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?

Author response:
Concerning the performance comparison on a machine with 8 GPU in tables 3-5,  This should have been more clearly stated in the manuscript. The node we were using has 8 GPUs, but we only required 1 GPU through the scheduler to run these experiments. We specified the full node configuration for reproducibility purposes on a similar machine on a public cluster. The speed with 8 GPUs would have been better than the one displayed, but we did not show it to stay within reasonable bounds for an average (academic) user without access to a dedicated, recent cluster. We clarify now in the manuscript that we use only 1 GPU.

Reviewer comment (Experiments):
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?

Author response:
In response to your question about Figure 14 in the Appendix: in the comparison reported in appendix H.1 we tried to match the exact hyperparameters used in TorchRL and RLlib libraries. However, this is possible only to a certain degree, as RLlib’s implementations do not provide full transparency and an in-depth knowledge of the codebase is needed for a better understanding. The cause for the slight reward mismatch could be due to this, as we are using a very simple training loop for TorchRL while RLlib could be performing some minor extra optimizations.  Additionally, we believe we can exclude an incorrect implementation for two reasons: (1) the PPO implementation used in this experiment is the same used throughout the other experiments in the paper, which shows that its performance matches the expected one on other simulators (2) the difference in rewards could be argued to be negligible (both implementations reach the “solved'' state for the task) and the sample efficiency profiles follow the same shape.

Reviewer comment (Evaluation):
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?

Author response:
Regarding the question of whether or not the statement in table 4 "Those data moves are unaccounted for in this table but can double the GAE runtime" applies only to TorchRL, the answer is no, we refer to the other libraries. All the libraries we compare against work directly with numpy (ie, the replay buffer stores numpy arrays, gae computation in numpy) and have to cast to pytorch once the batch is ready. Casting the vectors to pytorch takes extra time that is not accounted for in the table. Because TorchRL works directly with pytorch-backed data storage, the casting operation is not needed. However, we decided not to include the casting time in the table because it is not strictly part of the operation.

Reviewer comment (Presentation):
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?

Author response:
We agree more emphasis should be put on the library’s current limitations. Some trade-offs are due to design decisions, i.e. TorchRL is a Pytorch-first library, so JAX users would need to do conversions to use it. Modularity and generality also mean that the focus is not just on simplicity but we need to support complex production scenarios. FSDP is one case that we are currently improving (expected for Q1 2024). Better support for torch.compile is also needed and should appear in the same timeframe. We are also working on eliminating any overhead due to TensorDict conversions.  We are finally still actively improving the documentation.

Reviewer comment (Reproducibility):
Details of TorchRL’s communication class could be more clear. What defines a communication class? What part of downstream RL algorithms does it accelerate?

Author response:
Regarding the details of TorchRL’s communication class:  With “communication class” we are referring to the TensorDict data carrier itself. We acknowledge the poor choice of words. This observation also raises a valid point: the specific role in accelerating downstream RL algorithms could be introduced more explicitly, explaining how TensorDict features help coding RL algorithms. This suggestion aligns with other reviewers' feedback, it is now addressed in the manuscript.

Reviewer comment (Novelty):
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.

Author response:
Following the reviewer’s advice, we added in the appendix a more extensive comparison of TensorDict with other solutions across the RL landscape (incl. pytrees).

Reviewer comment (Experiments):
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\)?

Author response:
Additionally, we have also added a detailed benchmark of TensorDict overhead against PyTrees and other solutions in the appendix. We are actively working on resolving existing issues and we have a clear roadmap to address these, which is detailed in our repository.

Reviewer comment (Experiments):
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.

Author response:
Regarding efficiency comparisons with other pytorch-based libs, in table 4-5 we compared TorchRL throughput with SB3, Tianshou and RLLib, which all use PyTorch but not TensorDict.

Reviewer comment (Experiments):
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.

Author response:
CleanRL not in table 5: The data collection speed measures in table 5 were executed on Gym environments where the code execution was either numpy-based or another backend (eg, ALE). In these scenarios, what we measured was independent of the ML backend (Jax or PyTorch) but related to how data was collected across environments. We would not expect a different throughput with CleanRL, which uses the batched environment API from gym itself.

Reviewer comment (Reproducibility):
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?

Author response:
Tensor-based operations: In general, a TensorDict with an empty batch-size can contain tensors of any shape, and shape-related ops are not available (eg, split, view). If shape-related operations are required over a subset of tensors, there are currently two main solutions:
1. Organise tensors semantically: a root tensordict with an empty shape and a nested tensordict with a non-empty shape that accepts shape-related operations. This is how we approach the problem in the MARL API.
2. Isolate a set of tensors with `select`, reset the shape, and call a shape-related method. Depending on the use case there may be simpler ways of achieving the same result.

Reviewer comment (Experiments):
Are there computational efficiency results of TorchRL on accelerator-based simulators such as IsaacGym, Brax, and MJX?

Author response:
Results on accelerator-based simulators: The results in Table 3 are obtained with VMAS, an accelerator-based simulator. It is shown how TorchRL compares to RLLib in leveraging vectorization to obtain higher throughput. Other pytorch accelerator-based simulators will behave similarly (e.g., Isaac or habitat). We did not include these metrics in the paper for the sake of conciseness (ie, these simulators will perform at their best when executed within the library). Jax-based simulators are compatible with TorchRL too, but they will suffer from a consequent jax-to-torch io overhead.Our tutorials provide a simple example of how to code an accelerator-based environment using the TorchRL API.

Reviewer comment (Writing):
Typo in 3rd paragraph under Section 2, Subsection “Environment API, wrappers, and transformations” with the “Isaac” citation.

Author response:
Finally, We have fixed the typo.

Reviewer comment (Writing):
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?

Author response:
Regarding the confusing sentence in Section 2, we have now corrected the phrasing.
What we meant is that using lists of transforms rather than wrappers does not lower the amount of things one can do. A common worry with lists of transforms is that a transform that is not a wrapper cannot access the environment information for instance. TorchRL's transforms have a bidirectional registration mechanism that allows them to access properties of the parent env, thereby bringing flexibility without the cost.