# coding=utf-8

# Copyright 2020 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

#

# https://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.

"""Base class for Python metrics."""

from __future__ import absolute_import

from __future__ import division

from __future__ import print_function

import abc

from typing import Any, Optional, Sequence, Text, Union

from absl import logging

import numpy as np

import six

import tensorflow as tf # pylint: disable=g-explicit-tensorflow-version-import

from tf_agents.metrics import tf_metric

from tf_agents.trajectories import trajectory as traj

from tf_agents.typing import types

from tf_agents.utils import common

PyMetricType = types.ForwardRef('PyMetric') # pylint: disable=invalid-name

MetricType = Union[tf_metric.TFStepMetric, PyMetricType]

def run_summaries(

metrics: Sequence[PyMetricType],

session: Optional[tf.compat.v1.Session] = None,

):

"""Execute summary ops for py_metrics.

Args:

metrics: A list of py_metric.Base objects.

session: A TensorFlow session-like object. If it is not provided, it will

use the current TensorFlow session context manager.

Raises:

RuntimeError: If .tf_summaries() was not previously called on any of the

`metrics`.

AttributeError: If session is not provided and there is no default session

provided by a context manager.

"""

if session is None:

default_session = tf.compat.v1.get_default_session()

if default_session is None:

raise AttributeError(

'No TensorFlow session-like object was provided, and none '

"could be retrieved using 'tf.get_default_session()'."

)

session = default_session

for metric in metrics:

if metric.summary_op is None:

raise RuntimeError(

'metric.tf_summaries() must be called on py_metric '

'{} before attempting to run '

'summaries.'.format(metric.name)

)

summary_ops = [metric.summary_op for metric in metrics]

feed_dict = dict(

(metric.summary_placeholder, metric.result()) for metric in metrics

)

session.run(summary_ops, feed_dict=feed_dict)

@six.add_metaclass(abc.ABCMeta)

class PyMetric(tf.Module):

"""Defines the interface for metrics."""

def __init__(self, name: Text, prefix: Text = 'Metrics'):

"""Creates a metric."""

super(PyMetric, self).__init__(name)

self._prefix = prefix

self._summary_placeholder = None

self._summary_op = None

@property

def prefix(self) -> Text:

"""Prefix for the metric."""

return self._prefix

@abc.abstractmethod

def reset(self):

"""Resets internal stat gathering variables used to compute the metric."""

@abc.abstractmethod

def result(self) -> Any:

"""Evaluates the current value of the metric."""

def log(self):

tag = common.join_scope(self.prefix, self.name)

logging.info('%s', '{0} = {1}'.format(tag, self.result()))

def tf_summaries(

self,

train_step: types.Int = None,

step_metrics: Sequence[MetricType] = (),

) -> tf.Operation:

"""Build TF summary op and placeholder for this metric.

To execute the op, call py_metric.run_summaries.

Args:

train_step: Step counter for training iterations. If None, no metric is

generated against the global step.

step_metrics: Step values to plot as X axis in addition to global_step.

Returns:

The summary op.

Raises:

RuntimeError: If this method has already been called (it can only be

called once).

ValueError: If any item in step_metrics is not of type PyMetric or

tf_metric.TFStepMetric.

"""

if self.summary_op is not None:

raise RuntimeError('metric.tf_summaries() can only be called once.')

tag = common.join_scope(self.prefix, self.name)

summaries = []

summaries.append(

tf.compat.v2.summary.scalar(

name=tag, data=self.summary_placeholder, step=train_step

)

)

prefix = self.prefix

if prefix:

prefix += '_'

for step_metric in step_metrics:

# Skip plotting the metrics against itself.

if self.name == step_metric.name:

continue

step_tag = '{}vs_{}/{}'.format(prefix, step_metric.name, self.name)

if isinstance(step_metric, PyMetric):

step_tensor = step_metric.summary_placeholder

elif isinstance(step_metric, tf_metric.TFStepMetric):

step_tensor = step_metric.result()

else:

raise ValueError(

'step_metric is not PyMetric or TFStepMetric: {}'.format(

step_metric

)

)

summaries.append(

tf.compat.v2.summary.scalar(

name=step_tag, data=self.summary_placeholder, step=step_tensor

)

)

self._summary_op = tf.group(*summaries)

return self._summary_op

@property

def summary_placeholder(self) -> tf.compat.v1.placeholder:

"""TF placeholder to be used for the result of this metric."""

if self._summary_placeholder is None:

result = self.result()

if not isinstance(result, (np.ndarray, np.generic)):

result = np.array(result)

dtype = tf.as_dtype(result.dtype)

shape = result.shape

self._summary_placeholder = tf.compat.v1.placeholder(

dtype, shape=shape, name='{}_ph'.format(self.name)

)

return self._summary_placeholder

@property

def summary_op(self) -> tf.Operation:

"""TF summary op for this metric."""

return self._summary_op

@staticmethod

def aggregate(metrics: Sequence[PyMetricType]) -> types.Float:

"""Aggregates a list of metrics.

The default behaviour is to return the average of the metrics.

Args:

metrics: a list of metrics, of the same class.

Returns:

The result of aggregating this metric.

"""

return np.mean([metric.result() for metric in metrics])

def __call__(self, *args):

"""Method to update the metric contents.

To change the behavior of this function, override the call method.

Different subclasses might use this differently. For instance, the

PyStepMetric takes in a trajectory, while the CounterMetric takes no

parameters.

Args:

*args: See call method of subclass for specific arguments.

"""

self.call(*args)

class PyStepMetric(PyMetric):

"""Defines the interface for metrics that operate on trajectories."""

@abc.abstractmethod

def call(self, trajectory: traj.Trajectory):

"""Processes a trajectory to update the metric.

Args:

trajectory: A trajectory.Trajectory.

"""