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")

    rewards = []

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

        for _ in range(10):
            action, _ = trainer.model.predict(
                obs, deterministic=True
            )
            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}")


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

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

    main()