517 lines
20 KiB
Python
517 lines
20 KiB
Python
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")) |