Repository: pytorch/rl

Issue:
[BUG] missing lock for vecnorm transform
## Describe the bug

VecNorm goes into racing issues because updates to the shared tensors are not locked.

## To Reproduce

VecNorm tests are flacky because of racing operations between workers.

## Suggested solution

A lock should be introduced in vecnorm. For context, vecnorm contains tensors that are placed in shared memory (or equivalently on cuda). To make sure that all vecnorm instances have access to the same shared tensor, we create a vecnorm copy in main using the `EnvCreator` class. This original instance [creates the buffers](https://github.com/facebookresearch/rl/blob/2f2105e2d7f7ae844b740af44a0853c5b8a1f032/torchrl/envs/transforms/transforms.py#L1824) that will store the running stats in shared memory. Then, we link each vecnorm instance on each process to the copy on main by [passing a `state_dict` to each subprocess](https://github.com/facebookresearch/rl/blob/2f2105e2d7f7ae844b740af44a0853c5b8a1f032/torchrl/envs/env_creator.py#L128) (note that [we use `get_extra_state`](https://github.com/facebookresearch/rl/blob/2f2105e2d7f7ae844b740af44a0853c5b8a1f032/torchrl/envs/transforms/transforms.py#L1888-L1896) to pass a TensorDict that contains all the shared tensors). Have a look at [VecNorm docstrings](https://github.com/facebookresearch/rl/blob/2f2105e2d7f7ae844b740af44a0853c5b8a1f032/torchrl/envs/transforms/transforms.py#L1845-L1854) and [tests](https://github.com/facebookresearch/rl/blob/2f2105e2d7f7ae844b740af44a0853c5b8a1f032/test/test_transforms.py#L186) for more info.
Ideally, the solution should pass / get the lock in `get_extra_state` and `set_extra_state`.

## Checklist

- [x] I have checked that there is no similar issue in the repo (**required**)
- [x] I have read the [documentation](https://github.com/facebookresearch/rl/tree/main/docs/) (**required**)
- [x] I have provided a minimal working example to reproduce the bug (**required**)


Discussion:

Resolving patch:
diff --git a/torchrl/collectors/collectors.py b/torchrl/collectors/collectors.py
--- a/torchrl/collectors/collectors.py
+++ b/torchrl/collectors/collectors.py
@@ -72,6 +72,8 @@ def recursive_map_to_cpu(dictionary: OrderedDict) -> OrderedDict:
             k: recursive_map_to_cpu(item)
             if isinstance(item, OrderedDict)
             else item.cpu()
+            if isinstance(item, torch.Tensor)
+            else item
             for k, item in dictionary.items()
         }
     )
diff --git a/torchrl/envs/transforms/transforms.py b/torchrl/envs/transforms/transforms.py
--- a/torchrl/envs/transforms/transforms.py
+++ b/torchrl/envs/transforms/transforms.py
@@ -5,6 +5,7 @@
 
 from __future__ import annotations
 
+import multiprocessing as mp
 from copy import deepcopy, copy
 from typing import Any, List, Optional, OrderedDict, Sequence, Union
 from warnings import warn
@@ -1707,6 +1708,7 @@ def __init__(
         self,
         keys_in: Optional[Sequence[str]] = None,
         shared_td: Optional[TensorDictBase] = None,
+        lock: mp.Lock = (mp.Lock()),
         decay: float = 0.9999,
         eps: float = 1e-4,
     ) -> None:
@@ -1732,10 +1734,14 @@ def __init__(
                         f"with keys {shared_td.keys()}"
                     )
 
+        self.lock = lock
         self.decay = decay
         self.eps = eps
 
     def _call(self, tensordict: TensorDictBase) -> TensorDictBase:
+        if self.lock is not None:
+            self.lock.acquire()
+
         for key in self.keys_in:
             if key not in tensordict.keys():
                 continue
@@ -1746,6 +1752,10 @@ def _call(self, tensordict: TensorDictBase) -> TensorDictBase:
             )
 
             tensordict.set_(key, new_val)
+
+        if self.lock is not None:
+            self.lock.release()
+
         return tensordict
 
     def _init(self, tensordict: TensorDictBase, key: str) -> None:
