#!/usr/bin/env python3 """Soci Agent NN — Local Training Script Equivalent to notebooks/soci_agent_nn.ipynb but runs as a standalone script. Trains the SociAgentTransformer, exports to ONNX, and optionally pushes to HF Hub. Usage: python scripts/nn_train.py # Train from scratch (synthetic data) python scripts/nn_train.py --data data/nn_training # Train on collected + synthetic data python scripts/nn_train.py --push # Train and push to HF Hub python scripts/nn_train.py --epochs 50 --lr 1e-4 # Custom hyperparameters python scripts/nn_train.py --resume # Resume from existing weights Requires: pip install torch onnx onnxruntime numpy huggingface_hub """ from __future__ import annotations import argparse import json import logging import math import os import random import sys import time from collections import Counter from pathlib import Path import numpy as np logging.basicConfig( level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s", stream=sys.stdout, ) logger = logging.getLogger("nn_train") # ── Paths ──────────────────────────────────────────────────────────────── SCRIPT_DIR = Path(__file__).parent PROJECT_DIR = SCRIPT_DIR.parent MODEL_DIR = PROJECT_DIR / "models" DATA_DIR = PROJECT_DIR / "data" / "nn_training" SAMPLES_FILE = DATA_DIR / "collected_samples.jsonl" # ══════════════════════════════════════════════════════════════════════════ # 1. Domain Constants — must match the Soci simulation # ══════════════════════════════════════════════════════════════════════════ ACTION_TYPES = ["move", "work", "eat", "sleep", "talk", "exercise", "shop", "relax", "wander"] ACTION_TO_IDX = {a: i for i, a in enumerate(ACTION_TYPES)} NUM_ACTIONS = len(ACTION_TYPES) LOCATIONS = [ # Residential (17) "house_elena", "house_marcus", "house_helen", "house_diana", "house_kai", "house_priya", "house_james", "house_rosa", "house_yuki", "house_frank", "apartment_block_1", "apartment_block_2", "apartment_block_3", "apt_northeast", "apt_northwest", "apt_southeast", "apt_southwest", # Commercial (8) "cafe", "grocery", "bar", "restaurant", "bakery", "cinema", "diner", "pharmacy", # Work (5) "office", "office_tower", "factory", "school", "hospital", # Public (8) "park", "gym", "library", "church", "town_square", "sports_field", "street_north", "street_south", "street_east", "street_west", ] LOC_TO_IDX = {loc: i for i, loc in enumerate(LOCATIONS)} NUM_LOCATIONS = len(LOCATIONS) # Zone encoding LOC_ZONE = {} for _loc in LOCATIONS: if _loc.startswith(("house_", "apartment_", "apt_")): LOC_ZONE[_loc] = 0 elif _loc in ("cafe", "grocery", "bar", "restaurant", "bakery", "cinema", "diner", "pharmacy"): LOC_ZONE[_loc] = 1 elif _loc in ("office", "office_tower", "factory", "school", "hospital"): LOC_ZONE[_loc] = 2 else: LOC_ZONE[_loc] = 3 ACTION_NEEDS = { "work": {"purpose": 0.3}, "eat": {"hunger": 0.5}, "sleep": {"energy": 0.6}, "talk": {"social": 0.3}, "exercise": {"energy": -0.1, "fun": 0.2, "comfort": 0.1}, "shop": {"hunger": 0.1, "comfort": 0.1}, "relax": {"energy": 0.1, "fun": 0.2, "comfort": 0.2}, "wander": {"fun": 0.1}, "move": {}, } ACTION_DURATIONS = {"move": 1, "work": 4, "eat": 2, "sleep": 8, "talk": 2, "exercise": 3, "shop": 2, "relax": 2, "wander": 1} NEED_NAMES = ["hunger", "energy", "social", "purpose", "comfort", "fun"] PERSONALITY_NAMES = ["openness", "conscientiousness", "extraversion", "agreeableness", "neuroticism"] NUM_TIME_PERIODS = 7 FEATURE_DIM = 47 # ══════════════════════════════════════════════════════════════════════════ # 2. Personas — 20 Soci characters (from personas.yaml) # ══════════════════════════════════════════════════════════════════════════ PERSONAS = [ # House 1 — Elena & Lila (roommates) {"id": "elena", "name": "Elena Vasquez", "age": 34, "gender": "female", "occ": "software engineer", "O": 8, "C": 7, "E": 4, "A": 6, "N": 5, "home": "house_elena", "work": "office", "tags": ["freelance", "introvert", "tech"], "hangouts": ["cafe", "library"], # where she goes to think/work remotely "routine_bias": {}}, {"id": "lila", "name": "Lila Santos", "age": 33, "gender": "female", "occ": "artist", "O": 10, "C": 3, "E": 6, "A": 7, "N": 7, "home": "house_elena", "work": "library", "tags": ["creative", "emotional", "crush_elena"], "hangouts": ["park", "cafe", "library"], # paints outdoors, hangs near Elena "routine_bias": {"relax": 0.15, "wander": 0.10}}, # House 2 — Marcus & Zoe (siblings) {"id": "marcus", "name": "Marcus Chen", "age": 28, "gender": "male", "occ": "fitness trainer", "O": 5, "C": 8, "E": 9, "A": 7, "N": 3, "home": "house_marcus", "work": "gym", "tags": ["athletic", "extrovert", "community"], "hangouts": ["park", "sports_field", "cafe"], "routine_bias": {"exercise": 0.20, "talk": 0.10}}, {"id": "zoe", "name": "Zoe Chen-Williams", "age": 19, "gender": "female", "occ": "college student", "O": 8, "C": 4, "E": 8, "A": 6, "N": 7, "home": "house_marcus", "work": "library", "tags": ["student", "social_media", "young"], "hangouts": ["cafe", "cinema", "park", "town_square"], "routine_bias": {"talk": 0.15, "wander": 0.10}}, # House 3 — Helen & Alice (close friends) {"id": "helen", "name": "Helen Park", "age": 67, "gender": "female", "occ": "retired teacher", "O": 6, "C": 8, "E": 6, "A": 8, "N": 4, "home": "house_helen", "work": "library", "tags": ["retired", "bookworm", "widow"], "hangouts": ["library", "park", "bakery", "church"], "routine_bias": {"relax": 0.15}}, {"id": "alice", "name": "Alice Fontaine", "age": 58, "gender": "female", "occ": "retired accountant", "O": 5, "C": 8, "E": 6, "A": 8, "N": 3, "home": "house_helen", "work": "bakery", "tags": ["retired", "baker", "nurturing"], "hangouts": ["bakery", "grocery", "church"], "routine_bias": {"work": 0.10}}, # loves baking, spends extra time at bakery # House 4 — Diana & Marco (mother & son) {"id": "diana", "name": "Diana Novak", "age": 41, "gender": "female", "occ": "grocery store owner", "O": 4, "C": 9, "E": 5, "A": 6, "N": 7, "home": "house_diana", "work": "grocery", "tags": ["business_owner", "single_mother", "protective"], "hangouts": ["grocery"], # rarely leaves the store "routine_bias": {"work": 0.20}}, {"id": "marco", "name": "Marco Delgado", "age": 16, "gender": "male", "occ": "high school student", "O": 7, "C": 4, "E": 6, "A": 5, "N": 6, "home": "house_diana", "work": "school", "tags": ["student", "teen", "gamer"], "hangouts": ["park", "cinema", "cafe", "sports_field"], "routine_bias": {"relax": 0.10, "wander": 0.10}}, # House 5 — Kai (lives alone) {"id": "kai", "name": "Kai Okonkwo", "age": 22, "gender": "nonbinary", "occ": "barista", "O": 9, "C": 3, "E": 7, "A": 5, "N": 6, "home": "house_kai", "work": "cafe", "tags": ["musician", "creative", "dropout"], "hangouts": ["bar", "park", "town_square"], # plays music, socializes "routine_bias": {"relax": 0.10, "talk": 0.10}}, # House 6 — Priya & Nina (flatmates) {"id": "priya", "name": "Priya Sharma", "age": 38, "gender": "female", "occ": "doctor", "O": 7, "C": 9, "E": 5, "A": 8, "N": 6, "home": "house_priya", "work": "hospital", "tags": ["overworked", "caring", "guilt"], "hangouts": ["hospital", "pharmacy"], # rarely leaves work orbit "routine_bias": {"work": 0.25}}, # long hospital hours {"id": "nina", "name": "Nina Volkov", "age": 29, "gender": "female", "occ": "real estate agent", "O": 5, "C": 8, "E": 9, "A": 4, "N": 5, "home": "house_priya", "work": "office", "tags": ["ambitious", "networker", "suspicious"], "hangouts": ["cafe", "restaurant", "office_tower"], "routine_bias": {"talk": 0.15, "work": 0.10}}, # House 7 — James & Theo (housemates) {"id": "james", "name": "James O'Brien", "age": 55, "gender": "male", "occ": "bar owner", "O": 5, "C": 6, "E": 8, "A": 7, "N": 4, "home": "house_james", "work": "bar", "tags": ["social_hub", "divorced", "storyteller"], "hangouts": ["bar"], # his whole life revolves around the bar "routine_bias": {"talk": 0.20}}, {"id": "theo", "name": "Theo Blackwood", "age": 45, "gender": "male", "occ": "construction worker", "O": 3, "C": 7, "E": 4, "A": 5, "N": 5, "home": "house_james", "work": "factory", "tags": ["blue_collar", "stoic", "handy"], "hangouts": ["bar", "diner"], # bar after work "routine_bias": {"work": 0.15}}, # House 8 — Rosa & Omar {"id": "rosa", "name": "Rosa Martelli", "age": 62, "gender": "female", "occ": "restaurant owner", "O": 6, "C": 9, "E": 7, "A": 8, "N": 5, "home": "house_rosa", "work": "restaurant", "tags": ["nurturing", "italian", "community_mother"], "hangouts": ["restaurant", "grocery"], # buys ingredients, feeds everyone "routine_bias": {"work": 0.20, "eat": 0.05}}, {"id": "omar", "name": "Omar Hassan", "age": 50, "gender": "male", "occ": "taxi driver", "O": 6, "C": 6, "E": 7, "A": 7, "N": 4, "home": "house_rosa", "work": "restaurant", "tags": ["immigrant", "philosophical", "hardworking"], "hangouts": ["restaurant", "cafe", "park"], "routine_bias": {"wander": 0.15}}, # drives around town = wander # House 9 — Yuki & Devon (flatmates) {"id": "yuki", "name": "Yuki Tanaka", "age": 26, "gender": "female", "occ": "yoga instructor", "O": 8, "C": 6, "E": 5, "A": 9, "N": 3, "home": "house_yuki", "work": "gym", "tags": ["mindful", "calm", "empathetic"], "hangouts": ["park", "gym", "library"], # meditates in park "routine_bias": {"exercise": 0.15, "relax": 0.10}}, {"id": "devon", "name": "Devon Reeves", "age": 30, "gender": "male", "occ": "freelance journalist", "O": 9, "C": 5, "E": 6, "A": 4, "N": 6, "home": "house_yuki", "work": "office", "tags": ["investigative", "paranoid", "curious"], "hangouts": ["cafe", "bar", "library", "town_square"], # interviews, research "routine_bias": {"wander": 0.15, "talk": 0.10}}, # House 10 — Frank, George & Sam {"id": "frank", "name": "Frank Kowalski", "age": 72, "gender": "male", "occ": "retired mechanic", "O": 3, "C": 7, "E": 5, "A": 4, "N": 5, "home": "house_frank", "work": "bar", "tags": ["retired", "cantankerous", "creature_of_habit"], "hangouts": ["bar", "diner"], # same bar stool every night "routine_bias": {"relax": 0.15}}, {"id": "george", "name": "George Adeyemi", "age": 47, "gender": "male", "occ": "night shift security", "O": 4, "C": 7, "E": 3, "A": 6, "N": 4, "home": "house_frank", "work": "factory", "tags": ["night_shift", "widower", "observant"], "hangouts": ["park"], # naps in park during day "routine_bias": {}}, # schedule handled by night_shift tag {"id": "sam", "name": "Sam Nakamura", "age": 40, "gender": "nonbinary", "occ": "librarian", "O": 7, "C": 8, "E": 3, "A": 7, "N": 4, "home": "house_frank", "work": "library", "tags": ["quiet", "bookish", "inclusive"], "hangouts": ["library", "park", "cafe"], "routine_bias": {"work": 0.10, "relax": 0.05}}, ] # ══════════════════════════════════════════════════════════════════════════ # 3. Feature Encoding # ══════════════════════════════════════════════════════════════════════════ def _time_period(hour: int) -> int: if hour < 6: return 0 if hour < 9: return 1 if hour < 12: return 2 if hour < 14: return 3 if hour < 18: return 4 if hour < 22: return 5 return 6 def encode_features( persona: dict, hour: int, minute: int, day: int, needs: dict, mood: float, current_loc: str, num_people_here: int = 0, ) -> list[float]: """Encode agent state into 47-dim feature vector.""" f: list[float] = [] # Personality (5) f.append(persona.get("O", persona.get("openness", 5)) / 10.0) f.append(persona.get("C", persona.get("conscientiousness", 5)) / 10.0) f.append(persona.get("E", persona.get("extraversion", 5)) / 10.0) f.append(persona.get("A", persona.get("agreeableness", 5)) / 10.0) f.append(persona.get("N", persona.get("neuroticism", 5)) / 10.0) # Age (1) f.append(persona.get("age", 30) / 100.0) # Time cyclical (4) f.append(math.sin(2 * math.pi * hour / 24)) f.append(math.cos(2 * math.pi * hour / 24)) f.append(math.sin(2 * math.pi * minute / 60)) f.append(math.cos(2 * math.pi * minute / 60)) # Day (2) dow = ((day - 1) % 7) f.append(dow / 7.0) f.append(1.0 if dow >= 5 else 0.0) # Needs (6) for n in NEED_NAMES: f.append(needs.get(n, 0.5)) # Mood (1) f.append(max(-1.0, min(1.0, mood))) # Urgency (2) vals = [needs.get(n, 0.5) for n in NEED_NAMES] urgent_idx = int(np.argmin(vals)) f.append(urgent_idx / 5.0) f.append(1.0 if any(v < 0.15 for v in vals) else 0.0) # Location zone (1) zone = LOC_ZONE.get(current_loc, 3) f.append(zone / 3.0) # Home/work flags (2) home = persona.get("home", persona.get("home_location", "")) work = persona.get("work", persona.get("work_location", "")) f.append(1.0 if current_loc == home else 0.0) f.append(1.0 if current_loc == work else 0.0) # People density (1) f.append(min(num_people_here / 10.0, 1.0)) # Location type one-hot (6) loc_oh = [0.0] * 6 if current_loc.startswith(("house_", "apartment_", "apt_")): loc_oh[0] = 1.0 elif zone == 1: loc_oh[1] = 1.0 elif zone == 2: loc_oh[2] = 1.0 elif current_loc.startswith("street_"): loc_oh[4] = 1.0 else: loc_oh[3] = 1.0 if current_loc == home: loc_oh[5] = 1.0 f.extend(loc_oh) # Time period one-hot (7) tp = [0.0] * NUM_TIME_PERIODS tp[_time_period(hour)] = 1.0 f.extend(tp) # Last action one-hot (9) — random for synthetic, zeros for real last_action_oh = [0.0] * NUM_ACTIONS if random.random() < 0.8: last_action_oh[random.randint(0, NUM_ACTIONS - 1)] = 1.0 f.extend(last_action_oh) return f # ══════════════════════════════════════════════════════════════════════════ # 4. Synthetic Data Generator # ══════════════════════════════════════════════════════════════════════════ def _is_night_shift(persona: dict) -> bool: return "night_shift" in persona.get("tags", []) def _is_retired(persona: dict) -> bool: return "retired" in persona.get("tags", []) def _is_student(persona: dict) -> bool: return "student" in persona.get("tags", []) def _persona_hangout(persona: dict, fallbacks: list[str]) -> str: """Pick a location the persona naturally gravitates toward.""" hangouts = persona.get("hangouts", []) if hangouts and random.random() < 0.6: return random.choice(hangouts) return random.choice(fallbacks) def _apply_routine_bias(persona: dict, action: str | None) -> str | None: """Probabilistically override action based on persona routine_bias.""" bias = persona.get("routine_bias", {}) for biased_action, prob in bias.items(): if random.random() < prob: return biased_action return action def _generate_needs_for_persona(persona: dict, hour: int) -> dict: """Generate needs influenced by persona lifestyle, not purely random.""" needs = {} tags = persona.get("tags", []) is_night = _is_night_shift(persona) for n in NEED_NAMES: # Base: 15% chance critical, else moderate-to-full if random.random() < 0.15: needs[n] = round(random.uniform(0.0, 0.2), 2) else: needs[n] = round(random.uniform(0.2, 1.0), 2) # Persona-specific need tendencies if "overworked" in tags: # Priya: chronically low energy, low social needs["energy"] = round(min(needs["energy"], random.uniform(0.1, 0.5)), 2) needs["social"] = round(min(needs["social"], random.uniform(0.1, 0.5)), 2) if "athletic" in tags: # Marcus: high energy baseline, low fun without exercise needs["energy"] = round(max(needs["energy"], random.uniform(0.5, 0.9)), 2) if "emotional" in tags: # Lila: volatile needs swing = random.choice(NEED_NAMES) needs[swing] = round(random.uniform(0.0, 0.3), 2) if "creature_of_habit" in tags: # Frank: stable moderate needs for n in NEED_NAMES: needs[n] = round(needs[n] * 0.7 + 0.2, 2) if is_night: # George: energy inverted — tired during day, awake at night if 6 <= hour <= 18: needs["energy"] = round(min(needs["energy"], random.uniform(0.05, 0.35)), 2) else: needs["energy"] = round(max(needs["energy"], random.uniform(0.5, 0.9)), 2) if "student" in tags: # Students: higher social need, lower purpose needs["social"] = round(max(needs["social"], random.uniform(0.3, 0.7)), 2) needs["fun"] = round(max(needs["fun"], random.uniform(0.2, 0.5)), 2) if "nurturing" in tags or "community_mother" in tags: # Rosa, Alice: high comfort, purpose from feeding/helping others needs["purpose"] = round(max(needs["purpose"], random.uniform(0.4, 0.8)), 2) if "mindful" in tags: # Yuki: generally balanced, rarely critical for n in NEED_NAMES: needs[n] = round(max(needs[n], 0.2), 2) return needs def _mood_for_persona(persona: dict, needs: dict) -> float: """Generate mood influenced by personality and current needs.""" tags = persona.get("tags", []) # Base mood from needs average avg_need = sum(needs.values()) / len(needs) base_mood = (avg_need - 0.5) * 2 # maps 0-1 to -1..+1 # Neuroticism makes mood more volatile n_factor = persona.get("N", 5) / 10.0 volatility = random.uniform(-0.5, 0.5) * n_factor base_mood += volatility if "calm" in tags or "mindful" in tags: base_mood = base_mood * 0.6 + 0.2 # dampen toward positive if "emotional" in tags: base_mood += random.uniform(-0.4, 0.4) return round(max(-1.0, min(1.0, base_mood)), 2) def _starting_location(persona: dict, hour: int, is_weekend: bool) -> str: """Pick a realistic starting location based on time and persona.""" tags = persona.get("tags", []) is_night = _is_night_shift(persona) period = _time_period(hour) # Night shift workers: at work during night, home during day if is_night: if period in (0, 6): # late night / night — at work return persona["work"] elif period in (1, 2): # morning — heading home or sleeping return random.choice([persona["home"], persona["work"]]) else: # daytime — at home (sleeping) or park (napping) return random.choice([persona["home"], "park"] if random.random() < 0.7 else [persona["home"]]) # Normal schedule if period == 0: # late night — home return persona["home"] elif period == 1: # early morning — home or commuting return random.choice([persona["home"], persona["work"]]) elif period in (2, 4) and not is_weekend: # working hours if _is_retired(persona): return random.choice([persona["home"]] + persona.get("hangouts", ["park"])) if _is_student(persona): return random.choice([persona["work"], "library", persona["home"]]) return random.choice([persona["work"], persona["work"], persona["work"], _persona_hangout(persona, ["cafe"])]) elif period == 3: # lunch return random.choice([persona["work"], "cafe", "restaurant", "diner", "park"]) elif period == 5: # evening return random.choice([persona["home"], _persona_hangout(persona, ["bar", "cafe", "park"])]) elif period == 6: # night return random.choice([persona["home"], persona["home"], _persona_hangout(persona, ["bar"])]) return persona["home"] def generate_action_example(persona: dict) -> dict: """Generate one training example with persona-aware rule-based labels.""" hour = random.randint(0, 23) minute = random.choice([0, 15, 30, 45]) day = random.randint(1, 30) is_weekend = ((day - 1) % 7) >= 5 tags = persona.get("tags", []) is_night = _is_night_shift(persona) needs = _generate_needs_for_persona(persona, hour) mood = _mood_for_persona(persona, needs) current_loc = _starting_location(persona, hour, is_weekend) # --- Determine action using rule-based logic --- # Priority 1: Critical needs urgent = [(n, v) for n, v in needs.items() if v < 0.15] urgent.sort(key=lambda x: x[1]) action = None target_loc = current_loc duration = 1 if urgent: need_name = urgent[0][0] if need_name == "hunger": action = "eat" # Persona-aware eating locations eat_locs = ["cafe", "restaurant", "grocery", "bakery", "diner", persona["home"]] if "community_mother" in tags: # Rosa eats at her restaurant eat_locs = ["restaurant", persona["home"]] elif "baker" in tags: # Alice eats at bakery or home eat_locs = ["bakery", persona["home"]] target_loc = random.choice(eat_locs) duration = 2 elif need_name == "energy": action = "sleep" target_loc = persona["home"] duration = random.choice([4, 6, 8]) elif need_name == "social": action = "talk" social_locs = ["cafe", "bar", "park", "town_square", current_loc] if "social_hub" in tags: # James talks at his bar social_locs = ["bar", "bar", "restaurant", "park"] elif "networker" in tags: # Nina networks everywhere social_locs = ["cafe", "restaurant", "office", "office_tower"] target_loc = random.choice(social_locs) duration = 2 elif need_name == "purpose": action = "work" target_loc = persona["work"] duration = 4 elif need_name == "comfort": action = "relax" target_loc = random.choice([persona["home"], "park", "library"]) duration = 2 elif need_name == "fun": action = random.choice(["relax", "exercise", "wander"]) fun_locs = ["park", "gym", "cinema", "bar", "sports_field"] if "teen" in tags or "student" in tags: fun_locs = ["cinema", "park", "cafe", "sports_field", "town_square"] target_loc = random.choice(fun_locs) duration = 2 # Priority 2: Night shift inverted schedule (George) if action is None and is_night: period = _time_period(hour) if period in (0, 6): # night — George is at work action = "work" target_loc = persona["work"] duration = 4 elif period == 1: # early morning — heading home action = "move" target_loc = persona["home"] duration = 1 elif period in (2, 3): # day — sleeping if needs["energy"] < 0.6: action = "sleep" target_loc = persona["home"] duration = random.choice([4, 6, 8]) else: # Sometimes naps in park action = "relax" target_loc = random.choice([persona["home"], "park"]) duration = 2 elif period in (4, 5): # afternoon/evening — wake up, eat, prep for work r = random.random() if needs["hunger"] < 0.5: action = "eat" target_loc = random.choice(["diner", "restaurant", persona["home"]]) duration = 2 elif r < 0.3: action = "talk" target_loc = random.choice(["park", "cafe"]) duration = 2 else: action = "move" target_loc = persona["work"] duration = 1 # Priority 3: Persona-specific behavioral patterns if action is None: period = _time_period(hour) # Frank: same bar stool every evening/night if persona["id"] == "frank" and period in (5, 6): if random.random() < 0.7: action = "relax" target_loc = "bar" duration = 3 # Lila: gravitates toward Elena (crush) — seeks her hangouts elif persona["id"] == "lila" and random.random() < 0.15: action = random.choice(["wander", "talk", "relax"]) target_loc = random.choice(["house_elena", "cafe", "library", "office"]) duration = 2 # Rosa: spends mornings buying ingredients, cooks all day elif persona["id"] == "rosa" and period in (1, 2): if random.random() < 0.4: action = "shop" target_loc = "grocery" duration = 2 # Devon: investigative journalist, wanders and interviews elif persona["id"] == "devon" and period in (2, 4): if random.random() < 0.3: action = random.choice(["wander", "talk"]) target_loc = random.choice(["cafe", "bar", "town_square", "library", "park"]) duration = 2 # Omar: taxi driver — wanders the streets during work hours elif persona["id"] == "omar" and period in (2, 3, 4) and not is_weekend: if random.random() < 0.5: action = "wander" target_loc = random.choice(["street_north", "street_south", "street_east", "street_west", "town_square", "cafe", "restaurant"]) duration = 2 # Diana: barely leaves the grocery store on weekdays elif persona["id"] == "diana" and not is_weekend and period in (2, 3, 4): if random.random() < 0.7: action = "work" target_loc = "grocery" duration = 4 # Marcus: morning exercise is sacred elif persona["id"] == "marcus" and period == 1: if random.random() < 0.6: action = "exercise" target_loc = random.choice(["gym", "park", "sports_field"]) duration = 3 # Yuki: morning meditation/yoga elif persona["id"] == "yuki" and period == 1: if random.random() < 0.5: action = "exercise" target_loc = random.choice(["park", "gym"]) duration = 3 # Priority 4: Apply routine_bias override if action is None: biased = _apply_routine_bias(persona, None) if biased: action = biased target_loc = _persona_hangout(persona, ["park", "cafe", persona["home"]]) duration = 2 # Priority 5: General time-of-day patterns (fallback) if action is None: period = _time_period(hour) if period == 0: # Late night action = "sleep" target_loc = persona["home"] duration = 8 elif period == 1: # Early morning r = random.random() if needs["hunger"] < 0.5: action = "eat" target_loc = random.choice(["cafe", "bakery", persona["home"]]) duration = 2 elif r < 0.3 and persona["E"] >= 6: action = "exercise" target_loc = random.choice(["gym", "park", "sports_field"]) duration = 3 else: action = "move" target_loc = persona["work"] duration = 1 elif period in (2, 4): # Mid-morning / Afternoon if is_weekend: r = random.random() if _is_retired(persona): # Retired: relaxed weekend routine if r < 0.35: action = "relax" target_loc = _persona_hangout(persona, ["park", "library", persona["home"]]) elif r < 0.55: action = "talk" target_loc = _persona_hangout(persona, ["cafe", "park", "church"]) elif r < 0.7: action = "shop" target_loc = random.choice(["grocery", "pharmacy", "bakery"]) else: action = "wander" target_loc = random.choice(["park", "town_square", "street_north"]) duration = random.choice([2, 3]) elif _is_student(persona): # Students: social weekends if r < 0.3: action = "talk" target_loc = random.choice(["cafe", "park", "cinema", "town_square"]) elif r < 0.5: action = "relax" target_loc = random.choice(["cinema", "park", persona["home"]]) elif r < 0.65: action = "exercise" target_loc = random.choice(["gym", "park", "sports_field"]) elif r < 0.8: action = "wander" target_loc = random.choice(["town_square", "street_north", "street_south"]) else: action = "shop" target_loc = random.choice(["grocery", "pharmacy"]) duration = random.choice([2, 3]) else: if r < 0.25: action = "relax" target_loc = _persona_hangout(persona, ["park", "cafe", "library", persona["home"]]) elif r < 0.45 and persona["E"] >= 6: action = "talk" target_loc = _persona_hangout(persona, ["cafe", "park", "town_square"]) elif r < 0.6: action = "shop" target_loc = random.choice(["grocery", "pharmacy"]) elif r < 0.8: action = "exercise" target_loc = random.choice(["gym", "park", "sports_field"]) else: action = "wander" target_loc = random.choice(["park", "town_square", "street_north", "street_south"]) duration = random.choice([2, 3]) else: # Weekday work hours work_prob = 0.5 + persona["C"] * 0.05 # Business owners and doctors work even harder if "business_owner" in tags or persona["occ"] == "doctor": work_prob += 0.15 if _is_retired(persona): work_prob = 0.15 # retired people rarely "work" if random.random() < work_prob: action = "work" target_loc = persona["work"] duration = 4 else: action = random.choice(["wander", "relax", "talk"]) target_loc = _persona_hangout(persona, ["cafe", "park", "town_square"]) duration = 2 elif period == 3: # Midday / lunch if needs["hunger"] < 0.6: action = "eat" lunch_locs = ["cafe", "restaurant", "bakery", "diner", "park"] # People eat near their workplace if current_loc == persona["work"]: lunch_locs = ["cafe", "restaurant", "diner", "bakery"] target_loc = random.choice(lunch_locs) duration = 2 else: action = "relax" target_loc = random.choice(["park", "cafe"]) duration = 1 elif period == 5: # Evening r = random.random() social_bias = persona["E"] / 10.0 if r < social_bias * 0.5: action = "talk" evening_social = ["bar", "restaurant", "park", "cafe"] if "social_hub" in tags: evening_social = ["bar", "bar", "restaurant"] target_loc = random.choice(evening_social) duration = 2 elif r < 0.4: action = "eat" target_loc = random.choice(["restaurant", "bar", "diner", persona["home"]]) duration = 2 elif r < 0.55: action = "exercise" target_loc = random.choice(["gym", "park", "sports_field"]) duration = 3 elif r < 0.7: action = "relax" target_loc = _persona_hangout(persona, ["cinema", "bar", persona["home"], "library"]) duration = 2 else: action = "relax" target_loc = persona["home"] duration = 2 elif period == 6: # Night if needs["energy"] < 0.4: action = "sleep" target_loc = persona["home"] duration = 8 else: action = "relax" target_loc = persona["home"] duration = 2 # 30% chance of picking "move" if target != current if target_loc != current_loc and action != "move": if random.random() < 0.3: action = "move" duration = 1 # Retired and elderly people do shorter activities if _is_retired(persona) and duration > 3 and action not in ("sleep", "work"): duration = min(duration, 3) # Teens/students have shorter attention spans for non-social activities if _is_student(persona) and action in ("relax", "work") and random.random() < 0.3: duration = max(1, duration - 1) features = encode_features( persona=persona, hour=hour, minute=minute, day=day, needs=needs, mood=mood, current_loc=current_loc, num_people_here=random.randint(0, 8), ) return { "features": features, "action_idx": ACTION_TO_IDX[action], "target_loc_idx": LOC_TO_IDX.get(target_loc, 0), "duration": min(max(duration, 1), 8), } def generate_dataset(n: int) -> list[dict]: """Generate n synthetic training examples.""" data = [] for _ in range(n): persona = random.choice(PERSONAS) data.append(generate_action_example(persona)) return data # ══════════════════════════════════════════════════════════════════════════ # 5. Model Architecture — SociAgentTransformer # ══════════════════════════════════════════════════════════════════════════ def build_model(): """Build the SociAgentTransformer model.""" import torch import torch.nn as nn import torch.nn.functional as F class FeatureTokenizer(nn.Module): GROUPS = [ ("personality", 0, 6), ("time", 6, 12), ("needs", 12, 21), ("location", 21, 31), ("time_period", 31, 38), ("last_action", 38, 47), ] def __init__(self, d_model: int): super().__init__() self.projections = nn.ModuleList() for name, start, end in self.GROUPS: self.projections.append(nn.Sequential( nn.Linear(end - start, d_model), nn.LayerNorm(d_model), nn.GELU(), )) self.pos_embed = nn.Parameter(torch.randn(1, len(self.GROUPS), d_model) * 0.02) def forward(self, features): tokens = [] for i, (_, start, end) in enumerate(self.GROUPS): tokens.append(self.projections[i](features[:, start:end])) tokens = torch.stack(tokens, dim=1) return tokens + self.pos_embed class MoEFeedForward(nn.Module): def __init__(self, d_model, d_ff, num_experts=4, top_k=2): super().__init__() self.num_experts = num_experts self.top_k = top_k self.gate = nn.Linear(d_model, num_experts, bias=False) self.experts = nn.ModuleList([ nn.Sequential(nn.Linear(d_model, d_ff), nn.GELU(), nn.Linear(d_ff, d_model)) for _ in range(num_experts) ]) def forward(self, x): B, S, D = x.shape gate_probs = F.softmax(self.gate(x), dim=-1) top_k_probs, top_k_idx = gate_probs.topk(self.top_k, dim=-1) top_k_probs = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True) output = torch.zeros_like(x) for k in range(self.top_k): eidx = top_k_idx[:, :, k] w = top_k_probs[:, :, k].unsqueeze(-1) for e in range(self.num_experts): mask = (eidx == e).unsqueeze(-1) if mask.any(): output = output + mask.float() * w * self.experts[e](x) return output class TransformerBlock(nn.Module): def __init__(self, d_model, nhead, d_ff, num_experts=4, dropout=0.1): super().__init__() self.attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=True) self.norm1 = nn.LayerNorm(d_model) self.moe_ff = MoEFeedForward(d_model, d_ff, num_experts) self.norm2 = nn.LayerNorm(d_model) self.dropout = nn.Dropout(dropout) def forward(self, x): attn_out, _ = self.attn(x, x, x) x = self.norm1(x + self.dropout(attn_out)) ff_out = self.moe_ff(x) return self.norm2(x + self.dropout(ff_out)) class SociAgentTransformer(nn.Module): def __init__(self, d_model=128, nhead=8, num_layers=4, d_ff=256, num_experts=4, dropout=0.1): super().__init__() self.tokenizer = FeatureTokenizer(d_model) self.layers = nn.ModuleList([ TransformerBlock(d_model, nhead, d_ff, num_experts, dropout) for _ in range(num_layers) ]) self.cls_query = nn.Parameter(torch.randn(1, 1, d_model) * 0.02) self.cls_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=True) self.cls_norm = nn.LayerNorm(d_model) self.action_head = nn.Sequential( nn.Linear(d_model, d_model), nn.GELU(), nn.Dropout(dropout), nn.Linear(d_model, NUM_ACTIONS), ) self.location_head = nn.Sequential( nn.Linear(d_model + NUM_ACTIONS, d_model), nn.GELU(), nn.Dropout(dropout), nn.Linear(d_model, NUM_LOCATIONS), ) self.duration_head = nn.Sequential( nn.Linear(d_model + NUM_ACTIONS, d_model // 2), nn.GELU(), nn.Linear(d_model // 2, 1), ) def forward(self, features): tokens = self.tokenizer(features) for layer in self.layers: tokens = layer(tokens) B = features.shape[0] cls = self.cls_query.expand(B, -1, -1) cls_out, _ = self.cls_attn(cls, tokens, tokens) h = self.cls_norm(cls_out.squeeze(1)) action_logits = self.action_head(h) action_probs = F.softmax(action_logits.detach(), dim=-1) h_a = torch.cat([h, action_probs], dim=-1) location_logits = self.location_head(h_a) duration = torch.sigmoid(self.duration_head(h_a)) * 7.0 + 1.0 return { "action_logits": action_logits, "location_logits": location_logits, "duration": duration.squeeze(-1), } return SociAgentTransformer() # ══════════════════════════════════════════════════════════════════════════ # 6. Training # ══════════════════════════════════════════════════════════════════════════ def train( epochs: int = 30, batch_size: int = 512, lr: float = 3e-4, num_train: int = 100_000, num_val: int = 10_000, data_dir: str | None = None, resume: bool = False, ): """Full training pipeline: generate/load data, train, export ONNX.""" import torch import torch.nn as nn from torch.utils.data import Dataset, DataLoader DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") logger.info(f"Device: {DEVICE}") if DEVICE.type == "cuda": logger.info(f"GPU: {torch.cuda.get_device_name()}") MODEL_DIR.mkdir(parents=True, exist_ok=True) best_pt = MODEL_DIR / "soci_agent_best.pt" onnx_path = MODEL_DIR / "soci_agent.onnx" # ── Load / generate data ───────────────────────────────────────── collected = [] source_counts: dict[str, int] = {} # Load collected samples from live sim (if available) samples_file = Path(data_dir) / "collected_samples.jsonl" if data_dir else SAMPLES_FILE if samples_file.exists(): with open(samples_file) as f: for line in f: line = line.strip() if line: sample = json.loads(line) collected.append(sample) src = sample.get("source", "unknown") source_counts[src] = source_counts.get(src, 0) + 1 logger.info(f"Loaded {len(collected):,} collected samples — sources: {source_counts}") # Oversample LLM-sourced data 3x (higher quality than NN/routine) llm_sources = {"gemini", "claude", "groq"} llm_samples = [s for s in collected if s.get("source", "") in llm_sources] if llm_samples: logger.info(f"Oversampling {len(llm_samples):,} LLM-sourced samples (3x weight)") collected.extend(llm_samples * 2) # Generate synthetic data to fill up to target size total_target = num_train + num_val synthetic_needed = max(0, total_target - len(collected)) if synthetic_needed > 0: logger.info(f"Generating {synthetic_needed:,} synthetic samples...") random.seed(42) collected.extend(generate_dataset(synthetic_needed)) random.shuffle(collected) split = int(len(collected) * 0.9) train_data = collected[:split] val_data = collected[split:] # ── Dataset ────────────────────────────────────────────────────── class ActionDataset(Dataset): def __init__(self, data): self.features = torch.tensor([d["features"] for d in data], dtype=torch.float32) self.actions = torch.tensor([d["action_idx"] for d in data], dtype=torch.long) self.locations = torch.tensor([d["target_loc_idx"] for d in data], dtype=torch.long) self.durations = torch.tensor([d["duration"] for d in data], dtype=torch.float32) def __len__(self): return len(self.actions) def __getitem__(self, idx): return { "features": self.features[idx], "action": self.actions[idx], "location": self.locations[idx], "duration": self.durations[idx], } train_ds = ActionDataset(train_data) val_ds = ActionDataset(val_data) train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=0, pin_memory=(DEVICE.type == "cuda")) val_loader = DataLoader(val_ds, batch_size=1024, shuffle=False, num_workers=0, pin_memory=(DEVICE.type == "cuda")) logger.info(f"Train: {len(train_ds):,}, Val: {len(val_ds):,}") # ── Model ──────────────────────────────────────────────────────── model = build_model().to(DEVICE) total_params = sum(p.numel() for p in model.parameters()) logger.info(f"Model parameters: {total_params:,} ({total_params * 4 / 1024 / 1024:.1f} MB fp32)") if resume and best_pt.exists(): model.load_state_dict(torch.load(str(best_pt), map_location=DEVICE, weights_only=True)) logger.info(f"Resumed from {best_pt}") # ── Class weights ──────────────────────────────────────────────── action_counts = torch.zeros(NUM_ACTIONS) for d in train_data: action_counts[d["action_idx"]] += 1 action_weights = 1.0 / (action_counts + 1.0) action_weights = action_weights / action_weights.sum() * NUM_ACTIONS action_weights = action_weights.to(DEVICE) logger.info("Action distribution:") for idx in range(NUM_ACTIONS): count = int(action_counts[idx]) pct = count / len(train_data) * 100 logger.info(f" {ACTION_TYPES[idx]:>10s}: {count:6d} ({pct:.1f}%)") # ── Loss & optimizer ───────────────────────────────────────────── action_loss_fn = nn.CrossEntropyLoss(weight=action_weights) location_loss_fn = nn.CrossEntropyLoss() duration_loss_fn = nn.MSELoss() W_ACTION = 1.0 W_LOCATION = 0.5 W_DURATION = 0.2 optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs, eta_min=1e-6) logger.info(f"Training for {epochs} epochs, LR={lr}, batch_size={batch_size}") # ── Training loop ──────────────────────────────────────────────── best_val_acc = 0.0 history = {"train_loss": [], "val_loss": [], "val_action_acc": [], "val_loc_acc": []} for epoch in range(epochs): # Train model.train() total_loss = 0.0 n_batches = 0 for batch in train_loader: feat = batch["features"].to(DEVICE) out = model(feat) loss = ( W_ACTION * action_loss_fn(out["action_logits"], batch["action"].to(DEVICE)) + W_LOCATION * location_loss_fn(out["location_logits"], batch["location"].to(DEVICE)) + W_DURATION * duration_loss_fn(out["duration"], batch["duration"].to(DEVICE)) ) optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() total_loss += loss.item() n_batches += 1 scheduler.step() avg_train_loss = total_loss / n_batches # Validate model.eval() val_loss = 0.0 correct_action = 0 correct_loc = 0 total = 0 with torch.no_grad(): for batch in val_loader: feat = batch["features"].to(DEVICE) out = model(feat) loss = ( W_ACTION * action_loss_fn(out["action_logits"], batch["action"].to(DEVICE)) + W_LOCATION * location_loss_fn(out["location_logits"], batch["location"].to(DEVICE)) + W_DURATION * duration_loss_fn(out["duration"], batch["duration"].to(DEVICE)) ) val_loss += loss.item() pred_action = out["action_logits"].argmax(dim=-1) pred_loc = out["location_logits"].argmax(dim=-1) correct_action += (pred_action == batch["action"].to(DEVICE)).sum().item() correct_loc += (pred_loc == batch["location"].to(DEVICE)).sum().item() total += feat.shape[0] avg_val_loss = val_loss / len(val_loader) action_acc = correct_action / total if total > 0 else 0 loc_acc = correct_loc / total if total > 0 else 0 history["train_loss"].append(avg_train_loss) history["val_loss"].append(avg_val_loss) history["val_action_acc"].append(action_acc) history["val_loc_acc"].append(loc_acc) if action_acc > best_val_acc: best_val_acc = action_acc torch.save(model.state_dict(), str(best_pt)) if (epoch + 1) % 5 == 0 or epoch == 0: lr_now = scheduler.get_last_lr()[0] logger.info( f"Epoch {epoch+1:3d}/{epochs} | " f"Train: {avg_train_loss:.4f} | " f"Val: {avg_val_loss:.4f} | " f"Act Acc: {action_acc:.1%} | " f"Loc Acc: {loc_acc:.1%} | " f"LR: {lr_now:.2e}" ) logger.info(f"Best validation action accuracy: {best_val_acc:.1%}") # ── Per-action accuracy ────────────────────────────────────────── model.load_state_dict(torch.load(str(best_pt), map_location=DEVICE, weights_only=True)) model.eval() cm = np.zeros((NUM_ACTIONS, NUM_ACTIONS), dtype=int) with torch.no_grad(): for batch in val_loader: feat = batch["features"].to(DEVICE) out = model(feat) preds = out["action_logits"].argmax(dim=-1).cpu().numpy() labels = batch["action"].numpy() for p, l in zip(preds, labels): cm[l][p] += 1 logger.info("Per-action accuracy:") for i, action in enumerate(ACTION_TYPES): row_total = cm[i].sum() correct = cm[i][i] acc = correct / row_total if row_total > 0 else 0 logger.info(f" {action:>10s}: {acc:.1%} ({correct}/{row_total})") # ── Test scenarios ─────────────────────────────────────────────── import torch.nn.functional as F @torch.no_grad() def predict(persona, hour, minute, day, needs, mood, loc, num_people=0): features = encode_features(persona, hour, minute, day, needs, mood, loc, num_people) feat_t = torch.tensor([features], dtype=torch.float32, device=DEVICE) out = model(feat_t) action_probs = F.softmax(out["action_logits"][0] / 0.7, dim=-1) action_idx = action_probs.argmax().item() loc_idx = out["location_logits"][0].argmax().item() dur = max(1, min(8, round(out["duration"][0].item()))) return ACTION_TYPES[action_idx], LOCATIONS[loc_idx], dur, action_probs[action_idx].item() logger.info("Test scenarios:") a, l, d, c = predict(PERSONAS[0], 0, 30, 5, {"hunger": 0.5, "energy": 0.05, "social": 0.4, "purpose": 0.6, "comfort": 0.3, "fun": 0.3}, -0.3, "office") logger.info(f" Elena midnight exhausted at office: {a} -> {l} ({d} ticks, {c:.0%})") a, l, d, c = predict(PERSONAS[2], 12, 30, 3, {"hunger": 0.05, "energy": 0.7, "social": 0.5, "purpose": 0.6, "comfort": 0.5, "fun": 0.4}, 0.2, "gym", 5) logger.info(f" Marcus lunchtime starving at gym: {a} -> {l} ({d} ticks, {c:.0%})") a, l, d, c = predict(PERSONAS[8], 10, 0, 6, {"hunger": 0.6, "energy": 0.7, "social": 0.5, "purpose": 0.5, "comfort": 0.7, "fun": 0.4}, 0.5, "house_kai") logger.info(f" Kai Saturday morning at home: {a} -> {l} ({d} ticks, {c:.0%})") # George (night shift) — should sleep during the day george = [p for p in PERSONAS if p["id"] == "george"][0] a, l, d, c = predict(george, 11, 0, 3, {"hunger": 0.4, "energy": 0.15, "social": 0.5, "purpose": 0.7, "comfort": 0.5, "fun": 0.4}, -0.1, "house_frank") logger.info(f" George midday after night shift: {a} -> {l} ({d} ticks, {c:.0%})") # Frank — evening at the bar frank = [p for p in PERSONAS if p["id"] == "frank"][0] a, l, d, c = predict(frank, 20, 0, 4, {"hunger": 0.5, "energy": 0.4, "social": 0.3, "purpose": 0.6, "comfort": 0.5, "fun": 0.3}, 0.1, "bar") logger.info(f" Frank evening at the bar: {a} -> {l} ({d} ticks, {c:.0%})") # Priya — overworked at hospital priya = [p for p in PERSONAS if p["id"] == "priya"][0] a, l, d, c = predict(priya, 15, 0, 2, {"hunger": 0.3, "energy": 0.2, "social": 0.3, "purpose": 0.8, "comfort": 0.4, "fun": 0.2}, -0.2, "hospital") logger.info(f" Priya afternoon exhausted at hospital: {a} -> {l} ({d} ticks, {c:.0%})") # ── Export to ONNX ─────────────────────────────────────────────── logger.info("Exporting to ONNX...") model.cpu().eval() dummy = torch.randn(1, FEATURE_DIM) torch.onnx.export( model, dummy, str(onnx_path), input_names=["features"], output_names=["action_logits", "location_logits", "duration"], dynamic_axes={"features": {0: "batch"}}, opset_version=17, dynamo=False, ) # Verify ONNX import onnx onnx_model = onnx.load(str(onnx_path)) onnx.checker.check_model(onnx_model) onnx_size = onnx_path.stat().st_size / 1024 logger.info(f"ONNX exported: {onnx_path} ({onnx_size:.0f} KB)") # Benchmark ONNX import onnxruntime as ort session = ort.InferenceSession(str(onnx_path)) batch_input = np.random.randn(50, FEATURE_DIM).astype(np.float32) start = time.perf_counter() for _ in range(100): session.run(None, {"features": batch_input}) elapsed = (time.perf_counter() - start) / 100 logger.info(f"ONNX inference (50 agents): {elapsed*1000:.1f} ms per batch") # ── Save training stats ────────────────────────────────────────── stats = { "best_val_action_acc": best_val_acc, "epochs": epochs, "train_samples": len(train_ds), "val_samples": len(val_ds), "collected_samples": sum(source_counts.values()), "source_counts": source_counts, "model_size_kb": onnx_size, "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"), "history": history, } stats_path = MODEL_DIR / "training_stats.json" stats_path.write_text(json.dumps(stats, indent=2)) logger.info(f"Stats saved to {stats_path}") # ── Plot training graphs ────────────────────────────────────────── plot_training_graphs(stats_path) return best_val_acc def plot_training_graphs(stats_path: Path | str | None = None): """Plot training loss and accuracy curves from saved training stats. Saves the plot to models/training_graphs.png and displays it. """ import matplotlib matplotlib.use("Agg") # non-interactive backend as fallback import matplotlib.pyplot as plt stats_path = Path(stats_path) if stats_path else MODEL_DIR / "training_stats.json" if not stats_path.exists(): logger.error(f"No training stats found at {stats_path}") return stats = json.loads(stats_path.read_text()) history = stats.get("history", {}) train_loss = history.get("train_loss", []) val_loss = history.get("val_loss", []) val_action_acc = history.get("val_action_acc", []) val_loc_acc = history.get("val_loc_acc", []) if not train_loss: logger.error("No training history found in stats file") return epochs_range = list(range(1, len(train_loss) + 1)) fig, axes = plt.subplots(1, 3, figsize=(18, 5)) fig.suptitle( f"Soci Agent NN Training — {stats.get('timestamp', '?')} | " f"Best Action Acc: {stats.get('best_val_action_acc', 0):.1%}", fontsize=13, fontweight="bold", ) # Loss curves ax = axes[0] ax.plot(epochs_range, train_loss, label="Train Loss", color="#2196F3", linewidth=2) ax.plot(epochs_range, val_loss, label="Val Loss", color="#F44336", linewidth=2) ax.set_xlabel("Epoch") ax.set_ylabel("Loss") ax.set_title("Training & Validation Loss") ax.legend() ax.grid(True, alpha=0.3) ax.set_xlim(1, len(train_loss)) # Action accuracy ax = axes[1] ax.plot(epochs_range, [a * 100 for a in val_action_acc], label="Action Accuracy", color="#4CAF50", linewidth=2) best_epoch = int(np.argmax(val_action_acc)) + 1 best_acc = max(val_action_acc) * 100 ax.axhline(y=best_acc, color="#4CAF50", linestyle="--", alpha=0.4) ax.annotate(f"Best: {best_acc:.1f}% (epoch {best_epoch})", xy=(best_epoch, best_acc), fontsize=9, xytext=(best_epoch + 1, best_acc - 3), arrowprops=dict(arrowstyle="->", color="#4CAF50"), color="#4CAF50") ax.set_xlabel("Epoch") ax.set_ylabel("Accuracy (%)") ax.set_title("Action Prediction Accuracy") ax.legend() ax.grid(True, alpha=0.3) ax.set_xlim(1, len(train_loss)) # Location accuracy ax = axes[2] if val_loc_acc: ax.plot(epochs_range, [a * 100 for a in val_loc_acc], label="Location Accuracy", color="#FF9800", linewidth=2) best_loc_epoch = int(np.argmax(val_loc_acc)) + 1 best_loc = max(val_loc_acc) * 100 ax.axhline(y=best_loc, color="#FF9800", linestyle="--", alpha=0.4) ax.annotate(f"Best: {best_loc:.1f}% (epoch {best_loc_epoch})", xy=(best_loc_epoch, best_loc), fontsize=9, xytext=(best_loc_epoch + 1, best_loc - 3), arrowprops=dict(arrowstyle="->", color="#FF9800"), color="#FF9800") ax.set_xlabel("Epoch") ax.set_ylabel("Accuracy (%)") ax.set_title("Location Prediction Accuracy") ax.legend() ax.grid(True, alpha=0.3) ax.set_xlim(1, len(train_loss)) # Footer with training info footer = ( f"Train: {stats.get('train_samples', '?'):,} samples | " f"Val: {stats.get('val_samples', '?'):,} samples | " f"Collected: {stats.get('collected_samples', 0):,} | " f"Model: {stats.get('model_size_kb', 0):.0f} KB" ) fig.text(0.5, 0.01, footer, ha="center", fontsize=9, color="gray") plt.tight_layout(rect=[0, 0.03, 1, 0.95]) graph_path = MODEL_DIR / "training_graphs.png" fig.savefig(str(graph_path), dpi=150, bbox_inches="tight") logger.info(f"Training graphs saved to {graph_path}") # Try to display interactively try: import warnings with warnings.catch_warnings(): warnings.simplefilter("ignore") matplotlib.use("TkAgg") plt.show(block=False) plt.pause(0.5) except Exception: pass # headless environment, PNG saved is enough plt.close(fig) def _push_to_hub(best_pt, onnx_path, stats_path, repo_id, best_val_acc, epochs, num_train, base_url: str = "https://raymelius-soci2.hf.space"): """Upload model files to HuggingFace Hub, then trigger live reload.""" from huggingface_hub import HfApi, login token = os.environ.get("HF_TOKEN", "") if not token: logger.error("HF_TOKEN not set — cannot push. Export it: export HF_TOKEN=hf_...") return login(token=token) api = HfApi() api.create_repo(repo_id, exist_ok=True) # Config config = { "architecture": "SociAgentTransformer", "d_model": 128, "nhead": 8, "num_layers": 4, "d_ff": 256, "num_experts": 4, "feature_dim": FEATURE_DIM, "num_actions": NUM_ACTIONS, "num_locations": NUM_LOCATIONS, "action_types": ACTION_TYPES, "locations": LOCATIONS, "action_durations": ACTION_DURATIONS, "need_names": NEED_NAMES, "personality_names": PERSONALITY_NAMES, "best_val_action_acc": best_val_acc, "training_samples": num_train, "epochs": epochs, } config_path = MODEL_DIR / "config.json" config_path.write_text(json.dumps(config, indent=2)) for local, remote in [ (onnx_path, "soci_agent.onnx"), (best_pt, "soci_agent_best.pt"), (config_path, "config.json"), (stats_path, "training_stats.json"), ]: if local.exists(): api.upload_file( path_or_fileobj=str(local), path_in_repo=remote, repo_id=repo_id, commit_message=f"Train: acc={best_val_acc:.1%}, {epochs} epochs", ) logger.info(f"Uploaded {remote}") logger.info(f"Model pushed to https://huggingface.co/{repo_id}") # Trigger hot-reload on the live simulation try: import httpx resp = httpx.post(f"{base_url}/api/nn/reload", timeout=30.0) if resp.status_code == 200: logger.info(f"Live sim NN reloaded: {resp.json().get('message', 'ok')}") else: logger.warning(f"Could not reload live sim NN: HTTP {resp.status_code}") except Exception as e: logger.warning(f"Could not reach live sim for reload: {e}") # ══════════════════════════════════════════════════════════════════════════ # CLI # ══════════════════════════════════════════════════════════════════════════ def main(): parser = argparse.ArgumentParser( description="Soci Agent NN — Local Training Script", formatter_class=argparse.RawDescriptionHelpFormatter, epilog="""Examples: python scripts/nn_train.py # Train from scratch python scripts/nn_train.py --resume --epochs 50 # Continue training python scripts/nn_train.py --data data/nn_training # Use collected samples python scripts/nn_train.py --push # Push existing model to HF Hub python scripts/nn_train.py --graph # Show graphs from last training """, ) parser.add_argument("--epochs", type=int, default=30, help="Training epochs (default: 30)") parser.add_argument("--batch-size", type=int, default=512, help="Batch size (default: 512)") parser.add_argument("--lr", type=float, default=3e-4, help="Learning rate (default: 3e-4)") parser.add_argument("--train-samples", type=int, default=100_000, help="Number of synthetic training samples (default: 100000)") parser.add_argument("--val-samples", type=int, default=10_000, help="Number of validation samples (default: 10000)") parser.add_argument("--data", type=str, default=None, help="Path to directory with collected_samples.jsonl") parser.add_argument("--resume", action="store_true", help="Resume from existing weights in models/") parser.add_argument("--push", action="store_true", help="Push existing model to HuggingFace Hub (no training)") parser.add_argument("--graph", action="store_true", help="Display training graphs from last training run") parser.add_argument("--repo", default="RayMelius/soci-agent-nn", help="HF Hub repo ID (default: RayMelius/soci-agent-nn)") parser.add_argument("--url", default="https://raymelius-soci2.hf.space", help="Live simulation URL for hot-reload after push (default: HF Space)") args = parser.parse_args() # --graph: just display graphs and exit if args.graph: plot_training_graphs() return # --push: just push existing model to HF Hub and exit if args.push: stats_path = MODEL_DIR / "training_stats.json" best_pt = MODEL_DIR / "soci_agent_best.pt" onnx_path = MODEL_DIR / "soci_agent.onnx" if stats_path.exists(): stats = json.loads(stats_path.read_text()) best_val_acc = stats.get("best_val_action_acc", 0) ep = stats.get("epochs", 0) n_train = stats.get("train_samples", 0) else: best_val_acc, ep, n_train = 0, 0, 0 _push_to_hub(best_pt, onnx_path, stats_path, args.repo, best_val_acc, ep, n_train, base_url=args.url) return # Default: train train( epochs=args.epochs, batch_size=args.batch_size, lr=args.lr, num_train=args.train_samples, num_val=args.val_samples, data_dir=args.data, resume=args.resume, ) if __name__ == "__main__": main()