强化学习最佳实践

1pip install https://ark-public-example-cn-beijing.tos-cn-beijing.volces.com/ark-sdk/ark_sdk-0.2.11.tar.gz

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16ark init workspace ark_rl_project --template rl_search_mcp_demo
#初始化项目命令执行成功后工作区结构如下：
#<ark_rl_project>
#├── data
#│   └── search_dataset_dev_100.jsonl
#│   └── search_dataset.jsonl
#├── plugins
#│   ├── draft_rollout_arkitect.py
#│   └── rollout.py
#│   └── llm_grader.py
#│   └── test_utils.py
#├── job.py
#├── README.md
#└── requirements.txt
#└── test_faas.py
#└── ......

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
168
169
170
171
172
173
174
175MAX_STEPS = 6
TOOL_CALL_MAX_RETRY = 10
TOOL_CALL_TIMEOUT = 20.0  # 单次工具调用最大时长（秒
rate_limiter = AsyncLimiter(max_rate=20, time_period=1)

LOCAL_TEST_MODEL = "doubao-seed-2-0-mini-260428"

search_tool = {
    "type": "function",
    "function": {
        "name": "web_search",
        "description": "互联网搜索",
        "parameters": {
            "type": "object",
            "properties": {
                "Query": {
                    "type": "string",
                    "description": "搜索的关键词，1~100个字符(过长会截断)，不支持多词搜索",
                },
                "Count": {
                    "type": "integer",
                    "description": "搜索结果数量",
                    "default": 10,
                },
            },
            "required": ["Query"],
        },
    },
}

@rollout(
    name="code_sandbox_grader_single_grader",
    description="code_sandbox_grader",
    runtime=Runtime(
        instance=PluginInstance.CPU2MEM4,
        timeout=900,
    ),
)
@coze_monitor
async def demo_rollout(
    context: PluginContext,
    proxy: RolloutInferenceProxy,
    sample: ChatCompletionSample,
) -> Optional[RolloutResult]:
    """
    训练时，每个样本会调用 n_sample 次rollout函数，每次调用rollout函数当前仅支持返回一条轨迹（wrap_client_inplace的client最后一次调用上下文）
    """
    # NOTE: 由于search mcp的封装，无法修改tool定义，使用直接调用api的方式
    # from arkitect.core.component.tool.tool_pool import ToolPool, build_tool_pool
    # from arkitect.core.component.tool.builder import build_mcp_clients_from_config
    # mcp_clients, cleanup = build_mcp_clients_from_config(CONFIG_FILE_PATH)
    # tool_pool = build_tool_pool(list(mcp_clients.values()))
    # tools = [tool.model_dump() for tool in await tool_pool.list_tools()]
    tools = [search_tool]
    schema = search_tool["function"]["parameters"]
    search_api_key = os.getenv("SEARCH_API_KEY")
    assert search_api_key, "env SEARCH_API_KEY not set"
    async with httpx.AsyncClient(
        base_url="https://open.feedcoopapi.com",
        headers={
            "Authorization": f"Bearer {search_api_key}",
            "Content-Type": "application/json",
        },
    ) as tool_client:
        req = sample.model_dump()
        messages = req.pop("messages")
        client = AsyncArk(base_url=proxy.url, api_key=proxy.jwt_token)

        wrap_inplace_trace(client)
        wrap_async_client_inplace(client, proxy=proxy)

        step = 0

        while True:
            completion = await client.chat.completions.create(
                # model 字段仅在本地测试时生效
                model=LOCAL_TEST_MODEL,
                messages=messages,
                tools=tools,
                # NOTE: 为了避免训练效果受影响，被wrap后的client无法在此处指定采样参数或thinking type, 需要在样本或任务超参中指定
            )
            messages.append(completion.choices[0].message.model_dump())
            step += 1

            if step >= MAX_STEPS:
                break
            if completion.choices[0].finish_reason != "tool_calls":
                break

            tool_calls = completion.choices[0].message.tool_calls
            for tool_call in tool_calls or []:
                try:
                    name = tool_call.function.name
                    params = json.loads(tool_call.function.arguments)
                    assert name == "web_search", f"tool name not web_search: {name}"
                    jsonschema.validate(instance=params, schema=schema)
                except (
                    json.JSONDecodeError,
                    AssertionError,
                    jsonschema.ValidationError,
                ) as e:
                    # NOTE: extra字段用于传递额外的信息给reward函数，此处通知reward需要惩罚为-1分
                    logger.error(f"tool call format check failed: {repr(e)}")
                    return RolloutResult(
                        status=PluginStatus.SUCCESS,
                        extra={
                            "reward": -1,
                        },
                    )
                res = await exec_tool_call(
                    tool_client=tool_client,
                    name=name,
                    arguments=params,
                )
                messages.append(
                    {
                        "role": "tool",
                        "content": res,
                        "tool_call_id": tool_call.id,
                    }
                )

    return RolloutResult(
        status=PluginStatus.SUCCESS,
    )

