🏔️ Hill Climbing Predicate Search

function HillClimbing(initial_state): // Initialize with starting state current_state = initial_state current_score = evaluate(current_state) while True: // Generate all neighbors (add one predicate) neighbors = generate_neighbors(current_state) // Evaluate each neighbor best_neighbor = None best_score = current_score for neighbor in neighbors: score = evaluate(neighbor) if score < best_score: // Lower is better best_neighbor = neighbor best_score = score // If no improvement found, terminate if best_neighbor is None: break // Move to best neighbor current_state = best_neighbor current_score = best_score return current_state

def run_hill_climbing( initial_state: _S, check_goal: Callable[[_S], bool], get_successors: Callable[[_S], Iterator[Tuple[_A, _S, float]]], heuristic: Callable[[_S], float], early_termination_heuristic_thresh: Optional[float] = None, enforced_depth: int = 0, parallelize: bool = False, verbose: bool = True, timeout: float = float('inf') ) -> Tuple[List[_S], List[_A], List[float]]: """Enforced hill climbing local search. For each node, the best child node is always selected, if that child is an improvement over the node. If no children improve on the node, look at the children's children, etc., up to enforced_depth, where enforced_depth 0 corresponds to simple hill climbing. Terminate when no improvement can be found. early_termination_heuristic_thresh allows for searching until heuristic reaches a specified value. Lower heuristic is better. """ # Initialize search state cur_node: _HeuristicSearchNode[_S, _A] = _HeuristicSearchNode(initial_state, 0, 0) last_heuristic = heuristic(cur_node.state) heuristics = [last_heuristic] visited = {initial_state} if verbose: logging.info(f"Starting hill climbing at state {cur_node.state} with heuristic {last_heuristic}") start_time = time.perf_counter() while True: # Early termination: heuristic reached target value if early_termination_heuristic_thresh is not None \ and last_heuristic <= early_termination_heuristic_thresh: break # Goal check: found optimal solution if check_goal(cur_node.state): if verbose: logging.info("Terminating hill climbing, achieved goal") break # Search for best improvement best_heuristic = float("inf") best_child_node = None current_depth_nodes = [cur_node] all_best_heuristics = [] # Enforced depth search: look at children's children up to enforced_depth for depth in range(0, enforced_depth + 1): if verbose: logging.info(f"Searching for an improvement at depth {depth}") successors_at_depth = [] for parent in current_depth_nodes: for action, child_state, cost in get_successors(parent.state): # Timeout check if time.perf_counter() - start_time > timeout: raise TimeoutError() # Skip already visited states if child_state in visited: continue visited.add(child_state) # Create child node with path cost child_path_cost = parent.cumulative_cost + cost child_node = _HeuristicSearchNode( state=child_state, edge_cost=cost, cumulative_cost=child_path_cost, parent=parent, action=action) successors_at_depth.append(child_node) # Evaluate child heuristic (THIS IS WHERE J_surr IS CALCULATED!) if not parallelize: child_heuristic = heuristic(child_node.state) if child_heuristic < best_heuristic: best_heuristic = child_heuristic best_child_node = child_node # Parallel heuristic computation (if enabled) if parallelize: num_cpus = mp.cpu_count() fn = lambda n: (heuristic(n.state), n) with mp.Pool(processes=num_cpus) as p: for child_heuristic, child_node in p.map(fn, successors_at_depth): if child_heuristic < best_heuristic: best_heuristic = child_heuristic best_child_node = child_node all_best_heuristics.append(best_heuristic) # Found improvement at this depth if last_heuristic > best_heuristic: if verbose: logging.info(f"Found an improvement at depth {depth}") break # Continue to next depth level current_depth_nodes = successors_at_depth if verbose: logging.info(f"No improvement found at depth {depth}") # Termination conditions if best_child_node is None: if verbose: logging.info("Terminating hill climbing, no more successors") break if last_heuristic <= best_heuristic: if verbose: logging.info("Terminating hill climbing, could not improve score") break # Move to best child and continue heuristics.extend(all_best_heuristics) cur_node = best_child_node last_heuristic = best_heuristic if verbose: logging.info(f"Hill climbing reached new state {cur_node.state} with heuristic {last_heuristic}") # Return search path and results states, actions = _finish_plan(cur_node) assert len(states) == len(heuristics) return states, actions, heuristics

