code: def restore_and_get_checkpoint_manager(
      root_dir, agent, metrics, step_metric
  ):
      """Restores from `root_dir` and returns a function that writes checkpoints."""
      trackable_objects = {metric.name: metric for metric in metrics}
      trackable_objects[AGENT_CHECKPOINT_NAME] = agent
      trackable_objects[STEP_CHECKPOINT_NAME] = step_metric
      checkpoint = tf.train.Checkpoint(**trackable_objects)
      checkpoint_manager = tf.train.CheckpointManager(
          checkpoint=checkpoint, directory=root_dir, max_to_keep=5
      )
      latest = checkpoint_manager.latest_checkpoint
      if latest is not None:
          logging.info("Restoring checkpoint from %s.", latest)
          checkpoint.restore(latest)
          logging.info(
              "Successfully restored to step %s.", step_metric.result()
          )
      else:
          logging.info(
              "Did not find a pre-existing checkpoint. Starting from scratch."
          )
      return checkpoint_manager
  
file_path: agents/tf_agents/bandits/agents/examples/v2/trainer.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import os
  from absl import logging
  import tensorflow
  from tf_agents.bandits.replay_buffers import bandit_replay_buffer
  from tf_agents.drivers import dynamic_step_driver
  from tf_agents.eval import metric_utils
  from tf_agents.metrics import export_utils
  from tf_agents.metrics import tf_metrics
  from tf_agents.policies import policy_saver
label_desc: Use image node_id 4 for calling a global function with example usage: restore_and_get_checkpoint_manager(root_dir, agent, metrics, step_metric) and returns: checkpoint_manager