hangX/lerobot
1
0
Fork 0
forked from tangger/lerobot
Code Pull Requests Activity

eval.mp4 works!

Browse Source
This commit is contained in:
Cadene
2024-01-30 23:30:14 +00:00
parent 1144819c29
commit 1e52499490
4 changed files with 127 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

57
lerobot/scripts/eval.py
View File

@@ -4,10 +4,12 @@ import hydra
import imageio
import numpy as np
import torch
from tensordict import TensorDict
from termcolor import colored
from ..lib.envs import make_env
from ..lib.utils import set_seed
from lerobot.lib.envs.factory import make_env
from lerobot.lib.tdmpc import TDMPC
from lerobot.lib.utils import set_seed
def eval_agent(
@@ -21,32 +23,45 @@ def eval_agent(
episode_successes = []
episode_lengths = []
for i in range(num_episodes):
obs, done, ep_reward, t = env.reset(), False, 0, 0
td = env.reset()
obs = {}
obs["rgb"] = td["observation"]["camera"]
obs["state"] = td["observation"]["robot_state"]
done = False
ep_reward = 0
t = 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())
action = agent.act(obs, t0=t == 0, eval_mode=True, step=100000)
td = TensorDict({"action": action}, batch_size=[])
td = env.step(td)
reward = td["next", "reward"].item()
success = td["next", "success"].item()
done = td["next", "done"].item()
obs = {}
obs["rgb"] = td["next", "observation"]["camera"]
obs["state"] = td["next", "observation"]["robot_state"]
ep_reward += reward
if "success" in info and info["success"]:
if 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,
)
frame = env.render()
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)
frames = np.stack(frames) # .transpose(0, 3, 1, 2)
# TODO(rcadene): make fps configurable
imageio.mimsave(video_path, frames, fps=15)
return {
@@ -63,8 +78,20 @@ def eval(cfg: dict):
print(colored("Log dir:", "yellow", attrs=["bold"]), cfg.log_dir)
env = make_env(cfg)
agent = TDMPC(cfg)
# ckpt_path = "/home/rcadene/code/fowm/logs/xarm_lift/all/default/2/models/offline.pt"
ckpt_path = "/home/rcadene/code/fowm/logs/xarm_lift/all/default/2/models/final.pt"
agent.load(ckpt_path)
eval_metrics = eval_agent(env, agent, num_episodes=10, save_video=True)
eval_metrics = eval_agent(
env,
agent,
num_episodes=10,
save_video=True,
video_path=Path("tmp/2023_01_29_xarm_lift_final/eval.mp4"),
)
print(eval_metrics)
if __name__ == "__main__":