在这里插入代码片# 使用连续PPO算法构建智能体
import torch.optim.lr_scheduler as lr_scheduler

import math
import torch
import torch.nn.functional as F
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm
import utils.rl_utils as rl_utils
import random

def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)  # if you are using multi-GPU.
    # Ensure reproducibility in cudnn
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False



class PolicyNet(torch.nn.Module):
    def __init__(self, state_dim, hidden_dim, action_dim):
        super(PolicyNet, self).__init__()
        self.fc1 = torch.nn.Linear(state_dim, hidden_dim)
        self.fc2 = torch.nn.Linear(hidden_dim, action_dim)
        self._init_weights()

    def _init_weights(self):
        # 使用 Kaiming 均匀初始化方法对 全连接层 的权重进行初始化，并将偏置初始化为零
        nn.init.kaiming_uniform_(self.fc1.weight, nonlinearity='relu')
        nn.init.zeros_(self.fc1.bias)
        nn.init.kaiming_uniform_(self.fc2.weight, nonlinearity='relu')
        nn.init.zeros_(self.fc2.bias)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        return F.softmax(self.fc2(x), dim=1)


class ValueNet(torch.nn.Module):
    def __init__(self, state_dim, hidden_dim):
        super(ValueNet, self).__init__()
        self.fc1 = torch.nn.Linear(state_dim, hidden_dim)
        self.fc2 = torch.nn.Linear(hidden_dim, 1)
        self._init_weights()

    def _init_weights(self):
        # 使用 Kaiming 均匀初始化方法对 全连接层 的权重进行初始化，并将偏置初始化为零
        nn.init.kaiming_uniform_(self.fc1.weight, nonlinearity='relu')
        nn.init.zeros_(self.fc1.bias)
        nn.init.kaiming_uniform_(self.fc2.weight, nonlinearity='relu')
        nn.init.zeros_(self.fc2.bias)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        return self.fc2(x)


class PPO:
    ''' PPO算法,采用截断方式 '''
    def __init__(self, state_dim, hidden_dim, action_dim, actor_lr, critic_lr,
                 lmbda, epochs, eps, gamma, device):
        self.actor = PolicyNet(state_dim, hidden_dim, action_dim).to(device)
        self.critic = ValueNet(state_dim, hidden_dim).to(device)
        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(),
                                                lr=actor_lr)
        self.critic_optimizer = torch.optim.Adam(self.critic.parameters(),
                                                 lr=critic_lr)
        self.gamma = gamma
        self.lmbda = lmbda
        self.epochs = epochs  # 一条序列的数据用于训练轮数
        self.eps = eps  # PPO中截断范围的参数
        self.device = device

    def take_action(self, state):
        state = torch.tensor([state], dtype=torch.float).to(self.device)
        probs = self.actor(state)# 表示给定状态下各个动作的概率分布
        action_dist = torch.distributions.Categorical(probs)# 创建概率分布对象，表示抽中0的动作概率是probs[0],抽中1的概率概率是probs[1]
        action = action_dist.sample()
        return action.item()

    def update(self, transition_dict):
        states = torch.tensor(transition_dict['states'],
                              dtype=torch.float).to(self.device)
        actions = torch.tensor(transition_dict['actions']).view(-1, 1).to(
            self.device)
        rewards = torch.tensor(transition_dict['rewards'],
                               dtype=torch.float).view(-1, 1).to(self.device)
        next_states = torch.tensor(transition_dict['next_states'],
                                   dtype=torch.float).to(self.device)
        dones = torch.tensor(transition_dict['dones'],
                             dtype=torch.float).view(-1, 1).to(self.device)
        td_target = rewards + self.gamma * self.critic(next_states) * (1 -
                                                                       dones)
        td_delta = td_target - self.critic(states)
        advantage = rl_utils.compute_advantage(self.gamma, self.lmbda,
                                               td_delta.cpu()).to(self.device)
        old_log_probs = torch.log(self.actor(states).gather(1,# torch.log(...)：对上一步获取的概率值取对数。这在强化学习中常用，因为对数概率在计算损失和进行策略更新时更加稳定和有效
                                                            actions)).detach()

        for _ in range(self.epochs):
            log_probs = torch.log(self.actor(states).gather(1, actions))
            ratio = torch.exp(log_probs - old_log_probs)
            surr1 = ratio * advantage
            surr2 = torch.clamp(ratio, 1 - self.eps,
                                1 + self.eps) * advantage  # 截断
            actor_loss = torch.mean(-torch.min(surr1, surr2))  # PPO损失函数
            critic_loss = torch.mean(
                F.mse_loss(self.critic(states), td_target.detach()))
            self.actor_optimizer.zero_grad()
            self.critic_optimizer.zero_grad()
            actor_loss.backward()
            critic_loss.backward()
            self.actor_optimizer.step()
            self.critic_optimizer.step()


