<p>Author here.</p>

<blockquote>
  <p>What does the reward from the initial_inference represent?</p>
</blockquote>

<p>The initial inference "predicts" the last observed reward. This isn't actually used for anything, but makes our code simpler: The prediction head can simply always predict the immediately preceding reward. For the dynamics network, this would be the reward observed after applying the action that's given as an input to the dynamics network.</p>

<p>At the beginning of the game there is no last observed reward, so we just set it to 0.</p>

<p>The reward target computation in the pseudocode was indeed misaligned; I've just uploaded a new version to arXiv. </p>

<p>Where it used to say</p>

<pre class="lang-py prettyprint-override"><code>      if current_index &lt; len(self.root_values):
        targets.append((value, self.rewards[current_index],
                        self.child_visits[current_index]))
      else:
        # States past the end of games are treated as absorbing states.
        targets.append((0, 0, []))
</code></pre>

<p>It should be:</p>

<pre class="lang-py prettyprint-override"><code>      # For simplicity the network always predicts the most recently received
      # reward, even for the initial representation network where we already
      # know this reward.
      if current_index &gt; 0 and current_index &lt;= len(self.rewards):
        last_reward = self.rewards[current_index - 1]
      else:
        last_reward = 0

      if current_index &lt; len(self.root_values):
        targets.append((value, last_reward, self.child_visits[current_index]))
      else:
        # States past the end of games are treated as absorbing states.
        targets.append((0, last_reward, []))
</code></pre>

<p>Hope that helps!</p>
Is the reward value in MuZero's pseudocode misaligned?
<p><a href="https://arxiv.org/abs/1911.08265" rel="nofollow noreferrer">MuZero</a>, a deep reinforcement learning technique, was just released, and I've been trying to implement it by looking at its <a href="https://arxiv.org/src/1911.08265v1/anc/pseudocode.py" rel="nofollow noreferrer">pseudocode</a> and this <a href="https://medium.com/applied-data-science/how-to-build-your-own-muzero-in-python-f77d5718061a" rel="nofollow noreferrer">helpful tutorial</a> on Medium.</p>

<p>However, there's something confusing me about how rewards are handled during training in the pseudocode, and it would be great if someone could verify that I'm reading the code correctly, and if I am, explain why this training algorithm works.</p>

<p>Here's the training function (from the <a href="https://arxiv.org/src/1911.08265v1/anc/pseudocode.py" rel="nofollow noreferrer">pseudocode</a>):</p>

<pre class="lang-py prettyprint-override"><code>def update_weights(optimizer: tf.train.Optimizer, network: Network, batch,
                   weight_decay: float):
  loss = 0
  for image, actions, targets in batch:
    # Initial step, from the real observation.
    value, reward, policy_logits, hidden_state = network.initial_inference(
        image)
    predictions = [(1.0, value, reward, policy_logits)]

    # Recurrent steps, from action and previous hidden state.
    for action in actions:
      value, reward, policy_logits, hidden_state = network.recurrent_inference(
          hidden_state, action)
      predictions.append((1.0 / len(actions), value, reward, policy_logits))

      hidden_state = tf.scale_gradient(hidden_state, 0.5)

    for prediction, target in zip(predictions, targets):
      gradient_scale, value, reward, policy_logits = prediction
      target_value, target_reward, target_policy = target

      l = (
          scalar_loss(value, target_value) +
          scalar_loss(reward, target_reward) +
          tf.nn.softmax_cross_entropy_with_logits(
              logits=policy_logits, labels=target_policy))

      loss += tf.scale_gradient(l, gradient_scale)

  for weights in network.get_weights():
    loss += weight_decay * tf.nn.l2_loss(weights)

  optimizer.minimize(loss)
</code></pre>

<p>I'm interested in the <code>reward</code> in the loss, specifically. Note that the loss gets all of its values from the <code>predictions</code>. The first <code>reward</code> added to <code>predictions</code> is from the <code>network.initial_inference</code> function. Afterwards, there are <code>len(actions)</code> more rewards added to <code>predictions</code>, all of which come from the <code>network.recurrent_inference</code> function.</p>

<p>Based on the tutorial <code>initial_inference</code> and <code>recurrent_inference</code> functions are built out of 3 different functions:</p>

<ol>
<li><em>Prediction</em> Input: internal game state. Output: policy, value (predicted sum of best possible future rewards)</li>
<li><em>Dynamics</em> Input: internal state of a game, action. Output: reward from taking that action, new internal state of the game.</li>
<li><em>Representation</em> Input: external state of a game. Output: internal state of the game</li>
</ol>

<p>The <code>initial_inference</code> function takes in an external game state, uses the <code>representation</code> function to turn it into an internal state, and then uses the <code>prediction</code> function on that internal game state. It outputs the internal state, the policy, and value.</p>

<p>The <code>recurrent_inference</code> function takes in an internal game state and an action. It uses the <code>dynamics</code> function to get a new internal game state and reward from the old game state and action. It then applies the <code>prediction</code> function to the new internal game state to get a policy and value of that new internal state. Thus, the final output is a new internal state, a reward, a policy, and a value.</p>

<p>However, in the pseudocode, the <code>initial_inference</code> function <em>also returns a reward</em>.</p>

<p>My main problem: <strong>What does that reward represent?</strong></p>

