# coding=utf-8

# Copyright 2020 The TF-Agents Authors.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# https://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

"""Utilities for policies."""

from __future__ import absolute_import

from __future__ import division

from __future__ import print_function

import collections

from typing import Optional, Sequence, Text, Union

import tensorflow as tf

from tf_agents.specs import bandit_spec_utils

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 common

def get_num_actions_from_tensor_spec(

action_spec: types.BoundedTensorSpec,

) -> int:

"""Validates `action_spec` and returns number of actions.

`action_spec` must specify a scalar int32 or int64 with minimum zero.

Args:

action_spec: a `BoundedTensorSpec`.

Returns:

The number of actions described by `action_spec`.

Raises:

ValueError: if `action_spec` is not an bounded scalar int32 or int64 spec

with minimum 0.

"""

if not isinstance(action_spec, tensor_spec.BoundedTensorSpec):

raise ValueError(

'Action spec must be a `BoundedTensorSpec`; got {}'.format(

type(action_spec)

)

)

if action_spec.shape.rank != 0:

raise ValueError(

'Action spec must be a scalar; got shape{}'.format(action_spec.shape)

)

if action_spec.dtype not in (tf.int32, tf.int64):

raise ValueError(

'Action spec must be have dtype int32 or int64; got {}'.format(

action_spec.dtype

)

)

if action_spec.minimum != 0:

raise ValueError(

'Action spec must have minimum 0; got {}'.format(action_spec.minimum)

)

return action_spec.maximum + 1

class InfoFields(object):

"""Strings which can be used in the policy info fields."""

LOG_PROBABILITY = policy_step.CommonFields.LOG_PROBABILITY

# Mean of predicted rewards (per arm).

PREDICTED_REWARDS_MEAN = 'predicted_rewards_mean'

# For multi-objective policies, this field stores the scalarized predicted

# mean rewards that the policy uses to choose an arm, whereas

# `PREDICTED_REWARDS_MEAN` stores the predicted mean rewards for all

# objectives.

MULTIOBJECTIVE_SCALARIZED_PREDICTED_REWARDS_MEAN = (

'multiobjective_scalarized_predicted_rewards_mean'

)

# Optimistic estimates of predicted rewards (per arm).

PREDICTED_REWARDS_OPTIMISTIC = 'predicted_rewards_optimistic'

# Samples of predicted rewards (per arm).

PREDICTED_REWARDS_SAMPLED = 'predicted_rewards_sampled'

# Type of bandit policy (see enumerations in `BanditPolicyType`).

BANDIT_POLICY_TYPE = 'bandit_policy_type'

# Used to store the chosen action for a per-arm model.

CHOSEN_ARM_FEATURES = 'chosen_arm_features'

PolicyInfo = collections.namedtuple( # pylint: disable=invalid-name

'PolicyInfo',

(

InfoFields.LOG_PROBABILITY,

InfoFields.PREDICTED_REWARDS_MEAN,

InfoFields.MULTIOBJECTIVE_SCALARIZED_PREDICTED_REWARDS_MEAN,

InfoFields.PREDICTED_REWARDS_OPTIMISTIC,

InfoFields.PREDICTED_REWARDS_SAMPLED,

InfoFields.BANDIT_POLICY_TYPE,

),

)

# Set default empty tuple for all fields.

PolicyInfo.__new__.__defaults__ = ((),) * len(PolicyInfo._fields)

PerArmPolicyInfo = collections.namedtuple( # pylint: disable=invalid-name

'PerArmPolicyInfo',

(

InfoFields.LOG_PROBABILITY,

InfoFields.PREDICTED_REWARDS_MEAN,

InfoFields.MULTIOBJECTIVE_SCALARIZED_PREDICTED_REWARDS_MEAN,

InfoFields.PREDICTED_REWARDS_OPTIMISTIC,

InfoFields.PREDICTED_REWARDS_SAMPLED,

InfoFields.BANDIT_POLICY_TYPE,

InfoFields.CHOSEN_ARM_FEATURES,

),

)

# Set default empty tuple for all fields.

PerArmPolicyInfo.__new__.__defaults__ = ((),) * len(PerArmPolicyInfo._fields)

