code: def _compute_distributional_critic_loss(
      sampled_q_t_all: List[tf.Tensor],
      q_tm1_all: List[tf.Tensor],
      r_t_all: tf.Tensor,
      d_t: tf.Tensor,
      discount: float,
      num_samples: int,
  ):
      """Compute loss and sampled Q-values for distributional critics."""
      # Compute average logits by first reshaping them and normalizing them
      # across atoms.
      batch_size = r_t_all.get_shape()[0]
      # Cast the additional discount to match the environment discount dtype.
      discount = tf.cast(discount, dtype=d_t.dtype)
      critic_losses = []
      sampled_q_ts = []
      for idx, (
          sampled_q_t_distributions,
          q_tm1_distribution,
      ) in enumerate(zip(sampled_q_t_all, q_tm1_all)):
          # Compute loss for distributional critic for objective c
          sampled_logits = tf.reshape(
              sampled_q_t_distributions.logits,
              [num_samples, batch_size, -1],
          )  # [N, B, A]
          sampled_logprobs = tf.math.log_softmax(
              sampled_logits, axis=-1
          )
          averaged_logits = tf.reduce_logsumexp(
              sampled_logprobs, axis=0
          )
  
          # Construct the expected distributional value for bootstrapping.
          q_t_distribution = networks.DiscreteValuedDistribution(
              values=sampled_q_t_distributions.values,
              logits=averaged_logits,
          )
  
          # Compute critic distributional loss.
          critic_loss = losses.categorical(
              q_tm1_distribution,
              r_t_all[:, idx],
              discount * d_t,
              q_t_distribution,
          )
          critic_losses.append(tf.reduce_mean(critic_loss))
  
          # Compute Q-values of sampled actions and reshape to [N, B].
          sampled_q_ts.append(
              tf.reshape(
                  sampled_q_t_distributions.mean(), (num_samples, -1)
              )
          )
  
      critic_loss = tf.reduce_mean(critic_losses)
      sampled_q_t = tf.stack(sampled_q_ts, axis=-1)  # [N, B, C]
      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 1 for calling a global function with example usage: _compute_distributional_critic_loss(sampled_q_t_all, q_tm1_all, r_t_all, d_t, discount, num_samples) and returns: critic_loss,  sampled_q_t