torchrl 0.11.0__cp314-cp314-manylinux_2_28_aarch64.whl

benchmarks/ecosystem/vmas_rllib_vs_torchrl_sampling_performance.py

@@ -0,0 +1,231 @@
+# 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 os
+import pickle
+
+import time
+from pathlib import Path
+
+import numpy as np
+
+import ray
+
+import vmas
+from matplotlib import pyplot as plt
+from ray import tune
+
+from ray.rllib.agents.ppo import PPOTrainer
+from ray.rllib.algorithms.callbacks import DefaultCallbacks
+from ray.tune import register_env
+from torchrl._utils import logger as torchrl_logger
+from torchrl.collectors import SyncDataCollector
+from torchrl.envs.libs.vmas import VmasEnv
+from vmas import Wrapper
+
+
+def store_pickled_evaluation(name: str, evaluation: dict):
+    save_folder = f"{os.path.dirname(os.path.realpath(__file__))}"
+    file = f"{save_folder}/{name}.pkl"
+
+    pickle.dump(evaluation, open(file, "wb"))
+
+
+def load_pickled_evaluation(
+    name: str,
+):
+    save_folder = f"{os.path.dirname(os.path.realpath(__file__))}"
+    file = Path(f"{save_folder}/{name}.pkl")
+
+    if file.is_file():
+        return pickle.load(open(file, "rb"))
+    return None
+
+
+def run_vmas_torchrl(
+    scenario_name: str, n_envs: int, n_steps: int, device: str, seed: int = 0
+):
+    env = VmasEnv(
+        scenario_name,
+        device=device,
+        num_envs=n_envs,
+        continuous_actions=False,
+        seed=seed,
+    )
+
+    collector = SyncDataCollector(
+        env,
+        policy=None,
+        device=device,
+        frames_per_batch=n_envs * n_steps,
+        total_frames=n_envs * n_steps,
+    )
+
+    init_time = time.time()
+
+    for _data in collector:
+        pass
+
+    total_time = time.time() - init_time
+    collector.shutdown()
+    return total_time
+
+
+def run_vmas_rllib(
+    scenario_name: str, n_envs: int, n_steps: int, device: str, seed: int = 0
+):
+    class TimerCallback(DefaultCallbacks):
+        result_time = None
+
+        def on_train_result(
+            self,
+            *,
+            algorithm,
+            result: dict,
+            **kwargs,
+        ) -> None:
+            TimerCallback.result_time = (
+                result["timers"]["training_iteration_time_ms"]
+                - result["timers"]["learn_time_ms"]
+            )
+
+    def env_creator(config: dict):
+        env = vmas.make_env(
+            scenario=config["scenario_name"],
+            num_envs=config["num_envs"],
+            device=config["device"],
+            continuous_actions=False,
+            wrapper=Wrapper.RLLIB,
+        )
+        return env
+
+    if not ray.is_initialized():
+        ray.init()
+    register_env(scenario_name, lambda config: env_creator(config))
+
+    num_gpus = 0.5 if device == "cuda" else 0
+    num_gpus_per_worker = 0.5 if device == "cuda" else 0
+    tune.run(
+        PPOTrainer,
+        stop={"training_iteration": 1},
+        config={
+            "seed": seed,
+            "framework": "torch",
+            "env": scenario_name,
+            "train_batch_size": n_envs * n_steps,
+            "rollout_fragment_length": n_steps,
+            "sgd_minibatch_size": n_envs * n_steps,
+            "num_gpus": num_gpus,
+            "num_workers": 0,
+            "num_gpus_per_worker": num_gpus_per_worker,
+            "num_envs_per_worker": n_envs,
+            "batch_mode": "truncate_episodes",
+            "env_config": {
+                "device": device,
+                "num_envs": n_envs,
+                "scenario_name": scenario_name,
+                "max_steps": n_steps,
+            },
+            "callbacks": TimerCallback,
+        },
+    )
+    assert TimerCallback.result_time is not None
+    TimerCallback.result_time /= 1_000  # convert to seconds
+    return TimerCallback.result_time
+
+
+def run_comparison_torchrl_rllib(
+    scenario_name: str,
+    device: str,
+    n_steps: int = 100,
+    max_n_envs: int = 3000,
+    step_n_envs: int = 3,
+):
+    """
+
+    Args:
+        scenario_name (str): name of scenario to benchmark
+        device (str):  device to ron comparison on ("cpu" or "cuda")
+        n_steps (int):  number of environment steps
+        max_n_envs (int): the maximum number of parallel environments to test
+        step_n_envs (int): the step size in number of environments from 1 to max_n_envs
+
+    """
+    list_n_envs = np.linspace(1, max_n_envs, step_n_envs)
+
+    figure_name = f"VMAS_{scenario_name}_{n_steps}_{device}_steps_rllib_vs_torchrl"
+    figure_name_pkl = figure_name + f"_range_{1}_{max_n_envs}_num_{step_n_envs}"
+
+    evaluation = load_pickled_evaluation(figure_name_pkl)
+    if not evaluation:
+        evaluation = {}
+    for framework in ["TorchRL", "RLlib"]:
+        if framework not in evaluation.keys():
+            torchrl_logger.info(f"\nFramework {framework}")
+            vmas_times = []
+            for n_envs in list_n_envs:
+                n_envs = int(n_envs)
+                torchrl_logger.info(f"Running {n_envs} environments")
+                if framework == "TorchRL":
+                    vmas_times.append(
+                        (n_envs * n_steps)
+                        / run_vmas_torchrl(
+                            scenario_name=scenario_name,
+                            n_envs=n_envs,
+                            n_steps=n_steps,
+                            device=device,
+                        )
+                    )
+                else:
+                    vmas_times.append(
+                        (n_envs * n_steps)
+                        / run_vmas_rllib(
+                            scenario_name=scenario_name,
+                            n_envs=n_envs,
+                            n_steps=n_steps,
+                            device=device,
+                        )
+                    )
+                torchrl_logger.info(f"fps {vmas_times[-1]}s")
+            evaluation[framework] = vmas_times
+
+    store_pickled_evaluation(name=figure_name_pkl, evaluation=evaluation)
+
+    fig, ax = plt.subplots()
+    for key, item in evaluation.items():
+        ax.plot(
+            list_n_envs,
+            item,
+            label=key,
+        )
+
+    plt.xlabel("Number of batched environments", fontsize=14)
+    plt.ylabel("Frames per second", fontsize=14)
+    ax.legend(loc="upper left")
+
+    ax.set_title(
+        f"Execution time of '{scenario_name}' for {n_steps} steps on {device}.",
+        fontsize=8,
+    )
+
+    save_folder = os.path.dirname(os.path.realpath(__file__))
+    plt.savefig(f"{save_folder}/{figure_name}.pdf")
+
+
+if __name__ == "__main__":
+    # pip install matplotlib
+    # pip install "ray[rllib]"==2.1.0
+    # pip install torchrl
+    # pip install vmas
+    # pip install numpy==1.23.5
+
+    run_comparison_torchrl_rllib(
+        scenario_name="simple_spread",
+        device="cuda",
+        n_steps=100,
+        max_n_envs=30000,
+        step_n_envs=10,
+    )

