class_docstr: An agent that mixes a set of agents with a given mixture.
  
  For every data sample, the agent updates the sub-agent that was used to make
  the action choice in that sample. For this update to happen, the mixture agent
  needs to have the information on which sub-agent is "responsible" for the
  action. This information is in a policy info field `mixture_agent_id`.
  
  Note that this agent makes use of `tf.dynamic_partition`, and thus it is not
  compatible with XLA.
class_docstr_tok: ["An","agent","that","mixes","a","set","of","agents","with","a","given","mixture",".","For","every","data","sample",",","the","agent","updates","the","sub-agent","that","was","used","to","make","the","action","choice","in","that","sample",".","For","this","update","to","happen",",","the","mixture","agent","needs","to","have","the","information","on","which","sub-agent","is","``","responsible","''","for","the","action",".","This","information","is","in","a","policy","info","field","`","mixture_agent_id","`",".","Note","that","this","agent","makes","use","of","`","tf.dynamic_partition","`",",","and","thus","it","is","not","compatible","with","XLA","."]
code: def _train(self, experience, weights=None):
      del weights  # unused
      experience = self._as_trajectory(experience)
  
      reward, _ = nest_utils.flatten_multi_batched_nested_tensors(
          experience.reward, self._time_step_spec.reward
      )
      action, _ = nest_utils.flatten_multi_batched_nested_tensors(
          experience.action, self._action_spec
      )
      observation, _ = nest_utils.flatten_multi_batched_nested_tensors(
          experience.observation, self._time_step_spec.observation
      )
      (
          policy_choice,
          _,
      ) = nest_utils.flatten_multi_batched_nested_tensors(
          experience.policy_info[mixture_policy.MIXTURE_AGENT_ID],
          self._time_step_spec.reward,
      )
      (
          original_infos,
          _,
      ) = nest_utils.flatten_multi_batched_nested_tensors(
          experience.policy_info[mixture_policy.SUBPOLICY_INFO],
          self._original_info_spec,
      )
  
      partitioned_nested_infos = nest_utils.batch_nested_tensors(
          _dynamic_partition_of_nested_tensors(
              original_infos, policy_choice, self._num_agents
          )
      )
  
      partitioned_nested_rewards = [
          nest_utils.batch_nested_tensors(t)
          for t in _dynamic_partition_of_nested_tensors(
              reward, policy_choice, self._num_agents
          )
      ]
      partitioned_nested_actions = [
          nest_utils.batch_nested_tensors(t)
          for t in _dynamic_partition_of_nested_tensors(
              action, policy_choice, self._num_agents
          )
      ]
      partitioned_nested_observations = [
          nest_utils.batch_nested_tensors(t)
          for t in _dynamic_partition_of_nested_tensors(
              observation, policy_choice, self._num_agents
          )
      ]
      loss = 0
      for k in range(self._num_agents):
          per_policy_experience = trajectory.single_step(
              observation=partitioned_nested_observations[k],
              action=partitioned_nested_actions[k],
              policy_info=partitioned_nested_infos[k],
              reward=partitioned_nested_rewards[k],
              discount=tf.zeros_like(partitioned_nested_rewards[k]),
          )
          loss_info = self._agents[k].train(per_policy_experience)
          loss += loss_info.loss
      common.function_in_tf1()(self._update_mixture_distribution)(
          experience
      )
      return tf_agent.LossInfo(loss=(loss), extra=())
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import abc
  from typing import List, Optional, Sequence, Text
  import gin
  import tensorflow
  from tf_agents.agents import data_converter
  from tf_agents.agents import tf_agent
  from tf_agents.bandits.policies import mixture_policy
  from tf_agents.trajectories import trajectory
  from tf_agents.typing import types
  from tf_agents.utils import common
  from tf_agents.utils import nest_utils
label_desc: Use image node_id 4 for calling the MixtureAgent obj's underlying member method code with example usage: obj._train(experience, weights) and returns: tf_agent