feature(xjy): add the rnd-related features

[xiongjyu]

实验记录文档

https://ai.feishu.cn/wiki/CKOBwu0NJicxa8kkcvOcPQ3yn9H?from=from_copylink

总结

目标

在 Venture 等长期稀疏外部奖励环境下，UniZero/MuZero 仅靠外部回报往往难以形成有效探索轨迹。此研究旨在调优 RND（Random Network Distillation），尽量对齐 OpenAI 原始 RND 实现的网络结构与关键参数，通过充分利用内在奖励，用于提升 UniZero/MuZero 的探索能力。 

准备工作

测试 RND 的开源实现（pytorch版本） 并对核心细节进行消融

• Collector:

  • 橙色是原实验结果
  • 天蓝色是归一化内在奖励不使用return
  • 蓝色是优化RND网络不使用当下step的结果，使用下一步obs
  • 红色不使用adv，只使用内在奖励
  

• Learner
  

• 结论：

  1. 内在奖励归一化使用 累积折扣奖励是十分必要的。
  2. 仅使用内在奖励与外在奖励加权（不涉及内在优势）对性能影响十分大，但在 UniZero/MuZero上暂且先不考虑 内在优势。

核心实现细节

模型架构

• RND 网络：
  • predict 和 target 网络的 backbone 均使用3层卷积网络，统一使用 LeakyReLU 激活函数，相关参数如下：
    • hidden_size_list = [32,64,64]
    • kernel_size_list = [8,4,3]
    • stride_size_list = [4,2,1]
  • 由于 target 网络完全冻结，predict 网络需要拟合 target 网络，所以 predict 网络应该比 target 网络更复杂，泛化能力更强。因此 target 网络在 backbone之后加上一层 FC 网络； predict 网络在 backbone 之后加三层 FC 网络。
  • 输入归一化：收集大量的 obs 数据，通过running_mean_std 进行归一化，计算方式为 (obs-mean)/std，此外对归一化后的结果缩放到[-5, 5]之间。每次输入一批数据先归一化再进行后续处理，处理完后再将这批数据计入 running_mean_std中。
    • 对于 图片数据，归一化的统计数据的是原始 CHW 的输入。
    • 对于 文本数据，归一化的统计数据是 文本经过冻结的预训练编码模型的embedding结果。
  • 内在奖励归一化：定义一个 running_mean_std 统计数据进行归一化，计算方式为 r_int/std（因为内在奖励一定不小于0，所以不减去均值）。计入running_mean_std的统计数据实现了2种：
    • 基于即时的内在奖励：将 每次收集的内在奖励存入 running_mean_std。
    • 基于累积折扣的内在奖励和：将一整条轨迹（从第一步到当前步）的所有内在奖励进行折扣求和，将最后结果存入running_mean_std。由 RND部分的消融证明，这种方式十分必要。
  • 外在奖励归一化：实现了 2种归一化方式，一种是与输入归一化一致的通过running_mean_std来实现，计算方式为 (obs-mean)/std；一种是直接使用 sign()函数来判断奖励是正的还是负的，将奖励规范到{-1,0,1}。
  • 内在奖励权重的自适应调整: 由于初期所有状态都是新颖状态，导致初期的内在奖励普遍虚高（噪声很大），为了防止初期过高的内在奖励对 target_reward的影响过大，使其认为一开始探索的轨迹都是好轨迹，从而破坏其探索能力，实现了两种内在奖励权重的自适应变化方法，一种是线性增长，一种是cosine增长。
• 内外在奖励融合：原始的world_model 的 dynamic head 预测的 reward 通过仅包含外在奖励的 target_reward来训练；+RND方式下world_model的dynamic head预测的 reward 包含外在奖励与内在奖励的加权和，即 target_reward = 外在奖励 + 内在奖励 * w。
• 将 world_model 和 rnd 网络进行耦合，使用同一个优化器、相同 learning rate 与 weight decay 进行训练。

训练过程

