一、数据处理

用于训练 RNN 的 mini-batch 数据不同于通常的数据。 这些数据通常应按时间序列排列。 对于 DI-engine, 这个处理是在 collector 阶段完成的。 用户需要在配置文件中指定 learn_unroll_len 以确保序列数据的长度与算法匹配。 对于大多数情况, learn_unroll_len 应该等于 RNN 的历史长度(a.k.a 时间序列长度),但在某些情况下并非如此。比如,在 r2d2 中, 我们使用burn-in操作, 序列长度等于 learn_unroll_len + burnin_step 。 这里将在下一节中具体解释。

什么是数据处理?

数据处理指的是为循环神经网络(RNN)训练准备时间序列数据的过程。这个过程包括将收集到的数据组织成适当格式的小批量(mini-batches),这些批量数据将用于网络的训练。这一步骤通常发生在DI-engine的collector阶段,也就是数据收集和预处理发生的地方。用户需要在配置文件中指定 learn_unroll_len 以确保序列数据的长度与算法匹配。 对于大多数情况, learn_unroll_len 应该等于 RNN 的历史长度(a.k.a 时间序列长度),但在某些情况下并非如此。比如,在 r2d2 中, 我们使用burn-in操作, 序列长度等于 learn_unroll_len + burnin_step 。例如,如果你设置 learn_unroll_len = 10 和 burnin_step = 5,那么 RNN 实际接收的输入序列长度将是 15:前 5 步为 burn-in(用于预热隐藏状态),接下来的 10 步作为学习的一部分。这样设置可以帮助 RNN 在计算梯度和进行权重更新时,有一个更加准确的隐藏状态作为起点。
部分名词解释

  • mini-batches:在机器学习中,特别是在训练神经网络时,数据一般被分成小的批次进行处理,这些批次被称为 “mini-batch”。一个 mini-batch 包含了一组样本,这组样本用于执行单次迭代的前向传播和反向传播,以更新网络的权重。使用 mini-batches 而不是单个样本或整个数据集(后者称为 “batch” 或 “full-batch”)可以平衡计算效率和内存限制,有助于提高学习的稳定性和收敛速度。
  • collector阶段:在 DI-engine中,collector 阶段是指环境与智能体交互并收集经验数据的过程。在这个阶段,智能体根据其当前的策略执行操作,环境则返回新的状态、奖励和其他可能的信息,如是否达到终止状态。收集到的数据(经常被称为经验或转换)随后被用于训练智能体的模型,例如对策略或价值函数进行更新。

为什么要进行数据处理:

  1. 保持时间依赖性:RNN的核心优势是处理具有时间序列依赖性的数据,比如语言、视频帧、股票价格等。正确的数据处理确保了这些时间依赖性在训练数据中得以保留,使得模型能够学习到数据中的序列特征。

  2. 提高学习效率:通过将数据划分为与模型期望的序列长度匹配的批次,可以提高模型学习的效率。这样做可以确保网络在每次更新时都接收到足够的上下文信息。

  3. 适配算法要求:不同的RNN算法可能需要不同形式的输入数据。例如,标准的RNN只需要过去的信息,而一些变体如LSTM或GRU可能会处理更长的序列。特定的算法,如R2D2,还可能需要额外的步骤(如burn-in),以便更好地初始化网络状态。

  4. 处理不规则长度:在现实世界的数据集中,序列长度往往是不规则的。数据处理确保了每个mini-batch都有统一的序列长度,这通常通过截断过长的序列或填充过短的序列来实现。

  5. 优化内存和计算资源:通过将数据组织成具有固定时间步长的批次,可以更有效地利用GPU等计算资源,因为这些资源在处理固定大小的数据时通常更高效。

  6. 稳定学习过程:特别是在强化学习中,使用如n-step返回或经验回放的技术,可以帮助模型从环境反馈中学习,并减少方差,从而稳定学习过程。

    如何进行数据处理

    比如原始采样数据是 [x_1,x_2,x_3,x_4,x_5,x_6],每个x表示 [s_t,a_t,r_t,d_t,s_{t+1}](或者log_\pi(a_t|s_t) ,隐藏状态等),此时 n_sample = 6 。此时根据所需 RNN 的序列长度即 learn_unroll_len 有以下三种情况:

  1. n_sample >= learn_unroll_len 并且 n_sample 可以被 learn_unroll_len 除尽: 例如 n_sample=6 和 learn_unroll_len=3,数据将被排列为:[[x_1,x_2,x_3],[x_4,x_5,x_6]]

  2. n_sample >= learn_unroll_len 并且 n_sample 不可以被 learn_unroll_len 除尽: 默认情况下,残差数据将由上一个样本中的一部分数据填充,例如如果 n_sample=6 和 learn_unroll_len=4 ,数据将被排列为 [[x_1,x_2,x_3,x_4],[x_3,x_4,x_5,x_6]]

  3. n_sample < learn_unroll_len:例如如果 n_sample=6 和 learn_unroll_len=7,默认情况下,算法将使用 null_padding 方法,数据将被排列为 [[x_1,x_2,x_3,x_4,x_5,x_6,x_{null}]]x_{null}类似于x_6 但它的 done=True 和 reward=0。

