class_docstr: An environment wrapper that adds action masks to observations.
class_docstr_tok: ["An","environment","wrapper","that","adds","action","masks","to","observations","."]
docstr: Initializes a `BernoulliActionMaskTFEnvironment`.
  
  Args:
    original_environment: Instance of `BanditTFEnvironment`. This environment
      will be wrapped.
    action_constraint_join_fn: A function that joins the osbervation from the
      original environment with the generated masks.
    action_probability: The probability that any action in the action space is
      allower by the generated mask.
docstr_tok: ["Initializes","a","`","BernoulliActionMaskTFEnvironment","`",".","Args",":","original_environment",":","Instance","of","`","BanditTFEnvironment","`",".","This","environment","will","be","wrapped",".","action_constraint_join_fn",":","A","function","that","joins","the","osbervation","from","the","original","environment","with","the","generated","masks",".","action_probability",":","The","probability","that","any","action","in","the","action","space","is","allower","by","the","generated","mask","."]
code: def __init__(
      self,
      original_environment: bandit_tf_environment.BanditTFEnvironment,
      action_constraint_join_fn: Callable[
          [types.TensorSpec, types.TensorSpec], types.TensorSpec
      ],
      action_probability: float,
  ):
      """Initializes a `BernoulliActionMaskTFEnvironment`.
  
      Args:
        original_environment: Instance of `BanditTFEnvironment`. This environment
          will be wrapped.
        action_constraint_join_fn: A function that joins the osbervation from the
          original environment with the generated masks.
        action_probability: The probability that any action in the action space is
          allower by the generated mask.
      """
      self._original_environment = original_environment
      assert isinstance(
          original_environment,
          bandit_tf_environment.BanditTFEnvironment,
      ), "The wrapped environment needs to be a `BanditTFEnvironment`."
      self._action_constraint_join_fn = action_constraint_join_fn
      self._action_probability = action_probability
      self._batch_size = self._original_environment.batch_size
      action_spec = self._original_environment.action_spec()
      observation_spec_without_mask = (
          self._original_environment.time_step_spec().observation
      )
      self._num_actions = (
          policy_utilities.get_num_actions_from_tensor_spec(action_spec)
      )
  
      mask_spec = tf.TensorSpec([self._num_actions], dtype=tf.int32)
      joined_observation_spec = self._action_constraint_join_fn(
          observation_spec_without_mask, mask_spec
      )
      time_step_spec = ts.time_step_spec(joined_observation_spec)
  
      self._current_mask = tf.compat.v2.Variable(
          tf.ones([self.batch_size, self._num_actions], dtype=tf.int32)
      )
  
      super(BernoulliActionMaskTFEnvironment, self).__init__(
          time_step_spec=time_step_spec,
          action_spec=action_spec,
          batch_size=self._batch_size,
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from typing import Callable
  import gin
  import tensorflow
  import tensorflow_probability
  from tf_agents.bandits.environments import bandit_tf_environment
  from tf_agents.policies import utils
  from tf_agents.trajectories import time_step
  from tf_agents.typing import types
  from tf_agents.utils import common
label_desc: Use image node_id 1 to create a new BernoulliActionMaskTFEnvironment object from inherited base classes: bandit_tf_environment with example: obj = BernoulliActionMaskTFEnvironment(original_environment, action_constraint_join_fn, action_probability)