qwenjev / train /evaluate_local.py

v0.4.9: vision-aware pooling (RCA fix) + v4 namespace + eps0.45 collect config

66bc5b7 verified 28 days ago

6.44 kB

"""Local evaluation: run the hybrid agent on public games and report
per-game scores in the spirit of the official RHAE metric.
    per-level = min(human/agent, 1.0) ^ 2      (human baseline unknown locally;
                                                we report completion + actions)
    game      = level-weighted average (1-indexed weights)
    total     = average over games
Without official human baselines this is a RELATIVE tool: compare agent
variants against each other, never treat absolute numbers as leaderboard
predictions.
"""

from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

import yaml  # noqa: E402

from data.collect_selfplay import PUBLIC_GAMES_FALLBACK, grids_from, make_env, meta_from  # noqa: E402
from data.mock_env import step_env  # noqa: E402
from qwenjev.hybrid import HybridAgent, HybridConfig  # noqa: E402
from qwenjev.system1 import QwenJevModel, System1Config  # noqa: E402
from qwenjev.system2 import System2Config, System2Controller  # noqa: E402
from qwenjev.vision import TinyVLMConfig  # noqa: E402


def run_game(env, game_id: str, brain: HybridAgent, max_actions: int) -> dict:
    """Play one episode against any environment (mock or official).
    Lifecycle: WIN ends the episode; GAME_OVER triggers the hybrid brain's
    episode-boundary hook and a fresh reset, then play continues within the
    same action budget — identical budget for every agent variant, so
    best_level/actions stay comparable across runs."""
    env.reset()
    actions = 0
    resets = 0
    best_level = 0
    final_state = "NOT_FINISHED"
    started = time.time()
    for _ in range(max_actions):
        meta = meta_from(env)
        final_state = meta["state"]
        if final_state == "WIN":
            break
        if final_state == "GAME_OVER":
            brain.on_game_over(game_id, [f"game over at action {actions}"])
            env.reset()
            resets += 1
            continue
        grids = grids_from(env)
        best_level = max(best_level, meta["levels_completed"])
        cand = brain.decide(
            game_id, grids,
            levels_completed=meta["levels_completed"],
            win_levels=meta["win_levels"],
            available_actions=meta["available_actions"],
            game_state=meta["state"],
        )
        if cand.is_complex:
            step_env(env, cand.name, int(cand.x or 0), int(cand.y or 0))
        else:
            step_env(env, cand.name)
        actions += 1
        best_level = max(best_level, meta_from(env)["levels_completed"])
    return {
        "game_id": game_id,
        "win": final_state == "WIN",
        "final_state": final_state,
        "best_level": best_level,
        "resets": resets,
        "actions": actions,
        "seconds": round(time.time() - started, 1),
    }


