code: def _compute_critic_loss(
      sampled_q_t_all: tf.Tensor,
      q_tm1_all: tf.Tensor,
      r_t_all: tf.Tensor,
      d_t: tf.Tensor,
      discount: float,
      num_samples: int,
      num_critic_heads: int,
  ):
      """Compute loss and sampled Q-values for (non-distributional) critics."""
      # Reshape Q-value samples back to original batch dimensions and average
      # them to compute the TD-learning bootstrap target.
      batch_size = r_t_all.get_shape()[0]
      sampled_q_t = tf.reshape(
          sampled_q_t_all, (num_samples, batch_size, num_critic_heads)
      )  # [N,B,C]
      q_t = tf.reduce_mean(sampled_q_t, axis=0)  # [B, C]
  
      # Flatten q_t and q_tm1; necessary for trfl.td_learning
      q_t = tf.reshape(q_t, [-1])  # [B*C]
      q_tm1 = tf.reshape(q_tm1_all, [-1])  # [B*C]
  
      # Flatten r_t_all; necessary for trfl.td_learning
      r_t_all = tf.reshape(r_t_all, [-1])  # [B*C]
  
      # Broadcast and then flatten d_t, to match shape of q_t and q_tm1
      d_t = tf.tile(d_t, [num_critic_heads])  # [B*C]
      # Cast the additional discount to match the environment discount dtype.
      discount = tf.cast(discount, dtype=d_t.dtype)
  
      # Critic loss.
      critic_loss = trfl.td_learning(
          q_tm1, r_t_all, discount * d_t, q_t
      ).loss
      critic_loss = tf.reduce_mean(critic_loss)
      return critic_loss, sampled_q_t
  
imports: import dataclasses
  import time
  from typing import Callable, List, Optional, Sequence
  import acme
  from acme import types
  from acme.tf import losses
  from acme.tf import networks
  from acme.tf import savers
  from acme.tf import utils
  from acme.utils import counting
  from acme.utils import loggers
  import numpy
  import sonnet
  import tensorflow
  import trfl
label_desc: Use image node_id 2 for calling a global function with example usage: _compute_critic_loss(sampled_q_t_all, q_tm1_all, r_t_all, d_t, discount, num_samples, num_critic_heads) and returns: critic_loss,  sampled_q_t