🤖 AI

CoCoDA: Co-evolving Compositional DAG for Tool-Augmented Agents

CoCoDA is a framework that enables small language models to match or exceed larger teachers by co-evolving a planner and a growing tool library within a single compositional code DAG, which optimizes retrieval efficiency through symbolic unification and accelerates learning via a structure-aware reward mechanism.

Original authors: Ziyang Yu, Qiyue Li, Liang Zhao

Published 2026-05-12

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Ziyang Yu, Qiyue Li, Liang Zhao

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a very smart but small robot assistant (the "planner") that needs to solve complex problems, like doing advanced math or writing code. This robot has a limited amount of "brain space" (context window) to hold information at once. To help it, we give it a toolbox of skills (tools) it can use.

The problem with existing methods is like trying to carry a massive library of books in your pocket. As you add more useful books (tools) to help the robot, the pocket gets too heavy, and the robot can't fit enough of them in its brain to solve the problem. Also, the robot doesn't know how to organize these books; it just sees a messy pile of text, making it hard to find the right one quickly.

CoCoDA is a new system that solves this by changing how the toolbox is built and how the robot learns to use it. Here is how it works, using simple analogies:

1. The "Living Blueprint" (The Compositional DAG)

Instead of a messy pile of books, CoCoDA builds a living blueprint (a Directed Acyclic Graph, or DAG).

The Nodes: Think of these as individual tools. Some are basic tools (like a hammer or a screwdriver), and others are "super-tools" (like a pre-assembled drill press) that are made by combining the basic ones.
The Connections: The blueprint shows exactly how the super-tools are built from the basic ones. It's like a recipe book where the "Cake" recipe points to the "Mix Batter" and "Bake" steps.
The Benefit: Because the blueprint knows the structure, the robot doesn't need to read every single page of every book. It can look at the "Table of Contents" (signatures) and the "Recipe Steps" (dependencies) to instantly know which tools fit together.

2. The "Smart Filter" (Typed DAG Retrieval)

When the robot faces a new problem, it needs to pick the right tools. Old methods would read the description of every tool in the library, which is slow and uses up too much brain space.

CoCoDA uses a four-stage filter to narrow down the choices, like a security checkpoint:

Stage 1 (The ID Check): It checks the tool's "ID card" (its type signature). If the tool doesn't match the job (e.g., trying to use a hammer to cut paper), it's instantly rejected. This costs almost no brain power.
Stage 2 (The Description Scan): For the tools that passed the ID check, it reads their short descriptions to see if they sound relevant.
Stage 3 (The Rule Check): It checks the "rules" (pre-conditions and post-conditions). Does this tool require inputs it doesn't have? Will it produce the right output? If not, it's tossed out.
Stage 4 (The Example Match): Finally, it looks at real-world examples of the tool in action to pick the absolute best match.

The Magic: Because it filters out bad tools at the cheap, fast stages (1 and 2), it only spends expensive "brain power" reading the full details of the few tools that actually survived the cut. This keeps the robot's "pocket" from getting too heavy, even as the library grows huge.

3. The "Co-Evolution" (Learning Together)

This is the most unique part. Usually, the robot learns to use tools, but the tools stay the same. In CoCoDA, the robot and the toolbox grow up together.

Discovery: When the robot successfully solves a hard problem by chaining together several basic steps (like "add," "multiply," then "divide"), a "Teacher" (a larger, smarter AI) looks at that success.
Invention: The Teacher says, "Hey, that sequence of steps is useful! Let's bundle them into a new 'Super-Tool' and add it to the blueprint."
Reward: The robot gets a bonus score (reward) for using these new Super-Tools because they save it from having to do the basic steps over and over again.
The Loop: As the robot gets better, it finds more patterns to bundle. As the toolbox gets smarter, the robot can solve harder problems. They evolve in sync.

The Results

The paper tested this system on math problems, table analysis, and coding tasks.

The "Small vs. Big" Test: They used a small robot (8 billion "brain units") and compared it to a giant robot (32 billion units).
The Outcome: With CoCoDA, the small robot was able to match or even beat the giant robot on difficult math and logic puzzles.
Efficiency: The small robot didn't just get lucky; it learned to organize its knowledge so efficiently that it didn't need the giant brain size to do the heavy lifting.

In short, CoCoDA teaches a small AI to build its own organized, hierarchical library of skills, allowing it to punch way above its weight class by knowing exactly which tools to grab and how to combine them without getting overwhelmed.

Technical Summary: CoCoDA – Co-evolving Compositional DAG for Tool-Augmented Agents

1. Problem Formulation

Tool-augmented language models (LLMs) extend small planners with external executable skills, allowing them to offload computation and symbolic manipulation. However, scaling the tool library introduces a fundamental tension: the library must evolve with the planner to capture new reusable subroutines, yet retrieval from a growing library must remain within a fixed context budget.

Existing approaches fail to resolve this coupling. Static tool-use systems assume a fixed inventory, while skill-library methods typically store tools as flat text records or natural-language memories. These flat structures cause prompt costs to grow linearly with library size and obscure the typed, compositional structure of executable code. Conversely, code-library methods exploit structure but are not designed for LLM planners operating under token-budgeted retrieval. The core challenge is to design a mechanism that is simultaneously code-aware, retrieval-efficient, and learnable online, solving two coupled problems: how to grow the library with the policy, and how to retrieve from that growing library under a fixed context window.

2. Methodology: The CoCoDA Framework

CoCoDA addresses these challenges through a single code-native structure: a Compositional Code DAG (Directed Acyclic Graph). This framework co-evolves the planner ( $\pi_\theta$ ) and the tool library ( $L$ ) via two primary components: Typed DAG Retrieval and an online co-evolution procedure.