@@ -1885,11 +1895,20 @@ def build_td_for_shared_vecnorm(
             return td_select.memmap_()
         return td_select.share_memory_()
 
-    def get_extra_state(self) -> TensorDictBase:
-        return self._td
-
-    def set_extra_state(self, td: TensorDictBase) -> None:
-        if not td.is_shared():
+    def get_extra_state(self) -> OrderedDict:
+        return OrderedDict([("lock", self.lock), ("td", self._td)])
+
+    def set_extra_state(self, state: OrderedDict) -> None:
+        lock = state["lock"]
+        if lock is not None:
+            """
+            since locks can't be serialized, we have use cases for stripping them
+            for example in ParallelEnv, in which case keep the lock we already have
+            to avoid an updated tensor dict being sent between processes to erase locks
+            """
+            self.lock = lock
+        td = state["td"]
+        if td is not None and not td.is_shared():
             raise RuntimeError(
                 "Only shared tensordicts can be set in VecNorm transforms"
             )
diff --git a/torchrl/envs/vec_env.py b/torchrl/envs/vec_env.py
--- a/torchrl/envs/vec_env.py
+++ b/torchrl/envs/vec_env.py
@@ -10,6 +10,7 @@
 from copy import deepcopy
 from logging import warn
 from multiprocessing import connection
+from multiprocessing.synchronize import Lock as MpLock
 from time import sleep
 from typing import Callable, Optional, Sequence, Union, Any, List
 
@@ -798,6 +799,19 @@ def to(self, device: DEVICE_TYPING):
         return self
 
 
+def recursively_strip_locks_from_state_dict(state_dict: OrderedDict) -> OrderedDict:
+    return OrderedDict(
+        **{
+            k: recursively_strip_locks_from_state_dict(item)
+            if isinstance(item, OrderedDict)
+            else None
+            if isinstance(item, MpLock)
+            else item
+            for k, item in state_dict.items()
+        }
+    )
+
+
 def _run_worker_pipe_shared_mem(
     idx: int,
     parent_pipe: connection.Connection,
@@ -921,7 +935,7 @@ def _run_worker_pipe_shared_mem(
             child_pipe.send((msg, None))
 
         elif cmd == "state_dict":
-            state_dict = env.state_dict()
+            state_dict = recursively_strip_locks_from_state_dict(env.state_dict())
             msg = "state_dict"
             child_pipe.send((msg, state_dict))
 


Test patch:
diff --git a/test/test_collector.py b/test/test_collector.py
--- a/test/test_collector.py
+++ b/test/test_collector.py
@@ -483,8 +483,8 @@ def test_collector_vecnorm_envcreator():
 
     s = c.state_dict()
 
-    td1 = s["worker0"]["env_state_dict"]["worker3"]["_extra_state"].clone()
-    td2 = s["worker1"]["env_state_dict"]["worker0"]["_extra_state"].clone()
+    td1 = s["worker0"]["env_state_dict"]["worker3"]["_extra_state"]["td"].clone()
+    td2 = s["worker1"]["env_state_dict"]["worker0"]["_extra_state"]["td"].clone()
     assert (td1 == td2).all()
 
     next(c_iter)
@@ -492,8 +492,8 @@ def test_collector_vecnorm_envcreator():
 
     s = c.state_dict()
 
-    td3 = s["worker0"]["env_state_dict"]["worker3"]["_extra_state"].clone()
-    td4 = s["worker1"]["env_state_dict"]["worker0"]["_extra_state"].clone()
+    td3 = s["worker0"]["env_state_dict"]["worker3"]["_extra_state"]["td"].clone()
+    td4 = s["worker1"]["env_state_dict"]["worker0"]["_extra_state"]["td"].clone()
     assert (td3 == td4).all()
     assert (td1 != td4).any()
 
diff --git a/test/test_transforms.py b/test/test_transforms.py
--- a/test/test_transforms.py
+++ b/test/test_transforms.py
@@ -4,7 +4,6 @@
 # LICENSE file in the root directory of this source tree.
 import argparse
 from copy import copy
-from time import sleep
 
 import pytest
 import torch
@@ -47,118 +46,6 @@
 TIMEOUT = 10.0
 
 
-def _test_vecnorm_subproc(idx, queue_out: mp.Queue, queue_in: mp.Queue):
-    td = queue_in.get(timeout=TIMEOUT)
-    if _has_gym:
-        env = GymEnv("Pendulum-v1")
-    else:
-        env = ContinuousActionVecMockEnv()
-    t = VecNorm(shared_td=td)
-    env = TransformedEnv(env, t)
-    env.set_seed(1000 + idx)
-    tensordict = env.reset()
-    for _ in range(10):
-        env.rand_step(tensordict)
-    queue_out.put(True)
-    msg = queue_in.get(timeout=TIMEOUT)
-    assert msg == "all_done"
-    obs_sum = t._td.get("next_observation_sum").clone()
-    obs_ssq = t._td.get("next_observation_ssq").clone()
-    obs_count = t._td.get("next_observation_ssq").clone()
-    reward_sum = t._td.get("reward_sum").clone()
-    reward_ssq = t._td.get("reward_ssq").clone()
-    reward_count = t._td.get("reward_ssq").clone()
-
-    td_out = TensorDict(
-        {
-            "obs_sum": obs_sum,
-            "obs_ssq": obs_ssq,
-            "obs_count": obs_count,
-            "reward_sum": reward_sum,
-            "reward_ssq": reward_ssq,
-            "reward_count": reward_count,
-        },
-        [],
-    ).share_memory_()
-    queue_out.put(td_out)
-    msg = queue_in.get(timeout=TIMEOUT)
-    assert msg == "all_done"
-    queue_in.close()
-    queue_out.close()
-    del queue_in, queue_out
-
-
-@pytest.mark.parametrize("nprc", [2, 5])
-def test_vecnorm_parallel(nprc):
-    queues = []
-    prcs = []
-    if _has_gym:
-        td = VecNorm.build_td_for_shared_vecnorm(GymEnv("Pendulum-v1"))
-    else:
-        td = VecNorm.build_td_for_shared_vecnorm(
-            ContinuousActionVecMockEnv(),
-        )
-    for idx in range(nprc):
-        prc_queue_in = mp.Queue(1)
-        prc_queue_out = mp.Queue(1)
-        p = mp.Process(
-            target=_test_vecnorm_subproc,
-            args=(
-                idx,
-                prc_queue_in,
-                prc_queue_out,
-            ),
-        )
-        p.start()
-        prc_queue_out.put(td)
-        prcs.append(p)
-        queues.append((prc_queue_in, prc_queue_out))
-
-    dones = [queue[0].get(timeout=TIMEOUT) for queue in queues]
-    assert all(dones)
-    msg = "all_done"
-    for idx in range(nprc):
-        queues[idx][1].put(msg)
-
-    obs_sum = td.get("next_observation_sum").clone()
-    obs_ssq = td.get("next_observation_ssq").clone()
-    obs_count = td.get("next_observation_ssq").clone()
-    reward_sum = td.get("reward_sum").clone()
-    reward_ssq = td.get("reward_ssq").clone()
-    reward_count = td.get("reward_ssq").clone()
-
-    for idx in range(nprc):
-        td_out = queues[idx][0].get(timeout=TIMEOUT)
-        _obs_sum = td_out.get("obs_sum")
-        _obs_ssq = td_out.get("obs_ssq")
-        _obs_count = td_out.get("obs_count")
-        _reward_sum = td_out.get("reward_sum")
-        _reward_ssq = td_out.get("reward_ssq")
-        _reward_count = td_out.get("reward_count")
-        assert (obs_sum == _obs_sum).all()
-        assert (obs_ssq == _obs_ssq).all()
-        assert (obs_count == _obs_count).all()
-        assert (reward_sum == _reward_sum).all()
-        assert (reward_ssq == _reward_ssq).all()
-        assert (reward_count == _reward_count).all()
-
-        obs_sum, obs_ssq, obs_count, reward_sum, reward_ssq, reward_count = (
-            _obs_sum,
-            _obs_ssq,
-            _obs_count,
-            _reward_sum,
-            _reward_ssq,
-            _reward_count,
-        )
-
-    msg = "all_done"
-    for idx in range(nprc):
-        queues[idx][1].put(msg)
-    del queues
-    for p in prcs:
-        p.join()
-
-
 def _test_vecnorm_subproc_auto(idx, make_env, queue_out: mp.Queue, queue_in: mp.Queue):
     env = make_env()
     env.set_seed(1000 + idx)
@@ -215,15 +102,13 @@ def test_vecnorm_parallel_auto(nprc):
         prcs.append(p)
         queues.append((prc_queue_in, prc_queue_out))
 
-    td = list(make_env.state_dict().values())[0]
     dones = [queue[0].get() for queue in queues]
     assert all(dones)
     msg = "all_done"
     for idx in range(nprc):
         queues[idx][1].put(msg)
 
-    sleep(0.01)  # Further fix for flacky test: vecnorm should have a locking
-    # mechanism
+    td = make_env.state_dict()["_extra_state"]["td"]
 
     obs_sum = td.get("next_observation_sum").clone()
     obs_ssq = td.get("next_observation_ssq").clone()
@@ -303,7 +188,7 @@ def test_parallelenv_vecnorm():
     parallel_sd = parallel_env.state_dict()
     assert "worker0" in parallel_sd
     worker_sd = parallel_sd["worker0"]
-    td = list(worker_sd.values())[0]
+    td = worker_sd["_extra_state"]["td"]
     queue_out.put("start")
     msg = queue_in.get(timeout=TIMEOUT)
     assert msg == "first round"