foxhunt/scripts/compare-rl-state.py

#!/usr/bin/env python3
"""Compare raw RL-loop state dumps between two determinism runs.

Reads the dumps produced by
`IntegratedTrainer::dump_rl_state_for_debug` (env-gated on
`FOXHUNT_DETERMINISM_DEBUG_RL=1`) and prints a per-step
exact-bytewise-equality verdict for each of the 6 candidate groups
identified by Phase 2.4's mamba2 falsification:

  Group C (action sampling RNG):
    - action_prng_state (B u32)
    - iqn_prng_state    (B u32)

  Group D (LOB simulator + unit state):
    - pyramid_units_count, unit_entry_step/_price/_lots/_active/...
    - prev_realized_pnl, prev_position_lots, steps_since_done
    - trade_duration_emit, close_unit_index, unit_prev_pos_lots,
      unit_initial_r, unit_trail_distance

  Group F (reward shaping + RL outputs):
    - rewards, raw_rewards, outcome_ema, reward_abs, dones
    - actions, next_actions, log_pi_old, advantages, returns
    - CMDP/edge-decay per-batch: session_pnl, consec_loss,
      cooldown_remaining, ph_mu/_count/_m/_mmin/_stat

  Group E (full ISV array — host mapped-pinned, RL_SLOTS_END floats):
    - isv_full (757 f32)

Group A (CfC carry `attn_context_d` / `h_t_d`) and Group B (PER state
priority_tree / sample_prng / sample_indices) are dumped by the
mamba2 + per debug dumpers respectively — Phase 2.5 reuses those.

Verdict logic:
  - All groups EQUAL at all dumped steps → bug is elsewhere (Phase 2.6).
  - Group C diverges → action sampling RNG path is non-deterministic.
  - Group D diverges → LOB sim is the culprit.
  - Group E diverges → an ISV controller is non-deterministic.
  - Group F diverges → reward shaping / advantage / log_pi pipeline
                       is the culprit.
  - Multiple groups diverge → ranked by FIRST step they diverge at;
                              the earliest one is upstream.

Exit codes:
  0  both runs identical for all dumped buffers at all dumped steps
  1  divergence found; verdict printed naming the candidate group
  2  inputs missing or malformed

Usage:
  python3 scripts/compare-rl-state.py <run_a_dir> <run_b_dir>

Each dir contains `step_{0,1,2,3}_<name>.bin`.

Per `feedback_cpu_is_read_only`: this script ONLY reads device dumps;
it performs no compute that affects training.
"""
from __future__ import annotations

import argparse
import struct
import sys
from pathlib import Path


# Each entry: (logical name, on-disk basename, element struct format,
# element size in bytes). The groups partition the candidate space
# so the verdict can name the culprit class precisely.
GROUP_C_RNG = [
    ("action_prng_state", "action_prng_state", "I", 4),
    ("iqn_prng_state",    "iqn_prng_state",    "I", 4),
]
GROUP_D_LOBSIM = [
    ("pyramid_units_count", "pyramid_units_count", "i", 4),
    ("unit_entry_step",     "unit_entry_step",     "i", 4),
    ("unit_entry_price",    "unit_entry_price",    "f", 4),
    ("unit_lots",           "unit_lots",           "i", 4),
    ("unit_active",         "unit_active",         "B", 1),
    ("unit_initial_r",      "unit_initial_r",      "f", 4),
    ("unit_trail_distance", "unit_trail_distance", "f", 4),
    ("unit_prev_pos_lots",  "unit_prev_pos_lots",  "i", 4),
    ("close_unit_index",    "close_unit_index",    "i", 4),
    ("prev_realized_pnl",   "prev_realized_pnl",   "f", 4),
    ("prev_position_lots",  "prev_position_lots",  "i", 4),
    ("steps_since_done",    "steps_since_done",    "i", 4),
    ("trade_duration_emit", "trade_duration_emit", "f", 4),
]
GROUP_F_REWARDS = [
    ("rewards",             "rewards",             "f", 4),
    ("raw_rewards",         "raw_rewards",         "f", 4),
    ("outcome_ema",         "outcome_ema",         "f", 4),
    ("reward_abs",          "reward_abs",          "f", 4),
    ("dones",               "dones",               "f", 4),
    ("actions",             "actions",             "i", 4),
    ("next_actions",        "next_actions",        "i", 4),
    ("log_pi_old",          "log_pi_old",          "f", 4),
    ("advantages",          "advantages",          "f", 4),
    ("returns",             "returns",             "f", 4),
    # CMDP + edge-decay
    ("session_pnl",         "session_pnl",         "f", 4),
    ("consec_loss",         "consec_loss",         "f", 4),
    ("cooldown_remaining",  "cooldown_remaining",  "f", 4),
    ("ph_mu",               "ph_mu",               "f", 4),
    ("ph_count",            "ph_count",            "f", 4),
    ("ph_m",                "ph_m",                "f", 4),
    ("ph_mmin",             "ph_mmin",             "f", 4),
    ("ph_stat",             "ph_stat",             "f", 4),
]
GROUP_E_ISV = [
    ("isv_full",            "isv_full",            "f", 4),
]

