text: # 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.
  
  """Bandit piecewise linear stationary environment."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import gin
  import tensorflow as tf  # pylint: disable=g-explicit-tensorflow-version-import
  
  from tf_agents.bandits.environments import non_stationary_stochastic_environment as nsse
  from tf_agents.specs import tensor_spec
  from tf_agents.typing import types
  
  
  def _raise_batch_shape_error(tensor_name, batch_shape):
    raise ValueError('`{tensor_name}` must have batch shape with length 1; '
                     'got {batch_shape}.'.format(
                         tensor_name=tensor_name,
                         batch_shape=batch_shape))
  
  
  @gin.configurable
  class PiecewiseStationaryDynamics(nsse.EnvironmentDynamics):
    """A piecewise stationary environment dynamics.
  
    This is a piecewise stationary environment which computes rewards as:
  
    rewards(t) = observation(t) * observation_to_reward(i) + additive_reward(i)
  
    where t is the environment time (env_time) and i is the index of each piece.
    The environment time is incremented after the reward is computed while the
    piece index is incremented at the end of the time interval. The parameters
    observation_to_reward(i), additive_reward(i), and the length of interval, are
    drawn from given distributions at the beginning of each temporal interval.
    """
  
    def __init__(self,
                 observation_distribution: types.Distribution,
                 interval_distribution: types.Distribution,
                 observation_to_reward_distribution: types.Distribution,
                 additive_reward_distribution: types.Distribution):
      """Initialize the parameters of the piecewise dynamics.
  
      Args:
        observation_distribution: A distribution from tfp.distributions with shape
          `[batch_size, observation_dim]` Note that the values of `batch_size` and
          `observation_dim` are deduced from the distribution.
        interval_distribution: A scalar distribution from `tfp.distributions`. The
          value is casted to `int64` to update the time range.
        observation_to_reward_distribution: A distribution from
          `tfp.distributions` with shape `[observation_dim, num_actions]`. The
          value `observation_dim` must match the second dimension of
          `observation_distribution`.
        additive_reward_distribution: A distribution from `tfp.distributions` with
          shape `[num_actions]`. This models the non-contextual behavior of the
          bandit.
      """
      self._observation_distribution = observation_distribution
      self._interval_distribution = interval_distribution
      self._observation_to_reward_distribution = (
          observation_to_reward_distribution)
      self._additive_reward_distribution = additive_reward_distribution
  
      observation_batch_shape = observation_distribution.batch_shape
      reward_batch_shape = additive_reward_distribution.batch_shape
  
      if observation_batch_shape.rank != 2:
        _raise_batch_shape_error(
            'observation_distribution', observation_batch_shape)
  
      if reward_batch_shape.rank != 1:
        _raise_batch_shape_error(
            'additive_reward_distribution', reward_batch_shape)
  
      if additive_reward_distribution.dtype != tf.float32:
        raise ValueError('Reward  must have dtype float32; got {}'.format(
            self._reward.dtype))
  
      expected_observation_to_reward_shape = [
          tf.compat.dimension_value(
              self._observation_distribution.batch_shape[1:]),
          tf.compat.dimension_value(
              self._additive_reward_distribution.batch_shape[0])]
  
      observation_to_reward_shape = [
          tf.compat.dimension_value(x)
          for x in observation_to_reward_distribution.batch_shape]
  
      if (observation_to_reward_shape !=
          expected_observation_to_reward_shape):
        raise ValueError(
            'Observation to reward has {} as expected shape; got {}'.format(
                observation_to_reward_shape,
                expected_observation_to_reward_shape))
  
      self._current_interval = tf.compat.v2.Variable(
          tf.cast(interval_distribution.sample(), dtype=tf.int64),
          dtype=tf.int64, name='interval')
      self._current_observation_to_reward = tf.compat.v2.Variable(
          observation_to_reward_distribution.sample(),
          dtype=tf.float32,
          name='observation_to_reward')
      self._current_additive_reward = tf.compat.v2.Variable(
          additive_reward_distribution.sample(),
          dtype=tf.float32,
          name='additive_reward')
  
    @property
    def batch_size(self) -> types.Int:
      return tf.compat.dimension_value(
          self._observation_distribution.batch_shape[0])
  
    @property
    def observation_spec(self) -> types.TensorSpec:
      return tensor_spec.TensorSpec(
          shape=self._observation_distribution.batch_shape[1:],
          dtype=self._observation_distribution.dtype,
          name='observation_spec')
  
    @property
    def action_spec(self) -> types.BoundedTensorSpec:
      return tensor_spec.BoundedTensorSpec(
          shape=(),
          dtype=tf.int32,
          minimum=0,
          maximum=tf.compat.dimension_value(
              self._additive_reward_distribution.batch_shape[0]) - 1,
          name='action')
  
    def observation(self, unused_t) -> types.NestedTensor:
      return self._observation_distribution.sample()
  
    def reward(self,
               observation: types.NestedTensor,
               t: types.Int) -> types.NestedTensor:
      def same_interval_parameters():
        """Returns the parameters of the current piece.
  
        Returns:
          The pair of `tf.Tensor` `(observation_to_reward, additive_reward)`.
        """
        return [self._current_observation_to_reward,
                self._current_additive_reward]
  
      def new_interval_parameters():
        """Update and returns the piece parameters.
  
        Returns:
          The pair of `tf.Tensor` `(observation_to_reward, additive_reward)`.
        """
        tf.compat.v1.assign_add(
            self._current_interval,
            tf.cast(self._interval_distribution.sample(), dtype=tf.int64))
        tf.compat.v1.assign(self._current_observation_to_reward,
                            self._observation_to_reward_distribution.sample())
        tf.compat.v1.assign(self._current_additive_reward,
                            self._additive_reward_distribution.sample())
  
        return [self._current_observation_to_reward,
                self._current_additive_reward]
  
      observation_to_reward, additive_reward = tf.cond(
          t < self._current_interval,
          same_interval_parameters,
          new_interval_parameters)
  
      reward = (tf.matmul(observation, observation_to_reward) +
                tf.reshape(additive_reward, [1, -1]))
      return reward
  
    @gin.configurable
    def compute_optimal_reward(
        self, observation: types.NestedTensor) -> types.NestedTensor:
      deterministic_reward = tf.matmul(
          observation, self._current_observation_to_reward)
      optimal_action_reward = tf.reduce_max(deterministic_reward, axis=-1)
      return optimal_action_reward
  
    @gin.configurable
    def compute_optimal_action(
        self, observation: types.NestedTensor) -> types.NestedTensor:
      deterministic_reward = tf.matmul(
          observation, self._current_observation_to_reward)
      optimal_action = tf.argmax(
          deterministic_reward, axis=-1, output_type=tf.int32)
      return optimal_action
  
  
  @gin.configurable
  class PiecewiseStochasticEnvironment(nsse.NonStationaryStochasticEnvironment):
    """Implements a piecewise stationary linear environment."""
  
    def __init__(self,
                 observation_distribution: types.Distribution,
                 interval_distribution: types.Distribution,
                 observation_to_reward_distribution: types.Distribution,
                 additive_reward_distribution: types.Distribution):
      """Initialize the environment with the dynamics parameters.
  
      Args:
        observation_distribution: A distribution from `tfp.distributions` with
          shape `[batch_size, observation_dim]`. Note that the values of
          `batch_size` and `observation_dim` are deduced from the distribution.
        interval_distribution: A scalar distribution from `tfp.distributions`. The
          value is casted to `int64` to update the time range.
        observation_to_reward_distribution: A distribution from
          `tfp.distributions` with shape `[observation_dim, num_actions]`. The
          value `observation_dim` must match the second dimension of
          `observation_distribution`.
        additive_reward_distribution: A distribution from `tfp.distributions` with
          shape `[num_actions]`. This models the non-contextual behavior of the
          bandit.
      """
      super(PiecewiseStochasticEnvironment, self).__init__(
          PiecewiseStationaryDynamics(
              observation_distribution,
              interval_distribution,
              observation_to_reward_distribution,
              additive_reward_distribution))
  