| 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 |