[Refactor] Use _set_tuple for faster set (pytorch#1372)

Author: vmoens

torchrl/envs/utils.py

@@ -231,15 +231,15 @@ def _set_single_key(source, dest, key, clone=False):
             new_val = dest.get(k, None)
             if new_val is None:
                 new_val = _clone_no_keys(val)
-                dest.set(k, new_val)
-                # dest._set_tuple(k, new_val, inplace=False, validated=True)
+                # dest.set(k, new_val)
+                dest._set_str(k, new_val, inplace=False, validated=True)
             source = val
             dest = new_val
         else:
             if clone:
                 val = val.clone()
-            dest.set(k, val)
-            # dest._set_tuple(k, val, inplace=False, validated=True)
+            # dest.set(k, val)
+            dest._set_str(k, val, inplace=False, validated=True)
 
 
 def _set(source, dest, key, total_key, excluded):
@@ -257,13 +257,13 @@ def _set(source, dest, key, total_key, excluded):
                     _set(val, new_val, subkey, total_key, excluded) or non_empty_local
                 )
             if non_empty_local:
-                dest.set(key, new_val)
-                # dest._set_tuple(key, new_val, inplace=False, validated=True)
+                # dest.set(key, new_val)
+                dest._set_str(key, new_val, inplace=False, validated=True)
             non_empty = non_empty_local
         else:
             non_empty = True
-            dest.set(key, val)
-            # dest._set_tuple(key, val, inplace=False, validated=True)
+            # dest.set(key, val)
+            dest._set_str(key, val, inplace=False, validated=True)
     return non_empty
 
 

torchrl/envs/vec_env.py

@@ -20,6 +20,7 @@
 import torch
 
 from tensordict import TensorDict, unravel_key
+from tensordict._tensordict import _unravel_key_to_tuple
 from tensordict.tensordict import LazyStackedTensorDict, TensorDictBase
 from torch import multiprocessing as mp
 from torchrl._utils import _check_for_faulty_process, VERBOSE
@@ -761,16 +762,18 @@ def _step(self, tensordict: TensorDictBase) -> TensorDictBase:
         if self._single_task:
             # this is faster than update_ but won't work for lazy stacks
             for key in self.env_input_keys:
-                self.shared_tensordict_parent.set(
-                    key,
-                    tensordict.get(key),
-                    inplace=True,
-                )
-                # self.shared_tensordict_parent._set_tuple(
+                # self.shared_tensordict_parent.set(
                 #     key,
-                #     tensordict._get_tuple(key, None),
+                #     tensordict.get(key),
                 #     inplace=True,
                 # )
+                key = _unravel_key_to_tuple(key)
+                self.shared_tensordict_parent._set_tuple(
+                    key,
+                    tensordict._get_tuple(key, None),
+                    inplace=True,
+                    validated=True,
+                )
         else:
             self.shared_tensordict_parent.update_(
                 tensordict.select(*self.env_input_keys, strict=False)
@@ -1062,8 +1065,14 @@ def _run_worker_pipe_shared_mem(
             i += 1
             if local_tensordict is not None:
                 for key in env_input_keys:
-                    local_tensordict.set(key, shared_tensordict.get(key))
-                    # local_tensordict._set_tuple(key, shared_tensordict._get_tuple(key, None), inplace=False, validated=True)
+                    # local_tensordict.set(key, shared_tensordict.get(key))
+                    key = _unravel_key_to_tuple(key)
+                    local_tensordict._set_tuple(
+                        key,
+                        shared_tensordict._get_tuple(key, None),
+                        inplace=False,
+                        validated=True,
+                    )
             else:
                 local_tensordict = shared_tensordict.clone(recurse=False)
             local_tensordict = env._step(local_tensordict)

torchrl/objectives/common.py

@@ -412,8 +412,10 @@ def _param_getter(self, network_name):
                                 value_to_set = getattr(self, value_to_set).detach()
                             else:
                                 value_to_set = getattr(self, value_to_set)
-                        params.set(key, value_to_set)
-                        # params._set_tuple(key, value_to_set, inplace=False, validated=True)
+                        # params.set(key, value_to_set)
+                        params._set_tuple(
+                            key, value_to_set, inplace=False, validated=True
+                        )
                 return params
             else:
                 params = getattr(self, param_name)
@@ -448,9 +450,10 @@ def _target_param_getter(self, network_name):
                         value_to_set = getattr(
                             self, self.SEP.join(["_target_" + network_name, *key])
                         )
-                        # _set is faster bc is bypasses the checks
-                        target_params.set(key, value_to_set)
-                        # target_params._set_tuple(key, value_to_set, inplace=False, validated=True)
+                        # target_params.set(key, value_to_set)
+                        target_params._set_tuple(
+                            key, value_to_set, inplace=False, validated=True
+                        )
                 return target_params
             else:
                 params = getattr(self, param_name)