强化学习 | (5) RUDDER：A practical tutorial

项目地址

A step-by-step guide to applying RUDDER

在本教程中，我将向您展示如何逐步应用RUDDER以及如何使用PyTorch实现奖励重新分配模型。 您可以将其用作快速指南，以将RUDDER应用于您的RL设置，并预先评估RUDDER是否可以事先改善您的任务。 该代码可以在合理的时间内在通用CPU上运行。

RUDDER Blog, RUDDER Paper, Rudder Repo

Overview: How to apply RUDDER:

我建议按照以下步骤检查RUDDER是否可以解决您的RL问题并应用RUDDER：

1. 从您的环境中取一组样本（与环境交互）。

• 一个样本是一个 (状态，动作，奖励)序列/轨迹
• 样本的return（奖励之和）应该有差异/方差。
• 样本不必来自同一策略，您还可以包括人类示范等。

2. 在您的样本上训练LSTM（或您想用于奖励重新分配的任何模型类别）。

• 这部分是一个监督机器学习任务
• 训练的主要任务是预测序列结束时每个样本的return。
• 作为辅助任务，我们将训练模型以预测每个序列位置的最终return。 这将允许我们能够使用预测的差异来做贡献分析，并且在某些情况下，这也有助于将奖励推回更远的时间。 与主要任务相比，降低此任务的贡献。
• 如果您的环境还包括不在序列末尾的中间奖励，则可以使用本文所述的辅助任务通过这些中间奖励引入梯度，以简化学习。
• 我建议尝试使模型尽可能简单，因为这可以使我的实验中的奖励重新分配和贡献分析更加清晰。
  A. 从没有遗忘门和没有输出门的LSTM开始，单元输入仅接收前向连接，而门仅接收递归连接。
  B. 增加模型的复杂度，直到模型能够足够好地预测主要任务（在序列的末尾的return）为止。 您可以按此顺序包括输出门，遗忘门和全连接的门来完成此操作。

3. 检查reward 重新分布是否有意义

• 查看贡献分析的输出，例如 LSTM在每个序列位置的预测差异，并检查奖励是否重新分配到序列中的较早位置。
• 检查模型是否过拟合。 如果是这样，请增加样本数量或降低模型复杂性。
• 此时，您会感觉到RUDDER是否可以帮助您完成任务。

4. 设置Lessons Buffer以在agent的训练期间训练奖励重新分配模型。

• 从前面的观点中，您应该了解LSTM要正确学习必须具有多大的数据集。
• 我建议使用Lessons Buffer，有较高奖励重新分配模型损失的样本优先处理，并增加buffer中样本returns的方差。

5. 使用重新分配的奖励来训练agent，例如 通过PPO。

• 对于基于AC的方法，您可以替换奖励信号，即Critic将通过重新分配的奖励对其进行训练。
• 如果您的环境具有即时奖励，则可能需要将原始奖励信号与重新分配的奖励混合在一起。

以下各节将显示第1-3点的示例，他们可以作为预评估，如果RUDDER可以帮助您完成任务。

Importing packages

我将使用我仓库中的widis-lstm-tools v0.4来简化LSTM。 您可以通过下面的命令安装：

pip install git+https://github.com/widmi/widis-lstm-tools

import numpy as np
import torch
import tqdm
from torch.utils.data import Dataset
from matplotlib import pyplot as plt
from widis_lstm_tools.nn import LSTMLayer# Prepare some random generators for later
rnd_gen = np.random.RandomState(seed=123)
_ = torch.manual_seed(123)

Creating an example task with delayed reward

为简单起见，我们使用一个简单的一维环境，并使用随机策略与环境交互产生样本。可以随意用您的任务样本替换此环境。

在此环境中，agent从位置0开始，并且可以在每个时间步可以选择向右或向左移动，这将分别使agent的位置增加或减少+1或-1。 会离开位置范围[-6、6]的动作将被忽略，因此不同的状态数量限制为13。状态仅由当前位置组成。 agent在每个轨迹/情节可以执行总共50个动作，即每个游戏序列/轨迹的长度为50。

可视化：

