torchrl 0.11.0__cp314-cp314t-macosx_11_0_arm64.whl

benchmarks/test_replaybuffer_benchmark.py

@@ -0,0 +1,254 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+import argparse
+import functools
+
+import pytest
+import torch
+from tensordict import TensorDict
+
+from torchrl.data import (
+    LazyMemmapStorage,
+    LazyTensorStorage,
+    ListStorage,
+    ReplayBuffer,
+    TensorDictPrioritizedReplayBuffer,
+    TensorDictReplayBuffer,
+)
+from torchrl.data.replay_buffers import (
+    RandomSampler,
+    SamplerWithoutReplacement,
+    SliceSampler,
+)
+
+_TensorDictPrioritizedReplayBuffer = functools.partial(
+    TensorDictPrioritizedReplayBuffer, alpha=1, beta=0.9
+)
+# preserve the name of the class even after partial
+_TensorDictPrioritizedReplayBuffer.__name__ = TensorDictPrioritizedReplayBuffer.__name__
+
+
+class create_rb:
+    def __init__(self, rb, storage, sampler, populated, size=1_000_000):
+        self.storage = storage
+        self.rb = rb
+        self.sampler = sampler
+        self.populated = populated
+        self.size = size
+
+    def __call__(self):
+
+        kwargs = {"batch_size": 256}
+        if self.sampler is not None:
+            kwargs["sampler"] = self.sampler()
+        if self.storage is not None:
+            kwargs["storage"] = self.storage(self.size)
+
+        rb = self.rb(**kwargs)
+        data = TensorDict(
+            {
+                "a": torch.zeros(self.size, 5),
+                ("b", "c"): torch.zeros(self.size, 3, 32, 32, dtype=torch.uint8),
+            },
+            batch_size=[self.size],
+        )
+        if "sampler" in kwargs and isinstance(kwargs["sampler"], SliceSampler):
+            data["traj"] = torch.arange(self.size) // 123
+        if self.populated:
+            rb.extend(data)
+            return ((rb,), {})
+        else:
+            return ((rb, data), {})
+
+
+def populate(rb, td):
+    rb.extend(td)
+
+
+def sample(rb):
+    rb.sample()
+
+
+def iterate(rb):
+    next(rb)
+
+
+@pytest.mark.parametrize(
+    "rb,storage,sampler,size",
+    [
+        [TensorDictReplayBuffer, ListStorage, RandomSampler, 4000],
+        [TensorDictReplayBuffer, LazyMemmapStorage, RandomSampler, 10_000],
+        [TensorDictReplayBuffer, LazyTensorStorage, RandomSampler, 10_000],
+        [TensorDictReplayBuffer, ListStorage, SamplerWithoutReplacement, 4000],
+        [TensorDictReplayBuffer, LazyMemmapStorage, SamplerWithoutReplacement, 10_000],
+        [TensorDictReplayBuffer, LazyTensorStorage, SamplerWithoutReplacement, 10_000],
+        [
+            TensorDictReplayBuffer,
+            LazyMemmapStorage,
+            functools.partial(SliceSampler, num_slices=8, traj_key="traj"),
+            10_000,
+        ],
+        [
+            TensorDictReplayBuffer,
+            LazyTensorStorage,
+            functools.partial(SliceSampler, num_slices=8, traj_key="traj"),
+            10_000,
+        ],
+        [_TensorDictPrioritizedReplayBuffer, ListStorage, None, 4000],
+        [_TensorDictPrioritizedReplayBuffer, LazyMemmapStorage, None, 10_000],
+        [_TensorDictPrioritizedReplayBuffer, LazyTensorStorage, None, 10_000],
+    ],
+)
+def test_rb_sample(benchmark, rb, storage, sampler, size):
+    (rb,), _ = create_rb(
+        rb=rb,
+        storage=storage,
+        sampler=sampler,
+        populated=True,
+        size=size,
+    )()
+    torch.manual_seed(0)
+    benchmark(sample, rb)
+
+
+def infinite_iter(obj):
+    torch.manual_seed(0)
+    while True:
+        yield from iter(obj)
+
+
+@pytest.mark.parametrize(
+    "rb,storage,sampler,size",
+    [
+        [TensorDictReplayBuffer, ListStorage, RandomSampler, 4000],
+        [TensorDictReplayBuffer, LazyMemmapStorage, RandomSampler, 10_000],
+        [TensorDictReplayBuffer, LazyTensorStorage, RandomSampler, 10_000],
+        [TensorDictReplayBuffer, ListStorage, SamplerWithoutReplacement, 4000],
+        [TensorDictReplayBuffer, LazyMemmapStorage, SamplerWithoutReplacement, 10_000],
+        [TensorDictReplayBuffer, LazyTensorStorage, SamplerWithoutReplacement, 10_000],
+        [_TensorDictPrioritizedReplayBuffer, ListStorage, None, 4000],
+        [_TensorDictPrioritizedReplayBuffer, LazyMemmapStorage, None, 10_000],
+        [_TensorDictPrioritizedReplayBuffer, LazyTensorStorage, None, 10_000],
+    ],
+)
+def test_rb_iterate(benchmark, rb, storage, sampler, size):
+    (rb,), _ = create_rb(
+        rb=rb,
+        storage=storage,
+        sampler=sampler,
+        populated=True,
+        size=size,
+    )()
+    benchmark(iterate, infinite_iter(rb))
+
+
+@pytest.mark.parametrize(
+    "rb,storage,sampler,size",
+    [
+        [TensorDictReplayBuffer, ListStorage, RandomSampler, 400],
+        [TensorDictReplayBuffer, LazyMemmapStorage, RandomSampler, 400],
+        [TensorDictReplayBuffer, LazyTensorStorage, RandomSampler, 400],
+        [TensorDictReplayBuffer, ListStorage, SamplerWithoutReplacement, 400],
+        [TensorDictReplayBuffer, LazyMemmapStorage, SamplerWithoutReplacement, 400],
+        [TensorDictReplayBuffer, LazyTensorStorage, SamplerWithoutReplacement, 400],
+        [_TensorDictPrioritizedReplayBuffer, ListStorage, None, 400],
+        [_TensorDictPrioritizedReplayBuffer, LazyMemmapStorage, None, 400],
+        [_TensorDictPrioritizedReplayBuffer, LazyTensorStorage, None, 400],
+    ],
+)
+def test_rb_populate(benchmark, rb, storage, sampler, size):
+    benchmark.pedantic(
+        populate,
+        setup=create_rb(
+            rb=rb,
+            storage=storage,
+            sampler=sampler,
+            populated=False,
+            size=size,
+        ),
+        iterations=1,
+        rounds=50,
+    )
+
+
+class create_compiled_tensor_rb:
+    def __init__(
+        self, rb, storage, sampler, storage_size, data_size, iters, compilable=False
+    ):
+        self.storage = storage
+        self.rb = rb
+        self.sampler = sampler
+        self.storage_size = storage_size
+        self.data_size = data_size
+        self.iters = iters
+        self.compilable = compilable
+
+    def __call__(self):
+        kwargs = {}
+        if self.sampler is not None:
+            kwargs["sampler"] = self.sampler()
+        if self.storage is not None:
+            kwargs["storage"] = self.storage(
+                self.storage_size, compilable=self.compilable
+            )
+
+        rb = self.rb(batch_size=3, compilable=self.compilable, **kwargs)
+        data = torch.randn(self.data_size, 1)
+        return ((rb, data, self.iters), {})
+
+
+def extend_and_sample(rb, td, iters):
+    for _ in range(iters):
+        rb.extend(td)
+        rb.sample()
+
+
+def extend_and_sample_compiled(rb, td, iters):
+    @torch.compile
+    def fn(td):
+        rb.extend(td)
+        rb.sample()
+
+    for _ in range(iters):
+        fn(td)
+
+
+@pytest.mark.parametrize(
+    "rb,storage,sampler,storage_size,data_size,iters,compiled",
+    [
+        [ReplayBuffer, LazyTensorStorage, RandomSampler, 10_000, 10_000, 100, True],
+        [ReplayBuffer, LazyTensorStorage, RandomSampler, 10_000, 10_000, 100, False],
+        [ReplayBuffer, LazyTensorStorage, RandomSampler, 100_000, 10_000, 100, True],
+        [ReplayBuffer, LazyTensorStorage, RandomSampler, 100_000, 10_000, 100, False],
+        [ReplayBuffer, LazyTensorStorage, RandomSampler, 1_000_000, 10_000, 100, True],
+        [ReplayBuffer, LazyTensorStorage, RandomSampler, 1_000_000, 10_000, 100, False],
+    ],
+)
+def test_rb_extend_sample(
+    benchmark, rb, storage, sampler, storage_size, data_size, iters, compiled
+):
+    if compiled:
+        torch._dynamo.reset_code_caches()
+
+    benchmark.pedantic(
+        extend_and_sample_compiled if compiled else extend_and_sample,
+        setup=create_compiled_tensor_rb(
+            rb=rb,
+            storage=storage,
+            sampler=sampler,
+            storage_size=storage_size,
+            data_size=data_size,
+            iters=iters,
+            compilable=compiled,
+        ),
+        iterations=1,
+        warmup_rounds=10,
+        rounds=50,
+    )
+
+
+if __name__ == "__main__":
+    args, unknown = argparse.ArgumentParser().parse_known_args()
+    pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)

