# 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.

"""TensorFlow Policies API."""

from __future__ import absolute_import

from __future__ import division

from __future__ import print_function

import abc

from typing import Optional, Sequence, Text

import six

import tensorflow as tf

import tensorflow_probability as tfp

from tf_agents.distributions import reparameterized_sampling

from tf_agents.specs import tensor_spec

from tf_agents.trajectories import policy_step

from tf_agents.trajectories import time_step as ts

from tf_agents.trajectories import trajectory

from tf_agents.typing import types

from tf_agents.utils import common

from tf_agents.utils import nest_utils

tfd = tfp.distributions

@six.add_metaclass(abc.ABCMeta)

class TFPolicy(tf.Module):

"""Abstract base class for TF Policies.

The Policy represents a mapping from `time_steps` recieved from the

environment to `actions` that can be applied to the environment.

Agents expose two policies. A `policy` meant for deployment and evaluation,

and a `collect_policy` for collecting data from the environment. The

`collect_policy` is usually stochastic for exploring the environment better

and may log auxilliary information such as log probabilities required for

training as well. `Policy` objects can also be created directly by the users

without using an `Agent`.

The main methods of TFPolicy are:

* `action`: Maps a `time_step` from the environment to an action.

* `distribution`: Maps a `time_step` to a distribution over actions.

* `get_initial_state`: Generates the initial state for stateful policies, e.g.

RNN/LSTM policies.

Example usage:

```

env = SomeTFEnvironment()

policy = TFRandomPolicy(env.time_step_spec(), env.action_spec())

# Or policy = agent.policy or agent.collect_policy

policy_state = policy.get_initial_state(env.batch_size)

time_step = env.reset()

while not time_step.is_last():

policy_step = policy.action(time_step, policy_state)

time_step = env.step(policy_step.action)

policy_state = policy_step.state

# policy_step.info may contain side info for logging, such as action log

# probabilities.

```

Policies can be saved to disk as SavedModels (see policy_saver.py and

policy_loader.py) or as TF Checkpoints.

A `PyTFEagerPolicy` can be used to wrap a `TFPolicy` so that it works with

`PyEnvironment`s.

**NOTE**: For API consistency, subclasses are not allowed to override public

methods of `TFPolicy` class. Instead, they may implement the protected methods

including `_get_initial_state`, `_action`, and `_distribution`. This

public-calls-private convention allowed this base class to do things like

properly add `spec` and shape checks, which provide users an easier experience

when debugging their environments and networks.

For researchers, and those developing new Policies, the `TFPolicy` base class

constructor also accept a `validate_args` parameter. If `False`, this

disables all spec structure, dtype, and shape checks in the public methods of

these classes. It allows algorithm developers to iterate and try different

input and output structures without worrying about overly restrictive

requirements, or input and output states being in a certain format. However,

*disabling argument validation* can make it very hard to identify structural

input or algorithmic errors; and should not be done for final, or

production-ready, Policies. In addition to having implementations that may

disagree with specs, this mean that the resulting Policy may no longer

interact well with other parts of TF-Agents. Examples include impedance

mismatches with Actor/Learner APIs, replay buffers, and the model export

functionality in `PolicySaver.

"""

# TODO(b/127327645) Remove this attribute.

# This attribute allows subclasses to back out of automatic tf.function

# attribute inside TF1 (for autodeps).

_enable_functions = True

def __init__(

self,

time_step_spec: ts.TimeStep,

action_spec: types.NestedTensorSpec,

policy_state_spec: types.NestedTensorSpec = (),

info_spec: types.NestedTensorSpec = (),

clip: bool = True,

emit_log_probability: bool = False,

automatic_state_reset: bool = True,

observation_and_action_constraint_splitter: Optional[

types.Splitter

] = None,

validate_args: bool = True,

name: Optional[Text] = None,

):

