class_docstr: A REINFORCE Agent.
  
  Implements:
  
  REINFORCE algorithm from
  
  "Simple statistical gradient-following algorithms for connectionist
  reinforcement learning"
  Williams, R.J., 1992.
  https://www-anw.cs.umass.edu/~barto/courses/cs687/williams92simple.pdf
  
  REINFORCE with state-value baseline, where state-values are estimated with
  function approximation, from
  
  "Reinforcement learning: An introduction" (Sec. 13.4)
  Sutton, R.S. and Barto, A.G., 2018.
  http://incompleteideas.net/book/the-book-2nd.html
  
  The REINFORCE agent can be optionally provided with:
  - value_network: A `tf_agents.network.Network` which parameterizes state-value
    estimation as a neural network. The network will be called with
    call(observation, step_type) and returns a floating point state-values
    tensor.
  - value_estimation_loss_coef: Weight on the value prediction loss.
  
  If value_network and value_estimation_loss_coef are provided, advantages are
  computed as
    `advantages = (discounted accumulated rewards) - (estimated state-values)`
  and the overall learning objective becomes:
    `(total loss) =
      (policy gradient loss) +
      value_estimation_loss_coef * (squared error of estimated state-values)`
class_docstr_tok: ["A","REINFORCE","Agent",".","Implements",":","REINFORCE","algorithm","from","''","Simple","statistical","gradient-following","algorithms","for","connectionist","reinforcement","learning","''","Williams",",","R.J.",",","1992.","https",":","\/\/www-anw.cs.umass.edu\/~barto\/courses\/cs687\/williams92simple.pdf","REINFORCE","with","state-value","baseline",",","where","state-values","are","estimated","with","function","approximation",",","from","''","Reinforcement","learning",":","An","introduction","''","(","Sec",".","13.4",")","Sutton",",","R.S",".","and","Barto",",","A.G.",",","2018.","http",":","\/\/incompleteideas.net\/book\/the-book-2nd.html","The","REINFORCE","agent","can","be","optionally","provided","with",":","-","value_network",":","A","`","tf_agents.network.Network","`","which","parameterizes","state-value","estimation","as","a","neural","network",".","The","network","will","be","called","with","call","(","observation",",","step_type",")","and","returns","a","floating","point","state-values","tensor",".","-","value_estimation_loss_coef",":","Weight","on","the","value","prediction","loss",".","If","value_network","and","value_estimation_loss_coef","are","provided",",","advantages","are","computed","as","`","advantages","=","(","discounted","accumulated","rewards",")","-","(","estimated","state-values",")","`","and","the","overall","learning","objective","becomes",":","`","(","total","loss",")","=","(","policy","gradient","loss",")","+","value_estimation_loss_coef","*","(","squared","error","of","estimated","state-values",")","`"]
code: def total_loss(
      self,
      experience: traj.Trajectory,
      returns: types.Tensor,
      weights: types.Tensor,
      training: bool = False,
  ) -> tf_agent.LossInfo:
      # Ensure we see at least one full episode.
      time_steps = ts.TimeStep(
          experience.step_type,
          tf.zeros_like(experience.reward),
          tf.zeros_like(experience.discount),
          experience.observation,
      )
      is_last = experience.is_last()
      num_episodes = tf.reduce_sum(tf.cast(is_last, tf.float32))
      tf.debugging.assert_greater(
          num_episodes,
          0.0,
          message=(
              "No complete episode found. REINFORCE requires full episodes "
              "to compute losses."
          ),
      )
  
      # Mask out partial episodes at the end of each batch of time_steps.
      # NOTE: We use is_last rather than is_boundary because the last transition
      # is the transition with the last valid reward.  In other words, the
      # reward on the boundary transitions do not have valid rewards.  Since
      # REINFORCE is calculating a loss w.r.t. the returns (and not bootstrapping)
      # keeping the boundary transitions is irrelevant.
      valid_mask = tf.cast(experience.is_last(), dtype=tf.float32)
      valid_mask = tf.math.cumsum(valid_mask, axis=1, reverse=True)
      valid_mask = tf.cast(valid_mask > 0, dtype=tf.float32)
      if weights is not None:
          weights *= valid_mask
      else:
          weights = valid_mask
  
      advantages = returns
      value_preds = None
  
      if self._baseline:
          value_preds, _ = self._value_network(
              time_steps.observation,
              time_steps.step_type,
              training=True,
          )
          if self._debug_summaries:
              tf.compat.v2.summary.histogram(
                  name="value_preds",
                  data=value_preds,
                  step=self.train_step_counter,
              )
  
      advantages = self._advantage_fn(returns, value_preds)
      if self._debug_summaries:
          tf.compat.v2.summary.histogram(
              name="advantages",
              data=advantages,
              step=self.train_step_counter,
          )
  
      # TODO(b/126592060): replace with tensor normalizer.
      if self._normalize_returns:
          advantages = _standard_normalize(advantages, axes=(0, 1))
          if self._debug_summaries:
              tf.compat.v2.summary.histogram(
                  name="normalized_%s"
                  % ("advantages" if self._baseline else "returns"),
                  data=advantages,
                  step=self.train_step_counter,
              )
  
      nest_utils.assert_same_structure(time_steps, self.time_step_spec)
      policy_state = _get_initial_policy_state(
          self.collect_policy, time_steps
      )
      actions_distribution = self.collect_policy.distribution(
          time_steps, policy_state=policy_state
      ).action
  
      policy_gradient_loss = self.policy_gradient_loss(
          actions_distribution,
          experience.action,
          experience.is_boundary(),
          advantages,
          num_episodes,
          weights,
      )
  
      entropy_regularization_loss = self.entropy_regularization_loss(
          actions_distribution, weights
      )
  
      network_regularization_loss = tf.nn.scale_regularization_loss(
          self._actor_network.losses
      )
  
      total_loss = (
          policy_gradient_loss
          + network_regularization_loss
          + entropy_regularization_loss
      )
  
      losses_dict = {
          "policy_gradient_loss": policy_gradient_loss,
          "policy_network_regularization_loss": network_regularization_loss,
          "entropy_regularization_loss": entropy_regularization_loss,
          "value_estimation_loss": 0.0,
          "value_network_regularization_loss": 0.0,
      }
  
      value_estimation_loss = None
      if self._baseline:
          value_estimation_loss = self.value_estimation_loss(
              value_preds, returns, num_episodes, weights
          )
          value_network_regularization_loss = (
              tf.nn.scale_regularization_loss(
                  self._value_network.losses
              )
          )
          total_loss += (
              value_estimation_loss + value_network_regularization_loss
          )
          losses_dict["value_estimation_loss"] = value_estimation_loss
          losses_dict[
              "value_network_regularization_loss"
          ] = value_network_regularization_loss
  
      loss_info_extra = ReinforceAgentLossInfo(**losses_dict)
  
      losses_dict[
          "total_loss"
      ] = total_loss  # Total loss not in loss_info_extra.
  
      common.summarize_scalar_dict(
          losses_dict, self.train_step_counter, name_scope="Losses/"
      )
  
      return tf_agent.LossInfo(total_loss, loss_info_extra)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  from typing import Callable, Optional, Text
  import gin
  import numpy
  import tensorflow
  from tf_agents.agents import data_converter
  from tf_agents.agents import tf_agent
  from tf_agents.networks import network
  from tf_agents.policies import actor_policy
  from tf_agents.policies import greedy_policy
  from tf_agents.trajectories import time_step
  from tf_agents.trajectories import trajectory
  from tf_agents.typing import types
  from tf_agents.utils import common
  from tf_agents.utils import eager_utils
  from tf_agents.utils import nest_utils
  from tf_agents.utils import value_ops
label_desc: Use image node_id 4 for calling the ReinforceAgent obj's underlying member method code with example usage: obj.total_loss(experience, returns, weights, training) and returns: tf_agent