torchrl 0.10.1__cp313-cp313-win_amd64.whl → 0.11.0__cp313-cp313-win_amd64.whl

benchmarks/ecosystem/gym_env_throughput.py

@@ -27,7 +27,7 @@ from torchrl.collectors import (
 )
 from torchrl.envs import EnvCreator, GymEnv, ParallelEnv
 from torchrl.envs.libs.gym import gym_backend as gym_bc, set_gym_backend
-from torchrl.envs.utils import RandomPolicy
+from torchrl.modules import RandomPolicy
 
 if __name__ == "__main__":
     avail_devices = ("cpu",)

benchmarks/storage/benchmark_sample_latency_over_rpc.py

@@ -144,7 +144,7 @@ if __name__ == "__main__":
     rank = args.rank
     storage_type = args.storage
 
-    torchrl_logger.info(f"Rank: {rank}; Storage: {storage_type}")
+    torchrl_logger.debug(f"RANK: {rank}; Storage: {storage_type}")
 
     os.environ["MASTER_ADDR"] = "localhost"
     os.environ["MASTER_PORT"] = "29500"

benchmarks/test_collectors_benchmark.py

@@ -18,7 +18,7 @@ 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.envs.utils import RandomPolicy
+from torchrl.modules import RandomPolicy
 
 
 def single_collector_setup():

benchmarks/test_non_tensor_env_benchmark.py

@@ -0,0 +1,70 @@
+# 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.
+from __future__ import annotations
+
+import gc
+import time
+
+import pytest
+from tensordict import set_capture_non_tensor_stack
+from torchrl.envs import ParallelEnv, SerialEnv
+from torchrl.testing.mocking_classes import EnvWithMetadata
+
+
+def _rollout(env, n_steps: int, break_when_any_done: bool) -> None:
+    env.rollout(n_steps, break_when_any_done=break_when_any_done)
+
+
+@pytest.mark.parametrize("break_when_any_done", [True, False])
+@pytest.mark.parametrize(
+    "kind,use_buffers",
+    [
+        pytest.param("single", None, id="single"),
+        pytest.param("serial", False, id="serial-no-buffers"),
+        pytest.param("serial", True, id="serial-buffers"),
+        pytest.param("parallel", False, id="parallel-no-buffers"),
+        pytest.param("parallel", True, id="parallel-buffers"),
+    ],
+)
+@pytest.mark.parametrize("n_steps", [1000])
+def test_non_tensor_env_rollout_speed(
+    benchmark,
+    break_when_any_done: bool,
+    kind: str,
+    use_buffers: bool | None,
+    n_steps: int,
+):
+    """Benchmarks a single rollout, after a warmup rollout, for non-tensor stacking envs.
+
+    Mirrors `test/test_envs.py::TestNonTensorEnv`'s option matrix (single/serial/parallel,
+    break_when_any_done, use_buffers).
+    """
+    with set_capture_non_tensor_stack(False):
+        if kind == "single":
+            env = EnvWithMetadata()
+        elif kind == "serial":
+            env = SerialEnv(2, EnvWithMetadata, use_buffers=use_buffers)
+        elif kind == "parallel":
+            env = ParallelEnv(2, EnvWithMetadata, use_buffers=use_buffers)
+        else:
+            raise RuntimeError(f"Unknown kind={kind}")
+
+        env.set_seed(0)
+        env.reset()
+
+        try:
+            # Warmup run (not timed)
+            _rollout(env, n_steps=n_steps, break_when_any_done=break_when_any_done)
+
+            # Timed run(s)
+            benchmark(
+                _rollout, env, n_steps=n_steps, break_when_any_done=break_when_any_done
+            )
+        finally:
+            env.close(raise_if_closed=False)
+            del env
+            # Give multiprocessing envs a brief chance to terminate cleanly.
+            time.sleep(0.05)
+            gc.collect()

benchmarks/test_objectives_benchmarks.py

@@ -172,6 +172,8 @@ def _maybe_compile(fn, compile, td, fullgraph=FULLGRAPH, warmup=3):
 def test_dqn_speed(
     benchmark, backward, compile, n_obs=8, n_act=4, depth=3, ncells=128, batch=128
 ):
