DualFlow/models/uf-rl/超滤训练源码/DQN_train.py

import os
import time
import random
import numpy as np
import torch

import gymnasium as gym
from gymnasium import spaces
from stable_baselines3 import DQN
from stable_baselines3.common.monitor import Monitor
from stable_baselines3.common.vec_env import DummyVecEnv
from stable_baselines3.common.callbacks import BaseCallback

from DQN_env import UFParams, UFSuperCycleEnv


# ==== 定义强化学习超参数 ====
class DQNParams:
    """
    DQN 超参数定义类
    用于统一管理模型训练参数
    """
    # 学习率，控制神经网络更新步长
    learning_rate: float = 1e-4

    # 经验回放缓冲区大小（步数）
    buffer_size: int = 10000

    # 学习开始前需要收集的步数
    learning_starts: int = 200

    # 每次从经验池中采样的样本数量
    batch_size: int = 32

    # 折扣因子，越接近1越重视长期奖励
    gamma: float = 0.95

    # 每隔多少步训练一次
    train_freq: int = 4

    # 目标网络更新间隔
    target_update_interval: int = 2000

    # 初始探索率 ε
    exploration_initial_eps: float = 1.0

    # 从初始ε衰减到最终ε所占的训练比例
    exploration_fraction: float = 0.3

    # 最终探索率 ε
    exploration_final_eps: float = 0.02

    # 日志备注（用于区分不同实验）
    remark: str = "default"

class UFEpisodeRecorder:
    """记录episode中的决策和结果"""

    def __init__(self):
        self.episode_data = []
        self.current_episode = []

    def record_step(self, obs, action, reward, done, info):
        """记录单步信息"""
        step_data = {
            "obs": obs.copy(),
            "action": action.copy(),
            "reward": reward,
            "done": done,
            "info": info.copy() if info else {}
        }
        self.current_episode.append(step_data)

        if done:
            self.episode_data.append(self.current_episode)
            self.current_episode = []

    def get_episode_stats(self, episode_idx=-1):
        """获取episode统计信息"""
        if not self.episode_data:
            return {}

        episode = self.episode_data[episode_idx]
        total_reward = sum(step["reward"] for step in episode)
        avg_recovery = np.mean([step["info"].get("recovery", 0) for step in episode if "recovery" in step["info"]])
        feasible_steps = sum(1 for step in episode if step["info"].get("feasible", False))

        return {
            "total_reward": total_reward,
            "avg_recovery": avg_recovery,
            "feasible_steps": feasible_steps,
            "total_steps": len(episode)
        }


# ==== 定义强化学习训练回调器 ====
class UFTrainingCallback(BaseCallback):
    """
    强化学习训练回调，用于记录每一步的数据到 recorder。
    1. 不依赖环境内部 last_* 属性
    2. 使用环境接口提供的 obs、actions、rewards、dones、infos
    3. 自动处理 episode 结束时的统计
    """

    def __init__(self, recorder, verbose=0):
        super(UFTrainingCallback, self).__init__(verbose)
        self.recorder = recorder

    def _on_step(self) -> bool:
        try:
            new_obs = self.locals.get("new_obs")
            actions = self.locals.get("actions")
            rewards = self.locals.get("rewards")
            dones = self.locals.get("dones")
            infos = self.locals.get("infos")

            if len(new_obs) > 0:
                step_obs = new_obs[0]
                step_action = actions[0] if actions is not None else None
                step_reward = rewards[0] if rewards is not None else 0.0
                step_done = dones[0] if dones is not None else False
                step_info = infos[0] if infos is not None else {}

                # 打印当前 step 的信息
                if self.verbose:
                    print(f"[Step {self.num_timesteps}] 动作={step_action}, 奖励={step_reward:.3f}, Done={step_done}")

                # 记录数据
                self.recorder.record_step(
                    obs=step_obs,
                    action=step_action,
                    reward=step_reward,
                    done=step_done,
                    info=step_info,
                )

        except Exception as e:
            if self.verbose:
                print(f"[Callback Error] {e}")

        return True




class DQNTrainer:
    def __init__(self, env, params, callback=None):
        self.env = env
        self.params = params
        self.callback = callback
        self.log_dir = self._create_log_dir()
        self.model = self._create_model()

    def _create_log_dir(self):
        # 创建训练日志
        timestamp = time.strftime("%Y%m%d-%H%M%S")
        log_name = (
            f"DQN_lr{self.params.learning_rate}_buf{self.params.buffer_size}_bs{self.params.batch_size}"
            f"_gamma{self.params.gamma}_exp{self.params.exploration_fraction}"
            f"_{self.params.remark}_{timestamp}"
        )
        log_dir = os.path.join("./uf_dqn_tensorboard", log_name)
        os.makedirs(log_dir, exist_ok=True)
        return log_dir

    def _create_model(self):
        return DQN(
            policy="MlpPolicy",
            env=self.env,
            learning_rate=self.params.learning_rate,
            buffer_size=self.params.buffer_size,
            learning_starts=self.params.learning_starts,
            batch_size=self.params.batch_size,
            gamma=self.params.gamma,
            train_freq=self.params.train_freq,
            target_update_interval=1,
            tau=0.005,
            exploration_initial_eps=self.params.exploration_initial_eps,
            exploration_fraction=self.params.exploration_fraction,
            exploration_final_eps=self.params.exploration_final_eps,
            verbose=1,
            tensorboard_log=self.log_dir
        )

    def train(self, total_timesteps: int):
        if self.callback:
            self.model.learn(total_timesteps=total_timesteps, callback=self.callback)
        else:
            self.model.learn(total_timesteps=total_timesteps)
        print(f"模型训练完成，日志保存在：{self.log_dir}")

    def save(self, path=None):
        if path is None:
            path = os.path.join(self.log_dir, "dqn_model.zip")
        self.model.save(path)
        print(f"模型已保存到：{path}")

    def load(self, path):
        self.model = DQN.load(path, env=self.env)
        print(f"模型已从 {path} 加载")


def set_global_seed(seed: int):
    """固定全局随机种子，保证训练可复现"""
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)  # 如果使用GPU
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False


def train_uf_rl_agent(params: UFParams, total_timesteps: int = 10000, seed: int = 2025):
    set_global_seed(seed)
    recorder = UFEpisodeRecorder()
    callback = UFTrainingCallback(recorder, verbose=1)

    def make_env():
        env = UFSuperCycleEnv(params)
        env = Monitor(env)
        return env

    env = DummyVecEnv([make_env])

    dqn_params = DQNParams()
    trainer = DQNTrainer(env, dqn_params, callback=callback)
    trainer.train(total_timesteps)
    trainer.save()

    stats = callback.recorder.get_episode_stats()
    print(f"训练完成 - 总奖励: {stats.get('total_reward', 0):.2f}, 平均回收率: {stats.get('avg_recovery', 0):.3f}")

    return trainer.model


# 训练
if __name__ == "__main__":
    # 初始化参数
    params = UFParams()

    # 训练RL代理
    print("开始训练RL代理...")
    train_uf_rl_agent(params, total_timesteps=50000)