-
Notifications
You must be signed in to change notification settings - Fork 21
Expand file tree
/
Copy pathtest_rl_processing.py
More file actions
167 lines (146 loc) · 7.01 KB
/
Copy pathtest_rl_processing.py
File metadata and controls
167 lines (146 loc) · 7.01 KB
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
156
157
158
159
160
161
162
163
164
165
166
167
from typing import ( # Optional was already here, this is fine.
Any,
Dict,
List,
Optional,
Union,
)
import pytest
from eval_protocol.agent.models import StepData
from eval_protocol.models import EvaluateResult, Message as RewardKitMessage, StepOutput
from eval_protocol.rl_processing import RLDataAligner
class TestRLDataAligner:
def create_mock_step_data(
self,
system_step_index: int,
assistant_turn_index: Optional[Union[int, str]] = None,
content: str = "assistant action",
) -> StepData:
"""Helper to create a StepData instance for testing."""
obs_msgs = [RewardKitMessage(role="user", content="Hello")]
action = {"type": "text", "content": content}
res_msgs = [*obs_msgs, RewardKitMessage(role="assistant", content=content)]
step_info = {}
if assistant_turn_index is not None:
step_info["assistant_turn_index"] = assistant_turn_index
return StepData(
system_step_index=system_step_index,
observation_data=obs_msgs,
action_taken=action,
resulting_messages_history=res_msgs,
step_info=step_info,
)
def test_align_single_rollout_with_step_outputs(self):
aligner = RLDataAligner()
rollout_id = "rollout1"
# User's reward function output
user_eval_result = EvaluateResult(
score=0.8,
step_outputs=[
StepOutput(
step_index=0, base_reward=0.25, reason="First action good"
), # Matches assistant_turn_index 0
StepOutput(
step_index="turn_1", base_reward=0.75, reason="Second action better"
), # Matches assistant_turn_index "turn_1"
],
)
# System's collected StepData
# RLRolloutWorker should populate step_info with 'assistant_turn_index'
step_data_list = [
self.create_mock_step_data(system_step_index=0, assistant_turn_index=0, content="Action 1"),
self.create_mock_step_data(
system_step_index=1, assistant_turn_index="intermediate_tool_step"
), # No user reward for this
self.create_mock_step_data(system_step_index=2, assistant_turn_index="turn_1", content="Action 2"),
self.create_mock_step_data(
system_step_index=3, assistant_turn_index=2, content="Action 3"
), # No user reward for this
]
aligned_step_data = aligner.align_data_for_rl_processing(
current_eval_result=user_eval_result,
current_step_data_list=step_data_list,
rollout_id=rollout_id,
)
assert len(aligned_step_data) == 4
# Check base_rewards
assert aligned_step_data[0].base_reward == 0.25 # Matched step_index 0
assert aligned_step_data[1].base_reward is None # No matching step_index "intermediate_tool_step"
assert aligned_step_data[2].base_reward == 0.75 # Matched step_index "turn_1"
assert aligned_step_data[3].base_reward is None # No StepOutput for step_index 2
def test_align_single_rollout_no_step_outputs(self):
aligner = RLDataAligner()
rollout_id = "rollout2"
user_eval_result = EvaluateResult(score=0.9, reason="Overall score only")
step_data_list = [self.create_mock_step_data(system_step_index=0, assistant_turn_index=0)]
aligned_step_data = aligner.align_data_for_rl_processing(
current_eval_result=user_eval_result,
current_step_data_list=step_data_list,
rollout_id=rollout_id,
)
assert aligned_step_data[0].base_reward is None
def test_align_single_rollout_empty_step_outputs(self):
aligner = RLDataAligner()
rollout_id = "rollout3"
user_eval_result = EvaluateResult(score=0.7, step_outputs=[]) # Empty list
step_data_list = [self.create_mock_step_data(system_step_index=0, assistant_turn_index=0)]
aligned_step_data = aligner.align_data_for_rl_processing(
current_eval_result=user_eval_result,
current_step_data_list=step_data_list,
rollout_id=rollout_id,
)
assert aligned_step_data[0].base_reward is None
def test_align_step_output_index_not_in_step_data_info(self):
"""Test when a StepOutput.step_index has no corresponding assistant_turn_index in StepData."""
aligner = RLDataAligner()
rollout_id = "rollout4"
user_eval_result = EvaluateResult(
score=0.5,
step_outputs=[StepOutput(step_index="non_existent_turn", base_reward=1.0)],
)
step_data_list = [self.create_mock_step_data(system_step_index=0, assistant_turn_index=0)]
aligned_step_data = aligner.align_data_for_rl_processing(
current_eval_result=user_eval_result,
current_step_data_list=step_data_list,
rollout_id=rollout_id,
)
assert aligned_step_data[0].base_reward is None # No match
def test_align_step_data_missing_assistant_turn_index_in_info(self):
"""Test when StepData.step_info is missing the 'assistant_turn_index' key."""
aligner = RLDataAligner()
rollout_id = "rollout5"
user_eval_result = EvaluateResult(score=0.5, step_outputs=[StepOutput(step_index=0, base_reward=1.0)])
# Create StepData *without* 'assistant_turn_index' in step_info
step_data_list = [
StepData(
system_step_index=0,
observation_data=[],
action_taken={},
resulting_messages_history=[],
step_info={}, # No assistant_turn_index
)
]
aligned_step_data = aligner.align_data_for_rl_processing(
current_eval_result=user_eval_result,
current_step_data_list=step_data_list,
rollout_id=rollout_id,
)
assert aligned_step_data[0].base_reward is None # Cannot map
def test_align_preserves_other_step_data_fields(self):
aligner = RLDataAligner()
rollout_id = "rollout6"
user_eval_result = EvaluateResult(score=0.8, step_outputs=[StepOutput(step_index=0, base_reward=0.99)])
original_step_data = self.create_mock_step_data(system_step_index=0, assistant_turn_index=0)
original_step_data.policy_logprobs = {"logp": -0.1}
original_step_data.policy_value_estimate = 0.5
original_step_data.advantage = -0.05 # Should remain untouched by this aligner
step_data_list = [original_step_data]
aligned_step_data = aligner.align_data_for_rl_processing(
current_eval_result=user_eval_result,
current_step_data_list=step_data_list,
rollout_id=rollout_id,
)
assert aligned_step_data[0].base_reward == 0.99
assert aligned_step_data[0].policy_logprobs == {"logp": -0.1}
assert aligned_step_data[0].policy_value_estimate == 0.5
assert aligned_step_data[0].advantage == -0.05 # Ensure unrelated fields are not wiped