class_docstr: 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.
class_docstr_tok: ["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","."]
docstr: 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.
docstr_tok: ["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","."]
code: 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",
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import gin
  import tensorflow
  from tf_agents.bandits.environments import non_stationary_stochastic_environment
  from tf_agents.specs import tensor_spec
  from tf_agents.typing import types
label_desc: Use image node_id 1 to create a new PiecewiseStationaryDynamics object from inherited base classes: nsse with example: obj = PiecewiseStationaryDynamics(observation_distribution, interval_distribution, observation_to_reward_distribution, additive_reward_distribution)