scripts/gen_bc_dataset.py

#!/usr/bin/env python3
"""Sinh dataset BC từ oracle grid search (stage 2a — oracle-guided warm start).

Bối cảnh (16/07/2026): chuỗi 3 giả thuyết về Q|pot dính 0.53 đều bị bác;
smoking gun: agent chọn NGẪU NHIÊN trong tập pot combo (Q combo random
0.524 ≈ agent 0.522 ≈ blind 0.531), trong khi trần thật 0.869. Lối ra
(Documents/PoolCoach_Stage2_BC_WarmStart_Design.md): oracle grid search =
expert demonstrations rẻ → behavior cloning → PPO fine-tune.

Mỗi bàn: grid search Y HỆT oracle_controllability (phi ghost-ball mọi lỗ
khả thi × V0 × side × vert), label = action của combo (pot && !scratch) có
Q TỐT NHẤT — KỂ CẢ b2-potted (Q=1, khớp đúng reward env đang tối ưu).
Bàn 0 pot combo (hiếm; 0/100 ở oracle run 16/07): bỏ qua.

Action label đổi về [-1,1]^4 — nghịch đảo đúng map trong env.step:
    a0 = phi/180 - 1          (phi ∈ [0,360))
    a1 = (V0 - 0.5)/3.5*2 - 1 (V0 ∈ [0.5, 4.0])
    a2 = side/0.4,  a3 = vert/0.4

Chạy từ gốc repo (Numba JIT ~40s/worker lúc khởi động):
    python scripts/gen_bc_dataset.py --tables 2000 --seed 123 --workers 4   # v1 (~35 phút)
    python scripts/gen_bc_dataset.py --tables 20 --workers 2                # smoke
    python scripts/gen_bc_dataset.py --tables 20000 --seed 124 --workers 4  # run chính (qua đêm)

Seed 123 ≠ 42 (oracle run) để không trùng 100 bàn đã phân tích.

Output: data/bc_dataset_<N>_<seed>.npz — 2 BỘ LABEL trong 1 lần chạy
(smoke 20/07: b2-lucky chiếm ~50% argmax → lưu cả 2, train_bc chọn):
    obs (N,6) float32         — chuẩn hoá [0,1] y hệt env._obs
    actions (N,4), q (N,)     — argmax Q KỂ CẢ b2-potted (= reward env)
    actions_xb2 (N,4), q_xb2  — argmax Q LOẠI b2-lucky (điều bi "thật")
    n_pot (N,), b2_lucky (N,) + grid metadata
    combos (M,6) [phi,v0,side,vert,q,b2p], combo_row (M,) — MỌI pot combo,
        nhiên liệu cho relabel_bc_dataset.py (BC v1 fail G1 vì multi-modality
        → đổi label rule post-hoc, không cần re-simulate)
+ data/bc_dataset_<N>_<seed>_sanity.png (hist Q label + action components)
"""

from __future__ import annotations

import argparse
import multiprocessing as mp
import sys
import time
from pathlib import Path

sys.stdout.reconfigure(encoding="utf-8")
sys.stderr.reconfigure(encoding="utf-8")

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))
sys.path.insert(0, str(ROOT / "scripts"))

# Tái dùng grid-search core của oracle (checklist §9 — pattern import thẳng,
# không refactor script oracle đã ra kết quả 16/07 để tránh regression):
# _init_worker (env + JIT warmup per worker), _simulate_shot, sample_tables.
import oracle_controllability as oc


def _to_action(phi, v0, side, vert):
    """Nghịch đảo map của env.step → [-1,1]^4 (verify round-trip 20/07)."""
    import numpy as np

    return np.clip(
        np.array([phi / 180.0 - 1.0,
                  (v0 - 0.5) / 3.5 * 2.0 - 1.0,
                  side / 0.4,
                  vert / 0.4], dtype=np.float32),
        -1.0, 1.0)


