How to modify main.py to output solution schedule? about rl-job-shop-scheduling HOT 4 CLOSED

class CustomCallbacks(DefaultCallbacks):
    
    def __init__(self, legacy_callbacks_dict: Dict[str, callable] = None):
        super(CustomCallbacks, self).__init__(legacy_callbacks_dict)
        self.best_solution_makespan = float('inf')
        self.best_solution = None

    def on_episode_end(self, worker: "RolloutWorker", base_env: BaseEnv,
                       policies: Dict[PolicyID, Policy],
                       episode: MultiAgentEpisode, **kwargs):
        env = base_env.get_unwrapped()[0]
        if env.last_time_step != float('inf'):
            episode.custom_metrics['make_span'] = env.last_time_step
            if env.last_time_step < self.best_solution_makespan:
                self.best_solution_makespan = env.last_time_step
                self.best_solution = env.solution

print(f'best solution makespan {config['callbacks'].best_solution_makespan}')
print(f'best solution {config['callbacks'].best_solution}')

trainer = PPOTrainer(config=config)
while start_time + stop['time_total_s'] > time.time():
        result = trainer.train()

        result = wandb_tune._clean_log(result)
        log, config_update = _handle_result(result)
        wandb.log(log)
        # wandb.config.update(config_update, allow_val_change=True)

ray.shutdown()
print(f'best solution makespan {config['callbacks'].best_solution_makespan}')
print(f'best solution {config['callbacks'].best_solution}')