78c163f37be165673b040cff49e1ee98eb89ca0f,tf_agents/bandits/agents/examples/v2/trainer.py,,train,#Any#Any#Any#Any#Any#Any#Any#,104

Before Change

    metric_utils.log_metrics(metrics)
    for metric in metrics:
      tf.summary.scalar(
          metric.name, metric.result(), step=step_metric.result())
    checkpoint_manager.save()
    saver.save(os.path.join(root_dir, "policy_%d" % step_metric.result()))

After Change

      tf_metrics.AverageEpisodeLengthMetric(batch_size=environment.batch_size)
  ] + list(additional_metrics)

  if isinstance(environment.reward_spec(), dict):
    metrics += [tf_metrics.AverageReturnMultiMetric(
        reward_spec=environment.reward_spec(),
        batch_size=environment.batch_size)]
  else:
    metrics += [
        tf_metrics.AverageReturnMetric(batch_size=environment.batch_size)]

  if training_data_spec_transformation_fn is not None:
    add_batch_fn = lambda data: replay_buffer.add_batch(  // pylint: disable=g-long-lambda
        training_data_spec_transformation_fn(data))
  else:

Instances