class_docstr: Class to build GreedyMultiObjectiveNeuralPolicy objects.
class_docstr_tok: ["Class","to","build","GreedyMultiObjectiveNeuralPolicy","objects","."]
docstr: Builds a GreedyMultiObjectiveNeuralPolicy based on multiple networks.
  
  This policy takes an iterable of `tf_agents.Network`, each responsible for
  predicting a specific objective, along with a `Scalarizer` object to
  generate an action by maximizing the scalarized objective, i.e., the output
  of the `Scalarizer` applied to the multiple predicted objectives by the
  networks.
  
  Args:
    time_step_spec: A `TimeStep` spec of the expected time_steps.
    action_spec: A nest of `BoundedTensorSpec` representing the actions.
    scalarizer: A
      `tf_agents.bandits.multi_objective.multi_objective_scalarizer.Scalarizer`
      object that implements scalarization of multiple objectives into a
      single scalar reward.
    objective_networks: A Sequence of `tf_agents.network.Network` objects to
      be used by the policy. Each network will be called with
      call(observation, step_type) and is expected to provide a prediction for
      a specific objective for all actions.
    observation_and_action_constraint_splitter: A function used for masking
      valid/invalid actions with each state of the environment. The function
      takes in a full observation and returns a tuple consisting of 1) the
      part of the observation intended as input to the network and 2) the
      mask.  The mask should be a 0-1 `Tensor` of shape `[batch_size,
      num_actions]`. This function should also work with a `TensorSpec` as
      input, and should output `TensorSpec` objects for the observation and
      mask.
    accepts_per_arm_features: (bool) Whether the policy accepts per-arm
      features.
    emit_policy_info: (tuple of strings) what side information we want to get
      as part of the policy info. Allowed values can be found in
      `policy_utilities.PolicyInfo`.
    name: The name of this policy. All variables in this module will fall
      under that name. Defaults to the class name.
  
  Raises:
    NotImplementedError: If `action_spec` contains more than one
      `BoundedTensorSpec` or the `BoundedTensorSpec` is not valid.
    NotImplementedError: If `action_spec` is not a `BoundedTensorSpec` of type
      int32 and shape ().
    ValueError: If `objective_networks` has fewer than two networks.
    ValueError: If `accepts_per_arm_features` is true but `time_step_spec` is
      None.
