def exp_decay_schedule_builder(base_value, decay_steps, decay_rate, staircase, round_num):
  """Creates a learning rate schedule with exponential root decay.

  Args:
    base_value: The base value of the quantity to decay over time.
    decay_steps: A positive scalar that governs how much the value decays at a
      given round number.
    decay_rate: A float between 0 and 1 that governs how quickly the decay
      occurs.
    staircase: A boolean. If set to True, the decaying occurs in discrete
      intervals.

  Returns:
    A 1-arg callable that produces a decayed version of the base value when
      passed the current round_num.
  """
  if staircase:
    def exp_decay_fn(round_num):
      return base_value * pow(decay_rate, round_num // decay_steps)
  else:
    def exp_decay_fn(round_num):
      return base_value * pow(decay_rate, round_num / decay_steps)

  return exp_decay_fn(round_num)