Skip to content

Commit 8fb55ef

Browse files
[Docs] InitTracker cleanup (#2330)
1 parent cdcd9f3 commit 8fb55ef

File tree

1 file changed

+5
-3
lines changed

1 file changed

+5
-3
lines changed

torchrl/envs/transforms/transforms.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff 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:

0 commit comments

Comments
 (0)