Building an AI Agent for USACO Programming Problems

Introduction

This project implements an AI-powered competitive programming agent specifically designed for USACO (USA Computing Olympiad) problems. By leveraging large language models like Claude 3 Opus and GPT-4o, combined with a retrieval-augmented generation approach and human-in-the-loop feedback, the system helps solve algorithmic challenges and provides detailed explanations.

Inspiration

The original idea for using A.I. to solve competitive programming problems in this manner was pioneered by Quan Shi, Michael Tang, Karthik Narasimhan, and Shunyu Yao in their paper "Can Language Models Solve Olympiad Programming?". They proposed competitive programming problems as a way of benchmarking an LM and set up a framework with retrieval and human-in-the-loop to improve performance. This framework, particularly human-in-the-loop, was astoundingly effective and bolstered the performance of GPT-4 from 0% to 86.7% on a certain set of 15 problems.

The USACO past problems made a perfect dataset for this project for several reasons. Firstly, each problem has a consistent format and has been catalogued with a specific difficulty level. Second, each problem has an official solution with both logical reasoning/explanation in natural language as well as an example Java/C++ implementation. Finally, for each problem, there are 10-20 high-quality test cases that usually contain some tricky edge cases.

System Architecture

The system is built on a multi-stage approach:

1. Draft Generation: Claude 3 Opus creates an initial solution approach
2. Example Retrieval: Similar problems are retrieved from the USACO dataset
3. Solution Generation: A comprehensive solution is created with retrieved examples
4. Testing: The solution is executed against test cases
5. Human Feedback: Users can provide hints when solutions fail

def build_graph(draft_solver, solver):
    builder = StateGraph(State)

    builder.add_node("draft", draft_solver)
    builder.add_edge(START, "draft")
    builder.add_node("retrieve", retrieve_examples)
    builder.add_node("solve", solver)
    builder.add_node("evaluate", evaluate)
    builder.add_edge("draft", "retrieve")
    builder.add_edge("retrieve", "solve")
    builder.add_edge("solve", "evaluate")

    def control_edge(state: State):
        if state.get("status") == "success":
            return END
        return "solve"

    builder.add_conditional_edges("evaluate", control_edge, {END: END, "solve": "solve"})
    
    return builder.compile()

Example: Solving the "Cow Crossing" Problem

Here's how the system tackles a typical problem:

Please input the title of the problem: Cow Crossing

Please input the problem text:
Farmer John has N cows (1 ≤ N ≤ 100,000) that need to cross a river.
Each cow has a weight w_i and a swimming speed s_i...

The system first retrieves similar problems:

Retrieved examples:
<problem>
Farmer John has N cows (1 ≤ N ≤ 100,000), each with its own height h_i...
</problem>
<solution>
This is a greedy algorithm problem. We need to sort the cows by their height ratio...

It then generates a Python solution:

def min_crossing_time(cows):
    # Sort cows by w/s ratio (crossing time)
    cows.sort(key=lambda x: x[0]/x[1])
    
    total_time = 0
    for i, (w, s) in enumerate(cows):
        crossing_time = w / s
        total_time += crossing_time * (len(cows) - i)
    
    return int(total_time)

Challenges and Future Work

Several challenges emerged during development:

1. State management: Ensuring correct state transitions in the LangGraph workflow
2. Error handling: Properly capturing and displaying execution errors
3. Output formatting: Limiting verbose outputs from test cases

Future improvements include:

• Implementing a web interface for easier interactions
• Adding visualization tools for algorithms
• Expanding the retrieval database with more competition problems

Conclusion

This project demonstrates the potential of AI-assisted programming education. By combining modern LLMs with retrieval augmentation and human feedback, we've created a system that can tackle complex algorithmic problems and explain its reasoning process. While not perfect, it shows how AI can serve as both a teaching tool and a programming assistant for competitive programmers.

The code is available on GitHub.

---

This project was inspired by cutting-edge research and built with LangChain, LangGraph, and the latest LLMs from Anthropic and OpenAI.