def build_s1(cfg: dict):
    """System-1 from the yaml config; checkpoint loaded when present."""
    s1_cfg = cfg.get("system1", {})
    from pathlib import Path

    # v0.4.9: pool_vision auto-detection. The saved qwenjev_s1.pt config
    # carries the flag from v0.4.9 on; ABSENCE = v3-era checkpoint =
    # last-position pooling (heads built with matching H input).
    pool_vision = s1_cfg.get("pool_vision")
    if pool_vision is None:
        pool_vision = False
        try:
            import torch as _torch

            ck = (Path(cfg.get("out_dir", "checkpoints/qwenjev-s1"))
                  / "qwenjev_s1.pt")
            if ck.exists():
                payload = _torch.load(ck, map_location="cpu",
                                      weights_only=False)
                pool_vision = bool(
                    (payload.get("config") or {}).get("pool_vision", False))
        except Exception:
            pool_vision = False
    model_cfg = System1Config(
        backend=s1_cfg.get("backend", "tiny"),
        hf_model=s1_cfg.get("hf_model", "Qwen/Qwen3-VL-2B-Instruct"),
        max_leaf_tokens=int(s1_cfg.get("max_leaf_tokens", 512)),
        score_levels=int(s1_cfg.get("score_levels", 10)),
        leaf_batch_size=int(s1_cfg.get(
            "eval_leaf_batch", s1_cfg.get("leaf_batch_size", 32))),
        pool_vision=bool(pool_vision),
        tiny=TinyVLMConfig(**(s1_cfg.get("tiny") or {})),
    )
    s1 = QwenJevModel(model_cfg).build()
    ckpt = Path(cfg.get("out_dir", "checkpoints/qwenjev-s1")) / "qwenjev_s1.pt"
    if ckpt.exists():
        s1.load_heads(ckpt.parent)
        # stage B/C continuity: merge the SFT adapters into the fp16 base
        # so evaluation runs on the SAME hidden states the heads saw
        if model_cfg.backend == "hf":
            from qwenjev.hf_compat import resume_adapters

            hf_lora = ckpt.parent / "hf_lora"
            if hf_lora.exists():
                s1._backbone, _merged = resume_adapters(
                    s1._backbone, str(hf_lora), merge=True)
                print(f"stage B LoRA merged -> {hf_lora}")
        print(f"System-1 checkpoint loaded -> {ckpt}")
    return s1


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", default="configs/train.yaml")
    ap.add_argument("--games", nargs="*", default=None)
    ap.add_argument("--out", default="eval_report.json")
    args = ap.parse_args()
    cfg = yaml.safe_load(open(args.config))

    s1 = None
    try:
        s1 = build_s1(cfg)
    except Exception as exc:
        print(f"[warn] System-1 unavailable ({exc}); heuristic policy will run")

    s2 = None
    if cfg.get("system2", {}).get("enabled"):
        s2 = System2Controller(System2Config(endpoint=cfg["system2"].get("endpoint")))

    brain = HybridAgent(s1=s1, s2=s2, config=HybridConfig(game_time_limit_s=600))
    games = args.games or PUBLIC_GAMES_FALLBACK[:5]  # default: quick smoke set

    report = []
    max_actions = int(cfg.get("eval_max_actions", 1500))
    for game_id in games:
        try:
            env = make_env(game_id, None)
            res = run_game(env, game_id, brain, max_actions=max_actions)
            report.append(res)
            print(res)
        except Exception as exc:
            print(f"[warn] {game_id}: {exc}")

    Path(args.out).write_text(json.dumps(report, indent=2))
    wins = sum(1 for r in report if r["win"])
    print(f"\nsummary: {wins}/{len(report)} wins -> {args.out}")