sota-check/README.md

@@ -0,0 +1,35 @@
+# SOTA Performance checks
+
+This folder contains a `submitit-release-check.sh` file that executes all
+the training scripts using `sbatch` with the default configuration and long them
+into a common WandB project.
+
+This script is to be executed before every release to assess the performance of
+the various algorithms available in torchrl. The name of the project will include
+the specific commit of torchrl used to run the scripts (e.g. `torchrl-examples-check-<commit>`).
+
+## Usage
+
+To display the script usage, you can use the `--help` option:
+
+```bash
+./submitit-release-check.sh --help
+```
+
+## Setup
+
+The following setup should allow you to run the scripts:
+
+```bash
+export MUJOCO_GL=egl
+
+conda create -n rl-sota-bench python=3.10 -y 
+conda install anaconda::libglu -y
+pip3 install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu121
+pip3 install "gymnasium[atari,mujoco]" vmas tqdm wandb pygame "moviepy<2.0.0" imageio submitit hydra-core transformers
+
+cd /path/to/tensordict
+python setup.py develop
+cd /path/to/torchrl
+python setup.py develop
+```

sota-implementations/README.md

@@ -0,0 +1,142 @@
+# Examples
+
+We provide examples to train the following algorithms:
+- [CQL](../sota-implementations/cql/)
+- [DDPG](ddpg/ddpg.py)
+- [DQN](../sota-implementations/dqn/)
+- [Decision Transformers](../sota-implementations/decision_transformer)
+- [Discrete SAC](discrete_sac/discrete_sac.py)
+- [Dreamer](../sota-implementations/dreamer)
+- [IQL](iql/)
+- [Impala](impala/)
+- [PPO](../sota-implementations/ppo/)
+- [REDQ](redq/redq.py)
+- [SAC](sac/sac.py)
+- [TD3](../sota-implementations/td3/td3.py)
+- [Various multiagent examples](multiagent/)
+
+To run these examples, make sure you have installed hydra:
+```
+pip install hydra-core
+```
+
+Scripts can be run from the directory of interest using:
+```
+python sac.py
+```
+or similar. Hyperparameters can be easily changed by providing the arguments to hydra:
+```
+python sac.py collector.frames_per_batch=63
+```
+
+[//]: # (# Results)
+
+[//]: # ()
+[//]: # (Here we can see some results for the SAC and REDQ algorithm.)
+
+[//]: # (We average the results over 5 different seeds and plot the standard error.)
+
+[//]: # (## Gym's HalfCheetah-v4)
+
+[//]: # ()
+[//]: # (<p align="center">)
+
+[//]: # (<img src="media/halfcheetah_chart.png" width="600px">)
+
+[//]: # (</p>)
+
+[//]: # (To reproduce a single run:)
+
+[//]: # ()
+[//]: # (```)
+
+[//]: # (python sac/sac.py env.name="HalfCheetah-v4" env.task="" env.library="gym")
+
+[//]: # (```)
+
+[//]: # ()
+[//]: # (``` )
+
+[//]: # (python redq/redq.py env.name="HalfCheetah-v4" env.library="gymnasium")
+
+[//]: # (```)
+
+[//]: # ()
+[//]: # ()
+[//]: # (## dm_control's cheetah-run)
+
+[//]: # ()
+[//]: # (<p align="center">)
+
+[//]: # (<img src="media/cheetah_chart.png" width="600px">)
+
+[//]: # (</p>)
+
+[//]: # (To reproduce a single run:)
+
+[//]: # ()
+[//]: # (```)
+
+[//]: # (python sac/sac.py env.name="cheetah" env.task="run" env.library="dm_control")
+
+[//]: # (```)
+
+[//]: # ()
+[//]: # (``` )
+
+[//]: # (python redq/redq.py env.name="cheetah" env.task="run" env.library="dm_control")
+
+[//]: # (```)
+
+[//]: # ()
+[//]: # ([//]: # &#40;TODO: adapt these scripts&#41;)
+[//]: # ([//]: # &#40;## Gym's Ant-v4&#41;)
+[//]: # ()
+[//]: # ([//]: # &#40;&#41;)
+[//]: # ([//]: # &#40;<p align="center">&#41;)
+[//]: # ()
+[//]: # ([//]: # &#40;<img src="media/ant_chart.png" width="600px">&#41;)
+[//]: # ()
+[//]: # ([//]: # &#40;</p>&#41;)
+[//]: # ()
+[//]: # ([//]: # &#40;To reproduce a single run:&#41;)
+[//]: # ()
+[//]: # ([//]: # &#40;&#41;)
+[//]: # ([//]: # &#40;```&#41;)
+[//]: # ()
+[//]: # ([//]: # &#40;python sac/sac.py env.name="Ant-v4" env.task="" env.library="gym"&#41;)
+[//]: # ()
+[//]: # ([//]: # &#40;```&#41;)
+[//]: # ()
+[//]: # ([//]: # &#40;&#41;)
+[//]: # ([//]: # &#40;``` &#41;)
+[//]: # ()
+[//]: # ([//]: # &#40;python redq/redq.py env_name="Ant-v4" env_task="" env_library="gym"&#41;)
+[//]: # ()
+[//]: # ([//]: # &#40;```&#41;)
+[//]: # ()
+[//]: # ([//]: # &#40;&#41;)
+[//]: # ([//]: # &#40;## Gym's Walker2D-v4&#41;)
+[//]: # ()
+[//]: # ([//]: # &#40;&#41;)
+[//]: # ([//]: # &#40;<p align="center">&#41;)
+[//]: # ()
+[//]: # ([//]: # &#40;<img src="media/walker2d_chart.png" width="600px">&#41;)
+[//]: # ()
+[//]: # ([//]: # &#40;</p>&#41;)
+[//]: # ()
+[//]: # ([//]: # &#40;To reproduce a single run:&#41;)
+[//]: # ()
+[//]: # ([//]: # &#40;&#41;)
+[//]: # ([//]: # &#40;```&#41;)
+[//]: # ()
+[//]: # ([//]: # &#40;python sac/sac.py env_name="Walker2D-v4" env_task="" env_library="gym"&#41;)
+[//]: # ()
+[//]: # ([//]: # &#40;```&#41;)
+[//]: # ()
+[//]: # ([//]: # &#40;&#41;)
+[//]: # ([//]: # &#40;``` &#41;)
+[//]: # ()
+[//]: # ([//]: # &#40;python redq/redq.py env_name="Walker2D-v4" env_task="" env_library="gym"&#41;)
+[//]: # ()
+[//]: # ([//]: # &#40;```&#41;)

sota-implementations/a2c/README.md

@@ -0,0 +1,39 @@
+## Reproducing Advantage Actor Critic (A2C) Algorithm Results
+
+This repository contains scripts that enable training agents using the Advantage Actor Critic (A2C) Algorithm on MuJoCo and Atari environments. We follow the original paper [Asynchronous Methods for Deep Reinforcement Learning](https://arxiv.org/abs/1602.01783) by Mnih et al. (2016) to implement the A2C algorithm but fix the number of steps during the collection phase.
+
+
+## Examples Structure
+
+Please note that each example is independent of each other for the sake of simplicity. Each example contains the following files:
+
+1. **Main Script:** The definition of algorithm components and the training loop can be found in the main script  (e.g. a2c_atari.py).
+
+2. **Utils File:** A utility file is provided to contain various helper functions, generally to create the environment and the models (e.g. utils_atari.py).
+
+3. **Configuration File:** This file includes default hyperparameters specified in the original paper. Users can modify these hyperparameters to customize their experiments  (e.g. config_atari.yaml).
+
+
+## Running the Examples
+
+You can execute the A2C algorithm on Atari environments by running the following command:
+
+```bash
+python a2c_atari.py compile.compile=1 compile.cudagraphs=1
+```
+
+
+You can execute the A2C algorithm on MuJoCo environments by running the following command:
+
+```bash
+python a2c_mujoco.py compile.compile=1 compile.cudagraphs=1
+``` 
+
+## Runtimes
+
+Runtimes when executed on H100:
+
+| Environment | Eager     | Compile   | Compile+cudagraphs |
+|-------------|-----------|-----------|--------------------|
+| MUJOCO      | < 25 mins | < 23 mins | < 20 mins          |
+| ATARI       | < 85 mins | < 60 mins | < 45 mins          |