使用强化学习训练平衡摆杆#

本笔记本是 AI for Beginners Curriculum 的一部分，灵感来源于官方 PyTorch 教程和这个 Cartpole PyTorch 实现。

在这个示例中，我们将使用强化学习（RL）训练一个模型，使其能够在水平轨道上移动的推车上平衡摆杆。我们将使用 OpenAI Gym 环境来模拟摆杆。

注意: 你可以在本地运行本课程的代码（例如，通过 Visual Studio Code），此时模拟器会在新窗口中打开。如果在线运行代码，你可能需要对代码进行一些调整，具体说明见这里。

我们将从确保 Gym 已安装开始：

In [ ]:

import sys
!{sys.executable} -m pip install gym

现在让我们创建 CartPole 环境，并了解如何对其进行操作。一个环境具有以下属性：

动作空间 是我们在模拟的每一步中可以执行的所有可能动作的集合
观察空间 是我们能够获取的所有可能观察值的集合

In [ ]:

import gym

env = gym.make("CartPole-v1")

print(f"Action space: {env.action_space}")
print(f"Observation space: {env.observation_space}")

让我们看看模拟是如何运行的。以下循环会运行模拟，直到 env.step 不返回终止标志 done。我们将使用 env.action_space.sample() 随机选择动作，这意味着实验可能会非常快地失败（当小车的速度、位置或角度超出某些限制时，CartPole 环境会终止）。

模拟会在新窗口中打开。你可以多次运行代码，观察它的表现。

In [ ]:

env.reset()

done = False
total_reward = 0
while not done:
   env.render()
   obs, rew, done, info = env.step(env.action_space.sample())
   total_reward += rew
   print(f"{obs} -> {rew}")
print(f"Total reward: {total_reward}")

你会注意到观测值包含4个数字，它们是：

小车的位置
小车的速度
杆子的角度
杆子的旋转速度

rew 是我们在每一步中获得的奖励。在 CartPole 环境中，每进行一步模拟都会获得1分奖励，目标是最大化总奖励，也就是让小车平衡杆子而不倒下的时间尽可能长。

在强化学习中，我们的目标是训练一个策略 $\pi$，它会根据每个状态 $s$ 告诉我们应该采取的动作 $a$，本质上就是 $a = \pi(s)$。

如果你想要一个概率性的解决方案，可以将策略视为为每个动作返回一组概率，也就是说 $\pi(a|s)$ 表示在状态 $s$ 下采取动作 $a$ 的概率。

策略梯度方法#

在最简单的强化学习算法中，称为策略梯度，我们将训练一个神经网络来预测下一步的动作。

In [ ]:

import numpy as np
import matplotlib.pyplot as plt
import torch

num_inputs = 4
num_actions = 2

model = torch.nn.Sequential(
    torch.nn.Linear(num_inputs, 128, bias=False, dtype=torch.float32),
    torch.nn.ReLU(),
    torch.nn.Linear(128, num_actions, bias = False, dtype=torch.float32),
    torch.nn.Softmax(dim=1)
)

我们将通过运行许多实验来训练网络，并在每次运行后更新我们的网络。让我们定义一个函数来运行实验并返回结果（所谓的轨迹）——所有状态、动作（及其推荐概率）和奖励：

In [ ]:

def run_episode(max_steps_per_episode = 10000,render=False):    
    states, actions, probs, rewards = [],[],[],[]
    state = env.reset()
    for _ in range(max_steps_per_episode):
        if render:
            env.render()
        action_probs = model(torch.from_numpy(np.expand_dims(state,0)))[0]
        action = np.random.choice(num_actions, p=np.squeeze(action_probs.detach().numpy()))
        nstate, reward, done, info = env.step(action)
        if done:
            break
        states.append(state)
        actions.append(action)
        probs.append(action_probs.detach().numpy())
        rewards.append(reward)
        state = nstate
    return np.vstack(states), np.vstack(actions), np.vstack(probs), np.vstack(rewards)

您可以运行一个未训练网络的情节，并观察到总奖励（即情节长度）非常低：

In [ ]:

s, a, p, r = run_episode()
print(f"Total reward: {np.sum(r)}")

政策梯度算法的一个棘手方面是使用折扣奖励。其思想是我们在游戏的每一步计算总奖励的向量，并在此过程中使用某个系数 $gamma$ 对早期奖励进行折扣。我们还对结果向量进行归一化，因为我们将使用它作为权重来影响我们的训练：

In [ ]:

eps = 0.0001

def discounted_rewards(rewards,gamma=0.99,normalize=True):
    ret = []
    s = 0
    for r in rewards[::-1]:
        s = r + gamma * s
        ret.insert(0, s)
    if normalize:
        ret = (ret-np.mean(ret))/(np.std(ret)+eps)
    return ret

现在开始实际训练！我们将运行300次实验，每次实验中我们将执行以下步骤：

运行实验并收集轨迹数据。
计算所采取的动作与预测概率之间的差异（gradients）。差异越小，说明我们越确信采取了正确的动作。
计算折扣奖励，并将梯度乘以折扣奖励——这样可以确保高奖励的步骤对最终结果的影响比低奖励的步骤更大。
我们的神经网络的目标动作部分来自运行期间的预测概率，部分来自计算出的梯度。我们将使用alpha参数来确定梯度和奖励在多大程度上被考虑——这被称为强化学习算法的学习率。
最后，我们基于状态和目标动作训练网络，并重复这一过程。

In [ ]:

optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

def train_on_batch(x, y):
    x = torch.from_numpy(x)
    y = torch.from_numpy(y)
    optimizer.zero_grad()
    predictions = model(x)
    loss = -torch.mean(torch.log(predictions) * y)
    loss.backward()
    optimizer.step()
    return loss

