torchrl 0.11.0__cp314-cp314t-manylinux_2_28_aarch64.whl

benchmarks/benchmark_batched_envs.py

@@ -0,0 +1,104 @@
+# 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.
+"""
+Benchmarking different types of batched environments
+====================================================
+Compares runtime for different environments which allow performing operations in a batch.
+- SerialEnv executes the operations sequentially
+- ParallelEnv uses multiprocess parallelism
+- MultiThreadedEnv uses multithreaded parallelism and is based on envpool library.
+
+Run as "python benchmarks/benchmark_batched_envs.py"
+Requires pandas ("pip install pandas").
+
+"""
+
+import pandas as pd
+from torchrl._utils import logger as torchrl_logger
+
+pd.set_option("display.max_columns", 100)
+pd.set_option("display.width", 1000)
+import torch
+from torch.utils.benchmark import Timer
+from torchrl.envs import MultiThreadedEnv, ParallelEnv, SerialEnv
+from torchrl.envs.libs.gym import GymEnv
+
+N_STEPS = 1000
+
+
+def create_multithreaded(num_workers, device):
+    env = MultiThreadedEnv(num_workers=num_workers, env_name="Pendulum-v1")
+    # GPU doesn't lead to any speedup for MultiThreadedEnv, as the underlying library (envpool) works only on CPU
+    env = env.to(device=torch.device(device))
+    env.rollout(policy=None, max_steps=5)  # Warm-up
+    return env
+
+
+def factory():
+    return GymEnv("Pendulum-v1")
+
+
+def create_serial(num_workers, device):
+    env = SerialEnv(num_workers=num_workers, create_env_fn=factory)
+    env = env.to(device=torch.device(device))
+    env.rollout(policy=None, max_steps=5)  # Warm-up
+    return env
+
+
+def create_parallel(num_workers, device):
+    env = ParallelEnv(num_workers=num_workers, create_env_fn=factory)
+    env = env.to(device=torch.device(device))
+    env.rollout(policy=None, max_steps=5)  # Warm-up
+    return env
+
+
+def run_env(env):
+    env.rollout(policy=None, max_steps=N_STEPS)
+
+
+if __name__ == "__main__":
+    res = {}
+    devices = ["cpu"]
+    if torch.cuda.is_available():
+        devices.append("cuda")
+    for device in devices:
+        for num_workers in [1, 4, 16]:
+            torchrl_logger.info(f"With num_workers={num_workers}, {device}")
+            torchrl_logger.info("Multithreaded...")
+            env_multithreaded = create_multithreaded(num_workers, device)
+            res_multithreaded = Timer(
+                stmt="run_env(env)",
+                setup="from __main__ import run_env",
+                globals={"env": env_multithreaded},
+            )
+            time_multithreaded = res_multithreaded.blocked_autorange().mean
+
+            torchrl_logger.info("Serial...")
+            env_serial = create_serial(num_workers, device)
+            res_serial = Timer(
+                stmt="run_env(env)",
+                setup="from __main__ import run_env",
+                globals={"env": env_serial},
+            )
+            time_serial = res_serial.blocked_autorange().mean
+
+            torchrl_logger.info("Parallel...")
+            env_parallel = create_parallel(num_workers, device)
+            res_parallel = Timer(
+                stmt="run_env(env)",
+                setup="from __main__ import run_env",
+                globals={"env": env_parallel},
+            )
+            time_parallel = res_parallel.blocked_autorange().mean
+
+            res[f"num_workers_{num_workers}_{device}"] = {
+                "Serial, s": time_serial,
+                "Parallel, s": time_parallel,
+                "Multithreaded, s": time_multithreaded,
+            }
+    df = pd.DataFrame(res).round(3)
+    gain = 1 - df.loc["Multithreaded, s"] / df.loc["Parallel, s"]
+    df.loc["Gain, %", :] = (gain * 100).round(1)
+    df.to_csv("multithreaded_benchmark.csv")

benchmarks/conftest.py

