def punish_small_rl_headways(vehicles, rl_ids, headway_threshold, penalty_gain=1, penalty_exponent=1):
    """
    A reward function used to train rl vehicles to avoid small headways.

    :param vehicles {dict} - contains the state of all vehicles in the network (generally self.vehicles)
    :param rl_ids: {list} - list of ids for rl vehicles in the network (generally self.rl_ids)
    :param headway_threshold {float} - the maximum headway allowed for rl vehicles before being penalized
    :param penalty_gain {float} - sets the penalty for each rl vehicle between 0 and this value
    :param penalty_exponent {float} - used to allow exponential punishing of smaller headways
    :return: a (non-negative) penalty on rl vehicles whose headway is below the headway_threshold
    """
    headway_penalty = 0
    for veh_id in rl_ids:
        if vehicles[veh_id]["headway"] < headway_threshold:
            headway_penalty += (((headway_threshold - vehicles[veh_id]["headway"]) / headway_threshold)
                                ** penalty_exponent) * penalty_gain

    # in order to keep headway penalty (and thus reward function) positive
    max_headway_penalty = len(rl_ids) * penalty_gain

    return max_headway_penalty - headway_penalty