【DRL】最简单的策略梯度(Policy Gradient)算法

转载请注明出处：小锋学长生活大爆炸[xfxuezhagn.cn]

如果本文帮助到了你，欢迎[点赞、收藏、关注]哦~

代码来自OpenAI SpinningUp

该代码是策略梯度（Policy Gradient）方法最简实现，采样一批episode，计算每步动作的对数概率乘以整条轨迹的总奖励，作为损失函数优化策略网络，实现对策略的改进。适合用来理解PG算法的基本原理和流程。

# coding: utf-8

import torch
import torch.nn as nn
from torch.distributions.categorical import Categorical
from torch.optim import Adam
import numpy as np
import gym
import time
from gym.spaces import Discrete, Box
from tqdm import tqdm

# MLP网络构建函数，构造一个多层感知机（MLP），作为策略网络，输出每个动作的logits。
def mlp(sizes, activation=nn.Tanh, output_activation=nn.Identity):
    layers = []
    for j in range(len(sizes)-1):
        act = activation if j < len(sizes)-2 else output_activation
        layers += [nn.Linear(sizes[j], sizes[j+1]), act()]
    return nn.Sequential(*layers)

# train训练函数
def train(env_name='CartPole-v0', hidden_sizes=[32], lr=1e-2, epochs=50, batch_size=5000, render=False):
    # 环境初始化和检查：只支持连续状态空间（Box）和离散动作空间（Discrete）
    env = gym.make(env_name)
    assert isinstance(env.observation_space, Box), "此示例仅适用于具有连续状态空间的env"
    assert isinstance(env.action_space, Discrete), "此示例仅适用于具有离散动作空间的env"
    # observation_space是智能体能够从环境中观测到的所有可能信息的集合，它定义了智能体可以获取到的关于环境状态的部分或全部信息，这些信息是智能体做出决策的依据
    obs_dim = env.observation_space.shape[0] # 4: (小车位置、小车速度、杆子角度、杆子角速度)
    # action是智能体在某个状态下可以采取的具体行为或选择
    n_acts = env.action_space.n # 2: 左、右

    # 搭建策略网络。logits_net模块的输出可用于构造动作的对数概率和概率
    # 当谈论具有“logits”的分类分布时，是指每个结果的概率由logits的Softmax函数给出
    logits_net = mlp(sizes=[obs_dim]+hidden_sizes+[n_acts])

    # 输入状态obs，输出Categorical分布对象，用于采样动作
    def get_policy(obs):
        logits = logits_net(obs)
        # Category 是一个 PyTorch Distribution 对象，它包含了一些与概率分布相关的数学函数
        # 特别是，它有一个从分布中采样的方法和一个计算给定样本的对数概率的方法
        return Categorical(logits=logits)

    # 从策略分布(logits计算出的概率)中采样一个动作
    # 这个特定的get_action函数假设只提供一个obs，因此只提供一个integer的动作输出。
    # 这就是为什么它使用.item()，它用于获取仅具有一个元素的Tensor的内容。
    def get_action(obs):
        return get_policy(obs).sample().item()

    # 计算策略梯度损失，权重为每个episode的总回报
    # episode是智能体从环境的初始状态开始执行动作，直到到达终止状态的整个过程
    # 当插入正确的数据时，这种损失的梯度等于策略梯度
    # 正确的数据意味着在根据当前策略采取行动时收集的一组(state, action, weight)二元组
    # 其中state-action对的权重是它所属的事件的返回值
    # 与监督学习中典型意义上的损失函数相比，这里的损失函数有两点不同：
    #     1. 数据分布取决于参数：标准的损失函数通常定义在一个固定的数据分布上，它与要优化的参数无关。但这里的情况并非如此，这里的数据必须根据最近的策略进行采样
    #     2. 不衡量表现：标准的损失函数通常评估关心的性能指标，而在这里关心的是预期收益。当用当前参数生成的数据评估时，它的性能梯度为负。在梯度下降的第一步之后，与性能就不再有联系。这意味着，对于给定的一批数据，最小化这个loss函数并不能保证提高预期回报。在策略梯度中，你应该只关心平均回报。损失函数没有任何意义。
    def compute_loss(obs, act, weights):
        logp = get_policy(obs).log_prob(act)
        return -(logp * weights).mean()

    # Adam优化器初始化。策略参数就是logits_net的所有参数
    optimizer = Adam(logits_net.parameters(), lr=lr)

    def reward_to_go(rews):
        n = len(rews)
        rtgs = np.zeros_like(rews)
        for i in reversed(range(n)):
            rtgs[i] = rews[i] + (rtgs[i+1] if i+1 < n else 0)
        return rtgs
    
    # 用于训练策略
    def train_one_epoch():
        # 制作一些空列表进行日志记录
        batch_obs = []          # for observations
        batch_acts = []         # for actions
        batch_weights = []      # 策略梯度中R(tau)的权重
        batch_rets = []         # 衡量episode的回报
        batch_lens = []         # 衡量episode的步数

        # 重置特定于episode的变量
        obs = env.reset()       # 第一个观察来自启动时的分布
        done = False            # 来自env的信号，表明episode已结束
        ep_rews = []            # episode期间累积的奖励列表
        # 只展示第一个episode的渲染，不然训练很慢
        finished_rendering_this_epoch = False

        bar = tqdm(total=batch_size)
        # 与环境交互，采集一个batch(如5000步)的(s, a, r)数据，保存每一步的(观察、行动、奖励)
        while True:
            bar.update(1)
            # 渲染每个epoch的第一个episode
            if (not finished_rendering_this_epoch) and render: env.render()

            # 保存观察
            batch_obs.append(obs.copy())

            # 在环境中执行一个行动
            act = get_action(torch.as_tensor(obs, dtype=torch.float32))
            obs, rew, done, _ = env.step(act)

            # 保存行动和奖励
            batch_acts.append(act)
            ep_rews.append(rew)

            # 每个episode结束时，将该episode的总回报赋给该episode中所有step作为权重
            if done:
                # 如果episode已结束，则记录有关episode的信息
                ep_ret, ep_len = sum(ep_rews), len(ep_rews)
                batch_rets.append(ep_ret)
                batch_lens.append(ep_len)

                # 每个logprob(a|s)的权重为R(tau)
                # 将整个回合的总回报 ep_ret 平均分配给该回合的每一步
                # 认为回合中所有动作对最终结果的贡献相同，无论动作发生在早期还是晚期
                # 实现简单，但方差大（早期动作的结果受更多随机因素影响）
                # 可能导致梯度估计不准确，尤其是在长回合任务中
                batch_weights += [ep_ret] * ep_len
                # 另一种形式：优势函数的简化版
                # 使用未来回报作为每一步的权重
                # 认为每个动作只影响从该时刻开始的未来回报，而非整个回合
                # 降低了方差（去除了与当前动作无关的历史回报）
                # 更符合因果关系：动作只影响未来，不影响过去。
                # batch_weights += list(reward_to_go(ep_rews))



                # 重置特定于episode的变量
                obs, done, ep_rews = env.reset(), False, []

                # 此epoch下不会再渲染
                finished_rendering_this_epoch = True

                # 当采集到足够多数据时就中止
                if len(batch_obs) > batch_size: break
        
        # 用所有采集的数据进行一次策略网络更新（梯度下降）
        optimizer.zero_grad()
        batch_loss = compute_loss(obs=torch.as_tensor(batch_obs, dtype=torch.float32),
                                  act=torch.as_tensor(batch_acts, dtype=torch.int32),
                                  weights=torch.as_tensor(batch_weights, dtype=torch.float32)
                                  )
        batch_loss.backward()
        optimizer.step()
        # 返回本轮损失、回报和episode长度
        return batch_loss, batch_rets, batch_lens

    # 训练主循环。连续训练指定的轮数，每轮输出损失、平均回报、平均episode长度
    for i in range(epochs):
        batch_loss, batch_rets, batch_lens = train_one_epoch()
        print('epoch: %3d \t loss: %.3f \t return: %.3f \t ep_len: %.3f' % (i, batch_loss, np.mean(batch_rets), np.mean(batch_lens)))

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--env_name', '--env', type=str, default='CartPole-v0')
    parser.add_argument('--render', action='store_true')
    parser.add_argument('--lr', type=float, default=1e-2)
    args = parser.parse_args()
    print('\n使用最简单的策略梯度公式.\n')
    train(env_name=args.env_name, render=args.render, lr=args.lr)

运行：

python 1_simple_pg.py  --render

输出：