policy gradient代码 pytorch框架

原创

已于 2022-04-23 16:48:31 修改 · 3.9k 阅读

6 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch

于 2022-04-18 17:56:05 首次发布

相信你能这么百度，已经都知道什么是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