model_loading.py

"""
Stable-Baselines3 模型加载辅助工具。
"""

from pathlib import Path

from stable_baselines3 import DDPG, PPO, SAC, TD3
from stable_baselines3.common.save_util import load_from_zip_file


ALGO_LOADERS = {
    "sac": SAC,
    "ppo": PPO,
    "td3": TD3,
    "ddpg": DDPG,
}


def infer_model_algo(model_path):
    """
    根据 Stable-Baselines3 保存的元数据推断模型算法。

    当前仓库主要区分 PPO 与 SAC，TD3 / DDPG 仅做最佳努力识别。
    """
    data, _, _ = load_from_zip_file(model_path, device="cpu", load_data=True)
    policy_class = data.get("policy_class")

    if policy_class is None:
        raise ValueError(f"无法从模型中读取 policy_class: {model_path}")

    module_name = getattr(policy_class, "__module__", "").lower()
    class_name = getattr(policy_class, "__name__", str(policy_class))

    if "stable_baselines3.sac." in module_name or class_name == "SACPolicy":
        return "sac"

    if data.get("rollout_buffer_class") is not None:
        return "ppo"

    if "stable_baselines3.td3." in module_name or class_name == "TD3Policy":
        if data.get("policy_delay") == 1:
            return "ddpg"
        return "td3"

    return "unknown"


def load_saved_observation_space(model_path):
    """读取模型保存时记录的观察空间。"""
    data, _, _ = load_from_zip_file(model_path, device="cpu", load_data=True)
    return data.get("observation_space")


def find_model_candidates(expected_algo, search_root="models", limit=3):
    """在默认模型目录下查找同算法候选模型，便于给出恢复建议。"""
    candidates = []
    root = Path(search_root)

    if not root.exists():
        return candidates

    for path in sorted(root.rglob("*.zip")):
        try:
            algo = infer_model_algo(path)
        except (OSError, ValueError, KeyError, TypeError):
            continue

        if algo == expected_algo:
            candidates.append(str(path))

        if len(candidates) >= limit:
            break

    return candidates


def load_model_for_algo(model_path, env, algo):
    """校验模型算法后再加载，避免不匹配时出现难以理解的底层异常。"""
    expected_algo = algo.lower()

    if expected_algo not in ALGO_LOADERS:
        raise ValueError(f"不支持的算法: {algo}")

    actual_algo = infer_model_algo(model_path)
    if actual_algo != "unknown" and actual_algo != expected_algo:
        message = (
            f"模型算法不匹配：当前脚本期望加载 {expected_algo.upper()}，"
            f"但 `{model_path}` 看起来是 {actual_algo.upper()} 模型。"
        )

        suggestions = find_model_candidates(expected_algo)
        if suggestions:
            suggestion_lines = "\n".join(f"  - {candidate}" for candidate in suggestions)
            message += f"\n可尝试以下 {expected_algo.upper()} checkpoint：\n{suggestion_lines}"

        raise ValueError(message)

    return ALGO_LOADERS[expected_algo].load(model_path, env=env)