"""
visualize.py — OpenEnv Data Pipeline Debugger
Updated: added generate_replay_html() for the step-by-step Replay Dashboard.
Original generate_reward_chart() is fully preserved and unchanged.
"""
from __future__ import annotations
import json
import os
from typing import Any
# ──────────────────────────────────────────────────────────────────────────────
# Original generate_reward_chart() — unchanged
# ──────────────────────────────────────────────────────────────────────────────
def generate_reward_chart(
training_results: list[dict],
output_path: str = "training_results.html",
) -> str:
"""
Generate an interactive Chart.js reward-curve HTML file.
Parameters
----------
training_results : list[dict]
Each dict: {"episode": int, "score": float, "task": str}
output_path : str
Where to write the HTML file.
Returns
-------
str — absolute path of the written file.
"""
episodes = [r["episode"] for r in training_results]
scores = [round(r["score"], 4) for r in training_results]
tasks = [r.get("task", "unknown") for r in training_results]
# Colour per task tier
task_colors = {
"task_easy_schema_fix": "#1D9E75",
"task_medium_data_quality": "#378ADD",
"task_hard_pipeline_orchestration": "#EF9F27",
"task_veryhard_streaming_pipeline": "#D85A30",
"task_expert_multi_source_join": "#534AB7",
}
point_colors = [task_colors.get(t, "#888780") for t in tasks]
html = f"""
OpenEnv — Training Reward Curves
OpenEnv — Training Reward Curves
Curriculum training across {len(episodes)} episodes · {len(set(tasks))} task tiers
Best score
{max(scores):.4f}
Final score
{scores[-1]:.4f}
Total episodes
{len(episodes)}
Improvement
+{(scores[-1]-scores[0]):.3f}
"""
with open(output_path, "w", encoding="utf-8") as f:
f.write(html)
print(f"[visualize] Reward chart -> {os.path.abspath(output_path)}")
return os.path.abspath(output_path)
# ──────────────────────────────────────────────────────────────────────────────
# NEW generate_replay_html() — Replay & Step Debugger Dashboard
# ──────────────────────────────────────────────────────────────────────────────
def generate_replay_html(
episode_log: list, # list[StepRecord] from env.history
episode_num: int = 0,
task_id: str = "",
final_score: float | None = None,
output_path: str | None = None,
) -> str:
"""
Generate a standalone Replay & Step Debugger HTML dashboard from one episode.
Parameters
----------
episode_log : list[StepRecord]
Populated automatically by DataPipelineEnvironment.history after an episode.
episode_num : int
Episode number (for labelling only).
task_id : str
Task identifier string.
final_score : float | None
Override the final score shown in the stats bar.
output_path : str | None
Where to write the file. Defaults to replay_ep{N}.html.
Returns
-------
str — absolute path of the written file.
Usage
-----
>>> from visualize import generate_replay_html
>>> path = generate_replay_html(env.history, episode_num=69,
... task_id="task_hard_pipeline_orchestration")
>>> print(f"Replay saved → {path}")
"""
if output_path is None:
output_path = f"replay_ep{episode_num}.html"
steps_data = [s.to_dict() if hasattr(s, "to_dict") else s for s in episode_log]
total_steps = len(steps_data)
if total_steps == 0:
raise ValueError("episode_log is empty — run an episode first.")
last = steps_data[-1]
score = final_score if final_score is not None else last.get("cumulative_reward", 0)
task_label = task_id.replace("task_", "").replace("_", " ").title() if task_id else "Unknown task"
# Chip CSS class per action type
chip_map = {
"inspect": "chip-inspect",
"cast_column": "chip-cast",
"drop_nulls": "chip-drop",
"fill_nulls": "chip-fill",
"drop_duplicates": "chip-drop",
"filter_outliers": "chip-filter",
"rename_column": "chip-cast",
"reorder_stages": "chip-reorder",
"apply_business_rule": "chip-filter",
"validate": "chip-validate",
"submit": "chip-submit",
}
html = f"""
Replay — Episode {episode_num} · {task_label}
Replay — Episode {episode_num}
{task_label} · {total_steps} steps recorded
Task
{task_label}
Episode
{episode_num}
Final score
{score:.4f}
Steps used
— / {total_steps}
Episode timeline — click any dot to jumpStep 1 of {total_steps}
step 1step {total_steps // 2}step {total_steps}
Current step
1inspect
—
Reasoning—
Observation Summary—
Reward—
Cumulative—
Bugs left—
Reward curve
Action log (last 6 steps)
← → arrow keys also work
"""
with open(output_path, "w", encoding="utf-8") as f:
f.write(html)
print(f"[visualize] Replay dashboard -> {os.path.abspath(output_path)}")
return os.path.abspath(output_path)