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","."]
code: def _apply_action(self, action: types.Array) -> types.Array:
      if len(action) != self.batch_size:
          raise ValueError("Number of actions must match batch size.")
      global_obs = self._observation[
          GLOBAL_KEY
      ]  # pytype: disable=attribute-error  # trace-all-classes
      batch_size_range = list(range(self.batch_size))
      arm_obs = tf.nest.map_structure(
          lambda x: x[batch_size_range, action, :],
          self._observation[
              PER_ARM_KEY
          ],  # pytype: disable=attribute-error  # trace-all-classes
      )
  
      def _get_element_from_batch(structure, index):
          return tf.nest.map_structure(lambda x: x[index], structure)
  
      reward = np.stack(
          [
              # pytype: disable=wrong-arg-count  # trace-all-classes
              self._reward_fn(
                  _get_element_from_batch(global_obs, b),
                  _get_element_from_batch(arm_obs, b),
              )
              for b in batch_size_range
              # pytype: enable=wrong-arg-count
          ]
      )
      return reward
  
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 6 for calling the StationaryStochasticStructuredPyEnvironment obj's underlying member method code with example usage: obj._apply_action(action) and returns: reward, tf