GR00T-WholeBodyControl/gear_sonic/eval_exp.py

#!/usr/bin/env python3  # noqa: EXE001
# Copyright 2025 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import glob
import itertools
import json
import os
from pathlib import Path
import signal
import subprocess
import sys
import time
import hydra
from loguru import logger
import omegaconf
import wandb
import yaml

from gear_sonic.trl.callbacks import im_eval_callback
from gear_sonic.utils import config_utils

config_utils.register_rl_resolvers()


class CheckpointEvaluator:
    """Continuously monitors an experiment directory for new checkpoints and evaluates them sequentially."""

    def __init__(self, config):
        self.config = config
        self.experiment_dir = Path(config.experiment_dir)
        self.evaluated_checkpoints: set[str] = set()
        self.shutdown_flag = False
        self.last_evaluation_time = time.time()
        self.evaluation_timeout = config.get("evaluation_timeout", 24 * 3600)
        self.eval_frequency = config.get("eval_frequency", None)
        self.eval_last_n = config.get("eval_last_n", None)

        if not self.experiment_dir.exists():
            raise ValueError(f"Experiment directory does not exist: {self.experiment_dir}")

        self.find_evaluated_checkpoints()

        self.wandb_run_id = None
        self.wandb_project = None
        self.wandb_entity = None
        self._load_wandb_config()

        logger.info(f"Monitoring experiment directory: {self.experiment_dir}")
        logger.info(f"Scan interval: {config.scan_interval} seconds")
        logger.info(f"Evaluation timeout: {self.evaluation_timeout / 3600:.1f} hours")
        if self.wandb_run_id:
            logger.info(f"Wandb logging enabled: run_id={self.wandb_run_id}")

        self._backfill_wandb()

        signal.signal(signal.SIGINT, self._signal_handler)
        signal.signal(signal.SIGTERM, self._signal_handler)

    def _signal_handler(self, signum, frame):  # noqa: ARG002
        logger.info("Received shutdown signal. Stopping checkpoint monitoring...")
        self.shutdown_flag = True

    def _load_wandb_config(self):
        meta_path = self.experiment_dir / "meta.yaml"
        if meta_path.exists():
            with open(meta_path) as f:
                meta = yaml.safe_load(f)
            self.wandb_run_id = meta.get("wandb_run")

        config_path = self.experiment_dir / ".hydra" / "config.yaml"
        if not config_path.exists():
            config_path = self.experiment_dir / "config.yaml"
        if config_path.exists():
            try:
                with open(config_path) as f:
                    train_config = yaml.safe_load(f)
                wandb_cfg = train_config.get("wandb", {})
                self.wandb_project = train_config.get("project_name", "TRL_G1_Track")
                self.wandb_entity = wandb_cfg.get("wandb_entity", None)
            except Exception as e:  # noqa: BLE001
                logger.warning(f"Could not load training config for wandb: {e}")

        if self.wandb_project is None:
            self.wandb_project = "TRL_G1_Track"
        if self.wandb_entity is None:
            self.wandb_entity = None  # uses wandb default entity

    def _get_wandb_logged_steps(self) -> set[int]:
        logged_steps = set()
        if not self.wandb_run_id:
            return logged_steps

        try:
            api = wandb.Api(timeout=30)
            run = api.run(f"{self.wandb_entity}/{self.wandb_project}/{self.wandb_run_id}")
            hist = run.scan_history(
                keys=["eval/success/success_rate", "eval_step"],
                min_step=0,
                page_size=10000,
            )
            for row in hist:
                if row.get("eval/success/success_rate") is not None:
                    step = row.get("eval_step")
                    if step is not None:
                        logged_steps.add(int(step))
        except Exception as e:  # noqa: BLE001
            logger.warning(f"Could not query wandb for logged eval steps: {e}")

        return logged_steps

    def _backfill_wandb(self):
        if not self.wandb_run_id:
            self._load_wandb_config()
            if not self.wandb_run_id:
                logger.info("No wandb run ID available, skipping backfill")
                return

        eval_dir = self.experiment_dir / "eval"
        if not eval_dir.exists():
            return

        completed_steps = []
        for eval_subdir in sorted(eval_dir.iterdir()):
            if not eval_subdir.is_dir():
                continue
            try:
                step_num = int(eval_subdir.name)
            except ValueError:
                continue
            if (eval_subdir / "all_eval_finish.txt").exists():
                completed_steps.append((step_num, str(eval_subdir)))

        if not completed_steps:
            logger.info("No completed eval steps found on disk, nothing to backfill")
            return

        logged_steps = self._get_wandb_logged_steps()
        missing = [(step, path) for step, path in completed_steps if step not in logged_steps]

        if not missing:
            logger.info(
                f"All {len(completed_steps)} eval steps already logged to wandb, no backfill needed"
            )
            return

        logger.info(f"Backfilling {len(missing)}/{len(completed_steps)} eval steps to wandb")
        for eval_step, checkpoint_work_dir in missing:
            self._log_eval_to_wandb(eval_step, checkpoint_work_dir)

        logger.info("Backfill complete")

    def _log_eval_to_wandb(self, eval_step: int, checkpoint_work_dir: str):
        if not self.wandb_run_id:
            self._load_wandb_config()
            if not self.wandb_run_id:
                logger.warning("No wandb run ID found, skipping wandb logging")
                return

        eval_dir = Path(checkpoint_work_dir)
        if not eval_dir.exists():
            return

        try:
            wandb.init(
                id=self.wandb_run_id,
                project=self.wandb_project,
                entity=self.wandb_entity,
                resume="allow",
            )

            wandb.define_metric("eval_step")
            wandb.define_metric("eval/*", step_metric="eval_step")
            wandb.define_metric("videos_hard*", step_metric="eval_step")
            for subdir in sorted(eval_dir.iterdir()):
                if subdir.is_dir() and subdir.name != "train":
                    wandb.define_metric(f"{subdir.name}/*", step_metric="eval_step")

            all_metrics = {"eval_step": eval_step}

            for subdir in sorted(eval_dir.iterdir()):
                if not subdir.is_dir():
                    continue

                try:
                    metrics_file = subdir / "metrics_eval.json"
                    metrics_finish = subdir / "metrics_finish.txt"
                    if metrics_finish.exists() and metrics_file.exists():
                        self._log_metrics(eval_step, metrics_file)

                    render_finish = subdir / "render_finish.txt"
                    video_dir = subdir / "render_results"
                    if render_finish.exists() and video_dir.exists():
                        self._log_videos(eval_step, metrics_file, video_dir)
                except Exception as subdir_e:  # noqa: BLE001
                    logger.error(
                        f"Failed to log subdir {subdir.name} for step {eval_step}: {subdir_e}"
                    )

            wandb.log(all_metrics)
            wandb.finish()
            logger.info(f"Logged eval results to wandb for step {eval_step}")

        except Exception as e:  # noqa: BLE001
            logger.error(f"Failed to log to wandb for step {eval_step}: {e}")
            try:
                wandb.finish()
            except Exception as e:  # noqa: BLE001
                logger.error(f"Error finishing wandb: {e}")

    def _load_metrics(self, eval_step: int, metrics_file: Path) -> dict | None:
        try:
            with open(metrics_file) as f:
                metrics_eval = json.load(f)
        except json.JSONDecodeError:
            logger.error(f"Error loading {metrics_file}")
            return None

        log_keys = metrics_eval.pop("log_keys", None)

        file_size_mb = metrics_file.stat().st_size / 1024 / 1024
        if file_size_mb > 20:
            metrics_eval.pop("eval/all_metrics_dict", None)
            metrics_eval.pop("eval/failed_metrics_dict", None)
            logger.info(
                f"Skipping per-motion dicts for {metrics_file.parent.name} ({file_size_mb:.0f} MB > 20 MB)"
            )
        else:
            if "eval/all_metrics_dict" in metrics_eval:
                metrics_eval["eval/all_metrics_dict"] = im_eval_callback.create_html_table(
                    metrics_eval["eval/all_metrics_dict"]
                )
            if "eval/failed_metrics_dict" in metrics_eval:
                metrics_eval["eval/failed_metrics_dict"] = im_eval_callback.create_html_table(
                    metrics_eval["eval/failed_metrics_dict"]
                )

        for key in ["failed_keys", "failed_idxes"]:
            metrics_eval.pop(key, None)

        metrics_eval["eval_step"] = eval_step

        if log_keys is not None:
            metrics_eval = {f"{log_keys}/{k}": v for k, v in metrics_eval.items()}
            metrics_eval["eval_step"] = eval_step

        return metrics_eval

    def _log_metrics(self, eval_step: int, metrics_file: Path):
        metrics = self._load_metrics(eval_step, metrics_file)
        if metrics:
            wandb.log(metrics)

    def _log_videos(self, eval_step: int, metrics_file: Path, video_dir: Path):
        if not video_dir.exists():
            return

        log_keys = None
        if metrics_file.exists():
            try:
                with open(metrics_file) as f:
                    metrics = json.load(f)
                    log_keys = metrics.get("log_keys")
            except Exception as e:  # noqa: BLE001
                logger.error(f"Error getting log_keys from metrics file: {e}")

        video_files = sorted(
            [
                (i, f)
                for i, f in enumerate(sorted(video_dir.iterdir()))
                if f.is_file() and f.name.endswith(".mp4")
            ]
        )

        if not video_files:
            return

        prefix = f"videos_hard_{log_keys}" if log_keys else "videos_hard"
        wandb_videos = {
            f"{prefix}/{i:04d}": wandb.Video(str(video_file), format="mp4")
            for i, video_file in reversed(video_files)
        }
        wandb_videos["eval_step"] = eval_step
        wandb.log(wandb_videos)

    def find_evaluated_checkpoints(self):
        """Find all checkpoints that have been successfully evaluated."""
        eval_dir = self.experiment_dir / "eval"

        if not eval_dir.exists():
            logger.info("No eval directory found")
            return

        for eval_subdir in sorted(eval_dir.iterdir()):
            if eval_subdir.is_dir():
                metrics_finish_file = eval_subdir / "metrics_finish.txt"
                metrics_file = eval_subdir / "metrics_eval.json"
                render_finish_file = eval_subdir / "render_finish.txt"
                if (
                    metrics_finish_file.exists()
                    and metrics_file.exists()
                    and render_finish_file.exists()
                ):
                    try:
                        step_num = int(eval_subdir.name)
                        checkpoint_path = (
                            self.experiment_dir / f"model_step_{step_num:06d}.pt"
                        )
                        if checkpoint_path.exists():
                            self.evaluated_checkpoints.add(str(checkpoint_path))
                    except ValueError:
                        pass

        logger.info(f"Found {len(self.evaluated_checkpoints)} already evaluated checkpoints")

    def find_checkpoints(self) -> list[Path]:
        """Find all checkpoint files in the experiment directory."""
        checkpoint_pattern = str(self.experiment_dir / "model_step_*.pt")
        checkpoints = sorted(
            [Path(p) for p in glob.glob(checkpoint_pattern)],
            key=lambda p: int(p.stem.split("_")[-1]),
        )
        return checkpoints

    def is_checkpoint_ready(self, checkpoint_path: Path) -> bool:
        """Check if a checkpoint is ready for evaluation (not being written)."""
        checkpoint_ready_delay = self.config.get("checkpoint_ready_delay", 60)
        mtime = checkpoint_path.stat().st_mtime
        age = time.time() - mtime
        return age > checkpoint_ready_delay

    def evaluate_checkpoint(
        self,
        checkpoint_path: Path,
        mode: str = "metrics",
        work_dir: str = None,
        num_render_videos: int = None,
        eval_step: int = None,  # noqa: ARG002
        eval_dataset: str = None,
        eval_mode: str = None,
    ):
        """Evaluate a single checkpoint using eval_agent_trl.py."""
        checkpoint_str = str(checkpoint_path)
        success = False

        mode_finish_file = os.path.join(work_dir, f"{mode}_finish.txt")
        metrics_file = os.path.join(work_dir, "metrics_eval.json")
        skip = os.path.exists(mode_finish_file)
        if skip and mode == "metrics" and not os.path.exists(metrics_file):
            logger.info(f"[{mode}] Not skipping since metrics file not found: {metrics_file}")
            skip = False
        if skip:
            logger.info(
                f"[{mode}] Skipping evaluation for checkpoint: {checkpoint_path} because it has already been evaluated"  # noqa: E501
            )
            return True

        try:
            logger.info(f"[{mode}] Starting evaluation for checkpoint: {checkpoint_path}")

            eval_callbacks = self.config.get("eval_callbacks", "im_eval")

            if mode == "metrics":
                cmd = f"accelerate launch gear_sonic/eval_agent_trl.py +checkpoint={checkpoint_str} +headless=True ++eval_callbacks={eval_callbacks} ++run_eval_loop=False"  # noqa: E501
                cmd += f" ++num_envs={self.config.num_eval_envs}"
                cmd += f" ++eval_output_dir={work_dir}"
                if eval_mode is not None:
                    cmd += f" ++use_encoder={eval_mode}"
                cmd += " ++manager_env.commands.motion.motion_lib_cfg.multi_thread=False"
                cmd += " +manager_env/terminations=tracking/eval"
                if eval_dataset is not None:
                    cmd += (
                        f" +manager_env.commands.motion.motion_lib_cfg.motion_file={eval_dataset}"
                    )
                    cmd += f" +log_keys={Path(eval_dataset).name}_{eval_mode if eval_mode is not None else 'all'}"

            elif mode == "render":
                cmd = f"python -u gear_sonic/eval_agent_trl.py +checkpoint={checkpoint_str} +headless=True ++eval_callbacks={eval_callbacks} ++run_eval_loop=False"  # noqa: E501
                cmd += f" ++num_envs={num_render_videos}"
                cmd += f" ++metrics_file={metrics_file}"
                render_sort_by = self.config.get("render_sort_by", None)
                if render_sort_by is not None:
                    cmd += f" ++render_sort_by={render_sort_by}"
                cmd += f" ++manager_env.config.save_rendering_dir={work_dir}/render_results"
                cmd += " ++manager_env.config.render_results=True"
                cmd += " ++manager_env.config.env_spacing=10.0"
                cmd += " +manager_env/recorders=render"
                cmd += " ++manager_env.commands.motion.motion_lib_cfg.multi_thread=False"
                if eval_mode is not None:
                    cmd += f" ++use_encoder={eval_mode}"

                if eval_dataset is not None:
                    cmd += (
                        f" +manager_env.commands.motion.motion_lib_cfg.motion_file={eval_dataset}"
                    )

            extra_overrides = self.config.get("extra_overrides", [])
            for override in extra_overrides:
                cmd += f" {override}"

            logger.info(f"Running command: {cmd}")
            capture_output = self.config.get("capture_output", True)
            timeout_seconds = self.config.get("render_timeout", 3600) if mode == "render" else 21600
            proc = subprocess.Popen(
                cmd,
                shell=True,
                preexec_fn=os.setsid,
                stdout=subprocess.PIPE if capture_output else None,
                stderr=subprocess.PIPE if capture_output else None,
                text=True,
            )
            try:
                stdout_data, stderr_data = proc.communicate(timeout=timeout_seconds)
            except subprocess.TimeoutExpired:
                os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
                proc.communicate()
                logger.error(f"Evaluation timeout for checkpoint: {checkpoint_path}")
                return False
            result_returncode = proc.returncode
            result_stdout = stdout_data or ""
            result_stderr = stderr_data or ""

            success = result_returncode == 0 and os.path.exists(metrics_file)

            if mode == "render":
                found_videos = len(glob.glob(os.path.join(work_dir, "render_results", "*.mp4")))
                expected_videos = num_render_videos or self.config.get("num_render_videos", 64)
                if result_returncode == 0 and found_videos < expected_videos:
                    logger.warning(
                        f"[{mode}] Fewer videos than requested: found {found_videos}/{expected_videos} "
                        f"(OK if dataset has fewer unique motions)"
                    )
                elif found_videos == 0:
                    logger.error(f"[{mode}] No videos produced")
                    success = False

            if success:
                logger.info(f"[{mode}] Successfully evaluated checkpoint: {checkpoint_path}")
                self.last_evaluation_time = time.time()
            else:
                logger.error(f"[{mode}] Evaluation failed for checkpoint {checkpoint_path}")
                logger.error("=" * 20 + " stdout " + "=" * 20)
                logger.error(result_stdout)
                logger.error("=" * 20 + " stderr " + "=" * 20)
                logger.error(result_stderr)
                logger.error("=" * 20 + " end " + "=" * 20 + "\n")
                if not os.path.exists(metrics_file):
                    logger.error(f"[{mode}] Metrics file not found: {metrics_file}")

        except Exception as e:  # noqa: BLE001
            logger.error(f"Error evaluating checkpoint {checkpoint_path}: {e}")
            return False

        return success

    def run(self):
        """Main monitoring loop."""
        single_pass = self.config.get("single_pass", False)
        if single_pass:
            logger.info("Running in single-pass mode...")
        else:
            logger.info("Starting checkpoint monitoring loop...")
        eval_datasets = self.config.get("eval_datasets", None)
        eval_modes = self.config.get("eval_modes", [None])
        num_render_videos = self.config.get("num_render_videos", 64)
        num_test_render_videos = self.config.get("num_test_render_videos", 32)
        while not self.shutdown_flag:
            try:
                checkpoints = self.find_checkpoints()

                new_checkpoints = []
                for cp in checkpoints:
                    cp_str = str(cp)
                    if (
                        self.eval_frequency is not None
                        and int(cp.stem.split("_")[-1]) % self.eval_frequency != 0
                    ):
                        continue
                    if cp_str not in self.evaluated_checkpoints and self.is_checkpoint_ready(cp):
                        new_checkpoints.append(cp)

                if self.eval_last_n is not None and len(new_checkpoints) > self.eval_last_n:
                    skipped = len(new_checkpoints) - self.eval_last_n
                    new_checkpoints = new_checkpoints[-self.eval_last_n :]
                    logger.info(
                        f"eval_last_n={self.eval_last_n}: skipping {skipped} earlier checkpoints"
                    )

                if single_pass and not new_checkpoints:
                    logger.info("Single-pass mode: no new checkpoints to evaluate, exiting")
                    break

                evaluation_success_count = 0
                for checkpoint in new_checkpoints:
                    if self.shutdown_flag:
                        break

                    eval_step = int(checkpoint.stem.split("_")[-1])
                    checkpoint_work_dir = os.path.join(
                        self.experiment_dir, "eval", f"{eval_step:06d}"
                    )
                    os.makedirs(checkpoint_work_dir, exist_ok=True)
                    logger.info(f"Found new checkpoint: {checkpoint}")

                    success = True
                    metrics_success = True

                    for mode in ["metrics", "render"]:
                        mode_work_dir = checkpoint_work_dir + "/train"
                        mode_success = self.evaluate_checkpoint(
                            checkpoint,
                            mode=mode,
                            work_dir=mode_work_dir,
                            eval_step=eval_step,
                            num_render_videos=num_render_videos,
                        )
                        if mode_success:
                            with open(os.path.join(mode_work_dir, f"{mode}_finish.txt"), "w") as f:
                                f.write(f"{mode}_finish")
                        success = success and mode_success
                        if mode == "metrics":
                            metrics_success = metrics_success and mode_success

                    if eval_datasets is not None:
                        for eval_dataset, eval_mode in itertools.product(eval_datasets, eval_modes):
                            for mode in ["metrics", "render"]:
                                mode_work_dir = (
                                    checkpoint_work_dir
                                    + f"/{Path(eval_dataset).name}_{eval_mode if eval_mode is not None else 'all'}"
                                )
                                mode_success = self.evaluate_checkpoint(
                                    checkpoint,
                                    mode=mode,
                                    work_dir=mode_work_dir,
                                    eval_step=eval_step,
                                    eval_dataset=eval_dataset,
                                    num_render_videos=num_test_render_videos,
                                    eval_mode=eval_mode,
                                )
                                if mode_success:
                                    with open(
                                        os.path.join(mode_work_dir, f"{mode}_finish.txt"), "w"
                                    ) as f:
                                        f.write(f"{mode}_finish")
                                success = success and mode_success
                                if mode == "metrics":
                                    metrics_success = metrics_success and mode_success

                    if success:
                        with open(
                            os.path.join(checkpoint_work_dir, "all_eval_finish.txt"), "w"
                        ) as f:
                            f.write("all_eval_finish")
                        self._log_eval_to_wandb(eval_step, checkpoint_work_dir)
                    elif metrics_success:
                        logger.warning(
                            f"Render failed for step {eval_step}, logging metrics-only to W&B"
                        )
                        self._log_eval_to_wandb(eval_step, checkpoint_work_dir)

                    if success:
                        self.evaluated_checkpoints.add(str(checkpoint))
                        evaluation_success_count += 1
                        if eval_step >= self.config.max_train_steps:
                            logger.info(
                                f"Reached max train steps: {eval_step} >= {self.config.max_train_steps}. Shutting down..."  # noqa: E501
                            )
                            self.shutdown_flag = True
                            break

                if new_checkpoints:
                    logger.info(f"Evaluated {evaluation_success_count} new checkpoints")
                    logger.info(f"Total evaluated checkpoints: {len(self.evaluated_checkpoints)}")

                if single_pass:
                    logger.info(
                        f"Single-pass mode: evaluated {evaluation_success_count} checkpoint(s), exiting"
                    )
                    break

                time_since_last_eval = time.time() - self.last_evaluation_time
                if time_since_last_eval > self.evaluation_timeout:
                    logger.info(
                        f"No checkpoints evaluated in {time_since_last_eval / 3600:.1f} hours. Shutting down..."
                    )
                    self.shutdown_flag = True
                    break

                time.sleep(self.config.scan_interval)

            except KeyboardInterrupt:
                logger.info("Received keyboard interrupt. Shutting down...")
                break
            except Exception as e:  # noqa: BLE001
                logger.error(f"Error in monitoring loop: {e}")
                time.sleep(self.config.scan_interval)

        logger.info("Checkpoint monitoring stopped.")


