Add evaluate jit models to tensorboard

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_<step>.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/<experiment>, and full results are cached under
+    logs/go2_lab/<experiment>/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_<step>.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/<experiment>/robogauge_results/results_<step>.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_<step>.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()

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 \