pytorch · vmoens · Jan 31, 2024 · Jan 19, 2024 · Jan 31, 2024 · Jan 31, 2024
diff --git a/.github/unittest/linux_examples/scripts/run_test.sh b/.github/unittest/linux_examples/scripts/run_test.sh
@@ -114,6 +114,7 @@ python .github/unittest/helpers/coverage_run_parallel.py examples/dqn/dqn_atari.
   buffer.batch_size=10 \
   device=cuda:0 \
   loss.num_updates=1 \
+  logger.backend= \
   buffer.buffer_size=120
 python .github/unittest/helpers/coverage_run_parallel.py examples/cql/discrete_cql_online.py \
   collector.total_frames=48 \
@@ -256,6 +257,7 @@ python .github/unittest/helpers/coverage_run_parallel.py examples/dqn/dqn_atari.
   buffer.batch_size=10 \
   device=cuda:0 \
   loss.num_updates=1 \
+  logger.backend= \
   buffer.buffer_size=120
 python .github/unittest/helpers/coverage_run_parallel.py examples/redq/redq.py \
   num_workers=2 \

diff --git a/examples/bandits/dqn.py b/examples/bandits/dqn.py
@@ -7,11 +7,12 @@
 
 import torch
 import tqdm
-from torch import nn
 
+from tensordict.nn import TensorDictSequential
+from torch import nn
 from torchrl.envs.libs.openml import OpenMLEnv
 from torchrl.envs.utils import ExplorationType, set_exploration_type
-from torchrl.modules import DistributionalQValueActor, EGreedyWrapper, MLP, QValueActor
+from torchrl.modules import DistributionalQValueActor, EGreedyModule, MLP, QValueActor
 from torchrl.objectives import DistributionalDQNLoss, DQNLoss
 
 parser = argparse.ArgumentParser()
@@ -85,12 +86,14 @@
         actor(env.reset())
         loss = DQNLoss(actor, loss_function="smooth_l1", action_space=env.action_spec)
         loss.make_value_estimator(gamma=0.0)
-    policy = EGreedyWrapper(
+    policy = TensorDictSequential(
         actor,
-        eps_init=eps_greedy,
-        eps_end=0.0,
-        annealing_num_steps=n_steps,
-        spec=env.action_spec,
+        EGreedyModule(
+            eps_init=eps_greedy,
+            eps_end=0.0,
+            annealing_num_steps=n_steps,
+            spec=env.action_spec,
+        ),
     )
     optim = torch.optim.Adam(loss.parameters(), lr, weight_decay=wd)
 

diff --git a/examples/cql/discrete_cql_config.yaml b/examples/cql/discrete_cql_config.yaml
@@ -2,7 +2,7 @@
 env:
   name: CartPole-v1
   task: ""
-  backend: gym
+  backend: gymnasium
   n_samples_stats: 1000
   max_episode_steps: 200
   seed: 0
@@ -36,7 +36,7 @@ replay_buffer:
   prb: 0
   buffer_prefetch: 64
   size: 1_000_000
-  scratch_dir: ${env.exp_name}_${env.seed}
+  scratch_dir: null
 
 # Optimization
 optim:

diff --git a/examples/cql/discrete_cql_online.py b/examples/cql/discrete_cql_online.py
@@ -73,7 +73,7 @@ def main(cfg: "DictConfig"):  # noqa: F821
         batch_size=cfg.optim.batch_size,
         prb=cfg.replay_buffer.prb,
         buffer_size=cfg.replay_buffer.size,
-        buffer_scratch_dir=cfg.replay_buffer.scratch_dir,
+        scratch_dir=cfg.replay_buffer.scratch_dir,
         device="cpu",
     )
 

diff --git a/examples/cql/offline_config.yaml b/examples/cql/offline_config.yaml
@@ -5,7 +5,7 @@ env:
   library: gym
   n_samples_stats: 1000
   seed: 0
-  backend: gym  # D4RL uses gym so we make sure gymnasium is hidden
+  backend: gymnasium
 
 # torchrl_logger
 logger:

diff --git a/examples/cql/online_config.yaml b/examples/cql/online_config.yaml
@@ -6,7 +6,7 @@ env:
   seed: 0
   train_num_envs: 1
   eval_num_envs: 1
-  backend: gym
+  backend: gymnasium
 
 # Collector
 collector:

diff --git a/examples/cql/utils.py b/examples/cql/utils.py
@@ -121,7 +121,7 @@ def make_replay_buffer(
     batch_size,
     prb=False,
     buffer_size=1000000,
-    buffer_scratch_dir=None,
+    scratch_dir=None,
     device="cpu",
     prefetch=3,
 ):