if __name__ == "__main__":
    main()

	"""Local evaluation: run the hybrid agent on public games and report
	per-game scores in the spirit of the official RHAE metric.

	per-level = min(human/agent, 1.0) ^ 2 (human baseline unknown locally;
	we report completion + actions)
	game = level-weighted average (1-indexed weights)
	total = average over games

	Without official human baselines this is a RELATIVE tool: compare agent
	variants against each other, never treat absolute numbers as leaderboard
	predictions.
	"""

	from __future__ import annotations

	import argparse
	import json
	import sys
	import time
	from pathlib import Path

	sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

	import yaml # noqa: E402

	from data.collect_selfplay import PUBLIC_GAMES_FALLBACK, grids_from, make_env, meta_from # noqa: E402
	from data.mock_env import step_env # noqa: E402
	from qwenjev.hybrid import HybridAgent, HybridConfig # noqa: E402
	from qwenjev.system1 import QwenJevModel, System1Config # noqa: E402
	from qwenjev.system2 import System2Config, System2Controller # noqa: E402
	from qwenjev.vision import TinyVLMConfig # noqa: E402


	def run_game(env, game_id: str, brain: HybridAgent, max_actions: int) -> dict:
	"""Play one episode against any environment (mock or official).

	Lifecycle: WIN ends the episode; GAME_OVER triggers the hybrid brain's
	episode-boundary hook and a fresh reset, then play continues within the
	same action budget — identical budget for every agent variant, so
	best_level/actions stay comparable across runs."""
	env.reset()
	actions = 0
	resets = 0
	best_level = 0
	final_state = "NOT_FINISHED"
	started = time.time()
	for _ in range(max_actions):
	meta = meta_from(env)
	final_state = meta["state"]
	if final_state == "WIN":
	break
	if final_state == "GAME_OVER":
	brain.on_game_over(game_id, [f"game over at action {actions}"])
	env.reset()
	resets += 1
	continue
	grids = grids_from(env)
	best_level = max(best_level, meta["levels_completed"])
	cand = brain.decide(
	game_id, grids,
	levels_completed=meta["levels_completed"],
	win_levels=meta["win_levels"],
	available_actions=meta["available_actions"],
	game_state=meta["state"],
	)
	if cand.is_complex:
	step_env(env, cand.name, int(cand.x or 0), int(cand.y or 0))
	else:
	step_env(env, cand.name)
	actions += 1
	best_level = max(best_level, meta_from(env)["levels_completed"])
	return {
	"game_id": game_id,
	"win": final_state == "WIN",
	"final_state": final_state,
	"best_level": best_level,
	"resets": resets,
	"actions": actions,
	"seconds": round(time.time() - started, 1),
	}


	def build_s1(cfg: dict):
	"""System-1 from the yaml config; checkpoint loaded when present."""
	s1_cfg = cfg.get("system1", {})
	from pathlib import Path

	# v0.4.9: pool_vision auto-detection. The saved qwenjev_s1.pt config
	# carries the flag from v0.4.9 on; ABSENCE = v3-era checkpoint =
	# last-position pooling (heads built with matching H input).
	pool_vision = s1_cfg.get("pool_vision")
	if pool_vision is None:
	pool_vision = False
	try:
	import torch as _torch

	ck = (Path(cfg.get("out_dir", "checkpoints/qwenjev-s1"))
	/ "qwenjev_s1.pt")
	if ck.exists():
	payload = _torch.load(ck, map_location="cpu",
	weights_only=False)
	pool_vision = bool(
	(payload.get("config") or {}).get("pool_vision", False))
	except Exception:
	pool_vision = False
	model_cfg = System1Config(
	backend=s1_cfg.get("backend", "tiny"),
	hf_model=s1_cfg.get("hf_model", "Qwen/Qwen3-VL-2B-Instruct"),
	max_leaf_tokens=int(s1_cfg.get("max_leaf_tokens", 512)),
	score_levels=int(s1_cfg.get("score_levels", 10)),
	leaf_batch_size=int(s1_cfg.get(
	"eval_leaf_batch", s1_cfg.get("leaf_batch_size", 32))),
	pool_vision=bool(pool_vision),
	tiny=TinyVLMConfig(**(s1_cfg.get("tiny") or {})),
	)
	s1 = QwenJevModel(model_cfg).build()
	ckpt = Path(cfg.get("out_dir", "checkpoints/qwenjev-s1")) / "qwenjev_s1.pt"
	if ckpt.exists():
	s1.load_heads(ckpt.parent)
	# stage B/C continuity: merge the SFT adapters into the fp16 base
	# so evaluation runs on the SAME hidden states the heads saw
	if model_cfg.backend == "hf":
	from qwenjev.hf_compat import resume_adapters

	hf_lora = ckpt.parent / "hf_lora"
	if hf_lora.exists():
	s1._backbone, _merged = resume_adapters(
	s1._backbone, str(hf_lora), merge=True)
	print(f"stage B LoRA merged -> {hf_lora}")
	print(f"System-1 checkpoint loaded -> {ckpt}")
	return s1


	def main() -> None:
	ap = argparse.ArgumentParser()
	ap.add_argument("--config", default="configs/train.yaml")
	ap.add_argument("--games", nargs="*", default=None)
	ap.add_argument("--out", default="eval_report.json")
	args = ap.parse_args()
	cfg = yaml.safe_load(open(args.config))

	s1 = None
	try:
	s1 = build_s1(cfg)
	except Exception as exc:
	print(f"[warn] System-1 unavailable ({exc}); heuristic policy will run")

	s2 = None
	if cfg.get("system2", {}).get("enabled"):
	s2 = System2Controller(System2Config(endpoint=cfg["system2"].get("endpoint")))

	brain = HybridAgent(s1=s1, s2=s2, config=HybridConfig(game_time_limit_s=600))
	games = args.games or PUBLIC_GAMES_FALLBACK[:5] # default: quick smoke set

	report = []
	max_actions = int(cfg.get("eval_max_actions", 1500))
	for game_id in games:
	try:
	env = make_env(game_id, None)
	res = run_game(env, game_id, brain, max_actions=max_actions)
	report.append(res)
	print(res)
	except Exception as exc:
	print(f"[warn] {game_id}: {exc}")

	Path(args.out).write_text(json.dumps(report, indent=2))
	wins = sum(1 for r in report if r["win"])
	print(f"\nsummary: {wins}/{len(report)} wins -> {args.out}")


	if __name__ == "__main__":
	main()