def populate_policy_info(

arm_observations: types.Tensor,

chosen_actions: types.Tensor,

rewards_for_argmax: types.Tensor,

est_rewards: types.Tensor,

emit_policy_info: Sequence[Text],

accepts_per_arm_features: bool,

) -> PolicyInfo:

"""Populates policy info given all needed input.

Args:

arm_observations: In case the policy accepts per-arm feautures, this is a

Tensor with the per-arm features. Otherwise its value is unused.

chosen_actions: A Tensor with the indices of the chosen actions.

rewards_for_argmax: The sampled or optimistically boosted reward estimates

based on which the policy chooses the action greedily.

est_rewards: A Tensor with the rewards estimated by the model.

emit_policy_info: A set of policy info keys, specifying wich info fields to

populate

accepts_per_arm_features: (bool) Whether the policy accepts per-arm

features.

Returns:

A policy info.

"""

if accepts_per_arm_features:

# Saving the features for the chosen action to the policy_info.

chosen_arm_features = tf.nest.map_structure(

lambda t: tf.gather(params=t, indices=chosen_actions, batch_dims=1),

arm_observations,

)

policy_info = PerArmPolicyInfo(

predicted_rewards_optimistic=(

rewards_for_argmax

if InfoFields.PREDICTED_REWARDS_OPTIMISTIC in emit_policy_info

else ()

),

predicted_rewards_sampled=(

rewards_for_argmax

if InfoFields.PREDICTED_REWARDS_SAMPLED in emit_policy_info

else ()

),

predicted_rewards_mean=(

est_rewards

if InfoFields.PREDICTED_REWARDS_MEAN in emit_policy_info

else ()

),

chosen_arm_features=chosen_arm_features,

)

else:

policy_info = PolicyInfo(

predicted_rewards_optimistic=(

rewards_for_argmax

if InfoFields.PREDICTED_REWARDS_OPTIMISTIC in emit_policy_info

else ()

),

predicted_rewards_sampled=(

rewards_for_argmax

if InfoFields.PREDICTED_REWARDS_SAMPLED in emit_policy_info

else ()

),

predicted_rewards_mean=(

est_rewards

if InfoFields.PREDICTED_REWARDS_MEAN in emit_policy_info

else ()

),

)

return policy_info

class BanditPolicyType(object):

"""Enumeration of bandit policy types."""

# No bandit policy type specified.

UNKNOWN = 0

# Greedy decision made by bandit agent.

GREEDY = 1

# Random decision for exploration made by epsilon-greedy agent sampled from

# uniform distribution over actions.

UNIFORM = 2

# Decision made by Boltzmann exploration.

BOLTZMANN = 3

# Decision made by FALCON.

FALCON = 4

def create_bandit_policy_type_tensor_spec(

shape: types.Shape,

) -> types.BoundedTensorSpec:

"""Create tensor spec for bandit policy type."""

return tensor_spec.BoundedTensorSpec(

shape=shape,

dtype=tf.int32,

minimum=BanditPolicyType.UNKNOWN,

maximum=BanditPolicyType.FALCON,

)

@common.function

def masked_argmax(

input_tensor: types.Tensor,

mask: types.Tensor,

output_type: tf.DType = tf.int32,

) -> types.Tensor:

"""Computes the argmax where the allowed elements are given by a mask.

If a row of `mask` contains all zeros, then this method will return -1 for the

corresponding row of `input_tensor`.

Args:

input_tensor: Rank-2 Tensor of floats.

mask: 0-1 valued Tensor of the same shape as input.

output_type: Integer type of the output.

Returns:

A Tensor of rank 1 and type `output_type`, with the masked argmax of every

row of `input_tensor`.

"""

input_tensor.shape.assert_is_compatible_with(mask.shape)

neg_inf = tf.constant(-float('Inf'), input_tensor.dtype)

modified_input = tf.compat.v2.where(

tf.cast(mask, tf.bool), input_tensor, neg_inf

)

argmax_tensor = tf.argmax(modified_input, axis=-1, output_type=output_type)

# Replace results for invalid mask rows with -1.

reduce_mask = tf.cast(tf.reduce_max(mask, axis=1), tf.bool)

neg_one = tf.constant(-1, output_type)

return tf.compat.v2.where(reduce_mask, argmax_tensor, neg_one)

