sciworld/impl/sciworld_agent_loop.py

"""SciWorld AgentLoop for verl —— 把 AgentGym 的 HTTP 环境接进 verl 的多轮 rollout。

verl 的 multi_turn 是围绕 tool call 设计的; SciWorld 是"环境返回纯文本 observation"
的形态, 所以按 notes 的建议实现一个 AgentLoop 子类。

协议与我们自建评测/ROSE/OPD 完全一致(便于横向对比):
  system(通用 Qwen) + user(24 动作 REACT 说明) + assistant("OK. I'll follow…")
  + user(task_description + look around) + [assistant 动作 / user observation] * N
动作解析取 "Action:" 之后首行; 每条轨迹的 reward = 最终 score(0~1)。

数据里每行需要:
  item_id      : "sciworld_{games下标}"
  raw_prompt   : 开场三条消息(list[dict])
  agent_name   : "sciworld"
"""
import asyncio
import logging
import os
from typing import Any
from uuid import uuid4

import aiohttp

from verl.experimental.agent_loop.agent_loop import AgentLoopBase, AgentLoopOutput, register
from verl.utils.profiler import simple_timer
from verl.utils.rollout_trace import rollout_trace_op
from verl.workers.rollout.replica import TokenOutput

logger = logging.getLogger(__file__)
logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN"))

ENV_BASE = os.getenv("SCIWORLD_ENV_BASE", "http://127.0.0.1:36001")
ENV_TIMEOUT = float(os.getenv("SCIWORLD_ENV_TIMEOUT", "300"))


def parse_action(text: str) -> str:
    tail = text.split("Action:")[-1] if "Action:" in text else text
    for line in tail.strip().splitlines():
        line = line.strip()
        if line:
            return line.strip("`* ")
    return "look around"


