def run_episode(environment, agent, is_training=False):
  """Run a single episode."""

  timestep = environment.reset()

  while not timestep.last():
    action = agent.step(timestep, is_training)
    new_timestep = environment.step(action)

    if is_training:
      agent.update(timestep, action, new_timestep)

    timestep = new_timestep

  episode_return = environment.episode_return

  return episode_return