<p>In <a href="https://medium.com/applied-data-science/how-to-build-your-own-muzero-in-python-f77d5718061a" rel="nofollow noreferrer">the tutorial</a>, they just implicitly assume that the reward from the <code>initial_inference</code> function is 0. (See <a href="https://miro.medium.com/max/1514/1*GA72IpY7ZciGshmVvtl8kQ.png" rel="nofollow noreferrer">this image</a> from the tutorial.) So is that what's going on? Is there actually no reward, so the <code>initial_inference</code> just always returns a 0 for the reward?</p>

<p>Let's assume that that's the case.</p>

<p>Under this assumption, then, the first reward in the <code>predictions</code> list will be the 0 that the <code>initial_inference</code> function will return for the reward. Then, in the loss, this 0 will be compared with the first element of the <code>target</code> list.</p>

<p>Here's how the <code>target</code> is created:</p>

<pre class="lang-py prettyprint-override"><code>  def make_target(self, state_index: int, num_unroll_steps: int, td_steps: int,
                  to_play: Player):
    # The value target is the discounted root value of the search tree N steps
    # into the future, plus the discounted sum of all rewards until then.
    targets = []
    for current_index in range(state_index, state_index + num_unroll_steps + 1):
      bootstrap_index = current_index + td_steps
      if bootstrap_index &lt; len(self.root_values):
        value = self.root_values[bootstrap_index] * self.discount**td_steps
      else:
        value = 0

      for i, reward in enumerate(self.rewards[current_index:bootstrap_index]):
        value += reward * self.discount**i  # pytype: disable=unsupported-operands

      if current_index &lt; len(self.root_values):
        targets.append((value, self.rewards[current_index],
                        self.child_visits[current_index]))
      else:
        # States past the end of games are treated as absorbing states.
        targets.append((0, 0, []))
    return targets
</code></pre>

<p>The <code>targets</code> returned by this function become the <code>target</code> list in the <code>update_weights</code> function. So the first value in <code>targets</code> is <code>self.rewards[current_index]</code>. The <code>self.rewards</code> is a list of all of the rewards received while playing a game. The only time it is edited is within this function <code>apply</code>:</p>

<pre class="lang-py prettyprint-override"><code>  def apply(self, action: Action):
    reward = self.environment.step(action)
    self.rewards.append(reward)
    self.history.append(action)
</code></pre>

<p>The <code>apply</code> function is only called here:</p>

<pre class="lang-py prettyprint-override"><code># Each game is produced by starting at the initial board position, then
# repeatedly executing a Monte Carlo Tree Search to generate moves until the end
# of the game is reached.
def play_game(config: MuZeroConfig, network: Network) -&gt; Game:
  game = config.new_game()

  while not game.terminal() and len(game.history) &lt; config.max_moves:
    # At the root of the search tree we use the representation function to
    # obtain a hidden state given the current observation.
    root = Node(0)
    current_observation = game.make_image(-1)
    expand_node(root, game.to_play(), game.legal_actions(),
                network.initial_inference(current_observation))
    add_exploration_noise(config, root)

    # We then run a Monte Carlo Tree Search using only action sequences and the
    # model learned by the network.
    run_mcts(config, root, game.action_history(), network)
    action = select_action(config, len(game.history), root, network)
    game.apply(action)
    game.store_search_statistics(root)
  return game
</code></pre>

<p>To me, it looks like <em>every single time an action is taken, a reward is generated</em>. So the first reward in the <code>self.rewards</code> list should be the reward from taking the first action in the game.</p>

<p>The issue becomes clear if <code>current_index = 0</code> in <code>self.rewards[current_index]</code>. In this case, the <code>predictions</code> list will have a 0 for the first reward because it always does. However, the <code>targets</code> list, will have the reward given for completing the first action.</p>

<p>So, to me, <em>it seems like the rewards are misaligned.</em></p>

<p>If we continue, the second reward in the <code>predictions</code> list will be the reward from <code>recurrent_inference</code> for completing the <em>first</em> action. However, the second reward in the <code>targets</code> list will be the reward stored in the game for completing the <em>second</em> action.</p>

<p>So, overall, I have three questions that build on one another:</p>

<ol>
<li><strong>What does the reward from the <code>initial_inference</code> represent? (What is it?)</strong></li>
<li>If it is 0, and it's supposed to represent a reward, are the rewards between the <code>predictions</code> and <code>targets</code> misaligned? (i.e. should the second reward in <code>predictions</code> actually be matched with the first reward in <code>targets</code>?)</li>
<li>If they are misaligned, will the network still train and work correctly?</li>
</ol>

<p>(Another curiosity to note is that despite this misalignment (assuming there is misalignment), both the <code>predictions</code> and <code>targets</code> length do have the same length. The targets length is defined by the line <code>for current_index in range(state_index, state_index + num_unroll_steps + 1)</code> in the <code>make_target</code> function above. Above, we also computed that the length of <code>predictions</code> is <code>len(actions) + 1</code>. And <code>len(actions)</code> is defined by <code>g.history[i:i + num_unroll_steps]</code> in the <code>sample_batch</code> function (see <a href="https://arxiv.org/src/1911.08265v1/anc/pseudocode.py" rel="nofollow noreferrer">the pseudocode</a>). Thus, the length of both lists are the same.)</p>

<p><strong>What's going on?</strong></p>