class PolicyNetContinuous(torch.nn.Module):
    def __init__(self, state_dim, hidden_dim, action_dim):
        super(PolicyNetContinuous, self).__init__()
        self.fc1 = torch.nn.Linear(state_dim, hidden_dim)
        self.fc2 = torch.nn.Linear(hidden_dim,hidden_dim)
        self.fc3 = torch.nn.Linear(hidden_dim,hidden_dim)
        self.fc_mu = torch.nn.Linear(hidden_dim, action_dim)
        self.fc_std = torch.nn.Linear(hidden_dim, action_dim)
        # self._init_weights()

    def _init_weights(self):
        # 使用 Kaiming 均匀初始化方法对 全连接层 的权重进行初始化，并将偏置初始化为零
        nn.init.kaiming_uniform_(self.fc1.weight, nonlinearity='relu')
        nn.init.zeros_(self.fc1.bias)
        nn.init.kaiming_uniform_(self.fc2.weight, nonlinearity='relu')
        nn.init.zeros_(self.fc2.bias)
        nn.init.kaiming_uniform_(self.fc_mu.weight, nonlinearity='tanh')
        nn.init.zeros_(self.fc_mu.bias)
        nn.init.kaiming_uniform_(self.fc_std.weight, nonlinearity='relu')
        nn.init.zeros_(self.fc_std.bias)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = F.relu(self.fc3(x))
        mu = torch.tanh(self.fc_mu(x))# torch.tanh 是双曲正切函数，其输出范围是 [-1, 1]
        std = F.softplus(self.fc_std(x))# F.softplus 是一种平滑的 ReLU 变体，输出为正值，避免标准差为负数
        return mu, std # std 表示生成动作的标准差，用于控制智能体在当前状态下的探索程度

class ValueNet(torch.nn.Module):
    def __init__(self, state_dim, hidden_dim):
        super(ValueNet, self).__init__()
        self.fc1 = torch.nn.Linear(state_dim, hidden_dim)
        self.fc2 = torch.nn.Linear(hidden_dim, hidden_dim)
        self.fc3 = torch.nn.Linear(hidden_dim, 1)
        # self._init_weights()

    def _init_weights(self):
        # 使用 Kaiming 均匀初始化方法对 全连接层 的权重进行初始化，并将偏置初始化为零
        nn.init.kaiming_uniform_(self.fc1.weight, nonlinearity='relu')
        nn.init.zeros_(self.fc1.bias)
        nn.init.kaiming_uniform_(self.fc2.weight, nonlinearity='relu')
        nn.init.zeros_(self.fc2.bias)
        nn.init.kaiming_uniform_(self.fc3.weight, nonlinearity='relu')
        nn.init.zeros_(self.fc3.bias)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        return self.fc3(x)
    

