class_docstr: Stationary Stochastic Bandit environment with per-arm features.
docstr: Initializes the environment.
  
  In each round, global context is generated by global_context_sampling_fn,
  per-arm contexts are generated by arm_context_sampling_fn. The reward_fn
  function takes the concatenation of a global and a per-arm feature, and
  outputs a possibly random reward.
  In case `num_action_fn` is specified, the number of actions will be dynamic
  and a `num_actions` feature key indicates the number of actions in any given
  sample.
  
  Example:
    def global_context_sampling_fn():
      return np.random.randint(0, 10, [2])  # 2-dimensional global features.
  
    def arm_context_sampling_fn():
      return np.random.randint(-3, 4, [3])  # 3-dimensional arm features.
  
    def reward_fn(x):
      return sum(x)
  
    def num_actions_fn():
      return np.random.randint(2, 6)
  
    env = StationaryStochasticPerArmPyEnvironment(global_context_sampling_fn,
                                                  arm_context_sampling_fn,
                                                  5,
                                                  reward_fn,
                                                  num_actions_fn)
  
  Args:
    global_context_sampling_fn: A function that outputs a random 1d array or
      list of ints or floats. This output is the global context. Its shape and
      type must be consistent across calls.
    arm_context_sampling_fn: A function that outputs a random 1 array or list
      of ints or floats (same type as the output of
      `global_context_sampling_fn`). This output is the per-arm context. Its
      shape must be consistent across calls.
    max_num_actions: (int) the maximum number of actions in every sample. If
      `num_actions_fn` is not set, this many actions are available in every
      time step.
    reward_fn: A function that generates a reward when called with an
      observation.
    num_actions_fn: If set, it should be a function that outputs a single
      integer specifying the number of actions for a given time step. The
      value output by this function will be capped between 1 and
      `max_num_actions`. The number of actions will be encoded in the
      observation by the feature key `num_actions`.
    batch_size: The batch size.
    name: The name of this environment instance.
