Upload folder using huggingface_hub

app.py

@@ -1,14 +1,14 @@
 """
 RoboGen — HaptalAI Synthetic Robotics Dataset Generator
-Gradio Space: HaptalAI/robogen
-
-Step-by-step UI:
-  1 → Robot selection (card buttons)
-  2 → Task selection (dropdown)
-  3 → Parameter configuration (sliders + checkboxes with tooltips)
-  4 → Generate (progress bar)
-  5 → Results dashboard (quality score, band, failure breakdown)
-  6 → Email gate → Download zip (parquet + README)
+Gradio 5.9.1 / Python 3.11
+
+Step flow:
+  1  Robot selection (card-style radio)
+  2  Task dropdown
+  3  Parameter sliders + failure checkboxes
+  4  Generate button
+  5  Quality results dashboard
+  6  Email gate + zip download
 """
 
 from __future__ import annotations
@@ -16,18 +16,15 @@ from __future__ import annotations
 import os
 import sys
 import io
-import json
 import zipfile
 import tempfile
 import traceback
 from typing import Optional, Dict, List
 
-# Allow running as  python space/app.py  from repo root
-sys.path.insert(0, os.path.dirname(__file__))
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
 
 import gradio as gr
 import pandas as pd
-import numpy as np
 
 from generator import (
     generate_dataset,
@@ -40,25 +37,13 @@ from generator import (
 from readme_gen import generate_readme
 from airtable import log_email
 
-# ── Load CSS ──────────────────────────────────────────────────────────────────
+# ── CSS ───────────────────────────────────────────────────────────────────────
 
-_CSS_PATH = os.path.join(os.path.dirname(__file__), "style.css")
-with open(_CSS_PATH) as _f:
-    _CSS = _f.read()
+_here = os.path.dirname(os.path.abspath(__file__))
+with open(os.path.join(_here, "style.css")) as _f:
+    CSS = _f.read()
 
-# ── Robot display config ──────────────────────────────────────────────────────
-
-ROBOT_ICONS = {
-    "SO-100": "SO-100",
-    "SO-101": "SO-101",
-    "Koch":   "Koch",
-}
-
-ROBOT_DESCRIPTIONS = {
-    "SO-100": "Low-cost 6-DOF arm, community favourite",
-    "SO-101": "Upgraded SO-100 with refined kinematics",
-    "Koch":   "Koch arm — drawer & manipulation tasks",
-}
+# ── Constants ─────────────────────────────────────────────────────────────────
 
 TASK_LABELS = {
     "pick_and_place":    "Pick and Place",
@@ -74,9 +59,7 @@ FAILURE_LABELS = {
     "torque_saturation": "Torque Saturation",
 }
 
-# ── Default parameters per robot × task ──────────────────────────────────────
-
-DEFAULTS: Dict[str, Dict] = {
+DEFAULTS = {
     "SO-100": {"n_eps": 50, "success": 70, "fmin": 1.0, "fmax": 10.0},
     "SO-101": {"n_eps": 50, "success": 70, "fmin": 1.0, "fmax": 10.0},
     "Koch":   {"n_eps": 30, "success": 75, "fmin": 0.5, "fmax":  8.0},
@@ -84,39 +67,30 @@ DEFAULTS: Dict[str, Dict] = {
 
 # ── HTML helpers ──────────────────────────────────────────────────────────────
 
-def _make_results_html(result: Dict, robot: str, task: str) -> str:
-    score  = result["overall_score"]
-    band   = result["band"]
-    n_pass = result["n_passed"]
-    n_flag = result["n_flagged"]
-    n_eps  = result["n_episodes"]
+def _results_html(result: Dict, robot: str, task: str) -> str:
+    score    = result["overall_score"]
+    band     = result["band"]
+    n_pass   = result["n_passed"]
+    n_flag   = result["n_flagged"]
+    n_eps    = result["n_episodes"]
     mismatch = result["mean_mismatch"]
     fb       = result["failure_breakdown"]
     scorer   = result["scorer_used"]
-
-    band_cls  = band.lower()
+    band_cls = band.lower()
     band_desc = {
         "Clean":   "Trajectories are smooth and anomaly-free. Ready for policy training.",
         "Review":  "Some anomalies detected. Review flagged episodes before training.",
         "Flagged": "High anomaly rate. Best used for failure analysis and augmentation.",
     }.get(band, "")
-
-    # Failure bars
-    total_failures = sum(fb.values()) or 1
-    bar_html = ""
-    for key, count in sorted(fb.items(), key=lambda x: -x[1]):
-        label = FAILURE_LABELS.get(key, key)
-        pct   = count / total_failures * 100
-        bar_html += f"""
-        <div class="rg-failure-bar">
-          <span class="rg-failure-label">{label}</span>
-          <div class="rg-bar-track"><div class="rg-bar-fill" style="width:{pct:.0f}%"></div></div>
-          <span class="rg-bar-count">{count}</span>
-        </div>"""
-
-    task_label  = TASK_LABELS.get(task, task)
-    no_failures = "No failure episodes in dataset." if not fb else ""
-
+    total = sum(fb.values()) or 1
+    bars  = "".join(
+        f'<div class="rg-failure-bar">'
+        f'<span class="rg-failure-label">{FAILURE_LABELS.get(k,k)}</span>'
+        f'<div class="rg-bar-track"><div class="rg-bar-fill" style="width:{v/total*100:.0f}%"></div></div>'
+        f'<span class="rg-bar-count">{v}</span></div>'
+        for k, v in sorted(fb.items(), key=lambda x: -x[1])
+    )
+    task_label = TASK_LABELS.get(task, task)
     return f"""
 <div class="rg-results">
   <div class="rg-score-row">
@@ -129,63 +103,27 @@ def _make_results_html(result: Dict, robot: str, task: str) -> str:
       <div class="rg-band-desc">{band_desc}</div>
     </div>
   </div>
-
   <div class="rg-stat-grid">
-    <div class="rg-stat">
-      <div class="rg-stat-value">{n_eps}</div>
-      <div class="rg-stat-label">Total Episodes</div>
-    </div>
-    <div class="rg-stat">
-      <div class="rg-stat-value" style="color:var(--green)">{n_pass}</div>
-      <div class="rg-stat-label">Passed</div>
-    </div>
-    <div class="rg-stat">
-      <div class="rg-stat-value" style="color:var(--red)">{n_flag}</div>
-      <div class="rg-stat-label">Flagged</div>
-    </div>
-    <div class="rg-stat">
-      <div class="rg-stat-value">{mismatch:.3f}</div>
-      <div class="rg-stat-label">Mean Mismatch</div>
-    </div>
-    <div class="rg-stat">
-      <div class="rg-stat-value">{robot}</div>
-      <div class="rg-stat-label">Robot</div>
-    </div>
-    <div class="rg-stat">
-      <div class="rg-stat-value" style="font-size:0.9rem">{task_label}</div>
-      <div class="rg-stat-label">Task</div>
-    </div>
+    <div class="rg-stat"><div class="rg-stat-value">{n_eps}</div><div class="rg-stat-label">Total Episodes</div></div>
+    <div class="rg-stat"><div class="rg-stat-value" style="color:var(--green)">{n_pass}</div><div class="rg-stat-label">Passed</div></div>
+    <div class="rg-stat"><div class="rg-stat-value" style="color:var(--red)">{n_flag}</div><div class="rg-stat-label">Flagged</div></div>
+    <div class="rg-stat"><div class="rg-stat-value">{mismatch:.3f}</div><div class="rg-stat-label">Mean Mismatch</div></div>
+    <div class="rg-stat"><div class="rg-stat-value">{robot}</div><div class="rg-stat-label">Robot</div></div>
+    <div class="rg-stat"><div class="rg-stat-value" style="font-size:0.9rem">{task_label}</div><div class="rg-stat-label">Task</div></div>
   </div>
-
   <div class="rg-failure-section">
     <div class="rg-failure-title">Failure Type Breakdown</div>
-    {bar_html or no_failures}
+    {bars or "No failure episodes in dataset."}
   </div>
-
   <div class="rg-scorer-note">
     Scored by HaptalAI misalignment benchmark &middot; scorer: <code>{scorer}</code>
   </div>
-</div>
-"""
-
-
-# ── Download bundle builder ───────────────────────────────────────────────────
+</div>"""
 
-def _build_zip(
-    df:      pd.DataFrame,
-    result:  Dict,
-    robot:   str,
-    task:    str,
-    n_eps:   int,
-    success: float,
-    fmin:    float,
-    fmax:    float,
-    failures: List[str],
-) -> str:
-    """Annotate DF, write parquet + README into a temp zip, return zip path."""
-    df_annotated = annotate_quality_scores(df, result)
 
-    readme_md = generate_readme(
+def _build_zip(df, result, robot, task, n_eps, success, fmin, fmax, failures) -> str:
+    df_out = annotate_quality_scores(df, result)
+    readme = generate_readme(
         robot=robot, task=task, n_episodes=n_eps,
         success_rate=success / 100, force_min=fmin, force_max=fmax,
         failures=failures,
@@ -195,423 +133,285 @@ def _build_zip(
         failure_breakdown=result["failure_breakdown"],
         scorer_used=result["scorer_used"],
     )
-
-    tag     = f"{robot.replace('-', '')}_{task}"
-    zip_fd, zip_path = tempfile.mkstemp(suffix=".zip", prefix=f"robogen_{tag}_")
-    os.close(zip_fd)
-
-    with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_DEFLATED) as zf:
-        # Parquet
+    tag = f"{robot.replace('-','')}_{task}"
+    fd, path = tempfile.mkstemp(suffix=".zip", prefix=f"robogen_{tag}_")
+    os.close(fd)
+    with zipfile.ZipFile(path, "w", compression=zipfile.ZIP_DEFLATED) as zf:
         buf = io.BytesIO()
-        df_annotated.to_parquet(buf, index=False)
+        df_out.to_parquet(buf, index=False)
         zf.writestr(f"robogen_{tag}.parquet", buf.getvalue())
-        # README
-        zf.writestr("README.md", readme_md.encode("utf-8"))
+        zf.writestr("README.md", readme.encode("utf-8"))
+    return path
+
+
+# ── Event handlers (module level — Gradio 5 requirement) ─────────────────────
+
+def on_robot_select(robot: str):
+    if not robot:
+        return (
+            gr.update(visible=False),
+            gr.update(choices=[], value=None),
+            gr.update(visible=False),
+            gr.update(visible=False),
+            "",
+        )
+    tasks_raw  = TASKS_BY_ROBOT[robot]
+    tasks_disp = [(TASK_LABELS.get(t, t), t) for t in tasks_raw]
+    return (
+        gr.update(visible=True),
+        gr.update(choices=tasks_disp, value=tasks_raw[0]),
+        gr.update(visible=False),
+        gr.update(visible=False),
+        robot,
+    )
+
+
+def on_task_select(task: str, robot: str):
+    if not task or not robot:
+        return gr.update(visible=False), gr.update(visible=False), 50, 70, 1.0, 10.0
+    d  = DEFAULTS.get(robot, DEFAULTS["SO-100"])
+    fr = ROBOT_CONFIG[robot]["force_range"]
+    return (
+        gr.update(visible=True),
+        gr.update(visible=True),
+        d["n_eps"],
+        d["success"],
+        fr[0],
+        fr[1],
+    )
+
+
+def on_generate(robot, task, n_eps, success_pct, fmin, fmax, failures):
+    if not robot or not task:
+        return (
+            "Please complete steps 1 and 2 first.",
+            gr.update(visible=False), "",
+            gr.update(visible=False),
+            None, None,
+        )
+    if not failures:
+        failures = list(FAILURE_TYPES)
+    try:
+        df = generate_dataset(
+            robot=robot, task=task,
+            n_episodes=int(n_eps),
+            success_rate=float(success_pct) / 100,
+            force_min=float(fmin), force_max=float(fmax),
+            enabled_failures=list(failures),
+            seed=None,
+        )
+        result  = score_dataset(df)
+        panel   = _results_html(result, robot, task)
+        status  = (
+            f"Generated {len(df):,} rows across {result['n_episodes']} episodes — "
+            f"score **{result['overall_score']:.1f}/100** ({result['band']})"
+        )
+        return (
+            status,
+            gr.update(visible=True), panel,
+            gr.update(visible=True),
+            df, result,
+        )
+    except Exception:
+        return (
+            f"Generation failed:\n```\n{traceback.format_exc()}\n```",
+            gr.update(visible=False), "",
+            gr.update(visible=False),
+            None, None,
+        )
 
-    return zip_path
+
+def on_email_submit(email, robot, task, n_eps, success_pct, fmin, fmax, failures, df, result):
+    if not email or "@" not in email:
+        return "Please enter a valid email address.", gr.update(visible=False)
+    if df is None or result is None:
+        return "Generate a dataset first (Step 4).", gr.update(visible=False)
+    try:
+        ok, msg = log_email(
+            email=email.strip(), robot=robot, task=task,
+            n_episodes=int(n_eps),
+            quality_score=result["overall_score"],
+            band=result["band"],
+        )
+        if not ok:
+            print(f"[RoboGen] Airtable: {msg}")
+    except Exception as exc:
+        print(f"[RoboGen] Airtable exception: {exc}")
+    try:
+        path = _build_zip(
+            df=df, result=result, robot=robot, task=task,
+            n_eps=int(n_eps), success=float(success_pct),
+            fmin=float(fmin), fmax=float(fmax),
+            failures=list(failures),
+        )
+        return "Email confirmed. Your download is ready below.", gr.update(visible=True, value=path)
+    except Exception:
+        return (
+            f"Download preparation failed:\n```\n{traceback.format_exc()}\n```",
+            gr.update(visible=False),
+        )
 
 
-# ── Gradio app ────────────────────────────────────────────────────────────────
+# ── Build UI ──────────────────────────────────────────────────────────────────
 
-def build_app() -> gr.Blocks:
+with gr.Blocks(css=CSS, title="RoboGen") as demo:
 
-    with gr.Blocks(
-        css=_CSS,
-        theme=gr.themes.Base(
-            primary_hue="purple",
-            neutral_hue="slate",
-        ),
-        title="RoboGen — Synthetic Robotics Datasets",
-        analytics_enabled=False,
-    ) as demo:
+    robot_state  = gr.State("")
+    df_state     = gr.State(None)
+    result_state = gr.State(None)
 
-        # ── Persistent state ──────────────────────────────────────────────
-        robot_state  = gr.State("")
-        df_state     = gr.State(None)
-        result_state = gr.State(None)
+    gr.HTML("""
+    <div class="rg-header">
+      <div class="rg-logo">RoboGen</div>
+      <div class="rg-tagline">Synthetic robotics datasets, physics-accurate &amp; quality-scored</div>
+      <div class="rg-badge">LeRobot-format &nbsp;&middot;&nbsp; SO-100 / SO-101 / Koch &nbsp;&middot;&nbsp; HaptalAI</div>
+    </div>""")
 
-        # ── Header ────────────────────────────────────────────────────────
+    # ── Step 1 ────────────────────────────────────────────────────────────────
+    with gr.Group(elem_classes=["step-card"]):
         gr.HTML("""
-        <div class="rg-header">
-          <div class="rg-logo">RoboGen</div>
-          <div class="rg-tagline">Synthetic robotics datasets, physics-accurate &amp; quality-scored</div>
-          <div class="rg-badge">LeRobot-format &nbsp;&middot;&nbsp; SO-100 / SO-101 / Koch &nbsp;&middot;&nbsp; HaptalAI</div>
-        </div>
-        """)
-
-        # ────────────────────────────────────────────────────────────────────
-        # STEP 1 — Robot selection
-        # ────────────────────────────────────────────────────────────────────
-        with gr.Group(elem_classes=["step-card"]):
-            gr.HTML("""
-            <div class="step-header">
-              <span class="step-num">1</span>
-              <span class="step-title">Select Robot</span>
-            </div>""")
-
-            robot_select = gr.Radio(
-                choices=["SO-100", "Koch", "SO-101"],
-                value=None,
-                label="",
-                elem_classes=["robot-radio"],
-            )
+        <div class="step-header">
+          <span class="step-num">1</span>
+          <span class="step-title">Select Robot</span>
+        </div>""")
+        robot_select = gr.Radio(
+            choices=["SO-100", "Koch", "SO-101"],
+            value=None,
+            label="",
+            elem_classes=["robot-radio"],
+        )
 
-        # ────────────────────────────────────────────────────────────────────
-        # STEP 2 — Task selection
-        # ────────────────────────────────────────────────────────────────────
-        with gr.Group(visible=False, elem_classes=["step-card"]) as step2_grp:
-            gr.HTML("""
-            <div class="step-header">
-              <span class="step-num">2</span>
-              <span class="step-title">Select Task</span>
-            </div>""")
-
-            task_select = gr.Dropdown(
-                choices=[],
-                value=None,
-                label="Task",
-                interactive=True,
+    # ── Step 2 ────────────────────────────────────────────────────────────────
+    with gr.Group(visible=False, elem_classes=["step-card"]) as step2_grp:
+        gr.HTML("""
+        <div class="step-header">
+          <span class="step-num">2</span>
+          <span class="step-title">Select Task</span>
+        </div>""")
+        task_select = gr.Dropdown(choices=[], value=None, label="Task", interactive=True)
+
+    # ── Step 3 ────────────────────────────────────────────────────────────────
+    with gr.Group(visible=False, elem_classes=["step-card"]) as step3_grp:
+        gr.HTML("""
+        <div class="step-header">
+          <span class="step-num">3</span>
+          <span class="step-title">Configure Parameters</span>
+        </div>""")
+        with gr.Row():
+            n_episodes_slider = gr.Slider(
+                minimum=10, maximum=500, value=50, step=5,
+                label="Number of Episodes",
+                info="Total episodes in the dataset (10–500)",
             )
-
-        # ────────────────────────────────────────────────────────────────────
-        # STEP 3 — Parameters
-        # ────────────────────────────────────────────────────────────────────
-        with gr.Group(visible=False, elem_classes=["step-card"]) as step3_grp:
-            gr.HTML("""
-            <div class="step-header">
-              <span class="step-num">3</span>
-              <span class="step-title">Configure Parameters</span>
-            </div>""")
-
-            with gr.Row():
-                n_episodes_slider = gr.Slider(
-                    minimum=10, maximum=500, value=50, step=5,
-                    label="Number of Episodes",
-                    info="Total episodes in the dataset (10–500)",
-                )
-                success_slider = gr.Slider(
-                    minimum=0, maximum=100, value=70, step=5,
-                    label="Success Rate (%)",
-                    info="Fraction of episodes with successful trajectories",
-                )
-
-            with gr.Row():
-                force_min_slider = gr.Slider(
-                    minimum=0.1, maximum=10.0, value=1.0, step=0.1,
-                    label="Min Contact Force (N)",
-                    info="Lower bound of spring-damper contact force during grasping",
-                )
-                force_max_slider = gr.Slider(
-                    minimum=1.0, maximum=20.0, value=10.0, step=0.5,
-                    label="Max Contact Force (N)",
-                    info="Upper bound of contact force — higher = firmer grip",
-                )
-
-            gr.HTML("""
-            <div style="margin: 4px 0 8px;font-size:0.82rem;color:#8892a4;">
-              <b>Failure types to include</b> &nbsp;
-              <span style="font-style:italic;">
-                Grasp Slip — gripper opens mid-episode &nbsp;|&nbsp;
-                Velocity Spike — servo glitch (z&gt;6.5) &nbsp;|&nbsp;
-                Torque Saturation — joint hits angular limit
-              </span>
-            </div>""")
-
-            failure_check = gr.CheckboxGroup(
-                choices=["grasp_slip", "velocity_spike", "torque_saturation"],
-                value=["grasp_slip", "velocity_spike", "torque_saturation"],
-                label="",
-                elem_classes=["checkbox-group"],
+            success_slider = gr.Slider(
+                minimum=0, maximum=100, value=70, step=5,
+                label="Success Rate (%)",
+                info="Fraction of episodes with successful trajectories",
             )
-
-        # ────────────────────────────────────────────────────────────────────
-        # STEP 4 — Generate
-        # ────────────────────────────────────────────────────────────────────
-        with gr.Group(visible=False, elem_classes=["step-card"]) as step4_grp:
-            gr.HTML("""
-            <div class="step-header">
-              <span class="step-num">4</span>
-              <span class="step-title">Generate Dataset</span>
-            </div>""")
-
-            generate_btn = gr.Button(
-                "Generate Dataset",
-                elem_classes=["btn-generate"],
-                size="lg",
+        with gr.Row():
+            force_min_slider = gr.Slider(
+                minimum=0.1, maximum=10.0, value=1.0, step=0.1,
+                label="Min Contact Force (N)",
+                info="Lower bound of spring-damper contact force during grasping",
             )
-            gen_status = gr.Markdown("", elem_classes=["status-msg"])
-
-        # ────────────────────────────────────────────────────────────────────
-        # STEP 5 — Results dashboard
-        # ────────────────────────────────────────────────────────────────────
-        with gr.Group(visible=False, elem_classes=["step-card"]) as step5_grp:
-            gr.HTML("""
-            <div class="step-header">
-              <span class="step-num">5</span>
-              <span class="step-title">Quality Results</span>
-            </div>""")
-            results_html = gr.HTML("")
-
-        # ────────────────────────────────────────────────────────────────────
-        # STEP 6 — Email gate + Download
-        # ────────────────────────────────────────────────────────────────────
-        with gr.Group(visible=False, elem_classes=["step-card"]) as step6_grp:
-            gr.HTML("""
-            <div class="step-header">
-              <span class="step-num">6</span>
-              <span class="step-title">Download Dataset</span>
-            </div>
-            <div class="email-gate-note">
-              Enter your email to unlock the download. You'll receive occasional
-              updates on new robot configs and dataset improvements.
-            </div>""")
-
-            with gr.Row():
-                email_input = gr.Textbox(
-                    placeholder="you@example.com",
-                    label="Email",
-                    scale=4,
-                    max_lines=1,
-                )
-                email_btn = gr.Button(
-                    "Confirm →",
-                    elem_classes=["btn-primary"],
-                    scale=1,
-                )
-
-            email_status = gr.Markdown("", visible=True)
-
-            download_file = gr.File(
-                label="Download robogen_dataset.zip",
-                visible=False,
-                interactive=False,
+            force_max_slider = gr.Slider(
+                minimum=1.0, maximum=20.0, value=10.0, step=0.5,
+                label="Max Contact Force (N)",
+                info="Upper bound of contact force — higher = firmer grip",
             )
+        gr.HTML("""
+        <div style="margin:4px 0 8px;font-size:0.82rem;color:#8892a4;">
+          <b>Failure types to include</b> &nbsp;
+          <span style="font-style:italic;">
+            Grasp Slip — gripper opens mid-episode &nbsp;|&nbsp;
+            Velocity Spike — servo glitch (z&gt;6.5) &nbsp;|&nbsp;
+            Torque Saturation — joint hits angular limit
+          </span>
+        </div>""")
+        failure_check = gr.CheckboxGroup(
+            choices=["grasp_slip", "velocity_spike", "torque_saturation"],
+            value=["grasp_slip", "velocity_spike", "torque_saturation"],
+            label="",
+            elem_classes=["checkbox-group"],
+        )
 
-        # ════════════════════════════════════════════════════════════════════
-        # EVENT HANDLERS
-        # ════════════════════════════════════════════════════════════════════
-
-        # ── Step 1 → Step 2: Robot selected ──────────────────────────────
-        def on_robot_select(robot: str):
-            if not robot:
-                return (
-                    gr.update(visible=False),
-                    gr.update(choices=[], value=None),
-                    gr.update(visible=False),
-                    gr.update(visible=False),
-                    robot,
-                )
-            tasks_raw = TASKS_BY_ROBOT[robot]
-            tasks_disp = [(TASK_LABELS.get(t, t), t) for t in tasks_raw]
-            d = DEFAULTS.get(robot, DEFAULTS["SO-100"])
-            return (
-                gr.update(visible=True),                            # step2_grp
-                gr.update(choices=tasks_disp, value=tasks_raw[0]), # task_select
-                gr.update(visible=False),                           # step3_grp
-                gr.update(visible=False),                           # step4_grp
-                robot,                                              # robot_state
+    # ── Step 4 ────────────────────────────────────────────────────────────────
+    with gr.Group(visible=False, elem_classes=["step-card"]) as step4_grp:
+        gr.HTML("""
+        <div class="step-header">
+          <span class="step-num">4</span>
+          <span class="step-title">Generate Dataset</span>
+        </div>""")
+        generate_btn = gr.Button("Generate Dataset", elem_classes=["btn-generate"], size="lg")
+        gen_status   = gr.Markdown("", elem_classes=["status-msg"])
+
+    # ── Step 5 ────────────────────────────────────────────────────────────────
+    with gr.Group(visible=False, elem_classes=["step-card"]) as step5_grp:
+        gr.HTML("""
+        <div class="step-header">
+          <span class="step-num">5</span>
+          <span class="step-title">Quality Results</span>
+        </div>""")
+        results_html = gr.HTML("")
+
+    # ── Step 6 ────────────────────────────────────────────────────────────────
+    with gr.Group(visible=False, elem_classes=["step-card"]) as step6_grp:
+        gr.HTML("""
+        <div class="step-header">
+          <span class="step-num">6</span>
+          <span class="step-title">Download Dataset</span>
+        </div>
+        <div class="email-gate-note">
+          Enter your email to unlock the download. You'll receive occasional
+          updates on new robot configs and dataset improvements.
+        </div>""")
+        with gr.Row():
+            email_input = gr.Textbox(
+                placeholder="you@example.com", label="Email",
+                scale=4, max_lines=1,
             )
+            email_btn = gr.Button("Confirm", elem_classes=["btn-primary"], scale=1)
+        email_status = gr.Markdown("")
+        download_file = gr.File(label="Download robogen_dataset.zip", visible=False)
 
-        robot_select.change(
-            on_robot_select,
-            inputs=[robot_select],
-            outputs=[step2_grp, task_select, step3_grp, step4_grp, robot_state],
-        )
+    # ── Wire events ───────────────────────────────────────────────────────────
 
-        # ── Step 2 → Step 3: Task selected ───────────────────────────────
-        def on_task_select(task: str, robot: str):
-            if not task or not robot:
-                return (
-                    gr.update(visible=False),
-                    gr.update(visible=False),
-                    50, 70, 1.0, 10.0,
-                )
-            d = DEFAULTS.get(robot, DEFAULTS["SO-100"])
-            cfg_fr = ROBOT_CONFIG[robot]["force_range"]
-            return (
-                gr.update(visible=True),   # step3_grp
-                gr.update(visible=True),   # step4_grp
-                d["n_eps"],
-                d["success"],
-                cfg_fr[0],
-                cfg_fr[1],
-            )
+    robot_select.change(
+        fn=on_robot_select,
+        inputs=[robot_select],
+        outputs=[step2_grp, task_select, step3_grp, step4_grp, robot_state],
+        api_name=False,
+    )
 
-        task_select.change(
-            on_task_select,
-            inputs=[task_select, robot_state],
-            outputs=[
-                step3_grp, step4_grp,
-                n_episodes_slider, success_slider,
-                force_min_slider, force_max_slider,
-            ],
-        )
+    task_select.change(
+        fn=on_task_select,
+        inputs=[task_select, robot_state],
+        outputs=[step3_grp, step4_grp, n_episodes_slider, success_slider,
+                 force_min_slider, force_max_slider],
+        api_name=False,
+    )
 
-        # ── Step 4: Generate ─────────────────────────────────────────────
-        def on_generate(
-            robot, task, n_eps, success_pct, fmin, fmax, failures,
-            progress=gr.Progress(),
-        ):
-            if not robot or not task:
-                return (
-                    "Please complete steps 1 and 2 first.",
-                    gr.update(visible=False),
-                    gr.update(visible=False),
-                    gr.update(visible=False),
-                    None, None,
-                )
-            if not failures:
-                failures = list(FAILURE_TYPES)
-
-            try:
-                # ── Generation ──────────────────────────────────────────
-                def gen_progress(frac, msg):
-                    progress(frac * 0.65, desc=msg)
-
-                progress(0.0, desc="Generating episodes…")
-                df = generate_dataset(
-                    robot=robot, task=task,
-                    n_episodes=int(n_eps),
-                    success_rate=success_pct / 100,
-                    force_min=float(fmin), force_max=float(fmax),
-                    enabled_failures=list(failures),
-                    seed=None,
-                    progress_callback=gen_progress,
-                )
-
-                # ── Scoring ─────────────────────────────────────────────
-                progress(0.70, desc="Running quality checks…")
-
-                def score_progress(frac, msg):
-                    progress(0.70 + frac * 0.20, desc=msg)
-
-                result = score_dataset(df, progress_callback=score_progress)
-
-                progress(0.92, desc="Preparing results…")
-                results_panel = _make_results_html(result, robot, task)
-                progress(1.0, desc="Done")
-
-                status = (
-                    f"Generated {len(df):,} rows across {result['n_episodes']} episodes — "
-                    f"score **{result['overall_score']:.1f}/100** ({result['band']})"
-                )
-
-                return (
-                    status,
-                    gr.update(visible=True),   # step5_grp
-                    results_panel,             # results_html
-                    gr.update(visible=True),   # step6_grp
-                    df,                        # df_state
-                    result,                    # result_state
-                )
-
-            except Exception:
-                err = traceback.format_exc()
-                return (
-                    f"Generation failed:\n```\n{err}\n```",
-                    gr.update(visible=False),
-                    "",
-                    gr.update(visible=False),
-                    None, None,
-                )
-
-        generate_btn.click(
-            on_generate,
-            inputs=[
-                robot_state, task_select,
-                n_episodes_slider, success_slider,
-                force_min_slider, force_max_slider,
-                failure_check,
-            ],
-            outputs=[
-                gen_status,
-                step5_grp, results_html,
-                step6_grp,
-                df_state, result_state,
-            ],
-        )
+    generate_btn.click(
+        fn=on_generate,
+        inputs=[robot_state, task_select, n_episodes_slider, success_slider,
+                force_min_slider, force_max_slider, failure_check],
+        outputs=[gen_status, step5_grp, results_html, step6_grp, df_state, result_state],
+        api_name=False,
+    )
 
-        # ── Step 6: Email gate → unlock download ──────────────────────────
-        def on_email_submit(
-            email: str,
-            robot: str,
-            task: str,
-            n_eps: float,
-            success_pct: float,
-            fmin: float,
-            fmax: float,
-            failures: List[str],
-            df: Optional[pd.DataFrame],
-            result: Optional[Dict],
-        ):
-            if not email or "@" not in email:
-                return (
-                    "Please enter a valid email address.",
-                    gr.update(visible=False),
-                )
-            if df is None or result is None:
-                return (
-                    "Generate a dataset first (Step 4).",
-                    gr.update(visible=False),
-                )
-
-            # Fire Airtable (failure is non-blocking)
-            try:
-                ok, msg = log_email(
-                    email=email.strip(),
-                    robot=robot, task=task,
-                    n_episodes=int(n_eps),
-                    quality_score=result["overall_score"],
-                    band=result["band"],
-                )
-                if not ok:
-                    print(f"[RoboGen] Airtable log failed: {msg}")
-            except Exception as exc:
-                print(f"[RoboGen] Airtable exception: {exc}")
-
-            # Build download zip regardless of Airtable outcome
-            try:
-                zip_path = _build_zip(
-                    df=df, result=result,
-                    robot=robot, task=task,
-                    n_eps=int(n_eps), success=success_pct,
-                    fmin=float(fmin), fmax=float(fmax),
-                    failures=list(failures),
-                )
-                return (
-                    "Email confirmed. Your download is ready below.",
-                    gr.update(visible=True, value=zip_path),
-                )
-            except Exception:
-                err = traceback.format_exc()
-                return (
-                    f"Download preparation failed:\n```\n{err}\n```",
-                    gr.update(visible=False),
-                )
-
-        email_btn.click(
-            on_email_submit,
-            inputs=[
-                email_input,
-                robot_state, task_select,
+    email_btn.click(
+        fn=on_email_submit,
+        inputs=[email_input, robot_state, task_select,
                 n_episodes_slider, success_slider,
                 force_min_slider, force_max_slider,
-                failure_check,
-                df_state, result_state,
-            ],
-            outputs=[email_status, download_file],
-        )
-
-    return demo
+                failure_check, df_state, result_state],
+        outputs=[email_status, download_file],
+        api_name=False,
+    )
 
+# ── Launch ────────────────────────────────────────────────────────────────────
 
-# ── Entry point ───────────────────────────────────────────────────────────────
+demo.queue()
 
 if __name__ == "__main__":
-    app = build_app()
-    app.queue()
-    app.launch(
-        server_name="0.0.0.0",
-        server_port=int(os.environ.get("PORT", 7860)),
-        show_error=True,
-    )
+    demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))