72 lines
2.2 KiB
Python
72 lines
2.2 KiB
Python
from pathlib import Path
|
|
|
|
import hydra
|
|
import imageio
|
|
import numpy as np
|
|
import torch
|
|
from termcolor import colored
|
|
|
|
from ..lib.envs import make_env
|
|
from ..lib.utils import set_seed
|
|
|
|
|
|
def eval_agent(
|
|
env, agent, num_episodes: int, save_video: bool = False, video_path: Path = None
|
|
):
|
|
"""Evaluate a trained agent and optionally save a video."""
|
|
if save_video:
|
|
assert video_path is not None
|
|
assert video_path.suffix == ".mp4"
|
|
episode_rewards = []
|
|
episode_successes = []
|
|
episode_lengths = []
|
|
for i in range(num_episodes):
|
|
obs, done, ep_reward, t = env.reset(), False, 0, 0
|
|
ep_success = False
|
|
if save_video:
|
|
frames = []
|
|
while not done:
|
|
action = agent.act(obs, t0=t == 0, eval_mode=True, step=step)
|
|
obs, reward, done, info = env.step(action.cpu().numpy())
|
|
ep_reward += reward
|
|
if "success" in info and info["success"]:
|
|
ep_success = True
|
|
if save_video:
|
|
frame = env.render(
|
|
mode="rgb_array",
|
|
# TODO(rcadene): make height, width, camera_id configurable
|
|
height=384,
|
|
width=384,
|
|
camera_id=0,
|
|
)
|
|
frames.append(frame)
|
|
t += 1
|
|
episode_rewards.append(float(ep_reward))
|
|
episode_successes.append(float(ep_success))
|
|
episode_lengths.append(t)
|
|
if save_video:
|
|
frames = np.stack(frames).transpose(0, 3, 1, 2)
|
|
video_path.parent.mkdir(parents=True, exist_ok=True)
|
|
# TODO(rcadene): make fps configurable
|
|
imageio.mimsave(video_path, frames, fps=15)
|
|
return {
|
|
"episode_reward": np.nanmean(episode_rewards),
|
|
"episode_success": np.nanmean(episode_successes),
|
|
"episode_length": np.nanmean(episode_lengths),
|
|
}
|
|
|
|
|
|
@hydra.main(version_base=None, config_name="default", config_path="../configs")
|
|
def eval(cfg: dict):
|
|
assert torch.cuda.is_available()
|
|
set_seed(cfg.seed)
|
|
print(colored("Log dir:", "yellow", attrs=["bold"]), cfg.log_dir)
|
|
|
|
env = make_env(cfg)
|
|
|
|
eval_metrics = eval_agent(env, agent, num_episodes=10, save_video=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
eval()
|