如果agent处于位置2的状态，它将获得+1奖励。 我们将此任务转换为延迟奖励任务，在该任务中我们知道立即奖励会是什么样子。 为此，我们假设没有立即向agent显示奖励，而是在最后一个时间步显示了累积的奖励。 请注意，这不再是MDP，但是对于使用LSTM的奖励重新分配，这不会影响我们。

class Environment(Dataset):def __init__(self, n_samples: int, max_timestep: int, n_positions: int, rnd_gen: np.random.RandomState):"""Our simple 1D environment as PyTorch Dataset"""super(Environment, self).__init__()n_actions = 2    # 两个动作zero_position = int(np.ceil(n_positions / 2.)) # 0位置 7coin_position = zero_position + 2 #有奖励的位置# Generate random action sequencesactions = np.asarray(rnd_gen.randint(low=0, high=2, size=(n_samples, max_timestep)), dtype=np.int) #随机动作序列 (样本数/序列数，最大序列长度/时间步)actions_onehot = np.identity(n_actions, dtype=np.float32)[actions] #把动作索引转换为one-hot动作表示# Generate observations from action sequencesactions[:] = (actions * 2) - 1observations = np.full(fill_value=zero_position, shape=(n_samples, max_timestep), dtype=np.int) # 状态序列for t in range(max_timestep-1):action = actions[:, t]observations[:, t+1] = np.clip(observations[:, t] + action, 0, n_positions-1) #对于t时刻的状态 施加t时刻的动作 的得到t+1时刻的状态observations_onehot = np.identity(n_positions, dtype=np.float32)[observations] # 把状态转换为one-hot表示# Calculate rewards (sum over coin position for all timesteps)rewards = np.zeros(shape=(n_samples, max_timestep), dtype=np.float32) rewards[:, -1] = observations_onehot[:, :, coin_position].sum(axis=1) # 把每个轨迹的reward加起来 放在最后一个时间步上self.actions = actions_onehotself.observations = observations_onehotself.rewards = rewardsdef __len__(self):return self.rewards.shape[0]def __getitem__(self, idx):return self.observations[idx], self.actions[idx], self.rewards[idx]n_positions = 13
env = Environment(n_samples=1000, max_timestep=50, n_positions=13, rnd_gen=rnd_gen)
env_loader = torch.utils.data.DataLoader(env, batch_size=8, num_workers=4)

obs0, a0, r0 = env.__getitem__(3)  # 一个样本/轨迹 的状态、动作以及奖励序列
obs1, a1, r1 = env.__getitem__(25)
fig, axes = plt.subplots(3, 2, figsize=(8, 4.5), dpi=100)
axes[0, 0].plot(obs0.argmax(-1) - 6)
axes[0, 1].plot(obs1.argmax(-1) - 6)
axes[0, 0].set_ylim(-6, 6)
axes[0, 1].set_ylim(-6, 6)
axes[0, 0].axhline(2, linestyle='--', color='r')
axes[0, 1].axhline(2, linestyle='--', color='r')
axes[0, 0].xaxis.grid(True)
axes[0, 1].xaxis.grid(True)
axes[0, 0].set_title('observations (sample 1)')
axes[0, 1].set_title('observations (sample 2)')
axes[0, 0].set_xlabel('time (environment steps)')
axes[0, 1].set_xlabel('time (environment steps)')axes[1, 0].plot(a0.argmax(-1))
axes[1, 1].plot(a1.argmax(-1))
axes[1, 0].xaxis.grid(True)
axes[1, 1].xaxis.grid(True)
axes[1, 0].set_title('actions (sample 1)')
axes[1, 1].set_title('actions (sample 2)')
axes[1, 0].set_xlabel('time (environment steps)')
axes[1, 1].set_xlabel('time (environment steps)')axes[2, 0].plot(r0)
axes[2, 1].plot(r1)
axes[2, 0].xaxis.grid(True)
axes[2, 1].xaxis.grid(True)
axes[2, 0].set_title('original rewards (sample 1)')
axes[2, 1].set_title('original rewards (sample 2)')
axes[2, 0].set_xlabel('time (environment steps)')
axes[2, 1].set_xlabel('time (environment steps)')fig.tight_layout()

…我们可以看到我们创建了有延迟奖励的任务。 每当agent在位置2时，它会在序列结束时获得+1奖励。

