policy gradient代码 pytorch框架

相信你能这么百度，已经都知道什么是PG（policy gradient，PG）算法了，也已经装有pytorch了。所以，这里我不做任何对PG算法的讲解，因为我嘴笨可能讲不清楚。

这篇博客存在意义：网上都是牵扯到Categorical函数（分类分布）的代码，不用Categorical能不能写？答案：能！！灵感来自莫烦大佬的policy gradient tensorfollow版。

以下是PG玩cartpole游戏，训练只需10min，测试时游戏已经可以永远不倒。

GPU版（用的是GPU训练，所以你要先安装Cuda）

至此，代码如下，拿走不谢，复制即用，不行砍我！

#开发者：Bright Fang
#开发时间：2022/4/12 11:35
import torch
import torch.nn as nn
import torch.nn.functional as F
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
import numpy as np
import gym
LearningRate=0.01
Gamma=0.9#Gamma越大越容易收敛
Switch=0#训练、测试切换标志
env=gym.make('CartPole-v1')
env=env.unwrapped
state_number=env.observation_space.shape[0]
action_number=env.action_space.n
'''policygrandient第一步先建网络'''
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.in_to_y1=nn.Linear(state_number,20)
        self.in_to_y1.weight.data.normal_(0,0.1)
        self.y1_to_y2=nn.Linear(20,10)
        self.y1_to_y2.weight.data.normal_(0,0.1)
        self.out=nn.Linear(10,action_number)
        self.out.weight.data.normal_(0,0.1)
    def forward(self,inputstate):
        inputstate=self.in_to_y1(inputstate)
        inputstate=F.relu(inputstate)
        inputstate=self.y1_to_y2(inputstate)
        inputstate=torch.sigmoid(inputstate)
        act=self.out(inputstate)
        # return act
        return F.softmax(act,dim=-1)
class PG():
    def __init__(self):
        self.policy = Net().cuda()
        self.rewards,self.obs,self.acts = [],[],[]
        self.renderflag=False
        self.optimizer=torch.optim.Adam(self.policy.parameters(),lr=LearningRate)
    '''第二步 定义选择动作函数'''
    def choose(self,inputstate):
        inputstate=torch.FloatTensor(inputstate).cuda()
        probs=self.policy(inputstate).cpu().detach().numpy()
        action=np.random.choice(np.arange(action_number),p=probs)
        return action
    '''第三步 存储每一个回合的数据'''
    def store_transtion(self,s,a,r):
        self.obs.append(s)
        self.acts.append(a)
        self.rewards.append(r)
    '''第四步 学习'''
    def learn(self):
        # pass
        discounted_ep_r =np.zeros_like(self.rewards)
        running_add=0
        for t in reversed(range(0,len(self.rewards))):
            running_add=running_add*Gamma+self.rewards[t]
            discounted_ep_r[t]=running_add#例如，discounted_ep_r是1*87的列表，列表的第一个值为58，最后一个值为1
        #先减去平均数再除以标准差，就可对奖励归一化，奖励列表的中间段为0，最左为+2.1，最右为-1.9.
        discounted_ep_r-=np.mean(discounted_ep_r)
        discounted_ep_r/=np.std(discounted_ep_r)
        discounted_ep_rs_norm=discounted_ep_r
        self.optimizer.zero_grad()
        #把一个回合的状态、动作、奖励三个列表转为tensor
        self.obs=np.array(self.obs)
        state_tensor = torch.FloatTensor(self.obs).cuda()
        reward_tensor = torch.FloatTensor(discounted_ep_rs_norm).cuda()
        action_tensor = torch.LongTensor(self.acts).cuda()
        #我们可以用G值直接进行学习，但一般来说，对数据进行归一化处理后，训练效果会更好
        log_prob=torch.log(self.policy(state_tensor))#log_prob是拥有两个动作概率的张量，一个左动作概率，一个右动作概率
        selected_log_probs =reward_tensor * log_prob[np.arange(len(action_tensor)), action_tensor]#np.arange(len(action_tensor))是log_prob的索引，
        # action_tensor由0、1组成，于是log_prob[np.arange(len(action_tensor)), action_tensor]就可以取到我们已经选择了的动作的概率，是拥有一个动作概率的张量
        loss=-selected_log_probs.mean()
        loss.backward()
        self.optimizer.step()
        self.obs,self.acts,self.rewards=[],[],[]