"""Initialization of TFPolicy class.

Args:

time_step_spec: A `TimeStep` spec of the expected time_steps. Usually

provided by the user to the subclass.

action_spec: A nest of BoundedTensorSpec representing the actions. Usually

provided by the user to the subclass.

policy_state_spec: A nest of TensorSpec representing the policy_state.

Provided by the subclass, not directly by the user.

info_spec: A nest of TensorSpec representing the policy info. Provided by

the subclass, not directly by the user.

clip: Whether to clip actions to spec before returning them. Default

True. Most policy-based algorithms (PCL, PPO, REINFORCE) use unclipped

continuous actions for training.

emit_log_probability: Emit log-probabilities of actions, if supported. If

True, policy_step.info will have CommonFields.LOG_PROBABILITY set.

Please consult utility methods provided in policy_step for setting and

retrieving these. When working with custom policies, either provide a

dictionary info_spec or a namedtuple with the field 'log_probability'.

automatic_state_reset: If `True`, then `get_initial_policy_state` is used

to clear state in `action()` and `distribution()` for for time steps

where `time_step.is_first()`.

observation_and_action_constraint_splitter: A function used to process

observations with action constraints. These constraints can indicate,

for example, a mask of valid/invalid actions for a given state of the

environment. The function takes in a full observation and returns a

tuple consisting of 1) the part of the observation intended as input to

the network and 2) the constraint. An example

`observation_and_action_constraint_splitter` could be as simple as: ```

def observation_and_action_constraint_splitter(observation): return

observation['network_input'], observation['constraint'] ``` *Note*: when

using `observation_and_action_constraint_splitter`, make sure the

provided `q_network` is compatible with the network-specific half of the

output of the `observation_and_action_constraint_splitter`. In

particular, `observation_and_action_constraint_splitter` will be called

on the observation before passing to the network. If

`observation_and_action_constraint_splitter` is None, action constraints

are not applied.

validate_args: Python bool. Whether to verify inputs to, and outputs of,

functions like `action` and `distribution` against spec structures,

dtypes, and shapes. Research code may prefer to set this value to

`False` to allow iterating on input and output structures without being

hamstrung by overly rigid checking (at the cost of harder-to-debug

errors). See also `TFAgent.validate_args`.

name: A name for this module. Defaults to the class name.

"""

super(TFPolicy, self).__init__(name=name)

common.check_tf1_allowed()

common.tf_agents_gauge.get_cell('TFAPolicy').set(True)

common.assert_members_are_not_overridden(base_cls=TFPolicy, instance=self)

if not isinstance(time_step_spec, ts.TimeStep):

raise ValueError(

'The `time_step_spec` must be an instance of `TimeStep`, but is `{}`.'

.format(type(time_step_spec))

)

self._time_step_spec = tensor_spec.from_spec(time_step_spec)

self._action_spec = tensor_spec.from_spec(action_spec)

self._policy_state_spec = tensor_spec.from_spec(policy_state_spec)

self._emit_log_probability = emit_log_probability

self._validate_args = validate_args

if emit_log_probability:

log_probability_spec = tensor_spec.BoundedTensorSpec(

shape=(),

dtype=tf.float32,

maximum=0,

minimum=-float('inf'),

name='log_probability',

)

log_probability_spec = tf.nest.map_structure(

lambda _: log_probability_spec, action_spec

)

info_spec = policy_step.set_log_probability(

info_spec, log_probability_spec

) # pytype: disable=wrong-arg-types

self._info_spec = tensor_spec.from_spec(info_spec)

self._setup_specs()

self._clip = clip

self._action_fn = common.function_in_tf1(reduce_retracing=False)(

self._action

)

self._automatic_state_reset = automatic_state_reset

self._observation_and_action_constraint_splitter = (

observation_and_action_constraint_splitter

)

def _setup_specs(self):

self._policy_step_spec = policy_step.PolicyStep(

action=self._action_spec,

state=self._policy_state_spec,

info=self._info_spec,

)

self._trajectory_spec = trajectory.from_transition(

self._time_step_spec, self._policy_step_spec, self._time_step_spec

)

def variables(self) -> Sequence[tf.Variable]:

"""Returns the list of Variables that belong to the policy."""

# Ignore self._variables() in favor of using tf.Module's tracking.

return super(TFPolicy, self).variables

@property

def observation_and_action_constraint_splitter(self) -> types.Splitter:

return self._observation_and_action_constraint_splitter

@property

def validate_args(self) -> bool:

"""Whether `action` & `distribution` validate input and output args."""