benchmarks/requirements.txt

@@ -0,0 +1,7 @@
+pytest-benchmark
+tenacity
+safetensors
+tqdm
+pandas
+numpy
+matplotlib

benchmarks/storage/benchmark_sample_latency_over_rpc.py

@@ -0,0 +1,193 @@
+# 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.
+
+"""
+Sample latency benchmarking (using RPC)
+======================================
+A rough benchmark of sample latency using different storage types over the network using `torch.rpc`.
+Run this script with --rank=0 and --rank=1 flags set in separate processes - these ranks correspond to the trainer worker and buffer worker respectively, and both need to be initialised.
+e.g. to benchmark LazyMemmapStorage, run the following commands using either two separate shells or multiprocessing.
+    - python3 benchmark_sample_latency_over_rpc.py --rank=0 --storage=LazyMemmapStorage
+    - python3 benchmark_sample_latency_over_rpc.py --rank=1 --storage=LazyMemmapStorage
+This code is based on examples/distributed/distributed_replay_buffer.py.
+"""
+import argparse
+import os
+import pickle
+import sys
+import time
+import timeit
+from datetime import datetime
+
+import torch
+import torch.distributed.rpc as rpc
+from tensordict import TensorDict
+from torchrl._utils import logger as torchrl_logger
+from torchrl.data.replay_buffers import RemoteTensorDictReplayBuffer
+from torchrl.data.replay_buffers.samplers import RandomSampler
+from torchrl.data.replay_buffers.storages import (
+    LazyMemmapStorage,
+    LazyTensorStorage,
+    ListStorage,
+)
+from torchrl.data.replay_buffers.writers import RoundRobinWriter
+
+RETRY_LIMIT = 2
+RETRY_DELAY_SECS = 3
+REPLAY_BUFFER_NODE = "ReplayBuffer"
+TRAINER_NODE = "Trainer"
+TENSOR_SIZE = 3 * 86 * 86
+BUFFER_SIZE = 1001
+BATCH_SIZE = 256
+REPEATS = 1000
+
+storage_options = {
+    "LazyMemmapStorage": LazyMemmapStorage,
+    "LazyTensorStorage": LazyTensorStorage,
+    "ListStorage": ListStorage,
+}
+
+storage_arg_options = {
+    "LazyMemmapStorage": {"scratch_dir": "/tmp/", "device": torch.device("cpu")},
+    "LazyTensorStorage": {},
+    "ListStorage": {},
+}
+parser = argparse.ArgumentParser(
+    description="RPC Replay Buffer Example",
+    formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+)
+
+parser.add_argument(
+    "--rank",
+    type=int,
+    default=-1,
+    help="Node Rank [0 = Replay Buffer, 1 = Dummy Trainer, 2+ = Dummy Data Collector]",
+)
+
+parser.add_argument(
+    "--storage",
+    type=str,
+    default="LazyMemmapStorage",
+    help="Storage type [LazyMemmapStorage, LazyTensorStorage, ListStorage]",
+)
+
+
+class DummyTrainerNode:
+    def __init__(self) -> None:
+        self.id = rpc.get_worker_info().id
+        self.replay_buffer = self._create_replay_buffer()
+        self._ret = None
+
+    def train(self, batch_size: int) -> None:
+        start_time = timeit.default_timer()
+        ret = rpc.rpc_sync(
+            self.replay_buffer.owner(),
+            ReplayBufferNode.sample,
+            args=(self.replay_buffer, batch_size),
+        )
+        if storage_type == "ListStorage":
+            self._ret = ret[0]
+        else:
+            if self._ret is None:
+                self._ret = ret
+            else:
+                self._ret.update_(ret)
+        # make sure the content is read
+        self._ret["observation"] + 1
+        self._ret["next_observation"] + 1
+        return timeit.default_timer() - start_time
+
+    def _create_replay_buffer(self) -> rpc.RRef:
+        while True:
+            try:
+                replay_buffer_info = rpc.get_worker_info(REPLAY_BUFFER_NODE)
+                buffer_rref = rpc.remote(
+                    replay_buffer_info, ReplayBufferNode, args=(1000000,)
+                )
+                torchrl_logger.info(f"Connected to replay buffer {replay_buffer_info}")
+                return buffer_rref
+            except Exception:
+                torchrl_logger.info("Failed to connect to replay buffer")
+                time.sleep(RETRY_DELAY_SECS)
+
+
+class ReplayBufferNode(RemoteTensorDictReplayBuffer):
+    def __init__(self, capacity: int):
+        super().__init__(
+            storage=storage_options[storage_type](
+                max_size=capacity, **storage_arg_options[storage_type]
+            ),
+            sampler=RandomSampler(),
+            writer=RoundRobinWriter(),
+            collate_fn=lambda x: x,
+        )
+        tds = TensorDict(
+            {
+                "observation": torch.randn(
+                    BUFFER_SIZE,
+                    TENSOR_SIZE,
+                ),
+                "next_observation": torch.randn(
+                    BUFFER_SIZE,
+                    TENSOR_SIZE,
+                ),
+            },
+            batch_size=[BUFFER_SIZE],
+        )
+        self.extend(tds)
+
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+    rank = args.rank
+    storage_type = args.storage
+
+    torchrl_logger.debug(f"RANK: {rank}; Storage: {storage_type}")
+
+    os.environ["MASTER_ADDR"] = "localhost"
+    os.environ["MASTER_PORT"] = "29500"
+    os.environ["TORCH_DISTRIBUTED_DEBUG"] = "DETAIL"
+    options = rpc.TensorPipeRpcBackendOptions(
+        num_worker_threads=16, init_method="tcp://localhost:10002", rpc_timeout=120
+    )
+    if rank == 0:
+        # rank 0 is the trainer
+        rpc.init_rpc(
+            TRAINER_NODE,
+            rank=rank,
+            backend=rpc.BackendType.TENSORPIPE,
+            rpc_backend_options=options,
+        )
+        trainer = DummyTrainerNode()
+        results = []
+        for i in range(REPEATS):
+            result = trainer.train(batch_size=BATCH_SIZE)
+            if i == 0:
+                continue
+            results.append(result)
+            torchrl_logger.info(f"{i}, {results[-1]}")
+
+        with open(
+            f'./benchmark_{datetime.now().strftime("%d-%m-%Y%H:%M:%S")};batch_size={BATCH_SIZE};tensor_size={TENSOR_SIZE};repeat={REPEATS};storage={storage_type}.pkl',
+            "wb+",
+        ) as f:
+            pickle.dump(results, f)
+
+        tensor_results = torch.tensor(results)
+        torchrl_logger.info(f"Mean: {torch.mean(tensor_results)}")
+        breakpoint()
+    elif rank == 1:
+        # rank 1 is the replay buffer
+        # replay buffer waits passively for construction instructions from trainer node
+        rpc.init_rpc(
+            REPLAY_BUFFER_NODE,
+            rank=rank,
+            backend=rpc.BackendType.TENSORPIPE,
+            rpc_backend_options=options,
+        )
+        breakpoint()
+    else:
+        sys.exit(1)
+    rpc.shutdown()

