code: # Copyright 2018/2019 The RLgraph authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  from rlgraph import get_backend
  from rlgraph.utils.rlgraph_errors import RLGraphError
  from rlgraph.components.component import Component
  from rlgraph.components.explorations.epsilon_exploration import EpsilonExploration
  from rlgraph.components.common.noise_components import NoiseComponent
  from rlgraph.spaces import IntBox, FloatBox
  from rlgraph.spaces.space_utils import sanity_check_space
  from rlgraph.utils.util import convert_dtype
  from rlgraph.utils.decorators import rlgraph_api, graph_fn
  
  if get_backend() == "tf":
      import tensorflow as tf
  elif get_backend() == "pytorch":
      import torch
  
  
  class Exploration(Component):
      """
      A Component that can be plugged on top of a Policy's output to produce action choices.
      It includes noise and/or epsilon-based exploration options as well as an out-Socket to draw actions from
      the Policy's distribution - either by sampling or by deterministically choosing the max-likelihood value.
      """
      def __init__(self, epsilon_spec=None, noise_spec=None, scope="exploration", **kwargs):
          """
          Args:
              epsilon_spec (any): The spec or Component object itself to construct an EpsilonExploration Component.
              noise_spec (dict): The specification dict for a noise generator that adds noise to the NN's output.
          """
          super(Exploration, self).__init__(scope=scope, **kwargs)
  
          self.action_space = None  # The actual action space (may not have batch-rank, just the plain space)
          self.flat_action_space = None
  
          self.epsilon_exploration = None
          self.noise_component = None
  
          # For define-by-run sampling.
          self.sample_obj = None
  
          # Don't allow both epsilon and noise component
          if epsilon_spec and noise_spec:
              raise RLGraphError("Cannot use both epsilon exploration and a noise component at the same time.")
  
          # Add epsilon component.
          if epsilon_spec:
              self.epsilon_exploration = EpsilonExploration.from_spec(epsilon_spec)
              self.add_components(self.epsilon_exploration)
  
              # Define our interface.
              @rlgraph_api(component=self)
              def get_action(self, actions, time_step, use_exploration=True):
                  """
                  Action depends on time-step (e.g. epsilon-decay).
                  """
                  epsilon_decisions = self.epsilon_exploration.do_explore(actions, time_step)
                  return self._graph_fn_pick(use_exploration, epsilon_decisions, actions)
  
          # Add noise component.
          elif noise_spec:
              self.noise_component = NoiseComponent.from_spec(noise_spec)
              self.add_components(self.noise_component)
  
              @rlgraph_api(component=self)
              def get_action(self, actions, time_step=0, use_exploration=True):
                  """
                  Noise is added to the sampled action.
                  """
                  noise = self.noise_component.get_noise()
                  return self._graph_fn_add_noise(use_exploration, noise, actions)
  
          # Don't explore at all. Simple pass-through.
          else:
              @rlgraph_api(component=self)
              def get_action(self, actions, time_step=0, use_exploration=False):
                  """
                  Action is returned as is.
                  """
                  return actions
  
      def check_input_spaces(self, input_spaces, action_space=None):
          action_sample_space = input_spaces["actions"]
  
          if get_backend() == "tf":
              sanity_check_space(action_sample_space, must_have_batch_rank=True)
  
          assert action_space is not None
          self.action_space = action_space
          self.flat_action_space = action_space.flatten()
  
          if self.epsilon_exploration and self.noise_component:
              # Check again at graph creation? This is currently redundant to the check in __init__
              raise RLGraphError("Cannot use both epsilon exploration and a noise component at the same time.")
  
          if self.epsilon_exploration:
              sanity_check_space(self.action_space, must_have_categories=True, num_categories=(1, None),
                                 allowed_sub_types=[IntBox])
          elif self.noise_component:
              sanity_check_space(self.action_space, allowed_sub_types=[FloatBox])
  
      @graph_fn(flatten_ops=True, split_ops=True, add_auto_key_as_first_param=True)
      def _graph_fn_pick(self, key, use_exploration, epsilon_decisions, sample):
          """
          Exploration for discrete action spaces.
          Either pick a random action (if `use_exploration` and `epsilon_decision` are True),
              or return non-exploratory action.
  
          Args:
              use_exploration (DataOp): The master switch determining, whether to use exploration or not.
              epsilon_decisions (DataOp): The bool coming from the epsilon-exploration component specifying
                  whether to use exploration or not (per batch item).
              sample (DataOp): The output from a distribution's "sample_deterministic" OR "sample_stochastic".
  
          Returns:
              DataOp: The DataOp representing the action. This will match the shape of self.action_space.
          """
          if get_backend() == "tf":
              if use_exploration is False:
                  return sample
              else:
                  random_actions = tf.random_uniform(
                      shape=tf.shape(sample),
                      maxval=self.flat_action_space[key].num_categories,
                      dtype=convert_dtype("int")
                  )
  
                  return tf.where(
                      # `use_exploration` given as actual bool or as tensor?
                      condition=epsilon_decisions if use_exploration is True else tf.logical_and(
                          use_exploration, epsilon_decisions
                      ),
                      x=random_actions,
                      y=sample
                  )
          elif get_backend() == "pytorch":
              # N.b. different order versus TF because we dont want to execute the sampling below.
              if bool(use_exploration) is False:
                  return sample
  
              if self.sample_obj is None:
                  # Don't create new sample objects very time.
                  self.sample_obj = torch.distributions.Uniform(0, self.flat_action_space[key].num_categories)
  
              random_actions = self.sample_obj.sample(sample.shape).int()
              if bool(use_exploration) is True:
                  return torch.where(epsilon_decisions, random_actions, sample)
              else:
                  if not isinstance(use_exploration, torch.ByteTensor):
                      use_exploration = use_exploration.byte()
                  if not isinstance(epsilon_decisions, torch.ByteTensor):
                      epsilon_decisions = epsilon_decisions.byte()
                  return torch.where(use_exploration & epsilon_decisions, random_actions, sample)
  
      @graph_fn
      def _graph_fn_add_noise(self, use_exploration, noise, sample):
          """
          Noise for continuous action spaces.
          Return the action with added noise.
  
          Args:
              use_exploration (DataOp): The master switch determining, whether to add noise or not.
              noise (DataOp): The noise coming from the noise component.
              sample (DataOp): The output from a distribution's "sample_deterministic" or "sample_stochastic" API-method.
  
          Returns:
              DataOp: The DataOp representing the action. This will match the shape of self.action_space.
  
          """
          if get_backend() == "tf":
              return tf.cond(
                  use_exploration, true_fn=lambda: sample + noise, false_fn=lambda: sample
              )
          elif get_backend() == "pytorch":
              if use_exploration:
                  return sample + noise
              else:
                  return sample
  