现在，让我们看看RUDDER在这种情况下可以为我们做些什么。

Training a reward redistribution model

现在，我们将训练一个LSTM模型来执行此任务上的奖励重新分配。

Defining the model

我们的模型将由一个带有16个单元的LSTM层组成。 我们将逐时间步将拼接的动作和观察输入到网络中。 我们还可以使用单独的网络来预处理动作和观察/状态。 在我们的实验中，它有助于在拼接之前使动作和观测特征的数量相似。 在实践中，不一定总是也要提供观测值的增量（deltas），像我们在论文中做的那样，对于atari游戏。 在此示例中，我将省略它。

我们将从简化的LSTM模型开始，不需要忘记门和输出门。 此外，我们仅将输入门连接到循环连接，将单元输入连接到正向连接。 由于我们正在执行回归以预测return值，因此LSTM的输出函数将是线性的。

如果我们的模型不足以预测序列结束时的return，我们将不得不增加模型的复杂度，直到足够好为止。 您可以按此顺序添加结构包括输出门，遗忘门和全连接门来完成此操作。

class Net(torch.nn.Module):def __init__(self, n_positions, n_actions, n_lstm):super(Net, self).__init__()# This will create an LSTM layer where we will feed the concatenateself.lstm1 = LSTMLayer(in_features=n_positions+n_actions, out_features=n_lstm, inputformat='NLC',# cell input: initialize weights to forward inputs with xavier, disable connections to recurrent inputsw_ci=(torch.nn.init.xavier_normal_, False),# input gate: disable connections to forward inputs, initialize weights to recurrent inputs with xavierw_ig=(False, torch.nn.init.xavier_normal_),# output gate: disable all connection (=no forget gate) and disable biasw_og=False, b_og=False,# forget gate: disable all connection (=no forget gate) and disable biasw_fg=False, b_fg=False,# LSTM output activation is set to identity functiona_out=lambda x: x)# After the LSTM layer, we add a fully connected output layerself.fc_out = torch.nn.Linear(n_lstm, 1)def forward(self, observations, actions):# Process input sequence by LSTM# observations (batch_size, timesteps, obs_dim)# actions (batch_size, timesteps, act_dim)lstm_out, *_ = self.lstm1(torch.cat([observations, actions], dim=-1),return_all_seq_pos=True  # return predictions for all sequence positions)#print(lstm_out.shape) (batch_size, timesteps, lstm_units)net_out = self.fc_out(lstm_out) return net_out # (batch_size, timesteps, 1)# Create Network
device = 'cpu'
net = Net(n_positions=n_positions, n_actions=2, n_lstm=16)
_ = net.to(device)

Defining the loss function

如概述中所述，我们将训练两个任务：

1. 训练的主要任务是预测序列结束时每个样本的return。
2. 辅助任务是训练模型以预测每个序列位置的最终return。 这将使我们能够通过预测差异进行贡献分析。 在某些情况下，这有助于将奖励回推到更远的时间。 对于LSTM来说，预测每个序列位置的最终return通常要困难得多，因为它需要预测未知的未来。 但是，LSTM不需要在此任务上完美，因为正确的return分解仅取决于最后一个时间步长（主任务）的预测。

我们在序列中没有任何即时奖励，因此我们将仅使用这两个任务（主任务，辅助任务）。 我们将辅助任务的权重降低到0.5。 给定模型输出（=预测）和奖励序列，我们可以按以下方式计算损失：

def lossfunction(predictions, rewards): #prediction (batch_size, time_steps)#rewards (batch_size, time_steps)returns = rewards.sum(dim=1) # (batch_size) # Main task: predicting return at last timestepmain_loss = torch.mean(predictions[:, -1] - returns) ** 2# Auxiliary task: predicting final return at every timestep ([..., None] is for correct broadcasting)aux_loss = torch.mean(predictions[:, :] - returns[..., None]) ** 2 # returns[..., None] (batch_size,1)# Combine lossesloss = main_loss + aux_loss * 0.5return loss

Training the model