'''训练'''
if Switch==0:
    print("训练PG中...")
    f=PG()
    for i in range(2000):
        r=0
        observation=env.reset()
        while True:
            if f.renderflag: env.render()
            action=f.choose(observation)
            observation_,reward,done,info=env.step(action)
            #修改reward
            x, x_dot, theta, theta_dot = observation_
            r1 = (env.x_threshold - abs(x)) / env.x_threshold - 0.8
            r2 = (env.theta_threshold_radians - abs(theta)) / env.theta_threshold_radians - 0.5
            r3=3*r1+r2
            #你也可以不修改奖励，直接用reward，都能收敛
            f.store_transtion(observation,action,r3)
            r+=r3
            if done:
                f.learn()
                break
            observation=observation_
        print("\rEp: {} rewards: {}".format(i,r), end="")
        if i % 10 == 0 and i > 500:
            save_data = {
    
    'net': f.policy.state_dict(), 'opt': f.optimizer.state_dict(), 'i': i}
            torch.save(save_data, "D:\PyCharm 2019.3\mytorch_spacework\demo\model_PG.pth")
else:
    print("测试PG中...")
    c=PG()
    checkpoint = torch.load("D:\PyCharm 2019.3\mytorch_spacework\demo\model_PG.pth")
    c.policy.load_state_dict(checkpoint['net'])
    for j in range(10):
        state = env.reset()
        total_rewards = 0
        while True:
            env.render()
            state = torch.FloatTensor(state)
            action=c.choose(state)
            new_state, reward, done, info = env.step(action)  # 执行动作
            total_rewards += reward
            if done:
                print("Score", total_rewards)
                break
            state = new_state
    env.close()

不出意外，上述代码应该不能复制即用，我早已料到。因为：

1. 你们的D盘里没有PyCharm 2019.3\mytorch_spacework\demo文件夹，毕竟大家代码的存放地方都不一样。我的代码是放在这里的，
   所以，要么把我的D:\PyCharm 2019.3\mytorch_spacework\demo\改为你自己的当前代码所在位置，要么你就直接把

torch.save(save_data, "D:\PyCharm 2019.3\mytorch_spacework\demo\model_PG.pth")

改为

torch.save(save_data, "E:\model_PG.pth")

直接放在你的E盘里，别告诉我你的电脑没有E盘啊。
2. 并不是所有人都已经安装了Cuda的，所以我给出了CPU版PG算法玩cartpole的代码，如下：

CPU版

#开发者：Bright Fang
#开发时间：2022/4/12 11:35
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import gym
LearningRate=0.01
Gamma=0.9#Gamma越大越容易收敛
Switch=0#训练、测试切换标志
env=gym.make('CartPole-v1')
env=env.unwrapped
state_number=env.observation_space.shape[0]
action_number=env.action_space.n
'''policygrandient第一步先建网络'''
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.in_to_y1=nn.Linear(state_number,20)
        self.in_to_y1.weight.data.normal_(0,0.1)
        self.y1_to_y2=nn.Linear(20,10)
        self.y1_to_y2.weight.data.normal_(0,0.1)
        self.out=nn.Linear(10,action_number)
        self.out.weight.data.normal_(0,0.1)
    def forward(self,inputstate):
        inputstate=self.in_to_y1(inputstate)
        inputstate=F.relu(inputstate)
        inputstate=self.y1_to_y2(inputstate)
        inputstate=torch.sigmoid(inputstate)
        act=self.out(inputstate)
        # return act
        return F.softmax(act,dim=-1)