def evaluate_with_operators( candidate_predicates: FrozenSet[Predicate], low_level_trajs: List[LowLevelTrajectory], segmented_trajs: List[List[Segment]], strips_ops: List[STRIPSOperator], option_specs: List[OptionSpec] ) -> float: """Calculate J_surr for the given predicate set. This is the core of the surrogate objective function that estimates the expected number of nodes that would need to be created/expanded before finding a refinable skeleton during task planning. """ assert self.metric_name in ("num_nodes_created", "num_nodes_expanded") score = 0.0 seen_demos = 0 # For each demonstration trajectory for ll_traj, seg_traj in zip(low_level_trajs, segmented_trajs): if seen_demos >= CFG.grammar_search_max_demos: break if not ll_traj.is_demo: continue # Extract demonstration information demo_atoms_sequence = utils.segment_trajectory_to_atoms_sequence(seg_traj) seen_demos += 1 init_atoms = demo_atoms_sequence[0] goal = self._train_tasks[ll_traj.train_task_idx].goal # Ground operators and create planning heuristic objects = set(ll_traj.states[0]) dummy_nsrts = utils.ops_and_specs_to_dummy_nsrts(strips_ops, option_specs) ground_nsrts, reachable_atoms = task_plan_grounding( init_atoms, objects, dummy_nsrts, allow_noops=CFG.grammar_search_expected_nodes_allow_noops) heuristic = utils.create_task_planning_heuristic( CFG.sesame_task_planning_heuristic, init_atoms, goal, ground_nsrts, candidate_predicates | self._initial_predicates, objects) # Calculate expected planning time for this demo expected_planning_time = 0.0 refinable_skeleton_not_found_prob = 1.0 # Set skeleton budget if CFG.grammar_search_expected_nodes_max_skeletons == -1: max_skeletons = CFG.sesame_max_skeletons_optimized else: max_skeletons = CFG.grammar_search_expected_nodes_max_skeletons # Generate skeletons using task planning generator = task_plan(init_atoms, goal, ground_nsrts, reachable_atoms, heuristic, CFG.seed, CFG.grammar_search_task_planning_timeout, max_skeletons, use_visited_state_set=False) try: for idx, (_, plan_atoms_sequence, metrics) in enumerate(generator): assert goal.issubset(plan_atoms_sequence[-1]) # Calculate refinement probability using geometric distribution refinement_prob = self._get_refinement_prob(demo_atoms_sequence, plan_atoms_sequence) # Get number of nodes created/expanded for this skeleton assert self.metric_name in metrics num_nodes = metrics[self.metric_name] # Calculate probability this is the first refinable skeleton p_first_refinable = refinable_skeleton_not_found_prob * refinement_prob # Expected contribution: P(first refinable is k) * nodes_k expected_planning_time += p_first_refinable * num_nodes # Add backtracking cost for skeletons after the first if idx > 0: w = CFG.grammar_search_expected_nodes_backtracking_cost # 1000 expected_planning_time += p_first_refinable * w # Update probability that no skeleton yet is refinable refinable_skeleton_not_found_prob *= (1 - refinement_prob) except (PlanningTimeout, PlanningFailure): # If planning failed, add upper bound with full probability pass # Add worst-case upper bound for remaining probability ub = CFG.grammar_search_expected_nodes_upper_bound # 100000 expected_planning_time += refinable_skeleton_not_found_prob * ub # Accumulate score across all demos score += expected_planning_time return score @staticmethod def _get_refinement_prob( demo_atoms_sequence: Sequence[Set[GroundAtom]], plan_atoms_sequence: Sequence[Set[GroundAtom]] ) -> float: """Estimate probability that plan_atoms_sequence is refinable. Uses geometric distribution: p * (1-p)^|demo_len - plan_len| Assumes demonstrations are optimal with probability p. """ demo_len = len(demo_atoms_sequence) plan_len = len(plan_atoms_sequence) exponent = abs(demo_len - plan_len) p = CFG.grammar_search_expected_nodes_optimal_demo_prob # 0.99999 return p * (1 - p) ** exponent