def _label_table(args):
    """Grid search 1 bàn → 2 bộ label BC (gồm b2-lucky / loại b2-lucky).

    Smoke 20/07: b2-lucky (Q=1) chiếm ~50% label argmax — lưu CẢ HAI trong
    1 lần chạy để train_bc chọn (--labels b2|xb2), ablation cho luận văn.

    Trả (idx, obs6, act_b2, q_b2, act_xb2, q_xb2, n_pot, b2_lucky, n_sims);
    obs = None nếu bàn không có pot combo nào (bỏ qua bàn đó).
    Chạy trong worker — oc._ENV/_GRIDS đã được oc._init_worker dựng sẵn.
    """
    import numpy as np

    idx, cue_xy, b1_xy, b2_xy = args
    v0_grid, side_grid, vert_grid, _ = oc._GRIDS
    cue_xy, b1_xy, b2_xy = map(np.asarray, (cue_xy, b1_xy, b2_xy))

    # phi ứng viên: ghost-ball của mọi lỗ khả thi (cùng logic oracle/aim any)
    phis = [
        float(np.degrees(np.arctan2(d[1], d[0])) % 360.0)
        for d in oc._ENV._ghost_dirs_any(cue_xy, b1_xy)
    ]

    n_sims = n_pot = 0
    rows = []  # MỌI pot combo (phi, v0, side, vert, q, b2p) — lưu vào npz
               # để relabel_bc_dataset.py đổi label rule không cần re-simulate
    best_q, best_act, best_b2 = -1.0, None, False   # argmax KỂ CẢ b2 (= env)
    best_xq, best_xact = -1.0, None                 # argmax LOẠI b2-lucky
    for phi in phis:
        for v0 in v0_grid:
            for a in side_grid:
                for b in vert_grid:
                    n_sims += 1
                    potted, scratch, b2p, q = oc._simulate_shot(
                        cue_xy, b1_xy, b2_xy, phi, float(v0), float(a), float(b))
                    if not (potted and not scratch):
                        continue
                    n_pot += 1
                    rows.append((phi, float(v0), float(a), float(b),
                                 q, float(b2p)))
                    if q > best_q:
                        best_q, best_b2 = q, bool(b2p)
                        best_act = (phi, float(v0), float(a), float(b))
                    if not b2p and q > best_xq:
                        best_xq = q
                        best_xact = (phi, float(v0), float(a), float(b))

    if best_act is None:
        return idx, None, None, 0.0, None, 0.0, 0, False, n_sims, None

    # obs chuẩn hoá [0,1] — y hệt PositionPlayEnv._obs (vị trí BAN ĐẦU)
    w, l = oc._ENV.w, oc._ENV.l
    obs = np.clip(
        np.array([cue_xy[0] / w, cue_xy[1] / l,
                  b1_xy[0] / w, b1_xy[1] / l,
                  b2_xy[0] / w, b2_xy[1] / l], dtype=np.float32),
        0.0, 1.0)

    if best_xact is None:  # hiếm: MỌI pot combo đều b2-lucky → fallback b2
        best_xact, best_xq = best_act, best_q
    return (idx, obs, _to_action(*best_act), float(best_q),
            _to_action(*best_xact), float(best_xq), n_pot, best_b2, n_sims,
            np.array(rows, dtype=np.float32))


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--tables", type=int, default=2000)
    p.add_argument("--seed", type=int, default=123,
                   help="seed sample bàn; ≠ 42 (oracle run 16/07)")
    p.add_argument("--v0-steps", type=int, default=10,
                   help="số mức V0 trong [0.5, 4.0] (khớp oracle)")
    p.add_argument("--spin-steps", type=int, default=5,
                   help="số mức side/vert trong [-0.4, 0.4]; nên LẺ để có 0")
    p.add_argument("--workers", type=int, default=4)
    p.add_argument("--out", default=None,
                   help="mặc định data/bc_dataset_<N>_<seed>.npz")
    args = p.parse_args()

    import numpy as np

    out = (Path(args.out) if args.out
           else ROOT / "data" / f"bc_dataset_{args.tables}_{args.seed}.npz")
    out.parent.mkdir(parents=True, exist_ok=True)

    v0_grid = np.linspace(0.5, 4.0, args.v0_steps)
    side_grid = np.linspace(-0.4, 0.4, args.spin_steps)
    vert_grid = np.linspace(-0.4, 0.4, args.spin_steps)

    tables = oc.sample_tables(args.tables, args.seed)
    per_pocket = args.v0_steps * args.spin_steps ** 2
    print(f"== Sinh BC dataset: {args.tables} bàn, seed {args.seed}, "
          f"~{per_pocket} sim/lỗ khả thi (TB ~2.2 lỗ/bàn) ==")
    print(f"   grid: V0 {args.v0_steps} mức x side/vert {args.spin_steps} mức "
          f"— label = argmax Q (kể cả b2-potted)")
    print(f"   {args.workers} worker — JIT Numba ~40s lúc khởi động...\n")

    t0 = time.time()
    results = {}
    skipped = total_sims = 0
    step = max(1, args.tables // 100)
    with mp.Pool(args.workers, initializer=oc._init_worker,
                 initargs=(v0_grid, side_grid, vert_grid, 0.0)) as pool:
        for k, (idx, obs, act, q, act_x, q_x, n_pot, b2, n_sims,
                comb) in enumerate(
                pool.imap_unordered(_label_table, tables), 1):
            total_sims += n_sims
            if obs is None:
                skipped += 1
            else:
                results[idx] = (obs, act, q, act_x, q_x, n_pot, b2, comb)
            if k % step == 0 or k == len(tables):
                el = time.time() - t0
                eta = el / k * (len(tables) - k)
                print(f"  {k}/{len(tables)} bàn — {len(results)} label, "
                      f"{skipped} bỏ qua   [{el/60:.1f} phút, còn ~{eta/60:.1f}]")

    el = time.time() - t0
    print(f"\nXong {total_sims} sim trong {el/60:.1f} phút "
          f"({total_sims/el:.0f} sim/s)")
    if not results:
        print("KHÔNG có bàn nào pot được — kiểm tra lại grid/env!")
        return

    order = sorted(results)
    obs_arr = np.stack([results[i][0] for i in order])
    act_arr = np.stack([results[i][1] for i in order])
    q_arr = np.array([results[i][2] for i in order], dtype=np.float32)
    actx_arr = np.stack([results[i][3] for i in order])
    qx_arr = np.array([results[i][4] for i in order], dtype=np.float32)
    npot_arr = np.array([results[i][5] for i in order], dtype=np.int32)
    b2_arr = np.array([results[i][6] for i in order], dtype=np.int8)

    # MỌI pot combo — nhiên liệu cho relabel_bc_dataset.py (fail path §7:
    # đổi label rule là post-processing, không phải re-simulate 45 phút)
    combos = np.concatenate([results[i][7] for i in order])
    combo_row = np.concatenate([
        np.full(len(results[i][7]), r, dtype=np.int32)
        for r, i in enumerate(order)
    ])

    np.savez_compressed(out, obs=obs_arr, actions=act_arr, q=q_arr,
                        actions_xb2=actx_arr, q_xb2=qx_arr,
                        n_pot=npot_arr, b2_lucky=b2_arr,
                        combos=combos, combo_row=combo_row,
                        seed=np.int64(args.seed), v0_grid=v0_grid,
                        side_grid=side_grid, vert_grid=vert_grid)
    print(f"Dataset -> {out}  ({len(order)} sample, 2 bộ label b2+xb2, "
          f"{len(combos)} pot combo lưu kèm cho relabel)")

    # ------------------------------------------------------- sanity stats
    v0_lbl = 0.5 + (act_arr[:, 1] + 1.0) / 2.0 * 3.5
    wrap = float(np.mean(np.abs(act_arr[:, 0]) > 0.95))
    print(f"\n== Sanity (kỳ vọng theo design doc §8) ==")
    print(f"  Q label b2   : mean {q_arr.mean():.3f} | "
          f"median {np.median(q_arr):.3f} | "
          f"p10/p90 {np.percentile(q_arr, 10):.3f}/{np.percentile(q_arr, 90):.3f} | "
          f"b2 lucky (Q=1): {float(b2_arr.mean()):.1%}")
    print(f"  Q label xb2  : mean {qx_arr.mean():.3f} "
          f"(kỳ vọng ~0.87 = trần thật oracle 16/07) | "
          f"median {np.median(qx_arr):.3f}")
    print(f"  Q > 0.7 (b2) : {float(np.mean(q_arr > 0.7)):.1%}   |   "
          f"Q > 0.7 (xb2): {float(np.mean(qx_arr > 0.7)):.1%}")
    print(f"  pot combo/bàn: mean {npot_arr.mean():.1f} "
          f"(multi-modality material — design §7)")
    print(f"  V0 label     : mean {v0_lbl.mean():.2f} m/s | "
          f"mức max (4.0): {float(np.mean(v0_lbl > 3.9)):.1%} | "
          f"mức min (0.5): {float(np.mean(v0_lbl < 0.6)):.1%}")
    print(f"  |side|/|vert|: {np.abs(act_arr[:, 2]).mean():.2f} / "
          f"{np.abs(act_arr[:, 3]).mean():.2f} (0-1) | "
          f"spin != 0: {float(np.mean((act_arr[:, 2] != 0) | (act_arr[:, 3] != 0))):.1%}")
    print(f"  phi wraparound risk (|a0| > 0.95): {wrap:.1%} "
          f"(cao → cân nhắc label cos/sin, design §7)")

    # ------------------------------------------------------- sanity plot
    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    fig, axes = plt.subplots(2, 3, figsize=(14, 8))

    ax = axes[0][0]
    ax.hist(q_arr, bins=np.linspace(0, 1, 21), alpha=0.6,
            label=f"b2 (mean {q_arr.mean():.3f})")
    ax.hist(qx_arr, bins=np.linspace(0, 1, 21), alpha=0.6,
            label=f"xb2 (mean {qx_arr.mean():.3f})")
    ax.set_title("Q label — xb2 kỳ vọng mean ~0.87")
    ax.legend(loc="upper left")

    ax = axes[0][1]
    ax.hist((act_arr[:, 0] + 1.0) * 180.0, bins=36, alpha=0.75)
    ax.set_title("phi label (độ) — kỳ vọng trải đều")

    def bar_levels(ax, values, grid, title):
        lv, cnt = np.unique(np.round(values, 3), return_counts=True)
        width = 0.6 * (grid[1] - grid[0]) if len(grid) > 1 else 0.1
        ax.bar(lv, cnt, width=width, alpha=0.75)
        ax.set_title(title)

    bar_levels(axes[0][2], v0_lbl, v0_grid, "V0 label (m/s) — có đa dạng không?")
    bar_levels(axes[1][0], act_arr[:, 2] * 0.4, side_grid, "side label")
    bar_levels(axes[1][1], act_arr[:, 3] * 0.4, vert_grid, "vert label")

    ax = axes[1][2]
    ax.hist(npot_arr, bins=30, alpha=0.75)
    ax.set_title("pot combo / bàn")

    fig.suptitle(f"BC dataset sanity — {len(order)} sample, seed {args.seed}")
    fig.tight_layout()
    png = out.with_name(out.stem + "_sanity.png")
    fig.savefig(png, dpi=120)
    print(f"\nSanity plot -> {png}")
    print(f"\nBước kế: python scripts/train_bc.py --dataset {out.relative_to(ROOT)}")


if __name__ == "__main__":
    main()