158 lines
4.4 KiB
Python
158 lines
4.4 KiB
Python
import logging
|
|
import threading
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import hydra
|
|
import imageio
|
|
import numpy as np
|
|
import torch
|
|
import tqdm
|
|
from tensordict.nn import TensorDictModule
|
|
from torchrl.envs import EnvBase
|
|
|
|
from lerobot.common.datasets.factory import make_offline_buffer
|
|
from lerobot.common.envs.factory import make_env
|
|
from lerobot.common.logger import log_output_dir
|
|
from lerobot.common.policies.factory import make_policy
|
|
from lerobot.common.utils import init_logging, set_seed
|
|
|
|
|
|
def write_video(video_path, stacked_frames, fps):
|
|
imageio.mimsave(video_path, stacked_frames, fps=fps)
|
|
|
|
|
|
def eval_policy(
|
|
env: EnvBase,
|
|
policy: TensorDictModule = None,
|
|
num_episodes: int = 10,
|
|
max_steps: int = 30,
|
|
save_video: bool = False,
|
|
video_dir: Path = None,
|
|
fps: int = 15,
|
|
return_first_video: bool = False,
|
|
):
|
|
start = time.time()
|
|
sum_rewards = []
|
|
max_rewards = []
|
|
successes = []
|
|
threads = []
|
|
for i in tqdm.tqdm(range(num_episodes)):
|
|
tensordict = env.reset()
|
|
|
|
ep_frames = []
|
|
|
|
if save_video or (return_first_video and i == 0):
|
|
|
|
def rendering_callback(env, td=None):
|
|
ep_frames.append(env.render()) # noqa: B023
|
|
|
|
# render first frame before rollout
|
|
rendering_callback(env)
|
|
else:
|
|
rendering_callback = None
|
|
|
|
with torch.inference_mode():
|
|
rollout = env.rollout(
|
|
max_steps=max_steps,
|
|
policy=policy,
|
|
callback=rendering_callback,
|
|
auto_reset=False,
|
|
tensordict=tensordict,
|
|
auto_cast_to_device=True,
|
|
)
|
|
# print(", ".join([f"{x:.3f}" for x in rollout["next", "reward"][:,0].tolist()]))
|
|
ep_sum_reward = rollout["next", "reward"].sum()
|
|
ep_max_reward = rollout["next", "reward"].max()
|
|
ep_success = rollout["next", "success"].any()
|
|
sum_rewards.append(ep_sum_reward.item())
|
|
max_rewards.append(ep_max_reward.item())
|
|
successes.append(ep_success.item())
|
|
|
|
if save_video or (return_first_video and i == 0):
|
|
stacked_frames = np.stack(ep_frames)
|
|
|
|
if save_video:
|
|
video_dir.mkdir(parents=True, exist_ok=True)
|
|
video_path = video_dir / f"eval_episode_{i}.mp4"
|
|
thread = threading.Thread(
|
|
target=write_video,
|
|
args=(str(video_path), stacked_frames, fps),
|
|
)
|
|
thread.start()
|
|
threads.append(thread)
|
|
|
|
if return_first_video and i == 0:
|
|
first_video = stacked_frames.transpose(0, 3, 1, 2)
|
|
|
|
for thread in threads:
|
|
thread.join()
|
|
|
|
info = {
|
|
"avg_sum_reward": np.nanmean(sum_rewards),
|
|
"avg_max_reward": np.nanmean(max_rewards),
|
|
"pc_success": np.nanmean(successes) * 100,
|
|
"eval_s": time.time() - start,
|
|
"eval_ep_s": (time.time() - start) / num_episodes,
|
|
}
|
|
if return_first_video:
|
|
return info, first_video
|
|
return info
|
|
|
|
|
|
@hydra.main(version_base=None, config_name="default", config_path="../configs")
|
|
def eval_cli(cfg: dict):
|
|
eval(cfg, out_dir=hydra.core.hydra_config.HydraConfig.get().runtime.output_dir)
|
|
|
|
|
|
def eval(cfg: dict, out_dir=None):
|
|
if out_dir is None:
|
|
raise NotImplementedError()
|
|
|
|
init_logging()
|
|
|
|
if cfg.device == "cuda":
|
|
assert torch.cuda.is_available()
|
|
else:
|
|
logging.warning("Using CPU, this will be slow.")
|
|
|
|
torch.backends.cudnn.benchmark = True
|
|
torch.backends.cuda.matmul.allow_tf32 = True
|
|
set_seed(cfg.seed)
|
|
|
|
log_output_dir(out_dir)
|
|
|
|
logging.info("make_offline_buffer")
|
|
offline_buffer = make_offline_buffer(cfg)
|
|
|
|
logging.info("make_env")
|
|
env = make_env(cfg, transform=offline_buffer._transform)
|
|
|
|
if cfg.policy.pretrained_model_path:
|
|
policy = make_policy(cfg)
|
|
policy = TensorDictModule(
|
|
policy,
|
|
in_keys=["observation", "step_count"],
|
|
out_keys=["action"],
|
|
)
|
|
else:
|
|
# when policy is None, rollout a random policy
|
|
policy = None
|
|
|
|
metrics = eval_policy(
|
|
env,
|
|
policy=policy,
|
|
save_video=True,
|
|
video_dir=Path(out_dir) / "eval",
|
|
fps=cfg.env.fps,
|
|
max_steps=cfg.env.episode_length // cfg.n_action_steps,
|
|
num_episodes=cfg.eval_episodes,
|
|
)
|
|
print(metrics)
|
|
|
|
logging.info("End of eval")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
eval_cli()
|