Atomic Task Graph: A 7B Model That Beats GPT-4 ReAct on ALFWorld and WebShop Has Nothing to Do With the 7B. It Is the Control Framework.

# ATG: training-free, no fine-tuning, no external demonstrations
# Source: arXiv:2607.01942
# Design intent: represent planning + execution as an explicit DAG
# rather than a growing textual trajectory

# Node definition
node = {
    "id": "v_j",
    "tool": "search_api",          # f_j: the selected tool
    "input": "Beijing weather",     # i_j: input (may reference o_k from upstream nodes)
    "output": None,                 # o_j: filled in during execution
    "status": "pending",           # pending / running / succeeded / failed
    "error": None,
    "refinement_depth": 2,         # which level of recursive compilation produced this node
    "parent_node": "v_coarse"      # which parent was refined to produce this node
}

# Edge definition
edge = {
    "from": "v_j",    # o_j is used as i_k
    "to": "v_k"
}

# ATG: interface-preserving recursive graph compilation
# Source: arXiv:2607.01942 Section 4.1 (reconstructed from paper's algorithm)
# Design intent: expose subtask dependencies explicitly as a DAG at each level.
# Each refinement step preserves the parent node's external i/o interface,
# making the refinement history a usable basis for error tracing and localized repair.

from dataclasses import dataclass, field
from typing import Optional

@dataclass
class AtomicNode:
    """
    One node in the ATG. v_j = (i_j, f_j, o_j).
    An atomic node is one that can be directly executed as a single tool call.
    A non-atomic node needs further refinement.
    """
    node_id: str
    tool: Optional[str]                  # f_j: the tool to invoke
    input_spec: str                      # i_j: what inputs this node needs
    output_spec: str                     # o_j: what this node produces
    is_atomic: bool = False              # True when no further decomposition needed
    parent_id: Optional[str] = None     # which node was refined to produce this
    refinement_depth: int = 0           # depth at which this node was created
    # State tracking during execution:
    status: str = "pending"             # pending / running / succeeded / failed
    actual_input: Optional[dict] = None
    actual_output: Optional[dict] = None
    error_message: Optional[str] = None

@dataclass
class AtomicTaskGraph:
    nodes: dict = field(default_factory=dict)   # node_id → AtomicNode
    edges: list = field(default_factory=list)   # list of (from_id, to_id)
    refinement_history: list = field(default_factory=list)  # G_0, G_1, ..., G_final

def recursive_compile(task_description: str, tools: list, llm) -> AtomicTaskGraph:
    """
    Recursively compile a coarse task into an atomic task graph.
    ← THIS is the planning mechanism: no growing textual trajectory.
    Each LLM call only sees context relevant to the current node being refined.
    """
    atg = AtomicTaskGraph()

    # Initial coarse node: the whole task as one non-atomic unit
    root = AtomicNode(
        node_id="root",
        tool=None,
        input_spec=task_description,
        output_spec="complete_task_result",
        is_atomic=False,
        refinement_depth=0
    )
    atg.nodes["root"] = root
    atg.refinement_history.append({"depth": 0, "nodes": ["root"], "edges": []})

    nodes_to_refine = ["root"]
    depth = 1

    while nodes_to_refine:
        next_round = []
        for node_id in nodes_to_refine:
            parent = atg.nodes[node_id]

            # ← LOCALIZED CONTEXT: LLM only sees the current node's i/o spec
            # and the list of available tools. NOT the full task history.
            # This is why ATG reduces hallucinations: no accumulation of irrelevant context.
            prompt = f"""
            Task node to refine:
              Input: {parent.input_spec}
              Required output: {parent.output_spec}
              Available tools: {[t['name'] for t in tools]}

            Decompose this node into an atomic subgraph.
            Rules:
            1. The subgraph must consume EXACTLY the same inputs as this node (interface preservation)
            2. The subgraph must produce EXACTLY the same output type (interface preservation)
            3. Each child node must be either:
               a) Directly executable with one tool call (atomic)
               b) A coarser subtask to be refined in the next round
            4. Specify all input-output dependencies between child nodes.

            Return as JSON: {{children: [{{id, tool_or_none, input_spec, output_spec, is_atomic}}],
                            edges: [[from_id, to_id]]}}
            """

            refinement = llm.generate_json(prompt)

            # Replace parent node in graph with its subgraph
            # ← Interface preserved: subgraph collectively has same i/o as parent
            del atg.nodes[node_id]
            for child in refinement["children"]:
                child_node = AtomicNode(
                    node_id=child["id"],
                    tool=child.get("tool_or_none"),
                    input_spec=child["input_spec"],
                    output_spec=child["output_spec"],
                    is_atomic=child["is_atomic"],
                    parent_id=node_id,           # ← records which parent was refined
                    refinement_depth=depth
                )
                atg.nodes[child["id"]] = child_node
                if not child_node.is_atomic:
                    next_round.append(child["id"])  # needs further refinement

            for (from_id, to_id) in refinement["edges"]:
                atg.edges.append((from_id, to_id))

        # Record this round's graph in refinement history
        # ← refinement history is what enables localized repair later
        atg.refinement_history.append({
            "depth": depth,
            "nodes": list(atg.nodes.keys()),
            "edges": list(atg.edges)
        })
        nodes_to_refine = next_round
        depth += 1

    return atg