def has_bandit_policy_type(

info: Optional[Union[PolicyInfo, PerArmPolicyInfo]],

check_for_tensor: bool = False,

) -> bool:

"""Check if policy info has `bandit_policy_type` field/tensor."""

if not info:

return False

fields = getattr(info, '_fields', None)

has_field = fields is not None and InfoFields.BANDIT_POLICY_TYPE in fields

if has_field and check_for_tensor:

return isinstance(info.bandit_policy_type, tf.Tensor)

else:

return has_field

def has_chosen_arm_features(

info: Optional[Union[PolicyInfo, PerArmPolicyInfo]],

check_for_tensor: bool = False,

) -> bool:

"""Check if policy info has `chosen_arm_features` field/tensor."""

if not info:

return False

fields = getattr(info, '_fields', None)

has_field = fields is not None and InfoFields.CHOSEN_ARM_FEATURES in fields

if has_field and check_for_tensor:

return isinstance(info.chosen_arm_features, tf.Tensor)

else:

return has_field

def set_bandit_policy_type(

info: Optional[Union[PolicyInfo, PerArmPolicyInfo]],

bandit_policy_type: types.Tensor,

) -> Union[PolicyInfo, PerArmPolicyInfo]:

"""Sets the InfoFields.BANDIT_POLICY_TYPE on info to bandit_policy_type.

If policy `info` does not support InfoFields.BANDIT_POLICY_TYPE, this method

returns `info` as-is (without any modification).

Args:

info: Policy info on which to set bandit policy type.

bandit_policy_type: Tensor containing BanditPolicyType enums or TensorSpec

from `create_bandit_policy_type_tensor_spec()`.

Returns:

Policy info with modified field (if possible).

"""

if not info:

return PolicyInfo(bandit_policy_type=bandit_policy_type)

fields = getattr(info, '_fields', None)

if fields is not None and InfoFields.BANDIT_POLICY_TYPE in fields:

return info._replace(bandit_policy_type=bandit_policy_type)

try:

info[InfoFields.BANDIT_POLICY_TYPE] = bandit_policy_type

except TypeError:

pass

return info

@common.function

def bandit_policy_uniform_mask(

values: types.Tensor, mask: types.Tensor

) -> types.Tensor:

"""Set bandit policy type tensor to BanditPolicyType.UNIFORM based on mask.

Set bandit policy type `values` to BanditPolicyType.UNIFORM; returns tensor

where output[i] is BanditPolicyType.UNIFORM if mask[i] is True, otherwise it

is left as values[i].

Args:

values: Tensor containing `BanditPolicyType` enumerations.

mask: Tensor of the same shape as `values` with boolean flags indicating

values to set to `BanditPolicyType.UNIFORM`.

Returns:

Tensor containing `BanditPolicyType` enumerations with masked values.

"""

return tf.where(

mask, tf.fill(tf.shape(values), BanditPolicyType.UNIFORM), values

)

def get_model_index(arm_index: int, accepts_per_arm_features: bool) -> int:

"""Returns the model index for a specific arm.

The number of models depends on the observation format: If the policy accepts

per-arm features, there is only one single model used for every arm. Otherwise

there is a model for every arm.

Args:

arm_index: The index of the arm for which the model index is needed.

accepts_per_arm_features: (bool) Whether the policy works with per-arm

features.

Returns:

The index of the model for the arm requested.

"""

return 0 if accepts_per_arm_features else arm_index

def create_chosen_arm_features_info_spec(

observation_spec: types.NestedTensorSpec,

) -> types.NestedTensorSpec:

"""Creates the chosen arm features info spec from the arm observation spec."""

arm_spec = observation_spec[bandit_spec_utils.PER_ARM_FEATURE_KEY]

return tensor_spec.remove_outer_dims_nest(arm_spec, 1)

def check_no_mask_with_arm_features(

accepts_per_arm_features: bool,

observation_and_action_constraint_splitter: types.Splitter,

):

if accepts_per_arm_features and (

observation_and_action_constraint_splitter is not None

):

raise NotImplementedError(

'Action masking is not allowed for Bandit problems with arm features. '

'To implement a policy that handles variable number of actions, please '

'use the `{}` feature key.'.format(

bandit_spec_utils.NUM_ACTIONS_FEATURE_KEY

)

)