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