Repository: pytorch/rl

Issue:
[BUG] `SACLoss` cannot be loaded from `state_dict`
The SACLoss has a `target_entropy_buffer` which is instantiated not at init time but lazily https://github.com/pytorch/rl/blob/106368ffe627aada70ebb33ecf545d0f193c7757/torchrl/objectives/sac.py#L424

This is a problem as it will prevent reloading this loss from `state_dict`.

What will happen is that `state_dict` will likely contain this buffer and the reloaded loss will not (as the lazy step has not happened yet)

## Solution
We instantiate all buffers at init time. That seems the only way to me to support loading from `state_dict`


Discussion:
"the only way" is rarely true :) 
Can you describe the problem more in details, ie what you're trying to do (when do you save / when do you load etc)?
Sorry, here is the example
```python
import torch

from torchrl.data import UnboundedContinuousTensorSpec
from torchrl.modules import QValueActor, ProbabilisticActor, TanhDelta, ValueOperator
from tensordict.nn import TensorDictModule

from torchrl.objectives import SACLoss

if __name__ == "__main__":
    model = torch.nn.Linear(3, 4)  # does not work
    # model = torch.nn.Linear(1, 1) works
    actor_module = TensorDictModule(
        torch.nn.Linear(3, 4), in_keys=["obs"], out_keys=["logits"]
    )
    policy = ProbabilisticActor(
        module=actor_module,
        in_keys=["logits"],
        out_keys=["action"],
        distribution_class=TanhDelta,
    )
    value = ValueOperator(module=model, in_keys=["obs"], out_keys="value")

    loss = SACLoss(
        actor_network=policy,
        qvalue_network=value,
        action_spec=UnboundedContinuousTensorSpec(shape=(2,)),
    )
    _ = (
        loss.target_entropy
    )  # This instantiates a buffer. It is called during the forward pass. Here we call it to emulate a forward
    state = loss.state_dict()

    loss = SACLoss(
        actor_network=policy,
        qvalue_network=value,
        action_spec=UnboundedContinuousTensorSpec(shape=(2,)),
    )
    loss.load_state_dict(
        state
    )  # The loss.target_entropy was in the state_dict but cannot be found in the loss as it does not exist yet

```

```bash
Traceback (most recent call last):
  File "/Users/matbet/PycharmProjects/rl/prova.py", line 30, in <module>
    loss.load_state_dict(
  File "/Users/matbet/miniconda3/envs/torchrl/lib/python3.9/site-packages/torch/nn/modules/module.py", line 2041, in load_state_dict
    raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for SACLoss:
	Unexpected key(s) in state_dict: "target_entropy_buffer". 
```
one simple solution is calling `_ = self.target_entropy` in the init of the loss. Or transfer the whole logic of `self.target_entropy` in init like it is done in `DiscreteSACLoss`
Let me refactor that
We need this to be instantiated before the keys are set, because it reads the action key in the action spec. So that can't be done in `__init__`.
I wonder why we can make it in the discrete case though...
My suggestion is to make it happen in state_dict if it's not done already
In your example, I'm a bit confused by 
```python
    model = torch.nn.Linear(3, 4)  # does not work
    # model = torch.nn.Linear(1, 1) works
```
can you explain?

Resolving patch:
diff --git a/torchrl/objectives/sac.py b/torchrl/objectives/sac.py
--- a/torchrl/objectives/sac.py
+++ b/torchrl/objectives/sac.py
@@ -5,6 +5,7 @@
 import math
 import warnings
 from dataclasses import dataclass
+from functools import wraps
 from numbers import Number
 from typing import Dict, Optional, Tuple, Union
 
@@ -43,6 +44,15 @@
     FUNCTORCH_ERROR = err
 
 
