diff --git a/llm/llm_schelling/README.md b/llm/llm_schelling/README.md new file mode 100644 index 000000000..1cb272893 --- /dev/null +++ b/llm/llm_schelling/README.md @@ -0,0 +1,73 @@ +# LLM Schelling Segregation + +## Summary + +An LLM-powered reimplementation of Schelling's (1971) classic segregation model where agents +**reason in natural language** about their neighborhood before deciding to stay or move — +instead of following a fixed satisfaction threshold. + +## The Model + +Agents of two groups (A and B) are placed on a 5×5 grid. Each step: + +1. Each agent observes its Moore neighborhood (up to 8 neighbors on a torus grid) +2. It describes the neighborhood composition to the LLM +3. The LLM decides: `happy` (stay) or `unhappy` (move to a random empty cell) + +The simulation tracks happiness levels and a segregation index over time. The model stops +when all agents are happy. + +## What makes this different from classical Schelling + +Classical Schelling uses a fixed threshold rule — an agent moves if fewer than X% of its +neighbors share its group. The outcome is mathematically determined by that threshold. + +Here, agents **reason** at each step: + +> "I belong to Group A. My neighborhood has 3 Group A and 4 Group B neighbors. +> The mix is manageable — I feel comfortable here. I'll stay." + +This produces **qualitatively different dynamics**. The LLM agents weigh context, +not just ratios — and the result shows it. + +## Visualization + +**Initial state (Step 0) — random placement:** + +![Initial state — blue (Group A) and orange (Group B) agents distributed randomly](schelling_initial.png) + +Random mix across the 5×5 grid. No clustering, no preference expressed yet. Happiness and +segregation charts are empty. + +**After 2 LLM reasoning steps — stable integration:** + +![Step 2 — all agents happy, zero segregation, model stopped](schelling_dashboard.png) + +| Step | Happy | Unhappy | Segregation Index | What happened | +|------|-------|---------|-------------------|---------------| +| 0 | — | — | — | Random initial placement | +| 1 | ~16 | ~0 | 0.00 | LLM agents assess neighbors, most decide to stay | +| 2 | 19 | 0 | 0.00 | All agents happy — simulation stops | + +**Why this matters:** The classical Schelling model *always* produces segregation from even +mild preferences. The LLM version produced **zero segregation** — agents reasoned their way +to comfort in a diverse neighborhood. No hardcoded tolerance parameter, no rule relaxation. +The LLM considered context (neighborhood quality, stability, social mix) and decided the +mixed state was acceptable. This is something a fixed-threshold agent cannot do. + +## How to Run + +```bash +cp .env.example .env # fill in your API key +pip install -r requirements.txt +solara run app.py +``` + +## Supported LLM Providers + +Gemini, OpenAI, Anthropic, Ollama (local) — configured via `.env`. + +## Reference + +Schelling, T.C. (1971). Dynamic models of segregation. +*Journal of Mathematical Sociology*, 1(2), 143–186. diff --git a/llm/llm_schelling/app.py b/llm/llm_schelling/app.py new file mode 100644 index 000000000..d9716358c --- /dev/null +++ b/llm/llm_schelling/app.py @@ -0,0 +1,134 @@ +import os +import sys +from pathlib import Path + +from dotenv import load_dotenv + +load_dotenv(Path(__file__).parent / ".env") +os.environ["PYTHONIOENCODING"] = "utf-8" +if sys.stdout.encoding != "utf-8": + sys.stdout.reconfigure(encoding="utf-8") +if sys.stderr.encoding != "utf-8": + sys.stderr.reconfigure(encoding="utf-8") + +import matplotlib.pyplot as plt # noqa: E402 +import solara # noqa: E402 +from llm_schelling.model import LLMSchellingModel # noqa: E402 +from mesa.visualization import SolaraViz # noqa: E402 +from mesa.visualization.utils import update_counter # noqa: E402 + +GROUP_COLORS = {0: "#2196F3", 1: "#FF5722"} # Blue, Orange + +model_params = { + "width": { + "type": "SliderInt", + "value": 5, + "label": "Grid width", + "min": 5, + "max": 20, + "step": 1, + }, + "height": { + "type": "SliderInt", + "value": 5, + "label": "Grid height", + "min": 5, + "max": 20, + "step": 1, + }, + "density": { + "type": "SliderFloat", + "value": 0.8, + "label": "Population density", + "min": 0.1, + "max": 1.0, + "step": 0.05, + }, + "minority_fraction": { + "type": "SliderFloat", + "value": 0.4, + "label": "Minority fraction", + "min": 0.1, + "max": 0.5, + "step": 0.05, + }, +} + + +def SchellingGridPlot(model): + """Visualize the grid showing agent groups and happiness.""" + update_counter.get() + + width = model.grid.dimensions[0] + height = model.grid.dimensions[1] + + fig, ax = plt.subplots(figsize=(6, 6)) + ax.set_xlim(0, width) + ax.set_ylim(0, height) + ax.set_aspect("equal") + ax.set_xticks([]) + ax.set_yticks([]) + ax.set_title("Schelling Segregation (LLM)\nBlue=Group A, Orange=Group B, X=Unhappy") + + for agent in model.agents: + x, y = agent.pos + color = GROUP_COLORS[agent.group] + marker = "o" if agent.is_happy else "x" + ax.plot( + x + 0.5, + y + 0.5, + marker=marker, + color=color, + markersize=8, + markeredgewidth=2, + ) + + return solara.FigureMatplotlib(fig) + + +def SchellingStatsPlot(model): + """Combined happiness + segregation index chart.""" + update_counter.get() + + df = model.datacollector.get_model_vars_dataframe() + + fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(6, 5), facecolor="#1a1a2e") + fig.patch.set_facecolor("#1a1a2e") + + for ax in (ax1, ax2): + ax.set_facecolor("#16213e") + ax.tick_params(colors="white") + ax.xaxis.label.set_color("white") + ax.yaxis.label.set_color("white") + ax.title.set_color("white") + for spine in ax.spines.values(): + spine.set_edgecolor("#444") + + if not df.empty: + ax1.plot(df.index, df["happy"], color="#4CAF50", label="Happy", linewidth=2) + ax1.plot(df.index, df["unhappy"], color="#F44336", label="Unhappy", linewidth=2) + ax1.set_title("Agent Happiness", fontsize=11) + ax1.set_ylabel("Count", color="white") + ax1.legend(facecolor="#16213e", labelcolor="white", fontsize=8) + + if not df.empty and "segregation_index" in df.columns: + ax2.plot(df.index, df["segregation_index"], color="#2196F3", linewidth=2) + ax2.set_title("Segregation Index", fontsize=11) + ax2.set_ylabel("Index", color="white") + ax2.set_xlabel("Step", color="white") + + fig.tight_layout(pad=1.5) + return solara.FigureMatplotlib(fig) + + +model = LLMSchellingModel() + +page = SolaraViz( + model, + components=[ + SchellingGridPlot, + SchellingStatsPlot, + ], + model_params=model_params, + name="LLM Schelling Segregation", +) diff --git a/llm/llm_schelling/llm_schelling/__init__.py b/llm/llm_schelling/llm_schelling/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/llm/llm_schelling/llm_schelling/agent.py b/llm/llm_schelling/llm_schelling/agent.py new file mode 100644 index 000000000..b810d4c3a --- /dev/null +++ b/llm/llm_schelling/llm_schelling/agent.py @@ -0,0 +1,100 @@ +import logging + +from mesa_llm.llm_agent import LLMAgent +from mesa_llm.reasoning.reasoning import Reasoning + +logger = logging.getLogger(__name__) + + +class SchellingAgent(LLMAgent): + """ + An LLM-powered agent in the Schelling Segregation model. + + Unlike the classical Schelling model where agents move if fewer than + a fixed fraction of neighbors share their group, this agent reasons + about its neighborhood using an LLM and decides whether it feels + comfortable staying or wants to move. + + Attributes: + group (int): The agent's group identity (0 or 1). + is_happy (bool): Whether the agent is satisfied with its location. + """ + + def __init__(self, model, reasoning: type[Reasoning], group: int): + group_label = "Group A" if group == 0 else "Group B" + other_label = "Group B" if group == 0 else "Group A" + + system_prompt = f"""You are an agent in a social simulation. You belong to {group_label}. +You are deciding whether you feel comfortable in your current neighborhood. +Look at your neighbors: if too many belong to {other_label} and too few to {group_label}, +you may feel uncomfortable and want to move. +Respond with ONLY one word: 'happy' if you want to stay, or 'unhappy' if you want to move.""" + + super().__init__( + model=model, + reasoning=reasoning, + system_prompt=system_prompt, + vision=1, + internal_state=["group", "is_happy"], + ) + self.group = group + self.is_happy = True + + def step(self): + """Decide whether to move based on LLM reasoning about neighborhood.""" + obs = self.generate_obs() + + # Count neighbors by group + neighbors = list(self.cell.neighborhood.agents) if self.cell else [] + same = sum( + 1 for n in neighbors if hasattr(n, "group") and n.group == self.group + ) + total = len(neighbors) + different = total - same + + if total == 0: + self.is_happy = True + self.internal_state = [f"group:{self.group}", "is_happy:True"] + return + + group_label = "Group A" if self.group == 0 else "Group B" + other_label = "Group B" if self.group == 0 else "Group A" + + step_prompt = f"""You belong to {group_label}. +Your current neighborhood has: +- {same} neighbors from {group_label} (your group) +- {different} neighbors from {other_label} (the other group) +- {total} neighbors total + +Do you feel comfortable here, or do you want to move to a different location? +Respond with ONLY one word: 'happy' or 'unhappy'.""" + + plan = self.reasoning.plan(obs, step_prompt=step_prompt) + + # Parse LLM response + response_text = "" + try: + if hasattr(plan, "llm_plan") and plan.llm_plan: + for block in plan.llm_plan: + if hasattr(block, "text"): + response_text += block.text.lower() + except Exception as e: + logger.warning("Failed to parse LLM response: %s", e) + + self.is_happy = "unhappy" not in response_text + self.internal_state = [ + f"group:{self.group}", + f"is_happy:{self.is_happy}", + ] + + # If unhappy, move to a random empty cell + if not self.is_happy: + empty_cells = [ + cell + for cell in self.model.grid.all_cells + if len(list(cell.agents)) == 0 + ] + if empty_cells: + new_cell = self.model.random.choice(empty_cells) + self.cell = new_cell + self.pos = new_cell.coordinate diff --git a/llm/llm_schelling/llm_schelling/model.py b/llm/llm_schelling/llm_schelling/model.py new file mode 100644 index 000000000..e3cd9da4e --- /dev/null +++ b/llm/llm_schelling/llm_schelling/model.py @@ -0,0 +1,94 @@ +import mesa +from mesa.discrete_space import OrthogonalMooreGrid +from mesa_llm.reasoning.cot import CoTReasoning + +from .agent import SchellingAgent + + +class LLMSchellingModel(mesa.Model): + """ + An LLM-powered Schelling Segregation model. + + The classical Schelling (1971) segregation model shows that even mild + individual preferences for same-group neighbors lead to strong global + segregation. In the classical model, agents move if fewer than a fixed + threshold fraction of their neighbors share their group. + + This model replaces the threshold rule with LLM reasoning: agents + describe their neighborhood in natural language and decide whether they + feel comfortable staying or want to move. This produces richer dynamics + where the decision depends on framing, context, and reasoning — not + just a number. + + Reference: + Schelling, T.C. (1971). Dynamic models of segregation. + Journal of Mathematical Sociology, 1(2), 143-186. + + Args: + width (int): Grid width. + height (int): Grid height. + density (float): Fraction of cells that are occupied. + minority_fraction (float): Fraction of agents in group 1 (minority). + llm_model (str): LLM model string in 'provider/model' format. + rng: Random number generator. + """ + + def __init__( + self, + width: int = 5, + height: int = 5, + density: float = 0.8, + minority_fraction: float = 0.4, + llm_model: str = "cerebras/llama3.1-8b", + rng=None, + ): + super().__init__(rng=rng) + + self.grid = OrthogonalMooreGrid((width, height), torus=True, random=self.random) + + self.datacollector = mesa.DataCollector( + model_reporters={ + "happy": lambda m: sum(1 for a in m.agents if a.is_happy), + "unhappy": lambda m: sum(1 for a in m.agents if not a.is_happy), + "segregation_index": self._segregation_index, + } + ) + + # Place agents on grid + for cell in self.grid.all_cells: + if self.random.random() < density: + group = 1 if self.random.random() < minority_fraction else 0 + agent = SchellingAgent( + model=self, + reasoning=CoTReasoning, + group=group, + ) + agent.cell = cell + agent.pos = cell.coordinate + + self.running = True + self.datacollector.collect(self) + + def step(self): + """Advance the model by one step.""" + self.agents.shuffle_do("step") + self.datacollector.collect(self) + + # Stop if everyone is happy + if all(a.is_happy for a in self.agents): + self.running = False + + @staticmethod + def _segregation_index(model): + """ + Measure segregation as the average fraction of same-group neighbors. + Higher values indicate more segregation. + """ + scores = [] + for agent in model.agents: + neighbors = list(agent.cell.neighborhood.agents) + if not neighbors: + continue + same = sum(1 for n in neighbors if n.group == agent.group) + scores.append(same / len(neighbors)) + return sum(scores) / len(scores) if scores else 0.0 diff --git a/llm/llm_schelling/requirements.txt b/llm/llm_schelling/requirements.txt new file mode 100644 index 000000000..3302bc4d9 --- /dev/null +++ b/llm/llm_schelling/requirements.txt @@ -0,0 +1,2 @@ +mesa[viz]>=3.0 +mesa-llm>=0.1.0 diff --git a/llm/llm_schelling/schelling_dashboard.png b/llm/llm_schelling/schelling_dashboard.png new file mode 100644 index 000000000..a4e8abd2a Binary files /dev/null and b/llm/llm_schelling/schelling_dashboard.png differ diff --git a/llm/llm_schelling/schelling_initial.png b/llm/llm_schelling/schelling_initial.png new file mode 100644 index 000000000..ff43a7d8f Binary files /dev/null and b/llm/llm_schelling/schelling_initial.png differ