class_docstr: Exponentially Q-weighted policy.
  
  Given a stochastic policy and a critic, returns a (stochastic) policy which
  samples multiple actions from the underlying policy, computes the Q-values for
  each action, and chooses the final action among the sampled ones with
  probability proportional to the exponentiated Q values, tempered by
  a parameter beta.
class_docstr_tok: ["Exponentially","Q-weighted","policy",".","Given","a","stochastic","policy","and","a","critic",",","returns","a","(","stochastic",")","policy","which","samples","multiple","actions","from","the","underlying","policy",",","computes","the","Q-values","for","each","action",",","and","chooses","the","final","action","among","the","sampled","ones","with","probability","proportional","to","the","exponentiated","Q","values",",","tempered","by","a","parameter","beta","."]
code: def __call__(self, inputs: types.NestedTensor) -> tf.Tensor:
      # Inputs are of size [B, ...]. Here we tile them to be of shape [N, B, ...].
      tiled_inputs = tf2_utils.tile_nested(
          inputs, self._num_action_samples
      )
      shape = tf.shape(tree.flatten(tiled_inputs)[0])
      n, b = shape[0], shape[1]
      tf.debugging.assert_equal(
          n,
          self._num_action_samples,
          "Internal Error. Unexpected tiled_inputs shape.",
      )
      dummy_zeros_n_b = tf.zeros((n, b))
      # Reshape to [N * B, ...].
      merge = lambda x: snt.merge_leading_dims(x, 2)
      tiled_inputs = tree.map_structure(merge, tiled_inputs)
  
      tiled_actions = self._actor_network(tiled_inputs)
  
      # Compute Q-values and the resulting tempered probabilities.
      q = self._critic_network(tiled_inputs, tiled_actions)
      boltzmann_logits = q / self._beta
  
      boltzmann_logits = snt.split_leading_dim(
          boltzmann_logits, dummy_zeros_n_b, 2
      )
      # [B, N]
      boltzmann_logits = tf.transpose(boltzmann_logits, perm=(1, 0))
      # Resample one action per batch according to the Boltzmann distribution.
      action_idx = tfp.distributions.Categorical(
          logits=boltzmann_logits
      ).sample()
      # [B, 2], where the first column is 0, 1, 2,... corresponding to indices to
      # the batch dimension.
      action_idx = tf.stack((tf.range(b), action_idx), axis=1)
  
      tiled_actions = snt.split_leading_dim(
          tiled_actions, dummy_zeros_n_b, 2
      )
      action_dim = len(tiled_actions.get_shape().as_list())
      tiled_actions = tf.transpose(
          tiled_actions, perm=[1, 0] + list(range(2, action_dim))
      )
      # [B, ...]
      action_sample = tf.gather_nd(tiled_actions, action_idx)
  
      return action_sample
  
file_path: acme/acme/tf/networks/stochastic.py
imports: from acme import types
  from acme.tf import utils
  import sonnet
  import tensorflow
  import tensorflow_probability
  import tree
label_desc: Use image node_id 2 for calling the ExpQWeightedPolicy obj's underlying member method code with example usage: obj.__call__(inputs) and returns: action_sample