class_docstr: GumbelSoftmax distribution with temperature and logits.
  
  The implementation is almost identical to tfp.distributions.
  relaxed_onehot_categorical.RelaxedOneHotCategorical except for the following:
  
  1. Add mode() function to return mode of the underlying categorical
     distribution (There is no mode() defined in RelaxedOneHotCategorical)
  2. Add a convert_to_integer() function to convert the sample from non-integer
     to integer. Note that the sample function returns one_hot format of the
     discrete action that is different from regular distributions.
  3. log_prob() of RelaxedOneHotCategorical will return INF when the input is
     at boundary. In this implementation, we add a small epsilon to avoid
     getting NAN. In addition, when the input is discrete, we calculate log_prob
     using the underlying categorical distribution.
class_docstr_tok: ["GumbelSoftmax","distribution","with","temperature","and","logits",".","The","implementation","is","almost","identical","to","tfp.distributions",".","relaxed_onehot_categorical.RelaxedOneHotCategorical","except","for","the","following",":","1",".","Add","mode","(",")","function","to","return","mode","of","the","underlying","categorical","distribution","(","There","is","no","mode","(",")","defined","in","RelaxedOneHotCategorical",")","2",".","Add","a","convert_to_integer","(",")","function","to","convert","the","sample","from","non-integer","to","integer",".","Note","that","the","sample","function","returns","one_hot","format","of","the","discrete","action","that","is","different","from","regular","distributions",".","3.","log_prob","(",")","of","RelaxedOneHotCategorical","will","return","INF","when","the","input","is","at","boundary",".","In","this","implementation",",","we","add","a","small","epsilon","to","avoid","getting","NAN",".","In","addition",",","when","the","input","is","discrete",",","we","calculate","log_prob","using","the","underlying","categorical","distribution","."]
code: def convert_to_one_hot(self, samples):
      return tf.one_hot(
          tf.argmax(samples, axis=-1),
          self.distribution.event_size,
          dtype=self._output_dtype,
      )
  
file_path: agents/tf_agents/distributions/gumbel_softmax.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import tensorflow.compat.v2
  import tensorflow_probability
label_desc: Use image node_id 3 for calling the GumbelSoftmax obj's underlying member method code with example usage: obj.convert_to_one_hot(samples) and returns: tf