return self._validate_args

def get_initial_state(

self, batch_size: Optional[types.Int]

) -> types.NestedTensor:

"""Returns an initial state usable by the policy.

Args:

batch_size: Tensor or constant: size of the batch dimension. Can be None

in which case no dimensions gets added.

Returns:

A nested object of type `policy_state` containing properly

initialized Tensors.

"""

return self._get_initial_state(batch_size)

def _maybe_reset_state(self, time_step, policy_state):

if policy_state is (): # pylint: disable=literal-comparison

return policy_state

batch_size = tf.compat.dimension_value(time_step.discount.shape[0])

if batch_size is None:

batch_size = tf.shape(time_step.discount)[0]

# Make sure we call this with a kwarg as it may be wrapped in tf.function

# which would expect a tensor if it was not a kwarg.

zero_state = self.get_initial_state(batch_size=batch_size)

condition = time_step.is_first()

# When experience is a sequence we only reset automatically for the first

# time_step in the sequence as we can't easily generalize how the policy is

# unrolled over the sequence.

if nest_utils.get_outer_rank(time_step, self._time_step_spec) > 1:

condition = time_step.is_first()[:, 0, ...]

return nest_utils.where(condition, zero_state, policy_state)

def action(

self,

time_step: ts.TimeStep,

policy_state: types.NestedTensor = (),

seed: Optional[types.Seed] = None,

) -> policy_step.PolicyStep:

"""Generates next action given the time_step and policy_state.

Args:

time_step: A `TimeStep` tuple corresponding to `time_step_spec()`.

policy_state: A Tensor, or a nested dict, list or tuple of Tensors

representing the previous policy_state.

seed: Seed to use if action performs sampling (optional).

Returns:

A `PolicyStep` named tuple containing:

`action`: An action Tensor matching the `action_spec`.

`state`: A policy state tensor to be fed into the next call to action.

`info`: Optional side information such as action log probabilities.

Raises:

RuntimeError: If subclass __init__ didn't call super().__init__.

ValueError or TypeError: If `validate_args is True` and inputs or

outputs do not match `time_step_spec`, `policy_state_spec`,

or `policy_step_spec`.

"""

if self._enable_functions and getattr(self, '_action_fn', None) is None:

raise RuntimeError(

'Cannot find _action_fn. Did %s.__init__ call super?'

% type(self).__name__

)

if self._enable_functions:

action_fn = self._action_fn

else:

action_fn = self._action

if self._validate_args:

time_step = nest_utils.prune_extra_keys(self._time_step_spec, time_step)

policy_state = nest_utils.prune_extra_keys(

self._policy_state_spec, policy_state

)

nest_utils.assert_same_structure(

time_step,

self._time_step_spec,

message='time_step and time_step_spec structures do not match',

)

if common.safe_has_state(policy_state):

nest_utils.assert_same_structure(

policy_state,

self._policy_state_spec,

message=(

'policy_state and policy_state_spec structures do not match'

),

)

if self._automatic_state_reset:

policy_state = self._maybe_reset_state(time_step, policy_state)

step = action_fn(time_step=time_step, policy_state=policy_state, seed=seed)

def clip_action(action, action_spec):

if isinstance(action_spec, tensor_spec.BoundedTensorSpec):

return common.clip_to_spec(action, action_spec)

return action

if self._validate_args:

nest_utils.assert_same_structure(

step.action,

self._action_spec,

message='action and action_spec structures do not match',

)

if self._clip:

clipped_actions = tf.nest.map_structure(

clip_action, step.action, self._action_spec

)

step = step._replace(action=clipped_actions)

if self._validate_args:

nest_utils.assert_same_structure(

step,

self._policy_step_spec,

message='action output and policy_step_spec structures do not match',

)

def compare_to_spec(value, spec):

return value.dtype.is_compatible_with(spec.dtype)

compatibility = [

compare_to_spec(v, s)

for (v, s) in zip(

tf.nest.flatten(step.action), tf.nest.flatten(self.action_spec)

)

]

if not all(compatibility):

get_dtype = lambda x: x.dtype

action_dtypes = tf.nest.map_structure(get_dtype, step.action)

spec_dtypes = tf.nest.map_structure(get_dtype, self.action_spec)