In [ ]:

alpha = 1e-4

history = []
for epoch in range(300):
    states, actions, probs, rewards = run_episode()
    one_hot_actions = np.eye(2)[actions.T][0]
    gradients = one_hot_actions-probs
    dr = discounted_rewards(rewards)
    gradients *= dr
    target = alpha*np.vstack([gradients])+probs
    train_on_batch(states,target)
    history.append(np.sum(rewards))
    if epoch%100==0:
        print(f"{epoch} -> {np.sum(rewards)}")

plt.plot(history)

现在让我们运行带有渲染的剧集来查看结果：

In [ ]:

_ = run_episode(render=True)

希望你能看到，现在杆子已经可以很好地保持平衡了！

Actor-Critic 模型#

Actor-Critic 模型是策略梯度的进一步发展，在该模型中，我们构建一个神经网络来同时学习策略和估计奖励。这个网络将有两个输出（或者你可以将其视为两个独立的网络）：

Actor 将通过给出状态概率分布来推荐采取的行动，就像在策略梯度模型中一样。
Critic 将估算这些行动可能带来的奖励。它会返回在给定状态下未来的总估计奖励。

让我们定义这样一个模型：

In [ ]:

from itertools import count
import torch.nn.functional as F

In [ ]:

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
env = gym.make("CartPole-v1")

state_size = env.observation_space.shape[0]
action_size = env.action_space.n
lr = 0.0001

class Actor(torch.nn.Module):
    def __init__(self, state_size, action_size):
        super(Actor, self).__init__()
        self.state_size = state_size
        self.action_size = action_size
        self.linear1 = torch.nn.Linear(self.state_size, 128)
        self.linear2 = torch.nn.Linear(128, 256)
        self.linear3 = torch.nn.Linear(256, self.action_size)

    def forward(self, state):
        output = F.relu(self.linear1(state))
        output = F.relu(self.linear2(output))
        output = self.linear3(output)
        distribution = torch.distributions.Categorical(F.softmax(output, dim=-1))
        return distribution


class Critic(torch.nn.Module):
    def __init__(self, state_size, action_size):
        super(Critic, self).__init__()
        self.state_size = state_size
        self.action_size = action_size
        self.linear1 = torch.nn.Linear(self.state_size, 128)
        self.linear2 = torch.nn.Linear(128, 256)
        self.linear3 = torch.nn.Linear(256, 1)

    def forward(self, state):
        output = F.relu(self.linear1(state))
        output = F.relu(self.linear2(output))
        value = self.linear3(output)
        return value

我们需要稍微修改我们的 discounted_rewards 和 run_episode 函数：

In [ ]:

def discounted_rewards(next_value, rewards, masks, gamma=0.99):
    R = next_value
    returns = []
    for step in reversed(range(len(rewards))):
        R = rewards[step] + gamma * R * masks[step]
        returns.insert(0, R)
    return returns

def run_episode(actor, critic, n_iters):
    optimizerA = torch.optim.Adam(actor.parameters())
    optimizerC = torch.optim.Adam(critic.parameters())
    for iter in range(n_iters):
        state = env.reset()
        log_probs = []
        values = []
        rewards = []
        masks = []
        entropy = 0
        env.reset()

        for i in count():
            env.render()
            state = torch.FloatTensor(state).to(device)
            dist, value = actor(state), critic(state)

            action = dist.sample()
            next_state, reward, done, _ = env.step(action.cpu().numpy())

            log_prob = dist.log_prob(action).unsqueeze(0)
            entropy += dist.entropy().mean()

            log_probs.append(log_prob)
            values.append(value)
            rewards.append(torch.tensor([reward], dtype=torch.float, device=device))
            masks.append(torch.tensor([1-done], dtype=torch.float, device=device))

            state = next_state

            if done:
                print('Iteration: {}, Score: {}'.format(iter, i))
                break


        next_state = torch.FloatTensor(next_state).to(device)
        next_value = critic(next_state)
        returns = discounted_rewards(next_value, rewards, masks)

        log_probs = torch.cat(log_probs)
        returns = torch.cat(returns).detach()
        values = torch.cat(values)

        advantage = returns - values

        actor_loss = -(log_probs * advantage.detach()).mean()
        critic_loss = advantage.pow(2).mean()

        optimizerA.zero_grad()
        optimizerC.zero_grad()
        actor_loss.backward()
        critic_loss.backward()
        optimizerA.step()
        optimizerC.step()

现在我们将运行主要训练循环。我们将通过计算适当的损失函数并更新网络参数来使用手动网络训练过程：

In [ ]:


actor = Actor(state_size, action_size).to(device)
critic = Critic(state_size, action_size).to(device)
run_episode(actor, critic, n_iters=100)

In [ ]:

env.close()

要点#

在这个演示中，我们看到了两种强化学习算法：简单的策略梯度算法和更复杂的演员-评论家算法。可以看到，这些算法使用的是状态、动作和奖励的抽象概念，因此它们可以应用于非常不同的环境。

强化学习使我们能够仅通过观察最终奖励来学习解决问题的最佳策略。不需要标注数据集这一点使我们可以多次重复模拟来优化模型。然而，强化学习领域仍然存在许多挑战，如果你决定更多地关注这个有趣的人工智能领域，你将会进一步了解这些挑战。

免责声明：
本文档使用AI翻译服务 Co-op Translator 进行翻译。尽管我们努力确保翻译的准确性，但请注意，自动翻译可能包含错误或不准确之处。原始语言的文档应被视为权威来源。对于重要信息，建议使用专业人工翻译。我们不对因使用此翻译而产生的任何误解或误读承担责任。