import argparse import os import re import numpy as np from tqdm import tqdm # WLPlan Imports from wlplan.data import DomainDataset, ProblemDataset from wlplan.feature_generator import init_feature_generator from wlplan.planning import Atom, State, parse_domain, parse_problem ALL_DOMAINS = ["blocks", "gripper", "logistics", "visitall-from-everywhere"] SPLITS = ["train", "validation", "test-interpolation", "test-extrapolation"] def parse_traj_line_to_state(line, pred_map): """ Parses a line like "(on a b) (clear c)" into a wlplan State object. """ line = line.strip() if not line: return State([]) # Regex to find all (predicate arg1 arg2 ...) groups matches = re.findall(r"\(([\w-]+(?: [\w-]+)*)\)", line) atoms = [] for m in matches: parts = m.split() pred_name = parts[0] objs = parts[1:] if pred_name in pred_map: # Create Atom: (Predicate, [objects]) atoms.append(Atom(pred_map[pred_name], objs)) return State(atoms) def main(): parser = argparse.ArgumentParser() parser.add_argument("--data_dir", default="data") parser.add_argument("--output_dir", default="data/encodings/graphs") parser.add_argument("--model_dir", default="data/encodings/models") parser.add_argument("--iterations", type=int, default=2, help="WL iterations") parser.add_argument("--domain", type=str, default=None, help="Specific domain") args = parser.parse_args() # Make sure directories exists os.makedirs(args.output_dir, exist_ok=True) os.makedirs(args.model_dir, exist_ok=True) os.makedirs(os.path.join(args.data_dir, "pddl"), exist_ok=True) os.makedirs(os.path.join(args.data_dir, "states"), exist_ok=True) os.makedirs(os.path.join(args.data_dir, "plans"), exist_ok=True) os.makedirs(os.path.join(args.data_dir, "encodings"), exist_ok=True) os.makedirs(os.path.join(args.data_dir, "encodings", "graphs"), exist_ok=True) os.makedirs(os.path.join(args.data_dir, "encodings", "models"), exist_ok=True) domains_to_run = [args.domain] if args.domain else ALL_DOMAINS for domain_name in domains_to_run: print(f"\n=== Processing Domain: {domain_name} ===") # Paths domain_pddl = os.path.join(args.data_dir, "pddl", domain_name, "domain.pddl") train_states_dir = os.path.join(args.data_dir, "states", domain_name, "train") if not os.path.exists(domain_pddl): print(f" [Error] Domain PDDL not found: {domain_pddl}") continue # 1. Parse Domain try: wl_domain = parse_domain(domain_pddl) except Exception as e: print(f" [Error] Failed to parse domain: {e}") continue pred_map = {p.name: p for p in wl_domain.predicates} # 2. Initialize Feature Generator (ILG = Instance Learning Graph) feature_gen = init_feature_generator( feature_algorithm="wl", domain=wl_domain, graph_representation="ilg", iterations=args.iterations, pruning="none", multiset_hash=True, ) # 3. Collect Vocabulary (Train Split Only) print(" [WL] Collecting vocabulary from training data...") # We need to load a subset of training data to build the vocabulary. # Loading ALL training states might be slow, but it ensures full coverage. train_files = sorted( [f for f in os.listdir(train_states_dir) if f.endswith(".traj")] ) pddl_train_dir = os.path.join(args.data_dir, "pddl", domain_name, "train") wl_problems = [] for t_file in tqdm(train_files, desc="Parsing Train"): prob_name = t_file.replace(".traj", "") prob_pddl = os.path.join(pddl_train_dir, f"{prob_name}.pddl") traj_path = os.path.join(train_states_dir, t_file) if not os.path.exists(prob_pddl): continue try: wl_prob = parse_problem(domain_pddl, prob_pddl) # Read trajectory with open(traj_path, "r") as f: lines = f.readlines() # Parse states states = [parse_traj_line_to_state(line, pred_map) for line in lines] wl_problems.append(ProblemDataset(wl_prob, states)) except Exception: continue # Collect if not wl_problems: print(" [Error] No valid training data found.") continue full_train_ds = DomainDataset(wl_domain, wl_problems) feature_gen.collect(full_train_ds) print(f" [WL] Vocabulary Size: {feature_gen.get_n_features()}") # 4. Save Feature Generator (JSON) os.makedirs(args.model_dir, exist_ok=True) save_path = os.path.join(args.model_dir, f"{domain_name}_wl.json") feature_gen.save(save_path) print(f" [WL] Saved model to {save_path}") # 5. Embed All Splits for split in SPLITS: print(f" [WL] Embedding split: {split}") split_state_dir = os.path.join(args.data_dir, "states", domain_name, split) split_pddl_dir = os.path.join(args.data_dir, "pddl", domain_name, split) split_out_dir = os.path.join(args.output_dir, domain_name, split) os.makedirs(split_out_dir, exist_ok=True) if not os.path.exists(split_state_dir): continue traj_files = sorted([f for f in os.listdir(split_state_dir) if f.endswith(".traj")]) for t_file in tqdm(traj_files, desc=f"Embedding {split}"): prob_name = t_file.replace(".traj", "") prob_pddl = os.path.join(split_pddl_dir, f"{prob_name}.pddl") traj_path = os.path.join(split_state_dir, t_file) out_traj_path = os.path.join(split_out_dir, f"{prob_name}.npy") out_goal_path = os.path.join(split_out_dir, f"{prob_name}_goal.npy") if not os.path.exists(prob_pddl): continue try: # Parse Problem & States wl_prob = parse_problem(domain_pddl, prob_pddl) with open(traj_path, "r") as f: lines = f.readlines() states = [parse_traj_line_to_state(l, pred_map) for l in lines] # Embed Trajectory # feature_gen.embed returns a flattened list of vectors [v_s0, v_s1, ...] mini_ds = DomainDataset( wl_domain, [ProblemDataset(wl_prob, states)] ) embs = feature_gen.embed(mini_ds) # Use 'embs' directly, not 'embs[0]' traj_matrix = np.array(embs, dtype=np.float32) # [T, D] # Embed Goal # We create a dummy state containing only the goal atoms goal_atoms = list(wl_prob.positive_goals) goal_state = State(goal_atoms) goal_ds = DomainDataset( wl_domain, [ProblemDataset(wl_prob, [goal_state])] ) goal_embs = feature_gen.embed(goal_ds) # Use 'goal_embs[0]' (the first and only vector) goal_vec = np.array(goal_embs[0], dtype=np.float32) # [D] # Save np.save(out_traj_path, traj_matrix) np.save(out_goal_path, goal_vec) except Exception as e: print(f" Error embedding {prob_name}: {e}") if __name__ == "__main__": main()