DualFlow/models/uf-rl/uf_train/data_to_rl/get_reset_pool.py

import yaml
from pathlib import Path

from loader import load_all_units_cycles, resolve_path
from state_construction import build_chem_cycle_state, STATE_COLUMNS


def load_config(path: Path) -> dict:
    with open(path, "r", encoding="utf-8") as f:
        return yaml.safe_load(f)


def generate_reset_state_pool(config_path: Path):
    cfg = load_config(config_path)
    project_root = Path(cfg["Paths"]["project_root"])

    output_dir = resolve_path(
        project_root,
        cfg["Paths"]["data_to_rl"]["output_dir"]
    )
    output_dir.mkdir(parents=True, exist_ok=True)

    df = load_all_units_cycles(cfg)
    df = df[df["chem_cycle_valid"] == True]

    state_df = build_chem_cycle_state(df)

    # 严格完整性过滤
    state_df = state_df.dropna(subset=STATE_COLUMNS)

    state_df.to_csv(
        output_dir / "reset_state_pool.csv",
        index=False
    )