(self)
| 46 | pass |
| 47 | |
| 48 | def _on_training_end(self) -> None: |
| 49 | print(f'n_epoch: {self.n_epoch}, num_timesteps: {self.model.num_timesteps}') |
| 50 | # save time |
| 51 | time_elapsed = time.time() - self.model.start_time |
| 52 | wandb.log({ |
| 53 | 'time/n_epoch': self.n_epoch, |
| 54 | 'time/sec_per_epoch': time_elapsed / (self.n_epoch+1), |
| 55 | 'time/fps': (self.model.num_timesteps-self.model.start_num_timesteps) / time_elapsed, |
| 56 | 'time/train': self.model.t_train, |
| 57 | 'time/train_values': self.model.t_train_values, |
| 58 | 'time/rollout': self.model.t_rollout |
| 59 | }, step=self.model.num_timesteps) |
| 60 | wandb.log(self.model.train_debug, step=self.model.num_timesteps) |
| 61 | |
| 62 | # evaluate and save checkpoint |
| 63 | if (self.model.num_timesteps - self._last_time_eval) >= self._eval_step: |
| 64 | self._last_time_eval = self.model.num_timesteps |
| 65 | # evaluate |
| 66 | eval_video_path = (self._video_path / f'eval_{self.model.num_timesteps}.mp4').as_posix() |
| 67 | avg_ep_stat, ep_events = self.evaluate_policy(self.vec_env, self.model.policy, eval_video_path) |
| 68 | # log to wandb |
| 69 | wandb.log({f'video/{self.model.num_timesteps}': wandb.Video(eval_video_path)}, |
| 70 | step=self.model.num_timesteps) |
| 71 | wandb.log(avg_ep_stat, step=self.model.num_timesteps) |
| 72 | # save events |
| 73 | # eval_json_path = (video_path / f'event_{self.model.num_timesteps}.json').as_posix() |
| 74 | # with open(eval_json_path, 'w') as fd: |
| 75 | # json.dump(ep_events, fd, indent=4, sort_keys=False) |
| 76 | # wandb.save(eval_json_path) |
| 77 | |
| 78 | ckpt_path = (self._ckpt_dir / f'ckpt_{self.model.num_timesteps}.pth').as_posix() |
| 79 | self.model.save(ckpt_path) |
| 80 | wandb.save(f'./{ckpt_path}') |
| 81 | self.n_epoch += 1 |
| 82 | |
| 83 | # CONFIGHACK: curriculum |
| 84 | # num_zombies = {} |
| 85 | # for i in range(self.vec_env.num_envs): |
| 86 | # env_all_tasks = self.vec_env.get_attr('all_tasks',indices=i)[0] |
| 87 | # num_zombies[f'train/n_veh/{i}'] = env_all_tasks[0]['num_zombie_vehicles'] |
| 88 | # num_zombies[f'train/n_ped/{i}'] = env_all_tasks[0]['num_zombie_walkers'] |
| 89 | # if wandb.config['curriculum']: |
| 90 | # if avg_ep_stat['eval/route_completed_in_km'] > 1.0: |
| 91 | # # and avg_ep_stat['eval/red_light']>0: |
| 92 | # for env_task in env_all_tasks: |
| 93 | # env_task['num_zombie_vehicles'] += 10 |
| 94 | # env_task['num_zombie_walkers'] += 10 |
| 95 | # self.vec_env.set_attr('all_tasks', env_all_tasks, indices=i) |
| 96 | |
| 97 | # wandb.log(num_zombies, step=self.model.num_timesteps) |
| 98 | |
| 99 | def _on_rollout_end(self): |
| 100 | wandb.log({'time/rollout': self.model.t_rollout}, step=self.model.num_timesteps) |
nothing calls this directly
no test coverage detected