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

"""Learner implementation for Agents. Refer to the examples dir."""

import os

from typing import Any, Optional, Tuple

from absl import logging

import gin

import tensorflow.compat.v2 as tf

from tf_agents.agents import tf_agent

from tf_agents.specs import tensor_spec

from tf_agents.train import interval_trigger

from tf_agents.typing import types

from tf_agents.utils import common

TRAIN_DIR = 'train'

POLICY_SAVED_MODEL_DIR = 'policies'

COLLECT_POLICY_SAVED_MODEL_DIR = 'collect_policy'

GREEDY_POLICY_SAVED_MODEL_DIR = 'greedy_policy'

RAW_POLICY_SAVED_MODEL_DIR = 'policy'

POLICY_CHECKPOINT_DIR = 'checkpoints'

REPLAY_BUFFER_CHECKPOINT_DIR = 'replay_buffer_checkpoints'

ExperienceAndSampleInfo = Tuple[types.NestedTensor, Tuple[Any, ...]]

@gin.configurable

class Learner(tf.Module):

"""Manages all the learning details needed when training an agent.

These include:

* Using distribution strategies correctly

* Summaries

* Checkpoints

* Minimizing entering/exiting TF context:

Especially in the case of TPUs scheduling a single TPU program to

perform multiple train steps is critical for performance.

* Generalizes the train call to be done correctly across CPU, GPU, or TPU

executions managed by DistributionStrategies. This uses `strategy.run` and

then makes sure to do a reduce operation over the `LossInfo` returned by

the agent.

"""

def __init__(

self,

root_dir,

train_step,

agent,

experience_dataset_fn=None,

after_train_strategy_step_fn=None,

triggers=None,

checkpoint_interval=100000,

summary_interval=1000,

max_checkpoints_to_keep=3,

use_kwargs_in_agent_train=False,

strategy=None,

run_optimizer_variable_init=True,

use_reverb_v2=False,

direct_sampling=False,

experience_dataset_options=None,

strategy_run_options=None,

summary_root_dir=None,

):

"""Initializes a Learner instance.

Args:

root_dir: Main directory path where checkpoints, saved_models, and

summaries (if summary_dir is not specified) will be written to.

train_step: a scalar tf.int64 `tf.Variable` which will keep track of the

number of train steps. This is used for artifacts created like

summaries, or outputs in the root_dir.

agent: `tf_agent.TFAgent` instance to train with.

experience_dataset_fn: a function that will create an instance of a

tf.data.Dataset used to sample experience for training. Required for

using the Learner as is. Optional for subclass learners which take a new

iterator each time when `learner.run` is called.

after_train_strategy_step_fn: (Optional) callable of the form `fn(sample,

loss)` which can be used for example to update priorities in a replay

buffer where sample is pulled from the `experience_iterator` and loss is

a `LossInfo` named tuple returned from the agent. This is called after

every train step. It runs using `strategy.run(...)`.

triggers: List of callables of the form `trigger(train_step)`. After every

`run` call every trigger is called with the current `train_step` value

as an np scalar.

checkpoint_interval: Number of train steps in between checkpoints. Note

these are placed into triggers and so a check to generate a checkpoint

only occurs after every `run` call. Set to -1 to disable (this is not

recommended, because it means that if the pipeline gets preempted, all

previous progress is lost). This only takes care of the checkpointing

the training process. Policies must be explicitly exported through

triggers.

summary_interval: Number of train steps in between summaries. Note these

are placed into triggers and so a check to generate a checkpoint only

occurs after every `run` call.

max_checkpoints_to_keep: Maximum number of checkpoints to keep around.

These are used to recover from pre-emptions when training.

use_kwargs_in_agent_train: If True the experience from the replay buffer

is passed into the agent as kwargs. This requires samples from the RB to

be of the form `dict(experience=experience, kwarg1=kwarg1, ...)`. This

is useful if you have an agent with a custom argspec.

strategy: (Optional) `tf.distribute.Strategy` to use during training.

run_optimizer_variable_init: Specifies if the variables of the optimizer

are initialized before checkpointing. This should be almost always

`True` (default) to ensure that the state of the optimizer is

checkpointed properly. The initialization of the optimizer variables

happens by building the Tensorflow graph. This is done by calling a

`get_concrete_function` on the agent's `train` method which requires

passing some input. Since, no real data is available at this point we

use the batched form of `training_data_spec` to achieve this (standard

technique). The problem arises when the agent expects some agent

specific batching of the input. In this case, there is no _general_ way

at this point in the learner to batch the impacted specs properly. To

avoid breaking the code in these specific cases, we recommend turning

off initialization of the optimizer variables by setting the value of

this field to `False`.

use_reverb_v2: If True then we expect the dataset samples to return a

named_tuple with a data and an info field. If False we expect a

tuple(data, info).

direct_sampling: Do not use replay_buffer, but sample from offline dataset

directly.

experience_dataset_options: (Optional) `tf.distribute.InputOptions` passed

to `strategy.distribute_datasets_from_function`, used to control options

on how this dataset is distributed.

strategy_run_options: (Optional) `tf.distribute.RunOptions` passed to

`strategy.run`. This is passed to every strategy.run invocation by the

learner.

summary_root_dir: (Optional) Root directory path where summaries will be

written to.

"""

