geoguesser-env / env /scripts /render_trace.py
AdithyaSK's picture
AdithyaSK HF Staff
Deploy geoguesser environment with train/eval splits
ee44678 verified
Raw
History Blame Contribute Delete
9.49 kB
# SPDX-License-Identifier: BSD-3-Clause
"""Render a rollout trace as a single HTML page you can scroll through.
A JSONL trace answers "what was the reward"; it does not show you *why*. This
lays each episode out as a filmstrip: the view the model was looking at, what it
said, the action it chose, what the environment replied, and the running cost —
ending with the reveal and the score.
Usage:
python scripts/render_trace.py rollouts/anthropic_agentic
python scripts/render_trace.py rollouts/hf_agentic --out qwen.html
"""
from __future__ import annotations
import argparse
import base64
import html
import json
import pathlib
def _data_uri(path: pathlib.Path) -> str:
"""Inline an image so the report is a single portable file."""
if not path.exists():
return ""
mime = "image/png" if path.suffix == ".png" else "image/jpeg"
return f"data:{mime};base64,{base64.b64encode(path.read_bytes()).decode()}"
def _fmt_distance(km: float | None) -> str:
if km is None:
return "—"
return f"{km * 1000:.0f} m" if km < 10 else f"{km:,.0f} km"
def _verdict(km: float | None) -> str:
if km is None:
return "no guess"
if km < 0.025:
return "perfect"
if km < 25:
return "pinpoint"
if km < 200:
return "close"
if km < 1500:
return "right region"
return "wrong continent"
def render(trace_dir: pathlib.Path, out: pathlib.Path) -> None:
"""Write one HTML report for a trace directory."""
records = [
json.loads(line)
for line in (trace_dir / "trace.jsonl").read_text().splitlines()
if line.strip()
]
summary_path = trace_dir / "summary.json"
summary = json.loads(summary_path.read_text()) if summary_path.exists() else {}
episodes = []
for record in records:
turns = []
for row in record.get("trace", []):
image = ""
if row.get("image"):
image = _data_uri(trace_dir / row["image"])
action = row.get("action") or {}
kind = (
action.get("action", "reset")
if isinstance(action, dict)
else str(action)
)
args = (
", ".join(
f"{key}={value}" for key, value in action.items() if key != "action"
)
if isinstance(action, dict)
else ""
)
turns.append(
{
"turn": row.get("turn"),
"kind": kind,
"args": args,
"reply": row.get("reply") or "",
"feedback": row.get("feedback") or "",
"steps": row.get("steps_remaining"),
"cost": row.get("action_cost"),
"image": image,
"image_kind": row.get("image_kind", "none"),
}
)
episodes.append({"record": record, "turns": turns})
parts = [
_HEAD.replace(
"__TITLE__", html.escape(str(summary.get("model", trace_dir.name)))
)
]
parts.append(_summary_block(summary, records))
for episode in episodes:
parts.append(_episode_block(episode))
parts.append("</body></html>")
out.write_text("".join(parts))
print(f"wrote {out} ({out.stat().st_size / 1e6:.1f} MB, {len(episodes)} episodes)")
def _summary_block(summary: dict, records: list[dict]) -> str:
rows = "".join(
f"<div class='stat'><span>{html.escape(label)}</span><b>{html.escape(str(value))}</b></div>"
for label, value in [
("model", summary.get("model", "—")),
("mode", summary.get("mode", "—")),
("episodes", summary.get("episodes", len(records))),
("mean reward", f"{summary.get('mean_reward', 0):.3f}"),
("median distance", _fmt_distance(summary.get("median_distance_km"))),
("within 200 km", f"{summary.get('within_200km', 0)}/{len(records)}"),
("parsed", f"{summary.get('parsed', 0)}/{len(records)}"),
("wall clock", f"{summary.get('wall_clock_s', 0):.1f}s"),
]
)
return f"<header><h1>rollout trace</h1><div class='stats'>{rows}</div></header>"
def _episode_block(episode: dict) -> str:
record = episode["record"]
km = record.get("distance_km")
reward = record.get("reward") or 0.0
head = (
f"<h2>episode {record.get('task_index')} "
f"<span class='country'>{html.escape(str(record.get('country', '')))}</span>"
f"<span class='verdict v-{_verdict(km).replace(' ', '-')}'>{_verdict(km)}</span>"
f"<span class='num'>{_fmt_distance(km)}</span>"
f"<span class='num'>reward {reward:.3f}</span>"
f"<span class='num'>{record.get('turns', 0)} turns</span>"
f"<span class='num'>{record.get('latency_s', 0):.1f}s</span></h2>"
)
cards = []
for turn in episode["turns"]:
image = (
f"<img src='{turn['image']}' alt='observation at turn {turn['turn']}'>"
if turn["image"]
else "<div class='noimg'>no image</div>"
)
reply = html.escape(turn["reply"][:1400])
meta = []
if turn["steps"] is not None:
meta.append(f"{turn['steps']} steps left")
if turn["cost"] is not None:
meta.append(f"cost {turn['cost']:.2f}")
cards.append(
f"""<div class='turn {"map" if turn["image_kind"] == "map" else ""}'>
<div class='thumb'>{image}</div>
<div class='body'>
<div class='act'><b>{html.escape(str(turn["kind"]))}</b>
<span>{html.escape(turn["args"])}</span>
<span class='meta'>{html.escape(" · ".join(meta))}</span></div>
<div class='fb'>{html.escape(turn["feedback"])}</div>
{f"<details><summary>model reply</summary><pre>{reply}</pre></details>" if reply else ""}
</div>
</div>"""
)
return f"<section>{head}<div class='turns'>{''.join(cards)}</div></section>"
_HEAD = """<!doctype html>
<html lang="en"><head><meta charset="utf-8">
<meta name="viewport" content="width=device-width,initial-scale=1">
<title>__TITLE__ — rollout trace</title>
<style>
:root {
--ground:#eef1f4; --panel:#fff; --edge:#d3dae0; --ink:#1d2328;
--soft:#5c666e; --faint:#8b959c; --wire:#4a7d99; --pin:#c4332a; --good:#3f7a55;
}
@media (prefers-color-scheme:dark){ :root{
--ground:#12171b; --panel:#191f24; --edge:#2b343b; --ink:#e6ebef;
--soft:#a3aeb6; --faint:#6e7a83; --wire:#7fb3cc; --pin:#e8695e; --good:#79b892; } }
*{box-sizing:border-box}
body{margin:0;background:var(--ground);color:var(--ink);
font:14px/1.6 ui-monospace,"SF Mono",Menlo,monospace;padding:24px 18px 60px}
header{max-width:1180px;margin:0 auto 26px}
h1{font-size:19px;margin:0 0 12px;font-weight:600;letter-spacing:.02em}
.stats{display:flex;flex-wrap:wrap;gap:18px}
.stat{background:var(--panel);border:1px solid var(--edge);border-radius:4px;
padding:7px 12px;min-width:118px}
.stat span{display:block;font-size:9.5px;letter-spacing:.12em;
text-transform:uppercase;color:var(--faint)}
.stat b{font-size:15px;font-weight:500;font-variant-numeric:tabular-nums}
section{max-width:1180px;margin:0 auto 30px;background:var(--panel);
border:1px solid var(--edge);border-radius:5px;overflow:hidden}
h2{margin:0;padding:11px 15px;font-size:13px;font-weight:600;
border-bottom:1px solid var(--edge);display:flex;align-items:center;
gap:12px;flex-wrap:wrap}
h2 .country{color:var(--soft);font-weight:400}
h2 .num{margin-left:auto;font-weight:400;color:var(--soft);font-size:11.5px;
font-variant-numeric:tabular-nums}
h2 .num + .num{margin-left:0}
.verdict{font-size:10px;letter-spacing:.1em;text-transform:uppercase;
padding:2px 7px;border-radius:3px;border:1px solid var(--edge);color:var(--soft)}
.v-pinpoint,.v-perfect{color:var(--good);border-color:var(--good)}
.v-wrong-continent,.v-no-guess{color:var(--pin);border-color:var(--pin)}
.turns{display:flex;flex-direction:column}
.turn{display:flex;gap:14px;padding:12px 15px;border-bottom:1px solid var(--edge)}
.turn:last-child{border-bottom:none}
.turn.map{background:color-mix(in srgb,var(--wire) 7%,transparent)}
.thumb{flex:0 0 200px}
.thumb img{width:200px;border-radius:3px;display:block;border:1px solid var(--edge)}
.noimg{width:200px;height:96px;border:1px dashed var(--edge);border-radius:3px;
display:flex;align-items:center;justify-content:center;color:var(--faint);font-size:11px}
.body{flex:1;min-width:0}
.act b{color:var(--wire)}
.act span{color:var(--soft);font-size:12px;margin-left:6px}
.act .meta{float:right;color:var(--faint);font-size:11px}
.fb{margin-top:5px;color:var(--soft);font-size:12.5px}
details{margin-top:8px}
summary{cursor:pointer;font-size:11px;color:var(--faint)}
pre{white-space:pre-wrap;font-size:11.5px;background:var(--ground);
border:1px solid var(--edge);border-radius:3px;padding:9px;margin:7px 0 0;
max-height:260px;overflow:auto}
</style></head><body>
"""
def main() -> None:
"""Command-line entry point."""
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("trace_dir", type=pathlib.Path)
parser.add_argument("--out", type=pathlib.Path)
args = parser.parse_args()
out = args.out or (args.trace_dir / "trace.html")
render(args.trace_dir, out)
if __name__ == "__main__":
main()