@backoff.on_exception(
    backoff.expo,
    RolloutRetryException,
    max_time=60,
    jitter=backoff.full_jitter,  # full jitter
    base=2,
    factor=0.4,
    max_value=20,
)
@observe(
    process_inputs=lambda d: {key: d["kwargs"][key] for key in ["name", "arguments"]},
)  # type: ignore
async def exec_tool_call(
    tool_client: httpx.AsyncClient,
    name: str,
    arguments: dict,
) -> str | list[ChatCompletionContentPartParam]:
    # from arkitect.core.component.tool.utils import convert_to_chat_completion_content_part_param
    # if isinstance(tool_client, ToolPool):
    #     res = await tool_client.execute_tool(name, arguments)
    #     content = convert_to_chat_completion_content_part_param(res)
    #     return content
    async with rate_limiter:
        try:
            resp = await asyncio.wait_for(
                tool_client.post(
                    "search_api/web_search",
                    json={**arguments, "SearchType": "web"},
                ),
                timeout=TOOL_CALL_TIMEOUT,
            )
            data = resp.json()
            resp.raise_for_status()
            if data.get("Result"):
                return json.dumps(data, ensure_ascii=False)

            # NOTE: 可以选择一些错误返回RETRY进行rollout重试，重试达到上线后导致任务失败，默认exception为Failure会重试后拉挂任务
            raise RolloutRetryException(
                f"web_search tool call got resp without Result: {data}"
            )

        except asyncio.TimeoutError:
            raise RolloutRetryException(
                f"tool call timeout after {TOOL_CALL_TIMEOUT}s"
            )
        except httpx.HTTPStatusError as e:
            raise RolloutRetryException(
                f"tool call HTTPStatusError: {e.response.status_code} {e.response.text}"
            )

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
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223# NOTE: 避免thinking模型因为重复输出等问题影响训练
MAX_COMPLETION_TOKENS = 4096

SYSTEM_PROMPT = """你是一位严谨的内容评估专家。
你的核心任务是：依据提供的【问题】和一份【正确答案列表】（列表中的每个答案都被视为该问题的有效解答），来判断我提供的【预测答案】是否正确回答了该【问题】。你应该首先给出你的判断依据，然后给出你的判断结果（即“正确”或“错误”）。
判断标准如下：
1. 【预测答案】不需要与【正确答案列表】中的任何答案在字面上完全相同，但应该在语义上相同。
2. 【正确答案列表】中的每个答案都可以被视为问题的正确答案，【预测答案】应该至少与其中之一在语义上是相同的。
3. 你必须仔细阅读【问题】和【预测答案】，并仔细考虑【预测答案】是否真的正确回答了【问题】，不可以因为【预测答案】中包含了正确答案的字眼而认为【预测答案】正确。
4. 对于模型没有认真回答问题，只是通过列举很多答案骗分的情况，请你判断为错误。
"""

USER_PROMPT_TEMPLATE = """
###问题###
{}
###正确答案列表###
{}
###预测答案###
{}
现在请开始你的判断：
"""

@group_grader(
    name="llm_group_grader",
    description="llm_group_grader",
    runtime=Runtime(
        instance=PluginInstance.CPU1MEM2,
        max_concurrency=32,
        timeout=900,
    ),
)
async def llm_grader(
    context: PluginContext,
    sample: ChatCompletionSample,
    trajectories: List[Trajectory],
):
    if (
        not sample.extra
        or "answer" not in sample.extra
        or sample.messages[-1].role != "user"
    ):
        return GroupGraderResult(
            rewards=[],
            metrics={},
            status=PluginStatus.DISCARD,
            error="sample is not valid",
        )

    ark = AsyncArk()
    reward_list = []
    tasks = []

    @backoff.on_exception(
        wait_gen=backoff.expo,
        exception=(
            ArkAPIConnectionError,
            ArkRateLimitError,
            ArkInternalServerError,
            json.JSONDecodeError,
        ),
        max_time=100,
        jitter=backoff.full_jitter,  # full jitter
        base=2,
        factor=0.4,
        max_value=20,  # 退避重试最大时长（秒）
    )
    async def single_grader(traj: Trajectory) -> float:
        # 返回rollout过程中提前设置的reward
        if traj.extra and "reward" in traj.extra:
            return traj.extra["reward"]

        format_score = validate_tool_calls(traj)
        if format_score is not None:
            logger.info(
                f"prompt_id: {sample.extra.get('prompt_id')} eval_res: {format_score} format check failed. {traj.messages}"
            )
            return format_score
        resp = await ark.chat.completions.create(
            model="doubao-seed-2-0-lite-260428",
            max_completion_tokens=MAX_COMPLETION_TOKENS,
            temperature=0,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {
                    "role": "user",
                    # NOTE: 仅reward最后的总结，并未考虑历史tool_calls
                    "content": USER_PROMPT_TEMPLATE.format(
                        sample.messages[-1].content,
                        sample.extra["answer"],
                        traj.messages[-1].content,
                    ),
                },
            ],
            response_format={
                "type": "json_schema",
                "json_schema": {
                    "name": "评估结果",
                    "description": "根据你作为内容评估专家的角色和评估标准进行评估的结果",
                    "schema": {
                        "type": "object",
                        "properties": {
                            "rationale": {
                                "type": "string",
                                "description": "详细判断理由",
                            },
                            "judgement": {
                                "type": "string",
                                "enum": ["正确", "错误"],
                                "description": "判断结果",
                            },
                        },
                        "required": ["rationale", "judgement"],
                    },
                    "strict": True,
                },
            },
        )

        eval_res = json.loads(resp.choices[0].message.content)
        logger.info(f"prompt_id: {sample.extra.get('prompt_id')} eval_res: {eval_res}")
        if eval_res["judgement"] == "正确":
            return 1.0
        else:
            return 0.0

    for traj in trajectories:
        tasks.append(asyncio.create_task(single_grader(traj)))

    results = await asyncio.gather(*tasks, return_exceptions=True)
    for res in results:
        if isinstance(res, Exception):
            return GroupGraderResult(
                rewards=[], metrics={}, status=PluginStatus.FAILURE, error=str(res)
            )
        reward_list.append(res)

    return GroupGraderResult(
        rewards=reward_list, metrics={}, status=PluginStatus.SUCCESS, error=""
    )

def validate_tool_calls(traj: Trajectory) -> Optional[float]:
    # 如果最后一轮是tool_calls则视为参数不对或者超出轮次，0分惩罚
    if traj.messages[-1].tool_calls:
        for tool_call in traj.messages[-1].tool_calls:
            try:
                params = json.loads(tool_call.function.arguments)
                assert tool_call.function.name == "web_search", (
                    f"tool name not web_search: {tool_call.function.name}"
                )
                jsonschema.validate(
                    instance=params,
                    schema={
                        "type": "object",
                        "properties": {
                            "Query": {
                                "type": "string",
                            },
                            "Count": {
                                "type": "integer",
                            },
                        },
                        "required": ["Query"],
                    },
                )

            except Exception as e:
                return -1.0
        return 0.0

    if not traj.messages[-1].content:
        return -1.0

    # 如果模型没有执行工具直接回答，-1分惩罚
    if all([not msg.tool_calls for msg in traj.messages]):
        return -1.0
    return None

