Simulating Ants In Python With OpenAI's o1-preview

Since the release of ChatGPT, I have been testing generative AI models with odd programming tasks. One of my favorites is asking the AI to write an ant simulation in Python (Pygame). I gave OpenAI's latest "reasoning" model, o1-preview, the following instruction:

Can you program an ant simulation in Pygame? Please place 100 ants, an ant nest and 25 sugar piles randomly on the screen. The ants should move randomly on the screen, leaving a scent trail that disappears after a short time. When an ant detects a sugar pile, it will bring sugar to the ant nest. If an ant detects the scent trail of another ant, it will follow it.

The first attempts were not very successful. Either the ants barely moved or they got stuck when they encountered a scent trail. I kept pointing out the problems to o1-preview, but initially without much success. The problem of getting stuck on a scent trail was only solved when I gave o1-preview the following prompt:

The ants still get stuck when they encounter a scent trail. I have an idea to solve this: If an ant moves on a scent trail (by another ant) the scent trail vanishes. If an ant moving in one direction and leaves a scent trail, the ant which encounters the trail will follow it in the opposite direction.

Now the simulation worked almost as desired, one last prompt and I was happy:

This is great! Just make one change: An ant carrying sugar (and leaving a scent trail) should not be able to encounter scent trails until it unloads the sugar at the nest.

It took me a total of about 20 minutes to get o1-preview to write the simulation I wanted. The amazing result was a script of 229 lines of Python code without me writing a single line of code. Moreover, even without my guidance, the code is clearly structured, modular, and follows the rules of object-oriented programming. While this example may be somewhat artificial, it shows me once again that AI can create real value. State-of-the-art models like o1-preview are being trained on NVIDIA GPUs, the engines of this industrial revolution. Here is the final code for the simulation:

import pygame
from pygame.math import Vector2
import random
import time
import math

# Initialize pygame
pygame.init()

# Set display window
WINDOW_WIDTH = 1200
WINDOW_HEIGHT = 700
size = (WINDOW_WIDTH, WINDOW_HEIGHT)
screen = pygame.display.set_mode(size)
pygame.display.set_caption("Ant Simulation")

# Set FPS and Clock
FPS = 60
clock = pygame.time.Clock()

# Colors
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
ANT_COLOR = BLACK
ANT_CARRYING_COLOR = (255, 0, 0)  # Red when carrying sugar
SCENT_COLOR = (170, 160, 160)
SUGAR_COLOR = (255, 215, 0)  # Gold for sugar
NEST_COLOR = (139, 69, 19)  # Brown for nest

# Ant parameters
NUM_ANTS = 100
ANT_SPEED = 2.5
OLFACTION_RADIUS = 75

# Sugar parameters
NUM_SUGAR_PILES = 25
SUGAR_RADIUS = 10