@hydra.main(config_path="config", config_name="eval_exp", version_base="1.1")
def main(config: omegaconf.OmegaConf) -> None:
    """Main function to start checkpoint monitoring and evaluation."""
    os.chdir(hydra.utils.get_original_cwd())

    single_pass = config.get("single_pass", False)

    experiment_dir = Path(config.experiment_dir)
    if not experiment_dir.exists():
        parent_dir = experiment_dir.parent
        prefix = experiment_dir.name
        logger.info(
            f"Experiment directory doesn't exist, looking for prefix match: {prefix}* in {parent_dir}"
        )

        while True:
            if parent_dir.exists():
                matches = sorted(
                    [d for d in parent_dir.iterdir() if d.is_dir() and d.name.startswith(prefix)]
                )[::-1]
                if matches:
                    experiment_dir = None
                    for match in matches:
                        if (match / "meta.yaml").exists():
                            experiment_dir = match
                            logger.info(
                                f"Found matching directory with meta.yaml: {experiment_dir}"
                            )
                            break
                    if experiment_dir is None:
                        experiment_dir = matches[-1]
                        logger.info(
                            f"Found matching directory (no meta.yaml yet): {experiment_dir}"
                        )
                    config.experiment_dir = str(experiment_dir)
                    break

            if single_pass:
                logger.info("Single-pass mode: no matching directory found, exiting")
                return

            logger.info("No match found yet, waiting...")
            time.sleep(5)

    meta_file = os.path.join(config.experiment_dir, "meta.yaml")
    logger.info(f"Waiting for meta.yaml to exist: {meta_file}")

    if single_pass and not os.path.exists(meta_file):
        logger.info("Single-pass mode: meta.yaml not found, exiting")
        return
    while not os.path.exists(meta_file):
        time.sleep(1)
    meta = yaml.safe_load(open(meta_file))  # noqa: SIM115
    config.max_train_steps = meta["max_train_steps"]
    logger.info(f"Loaded meta: {meta}")

    hydra_log_path = os.path.join(config.experiment_dir, "eval_exp.log")
    logger.remove()
    logger.add(hydra_log_path, level="DEBUG")
    console_log_level = os.environ.get("LOGURU_LEVEL", "INFO").upper()
    logger.add(sys.stdout, level=console_log_level, colorize=True)

    evaluator = CheckpointEvaluator(config)
    evaluator.run()


if __name__ == "__main__":
    main()