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 _train(self, experience, weights=None):
      experience = self._as_trajectory(experience)
  
      # Add a mask to ensure we reset the return calculation at episode
      # boundaries. This is needed in cases where episodes are truncated before
      # reaching a terminal state. Note experience is a batch of trajectories
      # where reward=next_step.reward so the mask may look shifted at first.
      non_last_mask = tf.cast(
          tf.math.not_equal(
              experience.next_step_type, ts.StepType.LAST
          ),
          tf.float32,
      )
      discounts = non_last_mask * experience.discount * self._gamma
      returns = value_ops.discounted_return(
          experience.reward, discounts, time_major=False
      )
  
      if self._debug_summaries:
          tf.compat.v2.summary.histogram(
              name="rewards",
              data=experience.reward,
              step=self.train_step_counter,
          )
          tf.compat.v2.summary.histogram(
              name="discounts",
              data=experience.discount,
              step=self.train_step_counter,
          )
          tf.compat.v2.summary.histogram(
              name="returns", data=returns, step=self.train_step_counter
          )
  
      with tf.GradientTape() as tape:
          loss_info = self.total_loss(
              experience,
              tf.stop_gradient(returns),
              weights=weights,
              training=True,
          )
          tf.debugging.check_numerics(
              loss_info.loss, "Loss is inf or nan"
          )
      variables_to_train = self._actor_network.trainable_weights
      if self._baseline:
          variables_to_train += self._value_network.trainable_weights
      grads = tape.gradient(loss_info.loss, variables_to_train)
  
      grads_and_vars = list(zip(grads, variables_to_train))
      if self._gradient_clipping:
          grads_and_vars = eager_utils.clip_gradient_norms(
              grads_and_vars, self._gradient_clipping
          )
  
      if self._summarize_grads_and_vars:
          eager_utils.add_variables_summaries(
              grads_and_vars, self.train_step_counter
          )
          eager_utils.add_gradients_summaries(
              grads_and_vars, self.train_step_counter
          )
  
      self._optimizer.apply_gradients(grads_and_vars)
      self.train_step_counter.assign_add(1)
  
      return tf.nest.map_structure(tf.identity, loss_info)
  
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 3 for calling the ReinforceAgent obj's underlying member method code with example usage: obj._train(experience, weights) and returns: tf