| 9 | |
| 10 | class WandbCallback(BaseCallback): |
| 11 | def __init__(self, cfg, vec_env): |
| 12 | super(WandbCallback, self).__init__(verbose=1) |
| 13 | |
| 14 | # save_dir = Path.cwd() |
| 15 | # self._save_dir = save_dir |
| 16 | self._video_path = Path('video') |
| 17 | self._video_path.mkdir(parents=True, exist_ok=True) |
| 18 | self._ckpt_dir = Path('ckpt') |
| 19 | self._ckpt_dir.mkdir(parents=True, exist_ok=True) |
| 20 | |
| 21 | # wandb.init(project=cfg.wb_project, dir=save_dir, name=cfg.wb_runname) |
| 22 | wandb.init(project=cfg.wb_project, name=cfg.wb_name, notes=cfg.wb_notes, tags=cfg.wb_tags) |
| 23 | wandb.config.update(OmegaConf.to_container(cfg)) |
| 24 | |
| 25 | wandb.save('./config_agent.yaml') |
| 26 | wandb.save('.hydra/*') |
| 27 | |
| 28 | self.vec_env = vec_env |
| 29 | |
| 30 | self._eval_step = int(1e5) |
| 31 | self._buffer_step = int(1e5) |
| 32 | |
| 33 | def _init_callback(self): |
| 34 | self.n_epoch = 0 |