class_docstr: AlphaZero-style learning.
class_docstr_tok: ["AlphaZero-style","learning","."]
docstr: Do a step of SGD on the loss.
docstr_tok: ["Do","a","step","of","SGD","on","the","loss","."]
code: def _step(self) -> tf.Tensor:
      """Do a step of SGD on the loss."""
  
      inputs = next(self._iterator)
      o_t, _, r_t, d_t, o_tp1, extras = inputs.data
      pi_t = extras["pi"]
  
      with tf.GradientTape() as tape:
          # Forward the network on the two states in the transition.
          logits, value = self._network(o_t)
          _, target_value = self._network(o_tp1)
          target_value = tf.stop_gradient(target_value)
  
          # Value loss is simply on-policy TD learning.
          value_loss = tf.square(
              r_t + self._discount * d_t * target_value - value
          )
  
          # Policy loss distills MCTS policy into the policy network.
          policy_loss = tf.nn.softmax_cross_entropy_with_logits(
              logits=logits, labels=pi_t
          )
  
          # Compute gradients.
          loss = tf.reduce_mean(value_loss + policy_loss)
          gradients = tape.gradient(
              loss, self._network.trainable_variables
          )
  
      self._optimizer.apply(
          gradients, self._network.trainable_variables
      )
  
      return loss
  
file_path: acme/acme/agents/tf/mcts/learning.py
imports: from typing import List, Optional
  import acme
  from acme.tf import utils
  from acme.utils import counting
  from acme.utils import loggers
  import numpy
  import sonnet
  import tensorflow
label_desc: Use image node_id 2 for calling the AZLearner obj's underlying member method code with example usage: obj._step() and returns: loss