From 30643dfed8cf9cf0f9919ebe8c369fd9004087ca Mon Sep 17 00:00:00 2001 From: wty-yy Date: Fri, 26 Jun 2026 21:32:36 +0800 Subject: [PATCH] Add evaluate jit models to tensorboard --- scripts/eval_jit_models_to_tensorboard.py | 504 ++++++++++++++++++++++ scripts/run_models.sh | 13 + 2 files changed, 517 insertions(+) create mode 100755 scripts/eval_jit_models_to_tensorboard.py diff --git a/scripts/eval_jit_models_to_tensorboard.py b/scripts/eval_jit_models_to_tensorboard.py new file mode 100755 index 0000000..c392c06 --- /dev/null +++ b/scripts/eval_jit_models_to_tensorboard.py @@ -0,0 +1,504 @@ +#!/usr/bin/env python3 +"""Batch-submit exported JIT checkpoints to RoboGauge and log scores to TensorBoard. + +Overview: + The script discovers policy_jit_.pt files under rsl_rl experiment logs, + submits checkpoints through robogauge/scripts/client.py, caches full + RoboGauge results as YAML, and writes scores as RoboGauge/{key} TensorBoard + scalars using the checkpoint id as global_step. + +Quick Start: + python scripts/eval_jit_models_to_tensorboard.py + +Full Command: + python scripts/eval_jit_models_to_tensorboard.py \ + --source-root /home/user/go2_rl_robotlab/logs/rsl_rl \ + --experiment-glob "go2_moe_cts_v3*" \ + --log-root logs/go2_lab \ + --task-name go2_lab \ + --server-url http://127.0.0.1:9973 \ + --reuse-source-results \ + --clear-tensorboard + +Options: + --source-root: Directory containing rsl_rl experiment logs. + --experiment-glob: One or more experiment directory patterns under source root. + --experiments: Explicit experiment names that override the glob. + --log-root: TensorBoard and result-cache output root. + --task-name: RoboGauge task name sent to the evaluation server. + --server-url: RoboGauge server URL used by RoboGaugeClient. + --request-timeout: HTTP request timeout for client calls. + --retry-interval: Delay between server availability retries. + --poll-interval: Delay between task status checks. + --start-step: Minimum checkpoint step to include. + --end-step: Maximum checkpoint step to include. + --max-checkpoints: Limit the number of discovered checkpoints. + --overwrite: Re-submit checkpoints even when output result YAML exists. + --reuse-source-results: Reuse source-run result YAML before submitting. + --clear-tensorboard: Delete existing output event files before writing. + --keep-going: Continue after failed checkpoints. + --dry-run: Print planned work without writing logs or contacting the server. + +Notes: + Start robogauge/scripts/server.py before running checkpoints that cannot be + satisfied from cached YAML. Output event files are written under + logs/go2_lab/, and full results are cached under + logs/go2_lab//robogauge_results. +""" + +from __future__ import annotations + +import argparse +import re +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import yaml +from torch.utils.tensorboard import SummaryWriter + + +DEFAULT_SOURCE_ROOT = "/home/user/go2_rl_robotlab/logs/rsl_rl" +DEFAULT_LOG_ROOT = "logs/go2_lab" +CHECKPOINT_RE = re.compile(r"policy_jit_(\d+)\.pt$") + + +@dataclass(frozen=True) +class Checkpoint: + """Represent one exported JIT checkpoint and its source run context.""" + + experiment_name: str + run_dir: Path + model_path: Path + step: int + + @property + def source_result_path(self) -> Path: + """Return the source training run's expected RoboGauge result path. + + Returns: + Path to the source results_.yaml file. + """ + return self.run_dir / "robogauge_results" / f"results_{self.step}.yaml" + + +def parse_args() -> argparse.Namespace: + """Build and parse the command line interface for batch evaluation. + + Returns: + Parsed CLI arguments. + """ + parser = argparse.ArgumentParser( + description=( + "Evaluate exported policy_jit_*.pt files through robogauge/scripts/client.py " + "and write scores with tags matching rsl_rl's RoboGauge/{key} TensorBoard logs." + ) + ) + parser.add_argument("--source-root", type=Path, default=Path(DEFAULT_SOURCE_ROOT)) + parser.add_argument( + "--experiment-glob", + nargs="+", + default=["go2_moe_cts_v3*"], + help="One or more experiment directory glob patterns under --source-root.", + ) + parser.add_argument( + "--experiments", + nargs="+", + help="Explicit experiment directory names under --source-root. Overrides --experiment-glob.", + ) + parser.add_argument("--log-root", type=Path, default=Path(DEFAULT_LOG_ROOT)) + parser.add_argument("--task-name", default="go2_lab") + parser.add_argument("--server-url", default="http://127.0.0.1:9973") + parser.add_argument("--request-timeout", type=float, default=5.0) + parser.add_argument("--retry-interval", type=float, default=2.0) + parser.add_argument("--poll-interval", type=float, default=60.0) + parser.add_argument("--start-step", type=int) + parser.add_argument("--end-step", type=int) + parser.add_argument("--max-checkpoints", type=int) + parser.add_argument( + "--overwrite", + action="store_true", + help="Re-submit checkpoints even when logs/go2_lab//robogauge_results/results_.yaml exists.", + ) + parser.add_argument( + "--reuse-source-results", + action="store_true", + help="Reuse existing results from the source training run's robogauge_results before submitting to the server.", + ) + parser.add_argument( + "--clear-tensorboard", + action="store_true", + help="Delete existing TensorBoard event files in each output experiment directory before writing scalars.", + ) + parser.add_argument( + "--keep-going", + action="store_true", + help="Continue with later checkpoints after a failed evaluation.", + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Print discovered checkpoints and planned actions without contacting the server or writing logs.", + ) + return parser.parse_args() + + +def checkpoint_step(path: Path) -> int | None: + """Extract the checkpoint step from a policy_jit_.pt filename. + + Args: + path: Checkpoint path whose filename should match the export pattern. + + Returns: + Parsed checkpoint step, or None when the filename does not match. + """ + match = CHECKPOINT_RE.match(path.name) + if match is None: + return None + return int(match.group(1)) + + +def discover_experiment_dirs(source_root: Path, experiment_globs: list[str], experiments: list[str] | None) -> list[Path]: + """Discover experiment directories under the source root. + + Explicit experiment names take precedence over the glob. Missing explicit + directories raise an error so typos do not silently skip evaluations. + + Args: + source_root: Directory that contains rsl_rl experiment directories. + experiment_globs: Glob patterns used when explicit experiments are absent. + experiments: Optional experiment directory names to evaluate. + + Returns: + Existing experiment directories to scan. + """ + if experiments: + dirs = [source_root / name for name in experiments] + else: + matched_dirs = { + path + for experiment_glob in experiment_globs + for path in source_root.glob(experiment_glob) + if path.is_dir() + } + dirs = sorted(matched_dirs) + missing = [path for path in dirs if not path.is_dir()] + if missing: + names = "\n".join(f" - {path}" for path in missing) + raise FileNotFoundError(f"Experiment directories not found:\n{names}") + return dirs + + +def discover_checkpoints(args: argparse.Namespace) -> list[Checkpoint]: + """Find exported policy JIT checkpoints that match CLI filters. + + The discovered checkpoints are sorted by experiment name, step, and path so + repeated runs process checkpoints in a stable order. + + Args: + args: Parsed CLI arguments containing discovery roots and filters. + + Returns: + Ordered checkpoints selected for evaluation. + """ + source_root = args.source_root.expanduser().resolve() + experiment_dirs = discover_experiment_dirs(source_root, args.experiment_glob, args.experiments) + checkpoints: list[Checkpoint] = [] + + for experiment_dir in experiment_dirs: + for model_path in sorted(experiment_dir.glob("*/jit_models/policy_jit_*.pt")): + step = checkpoint_step(model_path) + if step is None: + continue + if args.start_step is not None and step < args.start_step: + continue + if args.end_step is not None and step > args.end_step: + continue + checkpoints.append( + Checkpoint( + experiment_name=experiment_dir.name, + run_dir=model_path.parents[1], + model_path=model_path.resolve(), + step=step, + ) + ) + + checkpoints.sort(key=lambda item: (item.experiment_name, item.step, str(item.model_path))) + if args.max_checkpoints is not None: + checkpoints = checkpoints[: args.max_checkpoints] + return checkpoints + + +def load_yaml(path: Path) -> dict[str, Any]: + """Load a YAML mapping from disk. + + Args: + path: YAML file to read. + + Returns: + Parsed YAML mapping. + """ + with path.open("r", encoding="utf-8") as file: + data = yaml.safe_load(file) + if not isinstance(data, dict): + raise ValueError(f"Expected a mapping in {path}, got {type(data).__name__}.") + return data + + +def save_yaml(path: Path, data: dict[str, Any]) -> None: + """Write a YAML mapping to disk. + + Args: + path: Output YAML path. + data: Mapping to serialize. + """ + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", encoding="utf-8") as file: + yaml.safe_dump(data, file, allow_unicode=True, sort_keys=False) + + +def result_scores(results: dict[str, Any], source: Path | str) -> dict[str, Any]: + """Return the scores mapping from a RoboGauge result payload. + + Args: + results: RoboGauge result payload. + source: Human-readable origin used in validation errors. + + Returns: + Mapping of score names to scalar values. + """ + scores = results.get("scores") + if not isinstance(scores, dict): + raise ValueError(f"RoboGauge results from {source} do not contain a 'scores' mapping.") + return scores + + +def log_scores(writer: SummaryWriter, results: dict[str, Any], step: int, source: Path | str) -> None: + """Write RoboGauge score scalars to TensorBoard. + + Scores are logged as RoboGauge/{key}, matching the training runner's + TensorBoard tag format. + + Args: + writer: TensorBoard writer for the output experiment directory. + results: RoboGauge result payload containing a scores mapping. + step: TensorBoard global_step taken from the checkpoint id. + source: Human-readable result origin used in validation errors. + """ + for key, value in result_scores(results, source).items(): + writer.add_scalar(f"RoboGauge/{key}", value, step) + writer.flush() + + +def clear_tensorboard_events(log_dir: Path) -> None: + """Delete TensorBoard event files from an experiment output directory. + + Args: + log_dir: Experiment output directory that may contain event files. + """ + for event_file in log_dir.glob("events.out.tfevents*"): + event_file.unlink() + + +def output_paths(log_root: Path, checkpoint: Checkpoint) -> tuple[Path, Path]: + """Build output paths for an experiment and checkpoint result. + + Args: + log_root: Root directory for TensorBoard logs and result caches. + checkpoint: Checkpoint whose experiment and step define the paths. + + Returns: + Experiment log directory and checkpoint result YAML path. + """ + experiment_log_dir = log_root / checkpoint.experiment_name + result_path = experiment_log_dir / "robogauge_results" / f"results_{checkpoint.step}.yaml" + return experiment_log_dir, result_path + + +def write_cached_result( + writer: SummaryWriter, + results: dict[str, Any], + step: int, + result_path: Path, + source: Path | str, +) -> None: + """Persist a result payload and write its score scalars. + + Args: + writer: TensorBoard writer for the output experiment directory. + results: RoboGauge result payload to cache and log. + step: TensorBoard global_step taken from the checkpoint id. + result_path: Output YAML path for the cached payload. + source: Human-readable result origin used in validation errors. + """ + save_yaml(result_path, results) + log_scores(writer, results, step, source) + + +def submit_and_wait( + client: Any, + checkpoint: Checkpoint, + task_name: str, + poll_interval: float, + retry_interval: float, +) -> dict[str, Any]: + """Submit one checkpoint through RoboGaugeClient and wait for results. + + The function polls the existing client until the submitted task returns a + finished result, preserving one-checkpoint-at-a-time queueing for the batch. + + Args: + client: RoboGaugeClient instance used for submission and monitoring. + checkpoint: Checkpoint to evaluate. + task_name: RoboGauge task name sent with the evaluation request. + poll_interval: Seconds to sleep between task status checks. + retry_interval: Seconds to sleep between server availability retries. + + Returns: + RoboGauge result payload returned by the server. + """ + task_id = client.submit_task( + model_path=str(checkpoint.model_path), + step=checkpoint.step, + task_name=task_name, + experiment_name=checkpoint.experiment_name, + wait_for_server=True, + retry_interval=retry_interval, + ) + if task_id is None: + raise RuntimeError(f"Failed to submit {checkpoint.model_path}.") + + while True: + client.monitor_tasks(wait_for_server=True, retry_interval=retry_interval) + response = client.response_data.pop(task_id, None) + if response is not None: + results = response.get("results") + if not isinstance(results, dict): + raise ValueError(f"Task {task_id} finished without a valid results mapping: {response}") + response_step = int(response.get("step", checkpoint.step)) + if response_step != checkpoint.step: + raise ValueError( + f"Task {task_id} returned step {response_step}, expected {checkpoint.step}." + ) + return results + + if task_id not in client.processing_ids: + raise RuntimeError(f"Task {task_id} ended without a finished result.") + + time.sleep(poll_interval) + + +def print_plan(checkpoints: list[Checkpoint], log_root: Path, overwrite: bool, reuse_source_results: bool) -> None: + """Print a compact summary of the batch evaluation plan. + + Args: + checkpoints: Checkpoints selected for this run. + log_root: Output root for TensorBoard logs and result caches. + overwrite: Whether cached output results will be ignored. + reuse_source_results: Whether source-run result YAML can satisfy work. + """ + counts: dict[str, int] = {} + for checkpoint in checkpoints: + counts[checkpoint.experiment_name] = counts.get(checkpoint.experiment_name, 0) + 1 + print("Discovered checkpoints:") + for experiment_name in sorted(counts): + print(f" - {experiment_name}: {counts[experiment_name]}") + print(f"Output TensorBoard root: {log_root}") + print(f"Overwrite output results: {overwrite}") + print(f"Reuse source results: {reuse_source_results}") + + +def main() -> None: + """Run checkpoint discovery, evaluation, caching, and TensorBoard logging.""" + args = parse_args() + log_root = args.log_root.expanduser().resolve() + checkpoints = discover_checkpoints(args) + print_plan(checkpoints, log_root, args.overwrite, args.reuse_source_results) + + if not checkpoints: + print("No checkpoints matched the requested filters.") + return + + if args.dry_run: + for checkpoint in checkpoints[:20]: + experiment_log_dir, result_path = output_paths(log_root, checkpoint) + action = "evaluate" + if result_path.exists() and not args.overwrite: + action = "use-output-cache" + elif args.reuse_source_results and checkpoint.source_result_path.exists(): + action = "reuse-source-result" + print(f"{action}: step={checkpoint.step} model={checkpoint.model_path} tb={experiment_log_dir}") + if len(checkpoints) > 20: + print(f"... {len(checkpoints) - 20} more checkpoints omitted from dry-run preview.") + return + + log_root.mkdir(parents=True, exist_ok=True) + writers: dict[str, SummaryWriter] = {} + cleared_experiments: set[str] = set() + client: Any | None = None + completed = 0 + failed = 0 + + try: + for index, checkpoint in enumerate(checkpoints, start=1): + experiment_log_dir, result_path = output_paths(log_root, checkpoint) + experiment_log_dir.mkdir(parents=True, exist_ok=True) + if args.clear_tensorboard and checkpoint.experiment_name not in cleared_experiments: + clear_tensorboard_events(experiment_log_dir) + cleared_experiments.add(checkpoint.experiment_name) + + writer = writers.get(checkpoint.experiment_name) + if writer is None: + writer = SummaryWriter(str(experiment_log_dir)) + writers[checkpoint.experiment_name] = writer + + prefix = f"[{index}/{len(checkpoints)}] {checkpoint.experiment_name} step {checkpoint.step}" + try: + if result_path.exists() and not args.overwrite: + results = load_yaml(result_path) + log_scores(writer, results, checkpoint.step, result_path) + print(f"{prefix}: logged cached output result.") + completed += 1 + continue + + if args.reuse_source_results and checkpoint.source_result_path.exists() and not args.overwrite: + results = load_yaml(checkpoint.source_result_path) + write_cached_result(writer, results, checkpoint.step, result_path, checkpoint.source_result_path) + print(f"{prefix}: reused source result and logged TensorBoard scalars.") + completed += 1 + continue + + if client is None: + from robogauge.scripts.client import RoboGaugeClient + + client = RoboGaugeClient(args.server_url, request_timeout=args.request_timeout) + client.wait_until_available(retry_interval=args.retry_interval) + + print(f"{prefix}: submitting {checkpoint.model_path}.") + results = submit_and_wait( + client=client, + checkpoint=checkpoint, + task_name=args.task_name, + poll_interval=args.poll_interval, + retry_interval=args.retry_interval, + ) + write_cached_result(writer, results, checkpoint.step, result_path, "server") + print(f"{prefix}: evaluation finished and logged.") + completed += 1 + + except Exception as exc: + failed += 1 + print(f"{prefix}: failed: {exc}") + if not args.keep_going: + raise + + finally: + for writer in writers.values(): + writer.close() + + print(f"Done. Completed: {completed}, failed: {failed}, output: {log_root}") + + +if __name__ == "__main__": + main() diff --git a/scripts/run_models.sh b/scripts/run_models.sh index b0f7387..7cdce10 100755 --- a/scripts/run_models.sh +++ b/scripts/run_models.sh @@ -2,6 +2,19 @@ source /data/user/wutianyang/Programs/miniforge3/bin/activate go2-gym +python robogauge/scripts/run.py \ + --task-name go2_lab \ + --model-path /home/user/go2_rl_robotlab/logs/rsl_rl/go2_moe_cts_v4.2/2026-06-08_14-58-39/exported/policy.pt \ + --experiment-name go2_moe_cts_v4.2_0.6984 \ + --stress-benchmark \ + --stress-terrain-names flat slope_fd slope_bd stairs_fd stairs_bd wave obstacle \ + --num-processes 70 \ + --seeds 0 1 2 \ + --search-seeds 0 1 2 3 4 \ + --frictions 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.0 \ + --compress-logs \ + --headless + python robogauge/scripts/run.py \ --task-name go2_moe \ --model-path /data/user/wutianyang/Coding/go2_rl_gym/mytest_merge_data/rem_cts/go2_moe_cts_expert_goal_137000_0.6745/policies/policy.pt \