这里以r2d2算法为例,在r2d2中,在方法 _get_train_sample 中通过调用函数 get_nstep_return_data 和 get_train_sample 获取按时序排列的数据。

def _get_train_sample(self, data: list) -> Union[None, List[Any]]:
    data = get_nstep_return_data(data, self._nstep, gamma=self._gamma)
    return get_train_sample(data, self._sequence_len)

代码段 def _get_train_sample(self, data: list) 是一个方法,它的作用是从收集到的数据中提取用于训练 RNN 的样本。这个方法会在两个步骤中处理数据:

  • N步返回计算(get_nstep_return_data): 这个函数接受原始的经验数据,然后计算所谓的 N 步返回值。N 步返回是一个在强化学习中用于临时差分(Temporal Difference, TD)学习的概念,它考虑了从当前状态开始的未来 N 步的累积奖励。计算这个值需要使用折现因子 gamma。这个步骤的目的是为了让智能体学习如何根据当前的行动预测未来的奖励,这是强化学习中价值函数估计的重要部分。

  • 训练样本获取(get_train_sample): 在得到 N 步返回值之后,这个函数进一步处理数据以生成训练样本。具体地,它会根据 self._sequence_len(即时间序列长度或者 RNN 的历史长度)来选择数据序列。这意味着每个训练样本将是一个具有 self._sequence_len 长度的数据序列,这对于训练 RNN 来说是必要的,因为 RNN 需要一定长度的历史来维护其内部状态(或记忆)。
    有关这两个数据处理功能的工作流程见下图:

二、初始化隐藏状态 (Hidden State)

RNN用于处理具有时间依赖性的信息。RNN的隐藏状态(Hidden State)是其记忆的一部分,它能够捕捉到前一时间步长的信息。这些信息对于预测下一个动作或状态非常关键。在此上下文中,初始化RNN的隐藏状态是一个重要的步骤,它确保了RNN在开始新的数据批次处理时具有正确的起始状态。
策略的 _learn_model 需要初始化 RNN。这些隐藏状态来自 _collect_model 保存的 prev_state。 用户需要通过 _process_transition 函数将这些状态添加到 _learn_model 输入数据字典中。

def _process_transition(self, obs: Any, model_output: dict, timestep: namedtuple) -> dict:
    transition = {
        'obs': obs,
        'action': model_output['action'],
        'prev_state': model_output['prev_state'], # add ``prev_state`` key here
        'reward': timestep.reward,
        'done': timestep.done,
    }
    return transition

此函数用于处理从环境中收集到的每个时间步长的观察值、模型输出和时间步长信息,并将它们整理为一个字典,称为transition,这个字典包括如下键值对:

  • 'obs': 当前观察值。
  • 'action': 模型输出的动作。
  • 'prev_state': 模型输出的前一个隐藏状态,这是RNN在当前步骤之前的内部状态。
  • 'reward': 当前时间步长的奖励。
  • 'done': 表示当前时间步长是否是序列末尾的标志。
    存储 prev_state
  • transition字典中,'prev_state'键存储模型输出的前一个隐藏状态,这个状态将被用来初始化_learn_model中RNN的隐藏状态。

然后在 _learn_model 前向函数中, 调用它的重置函数 ( 对应 HiddenStateWrapper 里面的重置函数) 以用来初始化 RNN 的 prev_state。

def _forward_learn(self, data: dict) -> Dict[str, Any]:
     # forward
     data = self._data_preprocess_learn(data)
     self._learn_model.train()
     self._learn_model.reset(data_id=None, state=data['prev_state'][0])

