class_docstr: Categorical distribution, for discrete integer actions (specification key: `categorical`).
  
  Args:
      temperature_mode ("predicted" | "global"): Whether to predict the temperature via a linear
          transformation of the state embedding, or to parametrize the temperature by a separate
          set of trainable weights
          (<span style="color:#00C000"><b>default</b></span>: default temperature of 1).
      skip_linear (bool): Whether to not add the implicit linear logits layer, requires suitable
          network output shape according to action space, not compatible with temperature_mode
          (<span style="color:#00C000"><b>default</b></span>: false).
      name (string): <span style="color:#0000C0"><b>internal use</b></span>.
      action_spec (specification): <span style="color:#0000C0"><b>internal use</b></span>.
      input_spec (specification): <span style="color:#0000C0"><b>internal use</b></span>.
class_docstr_tok: ["Categorical","distribution",",","for","discrete","integer","actions","(","specification","key",":","`","categorical","`",")",".","Args",":","temperature_mode","(","``","predicted","''","|","``","global","''",")",":","Whether","to","predict","the","temperature","via","a","linear","transformation","of","the","state","embedding",",","or","to","parametrize","the","temperature","by","a","separate","set","of","trainable","weights","(","<","span","style=","''","color",":","#","00C000","''",">","<","b",">","default","<","\/b",">","<","\/span",">",":","default","temperature","of","1",")",".","skip_linear","(","bool",")",":","Whether","to","not","add","the","implicit","linear","logits","layer",",","requires","suitable","network","output","shape","according","to","action","space",",","not","compatible","with","temperature_mode","(","<","span","style=","''","color",":","#","00C000","''",">","<","b",">","default","<","\/b",">","<","\/span",">",":","false",")",".","name","(","string",")",":","<","span","style=","''","color",":","#","0000C0","''",">","<","b",">","internal","use","<","\/b",">","<","\/span",">",".","action_spec","(","specification",")",":","<","span","style=","''","color",":","#","0000C0","''",">","<","b",">","internal","use","<","\/b",">","<","\/span",">",".","input_spec","(","specification",")",":","<","span","style=","''","color",":","#","0000C0","''",">","<","b",">","internal","use","<","\/b",">","<","\/span",">","."]
code: def parametrize(self, *, x, conditions):
      epsilon = tf_util.constant(value=util.epsilon, dtype="float")
      log_epsilon = tf_util.constant(
          value=np.log(util.epsilon), dtype="float"
      )
      log_two = tf_util.constant(value=np.log(2.0), dtype="float")
  
      # Action values
      if self.action_values is None:
          action_values = x
      else:
          action_values = self.action_values.apply(x=x)
      shape = (
          (-1,)
          + self.action_spec.shape
          + (self.action_spec.num_values,)
      )
      action_values = tf.reshape(tensor=action_values, shape=shape)
  
      # Softplus standard deviation
      if self.temperature_mode == "global":
          multiples = (tf.shape(input=x)[0],) + tuple(
              1 for _ in range(self.action_spec.rank + 1)
          )
          softplus_temperature = tf.tile(
              input=self.temperature, multiples=multiples
          )
      elif self.temperature_mode == "predicted":
          softplus_temperature = self.temperature.apply(x=x)
          shape = (-1,) + self.action_spec.shape + (1,)
          softplus_temperature = tf.reshape(
              tensor=softplus_temperature, shape=shape
          )
  
      if self.temperature_mode is None:
          # Logits
          logits = action_values
  
          # Implicit states value
          state_value = tf.reduce_logsumexp(
              input_tensor=logits, axis=-1
          )
  
      else:
          # Clip softplus_temperature for numerical stability (epsilon < 1.0, hence negative)
          softplus_temperature = tf.clip_by_value(
              t=softplus_temperature,
              clip_value_min=log_epsilon,
              clip_value_max=-log_epsilon,
          )
  
          # Softplus transformation (based on https://arxiv.org/abs/2007.06059)
          softplus_shift = tf_util.constant(value=0.2, dtype="float")
          temperature = (
              tf.nn.softplus(features=softplus_temperature)
              + softplus_shift
          ) / (log_two + softplus_shift)
  
          # Logits
          logits = action_values / temperature
  
          # Implicit states value
          temperature = tf.squeeze(input=temperature, axis=-1)
          state_value = temperature * tf.reduce_logsumexp(
              input_tensor=logits, axis=-1
          )
  
      # # Explicit states value and advantage-based action values
      # state_value = self.state_value.apply(x=x)
      # state_value = tf.reshape(tensor=state_value, shape=shape[:-1])
      # action_values = tf.expand_dims(input=state_value, axis=-1) + action_values
      # action_values -= tf.math.reduce_mean(input_tensor=action_values, axis=-1, keepdims=True)
  
      # Action masking, affects action_values/probabilities/logits but not state_value
      if self.config.enable_int_action_masking:
          min_float = tf.fill(
              dims=tf.shape(input=action_values),
              value=tf_util.get_dtype(type="float").min,
          )
          action_values = tf.where(
              condition=conditions["mask"], x=action_values, y=min_float
          )
          logits = tf.where(
              condition=conditions["mask"], x=logits, y=min_float
          )
  
      # Softmax for corresponding probabilities
      probabilities = tf.nn.softmax(logits=logits, axis=-1)
  
      # "Normalized" logits
      logits = tf.math.log(x=(probabilities + epsilon))
      # Unstable
      # logits = tf.nn.log_softmax(logits=logits, axis=-1)
      # Doesn't take masking into account
      # logits = action_values - tf.expand_dims(input=state_value, axis=-1) ... / temperature
  
      if self.temperature_mode is None:
          return TensorDict(
              probabilities=probabilities,
              logits=logits,
              action_values=action_values,
              state_value=state_value,
          )
      else:
          return TensorDict(
              probabilities=probabilities,
              temperature=temperature,
              logits=logits,
              action_values=action_values,
              state_value=state_value,
          )
  
file_path: tensorforce/tensorforce/core/distributions/categorical.py
imports: import numpy
  import tensorflow
  from tensorforce import TensorforceError, util
  from tensorforce.core import layer_modules, TensorDict, TensorSpec, TensorsSpec, tf_function, tf_util
  from tensorforce.core.distributions import Distribution
label_desc: Use image node_id 4 for calling the Categorical obj's underlying member method code with example usage: obj.parametrize() and returns: TensorDict, TensorDict