def decay(rho, beta, progress):
    """
    Decay function from Rosokha & Younge (2019).
    :param rho: trajectory shape, in [0,1]
    :param beta: value to decay
    :param progress: relative progress in interval [0,T]
    :return: decayed value
    """

    assert rho >= 0 and rho <= 1, 'Rho must be in [0,1]'
    assert progress >= 0 and progress <= 1, 't/T must be in [0,1]'

    # No decay
    if rho == 0:
        return beta

    # Slow onset decay
    elif rho > 0 and rho < 0.5:
        return beta * (1 - (progress) ** (1 / (2 * rho)))

    # Linear decay
    elif rho == 0.5:
        return beta * (1 - progress)

    # Fast onset decay
    elif rho > 0.5 and rho < 1:
        return beta * (1 - progress) ** (1 / (2 * (1 - rho)))

    # Immediate decay
    elif rho == 1:
        return 0