1. 策略初始化：初始化参数， 创建策略网络（包含world_model和 RND网络）。
2. 初始观测归一化: 在正式开始有策略的探索前，让 agent 在环境中完全随机地采样大量数据。收集它看到的所有状态，用来做 RND的观察归一化。
3. 数据收集：让 agent 根据其当前的策略网络 与环境进行交互，把这期间的所有经验都记录在 replay buffer 里。（目前不涉及 RND）
4. 策略网络训练：在 replay buffer 中随机采样一批数据，首先 通过 RND网络计算内在奖励，并与外在奖励进行融合得到新的 target_reward，然后计算这批数据经过 RND 网络的 MSE 结果作为 loss 与 world_model 的 loss 进行加权，最后一起反向传播并更新参数。计算内在奖励过程： 先对数据的输入进行归一化，然后通过 RND网络给出 predict 和 target网络的结果，计算 MSE值 作为内在奖励；再对内在奖励进行归一化。

监控指标

• 为了方便分析 训练过程的 rnd 网络变化， 补充了与 RND相关的一些监控指标，包含 内在奖励平均/最大/方差值；内外在奖励融合以后的加权结果的平均/最大/方差值，rnd loss以及内在奖励权重的变化。
• 在 evaluate 阶段，还对整个episode的每一步都统计了 内在奖励，绘制成一个折线图，并将内在奖励最大的12 帧画面渲染成图片方便进行详细的分析。

遇到的问题与解决方案

初始观测归一化过程采集数据过慢

原始的实现代码基于 MuZeroCollector 来实现随机采样数据，每次还需要经过 policy网络，然后在随机采样一个动作。由于这个过程需要采集大量数据，所以大量与policy网络的交互导致速度过慢。

• 解决方案：手动初始化N个环境，不创建collector/policy等，每次与环境交互得到obs以后随机采样一个action进行交互。

不同变体实验对比

MuZero + RND 的实验

Venture环境

• 长期实验：MuZero收集了3.6M envsteps; MuZero + RND 收集了 5M envsteps

  • collector

    • 红色是MuZero的实验
    • 绿色是MuZero+RND的实验；内在奖励归一化使用的累积折扣奖励，不使用自适应内在奖励权重
    • 橙色与蓝色的区别是否是否自适应的内在奖励权重
    

  • evaluator
    

  • learner
    
    

  • rnd_reward_model
    
    

UniZero + RND的实验

Venture环境

• Collector:

  • 简介：
    • Rnd loss权重都为1； 内在奖励权重都为0.1；
    • 蓝色曲线实验为内在奖励权重不使用自适应权重；红色为使用自适应权重，10k-20k权重由0->0.1
  • 分析：
    • 目前 只有使用自适应权重的结果收集到了 外在奖励为1的轨迹
    

• Learner:

  • 分析：
    • 蓝色曲线整体收敛很快，基本在10k-20k期间整个 world model都收敛了，也就丧失了探索能力，导致性能差；
    • 红色曲线整体来看探索能力一直存在；并且熟练较为缓慢
  

• RND:

  • 蓝色曲线由于一开始就使用了内在奖励，而训练前期，所有状态的新颖度都很高，导致内在奖励很高，模型会以为前期探索的每一步都是选择了最好的动作；但实际上外在奖励一直为0，导致收敛很快。
  • 红色曲线等训练了一段时间world_model， 并收集了部分数据以后开始使用内在奖励，这个时候内在奖励属于一个相对稳定的阶段，只有遇到很少访问的状态才会给一个较大的内在奖励；
  

Zork1环境

• Collector:

  • 简介：
    • Rnd loss权重都为1； 内在奖励权重都为0.5；
    • 蓝色曲线实验为内在奖励权重不使用自适应权重；红色为使用自适应权重，1k-2k权重由0->0.5
  

• Learner:

  • 分析：
    • 蓝色曲线整体收敛很快，基本在0k-1k期间整个 world model都收敛了，也就丧失了探索能力，导致性能差；
    • 红色曲线整体来看探索能力一直存在；并且收敛较为缓慢
  

• RND: