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",")","`"]
docstr: Computes the policy gradient loss.
  
  Args:
    actions_distribution: A possibly batched tuple of action distributions.
    actions: Tensor with a batch of actions.
    is_boundary: Tensor of booleans that indicate if the corresponding action
      was in a boundary trajectory and should be ignored.
    returns: Tensor with a return from each timestep, aligned on index. Works
      better when returns are normalized.
    num_episodes: Number of episodes contained in the training data.
    weights: Optional scalar or element-wise (per-batch-entry) importance
      weights.  May include a mask for invalid timesteps.
  
  Returns:
    policy_gradient_loss: A tensor that will contain policy gradient loss for
      the on-policy experience.
docstr_tok: ["Computes","the","policy","gradient","loss",".","Args",":","actions_distribution",":","A","possibly","batched","tuple","of","action","distributions",".","actions",":","Tensor","with","a","batch","of","actions",".","is_boundary",":","Tensor","of","booleans","that","indicate","if","the","corresponding","action","was","in","a","boundary","trajectory","and","should","be","ignored",".","returns",":","Tensor","with","a","return","from","each","timestep",",","aligned","on","index",".","Works","better","when","returns","are","normalized",".","num_episodes",":","Number","of","episodes","contained","in","the","training","data",".","weights",":","Optional","scalar","or","element-wise","(","per-batch-entry",")","importance","weights",".","May","include","a","mask","for","invalid","timesteps",".","Returns",":","policy_gradient_loss",":","A","tensor","that","will","contain","policy","gradient","loss","for","the","on-policy","experience","."]
code: def policy_gradient_loss(
      self,
      actions_distribution: types.NestedDistribution,
      actions: types.NestedTensor,
      is_boundary: types.Tensor,
      returns: types.Tensor,
      num_episodes: types.Int,
      weights: Optional[types.Tensor] = None,
  ) -> types.Tensor:
      """Computes the policy gradient loss.
  
      Args:
        actions_distribution: A possibly batched tuple of action distributions.
        actions: Tensor with a batch of actions.
        is_boundary: Tensor of booleans that indicate if the corresponding action
          was in a boundary trajectory and should be ignored.
        returns: Tensor with a return from each timestep, aligned on index. Works
          better when returns are normalized.
        num_episodes: Number of episodes contained in the training data.
        weights: Optional scalar or element-wise (per-batch-entry) importance
          weights.  May include a mask for invalid timesteps.
  
      Returns:
        policy_gradient_loss: A tensor that will contain policy gradient loss for
          the on-policy experience.
      """
      # TODO(b/126594799): Add class IndependentNested(tfd.Distribution) to handle
      # nests of independent distributions like this.
      action_log_prob = common.log_probability(
          actions_distribution, actions, self.action_spec
      )
  
      # Filter out transitions between end state of previous episode and start
      # state of next episode.
      valid_mask = tf.cast(~is_boundary, tf.float32)
      action_log_prob *= valid_mask
  
      action_log_prob_times_return = action_log_prob * returns
  
      if weights is not None:
          action_log_prob_times_return *= weights
  
      if self._debug_summaries:
          tf.compat.v2.summary.histogram(
              name="action_log_prob",
              data=action_log_prob,
              step=self.train_step_counter,
          )
          tf.compat.v2.summary.histogram(
              name="action_log_prob_times_return",
              data=action_log_prob_times_return,
              step=self.train_step_counter,
          )
  
      # Policy gradient loss is defined as the sum, over timesteps, of action
      #   log-probability times the cumulative return from that timestep onward.
      #   For more information, see (Williams, 1992).
      policy_gradient_loss = -tf.reduce_sum(
          input_tensor=action_log_prob_times_return
      )
  
      # We take the mean over episodes by dividing by num_episodes.
      policy_gradient_loss = policy_gradient_loss / num_episodes
  
      return policy_gradient_loss
  
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 5 for calling the ReinforceAgent obj's underlying member method code with example usage: obj.policy_gradient_loss(actions_distribution, actions, is_boundary, returns, num_episodes, weights) and returns: policy_gradient_loss