File size: 5,263 Bytes
078ba7c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
dff2db9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
from framework.controller import (
    AdaptiveController,
    DecodingParameters,
    FrameworkState,
)


def _make_state(lam=0.30, novelty=0.20, bias=0.40, quality=0.60):
    return FrameworkState(
        lambda_model=lam,
        lambda_human=0.15,
        aggregate_novelty=novelty,
        aggregate_bias=bias,
        aggregate_quality=quality,
    )


def test_reward_signal_eq_xxxiii():
    controller = AdaptiveController()
    # r_{t+1} = D(s_t) - D(s_{t+1})
    assert abs(controller.compute_reward(0.8, 0.5) - 0.3) < 1e-12
    assert controller.compute_reward(0.5, 0.8) < 0


def test_td_error_and_value_update_eq_xxxv_xxxvi():
    controller = AdaptiveController(gamma=0.95, alpha=0.1)
    s0 = _make_state(novelty=0.10)
    s1 = _make_state(novelty=0.30)

    # With an empty value table: δ = r + γ·0 - 0 = r
    reward = 0.25
    delta = controller.compute_td_error(reward, s0, s1)
    assert abs(delta - reward) < 1e-12

    # V(s0) ← V(s0) + α·δ (Eq xxxvi)
    v_after = controller.update_value(s0, delta)
    assert abs(v_after - 0.1 * reward) < 1e-12

    # Second TD error now sees the updated V(s0)
    delta2 = controller.compute_td_error(reward, s0, s1)
    assert abs(delta2 - (reward - v_after)) < 1e-12


def test_theta_update_moves_along_direction_and_clips_eq_xxxvii():
    controller = AdaptiveController(eta=0.5)
    theta = DecodingParameters(temperature=0.7, top_p=0.9, top_k=50)

    # Novelty deficit direction increases exploration
    direction = controller.CONTROL_DIRECTIONS["novelty_deficit"]
    updated = controller.apply_parameter_update(theta, td_error=0.4, direction=direction)
    assert updated.temperature > theta.temperature
    assert updated.top_p >= theta.top_p
    assert updated.top_k >= theta.top_k

    # Large positive TD error must clip at θ_max
    updated_max = controller.apply_parameter_update(theta, td_error=100.0, direction=direction)
    assert updated_max.temperature == controller.TEMP_MAX
    assert updated_max.top_p == controller.TOP_P_MAX
    assert updated_max.top_k == controller.TOP_K_MAX

    # Negative TD error reverses along c and clips at θ_min
    updated_min = controller.apply_parameter_update(theta, td_error=-100.0, direction=direction)
    assert updated_min.temperature == controller.TEMP_MIN
    assert updated_min.top_p == controller.TOP_P_MIN
    assert updated_min.top_k == controller.TOP_K_MIN


def test_control_direction_follows_dominant_gap():
    controller = AdaptiveController()
    state = _make_state(lam=0.30)

    dominant, direction = controller.select_control_direction(
        {"decay_mismatch": 0.02, "novelty_deficit": 0.30,
         "bias_excess": 0.01, "quality_deficit": 0.05},
        state,
    )
    assert dominant == "novelty_deficit"
    assert direction["temperature"] > 0

    dominant, direction = controller.select_control_direction(
        {"decay_mismatch": 0.02, "novelty_deficit": 0.01,
         "bias_excess": 0.40, "quality_deficit": 0.05},
        state,
    )
    assert dominant == "bias_excess"
    assert direction["temperature"] < 0


def test_closed_loop_episode_bootstrap_and_transition():
    controller = AdaptiveController(eta=0.2, gamma=0.95, alpha=0.1)
    theta_0 = DecodingParameters(temperature=0.7, top_p=0.9, top_k=50)
    s0 = _make_state(novelty=0.10)

    d0 = {"total_distance": 0.6,
          "gap_vector": {"decay_mismatch": 0.0, "novelty_deficit": 0.25,
                         "bias_excess": 0.0, "quality_deficit": 0.05}}
    step1 = controller.initial_step(s0, d0, theta_0)
    # Bootstrap: δ_0 = D(s_0) with an empty value table
    assert abs(step1.td_error - 0.6) < 1e-12
    assert step1.theta_after.temperature > theta_0.temperature

    s1 = _make_state(novelty=0.30)
    d1 = {"total_distance": 0.35,
          "gap_vector": {"decay_mismatch": 0.0, "novelty_deficit": 0.05,
                         "bias_excess": 0.0, "quality_deficit": 0.05}}
    step2 = controller.transition_step(s0, s1, d0, d1, step1.theta_after)

    # r_1 = 0.6 - 0.35 (Eq xxxiii); V updated by α·δ (Eq xxxvi)
    assert abs(step2.reward - 0.25) < 1e-12
    assert step2.value_after != step2.value_before
    assert step2.distance_after == 0.35


def test_controller_accepts_cli_metric_keys():
    controller = AdaptiveController()

    recommendation = controller.recommend_parameters(
        {
            "novelty": 0.2,
            "self_bleu": 60.0,
            "fallback_quality": 0.8,
            "bias_proxy": 0.1,
            "perplexity": 40.0,
        },
        current_temp=0.7,
        current_top_p=0.9,
    )

    assert "temperature=0.85" in recommendation
    assert "top_p=0.96" in recommendation


def test_controller_clamps_recommendations():
    controller = AdaptiveController()

    recommendation = controller.recommend_parameters(
        {
            "Novelty": 0.0,
            "Self-BLEU": 100.0,
            "Fallback Quality": 0.0,
            "Bias Proxy": 1.0,
            "Perplexity": 999.0,
        },
        current_temp=1.2,
        current_top_p=0.98,
    )

    assert "temperature=1.20" in recommendation
    assert "top_p=0.98" in recommendation


if __name__ == "__main__":
    test_controller_accepts_cli_metric_keys()
    test_controller_clamps_recommendations()