modelscope--ms-swift/docs/source/Instruction/GRPO/DeveloperGuide/gym_env.md

GYM环境训练

GYM 风格的环境训练把"模型 → 环境 → 奖励"这条链路封装成一个抽象接口，让 LLM 像 Agent 一样与环境进行多轮交互，每一步的奖励直接由环境给出，无需再单独写 reward 函数从轨迹里反推。本文先介绍接口，再用一个完整的自定义示例（FrozenLake）说明如何接入训练。

Gym 接口

GYM 源自 Gymnasium库。在 ms-swift 中我们定义了如下接口：

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）等。完整示例参见 入参示例。

如果需要在每轮 rollout 之间额外控制对话历史（例如动态压缩、注入额外提示），推荐直接继承 MultiTurnScheduler 并实现 on_trajectory_start / on_turn_end hook，或重写 step / run 方法，详见多轮训练文档。

启动训练

使用内置的 gym_scheduler 把 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 模式:

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 模式

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 环境

FrozenLake 是 OpenAI Gym 中的经典任务：智能体从起点出发，需要穿过一片冰湖到达终点，途中要避开冰窟。原始环境如上图所示。下面以纯文本版本（把上图网格直接渲染成 ASCII 字符）为例。

以下完整代码参考完整代码：frozen_lake_plugin。

1. 定义 Env

每条数据派生一张随机 4x4 地图（随机洞 + 随机 S/G 位置，BFS 校验保证可解）。单元含义：S 起点 / G 终点 / H 冰窟（踩到=失败）/ F 安全冰面 / P 玩家当前位置。

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 完成了控制逻辑：

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 该文件，注册随之生效：

# 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来控制地图生成的随机性：

{"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：

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

OpenEnv 环境训练

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。

1. 启动 OpenEnv 服务器

安装 OpenEnv 和 Sudoku 环境包（textarena 和 nltk 会作为依赖自动安装）：

pip install openenv
pip install git+https://huggingface.co/spaces/openenv/sudoku

使用提供的启动脚本启动本地服务器（默认端口 8000）。MAX_CONCURRENT_ENVS 需 ≥ 训练时的 num_generations：

TEXTARENA_ENV_ID=Sudoku-v0 MAX_CONCURRENT_ENVS=8 python examples/train/grpo/plugin/openenv/start_sudoku_server.py

数据集中将 base_url 指向本地服务器地址：

{"messages":[{"role":"user","content":"Play"}],"env_config":{"name":"openenv","base_url":"http://127.0.0.1:8000"}}

2. 自定义 Scheduler

继承 OpenEnvScheduler，实现 Sudoku 专用的动作解析、观察格式化和多组件奖励：

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 示例）：

奖励组件	计算方式	作用
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 托管地址：

{"messages":[{"role":"user","content":"Play"}],"env_config":{"name":"openenv","base_url":"http://127.0.0.1:8000"}}

5. 注册 Scheduler

sudoku_scheduler.py 末尾已包含注册代码，通过 --external_plugins 加载即可：

# sudoku_scheduler.py 末尾
from swift.rollout.multi_turn import multi_turns
multi_turns['sudoku_scheduler'] = SudokuScheduler

6. 启动训练

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

注意事项

• 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