from stable_baselines3.common.callbacks import BaseCallback

class SuccessRateCallback(BaseCallback):
    def __init__(self, verbose=0):
        super(SuccessRateCallback, self).__init__(verbose)
        self.unsafe_count = 0  

    def _on_step(self):
        if self.locals["dones"][0]:
            if not self.locals["infos"][0]["is_success"]:
                self.unsafe_count += 1
            
            self.logger.record("unsafe_episodes", self.unsafe_count)
        
        return True
    
    