| 333 | self.num_training_map_level += 1 |
| 334 | |
| 335 | def save(self): |
| 336 | |
| 337 | save_dir = f'./pytorch_models/{current_time}' |
| 338 | os.makedirs(save_dir, exist_ok=True) |
| 339 | |
| 340 | |
| 341 | torch.save(self.policy_net.state_dict(), os.path.join(save_dir, 'policy_net.pth')) |
| 342 | torch.save(self.value_net.state_dict(), os.path.join(save_dir, 'value_net.pth')) |
| 343 | torch.save(self.Q_net1.state_dict(), os.path.join(save_dir, 'Q_net1.pth')) |
| 344 | torch.save(self.Q_net2.state_dict(), os.path.join(save_dir, 'Q_net2.pth')) |
| 345 | |
| 346 | print("====================================") |
| 347 | print(f"Models have been saved in {save_dir}...") |
| 348 | print("====================================") |
| 349 | |
| 350 | def load(self): |
| 351 | self.policy_net.load_state_dict(torch.load('./pytorch_models/policy_net.pth')) |