[Minor] Cleanup

Author: vmoens

sota-check/README.md

@@ -25,7 +25,7 @@ export MUJOCO_GL=egl
 
 conda create -n rl-sota-bench python=3.10 -y 
 conda install anaconda::libglu -y
-pip3 install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu121
+pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu121
 pip3 install "gymnasium[accept-rom-license,atari,mujoco]" vmas tqdm wandb pygame moviepy imageio submitit hydra-core transformers
 
 cd /path/to/tensordict

test/test_env.py

@@ -2549,7 +2549,8 @@ def _step(self, tensordict):
                     "reward": action.sum().unsqueeze(0),
                     **self.full_done_spec.zero(),
                     "observation": obs,
-                }
+                },
+                batch_size=[],
             )
 
     torch.manual_seed(0)

torchrl/envs/batched_envs.py

@@ -433,9 +433,6 @@ def _check_for_empty_spec(specs: CompositeSpec):
                 def map_device(key, value, device_map=device_map):
                     return value.to(device_map[key])
 
-                # self._env_tensordict.named_apply(
-                #     map_device, nested_keys=True, filter_empty=True
-                # )
                 self._env_tensordict.named_apply(
                     map_device,
                     nested_keys=True,
@@ -809,11 +806,6 @@ def select_and_clone(name, tensor):
             if name in selected_output_keys:
                 return tensor.clone()
 
-        # out = self.shared_tensordict_parent.named_apply(
-        #     select_and_clone,
-        #     nested_keys=True,
-        #     filter_empty=True,
-        # )
         out = self.shared_tensordict_parent.named_apply(
             select_and_clone,
             nested_keys=True,
@@ -1208,14 +1200,12 @@ def step_and_maybe_reset(
                 if x.device != device
                 else x.clone(),
                 device=device,
-                # filter_empty=True,
             )
             tensordict_ = tensordict_._fast_apply(
                 lambda x: x.to(device, non_blocking=self.non_blocking)
                 if x.device != device
                 else x.clone(),
                 device=device,
-                # filter_empty=True,
             )
         else:
             next_td = next_td.clone().clear_device_()
@@ -1271,7 +1261,6 @@ def select_and_clone(name, tensor):
         out = next_td.named_apply(
             select_and_clone,
             nested_keys=True,
-            # filter_empty=True,
         )
         if out.device != device:
             if device is None:
@@ -1357,7 +1346,6 @@ def select_and_clone(name, tensor):
         out = self.shared_tensordict_parent.named_apply(
             select_and_clone,
             nested_keys=True,
-            # filter_empty=True,
         )
         del out["next"]
 
@@ -1495,7 +1483,6 @@ def _run_worker_pipe_shared_mem(
         def look_for_cuda(tensor, has_cuda=has_cuda):
             has_cuda[0] = has_cuda[0] or tensor.is_cuda
 
-        # shared_tensordict.apply(look_for_cuda, filter_empty=True)
         shared_tensordict.apply(look_for_cuda)
         has_cuda = has_cuda[0]
     else:
@@ -1685,9 +1672,5 @@ def look_for_cuda(tensor, has_cuda=has_cuda):
                 child_pipe.send(("_".join([cmd, "done"]), None))
 
 
-def _filter_empty(tensordict):
-    return tensordict.select(*tensordict.keys(True, True))
-
-
 # Create an alias for possible imports
 _BatchedEnv = BatchedEnvBase

torchrl/objectives/common.py

@@ -252,7 +252,6 @@ def _compare_and_expand(param):
                     return param._apply_nest(
                         _compare_and_expand,
                         batch_size=[expand_dim, *param.shape],
-                        filter_empty=False,
                         call_on_nested=True,
                     )
                 if not isinstance(param, nn.Parameter):

torchrl/objectives/ppo.py

@@ -455,7 +455,7 @@ def get_entropy_bonus(self, dist: d.Distribution) -> torch.Tensor:
             entropy = dist.entropy()
         except NotImplementedError:
             x = dist.rsample((self.samples_mc_entropy,))
-            entropy = -dist.log_prob(x)
+            entropy = -dist.log_prob(x).mean(0)
         return entropy.unsqueeze(-1)
 
     def _log_weight(
@@ -1036,7 +1036,7 @@ def forward(self, tensordict: TensorDictBase) -> TensorDict:
             td_out.set("loss_entropy", -self.entropy_coef * entropy.mean())
 
         if self.critic_coef:
-            loss_critic = self.loss_critic(tensordict)
+            loss_critic = self.loss_critic(tensordict_copy)
             td_out.set("loss_critic", loss_critic.mean())
 
         return td_out