docstr_tok: ["Builds","a","GreedyMultiObjectiveNeuralPolicy","based","on","multiple","networks",".","This","policy","takes","an","iterable","of","`","tf_agents.Network","`",",","each","responsible","for","predicting","a","specific","objective",",","along","with","a","`","Scalarizer","`","object","to","generate","an","action","by","maximizing","the","scalarized","objective",",","i.e.",",","the","output","of","the","`","Scalarizer","`","applied","to","the","multiple","predicted","objectives","by","the","networks",".","Args",":","time_step_spec",":","A","`","TimeStep","`","spec","of","the","expected","time_steps",".","action_spec",":","A","nest","of","`","BoundedTensorSpec","`","representing","the","actions",".","scalarizer",":","A","`","tf_agents.bandits.multi_objective.multi_objective_scalarizer.Scalarizer","`","object","that","implements","scalarization","of","multiple","objectives","into","a","single","scalar","reward",".","objective_networks",":","A","Sequence","of","`","tf_agents.network.Network","`","objects","to","be","used","by","the","policy",".","Each","network","will","be","called","with","call","(","observation",",","step_type",")","and","is","expected","to","provide","a","prediction","for","a","specific","objective","for","all","actions",".","observation_and_action_constraint_splitter",":","A","function","used","for","masking","valid\/invalid","actions","with","each","state","of","the","environment",".","The","function","takes","in","a","full","observation","and","returns","a","tuple","consisting","of","1",")","the","part","of","the","observation","intended","as","input","to","the","network","and","2",")","the","mask",".","The","mask","should","be","a","0-1","`","Tensor","`","of","shape","`","[","batch_size",",","num_actions","]","`",".","This","function","should","also","work","with","a","`","TensorSpec","`","as","input",",","and","should","output","`","TensorSpec","`","objects","for","the","observation","and","mask",".","accepts_per_arm_features",":","(","bool",")","Whether","the","policy","accepts","per-arm","features",".","emit_policy_info",":","(","tuple","of","strings",")","what","side","information","we","want","to","get","as","part","of","the","policy","info",".","Allowed","values","can","be","found","in","`","policy_utilities.PolicyInfo","`",".","name",":","The","name","of","this","policy",".","All","variables","in","this","module","will","fall","under","that","name",".","Defaults","to","the","class","name",".","Raises",":","NotImplementedError",":","If","`","action_spec","`","contains","more","than","one","`","BoundedTensorSpec","`","or","the","`","BoundedTensorSpec","`","is","not","valid",".","NotImplementedError",":","If","`","action_spec","`","is","not","a","`","BoundedTensorSpec","`","of","type","int32","and","shape","(",")",".","ValueError",":","If","`","objective_networks","`","has","fewer","than","two","networks",".","ValueError",":","If","`","accepts_per_arm_features","`","is","true","but","`","time_step_spec","`","is","None","."]
code: def __init__(
      self,
      time_step_spec: Optional[ts.TimeStep],
      action_spec: Optional[types.NestedBoundedTensorSpec],
      scalarizer: multi_objective_scalarizer.Scalarizer,
      objective_networks: Sequence[Network],
      observation_and_action_constraint_splitter: types.Splitter = None,
      accepts_per_arm_features: bool = False,
      emit_policy_info: Tuple[Text, ...] = (),
      name: Optional[Text] = None,
  ):
      """Builds a GreedyMultiObjectiveNeuralPolicy based on multiple networks.
  
      This policy takes an iterable of `tf_agents.Network`, each responsible for
      predicting a specific objective, along with a `Scalarizer` object to
      generate an action by maximizing the scalarized objective, i.e., the output
      of the `Scalarizer` applied to the multiple predicted objectives by the
      networks.
  
      Args:
        time_step_spec: A `TimeStep` spec of the expected time_steps.
        action_spec: A nest of `BoundedTensorSpec` representing the actions.
        scalarizer: A
          `tf_agents.bandits.multi_objective.multi_objective_scalarizer.Scalarizer`
          object that implements scalarization of multiple objectives into a
          single scalar reward.
        objective_networks: A Sequence of `tf_agents.network.Network` objects to
          be used by the policy. Each network will be called with
          call(observation, step_type) and is expected to provide a prediction for
          a specific objective for all actions.
        observation_and_action_constraint_splitter: A function used for masking
          valid/invalid actions with each state of the environment. The function
          takes in a full observation and returns a tuple consisting of 1) the
          part of the observation intended as input to the network and 2) the
          mask.  The mask should be a 0-1 `Tensor` of shape `[batch_size,
          num_actions]`. This function should also work with a `TensorSpec` as
          input, and should output `TensorSpec` objects for the observation and
          mask.
        accepts_per_arm_features: (bool) Whether the policy accepts per-arm
          features.
        emit_policy_info: (tuple of strings) what side information we want to get
          as part of the policy info. Allowed values can be found in
          `policy_utilities.PolicyInfo`.
        name: The name of this policy. All variables in this module will fall
          under that name. Defaults to the class name.
  
      Raises:
        NotImplementedError: If `action_spec` contains more than one
          `BoundedTensorSpec` or the `BoundedTensorSpec` is not valid.
        NotImplementedError: If `action_spec` is not a `BoundedTensorSpec` of type
          int32 and shape ().
        ValueError: If `objective_networks` has fewer than two networks.
        ValueError: If `accepts_per_arm_features` is true but `time_step_spec` is
          None.
      """
      policy_utilities.check_no_mask_with_arm_features(
          accepts_per_arm_features,
          observation_and_action_constraint_splitter,
      )
      flat_action_spec = tf.nest.flatten(action_spec)
      if len(flat_action_spec) > 1:
          raise NotImplementedError(
              "action_spec can only contain a single BoundedTensorSpec."
          )
  
      action_spec = flat_action_spec[0]
      if (
          not tensor_spec.is_bounded(action_spec)
          or not tensor_spec.is_discrete(action_spec)
          or action_spec.shape.rank > 1
          or action_spec.shape.num_elements() != 1
      ):
          raise NotImplementedError(
              "action_spec must be a BoundedTensorSpec of type int32 and shape (). "
              "Found {}.".format(action_spec)
          )
      self._expected_num_actions = (
          action_spec.maximum - action_spec.minimum + 1
      )
      self._action_offset = action_spec.minimum
      policy_state_spec = []
      for network in objective_networks:
          policy_state_spec.append(network.state_spec)
          network.create_variables()
      self._objective_networks = objective_networks
      self._scalarizer = scalarizer
      self._num_objectives = len(self._objective_networks)
      if self._num_objectives < 2:
          raise ValueError(
              "Number of objectives should be at least two, but found to be {}".format(
                  self._num_objectives
              )
          )
  
      self._emit_policy_info = emit_policy_info
      predicted_rewards_mean = ()
      if (
          policy_utilities.InfoFields.PREDICTED_REWARDS_MEAN
          in emit_policy_info
      ):
          predicted_rewards_mean = tensor_spec.TensorSpec(
              [self._num_objectives, self._expected_num_actions]
          )
      scalarized_predicted_rewards_mean = ()
      if (
          policy_utilities.InfoFields.MULTIOBJECTIVE_SCALARIZED_PREDICTED_REWARDS_MEAN
          in emit_policy_info
      ):
          scalarized_predicted_rewards_mean = tensor_spec.TensorSpec(
              [self._expected_num_actions]
          )
      bandit_policy_type = ()
      if (
          policy_utilities.InfoFields.BANDIT_POLICY_TYPE
          in emit_policy_info
      ):
          bandit_policy_type = (
              policy_utilities.create_bandit_policy_type_tensor_spec(
                  shape=[1]
              )
          )
      if accepts_per_arm_features:
          if time_step_spec is None:
              raise ValueError(
                  "time_step_spec should not be None for per-arm-features policies, "
                  "but found to be."
              )
          # The features for the chosen arm is saved to policy_info.
          chosen_arm_features_info = (
              policy_utilities.create_chosen_arm_features_info_spec(
                  time_step_spec.observation
              )
          )
          info_spec = policy_utilities.PerArmPolicyInfo(
              predicted_rewards_mean=predicted_rewards_mean,
              multiobjective_scalarized_predicted_rewards_mean=scalarized_predicted_rewards_mean,
              bandit_policy_type=bandit_policy_type,
              chosen_arm_features=chosen_arm_features_info,
          )
      else:
          info_spec = policy_utilities.PolicyInfo(
              predicted_rewards_mean=predicted_rewards_mean,
              multiobjective_scalarized_predicted_rewards_mean=scalarized_predicted_rewards_mean,
              bandit_policy_type=bandit_policy_type,
          )
  
      self._accepts_per_arm_features = accepts_per_arm_features
  
      super(GreedyMultiObjectiveNeuralPolicy, self).__init__(
          time_step_spec,
          action_spec,
          policy_state_spec=policy_state_spec,
          clip=False,
          info_spec=info_spec,
          emit_log_probability="log_probability" in emit_policy_info,
          observation_and_action_constraint_splitter=(
              observation_and_action_constraint_splitter
          ),
          name=name,
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import copy
  from typing import List, Optional, Sequence, Text, Tuple
  import gin
  import tensorflow
  import tensorflow_probability
  from tf_agents.bandits.multi_objective import multi_objective_scalarizer
  from tf_agents.bandits.networks import heteroscedastic_q_network
  from tf_agents.bandits.policies import constraints
  from tf_agents.bandits.specs import utils
  from tf_agents.networks.network import Network
  from tf_agents.policies import tf_policy
  from tf_agents.policies import utils
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import policy_step
  from tf_agents.trajectories import time_step
  from tf_agents.typing import types
  from tf_agents.utils import common
label_desc: Use image node_id 1 to create a new GreedyMultiObjectiveNeuralPolicy object from inherited base classes: tf_policy with example: obj = GreedyMultiObjectiveNeuralPolicy(time_step_spec, action_spec, scalarizer, objective_networks, observation_and_action_constraint_splitter, accepts_per_arm_features, emit_policy_info, name)