| """Multi-GPU closed-loop ALFWorld evaluation from immutable best checkpoints.""" |
| from __future__ import annotations |
|
|
| import argparse |
| from collections import Counter, defaultdict |
| import heapq |
| import json |
| import os |
| from pathlib import Path |
| import random |
| import shutil |
| import sys |
| import uuid |
|
|
| PROJECT = Path("/share/shenmingxiang/closing-square") |
| WS = PROJECT / "verl_workspace" |
| sys.path.insert(0, str(WS / "third_party" / "Agent-R1")) |
| sys.path.insert(0, str(PROJECT)) |
|
|
| import pandas as pd |
| import torch |
| import torch.distributed as dist |
| from peft import PeftModel |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
| from recipes.alfworld.env.alfworld_wrapper import AlfworldTextworldEnv |
| from src.prompts import student_prompt, teacher_prompt |
| from src.prepare_structured_hints import build_hint_fields, render_hint, _load_traj |
|
|
| DEFAULT_MODEL = PROJECT / "models" / "Qwen2.5-3B-Instruct" |
| DATA_ROOT = WS / "data" / "alfworld" |
|
|
|
|
| def set_seed(seed: int) -> None: |
| random.seed(seed) |
| torch.manual_seed(seed) |
| torch.cuda.manual_seed_all(seed) |
|
|
|
|
| @torch.no_grad() |
| def constrained_action(model, tokenizer, prompt: str, actions: list[str], device, max_prompt_length: int) -> str: |
| old_side = tokenizer.truncation_side |
| tokenizer.truncation_side = "left" |
| encoded = tokenizer( |
| prompt, return_tensors="pt", add_special_tokens=True, |
| truncation=True, max_length=max_prompt_length, |
| ) |
| tokenizer.truncation_side = old_side |
| input_ids = encoded["input_ids"].to(device) |
| attention_mask = encoded["attention_mask"].to(device) |
| prompt_length = input_ids.shape[1] |
| token_to_action: dict[tuple[int, ...], str] = {} |
| sequences: list[tuple[int, ...]] = [] |
| for action in actions: |
| tokens = tuple(tokenizer(action, add_special_tokens=False)["input_ids"]) |
| token_to_action.setdefault(tokens, action) |
| sequences.append(tokens + (tokenizer.eos_token_id,)) |
| max_new_tokens = max(len(sequence) for sequence in sequences) |
|
|
| def allowed_tokens(_batch_id: int, generated: torch.Tensor) -> list[int]: |
| prefix = tuple(generated[prompt_length:].tolist()) |
| allowed = { |
| sequence[len(prefix)] |
| for sequence in sequences |
| if len(sequence) > len(prefix) and sequence[: len(prefix)] == prefix |
| } |
| return sorted(allowed) if allowed else [tokenizer.eos_token_id] |
|
|
| generated = model.generate( |
| input_ids=input_ids, |
| attention_mask=attention_mask, |
| do_sample=False, |
| temperature=1.0, |
| top_p=1.0, |
| top_k=50, |
| repetition_penalty=1.0, |
| num_beams=1, |
| max_new_tokens=max_new_tokens, |
| eos_token_id=tokenizer.eos_token_id, |
| pad_token_id=tokenizer.pad_token_id, |
| prefix_allowed_tokens_fn=allowed_tokens, |
| use_cache=True, |
| ) |
| continuation = generated[0, prompt_length:].tolist() |
| if tokenizer.eos_token_id in continuation: |
| continuation = continuation[: continuation.index(tokenizer.eos_token_id)] |
| key = tuple(continuation) |
| if key not in token_to_action: |
| raise RuntimeError(f"constrained decoding returned unknown token path: {continuation}") |
| return token_to_action[key] |
|
|
|
|
| def run_episode( |
| model, |
| tokenizer, |
| env, |
| item, |
| device, |
| max_steps: int, |
| max_prompt_length: int, |
| prompt_mode: str, |
| ): |
| obs, info = env.reset_with_info(item["game"]) |
| hint = item.get("H", "") |
| if prompt_mode == "teacher" and not hint: |
| expert_policy = [str(value) for value in info.get("policy_commands", [])] |
| hint = render_hint( |
| build_hint_fields(item["task_family"], _load_traj(item["game"]), expert_policy) |
| ) |
| history: list[str] = [] |
| done = False |
| steps = 0 |
| repeated = 0 |
| last_action = None |
| trajectory = [] |
| for _ in range(max_steps): |
| admissible = [str(value) for value in info.get("admissible_commands", [])] |
| if done or not admissible: |
| break |
| if prompt_mode == "teacher": |
| prompt = teacher_prompt(item["goal"], str(obs), admissible, hint, history) |
| else: |
| prompt = student_prompt(item["goal"], str(obs), admissible, history) |
| action = constrained_action(model, tokenizer, prompt, admissible, device, max_prompt_length) |
| repeated += int(action == last_action) |
| next_obs, _, done, info = env.step(action) |
| trajectory.append( |
| { |
| "step": steps, |
| "observation": str(obs), |
| "admissible": admissible, |
| "action": action, |
| "response": str(next_obs), |
| "success": bool(info.get("success", False)) if info else False, |
| } |
| ) |
| history.append(f"Action taken: {action}\nEnvironment response: {next_obs}") |
| obs = next_obs |
| last_action = action |
| steps += 1 |
| return { |
| "success": bool(info.get("success", False)) if info else False, |
| "steps": steps, |
| "repeated_actions": repeated, |
| "H": hint or None, |
| "trajectory": trajectory, |
| } |
|
|
|
|
| def atomic_json(path: Path, value, token: str) -> None: |
| temp = path.parent / f".{path.name}.tmp.{token}" |
| with temp.open("w", encoding="utf-8") as handle: |
| json.dump(value, handle, ensure_ascii=False, indent=2) |
| handle.flush() |
| os.fsync(handle.fileno()) |
| os.replace(temp, path) |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--name", required=True) |
| parser.add_argument("--adapter", default="none") |
| parser.add_argument("--model", type=Path, default=DEFAULT_MODEL) |
| parser.add_argument("--output", type=Path, required=True) |
| parser.add_argument("--max_steps", type=int, default=50) |
| parser.add_argument("--max_prompt_length", type=int, default=1536) |
| parser.add_argument("--seed", type=int, default=42) |
| parser.add_argument("--limit_per_split", type=int, default=0) |
| parser.add_argument("--limit_per_family", type=int, default=0) |
| parser.add_argument( |
| "--dataset_mode", choices=("standard", "pilot", "sft_val"), default="standard" |
| ) |
| parser.add_argument("--prompt_mode", choices=("student", "teacher"), default="student") |
| parser.add_argument( |
| "--hint_data", |
| type=Path, |
| default=WS / "data" / "privileged_plan_sft.jsonl", |
| ) |
| args = parser.parse_args() |
|
|
| world_size = int(os.environ.get("WORLD_SIZE", "1")) |
| distributed = world_size > 1 |
| if distributed: |
| dist.init_process_group("nccl") |
| rank = dist.get_rank() |
| local_rank = int(os.environ["LOCAL_RANK"]) |
| torch.cuda.set_device(local_rank) |
| else: |
| rank = 0 |
| local_rank = 0 |
| is_main = rank == 0 |
| device = torch.device("cuda", local_rank) |
| set_seed(args.seed) |
|
|
| items = [] |
| if args.dataset_mode == "sft_val": |
| by_game = {} |
| with args.hint_data.open("r", encoding="utf-8") as source: |
| for line in source: |
| row = json.loads(line) |
| if row["sft_split"] != "val": |
| continue |
| by_game.setdefault( |
| row["game_id"], |
| { |
| "split": "sft_val", |
| "game": row["game_id"], |
| "goal": row["goal"], |
| "task_family": row["task_family"], |
| "H": row["H"], |
| }, |
| ) |
| by_family = defaultdict(list) |
| for item in by_game.values(): |
| by_family[item["task_family"]].append(item) |
| rng = random.Random(args.seed) |
| for family in sorted(by_family): |
| family_items = sorted(by_family[family], key=lambda value: value["game"]) |
| rng.shuffle(family_items) |
| if args.limit_per_family: |
| family_items = family_items[: args.limit_per_family] |
| items.extend(family_items) |
| else: |
| source_root = DATA_ROOT / "pilot" / "pick_and_place" if args.dataset_mode == "pilot" else DATA_ROOT |
| for split in ("valid_seen", "valid_unseen"): |
| frame = pd.read_parquet(source_root / f"{split}.parquet") |
| if args.limit_per_split: |
| frame = frame.head(args.limit_per_split) |
| for _, row in frame.iterrows(): |
| extra = dict(row["extra_info"]) |
| items.append( |
| { |
| "split": split, |
| "game": str(extra["game_relative_path"]), |
| "goal": str(extra.get("goal_text", "")), |
| "task_family": str(extra["task_family"]), |
| } |
| ) |
| indexed_items = [(index, item) for index, item in enumerate(items) if index % world_size == rank] |
| if is_main: |
| print( |
| f"eval name={args.name} episodes={len(items)} world_size={world_size} " |
| f"adapter={args.adapter}", |
| flush=True, |
| ) |
|
|
| tokenizer = AutoTokenizer.from_pretrained(args.model) |
| tokenizer.pad_token = tokenizer.eos_token |
| tokenizer.padding_side = "right" |
| base = AutoModelForCausalLM.from_pretrained( |
| args.model, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True |
| ).to(device) |
| if args.adapter.lower() != "none": |
| adapter = Path(args.adapter).resolve() |
| model = PeftModel.from_pretrained(base, adapter) |
| else: |
| adapter = None |
| model = base |
| model.eval().requires_grad_(False) |
| model.config.use_cache = True |
|
|
| token_holder = [f"{os.getpid()}.{uuid.uuid4().hex}" if is_main else ""] |
| if distributed: |
| dist.broadcast_object_list(token_holder, src=0) |
| run_token = token_holder[0] |
| part_dir = args.output.parent / f".{args.output.name}.parts.{run_token}" |
| if is_main: |
| args.output.parent.mkdir(parents=True, exist_ok=True) |
| part_dir.mkdir(parents=False, exist_ok=False) |
| if distributed: |
| dist.barrier() |
|
|
| env = AlfworldTextworldEnv(data_root=str(DATA_ROOT), max_episode_steps=args.max_steps) |
| temp_part = part_dir / f".rank_{rank:02d}.jsonl.tmp" |
| final_part = part_dir / f"rank_{rank:02d}.jsonl" |
| with temp_part.open("w", encoding="utf-8") as handle: |
| for local_index, (global_index, item) in enumerate(indexed_items, start=1): |
| try: |
| result = run_episode( |
| model, tokenizer, env, item, device, args.max_steps, |
| args.max_prompt_length, args.prompt_mode, |
| ) |
| error = None |
| except Exception as exception: |
| result = { |
| "success": False, |
| "steps": 0, |
| "repeated_actions": 0, |
| "H": None, |
| "trajectory": [], |
| } |
| error = f"{type(exception).__name__}: {exception}" |
| finally: |
| env._close_env() |
| record = {"record_index": global_index, **item, **result, "error": error} |
| handle.write(json.dumps(record, ensure_ascii=False) + "\n") |
| if local_index % 5 == 0: |
| print(f"eval rank={rank} completed={local_index}/{len(indexed_items)}", flush=True) |
| handle.flush() |
| os.fsync(handle.fileno()) |
| os.replace(temp_part, final_part) |
| if distributed: |
| dist.barrier() |
|
|
| if is_main: |
| def lines(path: Path): |
| with path.open("r", encoding="utf-8") as source: |
| for line in source: |
| value = json.loads(line) |
| yield value["record_index"], line, value |
|
|
| streams = [lines(path) for path in sorted(part_dir.glob("rank_*.jsonl"))] |
| temp_output = args.output.parent / f".{args.output.name}.tmp.{run_token}" |
| merged_records = [] |
| with temp_output.open("w", encoding="utf-8") as output_handle: |
| for _, line, value in heapq.merge(*streams, key=lambda triple: triple[0]): |
| output_handle.write(line) |
| merged_records.append(value) |
| output_handle.flush() |
| os.fsync(output_handle.fileno()) |
| os.replace(temp_output, args.output) |
|
|
| groups = defaultdict(lambda: Counter(total=0, success=0, errors=0, steps=0, repeated=0)) |
| for record in merged_records: |
| for key in (record["split"], f"{record['split']}::{record['task_family']}", "pooled"): |
| groups[key]["total"] += 1 |
| groups[key]["success"] += int(record["success"]) |
| groups[key]["errors"] += int(record["error"] is not None) |
| groups[key]["steps"] += int(record["steps"]) |
| groups[key]["repeated"] += int(record["repeated_actions"]) |
| summary_groups = {} |
| for key, counts in sorted(groups.items()): |
| total = max(counts["total"], 1) |
| summary_groups[key] = { |
| **dict(counts), |
| "success_rate": counts["success"] / total, |
| "mean_steps": counts["steps"] / total, |
| "mean_repeated_actions": counts["repeated"] / total, |
| } |
| summary = { |
| "name": args.name, |
| "adapter": str(adapter) if adapter else None, |
| "base_model": str(args.model.resolve()), |
| "seed": args.seed, |
| "dataset_mode": args.dataset_mode, |
| "prompt_mode": args.prompt_mode, |
| "decoding": "greedy token-trie constrained to admissible commands", |
| "groups": summary_groups, |
| } |
| summary_path = args.output.with_suffix(args.output.suffix + ".summary.json") |
| atomic_json(summary_path, summary, run_token) |
| print(json.dumps(summary, ensure_ascii=False, indent=2), flush=True) |
| if distributed: |
| dist.barrier() |
| if is_main: |
| shutil.rmtree(part_dir) |
| if distributed: |
| dist.destroy_process_group() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
