Change internal eval callback to log to jsrl

steventango · Jul 20, 2023 · f85abf9 · f85abf9
1 parent f127c89
commit f85abf9
Showing 1 changed file with 15 additions and 10 deletions.
diff --git a/src/jsrl/jsrl.py b/src/jsrl/jsrl.py
@@ -164,22 +164,27 @@ def _init_callback(
             :return: A hybrid callback calling `callback` and performing evaluation.
             """
             callback = super()._init_callback(callback, progress_bar)
+            eval_callback = EvalCallback(
+                self.env,
+                callback_after_eval=JSRLAfterEvalCallback(
+                    self.policy,
+                    self.logger,
+                    verbose=self.verbose,
+                ),
+                eval_freq=self.policy.eval_freq,
+                n_eval_episodes=self.policy.n_eval_episodes,
+            )
             callback = CallbackList(
                 [
                     callback,
-                    EvalCallback(
-                        self.env,
-                        callback_after_eval=JSRLAfterEvalCallback(
-                            self.policy,
-                            self.logger,
-                            verbose=self.verbose,
-                        ),
-                        eval_freq=self.policy.eval_freq,
-                        n_eval_episodes=self.policy.n_eval_episodes,
-                    ),
+                    eval_callback,
                 ]
             )
             callback.init_callback(self)
+            default_record = eval_callback.logger.record
+            eval_callback.logger.record = lambda key, *args, **kwargs: default_record(
+                key.replace("eval/", "jsrl/"), *args, **kwargs
+            )
             return callback
 
         def predict(