raise TypeError(

"Policy produced an action with a dtype that doesn't "

'match its action_spec. Got action:\n %s\n with '

'action_spec:\n %s' % (action_dtypes, spec_dtypes)

)

return step

def distribution(

self, time_step: ts.TimeStep, policy_state: types.NestedTensor = ()

) -> policy_step.PolicyStep:

"""Generates the distribution over next actions given the time_step.

Args:

time_step: A `TimeStep` tuple corresponding to `time_step_spec()`.

policy_state: A Tensor, or a nested dict, list or tuple of Tensors

representing the previous policy_state.

Returns:

A `PolicyStep` named tuple containing:

`action`: A tf.distribution capturing the distribution of next actions.

`state`: A policy state tensor for the next call to distribution.

`info`: Optional side information such as action log probabilities.

Raises:

ValueError or TypeError: If `validate_args is True` and inputs or

outputs do not match `time_step_spec`, `policy_state_spec`,

or `policy_step_spec`.

"""

if self._validate_args:

time_step = nest_utils.prune_extra_keys(self._time_step_spec, time_step)

policy_state = nest_utils.prune_extra_keys(

self._policy_state_spec, policy_state

)

nest_utils.assert_same_structure(

time_step,

self._time_step_spec,

message='time_step and time_step_spec structures do not match',

)

nest_utils.assert_same_structure(

policy_state,

self._policy_state_spec,

message='policy_state and policy_state_spec structures do not match',

)

if self._automatic_state_reset:

policy_state = self._maybe_reset_state(time_step, policy_state)

step = self._distribution(time_step=time_step, policy_state=policy_state)

if self.emit_log_probability:

# This here is set only for compatibility with info_spec in constructor.

info = policy_step.set_log_probability(

step.info,

tf.nest.map_structure(

lambda _: tf.constant(0.0, dtype=tf.float32),

policy_step.get_log_probability(self._info_spec),

),

)

step = step._replace(info=info)

if self._validate_args:

nest_utils.assert_same_structure(

step,

self._policy_step_spec,

message=(

'distribution output and policy_step_spec structures do not match'

),

)

return step

def update(

self,

policy,

tau: float = 1.0,

tau_non_trainable: Optional[float] = None,

sort_variables_by_name: bool = False,

) -> tf.Operation:

"""Update the current policy with another policy.

This would include copying the variables from the other policy.

Args:

policy: Another policy it can update from.

tau: A float scalar in [0, 1]. When tau is 1.0 (the default), we do a hard

update. This is used for trainable variables.

tau_non_trainable: A float scalar in [0, 1] for non_trainable variables.

If None, will copy from tau.

sort_variables_by_name: A bool, when True would sort the variables by name

before doing the update.

Returns:

An TF op to do the update.

"""

if self.variables():

return common.soft_variables_update(

policy.variables(),

self.variables(),

tau=tau,

tau_non_trainable=tau_non_trainable,

sort_variables_by_name=sort_variables_by_name,

)

else:

return tf.no_op()

@property

def emit_log_probability(self) -> bool:

"""Whether this policy instance emits log probabilities or not."""

return self._emit_log_probability

@property

def time_step_spec(self) -> ts.TimeStep:

"""Describes the `TimeStep` tensors returned by `step()`.

Returns:

A `TimeStep` namedtuple with `TensorSpec` objects instead of Tensors,

which describe the shape, dtype and name of each tensor returned by

`step()`.

"""

return self._time_step_spec

@property

def action_spec(self) -> types.NestedTensorSpec:

"""Describes the TensorSpecs of the Tensors expected by `step(action)`.

`action` can be a single Tensor, or a nested dict, list or tuple of

Tensors.

Returns:

An single BoundedTensorSpec, or a nested dict, list or tuple of

`BoundedTensorSpec` objects, which describe the shape and

dtype of each Tensor expected by `step()`.

"""

return self._action_spec

@property

def policy_state_spec(self) -> types.NestedTensorSpec:

"""Describes the Tensors expected by `step(_, policy_state)`.

`policy_state` can be an empty tuple, a single Tensor, or a nested dict,

list or tuple of Tensors.

Returns:

An single TensorSpec, or a nested dict, list or tuple of

`TensorSpec` objects, which describe the shape and

dtype of each Tensor expected by `step(_, policy_state)`.

"""