class PPOContinuous:
    ''' 处理连续动作的PPO算法 '''
    def __init__(self, max_episode,state_dim, hidden_dim, action_dim, actor_lr, critic_lr,
                 lmbda, epochs, eps, gamma, device):
        self.actor = PolicyNetContinuous(state_dim, hidden_dim,
                                         action_dim).to(device)
        self.critic = ValueNet(state_dim, hidden_dim).to(device)
        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(),
                                                lr=actor_lr)
        self.critic_optimizer = torch.optim.Adam(self.critic.parameters(),
                                                 lr=critic_lr)
        
        # 学习率调度器
        lambda_func = lambda epoch: 1 - epoch / 1000  # 线性衰减，假设训练 1000 episode 完成衰减
        lambda_func = lambda episode: 0.5 * (1 + math.cos(math.pi * episode / max_episode)) # 余弦退火 
        self.actor_scheduler = lr_scheduler.LambdaLR(self.actor_optimizer, lr_lambda=lambda_func)
        self.critic_scheduler = lr_scheduler.LambdaLR(self.critic_optimizer, lr_lambda=lambda_func)

        self.gamma = gamma
        self.lmbda = lmbda
        self.epochs = epochs
        self.eps = eps
        self.device = device
        self.max_delta = 0.001
    def update_schedulers(self, episode):
        """在每个 episode 后更新调度器"""
        self.actor_scheduler.step(episode)
        self.critic_scheduler.step(episode)

    def take_action(self, state):
        state = torch.tensor([state], dtype=torch.float).to(self.device)
        mu, sigma = self.actor(state)# 输出均值和标准差
        # mu = torch.tanh(mu) * self.max_delta  # 缩放到 [-0.00135, 0.00135] 范围内
        # sigma = torch.clamp(sigma+1e-6, min=1e-6, max=self.max_delta)  # 将标准差裁剪到合适范围
        action_dist = torch.distributions.Normal(mu, sigma)#  # 定义正态分布
        action = action_dist.sample()# 从正态分布中采样
        action = action * 1e-3    
        action = torch.clamp(action, min=-self.max_delta, max=self.max_delta)
        
        # 转为list
        action = action.squeeze().tolist()
        return action# 返回动作

    def update(self, transition_dict):
        states = torch.tensor(transition_dict['states'],
                              dtype=torch.float).to(self.device)
        actions = torch.tensor(transition_dict['actions'],
                               dtype=torch.float).to(self.device)
        rewards = torch.tensor(transition_dict['rewards'],
                               dtype=torch.float).view(-1, 1).to(self.device)
        next_states = torch.tensor(transition_dict['next_states'],
                                   dtype=torch.float).to(self.device)
        dones = torch.tensor(transition_dict['dones'],
                             dtype=torch.float).view(-1, 1).to(self.device)
        # rewards = (rewards + 250.0) / 250.0  # 和TRPO一样,对奖励进行修改,方便训练
        # rewards = (rewards - rewards.mean()) / (rewards.std() + 1e-8)  # 和TRPO一样,对奖励进行修改,方便训练
        td_target = rewards + self.gamma * self.critic(next_states) * (1 -
                                                                       dones)
        td_delta = td_target - self.critic(states)
        advantage = rl_utils.compute_advantage(self.gamma, self.lmbda,
                                               td_delta.cpu()).to(self.device)
        # advantage = (advantage - advantage.mean()) / (advantage.std() + 1e-8)  # 标准化优势函数
        mu, std = self.actor(states)
        action_dists = torch.distributions.Normal(mu.detach(), std.detach())# 创建一个正态分布对象 action_dists
        # 动作是正态分布
        old_log_probs = action_dists.log_prob(actions)# 旧策略下的动作对数概率 old_log_probs

        for i in range(self.epochs):
            mu, std = self.actor(states)
            action_dists = torch.distributions.Normal(mu, std)
            log_probs = action_dists.log_prob(actions)
            ratio = torch.exp(log_probs - old_log_probs)
            surr1 = ratio * advantage
            surr2 = torch.clamp(ratio, 1 - self.eps, 1 + self.eps) * advantage
            actor_loss = torch.mean(-torch.min(surr1, surr2))
            critic_loss = torch.mean(
                F.mse_loss(self.critic(states), td_target.detach()))
            self.actor_optimizer.zero_grad()
            self.critic_optimizer.zero_grad()
            actor_loss.backward()
            critic_loss.backward()
            torch.nn.utils.clip_grad_norm_(self.actor.parameters(), max_norm=0.5)  # 添加梯度裁剪
            torch.nn.utils.clip_grad_norm_(self.critic.parameters(), max_norm=0.5)
            self.actor_optimizer.step()
            self.critic_optimizer.step()
        # 更新学习率
        self.actor_scheduler.step()
        self.critic_scheduler.step()
    
    def save(self,episode,moudle_dir):
        torch.save(self.actor.state_dict(), f'{moudle_dir}/{episode}PPO_actor_dec.pth')
        torch.save(self.critic.state_dict(), f'{moudle_dir}/{episode}PPO_critic_dec.pth')
        print('...save model...')

    def load(self,moudle_dir):
        self.actor.load_state_dict(torch.load(f'{moudle_dir}/PPO_actor_dec.pth'))
        self.critic.load_state_dict(torch.load(f'{moudle_dir}/PPO_critic_dec.pth'))
        print('...load...')