@@ -133,7 +133,7 @@ def make_replay_buffer(
             prefetch=prefetch,
             storage=LazyMemmapStorage(
                 buffer_size,
-                scratch_dir=buffer_scratch_dir,
+                scratch_dir=scratch_dir,
                 device=device,
             ),
             batch_size=batch_size,
@@ -144,7 +144,7 @@ def make_replay_buffer(
             prefetch=prefetch,
             storage=LazyMemmapStorage(
                 buffer_size,
-                scratch_dir=buffer_scratch_dir,
+                scratch_dir=scratch_dir,
                 device=device,
             ),
             batch_size=batch_size,
@@ -320,7 +320,6 @@ def make_discrete_loss(loss_cfg, model):
         model,
         loss_function=loss_cfg.loss_function,
         delay_value=True,
-        gamma=loss_cfg.gamma,
     )
     loss_module.make_value_estimator(gamma=loss_cfg.gamma)
     target_net_updater = SoftUpdate(loss_module, tau=loss_cfg.tau)

diff --git a/examples/ddpg/config.yaml b/examples/ddpg/config.yaml
@@ -21,7 +21,7 @@ collector:
 replay_buffer:
   size: 1000000
   prb: 0 # use prioritized experience replay
-  scratch_dir: ${torchrl_logger.exp_name}_${env.seed}
+  scratch_dir: null
 
 # optimization
 optim:

diff --git a/examples/ddpg/ddpg.py b/examples/ddpg/ddpg.py
@@ -74,7 +74,7 @@ def main(cfg: "DictConfig"):  # noqa: F821
         batch_size=cfg.optim.batch_size,
         prb=cfg.replay_buffer.prb,
         buffer_size=cfg.replay_buffer.size,
-        buffer_scratch_dir=cfg.replay_buffer.scratch_dir,
+        scratch_dir=cfg.replay_buffer.scratch_dir,
         device="cpu",
     )
 

diff --git a/examples/ddpg/utils.py b/examples/ddpg/utils.py
@@ -119,7 +119,7 @@ def make_replay_buffer(
     batch_size,
     prb=False,
     buffer_size=1000000,
-    buffer_scratch_dir=None,
+    scratch_dir=None,
     device="cpu",
     prefetch=3,
 ):
@@ -131,7 +131,7 @@ def make_replay_buffer(
             prefetch=prefetch,
             storage=LazyMemmapStorage(
                 buffer_size,
-                scratch_dir=buffer_scratch_dir,
+                scratch_dir=scratch_dir,
                 device=device,
             ),
             batch_size=batch_size,
@@ -142,7 +142,7 @@ def make_replay_buffer(
             prefetch=prefetch,
             storage=LazyMemmapStorage(
                 buffer_size,
-                scratch_dir=buffer_scratch_dir,
+                scratch_dir=scratch_dir,
                 device=device,
             ),
             batch_size=batch_size,

diff --git a/examples/decision_transformer/dt_config.yaml b/examples/decision_transformer/dt_config.yaml
@@ -36,7 +36,7 @@ replay_buffer:
   stacked_frames: 20
   buffer_prefetch: 64
   capacity: 1_000_000
-  buffer_scratch_dir:
+  scratch_dir:
   device: cpu
   prefetch: 3
 

diff --git a/examples/decision_transformer/odt_config.yaml b/examples/decision_transformer/odt_config.yaml
@@ -36,7 +36,7 @@ replay_buffer:
   stacked_frames: 20
   buffer_prefetch: 64
   capacity: 1_000_000
-  buffer_scratch_dir:
+  scratch_dir:
   device: cuda:0
   prefetch: 3
 

diff --git a/examples/decision_transformer/utils.py b/examples/decision_transformer/utils.py
@@ -296,7 +296,7 @@ def make_online_replay_buffer(offline_buffer, rb_cfg, reward_scaling=0.001):
     )
     storage = LazyMemmapStorage(
         max_size=rb_cfg.capacity,
-        scratch_dir=rb_cfg.buffer_scratch_dir,
+        scratch_dir=rb_cfg.scratch_dir,
         device=rb_cfg.device,
     )
 

diff --git a/examples/discrete_sac/config.yaml b/examples/discrete_sac/config.yaml
@@ -22,7 +22,7 @@ collector:
 replay_buffer:
   prb: 0 # use prioritized experience replay
   size: 1000000
-  scratch_dir: ${torchrl_logger.exp_name}_${env.seed}
+  scratch_dir: null
 
 # optim
 optim:

diff --git a/examples/discrete_sac/discrete_sac.py b/examples/discrete_sac/discrete_sac.py
@@ -74,7 +74,7 @@ def main(cfg: "DictConfig"):  # noqa: F821
         batch_size=cfg.optim.batch_size,
         prb=cfg.replay_buffer.prb,
         buffer_size=cfg.replay_buffer.size,
-        buffer_scratch_dir=cfg.replay_buffer.scratch_dir,
+        scratch_dir=cfg.replay_buffer.scratch_dir,
         device="cpu",
     )
 