if checkpoint_interval < 0:

logging.warning(

'Warning: checkpointing the training process is manually disabled.'

'This means training progress will NOT be automatically restored '

'if the job gets preempted.'

)

self._train_dir = os.path.join(root_dir, TRAIN_DIR)

summary_root_dir = (

root_dir if summary_root_dir is None else summary_root_dir

)

self._summary_dir = os.path.join(summary_root_dir, TRAIN_DIR)

self._use_reverb_v2 = use_reverb_v2

self._direct_sampling = direct_sampling

if summary_interval:

self.train_summary_writer = tf.compat.v2.summary.create_file_writer(

self._summary_dir, flush_millis=10000

)

else:

self.train_summary_writer = tf.summary.create_noop_writer()

self.train_step = train_step

self._agent = agent

self.use_kwargs_in_agent_train = use_kwargs_in_agent_train

self.strategy = strategy or tf.distribute.get_strategy()

dataset = None

if experience_dataset_fn:

with self.strategy.scope():

dataset = self.strategy.distribute_datasets_from_function(

lambda _: experience_dataset_fn(),

options=experience_dataset_options,

)

self._experience_iterator = iter(dataset)

self.after_train_strategy_step_fn = after_train_strategy_step_fn

self.triggers = triggers or []

# Prevent autograph from going into the agent.

self._agent.train = tf.autograph.experimental.do_not_convert(agent.train)

self._strategy_run_options = strategy_run_options

checkpoint_dir = os.path.join(self._train_dir, POLICY_CHECKPOINT_DIR)

with self.strategy.scope():

agent.initialize()

if run_optimizer_variable_init:

# Force a concrete function creation inside of the strategy scope to

# ensure that all variables, including optimizer slot variables, are

# created. This has to happen before the checkpointer is created.

if dataset is not None:

if use_reverb_v2:

batched_specs = dataset.element_spec.data

else:

# Assumes (experience, sample_info) = next(iterator)

batched_specs, _ = dataset.element_spec

else:

batched_specs = tensor_spec.add_outer_dims_nest(

self._agent.training_data_spec,

(None, self._agent.train_sequence_length),

)

if self.use_kwargs_in_agent_train:

batched_specs = dict(experience=batched_specs)

@common.function

def _create_variables(specs):

# TODO(b/170516529): Each replica has to be in the same graph.

# This can be ensured by placing the `strategy.run(...)` call inside

# the `tf.function`.

if self.use_kwargs_in_agent_train:

return self.strategy.run(

self._agent.train,

kwargs=specs,

options=self._strategy_run_options,

)

return self.strategy.run(

self._agent.train,

args=(specs,),

options=self._strategy_run_options,

)

_create_variables.get_concrete_function(batched_specs)

else:

# TODO(b/186052656) Update clients.

logging.warning('run_optimizer_variable_init = False is Deprecated')

self._checkpointer = common.Checkpointer(

checkpoint_dir,

max_to_keep=max_checkpoints_to_keep,

agent=self._agent,

train_step=self.train_step,

)

self._checkpointer.initialize_or_restore() # pytype: disable=attribute-error

for trigger in self.triggers:

if hasattr(trigger, 'set_start'):

trigger.set_start(self.train_step.numpy())

self.triggers.append(self._get_checkpoint_trigger(checkpoint_interval))

self.summary_interval = tf.constant(summary_interval, dtype=tf.int64)

@property

def train_step_numpy(self):

"""The current train_step.

Returns:

The current `train_step`. Note this will return a scalar numpy array which

holds the `train_step` value when this was called.

"""

return self.train_step.numpy()

def _get_checkpoint_trigger(self, checkpoint_interval):

if checkpoint_interval <= 0:

return lambda _, force_trigger=False: None

save_fn = lambda: self._checkpointer.save(self.train_step)

return interval_trigger.IntervalTrigger(

checkpoint_interval, save_fn, start=self.train_step.numpy()

)

def run(self, iterations=1, iterator=None, parallel_iterations=10):

"""Runs `iterations` iterations of training.

Args:

iterations: Number of train iterations to perform per call to run. The

iterations will be evaluated in a tf.while loop created by autograph.

Final aggregated losses will be returned.

iterator: The iterator to the dataset to use for training. If not

specified, `self._experience_iterator` is used.

parallel_iterations: Maximum number of train iterations to allow running

in parallel. This value is forwarded directly to the training tf.while

loop.

Returns:

The total loss computed before running the final step.

"""

assert iterations >= 1, (

'Iterations must be greater or equal to 1, was %d' % iterations

)

def _summary_record_if():

if self.summary_interval:

return tf.math.equal(

self.train_step % tf.constant(self.summary_interval), 0

)

else:

return tf.constant(False)

with self.train_summary_writer.as_default(), common.soft_device_placement(), tf.compat.v2.summary.record_if(

_summary_record_if

), self.strategy.scope():

iterator = iterator or self._experience_iterator

loss_info = self._train(

tf.constant(iterations), iterator, parallel_iterations

)

train_step_val = self.train_step.numpy()

for trigger in self.triggers:

trigger(train_step_val)

return loss_info

# Use tf.config.experimental_run_functions_eagerly(True) if you want to

# disable use of tf.function.

@common.function(autograph=True)

def _train(self, iterations, iterator, parallel_iterations):

# Call run explicitly once to get loss info shape for autograph. Because the

# for loop below will get converted to a `tf.while_loop` by autograph we

# need the shape of loss info to be well defined.

loss_info = self.single_train_step(iterator)

for _ in tf.range(iterations - 1):

tf.autograph.experimental.set_loop_options(

parallel_iterations=parallel_iterations

)

loss_info = self.single_train_step(iterator)

def _reduce_loss(loss):

rank = None

if isinstance(loss, tf.distribute.DistributedValues):

# If loss is distributed get the rank from the first replica.

rank = loss.values[0].shape.rank

elif tf.is_tensor(loss):

rank = loss.shape.rank

axis = None

if rank:

axis = tuple(range(0, rank))

return self.strategy.reduce(tf.distribute.ReduceOp.SUM, loss, axis=axis)

# We assume all data can be reduced in the loss_info. This means no

# string dtypes are currently allowed as LossInfo Fields.

reduced_loss_info = tf.nest.map_structure(_reduce_loss, loss_info)

