[BugFix] Upgrade tensordict deps (#953)

Author: vmoens

torchrl/collectors/collectors.py

@@ -653,7 +653,7 @@ def rollout(self) -> TensorDictBase:
                     # unlock the output tensordict to allow for new keys to be written
                     # these will be missed during the sync but at least we won't get an error during the update
                     is_shared = self._tensordict_out.is_shared()
-                    self._tensordict_out.unlock()
+                    self._tensordict_out.unlock_()
                     self._tensordict_out[..., j] = self._tensordict
                     if is_shared:
                         self._tensordict_out.share_memory_()

torchrl/envs/common.py

@@ -379,15 +379,15 @@ def step(self, tensordict: TensorDictBase) -> TensorDictBase:
         # sanity check
         self._assert_tensordict_shape(tensordict)
 
-        tensordict.lock()  # make sure _step does not modify the tensordict
+        tensordict.lock_()  # make sure _step does not modify the tensordict
         tensordict_out = self._step(tensordict)
         if tensordict_out is tensordict:
             raise RuntimeError(
                 "EnvBase._step should return outplace changes to the input "
                 "tensordict. Consider emptying the TensorDict first (e.g. tensordict.empty() or "
                 "tensordict.select()) inside _step before writing new tensors onto this new instance."
             )
-        tensordict.unlock()
+        tensordict.unlock_()
 
         obs_keys = self.observation_spec.keys(nested_keys=False)
         # we deliberately do not update the input values, but we want to keep track of

torchrl/envs/libs/vmas.py

@@ -238,7 +238,7 @@ def _reset(
                 device=self.device,
             )
 
-            if infos is not None:
+            if agent_info is not None:
                 agent_td.set("info", agent_info)
             if dones is not None:
                 agent_td.set("done", dones[i])
@@ -277,7 +277,7 @@ def _step(
                 device=self.device,
             )
 
-            if infos is not None:
+            if agent_info is not None:
                 agent_td.set("info", agent_info)
             agent_tds.append(agent_td)
 

torchrl/envs/transforms/transforms.py

@@ -36,8 +36,16 @@
     from torchvision.transforms.functional import center_crop
 
     try:
-        from torchvision.transforms.functional import resize
+        from torchvision.transforms.functional import InterpolationMode, resize
+
+        def interpolation_fn(interpolation):  # noqa: D103
+            return InterpolationMode(interpolation)
+
     except ImportError:
+
+        def interpolation_fn(interpolation):  # noqa: D103
+            return interpolation
+
         from torchvision.transforms.functional_tensor import resize
 
     _has_tv = True
@@ -65,6 +73,14 @@ def new_fun(self, observation_spec):
 
 
 def _apply_to_composite_inv(function):
+    # Changes the input_spec following a transform function.
+    # The usage is: if an env expects a certain input (e.g. a double tensor)
+    # but the input has to be transformed (e.g. it is float), this function will
+    # modify the spec to get a spec that from the outside matches what is given
+    # (ie a float).
+    # Now since EnvBase.step ignores new inputs (ie the root level of the
+    # tensor is not updated) an out_key that does not match the in_key has
+    # no effect on the spec.
     def new_fun(self, input_spec):
         if isinstance(input_spec, CompositeSpec):
             d = input_spec._specs
@@ -996,7 +1012,7 @@ def __init__(
         super().__init__(in_keys=in_keys, out_keys=out_keys)
         self.w = int(w)
         self.h = int(h)
-        self.interpolation = interpolation
+        self.interpolation = interpolation_fn(interpolation)
 
     def _apply_transform(self, observation: torch.Tensor) -> torch.Tensor:
         # flatten if necessary

torchrl/objectives/value/advantages.py

@@ -7,7 +7,7 @@
 from typing import List, Optional, Tuple, Union
 
 import torch
-from tensordict.nn import dispatch_kwargs
+from tensordict.nn import dispatch
 from tensordict.tensordict import TensorDictBase
 from torch import nn, Tensor
 
@@ -95,7 +95,7 @@ def is_functional(self):
         )
 
     @_self_set_grad_enabled
-    @dispatch_kwargs
+    @dispatch
     def forward(
         self,
         tensordict: TensorDictBase,
@@ -269,7 +269,7 @@ def is_functional(self):
         )
 
     @_self_set_grad_enabled
-    @dispatch_kwargs
+    @dispatch
     def forward(
         self,
         tensordict: TensorDictBase,
@@ -461,7 +461,7 @@ def is_functional(self):
         )
 
     @_self_set_grad_enabled
-    @dispatch_kwargs
+    @dispatch
     def forward(
         self,
         tensordict: TensorDictBase,