PyTorch策略梯度算法详情

import gym
import torch
import matplotlib.pyplot as plt
env = gym.make('CartPole-v0')

n_state = env.observation_space.shape[0]
print(n_state)

n_action = env.action_space.n
print(n_action)

def run_episode(env, weight):
    state = env.reset()
    grads = []
    total_reward = 0
    is_done = False
    while not is_done:
        state = torch.from_numpy(state).float()
        # 根据当前状态和输入权重计算两个动作的概率 probs
        z = torch.matmul(state, weight)
        probs = torch.nn.Softmax(dim=0)(z)
        # 根据结果概率采样一个动作 action
        action = int(torch.bernoulli(probs[1]).item())
        # 以概率作为输入计算 softmax 函数的导数 d_softmax
        d_softmax = torch.diag(probs) - probs.view(-1, 1) * probs
        # 计算与策略相关的对数导数d_log
        d_log = d_softmax[action] / probs[action]
        # 计算权重的梯度grad
        grad = state.view(-1, 1) * d_log
        grads.append(grad)
        state, reward, is_done, _ = env.step(action)
        total_reward += reward
        if is_done:
            break
    return total_reward, grads

n_episode = 1000
weight = torch.rand(n_state, n_action)
total_rewards = []

learning_rate = 0.001

for e in range(n_episode):
    total_reward, gradients = run_episode(env, weight)
    print('Episode {}: {}'.format(e + 1, total_reward))
    for i, gradient in enumerate(gradients):
        weight += learning_rate * gradient * (total_reward - i)
    total_rewards.append(total_reward)

print('Average total reward over {} episode: {}'.format(n_episode, sum(total_rewards)/n_episode))

plt.plot(total_rewards)
plt.xlabel('Episode')
plt.ylabel('Reward')
plt.show()

n_episode_eval = 1000
total_rewards_eval = []
for e in range(n_episode_eval):
    total_reward, _ = run_episode(env, weight)
    print('Episode {}: {}'.format(e+1, total_reward))
    total_rewards_eval.append(total_reward)

print('Average total reward over {} episode: {}'.format(n_episode_eval, sum(total_rewards_eval)/n_episode_eval))
# Average total reward over 1000 episode: 200