File size: 6,131 Bytes
87c8483 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 | import json
from pathlib import Path
from transformers import TrainerCallback
import wandb
from gear_sonic.trl.callbacks.im_eval_callback import create_html_table
from gear_sonic.trl.utils.common import wandb_run_exists
class ReadEvalCallback(TrainerCallback):
"""Callback to read evaluation metrics and log them to wandb."""
def __init__(self, eval_dir: str, check_interval: int = 1):
"""
Initialize the callback.
Args:
experiment_dir: Path to the experiment directory where eval results are stored
check_interval: How often to check for new evaluation results (in log steps)
"""
super().__init__()
self.eval_dir = Path(eval_dir)
self.check_interval = check_interval
self.last_check_step = 0
self.state = None
def on_log(self, args, state, control, logs=None, **kwargs):
del args, control, logs
return # Disable this callback for now. We will let eval script to handle the wandb logging.
def find_next_unread_checkpoint(self, current_eval_step: int = 0, mode: str = "metrics"):
"""
Find the next unread checkpoint folder that has finished evaluation.
Args:
current_eval_step: Current evaluation step to start searching from
Returns:
tuple: (eval_step, eval_step_dir) or (None, None) if not found
"""
if not self.eval_dir.exists():
return []
# Find all evaluation directories with 6-digit zero-padded names
eval_step_dirs = []
for d in self.eval_dir.iterdir():
if d.is_dir() and d.name.isdigit():
eval_step = int(d.name)
if eval_step > current_eval_step:
eval_step_dirs.append((eval_step, d))
if not eval_step_dirs:
return []
# Sort by step number to get the next one in sequence
eval_step_dirs.sort(key=lambda x: x[0])
# Find the first one that has finished evaluation
for eval_step, eval_step_dir in eval_step_dirs:
metrics_finish_file = eval_step_dir / "all_eval_finish.txt"
if metrics_finish_file.exists():
eval_dir_res = []
for f in eval_step_dir.iterdir():
if f.is_dir():
mode_finish_file = f / f"{mode}_finish.txt"
if mode_finish_file.exists():
eval_dir_res.append((eval_step, f))
return eval_dir_res
return []
def _check_and_log_eval_results(self, eval_step, eval_step_dir):
"""Check for new evaluation results and log them to wandb."""
metrics_file = eval_step_dir / "metrics_eval.json"
# Read and log the metrics
with open(metrics_file) as f:
try:
metrics_eval = json.load(f)
except json.JSONDecodeError:
print(f"Error loading metrics_eval.json for step {eval_step}")
return
if "log_keys" in metrics_eval:
self.log_keys = metrics_eval["log_keys"]
else:
self.log_keys = None
# Add eval_step if not already present
if "eval_step" not in metrics_eval:
metrics_eval["eval_step"] = eval_step
metrics_eval["eval/all_metrics_dict"][
"sampling_prob"
] = self.env._motion_lib._sampling_prob.cpu().numpy()
metrics_eval["eval/failed_metrics_dict"][
"sampling_prob"
] = self.env._motion_lib._sampling_prob.cpu().numpy()[metrics_eval["failed_idxes"]]
metrics_eval["eval/all_metrics_dict"] = create_html_table(
metrics_eval["eval/all_metrics_dict"]
)
metrics_eval["eval/failed_metrics_dict"] = create_html_table(
metrics_eval["eval/failed_metrics_dict"]
)
# Log to wandb
if self.accelerator.is_main_process and wandb_run_exists():
if self.log_keys is not None:
metrics_eval = {f"{self.log_keys}/{k}": v for k, v in metrics_eval.items()}
wandb.log(metrics_eval)
for key in ["failed_keys", "failed_idxes"]:
# ZL: why do we need to do this?
if key in metrics_eval:
del metrics_eval[key]
print(f"Logged evaluation metrics for step {eval_step} {eval_step_dir}")
def _check_and_log_eval_render_results(self, eval_step, eval_step_dir):
"""Check for new evaluation results and log them to wandb."""
video_dir = eval_step_dir / "render_results"
metrics_file = eval_step_dir / "metrics_eval.json"
if not video_dir.exists():
print(f"No render_results directory for step {eval_step}, skipping render logging")
return
# Read and log the metrics
with open(metrics_file) as f:
try:
metrics_eval = json.load(f)
except json.JSONDecodeError:
print(f"Error loading metrics_eval.json for step {eval_step}")
return
if "log_keys" in metrics_eval:
self.log_keys = metrics_eval["log_keys"]
else:
self.log_keys = None
video_files = []
for i, video_file in enumerate(sorted(video_dir.iterdir())):
if video_file.is_file() and video_file.name.endswith(".mp4"):
video_files.append((i, video_file))
if self.log_keys is not None:
wandb_videos = {
f"videos_hard_{self.log_keys}/{i: 04d}": wandb.Video(str(video_file), format="mp4")
for i, video_file in reversed(video_files)
}
else:
wandb_videos = {
f"videos_hard/{i: 04d}": wandb.Video(str(video_file), format="mp4")
for i, video_file in reversed(video_files)
}
wandb_videos["eval_step"] = eval_step
# Log to wandb
if self.accelerator.is_main_process and wandb_run_exists():
wandb.log(wandb_videos)
print(f"Logged rendered video for step {eval_step}")
|