这是策略的学习模型的前向传播函数。在该函数中,RNN的隐藏状态通过调用模型的reset函数来初始化。这是在每个训练批次开始时完成的,以确保RNN从正确的状态开始学习。

reset函数是HiddenStateWrapper类的一部分,它负责重置RNN的内部状态。

Burn-in(in R2D2)

Burn-in的概念来自 R2D2 (Recurrent Experience Replay In Distributed Reinforcement Learning)论文。在 R2D2 算法中,由于 LSTM 需要处理时间序列数据,因此它需要有一个合理的初始隐藏状态。burn-in 期是为了让 LSTM 有机会通过处理一些初始的序列数据来”预热”,这样它就能够在正式开始学习之前构建起一些有意义的内部状态。论文指出在使用 LSTM 时,最基础的方式是:

1.将完整的 episode 轨迹切分为很多序列样本。在每个序列样本的初始时刻,使用全部为0的 tensor 作为 RNN 网络的初始化 hidden state。

2.使用完整的 episode 轨迹用于 RNN 训练。

对于第一种方法,由于每个序列样本的初始时刻的 hidden state 应该包含之前时刻的信息,这里简单使用全为0的 Tensor 带来很大的 bias 对于第二种方法,往往在不同环境上,完整的一个episode的长度是变化的,很难直接用于 RNN 的训练。

Burn-in 给予 RNN 网络一个 burn-in period。 即使用 replay sequence 的前面一部分数据产生一个开始的隐藏状态 (hidden state),然后仅在 replay sequence 的后面一部分数据上更新 RNN 网络。

在 DI-engine 中,r2d2 使用 n-step td error, 即, self._nstep 是 n 的数量。 sequence length = burnin_step + learn_unroll_len. 所以在配置文件中, learn_unroll_len 应该设置为 sequence length - burnin_step。

在此设置中,原始展开的 obs 序列被拆分为 burnin_nstep_obs , main_obs 和 marget_obs。 burnin_nstep_obs 是 用于计算 RNN 的初始隐藏状态,用便未来用于计算 q_value、target_q_value 和 target_q_action。 main_obs 用于计算 q_value。在下面的代码中, [bs:-self._nstep] 表示使用来自的数据 bs 时间步长到 sequence length - self._nstep 时间步长。 target_obs 用于计算 target_q_value。
这个数据处理可以通过下面的代码来实现:

data['action'] = data['action'][bs:-self._nstep]
data['reward'] = data['reward'][bs:-self._nstep]

data['burnin_nstep_obs'] = data['obs'][:bs + self._nstep]
data['main_obs'] = data['obs'][bs:-self._nstep]
data['target_obs'] = data['obs'][bs + self._nstep:]

具体到代码的分析:

  • data[‘action’] 和 data[‘reward’] 的切片 [bs:-self._nstep] 表示从序列的第 bs 个时间步开始取,直到倒数第 self._nstep 个时间步结束。bs 代表burn-in步数,这样可以保证在计算TD误差时,动作和奖励与预测的Q值对应的时间步是一致的。

  • data[‘burnin_nstep_obs’] 存储了用于RNN初始化隐藏状态计算的观测序列。这里的 [:bs + self._nstep] 切片意味着取序列的开始直到 bs + self._nstep 时间步。这部分数据不会用于梯度更新,只用于生成RNN的初始状态。

  • data[‘main_obs’] 代表用于实际计算Q值的主观测序列。通过切片 [bs:-self._nstep] 获取,这样确保了主观测序列排除了用于burn-in的部分和用于计算目标Q值的最后几步(因为需要n步TD误差)。

  • data[‘target_obs’] 存储了用于计算目标Q值的观测序列。切片 [bs + self._nstep:] 意味着从burn-in期加上n步之后开始取,直到序列结束。

在R2D2算法中使用burn-in技术时,需要在每个时间步更新并保存RNN的隐藏状态,以便在学习阶段能够使用正确的状态来初始化RNN。由于RNN的隐藏状态是基于时间序列的,因此必须采用适当的方法来处理这一点。

