text: # # Unity ML-Agents Toolkit
  from typing import Dict, Any, Optional, List
  import os
  import attr
  from mlagents.trainers.training_status import GlobalTrainingStatus, StatusType
  from mlagents_envs.logging_util import get_logger
  
  logger = get_logger(__name__)
  
  
  @attr.s(auto_attribs=True)
  class ModelCheckpoint:
      steps: int
      file_path: str
      reward: Optional[float]
      creation_time: float
  
  
  class ModelCheckpointManager:
      @staticmethod
      def get_checkpoints(behavior_name: str) -> List[Dict[str, Any]]:
          checkpoint_list = GlobalTrainingStatus.get_parameter_state(
              behavior_name, StatusType.CHECKPOINTS
          )
          if not checkpoint_list:
              checkpoint_list = []
              GlobalTrainingStatus.set_parameter_state(
                  behavior_name, StatusType.CHECKPOINTS, checkpoint_list
              )
          return checkpoint_list
  
      @staticmethod
      def remove_checkpoint(checkpoint: Dict[str, Any]) -> None:
          """
          Removes a checkpoint stored in checkpoint_list.
          If checkpoint cannot be found, no action is done.
  
          :param checkpoint: A checkpoint stored in checkpoint_list
          """
          file_path: str = checkpoint["file_path"]
          if os.path.exists(file_path):
              os.remove(file_path)
              logger.info(f"Removed checkpoint model {file_path}.")
          else:
              logger.info(f"Checkpoint at {file_path} could not be found.")
          return
  
      @classmethod
      def _cleanup_extra_checkpoints(
          cls, checkpoints: List[Dict], keep_checkpoints: int
      ) -> List[Dict]:
          """
          Ensures that the number of checkpoints stored are within the number
          of checkpoints the user defines. If the limit is hit, checkpoints are
          removed to create room for the next checkpoint to be inserted.
  
          :param behavior_name: The behavior name whose checkpoints we will mange.
          :param keep_checkpoints: Number of checkpoints to record (user-defined).
          """
          while len(checkpoints) > keep_checkpoints:
              if keep_checkpoints <= 0 or len(checkpoints) == 0:
                  break
              ModelCheckpointManager.remove_checkpoint(checkpoints.pop(0))
          return checkpoints
  
      @classmethod
      def add_checkpoint(
          cls, behavior_name: str, new_checkpoint: ModelCheckpoint, keep_checkpoints: int
      ) -> None:
          """
          Make room for new checkpoint if needed and insert new checkpoint information.
          :param behavior_name: Behavior name for the checkpoint.
          :param new_checkpoint: The new checkpoint to be recorded.
          :param keep_checkpoints: Number of checkpoints to record (user-defined).
          """
          new_checkpoint_dict = attr.asdict(new_checkpoint)
          checkpoints = cls.get_checkpoints(behavior_name)
          checkpoints.append(new_checkpoint_dict)
          cls._cleanup_extra_checkpoints(checkpoints, keep_checkpoints)
          GlobalTrainingStatus.set_parameter_state(
              behavior_name, StatusType.CHECKPOINTS, checkpoints
          )
  
      @classmethod
      def track_final_checkpoint(
          cls, behavior_name: str, final_checkpoint: ModelCheckpoint
      ) -> None:
          """
          Ensures number of checkpoints stored is within the max number of checkpoints
          defined by the user and finally stores the information about the final
          model (or intermediate model if training is interrupted).
          :param behavior_name: Behavior name of the model.
          :param final_checkpoint: Checkpoint information for the final model.
          """
          final_model_dict = attr.asdict(final_checkpoint)
          GlobalTrainingStatus.set_parameter_state(
              behavior_name, StatusType.FINAL_CHECKPOINT, final_model_dict
          )
  