@@ -0,0 +1,91 @@
+# 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 time
+import warnings
+from collections import defaultdict
+
+import pytest
+from torchrl._utils import logger as torchrl_logger
+
+CALL_TIMES = defaultdict(float)
+
+
+def pytest_sessionfinish(maxprint=50):
+    out_str = """
+Call times:
+===========
+"""
+    keys = list(CALL_TIMES.keys())
+    if len(keys) > 1:
+        maxchar = max(*[len(key) for key in keys])
+    elif len(keys) == 1:
+        maxchar = len(keys[0])
+    else:
+        return
+    for i, (key, item) in enumerate(
+        sorted(CALL_TIMES.items(), key=lambda x: x[1], reverse=True)
+    ):
+        spaces = "  " + " " * (maxchar - len(key))
+        out_str += f"\t{key}{spaces}{item: 4.4f}s\n"
+        if i == maxprint - 1:
+            break
+    torchrl_logger.info(out_str)
+
+
+@pytest.fixture(autouse=True)
+def measure_duration(request: pytest.FixtureRequest):
+    start_time = time.time()
+
+    def fin():
+        duration = time.time() - start_time
+        name = request.node.name
+        class_name = request.cls.__name__ if request.cls else None
+        name = name.split("[")[0]
+        if class_name is not None:
+            name = "::".join([class_name, name])
+        file = os.path.basename(request.path)
+        name = f"{file}::{name}"
+        CALL_TIMES[name] = CALL_TIMES[name] + duration
+
+    request.addfinalizer(fin)
+
+
+def pytest_addoption(parser):
+    parser.addoption("--rank", action="store")
+
+
+@pytest.fixture(scope="session", autouse=True)
+def set_warnings() -> None:
+    warnings.filterwarnings(
+        "ignore",
+        category=UserWarning,
+        message=r"Lazy modules are a new feature under heavy development",
+    )
+    warnings.filterwarnings(
+        "ignore",
+        category=UserWarning,
+        message=r"Couldn't cast the policy onto the desired device on remote process",
+    )
+    warnings.filterwarnings(
+        "ignore",
+        category=DeprecationWarning,
+        message=r"Deprecated call to `pkg_resources.declare_namespace",
+    )
+    warnings.filterwarnings(
+        "ignore",
+        category=DeprecationWarning,
+        message=r"Using or importing the ABCs",
+    )
+    warnings.filterwarnings(
+        "ignore",
+        category=DeprecationWarning,
+        message=r"Please use `coo_matrix` from the `scipy.sparse` namespace",
+    )
+    warnings.filterwarnings(
+        "ignore",
+        category=DeprecationWarning,
+        message=r"jax.tree_util.register_keypaths is deprecated|jax.ShapedArray is deprecated",
+    )

benchmarks/ecosystem/gym_env_throughput.py

