class_docstr: The critic network described in [1], predicting values for policy learning.
  
  Contains a copy of itself (EMA net) for weight regularization.
  The EMA net is updated after each train step via EMA (using the `ema_decay`
  parameter and the actual critic's weights). The EMA net is NOT used for target
  computations (we use the actual critic for that), its only purpose is to compute a
  weights regularizer term for the critic's loss such that the actual critic does not
  move too quickly.
class_docstr_tok: ["The","critic","network","described","in","[","1","]",",","predicting","values","for","policy","learning",".","Contains","a","copy","of","itself","(","EMA","net",")","for","weight","regularization",".","The","EMA","net","is","updated","after","each","train","step","via","EMA","(","using","the","`","ema_decay","`","parameter","and","the","actual","critic","'s","weights",")",".","The","EMA","net","is","NOT","used","for","target","computations","(","we","use","the","actual","critic","for","that",")",",","its","only","purpose","is","to","compute","a","weights","regularizer","term","for","the","critic","'s","loss","such","that","the","actual","critic","does","not","move","too","quickly","."]
docstr: Performs a forward pass through the critic network.
  
  Args:
      h: The deterministic hidden state of the sequence model. [B, dim(h)].
      z: The stochastic discrete representations of the original
          observation input. [B, num_categoricals, num_classes].
      use_ema: Whether to use the EMA-copy of the critic instead of the actual
          critic to perform this computation.
docstr_tok: ["Performs","a","forward","pass","through","the","critic","network",".","Args",":","h",":","The","deterministic","hidden","state","of","the","sequence","model",".","[","B",",","dim","(","h",")","]",".","z",":","The","stochastic","discrete","representations","of","the","original","observation","input",".","[","B",",","num_categoricals",",","num_classes","]",".","use_ema",":","Whether","to","use","the","EMA-copy","of","the","critic","instead","of","the","actual","critic","to","perform","this","computation","."]
code: def call(self, h, z, use_ema):
      """Performs a forward pass through the critic network.
  
      Args:
          h: The deterministic hidden state of the sequence model. [B, dim(h)].
          z: The stochastic discrete representations of the original
              observation input. [B, num_categoricals, num_classes].
          use_ema: Whether to use the EMA-copy of the critic instead of the actual
              critic to perform this computation.
      """
      # Flatten last two dims of z.
      assert len(z.shape) == 3
      z_shape = tf.shape(z)
      z = tf.reshape(z, shape=(z_shape[0], -1))
      assert len(z.shape) == 2
      out = tf.concat([h, z], axis=-1)
      out.set_shape(
          [
              None,
              (
                  get_num_z_categoricals(self.model_size)
                  * get_num_z_classes(self.model_size)
                  + get_gru_units(self.model_size)
              ),
          ]
      )
  
      if not use_ema:
          # Send h-cat-z through MLP.
          out = self.mlp(out)
          # Return expected return OR (expected return, probs of bucket values).
          return self.return_layer(out)
      else:
          out = self.mlp_ema(out)
          return self.return_layer_ema(out)
  
file_path: ray/python/ray/rllib/algorithms/dreamerv3/tf/models/critic_network.py
imports: from typing import Optional
  from ray.rllib.algorithms.dreamerv3.tf.models.components.mlp import MLP
  from ray.rllib.algorithms.dreamerv3.tf.models.components.reward_predictor_layer import RewardPredictorLayer
  from ray.rllib.algorithms.dreamerv3.utils import get_gru_units, get_num_z_categoricals, get_num_z_classes
  from ray.rllib.utils.framework import try_import_tf
label_desc: Use image node_id 2 for calling the CriticNetwork obj's underlying member method code with example usage: obj.call(h, z, use_ema) and returns: self, self