class_docstr: A layer outputting reward predictions using K bins and two-hot encoding.
  
  This layer is used in two models in DreamerV3: The reward predictor of the world
  model and the value function. K is 255 by default (see [1]) and doesn't change
  with the model size.
  
  Possible predicted reward/values range from symexp(-20.0) to symexp(20.0), which
  should cover any possible environment. Outputs of this layer are generated by
  generating logits/probs via a single linear layer, then interpreting the probs
  as weights for a weighted average of the different possible reward (binned) values.
class_docstr_tok: ["A","layer","outputting","reward","predictions","using","K","bins","and","two-hot","encoding",".","This","layer","is","used","in","two","models","in","DreamerV3",":","The","reward","predictor","of","the","world","model","and","the","value","function",".","K","is","255","by","default","(","see","[","1","]",")","and","does","n't","change","with","the","model","size",".","Possible","predicted","reward\/values","range","from","symexp","(","-20.0",")","to","symexp","(","20.0",")",",","which","should","cover","any","possible","environment",".","Outputs","of","this","layer","are","generated","by","generating","logits\/probs","via","a","single","linear","layer",",","then","interpreting","the","probs","as","weights","for","a","weighted","average","of","the","different","possible","reward","(","binned",")","values","."]
docstr: Computes the expected reward using N equal sized buckets of possible values.
  
  Args:
      inputs: The input tensor for the layer, which computes the reward bucket
          weights (logits). [B, dim].
  
  Returns:
      A tuple consisting of the expected rewards and the logits that parameterize
      the tfp `FiniteDiscrete` distribution object. To get the individual bucket
      probs, do `[FiniteDiscrete object].probs`.
docstr_tok: ["Computes","the","expected","reward","using","N","equal","sized","buckets","of","possible","values",".","Args",":","inputs",":","The","input","tensor","for","the","layer",",","which","computes","the","reward","bucket","weights","(","logits",")",".","[","B",",","dim","]",".","Returns",":","A","tuple","consisting","of","the","expected","rewards","and","the","logits","that","parameterize","the","tfp","`","FiniteDiscrete","`","distribution","object",".","To","get","the","individual","bucket","probs",",","do","`","[","FiniteDiscrete","object","]",".probs","`","."]
code: def call(self, inputs):
      """Computes the expected reward using N equal sized buckets of possible values.
  
      Args:
          inputs: The input tensor for the layer, which computes the reward bucket
              weights (logits). [B, dim].
  
      Returns:
          A tuple consisting of the expected rewards and the logits that parameterize
          the tfp `FiniteDiscrete` distribution object. To get the individual bucket
          probs, do `[FiniteDiscrete object].probs`.
      """
      # Compute the `num_buckets` weights.
      assert len(inputs.shape) == 2
      logits = tf.cast(self.reward_buckets_layer(inputs), tf.float32)
      # out=[B, `num_buckets`]
  
      # Compute the expected(!) reward using the formula:
      # `softmax(Linear(x))` [vectordot] `possible_outcomes`, where
      # `possible_outcomes` is the even-spaced (binned) encoding of all possible
      # symexp'd reward/values.
      # [2]: "The mean of the reward predictor pφ(ˆrt | zˆt) is used as reward
      # sequence rˆ1:H."
      probs = tf.nn.softmax(logits)
      possible_outcomes = tf.linspace(
          self.lower_bound,
          self.upper_bound,
          self.num_buckets,
      )
      # probs=possible_outcomes=[B, `num_buckets`]
  
      # Simple vector dot product (over last dim) to get the mean reward
      # weighted sum, where all weights sum to 1.0.
      expected_rewards = tf.reduce_sum(
          probs * possible_outcomes, axis=-1
      )
      # expected_rewards=[B]
  
      return expected_rewards, logits
  
imports: from ray.rllib.utils.framework import try_import_tf
label_desc: Use image node_id 2 for calling the RewardPredictorLayer obj's underlying member method code with example usage: obj.call(inputs) and returns: expected_rewards,  logits