近端策略优化PPO-算法简析-paddle2.0

import argparseimport picklefrom collections import namedtuplefrom itertools import countimport timeimport os, timeimport numpy as npimport matplotlib.pyplot as pltimport gymimport paddleimport paddle.nn as nnimport paddle.nn.functional as Fimport paddle.optimizer as optimfrom paddle.distribution import Normal, Categoricalfrom paddle.io import RandomSampler, BatchSampler, Datasetfrom visualdl import LogWriter

# Parametersenv_name = 'MountainCar-v0'gamma = 0.99render = Falseseed = 1log_interval = 10env = gym.make(env_name).unwrapped
num_state = env.observation_space.shape[0]
num_action = env.action_space.n

env.seed(seed)
Transition = namedtuple('Transition', ['state', 'action', 'action_prob', 'reward', 'next_state'])

class Actor(nn.Layer):
    def __init__(self):
        super(Actor, self).__init__()
        self.fc1 = nn.Linear(num_state, 128)
        self.action_head = nn.Linear(128, num_action)    def forward(self, x):
        x = F.relu(self.fc1(x))
        action_prob = F.softmax(self.action_head(x), axis=1)        return action_probclass Critic(nn.Layer):
    def __init__(self):
        super(Critic, self).__init__()
        self.fc1 = nn.Linear(num_state, 128)
        self.state_value = nn.Linear(128, 1)    def forward(self, x):
        x = F.relu(self.fc1(x))
        value = self.state_value(x)        return value

# init with datasetclass RandomDataset(Dataset):
    def __init__(self, num_samples):
        self.num_samples = num_samples    def __getitem__(self, idx):
        pass

    def __len__(self):
        return self.num_samples

class PPO():
    clip_param = 0.2
    max_grad_norm = 0.5
    ppo_update_time = 10
    buffer_capacity = 8000
    batch_size = 64
    ## 初始化参数
    def __init__(self):
        super(PPO, self).__init__()
        self.actor_net = Actor()
        self.critic_net = Critic()
        self.buffer = []
        self.counter = 0
        self.training_step = 0
        self.writer = LogWriter('./exp')

        clip = nn.ClipGradByNorm(self.max_grad_norm)
        self.actor_optimizer = optim.Adam(parameters = self.actor_net.parameters(),learning_rate= 1e-3, grad_clip=clip)
        self.critic_net_optimizer = optim.Adam(parameters = self.critic_net.parameters(), learning_rate=3e-3,grad_clip=clip)        if not os.path.exists('./param'):
            os.makedirs('./param/net_param')
            os.makedirs('./param/img')    # 选择动作
    def select_action(self, state):
        state = paddle.to_tensor(state,dtype="float32").unsqueeze(0)        with paddle.no_grad():
            action_prob = self.actor_net(state)
        dist = Categorical(action_prob)
        action = dist.sample([1]).squeeze(0)
        action = action.cpu().numpy()[0]        return action, action_prob[:, int(action)].numpy()[0]    # 评估值
    def get_value(self, state):
        state = paddle.to_tensor(state)        with paddle.no_grad():
            value = self.critic_net(state)        return value.numpy()    def save_param(self):
        paddle.save(self.actor_net.state_dict(), './param/net_param/actor_net' + str(time.time())[:10] +'.param')
        paddle.save(self.critic_net.state_dict(), './param/net_param/critic_net' + str(time.time())[:10] +'.param')    def store_transition(self, transition):
        self.buffer.append(transition)
        self.counter += 1

    def update(self, i_ep):
        state = paddle.to_tensor([t.state for t in self.buffer], dtype="float32")
        action = paddle.to_tensor([t.action for t in self.buffer], dtype="int64").reshape([-1, 1])
        reward = [t.reward for t in self.buffer]        # update: don't need next_state
 
        old_action_prob = paddle.to_tensor([t.action_prob for t in self.buffer], dtype="float32").reshape([-1, 1])

        R = 0
        Gt = []        for r in reward[::-1]:
            R = r + gamma * R
            Gt.insert(0, R)
        Gt = paddle.to_tensor(Gt, dtype="float32")        # print("The agent is updateing....")
        for i in range(self.ppo_update_time):            for index in BatchSampler(sampler=RandomSampler(RandomDataset(len(self.buffer))), batch_size=self.batch_size, drop_last=False):                if self.training_step % 1000 == 0:                    print('I_ep {} ，train {} times'.format(i_ep, self.training_step))
                    self.save_param()

                index = paddle.to_tensor(index)
                Gt_index = paddle.index_select(x=Gt, index=index).reshape([-1, 1])                
                # V = self.critic_net(state[index])
                V = self.critic_net(paddle.index_select(state,index))
                
                delta = Gt_index - V
                advantage = delta.detach()                # epoch iteration, PPO core!!!                
                action_prob = self.actor_net(paddle.index_select(state,index))  # new policy
                action_prob = paddle.concat([action_prob[i][int(paddle.index_select(action,index)[i])] for i in range(len(action_prob))]).reshape([-1,1])

                ratio = (action_prob / paddle.index_select(old_action_prob,index))
                surr1 = ratio * advantage
                surr2 = paddle.clip(ratio, 1 - self.clip_param, 1 + self.clip_param) * advantage                # update actor network
                surr = paddle.concat([surr1,surr2],1) 
                action_loss = -paddle.min(surr,1).mean()  # MAX->MIN desent
                self.writer.add_scalar('loss/action_loss', action_loss, self.training_step)
                self.actor_optimizer.clear_grad()
                action_loss.backward()
                self.actor_optimizer.step()                # update critic network
                value_loss = F.mse_loss(Gt_index, V)
                self.writer.add_scalar('loss/value_loss', value_loss, self.training_step)
                self.critic_net_optimizer.clear_grad()
                value_loss.backward()
                self.critic_net_optimizer.step()
                self.training_step += 1

        del self.buffer[:]  # clear experiencedef main():
    agent = PPO()    for i_epoch in range(1000):
        state = env.reset()        if render: env.render()        for t in count():
            action, action_prob = agent.select_action(state)
            next_state, reward, done, _ = env.step(action)
            trans = Transition(state, action, action_prob, reward, next_state)            if render: env.render()
            agent.store_transition(trans)
            state = next_state            if done :                if len(agent.buffer) >= agent.batch_size: agent.update(i_epoch)
                agent.writer.add_scalar('Steptime/steptime', t, i_epoch)                # print("Number of steps to achieve the goal:{} , Steptime:{}".format(t,i_epoch))
                breakif __name__ == '__main__':
    main()    print("end")