# ATG: dependency-aware execution with pre-execution thought experiment
# and minimal necessary subgraph repair
# Source: arXiv:2607.01942 Sections 4.2 and 4.3
# Design intent: exploit the explicit DAG for parallel scheduling and
# catch failures early (thought experiment) rather than after expensive environment calls.

import asyncio

async def thought_experiment(atg: AtomicTaskGraph, llm) -> dict:
    """
    Pre-execution simulation: check the compiled graph BEFORE touching the environment.
    ← Detects 24.6% risky plans on ALFWorld with >74% precision.
    ← Catches: invalid tool assignments, missing intermediate steps,
       interface mismatches, implausible execution paths.
    ← Cost: one LLM call. Cost of missing this: failed environment interactions.
    """
    prompt = f"""
    Pre-execution validation. Simulate the following task graph internally:
    Nodes: {[{n.node_id: {'tool': n.tool, 'input': n.input_spec, 'output': n.output_spec}}
             for n in atg.nodes.values()]}
    Edges (dependencies): {atg.edges}

    Check 

for:
    1. Invalid tool selections (tool exists but wrong for this input type)
    2. Missing intermediate steps (an output cannot produce the next node's required input)
    3. Interface mismatches (output type of edge source ≠ input type of edge destination)
    4. Implausible execution paths (correct tools in wrong order)

    For each issue found, report: {{failing_node_id, issue_type, description}}
    If no issues found, return: {{issues: []}}
    """
    result = llm.generate_json(prompt)
    return result  # {"issues": [...]} or {"issues": []}

async def execute_node(node: AtomicNode, resolved_inputs: dict, environment) -> dict:
    """Execute one atomic tool call with resolved inputs."""
    try:
        actual_output = await environment.call_tool(node.tool, resolved_inputs)
        node.status = "succeeded"
        node.actual_output = actual_output
        return actual_output
    except Exception as e:
        node.status = "failed"
        node.error_message = str(e)
        raise

