add tqdm (#101)

dongminlee94 · Aug 1, 2022 · 1724105 · 1724105
1 parent f25fd7d
commit 1724105
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/meta_rl/pearl/algorithm/meta_learner.py b/src/meta_rl/pearl/algorithm/meta_learner.py
@@ -11,6 +11,7 @@
 import torch
 from gym.envs.mujoco.half_cheetah import HalfCheetahEnv
 from torch.utils.tensorboard import SummaryWriter
+from tqdm import tqdm
 
 from meta_rl.pearl.algorithm.buffers import MultiTaskReplayBuffer
 from meta_rl.pearl.algorithm.sac import SAC
@@ -171,7 +172,8 @@ def meta_train(self) -> None:
 
  # 첫번째 반복단계에 한해 모든 메타-트레이닝 태스크에 대한 경로를 수집하여 리플레이 버퍼에 저장
  if iteration == 0:
- for index in self.train_tasks:
+ print("Collecting initial pool of data for train and eval")
+ for index in tqdm(self.train_tasks):
  self.env.reset_task(index)
  self.collect_train_data(
  task_index=index,