From 39f1ccf378390dfe78f6048e6efb6bfcf38dd4b7 Mon Sep 17 00:00:00 2001 From: xhluca Date: Thu, 7 May 2026 12:20:06 -0400 Subject: [PATCH] Add missing extract / prepare scripts for task-intent generation The pipeline section of the README jumps straight from `a3-explore` to `scripts/generate_task_intents.py`, but the latter consumes prepared prompts under `outputs/task_intents/prompts//` that the released code did not provide a way to build. This adds the two intermediate scripts and documents the full Step 2 in the README, including which part of the exploration trajectory is fed to the Task Designer and how `step_num` is selected (random sample of N steps per trajectory, dropping step 0). - scripts/extract_chat_messages.py: agentlab `step_*.pkl.gz` -> per-step JSON with chat messages, goal, screenshot path, etc. - scripts/prepare_tasks_intents_prompts.py: per-trajectory random step sample + appended TASK_INTENT_PROMPT_TEMPLATE user turn, written as `task_.step_.json` matching `generate_task_intents.py`'s expected input layout. - README.md: expand Step 2 to spell out the four sub-steps and the sampling semantics. Co-Authored-By: Claude Opus 4.7 (1M context) --- README.md | 28 ++- scripts/extract_chat_messages.py | 211 ++++++++++++++++++ scripts/prepare_tasks_intents_prompts.py | 267 +++++++++++++++++++++++ 3 files changed, 505 insertions(+), 1 deletion(-) create mode 100644 scripts/extract_chat_messages.py create mode 100644 scripts/prepare_tasks_intents_prompts.py diff --git a/README.md b/README.md index c89bf9a..c422340 100644 --- a/README.md +++ b/README.md @@ -53,10 +53,36 @@ python scripts/create_personas.py ### Step 2: Generate task intents (via exploration) ```bash +# 2a. Run the exploration agent. Trajectories are saved as agentlab pickles +# under $AGENTLAB_EXP_ROOT///step_*.pkl.gz. a3-explore -python scripts/generate_task_intents.py + +# 2b. Extract chat messages from each step pickle into a parallel JSON tree +# at outputs/chat_messages///step_*.json. +python scripts/extract_chat_messages.py --find-latest + +# 2c. For each trajectory, randomly sample N steps (default 3, skipping step 0), +# append the TASK_INTENT_PROMPT_TEMPLATE as a final user turn, and write each +# prompt as outputs/task_intents/prompts//task_.step_.json. +python scripts/prepare_tasks_intents_prompts.py --find-latest + +# 2d. Send each prepared prompt to the Task Designer LLM. Completions land in +# outputs/task_intents/completions///. +python scripts/generate_task_intents.py \ + --exploration-model \ + --model ``` +**What gets passed to the Task Designer:** the *full chat-message history* of a +single exploration step (system prompt, goal, every prior assistant action and +observation up to that step) with one extra user turn appended that contains +`TASK_INTENT_PROMPT_TEMPLATE`. The step is selected by uniform random sampling +of `step_*.json` files in the trajectory, after dropping `step_0` (the initial +observation, which has no agent actions yet). The number of samples per +trajectory is controlled by `--num_samples` (default `3`); the number of intents +requested per prompt is `--num_intents` (default `2`). Sampling is seeded by +`--seed` (default `42`) for reproducibility. + ### Step 3: Create A3-Synth task configs ```bash python scripts/create_synth_configs.py diff --git a/scripts/extract_chat_messages.py b/scripts/extract_chat_messages.py new file mode 100644 index 0000000..0b591e1 --- /dev/null +++ b/scripts/extract_chat_messages.py @@ -0,0 +1,211 @@ +""" +Given a directory of agentlab results (e.g. agentlab_results/2025-09-17_01-28-02_genericagent-qwen-qwen3-32b-on-exploration-exploration), +extract all the chat messages into json files matching each of the pickles. Save them into a new directory called chat_messages. +""" + +import pickle +import gzip +import orjson +import argparse +import re +from pathlib import Path +from tqdm import tqdm + + +def find_latest_results_dir(base_dir, agent_name, benchmark="exploration"): + """ + Find the latest results directory matching the agent name and benchmark. + + Directory naming pattern: YYYY-MM-DD_HH-MM-SS_genericagent-{agent}-on-{benchmark}-{suffix} + """ + base_path = Path(base_dir) + if not base_path.exists(): + print(f"Base directory {base_dir} does not exist") + return None + + # Pattern to match directories with timestamp and keywords + # Looking for: on-{benchmark} and containing agent_name + matching_dirs = [] + + for d in base_path.iterdir(): + if not d.is_dir(): + continue + name = d.name + # Check if it contains the benchmark pattern (e.g., "on-exploration") + if f"on-{benchmark}" not in name: + continue + # Check if it contains the agent name + if agent_name not in name: + continue + # Check if it starts with a date pattern + if not re.match(r"^\d{4}-\d{2}-\d{2}_\d{2}-\d{2}-\d{2}_", name): + continue + matching_dirs.append(d) + + if not matching_dirs: + print(f"No directories found matching agent '{agent_name}' and benchmark '{benchmark}'") + return None + + # Sort by name (timestamp is at the start, so lexicographic sort works) + matching_dirs.sort(key=lambda x: x.name, reverse=True) + latest = matching_dirs[0] + + print(f"Found {len(matching_dirs)} matching directories") + print(f"Selected latest: {latest.name}") + + return latest + + +def extract_chat_messages(results_dir, output_dir, force_reprocess=False): + """Extract chat messages from all pickle files in the results directory""" + results_path = Path(results_dir) + + if not results_path.exists(): + print(f"Results directory {results_dir} does not exist") + return + + # Create chat_messages output directory + output_dir = Path(output_dir) + output_dir.mkdir(exist_ok=True) + + # Create subdirectory matching the results directory name + output_subdir = output_dir / results_path.name + output_subdir.mkdir(exist_ok=True) + + print(f"Extracting chat messages from {results_dir}") + print(f"Output directory: {output_subdir}") + + # First pass: count all pickle files for progress bar + all_pickle_files = [] + subdirs = [d for d in results_path.iterdir() if d.is_dir()] + + print("Scanning for pickle files...") + for subdir in tqdm(subdirs, desc="Scanning directories"): + step_files = list(subdir.glob("step_*.pkl.gz")) + step_files.sort(key=lambda x: int(x.stem.split("_")[1].split(".")[0])) + all_pickle_files.extend([(subdir, pickle_file) for pickle_file in step_files]) + + print(f"Found {len(all_pickle_files)} pickle files to process") + + # Process each pickle file with progress bar + processed_count = 0 + skipped_count = 0 + + for subdir, pickle_file in tqdm(all_pickle_files, desc="Extracting messages"): + # Create corresponding output directory + sub_output_dir = output_subdir / subdir.name + sub_output_dir.mkdir(exist_ok=True) + + # Check if output file already exists (caching) + json_filename = pickle_file.stem.replace(".pkl", ".json") + json_path = sub_output_dir / json_filename + + if json_path.exists() and not force_reprocess: + # Check if output is newer than input (basic cache validation) + if json_path.stat().st_mtime > pickle_file.stat().st_mtime: + skipped_count += 1 + continue + + # Extract chat messages + pickle_file_path = Path(pickle_file).resolve() + try: + with gzip.open(pickle_file_path, "rb") as f: + data = pickle.load(f) + + # Extract chat messages from the obs + if "chat_messages" not in data.agent_info: + messages = [] + chat_model_args = {} + goal = "" + else: + messages = data.agent_info["chat_messages"].to_openai() + chat_model_args = data.agent_info.extra_info.get("chat_model_args", {}) + goal = data.obs.get("goal", "") + + except Exception as e: + tqdm.write(f"Error processing {pickle_file_path}: {e}") + continue + + pickle_file_path = Path(pickle_file).resolve() + image_file_path = ( + str(pickle_file_path) + .replace(".pkl.gz", ".png") + .replace("step", "screenshot_step") + ) + + extracted = { + "pickle_path": str(pickle_file_path), + "image_path": str(image_file_path), + "goal": goal, + "chat_model_args": chat_model_args, + "step_num": data.step, + "messages": messages, + } + # Save to JSON file with matching name + with open(json_path, "w") as f: + f.write(orjson.dumps(extracted).decode()) + + processed_count += 1 + + print(f"Processed {processed_count} pickle files") + print(f"Skipped {skipped_count} files (already up-to-date)") + print(f"Chat messages saved to {output_subdir}") + + +def main(): + parser = argparse.ArgumentParser( + description="Extract chat messages from agentlab results" + ) + parser.add_argument( + "--path", + "-p", + default="", + help="Path to the agentlab results directory", + ) + parser.add_argument( + "--find-latest", + "-l", + metavar="AGENT", + help="Find the latest results directory for the specified agent name (e.g., 'gemini-3-flash-preview')", + ) + parser.add_argument( + "--benchmark", + "-b", + default="exploration", + help="Benchmark name to match when using --find-latest (default: exploration)", + ) + parser.add_argument( + "--base-dir", + default="agentlab_results", + help="Base directory to search in when using --find-latest (default: agentlab_results)", + ) + parser.add_argument( + "--force", + "-f", + action="store_true", + help="Force reprocessing of all files, ignoring cache", + ) + parser.add_argument( + "--output_dir", + "-o", + default="outputs/chat_messages", + help="Directory to save extracted chat messages", + ) + + args = parser.parse_args() + + # Determine results directory + if args.find_latest: + results_dir = find_latest_results_dir(args.base_dir, args.find_latest, args.benchmark) + if results_dir is None: + return + elif args.path: + results_dir = args.path + else: + parser.error("Either --path or --find-latest must be specified") + + extract_chat_messages(results_dir, args.output_dir, force_reprocess=args.force) + + +if __name__ == "__main__": + main() diff --git a/scripts/prepare_tasks_intents_prompts.py b/scripts/prepare_tasks_intents_prompts.py new file mode 100644 index 0000000..9d1b2c0 --- /dev/null +++ b/scripts/prepare_tasks_intents_prompts.py @@ -0,0 +1,267 @@ +""" +Script to prepare prompts for task generation from agentlab chat messages. +Randomly samples n (default 4) step_*.json files from each trajectory and adds a user message +requesting the model to come up with a task following the original instructions. +""" + +import datetime +import random +import re +import argparse +from pathlib import Path +import shutil +from tqdm import tqdm + +import orjson + +from agent_as_annotators.utils import load_all_model_configs + + +def find_latest_exp_dir(base_dir, agent_name, benchmark="exploration"): + """ + Find the latest experiment directory matching the agent name and benchmark. + + Directory naming pattern: YYYY-MM-DD_HH-MM-SS_genericagent-{agent}-on-{benchmark}-{suffix} + """ + base_path = Path(base_dir) + if not base_path.exists(): + print(f"Base directory {base_dir} does not exist") + return None + + # Pattern to match directories with timestamp and keywords + matching_dirs = [] + + for d in base_path.iterdir(): + if not d.is_dir(): + continue + name = d.name + # Check if it contains the benchmark pattern (e.g., "on-exploration") + if f"on-{benchmark}" not in name: + continue + # Check if it contains the agent name + if agent_name not in name: + continue + # Check if it starts with a date pattern + if not re.match(r"^\d{4}-\d{2}-\d{2}_\d{2}-\d{2}-\d{2}_", name): + continue + matching_dirs.append(d) + + if not matching_dirs: + print(f"No directories found matching agent '{agent_name}' and benchmark '{benchmark}'") + return None + + # Sort by name (timestamp is at the start, so lexicographic sort works) + matching_dirs.sort(key=lambda x: x.name, reverse=True) + latest = matching_dirs[0] + + print(f"Found {len(matching_dirs)} matching directories") + print(f"Selected latest: {latest.name}") + + return latest + +from agent_as_annotators.prompts import ( + TASK_INTENT_PROMPT_TEMPLATE, + WEBARENA_ANNOTATOR_INSTRUCTIONS +) + +def is_trajectory_dir(d): + if d.name.startswith('.'): + return False + if not d.is_dir(): + return False + if d.name.startswith('_'): + return False + if not "_on_" in d.name: + return False + return True + +def parse_trajectory_dir_name(d): + # first split by . to get .{num}_{seed} + before, _, num_seed = d.name.rpartition('.') + num, _, seed = num_seed.partition('_') + # get YYYY-MM-DD_HH-MM-SS_--on- + date_str, _, rest = before.partition('_') + time_str, _, rest = rest.partition('_') + agent, _, task = rest.partition('-on-') + return { + 'date': date_str, + 'time': time_str, + 'agent': agent, + 'task_name': task, + 'task_num': int(num), + 'seed': int(seed) + } + +def get_trajectory_dirs(exp_dir): + return [d for d in exp_dir.iterdir() if is_trajectory_dir(d)] + +def get_step_number(step_file): + return int(step_file.stem.split('_')[1].split('.')[0]) + +def main(exp_dir, output_dir, num_samples=2, seed=42, num_intents=2): + exp_dir = Path(exp_dir) + output_dir = Path(output_dir) + trajectory_dirs = get_trajectory_dirs(exp_dir) + + print(f"Found {len(trajectory_dirs)} trajectory directories in {exp_dir}") + + # if the output file exists, remove it to start fresh, but prompt the user + if output_dir.exists(): + response = input(f"Output directory '{output_dir}' already exists. (d)elete / (c)ontinue? ") + if response.lower() == 'd': + # delete the output directory + shutil.rmtree(output_dir) + elif response.lower() == 'c': + print("Continuing...") + else: + print("Invalid response. Exiting...") + exit(0) + + output_dir.mkdir(parents=True, exist_ok=True) + random.seed(seed) # For reproducibility + + fails = 0 + for traj_dir in tqdm(trajectory_dirs, desc="Processing trajectories"): + traj_info = parse_trajectory_dir_name(traj_dir) + step_files = list(traj_dir.glob("step_*.json")) + step_files.sort(key=get_step_number) + step_files = step_files[1:] # remove step 0 + + if len(step_files) == 0: + fails += 1 + print(f"No step files found in {traj_dir}, skipping.") + continue + + # Sample up to 4 random step files + sampled_files = random.sample(step_files, min(num_samples, len(step_files))) + + for step_file in sampled_files: + task_num = traj_info['task_num'] + step_num = get_step_number(step_file) + + save_fname = f"task_{task_num}.step_{step_num}.json" + save_path = output_dir / save_fname + + if save_path.exists(): + tqdm.write(f"Skipping {save_path} because it already exists") + continue + with open(step_file, 'rb') as f: + step_data = orjson.loads(f.read()) + + step_data['task_num'] = task_num + step_data['step_num'] = step_num + step_data['traj_info'] = traj_info # Add trajectory info to step data + + # Add user message requesting task generation + user_message = { + "role": "user", + "content": TASK_INTENT_PROMPT_TEMPLATE.format( + annotator_instructions=WEBARENA_ANNOTATOR_INSTRUCTIONS, + num_intents=num_intents + ) + } + + step_data['messages'].append(user_message) + + # add to messages + + with open(save_path, 'w') as f: + f.write(orjson.dumps(step_data).decode('utf-8')) + + if fails > 0: + print(f"Failed to process {fails} trajectories") + + print(f"Prepared prompts saved in {output_dir}") + +if __name__ == "__main__": + this_file = Path(__file__) + model_configs_path = this_file.parent.parent / "configs" / "model_configs.json" + shorthands_to_configs = load_all_model_configs(model_configs_path) + + parser = argparse.ArgumentParser(description="Prepare prompts for task generation from chat messages") + parser.add_argument( + "--exp_dir", + type=str, + default="", + help="Directory containing trajectory subdirectories with step_*.json files", + ) + parser.add_argument( + "--find-latest", + "-l", + metavar="AGENT", + help="Find the latest exp directory for the specified agent name (e.g., 'gemini-3-flash-preview')", + ) + parser.add_argument( + "--benchmark", + "-b", + default="exploration", + help="Benchmark name to match when using --find-latest (default: exploration)", + ) + parser.add_argument( + "--base-dir", + default="outputs/chat_messages", + help="Base directory to search in when using --find-latest (default: outputs/chat_messages)", + ) + parser.add_argument( + "--output-dir-base", + type=str, + default="outputs/task_intents/prompts", + help="Base directory to save the prepared prompts. The final output directory will be //", + ) + parser.add_argument( + "-e", + "--exploration-model", + type=str, + default=None, + choices=list(shorthands_to_configs.keys()), + help="Model used for exploration (used for output directory naming). If --find-latest is provided, this is inferred automatically.", + ) + parser.add_argument( + "--num_samples", + type=int, + default=3, + help="Number of step files to sample from each trajectory", + ) + parser.add_argument( + "--num_intents", + type=int, + default=2, + help="Number of intents to request from the model per prompt", + ) + parser.add_argument( + "--seed", + type=int, + default=42, + help="Random seed for sampling", + ) + args = parser.parse_args() + + # Determine exp directory + if args.find_latest: + exp_dir = find_latest_exp_dir(args.base_dir, args.find_latest, args.benchmark) + if exp_dir is None: + exit(1) + elif args.exp_dir: + exp_dir = args.exp_dir + else: + parser.error("Either --exp_dir or --find-latest must be specified") + + # Determine exploration model: --find-latest supersedes -e + if args.find_latest and args.exploration_model is None: + # Use --find-latest value as the exploration model + exploration_model = args.find_latest + if exploration_model not in shorthands_to_configs: + parser.error(f"Model '{exploration_model}' from --find-latest not found in model configs. " + f"Available: {list(shorthands_to_configs.keys())}") + elif args.exploration_model is not None: + exploration_model = args.exploration_model + else: + parser.error("Either --find-latest or -e/--exploration-model must be specified") + + # Build output directory with model name subdirectory (matching generate_task_intents.py convention) + exploration_model_name = shorthands_to_configs[exploration_model]["kwargs"]["model_name"] + exploration_model_dirname = exploration_model_name.replace('/', '_') + output_dir = Path(args.output_dir_base) / exploration_model_dirname + print(f"Output directory: '{output_dir}'") + + main(exp_dir, output_dir, num_samples=args.num_samples, seed=args.seed, num_intents=args.num_intents) \ No newline at end of file