code: def generalized_lambda_returns(
      rewards,
      pcontinues,
      values,
      bootstrap_value,
      lambda_=1,
      name="generalized_lambda_returns",
  ):
      """Computes lambda-returns along a batch of (chunks of) trajectories.
  
      For lambda=1 these will be multistep returns looking ahead from each
      state to the end of the chunk, where bootstrap_value is used. If you pass an
      entire trajectory and zeros for bootstrap_value, this is just the Monte-Carlo
      return / TD(1) target.
  
      For lambda=0 these are one-step TD(0) targets.
  
      For inbetween values of lambda these are lambda-returns / TD(lambda) targets,
      except that traces are always cut off at the end of the chunk, since we can't
      see returns beyond then. If you pass an entire trajectory with zeros for
      bootstrap_value though, then they're plain TD(lambda) targets.
  
      lambda can also be a tensor of values in [0, 1], determining the mix of
      bootstrapping vs further accumulation of multistep returns at each timestep.
      This can be used to implement Retrace and other algorithms. See
      `sequence_ops.multistep_forward_view` for more info on this. Another way to
      think about the end-of-chunk cutoff is that lambda is always effectively zero
      on the timestep after the end of the chunk, since at the end of the chunk we
      rely entirely on bootstrapping and can't accumulate returns looking further
      into the future.
  
      The sequences in the tensors should be aligned such that an agent in a state
      with value `V` 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`.
  
      Subtracting `values` from these lambda-returns will yield estimates of the
      advantage function which can be used for both the policy gradient loss and
      the baseline value function loss in A3C / GAE.
  
      Args:
        rewards: 2-D Tensor with shape `[T, B]`.
        pcontinues: 2-D Tensor with shape `[T, B]`.
        values: 2-D Tensor containing estimates of the state values for timesteps
          0 to `T-1`. Shape `[T, B]`.
        bootstrap_value: 1-D Tensor containing an estimate of the value of the
          final state at time `T`, used for bootstrapping the target n-step
          returns. Shape `[B]`.
        lambda_: an optional scalar or 2-D Tensor with shape `[T, B]`.
        name: Customises the name_scope for this op.
  
      Returns:
        2-D Tensor with shape `[T, B]`
      """
      values.get_shape().assert_has_rank(2)
      rewards.get_shape().assert_has_rank(2)
      pcontinues.get_shape().assert_has_rank(2)
      bootstrap_value.get_shape().assert_has_rank(1)
      scoped_values = [
          rewards,
          pcontinues,
          values,
          bootstrap_value,
          lambda_,
      ]
      with tf.name_scope(name, values=scoped_values):
          if lambda_ == 1:
              # This is actually equivalent to the branch below, just an optimisation
              # to avoid unnecessary work in this case:
              return sequence_ops.scan_discounted_sum(
                  rewards,
                  pcontinues,
                  initial_value=bootstrap_value,
                  reverse=True,
                  back_prop=False,
                  name="multistep_returns",
              )
          else:
              v_tp1 = tf.concat(
                  axis=0,
                  values=[
                      values[1:, :],
                      tf.expand_dims(bootstrap_value, 0),
                  ],
              )
              # `back_prop=False` prevents gradients flowing into values and
              # bootstrap_value, which is what you want when using the bootstrapped
              # lambda-returns in an update as targets for values.
              return sequence_ops.multistep_forward_view(
                  rewards,
                  pcontinues,
                  v_tp1,
                  lambda_,
                  back_prop=False,
                  name="generalized_lambda_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 2 for calling a global function with example usage: generalized_lambda_returns(rewards, pcontinues, values, bootstrap_value, lambda_, name) and returns: sequence_ops, sequence_ops