diff --git a/torchrl/_utils.py b/torchrl/_utils.py index f2ce0cf520e..0b4dd03a636 100644 --- a/torchrl/_utils.py +++ b/torchrl/_utils.py @@ -117,10 +117,15 @@ def print(prefix=None) -> str: # noqa: T202 return "\n".join(string) @classmethod - def todict(cls, percall=True): + def todict(cls, percall=True, prefix=None): + def _make_key(key): + if prefix: + return f"{prefix}/{key}" + return key + if percall: - return {key: val[0] for key, val in cls._REG.items()} - return {key: val[1] for key, val in cls._REG.items()} + return {_make_key(key): val[0] for key, val in cls._REG.items()} + return {_make_key(key): val[1] for key, val in cls._REG.items()} @staticmethod def erase():