return reduced_loss_info

def single_train_step(self, iterator):

sample = next(iterator)

if self._direct_sampling:

experience, sample_info = sample, None

elif self._use_reverb_v2:

experience, sample_info = sample.data, sample.info

else:

experience, sample_info = sample

if self.use_kwargs_in_agent_train:

loss_info = self.strategy.run(

self._agent.train,

kwargs=experience,

options=self._strategy_run_options,

)

else:

loss_info = self.strategy.run(

self._agent.train,

args=(experience,),

options=self._strategy_run_options,

)

if self.after_train_strategy_step_fn:

if self.use_kwargs_in_agent_train:

self.strategy.run(

self.after_train_strategy_step_fn,

kwargs=dict(

experience=(experience, sample_info), loss_info=loss_info

),

options=self._strategy_run_options,

)

else:

self.strategy.run(

self.after_train_strategy_step_fn,

args=((experience, sample_info), loss_info),

options=self._strategy_run_options,

)

return loss_info

def loss(

self,

experience_and_sample_info: Optional[ExperienceAndSampleInfo] = None,

reduce_op: tf.distribute.ReduceOp = tf.distribute.ReduceOp.SUM,

) -> tf_agent.LossInfo:

"""Computes loss for the experience.

Since this calls agent.loss() it does not update gradients or

increment the train step counter. Networks are called with `training=False`

so statistics like batch norm are not updated.

Args:

experience_and_sample_info: A batch of experience and sample info. If not

specified, `next(self._experience_iterator)` is used.

reduce_op: a `tf.distribute.ReduceOp` value specifying how loss values

should be aggregated across replicas.

Returns:

The total loss computed.

"""

def _summary_record_if():

return tf.math.equal(

self.train_step % tf.constant(self.summary_interval), 0

)

with self.train_summary_writer.as_default(), common.soft_device_placement(), tf.compat.v2.summary.record_if(

_summary_record_if

), self.strategy.scope():

if experience_and_sample_info is None:

sample = next(self._experience_iterator)

if self._direct_sampling:

experience_and_sample_info = (sample, None)

elif self._use_reverb_v2:

experience_and_sample_info = (sample.data, sample.info)

else:

experience_and_sample_info = sample

loss_info = self._loss(experience_and_sample_info, reduce_op)

return loss_info

# Use tf.config.experimental_run_functions_eagerly(True) if you want to

# disable use of tf.function.

@common.function(autograph=True)

def _loss(

self,

experience_and_sample_info: ExperienceAndSampleInfo,

reduce_op: tf.distribute.ReduceOp,

) -> tf_agent.LossInfo:

(experience, sample_info) = experience_and_sample_info

if self.use_kwargs_in_agent_train:

loss_info = self.strategy.run(self._agent.loss, kwargs=experience)

else:

loss_info = self.strategy.run(self._agent.loss, args=(experience,))

if self.after_train_strategy_step_fn:

if self.use_kwargs_in_agent_train:

self.strategy.run(

self.after_train_strategy_step_fn,

kwargs=dict(

experience=(experience, sample_info), loss_info=loss_info

),

options=self._strategy_run_options,

)

else:

self.strategy.run(

self.after_train_strategy_step_fn,

args=((experience, sample_info), loss_info),

options=self._strategy_run_options,

)

def _reduce_loss(loss):

rank = None

if isinstance(loss, tf.distribute.DistributedValues):

rank = loss.values[0].shape.rank

elif tf.is_tensor(loss):

rank = loss.shape.rank

axis = None

if rank:

axis = tuple(range(0, rank))

return self.strategy.reduce(reduce_op, loss, axis=axis)

# We assume all data can be reduced in the loss_info. This means no

# string dtypes are currently allowed as LossInfo Fields.

reduced_loss_info = tf.nest.map_structure(_reduce_loss, loss_info)

return reduced_loss_info