import tensorflow as tf



def load_agent_ckpt(ckpt_dir, tf_agent, global_step=None):
  if global_step is None:
    global_step = tf.compat.v1.train.get_or_create_global_step()
  train_checkpointer = common.Checkpointer(
      ckpt_dir=ckpt_dir, agent=tf_agent, global_step=global_step)
  train_checkpointer.initialize_or_restore().assert_existing_objects_matched()
