Repository: pytorch/rl

Issue:
[BUG] Use state_dict instead of named_parameters for updating policy in subprocess?
Update policy weights do not consider variables that are registered as buffer? For example, when using batch normalization, there are variables like running_mean, but they won't show up on named_parameters(), so the update_policy_weights won't correctly update the policy in the subprocess when using the collectors?


Discussion:
Good point, we could do that! 
Let me fix this

Resolving patch:
diff --git a/torchrl/collectors/collectors.py b/torchrl/collectors/collectors.py
--- a/torchrl/collectors/collectors.py
+++ b/torchrl/collectors/collectors.py
@@ -195,17 +195,6 @@ def _get_policy_and_device(
             observation_spec (TensorSpec, optional): spec of the observations
 
         """
-        # if create_env_fn is not None:
-        #     if create_env_kwargs is None:
-        #         create_env_kwargs = {}
-        #     self.create_env_fn = create_env_fn
-        #     if isinstance(create_env_fn, EnvBase):
-        #         env = create_env_fn
-        #     else:
-        #         env = self.create_env_fn(**create_env_kwargs)
-        # else:
-        #     env = None
-
         if policy is None:
             if not hasattr(self, "env") or self.env is None:
                 raise ValueError(
@@ -271,15 +260,24 @@ def _get_policy_and_device(
             )
 
         device = torch.device(device) if device is not None else policy_device
-        if device is None:
-            device = torch.device("cpu")
         get_weights_fn = None
         if policy_device != device:
-            get_weights_fn = policy.state_dict
-            policy = deepcopy(policy).requires_grad_(False).to(device)
-            if device == torch.device("cpu"):
-                policy.share_memory()
-        return policy, device, get_weights_fn
+            param_and_buf = dict(policy.named_parameters())
+            param_and_buf.update(dict(policy.named_buffers()))
+
+            def get_weights_fn(param_and_buf=param_and_buf):
+                return TensorDict(param_and_buf, []).apply(lambda x: x.data)
+
+            policy_cast = deepcopy(policy).requires_grad_(False).to(device)
+            # here things may break bc policy.to("cuda") gives us weights on cuda:0 (same
+            # but different)
+            try:
+                device = next(policy_cast.parameters()).device
+            except StopIteration:  # noqa
+                pass
+        else:
+            policy_cast = policy
+        return policy_cast, device, get_weights_fn
 
     def update_policy_weights_(
         self, policy_weights: Optional[TensorDictBase] = None
@@ -809,7 +807,8 @@ def rollout(self) -> TensorDictBase:
                 if self._frames < self.init_random_frames:
                     self.env.rand_step(self._tensordict)
                 else:
-                    self.env.step(self.policy(self._tensordict))
+                    self.policy(self._tensordict)
+                    self.env.step(self._tensordict)
 
                 # we must clone all the values, since the step / traj_id updates are done in-place
                 try:
@@ -1125,9 +1124,9 @@ def device_err_msg(device_name, devices_list):
             )
             self._policy_dict[_device] = _policy
             if isinstance(_policy, nn.Module):
-                self._policy_weights_dict[_device] = TensorDict(
-                    dict(_policy.named_parameters()), []
-                )
+                param_dict = dict(_policy.named_parameters())
+                param_dict.update(_policy.named_buffers())
+                self._policy_weights_dict[_device] = TensorDict(param_dict, [])
             else:
                 self._policy_weights_dict[_device] = TensorDict({}, [])
 
@@ -1209,7 +1208,7 @@ def update_policy_weights_(self, policy_weights=None) -> None:
                     policy_weights
                 )
             elif self._get_weights_fn_dict[_device] is not None:
-                self._policy_dict[_device].load_state_dict(
+                self._policy_weights_dict[_device].update_(
                     self._get_weights_fn_dict[_device]()
                 )
 
@@ -1977,9 +1976,17 @@ def _main_async_collector(
 ) -> None:
     pipe_parent.close()
     # init variables that will be cleared when closing
-    tensordict = data = d = data_in = dc = dc_iter = None
-
-    dc = SyncDataCollector(
+    tensordict = data = d = data_in = inner_collector = dc_iter = None
+
+    # send the policy to device
+    try:
+        policy = policy.to(device)
+    except Exception:
+        warnings.warn(
+            "Couldn't cast the policy onto the desired device on remote process. "
+            "If your policy is not a nn.Module instance you can probably ignore this warning."
+        )
+    inner_collector = SyncDataCollector(
         create_env_fn,
         create_env_kwargs=create_env_kwargs,
         policy=policy,
@@ -1998,7 +2005,7 @@ def _main_async_collector(
     )
     if verbose:
         print("Sync data collector created")
-    dc_iter = iter(dc)
+    dc_iter = iter(inner_collector)
     j = 0
     pipe_child.send("instantiated")
 
@@ -2046,9 +2053,9 @@ def _main_async_collector(
                 continue
         if msg in ("continue", "continue_random"):
             if msg == "continue_random":
-                dc.init_random_frames = float("inf")
+                inner_collector.init_random_frames = float("inf")
             else:
-                dc.init_random_frames = -1
+                inner_collector.init_random_frames = -1
 
             d = next(dc_iter)
             if pipe_child.poll(_MIN_TIMEOUT):
@@ -2084,14 +2091,14 @@ def _main_async_collector(
                 continue
 
         elif msg == "update":
-            dc.update_policy_weights_()
+            inner_collector.update_policy_weights_()
             pipe_child.send((j, "updated"))
             has_timed_out = False
             continue
 
         elif msg == "seed":
             data_in, static_seed = data_in
-            new_seed = dc.set_seed(data_in, static_seed=static_seed)
+            new_seed = inner_collector.set_seed(data_in, static_seed=static_seed)
             torch.manual_seed(data_in)
             np.random.seed(data_in)
             pipe_child.send((new_seed, "seeded"))
@@ -2099,12 +2106,12 @@ def _main_async_collector(
             continue
 
         elif msg == "reset":
-            dc.reset()
+            inner_collector.reset()
             pipe_child.send((j, "reset"))
             continue
 
         elif msg == "state_dict":
-            state_dict = dc.state_dict()
+            state_dict = inner_collector.state_dict()
             # send state_dict to cpu first
             state_dict = recursive_map_to_cpu(state_dict)
             pipe_child.send((state_dict, "state_dict"))
@@ -2113,15 +2120,15 @@ def _main_async_collector(
 
         elif msg == "load_state_dict":
             state_dict = data_in
-            dc.load_state_dict(state_dict)
+            inner_collector.load_state_dict(state_dict)
             pipe_child.send((j, "loaded"))
             has_timed_out = False
             continue
 
         elif msg == "close":
             del tensordict, data, d, data_in
-            dc.shutdown()
-            del dc, dc_iter
+            inner_collector.shutdown()
+            del inner_collector, dc_iter
             pipe_child.send("closed")
             if verbose:
                 print(f"collector {idx} closed")


Test patch:
diff --git a/test/test_collector.py b/test/test_collector.py
--- a/test/test_collector.py
+++ b/test/test_collector.py
@@ -33,7 +33,7 @@
 )
 from torchrl.collectors.utils import split_trajectories
 from torchrl.data import CompositeSpec, UnboundedContinuousTensorSpec
-from torchrl.envs import EnvCreator, ParallelEnv, SerialEnv, StepCounter
+from torchrl.envs import EnvBase, EnvCreator, ParallelEnv, SerialEnv, StepCounter
 from torchrl.envs.libs.gym import _has_gym, GymEnv
 from torchrl.envs.transforms import TransformedEnv, VecNorm
 from torchrl.modules import Actor, LSTMNet, OrnsteinUhlenbeckProcessWrapper, SafeModule
@@ -1327,6 +1327,110 @@ def test_reset_heterogeneous_envs():
     ).all()
 
 
+@pytest.mark.skipif(not torch.cuda.device_count(), reason="No casting if no cuda")
+class TestUpdateParams:
+    class DummyEnv(EnvBase):
+        def __init__(self, device, batch_size=[]):  # noqa: B006
+            super().__init__(batch_size=batch_size, device=device)
+            self.state = torch.zeros(self.batch_size, device=device)
+            self.output_spec = CompositeSpec(
+                observation=CompositeSpec(
+                    state=UnboundedContinuousTensorSpec(shape=(), device=device)
+                )
+            )
+            self.action_spec = UnboundedContinuousTensorSpec(
+                shape=batch_size, device=device
+            )
+            self.reward_spec = UnboundedContinuousTensorSpec(
+                shape=(*batch_size, 1), device=device
+            )
+
+        def _step(
+            self,
+            tensordict,
+        ):
+            action = tensordict.get("action")
+            self.state += action
+            return TensorDict(
+                {
+                    "next": {
+                        "state": self.state.clone(),
+                        "reward": self.reward_spec.zero(),
+                        "done": self.done_spec.zero(),
+                    }
+                },
+                self.batch_size,
+            )
+
+        def _reset(self, tensordict=None):
+            self.state.zero_()
+            return TensorDict({"state": self.state.clone()}, self.batch_size)
+
+        def _set_seed(self, seed):
+            return seed
+
+    class Policy(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.param = nn.Parameter(torch.zeros(()))
+            self.register_buffer("buf", torch.zeros(()))
+            self.in_keys = []
+            self.out_keys = ["action"]
+
+        def forward(self, td):
+            td["action"] = (self.param + self.buf).expand(td.shape)
+            return td
+
+    @pytest.mark.parametrize(
+        "collector", [MultiSyncDataCollector, MultiaSyncDataCollector]
+    )
+    @pytest.mark.parametrize("give_weights", [True, False])
+    @pytest.mark.parametrize(
+        "policy_device,env_device",
+        [
+            ["cpu", "cuda"],
+            ["cuda", "cpu"],
+            ["cpu", "cuda:0"],
+            ["cuda:0", "cpu"],
+            ["cuda", "cuda:0"],
+            ["cuda:0", "cuda"],
+        ],
+    )
+    def test_param_sync(self, give_weights, collector, policy_device, env_device):
+        policy = TestUpdateParams.Policy().to(policy_device)
+
+        env = EnvCreator(lambda: TestUpdateParams.DummyEnv(device=env_device))
+        device = env().device
+        env = [env]
+        col = collector(
+            env, policy, device=device, total_frames=200, frames_per_batch=10
+        )
+        try:
+            for i, data in enumerate(col):
+                if i == 0:
+                    assert (data["action"] == 0).all()
+                    # update policy
+                    policy.param.data += 1
+                    policy.buf.data += 2
+                    if give_weights:
+                        d = dict(policy.named_parameters())
+                        d.update(policy.named_buffers())
+                        p_w = TensorDict(d, [])
+                    else:
+                        p_w = None
+                    col.update_policy_weights_(p_w)
+                elif i == 20:
+                    if (data["action"] == 1).all():
+                        raise RuntimeError("Failed to update buffer")
+                    elif (data["action"] == 2).all():
+                        raise RuntimeError("Failed to update params")
+                    elif (data["action"] == 0).all():
+                        raise RuntimeError("Failed to update params and buffers")
+                    assert (data["action"] == 3).all()
+        finally:
+            col.shutdown()
+
+
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()
     pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)