class_docstr: Predict the RSSM's z^(t+1), given h(t), z^(t), and a(t).
  
  Disagreement (stddev) between the N networks in this model on what the next z^ would
  be are used to produce intrinsic rewards for enhanced, curiosity-based exploration.
  
  TODO
code: def compute_intrinsic_rewards(self, h, z, a):
      forward_train_outs = self.forward_train(a=a, h=h, z=z)
      B = tf.shape(h)[0]
  
      # Intrinsic rewards are computed as:
      # Stddev (between the different nets) of the 32x32 discrete, stochastic
      # probabilities. Meaning that if the larger the disagreement
      # (stddev) between the nets on what the probabilities for the different
      # classes should be, the higher the intrinsic reward.
      z_predicted_probs_N_B = forward_train_outs[
          "z_predicted_probs_N_HxB"
      ]
      N = len(z_predicted_probs_N_B)
      z_predicted_probs_N_B = tf.stack(z_predicted_probs_N_B, axis=0)
      # Flatten z-dims (num_categoricals x num_classes).
      z_predicted_probs_N_B = tf.reshape(
          z_predicted_probs_N_B, shape=(N, B, -1)
      )
  
      # Compute stddevs over all disagree nets (axis=0).
      # Mean over last axis ([num categoricals] x [num classes] folded axis).
      stddevs_B_mean = tf.reduce_mean(
          tf.math.reduce_std(z_predicted_probs_N_B, axis=0),
          axis=-1,
      )
      # TEST:
      stddevs_B_mean -= tf.reduce_mean(stddevs_B_mean)
      # END TEST
      return {
          "rewards_intrinsic": stddevs_B_mean
          * self.intrinsic_rewards_scale,
          "forward_train_outs": forward_train_outs,
      }
  
file_path: ray/python/ray/rllib/algorithms/dreamerv3/tf/models/disagree_networks.py
imports: from ray.rllib.algorithms.dreamerv3.tf.models.components.mlp import MLP
  from ray.rllib.algorithms.dreamerv3.tf.models.components.representation_layer import RepresentationLayer
  from ray.rllib.utils.framework import try_import_tf, try_import_tfp
label_desc: Use image node_id 3 for calling the DisagreeNetworks obj's underlying member method code with example usage: obj.compute_intrinsic_rewards(h, z, a) and returns: dict