class_docstr: A policy that chooses from a set of policies to decide the action.
class_docstr_tok: ["A","policy","that","chooses","from","a","set","of","policies","to","decide","the","action","."]
code: def _action(self, time_step, policy_state, seed=None):
      first_obs = tf.nest.flatten(time_step.observation)[0]
      batch_size = (
          tf.compat.dimension_value(first_obs.shape[0])
          or tf.shape(first_obs)[0]
      )
      policy_choice = self._mixture_distribution.sample(batch_size)
      policy_steps = [
          policy.action(time_step, policy_state)
          for policy in self._policies
      ]
      policy_actions = nest_utils.stack_nested_tensors(
          [step.action for step in policy_steps], axis=-1
      )
      policy_infos = nest_utils.stack_nested_tensors(
          [step.info for step in policy_steps], axis=-1
      )
  
      expanded_choice = tf.expand_dims(policy_choice, axis=-1)
      mixture_action = tf.nest.map_structure(
          lambda t: tf.gather(t, policy_choice, batch_dims=1),
          policy_actions,
      )
  
      expanded_mixture_info = tf.nest.map_structure(
          lambda t: tf.gather(
              t, expanded_choice, batch_dims=1, axis=-1
          ),
          policy_infos,
      )
      mixture_info = tf.nest.map_structure(
          lambda t: tf.squeeze(t, axis=-1), expanded_mixture_info
      )
      return policy_step.PolicyStep(
          mixture_action,
          policy_state,
          {
              MIXTURE_AGENT_ID: policy_choice,
              SUBPOLICY_INFO: mixture_info,
          },
      )
  
file_path: agents/tf_agents/bandits/policies/mixture_policy.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from typing import Optional, Sequence, Text
  import gin
  import tensorflow
  import tensorflow_probability
  from tf_agents.policies import tf_policy
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import policy_step
  from tf_agents.typing import types
  from tf_agents.utils import nest_utils
label_desc: Use image node_id 4 for calling the MixturePolicy obj's underlying member method code with example usage: obj._action(time_step, policy_state, seed) and returns: policy_step