optimizer = torch.optim.Adam(net.parameters(), lr=1e-3, weight_decay=1e-5)update = 0
n_updates = 10000 #batch-GD 更新次数
running_loss = 100.
progressbar = tqdm.tqdm(total=n_updates)
while update < n_updates:for data in env_loader: # 监督学习设置 训练数据以batch的形式提供# Get samplesobservations, actions, rewards = data 
# print(observations.shape) (batch_size, time_steps,obs_dim)
# print(actions.shape) (batch_size, time_steps,act_dim)
# print(rewards.shape) (batch_size, time_steps)observations, actions, rewards = observations.to(device), actions.to(device), rewards.to(device)# Reset gradientsoptimizer.zero_grad()# Get outputs for networkoutputs = net(observations=observations, actions=actions)  #(batch_size, time_steps, 1)# Calculate loss, do backward pass, and updateloss = lossfunction(outputs[..., 0], rewards) #计算lossloss.backward() #计算梯度running_loss = running_loss*0.99 + loss*0.01 #滑动平均计算 最近几次更新过程的平均lossoptimizer.step() #更新参数update += 1progressbar.set_description(f"Loss: {running_loss:8.4f}")progressbar.update(1)progressbar.close()

Computing the reward redistribution

现在，我们可以将奖励重新分配计算为训练模型的预测差异。 在最佳情况下，奖励将重新分配给导致奖励和return的动作。 实际上，它将分配给模型用来预测return的东西，比如重要的状态。 这是辅助任务帮助将奖励进一步移回此处的地方。 当然，LSTM预测越好，这种奖励重新分配就越好。

现在使用训练好的模型来计算奖励重分布。

# Load 2 samples
obs0, a0, r0 = env.__getitem__(3)
obs1, a1, r1 = env.__getitem__(25)# Apply our reward redistribution model to the samples
observations = torch.stack([torch.Tensor(obs0).to(device), torch.Tensor(obs1).to(device)], dim=0)
actions = torch.stack([torch.Tensor(a0).to(device), torch.Tensor(a1).to(device)], dim=0)
rewards = torch.stack([torch.Tensor(r0).to(device), torch.Tensor(r1).to(device)], dim=0)
predictions = net(observations=observations.to(device), actions=actions.to(device))[..., 0]
#print(predictions.shape) # (batch_size, time_steps)# Use the differences of predictions as redistributed reward 重分布的reward是 当前时刻的预测-前一时刻的预测
redistributed_reward = predictions[:, 1:] - predictions[:, :-1]
#print(redistributed_reward.shape) #(batch_size, timesteps-1)
# For the first timestep we will take (0-predictions[:, :1]) as redistributed reward
redistributed_reward = torch.cat([predictions[:, :1], redistributed_reward], dim=1) #第一个时间步的重分布reward=第一时刻的预测

我们还可以强制要求新序列和原始序列具有相同的return。 我们将通过计算预测误差并将误差的校正平均分布在所有序列位置上来实现：

# Calculate prediction error
returns = rewards.sum(dim=1)
predicted_returns = redistributed_reward.sum(dim=1)
prediction_error = returns - predicted_returns# Distribute correction for prediction error equally over all sequence positions
redistributed_reward += prediction_error[:, None] / redistributed_reward.shape[1]

Inspecting the reward redistribution

最后，让我们看一下两个奖励重新分配的示例：