ALL_GROUPS = [
    ("C (action RNG)",      GROUP_C_RNG),
    ("D (LOB sim)",         GROUP_D_LOBSIM),
    ("E (ISV controllers)", GROUP_E_ISV),
    ("F (reward shaping)",  GROUP_F_REWARDS),
]


def read_typed(path: Path, fmt: str, elem_size: int) -> list:
    """Read a raw little-endian binary file as `fmt`-typed elements."""
    data = path.read_bytes()
    if len(data) % elem_size != 0:
        raise ValueError(
            f"{path}: size {len(data)} not a multiple of {elem_size}"
        )
    n = len(data) // elem_size
    return list(struct.unpack(f"<{n}{fmt}", data))


def compare_lists(a: list, b: list, label: str) -> tuple[bool, int, str]:
    """Bytewise exact equality. Returns (equal, first_diff_idx, summary)."""
    if len(a) != len(b):
        return False, -1, f"{label}: length mismatch A={len(a)} B={len(b)}"
    diffs: list[int] = []
    max_abs_delta = 0.0
    for i, (x, y) in enumerate(zip(a, b)):
        if x != y:
            diffs.append(i)
            try:
                ad = abs(float(x) - float(y))
                if ad > max_abs_delta:
                    max_abs_delta = ad
            except (TypeError, ValueError):
                pass
        if len(diffs) >= 5:
            break
    if not diffs:
        return True, -1, f"{label}: EQUAL ({len(a)} elements)"
    sample = ", ".join(
        f"[{i}]: A={a[i]!r} B={b[i]!r}" for i in diffs[:3]
    )
    return False, diffs[0], (
        f"{label}: DIVERGE at idx {diffs[0]} "
        f"max|Δ|={max_abs_delta:.6g} "
        f"(showing up to 3): {sample}"
    )


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("run_a", type=Path)
    parser.add_argument("run_b", type=Path)
    args = parser.parse_args()

    if not args.run_a.is_dir() or not args.run_b.is_dir():
        print(
            f"ERROR: run dirs missing: {args.run_a} or {args.run_b}",
            file=sys.stderr,
        )
        return 2

    # Per-step verdict. Tracks the FIRST step that diverges and which
    # group diverged there.
    overall_any_diverge = False
    # Map (step, group_label) -> list of (name, ok, summary).
    per_step_group_results: dict[tuple[int, str], list] = {}
    first_diverge_step: int | None = None
    first_diverge_groups: set[str] = set()

    for step in (0, 1, 2, 3):
        print(f"=== step {step} ===")
        any_present = False
        step_had_diverge = False
        step_diverged_groups: set[str] = set()

        for group_label, members in ALL_GROUPS:
            print(f"  -- group {group_label} --")
            group_results: list[tuple[str, bool, str]] = []
            group_any_present = False
            for name, basename, fmt, elem_size in members:
                pa = args.run_a / f"step_{step}_{basename}.bin"
                pb = args.run_b / f"step_{step}_{basename}.bin"
                if not (pa.is_file() and pb.is_file()):
                    continue
                group_any_present = True
                any_present = True
                try:
                    a = read_typed(pa, fmt, elem_size)
                    b = read_typed(pb, fmt, elem_size)
                except ValueError as e:
                    print(f"    ERROR reading {name}: {e}")
                    return 2
                equal, _, summary = compare_lists(a, b, name)
                print(f"    {summary}")
                group_results.append((name, equal, summary))
                if not equal:
                    overall_any_diverge = True
                    step_had_diverge = True
                    step_diverged_groups.add(group_label)

            if not group_any_present:
                print("    SKIP — no buffers dumped for this group/step")
            per_step_group_results[(step, group_label)] = group_results

        if not any_present:
            print("  SKIP — no buffers dumped for this step")
            continue

        if step_had_diverge and first_diverge_step is None:
            first_diverge_step = step
            first_diverge_groups = step_diverged_groups

    print()
    if not overall_any_diverge:
        print(
            "VERDICT: all dumped RL-loop buffers match across all dumped\n"
            "steps. The bug is NOT in any of the 6 candidate groups (CfC\n"
            "carry / PER state / action RNG / LOB sim / ISV / reward\n"
            "shaping). Phase 2.6 dispatch needed — STOP per\n"
            "feedback_systematic_debugging."
        )
        return 0

    assert first_diverge_step is not None
    print(f"VERDICT (first divergent step = {first_diverge_step}):")
    print(f"  diverged group(s): {sorted(first_diverge_groups)}")
    print()
    print("Per-group culprit table:")
    for group_label, _ in ALL_GROUPS:
        results = per_step_group_results.get((first_diverge_step, group_label), [])
        any_div = any(not eq for _, eq, _ in results)
        if any_div:
            divnames = [name for name, eq, _ in results if not eq]
            print(f"  Group {group_label}: DIVERGE in {divnames}")
        else:
            present = len(results) > 0
            tag = "EQUAL" if present else "(no dump)"
            print(f"  Group {group_label}: {tag}")
    print()
    # Print interpretive hint based on which group(s) diverged FIRST.
    print("Next-step interpretation:")
    if "C (action RNG)" in first_diverge_groups:
        print("  Group C (action RNG) — likely culprit if the only group.")
        print("    Fix: xorshift32 update in rl_action_kernel must read")
        print("    same upstream Q/π logits across runs. If upstream Q is")
        print("    determined to be EQUAL but action_prng_state diverges,")
        print("    the PRNG was perturbed by branchy logic (e.g.,")
        print("    conditional advance based on action choice).")
    if "D (LOB sim)" in first_diverge_groups:
        print("  Group D (LOB sim) — LOB sim is producing different state.")
        print("    Inspect kernels in crates/ml-backtesting/src/sim/ for")
        print("    atomics, parallel reductions, order-dependent writes.")
    if "E (ISV controllers)" in first_diverge_groups:
        print("  Group E (ISV controllers) — one controller is non-det.")
        print("    Compare ISV slot ranges: popart_sigma_welford (725),")
        print("    kelly avg_win/loss_ema, wr_ema, ph_mean / edge-decay.")
        print("    Look for parallel reductions in ema_update_*,")
        print("    rl_kelly_*, rl_popart_*. Apply Phase-2-PER-rebuild fix")
        print("    pattern (single-block + __syncthreads + fixed order).")
    if "F (reward shaping)" in first_diverge_groups:
        print("  Group F (reward shaping) — rewards / advantages / log_pi")
        print("    diverge. Check rl_fused_reward_pipeline.cu,")
        print("    compute_advantage_return, apply_reward_scale,")
        print("    log_pi_at_action for non-det parallel writes.")
        print("    If `rewards` diverges with EQUAL action_prng + LOB state,")
        print("    the reward kernel itself is non-deterministic.")
        print("    If `actions` diverges → upstream (action sampling) bug.")
        print("    If `advantages` diverges with EQUAL `rewards` →")
        print("    `compute_advantage_return` reduction is non-det.")
    return 1


if __name__ == "__main__":
    sys.exit(main())