apis: tensorflow.cond tensorflow.shape torch.where torch.distributions.Uniform tensorflow.logical_and
api_extract: [(30, 'rlgraph.get_backend', 'get_backend', ([], {}), False, 'from rlgraph import get_backend\n'), (119, 'rlgraph.utils.decorators.graph_fn', 'graph_fn', ([], {'flatten_ops': '(True)', 'split_ops': '(True)', 'add_auto_key_as_first_param': '(True)'}), False, 'from rlgraph.utils.decorators import rlgraph_api, graph_fn\n'), (32, 'rlgraph.get_backend', 'get_backend', ([], {}), False, 'from rlgraph import get_backend\n'), (61, 'rlgraph.utils.rlgraph_errors.RLGraphError', 'RLGraphError', (['"""Cannot use both epsilon exploration and a noise component at the same time."""'], {}), False, 'from rlgraph.utils.rlgraph_errors import RLGraphError\n'), (65, 'rlgraph.components.explorations.epsilon_exploration.EpsilonExploration.from_spec', 'EpsilonExploration.from_spec', (['epsilon_spec'], {}), False, 'from rlgraph.components.explorations.epsilon_exploration import EpsilonExploration\n'), (69, 'rlgraph.utils.decorators.rlgraph_api', 'rlgraph_api', ([], {'component': 'self'}), False, 'from rlgraph.utils.decorators import rlgraph_api, graph_fn\n'), (102, 'rlgraph.get_backend', 'get_backend', ([], {}), False, 'from rlgraph import get_backend\n'), (103, 'rlgraph.spaces.space_utils.sanity_check_space', 'sanity_check_space', (['action_sample_space'], {'must_have_batch_rank': '(True)'}), False, 'from rlgraph.spaces.space_utils import sanity_check_space\n'), (111, 'rlgraph.utils.rlgraph_errors.RLGraphError', 'RLGraphError', (['"""Cannot use both epsilon exploration and a noise component at the same time."""'], {}), False, 'from rlgraph.utils.rlgraph_errors import RLGraphError\n'), (114, 'rlgraph.spaces.space_utils.sanity_check_space', 'sanity_check_space', (['self.action_space'], {'must_have_categories': '(True)', 'num_categories': '(1, None)', 'allowed_sub_types': '[IntBox]'}), False, 'from rlgraph.spaces.space_utils import sanity_check_space\n'), (135, 'rlgraph.get_backend', 'get_backend', ([], {}), False, 'from rlgraph import get_backend\n'), (187, 'rlgraph.get_backend', 'get_backend', ([], {}), False, 'from rlgraph import get_backend\n'), (188, 'tensorflow.cond', 'tf.cond', (['use_exploration'], {'true_fn': '(lambda : sample + noise)', 'false_fn': '(lambda : sample)'}), True, 'import tensorflow as tf\n'), (79, 'rlgraph.components.common.noise_components.NoiseComponent.from_spec', 'NoiseComponent.from_spec', (['noise_spec'], {}), False, 'from rlgraph.components.common.noise_components import NoiseComponent\n'), (82, 'rlgraph.utils.decorators.rlgraph_api', 'rlgraph_api', ([], {'component': 'self'}), False, 'from rlgraph.utils.decorators import rlgraph_api, graph_fn\n'), (92, 'rlgraph.utils.decorators.rlgraph_api', 'rlgraph_api', ([], {'component': 'self'}), False, 'from rlgraph.utils.decorators import rlgraph_api, graph_fn\n'), (117, 'rlgraph.spaces.space_utils.sanity_check_space', 'sanity_check_space', (['self.action_space'], {'allowed_sub_types': '[FloatBox]'}), False, 'from rlgraph.spaces.space_utils import sanity_check_space\n'), (153, 'rlgraph.get_backend', 'get_backend', ([], {}), False, 'from rlgraph import get_backend\n'), (191, 'rlgraph.get_backend', 'get_backend', ([], {}), False, 'from rlgraph import get_backend\n'), (160, 'torch.distributions.Uniform', 'torch.distributions.Uniform', (['(0)', 'self.flat_action_space[key].num_categories'], {}), False, 'import torch\n'), (164, 'torch.where', 'torch.where', (['epsilon_decisions', 'random_actions', 'sample'], {}), False, 'import torch\n'), (170, 'torch.where', 'torch.where', (['(use_exploration & epsilon_decisions)', 'random_actions', 'sample'], {}), False, 'import torch\n'), (140, 'tensorflow.shape', 'tf.shape', (['sample'], {}), True, 'import tensorflow as tf\n'), (142, 'rlgraph.utils.util.convert_dtype', 'convert_dtype', (['"""int"""'], {}), False, 'from rlgraph.utils.util import convert_dtype\n'), (147, 'tensorflow.logical_and', 'tf.logical_and', (['use_exploration', 'epsilon_decisions'], {}), True, 'import tensorflow as tf\n')]