class PG():
    def __init__(self):
        self.policy = Net()
        self.rewards,self.obs,self.acts = [],[],[]
        self.renderflag=False
        self.optimizer=torch.optim.Adam(self.policy.parameters(),lr=LearningRate)
    '''第二步 定义选择动作函数'''
    def choose(self,inputstate):
        inputstate=torch.FloatTensor(inputstate)
        probs=self.policy(inputstate).detach().numpy()
        action=np.random.choice(np.arange(action_number),p=probs)
        return action
    '''第三步 存储每一个回合的数据'''
    def store_transtion(self,s,a,r):
        self.obs.append(s)
        self.acts.append(a)
        self.rewards.append(r)
    '''第四步 学习'''
    def learn(self):
        # pass
        discounted_ep_r =np.zeros_like(self.rewards)
        running_add=0
        for t in reversed(range(0,len(self.rewards))):
            running_add=running_add*Gamma+self.rewards[t]
            discounted_ep_r[t]=running_add#例如，discounted_ep_r是1*87的列表，列表的第一个值为58，最后一个值为1
        #先减去平均数再除以标准差，就可对奖励归一化，奖励列表的中间段为0，最左为+2.1，最右为-1.9.
        discounted_ep_r-=np.mean(discounted_ep_r)
        discounted_ep_r/=np.std(discounted_ep_r)
        discounted_ep_rs_norm=discounted_ep_r
        self.optimizer.zero_grad()
        #把一个回合的状态、动作、奖励三个列表转为tensor
        self.obs=np.array(self.obs)
        state_tensor = torch.FloatTensor(self.obs)
        reward_tensor = torch.FloatTensor(discounted_ep_rs_norm)
        action_tensor = torch.LongTensor(self.acts)
        #我们可以用G值直接进行学习，但一般来说，对数据进行归一化处理后，训练效果会更好
        log_prob=torch.log(self.policy(state_tensor))#log_prob是拥有两个动作概率的张量，一个左动作概率，一个右动作概率
        selected_log_probs =reward_tensor * log_prob[np.arange(len(action_tensor)), action_tensor]#np.arange(len(action_tensor))是log_prob的索引，
        # action_tensor由0、1组成，于是log_prob[np.arange(len(action_tensor)), action_tensor]就可以取到我们已经选择了的动作的概率，是拥有一个动作概率的张量
        loss=-selected_log_probs.mean()
        loss.backward()
        self.optimizer.step()
        self.obs,self.acts,self.rewards=[],[],[]
'''训练'''
if Switch==0:
    print("训练PG中...")
    f=PG()
    for i in range(2000):
        r=0
        observation=env.reset()
        while True:
            if f.renderflag: env.render()
            action=f.choose(observation)
            observation_,reward,done,info=env.step(action)
            #修改reward
            x, x_dot, theta, theta_dot = observation_
            r1 = (env.x_threshold - abs(x)) / env.x_threshold - 0.8
            r2 = (env.theta_threshold_radians - abs(theta)) / env.theta_threshold_radians - 0.5
            r3=3*r1+r2
            #你也可以不修改奖励，直接用reward，都能收敛
            f.store_transtion(observation,action,r3)
            r+=r3
            if done:
                f.learn()
                break
            observation=observation_
        print("\rEp: {} rewards: {}".format(i,r), end="")
        if i % 10 == 0 and i > 500:
            save_data = {
    
    'net': f.policy.state_dict(), 'opt': f.optimizer.state_dict(), 'i': i}
            torch.save(save_data, "E:\model_PG.pth")
else:
    print("测试PG中...")
    c=PG()
    checkpoint = torch.load("E:\model_PG.pth")
    c.policy.load_state_dict(checkpoint['net'])
    for j in range(10):
        state = env.reset()
        total_rewards = 0
        while True:
            env.render()
            state = torch.FloatTensor(state)
            action=c.choose(state)
            new_state, reward, done, info = env.step(action)  # 执行动作
            total_rewards += reward
            if done:
                print("Score", total_rewards)
                break
            state = new_state
    env.close()

代码用法：
先把Switch标志为赋为0，先训练，训练个5-10min就直接停止训练（不要等了，如果让它自然地训练结束会等到猴年马月的），因为神经网络的参数已经被我们保存在了model_PG.pth里。然后，把Switch标志为赋为1，就可以看到训练的效果了。