@@ -0,0 +1,321 @@
+# 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.
+
+"""This script executes some envs across the Gym library with the explicit scope of testing the throughput using the various TorchRL components.
+
+We test:
+- gym async envs embedded in a TorchRL's GymEnv wrapper,
+- ParallelEnv with regular GymEnv instances,
+- Data collector
+- Multiprocessed data collectors with parallel envs.
+
+The tests are executed with various number of cpus, and on different devices.
+
+"""
+import time
+
+# import myosuite  # noqa: F401
+import torch
+import tqdm
+from torchrl._utils import timeit
+from torchrl.collectors import (
+    MultiaSyncDataCollector,
+    MultiSyncDataCollector,
+    SyncDataCollector,
+)
+from torchrl.envs import EnvCreator, GymEnv, ParallelEnv
+from torchrl.envs.libs.gym import gym_backend as gym_bc, set_gym_backend
+from torchrl.modules import RandomPolicy
+
+if __name__ == "__main__":
+    avail_devices = ("cpu",)
+    if torch.cuda.is_available():
+        avail_devices = avail_devices + ("cuda:0",)
+
+    for envname in [
+        "CartPole-v1",
+        "HalfCheetah-v4",
+        "myoHandReachRandom-v0",
+        "ALE/Breakout-v5",
+    ]:
+        # the number of collectors won't affect the resources, just impacts how the envs are split in sub-sub-processes
+        for num_workers, num_collectors in zip((32, 64, 8, 16), (8, 8, 2, 4)):
+            with open(f"{envname}_{num_workers}.txt".replace("/", "-"), "w+") as log:
+                if "myo" in envname:
+                    gym_backend = "gym"
+                else:
+                    gym_backend = "gymnasium"
+
+                total_frames = num_workers * 10_000
+
+                # pure gym
+                def make(envname=envname, gym_backend=gym_backend):
+                    with set_gym_backend(gym_backend):
+                        return gym_bc().make(envname)
+
+                with set_gym_backend(gym_backend):
+                    env = gym_bc().vector.AsyncVectorEnv(
+                        [make for _ in range(num_workers)]
+                    )
+                env.reset()
+                global_step = 0
+                times = []
+                start = time.time()
+                for _ in tqdm.tqdm(range(total_frames // num_workers)):
+                    env.step(env.action_space.sample())
+                    global_step += num_workers
+                env.close()
+                log.write(
+                    f"pure gym: {num_workers * 10_000 / (time.time() - start): 4.4f} fps\n"
+                )
+                log.flush()
+
+                # regular parallel env
+                for device in avail_devices:
+
+                    def make(envname=envname, gym_backend=gym_backend):
+                        with set_gym_backend(gym_backend):
+                            return GymEnv(envname, device="cpu")
+
+                    # env_make = EnvCreator(make)
+                    penv = ParallelEnv(num_workers, EnvCreator(make), device=device)
+                    with torch.inference_mode():
+                        # warmup
+                        penv.rollout(2)
+                        pbar = tqdm.tqdm(total=num_workers * 10_000)
+                        t0 = time.time()
+                        data = None
+                        for _ in range(100):
+                            data = penv.rollout(
+                                100, break_when_any_done=False, out=data
+                            )
+                            pbar.update(100 * num_workers)
+                    log.write(
+                        f"penv {device}: {num_workers * 10_000 / (time.time() - t0): 4.4f} fps\n"
+                    )
+                    log.flush()
+                    penv.close()
+                    timeit.print()
+                    del penv
+
+                for device in avail_devices:
+
+                    def make(envname=envname, gym_backend=gym_backend):
+                        with set_gym_backend(gym_backend):
+                            return GymEnv(envname, device="cpu")
+
+                    env_make = EnvCreator(make)
+                    # penv = SerialEnv(num_workers, env_make)
+                    penv = ParallelEnv(num_workers, env_make, device=device)
+                    collector = SyncDataCollector(
+                        penv,
+                        RandomPolicy(penv.action_spec),
+                        frames_per_batch=1024,
+                        total_frames=num_workers * 10_000,
+                        device=device,
+                    )
+                    pbar = tqdm.tqdm(total=num_workers * 10_000)
+                    total_frames = 0
+                    t0 = time.time()
+                    for data in collector:
+                        total_frames += data.numel()
+                        pbar.update(data.numel())
+                        pbar.set_description(
+                            f"single collector + torchrl penv: {total_frames / (time.time() - t0): 4.4f} fps"
+                        )
+                    log.write(
+                        f"single collector + torchrl penv {device}: {total_frames / (time.time() - t0): 4.4f} fps\n"
+                    )
+                    log.flush()
+                    collector.shutdown()
+                    del collector
+
+                for device in avail_devices:
+                    # gym parallel env
+                    def make_env(
+                        envname=envname,
+                        num_workers=num_workers,
+                        gym_backend=gym_backend,
+                        device=device,
+                    ):
+                        with set_gym_backend(gym_backend):
+                            penv = GymEnv(envname, num_envs=num_workers, device=device)
+                        return penv
+
+                    penv = make_env()
+                    # warmup
+                    penv.rollout(2)
+                    pbar = tqdm.tqdm(total=num_workers * 10_000)
+                    t0 = time.time()
+                    for _ in range(100):
+                        data = penv.rollout(100, break_when_any_done=False)
+                        pbar.update(100 * num_workers)
+                    log.write(
+                        f"gym penv {device}: {num_workers * 10_000 / (time.time() - t0): 4.4f} fps\n"
+                    )
+                    log.flush()
+                    penv.close()
+                    del penv
+
+                for device in avail_devices:
+                    # async collector
+                    # + torchrl parallel env
+                    def make_env(envname=envname, gym_backend=gym_backend):
+                        with set_gym_backend(gym_backend):
+                            return GymEnv(envname, device="cpu")
+
+                    penv = ParallelEnv(
+                        num_workers // num_collectors,
+                        EnvCreator(make_env),
+                        device=device,
+                    )
+                    collector = MultiaSyncDataCollector(
+                        [penv] * num_collectors,
+                        policy=RandomPolicy(penv.action_spec),
+                        frames_per_batch=1024,
+                        total_frames=num_workers * 10_000,
+                        device=device,
+                    )
+                    pbar = tqdm.tqdm(total=num_workers * 10_000)
+                    total_frames = 0
+                    for i, data in enumerate(collector):
+                        if i == num_collectors:
+                            t0 = time.time()
+                        if i >= num_collectors:
+                            total_frames += data.numel()
+                            pbar.update(data.numel())
+                            pbar.set_description(
+                                f"collector + torchrl penv: {total_frames / (time.time() - t0): 4.4f} fps"
+                            )
+                    log.write(
+                        f"async collector + torchrl penv {device}: {total_frames / (time.time() - t0): 4.4f} fps\n"
+                    )
+                    log.flush()
+                    collector.shutdown()
+                    del collector
+
+                for device in avail_devices:
+                    # async collector
+                    # + gym async env
+                    def make_env(
+                        envname=envname,
+                        num_workers=num_workers,
+                        gym_backend=gym_backend,
+                    ):
+                        with set_gym_backend(gym_backend):
+                            penv = GymEnv(envname, num_envs=num_workers, device="cpu")
+                        return penv
+
+                    penv = EnvCreator(
+                        lambda num_workers=num_workers // num_collectors: make_env(
+                            num_workers=num_workers
+                        )
+                    )
+                    collector = MultiaSyncDataCollector(
+                        [penv] * num_collectors,
+                        policy=RandomPolicy(penv().action_spec),
+                        frames_per_batch=1024,
+                        total_frames=num_workers * 10_000,
+                        num_sub_threads=num_workers // num_collectors,
+                        device=device,
+                    )
+                    pbar = tqdm.tqdm(total=num_workers * 10_000)
+                    total_frames = 0
+                    for i, data in enumerate(collector):
+                        if i == num_collectors:
+                            t0 = time.time()
+                        if i >= num_collectors:
+                            total_frames += data.numel()
+                            pbar.update(data.numel())
+                            pbar.set_description(
+                                f"{i} collector + gym penv: {total_frames / (time.time() - t0): 4.4f} fps"
+                            )
+                    log.write(
+                        f"async collector + gym penv {device}: {total_frames / (time.time() - t0): 4.4f} fps\n"
+                    )
+                    log.flush()
+                    collector.shutdown()
+                    del collector
+
+                for device in avail_devices:
+                    # sync collector
+                    # + torchrl parallel env
+                    def make_env(envname=envname, gym_backend=gym_backend):
+                        with set_gym_backend(gym_backend):
+                            return GymEnv(envname, device="cpu")
+
+                    penv = ParallelEnv(
+                        num_workers // num_collectors,
+                        EnvCreator(make_env),
+                        device=device,
+                    )
+                    collector = MultiSyncDataCollector(
+                        [penv] * num_collectors,
+                        policy=RandomPolicy(penv.action_spec),
+                        frames_per_batch=1024,
+                        total_frames=num_workers * 10_000,
+                        device=device,
+                    )
+                    pbar = tqdm.tqdm(total=num_workers * 10_000)
+                    total_frames = 0
+                    for i, data in enumerate(collector):
+                        if i == num_collectors:
+                            t0 = time.time()
+                        if i >= num_collectors:
+                            total_frames += data.numel()
+                            pbar.update(data.numel())
+                            pbar.set_description(
+                                f"collector + torchrl penv: {total_frames / (time.time() - t0): 4.4f} fps"
+                            )
+                    log.write(
+                        f"sync collector + torchrl penv {device}: {total_frames / (time.time() - t0): 4.4f} fps\n"
+                    )
+                    log.flush()
+                    collector.shutdown()
+                    del collector
+
+                for device in avail_devices:
+                    # sync collector
+                    # + gym async env
+                    def make_env(
+                        envname=envname,
+                        num_workers=num_workers,
+                        gym_backend=gym_backend,
+                    ):
+                        with set_gym_backend(gym_backend):
+                            penv = GymEnv(envname, num_envs=num_workers, device="cpu")
+                        return penv
+
+                    penv = EnvCreator(
+                        lambda num_workers=num_workers // num_collectors: make_env(
+                            num_workers=num_workers
+                        )
+                    )
+                    collector = MultiSyncDataCollector(
+                        [penv] * num_collectors,
+                        policy=RandomPolicy(penv().action_spec),
+                        frames_per_batch=1024,
+                        total_frames=num_workers * 10_000,
+                        num_sub_threads=num_workers // num_collectors,
+                        device=device,
+                    )
+                    pbar = tqdm.tqdm(total=num_workers * 10_000)
+                    total_frames = 0
+                    for i, data in enumerate(collector):
+                        if i == num_collectors:
+                            t0 = time.time()
+                        if i >= num_collectors:
+                            total_frames += data.numel()
+                            pbar.update(data.numel())
+                            pbar.set_description(
+                                f"{i} collector + gym penv: {total_frames / (time.time() - t0): 4.4f} fps"
+                            )
+                    log.write(
+                        f"sync collector + gym penv {device}: {total_frames / (time.time() - t0): 4.4f} fps\n"
+                    )
+                    log.flush()
+                    collector.shutdown()
+                    del collector
+    exit()