class_docstr: Stationary Stochastic Bandit environment with per-arm features.
code: def _apply_action(self, action: np.ndarray) -> types.Array:
      if action.shape[0] != 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 = range(self.batch_size)
      arm_obs = self._observation[PER_ARM_KEY][
          batch_size_range, action, :
      ]  # pytype: disable=attribute-error  # trace-all-classes
      reward = np.stack(
          [
              self._reward_fn(
                  np.concatenate((global_obs[b, :], arm_obs[b, :]))
              )
              for b in batch_size_range
          ]
      )
      return reward
  
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 5 for calling the StationaryStochasticPerArmPyEnvironment obj's underlying member method code with example usage: obj._apply_action(action) and returns: reward