收集模型(self._collect_model

当调用self._collect_modelforward方法进行数据收集时,通常设置inference=True。在这个模式下,每次只处理单个时间步的数据,这样就能在每个时间步获取到RNN的当前隐藏状态(prev_state)。

学习模型(self._learn_model

在学习阶段,调用self._learn_modelforward方法时则设置inference=False。在非推理模式下,传入的是一系列数据,而输出的prev_state字段仅代表序列中最后一个时间步的隐藏状态。

保存特定时间步的隐藏状态

为了保存除了最后一个时间步之外的特定时间步的隐藏状态,可以通过saved_hidden_state_timesteps参数来指定。这个参数是一个列表,表示需要保存隐藏状态的时间步。

在R2D2中的应用

在R2D2的实现中,通过指定saved_hidden_state_timesteps=[self._burnin_step, self._burnin_step + self._nstep],可以在调用网络的forward方法后,保存burnin_outputburnin_output_target中的特定时间步的隐藏状态。这些保存的隐藏状态将被用在后续的计算过程中,例如计算Q值(q_value)、目标Q值(target_q_value)和目标动作(target_q_action)。

示例代码

def _forward_learn(self, data: dict) -> Dict[str, Any]:
    # forward
    data = self._data_preprocess_learn(data)
    self._learn_model.train()
    self._target_model.train()
    # use the hidden state in timestep=0
    self._learn_model.reset(data_id=None, state=data['prev_state'][0])
    self._target_model.reset(data_id=None, state=data['prev_state'][0])

    if len(data['burnin_nstep_obs']) != 0:
        with torch.no_grad():
            inputs = {'obs': data['burnin_nstep_obs'], 'enable_fast_timestep': True}
            burnin_output = self._learn_model.forward(
                inputs, saved_hidden_state_timesteps=[self._burnin_step, self._burnin_step + self._nstep]
            )
            burnin_output_target = self._target_model.forward(
                inputs, saved_hidden_state_timesteps=[self._burnin_step, self._burnin_step + self._nstep]
            )

    self._learn_model.reset(data_id=None, state=burnin_output['saved_hidden_state'][0])
    inputs = {'obs': data['main_obs'], 'enable_fast_timestep': True}
    q_value = self._learn_model.forward(inputs)['logit']
    self._learn_model.reset(data_id=None, state=burnin_output['saved_hidden_state'][1])
    self._target_model.reset(data_id=None, state=burnin_output_target['saved_hidden_state'][1])

    next_inputs = {'obs': data['target_obs'], 'enable_fast_timestep': True}
    with torch.no_grad():
        target_q_value = self._target_model.forward(next_inputs)['logit']
        # argmax_action double_dqn
        target_q_action = self._learn_model.forward(next_inputs)['action']

这段代码是R2D2算法中学习过程的一个实现示例,它展示了如何在强化学习模型中结合使用burn-in技术和RNN。代码主要分为几个部分:

  1. 数据预处理:
    `self._data_preprocess_learn(data) 可能对输入数据进行了一些标准化处理或其他预处理步骤。

  2. 模型设置为训练模式:
    `self._learn_model.train()self._target_model.train() 将两个模型设置为训练模式,这在PyTorch中会启用梯度计算。

  3. 重置模型状态:
    使用 self._learn_model.reset()self._target_model.reset() 函数将两个模型的隐藏状态初始化为data['prev_state'][0],即序列中第一个时间步的隐藏状态。

  4. 执行burn-in过程:
    如果 data['burnin_nstep_obs'] 非空,即存在burn-in序列,则执行burn-in过程。在 torch.no_grad() 上下文中,这意味着在这个过程中不会计算梯度。burnin_outputburnin_output_target 分别保存了学习模型和目标模型在burn-in步骤后的输出,包括保存的隐藏状态。

  5. 计算Q值:
    然后代码通过 self._learn_model.reset() 将学习模型的隐藏状态设置为burn-in阶段最后的状态,准备计算主要观测序列的Q值。q_value 是根据主观测序列(data['main_obs'])计算出的Q值。

  6. 计算目标Q值和目标动作:
    接下来,使用 self._learn_model.reset()self._target_model.reset() 将两个模型的隐藏状态设置为burn-in阶段计算的第二个状态。这个状态用于计算下一步观测序列(data['target_obs'])的目标Q值(target_q_value)和目标动作(target_q_action),这些值用于更新Q网络的参数。

这段代码的关键在于,它使用了一个burn-in序列来生成RNN的初始隐藏状态,并将这个状态用于计算后续序列的Q值。这样做可以缓解由于隐藏状态初始化为零带来的偏差问题。此外,代码还展示了如何处理不同时间步的隐藏状态,这对于RNN在处理时间序列数据时至关重要。在强化学习中,这种处理方式有助于提高算法的稳定性和性能。