def main():
    import asyncio

    sample = ChatCompletionSample(
        **{
            "messages": [
                {
                    "role": "user",
                    "content": "通过景栗科技的私域运营服务和与薪勤科技的产品共创，哪两个公司在各自的领域实现了用户增长或应用上架？",
                }
            ],
            "thinking": {"type": "enabled"},
            "extra": {"answer": "景栗科技和薪勤科技"},
        }
    )
    trajectories = [
        Trajectory(
            **{
                "messages": [
                    {
                        "role": "assistant",
                        "content": "景栗科技的私域运营服务和薪勤科技的产品共创实现了用户增长或应用上架。",
                    }
                ],
                "finish_reason": "stop",
                "usage": {},
            }
        ),
        Trajectory(
            **{
                "messages": [
                    {
                        "role": "assistant",
                        "content": "景栗科技",
                    }
                ],
                "finish_reason": "stop",
                "usage": {},
            }
        ),
    ]
    res = asyncio.run(llm_grader({}, sample, trajectories))
    logger.info(f"grader done with result: {res}")

if __name__ == "__main__":
    main()

1 python job.py

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@rollout(
    name="demo_rollout",
    runtime=Runtime(
        instance=PluginInstance.CPU1MEM2,
        max_concurrency=100,
        min_replicas=1,
        max_replicas=10,
        timeout=900,
    ),
)
@coze_monitor
def demo_rollout(  # sync的函数，会由@rollout转换成线程并发的async函数来提升运行效率
    context: PluginContext,
    proxy: RolloutInferenceProxy,
    sample: ChatCompletionSample,
) -> Optional[RolloutResult]:
    """
    训练时，每个样本会调用 n_sample 次rollout函数，每次调用rollout函数当前仅支持返回一条轨迹（wrap_client_inplace的client最后一次调用上下文）
    """
    # 使用训练感知客户端 - 仅此一行不同！
    client = Ark(base_url=proxy.url, api_key=proxy.jwt_token)
    # 可以使用openai的client
    # from openai import OpenAI
    # from ark_sdk.core.plugin.rollout.proxy import wrap_client_inplace
    # client = OpenAI(base_url=proxy.url, api_key=proxy.jwt_token)
    wrap_inplace_trace(client)
    wrap_client_inplace(client, proxy=proxy)
    req = sample.model_dump()
    messages = req.pop("messages")
    tools = (req.pop("tools") or []) + rollout_tools
    step = 0

    while True:
        # 步骤2: 发起模型请求，由于模型在收到工具执行结果后仍然可能有工具调用意愿，因此需要多次请求
        completion = client.chat.completions.create(
            # model 字段仅在本地测试时生效
            model=LOCAL_TEST_MODEL,
            messages=messages,
            tools=tools,
        )
        # 注意：训练逻辑已经自动处理，无需手动判断模式或调用process_completion！
        messages.append(completion.choices[0].message.model_dump())
        step += 1
        logger.info(f"completion: {completion.choices[0].message.model_dump()}")

        if step >= 30:
            break

        if completion.choices[0].finish_reason != "tool_calls":
            # 模型最终总结，没有调用工具意愿
            break
        tool_calls = completion.choices[0].message.tool_calls
        for tool_call in tool_calls or []:
            tool_name = tool_call.function.name
            if tool_name == "get_current_weather":
                # 步骤 3：调用外部工具
                try:
                    args = json.loads(tool_call.function.arguments)
                    tool_result = get_current_weather(**args)
                except Exception as e:
                    logger.error(f"get_current_weather error: {e}")

                    # 重试，rollout请求，超过重试次数会DISCARD掉本条数据（包裹所有n smaple）
                    # return RolloutResult(status=PluginStatus.RETRY, error=str(e))

                    # 丢弃，丢弃本样本（包裹所有n smaple）
                    # return RolloutResult(status=PluginStatus.DISCARD, error="")

                    # 失败，返回FAILURE重试本样本，超过重试次数会拉挂任务。如果不处理exception，默认会返回FAILURE
                    # return RolloutResult(status=PluginStatus.FAILURE, error=str(e))

                    # extra字段用于传递额外的信息给reward函数，此处通知reward需要惩罚为-1分
                    return RolloutResult(
                        status=PluginStatus.SUCCESS,
                        extra={
                            "reward": -1,
                        },
                    )

                # 步骤 4：回填工具结果，并获取模型总结回复
                messages.append(
                    {
                        "role": "tool",
                        "content": tool_result,
                        "tool_call_id": tool_call.id,
                    }
                )

    # 失败，丢弃本样本（包裹所有n smaple）
    # return RolloutResult(status=PluginStatus.DISCARD, error="")
    # 默认return None则视为rollout成功
    return None

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@rollout(
    runtime=Runtime(
        instance=PluginInstance.CPU1MEM2,
        max_concurrency=100,
        min_replicas=1,
        max_replicas=10,
    ),
)
@coze_monitor
def demo_rollout(
    context: PluginContext,
    proxy: RolloutInferenceProxy,
    sample: ChatCompletionSample,
) -> Optional[RolloutResult]:
    client = Ark(base_url=proxy.url, api_key=proxy.jwt_token)
    # 可选：添加cozeloop trace，用于debug监控模型调用输入输出
    wrap_inplace_trace(client)
    req = sample.model_dump()
    messages = req.pop("messages")
    tools = (req.pop("tools") or []) + rollout_tools
    while True:
        # NOTE:强化学习场景特殊逻辑
        # 请求前发给 proxy 检查历史信息、更新内部状态（当 message 更改历史时，proxy 会自动截断内部维护的 token ids）。
        proxy.update_state_from_messages([ChatCompletionMessage(**msg) for msg in messages])
        # 步骤2: 发起模型请求，由于模型在收到工具执行结果后仍然可能有工具调用意愿，因此需要多次请求
        completion = client.chat.completions.create(
            # model 字段仅在本地测试时生效
            model=LOCAL_TEST_MODEL,
            messages=messages,
            tools=tools,
            # NOTE:强化学习场景特殊逻辑
            extra_body=proxy.get_extra_body(),
            extra_headers=proxy.headers,
        )
        assert isinstance(completion, ChatCompletion)
        # NOTE:强化学习场景特殊逻辑
        proxy.process_completion(ChatCompletionResponse(**completion.model_dump()))
        logger.info(f"completion: {completion.choices[0].message.model_dump()}")
        messages.append(completion.choices[0].message.model_dump())
        if completion.choices[0].finish_reason != "tool_calls":
            # 模型最终总结，没有调用工具意愿
            break
        tool_calls = completion.choices[0].message.tool_calls
        for tool_call in tool_calls or []:
            tool_name = tool_call.function.name
            if tool_name == "get_current_weather":
                # 步骤 3：调用外部工具
                args = json.loads(tool_call.function.arguments)
                tool_result = get_current_weather(**args)
                # 步骤 4：回填工具结果，并获取模型总结回复
                messages.append(
                    {
                        "role": "tool",
                        "content": tool_result,
                        "tool_call_id": tool_call.id,
                    }
                )
    return

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@group_grader(
    name="randaom_reward",
    runtime=Runtime(
        instance=PluginInstance.CPU1MEM2,
        max_concurrency=100,
        timeout=300,
    ),
)
def random_reward_fn(
    context: PluginContext,
    sample: ChatCompletionSample,
    trajectories: List[Trajectory],
) -> GroupGraderResult:
    """
    奖励函数：返回随机奖励

    参数:
    - trajectories: 完整的对话历史
    - sample: 样本数据，包含标准答案的字典

    返回:
    - list[float]: 奖励分数列表，每个分数对应一个候选回复（1.0表示完全匹配，0.0表示不匹配）

    依赖:
    - 数据集里的字典字段 extra 内需要携带 answer 字段。
    """
    rewards = [
        t.extra["reward"] if (t.extra and "reward" in t.extra) else random.random()
        for t in trajectories
    ]
    return GroupGraderResult(
        rewards=rewards, status=PluginStatus.SUCCESS, error="", metrics={}
    )