【强化学习】MuZero的解释与实现

本文是我自制的强化学习框架 SimpleDistributedRL 的评论文章。

上一篇：AlphaZero
下一个：随机 MuZero

这一次，我将解释 MuZero，它是 AlphaZero 的继任者。

没有零

AlphaZero 在搜索树时遇到了使用游戏规则的问题。
（前进一招后，需要知道板面的状态）
在这种情况下，AlphaZero 只能在规则已知的环境中使用，因此 MuZero 是一种经过扩展的算法，因此它也可以在马尔可夫决策过程（MDP）环境中使用，这是强化学习中通常假设的。
（下文在本文中区分环境时，将马尔可夫决策过程（MDP）环境称为雅达利环境，将围棋和将棋称为棋盘游戏环境。）

通过学习游戏规则本身，即使在游戏规则未知的 MDP 环境中，MuZero 也能进行学习。

参考

・MuZero 实施评论（用于 Breakout）| Mendako 从任何角度
・如何构建 MuZero AI
・MuZero：演练（第 1/3 部分）
・通过学习模型进行规划，掌握雅达利、围棋、国际象棋和将棋
・MuZero 纸
・论文伪代码：https://arxiv.org/src/1911.08265v2/anc/pseudocode.py

1. MDP环境的表示和学习

简单地说，MDP 是一个模型，其中下一个状态由当前状态和动作决定。 （详情见上篇文章）

MuZero 添加了两个新的神经网络来表示 MDP 环境。
我认为如果你先看整个网络会更容易理解，所以让我们先看看每个网络以及它是如何学习的。

1. 代表网络

表示网络负责将从环境中观察到的状态编码为在算法中使用的状态。
编码状态本身没有特殊意义。
（从机器学习的角度来看，它具有很强的特征提取意义）

在本文中，我们将从环境中观察到的状态写为 $o$，将编码后的状态写为 $s$，以便于区分。

2.动态网络

动态网络是一个从当前状态和动作预测下一个状态和立即奖励的网络。

*桌游环境没有即时奖励预测

3.预测网络

几乎与 AlphaZero 的 PV 网络相同。 （我只是为了匹配文章而改了名字）
唯一的区别是输入是表示网络编码状态$s$。

网络学习

至于学习，上面三个网络是同时训练的。
学习样本采集图如下。

符号是 u 代表即时奖励，z 代表评估值（AlphaZero 的输赢），a 代表行动，π 代表 MCTS 结果。 （如何计算 z 将在后面介绍）
我们需要学习下一个状态，因此我们将展开步骤保存为 1 个批次。

基于此的学习如下。

训练损失分为三种，每种计算的值之和就是最终损失。

$$
l_t(\theta) = \sum_{k=0}^K{ l^p(\pi_{t+k}, p_t^k)} + \sum_{k=0}^K{ l^v(z_{ t+k}, v_t^k) } + \sum_{k=1}^K{ l^r(u_{t+k}, r_t^k) } + c||\theta||^2
$$

*计算公式为MuZero 的下一篇方法论文用作参考。看来$l^r$ 是错误的，它从$k=1$ 开始。
不过好像L2正则化项写的不是很清楚，所以按照MuZero加了。

对于每个损失，$l^p$ 使用交叉熵计算损失，类似于 AlphaZero。
$l^v$ 和 $l^v$ ，但棋盘游戏环境是 MSE。
Atari 环境在将数据转换为分类分布后根据交叉熵计算损失。 （转换为分类分布在后面介绍）
转换为分类分布的原因是为了匹配三种损失类型的尺度。
最后一项将是 L2 正则化项。

2.MuZero MCTS

基本和AlphaZero一样，只是下降到叶子节点，返回棋盘的值，展开。

以下是 MCTS 的详细故事。

1. 结束状态
   与 AlphaZero 不同，MuZero 没有结束状态，因此它将永远搜索，直到到达叶节点。
   暂时来说，学习结束后的状态，似乎被学习为吸收状态。
   参考：吸收马尔可夫过程 | 英文维基百科

2. 传播价值
   正如我们在 AlphaZero 实施中所做的那样，Atari 环境会打折并传播奖励。
   （在棋盘游戏环境中按原样传播）

3. 奖励逆转
   我已经这样做了，但是在棋盘游戏环境中，下一步将轮到对手，因此奖励将被反转。
   （雅达利环境下无翻转）

4. 根状态下的狄利克雷噪声
   在不改变 AlphaZero 的情况下，添加 Dirichlet 噪声以鼓励探索。

5. 有效动作
   有效操作仅适用于可以从环境中获取的根状态。
   我们在树搜索期间不执行任何有效操作，让网络学习。