return self._policy_state_spec

@property

def info_spec(self) -> types.NestedTensorSpec:

"""Describes the Tensors emitted as info by `action` and `distribution`.

`info` can be an empty tuple, a single Tensor, or a nested dict,

list or tuple of Tensors.

Returns:

An single TensorSpec, or a nested dict, list or tuple of

`TensorSpec` objects, which describe the shape and

dtype of each Tensor expected by `step(_, policy_state)`.

"""

return self._info_spec

@property

def policy_step_spec(self) -> policy_step.PolicyStep:

"""Describes the output of `action()`.

Returns:

A nest of TensorSpec which describe the shape and dtype of each Tensor

emitted by `action()`.

"""

return self._policy_step_spec

# TODO(kbanoop, ebrevdo): Should this be collect_data_spec to mirror agents?

@property

def trajectory_spec(self) -> trajectory.Trajectory:

"""Describes the Tensors written when using this policy with an environment.

Returns:

A `Trajectory` containing all tensor specs associated with the

observation_spec, action_spec, policy_state_spec, and info_spec of

this policy.

"""

return self._trajectory_spec

@property

def collect_data_spec(self) -> trajectory.Trajectory:

"""Describes the Tensors written when using this policy with an environment.

Returns:

A nest of TensorSpec which describe the shape and dtype of each Tensor

required to train the agent which generated this policy.

"""

return self._trajectory_spec

# Subclasses MAY optionally override _action.

def _action(

self,

time_step: ts.TimeStep,

policy_state: types.NestedTensor,

seed: Optional[types.Seed] = None,

) -> policy_step.PolicyStep:

"""Implementation of `action`.

Args:

time_step: A `TimeStep` tuple corresponding to `time_step_spec()`.

policy_state: A Tensor, or a nested dict, list or tuple of Tensors

representing the previous policy_state.

seed: Seed to use if action performs sampling (optional).

Returns:

A `PolicyStep` named tuple containing:

`action`: An action Tensor matching the `action_spec`.

`state`: A policy state tensor to be fed into the next call to action.

`info`: Optional side information such as action log probabilities.

"""

seed_stream = tfp.util.SeedStream(seed=seed, salt='tf_agents_tf_policy')

distribution_step = self._distribution(time_step, policy_state) # pytype: disable=wrong-arg-types

actions = tf.nest.map_structure(

lambda d: reparameterized_sampling.sample(d, seed=seed_stream()),

distribution_step.action,

)

info = distribution_step.info

if self.emit_log_probability:

try:

log_probability = tf.nest.map_structure(

lambda a, d: d.log_prob(a), actions, distribution_step.action

)

info = policy_step.set_log_probability(info, log_probability)

except:

raise TypeError(

'%s does not support emitting log-probabilities.'

% type(self).__name__

)

return distribution_step._replace(action=actions, info=info)

## Subclasses MUST implement these.

def _distribution(

self, time_step: ts.TimeStep, policy_state: types.NestedTensorSpec

) -> policy_step.PolicyStep:

"""Implementation of `distribution`.

Args:

time_step: A `TimeStep` tuple corresponding to `time_step_spec()`.

policy_state: A Tensor, or a nested dict, list or tuple of Tensors

representing the previous policy_state.

Returns:

A `PolicyStep` named tuple containing:

`action`: A (optionally nested) of tfp.distribution.Distribution

capturing the distribution of next actions.

`state`: A policy state tensor for the next call to distribution.

`info`: Optional side information such as action log probabilities.

"""

raise NotImplementedError()

# Subclasses MAY optionally overwrite _get_initial_state.

def _get_initial_state(self, batch_size: int) -> types.NestedTensor:

"""Returns the initial state of the policy network.

Args:

batch_size: A constant or Tensor holding the batch size. Can be None, in

which case the state will not have a batch dimension added.

Returns:

A nest of zero tensors matching the spec of the policy network state.

"""

return tensor_spec.zero_spec_nest(

self._policy_state_spec,

outer_dims=None if batch_size is None else [batch_size],

)