benchmarks/test_collectors_benchmark.py

@@ -0,0 +1,240 @@
+# 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 time
+
+import pytest
+import torch.cuda
+import tqdm
+
+from torchrl.collectors import (
+    MultiaSyncDataCollector,
+    MultiSyncDataCollector,
+    SyncDataCollector,
+)
+from torchrl.data import LazyTensorStorage, ReplayBuffer
+from torchrl.data.utils import CloudpickleWrapper
+from torchrl.envs import EnvCreator, GymEnv, ParallelEnv, StepCounter, TransformedEnv
+from torchrl.envs.libs.dm_control import DMControlEnv
+from torchrl.modules import RandomPolicy
+
+
+def single_collector_setup():
+    device = "cuda:0" if torch.cuda.device_count() else "cpu"
+    env = TransformedEnv(DMControlEnv("cheetah", "run", device=device), StepCounter(50))
+    c = SyncDataCollector(
+        env,
+        RandomPolicy(env.action_spec),
+        total_frames=-1,
+        frames_per_batch=100,
+        device=device,
+    )
+    c = iter(c)
+    for i, _ in enumerate(c):
+        if i == 10:
+            break
+    return ((c,), {})
+
+
+def sync_collector_setup():
+    device = "cuda:0" if torch.cuda.device_count() else "cpu"
+    env = EnvCreator(
+        lambda: TransformedEnv(
+            DMControlEnv("cheetah", "run", device=device), StepCounter(50)
+        )
+    )
+    c = MultiSyncDataCollector(
+        [env, env],
+        RandomPolicy(env().action_spec),
+        total_frames=-1,
+        frames_per_batch=100,
+        device=device,
+    )
+    c = iter(c)
+    for i, _ in enumerate(c):
+        if i == 10:
+            break
+    return ((c,), {})
+
+
+def async_collector_setup():
+    device = "cuda:0" if torch.cuda.device_count() else "cpu"
+    env = EnvCreator(
+        lambda: TransformedEnv(
+            DMControlEnv("cheetah", "run", device=device), StepCounter(50)
+        )
+    )
+    c = MultiaSyncDataCollector(
+        [env, env],
+        RandomPolicy(env().action_spec),
+        total_frames=-1,
+        frames_per_batch=100,
+        device=device,
+    )
+    c = iter(c)
+    for i, _ in enumerate(c):
+        if i == 10:
+            break
+    return ((c,), {})
+
+
+def single_collector_setup_pixels():
+    device = "cuda:0" if torch.cuda.device_count() else "cpu"
+    # env = TransformedEnv(
+    #     DMControlEnv("cheetah", "run", device=device, from_pixels=True), StepCounter(50)
+    # )
+    env = TransformedEnv(GymEnv("ALE/Pong-v5"), StepCounter(50))
+    c = SyncDataCollector(
+        env,
+        RandomPolicy(env.action_spec),
+        total_frames=-1,
+        frames_per_batch=100,
+        device=device,
+    )
+    c = iter(c)
+    for i, _ in enumerate(c):
+        if i == 10:
+            break
+    return ((c,), {})
+
+
+def sync_collector_setup_pixels():
+    device = "cuda:0" if torch.cuda.device_count() else "cpu"
+    env = EnvCreator(
+        lambda: TransformedEnv(
+            # DMControlEnv("cheetah", "run", device=device, from_pixels=True),
+            GymEnv("ALE/Pong-v5"),
+            StepCounter(50),
+        )
+    )
+    c = MultiSyncDataCollector(
+        [env, env],
+        RandomPolicy(env().action_spec),
+        total_frames=-1,
+        frames_per_batch=100,
+        device=device,
+    )
+    c = iter(c)
+    for i, _ in enumerate(c):
+        if i == 10:
+            break
+    return ((c,), {})
+
+
+def async_collector_setup_pixels():
+    device = "cuda:0" if torch.cuda.device_count() else "cpu"
+    env = EnvCreator(
+        lambda: TransformedEnv(
+            # DMControlEnv("cheetah", "run", device=device, from_pixels=True),
+            GymEnv("ALE/Pong-v5"),
+            StepCounter(50),
+        )
+    )
+    c = MultiaSyncDataCollector(
+        [env, env],
+        RandomPolicy(env().action_spec),
+        total_frames=-1,
+        frames_per_batch=100,
+        device=device,
+    )
+    c = iter(c)
+    for i, _ in enumerate(c):
+        if i == 10:
+            break
+    return ((c,), {})
+
+
+def execute_collector(c):
+    # will run for 9 iterations (1 during setup)
+    next(c)
+
+
+def test_single(benchmark):
+    (c,), _ = single_collector_setup()
+    benchmark(execute_collector, c)
+
+
+def test_sync(benchmark):
+    (c,), _ = sync_collector_setup()
+    benchmark(execute_collector, c)
+
+
+def test_async(benchmark):
+    (c,), _ = async_collector_setup()
+    benchmark(execute_collector, c)
+
+
+@pytest.mark.skipif(not torch.cuda.device_count(), reason="no rendering without cuda")
+def test_single_pixels(benchmark):
+    (c,), _ = single_collector_setup_pixels()
+    benchmark(execute_collector, c)
+
+
+@pytest.mark.skipif(not torch.cuda.device_count(), reason="no rendering without cuda")
+def test_sync_pixels(benchmark):
+    (c,), _ = sync_collector_setup_pixels()
+    benchmark(execute_collector, c)
+
+
+@pytest.mark.skipif(not torch.cuda.device_count(), reason="no rendering without cuda")
+def test_async_pixels(benchmark):
+    (c,), _ = async_collector_setup_pixels()
+    benchmark(execute_collector, c)
+
+
+class TestRBGCollector:
+    @pytest.mark.parametrize(
+        "n_col,n_wokrers_per_col",
+        [
+            [2, 2],
+            [4, 2],
+            [8, 2],
+            [16, 2],
+            [2, 1],
+            [4, 1],
+            [8, 1],
+            [16, 1],
+        ],
+    )
+    def test_multiasync_rb(self, n_col, n_wokrers_per_col):
+        make_env = EnvCreator(lambda: GymEnv("ALE/Pong-v5"))
+        if n_wokrers_per_col > 1:
+            make_env = ParallelEnv(n_wokrers_per_col, make_env)
+            env = make_env
+            policy = RandomPolicy(env.action_spec)
+        else:
+            env = make_env()
+            policy = RandomPolicy(env.action_spec)
+
+        storage = LazyTensorStorage(10_000)
+        rb = ReplayBuffer(storage=storage)
+        rb.extend(env.rollout(2, policy).reshape(-1))
+        rb.append_transform(CloudpickleWrapper(lambda x: x.reshape(-1)), invert=True)
+
+        fpb = n_wokrers_per_col * 100
+        total_frames = n_wokrers_per_col * 100_000
+        c = MultiaSyncDataCollector(
+            [make_env] * n_col,
+            policy,
+            frames_per_batch=fpb,
+            total_frames=total_frames,
+            replay_buffer=rb,
+        )
+        frames = 0
+        pbar = tqdm.tqdm(total=total_frames - (n_col * fpb))
+        for i, _ in enumerate(c):
+            if i == n_col:
+                t0 = time.time()
+            if i >= n_col:
+                frames += fpb
+            if i > n_col:
+                fps = frames / (time.time() - t0)
+                pbar.update(fpb)
+                pbar.set_description(f"fps: {fps: 4.4f}")
+
+
+if __name__ == "__main__":
+    args, unknown = argparse.ArgumentParser().parse_known_args()
+    pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)