6. Q值归一化
   PUCT 中使用的 Q 值在 AlphaZero 中被假定为 [0,1]。
   由于 MuZero 可能取 0 到 1 以外的值，因此使用过去观察到的最小和最大 Q 值将其归一化为 0 到 1 之间的值。

   3. 其他故事

   1. MCTS 后的动作选择

   动作选择概率由以下公式表示，类似于 AlphaZero。

   $$ p(a) = \frac{N(a)^{1/{\tau}}}{\sum_b N(b)^{1/{\tau}} } $$

   $N$ 是您选择动作的次数，$\tau$ 是温度参数。
   棋盘游戏环境为前30步$\tau=1$（由概率根据步数决定），之后$\tau\rightarrow 1$（选择最大步数）。
   在 Atari 环境中，前 500k 步使用 $\tau=1$，前 750k 步使用 $\tau=0.5$，之后使用 $\tau=0.25$。

   2.状态值（z）的计算

   关于收集样本时出现的状态值 z 的计算，我们在 Atari 环境中使用 N 步折扣奖励。
   （与 Q-learning 中的 Multi-step Bootstrapping 相同）

   $$ z_t = u_{t+1} + \gamma u_{t+2} + ... + \gamma^{n-1} u_{t+m} + \gamma^n v_{t+n}$ $

   $u$ 是即时奖励，$\gamma$ 是贴现率，$v$ 是预测值。
   然而，该框架使用蒙特卡洛方法，将 N 扩展到剧集的结尾。

   3.重新调整价值和奖励

   对 Value 和 Reward 的目标值应用重新缩放函数并对值进行四舍五入。 （R2D2 方法）
   缩放功能如下。

   $$
   h(u) = 符号(u)(\sqrt{|u|+1}-1) + \epsilon u
   $$

   $sign$ 是符号函数，$\epsilon$ 是一个常数（使用 0.001）。

   4、奖励类别分配

   Atari 环境跨类别分配奖励以匹配损失规模。
   作为一种分类方法，十进制数对两个相邻的整数进行加权。
   例如，3.7 相对于两个相邻整数 (3,4) 的权重为 (0.3, 0.7)。

   

   这可以通过乘以权重来恢复。 （$3 乘以 0.3 + 4 乘以 0.7 = 3.7$）
   （它是十进制单热编码吗？）

   5. 样品优先

   用于训练的样本是根据它们的优先级选择的。
   这与 Rainbow 中使用的优先体验重放技术相同。
   详情请参阅链接。

   只有优先级计算不同，通过以下公式计算。

   $$p^i = |v_i - z_i|$$

   $v$ 是预测状态值，$z$ 是 n 步奖励。
   此外，Prioritized Experience Replay 的超参数似乎设置为 $\alpha=\beta=1$。

   6.梯度缩放

   扩展步骤之后的梯度对反向传播进行缩放，以便幅度对齐。
   缩放方法如下。 （来自伪代码）

   def scale_gradient(tensor, scale):
     """Scales the gradient for the backward pass."""
     return tensor * scale + tf.stop_gradient(tensor) * (1 - scale)
   

   应用要点如下。

   • 展开步骤后比例损失 1/K。 （K是部署步骤数）
   • 将动态网络中编码的下一个状态（梯度）缩放 1/2。

   以下是伪代码中仅学习部分的摘录。 （为清楚起见进行了编辑）

   def update_weights(optimizer, network, batch, weight_decay):
     loss = 0
   
     # バッチループ
     for init_state, actions, targets in batch:
       
       # --- 1st step
       # representation + prediction より P, V, hidden_state を出す
       # (reward は予測できないので多分なしだと思う…)
       # 最初はスケーリングなし（スケール値=1.0）
       value, reward, policy, hidden_state = network.initial_inference(init_state)
       predictions = [(1.0, value, reward, policy)]
   
       # --- 展開ステップ
       for action in actions:
         # dynamics + prediction より P, V, reward, next_hidden_state を出す
         value, reward, policy, hidden_state = network.recurrent_inference(hidden_state, action)
   
         # スケール値は 1/K
         predictions.append((1.0 / len(actions), value, reward, policy))
   
         # hidden_state は 0.5 でスケールする
         hidden_state = scale_gradient(hidden_state, 0.5)
       
       # --- 実際に勾配を計算
       for prediction, target in zip(predictions, targets):
         gradient_scale, value, reward, policy = prediction
         target_value, target_reward, target_policy = target
         
         # 各loss、scalar_loss は MSE(ボードゲーム) or cross_entropy(Atari)
         value_loss = scalar_loss(value, target_value)
         reward_loss = scalar_loss(reward, target_reward)
         policy_loss = cross_entropy(policy, target_policy)
         l = (value_loss + reward_loss + policy_loss)
       
         # backwardは各スケール値で
         loss += scale_gradient(l, gradient_scale)
     
     # L2正則化項を追加
     for weights in network.get_weights():
       loss += weight_decay * tf.nn.l2_loss(weights)
   
     optimizer.minimize(loss)
   
   

   8.编码后的缩放状态s

   编码后的状态 s 从 0 归一化为 1，以匹配 one-hot 编码动作的规模。

   9.重新分析

   用于训练的样本与使用旧参数训练的样本混合。
   因此，Reanalyze 是一种在使用最新的网络进行学习之前再次执行 MCTS，更新训练数据中的 $\pi$ 和 $z$，然后进行学习的方法。
   这允许重复使用过去的样本，从而提高样本效率。

   然而，由于框架上实现的复杂性以及需要在 Trainer 端执行 MCTS，实现被推迟。

   10. 学习率

   虽然没有写在论文里，但看伪代码，似乎和AlphaZero有一点不同。
   在 AlphaZero 中，使用固定值逐步改变（围棋中 0、300,000 和 500,000 步后为 0.02、0.002、0.0002），但 MuZero指数衰减（维基百科）看来学习率的调度是合理的。

   $$
   \alpha = \alpha_{init} \lambda^{T/ \alpha_t}
   $$

   $\alpha$ 是学习率，$\alpha_{init}$ 是初始学习率，$\lambda$ 是衰减率，$T$ 是总步数，$\alpha_t$ 是总步数腐烂的步骤。
   既然符号和术语合适，我也会用代码写出来

   lr_init = 0.05          # 初期学習率
   lr_decay_rate = 0.1     # 減衰率
   lr_decay_steps = 350e3  # 減衰総ステップ数
   T = 総ステップ数
   
   lr = lr_init * lr_decay_rate ** (T / lr_decay_steps)
   

   4.网络结构

   基本原理与 AlphaZero 相同。
   可以看到，内容多为 ResNet，所以在具有类图像特征的环境中似乎比较强。

   • 表情网络（棋盘游戏环境）

   

   • 表达网络（Atari 环境）

   它似乎从 96x96 缩小到 6x6 大小。 （剩余块相同）

   

   • 动态网络

   只写和AlphaZero一样，所以我没有信心。
   关于Reward的输出，我一开始是加在ResBlock下的，但是loss发散了，学习也不太顺利。
   因此，它们直接连接，如图所示。
   这样一来，下一个状态的预测就会有和表达网络一样的结构，而Reward的预测也会有和Policy/Value一样的结构，所以我觉得大概是对的。

   

   *在棋盘游戏环境中，不会有奖励输出。

   • 预测网络

   

   * 在棋盘游戏环境中，Value 端的输出将为 tanh。

   执行

   仅摘录相关部分。
   完整的代码是这里是。

   配置

   这是一个超参数。
   值是 Atari 环境中的参数。

   @dataclass
   class Config(DiscreteActionConfig):
       num_simulations: int = 50
       batch_size: int = 1024
       discount: float = 0.997
       
       # 学習率
       lr_init: float = 0.05
       lr_decay_rate: float = 0.1
       lr_decay_steps: int = 350_000
   
       # カテゴリ化する範囲
       v_min: int = -300
       v_max: int = 300
   
       # policyの温度パラメータのリスト
       policy_tau_schedule: List[dict] = [
           {"step": 0, "tau": 1.0},
           {"step": 500_000, "tau": 0.5},
           {"step": 750_000, "tau": 0.25},
       ]
   
       # td_steps: int = 10   # multisteps
       unroll_steps: int = 5  # unroll_steps
   
       # Root prior exploration noise.
       root_dirichlet_alpha: float = 0.3
       root_exploration_fraction: float = 0.25
   
       # PUCT
       c_base: float = 19652
       c_init: float = 1.25
   
       # Priority Experience Replay
       capacity: int = 1_000_000
       memory_name: str = "ProportionalMemory"
       memory_warmup_size: int = 1000
       memory_alpha: float = 1.0
       memory_beta_initial: float = 1.0
       memory_beta_steps: int = 0
   
   

   网络

   我会写输入雅达利环境的案例（处理后的96,96灰度图）。
   省略了每一层的详细参数。 （详见 git 代码）

   # 残差ブロック(AlphaZeroと同じ)
   class _ResidualBlock(keras.Model):
       def __init__(self, filters):
           super().__init__()
   
           self.conv1 = kl.Conv2D(filters=filters, kernel_size=(3,3))
           self.bn1 = kl.BatchNormalization()
           self.relu1 = kl.LeakyReLU()
           self.conv2 = kl.Conv2D(filters=filters, kernel_size=(3,3))
           self.bn2 = kl.BatchNormalization()
           self.relu2 = kl.LeakyReLU()
   
       def call(self, x):
           x1 = self.conv1(x)
           x1 = self.bn1(x1)
           x1 = self.relu1(x1)
           x1 = self.conv2(x1)
           x1 = self.bn2(x1)
           x = x + x1
           x = self.relu2(x)
           return x
   
   # --- 表現ネットワーク
   class _RepresentationNetwork(keras.Model):
       def __init__(self, config: Config):
           super().__init__()
   
           input_shape = (96, 96)  # atari image
           in_layer = c = kl.Input(shape=input_shape)
           c = kl.Reshape(input_shape + (1,))(c)  # (w, h) -> (w, h, 1)
   
           # ダウンサンプリング
           c = kl.Conv2D(128, kernel_size=3, strides=2, activation="relu")
           c = _ResidualBlock(128)(c)
           c = _ResidualBlock(128)(c)
           c = kl.Conv2D(256, kernel_size=3, strides=2, activation="relu")
           c = _ResidualBlock(256)(c)
           c = _ResidualBlock(256)(c)
           c = _ResidualBlock(256)(c)
           c = kl.AveragePooling2D(pool_size=3, strides=2)(c)
           c = _ResidualBlock(256)(c)
           c = _ResidualBlock(256)(c)
           c = _ResidualBlock(256)(c)
           c = kl.AveragePooling2D(pool_size=3, strides=2)(c)
           self.model = keras.Model(in_state, c)
   
       def call(self, state):
           x = self.model(state)
   
           # 隠れ状態はアクションとスケールを合わせるため0-1で正規化(一応batch毎)
           batch, h, w, d = x.shape
           s_min = tf.reduce_min(tf.reshape(x, (batch, -1)), axis=1, keepdims=True)
           s_max = tf.reduce_max(tf.reshape(x, (batch, -1)), axis=1, keepdims=True)
           s_min = s_min * tf.ones((batch, h * w * d), dtype=tf.float32)
           s_max = s_max * tf.ones((batch, h * w * d), dtype=tf.float32)
           s_min = tf.reshape(s_min, (batch, h, w, d))
           s_max = tf.reshape(s_max, (batch, h, w, d))
           epsilon = 1e-4  # div0 回避
           x = (x - s_min + epsilon) / tf.maximum((s_max - s_min), 2 * epsilon)
   
           return x
   
   # --- ダイナミクスネットワーク
   class _DynamicsNetwork(keras.Model):
       def __init__(self, config: Config, input_shape):
           super().__init__()
           self.action_num = config.action_num
           v_num = config.v_max - config.v_min + 1
   
           # input_shapeは表現ネットワーク後のshapeなので(6, 6, 256)
           h, w, ch = input_shape
   
           # hidden_state + action_space
           in_state = c = kl.Input(shape=(h, w, ch + self.action_num))
   
           # AlphaZeroブロック(詳細はAlphaZeroの記事を参照)
           c1 = AlphaZeroImageBlock(n_blocks=15)(c)
   
           # reward
           c2 = kl.Conv2D(1, kernel_size=1)(c)
           c2 = kl.BatchNormalization()(c2)
           c2 = kl.LeakyReLU()(c2)
           c2 = kl.Flatten()(c2)
           reward = kl.Dense(v_num, activation="softmax")(c2)
   
           # 出力は hidden_state, reward(category)
           self.model = keras.Model(in_state, [c1, reward])
   
       def call(self, hidden_state, action):
           batch_size, h, w, _ = hidden_state.shape
   
           # --- actionをイメージ化する
           action_image = tf.one_hot(action, self.action_num)  # (batch, action)
           action_image = tf.repeat(action_image, repeats=h * w, axis=1)  # (batch, action * h * w)
           action_image = tf.reshape(action_image, (batch_size, self.action_num, h, w))  # (batch, action, h, w)
           action_image = tf.transpose(action_image, perm=[0, 2, 3, 1])  # (batch, h, w, action)
   
           # --- hidden_stateの最後にアクション層を追加
           in_state = tf.concat([hidden_state, action_image], axis=3)
           x, reward_category = self.model(in_state)
   
           # 隠れ状態はアクションとスケールを合わせるため0-1で正規化(一応batch毎)
           batch, h, w, d = x.shape
           s_min = tf.reduce_min(tf.reshape(x, (batch, -1)), axis=1, keepdims=True)
           s_max = tf.reduce_max(tf.reshape(x, (batch, -1)), axis=1, keepdims=True)
           s_min = s_min * tf.ones((batch, h * w * d), dtype=tf.float32)
           s_max = s_max * tf.ones((batch, h * w * d), dtype=tf.float32)
           s_min = tf.reshape(s_min, (batch, h, w, d))
           s_max = tf.reshape(s_max, (batch, h, w, d))
           epsilon = 1e-4  # div0 回避
           x = (x - s_min + epsilon) / tf.maximum((s_max - s_min), 2 * epsilon)
   
           return x, reward_category
   
   
   # --- 予測ネットワーク
   class _PredictionNetwork(keras.Model):
       def __init__(self, config: Config, input_shape):
           super().__init__()
   
           v_num = config.v_max - config.v_min + 1
   
           # input_shapeは表現ネットワーク後のshapeなので(6, 6, 256)
           in_layer = c = kl.Input(shape=input_shape)
   
           # --- policy
           c1 = kl.Conv2D(2, kernel_size=(1, 1))(c)
           c1 = kl.BatchNormalization()(c1)
           c1 = kl.LeakyReLU()(c1)
           c1 = kl.Flatten()(c1)
           policy = kl.Dense(config.action_num, activation="softmax")(c1)
   
           # --- value
           c2 = kl.Conv2D(1, kernel_size=(1, 1))(c)
           c2 = kl.BatchNormalization()(c2)
           c2 = kl.LeakyReLU()(c2)
           c2 = kl.Flatten()(c2)
           value = kl.Dense(v_num, activation="softmax")(c2)
   
           self.model = keras.Model(in_layer, [policy, value])
   
       def call(self, state):
           return self.model(state)
   
   

   参数

   代码省略，因为它与 AlphaZero 几乎相同。
   此外，它还存储了 MCTS 期间用于 Q 值归一化的最大和最小 Q 值。

   远程内存

   我只是使用我在 Rainbow 的 Prioritized Experience Replay 中实现的内容。
   代码省略。

   分类功能

   def float_category_encode(val: float, v_min: int, v_max: int) -> List[float]:
       category = [0.0 for _ in range(v_max - v_min + 1)]
       low_int = math.floor(val)
       high_int = low_int + 1
       weight = val - low_int
       category[int(low_int - v_min)] = 1 - weight
       category[int(high_int - v_min)] = weight
       return category
   
   def float_category_decode(category: List[float], v_min: int) -> float:
       n = 0
       for i, w in enumerate(category):
           n += (i + v_min) * w
       return n
   
   

   工作人员

   流程与 AlphaZero 相同，但细节有所改变。

   class Worker(DiscreteActionWorker):
       def __init__(self, *args):
           super().__init__(*args)
   
           # 温度パラメータを扱いやすいように変更
           self.policy_tau_schedule = {}
           for tau_list in self.config.policy_tau_schedule:
               self.policy_tau_schedule[tau_list["step"]] = tau_list["tau"]
           self.policy_tau = self.policy_tau_schedule[0]
   
           self.total_step = 0
   
       def call_on_reset(self, state: np.ndarray, invalid_actions: List[int]) -> None:
           self.step = 0
           self.history = []
   
           self.N = {}  # 訪問回数(s,a)
           self.W = {}  # 累計報酬(s,a)
           self.Q = {}
   
       def _init_state(self, state_str):
           if state_str not in self.N:
               self.N[state_str] = [0 for _ in range(self.config.action_num)]
               self.W[state_str] = [0 for _ in range(self.config.action_num)]
               self.Q[state_str] = [0 for _ in range(self.config.action_num)]
   
       def call_policy(self, state: np.ndarray, invalid_actions: List[int]) -> int:
           # --- 表現ネットワークから初期状態を取得
           s0 = self.parameter.representation_network(state[np.newaxis, ...])
           s0_str = s0.ref()
           
           # --- シミュレーション
           for _ in range(self.config.num_simulations):
               self._simulation(s0, s0_str, invalid_actions)
   
           # --- 確率に比例したアクションを選択
           if not self.training:
               self.policy_tau = 0  # 評価時は決定的に
           if self.policy_tau == 0:
               # 温度パラメータ0は決定的
               counts = np.asarray(self.N[s0_str])
               action = random.choice(np.where(counts == counts.max())[0])
           else:
               # 確率的
               step_policy = np.array(
                   [self.N[s0_str][a] ** (1 / policy_tau) for a in range(self.config.action_num)]
               )
               step_policy /= step_policy.sum()
               action = random_choice_by_probs(step_policy)
   
           # 温度パラメータのschedule check
           if self.total_step in self.policy_tau_schedule:
               self.policy_tau = self.policy_tau_schedule[self.total_step]
   
           # 学習用のpolicyはtau=1
           N = sum(self.N[self.s0_str])
           self.step_policy = [self.N[self.s0_str][a] / N for a in range(self.config.action_num)]
   
           # サンプル用に保存
           self.state = state
           self.action = int(action)
           self.state_v = self.parameter.V[s0_str]
           return self.action
   
       # --- シミュレーション（1step,再帰,次の報酬を返す）
       def _simulation(self, state, state_str, invalid_actions, depth: int = 0):
           if depth >= 99999:  # for safety
               return 0
   
           # PVを予測
           self._init_state(state_str)
           self.parameter.pred_PV(state, state_str)
   
           # actionを選択
           puct_list = self._calc_puct(state_str, invalid_actions, depth == 0)
           action = random.choice(np.where(puct_list == np.max(puct_list))[0])
   
           # 次の状態を取得
           n_state, reward_category = self.parameter.dynamics_network(state, [action])
           n_state_str = n_state.ref()
           reward = float_category_decode(reward_category.numpy()[0], self.config.v_min)
           enemy_turn = self.config.env_player_num > 1  # 2player以上は相手番と決め打ち
   
           if self.N[state_str][action] == 0:
               # leaf node ならロールアウト
               self.parameter.pred_PV(n_state, n_state_str)
               n_reward = self.parameter.V[n_state_str]
           else:
               # 子ノードに降りる(展開)
               n_reward = self._simulation(n_state, n_state_str, [], depth + 1)
   
           # 次が相手のターンなら、報酬は最小になってほしいので-をかける
           if enemy_turn:
               n_reward = -n_reward
   
           # 割引報酬
           reward = reward + self.config.discount * n_reward
   
           self.N[state_str][action] += 1
           self.W[state_str][action] += reward
           self.Q[state_str][action] = self.W[state_str][action] / self.N[state_str][action]
   
           self.parameter.q_min = min(self.parameter.q_min, self.Q[state_str][action])
           self.parameter.q_max = max(self.parameter.q_max, self.Q[state_str][action])
   
           return reward
   
       def _calc_puct(self, state_str, invalid_actions, is_root):
           AlphaZeroとほぼ同じなので省略。
           追加要素は、Q値の正規化(MinMax)だけです。
   
       def call_on_step(
           self,
           next_state: np.ndarray,
           reward: float,
           done: bool,
           next_invalid_actions: List[int],
       ):
           self.step += 1
           self.total_step += 1
   
           if not self.training:
               return {}
   
           self.history.append(
               {
                   "state": self.state,
                   "action": self.action,
                   "policy": self.step_policy,
                   "reward": reward,
                   "state_v": self.state_v,
               }
           )
   
           # 終了時に割引報酬を計算してサンプルを送る
           if done:
               zero_category = float_category_encode(0, self.config.v_min, self.config.v_max)
   
               # calc MC reward
               reward = 0
               for h in reversed(self.history):
                   reward = h["reward"] + self.config.discount * reward
                   h["discount_reward"] = reward
   
               # batch create
               for idx in range(len(self.history)):
   
                   # --- policies
                   policies = [
                       [1 / self.config.action_num] * self.config.action_num for _ in range(self.config.unroll_steps + 1)
                   ]
                   for i in range(self.config.unroll_steps + 1):
                       if idx + i >= len(self.history):
                           break
                       policies[i] = self.history[idx + i]["policy"]
   
                   # --- values
                   values = [zero_category for _ in range(self.config.unroll_steps + 1)]
                   priority = 0
                   for i in range(self.config.unroll_steps + 1):
                       if idx + i >= len(self.history):
                           break
                       v = self.history[idx + i]["discount_reward"]
                       v = rescaling(v)
                       priority += v - self.history[idx + i]["state_v"]
                       values[i] = float_category_encode(v, self.config.v_min, self.config.v_max)
                   priority /= self.config.unroll_steps + 1
   
                   # --- actions
                   actions = [random.randint(0, self.config.action_num - 1) for _ in range(self.config.unroll_steps)]
                   for i in range(self.config.unroll_steps):
                       if idx + i >= len(self.history):
                           break
                       actions[i] = self.history[idx + i]["action"]
   
                   # --- rewards
                   rewards = [zero_category for _ in range(self.config.unroll_steps)]
                   for i in range(self.config.unroll_steps):
                       if idx + i >= len(self.history):
                           break
                       r = self.history[idx + i]["reward"]
                       r = rescaling(r)
                       rewards[i] = float_category_encode(r, self.config.v_min, self.config.v_max)
   
                   self.remote_memory.add(
                       {
                           "state": self.history[idx]["state"],
                           "actions": actions,
                           "policies": policies,
                           "values": values,
                           "rewards": rewards,
                       },
                       priority,
                   )
           return {}
   
   

   培训师

   这是学习部分。
   同时训练三个网络。

   def _scale_gradient(tensor, scale):
       """ Scales the gradient for the backward pass. """
       return tensor * scale + tf.stop_gradient(tensor) * (1 - scale)
   
   class Trainer(RLTrainer):
       def __init__(self, *args):
           super().__init__(*args)
   
           self.optimizer = keras.optimizers.Adam()
           # バッチ毎に出力
           self.cross_entropy_loss = keras.losses.CategoricalCrossentropy(axis=1, reduction=keras.losses.Reduction.NONE)
           self.train_count = 0
   
       def train(self):
           if self.remote_memory.length() < self.config.memory_warmup_size:
               return {}
           indices, batchs, weights = self.remote_memory.sample(self.train_count, self.config.batch_size)
   
           # (batch, dict, val) -> (batch, val)
           states = batchsよりデータ変換
   
           # (batch, dict, steps, val) -> (steps, batch, val)
           actions_list = batchsよりデータ変換
           policies_list = batchsよりデータ変換
           values_list = batchsよりデータ変換
           rewards_list = batchsよりデータ変換
   
           with tf.GradientTape() as tape:
               # --- 1st step
               hidden_states = self.parameter.representation_network(states)
               p_pred, v_pred = self.parameter.prediction_network(hidden_states)
   
               # loss
               policy_loss = _scale_gradient(self.cross_entropy_loss(policies_list[0], p_pred), 1.0)
               value_loss = _scale_gradient(self.cross_entropy_loss(values_list[0], v_pred), 1.0)
               reward_loss = tf.constant([0] * self.config.batch_size, dtype=tf.float32)
   
               # --- unroll steps
               gradient_scale = 1 / self.config.unroll_steps
               for t in range(self.config.unroll_steps):
                   # pred
                   n_hidden_states, p_rewards = self.parameter.dynamics_network(hidden_states, actions_list[t])
                   p_pred, v_pred = self.parameter.prediction_network(hidden_states)
   
                   # loss
                   value_loss += _scale_gradient(self.cross_entropy_loss(values_list[t + 1], v_pred), gradient_scale)
                   policy_loss += _scale_gradient(self.cross_entropy_loss(policies_list[t + 1], p_pred), gradient_scale)
                   reward_loss += _scale_gradient(self.cross_entropy_loss(rewards_list[t], p_rewards), gradient_scale)
   
                   hidden_states = _scale_gradient(n_hidden_states, 0.5)
   
               loss = tf.reduce_mean((value_loss + policy_loss + reward_loss) * weights)
   
               # 各ネットワークの正則化項を加える
               loss += tf.reduce_sum(self.parameter.representation_network.losses)
               loss += tf.reduce_sum(self.parameter.prediction_network.losses)
               loss += tf.reduce_sum(self.parameter.dynamics_network.losses)
   
           priorities = value_loss.numpy()
   
           # lr
           lr = self.config.lr_init * self.config.lr_decay_rate ** (self.train_count / self.config.lr_decay_steps)
           self.optimizer.learning_rate = lr
   
           variables = [
               self.parameter.representation_network.trainable_variables,
               self.parameter.prediction_network.trainable_variables,
               self.parameter.dynamics_network.trainable_variables,
           ]
           grads = tape.gradient(loss, variables)
           for i in range(len(variables)):
               self.optimizer.apply_gradients(zip(grads[i], variables[i]))
   
           self.train_count += 1
   
           # memory update
           self.remote_memory.update(indices, batchs, priorities)
   
           # 学習したらキャッシュは削除
           self.parameter.reset_cache()
           return {}
   
   

   学习

   虽然是在学习，但是由于缺乏spec，一直没有得到好的结果。
   （做完文章，在这里停了很久，所以放弃了，优先发帖）
   在 Mendako 的博客中，感觉他在 GCP 的 24-vCPU/128GB RAM/GPU T4 上学习了 48 小时后终于开始学习了。
   我的印象是，包括环境在内的学习是相当困难的......

   我已经确认我可以在一个非常简单的环境中学习。
   我用来检查的代码如下。

   import numpy as np
   import srl
   from srl.envs import grid
   from srl.rl.models.alphazero_image_block import AlphaZeroImageBlock
   from srl.runner import sequence
   
   
   def main():
       rl_config = srl.rl.muzero.Config(
           num_simulations=10,
           discount=0.9,
           batch_size=16,
           memory_warmup_size=200,
           memory_name="ReplayMemory",
           lr_init=0.002,
           lr_decay_steps=10_000,
           v_min=-2,
           v_max=2,
           unroll_steps=1,
           input_image_block=AlphaZeroImageBlock,
           input_image_block_kwargs={"n_blocks": 1, "filters": 16},
           dynamics_blocks=1,
           enable_rescale=False,
           weight_decay=0,
       )
       rl_config.processors = [grid.LayerProcessor()]
       env_config = srl.envs.Config("EasyGrid")
       config = sequence.Config(env_config, rl_config)
   
       config.model_summary()
   
       # --- 学習ループ
       parameter, memory, history = sequence.train(config, max_episodes=100)
       history.plot()
       history.plot_info("train", "loss")
   
       # --- evaluate
       rewards = sequence.evaluate(config, parameter, max_episodes=10, print_progress=True)
       print("mean", np.mean(rewards))
   
       # --- rendering
       rewards, render = sequence.render(config, parameter, enable_animation=True, print_progress=True)
       render.create_anime().save("_qiita1.gif")
   
   
   if __name__ == "__main__":
       main()
   
   

   • 损失

   

   • 最终结果

   

   细节

   ### env: EasyGrid, max episodes: 1, max steps: -1, timeout:  -1.00s
   ### 0, action 2, rewards [0.], next 0
   env   None
   work0 None
   ......
   .   G.
   . . X.
   .P   .
   ......
   
   
   V_net: 0.09338
    ←  :  30.0% (      3)(N),   0.08279(Q),   0.82895(PUCT),   0.29893(P),   0.09176(V),  -0.00000(reward)
    ↓  :  20.0% (      2)(N),   0.08392(Q),   0.80924(PUCT),   0.20884(P),   0.06570(V),  -0.00000(reward)
   *→  :  30.0% (      3)(N),   0.08174(Q),   0.73965(PUCT),   0.20910(P),   0.11388(V),  -0.00034(reward)
    ↑  :  20.0% (      2)(N),   0.08661(Q),   0.90842(PUCT),   0.28312(P),   0.10653(V),  -0.00000(reward)
   
   ### 1, action 3, rewards [0.], next 0
   env   {}
   work0 {}
   ......
   .   G.
   . . X.
   . P  .
   ......
   
   
   V_net: 0.12717
    ←  :  20.0% (      2)(N),   0.10917(Q),   0.91028(PUCT),   0.27648(P),   0.10678(V),  -0.00000(reward)
    ↓  :  30.0% (      3)(N),   0.10386(Q),   0.74972(PUCT),   0.20876(P),   0.13161(V),  -0.00000(reward)
    →  :  20.0% (      2)(N),   0.12856(Q),   0.87763(PUCT),   0.24479(P),   0.18315(V),  -0.00088(reward)
   *↑  :  30.0% (      3)(N),   0.13049(Q),   0.82274(PUCT),   0.26996(P),   0.16563(V),  -0.00001(reward)
   
   ### 2, action 0, rewards [0.], next 0
   env   {}
   work0 {}
   ......
   .   G.
   . . X.
   . P  .
   ......
   
   
   V_net: 0.12717
   *←  :  30.0% (      3)(N),   0.10285(Q),   0.81619(PUCT),   0.27648(P),   0.10678(V),  -0.00000(reward)
    ↓  :  30.0% (      3)(N),   0.11080(Q),   0.75298(PUCT),   0.20876(P),   0.13161(V),  -0.00000(reward)
    →  :  20.0% (      2)(N),   0.14924(Q),   0.88735(PUCT),   0.24479(P),   0.18315(V),  -0.00088(reward)
    ↑  :  20.0% (      2)(N),   0.14727(Q),   0.91960(PUCT),   0.26996(P),   0.16563(V),  -0.00001(reward)
   
   ### 3, action 3, rewards [0.], next 0
   env   {}
   work0 {}
   ......
   .   G.
   . . X.
   .P   .
   ......
   
   
   V_net: 0.09338
    ←  :  30.0% (      3)(N),   0.09546(Q),   0.83491(PUCT),   0.29893(P),   0.09176(V),  -0.00000(reward)
    ↓  :  20.0% (      2)(N),   0.08227(Q),   0.80846(PUCT),   0.20884(P),   0.06570(V),  -0.00000(reward)
    →  :  10.0% (      1)(N),   0.10214(Q),   0.95598(PUCT),   0.20910(P),   0.11388(V),  -0.00034(reward)
   *↑  :  40.0% (      4)(N),   0.10538(Q),   0.76796(PUCT),   0.28312(P),   0.10653(V),  -0.00000(reward)
   
   ### 4, action 3, rewards [0.], next 0
   env   {}
   work0 {}
   ......
   .   G.
   .P. X.
   .    .
   ......
   
   
   V_net: 0.10793
    ←  :  30.0% (      3)(N),   0.08260(Q),   0.83014(PUCT),   0.30023(P),   0.08683(V),  -0.00000(reward)
    ↓  :  20.0% (      2)(N),   0.07308(Q),   0.79139(PUCT),   0.19917(P),   0.04500(V),  -0.00000(reward)
    →  :  20.0% (      2)(N),   0.08945(Q),   0.81848(PUCT),   0.21388(P),   0.09458(V),  -0.00004(reward)
   *↑  :  30.0% (      3)(N),   0.10137(Q),   0.82562(PUCT),   0.28672(P),   0.11268(V),  -0.00000(reward)
   
   ### 5, action 0, rewards [0.], next 0
   env   {}
   work0 {}
   ......
   .P  G.
   . . X.
   .    .
   ......
   
   
   V_net: 0.12365
   *←  :  30.0% (      3)(N),   0.10337(Q),   0.82428(PUCT),   0.28441(P),   0.10541(V),  -0.00000(reward)
    ↓  :  20.0% (      2)(N),   0.07933(Q),   0.81485(PUCT),   0.21473(P),   0.07611(V),  -0.00000(reward)
    →  :  20.0% (      2)(N),   0.17113(Q),   0.84337(PUCT),   0.20362(P),   0.26590(V),  -0.00030(reward)
    ↑  :  30.0% (      3)(N),   0.12231(Q),   0.84586(PUCT),   0.29723(P),   0.13563(V),  -0.00000(reward)
   
   ### 6, action 2, rewards [0.], next 0
   env   {}
   work0 {}
   ......
   .P  G.
   . . X.
   .    .
   ......
   
   
   V_net: 0.12365
    ←  :  30.0% (      3)(N),   0.10337(Q),   0.82428(PUCT),   0.28441(P),   0.10541(V),  -0.00000(reward)
    ↓  :  10.0% (      1)(N),   0.06850(Q),   0.95128(PUCT),   0.21473(P),   0.07611(V),  -0.00000(reward)
   *→  :  30.0% (      3)(N),   0.20643(Q),   0.79287(PUCT),   0.20362(P),   0.26590(V),  -0.00030(reward)
    ↑  :  30.0% (      3)(N),   0.13197(Q),   0.85040(PUCT),   0.29723(P),   0.13563(V),  -0.00000(reward)
   
   ### 7, action 2, rewards [0.], next 0
   env   {}
   work0 {}
   ......
   . P G.
   . . X.
   .    .
   ......
   
   
   V_net: 0.29224
    ←  :  20.0% (      2)(N),   0.11201(Q),   0.93680(PUCT),   0.29559(P),   0.14662(V),  -0.00000(reward)
    ↓  :  10.0% (      1)(N),   0.23912(Q),   0.98168(PUCT),   0.18953(P),   0.26568(V),  -0.00000(reward)
   *→  :  40.0% (      4)(N),   0.27038(Q),   0.80180(PUCT),   0.22780(P),   0.67761(V),  -0.00011(reward)
    ↑  :  30.0% (      3)(N),   0.17119(Q),   0.85882(PUCT),   0.28709(P),   0.28577(V),  -0.00000(reward)
   
   ### 8, action 2, rewards [0.], next 0
   env   {}
   work0 {}
   ......
   .  PG.
   . . X.
   .    .
   ......
   
   
   V_net: 0.53148
    ←  :   0.0% (      0)(N),   0.00000(Q),   1.28810(PUCT),   0.20068(P),   0.10117(V),   0.00027(reward)
    ↓  :  10.0% (      1)(N),  -0.08549(Q),   0.71013(PUCT),   0.12940(P),  -0.09500(V),   0.00001(reward)
   *→  :  70.0% (      7)(N),   1.05088(Q),   1.22495(PUCT),   0.47794(P),   0.02293(V),   0.98078(reward)
    ↑  :  20.0% (      2)(N),   0.28879(Q),   0.88336(PUCT),   0.19198(P),   0.40778(V),   0.05570(reward)
   
   ### 9, action 2, rewards [1.], done(env), next 0
   env   {}
   work0 {}
   ......
   .   P.
   . . X.
   .    .
   ......
   
   

   奖励很好学。
   由于损失还不稳定，看来学习还是不够的。

   MuZero 的问题

   正如论文中提到的，到下一个状态的转换在动态网络中是决定性的。
   这并不准确，因为马尔可夫决策过程是概率性的。
   解决这个问题的方法是下一个技术，Stochastic MuZero。
   纸：使用学习模型在随机环境中进行规划

   我会在不久的将来写一篇文章。

原创声明：本文系作者授权爱码网发表，未经许可，不得转载;

原文地址：https://www.likecs.com/show-308623143.html