import os import torch from pathlib import Path import numpy as np import gymnasium as gym from gymnasium import spaces from typing import Dict, Tuple, Optional import torch import torch.nn as nn from dataclasses import dataclass, asdict from UF_resistance_models import ResistanceIncreaseModel, ResistanceDecreaseModel # 导入模型类 import copy # ======================= # 膜运行参数类：定义膜的基础运行参数 # ======================= @dataclass class UFParams: # —— 膜动态运行参数 —— q_UF: float = 360.0 # 过滤进水流量（m^3/h） TMP0: float = 0.03 # 初始跨膜压差 temp: float = 25.0 # 水温，摄氏度 # —— 膜阻力模型参数 —— nuK: float =4.92e+01 # 过滤阶段膜阻力增长模型参数 slope: float = 3.44e-01 # 全周期不可逆污染阻力增长斜率 power: float = 1.032 # 全周期不可逆污染阻力增长幂次 tau_bw_s: float = 30.0 # 物洗时长影响时间尺度 gamma_t: float = 1.0 # 物洗时长作用指数 ceb_removal: float = 150 # CEB去除膜阻力 # —— 膜运行约束参数 —— global_TMP_limit: float = 0.08 # TMP硬上限（MPa） TMP0_max: float = 0.035 # 初始TMP上限（MPa） TMP0_min: float = 0.01 # 初始TMP下限（MPa） q_UF_max: float = 400.0 # 进水流量上限（m^3/h） q_UF_min: float = 250.0 # 进水流量上限（m^3/h） temp_max: float = 40.0 # 温度上限（摄氏度） temp_min: float = 10.0 # 温度下限（摄氏度） nuK_max: float = 6e+01 # 物理周期总阻力增速上限（m^-1/s） nuK_min: float = 3e+01 # 物理周期总阻力增速下限（m^-1/s） slope_max: float = 10 # 化学周期长期阻力增速斜率上限 slope_min: float = 0.1 # 化学周期长期阻力增速斜率下限 power_max: float = 1.3 # 化学周期长期阻力增速幂次上限 power_min: float = 0.8 # 化学周期长期阻力增速幂次下限 ceb_removal_max: float = 150 # CEB去除阻力(已缩放)上限（m^-1） ceb_removal_min: float = 100 # CEB去除阻力(已缩放)下限（m^-1） # —— 反洗参数（固定） —— q_bw_m3ph: float = 1000.0 # 物理反洗流量（m^3/h） # —— CEB参数 —— T_ceb_interval_h: float = 60.0 # 固定每 k 小时做一次CEB v_ceb_m3: float = 30.0 # CEB用水体积（m^3） t_ceb_s: float = 40 * 60.0 # CEB时长（s） # —— 搜索范围（秒） —— L_min_s: float = 3800.0 # 过滤时长下限（s） L_max_s: float = 6000.0 # 过滤时长上限（s） t_bw_min_s: float = 40.0 # 物洗时长下限（s） t_bw_max_s: float = 60.0 # 物洗时长上限（s） # —— 网格 —— L_step_s: float = 60.0 # 过滤时长步长（s） t_bw_step_s: float = 5.0 # 物洗时长步长（s） # —— 奖励函数参数 —— k_rec = 5.0 # 回收率敏感度 k_res = 10.0 # 残余污染敏感度 rec_low, rec_high = 0.92, 0.99 rr0 = 0.08 # ======================= # 辅助函数：转换膜阻力与跨膜压差 # ======================= def xishan_viscosity(temp): # temp: 水温，单位摄氏度 """ 锡山水厂 PLC水温校正因子经验公式（25摄氏度标准） 返回温度修正后的水粘度（纯水修正），TODO：水厂水质与纯水相差较大，对粘度有一定影响 """ x = (temp + 273.15) / 300 factor = 890 / (280.68 * x ** -1.9 + 511.45 * x ** -7.7 + 61.131 * x ** -19.6 + 0.45903 * x ** -40) mu = 0.00089 / factor return mu def _calculate_resistance(tmp, q_UF, temp): """ 计算超滤膜阻力 R = TMP / (J * μ) 返回缩小1e10的膜阻力（超滤原膜阻力量级为1e12，过大的绝对值容易导致平稳拟合） """ A = 128 * 40 # m²，有效膜面积 mu = xishan_viscosity(temp) # 温度修正后的水粘度 TMP_Pa = tmp * 1e6 # 跨膜压差 MPa -> Pa J = q_UF / A / 3600 # 通量 m³/h -> m³/(m²·s) if J <= 0 or mu <= 0: return np.nan R = TMP_Pa / (J * mu) / 1e10 # 缩放膜阻力 return float(R) def _calculate_tmp(R, q_UF, temp): """ 还原超滤跨膜压差 TMP """ A = 128 * 40 # m²，有效膜面积 mu = xishan_viscosity(temp) # 温度修正后的水粘度 J = q_UF / A / 3600 # 通量 m³/h -> m³/(m²·s) TMP_Pa = R * J * mu * 1e10 tmp = TMP_Pa / 1e6 return float(tmp) # ======================= # 环境体模型加载函数 # ======================= def load_resistance_models(): """加载阻力变化模型，仅在首次调用时执行""" global resistance_model_fp, resistance_model_bw # 如果全局模型已存在，则直接返回 if "resistance_model_fp" in globals() and resistance_model_fp is not None: return resistance_model_fp, resistance_model_bw print("🔄 Loading resistance models...") # 初始化模型 resistance_model_fp = ResistanceIncreaseModel() resistance_model_bw = ResistanceDecreaseModel() # 取得当前脚本所在目录（即 rl_dqn_env.py 或 check_initial_state.py 同目录） base_dir = Path(__file__).resolve().parent # 构造模型路径 fp_path = base_dir / "resistance_model_fp.pth" bw_path = base_dir / "resistance_model_bw.pth" # 检查文件存在性 assert fp_path.exists(), f"缺少 {fp_path.name}" assert bw_path.exists(), f"缺少 {bw_path.name}" # 加载权重 resistance_model_fp.load_state_dict(torch.load(fp_path, map_location="cpu")) resistance_model_bw.load_state_dict(torch.load(bw_path, map_location="cpu")) # 设置推理模式 resistance_model_fp.eval() resistance_model_bw.eval() print("✅ Resistance models loaded successfully from current directory.") return resistance_model_fp, resistance_model_bw # ======================= # 环境体模型模拟函数 # ======================= def _delta_resistance(p, L_h: float) -> float: """ 过滤时段膜阻力上升量：调用 resistance_model_fp.pth 模型 """ return resistance_model_fp(p, L_h) def phi_bw_of(p, R0: float, R_end: float, L_h_start: float, L_h_next_start: float, t_bw_s: float) -> float: """ 物理冲洗去除膜阻力值：调用 resistance_model_bw 模型 """ return resistance_model_bw(p, R0, R_end, L_h_start, L_h_next_start, t_bw_s) def _v_bw_m3(p, t_bw_s: float) -> float: """ 物理反洗水耗 """ return float(p.q_bw_m3ph * (float(t_bw_s) / 3600.0)) def simulate_one_supercycle(p: UFParams, L_s: float, t_bw_s: float): """ 模拟一个超级周期（多次物理反洗 + 一次化学反洗） 返回: (info, next_params) """ L_h = float(L_s) / 3600.0 # 小周期过滤时间(h) tmp = p.TMP0 R0 = _calculate_resistance(p.TMP0, p.q_UF, p.temp) max_tmp_during_filtration = tmp min_tmp_during_filtration = tmp max_residual_increase = 0.0 t_small_cycle_h = (L_s + t_bw_s) / 3600.0 k_bw_per_ceb = int(np.floor(p.T_ceb_interval_h / t_small_cycle_h)) if k_bw_per_ceb < 1: k_bw_per_ceb = 1 energy_lookup = { 3600: 0.1034, 3660: 0.1031, 3720: 0.1029, 3780: 0.1026, 3840: 0.1023, 3900: 0.1021, 3960: 0.1019, 4020: 0.1017, 4080: 0.1015, 4140: 0.1012, 4200: 0.1011 } # --- 循环模拟物理反洗 --- for idx in range(k_bw_per_ceb): tmp_run_start = tmp q_UF = p.q_UF temp = p.temp R_run_start = _calculate_resistance(tmp_run_start, q_UF, temp) d_R = _delta_resistance(p, L_s) R_peak = R_run_start + d_R tmp_peak = _calculate_tmp(R_peak, q_UF, temp) max_tmp_during_filtration = max(max_tmp_during_filtration, tmp_peak) min_tmp_during_filtration = min(min_tmp_during_filtration, tmp_run_start) # 物洗膜阻力减小 L_h_start = (L_s + t_bw_s) / 3600.0 * idx L_h_next_start = (L_s + t_bw_s) / 3600.0 * (idx + 1) reversible_R = phi_bw_of(p, R_run_start, R_peak, L_h_start, L_h_next_start, t_bw_s) R_after_bw = R_peak - reversible_R tmp_after_bw = _calculate_tmp(R_after_bw, q_UF, temp) residual_inc = tmp_after_bw - tmp_run_start max_residual_increase = max(max_residual_increase, residual_inc) tmp = tmp_after_bw # --- CEB反洗 --- R_after_ceb = R_peak - p.ceb_removal tmp_after_ceb = _calculate_tmp(R_after_ceb, q_UF, temp) # ============================================================ # 生成本周期指标 # ============================================================ # --- 体积与能耗 --- V_feed_super = k_bw_per_ceb * p.q_UF * L_h V_loss_super = k_bw_per_ceb * _v_bw_m3(p, t_bw_s) + p.v_ceb_m3 V_net = max(0.0, V_feed_super - V_loss_super) recovery = max(0.0, V_net / max(V_feed_super, 1e-12)) T_super_h = k_bw_per_ceb * (L_s + t_bw_s) / 3600.0 + p.t_ceb_s / 3600.0 daily_prod_time_h = k_bw_per_ceb * L_h / T_super_h * 24.0 closest_L = min(energy_lookup.keys(), key=lambda x: abs(x - L_s)) ton_water_energy = energy_lookup[closest_L] #TODO:需确认新过滤时间范围下的吨水电耗 # --- 信息输出 --- info = { "q_UF": p.q_UF, "temp": p.temp, "recovery": recovery, "V_feed_super_m3": V_feed_super, "V_loss_super_m3": V_loss_super, "V_net_super_m3": V_net, "supercycle_time_h": T_super_h, "max_TMP_during_filtration": max_tmp_during_filtration, "min_TMP_during_filtration": min_tmp_during_filtration, "global_TMP_limit":p.global_TMP_limit, "max_residual_increase_per_run": max_residual_increase, "R0": R0, "R_after_ceb": R_after_ceb, "TMP0":p.TMP0, "TMP_after_ceb": tmp_after_ceb, "daily_prod_time_h": daily_prod_time_h, "ton_water_energy_kWh_per_m3": ton_water_energy, "k_bw_per_ceb": k_bw_per_ceb } # ============================================================ # 状态更新：生成 next_params（新状态） # ============================================================ next_params = copy.deepcopy(p) # 更新跨膜压差（TMP） next_params.TMP0 = tmp_after_ceb # 可选参数（当前保持不变，未来可扩展更新逻辑） next_params.slope = p.slope next_params.power = p.power next_params.ceb_removal = p.ceb_removal next_params.nuK = p.nuK next_params.q_UF = p.q_UF next_params.temp = p.temp return info, next_params def calculate_reward(p: UFParams, info: dict) -> float: """ TMP不参与奖励计算，仅考虑回收率与残余污染比例之间的权衡。 满足： - 当 recovery=0.97, residual_ratio=0.1 → reward = 0 - 当 recovery=0.90, residual_ratio=0.0 → reward = 0 - 在两者之间平衡（如 recovery≈0.94, residual_ratio≈0.05）→ reward > 0 """ recovery = info["recovery"] residual_ratio = (info["R_after_ceb"] - info["R0"]) / info["R0"] # 回收率奖励（在 [rec_low, rec_high] 内平滑上升） rec_norm = (recovery - p.rec_low) / (p.rec_high - p.rec_low) rec_reward = np.clip(np.tanh(p.k_rec * (rec_norm - 0.5)), -1, 1) # 残余比惩罚（超过rr0时快速变为负值） res_penalty = -np.tanh(p.k_res * (residual_ratio / p.rr0 - 1)) # 组合逻辑：权衡二者 total_reward = rec_reward + res_penalty # 再平移使指定点为零： # recovery=0.97, residual=0.1 → 0 # recovery=0.90, residual=0.0 → 0 # 经验上，这两点几乎对称，因此无需额外线性偏移 # 若希望严格归零，可用线性校正： total_reward -= 0.0 return total_reward def is_dead_cycle(info: dict) -> bool: """ 判断当前循环是否为成功循环（True）或失败循环（False） 失败条件： 1. 最大TMP超过设定上限； 2. 回收率低于75%； 3. 化学反冲洗后膜阻力上升超过10%。 参数： info: dict simulate_one_supercycle() 返回的指标字典，需包含： - max_TMP_during_filtration - recovery - R_after_ceb - R_run_start - TMP_limit（如果有定义） 返回： bool: True 表示成功循环，False 表示失败循环。 """ TMP_limit = info.get("global_TMP_limit", 0.08) # 默认硬约束上限 max_tmp = info.get("max_TMP_during_filtration", 0) recovery = info.get("recovery", 1.0) R_after_ceb = info.get("R_after_ceb", 0) R0 = info.get("R0", 1e-6) # 判断条件 if max_tmp > TMP_limit: return False if recovery < 0.75: return False if (R_after_ceb - R0) / R0 > 0.1: return False return True class UFSuperCycleEnv(gym.Env): """超滤系统环境（超级周期级别决策）""" metadata = {"render_modes": ["human"]} def __init__(self, base_params, resistance_models=None, max_episode_steps: int = 15): super(UFSuperCycleEnv, self).__init__() self.base_params = base_params self.current_params = copy.deepcopy(base_params) self.max_episode_steps = max_episode_steps self.current_step = 0 if resistance_models is None: self.resistance_model_fp, self.resistance_model_bw = load_resistance_models() else: self.resistance_model_fp, self.resistance_model_bw = resistance_models # 计算离散动作空间 self.L_values = np.arange( self.base_params.L_min_s, self.base_params.L_max_s, self.base_params.L_step_s ) self.t_bw_values = np.arange( self.base_params.t_bw_min_s, self.base_params.t_bw_max_s + self.base_params.t_bw_step_s, self.base_params.t_bw_step_s ) self.num_L = len(self.L_values) self.num_bw = len(self.t_bw_values) # 单一离散动作空间 self.action_space = spaces.Discrete(self.num_L * self.num_bw) # 状态空间，归一化在 _get_obs 中处理 self.observation_space = spaces.Box( low=np.zeros(8, dtype=np.float32), high=np.ones(8, dtype=np.float32), dtype=np.float32 ) # 初始化环境 self.reset(seed=None) def generate_initial_state(self): """ 随机生成一个初始状态，不进行死状态判断 """ self.current_params.TMP0 = np.random.uniform( self.current_params.TMP0_min, self.current_params.TMP0_max ) self.current_params.q_UF = np.random.uniform( self.current_params.q_UF_min, self.current_params.q_UF_max ) self.current_params.temp = np.random.uniform( self.current_params.temp_min, self.current_params.temp_max ) self.current_params.R0 = _calculate_resistance( self.current_params.TMP0, self.current_params.q_UF, self.current_params.temp ) self.current_params.nuK = np.random.uniform( self.current_params.nuK_min, self.current_params.nuK_max ) self.current_params.slope = np.random.uniform( self.current_params.slope_min, self.current_params.slope_max ) self.current_params.power = np.random.uniform( self.current_params.power_min, self.current_params.power_max ) self.current_params.ceb_removal = np.random.uniform( self.current_params.ceb_removal_min, self.current_params.ceb_removal_max ) return self._get_state_copy() def reset(self, seed=None, options=None, max_attempts: int = 200): super().reset(seed=seed) attempts = 0 while attempts < max_attempts: attempts += 1 self.generate_initial_state() # 生成随机初始状态 if self.check_dead_initial_state(max_steps=getattr(self, "max_episode_steps", 15), L_s=3800, t_bw_s=60): # True 表示可行，退出循环 break else: # 超过最大尝试次数仍未生成可行状态 raise RuntimeError(f"在 {max_attempts} 次尝试后仍无法生成可行初始状态。") # 初始化步数、动作、最大 TMP self.current_step = 0 self.last_action = (self.base_params.L_min_s, self.base_params.t_bw_min_s) self.max_TMP_during_filtration = self.current_params.TMP0 return self._get_obs(), {} def check_dead_initial_state(self, max_steps: int = None, L_s: int = 4900, t_bw_s: int = 50) -> bool: """ 判断当前环境生成的初始状态是否为可行（non-dead）。 使用最保守策略连续模拟 max_steps 次： 若任意一次 is_dead_cycle(info) 返回 False，则视为必死状态。 参数： max_steps: 模拟步数，默认使用 self.max_episode_steps L_s: 过滤时长(s)，默认 3800 t_bw_s: 物理反洗时长(s)，默认 60 返回： bool: True 表示可行状态（non-dead），False 表示必死状态 """ if max_steps is None: max_steps = getattr(self, "max_episode_steps", 15) # 生成初始状态 self.generate_initial_state() if not hasattr(self, "current_params"): raise AttributeError("generate_initial_state() 未设置 current_params。") import copy curr_p = copy.deepcopy(self.current_params) # 逐步模拟 for step in range(max_steps): try: info, next_params = simulate_one_supercycle(curr_p, L_s, t_bw_s) except Exception: # 异常即视为不可行 return False if not is_dead_cycle(info): # 任意一次失败即为必死状态 return False curr_p = next_params return True def _get_state_copy(self): return copy.deepcopy(self.current_params) def _get_obs(self): """ 构建当前环境归一化状态向量 """ # === 1. 从 current_params 读取动态参数 === TMP0 = self.current_params.TMP0 q_UF = self.current_params.q_UF temp = self.current_params.temp # === 2. 计算本周期初始膜阻力 === R0 = _calculate_resistance(TMP0, q_UF, temp) # === 3. 从 current_params 读取膜阻力增长模型参数 === nuk = self.current_params.nuK slope = self.current_params.slope power = self.current_params.power ceb_removal = self.current_params.ceb_removal # === 4. 从 current_params 动态读取上下限 === TMP0_min, TMP0_max = self.current_params.TMP0_min, self.current_params.TMP0_max q_UF_min, q_UF_max = self.current_params.q_UF_min, self.current_params.q_UF_max temp_min, temp_max = self.current_params.temp_min, self.current_params.temp_max nuK_min, nuK_max = self.current_params.nuK_min, self.current_params.nuK_max slope_min, slope_max = self.current_params.slope_min, self.current_params.slope_max power_min, power_max = self.current_params.power_min, self.current_params.power_max ceb_min, ceb_max = self.current_params.ceb_removal_min, self.current_params.ceb_removal_max # === 5. 归一化计算（clip防止越界） === TMP0_norm = np.clip((TMP0 - TMP0_min) / (TMP0_max - TMP0_min), 0, 1) q_UF_norm = np.clip((q_UF - q_UF_min) / (q_UF_max - q_UF_min), 0, 1) temp_norm = np.clip((temp - temp_min) / (temp_max - temp_min), 0, 1) # R0 不在 current_params 中定义上下限，设定经验范围 R0_norm = np.clip((R0 - 100.0) / (800.0 - 100.0), 0, 1) short_term_norm = np.clip((nuk - nuK_min) / (nuK_max - nuK_min), 0, 1) long_term_slope_norm = np.clip((slope - slope_min) / (slope_max - slope_min), 0, 1) long_term_power_norm = np.clip((power - power_min) / (power_max - power_min), 0, 1) ceb_removal_norm = np.clip((ceb_removal - ceb_min) / (ceb_max - ceb_min), 0, 1) # === 6. 构建观测向量 === obs = np.array([ TMP0_norm, q_UF_norm, temp_norm, R0_norm, short_term_norm, long_term_slope_norm, long_term_power_norm, ceb_removal_norm ], dtype=np.float32) return obs def _get_action_values(self, action): """ 将动作还原为实际时长 """ L_idx = action // self.num_bw t_bw_idx = action % self.num_bw return self.L_values[L_idx], self.t_bw_values[t_bw_idx] def step(self, action): self.current_step += 1 L_s, t_bw_s = self._get_action_values(action) L_s = np.clip(L_s, self.base_params.L_min_s, self.base_params.L_max_s) t_bw_s = np.clip(t_bw_s, self.base_params.t_bw_min_s, self.base_params.t_bw_max_s) # 模拟超级周期 info, next_params = simulate_one_supercycle(self.current_params, L_s, t_bw_s) # 根据 info 判断是否成功 feasible = is_dead_cycle(info) # True 表示成功循环，False 表示失败 if feasible: reward = calculate_reward(self.current_params, info) self.current_params = next_params terminated = False else: reward = -10 terminated = True truncated = self.current_step >= self.max_episode_steps self.last_action = (L_s, t_bw_s) next_obs = self._get_obs() info["feasible"] = feasible info["step"] = self.current_step return next_obs, reward, terminated, truncated, info