class_docstr: Implements Stationary Stochastic Bandit environments.
docstr: Initializes a Stationary Stochastic Bandit environment.
  
  In each round, context is generated by context_sampling_fn, this context is
  passed through a reward_function for each arm.
  
  Example:
    def context_sampling_fn():
      return np.random.randint(0, 10, [1, 2])  # 2-dim ints between 0 and 10
  
    def reward_fn1(x):
      return x[0]
    def reward_fn2(x):
      return x[1]
    reward_fns = [reward_fn1, reward_fn2]  # Two arms
  
    env = StationaryStochasticPyEnvironment(context_sampling_fn,
                                            reward_fns)
  
  Args:
    context_sampling_fn: A function that outputs a random 2d array or list of
      ints or floats, where the first dimension is batch size.
    reward_fns: A function that generates a (perhaps non-scalar) reward when
      called with an observation.
    constraint_fns: A function that generates a (perhaps non-scalar)
      constraint metric when called with an observation.
    batch_size: The batch size. Must match the outer dimension of the output
      of context_sampling_fn.
    name: The name of this environment instance.
docstr_tok: ["Initializes","a","Stationary","Stochastic","Bandit","environment",".","In","each","round",",","context","is","generated","by","context_sampling_fn",",","this","context","is","passed","through","a","reward_function","for","each","arm",".","Example",":","def","context_sampling_fn","(",")",":","return","np.random.randint","(","0",",","10",",","[","1",",","2","]",")","#","2-dim","ints","between","0","and","10","def","reward_fn1","(","x",")",":","return","x","[","0","]","def","reward_fn2","(","x",")",":","return","x","[","1","]","reward_fns","=","[","reward_fn1",",","reward_fn2","]","#","Two","arms","env","=","StationaryStochasticPyEnvironment","(","context_sampling_fn",",","reward_fns",")","Args",":","context_sampling_fn",":","A","function","that","outputs","a","random","2d","array","or","list","of","ints","or","floats",",","where","the","first","dimension","is","batch","size",".","reward_fns",":","A","function","that","generates","a","(","perhaps","non-scalar",")","reward","when","called","with","an","observation",".","constraint_fns",":","A","function","that","generates","a","(","perhaps","non-scalar",")","constraint","metric","when","called","with","an","observation",".","batch_size",":","The","batch","size",".","Must","match","the","outer","dimension","of","the","output","of","context_sampling_fn",".","name",":","The","name","of","this","environment","instance","."]
code: def __init__(
      self,
      context_sampling_fn: Callable[[], np.ndarray],
      reward_fns: Sequence[Callable[[np.ndarray], Sequence[float]]],
      constraint_fns: Optional[
          Sequence[Callable[[np.ndarray], Sequence[float]]]
      ] = None,
      batch_size: Optional[int] = 1,
      name: Optional[Text] = "stationary_stochastic",
  ):
      """Initializes a Stationary Stochastic Bandit environment.
  
      In each round, context is generated by context_sampling_fn, this context is
      passed through a reward_function for each arm.
  
      Example:
        def context_sampling_fn():
          return np.random.randint(0, 10, [1, 2])  # 2-dim ints between 0 and 10
  
        def reward_fn1(x):
          return x[0]
        def reward_fn2(x):
          return x[1]
        reward_fns = [reward_fn1, reward_fn2]  # Two arms
  
        env = StationaryStochasticPyEnvironment(context_sampling_fn,
                                                reward_fns)
  
      Args:
        context_sampling_fn: A function that outputs a random 2d array or list of
          ints or floats, where the first dimension is batch size.
        reward_fns: A function that generates a (perhaps non-scalar) reward when
          called with an observation.
        constraint_fns: A function that generates a (perhaps non-scalar)
          constraint metric when called with an observation.
        batch_size: The batch size. Must match the outer dimension of the output
          of context_sampling_fn.
        name: The name of this environment instance.
      """
      self._context_sampling_fn = context_sampling_fn
      self._reward_fns = reward_fns
      self._num_actions = len(reward_fns)
      self._constraint_fns = constraint_fns
      self._batch_size = batch_size
  
      action_spec = array_spec.BoundedArraySpec(
          shape=(),
          dtype=np.int32,
          minimum=0,
          maximum=self._num_actions - 1,
          name="action",
      )
  
      example_observation = self._context_sampling_fn()
      observation_spec = array_spec.ArraySpec.from_array(
          example_observation[0]
      )
      if example_observation.shape[0] != batch_size:
          raise ValueError(
              "The outer dimension of the observations should match the batch size."
          )
  
      # Figure out the reward spec.
      # If we have constraints, the reward_spec will be a nested dict with keys:
      # 'reward' and 'constraint' (defined in tf_agents.bandits.specs.utils).
      example_reward = np.asarray(reward_fns[0](example_observation[0]))
      reward_spec = array_spec.ArraySpec(
          example_reward.shape, np.float32, name="reward"
      )
      if self._constraint_fns is not None:
          example_constraint = np.asarray(
              constraint_fns[0](example_observation[0])
          )
          constraint_spec = array_spec.ArraySpec(
              example_constraint.shape, np.float32, name="constraint"
          )
          reward_spec = {
              bandits_spec_utils.REWARD_SPEC_KEY: reward_spec,
              bandits_spec_utils.CONSTRAINTS_SPEC_KEY: constraint_spec,
          }
  
      super(StationaryStochasticPyEnvironment, self).__init__(
          observation_spec, action_spec, reward_spec, name=name
      )
  
file_path: agents/tf_agents/bandits/environments/stationary_stochastic_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 StationaryStochasticPyEnvironment object from inherited base classes: bandit_py_environment with example: obj = StationaryStochasticPyEnvironment(context_sampling_fn, reward_fns, constraint_fns, batch_size, name)