2.1 Compositional Tool DAG Structure

The library $L$ is organized as a DAG $G=(V, E)$ where:

Nodes ( $V$ ): Partitioned into primitive nodes (atomic executable functions) and composite nodes (macro-actions invoking existing tools).
Edges ( $E$ ): Encode invocation dependencies, stratifying the library by abstraction rather than surface similarity.
Node Records ( $I(v)$ ): Each node stores a four-level record:
1. $L_1$ : Typed signature (symbolic).
2. $L_2$ : Natural language description.
3. $L_3$ : Pre/post-condition specifications and complexity.
4. $L_4$ : Worked input-output examples.

2.2 Typed DAG Retrieval

At inference time, CoCoDA employs a coarse-to-fine retrieval cascade that prunes candidates symbolically before incurring expensive LLM costs. Given a query-induced subgoal, the process constructs a nested sequence of survivor sets $S_0 \supseteq S_1 \supseteq S_2 \supseteq S_3 \supseteq S_4$ :

$L_1$ (Symbolic Pruning): Filters tools by unifying the subgoal signature with the tool's typed signature. This stage is symbolic and incurs zero prompt-token cost.
$L_2$ (Semantic Expansion): Ranks survivors by semantic relevance using descriptions. Crucially, this stage can climb to a subsuming composite or descend to a primitive based on the DAG structure, a capability flat text retrieval lacks.
$L_3$ (Specification Filtering): Verifies pre/post-conditions and complexity constraints. Surviving composites are certified as macro-actions equivalent to their primitive sub-trajectories.
$L_4$ (Example Disambiguation): Reranks survivors using worked examples to break ties.

This cascade ensures that expensive LLM-billed records are materialized only for progressively smaller candidate sets, keeping retrieval cost sublinear in library size.

2.3 Online Co-Evolution

The framework updates the planner and library simultaneously based on successful trajectories:

Library Growth: Successful trajectories are passed to a fixed teacher abstractor ( $M_T$ ) which proposes composite tools. A proposal is committed only if it preserves acyclicity and satisfies dependency/specification checks (via INSERTTOOL).
Planner Update: The planner is updated using a DAG-induced reward. Beyond the standard result reward ( $R_{res}$ $R_{r es}$ ), a compositional reward ( $R_{comp}$ $R_{co m p}$ ) credits tools by their "saved-call" count ( $\Phi$ $Φ$ ), defined as the number of primitive calls a composite replaces.
- $R(\tau) = R_{res}(\tau) + \lambda \sum \Phi(t_i)$ .
- This reward encourages the planner to route through deeper composites, effectively shrinking the number of tool calls ( $T(\tau)$ ) required for a task.

3. Key Contributions

Problem Formulation: The authors frame tool-library learning for small LLMs as a joint optimization problem with explicit retrieval-cost and context-window constraints, addressing the coupled growth of policy and library.
Typed DAG Retrieval: A code-native retrieval cascade that utilizes signatures, invocation edges, specifications, and examples to keep prompt cost sublinear in library size.
Co-Evolution Procedure: An online mechanism that inserts validated composite tools and trains the planner with a structural "saved-call" reward, encouraging the exploitation of newly learned abstractions.
Theoretical Guarantees: The paper provides proofs for:
- Retrieval cost reduction and sublinear retrieval time.
- Compositional advantage under the shaped reward (proving that using composites yields strictly higher rewards than equivalent primitive sequences).
- Monotone co-evolution under conservative updates.
- DAG well-formedness (acyclicity and depth bounds).

4. Experimental Results

CoCoDA was evaluated on six benchmarks across mathematical reasoning (GSM8K, MATH), tabular analysis (WikiTableQuestions, FinQA), and code tasks (EvalPlus, MBPP).

Performance: An 8B student model trained with CoCoDA matched or exceeded a 32B teacher model on GSM8K (93.67 vs. 93.40) and MATH (63.18 vs. 61.62).
Scaling: Gains over vanilla Chain-of-Thought (CoT) were most significant at smaller scales (e.g., +10.7 on GSM8K for the 0.6B model).
Baselines: CoCoDA consistently outperformed strong tool-use baselines (ReAct, ToRL, ReTool) and library-learning methods (CREATOR, TroVE). Notably, library-construction baselines without a tight student-teacher loop underperformed plain CoT on math and code tasks, confirming the necessity of co-evolution.
Efficiency: As library size grew to 1,600 nodes, CoCoDA's retrieval cost was approximately 11.4 $\times$ lower than flat retrieval and 4.7 $\times$ lower than text-hierarchical RAG.
Ablation: Removing the Compositional Tool DAG (CTD) caused the largest performance drop (2–4 points), highlighting the importance of dependency edges. Disabling Typed DAG Retrieval (TDR) hurt tabular tasks most, while removing the graph-aware reward (GAR) uniformly reduced performance.

5. Significance and Claims

The paper claims that CoCoDA resolves the tension between library growth and context constraints by leveraging the unique structural properties of executable code. By representing the library as a compositional DAG, the framework enables:

Sublinear Retrieval: The ability to surface complex, reusable abstractions without linearly increasing context costs.
Structural Learning: A reward signal that explicitly credits the planner for compressing primitive sequences into composites, driving the co-evolution of policy and library.
Scalability: The ability for small models (8B) to achieve the performance of much larger models (32B) by effectively utilizing an external, evolving procedural memory.

The authors conclude that this approach offers a viable path for small language models to offload computation and symbolic manipulation efficiently, provided the library co-evolves with the policy through a structure that respects the executable nature of the tools. Future work is noted to extend this to multimodal settings and multi-agent scenarios.