@register("sciworld")
class SciWorldAgentLoop(AgentLoopBase):
    """多轮 SciWorld: 学生出动作 -> HTTP 环境返回 observation -> 拼回对话, 直到 done 或轮数上限。"""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.prompt_length = self.rollout_config.prompt_length
        self.response_length = self.rollout_config.response_length
        mt = getattr(self.rollout_config, "multi_turn", None)
        self.max_turns = int(getattr(mt, "max_user_turns", 0) or 20)
        self.max_obs_len = int(getattr(mt, "max_tool_response_length", 0) or 2048)

    async def _post(self, session, path, payload, retries=3):
        """环境侧任何接口都可能返回 {"error": ...} 而非期望字段, 这里统一重试。"""
        last = None
        for attempt in range(retries):
            try:
                async with session.post(f"{ENV_BASE}/{path}", json=payload,
                                        timeout=aiohttp.ClientTimeout(total=ENV_TIMEOUT)) as r:
                    js = await r.json()
                if isinstance(js, dict) and "error" in js:
                    last = js["error"]
                    await asyncio.sleep(1.0 + attempt)
                    continue
                return js
            except Exception as e:
                last = repr(e)
                await asyncio.sleep(1.0 + attempt)
        logger.warning("sciworld %s failed after %d retries: %s", path, retries, last)
        return {"error": last}

    @rollout_trace_op
    async def run(self, sampling_params: dict[str, Any], priority: int = 0, **kwargs) -> AgentLoopOutput:
        priority = int(priority)
        try:
            with open("/home/haojinw2/agentloop_trace.log", "a") as _f:
                _f.write(f"RUN kwargs={sorted(kwargs.keys())}\n")
        except Exception:
            pass
        messages = list(kwargs["raw_prompt"])
        item_id = kwargs.get("item_id", "sciworld_0")
        data_idx = int(str(item_id).split("_")[-1])
        metrics = {}

        self._assert_mm_supported(False)
        prompt_ids = await self.ct_build_initial_tokens(messages)
        runtime_ids = list(prompt_ids)
        min_gs, max_gs = None, None      # 记录每轮生成所用的权重版本(异步框架用来算 staleness)
        response_mask: list[int] = []
        response_logprobs: list[float] = []
        score = 0.0
        turns = 0

        async with aiohttp.ClientSession() as session:
            try:
                with open("/home/haojinw2/agentloop_trace.log", "a") as _f:
                    _f.write(f"CREATE item={item_id}\n")
            except Exception:
                pass
            created = await self._post(session, "create", {}, retries=8)
            if "id" not in created:
                await asyncio.sleep(20)           # 等看护重启环境服务
                created = await self._post(session, "create", {}, retries=8)
            if "id" not in created:
                logger.warning("sciworld create failed, skip this trajectory: %s", created)
                return AgentLoopOutput(
                    prompt_ids=prompt_ids, response_ids=[], response_mask=[],
                    reward_score=0.0, num_turns=0, metrics={"num_preempted": -1},
                    extra_fields={"sciworld_score": 0.0, "sciworld_turns": 0,
                                  "min_global_steps": 0, "max_global_steps": 0,
                                  "sciworld_env_error": 1})
            env_id = created["id"]
            try:
                obs = await self._post(session, "reset", {"id": env_id, "data_idx": data_idx})
                if "observation" not in obs and "task_description" not in obs:
                    logger.warning("sciworld reset failed: %s", obs)
                    obs = {"task_description": "", "observation": "", "done": True, "score": 0}
                first = (obs.get("task_description", "") + "\n" + obs.get("observation", "")).strip()
                convo = list(messages) + [{"role": "user", "content": first}]
                merge_result, response_mask, response_logprobs = await self.ct_merge_non_assistant_msg(
                    messages, convo, runtime_ids, response_mask, response_logprobs)
                runtime_ids = merge_result.token_ids
                done = bool(obs.get("done", False))
                score = float(obs.get("score", 0) or 0) / 100.0

                while not done and turns < self.max_turns:
                    if len(runtime_ids) - len(prompt_ids) >= self.response_length:
                        break
                    with simple_timer("generate_sequences", metrics):
                        rid = (f"det-{priority}"
                               if getattr(self.rollout_config, "full_determinism", False) else uuid4().hex)
                        out: TokenOutput = await self.server_manager.generate(
                            request_id=rid, prompt_ids=runtime_ids,
                            sampling_params=sampling_params, priority=priority)
                    gs = (out.extra_fields or {}).get("global_steps")
                    if gs is not None:
                        min_gs = gs if min_gs is None else min(min_gs, gs)
                        max_gs = gs if max_gs is None else max(max_gs, gs)
                    merge_result, response_mask, response_logprobs = await self.ct_merge_assistant_token(
                        runtime_ids, out.token_ids, response_mask, response_logprobs,
                        assistant_logprobs=getattr(out, "log_probs", None))
                    runtime_ids = merge_result.token_ids
                    action_text = self.tokenizer.decode(out.token_ids, skip_special_tokens=True)
                    turns += 1

                    try:
                        st = await self._post(session, "step",
                                              {"id": env_id, "action": parse_action(action_text)})
                        obs_text = str(st.get("observation", ""))[: self.max_obs_len]
                        done = bool(st.get("done", False))
                        score = max(score, float(st.get("score", 0) or 0) / 100.0)
                    except Exception as e:  # 环境异常 = 该轨迹终止, 与 AgentGym 一致
                        logger.warning("sciworld step failed: %s", e)
                        obs_text, done = f"Error: {e}", True

                    prev = list(convo)
                    convo = prev + [{"role": "assistant", "content": action_text},
                                    {"role": "user", "content": obs_text}]
                    if not done:
                        merge_result, response_mask, response_logprobs = await self.ct_merge_non_assistant_msg(
                            prev + [{"role": "assistant", "content": action_text}], convo,
                            runtime_ids, response_mask, response_logprobs)
                        runtime_ids = merge_result.token_ids
            finally:
                try:
                    await self._post(session, "close", {"id": env_id})
                except Exception:
                    pass

        response_ids = runtime_ids[len(prompt_ids):]
        response_mask = response_mask[: len(response_ids)]
        if len(response_mask) < len(response_ids):
            response_mask += [0] * (len(response_ids) - len(response_mask))
        metrics.setdefault("num_preempted", -1)

        return AgentLoopOutput(
            prompt_ids=prompt_ids,
            response_ids=response_ids[: self.response_length],
            response_mask=response_mask[: self.response_length],
            response_logprobs=(response_logprobs[: self.response_length] if response_logprobs else None),
            reward_score=score,
            num_turns=turns,
            metrics=metrics,
            extra_fields={"sciworld_score": score, "sciworld_turns": turns,
                          "min_global_steps": min_gs if min_gs is not None else 0,
                          "max_global_steps": max_gs if max_gs is not None else 0},
        )