text: # coding=utf-8
  # Copyright 2018 The TF-Agents Authors.
  #
  # 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.
  
  """Wraps a python metric as a TF metric."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import contextlib
  import threading
  
  import tensorflow as tf
  from tf_agents.metrics import tf_metric
  
  nest = tf.contrib.framework.nest
  
  
  @contextlib.contextmanager
  def _check_not_called_concurrently(lock):
    """Checks the returned context is not executed concurrently with any other."""
    if not lock.acquire(False):  # Non-blocking.
      raise RuntimeError(
          'Detected concurrent execution of TFPyMetric ops.')
    try:
      yield
    finally:
      lock.release()
  
  
  class TFPyMetric(tf_metric.TFStepMetric):
    """Wraps a python metric as a TF metric."""
  
    def __init__(self, py_metric, name=None, dtype=tf.float32):
      """Creates a TF metric given a py metric to wrap.
  
      Args:
        py_metric: A batched python metric to wrap.
        name: Name of the metric.
        dtype: Data type of the metric.
      """
      name = name or py_metric.name
      super(TFPyMetric, self).__init__(name=name)
      self._py_metric = py_metric
      self._dtype = dtype
      self._lock = threading.Lock()
  
    def build(self, *args, **kwargs):
      pass
  
    def call(self, trajectory):
      """Update the value of the metric using trajectory.
  
      The trajectory can be either batched or un-batched depending on
      the expected inputs for the py_metric being wrapped.
  
      Args:
        trajectory: A tf_agents.trajectory.Trajectory.
  
      Returns:
        The arguments, for easy chaining.
      """
      def _call(*flattened_trajectories):
        with _check_not_called_concurrently(self._lock):
          packed_trajectories = nest.pack_sequence_as(
              structure=(trajectory),
              flat_sequence=flattened_trajectories)
          return self._py_metric(packed_trajectories)
  
      flattened_trajectories = nest.flatten(trajectory)
      metric_op = tf.py_func(
          _call,
          flattened_trajectories,
          [],
          stateful=True,
          name='metric_call_py_func')
  
      with tf.control_dependencies([metric_op]):
        return nest.map_structure(tf.identity, trajectory)
  
    def result(self):
      def _result():
        with _check_not_called_concurrently(self._lock):
          return  self._py_metric.result()
  
      return tf.py_func(
          _result,
          [],
          self._dtype,
          stateful=True,
          name='metric_result_py_func')
  
    def reset(self):
      def _reset():
        with _check_not_called_concurrently(self._lock):
          return  self._py_metric.reset()
  
      return tf.py_func(
          _reset, [], [],
          stateful=True,
          name='metric_reset_py_func')
  