diff --git a/examples/discrete_sac/utils.py b/examples/discrete_sac/utils.py
@@ -120,14 +120,14 @@ def make_replay_buffer(
     batch_size,
     prb=False,
     buffer_size=1000000,
-    buffer_scratch_dir=None,
+    scratch_dir=None,
     device="cpu",
     prefetch=3,
 ):
     with (
         tempfile.TemporaryDirectory()
-        if buffer_scratch_dir is None
-        else nullcontext(buffer_scratch_dir)
+        if scratch_dir is None
+        else nullcontext(scratch_dir)
     ) as scratch_dir:
         if prb:
             replay_buffer = TensorDictPrioritizedReplayBuffer(

diff --git a/examples/distributed/collectors/multi_nodes/ray_train.py b/examples/distributed/collectors/multi_nodes/ray_train.py
@@ -36,7 +36,7 @@
 if __name__ == "__main__":
 
     # 1. Define Hyperparameters
-    device = "cpu"  # if not torch.has_cuda else "cuda:0"
+    device = "cpu"  # if not torch.cuda.device_count() else "cuda:0"
     num_cells = 256
     max_grad_norm = 1.0
     frame_skip = 1

diff --git a/examples/iql/utils.py b/examples/iql/utils.py
@@ -125,7 +125,7 @@ def make_replay_buffer(
     batch_size,
     prb=False,
     buffer_size=1000000,
-    buffer_scratch_dir=None,
+    scratch_dir=None,
     device="cpu",
     prefetch=3,
 ):
@@ -137,7 +137,7 @@ def make_replay_buffer(
             prefetch=prefetch,
             storage=LazyMemmapStorage(
                 buffer_size,
-                scratch_dir=buffer_scratch_dir,
+                scratch_dir=scratch_dir,
                 device=device,
             ),
             batch_size=batch_size,
@@ -148,7 +148,7 @@ def make_replay_buffer(
             prefetch=prefetch,
             storage=LazyMemmapStorage(
                 buffer_size,
-                scratch_dir=buffer_scratch_dir,
+                scratch_dir=scratch_dir,
                 device=device,
             ),
             batch_size=batch_size,

diff --git a/examples/multiagent/iql.py b/examples/multiagent/iql.py
@@ -7,7 +7,7 @@
 import hydra
 import torch
 
-from tensordict.nn import TensorDictModule
+from tensordict.nn import TensorDictModule, TensorDictSequential
 from torch import nn
 from torchrl._utils import logger as torchrl_logger
 from torchrl.collectors import SyncDataCollector
@@ -17,7 +17,7 @@
 from torchrl.envs import RewardSum, TransformedEnv
 from torchrl.envs.libs.vmas import VmasEnv
 from torchrl.envs.utils import ExplorationType, set_exploration_type
-from torchrl.modules import EGreedyWrapper, QValueModule, SafeSequential
+from torchrl.modules import EGreedyModule, QValueModule, SafeSequential
 from torchrl.modules.models.multiagent import MultiAgentMLP
 from torchrl.objectives import DQNLoss, SoftUpdate, ValueEstimators
 from utils.logging import init_logging, log_evaluation, log_training
@@ -31,7 +31,7 @@ def rendering_callback(env, td):
 @hydra.main(version_base="1.1", config_path=".", config_name="iql")
 def train(cfg: "DictConfig"):  # noqa: F821
     # Device
-    cfg.train.device = "cpu" if not torch.has_cuda else "cuda:0"
+    cfg.train.device = "cpu" if not torch.cuda.device_count() else "cuda:0"
     cfg.env.device = cfg.train.device
 
     # Seeding
@@ -96,13 +96,15 @@ def train(cfg: "DictConfig"):  # noqa: F821
     )
     qnet = SafeSequential(module, value_module)
 
-    qnet_explore = EGreedyWrapper(
+    qnet_explore = TensorDictSequential(
         qnet,
-        eps_init=0.3,
-        eps_end=0,
-        annealing_num_steps=int(cfg.collector.total_frames * (1 / 2)),
-        action_key=env.action_key,
-        spec=env.unbatched_action_spec,
+        EGreedyModule(
+            eps_init=0.3,
+            eps_end=0,
+            annealing_num_steps=int(cfg.collector.total_frames * (1 / 2)),
+            action_key=env.action_key,
+            spec=env.unbatched_action_spec,
+        ),
     )
 
     collector = SyncDataCollector(

diff --git a/examples/multiagent/maddpg_iddpg.py b/examples/multiagent/maddpg_iddpg.py
@@ -36,7 +36,7 @@ def rendering_callback(env, td):
 @hydra.main(version_base="1.1", config_path=".", config_name="maddpg_iddpg")
 def train(cfg: "DictConfig"):  # noqa: F821
     # Device
-    cfg.train.device = "cpu" if not torch.has_cuda else "cuda:0"
+    cfg.train.device = "cpu" if not torch.cuda.device_count() else "cuda:0"
     cfg.env.device = cfg.train.device
 
     # Seeding

diff --git a/examples/multiagent/mappo_ippo.py b/examples/multiagent/mappo_ippo.py
@@ -31,7 +31,7 @@ def rendering_callback(env, td):
 @hydra.main(version_base="1.1", config_path=".", config_name="mappo_ippo")
 def train(cfg: "DictConfig"):  # noqa: F821
     # Device
-    cfg.train.device = "cpu" if not torch.has_cuda else "cuda:0"
+    cfg.train.device = "cpu" if not torch.cuda.device_count() else "cuda:0"
     cfg.env.device = cfg.train.device
 
     # Seeding

diff --git a/examples/multiagent/qmix_vdn.py b/examples/multiagent/qmix_vdn.py
@@ -7,7 +7,7 @@
 import hydra
 import torch
 
-from tensordict.nn import TensorDictModule
+from tensordict.nn import TensorDictModule, TensorDictSequential
 from torch import nn
 from torchrl._utils import logger as torchrl_logger
 from torchrl.collectors import SyncDataCollector
@@ -17,7 +17,7 @@
 from torchrl.envs import RewardSum, TransformedEnv
 from torchrl.envs.libs.vmas import VmasEnv
 from torchrl.envs.utils import ExplorationType, set_exploration_type
-from torchrl.modules import EGreedyWrapper, QValueModule, SafeSequential
+from torchrl.modules import EGreedyModule, QValueModule, SafeSequential
 from torchrl.modules.models.multiagent import MultiAgentMLP, QMixer, VDNMixer
 from torchrl.objectives import SoftUpdate, ValueEstimators
 from torchrl.objectives.multiagent.qmixer import QMixerLoss
@@ -31,7 +31,7 @@ def rendering_callback(env, td):
 @hydra.main(version_base="1.1", config_path=".", config_name="qmix_vdn")
 def train(cfg: "DictConfig"):  # noqa: F821
     # Device
-    cfg.train.device = "cpu" if not torch.has_cuda else "cuda:0"
+    cfg.train.device = "cpu" if not torch.cuda.device_count() else "cuda:0"
     cfg.env.device = cfg.train.device
 
     # Seeding
@@ -96,13 +96,15 @@ def train(cfg: "DictConfig"):  # noqa: F821
     )
     qnet = SafeSequential(module, value_module)
 
-    qnet_explore = EGreedyWrapper(
+    qnet_explore = TensorDictSequential(
         qnet,
-        eps_init=0.3,
-        eps_end=0,
-        annealing_num_steps=int(cfg.collector.total_frames * (1 / 2)),
-        action_key=env.action_key,
-        spec=env.unbatched_action_spec,
+        EGreedyModule(
+            eps_init=0.3,
+            eps_end=0,
+            annealing_num_steps=int(cfg.collector.total_frames * (1 / 2)),
+            action_key=env.action_key,
+            spec=env.unbatched_action_spec,
+        ),
     )
 
     if cfg.loss.mixer_type == "qmix":

diff --git a/examples/multiagent/sac.py b/examples/multiagent/sac.py
@@ -33,7 +33,7 @@ def rendering_callback(env, td):
 @hydra.main(version_base="1.1", config_path=".", config_name="sac")
 def train(cfg: "DictConfig"):  # noqa: F821
     # Device
-    cfg.train.device = "cpu" if not torch.has_cuda else "cuda:0"
+    cfg.train.device = "cpu" if not torch.cuda.device_count() else "cuda:0"
     cfg.env.device = cfg.train.device
 
     # Seeding