[BugFix] Make VecNorm Transform pickable (#1596)

Author: albertbou92

test/test_transforms.py

@@ -6,6 +6,7 @@
 import argparse
 
 import itertools
+import pickle
 import sys
 from copy import copy
 from functools import partial
@@ -7327,6 +7328,15 @@ def test_vecnorm_rollout(self, parallel, thr=0.2, N=200):
             env_t.close()
         self.SEED = 0
 
+    def test_pickable(self):
+
+        transform = VecNorm()
+        serialized = pickle.dumps(transform)
+        transform2 = pickle.loads(serialized)
+        assert transform.__dict__.keys() == transform2.__dict__.keys()
+        for key in sorted(transform.__dict__.keys()):
+            assert isinstance(transform.__dict__[key], type(transform2.__dict__[key]))
+
 
 def test_added_transforms_are_in_eval_mode_trivial():
     base_env = ContinuousActionVecMockEnv()

torchrl/envs/transforms/transforms.py

@@ -11,7 +11,7 @@
 from copy import copy
 from functools import wraps
 from textwrap import indent
-from typing import Any, List, Optional, OrderedDict, Sequence, Tuple, Union
+from typing import Any, Dict, List, Optional, OrderedDict, Sequence, Tuple, Union
 
 import numpy as np
 
@@ -4337,6 +4337,20 @@ def __repr__(self) -> str:
             f"eps={self.eps:4.4f}, keys={self.in_keys})"
         )
 
+    def __getstate__(self) -> Dict[str, Any]:
+        state = self.__dict__.copy()
+        _lock = state.pop("lock", None)
+        if _lock is not None:
+            state["lock_placeholder"] = None
+        return state
+
+    def __setstate__(self, state: Dict[str, Any]):
+        if "lock_placeholder" in state:
+            state.pop("lock_placeholder")
+            _lock = mp.Lock()
+            state["lock"] = _lock
+        self.__dict__.update(state)
+
 
 class RewardSum(Transform):
     """Tracks episode cumulative rewards.