[BugFix]: gradient propagation in advantage estimates (#322)

Author: vmoens

test/test_cost.py

@@ -1573,13 +1573,9 @@ def test_hold_out():
     assert y.requires_grad
 
     # exception
-    with pytest.raises(
-        RuntimeError,
-        match="hold_out_net requires the network parameter set to be non-empty.",
-    ):
-        net = torch.nn.Sequential()
-        with hold_out_net(net):
-            pass
+    net = torch.nn.Sequential()
+    with hold_out_net(net):
+        pass
 
 
 @pytest.mark.parametrize("mode", ["hard", "soft"])

torchrl/objectives/costs/utils.py

@@ -263,9 +263,7 @@ def __init__(self, network: nn.Module) -> None:
         try:
             self.p_example = next(network.parameters())
         except StopIteration:
-            raise RuntimeError(
-                "hold_out_net requires the network parameter set to be " "non-empty."
-            )
+            self.p_example = torch.tensor([])
         self._prev_state = []
 
     def __enter__(self) -> None:

torchrl/objectives/returns/advantages.py

@@ -29,6 +29,8 @@
 
 __all__ = ["GAE", "TDLambdaEstimate", "TDEstimate"]
 
+from ..costs.utils import hold_out_net
+
 
 class TDEstimate:
     """Temporal Difference estimate of advantage function.
@@ -89,7 +91,7 @@ def __call__(
             if self.average_rewards:
                 reward = reward - reward.mean()
                 reward = reward / reward.std().clamp_min(1e-4)
-                tensordict.set_(
+                tensordict.set(
                     "reward", reward
                 )  # we must update the rewards if they are used later in the code
 
@@ -106,12 +108,19 @@ def __call__(
             self.value_network(tensordict, **kwargs)
             value = tensordict.get(self.value_key)
 
-        with torch.set_grad_enabled(False):
+        with hold_out_net(self.value_network):
+            # we may still need to pass gradient, but we don't want to assign grads to
+            # value net params
             step_td = step_tensordict(tensordict)
             if target_params is not None:
+                # we assume that target parameters are not differentiable
                 kwargs["params"] = target_params
+            elif "params" in kwargs:
+                kwargs["params"] = [param.detach() for param in kwargs["params"]]
             if target_buffers is not None:
                 kwargs["buffers"] = target_buffers
+            elif "buffers" in kwargs:
+                kwargs["buffers"] = [buffer.detach() for buffer in kwargs["buffers"]]
             self.value_network(step_td, **kwargs)
             next_value = step_td.get(self.value_key)
 
@@ -190,7 +199,7 @@ def __call__(
             if self.average_rewards:
                 reward = reward - reward.mean()
                 reward = reward / reward.std().clamp_min(1e-4)
-                tensordict.set_(
+                tensordict.set(
                     "reward", reward
                 )  # we must update the rewards if they are used later in the code
 
@@ -209,12 +218,19 @@ def __call__(
             self.value_network(tensordict, **kwargs)
             value = tensordict.get(self.value_key)
 
-        with torch.set_grad_enabled(False):
+        with hold_out_net(self.value_network):
+            # we may still need to pass gradient, but we don't want to assign grads to
+            # value net params
             step_td = step_tensordict(tensordict)
             if target_params is not None:
+                # we assume that target parameters are not differentiable
                 kwargs["params"] = target_params
+            elif "params" in kwargs:
+                kwargs["params"] = [param.detach() for param in kwargs["params"]]
             if target_buffers is not None:
                 kwargs["buffers"] = target_buffers
+            elif "buffers" in kwargs:
+                kwargs["buffers"] = [buffer.detach() for buffer in kwargs["buffers"]]
             self.value_network(step_td, **kwargs)
             next_value = step_td.get(self.value_key)
 
@@ -295,7 +311,7 @@ def __call__(
             if self.average_rewards:
                 reward = reward - reward.mean()
                 reward = reward / reward.std().clamp_min(1e-4)
-                tensordict.set_(
+                tensordict.set(
                     "reward", reward
                 )  # we must update the rewards if they are used later in the code
 
@@ -312,12 +328,19 @@ def __call__(
             self.value_network(tensordict, **kwargs)
             value = tensordict.get("state_value")
 
-        with torch.set_grad_enabled(False):
+        with hold_out_net(self.value_network):
+            # we may still need to pass gradient, but we don't want to assign grads to
+            # value net params
             step_td = step_tensordict(tensordict)
             if target_params is not None:
+                # we assume that target parameters are not differentiable
                 kwargs["params"] = target_params
+            elif "params" in kwargs:
+                kwargs["params"] = [param.detach() for param in kwargs["params"]]
             if target_buffers is not None:
                 kwargs["buffers"] = target_buffers
+            elif "buffers" in kwargs:
+                kwargs["buffers"] = [buffer.detach() for buffer in kwargs["buffers"]]
             self.value_network(step_td, **kwargs)
             next_value = step_td.get("state_value")
             done = tensordict.get("done")