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: Initializes a CriticNetwork instance.
  
  Args:
      model_size: The "Model Size" used according to [1] Appendinx B.
         Use None for manually setting the different network sizes.
      num_buckets: The number of buckets to create. Note that the number of
          possible symlog'd outcomes from the used distribution is
          `num_buckets` + 1:
          lower_bound --bucket-- o[1] --bucket-- o[2] ... --bucket-- upper_bound
          o=outcomes
          lower_bound=o[0]
          upper_bound=o[num_buckets]
      lower_bound: The symlog'd lower bound for a possible reward value.
          Note that a value of -20.0 here already allows individual (actual env)
          rewards to be as low as -400M. Buckets will be created between
          `lower_bound` and `upper_bound`.
      upper_bound: The symlog'd upper bound for a possible reward value.
          Note that a value of +20.0 here already allows individual (actual env)
          rewards to be as high as 400M. Buckets will be created between
          `lower_bound` and `upper_bound`.
      ema_decay: The weight to use for updating the weights of the critic's copy
          vs the actual critic. After each training update, the EMA copy of the
          critic gets updated according to:
          ema_net=(`ema_decay`*ema_net) + (1.0-`ema_decay`)*critic_net
          The EMA copy of the critic is used inside the critic loss function only
          to produce a regularizer term against the current critic's weights, NOT
          to compute any target values.
docstr_tok: ["Initializes","a","CriticNetwork","instance",".","Args",":","model_size",":","The","``","Model","Size","''","used","according","to","[","1","]","Appendinx","B",".","Use","None","for","manually","setting","the","different","network","sizes",".","num_buckets",":","The","number","of","buckets","to","create",".","Note","that","the","number","of","possible","symlog","'d","outcomes","from","the","used","distribution","is","`","num_buckets","`","+","1",":","lower_bound","--","bucket","--","o","[","1","]","--","bucket","--","o","[","2","]","...","--","bucket","--","upper_bound","o=outcomes","lower_bound=o","[","0","]","upper_bound=o","[","num_buckets","]","lower_bound",":","The","symlog","'d","lower","bound","for","a","possible","reward","value",".","Note","that","a","value","of","-20.0","here","already","allows","individual","(","actual","env",")","rewards","to","be","as","low","as","-400M",".","Buckets","will","be","created","between","`","lower_bound","`","and","`","upper_bound","`",".","upper_bound",":","The","symlog","'d","upper","bound","for","a","possible","reward","value",".","Note","that","a","value","of","+20.0","here","already","allows","individual","(","actual","env",")","rewards","to","be","as","high","as","400M",".","Buckets","will","be","created","between","`","lower_bound","`","and","`","upper_bound","`",".","ema_decay",":","The","weight","to","use","for","updating","the","weights","of","the","critic","'s","copy","vs","the","actual","critic",".","After","each","training","update",",","the","EMA","copy","of","the","critic","gets","updated","according","to",":","ema_net=","(","`","ema_decay","`","*","ema_net",")","+","(","1.0-","`","ema_decay","`",")","*","critic_net","The","EMA","copy","of","the","critic","is","used","inside","the","critic","loss","function","only","to","produce","a","regularizer","term","against","the","current","critic","'s","weights",",","NOT","to","compute","any","target","values","."]
code: def __init__(
      self,
      *,
      model_size: Optional[str] = "XS",
      num_buckets: int = 255,
      lower_bound: float = -20.0,
      upper_bound: float = 20.0,
      ema_decay: float = 0.98,
  ):
      """Initializes a CriticNetwork instance.
  
      Args:
          model_size: The "Model Size" used according to [1] Appendinx B.
             Use None for manually setting the different network sizes.
          num_buckets: The number of buckets to create. Note that the number of
              possible symlog'd outcomes from the used distribution is
              `num_buckets` + 1:
              lower_bound --bucket-- o[1] --bucket-- o[2] ... --bucket-- upper_bound
              o=outcomes
              lower_bound=o[0]
              upper_bound=o[num_buckets]
          lower_bound: The symlog'd lower bound for a possible reward value.
              Note that a value of -20.0 here already allows individual (actual env)
              rewards to be as low as -400M. Buckets will be created between
              `lower_bound` and `upper_bound`.
          upper_bound: The symlog'd upper bound for a possible reward value.
              Note that a value of +20.0 here already allows individual (actual env)
              rewards to be as high as 400M. Buckets will be created between
              `lower_bound` and `upper_bound`.
          ema_decay: The weight to use for updating the weights of the critic's copy
              vs the actual critic. After each training update, the EMA copy of the
              critic gets updated according to:
              ema_net=(`ema_decay`*ema_net) + (1.0-`ema_decay`)*critic_net
              The EMA copy of the critic is used inside the critic loss function only
              to produce a regularizer term against the current critic's weights, NOT
              to compute any target values.
      """
      super().__init__(name="critic")
  
      self.model_size = model_size
      self.ema_decay = ema_decay
  
      # "Fast" critic network(s) (mlp + reward-pred-layer). This is the network
      # we actually train with our critic loss.
      # IMPORTANT: We also use this to compute the return-targets, BUT we regularize
      # the critic loss term such that the weights of this fast critic stay close
      # to the EMA weights (see below).
      self.mlp = MLP(
          model_size=self.model_size,
          output_layer_size=None,
      )
      self.return_layer = RewardPredictorLayer(
          num_buckets=num_buckets,
          lower_bound=lower_bound,
          upper_bound=upper_bound,
      )
  
      # Weights-EMA (EWMA) containing networks for critic loss (similar to a
      # target net, BUT not used to compute anything, just for the
      # weights regularizer term inside the critic loss).
      self.mlp_ema = MLP(
          model_size=self.model_size,
          output_layer_size=None,
          trainable=False,
      )
      self.return_layer_ema = RewardPredictorLayer(
          num_buckets=num_buckets,
          lower_bound=lower_bound,
          upper_bound=upper_bound,
          trainable=False,
      )
  
      # Trace self.call.
      dl_type = (
          tf.keras.mixed_precision.global_policy().compute_dtype
          or tf.float32
      )
      self.call = tf.function(
          input_signature=[
              tf.TensorSpec(
                  shape=[None, get_gru_units(model_size)], dtype=dl_type
              ),
              tf.TensorSpec(
                  shape=[
                      None,
                      get_num_z_categoricals(model_size),
                      get_num_z_classes(model_size),
                  ],
                  dtype=dl_type,
              ),
              tf.TensorSpec(shape=[], dtype=tf.bool),
          ]
      )(self.call)
  
file_path: 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 1 to create a new CriticNetwork object from inherited base classes: tf.keras with example: obj = CriticNetwork()