class_docstr: An environment wrapper that adds action masks to observations.
class_docstr_tok: ["An","environment","wrapper","that","adds","action","masks","to","observations","."]
code: def _observe(self):
      # pylint: disable=protected-access
      original_observation = self._original_environment._observe()
      mask = tfd.Bernoulli(self._action_probability).sample(
          sample_shape=[self._batch_size, self._num_actions]
      )
      mask = _maybe_add_one_action(mask)
      tf.compat.v1.assign(self._current_mask, mask)
  
      return self._action_constraint_join_fn(original_observation, mask)
  
file_path: agents/tf_agents/bandits/environments/bernoulli_action_mask_tf_environment.py
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 5 for calling the BernoulliActionMaskTFEnvironment obj's underlying member method code with example usage: obj._observe() and returns: self