wehub-resource-sync
361 lines
15 KiB
Markdown
361 lines
15 KiB
Markdown
# GYM环境训练
|
||
|
||
GYM 风格的环境训练把"模型 → 环境 → 奖励"这条链路封装成一个抽象接口，让 LLM 像 Agent 一样与环境进行多轮交互，每一步的奖励直接由环境给出，无需再单独写 reward 函数从轨迹里反推。本文先介绍接口，再用一个完整的自定义示例（FrozenLake）说明如何接入训练。
|
||
|
||
## Gym 接口
|
||
|
||
GYM 源自 [Gymnasium库](https://github.com/Farama-Foundation/Gymnasium)。在 ms-swift 中我们定义了如下接口：
|
||
|
||
```python
|
||
class Env(ABC):
|
||
|
||
def __init__(self, env_config):
|
||
"""env_config 来自数据集每行的 env_config 列，可承载初始化参数"""
|
||
self.env_config = env_config
|
||
|
||
@abstractmethod
|
||
async def reset(self, config: RolloutInferRequest) -> Tuple[str, Dict[str, Any], str]:
|
||
"""
|
||
Returns:
|
||
- observation: 作为首轮 user 消息发送给模型
|
||
- info: 调试/日志信息，记录到 completions.jsonl
|
||
- system_message: 本条轨迹的 system prompt
|
||
"""
|
||
pass
|
||
|
||
@abstractmethod
|
||
async def step(self, action: Messages) -> Tuple[str, float, bool, Dict[str, Any]]:
|
||
"""
|
||
Args:
|
||
action: 截止当前的完整对话消息，最后一条即模型最新回复
|
||
Returns:
|
||
- next_observation: 下一轮 user 消息
|
||
- reward: 当前 step 奖励
|
||
- done: 轨迹是否结束
|
||
- info: 调试/日志信息
|
||
"""
|
||
pass
|
||
|
||
@abstractmethod
|
||
async def close(self):
|
||
"""释放资源"""
|
||
pass
|
||
```
|
||
|
||
`reset` 接收到的 `RolloutInferRequest` 包含数据集行的 `messages`、`data_dict`（额外列，包括 `env_config`）等。完整示例参见 [入参示例](./multi_turn.md#多轮规划器-multiturnscheduler)。
|
||
|
||
> 如果需要在每轮 rollout 之间额外控制对话历史（例如动态压缩、注入额外提示），推荐直接继承 `MultiTurnScheduler` 并实现 `on_trajectory_start` / `on_turn_end` hook，或重写 `step` / `run` 方法，详见[多轮训练文档](./multi_turn.md#自定义多轮交互逻辑)。
|
||
|
||
## 启动训练
|
||
|
||
使用内置的 [gym_scheduler](https://github.com/modelscope/ms-swift/blob/main/swift/rollout/multi_turn.py) 把 env 串到多轮 rollout 中。
|
||
|
||
`GYMScheduler` 基于通用 hook 协议实现：
|
||
- 继承 `MultiTurnScheduler`，无需自定义 `run` 方法
|
||
- 实现 `on_trajectory_start`（调用 `env.reset`）和 `on_turn_end`（调用 `env.step`）
|
||
- 同时适用于 server mode（`run()`）和 colocate mode（`run_multi_turn()`）
|
||
|
||
用户自定义的 env 通过 `--external_plugins your_plugin.py` 加载，plugin 里执行 `envs['my_env'] = MyEnv` 完成注册（下文 FrozenLake 示例完整演示）。
|
||
|
||
**Colocate 模式**:
|
||
|
||
```bash
|
||
megatron rlhf \
|
||
--rlhf_type grpo \
|
||
--vllm_mode colocate \
|
||
--external_plugins examples/megatron/grpo/multi_turn/frozen_lake_plugin.py \
|
||
--multi_turn_scheduler gym_scheduler \
|
||
--gym_env frozen_lake \
|
||
--use_gym_env true \
|
||
--max_turns 10 \
|
||
...
|
||
|
||
# swift rlhf 同理
|
||
```
|
||
|
||
|
||
**Server 模式**
|
||
|
||
```bash
|
||
swift rollout \
|
||
--model xxx \
|
||
--use_gym_env true \
|
||
--external_plugins examples/megatron/grpo/multi_turn/frozen_lake_plugin.py \
|
||
--multi_turn_scheduler gym_scheduler \
|
||
--gym_env frozen_lake \
|
||
--max_turns 10
|
||
|
||
# trainer 侧需要加 --vllm_server_pass_dataset true，把 env_config 等额外列透传给 rollout 端
|
||
megatron rlhf --vllm_mode server --vllm_server_pass_dataset true ...
|
||
# or swift rlhf --vllm_mode server --vllm_server_pass_dataset true ...
|
||
```
|
||
|
||
|
||
环境选择有两种方式：
|
||
- 通过 `--gym_env env_name` 全局指定（同一脚本里所有 prompt 共用一个 env）；
|
||
- 在每行数据的 `env_config.name` 中指定（适用于多环境混合场景，每条数据可指向不同 env，会覆盖 `--gym_env`）。
|
||
|
||
## 示例：从零写一个 FrozenLake 环境
|
||
|
||
<img src="https://gymnasium.farama.org/_images/frozen_lake.gif" width="220" alt="FrozenLake 环境示意图（来源：Gymnasium 官方文档）" />
|
||
|
||
[FrozenLake](https://gymnasium.farama.org/environments/toy_text/frozen_lake/) 是 OpenAI Gym 中的经典任务：智能体从起点出发，需要穿过一片冰湖到达终点，途中要避开冰窟。原始环境如上图所示。下面以纯文本版本（把上图网格直接渲染成 ASCII 字符）为例。
|
||
|
||
以下完整代码参考完整代码：[frozen_lake_plugin](https://github.com/modelscope/ms-swift/blob/main/examples/megatron/grpo/multi_turn/frozen_lake_plugin.py)。
|
||
|
||
**1. 定义 Env**
|
||
|
||
每条数据派生一张随机 4x4 地图（随机洞 + 随机 S/G 位置，BFS 校验保证可解）。单元含义：`S` 起点 / `G` 终点 / `H` 冰窟（踩到=失败）/ `F` 安全冰面 / `P` 玩家当前位置。
|
||
|
||
```python
|
||
class FrozenLakeEnv(Env):
|
||
def __init__(self, env_config):
|
||
super().__init__(env_config)
|
||
self.size = int(env_config.get('size', 4))
|
||
self.p = float(env_config.get('p', 0.8))
|
||
seed = env_config.get('seed')
|
||
self.seed = int(seed) if seed is not None else None
|
||
|
||
async def reset(self, config: RolloutInferRequest):
|
||
self.grid = generate_random_map(size=self.size, p=self.p, seed=self.seed)
|
||
...
|
||
return observation, {'seed': self.seed}, SYSTEM_PROMPT
|
||
|
||
async def step(self, action: Messages):
|
||
move = _parse_action(action[-1]['content']) # <action>up|down|left|right</action>
|
||
# 推进一格、判断 G / H；外层 max_turns 由 scheduler 兜底
|
||
if cell == 'G': return obs, 1.0, True, {'status': 'goal'}
|
||
if cell == 'H': return obs, 0.0, True, {'status': 'hole'}
|
||
...
|
||
```
|
||
|
||
**2. GYMScheduler 的 hook 实现**
|
||
|
||
框架内置的 `GYMScheduler` 基于多轮 hook 完成了控制逻辑：
|
||
|
||
```python
|
||
class GYMScheduler(MultiTurnScheduler):
|
||
def on_trajectory_start(self, requests):
|
||
# 为每个请求创建 env，调用 env.reset，注入初始 observation
|
||
for req in requests:
|
||
env = self._create_env(req.data_dict.get('env_config', {}))
|
||
observation, info, system_message = env.reset(req)
|
||
req.messages = [system_msg, user_msg(observation)]
|
||
self._envs[req.uuid] = env
|
||
|
||
def on_turn_end(self, req, response_choice, current_turn):
|
||
# 调用 env.step，累积 reward，返回 done + rollout_infos
|
||
next_obs, reward, done, info = env.step(deepcopy(req.messages))
|
||
self._total_rewards[req.uuid] += reward
|
||
return {
|
||
'done': done,
|
||
'rollout_infos': {
|
||
'total_reward': self._total_rewards[req.uuid],
|
||
'step_rewards': [...],
|
||
}
|
||
}
|
||
|
||
def step(self, req, response_choice, current_turn):
|
||
# 注入下一帧 observation 到 user message
|
||
if self._pending_obs.get(req.uuid):
|
||
req.messages.append({'role': 'user', 'content': next_obs})
|
||
return {'infer_request': req}
|
||
```
|
||
|
||
用户只需实现 Env 接口，无需关心多轮控制细节。
|
||
|
||
**3. 注册**
|
||
|
||
将 env 类挂到 swift 的 `envs` 注册表里。`--external_plugins` 在训练启动时会 import 该文件，注册随之生效：
|
||
|
||
```python
|
||
# examples/megatron/grpo/multi_turn/frozen_lake_plugin.py
|
||
from swift.rollout.gym_env import Env, envs
|
||
|
||
class FrozenLakeEnv(Env):
|
||
...
|
||
|
||
envs['frozen_lake'] = FrozenLakeEnv
|
||
```
|
||
|
||
**4. 准备数据集**
|
||
|
||
数据集在这里仅作占位符处理，数据构造由环境生成，和 `env_config.seed`来控制地图生成的随机性：
|
||
|
||
```json
|
||
{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":0}}
|
||
{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":1}}
|
||
...
|
||
{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":127}}
|
||
```
|
||
|
||
**5. （可选）叠加自定义 reward**
|
||
|
||
设置 `--use_gym_env true` 后，env 给出的 `total_reward` 会自动作为一路奖励参与训练，无需再写 reward 函数。如果想在此之外再叠加自定义信号（如格式/长度等），通过 `--reward_funcs` 传入即可，gym 奖励会作为额外一列与 reward_funcs 拼在一起，由 `--reward_weights` 统一加权。例如同时启用一个格式校验 reward：
|
||
|
||
```bash
|
||
megatron rlhf ... --use_gym_env true --reward_funcs format --reward_weights 0.2 1.0
|
||
# reward_weights 末位对应 gym 的 total_reward
|
||
```
|
||
|
||
**6. 训练**
|
||
|
||
运行脚本参考：[`examples/megatron/grpo/multi_turn/frozen_lake.sh`](https://github.com/modelscope/ms-swift/blob/main/examples/megatron/grpo/multi_turn/frozen_lake.sh)
|
||
|
||
|
||
## OpenEnv 环境训练
|
||
|
||
[OpenEnv](https://github.com/huggingface/openenv) 是 HuggingFace 开源的 Agentic RL 环境框架，通过 WebSocket 与环境服务器交互。与上文 FrozenLake 的本地 `Env` 接口不同，OpenEnv 将环境逻辑放在独立的服务进程中，swift 通过 `OpenEnvScheduler` + `OpenEnvWrapper` 与之通信。
|
||
|
||
### 架构对比
|
||
|
||
| 特性 | 内置 Gym (`GYMScheduler`) | OpenEnv (`OpenEnvScheduler`) |
|
||
|------|--------------------------|------------------------------|
|
||
| 环境运行位置 | 训练进程内（Python 对象） | 独立服务器（WebSocket 通信） |
|
||
| 环境接口 | 继承 `Env`，实现 `reset/step/close` | 服务器提供 HTTP/WebSocket API |
|
||
| 注册方式 | `--external_plugins` + `envs` 注册表 | `--external_plugins` + `multi_turns` 注册表 |
|
||
| 适用场景 | 轻量本地环境（FrozenLake 等） | 复杂服务端环境（TextArena、CARLA 等） |
|
||
| 并发控制 | 无需 | 内置 Semaphore 限制并发连接 |
|
||
|
||
### OpenEnvScheduler
|
||
|
||
`OpenEnvScheduler` 继承 `GYMScheduler`，将本地 `Env` 替换为 `OpenEnvWrapper`（WebSocket 客户端）。核心设计：
|
||
|
||
- **`_create_env`**：创建 `OpenEnvWrapper`，连接 OpenEnv 服务器
|
||
- **`on_trajectory_start`**：为每个请求创建 wrapper，调用 `reset()`，用 Semaphore 限制并发（默认 4）
|
||
- **`on_turn_end`**：解析模型输出，调用 `wrapper.step()`，累积奖励
|
||
- **`parse_action`**（可覆盖）：将模型文本解析为 action dict，默认 `json.loads`
|
||
- **`format_observation`**（可覆盖）：将服务器返回的 observation 格式化为字符串，默认 `json.dumps`
|
||
|
||
用户通过继承 `OpenEnvScheduler` 并覆盖 `parse_action`、`format_observation`、`on_trajectory_start`、`on_turn_end` 来适配具体环境。
|
||
|
||
### 示例：Sudoku 环境
|
||
|
||
以 TextArena Sudoku 为例，模型需要通过 `[row col number]` 格式下棋，在 9x9 数独棋盘上填入数字。完整代码参考：[sudoku_scheduler.py](https://github.com/modelscope/ms-swift/blob/main/examples/train/grpo/plugin/openenv/sudoku_scheduler.py)。
|
||
|
||
**1. 启动 OpenEnv 服务器**
|
||
|
||
安装 OpenEnv 和 Sudoku 环境包（textarena 和 nltk 会作为依赖自动安装）：
|
||
|
||
```bash
|
||
pip install openenv
|
||
pip install git+https://huggingface.co/spaces/openenv/sudoku
|
||
```
|
||
|
||
使用提供的启动脚本启动本地服务器（默认端口 8000）。`MAX_CONCURRENT_ENVS` 需 ≥ 训练时的 `num_generations`：
|
||
|
||
```bash
|
||
TEXTARENA_ENV_ID=Sudoku-v0 MAX_CONCURRENT_ENVS=8 python examples/train/grpo/plugin/openenv/start_sudoku_server.py
|
||
```
|
||
|
||
数据集中将 `base_url` 指向本地服务器地址：
|
||
|
||
```json
|
||
{"messages":[{"role":"user","content":"Play"}],"env_config":{"name":"openenv","base_url":"http://127.0.0.1:8000"}}
|
||
```
|
||
|
||
**2. 自定义 Scheduler**
|
||
|
||
继承 `OpenEnvScheduler`，实现 Sudoku 专用的动作解析、观察格式化和多组件奖励：
|
||
|
||
```python
|
||
from swift.rollout.multi_turn import OpenEnvScheduler
|
||
|
||
class SudokuScheduler(OpenEnvScheduler):
|
||
def __init__(self, *args, **kwargs):
|
||
super().__init__(*args, **kwargs)
|
||
self._last_content_len = {} # 内容差分跟踪
|
||
|
||
async def on_trajectory_start(self, requests):
|
||
# 创建环境、解析棋盘、生成 hints
|
||
# hints 包括「保证正确的走法」和候选数字列表
|
||
...
|
||
|
||
async def on_turn_end(self, infer_request, response_choice, current_turn):
|
||
# 解析 [row col number]，step 环境
|
||
# 计算 5 路奖励：空格选择 / 合法移动 / 重复惩罚 / 进度 / 正确
|
||
# 返回更新后的棋盘 + hints 作为下一轮观察
|
||
...
|
||
|
||
def parse_action(self, text):
|
||
import re
|
||
match = re.search(r'\[\s*(\d+)\s+(\d+)\s+(\d+)\s*\]', text)
|
||
if match:
|
||
row, col, num = match.groups()
|
||
return {"message": f"[{row} {col} {num}]"}
|
||
return {"message": "[1 1 1]"}
|
||
```
|
||
|
||
**多组件奖励系统**（参考 [TRL Sudoku 示例](https://github.com/huggingface/trl/blob/main/examples/notebooks/openenv_sudoku_grpo.ipynb)）：
|
||
|
||
| 奖励组件 | 计算方式 | 作用 |
|
||
|---------|---------|------|
|
||
| `empty_cell_reward` | 目标是空格 +1 / 覆盖已有 -1 | 引导模型选择合法位置 |
|
||
| `valid_move_reward` | 合法新走法 +1 / 警告 -0.5 / 无效 0 | 鼓励合法操作 |
|
||
| `repetition_reward` | 重复走法指数惩罚（-2^n，上限 -10） | 避免重复 |
|
||
| `progress_reward` | (已填充 - 初始) / (81 - 初始) | 衡量解题进度 |
|
||
| `correct_reward` | 环境返回的二值奖励 | 完全解出 |
|
||
|
||
组合奖励 = 各组件均值之和，提供比单一二值奖励更密集的学习信号。
|
||
|
||
**3. Hints 系统**
|
||
|
||
每轮交互中，scheduler 解析当前棋盘状态，为模型提供提示：
|
||
|
||
- **GUARANTEED MOVES**：只有一个候选数字的空格（可直接填入）
|
||
- **Other options**：2-3 个候选数字的空格
|
||
- **MOVES ALREADY TRIED**：已尝试过的走法（避免重复）
|
||
|
||
**4. 准备数据集**
|
||
|
||
数据集仅作占位符，实际棋盘由环境服务器生成。`base_url` 指向 OpenEnv 托管地址：
|
||
|
||
```json
|
||
{"messages":[{"role":"user","content":"Play"}],"env_config":{"name":"openenv","base_url":"http://127.0.0.1:8000"}}
|
||
```
|
||
|
||
**5. 注册 Scheduler**
|
||
|
||
`sudoku_scheduler.py` 末尾已包含注册代码，通过 `--external_plugins` 加载即可：
|
||
|
||
```python
|
||
# sudoku_scheduler.py 末尾
|
||
from swift.rollout.multi_turn import multi_turns
|
||
multi_turns['sudoku_scheduler'] = SudokuScheduler
|
||
```
|
||
|
||
**6. 启动训练**
|
||
|
||
```bash
|
||
swift rlhf \
|
||
--rlhf_type grpo \
|
||
--model Qwen/Qwen3.5-4B \
|
||
--dataset examples/train/grpo/plugin/openenv/sudoku.jsonl \
|
||
--external_plugins examples/train/grpo/plugin/openenv/sudoku_scheduler.py \
|
||
--enable_thinking false \
|
||
--max_completion_length 256 \
|
||
--use_gym_env true \
|
||
--multi_turn_scheduler sudoku_scheduler \
|
||
--max_turns 20 \
|
||
--use_vllm true \
|
||
--vllm_mode colocate \
|
||
...
|
||
```
|
||
|
||
运行脚本参考：[`examples/train/grpo/plugin/openenv/run_grpo_sudoku.sh`](https://github.com/modelscope/ms-swift/blob/main/examples/train/grpo/plugin/openenv/run_grpo_sudoku.sh)
|
||
|
||
### 注意事项
|
||
|
||
- **vLLM 模式**：以上示例使用 `--vllm_mode colocate`，vLLM 与训练共享 GPU。若使用 `--vllm_mode server`，需额外启动 `swift rollout` 作为 vLLM 服务器，且 `--multi_turn_scheduler` 和 `--max_turns` 参数应传给 `swift rlhf` 而非 `swift rollout`。
|
||
- **并发会话数**：`start_sudoku_server.py` 的 `MAX_CONCURRENT_ENVS` 需 ≥ 训练时的 `num_generations`。默认的 `python -m textarena_env.server.app` 只支持 1 个并发会话。
|
||
- **enable_thinking**：Sudoku 等环境不需要 CoT 推理，建议设置 `--enable_thinking false` 以减少 token 消耗。
|
||
- **同步 I/O**：`OpenEnvWrapper` 的 `reset()`/`step()` 是同步 WebSocket 调用。`OpenEnvScheduler` 的子类应使用 `asyncio.to_thread()` 包装这些调用以避免阻塞事件循环。
|
||
|
||
|
||
参考资料:
|
||
|
||
- https://gymnasium.farama.org/environments/toy_text/frozen_lake/
|
||
- https://github.com/alibaba/ROLL/tree/main/roll/pipeline/agentic/env/frozen_lake
|
||
- https://github.com/huggingface/openenv
|
||
- https://github.com/huggingface/trl/blob/main/examples/notebooks/openenv_sudoku_grpo.ipynb
|