# Nest parameters
NEST_POSITION = (WINDOW_WIDTH // 2, WINDOW_HEIGHT // 2)
NEST_RADIUS = 20


# Ant class
class Ant:
    def __init__(self, x, y, speed, olfaction):
        self.x = x
        self.y = y
        self.speed = speed
        self.olfaction = olfaction
        self.carrying_sugar = False
        # Initialize scent_trail to None
        self.scent_trail = None
        # Random initial direction
        angle = random.uniform(0, 2 * math.pi)
        self.direction = Vector2(math.cos(angle), math.sin(angle))
        self.id = id(self)  # Unique identifier for the ant

    def update(self, all_scents, scents_to_remove):
        if not self.carrying_sugar:
            self.check_scent_overlap(all_scents, scents_to_remove)
        self.move()
        self.check_sugar()
        self.check_nest()

    def move(self):
        # Update position
        self.x += self.direction.x * self.speed
        self.y += self.direction.y * self.speed

        # Reflect at borders
        if self.x <= 0 or self.x >= WINDOW_WIDTH:
            self.direction.x *= -1
            self.x = max(0, min(self.x, WINDOW_WIDTH))
        if self.y <= 0 or self.y >= WINDOW_HEIGHT:
            self.direction.y *= -1
            self.y = max(0, min(self.y, WINDOW_HEIGHT))

        # Leave scent trail only if carrying sugar
        if self.carrying_sugar:
            self.leave_scent()

    def check_scent_overlap(self, all_scents, scents_to_remove):
        ant_pos = Vector2(self.x, self.y)
        for scent in all_scents:
            scent_pos = Vector2(scent['x'], scent['y'])
            if ant_pos.distance_to(scent_pos) < 5:
                if scent['ant_id'] != self.id:
                    # Remove the scent trail point
                    scents_to_remove.append(scent)
                    # Set ant's direction to opposite of scent's direction
                    self.direction = -scent['direction']
                    if self.direction.length() == 0:
                        # Assign a random direction if zero length
                        angle = random.uniform(0, 2 * math.pi)
                        self.direction = Vector2(math.cos(angle), math.sin(angle))
                    else:
                        self.direction = self.direction.normalize()
                    return  # Only process one scent at a time

        # If not carrying sugar and no scent encountered, move randomly
        self.random_move()

    def random_move(self):
        # Slightly adjust the current direction
        angle_change = random.uniform(-0.3, 0.3)
        current_angle = math.atan2(self.direction.y, self.direction.x)
        new_angle = current_angle + angle_change
        self.direction = Vector2(math.cos(new_angle), math.sin(new_angle))
        # Ensure direction is normalized
        if self.direction.length() == 0:
            angle = random.uniform(0, 2 * math.pi)
            self.direction = Vector2(math.cos(angle), math.sin(angle))

    def leave_scent(self):
        self.scent_trail = {
            'x': self.x,
            'y': self.y,
            'time': time.time(),
            'ant_id': self.id,
            'direction': self.direction.copy()
        }

    def check_sugar(self):
        if not self.carrying_sugar:
            ant_pos = Vector2(self.x, self.y)
            for sugar in sugar_piles:
                sugar_pos = Vector2(sugar.x, sugar.y)
                if ant_pos.distance_to(sugar_pos) < SUGAR_RADIUS:
                    self.carrying_sugar = True
                    # Move towards nest
                    nest_pos = Vector2(NEST_POSITION)
                    self.direction = (nest_pos - ant_pos).normalize()
                    break

    def check_nest(self):
        if self.carrying_sugar:
            ant_pos = Vector2(self.x, self.y)
            nest_pos = Vector2(NEST_POSITION)
            if ant_pos.distance_to(nest_pos) < NEST_RADIUS:
                self.carrying_sugar = False  # Drop off sugar at the nest
                # Reset scent_trail after dropping off sugar
                self.scent_trail = None
                # After dropping sugar, pick a random direction
                angle = random.uniform(0, 2 * math.pi)
                self.direction = Vector2(math.cos(angle), math.sin(angle))

    def draw(self):
        ant_color = ANT_COLOR if not self.carrying_sugar else ANT_CARRYING_COLOR
        pygame.draw.circle(screen, ant_color, (int(self.x), int(self.y)), 6)


# SugarPile class
class SugarPile:
    def __init__(self, x, y):
        self.x = x
        self.y = y

    def draw(self):
        pygame.draw.circle(screen, SUGAR_COLOR, (int(self.x), int(self.y)), SUGAR_RADIUS)


# Nest class
class Nest:
    def __init__(self, x, y):
        self.x = x
        self.y = y

    def draw(self):
        pygame.draw.circle(screen, NEST_COLOR, (int(self.x), int(self.y)), NEST_RADIUS)


# Initialize ants
ants = []
for _ in range(NUM_ANTS):
    ants.append(Ant(
        random.randint(NEST_POSITION[0] - 10, NEST_POSITION[0] + 10),
        random.randint(NEST_POSITION[1] - 10, NEST_POSITION[1] + 10),
        ANT_SPEED, OLFACTION_RADIUS))

# Initialize sugar piles
sugar_piles = []
while len(sugar_piles) < NUM_SUGAR_PILES:
    x = random.randint(50, WINDOW_WIDTH - 50)
    y = random.randint(50, WINDOW_HEIGHT - 50)
    # Ensure sugar is not too close to the nest
    if Vector2(x, y).distance_to(Vector2(NEST_POSITION)) > NEST_RADIUS + 100:
        sugar_piles.append(SugarPile(x, y))

# Initialize nest
nest = Nest(*NEST_POSITION)

# Main game loop
running = True
all_scents = []  # Global list of scents
while running:
    clock.tick(FPS)
    screen.fill(WHITE)
    scents_to_remove = []  # List of scents to remove

    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False

    # Update and draw ants
    for ant in ants:
        ant.update(all_scents, scents_to_remove)
        ant.draw()
        # If ant has a scent_trail, add it to global scents
        if ant.scent_trail is not None:
            all_scents.append(ant.scent_trail)

    # Remove old scents and scents that were encountered
    all_scents = [scent for scent in all_scents if time.time() - scent['time'] < 10 and scent not in scents_to_remove]

    # Draw sugar piles
    for sugar in sugar_piles:
        sugar.draw()

    # Draw nest
    nest.draw()

    # Draw scent trails
    for scent in all_scents:
        pygame.draw.circle(screen, SCENT_COLOR, (int(scent['x']), int(scent['y'])), 2)

    pygame.display.flip()

pygame.quit()