Re-introducing recorder_log_keys (#225)

Author: vmoens

torchrl/collectors/collectors.py

@@ -40,7 +40,7 @@
 
 _TIMEOUT = 1.0
 _MIN_TIMEOUT = 1e-3  # should be several orders of magnitude inferior wrt time spent collecting a trajectory
-_MAX_IDLE_COUNT = int(os.environ.get("MAX_IDLE_COUNT", 10))
+_MAX_IDLE_COUNT = int(os.environ.get("MAX_IDLE_COUNT", 1000))
 
 
 class RandomPolicy:

torchrl/trainers/helpers/recorder.py

@@ -3,7 +3,8 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
+from typing import Any
 
 
 @dataclass
@@ -19,3 +20,5 @@ class RecorderConfig:
     # number of batch collections in between two collections of validation rollouts. Default=1000.
     record_frames: int = 1000
     # number of steps in validation rollouts. " "Default=1000.
+    recorder_log_keys: Any = field(default_factory=lambda: ["reward"])
+    # Keys to log in the recorder

torchrl/trainers/helpers/trainers.py

@@ -248,6 +248,7 @@ def make_trainer(
             policy_exploration=policy_exploration,
             recorder=recorder,
             record_interval=cfg.record_interval,
+            log_keys=cfg.recorder_log_keys,
         )
         trainer.register_op(
             "post_steps_log",
@@ -262,7 +263,7 @@ def make_trainer(
             record_interval=cfg.record_interval,
             exploration_mode="random",
             suffix="exploration",
-            out_key="r_evaluation_exploration",
+            out_keys={"reward": "r_evaluation_exploration"},
         )
         trainer.register_op(
             "post_steps_log",