async def dependency_aware_execute(atg: AtomicTaskGraph, environment, llm):
    """
    Execute ATG nodes in topological order.
    Independent branches run in PARALLEL.
    ← THIS is why ATG uses 18.36 steps vs ReAct 31.42 on ALFWorld.
       Not fewer total operations — fewer sequential steps because parallel branches
       count as one step. The explicit dependency graph enables this directly.
    """
    resolved_outputs = {}  # node_id → actual output value
    in_flight = set()

    # Build dependency counts for topological ordering
    dependency_count = {nid: 0 for nid in atg.nodes}
    predecessors = {nid: [] for nid in atg.nodes}
    for (from_id, to_id) in atg.edges:
        dependency_count[to_id] += 1
        predecessors[to_id].append(from_id)

    # Queue: all nodes with no dependencies are initially ready
    ready_queue = [nid for nid, count in dependency_count.items() if count == 0]
    pending = asyncio.Queue()
    for nid in ready_queue:
        await pending.put(nid)

    while not pending.empty() or in_flight:
        # Collect all currently executable nodes and launch them in parallel
        # ← THIS is the parallel scheduling step:
        # Multiple nodes can be simultaneously ready → all launched at once
        batch = []
        while not pending.empty():
            batch.append(await pending.get())

        if batch:
            tasks = []
            for nid in batch:
                node = atg.nodes[nid]
                in_flight.add(nid)
                # Resolve inputs from predecessor outputs
                inputs = {}
                for pred_id in predecessors[nid]:
                    inputs[pred_id] = resolved_outputs.get(pred_id)
                tasks.append(execute_node(node, inputs, environment))

            # ← All tasks in batch execute concurrently (parallel branches)
            results = await asyncio.gather(*tasks, return_exceptions=True)

            for nid, result in zip(batch, results):
                in_flight.discard(nid)
                if isinstance(result, Exception):
                    await minimal_subgraph_repair(atg, nid, environment, llm)
                else:
                    resolved_outputs[nid] = result
                    # Check if any successors are now ready
                    for (from_id, to_id) in atg.edges:
                        if from_id == nid:
                            dependency_count[to_id] -= 1
                            if dependency_count[to_id] == 0:
                                await pending.put(to_id)

async def minimal_subgraph_repair(atg: AtomicTaskGraph, failed_node_id: str,
                                   environment, llm):
    """
    Repair ONLY the affected subgraph. Freeze everything else.
    ← Ablation: removing this costs -7.72 pts on ALFWorld (Mistral-7B),
       -6.48 on ScienceWorld. The most impactful single component.
    """
    failed_node = atg.nodes[failed_node_id]

    # Step 1: Localize failure via refinement history
    # Trace the failed node back to its origin in G_0 → G_1 → ... → G_final
    # Find lowest common historical ancestor (a_f) covering the failed region
    a_f = find_lowest_common_ancestor(
        failed_nodes=[failed_node_id],
        refinement_history=atg.refinement_history
    )
    # ← a_f is the original planning scope where the error was introduced
    # ← It provides the natural boundary for localized repair

    # Step 2: Identify affected downstream nodes
    affected = {failed_node_id}
    for (from_id, to_id) in atg.edges:
        if from_id in affected:
            affected.add(to_id)  # downstream nodes also need repair

    # Step 3: Freeze validated regions
    for nid, node in atg.nodes.items():
        if nid not in affected and node.status == "succeeded":
            node.status = "frozen"  # ← These nodes will NOT be re-executed
            # Their outputs in resolved_outputs are preserved and reusable

    # Step 4: Repair only the affected subgraph (bounded by a_f's interface)
    repair_prompt = f"""
    A failure occurred at node {failed_node_id}.
    Error: {failed_node.error_message}
    Ancestor context (from refinement history): {a_f}
    Affected subgraph nodes: {list(affected)}

    Repair by:
    - Replacing incorrect tool selections
    - Inserting missing intermediate nodes
    - Adjusting local dependencies
    Constraints:
    - The repaired subgraph MUST preserve the same external input-output interface
    - Nodes outside the affected set are frozen and must not be changed
    """
    repaired_subgraph = llm.generate_json(repair_prompt)
    # Reintegrate repaired subgraph into main ATG
    apply_repair(atg, affected, repaired_subgraph)

    # Step 5: Resume execution from repaired subgraph
    # (frozen nodes not re-executed, only repaired/pending nodes run)

def find_lowest_common_ancestor(failed_nodes, refinement_history):
    """
    Use the recorded graph evolution to find the smallest ancestor
    node from which all failed nodes were derived.
    ← This is the 'refinement history' from compilation time, now paying off at repair time.
    """
    # Walk up refinement history levels until a node containing all failed nodes is found
    for level in reversed(refinement_history):
        for node_id in level["nodes"]:
            # Check if this ancestor covers all failed nodes
            if all(is_descendant(fn, node_id, refinement_history) for fn in failed_nodes):
                return node_id
    return "root"  # worst case: repair from root (rare)

