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))
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]`
  
trfl/value_ops.py