File size: 840 Bytes
037569d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 | from stable_baselines3.common.callbacks import BaseCallback
class MetricsCallback(BaseCallback):
def __init__(self, verbose=0):
super(MetricsCallback, self).__init__(verbose)
self.success_count = 0
self.episode_count = 0
def _on_step(self) -> bool:
if self.locals["dones"][0]:
self.episode_count += 1
info = self.locals["infos"][0]
if "distance" in info:
self.logger.record("custom/final_distance", info["distance"])
if "is_success" in info:
if info["is_success"]:
self.success_count += 1
success_rate = self.success_count / self.episode_count
self.logger.record("custom/success_rate", success_rate)
return True |