code: def td_lambda(
      state_values,
      rewards,
      pcontinues,
      bootstrap_value,
      lambda_=1,
      name="BaselineLoss",
  ):
      """Constructs a TensorFlow graph computing the L2 loss for sequences.
  
      This loss learns the baseline for advantage actor-critic models. Gradients
      for this loss flow through each tensor in `state_values`, but no other
      input tensors. The baseline is regressed towards the n-step bootstrapped
      returns given by the reward/pcontinue sequence.
  
      This function is designed for batches of sequences of data. Tensors are
      assumed to be time major (i.e. the outermost dimension is time, the second
      outermost dimension is the batch dimension). We denote the sequence length
      in the shapes of the arguments with the variable `T`, the batch size with
      the variable `B`, neither of which needs to be known at construction time.
      Index `0` of the time dimension is assumed to be the start of the sequence.
  
      `rewards` and `pcontinues` are the sequences of data taken directly from the
      environment, possibly modulated by a discount. `state_values` are the
      sequences of (typically learnt) estimates of the values of the states
      visited along a batch of trajectories.
  
      The sequences in the tensors should be aligned such that an agent in a state
      with value `V` that takes an action transitions into another state
      with value `V'`, receiving reward `r` and pcontinue `p`. Then `V`, `r`
      and `p` are all at the same index `i` in the corresponding tensors. `V'` is
      at index `i+1`, or in the `bootstrap_value` tensor if `i == T`.
  
      See "High-dimensional continuous control using generalized advantage
      estimation" by Schulman, Moritz, Levine et al.
      (https://arxiv.org/abs/1506.02438).
  
      Args:
        state_values: 2-D Tensor of state-value estimates with shape `[T, B]`.
        rewards: 2-D Tensor with shape `[T, B]`.
        pcontinues: 2-D Tensor with shape `[T, B]`.
        bootstrap_value: 1-D Tensor with shape `[B]`.
        lambda_: an optional scalar or 2-D Tensor with shape `[T, B]`.
        name: Customises the name_scope for this op.
  
      Returns:
        A namedtuple with fields:
  
        * `loss`: a tensor containing the batch of losses, shape `[B]`.
        * `extra`: a namedtuple with fields:
            * temporal_differences, Tensor of shape `[T, B]`
            * discounted_returns, Tensor of shape `[T, B]`
      """
      scoped_values = [
          state_values,
          rewards,
          pcontinues,
          bootstrap_value,
      ]
      with tf.name_scope(name, values=scoped_values):
          discounted_returns = generalized_lambda_returns(
              rewards,
              pcontinues,
              state_values,
              bootstrap_value,
              lambda_,
          )
          temporal_differences = discounted_returns - state_values
          loss = 0.5 * tf.reduce_sum(
              tf.square(temporal_differences), axis=0, name="l2_loss"
          )
  
          return base_ops.LossOutput(
              loss,
              TDLambdaExtra(
                  temporal_differences=temporal_differences,
                  discounted_returns=discounted_returns,
              ),
          )
  
file_path: trfl/trfl/value_ops.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  import tensorflow.compat.v1
  from trfl import base_ops
  from trfl import sequence_ops
label_desc: Use image node_id 3 for calling a global function with example usage: td_lambda(state_values, rewards, pcontinues, bootstrap_value, lambda_, name) and returns: base_ops