def update_beta(beta_schedule: str, beta: float, rho: float,
                step: int) -> float:
    """Update beta based on schedule and training step.

    Args:
        beta_schedule (str): Schedule for beta update.
        beta (float): Initial beta.
        rho (float): Schedule decay parameter.
        step (int): Current training iteration.

    Returns:
        Updated beta as per input schedule.
    """
    if beta_schedule == "linear_decay":
        return beta * ((1.0 - rho)**step)
    return beta