import os os.environ['MKL_THREADING_LAYER'] = 'GNU' os.environ["OMP_NUM_THREADS"] = "8" os.environ["MKL_NUM_THREADS"] = "8" import torch import numpy as np from torch.nn import Module import gymnasium as gym from gym_env_250326 import CustomEnv from stable_baselines3 import PPO from stable_baselines3.common.vec_env import SubprocVecEnv from stable_baselines3.common.vec_env import DummyVecEnv from sb3_contrib import RecurrentPPO from torch.utils.tensorboard import SummaryWriter import pickle # 自定义模块导入 from torch import nn from stable_baselines3.common.policies import ActorCriticPolicy from stable_baselines3.common.utils import obs_as_tensor from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback from stable_baselines3.common.buffers import RolloutBuffer from gymnasium import spaces import types torch.backends.cudnn.enabled = False current_dir = os.path.dirname(os.path.abspath("__file__")) parent_dir = os.path.abspath(os.path.join(current_dir, os.pardir)) class Sin(Module): def __init__(self): super().__init__() def forward(self, x): return torch.sin(x) class RewardAwareEnvironmentWrapper(gym.Wrapper): """ 环境包装器，跟踪奖励历史并传递给特征提取器 """ def __init__(self, env): super().__init__(env) self.reward_history = [] self.max_reward_history = 60 def reset(self, **kwargs): self.reward_history = [] return self.env.reset(**kwargs) def step(self, action): # 修复：适配新的Gymnasium API (5个返回值) result = self.env.step(action) # 检查返回值的数量以兼容不同版本 if len(result) == 5: # 新版本 Gymnasium: obs, reward, terminated, truncated, info obs, reward, terminated, truncated, info = result done = terminated or truncated # 合并terminated和truncated为done else: # 旧版本 Gym: obs, reward, done, info obs, reward, done, info = result # 记录奖励历史 self.reward_history.append(reward) if len(self.reward_history) > self.max_reward_history: self.reward_history.pop(0) # 将奖励历史添加到info中，供特征提取器使用 info['reward_history'] = self.reward_history.copy() info['current_reward'] = reward # 返回与原环境相同格式的值 if len(result) == 5: return obs, reward, terminated, truncated, info else: return obs, reward, done, info class MultiTimeScaleLSTMExtractor(nn.Module): def __init__(self, observation_space, features_dim=32): super().__init__() self.n_obs = observation_space.shape[0] # 总共14个观测量 self.delayed_indices = list(range(0, 8)) self.current_indices = list(range(2, 8)) self.leading_indices = list(range(8, self.n_obs)) # self.delayed_indices = [] # self.current_indices = list(range(0, 6)) # self.leading_indices = list(range(6, self.n_obs)) if len(self.leading_indices) > 0: self.leading_seq_length = 30 self.leading_lstm = nn.LSTM( input_size=len(self.leading_indices), hidden_size=16, num_layers=1, batch_first=True, dropout=0.0 ) self.leading_mlp = nn.Sequential( nn.Linear(16, 16), Sin() ) # LSTM分支 - 处理时间延迟观测量 if len(self.delayed_indices) > 0: self.delayed_seq_length = 60 self.delayed_lstm = nn.LSTM( input_size=len(self.delayed_indices), hidden_size=8, num_layers=1, batch_first=True, dropout=0.0 ) self.delayed_mlp = nn.Sequential( nn.Linear(8, 8), Sin() ) # MLP分支 - 处理当前观测量 if len(self.current_indices) > 0: current_obs_count = len(self.current_indices) self.current_mlp = nn.Sequential( nn.Linear(current_obs_count, 16), Sin(), ) # 奖励历史LSTM - 新增 self.reward_seq_length = 30 self.reward_lstm = nn.LSTM( input_size=1, # 奖励是标量 hidden_size=8, num_layers=1, batch_first=True ) self.reward_mlp = nn.Sequential( nn.Linear(8, 8), Sin() ) # 简化注意力机制 - 降低复杂度 attention_dim = 16 # 统一注意力维度 # 将不同分支的输出投影到统一维度 self.leading_proj = nn.Linear(16, attention_dim) if len(self.leading_indices) > 0 else None self.delayed_proj = nn.Linear(8, attention_dim) if len(self.delayed_indices) > 0 else None self.current_proj = nn.Linear(16, attention_dim) if len(self.current_indices) > 0 else None self.reward_proj = nn.Linear(8, attention_dim) # 奖励分支投影 # 时间注意力机制 - 学习不同时间尺度的重要性 self.temporal_attention = nn.MultiheadAttention( embed_dim=attention_dim, num_heads=2, batch_first=True ) # 融合层 - 将所有分支的特征融合 num_branches = sum([len(self.leading_indices) > 0, len(self.delayed_indices) > 0, len(self.current_indices) > 0]) + 1 combined_size = num_branches * attention_dim # 每个分支16维 self.fusion = nn.Sequential( nn.Linear(combined_size, features_dim), # 直接输出到目标维度 Sin() # 只用一层 ) self.features_dim = features_dim # 添加记忆缓冲区 self.leading_memory = None self.delayed_memory = None self.reward_memory = None # 当前奖励存储（用于传递给update_reward_memory） self.current_reward = 0.0 def update_reward_memory(self, reward, batch_size, device): """更新奖励记忆""" reward_tensor = torch.full((batch_size, 1), reward, device=device, dtype=torch.float32) if self.reward_memory is None or self.reward_memory.shape[0] != batch_size: self.reward_memory = torch.zeros( (batch_size, self.reward_seq_length, 1), device=device, dtype=torch.float32 ) for i in range(self.reward_seq_length): self.reward_memory[:, i, :] = reward_tensor else: self.reward_memory = torch.roll(self.reward_memory, shifts=-1, dims=1) self.reward_memory[:, -1, :] = reward_tensor def forward(self, observations): # 处理观测值，创建或更新记忆缓冲区 if len(observations.shape) == 2: batch_size, n_obs = observations.shape # 管理超前信号记忆 if self.leading_memory is None or self.leading_memory.shape[0] != batch_size: self.leading_memory = torch.zeros( (batch_size, self.leading_seq_length, len(self.leading_indices)), device=observations.device, dtype=observations.dtype ) for i in range(self.leading_seq_length): self.leading_memory[:, i, :] = observations[:, self.leading_indices] else: self.leading_memory = torch.roll(self.leading_memory, shifts=-1, dims=1) self.leading_memory[:, -1, :] = observations[:, self.leading_indices] # 管理滞后信号记忆 if self.delayed_memory is None or self.delayed_memory.shape[0] != batch_size: self.delayed_memory = torch.zeros( (batch_size, self.delayed_seq_length, len(self.delayed_indices)), device=observations.device, dtype=observations.dtype ) for i in range(self.delayed_seq_length): self.delayed_memory[:, i, :] = observations[:, self.delayed_indices] else: self.delayed_memory = torch.roll(self.delayed_memory, shifts=-1, dims=1) self.delayed_memory[:, -1, :] = observations[:, self.delayed_indices] # 管理奖励记忆 - 使用存储的当前奖励 self.update_reward_memory(self.current_reward, batch_size, observations.device) features = [] # 处理超前观测量 if len(self.leading_indices) > 0: _, (leading_hidden, _) = self.leading_lstm(self.leading_memory) leading_features = self.leading_mlp(leading_hidden[-1]) # 取最后一层的隐状态 leading_features = self.leading_proj(leading_features) # 投影到统一维度 features.append(leading_features) # 处理滞后观测量 if len(self.delayed_indices) > 0: _, (delayed_hidden, _) = self.delayed_lstm(self.delayed_memory) delayed_features = self.delayed_mlp(delayed_hidden[-1]) delayed_features = self.delayed_proj(delayed_features) # 投影到统一维度 features.append(delayed_features) # 处理当前观测量 if len(self.current_indices) > 0: current_obs = observations[:, self.current_indices] current_features = self.current_mlp(current_obs) current_features = self.current_proj(current_features) # 投影到统一维度 features.append(current_features) # 处理奖励历史特征 if self.reward_memory is not None: _, (reward_hidden, _) = self.reward_lstm(self.reward_memory) reward_features = self.reward_mlp(reward_hidden[-1]) reward_features = self.reward_proj(reward_features) features.append(reward_features) # 应用时间注意力机制 if len(features) > 1: # 将特征重新排列为注意力机制的输入格式 stacked_features = torch.stack(features, dim=1) # [batch_size, num_branches, feature_dim] attended_features, _ = self.temporal_attention( stacked_features, stacked_features, stacked_features ) # 修复：使用reshape而不是view，或使用contiguous().view() combined_features = attended_features.reshape(attended_features.shape[0], -1) else: combined_features = torch.cat(features, dim=1) return self.fusion(combined_features) class MlpExtractor(nn.Module): """ 自定义的MLP特征提取器，添加了SB3所需的forward_actor和forward_critic方法 """ def __init__(self, feature_dim, latent_dim_pi=16, latent_dim_vf=16): super().__init__() self.latent_dim_pi = latent_dim_pi self.latent_dim_vf = latent_dim_vf # 创建actor和critic网络 self.policy_net = nn.Sequential( nn.Linear(feature_dim, 32), Sin(), nn.Linear(32, latent_dim_pi), Sin() ) self.value_net = nn.Sequential( nn.Linear(feature_dim, 32), Sin(), nn.Linear(32, latent_dim_vf), Sin() ) def forward(self, features): """同时提取actor和critic特征""" return self.policy_net(features), self.value_net(features) def forward_actor(self, features): """仅提取actor特征""" return self.policy_net(features) def forward_critic(self, features): """仅提取critic特征""" return self.value_net(features) class CustomActorCriticPolicy(ActorCriticPolicy): def __init__(self, observation_space, action_space, lr_schedule, **kwargs): # 移除网络相关的关键字参数 features_extractor_kwargs = kwargs.pop("features_extractor_kwargs", {}) features_extractor_kwargs["features_dim"] = 32 super().__init__( observation_space, action_space, lr_schedule, net_arch=[], # 使用空列表而不是None activation_fn=Sin, features_extractor_class=MultiTimeScaleLSTMExtractor, features_extractor_kwargs=features_extractor_kwargs, **kwargs ) def _build_mlp_extractor(self): # 使用自定义的MlpExtractor替代默认的 features_dim = self.features_extractor.features_dim self.mlp_extractor = MlpExtractor( feature_dim=features_dim, latent_dim_pi=16, latent_dim_vf=16 ) class RewardTrackingPPO(PPO): """ 扩展PPO以传递奖励信息给特征提取器 """ def collect_rollouts( self, env: GymEnv, callback: MaybeCallback, rollout_buffer: RolloutBuffer, n_rollout_steps: int, ) -> bool: """ 重写collect_rollouts方法以传递奖励信息 """ # 在每次收集rollout之前重置奖励 if hasattr(self.policy.features_extractor, 'current_reward'): self.policy.features_extractor.current_reward = 0.0 assert self._last_obs is not None, "No previous observation was provided" # Switch to eval mode (this affects batch norm / dropout) self.policy.set_training_mode(False) n_steps = 0 rollout_buffer.reset() # Sample new weights for the state dependent exploration if self.use_sde: self.policy.reset_noise(env.num_envs) callback.on_rollout_start() while n_steps < n_rollout_steps: if self.use_sde and self.sde_sample_freq > 0 and n_steps % self.sde_sample_freq == 0: # Sample a new noise matrix self.policy.reset_noise(env.num_envs) with torch.no_grad(): # Convert to pytorch tensor or to TensorDict obs_tensor = obs_as_tensor(self._last_obs, self.device) actions, values, log_probs = self.policy(obs_tensor) actions = actions.cpu().numpy() # Rescale and perform action clipped_actions = actions if isinstance(self.action_space, spaces.Box): if self.policy.squash_output: # Unscale the actions to match env bounds # if they were previously squashed (scaled in [-1, 1]) clipped_actions = self.policy.unscale_action(clipped_actions) else: # Otherwise, clip the actions to avoid out of bound error # as we are sampling from an unbounded Gaussian distribution clipped_actions = np.clip(actions, self.action_space.low, self.action_space.high) new_obs, rewards, dones, infos = env.step(clipped_actions) # 更新特征提取器中的奖励信息 if hasattr(self.policy.features_extractor, 'current_reward'): # 如果是向量化环境，取第一个环境的奖励 reward_to_update = rewards[0] if isinstance(rewards, np.ndarray) else rewards self.policy.features_extractor.current_reward = float(reward_to_update) self.num_timesteps += env.num_envs # Give access to local variables callback.on_step() if callback.on_step() is False: return False self._update_info_buffer(infos, dones) n_steps += 1 if isinstance(self.action_space, spaces.Discrete): # Reshape in case of discrete action actions = actions.reshape(-1, 1) # Handle timeout by bootstraping with value function # see GitHub issue #633 for idx, done in enumerate(dones): if ( done and infos[idx].get("terminal_observation") is not None and infos[idx].get("TimeLimit.truncated", False) ): terminal_obs = self.policy.obs_to_tensor(infos[idx]["terminal_observation"])[0] with torch.no_grad(): terminal_value = self.policy.predict_values(terminal_obs)[0] rewards[idx] += self.gamma * terminal_value rollout_buffer.add( self._last_obs, actions, rewards, self._last_episode_starts, values, log_probs, ) self._last_obs = new_obs self._last_episode_starts = dones with torch.no_grad(): # Compute value for the last timestep values = self.policy.predict_values(obs_as_tensor(new_obs, self.device)) rollout_buffer.compute_returns_and_advantage(last_values=values, dones=dones) callback.on_rollout_end() return True def _update_reward_in_extractor(self, reward): """ 更新特征提取器中的当前奖励 """ if hasattr(self.policy.features_extractor, 'current_reward'): # 修复：正确处理NumPy数组和标量值 if isinstance(reward, np.ndarray): # 如果是数组，取第一个元素（或者平均值，根据需要） reward_value = float(reward.item()) if reward.size == 1 else float(reward[0]) else: # 如果是标量，直接转换 reward_value = float(reward) self.policy.features_extractor.current_reward = reward_value if __name__ == '__main__': # 包装环境以跟踪奖励历史 base_env = CustomEnv(device_id=0) vec_env = RewardAwareEnvironmentWrapper(base_env) name = "d1a3o14_cloak_lstm" # model = PPO.load(os.path.join(parent_dir, "models", "250729", "d1a3o12_cloak_lstm.zip"), env=vec_env, device=torch.device("cuda:0")) model = RewardTrackingPPO( policy=CustomActorCriticPolicy, env=vec_env, device=torch.device("cuda:0"), n_steps=1024, batch_size=128, learning_rate=5e-4, gamma=0.99, gae_lambda=0.95, clip_range=0.2, ent_coef=0.01, max_grad_norm=0.5, verbose=0) writer = SummaryWriter(log_dir=os.path.join(parent_dir, "tensorboard", name)) max_reward = 0 history_data = [] for i in range(500): model.learn(total_timesteps=1200) test_env = model.get_env() test_obs = test_env.reset() list_reward = [] episolde_data = {'actions': [], 'observations': [], 'rewards': []} for step in range(200): test_action, _states = model.predict(observation=test_obs) # 修复：处理环境step的返回值 result = test_env.step(test_action) if len(result) == 5: test_obs, test_rewards, terminated, truncated, info = result test_dones = terminated or truncated else: test_obs, test_rewards, test_dones, info = result # 更新特征提取器中的奖励信息 model._update_reward_in_extractor(test_rewards) list_reward.append(test_rewards) episolde_data['actions'].append(test_action[0, :]) episolde_data['observations'].append(np.array(test_obs)) episolde_data['rewards'].append(test_rewards) history_data.append(episolde_data) avg_reward = np.mean(list_reward[-100:]) writer.add_scalar('Reward', np.mean(avg_reward), i) if avg_reward > max_reward: max_reward = avg_reward model.save(os.path.join(parent_dir, "models", "250729", name + ".zip"))