Small fix and improve logging message

This commit is contained in:
Cadene
2024-02-27 11:44:26 +00:00
parent 21670dce90
commit 7df542445c
5 changed files with 37 additions and 16 deletions

View File

@@ -5,6 +5,7 @@ import hydra
import imageio
import numpy as np
import torch
import tqdm
from tensordict.nn import TensorDictModule
from termcolor import colored
from torchrl.envs import EnvBase
@@ -32,7 +33,7 @@ def eval_policy(
max_rewards = []
successes = []
threads = []
for i in range(num_episodes):
for i in tqdm.tqdm(range(num_episodes)):
tensordict = env.reset()
ep_frames = []