text: from stable_baselines import PPO2
  from stable_baselines.common.policies import MlpPolicy
  from env.create_vec_env import create_vec_env
  from agents.load_ppo_params import load_ppo_params
  import os
  import time
  from stable_baselines.common.callbacks import BaseCallback, EveryNTimesteps
  import glob
  
  
  class SaveOnNumberTimeSteps(BaseCallback):
      """
      Callback for saving a model (the check is done every ``check_freq`` steps)
  
      :param log_dir: (str) Path to the folder where the model will be saved.
      :param verbose: (int)
      """
      def __init__(self, log_dir: str, verbose=1):
          super(SaveOnNumberTimeSteps, self).__init__(verbose)
          self.save_path = log_dir
  
      def _init_callback(self) -> None:
          os.makedirs(self.save_path, exist_ok=True)
          self.policy_number = 0
  
      def _on_step(self) -> bool:
          save_path_current = self.save_path + 'policy' + str(self.policy_number)
          self.model.save(save_path_current)
          self.policy_number += 1
          return True
  
  
  def train_ppo(replication=10, replication_start=0, save_freq=100000, time_steps=1000000, log_name=None, env=None):
      model_params = load_ppo_params()
      for _ in range(replication_start, replication):
          log_dir = "./logs/" + log_name + "/" + str(_) + "/"
          env.set_attr('filename', log_dir)
          model = PPO2(MlpPolicy, env, verbose=0, tensorboard_log=None, **model_params, seed=_)
          log_dir = "./logs/" + log_name + "/" + str(_) + "/" + "policy/"
          # Create the callback: save model every x steps
          checkpoint_on_event = SaveOnNumberTimeSteps(log_dir=log_dir)
          event_callback = EveryNTimesteps(n_steps=save_freq, callback=checkpoint_on_event)
          model.learn(total_timesteps=time_steps, tb_log_name="tensorboard", callback=event_callback)
          log_dir = log_dir + '/' + 'policy_last'
          model.save(log_dir)
          del model
          del event_callback
          del checkpoint_on_event
          del log_dir
          time.sleep(5)
      del model_params
  
  
  def evaluate_ppo(time_steps=1000000, num_envs=20, log_name=None, env=None, safe_states_eval=False, start=0):
      os.mkdir("./logs/" + log_name + "/evaluation/")
      paths = glob.glob(os.path.join('./logs/' + log_name, "*", "*", "policy_last*"))
      y = start
      num_steps = int(time_steps / num_envs)
      if safe_states_eval:
          env.set_attr('safe_states', True)
      for evaluate_replication in paths:
          log_dir = "./logs/" + log_name + "/evaluation/" + str(y) + "/"
          os.mkdir(log_dir)
          env.set_attr('filename', log_dir)
          # load_policy
          policy = PPO2.load(evaluate_replication)
          obs = env.reset()
          for i in range(num_steps):
              actions, _states = policy.predict(obs)
              obs, rewards, done, info = env.step(actions)
          del policy
          time.sleep(5)
          y += 1
      del paths
      del y
      del num_steps
  
  
  def benchmark_ppo(supervisor='baseline', safe_info=True, safe_states_training=False, safe_states_eval=False,
                    safety_distance_logical=.25, replication=10, replication_start=0, worker_start=1000, save_freq=100000,
                    time_steps_train=1000000, time_steps_eval=150000, lidar_num_bins=15, num_envs=20, cbf_quadratic=False,
                    cbf_learn_online=True, rewards=None, log_name=None, search_space='random', env=None,
                    scale_reward=False):
      if env is None:
          env = create_vec_env(visibility=17.5, safe_info=safe_info, safe_states=safe_states_training,
                               safety_distance_logical=safety_distance_logical, repl=replication_start,
                               worker_start=worker_start, supervisor=supervisor,
                               lidar_num_bins=lidar_num_bins, num_envs=num_envs, cbf_quadratic=cbf_quadratic,
                               cbf_learn_online=cbf_learn_online, rewards=rewards, log_name=log_name,
                               num_prior_data_cbf='2k', search_space=search_space, scale_reward=scale_reward)
      train_ppo(replication=replication, replication_start=replication_start, save_freq=save_freq,
                time_steps=time_steps_train, log_name=log_name, env=env)
      evaluate_ppo(time_steps=time_steps_eval, num_envs=num_envs, log_name=log_name, safe_states_eval=safe_states_eval,
                   env=env, start=replication_start)
      env.close()
  