File tree Expand file tree Collapse file tree 1 file changed +5
-3
lines changed Expand file tree Collapse file tree 1 file changed +5
-3
lines changed Original file line number Diff line number Diff line change @@ -6420,6 +6420,7 @@ class InitTracker(Transform):
64206420
64216421 Args:
64226422 init_key (NestedKey, optional): the key to be used for the tracker entry.
6423+ In case of multiple _reset flags, this key is used as the leaf replacement for each.
64236424
64246425 Examples:
64256426 >>> from torchrl.envs.libs.gym import GymEnv
@@ -6433,11 +6434,12 @@ class InitTracker(Transform):
64336434
64346435 """
64356436
6436- def __init__ (self , init_key : NestedKey = "is_init" ):
6437+ def __init__ (self , init_key : str = "is_init" ):
64376438 if not isinstance (init_key , str ):
6438- raise ValueError ("init_key can only be of type str." )
6439+ raise ValueError (
6440+ "init_key can only be of type str as it will be the leaf key associated to each reset flag."
6441+ )
64396442 self .init_key = init_key
6440- self .reset_key = "_reset"
64416443 super ().__init__ ()
64426444
64436445 def set_container (self , container : Union [Transform , EnvBase ]) -> None :
You can’t perform that action at this time.
0 commit comments