code: def epsilon_greedy(action_values, epsilon, legal_actions_mask=None):
      """Computes an epsilon-greedy distribution over actions.
  
      This returns a categorical distribution over a discrete action space. It is
      assumed that the trailing dimension of `action_values` is of length A, i.e.
      the number of actions. It is also assumed that actions are 0-indexed.
  
      This policy does the following:
  
      - With probability 1 - epsilon, take the action corresponding to the highest
      action value, breaking ties uniformly at random.
      - With probability epsilon, take an action uniformly at random.
  
      Args:
        action_values: A Tensor of action values with any rank >= 1 and dtype float.
          Shape can be flat ([A]), batched ([B, A]), a batch of sequences
          ([T, B, A]), and so on.
        epsilon: A scalar Tensor (or Python float) with value between 0 and 1.
        legal_actions_mask: An optional one-hot tensor having the shame shape and
          dtypes as `action_values`, defining the legal actions:
          legal_actions_mask[..., a] = 1 if a is legal, 0 otherwise.
          If not provided, all actions will be considered legal and
          `tf.ones_like(action_values)`.
  
      Returns:
        policy: tfp.distributions.Categorical distribution representing the policy.
      """
      with tf.name_scope(
          "epsilon_greedy", values=[action_values, epsilon]
      ):
          # Convert inputs to Tensors if they aren't already.
          action_values = tf.convert_to_tensor(action_values)
          epsilon = tf.convert_to_tensor(
              epsilon, dtype=action_values.dtype
          )
  
          # We compute the action space dynamically.
          num_actions = tf.cast(
              tf.shape(action_values)[-1], action_values.dtype
          )
  
          # Dithering action distribution.
          if legal_actions_mask is None:
              dither_probs = (
                  1 / num_actions * tf.ones_like(action_values)
              )
          else:
              dither_probs = (
                  1
                  / tf.reduce_sum(
                      legal_actions_mask, axis=-1, keepdims=True
                  )
                  * legal_actions_mask
              )
  
          # Greedy action distribution, breaking ties uniformly at random.
          max_value = tf.reduce_max(
              action_values, axis=-1, keepdims=True
          )
          greedy_probs = tf.cast(
              tf.equal(action_values, max_value), action_values.dtype
          )
          greedy_probs /= tf.reduce_sum(
              greedy_probs, axis=-1, keepdims=True
          )
  
          # Epsilon-greedy action distribution.
          probs = epsilon * dither_probs + (1 - epsilon) * greedy_probs
  
          # Make the policy object.
          policy = tfp.distributions.Categorical(probs=probs)
  
      return policy
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import tensorflow.compat.v1
  import tensorflow_probability
label_desc: Use image node_id 1 for calling a global function with example usage: epsilon_greedy(action_values, epsilon, legal_actions_mask) and returns: policy