DualFlow/models/uf-rl/uf_train/rl_model/DQN/run_dqn_train.py

"""
DQN 超滤强化学习训练与测试主脚本（工程化优化版）
"""

import os
import sys
import random
from pathlib import Path
import numpy as np
import torch

# ============================================================
# 1. 导入模块
# ============================================================
CURRENT_DIR = Path(__file__).resolve().parent

PROJECT_ROOT = CURRENT_DIR.parents[2]     # uf_train  # uf-rl


# ---------- 数据 ----------
from uf_train.data_to_rl.data_splitter import ResetStatePoolLoader

# ---------- 阻力模型 ----------
from uf_train.env.uf_resistance_models_load import load_resistance_models
from uf_train.env.uf_physics import UFPhysicsModel
from uf_train.env.env_params import  UFPhysicsParams, UFStateBounds, UFRewardParams, UFActionSpec
from uf_train.env.uf_env import UFSuperCycleEnv

from uf_train.env.env_visual import UFEpisodeRecorder, UFTrainingCallback

from uf_train.rl_model.DQN.dqn_params import DQNParams
from uf_train.rl_model.DQN.dqn_trainer import DQNTrainer


# ---------- SB3 ----------
from stable_baselines3.common.monitor import Monitor
from stable_baselines3.common.vec_env import DummyVecEnv





# ============================================================
#  随机种子
# ============================================================
def set_global_seed(seed: int):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

    print(f"[Seed] Global random seed = {seed}")


# ============================================================
# 4. Reset State Pool 加载与划分
# ============================================================
def load_reset_state_pools():
    loader = ResetStatePoolLoader(
        csv_path=RESET_STATE_CSV,
        train_ratio=0.8,
        shuffle=True,
        random_state=RANDOM_SEED,
    )

    train_pool, val_pool = loader.split()

    print("[Data] Reset state pool loaded")
    print(f"       Train pool size: {len(train_pool)}")
    print(f"       Val   pool size: {len(val_pool)}")

    return train_pool, val_pool


# ============================================================
# 5. 环境构造函数
# ============================================================
def make_env(
    physics: UFPhysicsModel,
    reward_params: UFRewardParams,
    action_spec: UFActionSpec,
    statebounds: UFStateBounds,
    reset_state_pool,
    seed: int,
):
    def _init():
        env = UFSuperCycleEnv(
            physics=physics,
            reward_params=reward_params,
            action_spec=action_spec,
            statebounds=statebounds,
            real_state_pool=reset_state_pool,
            RANDOM_SEED=seed,
        )
        env.action_space.seed(seed)
        env.observation_space.seed(seed)
        return Monitor(env)
    return _init



# ============================================================
# 6. 主流程
# ============================================================
def main():
    # ---------- Seed ----------
    set_global_seed(RANDOM_SEED)

    # ---------- Reset states ----------
    train_pool, val_pool = load_reset_state_pools()

    # ---------- Resistance models ----------
    phys_params = UFPhysicsParams()
    res_fp, res_bw = load_resistance_models(phys_params)

    # ---------- Physics ----------
    physics_model = UFPhysicsModel(
        phys_params=phys_params,
        resistance_model_fp=res_fp,
        resistance_model_bw=res_bw,
    )

    # ---------- RL specs ----------
    reward_params = UFRewardParams()
    action_spec = UFActionSpec()
    state_bounds = UFStateBounds()

    # ---------- Environments ----------
    train_env = DummyVecEnv([
        make_env(
            physics_model,
            reward_params,
            action_spec,
            state_bounds,
            train_pool,
            RANDOM_SEED,
        )
    ])

    val_env = DummyVecEnv([
        make_env(
            physics_model,
            reward_params,
            action_spec,
            state_bounds,
            val_pool,
            RANDOM_SEED,
        )
    ])

    # ---------- Callback ----------
    recorder = UFEpisodeRecorder()
    callback = UFTrainingCallback(recorder, verbose=1)

    # ---------- Trainer ----------
    dqn_params = DQNParams(remark="uf_dqn_real_reset")
    trainer = DQNTrainer(
        env=train_env,
        params=dqn_params,
        callback=callback,
        PROJECT_ROOT=PROJECT_ROOT
    )


    # ---------- Training ----------
    print("\n Start training")
    trainer.train(total_timesteps=TOTAL_TIMESTEPS)
    trainer.save()

    # ========================================================
    # 验证
    # ========================================================
    print("\n[Eval] Start validation rollout")

    TMP0_min = 0.01
    TMP0_max = 0.08

    rewards = []

    # ---------- 用于可视化的容器（只记录第一个 episode） ----------
    vis_tmp_series = []
    vis_action_series = []

    for ep_idx in range(len(val_pool)):
        obs = val_env.reset()
        episode_reward = 0.0

        for step in range(10):
            # ====================================================
            # 可视化：只记录第一个 validation episode
            # ====================================================
            if ep_idx == 0:
                TMP0_norm = obs[0]
                TMP0 = (
                        TMP0_norm * (TMP0_max - TMP0_min)
                        + TMP0_min
                )
                vis_tmp_series.append(TMP0)

            # ---------------- 策略决策 ----------------
            action, _ = trainer.model.predict(
                obs, deterministic=True
            )

            if ep_idx == 0:
                vis_action_series.append(action[0])

            # ---------------- 环境推进 ----------------
            obs, reward, done, _ = val_env.step(action)
            episode_reward += reward[0]

            if done:
                break

        rewards.append(episode_reward)

    # ========================================================
    # 验证结果保存
    # ========================================================
    rewards = np.asarray(rewards)

    save_path = Path(trainer.log_dir) / "val_rewards.npy"
    np.save(save_path, rewards)

    print(f"[Eval] Saved to {save_path}")
    print(f"[Eval] Mean reward = {rewards.mean():.3f}")

    # ========================================================
    # 可视化（第一个 validation episode）
    # ========================================================
    import matplotlib.pyplot as plt

    vis_tmp_series = np.asarray(vis_tmp_series)
    vis_action_series = np.asarray(vis_action_series)
    steps = np.arange(len(vis_tmp_series))

    # ---------- TMP 曲线 ----------
    plt.figure()
    plt.plot(steps, vis_tmp_series, marker="o")
    plt.axhline(
        TMP0_max,
        linestyle="--",
        label="TMP Upper Limit"
    )
    plt.xlabel("Step")
    plt.ylabel("TMP (MPa)")
    plt.title("Validation Episode TMP Evolution")
    plt.legend()
    plt.grid(True)
    plt.show()

    # ---------- Action 曲线 ----------
    plt.figure()
    plt.plot(steps, vis_action_series, marker="o")
    plt.xlabel("Step")
    plt.ylabel("Action")
    plt.title("Validation Episode Action Output")
    plt.grid(True)
    plt.show()


# ============================================================
# 入口
# ============================================================
if __name__ == "__main__":
    # ============================================================
    # 2. 全局配置
    # ============================================================
    RANDOM_SEED = 2025
    TOTAL_TIMESTEPS = 1500000

    RESET_STATE_CSV = (
            PROJECT_ROOT
            / "datasets/rl_ready/output/reset_state_pool.csv"
    )

    main()