cyy/RoboGauge
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add evaluate jit models to tensorboard

Browse Source
This commit is contained in:
wty-yy
2026-06-26 21:32:36 +08:00
parent 9c172a0d39
commit 30643dfed8
2 changed files with 517 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

504
scripts/eval_jit_models_to_tensorboard.py Executable file
View File

@@ -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()