[BugFix] Fix tutos (#1648)

Author: Vincent Moens

torchrl/collectors/collectors.py

@@ -567,7 +567,9 @@ def __init__(
             self.policy_weights = TensorDict({}, [])
 
         self.env: EnvBase = self.env.to(self.device)
-        self.max_frames_per_traj = max_frames_per_traj
+        self.max_frames_per_traj = (
+            int(max_frames_per_traj) if max_frames_per_traj is not None else 0
+        )
         if self.max_frames_per_traj is not None and self.max_frames_per_traj > 0:
             # let's check that there is no StepCounter yet
             for key in self.env.output_spec.keys(True, True):
@@ -595,9 +597,13 @@ def __init__(
                     f"This means {frames_per_batch - remainder} additional frames will be collected."
                     "To silence this message, set the environment variable RL_WARNINGS to False."
                 )
-        self.total_frames = total_frames
+        self.total_frames = (
+            int(total_frames) if total_frames != float("inf") else total_frames
+        )
         self.reset_at_each_iter = reset_at_each_iter
-        self.init_random_frames = init_random_frames
+        self.init_random_frames = (
+            int(init_random_frames) if init_random_frames is not None else 0
+        )
         if (
             init_random_frames is not None
             and init_random_frames % frames_per_batch != 0
@@ -620,7 +626,7 @@ def __init__(
                 f" ({-(-frames_per_batch // self.n_env) * self.n_env})."
                 "To silence this message, set the environment variable RL_WARNINGS to False."
             )
-        self.requested_frames_per_batch = frames_per_batch
+        self.requested_frames_per_batch = int(frames_per_batch)
         self.frames_per_batch = -(-frames_per_batch // self.n_env)
         self.exploration_type = (
             exploration_type if exploration_type else DEFAULT_EXPLORATION_TYPE
@@ -1234,11 +1240,15 @@ def device_err_msg(device_name, devices_list):
                     f"This means {frames_per_batch - remainder} additional frames will be collected."
                     "To silence this message, set the environment variable RL_WARNINGS to False."
                 )
-        self.total_frames = total_frames
+        self.total_frames = (
+            int(total_frames) if total_frames != float("inf") else total_frames
+        )
         self.reset_at_each_iter = reset_at_each_iter
         self.postprocs = postproc
-        self.max_frames_per_traj = max_frames_per_traj
-        self.requested_frames_per_batch = frames_per_batch
+        self.max_frames_per_traj = (
+            int(max_frames_per_traj) if max_frames_per_traj is not None else 0
+        )
+        self.requested_frames_per_batch = int(frames_per_batch)
         self.reset_when_done = reset_when_done
         if split_trajs is None:
             split_trajs = False
@@ -1247,7 +1257,9 @@ def device_err_msg(device_name, devices_list):
                 "Cannot split trajectories when reset_when_done is False."
             )
         self.split_trajs = split_trajs
-        self.init_random_frames = init_random_frames
+        self.init_random_frames = (
+            int(init_random_frames) if init_random_frames is not None else 0
+        )
         self.update_at_each_batch = update_at_each_batch
         self.exploration_type = exploration_type
         self.frames_per_worker = np.inf

torchrl/envs/transforms/r3m.py

@@ -302,7 +302,7 @@ def _init(self):
         transforms.append(resize)
 
         # R3M
-        if out_keys is None:
+        if out_keys in (None, []):
             if stack_images:
                 out_keys = ["r3m_vec"]
             else:

torchrl/envs/transforms/vip.py

@@ -2,7 +2,6 @@
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
-
 from typing import List, Optional, Union
 
 import torch
@@ -277,7 +276,7 @@ def _init(self):
         transforms.append(resize)
 
         # VIP
-        if out_keys is None:
+        if out_keys in (None, []):
             if stack_images:
                 out_keys = ["vip_vec"]
             else:

tutorials/sphinx-tutorials/coding_dqn.py

@@ -390,7 +390,7 @@ def get_replay_buffer(buffer_size, n_optim, batch_size):
 
 
 def get_collector(
-    obs_norm_sd,
+    stats,
     num_collectors,
     actor_explore,
     frames_per_batch,
@@ -399,7 +399,7 @@ def get_collector(
 ):
     data_collector = MultiaSyncDataCollector(
         [
-            make_env(parallel=True, obs_norm_sd=obs_norm_sd),
+            make_env(parallel=True, obs_norm_sd=stats),
         ]
         * num_collectors,
         policy=actor_explore,
@@ -566,7 +566,12 @@ def get_loss_module(actor, gamma):
 loss_module, target_net_updater = get_loss_module(actor, gamma)
 
 collector = get_collector(
-    stats, num_collectors, actor_explore, frames_per_batch, total_frames, device
+    stats=stats,
+    num_collectors=num_collectors,
+    actor_explore=actor_explore,
+    frames_per_batch=frames_per_batch,
+    total_frames=total_frames,
+    device=device,
 )
 optimizer = torch.optim.Adam(
     loss_module.parameters(), lr=lr, weight_decay=wd, betas=betas

tutorials/sphinx-tutorials/pendulum.py

@@ -652,6 +652,12 @@ class SinTransform(Transform):
     def _apply_transform(self, obs: torch.Tensor) -> None:
         return obs.sin()
 
+    # The transform must also modify the data at reset time
+    def _reset(
+        self, tensordict: TensorDictBase, tensordict_reset: TensorDictBase
+    ) -> TensorDictBase:
+        return self._call(tensordict_reset)
+
     # _apply_to_composite will execute the observation spec transform across all
     # in_keys/out_keys pairs and write the result in the observation_spec which
     # is of type ``Composite``
@@ -670,6 +676,12 @@ class CosTransform(Transform):
     def _apply_transform(self, obs: torch.Tensor) -> None:
         return obs.cos()
 
+    # The transform must also modify the data at reset time
+    def _reset(
+        self, tensordict: TensorDictBase, tensordict_reset: TensorDictBase
+    ) -> TensorDictBase:
+        return self._call(tensordict_reset)
+
     # _apply_to_composite will execute the observation spec transform across all
     # in_keys/out_keys pairs and write the result in the observation_spec which
     # is of type ``Composite``