class_docstr: Class to build GreedyMultiObjectiveNeuralPolicy objects.
class_docstr_tok: ["Class","to","build","GreedyMultiObjectiveNeuralPolicy","objects","."]
docstr: Predict the objectives using the policy's objective networks.
  
  Args:
    observation: The observation whose objectives are to be predicted.
    step_type: The `tf_agents.trajectories.time_step.StepType` for the input
      observation.
    policy_state: The states for the policy's objective networks.
  
  Returns:
    A tuple of a rank-3 tensor for the predicted objectives shaped as
    [batch_size, num_of_objectives, num_of_actions] and a list of updated
    states, one for each objective network.
  
  Raises:
    ValueError: If the output of any objective network is not a rank-2 tensor.
    ValueError: If the output size of any objective network does not match the
      expected number of actions.
docstr_tok: ["Predict","the","objectives","using","the","policy","'s","objective","networks",".","Args",":","observation",":","The","observation","whose","objectives","are","to","be","predicted",".","step_type",":","The","`","tf_agents.trajectories.time_step.StepType","`","for","the","input","observation",".","policy_state",":","The","states","for","the","policy","'s","objective","networks",".","Returns",":","A","tuple","of","a","rank-3","tensor","for","the","predicted","objectives","shaped","as","[","batch_size",",","num_of_objectives",",","num_of_actions","]","and","a","list","of","updated","states",",","one","for","each","objective","network",".","Raises",":","ValueError",":","If","the","output","of","any","objective","network","is","not","a","rank-2","tensor",".","ValueError",":","If","the","output","size","of","any","objective","network","does","not","match","the","expected","number","of","actions","."]
code: def _predict(
      self,
      observation: types.NestedSpecTensorOrArray,
      step_type: types.SpecTensorOrArray,
      policy_state: Sequence[types.TensorSpec],
  ) -> Tuple[tf.Tensor, List[types.TensorSpec]]:
      """Predict the objectives using the policy's objective networks.
  
      Args:
        observation: The observation whose objectives are to be predicted.
        step_type: The `tf_agents.trajectories.time_step.StepType` for the input
          observation.
        policy_state: The states for the policy's objective networks.
  
      Returns:
        A tuple of a rank-3 tensor for the predicted objectives shaped as
        [batch_size, num_of_objectives, num_of_actions] and a list of updated
        states, one for each objective network.
  
      Raises:
        ValueError: If the output of any objective network is not a rank-2 tensor.
        ValueError: If the output size of any objective network does not match the
          expected number of actions.
      """
      if not common.safe_has_state(policy_state):
          policy_state = [()] * self._num_objectives
      predicted_objective_values = []
      updated_policy_state = []
      for idx, (network, state) in enumerate(
          zip(self._objective_networks, policy_state)
      ):
          prediction, state = network(observation, step_type, state)
          updated_policy_state.append(copy.deepcopy(state))
          predicted_value = (
              prediction.q_value_logits
              if isinstance(
                  network,
                  heteroscedastic_q_network.HeteroscedasticQNetwork,
              )
              else prediction
          )
          if predicted_value.shape.rank != 2:
              raise ValueError(
                  "Prediction from network {} shoud be a rank-2 tensor, "
                  " but has shape {}".format(idx, predicted_value.shape)
              )
          if (
              predicted_value.shape[1] is not None
              and predicted_value.shape[1] != self._expected_num_actions
          ):
              raise ValueError(
                  "The number of actions ({}) does not match objective network {}"
                  " output size ({}).".format(
                      self._expected_num_actions,
                      idx,
                      predicted_value.shape[1],
                  )
              )
          predicted_objective_values.append(predicted_value)
      # Stack the list of predicted objective tensors into a rank-3 tensor shaped
      # as [batch_size, num_of_objectives, num_of_actions].
      predicted_objectives_tensor = tf.stack(
          predicted_objective_values, axis=1
      )
      return predicted_objectives_tensor, updated_policy_state
  
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 4 for calling the GreedyMultiObjectiveNeuralPolicy obj's underlying member method code with example usage: obj._predict(observation, step_type, policy_state) and returns: predicted_objectives_tensor,  updated_policy_state