File size: 4,662 Bytes
05feb67
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
import re

with open("env/rewards.py", "r") as f:
    code = f.read()

# 1. Change the signature of `calculate_step_reward`
old_sig = """    def calculate_step_reward(
        self,
        state: dict[str, Any] | None,
        action: str,
        result: dict[str, Any] | None,
        original_config: str | None = None,
        fixed_config: str | None = None,
        error_message: str | None = None,
        expected_config: str | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> float:"""

new_sig = """    def calculate_step_reward(
        self,
        state: dict[str, Any] | None,
        action: str,
        result: dict[str, Any] | None,
        original_config: str | None = None,
        fixed_config: str | None = None,
        error_message: str | None = None,
        expected_config: str | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> tuple[float, dict[str, float]]:"""

code = code.replace(old_sig, new_sig)

# 2. Change the return and body of `calculate_step_reward`
old_body = """        reward = 0.0
        reward += self._progress_reward(action, result)
        reward += self._execution_reward(result)
        reward += self._quality_reward(
            action=action,
            current_config=current_config,
            expected_config=expected_config,
            original_config=original_config,
            error_message=error_message,
            result=result,
            metadata=metadata,
        )
        reward += self._penalty_reward(state=state, result=result, current_config=current_config)

        return round(self._clamp_01(reward), 4)"""

new_body = """        prog = self._progress_reward(action, result)
        exec = self._execution_reward(result)
        det_score, hide_score, llm_score, qual = self._quality_reward(
            action=action,
            current_config=current_config,
            expected_config=expected_config,
            original_config=original_config,
            error_message=error_message,
            result=result,
            metadata=metadata,
        )
        pen = self._penalty_reward(state=state, result=result, current_config=current_config)
        
        reward = prog + exec + qual + pen
        reward_clamped = round(self._clamp_01(reward), 4)
        
        components = {
            "progress": prog,
            "execution": exec,
            "deterministic": float(det_score or 0.0),
            "hidden": float(hide_score or 0.0),
            "llm_judge": float(llm_score or 0.0),
            "penalty": pen,
            "total": reward_clamped
        }
        return reward_clamped, components"""

code = code.replace(old_body, new_body)

# 3. Change `_quality_reward` return signature and body
old_qual_sig = """    def _quality_reward(
        self,
        action: str,
        current_config: str,
        expected_config: str,
        original_config: str,
        error_message: str,
        result: dict[str, Any],
        metadata: dict[str, Any],
    ) -> float:"""

new_qual_sig = """    def _quality_reward(
        self,
        action: str,
        current_config: str,
        expected_config: str,
        original_config: str,
        error_message: str,
        result: dict[str, Any],
        metadata: dict[str, Any],
    ) -> tuple[float, float, float, float]:"""

code = code.replace(old_qual_sig, new_qual_sig)

old_qual_early_return = """        if not current_config or not expected_config:
            return 0.0010101"""
new_qual_early_return = """        if not current_config or not expected_config:
            return 0.0, 0.0, 0.0, 0.0010101"""
code = code.replace(old_qual_early_return, new_qual_early_return)

old_qual_return = """        quality_reward = 0.0
        quality_reward += self.QUALITY_WEIGHTS["deterministic"] * self._clamp_01(deterministic_score)
        quality_reward += self.QUALITY_WEIGHTS["hidden"] * self._clamp_01(hidden_pass_rate or 0.0)
        quality_reward += self.QUALITY_WEIGHTS["llm"] * self._clamp_01(llm_average)

        return quality_reward"""

new_qual_return = """        quality_reward = 0.0
        quality_reward += self.QUALITY_WEIGHTS["deterministic"] * self._clamp_01(deterministic_score)
        quality_reward += self.QUALITY_WEIGHTS["hidden"] * self._clamp_01(hidden_pass_rate or 0.0)
        quality_reward += self.QUALITY_WEIGHTS["llm"] * self._clamp_01(llm_average)

        return self._clamp_01(deterministic_score), self._clamp_01(hidden_pass_rate or 0.0), self._clamp_01(llm_average), quality_reward"""

code = code.replace(old_qual_return, new_qual_return)

with open("env/rewards.py", "w") as f:
    f.write(code)