class_docstr: Stationary Stochastic Bandit environment with structured features.
  
  This environment can generate global and per-arm observations of any nested
  structure.
class_docstr_tok: ["Stationary","Stochastic","Bandit","environment","with","structured","features",".","This","environment","can","generate","global","and","per-arm","observations","of","any","nested","structure","."]
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 two feature generating functions should output a single observation, not
  including either the batch_size or the number of actions.
  
  The reward_fn function takes a global and a per-arm feature, and outputs a
  possibly random reward.
  
  Example:
    def global_context_sampling_fn():
      return np.random.randint(0, 10, [2])  # 2-dimensional global features.
  
    def arm_context_sampling_fn():
      return {'armf1': np.random.randint(-3, 4, [3]),    # A dictionary of
              'armf2': np.random.randint(0, 2, [4, 5])}  # arm features.
  
    def reward_fn(global, arm):
      return sum(global) + arm['armf1'][0] + arm['armf2'][3, 3]
  
    env = StationaryStochasticPyEnvironment(global_context_sampling_fn,
                                            arm_context_sampling_fn,
                                            5,
                                            reward_fn,
                                            batch_size=5)
  
  Args:
    global_context_sampling_fn: A function that outputs a possibly nested
      structure of features. This output is the global context. Its shapes and
      types must be consistent accross calls.
    arm_context_sampling_fn: A function that outputs a possibly nested
      structure of features. This output is the per-arm context. Its shapes
      must be consistent accross calls.
    num_actions: (int) the number of actions in every sample.
    reward_fn: A function that generates a reward when called with a global
      and a per-arm observation.
    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","two","feature","generating","functions","should","output","a","single","observation",",","not","including","either","the","batch_size","or","the","number","of","actions",".","The","reward_fn","function","takes","a","global","and","a","per-arm","feature",",","and","outputs","a","possibly","random","reward",".","Example",":","def","global_context_sampling_fn","(",")",":","return","np.random.randint","(","0",",","10",",","[","2","]",")","#","2-dimensional","global","features",".","def","arm_context_sampling_fn","(",")",":","return","{","'armf1","'",":","np.random.randint","(","-3",",","4",",","[","3","]",")",",","#","A","dictionary","of","'armf2","'",":","np.random.randint","(","0",",","2",",","[","4",",","5","]",")","}","#","arm","features",".","def","reward_fn","(","global",",","arm",")",":","return","sum","(","global",")","+","arm","[","'armf1","'","]","[","0","]","+","arm","[","'armf2","'","]","[","3",",","3","]","env","=","StationaryStochasticPyEnvironment","(","global_context_sampling_fn",",","arm_context_sampling_fn",",","5",",","reward_fn",",","batch_size=5",")","Args",":","global_context_sampling_fn",":","A","function","that","outputs","a","possibly","nested","structure","of","features",".","This","output","is","the","global","context",".","Its","shapes","and","types","must","be","consistent","accross","calls",".","arm_context_sampling_fn",":","A","function","that","outputs","a","possibly","nested","structure","of","features",".","This","output","is","the","per-arm","context",".","Its","shapes","must","be","consistent","accross","calls",".","num_actions",":","(","int",")","the","number","of","actions","in","every","sample",".","reward_fn",":","A","function","that","generates","a","reward","when","called","with","a","global","and","a","per-arm","observation",".","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],
      num_actions: int,
      reward_fn: Callable[[types.Array], Sequence[float]],
      batch_size: Optional[int] = 1,
      name: Optional[Text] = "stationary_stochastic_structured",
  ):
      """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 two feature generating functions should output a single observation, not
      including either the batch_size or the number of actions.
  
      The reward_fn function takes a global and a per-arm feature, and outputs a
      possibly random reward.
  
      Example:
        def global_context_sampling_fn():
          return np.random.randint(0, 10, [2])  # 2-dimensional global features.
  
        def arm_context_sampling_fn():
          return {'armf1': np.random.randint(-3, 4, [3]),    # A dictionary of
                  'armf2': np.random.randint(0, 2, [4, 5])}  # arm features.
  
        def reward_fn(global, arm):
          return sum(global) + arm['armf1'][0] + arm['armf2'][3, 3]
  
        env = StationaryStochasticPyEnvironment(global_context_sampling_fn,
                                                arm_context_sampling_fn,
                                                5,
                                                reward_fn,
                                                batch_size=5)
  
      Args:
        global_context_sampling_fn: A function that outputs a possibly nested
          structure of features. This output is the global context. Its shapes and
          types must be consistent accross calls.
        arm_context_sampling_fn: A function that outputs a possibly nested
          structure of features. This output is the per-arm context. Its shapes
          must be consistent accross calls.
        num_actions: (int) the number of actions in every sample.
        reward_fn: A function that generates a reward when called with a global
          and a per-arm observation.
        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._num_actions = num_actions
      self._reward_fn = reward_fn
      self._batch_size = batch_size
  
      global_example = global_context_sampling_fn()
      arm_example = arm_context_sampling_fn()
      observation_spec = {
          GLOBAL_KEY: tf.nest.map_structure(
              array_spec.ArraySpec.from_array, global_example
          ),
          PER_ARM_KEY: array_spec.add_outer_dims_nest(
              tf.nest.map_structure(
                  array_spec.ArraySpec.from_array, arm_example
              ),
              (num_actions,),
          ),
      }
  
      action_spec = array_spec.BoundedArraySpec(
          shape=(),
          dtype=np.int32,
          minimum=0,
          maximum=num_actions - 1,
          name="action",
      )
  
      super(StationaryStochasticStructuredPyEnvironment, self).__init__(
          observation_spec, action_spec, name=name
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from typing import Callable, Optional, Sequence, Text
  import gin
  import numpy
  import tensorflow
  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
  from tf_agents.utils import nest_utils
label_desc: Use image node_id 1 to create a new StationaryStochasticStructuredPyEnvironment object from inherited base classes: bandit_py_environment with example: obj = StationaryStochasticStructuredPyEnvironment(global_context_sampling_fn, arm_context_sampling_fn, num_actions, reward_fn, batch_size, name)