code: def compute_cross_entropy_loss(
      sampled_actions: tf.Tensor,
      normalized_weights: tf.Tensor,
      online_action_distribution: tfp.distributions.Distribution,
  ) -> tf.Tensor:
      """Compute cross-entropy online and the reweighted target policy.
  
      Args:
        sampled_actions: samples used in the Monte Carlo integration in the policy
          loss. Expected shape is [N, B, ...], where N is the number of sampled
          actions and B is the number of sampled states.
        normalized_weights: target policy multiplied by the exponentiated Q values
          and normalized; expected shape is [N, B].
        online_action_distribution: policy to be optimized.
  
      Returns:
        loss_policy_gradient: the cross-entropy loss that, when differentiated,
          produces the policy gradient.
      """
  
      # Compute the M-step loss.
      log_prob = online_action_distribution.log_prob(sampled_actions)
  
      # Compute the weighted average log-prob using the normalized weights.
      loss_policy_gradient = -tf.reduce_sum(
          log_prob * normalized_weights, axis=0
      )
  
      # Return the mean loss over the batch of states.
      return tf.reduce_mean(loss_policy_gradient, axis=0)
  
file_path: acme/acme/tf/losses/mpo.py
imports: from typing import Dict, Tuple, Union
  import sonnet
  import tensorflow
  import tensorflow_probability
label_desc: Use image node_id 3 for calling a global function with example usage: compute_cross_entropy_loss(sampled_actions, normalized_weights, online_action_distribution) and returns: tf