"""The composed society at LLM scale (C3) — E11 re-instantiated in a population of LoRA agents. A population of `N` agents (LoRA adapters on a shared frozen base) evolves for `G` non-overlapping generations under the four operators the paper composes: grounded evaluation, directed recombination (sex), diversity-preserving selection, and retraining (mutation). The single grounding knob acts in the *evaluation* channel, exactly as in E11: selection scores each agent by ``g * verifier_fitness + (1 - g) * conformity``, where conformity is agreement with the population's own modal answer. The inheritance channel is identical in every arm and deliberately ungrounded — each child is a fresh LoRA distilled from its source model's *own answers* (self-consumption made literal), so knowledge survives only through the data channel. Arms (four-arm ablation, mirroring E11): ``full`` / ``no_grounding`` (g=0; the verifier never enters that arm's loop — offspring screening also falls back to conformity) / ``no_sex`` (children are redistilled copies of selected parents) / ``no_diversity`` (plain top-P selection). Pure, testable pieces live at module top (consensus, conformity, behavioural distance, quality-diversity selection, complementary pairing); the GPU loop is :func:`run_society_experiment`. ``python -m llm.experiment configs/llm/society_smoke.yaml``. """ from __future__ import annotations import shutil from pathlib import Path import numpy as np import pandas as pd from .tasks import Task, make_tasks, _normalise # Pure operators live in society_ops (shared with the v2 loop); re-exported here so the v1 code # path and its tests are unchanged. from .society_ops import ( # noqa: E402,F401 answer_of as _answer_of, arm_settings as _arm_settings_v2, behavioural_distance, complementary_pairs, conformity_scores, consensus_answers, fitness_of as _fitness, select_parents, ) def arm_settings(arm: str, g: float) -> dict: """v1 four-arm switches (boolean ``sex``), resolved from the shared table.""" s = _arm_settings_v2(arm, g) if s["sex"] == "linear": raise ValueError("sex_linear is a v2 arm (kind: llm_society_v2)") return {"g": s["g"], "sex": s["sex"] is not None, "diversity": s["diversity"]} # ---------------------------------------------------------------------------- the GPU loop def run_society_experiment(cfg: dict) -> pd.DataFrame: """Run the society loop for every configured arm; return tidy long-form rows. Config keys (with defaults): ``agents`` (6), ``generations`` (8), ``arms`` (all four), ``g`` (0.5), ``lam`` (0.3), ``n_test``/``n_val``/``n_conf``/``n_inherit`` pool sizes, ``n_candidates`` (6) offspring screened per pair, ``epochs`` (2) child SFT epochs, ``spec_train``/``spec_epochs`` gen-0 specialist budget, ``hard`` (False), ``keep_all_adapters`` (False). """ import torch from .evaluate import generate, load_model from .merge import load_specialists from .specialise import train_lora_on_tasks, train_specialist name = cfg["experiment"] base = cfg["base_model"] fams = list(cfg.get("families", ["lists", "strings", "arith"])) N = int(cfg.get("agents", 6)) G = int(cfg.get("generations", 8)) arms = list(cfg.get("arms", ["full", "no_grounding", "no_sex", "no_diversity"])) g_val = float(cfg.get("g", 0.5)) lam = float(cfg.get("lam", 0.3)) n_test = int(cfg.get("n_test", 40)) n_val = int(cfg.get("n_val", 30)) n_conf = int(cfg.get("n_conf", 60)) n_inherit = int(cfg.get("n_inherit", 240)) n_cand = int(cfg.get("n_candidates", 6)) epochs = int(cfg.get("epochs", 2)) elitism = int(cfg.get("elitism", 0)) n_parents = int(cfg.get("n_parents", max(2, N // 2))) spec_train = int(cfg.get("spec_train", 300)) spec_epochs = int(cfg.get("spec_epochs", 2)) hard = bool(cfg.get("hard", False)) keep_all = bool(cfg.get("keep_all_adapters", False)) seed = int(cfg["seed"]) lora = cfg.get("lora", {}) r, alpha = int(lora.get("r", 16)), int(lora.get("alpha", 32)) root = Path(cfg.get("adapters_dir", "models/llm")) / "society" / f"{name}_s{seed}" # fixed pools: test (reporting only), val (grounded selection signal + offspring screening) test = sum([make_tasks(f, n_test, seed=1000 + i, hard=hard) for i, f in enumerate(fams)], []) val = sum([make_tasks(f, n_val, seed=3000 + i, hard=hard) for i, f in enumerate(fams)], []) rng = np.random.default_rng(seed) # gen-0 founders: light per-family specialists, shared across arms (same starting population) founders = [] for i in range(N): fam = fams[i % len(fams)] d = root / "founders" / f"agent{i}_{fam}" if not (d / "adapter_config.json").exists(): train_specialist(base, fam, str(d), n_train=spec_train, epochs=spec_epochs, seed=seed * 100 + i, hard=hard, r=r, alpha=alpha) founders.append(str(d)) rows: list[dict] = [] for arm in arms: s = arm_settings(arm, g_val) agents = list(founders) for t in range(G): conf_pool = sum([make_tasks(f, n_conf // len(fams), seed=seed * 7919 + t * 13 + i, hard=hard) for i, f in enumerate(fams)], []) conf_prompts = [x.prompt for x in conf_pool] # ---- produce & score: one base, all agents attached as adapters model, tok = load_specialists(base, agents) fit_test, fit_val, conf_outs = [], [], [] for i in range(N): model.set_adapter(f"a{i}") fit_test.append(_fitness(generate(model, tok, [x.prompt for x in test]), test, fams)) fit_val.append(_fitness(generate(model, tok, [x.prompt for x in val]), val, fams)) conf_outs.append(generate(model, tok, conf_prompts)) consensus = consensus_answers(conf_outs) conf = conformity_scores(conf_outs, consensus) dist = behavioural_distance(conf_outs) cons_acc = float(np.mean([verify(c, x) for c, x in zip(consensus, conf_pool)])) fitness = np.array([fv["overall"] for fv in fit_val]) scores = s["g"] * fitness + (1.0 - s["g"]) * conf parents = select_parents(scores, dist, n_parents, diversity=s["diversity"], lam=lam) # ---- rows (reporting uses the verifier in every arm; the loop does not) for i in range(N): base_row = {"experiment": name, "arm": arm, "seed": seed, "generation": t, "agent": i, "selected": i in parents, "conformity": float(conf[i]), "score": float(scores[i])} for k, v in fit_test[i].items(): rows.append({**base_row, "metric": f"test_{k}", "value": v}) rows.append({**base_row, "metric": "val_overall", "value": float(fitness[i])}) rows.append({"experiment": name, "arm": arm, "seed": seed, "generation": t, "agent": -1, "selected": False, "conformity": float("nan"), "score": float("nan"), "metric": "consensus_acc", "value": cons_acc}) rows.append({"experiment": name, "arm": arm, "seed": seed, "generation": t, "agent": -1, "selected": False, "conformity": float("nan"), "score": float("nan"), "metric": "diversity_behav", "value": float(dist[np.triu_indices(N, 1)].mean())}) if t == G - 1: del model torch.cuda.empty_cache() break # ---- breed: pick each child's source (merged offspring, or a copied parent) inherit_pool = sum([make_tasks(f, n_inherit // len(fams), seed=seed * 104729 + t * 17 + i, hard=hard) for i, f in enumerate(fams)], []) n_bred = N - elitism child_sources: list[dict] = [] # per child: parent names + weights if s["sex"]: from .directed import sample_merge_weights pairs = complementary_pairs(parents, dist, n_bred) for c, (pa, pb) in enumerate(pairs): w = sample_merge_weights(2, n_cand, rng) best_i, best_v = 0, -np.inf for ci in range(n_cand): cname = f"g{t}c{c}k{ci}" model.add_weighted_adapter([f"a{pa}", f"a{pb}"], w[ci].tolist(), cname, combination_type="linear") model.set_adapter(cname) if s["g"] > 0: # grounded screening: verifier on val v = _fitness(generate(model, tok, [x.prompt for x in val]), val, fams)["overall"] else: # ungrounded screening: conformity only outs_c = generate(model, tok, conf_prompts) v = float(np.mean([_normalise(o) == cc for o, cc in zip(outs_c, consensus)])) if v > best_v: best_i, best_v = ci, v model.set_adapter(f"a{pa}") # never delete the active adapter model.delete_adapter(cname) cname = f"g{t}c{c}win" model.add_weighted_adapter([f"a{pa}", f"a{pb}"], w[best_i].tolist(), cname, combination_type="linear") model.set_adapter(cname) answers = generate(model, tok, [x.prompt for x in inherit_pool]) model.set_adapter(f"a{pa}") model.delete_adapter(cname) child_sources.append({"parents": (pa, pb), "answers": answers}) else: for c in range(n_bred): p = parents[c % len(parents)] model.set_adapter(f"a{p}") answers = generate(model, tok, [x.prompt for x in inherit_pool]) child_sources.append({"parents": (p, p), "answers": answers}) del model torch.cuda.empty_cache() # ---- reproduce: fresh LoRA per child, trained on its source's own answers. # Elites (overlapping generations): the top-scoring parents survive as unmodified # copies, applied identically in every arm — reproduction here is lossy distillation, # so without a survivor the mutation load erases the champion each generation. new_agents = [] for e in range(elitism): d = root / arm / f"gen{t + 1}" / f"elite{e}" shutil.copytree(agents[parents[e]], d, dirs_exist_ok=True) new_agents.append(str(d)) rows.append({"experiment": name, "arm": arm, "seed": seed, "generation": t + 1, "agent": e, "selected": False, "conformity": float("nan"), "score": float("nan"), "metric": "parents", "value": float(parents[e] * 100 + parents[e])}) for c, src in enumerate(child_sources, start=elitism): data = [Task(x.family, x.prompt, _answer_of(a)) for x, a in zip(inherit_pool, src["answers"]) if _answer_of(a)] d = root / arm / f"gen{t + 1}" / f"agent{c}" if len(data) >= 8: train_lora_on_tasks(base, data, str(d), epochs=epochs, seed=seed * 31 + t * N + c, r=r, alpha=alpha) else: # terminal degeneration: copy the source # Reason: a fully-degenerate source emits no usable answers; crashing would kill # a long sweep (cf. the neural terminal-collapse sentinel) — inherit unchanged. shutil.copytree(agents[src["parents"][0]], d, dirs_exist_ok=True) new_agents.append(str(d)) rows.append({"experiment": name, "arm": arm, "seed": seed, "generation": t + 1, "agent": c, "selected": False, "conformity": float("nan"), "score": float("nan"), "metric": "parents", "value": float(src["parents"][0] * 100 + src["parents"][1])}) if not keep_all and t > 0: # disk hygiene: drop generation t for d in agents: # (founders at t=0 are kept — shared) shutil.rmtree(d, ignore_errors=True) agents = new_agents return pd.DataFrame(rows)