# Let's plot our samples with the redistributed rewards:
redistributed_reward = redistributed_reward.cpu().detach().numpy()
rr0, rr1 = redistributed_reward[0], redistributed_reward[1]fig, axes = plt.subplots(4, 2, figsize=(8, 6), dpi=100)
axes[0, 0].plot(obs0.argmax(-1) - 6)
axes[0, 1].plot(obs1.argmax(-1) - 6)
axes[0, 0].set_ylim(-6, 6)
axes[0, 1].set_ylim(-6, 6)
axes[0, 0].axhline(2, linestyle='--', color='r')
axes[0, 1].axhline(2, linestyle='--', color='r')
axes[0, 0].xaxis.grid(True)
axes[0, 1].xaxis.grid(True)
axes[0, 0].set_title('observations (sample 1)')
axes[0, 1].set_title('observations (sample 2)')
axes[0, 0].set_xlabel('time (environment steps)')
axes[0, 1].set_xlabel('time (environment steps)')axes[1, 0].plot(a0.argmax(-1))
axes[1, 1].plot(a1.argmax(-1))
axes[1, 0].xaxis.grid(True)
axes[1, 1].xaxis.grid(True)
axes[1, 0].set_title('actions (sample 1)')
axes[1, 1].set_title('actions (sample 2)')
axes[1, 0].set_xlabel('time (environment steps)')
axes[1, 1].set_xlabel('time (environment steps)')axes[2, 0].plot(r0)
axes[2, 1].plot(r1)
axes[2, 0].xaxis.grid(True)
axes[2, 1].xaxis.grid(True)
axes[2, 0].set_title('original rewards (sample 1)')
axes[2, 1].set_title('original rewards (sample 2)')
axes[2, 0].set_xlabel('time (environment steps)')
axes[2, 1].set_xlabel('time (environment steps)')axes[3, 0].plot(rr0)
axes[3, 1].plot(rr1)
axes[3, 0].xaxis.grid(True)
axes[3, 1].xaxis.grid(True)
axes[3, 0].set_title('redistributed rewards (sample 1)')
axes[3, 1].set_title('redistributed rewards (sample 2)')
axes[3, 0].set_xlabel('time (environment steps)')
axes[3, 1].set_xlabel('time (environment steps)')fig.tight_layout()

我们可以看到，原本只在结束时才获得的奖励被转移到使agent更接近位置2的动作上。

这意味着我们不会恢复agent因位置2而获得+1奖励的环境，而是要获得延迟时间更短的奖励的环境，因为趋向位置2的动作会立即得到奖励…

并且我们通过监督的ML方法实现了所有这些目标！

A note on LSTM training

查看细胞状态和激活情况，以检查您的LSTM是在做有用的事情还是只是过拟合 ，例如 序列长度或某处饱和。

对于之前的图中的样本1，训练后的LSTM状态看起来像这样：

# Plost LSTM states for first sample of last processed minibatch
net.lstm1.plot_internals(filename=None, show_plot=True, mb_index=0, fdict=dict(figsize=(8, 8), dpi=100))

A note on contribution analysis

我们发现预测差异是我们实验中非常适合进行贡献分析的方法。 主要归因于三个不错的属性：

1. 奖励重新分配模型已经在未来进行了一些平均（因为它必须在序列的早期预测最终的return）。 这使得学习基础RL代理/方法变得更加容易。 但是，这也是该方法的缺点，因为现在奖励重新分配模型不仅必须学会预测最后一个序列位置的return。 其他贡献分析方法，例如积分梯度或逐层相关性传播，不需要奖励重新分配模型来学习每个序列位置的预测，因为它们可以基于最后一个序列位置的预测执行贡献分析。 但是，这也意味着，如果贡献分析来自完成的情节序列的结尾，则该代理可能会收到针对将来发生的事件的奖励。
2. 经过训练以预测最后一个序列位置return的奖励重新分配模型将使用来自任意序列位置的信息进行预测。 但是，我们希望将奖励尽可能早地重新分配给序列，以减少动作和相应奖励之间的延迟。 这意味着我们希望奖励重新分配模型尽可能早地检测序列中对预测return有用的信息。 而这正是这项辅助任务所要做的。 例如：通过计算球队进球的时间来预测足球比赛的结果很容易。 因此，该模型可能会将这些突出事件用于其预测。 由于我们不仅对分配奖励给这些事件感兴趣，而且还对导致事件的动作（导致目标的一些良好行为）感兴趣，因此我们必须推动模型尽早进行预测。
3. 使用每个序列位置的预测差异意味着我们可以轻松地对不完整序列进行贡献分析。 因此，我们无需等到情节序列完成就可以执行贡献分析并使用重新分配的奖励来训练我们的RL agent/方法。 在实践中，我们可以继续在后台的Lessons Buffer中训练完整情节的奖励重新分配模型，而我们使用相同的奖励重新分配模型对从环境中采样的新序列执行在线奖励重新分配。 这允许以在线方式训练具有重新分配的奖励的agent。

What is left to do?

此时，您可能已经能够评估RUDDER是否对您的任务有所帮助。 然后，您必须设置一个Lessons Buffer，以在策略更改时训练您的奖励重新分配模型，并向RL agent/方法提供重新分配的奖励，如概述中第4点和第5点所述。