wangzixiang
/
lerobot
mirror of https://github.com/huggingface/lerobot.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix end-to-end aloha

This commit is contained in:
Cadene 2024-04-10 14:26:30 +00:00
parent a18bcb39a7
commit 48ec479660
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
lerobot/scripts/eval.py
View File

@ -249,7 +249,7 @@ def eval_policy(
threads.append(thread) threads.append(thread)
episode_counter += 1 episode_counter += 1
videos = batch_stacked_frames.transpose(0, 3, 1, 2) videos = einops.rearrange(batch_stacked_frames, "b t h w c -> b t c h w")
for thread in threads: for thread in threads:
thread.join() thread.join()
@ -328,6 +328,9 @@ def eval(cfg: dict, out_dir=None, stats_path=None):
# Save info # Save info
with open(Path(out_dir) / "eval_info.json", "w") as f: with open(Path(out_dir) / "eval_info.json", "w") as f:
# remove pytorch tensors which are not serializable to save the evaluation results only
del info["episodes"]
del info["videos"]
json.dump(info, f, indent=2) json.dump(info, f, indent=2)
logging.info("End of eval") logging.info("End of eval")