docstr_tok: ["Initializes","the","environment",".","In","each","round",",","global","context","is","generated","by","global_context_sampling_fn",",","per-arm","contexts","are","generated","by","arm_context_sampling_fn",".","The","reward_fn","function","takes","the","concatenation","of","a","global","and","a","per-arm","feature",",","and","outputs","a","possibly","random","reward",".","In","case","`","num_action_fn","`","is","specified",",","the","number","of","actions","will","be","dynamic","and","a","`","num_actions","`","feature","key","indicates","the","number","of","actions","in","any","given","sample",".","Example",":","def","global_context_sampling_fn","(",")",":","return","np.random.randint","(","0",",","10",",","[","2","]",")","#","2-dimensional","global","features",".","def","arm_context_sampling_fn","(",")",":","return","np.random.randint","(","-3",",","4",",","[","3","]",")","#","3-dimensional","arm","features",".","def","reward_fn","(","x",")",":","return","sum","(","x",")","def","num_actions_fn","(",")",":","return","np.random.randint","(","2",",","6",")","env","=","StationaryStochasticPerArmPyEnvironment","(","global_context_sampling_fn",",","arm_context_sampling_fn",",","5",",","reward_fn",",","num_actions_fn",")","Args",":","global_context_sampling_fn",":","A","function","that","outputs","a","random","1d","array","or","list","of","ints","or","floats",".","This","output","is","the","global","context",".","Its","shape","and","type","must","be","consistent","across","calls",".","arm_context_sampling_fn",":","A","function","that","outputs","a","random","1","array","or","list","of","ints","or","floats","(","same","type","as","the","output","of","`","global_context_sampling_fn","`",")",".","This","output","is","the","per-arm","context",".","Its","shape","must","be","consistent","across","calls",".","max_num_actions",":","(","int",")","the","maximum","number","of","actions","in","every","sample",".","If","`","num_actions_fn","`","is","not","set",",","this","many","actions","are","available","in","every","time","step",".","reward_fn",":","A","function","that","generates","a","reward","when","called","with","an","observation",".","num_actions_fn",":","If","set",",","it","should","be","a","function","that","outputs","a","single","integer","specifying","the","number","of","actions","for","a","given","time","step",".","The","value","output","by","this","function","will","be","capped","between","1","and","`","max_num_actions","`",".","The","number","of","actions","will","be","encoded","in","the","observation","by","the","feature","key","`","num_actions","`",".","batch_size",":","The","batch","size",".","name",":","The","name","of","this","environment","instance","."]
code: def __init__(
      self,
      global_context_sampling_fn: Callable[[], types.Array],
      arm_context_sampling_fn: Callable[[], types.Array],
      max_num_actions: int,
      reward_fn: Callable[[types.Array], Sequence[float]],
      num_actions_fn: Optional[Callable[[], int]] = None,
      batch_size: Optional[int] = 1,
      name: Optional[Text] = "stationary_stochastic_per_arm",
  ):
      """Initializes the environment.
  
      In each round, global context is generated by global_context_sampling_fn,
      per-arm contexts are generated by arm_context_sampling_fn. The reward_fn
      function takes the concatenation of a global and a per-arm feature, and
      outputs a possibly random reward.
      In case `num_action_fn` is specified, the number of actions will be dynamic
      and a `num_actions` feature key indicates the number of actions in any given
      sample.
  
      Example:
        def global_context_sampling_fn():
          return np.random.randint(0, 10, [2])  # 2-dimensional global features.
  
        def arm_context_sampling_fn():
          return np.random.randint(-3, 4, [3])  # 3-dimensional arm features.
  
        def reward_fn(x):
          return sum(x)
  
        def num_actions_fn():
          return np.random.randint(2, 6)
  
        env = StationaryStochasticPerArmPyEnvironment(global_context_sampling_fn,
                                                      arm_context_sampling_fn,
                                                      5,
                                                      reward_fn,
                                                      num_actions_fn)
  
      Args:
        global_context_sampling_fn: A function that outputs a random 1d array or
          list of ints or floats. This output is the global context. Its shape and
          type must be consistent across calls.
        arm_context_sampling_fn: A function that outputs a random 1 array or list
          of ints or floats (same type as the output of
          `global_context_sampling_fn`). This output is the per-arm context. Its
          shape must be consistent across calls.
        max_num_actions: (int) the maximum number of actions in every sample. If
          `num_actions_fn` is not set, this many actions are available in every
          time step.
        reward_fn: A function that generates a reward when called with an
          observation.
        num_actions_fn: If set, it should be a function that outputs a single
          integer specifying the number of actions for a given time step. The
          value output by this function will be capped between 1 and
          `max_num_actions`. The number of actions will be encoded in the
          observation by the feature key `num_actions`.
        batch_size: The batch size.
        name: The name of this environment instance.
      """
      self._global_context_sampling_fn = global_context_sampling_fn
      self._arm_context_sampling_fn = arm_context_sampling_fn
      self._max_num_actions = max_num_actions
      self._reward_fn = reward_fn
      self._batch_size = batch_size
      self._num_actions_fn = num_actions_fn
  
      observation_spec = {
          GLOBAL_KEY: array_spec.ArraySpec.from_array(
              global_context_sampling_fn()
          ),
          PER_ARM_KEY: array_spec.add_outer_dims_nest(
              array_spec.ArraySpec.from_array(
                  arm_context_sampling_fn()
              ),
              (max_num_actions,),
          ),
      }
      if self._num_actions_fn is not None:
          num_actions_spec = array_spec.BoundedArraySpec(
              shape=(),
              dtype=np.dtype(type(self._num_actions_fn())),
              minimum=1,
              maximum=max_num_actions,
          )
          observation_spec.update({NUM_ACTIONS_KEY: num_actions_spec})
  
      action_spec = array_spec.BoundedArraySpec(
          shape=(),
          dtype=np.int32,
          minimum=0,
          maximum=max_num_actions - 1,
          name="action",
      )
  
      super(StationaryStochasticPerArmPyEnvironment, self).__init__(
          observation_spec, action_spec, name=name
      )
  
file_path: agents/tf_agents/bandits/environments/stationary_stochastic_per_arm_py_environment.py
imports: from typing import Callable, Optional, Sequence, Text
  import gin
  import numpy
  from tf_agents.bandits.environments import bandit_py_environment
  from tf_agents.bandits.specs import utils
  from tf_agents.specs import array_spec
  from tf_agents.typing import types
label_desc: Use image node_id 1 to create a new StationaryStochasticPerArmPyEnvironment object from inherited base classes: bandit_py_environment with example: obj = StationaryStochasticPerArmPyEnvironment(global_context_sampling_fn, arm_context_sampling_fn, max_num_actions, reward_fn, num_actions_fn, batch_size, name)