[Doc, BugFix] Fix tutos errors (#817)

Author: vmoens

tutorials/sphinx-tutorials/coding_dqn.py

@@ -378,7 +378,8 @@ def make_model():
         frames.append(current_frames)
 
     if data["done"].any():
-        traj_lengths.append(data["step_count"][data["done"]].float().mean().item())
+        done = data["done"].squeeze(-1)
+        traj_lengths.append(data["step_count"][done].float().mean().item())
 
     # check that we have enough data to start training
     if sum(frames) > init_random_frames:
@@ -612,7 +613,8 @@ def make_model():
         frames.append(current_frames)
 
     if data["done"].any():
-        traj_lengths.append(data["step_count"][data["done"]].float().mean().item())
+        done = data["done"].squeeze(-1)
+        traj_lengths.append(data["step_count"][done].float().mean().item())
 
     if sum(frames) > init_random_frames:
         for _ in range(n_optim):

tutorials/sphinx-tutorials/torchrl_demo.py

@@ -330,6 +330,7 @@
     Compose,
     NoopResetEnv,
     ObservationNorm,
+    StepCounter,
     ToTensorImage,
     TransformedEnv,
 )
@@ -358,7 +359,7 @@
     lambda: GymEnv("Pendulum-v1", frame_skip=3, from_pixels=True, pixels_only=False),
 )
 env = TransformedEnv(
-    base_env, Compose(NoopResetEnv(3), ToTensorImage())
+    base_env, Compose(StepCounter(), ToTensorImage())
 )  # applies transforms on batch of envs
 env.append_transform(ObservationNorm(in_keys=["pixels"], loc=2, scale=1))
 env.reset()
@@ -587,9 +588,9 @@
 for i in range(max_steps):
     actor(tensordict)
     tensordicts[i] = env.step(tensordict)
-    tensordict = step_mdp(tensordict)  # roughly equivalent to obs = next_obs
-    if env.is_done:
+    if tensordict["done"].any():
         break
+    tensordict = step_mdp(tensordict)  # roughly equivalent to obs = next_obs
 
 tensordicts_prealloc = tensordicts.clone()
 print("total steps:", i)
@@ -607,9 +608,9 @@
 for _ in range(max_steps):
     actor(tensordict)
     tensordicts.append(env.step(tensordict))
-    tensordict = step_mdp(tensordict)  # roughly equivalent to obs = next_obs
-    if env.is_done:
+    if tensordict["done"].any():
         break
+    tensordict = step_mdp(tensordict)  # roughly equivalent to obs = next_obs
 tensordicts_stack = torch.stack(tensordicts, 0)
 print("total steps:", i)
 print(tensordicts_stack)