scripts/check_predictor_training_state.py

#!/usr/bin/env python3
"""Validate a Predictor training_latest.pt completion state."""

from __future__ import annotations

import argparse
from pathlib import Path

import torch


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("path", type=Path)
    parser.add_argument("--expected_step", type=int, required=True)
    parser.add_argument("--supervision_mode", choices=("onpolicy", "offline_ffff"))
    args = parser.parse_args()

    if not args.path.is_file():
        raise FileNotFoundError(args.path)
    state = torch.load(args.path, map_location="cpu", weights_only=False)
    actual = int(state.get("global_step", -1))
    if actual != args.expected_step:
        raise ValueError(f"{args.path}: global_step={actual}, expected {args.expected_step}")
    if args.supervision_mode is not None:
        saved = str(state.get("supervision_mode"))
        if saved != args.supervision_mode:
            raise ValueError(
                f"{args.path}: supervision_mode={saved!r}, "
                f"expected {args.supervision_mode!r}"
            )
    print(
        f"valid training state: path={args.path} step={actual} "
        f"mode={state.get('supervision_mode', 'stage1')}"
    )


if __name__ == "__main__":
    main()