[BugFix] Make casting to 'meta' device uniform across cost modules (#…

…1715)
pytorch · Nov 27, 2023 · 38d9cb7 · 38d9cb7
1 parent 0f93943
commit 38d9cb7
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 16 deletions.
diff --git a/torchrl/envs/transforms/rlhf.py b/torchrl/envs/transforms/rlhf.py
@@ -112,7 +112,7 @@ def __init__(
 
         # check that the model has parameters
         params = TensorDict.from_module(actor)
-        with params.apply(_stateless_param).to_module(actor):
+        with params.apply(_stateless_param, device="meta").to_module(actor):
             # copy a stateless actor
             self.__dict__["functional_actor"] = deepcopy(actor)
         # we need to register these params as buffer to have `to` and similar

diff --git a/torchrl/objectives/common.py b/torchrl/objectives/common.py
@@ -289,9 +289,9 @@ def _compare_and_expand(param):
 
         # set the functional module: we need to convert the params to non-differentiable params
         # otherwise they will appear twice in parameters
-        with params.apply(_make_meta_params, device=torch.device("meta")).to_module(
-            module
-        ):
+        with params.apply(
+            self._make_meta_params, device=torch.device("meta")
+        ).to_module(module):
             # avoid buffers and params being exposed
             self.__dict__[module_name] = deepcopy(module)
 
@@ -435,6 +435,16 @@ def make_value_estimator(self, value_type: ValueEstimators = None, **hyperparams
 
         return self
 
+    @staticmethod
+    def _make_meta_params(param):
+        is_param = isinstance(param, nn.Parameter)
+
+        pd = param.detach().to("meta")
+
+        if is_param:
+            pd = nn.Parameter(pd, requires_grad=False)
+        return pd
+
 
 class _make_target_param:
     def __init__(self, clone):
@@ -446,13 +456,3 @@ def __call__(self, x):
                 x.data.clone() if self.clone else x.data, requires_grad=False
             )
         return x.data.clone() if self.clone else x.data
-
-
-def _make_meta_params(param):
-    is_param = isinstance(param, nn.Parameter)
-
-    pd = param.detach().to("meta")
-
-    if is_param:
-        pd = nn.Parameter(pd, requires_grad=False)
-    return pd
diff --git a/torchrl/objectives/ddpg.py b/torchrl/objectives/ddpg.py
@@ -198,7 +198,7 @@ def __init__(
 
         actor_critic = ActorCriticWrapper(actor_network, value_network)
         params = TensorDict.from_module(actor_critic)
-        params_meta = params.detach().to("meta")
+        params_meta = params.apply(self._make_meta_params, device=torch.device("meta"))
         with params_meta.to_module(actor_critic):
             self.actor_critic = deepcopy(actor_critic)
 

diff --git a/torchrl/objectives/multiagent/qmixer.py b/torchrl/objectives/multiagent/qmixer.py
@@ -213,7 +213,9 @@ def __init__(
 
         global_value_network = SafeSequential(local_value_network, mixer_network)
         params = TensorDict.from_module(global_value_network)
-        with params.detach().to("meta").to_module(global_value_network):
+        with params.apply(
+            self._make_meta_params, device=torch.device("meta")
+        ).to_module(global_value_network):
             self.global_value_network = deepcopy(global_value_network)
 
         self.convert_to_functional(