+def _delezify(func):
+    @wraps(func)
+    def new_func(self, *args, **kwargs):
+        self.target_entropy
+        return func(self, *args, **kwargs)
+
+    return new_func
+
+
 class SACLoss(LossModule):
     """TorchRL implementation of the SAC loss.
 
@@ -371,7 +381,6 @@ def __init__(
 
         self._target_entropy = target_entropy
         self._action_spec = action_spec
-        self.target_entropy_buffer = None
         if self._version == 1:
             self.actor_critic = ActorCriticWrapper(
                 self.actor_network, self.value_network
@@ -384,48 +393,54 @@ def __init__(
         if self._version == 1:
             self._vmap_qnetwork00 = vmap(qvalue_network)
 
+    @property
+    def target_entropy_buffer(self):
+        return self.target_entropy
+
     @property
     def target_entropy(self):
-        target_entropy = self.target_entropy_buffer
-        if target_entropy is None:
-            delattr(self, "target_entropy_buffer")
-            target_entropy = self._target_entropy
-            action_spec = self._action_spec
-            actor_network = self.actor_network
-            device = next(self.parameters()).device
-            if target_entropy == "auto":
-                action_spec = (
-                    action_spec
-                    if action_spec is not None
-                    else getattr(actor_network, "spec", None)
-                )
-                if action_spec is None:
-                    raise RuntimeError(
-                        "Cannot infer the dimensionality of the action. Consider providing "
-                        "the target entropy explicitely or provide the spec of the "
-                        "action tensor in the actor network."
-                    )
-                if not isinstance(action_spec, CompositeSpec):
-                    action_spec = CompositeSpec({self.tensor_keys.action: action_spec})
-                if (
-                    isinstance(self.tensor_keys.action, tuple)
-                    and len(self.tensor_keys.action) > 1
-                ):
-                    action_container_shape = action_spec[
-                        self.tensor_keys.action[:-1]
-                    ].shape
-                else:
-                    action_container_shape = action_spec.shape
-                target_entropy = -float(
-                    action_spec[self.tensor_keys.action]
-                    .shape[len(action_container_shape) :]
-                    .numel()
+        target_entropy = self._buffers.get("_target_entropy", None)
+        if target_entropy is not None:
+            return target_entropy
+        target_entropy = self._target_entropy
+        action_spec = self._action_spec
+        actor_network = self.actor_network
+        device = next(self.parameters()).device
+        if target_entropy == "auto":
+            action_spec = (
+                action_spec
+                if action_spec is not None
+                else getattr(actor_network, "spec", None)
+            )
+            if action_spec is None:
+                raise RuntimeError(
+                    "Cannot infer the dimensionality of the action. Consider providing "
+                    "the target entropy explicitely or provide the spec of the "
+                    "action tensor in the actor network."
                 )
-            self.register_buffer(
-                "target_entropy_buffer", torch.tensor(target_entropy, device=device)
+            if not isinstance(action_spec, CompositeSpec):
+                action_spec = CompositeSpec({self.tensor_keys.action: action_spec})
+            if (
+                isinstance(self.tensor_keys.action, tuple)
+                and len(self.tensor_keys.action) > 1
+            ):
+             
+                action_container_shape = action_spec[self.tensor_keys.action[:-1]].shape
+            else:
+                action_container_shape = action_spec.shape
+            target_entropy = -float(
+                action_spec[self.tensor_keys.action]
+                .shape[len(action_container_shape) :]
+                .numel()
             )
-            return self.target_entropy_buffer
-        return target_entropy
+        delattr(self, "_target_entropy")
+        self.register_buffer(
+            "_target_entropy", torch.tensor(target_entropy, device=device)
+        )
+        return self._target_entropy
+
+    state_dict = _delezify(LossModule.state_dict)
+    load_state_dict = _delezify(LossModule.load_state_dict)
 
     def _forward_value_estimator_keys(self, **kwargs) -> None:
         if self._value_estimator is not None:


Test patch:
diff --git a/test/test_cost.py b/test/test_cost.py
--- a/test/test_cost.py
+++ b/test/test_cost.py
@@ -3260,6 +3260,49 @@ def test_sac_notensordict(
         assert loss_actor == loss_val_td["loss_actor"]
         assert loss_alpha == loss_val_td["loss_alpha"]
 
+    def test_state_dict(self, version):
+        if version == 1:
+            pytest.skip("Test not implemented for version 1.")
+        model = torch.nn.Linear(3, 4)
+        actor_module = TensorDictModule(model, in_keys=["obs"], out_keys=["logits"])
+        policy = ProbabilisticActor(
+            module=actor_module,
+            in_keys=["logits"],
+            out_keys=["action"],
+            distribution_class=TanhDelta,
+        )
+        value = ValueOperator(module=model, in_keys=["obs"], out_keys="value")
+
+        loss = SACLoss(
+            actor_network=policy,
+            qvalue_network=value,
+            action_spec=UnboundedContinuousTensorSpec(shape=(2,)),
+        )
+        state = loss.state_dict()
+
+        loss = SACLoss(
+            actor_network=policy,
+            qvalue_network=value,
+            action_spec=UnboundedContinuousTensorSpec(shape=(2,)),
+        )
+        loss.load_state_dict(state)
+
+        # with an access in between
+        loss = SACLoss(
+            actor_network=policy,
+            qvalue_network=value,
+            action_spec=UnboundedContinuousTensorSpec(shape=(2,)),
+        )
+        loss.target_entropy
+        state = loss.state_dict()
+
+        loss = SACLoss(
+            actor_network=policy,
+            qvalue_network=value,
+            action_spec=UnboundedContinuousTensorSpec(shape=(2,)),
+        )
+        loss.load_state_dict(state)
+
 
 @pytest.mark.skipif(
     not _has_functorch, reason=f"functorch not installed: {FUNCTORCH_ERR}"