Round 0 (coarse, 1 node):
  v_root = "put a clean sponge on the counter"
  Non-atomic: decompose

Round 1 (medium, context = only v_root's i/o spec):
  v_A = (task, navigate_to_sponge_location, sponge_location)
  v_B = (sponge_location, pick_up_sponge, held_sponge)
  v_C = (held_sponge, navigate_to_sink, at_sink)
  v_D = (at_sink + held_sponge, clean_sponge, clean_held_sponge)
  v_E = (clean_held_sponge, navigate_to_counter, at_counter_with_sponge)
  v_F = (at_counter_with_sponge, place_on_surface, task_complete)
  Edges: A→B→C→D→E→F; A→C (navigation shares context)

  v_A: still non-atomic (navigate requires find + examine)
  v_B, v_C, v_D, v_E, v_F: atomic (single tool calls)

Round 2 (fine, context = only v_A's i/o spec):
  v_A1 = (task, find_object, object_candidates)
  v_A2 = (object_candidates, examine_receptacles, sponge_with_location)
  Edges: A1→A2→B (reintegrates into main graph preserving v_A's interface)

Refinement history recorded: G_0 (root) → G_1 (A-F) → G_2 (A1,A2,B-F)

Internal simulation of G_2:
  Check: find_object → examine_receptacles: output type matches? YES
  Check: examine_receptacles → pick_up_sponge: sponge_with_location feeds held_sponge? YES
  Check: clean_sponge tool: requires sink nearby? YES (v_C ensures at_sink before v_D)
  Check: place_on_surface: requires clean_held_sponge + at_counter? YES

Issues found: 0
Risky plan: NO
Proceed to execution

Step 1 (parallel batch 1):
  v_A1 = find_object("sponge") → ["sponge_1 in bathtub", "sponge_2 in kitchen sink"]
  [Only v_A1 ready — single node, no other independent branches]

Step 2 (parallel batch 2):
  v_A2 = examine_receptacles(["sponge_1 in bathtub", "sponge_2 in kitchen sink"])
       → sponge_1, bathtub, location_coords

Step 3 (parallel batch 3):
  v_B = pick_up_sponge(sponge_1, bathtub) → held_sponge_1
  [Only v_B ready]

Step 4 (parallel batch 4):
  v_C = navigate_to_sink() → at_sink_status
  [Only v_C ready]

Step 5 (parallel batch 5):
  v_D = clean_sponge(held_sponge_1, at_sink) → clean_held_sponge_1
  [Only v_D ready]

Step 6 (parallel batch 6):
  v_E = navigate_to_counter() → at_counter_status
  [Only v_E ready]

Step 7 (parallel batch 7):
  v_F = place_on_surface(clean_held_sponge_1, at_counter) → SUCCESS

Total execution steps: 7
← ATG avg: 18.36 steps on ALFWorld vs ReAct 31.42
  (7 is fewer than average; this task has sequential structure)
  Tasks with independent branches exploit parallel execution more aggressively

Assume v_D fails: clean_sponge fails because sink is occupied

State at failure:
  v_A1: succeeded (frozen), output preserved
  v_A2: succeeded (frozen), output preserved
  v_B: succeeded (frozen), held_sponge_1 preserved
  v_C: succeeded (frozen), at_sink preserved
  v_D: FAILED — error: "sink occupied"

Failure localization:
  Failed node: v_D
  Trace refinement history: v_D was created in Round 1 from v_root
  Lowest common ancestor: v_D itself (no siblings failed)
  Affected: {v_D, v_E, v_F} (downstream of v_D)

Repair:
  v_A1, v_A2, v_B, v_C: FROZEN (outputs preserved, not re-executed)
  Repair subgraph {v_D, v_E, v_F}:
    New v_D' = wait_for_sink() + clean_sponge (adds waiting step)
    Or: v_D' = find_alternative_sink() + navigate + clean
  External interface preserved: v_D' still outputs clean_held_sponge

Resume from v_D' only: 3 additional steps
vs global restart: 7 steps from scratch

Repair cost: 3 additional steps
Global restart cost: 7 steps
Savings: 57% reduction in additional steps