基于策略梯度(Policy Gradient)来序贯决策(sequential decision making)任务

798 阅读5分钟

sequential decision making

在人工智能领域,有一类重要的任务和人生选择很相似,即序贯决策(sequential decision making)任务。决策和预测任务不同,决策往往会带来“后果”,因此决策者需要为未来负责,在未来的时间点做出进一步的决策。

数据集简介

这里的数据集采用CartPole,推杆小车是一个经典的强化学习环境,通常用于测试和开发基本的强化学习算法。它模拟了一个由一个可移动车(cart)和一个悬挂杆(pole)组成的系统。

智能体的任务是控制车的水平运动,使得杆子保持竖直直到杆子倾斜超过一定角度或者车移动超出一定范围为止。具体特点包括:

状态空间由四个连续值组成,分别是车的位置、车的速度、杆子的角度以及杆子顶端的速度。

动作空间智能体可以选择向左或向右推车,这是一个离散的动作空间。

奖励(Reward):在每个时间步,智能体获得一个奖励 +1,直到杆子倾斜超过规定角度或者车移动超出规定范围,此时奖励为 0,并且回合结束。

这个环境简单直观,但对于测试强化学习算法的基本能力非常有效,尤其是在学习如何保持杆子竖直这种简单任务上。实现了一个简单的策略梯度算法,用于训练智能体在 CartPole 环境中学习最佳策略。这种方法通过最大化未来预期奖励的策略梯度来更新策略网络,以使智能体学会执行最优动作序列来最大化长期累积奖励。

熟悉这个数据集还是很有帮助的,很多强化学习的demo都是用这个数据集进来进行展示。

构建梯度策略

这里使用了最简单的网络核心构建方式, 这行代码调用了之前定义的 mlp 函数,用于构建策略网络的核心部分。mlp 函数接受一个列表 sizes,指定了神经网络的结构,包括输入维度(obs_dim,即观测空间的维度)、隐藏层大小(hidden_sizes)、输出维度(n_acts,即动作空间的大小)。这个策略网络是一个多层感知机(MLP),用于输出动作的 logits。

# 构建策略网络的核心
logits_net = mlp(sizes=[obs_dim]+hidden_sizes+[n_acts])

# 定义获取策略分布的函数
def get_policy(obs):
    logits = logits_net(obs)
    return Categorical(logits=logits)

# 定义动作选择函数(从策略中采样得到整数动作)
def get_action(obs):
    return get_policy(obs).sample().item()

# 定义损失函数,其梯度用于策略梯度更新
def compute_loss(obs, act, weights):
    logp = get_policy(obs).log_prob(act)
    return -(logp * weights).mean()

# 创建优化器
optimizer = Adam(logits_net.parameters(), lr=lr)

完整代码

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

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)

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(env_name='CartPole-v0', hidden_sizes=[32], lr=1e-2, 
          epochs=50, batch_size=5000, render=False):

    # 创建环境,检查空间,获取观测和动作的维度
    env = gym.make(env_name)
    assert isinstance(env.observation_space, Box), \
        "此示例仅适用于具有连续状态空间的环境。"
    assert isinstance(env.action_space, Discrete), \
        "此示例仅适用于具有离散动作空间的环境。"

    obs_dim = env.observation_space.shape[0]
    n_acts = env.action_space.n

    # 构建策略网络的核心
    logits_net = mlp(sizes=[obs_dim]+hidden_sizes+[n_acts])

    # 定义获取策略分布的函数
    def get_policy(obs):
        logits = logits_net(obs)
        return Categorical(logits=logits)

    # 定义动作选择函数(从策略中采样得到整数动作)
    def get_action(obs):
        return get_policy(obs).sample().item()

    # 定义损失函数,其梯度用于策略梯度更新
    def compute_loss(obs, act, weights):
        logp = get_policy(obs).log_prob(act)
        return -(logp * weights).mean()

    # 创建优化器
    optimizer = Adam(logits_net.parameters(), lr=lr)

    # 用于训练策略的函数
    def train_one_epoch():
        # 用于记录日志的空列表
        batch_obs = []          # 观测值
        batch_acts = []         # 动作
        batch_weights = []      # 策略梯度中的奖励到达加权
        batch_rets = []         # 每个回合的累积奖励
        batch_lens = []         # 每个回合的步数

        # 重置每个回合的变量
        obs = env.reset()       # 初始观测来自初始分布
        done = False            # 环境信号,表明回合结束
        ep_rews = []            # 回合中累积奖励的列表

        # 渲染每个回合的第一帧(如果需要)
        finished_rendering_this_epoch = False

        # 使用当前策略在环境中收集经验
        while True:

            # 渲染环境
            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)

            if done:
                # 如果回合结束,记录回合信息
                ep_ret, ep_len = sum(ep_rews), len(ep_rews)
                batch_rets.append(ep_ret)
                batch_lens.append(ep_len)

                # 每个动作logprob(a_t|s_t)的权重是从t开始的未来奖励到达
                batch_weights += list(reward_to_go(ep_rews))

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

                # 不再在本轮中渲染
                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()
        return batch_loss, batch_rets, batch_lens

    # 训练循环
    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)