""" Graph BPE (Byte-Pair Encoding) tokenizer for Instance Learning Graphs. Implements a self-contained BPE algorithm on graph structures, inspired by the GraphBPE paper (Shen & Poczos, ICML 2024 AI4Science Workshop). The original GraphBPE repo (https://github.com/A-Chicharito-S/GraphBPE) is specialized for molecular graphs with heavy dependencies (torch_geometric, hydra, hgraph2graph). This implementation adapts the core BPE concept for planning ILGs without those dependencies. Algorithm: 1. Initialize vocabulary: each unique node label is a token. 2. Iteratively merge the most frequent adjacent token pair into a new token. 3. After training, tokenize a graph by applying learned merges. 4. Create histogram embedding over the final vocabulary. Reference: Shen, Y., & Poczos, B. (2024). GraphBPE: Molecular Graphs Meet Byte-Pair Encoding. ICML 2024 AI for Science Workshop. """ import json import logging import os import re import sys from collections import Counter, defaultdict import numpy as np from tqdm import tqdm from code.tokenization.base import TokenizationStrategy logger = logging.getLogger(__name__) _PREDICATE_REGEX = re.compile(r"\(([\w-]+(?: [\w-]+)*)\)") def _progress_enabled() -> bool: return bool(sys.stdout.isatty()) class GraphBPETokenizer(TokenizationStrategy): """ Graph BPE tokenizer: learns frequent substructure patterns via iterative merging of adjacent node labels. Vocabulary is built by iteratively merging the most frequent pair of adjacent labels in the training set. At inference, learned merges are applied to produce a token sequence, then a histogram is created. """ def __init__(self, vocab_size: int = 1000, num_iterations: int = 100): super().__init__(name="GraphBPE") self.vocab_size = vocab_size self.num_iterations = num_iterations # Learned merge rules: [(token_a, token_b, merged_token), ...] self._merges: list[tuple[str, str, str]] = [] # Final vocabulary: {token_string -> index} self._vocabulary: dict[str, int] | None = None self._domain_info: dict[str, int] | None = None def fit( self, domain_pddl_path: str, train_states_dir: str, train_pddl_dir: str, ) -> None: """ Learn BPE merge rules from training data. 1. Extract initial token sequences from all training ILGs. 2. Iteratively merge the most frequent adjacent pair. 3. Build final vocabulary from all remaining tokens. """ import pddl domain = pddl.parse_domain(domain_pddl_path) self._domain_info = {p.name.lower(): p.arity for p in domain.predicates} # 1. Extract all labeled graphs from training data all_graphs: list[dict] = [] train_files = sorted( [f for f in os.listdir(train_states_dir) if f.endswith(".traj")] ) for t_file in tqdm( train_files, desc=f" [{self.name}] Parsing graphs", disable=(not _progress_enabled()), ): prob_name = t_file.replace(".traj", "") prob_pddl = os.path.join(train_pddl_dir, f"{prob_name}.pddl") traj_path = os.path.join(train_states_dir, t_file) if not os.path.exists(prob_pddl): continue try: problem = pddl.parse_problem(prob_pddl) objects = sorted( {o.name for o in problem.objects} | {o.name for o in domain.constants} ) goal_atoms = self._extract_goal_atoms(problem) with open(traj_path, "r") as f: lines = f.readlines() for line in lines: state_atoms = _PREDICATE_REGEX.findall(line.strip()) graph = self._build_labeled_graph(state_atoms, goal_atoms, objects) all_graphs.append(graph) except Exception: continue if not all_graphs: raise RuntimeError("No graphs collected during GraphBPE fit.") # 2. Run BPE merge iterations logger.info( f"[{self.name}] Running BPE on {len(all_graphs)} graphs, " f"target vocab={self.vocab_size}, max_iter={self.num_iterations}" ) self._run_bpe(all_graphs) # 3. Build final vocabulary all_tokens: set[str] = set() for graph in all_graphs: for label in graph["labels"].values(): all_tokens.add(label) sorted_tokens = sorted(all_tokens) self._vocabulary = {t: i for i, t in enumerate(sorted_tokens)} self.embedding_dim = len(self._vocabulary) self._is_fitted = True logger.info( f"[{self.name}] Fitted: {len(self._merges)} merges, " f"vocabulary size: {self.embedding_dim}" ) def _extract_goal_atoms(self, problem) -> list[str]: """Extract goal atoms from a pddl Problem object.""" import pddl.logic.predicates goals = [] def visit(node): if isinstance(node, pddl.logic.predicates.Predicate): args = [ t.name if hasattr(t, "name") else str(t) for t in node.terms ] goals.append(f"{node.name} {' '.join(args)}") elif hasattr(node, "operands"): for op in node.operands: visit(op) elif hasattr(node, "_operands"): for op in node._operands: visit(op) visit(problem.goal) return goals def _build_labeled_graph( self, state_atoms: list[str], goal_atoms: list[str], objects: list[str], ) -> dict: """ Build a graph with labeled nodes and edges. Returns dict with: - 'labels': {obj_name -> initial_label} - 'adj': {obj_name -> set of neighbor obj_names} """ # Initial labels: based on predicates involving each object obj_state_preds: dict[str, list[str]] = defaultdict(list) obj_goal_preds: dict[str, list[str]] = defaultdict(list) adj: dict[str, set[str]] = {obj: set() for obj in objects} # Process state atoms for atom_str in state_atoms: parts = atom_str.split() if not parts: continue pred = parts[0].lower() args = [a.lower() for a in parts[1:]] for arg in args: if arg in adj: obj_state_preds[arg].append(f"s:{pred}") if len(args) >= 2: for i in range(len(args)): for j in range(i + 1, len(args)): if args[i] in adj and args[j] in adj: adj[args[i]].add(args[j]) adj[args[j]].add(args[i]) # Process goal atoms for atom_str in goal_atoms: parts = atom_str.split() if not parts: continue pred = parts[0].lower() args = [a.lower() for a in parts[1:]] for arg in args: if arg in adj: obj_goal_preds[arg].append(f"g:{pred}") if len(args) >= 2: for i in range(len(args)): for j in range(i + 1, len(args)): if args[i] in adj and args[j] in adj: adj[args[i]].add(args[j]) adj[args[j]].add(args[i]) # Create initial labels (sorted predicate sets for determinism) labels = {} for obj in objects: s_preds = sorted(set(obj_state_preds.get(obj, []))) g_preds = sorted(set(obj_goal_preds.get(obj, []))) label_parts = s_preds + g_preds labels[obj] = "|".join(label_parts) if label_parts else "EMPTY" return {"labels": labels, "adj": adj} def _run_bpe(self, graphs: list[dict]) -> None: """ Run BPE merge iterations on the graphs. At each step: 1. Count all adjacent label pair frequencies across all graphs. 2. Pick the most frequent pair. 3. Merge: for every edge where both endpoints have these labels, combine them into one super-node with merged label. """ self._merges = [] for iteration in range(self.num_iterations): # Count adjacent pairs across all graphs pair_counts: Counter = Counter() for graph in graphs: labels = graph["labels"] adj = graph["adj"] for node, neighbors in adj.items(): if node not in labels: continue node_label = labels[node] for neighbor in neighbors: if neighbor not in labels: continue nbr_label = labels[neighbor] # Canonical ordering pair = tuple(sorted([node_label, nbr_label])) pair_counts[pair] += 1 if not pair_counts: break # Find most frequent pair best_pair, best_count = pair_counts.most_common(1)[0] if best_count < 2: break # No frequent pairs left merged_label = f"({best_pair[0]}+{best_pair[1]})" self._merges.append((best_pair[0], best_pair[1], merged_label)) # Apply merge to all graphs for graph in graphs: self._apply_merge(graph, best_pair[0], best_pair[1], merged_label) if len(self._vocabulary or {}) + len(self._merges) >= self.vocab_size: break def _apply_merge( self, graph: dict, label_a: str, label_b: str, merged_label: str, ) -> None: """ Apply a single merge to a graph. Finds all edges where one endpoint has label_a and the other label_b, contracts the first such edge found, and repeats. """ labels = graph["labels"] adj = graph["adj"] # Find all edges eligible for this merge merged = True while merged: merged = False for node in list(labels.keys()): if node not in labels: continue if labels[node] not in (label_a, label_b): continue for neighbor in list(adj.get(node, set())): if neighbor not in labels: continue # Check if this edge matches the pair pair = tuple(sorted([labels[node], labels[neighbor]])) if pair == tuple(sorted([label_a, label_b])): # Contract: keep 'node', remove 'neighbor' labels[node] = merged_label # Transfer neighbor's edges to node for n_nbr in adj.get(neighbor, set()): if n_nbr != node and n_nbr in adj: adj[node].add(n_nbr) adj[n_nbr].discard(neighbor) adj[n_nbr].add(node) # Remove neighbor del labels[neighbor] adj[node].discard(neighbor) if neighbor in adj: del adj[neighbor] merged = True break if merged: break def _apply_merges_to_graph(self, graph: dict) -> dict: """Apply all learned merges to a new graph (inference).""" import copy g = copy.deepcopy(graph) for label_a, label_b, merged_label in self._merges: self._apply_merge(g, label_a, label_b, merged_label) return g def transform_state( self, state_atoms: list[str], goal_atoms: list[str], objects: list[str], ) -> np.ndarray: """ Tokenize an ILG using learned BPE merges and create histogram. """ self._check_fitted() # Parse atoms parsed_state = [] for a in state_atoms: matches = _PREDICATE_REGEX.findall(a) parsed_state.extend(matches) parsed_goal = [] for a in goal_atoms: a_clean = a.replace("(", "").replace(")", "").strip() if a_clean: parsed_goal.append(a_clean) # Build graph and apply merges graph = self._build_labeled_graph(parsed_state, parsed_goal, objects) merged_graph = self._apply_merges_to_graph(graph) # Create histogram embedding = np.zeros(self.embedding_dim, dtype=np.float32) for label in merged_graph["labels"].values(): if label in self._vocabulary: embedding[self._vocabulary[label]] += 1.0 return embedding def transform_goal( self, goal_atoms: list[str], objects: list[str], ) -> np.ndarray: """Embed a goal by treating it as a state with no current atoms.""" return self.transform_state([], goal_atoms, objects) def get_embedding_dim(self) -> int: self._check_fitted() return self.embedding_dim def save_vocabulary(self, filepath: str) -> None: """Save merges and vocabulary to JSON.""" self._check_fitted() data = { "vocab_size": self.vocab_size, "num_iterations": self.num_iterations, "merges": self._merges, "vocabulary": self._vocabulary, "embedding_dim": self.embedding_dim, } with open(filepath, "w") as f: json.dump(data, f, indent=2) logger.info(f"[{self.name}] Saved vocabulary to {filepath}") def load_vocabulary(self, filepath: str) -> None: """Load merges and vocabulary from JSON.""" with open(filepath, "r") as f: data = json.load(f) self.vocab_size = data["vocab_size"] self.num_iterations = data["num_iterations"] self._merges = [tuple(m) for m in data["merges"]] self._vocabulary = data["vocabulary"] self.embedding_dim = data["embedding_dim"] self._is_fitted = True logger.info( f"[{self.name}] Loaded: {len(self._merges)} merges, " f"vocab size: {self.embedding_dim}" )