+    if compile == "reduce-overhead" and backward is not None:
+        pytest.skip("reduce-overhead with backward causes segfaults in CI")
     if compile:
         torch._dynamo.reset_code_caches()
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
@@ -227,6 +229,8 @@ def test_dqn_speed(
 def test_ddpg_speed(
     benchmark, backward, compile, n_obs=8, n_act=4, ncells=128, batch=128, n_hidden=64
 ):
+    if compile == "reduce-overhead" and backward is not None:
+        pytest.skip("reduce-overhead with backward causes segfaults in CI")
     if compile:
         torch._dynamo.reset_code_caches()
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
@@ -303,6 +307,8 @@ def test_ddpg_speed(
 def test_sac_speed(
     benchmark, backward, compile, n_obs=8, n_act=4, ncells=128, batch=128, n_hidden=64
 ):
+    if compile == "reduce-overhead" and backward is not None:
+        pytest.skip("reduce-overhead with backward causes segfaults in CI")
     if compile:
         torch._dynamo.reset_code_caches()
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
@@ -393,6 +399,8 @@ def test_sac_speed(
 def test_redq_speed(
     benchmark, backward, compile, n_obs=8, n_act=4, ncells=128, batch=128, n_hidden=64
 ):
+    if compile == "reduce-overhead" and backward is not None:
+        pytest.skip("reduce-overhead with backward causes segfaults in CI")
     if compile:
         torch._dynamo.reset_code_caches()
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
@@ -484,6 +492,8 @@ def test_redq_speed(
 def test_redq_deprec_speed(
     benchmark, backward, compile, n_obs=8, n_act=4, ncells=128, batch=128, n_hidden=64
 ):
+    if compile == "reduce-overhead" and backward is not None:
+        pytest.skip("reduce-overhead with backward causes segfaults in CI")
     if compile:
         torch._dynamo.reset_code_caches()
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
@@ -573,6 +583,8 @@ def test_redq_deprec_speed(
 def test_td3_speed(
     benchmark, backward, compile, n_obs=8, n_act=4, ncells=128, batch=128, n_hidden=64
 ):
+    if compile == "reduce-overhead" and backward is not None:
+        pytest.skip("reduce-overhead with backward causes segfaults in CI")
     if compile:
         torch._dynamo.reset_code_caches()
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
@@ -667,6 +679,8 @@ def test_td3_speed(
 def test_cql_speed(
     benchmark, backward, compile, n_obs=8, n_act=4, ncells=128, batch=128, n_hidden=64
 ):
+    if compile == "reduce-overhead" and backward is not None:
+        pytest.skip("reduce-overhead with backward causes segfaults in CI")
     if compile:
         torch._dynamo.reset_code_caches()
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
@@ -763,6 +777,8 @@ def test_a2c_speed(
     batch=128,
     T=10,
 ):
+    if compile == "reduce-overhead" and backward is not None:
+        pytest.skip("reduce-overhead with backward causes segfaults in CI")
     if compile:
         torch._dynamo.reset_code_caches()
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
@@ -866,6 +882,8 @@ def test_ppo_speed(
     batch=128,
     T=10,
 ):
+    if compile == "reduce-overhead" and backward is not None:
+        pytest.skip("reduce-overhead with backward causes segfaults in CI")
     if compile:
         torch._dynamo.reset_code_caches()
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
@@ -969,6 +987,8 @@ def test_reinforce_speed(
     batch=128,
     T=10,
 ):
+    if compile == "reduce-overhead" and backward is not None:
+        pytest.skip("reduce-overhead with backward causes segfaults in CI")
     if compile:
         torch._dynamo.reset_code_caches()
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
@@ -1072,6 +1092,8 @@ def test_iql_speed(
     batch=128,
     T=10,
 ):
+    if compile == "reduce-overhead" and backward is not None:
+        pytest.skip("reduce-overhead with backward causes segfaults in CI")
     if compile:
         torch._dynamo.reset_code_caches()
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

sota-implementations/a2c/a2c_atari.py

@@ -22,7 +22,7 @@ def main(cfg: DictConfig):  # noqa: F821
     from tensordict import from_module
     from tensordict.nn import CudaGraphModule
 
-    from torchrl._utils import timeit
+    from torchrl._utils import get_available_device, timeit
     from torchrl.collectors import SyncDataCollector
     from torchrl.data import LazyTensorStorage, TensorDictReplayBuffer
     from torchrl.data.replay_buffers.samplers import SamplerWithoutReplacement
@@ -33,11 +33,9 @@ def main(cfg: DictConfig):  # noqa: F821
     from torchrl.record.loggers import generate_exp_name, get_logger
     from utils_atari import eval_model, make_parallel_env, make_ppo_models
 
-    device = cfg.loss.device
-    if not device:
-        device = torch.device("cpu" if not torch.cuda.is_available() else "cuda:0")
-    else:
-        device = torch.device(device)
+    device = (
+        torch.device(cfg.loss.device) if cfg.loss.device else get_available_device()
+    )
 
     # Correct for frame_skip
     frame_skip = 4
@@ -76,8 +74,8 @@ def main(cfg: DictConfig):  # noqa: F821
         actor_network=actor,
         critic_network=critic,
         loss_critic_type=cfg.loss.loss_critic_type,
-        entropy_coef=cfg.loss.entropy_coef,
-        critic_coef=cfg.loss.critic_coef,
+        entropy_coeff=cfg.loss.entropy_coeff,
+        critic_coeff=cfg.loss.critic_coeff,
     )
 
     # use end-of-life as done key

sota-implementations/a2c/a2c_mujoco.py

@@ -23,7 +23,7 @@ def main(cfg: DictConfig):  # noqa: F821
     from tensordict import from_module
     from tensordict.nn import CudaGraphModule
 
-    from torchrl._utils import timeit
+    from torchrl._utils import get_available_device, timeit
     from torchrl.collectors import SyncDataCollector
     from torchrl.data import LazyTensorStorage, TensorDictReplayBuffer
     from torchrl.data.replay_buffers.samplers import SamplerWithoutReplacement
@@ -36,11 +36,9 @@ def main(cfg: DictConfig):  # noqa: F821
 
     # Define paper hyperparameters
 
-    device = cfg.loss.device
-    if not device:
-        device = torch.device("cpu" if not torch.cuda.is_available() else "cuda:0")
-    else:
-        device = torch.device(device)
+    device = (
+        torch.device(cfg.loss.device) if cfg.loss.device else get_available_device()
+    )
 
     num_mini_batches = cfg.collector.frames_per_batch // cfg.loss.mini_batch_size
     total_network_updates = (
@@ -77,8 +75,8 @@ def main(cfg: DictConfig):  # noqa: F821
         actor_network=actor,
         critic_network=critic,
         loss_critic_type=cfg.loss.loss_critic_type,
-        entropy_coef=cfg.loss.entropy_coef,
-        critic_coef=cfg.loss.critic_coef,
+        entropy_coeff=cfg.loss.entropy_coeff,
+        critic_coeff=cfg.loss.critic_coeff,
     )
 
     # Create optimizers

sota-implementations/cql/cql_offline.py

@@ -18,7 +18,7 @@ import numpy as np
 import torch
 import tqdm
 from tensordict.nn import CudaGraphModule
-from torchrl._utils import timeit
+from torchrl._utils import get_available_device, timeit
 from torchrl.envs.utils import ExplorationType, set_exploration_type
 from torchrl.objectives import group_optimizers
 from torchrl.record.loggers import generate_exp_name, get_logger
@@ -55,13 +55,9 @@ def main(cfg: DictConfig):  # noqa: F821
     # Set seeds
     torch.manual_seed(cfg.env.seed)
     np.random.seed(cfg.env.seed)
-    device = cfg.optim.device
-    if device in ("", None):
-        if torch.cuda.is_available():
-            device = "cuda:0"
-        else:
-            device = "cpu"
-    device = torch.device(device)
+    device = (
+        torch.device(cfg.optim.device) if cfg.optim.device else get_available_device()
+    )
 
     # Create replay buffer
     replay_buffer = make_offline_replay_buffer(cfg.replay_buffer)

sota-implementations/cql/cql_online.py

@@ -21,7 +21,7 @@ import torch
 import tqdm
 from tensordict import TensorDict
 from tensordict.nn import CudaGraphModule
-from torchrl._utils import timeit
+from torchrl._utils import get_available_device, timeit
 from torchrl.envs.utils import ExplorationType, set_exploration_type
 from torchrl.objectives import group_optimizers
 from torchrl.record.loggers import generate_exp_name, get_logger
@@ -60,13 +60,9 @@ def main(cfg: DictConfig):  # noqa: F821
     # Set seeds
     torch.manual_seed(cfg.env.seed)
     np.random.seed(cfg.env.seed)
-    device = cfg.optim.device
-    if device in ("", None):
-        if torch.cuda.is_available():
-            device = "cuda:0"
-        else:
-            device = "cpu"
-    device = torch.device(device)
+    device = (
+        torch.device(cfg.optim.device) if cfg.optim.device else get_available_device()
+    )
 
     # Create env
     train_env, eval_env = make_environment(

sota-implementations/cql/discrete_cql_offline.py

@@ -18,7 +18,7 @@ import numpy as np
 import torch
 import tqdm
 from tensordict.nn import CudaGraphModule
-from torchrl._utils import timeit
+from torchrl._utils import get_available_device, timeit
 from torchrl.envs.utils import ExplorationType, set_exploration_type
 from torchrl.record.loggers import generate_exp_name, get_logger
 from utils import (
@@ -36,10 +36,9 @@ torch.set_float32_matmul_precision("high")
 
 @hydra.main(version_base="1.1", config_path="", config_name="discrete_offline_config")
 def main(cfg):  # noqa: F821
-    device = cfg.optim.device
-    if device in ("", None):
-        device = "cuda:0" if torch.cuda.is_available() else "cpu"
-    device = torch.device(device)
+    device = (
+        torch.device(cfg.optim.device) if cfg.optim.device else get_available_device()
+    )
 
     # Create logger
     exp_name = generate_exp_name("DiscreteCQL", cfg.logger.exp_name)

sota-implementations/cql/discrete_cql_online.py

@@ -20,7 +20,7 @@ import torch
 import torch.cuda
 import tqdm
 from tensordict.nn import CudaGraphModule
-from torchrl._utils import timeit
+from torchrl._utils import get_available_device, timeit
 from torchrl.envs.utils import ExplorationType, set_exploration_type
 from torchrl.record.loggers import generate_exp_name, get_logger
 from utils import (
@@ -38,13 +38,9 @@ torch.set_float32_matmul_precision("high")
 
 @hydra.main(version_base="1.1", config_path="", config_name="discrete_online_config")
 def main(cfg: DictConfig):  # noqa: F821
-    device = cfg.optim.device
-    if device in ("", None):
-        if torch.cuda.is_available():
-            device = "cuda:0"
-        else:
-            device = "cpu"
-    device = torch.device(device)
+    device = (
+        torch.device(cfg.optim.device) if cfg.optim.device else get_available_device()
+    )
 
     # Create logger
     exp_name = generate_exp_name("DiscreteCQL", cfg.logger.exp_name)

sota-implementations/crossq/crossq.py

@@ -21,7 +21,7 @@ import torch.cuda
 import tqdm
 from tensordict import TensorDict
 from tensordict.nn import CudaGraphModule
-from torchrl._utils import timeit
+from torchrl._utils import get_available_device, timeit
 from torchrl.envs.utils import ExplorationType, set_exploration_type
 from torchrl.objectives import group_optimizers
 from torchrl.record.loggers import generate_exp_name, get_logger
@@ -40,13 +40,11 @@ torch.set_float32_matmul_precision("high")
 
 @hydra.main(version_base="1.1", config_path=".", config_name="config")
 def main(cfg: DictConfig):  # noqa: F821
-    device = cfg.network.device
-    if device in ("", None):
-        if torch.cuda.is_available():
-            device = torch.device("cuda:0")
-        else:
-            device = torch.device("cpu")
-    device = torch.device(device)
+    device = (
+        torch.device(cfg.network.device)
+        if cfg.network.device
+        else get_available_device()
+    )
 
     # Create logger
     exp_name = generate_exp_name("CrossQ", cfg.logger.exp_name)

sota-implementations/ddpg/ddpg.py

@@ -21,7 +21,7 @@ import torch.cuda
 import tqdm
 from tensordict import TensorDict
 from tensordict.nn import CudaGraphModule
-from torchrl._utils import timeit
+from torchrl._utils import get_available_device, timeit
 from torchrl.envs.utils import ExplorationType, set_exploration_type
 from torchrl.objectives import group_optimizers
 from torchrl.record.loggers import generate_exp_name, get_logger
@@ -39,21 +39,14 @@ from utils import (
 
 @hydra.main(version_base="1.1", config_path="", config_name="config")
 def main(cfg: DictConfig):  # noqa: F821
-    device = cfg.optim.device
-    if device in ("", None):
-        if torch.cuda.is_available():
-            device = "cuda:0"
-        else:
-            device = "cpu"
-    device = torch.device(device)
-
-    collector_device = cfg.collector.device
-    if collector_device in ("", None):
-        if torch.cuda.is_available():
-            collector_device = "cuda:0"
-        else:
-            collector_device = "cpu"
-    collector_device = torch.device(collector_device)
+    device = (
+        torch.device(cfg.optim.device) if cfg.optim.device else get_available_device()
+    )
+    collector_device = (
+        torch.device(cfg.collector.device)
+        if cfg.collector.device
+        else get_available_device()
+    )
 
     # Create logger
     exp_name = generate_exp_name("DDPG", cfg.logger.exp_name)

sota-implementations/decision_transformer/dt.py

@@ -17,7 +17,7 @@ import torch
 import tqdm
 from tensordict import TensorDict
 from tensordict.nn import CudaGraphModule
-from torchrl._utils import logger as torchrl_logger, timeit
+from torchrl._utils import get_available_device, logger as torchrl_logger, timeit
 from torchrl.envs.libs.gym import set_gym_backend
 from torchrl.envs.utils import ExplorationType, set_exploration_type
 from torchrl.modules.tensordict_module import DecisionTransformerInferenceWrapper
@@ -38,13 +38,9 @@ from utils import (
 def main(cfg: DictConfig):  # noqa: F821
     set_gym_backend(cfg.env.backend).set()
 
-    model_device = cfg.optim.device
-    if model_device in ("", None):
-        if torch.cuda.is_available():
-            model_device = "cuda:0"
-        else:
-            model_device = "cpu"
-    model_device = torch.device(model_device)
+    model_device = (
+        torch.device(cfg.optim.device) if cfg.optim.device else get_available_device()
+    )
 
     # Set seeds
     torch.manual_seed(cfg.env.seed)

sota-implementations/decision_transformer/online_dt.py

@@ -15,7 +15,7 @@ import numpy as np
 import torch
 import tqdm
 from tensordict.nn import CudaGraphModule
-from torchrl._utils import logger as torchrl_logger, timeit
+from torchrl._utils import get_available_device, logger as torchrl_logger, timeit
 from torchrl.envs.libs.gym import set_gym_backend
 from torchrl.envs.utils import ExplorationType, set_exploration_type
 from torchrl.modules.tensordict_module import DecisionTransformerInferenceWrapper
@@ -36,13 +36,9 @@ from utils import (
 def main(cfg: DictConfig):  # noqa: F821
     set_gym_backend(cfg.env.backend).set()
 
-    model_device = cfg.optim.device
-    if model_device in ("", None):
-        if torch.cuda.is_available():
-            model_device = "cuda:0"
-        else:
-            model_device = "cpu"
-    model_device = torch.device(model_device)
+    model_device = (
+        torch.device(cfg.optim.device) if cfg.optim.device else get_available_device()
+    )
 
     # Set seeds
     torch.manual_seed(cfg.env.seed)

sota-implementations/dqn/dqn_atari.py

@@ -17,7 +17,7 @@ import torch.nn
 import torch.optim
 import tqdm
 from tensordict.nn import CudaGraphModule, TensorDictSequential
-from torchrl._utils import timeit
+from torchrl._utils import get_available_device, timeit
 from torchrl.collectors import SyncDataCollector
 from torchrl.data import LazyMemmapStorage, TensorDictReplayBuffer
 from torchrl.envs import ExplorationType, set_exploration_type
@@ -33,13 +33,7 @@ torch.set_float32_matmul_precision("high")
 @hydra.main(config_path="", config_name="config_atari", version_base="1.1")
 def main(cfg: DictConfig):  # noqa: F821
 
-    device = cfg.device
-    if device in ("", None):
-        if torch.cuda.is_available():
-            device = "cuda:0"
-        else:
-            device = "cpu"
-    device = torch.device(device)
+    device = torch.device(cfg.device) if cfg.device else get_available_device()
 
     # Correct for frame_skip
     frame_skip = 4

sota-implementations/dqn/dqn_cartpole.py

@@ -12,7 +12,7 @@ import torch.nn
 import torch.optim
 import tqdm
 from tensordict.nn import CudaGraphModule, TensorDictSequential
-from torchrl._utils import timeit
+from torchrl._utils import get_available_device, timeit
 from torchrl.collectors import SyncDataCollector
 from torchrl.data import LazyTensorStorage, TensorDictReplayBuffer
 from torchrl.envs import ExplorationType, set_exploration_type
@@ -28,13 +28,7 @@ torch.set_float32_matmul_precision("high")
 @hydra.main(config_path="", config_name="config_cartpole", version_base="1.1")
 def main(cfg: DictConfig):  # noqa: F821
 
-    device = cfg.device
-    if device in ("", None):
-        if torch.cuda.is_available():
-            device = "cuda:0"
-        else:
-            device = "cpu"
-    device = torch.device(device)
+    device = torch.device(cfg.device) if cfg.device else get_available_device()
 
     # Make the components
     model = make_dqn_model(cfg.env.env_name, device=device)