deploy v5 (long mode, a100-large)

Dockerfile

@@ -0,0 +1,31 @@
+FROM python:3.11-slim
+
+ENV PYTHONUNBUFFERED=1 \
+    PIP_NO_CACHE_DIR=1 \
+    HF_HOME=/app/.cache \
+    TRANSFORMERS_CACHE=/app/.cache \
+    HF_HUB_ENABLE_HF_TRANSFER=1 \
+    TOKENIZERS_PARALLELISM=false
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        git build-essential ca-certificates curl && \
+    rm -rf /var/lib/apt/lists/*
+
+WORKDIR /app
+
+# Step 1: Install PyTorch with CUDA 12.1 support FIRST (order matters for unsloth).
+# The extra-index-url provides cu121 wheels; the default PyPI would give CPU-only.
+RUN pip install --upgrade pip && \
+    pip install \
+        torch==2.4.1 \
+        torchvision==0.19.1 \
+        torchaudio==2.4.1 \
+        --index-url https://download.pytorch.org/whl/cu121
+
+# Step 2: Install the rest of the dependencies (no torch here — already installed).
+COPY requirements.txt .
+RUN pip install -r requirements.txt
+
+COPY . .
+
+CMD ["python", "-u", "train_hf.py"]

README.md

@@ -1,10 +1,41 @@
 ---
-title: Daedalus Training Space
-emoji: 🌖
+title: DAEDALUS Training
+emoji: 🏛
 colorFrom: indigo
-colorTo: gray
+colorTo: purple
 sdk: docker
 pinned: false
+hardware: a100-large
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# DAEDALUS Training Space (v5)
+
+One-shot Docker Space that trains the DAEDALUS mechanism designer using
+Unsloth + Qwen2.5-0.5B-Instruct + two-phase SFT → GRPO:
+
+1. **Load** `unsloth/Qwen2.5-0.5B-Instruct-bnb-4bit` (4-bit, ~2× faster startup).
+2. **Attach a single LoRA** (rank 16, all attn + MLP, `lora_dropout=0` for Unsloth fast path).
+3. **SFT** — teach the JSON output schema on synthetic `(observation, valid_mechanism)` pairs.
+4. **GRPO** — five reward signals on the SAME LoRA:
+   - `reward_format`    — schema coverage ∈ [−1, 1]
+   - `reward_welfare`   — social welfare ratio W ∈ [0, 1]
+   - `reward_fairness`  — 1 − Gini(surplus) ∈ [0, 1]
+   - `reward_stability` — 1 − 3σ(welfare) ∈ [0, 1]
+   - `reward_composite` — full R = W × F × P × S × anti_collusion
+5. **Merge & push** the full 16-bit model to `Laksh718/daedalus-designer`.
+
+Sentinel line in container logs: `[grpo v5] five-reward single-adapter`
+
+| `TRAIN_MODE` | SFT examples | GRPO steps | A100 time | Use case                      |
+|--------------|--------------|------------|-----------|-------------------------------|
+| `smoke`      | 24           | 4          | ~3 min    | shake out build errors        |
+| `short`      | 320          | 120        | ~10 min   | quick sanity check            |
+| `long`       | 800          | 300        | ~30 min   | solid model (default)         |
+| `full`       | 2000         | 500        | ~75 min   | best quality                  |
+
+Load the trained model:
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+model = AutoModelForCausalLM.from_pretrained("Laksh718/daedalus-designer")
+tok   = AutoTokenizer.from_pretrained("Laksh718/daedalus-designer")
+```

daedalus/__init__.py

@@ -0,0 +1,47 @@
+"""DAEDALUS Environment Package - Mechanism Design via Adversarial RL.
+
+Public API:
+
+OpenEnv-compliant (preferred for training, evaluation, judges):
+    from daedalus import (
+        DaedalusOpenEnv, DaedalusAction, DaedalusObservation, 
+        DaedalusState, DaedalusEnvClient
+    )
+
+Legacy tuple-returning API (for the demo dashboard and old notebooks):
+    from daedalus import DaedalusEnvironment
+"""
+
+from .env import DaedalusEnvironment
+from .models import MechanismConfig, MarketOutcome, AgentState
+
+try:
+    from .openenv_env import DaedalusOpenEnv
+    from .openenv_client import DaedalusEnvClient
+    from .openenv_models import (
+        DaedalusAction,
+        DaedalusObservation,
+        DaedalusState,
+    )
+
+    _OPENENV_AVAILABLE = True
+except Exception:  # noqa: BLE001 - openenv-core not installed
+    _OPENENV_AVAILABLE = False
+
+__version__ = "2.0.0"
+
+__all__ = [
+    "DaedalusEnvironment",
+    "MechanismConfig",
+    "MarketOutcome",
+    "AgentState",
+]
+
+if _OPENENV_AVAILABLE:
+    __all__ += [
+        "DaedalusOpenEnv",
+        "DaedalusEnvClient",
+        "DaedalusAction",
+        "DaedalusObservation",
+        "DaedalusState",
+    ]

daedalus/agents.py

@@ -0,0 +1,266 @@
+"""
+DAEDALUS Sub-Agents — Strategic agent implementations.
+Each agent type exploits a different aspect of the mechanism.
+"""
+import math
+import random
+from typing import List, Optional
+from .models import AgentState, MechanismConfig
+
+
+class SubAgent:
+    """Base class for all sub-agents."""
+
+    def __init__(self, state: AgentState):
+        self.state = state
+
+    def generate_valuation(self) -> float:
+        """Draw a new private valuation from the agent's distribution."""
+        raise NotImplementedError
+
+    def compute_bid(self, mech: MechanismConfig, n_active: int, history: dict) -> float:
+        """Compute bid given the current mechanism and market history."""
+        raise NotImplementedError
+
+    def adapt(self, won: bool, mech: MechanismConfig, history: dict):
+        """Adapt strategy based on outcome."""
+        pass
+
+    def check_participation(self, mech: MechanismConfig) -> bool:
+        """Check if agent still wants to participate."""
+        return self.state.active
+
+
+class TruthfulBidder(SubAgent):
+    """Bids true private valuation. No strategic model."""
+
+    def generate_valuation(self) -> float:
+        self.state.valuation = 0.3 + random.random() * 0.6
+        return self.state.valuation
+
+    def compute_bid(self, mech: MechanismConfig, n_active: int, history: dict) -> float:
+        # Truthful: bid = valuation (dominant strategy under second-price)
+        self.state.bid = self.state.valuation
+        return self.state.bid
+
+    def check_participation(self, mech: MechanismConfig) -> bool:
+        return self.state.valuation >= mech.reserve_price
+
+
+class BidShader(SubAgent):
+    """
+    First-order strategic agent. Shades bid below valuation.
+    In first-price auctions with n symmetric bidders drawing uniformly on [0,1],
+    the Bayes-Nash equilibrium bid is v × (n-1)/n.
+    """
+
+    def generate_valuation(self) -> float:
+        self.state.valuation = 0.4 + random.random() * 0.5
+        return self.state.valuation
+
+    def compute_bid(self, mech: MechanismConfig, n_active: int, history: dict) -> float:
+        v = self.state.valuation
+
+        if mech.auction_type == "first_price":
+            # BNE shading: bid = v × (n-1)/n
+            shade_factor = (n_active - 1) / max(n_active, 2)
+
+            # Adapt based on revealed information
+            if mech.reveal_clearing_price:
+                shade_factor -= 0.03  # Can calibrate more precisely
+            if mech.reveal_competing_bids:
+                shade_factor -= 0.05  # Full info enables strategic shading
+
+            # Learn from price history
+            prices = history.get("clearing_prices", [])
+            if len(prices) > 2:
+                avg_price = sum(prices[-5:]) / min(len(prices), 5)
+                if avg_price < v * 0.7:
+                    shade_factor -= 0.05  # Shade more aggressively
+
+            shade_factor = max(0.5, min(shade_factor, 0.95))
+            self.state.shade_factor = 1 - shade_factor
+            self.state.bid = v * shade_factor
+        else:
+            # Second-price / VCG: truthful is dominant, minor bounded rationality noise
+            self.state.shade_factor = 0.02 + random.random() * 0.03
+            self.state.bid = v * (1 - self.state.shade_factor)
+
+        return self.state.bid
+
+    def adapt(self, won: bool, mech: MechanismConfig, history: dict):
+        if won:
+            self.state.shade_factor = min(self.state.shade_factor + 0.01, 0.35)
+        else:
+            self.state.shade_factor = max(self.state.shade_factor - 0.005, 0.02)
+
+    def check_participation(self, mech: MechanismConfig) -> bool:
+        return self.state.valuation >= mech.reserve_price * 1.1
+
+
+class Colluder(SubAgent):
+    """
+    Coalition agent that coordinates with a partner.
+    One bids high, partner bids very low, they rotate winning.
+    Exploits information transparency for cartel enforcement.
+    """
+
+    def generate_valuation(self) -> float:
+        self.state.valuation = 0.35 + random.random() * 0.5
+        return self.state.valuation
+
+    def compute_bid(self, mech: MechanismConfig, n_active: int, history: dict) -> float:
+        v = self.state.valuation
+        should_win = self.state.collusion_turn
+        reserve = mech.reserve_price
+
+        # Check if collusion is viable
+        can_collude = mech.coalition_policy == "allow" or (
+            mech.coalition_policy == "restrict" and random.random() > 0.5
+        )
+        penalty_risk = 0.4 if mech.collusion_penalty > 1 else 0.0
+
+        if can_collude and random.random() > penalty_risk:
+            if mech.reveal_winner_identity:
+                # Strong collusion when winner is revealed (cartel can enforce)
+                self.state.bid = reserve + 0.01 if should_win else reserve * 0.3
+            else:
+                if should_win:
+                    target = reserve + 0.02 + random.random() * 0.05
+                    self.state.bid = min(v, max(target, v * 0.5))
+                else:
+                    self.state.bid = reserve * 0.5
+        else:
+            # Compete honestly when collusion is too risky
+            factor = 0.85 if mech.auction_type == "first_price" else 0.98
+            self.state.bid = v * factor
+
+        return self.state.bid
+
+    def adapt(self, won: bool, mech: MechanismConfig, history: dict):
+        # If facing high collusion penalties, gradually abandon collusion
+        if mech.collusion_penalty > 1.5:
+            self.state.shade_factor = min(self.state.shade_factor + 0.02, 0.5)
+
+    def check_participation(self, mech: MechanismConfig) -> bool:
+        return self.state.valuation >= mech.reserve_price * 0.8
+
+
+class StrategicDropout(SubAgent):
+    """
+    Agent with participation threshold. Exits if expected surplus
+    falls below outside option (reserve utility).
+    """
+
+    def generate_valuation(self) -> float:
+        self.state.valuation = 0.15 + random.random() * 0.45
+        return self.state.valuation
+
+    def compute_bid(self, mech: MechanismConfig, n_active: int, history: dict) -> float:
+        v = self.state.valuation
+        expected_surplus = v - mech.reserve_price - 0.05
+
+        # Update cumulative surplus tracker
+        if expected_surplus > 0:
+            self.state.cumulative_surplus += expected_surplus * 0.1
+        else:
+            self.state.cumulative_surplus -= 0.02
+
+        # Check dropout condition
+        if (self.state.cumulative_surplus < -self.state.dropout_threshold or
+                mech.reserve_price > v * 0.9):
+            self.state.active = False
+            self.state.bid = 0.0
+            return 0.0
+
+        factor = 0.9 if mech.auction_type == "first_price" else 1.0
+        self.state.bid = v * factor
+        return self.state.bid
+
+    def adapt(self, won: bool, mech: MechanismConfig, history: dict):
+        # Re-entry: if conditions improve, come back
+        if not self.state.active:
+            participation = history.get("participation_rates", [])
+            if len(participation) > 3:
+                recent_avg = sum(participation[-3:]) / 3
+                if recent_avg > 0.6 and mech.reserve_price < 0.3:
+                    self.state.active = True
+                    self.state.cumulative_surplus = 0.0
+
+    def check_participation(self, mech: MechanismConfig) -> bool:
+        return self.state.active
+
+
+class BudgetExploiter(SubAgent):
+    """
+    Budget-constrained agent. Submits strategic bids within budget.
+    Tests payment rule edge cases and attempts to exhaust market liquidity.
+    """
+
+    def generate_valuation(self) -> float:
+        self.state.valuation = 0.2 + random.random() * 0.4
+        return self.state.valuation
+
+    def compute_bid(self, mech: MechanismConfig, n_active: int, history: dict) -> float:
+        if self.state.budget <= 0:
+            self.state.active = False
+            self.state.bid = 0.0
+            return 0.0
+
+        v = self.state.valuation
+        aggressiveness = 0.6 if self.state.budget > 1.5 else 0.3
+        self.state.bid = min(v * aggressiveness, self.state.budget * 0.3)
+        return self.state.bid
+
+    def adapt(self, won: bool, mech: MechanismConfig, history: dict):
+        # Budget gets deducted upon winning in the env step
+        if self.state.budget <= 0:
+            self.state.active = False
+
+    def check_participation(self, mech: MechanismConfig) -> bool:
+        return self.state.active and self.state.budget > 0
+
+
+# ── Factory ──────────────────────────────────────────
+AGENT_CLASSES = {
+    "truthful": TruthfulBidder,
+    "shader": BidShader,
+    "colluder": Colluder,
+    "dropout": StrategicDropout,
+    "exploiter": BudgetExploiter,
+}
+
+
+def create_default_population(stage: int = 0) -> List[SubAgent]:
+    """
+    Create an 8-agent population based on the curriculum stage.
+    Stage 0 is easiest (truthful baseline); Stage 4 is full adversarial.
+    """
+    # Base: 8 truthful agents
+    agents = [
+        TruthfulBidder(AgentState(agent_id=i, agent_type="truthful"))
+        for i in range(8)
+    ]
+
+    # Stage 1: Add 2 Shaders
+    if stage >= 1:
+        agents[2] = BidShader(AgentState(agent_id=2, agent_type="shader"))
+        agents[3] = BidShader(AgentState(agent_id=3, agent_type="shader"))
+
+    # Stage 2: Add 2 strategic Dropouts
+    if stage >= 2:
+        agents[6] = StrategicDropout(AgentState(agent_id=6, agent_type="dropout", dropout_threshold=0.08))
+        # Replacement at index 1
+        agents[1] = StrategicDropout(AgentState(agent_id=1, agent_type="dropout", dropout_threshold=0.12))
+
+    # Stage 3: Add 2 Colluders
+    if stage >= 3:
+        agents[4] = Colluder(AgentState(agent_id=4, agent_type="colluder", partner_id=5, collusion_turn=False))
+        agents[5] = Colluder(AgentState(agent_id=5, agent_type="colluder", partner_id=4, collusion_turn=True))
+
+    # Stage 4: Add 2 Budget Exploiters
+    if stage >= 4:
+        agents[7] = BudgetExploiter(AgentState(agent_id=7, agent_type="exploiter", budget=2.5))
+        agents[0] = BudgetExploiter(AgentState(agent_id=0, agent_type="exploiter", budget=2.0))
+
+    return agents

daedalus/env.py

@@ -0,0 +1,327 @@
+"""
+DAEDALUS Environment — OpenEnv-compliant RL environment.
+Trains an LLM to design market mechanisms against adversarial sub-agents.
+
+OpenEnv API:
+    env = DaedalusEnvironment()
+    obs = env.reset()
+    obs, reward, done, info = env.step(action)
+    state = env.state()
+"""
+import json
+import math
+import random
+from typing import Tuple, Dict, Any, List, Optional
+
+from .models import MechanismConfig, MarketOutcome, Observation, AgentState
+from .agents import (
+    SubAgent, create_default_population,
+    TruthfulBidder, BidShader, Colluder, StrategicDropout, BudgetExploiter,
+)
+from .rewards import (
+    compute_welfare_ratio, compute_fairness, compute_gini,
+    compute_participation_rate, compute_stability,
+    compute_collusion_signal, compute_composite_reward,
+)
+
+
+class DaedalusEnvironment:
+    """
+    DAEDALUS: Mechanism Design via Adversarial RL
+
+    An OpenEnv-compliant environment where an LLM agent designs
+    market mechanisms while strategic sub-agents probe for exploits.
+    Reward is the product of welfare, fairness, participation, and stability.
+
+    Themes:
+        #1 Multi-Agent Interactions (primary)
+        #4 Self-Improvement via curriculum (secondary)
+    """
+
+    def __init__(
+        self,
+        n_agents: int = 8,
+        episode_length: int = 50,
+        rounds_per_step: int = 5,
+        curriculum_stage: int = 0,
+    ):
+        self.n_agents = n_agents
+        self.episode_length = episode_length
+        self.rounds_per_step = rounds_per_step
+        self.curriculum_stage = curriculum_stage
+
+        # Internal state
+        self.agents: List[SubAgent] = []
+        self.mechanism = MechanismConfig()
+        self.round = 0
+        self.history: List[MarketOutcome] = []
+        self.welfare_history: List[float] = []
+        self.participation_history: List[float] = []
+        self.clearing_prices: List[float] = []
+        self.collusion_rotation = 0
+        self.done = False
+
+    def reset(self) -> dict:
+        """
+        Start a fresh episode.
+        Samples new population, resets mechanism to default, clears history.
+
+        Returns:
+            Observation dict for the LLM agent.
+        """
+        self.agents = create_default_population(stage=self.curriculum_stage)
+        self.mechanism = MechanismConfig()
+        self.round = 0
+        self.history = []
+        self.welfare_history = []
+        self.participation_history = []
+        self.clearing_prices = []
+        self.collusion_rotation = 0
+        self.done = False
+
+        # Generate initial valuations
+        for agent in self.agents:
+            agent.generate_valuation()
+
+        # Run warm-up rounds with default mechanism
+        for _ in range(3):
+            self._run_market_round()
+
+        return self._get_observation().to_dict()
+
+    def step(self, action: Any) -> Tuple[dict, float, bool, dict]:
+        """
+        Apply a mechanism configuration action and run market rounds.
+
+        Args:
+            action: Either a MechanismConfig, dict, or JSON string.
+
+        Returns:
+            (observation, reward, done, info)
+        """
+        if self.done:
+            return self._get_observation().to_dict(), 0.0, True, {"error": "Episode finished"}
+
+        # Parse action
+        if isinstance(action, MechanismConfig):
+            self.mechanism = action
+        elif isinstance(action, dict):
+            self.mechanism = MechanismConfig.from_dict(action)
+        elif isinstance(action, str):
+            try:
+                self.mechanism = MechanismConfig.from_json(action)
+            except (json.JSONDecodeError, TypeError):
+                return self._get_observation().to_dict(), 0.0, False, {"error": "Invalid JSON"}
+        else:
+            return self._get_observation().to_dict(), 0.0, False, {"error": "Invalid action type"}
+
+        self.round += 1
+
+        # Run market rounds
+        round_outcomes = []
+        for _ in range(self.rounds_per_step):
+            outcome = self._run_market_round()
+            round_outcomes.append(outcome)
+
+        # Compute reward from post-adaptation rounds (rounds 3-5)
+        eval_outcomes = round_outcomes[2:] if len(round_outcomes) >= 3 else round_outcomes
+        avg_reward = sum(o.composite_reward for o in eval_outcomes) / max(len(eval_outcomes), 1)
+
+        # Check termination
+        self.done = self.round >= self.episode_length
+        if self._check_collapse():
+            self.done = True
+
+        # Build info
+        latest = round_outcomes[-1] if round_outcomes else MarketOutcome()
+        info = {
+            "round": self.round,
+            "mechanism": self.mechanism.to_dict(),
+            "welfare_ratio": latest.welfare_ratio,
+            "gini_coefficient": latest.gini_coefficient,
+            "participation_rate": latest.participation_rate,
+            "stability_score": latest.stability_score,
+            "composite_reward": avg_reward,
+            "collusion_signal": latest.collusion_signal,
+            "active_agents": sum(1 for a in self.agents if a.state.active),
+            "dropout_count": latest.dropout_count,
+        }
+
+        return self._get_observation().to_dict(), avg_reward, self.done, info
+
+    def state(self) -> dict:
+        """Return current observable state (POMDP observation)."""
+        return self._get_observation().to_dict()
+
+    def _run_market_round(self) -> MarketOutcome:
+        """Execute one auction round and return outcome."""
+        mech = self.mechanism
+        active_agents = [a for a in self.agents if a.state.active]
+        n_active = len(active_agents)
+        n_total = len(self.agents)
+
+        # Generate new valuations
+        for agent in active_agents:
+            agent.generate_valuation()
+
+        # Collect bids
+        history = {
+            "clearing_prices": self.clearing_prices,
+            "participation_rates": self.participation_history,
+        }
+        for agent in self.agents:
+            if agent.state.active:
+                agent.compute_bid(mech, n_active, history)
+            else:
+                agent.state.bid = 0.0
+
+        # Filter valid bids (above reserve)
+        valid_bidders = [
+            a for a in active_agents
+            if a.state.bid >= mech.reserve_price
+        ]
+        valid_bidders.sort(key=lambda a: a.state.bid, reverse=True)
+
+        # Allocation & Payment
+        winner = None
+        payment = 0.0
+        winner_valuation = 0.0
+
+        if valid_bidders:
+            winner = valid_bidders[0]
+            winner_valuation = winner.state.valuation
+
+            if mech.auction_type == "first_price":
+                payment = winner.state.bid
+            elif mech.auction_type in ("second_price", "vcg"):
+                payment = valid_bidders[1].state.bid if len(valid_bidders) > 1 else mech.reserve_price
+
+            winner.state.surplus = winner.state.valuation - payment
+            winner.state.wins += 1
+
+            # Budget deduction for exploiter
+            if winner.state.agent_type == "exploiter":
+                winner.state.budget -= payment
+
+            # Collusion penalty
+            if mech.collusion_penalty > 0 and winner.state.agent_type == "colluder":
+                partner = self._get_partner(winner)
+                if partner and partner.state.active and partner.state.bid < mech.reserve_price * 0.8:
+                    penalty_amount = payment * mech.collusion_penalty * 0.3
+                    winner.state.surplus -= penalty_amount
+
+        self.collusion_rotation += 1
+        # Toggle collusion turns
+        for a in self.agents:
+            if a.state.agent_type == "colluder":
+                a.state.collusion_turn = not a.state.collusion_turn
+
+        # Adapt agents
+        for agent in self.agents:
+            won = winner is not None and agent.state.agent_id == winner.state.agent_id
+            agent.adapt(won, mech, history)
+
+        # Compute metrics
+        all_valuations = [a.state.valuation for a in self.agents]
+        surplus_dist = [max(a.state.surplus, 0) for a in active_agents]
+        bids = [a.state.bid for a in active_agents if a.state.bid > 0]
+        wins = [a.state.wins for a in self.agents]
+
+        welfare = compute_welfare_ratio(winner_valuation, all_valuations)
+        fairness = compute_fairness(surplus_dist)
+        participation = compute_participation_rate(n_active, n_total)
+        self.welfare_history.append(welfare)
+        self.participation_history.append(participation)
+        self.clearing_prices.append(payment)
+        stability = compute_stability(self.welfare_history)
+        collusion = compute_collusion_signal(bids, wins)
+
+        # Anti-collusion multiplier
+        collusion_mult = max(0.5, 1.0 - collusion * 0.5) if collusion > 0.5 else 1.0
+
+        composite = compute_composite_reward(
+            welfare, fairness, participation, stability,
+            collusion_penalty=collusion_mult,
+        )
+
+        outcome = MarketOutcome(
+            welfare_ratio=welfare,
+            gini_coefficient=compute_gini(surplus_dist),
+            participation_rate=participation,
+            clearing_price_mean=payment,
+            clearing_price_std=0.0,
+            dropout_count=n_total - n_active,
+            collusion_signal=collusion,
+            stability_score=stability,
+            composite_reward=composite,
+        )
+        self.history.append(outcome)
+        return outcome
+
+    def _get_observation(self) -> Observation:
+        """Build observation for the LLM agent."""
+        active = [a for a in self.agents if a.state.active]
+        bids = [a.state.bid for a in active if a.state.bid > 0]
+
+        # Population proxies (observable aggregate statistics)
+        proxies = {
+            "active_count": len(active),
+            "total_agents": len(self.agents),
+            "bid_mean": sum(bids) / max(len(bids), 1),
+            "bid_std": (sum((b - sum(bids)/max(len(bids),1))**2 for b in bids) / max(len(bids),1)) ** 0.5 if bids else 0,
+            "bid_correlation": compute_collusion_signal(
+                bids, [a.state.wins for a in self.agents]
+            ),
+            "rotation_entropy": self._compute_rotation_entropy(),
+            "dropout_rate": 1 - len(active) / max(len(self.agents), 1),
+            "budget_exhaustion_count": sum(
+                1 for a in self.agents
+                if a.state.agent_type == "exploiter" and a.state.budget <= 0
+            ),
+            "clearing_price_trend": self._compute_trend(self.clearing_prices),
+        }
+
+        return Observation(
+            mechanism_config=self.mechanism.to_dict(),
+            market_outcomes=[o.to_dict() for o in self.history[-20:]],
+            population_proxies=proxies,
+            round_number=self.round,
+            episode_length=self.episode_length,
+            curriculum_stage=self.curriculum_stage,
+        )
+
+    def _get_partner(self, agent: SubAgent) -> Optional[SubAgent]:
+        """Get collusion partner for a colluder agent."""
+        if agent.state.partner_id is not None:
+            for a in self.agents:
+                if a.state.agent_id == agent.state.partner_id:
+                    return a
+        return None
+
+    def _compute_rotation_entropy(self) -> float:
+        """Compute entropy of winner distribution (cartel detection)."""
+        wins = [a.state.wins for a in self.agents]
+        total = sum(wins)
+        if total <= 0:
+            return 1.0
+        entropy = 0.0
+        for w in wins:
+            if w > 0:
+                p = w / total
+                entropy -= p * math.log2(p)
+        max_entropy = math.log2(len(wins)) if len(wins) > 1 else 1
+        return entropy / max_entropy if max_entropy > 0 else 1.0
+
+    def _compute_trend(self, values: List[float], window: int = 5) -> float:
+        """Compute first derivative (trend) of recent values."""
+        if len(values) < 2:
+            return 0.0
+        recent = values[-window:]
+        if len(recent) < 2:
+            return 0.0
+        return (recent[-1] - recent[0]) / len(recent)
+
+    def _check_collapse(self) -> bool:
+        """Check if mechanism has collapsed (zero participation)."""
+        active = sum(1 for a in self.agents if a.state.active)
+        return active == 0

daedalus/models.py

@@ -0,0 +1,163 @@
+"""
+DAEDALUS Data Models — Typed dataclasses for mechanism design environment.
+"""
+from dataclasses import dataclass, field
+from typing import Literal, List, Optional
+import json
+
+
+@dataclass
+class MechanismConfig:
+    """Full mechanism specification — the designer's action."""
+    auction_type: Literal["first_price", "second_price", "vcg"] = "second_price"
+    reveal_reserve: bool = False
+    reveal_competing_bids: bool = False
+    reveal_winner_identity: bool = False
+    reveal_clearing_price: bool = True
+    reveal_bid_distribution: bool = False
+    reserve_price: float = 0.10
+    shill_penalty: float = 0.0
+    withdrawal_penalty: float = 0.0
+    collusion_penalty: float = 0.0
+    coalition_policy: Literal["allow", "restrict", "penalize_suspected", "penalize_confirmed"] = "allow"
+
+    def to_dict(self):
+        return {
+            "auction_type": self.auction_type,
+            "reveal_reserve": self.reveal_reserve,
+            "reveal_competing_bids": self.reveal_competing_bids,
+            "reveal_winner_identity": self.reveal_winner_identity,
+            "reveal_clearing_price": self.reveal_clearing_price,
+            "reveal_bid_distribution": self.reveal_bid_distribution,
+            "reserve_price": self.reserve_price,
+            "shill_penalty": self.shill_penalty,
+            "withdrawal_penalty": self.withdrawal_penalty,
+            "collusion_penalty": self.collusion_penalty,
+            "coalition_policy": self.coalition_policy,
+        }
+
+    @classmethod
+    def from_dict(cls, d: dict) -> "MechanismConfig":
+        return cls(**{k: v for k, v in d.items() if k in cls.__dataclass_fields__})
+
+    @classmethod
+    def from_json(cls, s: str) -> "MechanismConfig":
+        return cls.from_dict(json.loads(s))
+
+    def to_json(self) -> str:
+        return json.dumps(self.to_dict(), indent=2)
+
+
+@dataclass
+class MarketOutcome:
+    """Aggregate market outcome from one step (5 market rounds)."""
+    welfare_ratio: float = 0.0
+    gini_coefficient: float = 0.0
+    participation_rate: float = 1.0
+    clearing_price_mean: float = 0.0
+    clearing_price_std: float = 0.0
+    dropout_count: int = 0
+    collusion_signal: float = 0.0
+    stability_score: float = 1.0
+    composite_reward: float = 0.0
+
+    def to_dict(self):
+        return {
+            "welfare_ratio": self.welfare_ratio,
+            "gini_coefficient": self.gini_coefficient,
+            "participation_rate": self.participation_rate,
+            "clearing_price_mean": self.clearing_price_mean,
+            "clearing_price_std": self.clearing_price_std,
+            "dropout_count": self.dropout_count,
+            "collusion_signal": self.collusion_signal,
+            "stability_score": self.stability_score,
+            "composite_reward": self.composite_reward,
+        }
+
+
+@dataclass
+class AgentState:
+    """Internal state of a sub-agent (hidden from the designer)."""
+    agent_id: int = 0
+    agent_type: Literal["truthful", "shader", "colluder", "dropout", "exploiter"] = "truthful"
+    valuation: float = 0.0
+    bid: float = 0.0
+    surplus: float = 0.0
+    active: bool = True
+    budget: float = float("inf")
+    shade_factor: float = 0.0
+    wins: int = 0
+    # Colluder-specific
+    partner_id: Optional[int] = None
+    collusion_turn: bool = False
+    # Dropout-specific
+    dropout_threshold: float = 0.08
+    cumulative_surplus: float = 0.0
+
+
+@dataclass
+class Observation:
+    """What the designer agent sees at each timestep."""
+    mechanism_config: dict = field(default_factory=dict)
+    market_outcomes: List[dict] = field(default_factory=list)
+    population_proxies: dict = field(default_factory=dict)
+    round_number: int = 0
+    episode_length: int = 50
+    curriculum_stage: int = 0
+
+    def to_dict(self):
+        return {
+            "mechanism_config": self.mechanism_config,
+            "market_outcomes": self.market_outcomes,
+            "population_proxies": self.population_proxies,
+            "round_number": self.round_number,
+            "episode_length": self.episode_length,
+            "curriculum_stage": self.curriculum_stage,
+        }
+
+    def to_prompt(self) -> str:
+        """Convert observation to natural language prompt for LLM."""
+        lines = [
+            "You are a mechanism designer. Analyze the current market state and propose an optimal mechanism configuration.",
+            "",
+            f"Round: {self.round_number} / {self.episode_length}",
+            f"Curriculum Stage: {self.curriculum_stage}",
+            "",
+            "Current Mechanism:",
+            f"  Auction Type: {self.mechanism_config.get('auction_type', 'second_price')}",
+            f"  Reserve Price: {self.mechanism_config.get('reserve_price', 0.1):.3f}",
+            f"  Coalition Policy: {self.mechanism_config.get('coalition_policy', 'allow')}",
+            "",
+            "Recent Market Outcomes:",
+        ]
+
+        for i, outcome in enumerate(self.market_outcomes[-5:]):
+            lines.append(
+                f"  Round {self.round_number - len(self.market_outcomes) + i + 1}: "
+                f"W={outcome.get('welfare_ratio', 0):.3f} "
+                f"F={1 - outcome.get('gini_coefficient', 0):.3f} "
+                f"P={outcome.get('participation_rate', 1):.3f} "
+                f"R={outcome.get('composite_reward', 0):.3f}"
+            )
+
+        proxies = self.population_proxies
+        if proxies:
+            lines.extend([
+                "",
+                "Population Signals:",
+                f"  Active Bidders: {proxies.get('active_count', 8)}",
+                f"  Bid Correlation: {proxies.get('bid_correlation', 0):.3f}",
+                f"  Winner Rotation Entropy: {proxies.get('rotation_entropy', 1):.3f}",
+                f"  Dropout Rate: {proxies.get('dropout_rate', 0):.3f}",
+            ])
+
+        lines.extend([
+            "",
+            "Respond with a JSON mechanism configuration:",
+            '{"auction_type": "first_price|second_price|vcg", "reserve_price": float, '
+            '"reveal_reserve": bool, "reveal_competing_bids": bool, "reveal_winner_identity": bool, '
+            '"reveal_clearing_price": bool, "reveal_bid_distribution": bool, '
+            '"shill_penalty": float, "withdrawal_penalty": float, "collusion_penalty": float, '
+            '"coalition_policy": "allow|restrict|penalize_suspected|penalize_confirmed"}',
+        ])
+        return "\n".join(lines)

daedalus/openenv_client.py

@@ -0,0 +1,25 @@
+"""
+DAEDALUS OpenEnv Client — Typed interface for the DAEDALUS environment.
+Allows easy connection to the hosted HF Space.
+"""
+from typing import Optional
+from openenv.core.env_client import EnvClient
+from .openenv_models import DaedalusAction, DaedalusObservation, DaedalusState
+
+class DaedalusEnvClient(EnvClient[DaedalusAction, DaedalusObservation, DaedalusState]):
+    """
+    Typed client for the DAEDALUS Mechanism Design environment.
+    
+    Usage:
+        client = DaedalusEnvClient(base_url="https://kabilesh-c-daedalus-env.hf.space")
+        obs = client.reset(seed=42).sync()
+        obs = client.step(DaedalusAction(auction_type="second_price", ...)).sync()
+    """
+    def __init__(self, base_url: str, token: Optional[str] = None):
+        super().__init__(
+            base_url=base_url,
+            token=token,
+            action_type=DaedalusAction,
+            observation_type=DaedalusObservation,
+            state_type=DaedalusState
+        )

daedalus/openenv_env.py

@@ -0,0 +1,187 @@
+"""
+DAEDALUS OpenEnv-compliant environment.
+
+This is the canonical class to use. It inherits from
+`openenv.core.env_server.interfaces.Environment`, exposes the standard
+Gymnasium-style API (`reset`, `step`, `state`), and returns typed
+Pydantic observations whose `.reward` and `.done` fields replace the
+legacy `(obs, reward, done, info)` tuple.
+
+The legacy `DaedalusEnvironment` (in `daedalus.env`) is preserved for
+the demo dashboard and earlier notebooks so existing code keeps working.
+This class wraps the legacy env so there is exactly one source of truth
+for the auction / sub-agent / reward logic.
+"""
+
+from __future__ import annotations
+
+import random
+import uuid
+from typing import Any, Optional
+
+from openenv.core.env_server.interfaces import Environment
+
+from .env import DaedalusEnvironment as _LegacyEnv
+from .openenv_models import DaedalusAction, DaedalusObservation, DaedalusState
+from .rubrics import DaedalusCompositeRubric
+
+
+class DaedalusOpenEnv(Environment[DaedalusAction, DaedalusObservation, DaedalusState]):
+    """
+    OpenEnv-compliant DAEDALUS mechanism design environment.
+
+    Theme #1 - Multi-Agent Interactions (primary):
+        An LLM designer agent picks auction rules; a population of
+        adaptive sub-agents (truthful bidders, shaders, colluders,
+        strategic dropouts, budget exploiters) probe the mechanism for
+        exploits. The composite reward
+            R = welfare * fairness * participation * stability
+        forces multi-objective tradeoffs that single-objective gaming
+        cannot satisfy.
+
+    Theme #4 - Self-Improvement (secondary):
+        A 5-stage curriculum gradually mixes harder sub-agent types
+        into the population so the designer self-improves against
+        progressively stronger adversaries.
+
+    Usage::
+
+        from daedalus import DaedalusOpenEnv, DaedalusAction
+
+        env = DaedalusOpenEnv(n_agents=8, episode_length=20)
+        obs = env.reset(seed=42)
+
+        action = DaedalusAction(
+            auction_type="second_price",
+            reserve_price=0.15,
+            collusion_penalty=1.5,
+            coalition_policy="penalize_suspected",
+        )
+        obs = env.step(action)
+        print(obs.reward, obs.done)
+
+        print(env.state.step_count, env.state.n_active_agents)
+    """
+
+    SUPPORTS_CONCURRENT_SESSIONS = False
+
+    def __init__(
+        self,
+        n_agents: int = 8,
+        episode_length: int = 50,
+        rounds_per_step: int = 5,
+        curriculum_stage: int = 0,
+        **kwargs: Any,
+    ):
+        super().__init__(rubric=DaedalusCompositeRubric(), **kwargs)
+        self._env = _LegacyEnv(
+            n_agents=n_agents,
+            episode_length=episode_length,
+            rounds_per_step=rounds_per_step,
+            curriculum_stage=curriculum_stage,
+        )
+        self._episode_id: Optional[str] = None
+        self._last_reward: float = 0.0
+
+    # -------------------------------------------------------------------
+    # OpenEnv abstract API
+    # -------------------------------------------------------------------
+    def reset(
+        self,
+        seed: Optional[int] = None,
+        episode_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> DaedalusObservation:
+        if seed is not None:
+            random.seed(seed)
+        self._episode_id = episode_id or str(uuid.uuid4())
+        self._last_reward = 0.0
+        obs_dict = self._env.reset()
+        return self._make_observation(obs_dict, reward=None, done=False)
+
+    def step(
+        self,
+        action: DaedalusAction,
+        timeout_s: Optional[float] = None,
+        **kwargs: Any,
+    ) -> DaedalusObservation:
+        action_dict = self._action_to_dict(action)
+        obs_dict, reward, done, info = self._env.step(action_dict)
+        self._last_reward = float(reward)
+        observation = self._make_observation(
+            obs_dict,
+            reward=float(reward),
+            done=bool(done),
+            info=info,
+        )
+        return observation
+
+    @property
+    def state(self) -> DaedalusState:
+        active = sum(1 for a in self._env.agents if a.state.active)
+        total = len(self._env.agents)
+        cumulative_welfare = sum(self._env.welfare_history)
+        return DaedalusState(
+            episode_id=self._episode_id,
+            step_count=self._env.round,
+            n_active_agents=active,
+            n_total_agents=total,
+            cumulative_welfare=cumulative_welfare,
+            last_composite_reward=self._last_reward,
+            curriculum_stage=self._env.curriculum_stage,
+        )
+
+    # -------------------------------------------------------------------
+    # Convenience: legacy tuple-style step for demo / notebook code.
+    # NOT part of the OpenEnv contract; kept so existing scripts keep
+    # working without rewriting them.
+    # -------------------------------------------------------------------
+    def step_tuple(self, action: Any):
+        if isinstance(action, DaedalusAction):
+            action_dict = self._action_to_dict(action)
+        else:
+            action_dict = action
+        return self._env.step(action_dict)
+
+    # -------------------------------------------------------------------
+    # Helpers
+    # -------------------------------------------------------------------
+    @staticmethod
+    def _action_to_dict(action: Any) -> dict:
+        if isinstance(action, DaedalusAction):
+            return action.model_dump(exclude={"metadata"})
+        if isinstance(action, dict):
+            return action
+        if hasattr(action, "to_dict"):
+            return action.to_dict()
+        raise TypeError(
+            f"DaedalusOpenEnv.step expected DaedalusAction or dict, "
+            f"got {type(action).__name__}"
+        )
+
+    def _make_observation(
+        self,
+        obs_dict: dict,
+        *,
+        reward: Optional[float],
+        done: bool,
+        info: Optional[dict] = None,
+    ) -> DaedalusObservation:
+        # Copy so we don't mutate the caller's dict; add derived keys rubrics need.
+        metadata = dict(info or {})
+        if "gini_coefficient" in metadata and "fairness" not in metadata:
+            metadata["fairness"] = max(0.0, 1.0 - float(metadata["gini_coefficient"]))
+        return DaedalusObservation(
+            mechanism_config=obs_dict.get("mechanism_config", {}),
+            market_outcomes=obs_dict.get("market_outcomes", []),
+            population_proxies=obs_dict.get("population_proxies", {}),
+            round_number=int(obs_dict.get("round_number", 0)),
+            episode_length=int(obs_dict.get("episode_length", 50)),
+            curriculum_stage=int(obs_dict.get("curriculum_stage", 0)),
+            reward=reward,
+            done=done,
+            metadata=metadata,
+        )
+
+    def close(self) -> None:
+        self._env = None

daedalus/openenv_models.py

@@ -0,0 +1,81 @@
+"""
+DAEDALUS OpenEnv-compliant data models.
+
+These Pydantic models subclass the official OpenEnv types from
+`openenv.core.env_server.types`, so the environment satisfies the
+hackathon's "Use OpenEnv (latest release). Build on top of the
+framework" requirement properly.
+
+`DaedalusObservation.reward` and `.done` are inherited from the OpenEnv
+`Observation` base, so a single object captures everything a Gymnasium-
+style step would normally return as a tuple.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Literal
+
+from pydantic import Field
+
+from openenv.core.env_server.types import Action, Observation, State
+
+
+AuctionType = Literal["first_price", "second_price", "vcg"]
+CoalitionPolicy = Literal[
+    "allow",
+    "restrict",
+    "penalize_suspected",
+    "penalize_confirmed",
+]
+
+
+class DaedalusAction(Action):
+    """The mechanism the LLM designer is proposing for the next 5-round window."""
+
+    auction_type: AuctionType = Field(
+        default="second_price",
+        description="Pricing rule for clearing the auction.",
+    )
+    reserve_price: float = Field(
+        default=0.10, ge=0.0, le=0.9,
+        description="Reserve price below which bids are rejected.",
+    )
+
+    reveal_reserve: bool = Field(default=False)
+    reveal_competing_bids: bool = Field(default=False)
+    reveal_winner_identity: bool = Field(default=False)
+    reveal_clearing_price: bool = Field(default=True)
+    reveal_bid_distribution: bool = Field(default=False)
+
+    shill_penalty: float = Field(default=0.0, ge=0.0, le=3.0)
+    withdrawal_penalty: float = Field(default=0.0, ge=0.0, le=3.0)
+    collusion_penalty: float = Field(default=0.0, ge=0.0, le=3.0)
+
+    coalition_policy: CoalitionPolicy = Field(default="allow")
+
+
+class DaedalusObservation(Observation):
+    """
+    What the LLM designer agent sees at every step.
+
+    `reward` and `done` are inherited from the OpenEnv `Observation`
+    base, so this single object replaces the Gym-style
+    `(obs, reward, done, info)` tuple.
+    """
+
+    mechanism_config: Dict[str, Any] = Field(default_factory=dict)
+    market_outcomes: List[Dict[str, Any]] = Field(default_factory=list)
+    population_proxies: Dict[str, Any] = Field(default_factory=dict)
+    round_number: int = Field(default=0, ge=0)
+    episode_length: int = Field(default=50, ge=1)
+    curriculum_stage: int = Field(default=0, ge=0)
+
+
+class DaedalusState(State):
+    """Server-side state, exposed via Environment.state for introspection."""
+
+    n_active_agents: int = Field(default=0, ge=0)
+    n_total_agents: int = Field(default=0, ge=0)
+    cumulative_welfare: float = Field(default=0.0)
+    last_composite_reward: float = Field(default=0.0)
+    curriculum_stage: int = Field(default=0, ge=0)

daedalus/rewards.py

@@ -0,0 +1,180 @@
+"""
+DAEDALUS Reward Functions — Multiple independent reward signals.
+Multiplicative composite prevents single-objective gaming.
+"""
+import math
+from typing import List, Dict
+
+
+def compute_welfare_ratio(winner_valuation: float, all_valuations: List[float]) -> float:
+    """
+    W(t) — Social Welfare Ratio
+    Total allocated utility / theoretical maximum welfare.
+    Are goods going to agents who value them most?
+
+    Range: [0, 1]
+    """
+    if not all_valuations:
+        return 0.0
+    theoretical_max = max(all_valuations)
+    if theoretical_max <= 0:
+        return 0.0
+    return min(winner_valuation / theoretical_max, 1.0)
+
+
+def compute_gini(values: List[float]) -> float:
+    """Compute Gini coefficient of a distribution."""
+    if not values:
+        return 0.0
+    sorted_vals = sorted(values)
+    n = len(sorted_vals)
+    total = sum(sorted_vals)
+    if total <= 0:
+        return 0.0
+
+    numerator = 0.0
+    for i, v in enumerate(sorted_vals):
+        numerator += (2 * (i + 1) - n - 1) * v
+    return numerator / (n * total)
+
+
+def compute_fairness(payment_distribution: List[float]) -> float:
+    """
+    F(t) — Fairness Score
+    1 - Gini coefficient of payment/surplus distribution.
+    Prevents one agent from capturing all surplus.
+
+    Range: [0, 1]
+    """
+    gini = compute_gini(payment_distribution)
+    return max(0.0, 1.0 - gini)
+
+
+def compute_participation_rate(active_agents: int, total_agents: int) -> float:
+    """
+    P(t) — Participation Rate
+    Fraction of eligible agents who submitted bids.
+    Near zero drives composite reward to zero regardless of other terms.
+
+    Range: [0, 1]
+    """
+    if total_agents <= 0:
+        return 0.0
+    return active_agents / total_agents
+
+
+def compute_stability(welfare_history: List[float], window: int = 5) -> float:
+    """
+    S(t) — Stability Bonus
+    1 - normalized standard deviation of welfare over last N rounds.
+    Rewards consistent performance over high-variance mechanisms.
+
+    Range: [0, 1]
+    """
+    if len(welfare_history) < window:
+        return 1.0  # Give benefit of the doubt early on
+
+    recent = welfare_history[-window:]
+    mean = sum(recent) / len(recent)
+    variance = sum((x - mean) ** 2 for x in recent) / len(recent)
+    std_dev = math.sqrt(variance)
+    return max(0.0, 1.0 - std_dev * 3)  # Normalize
+
+
+def compute_collusion_signal(bids: List[float], wins: List[int]) -> float:
+    """
+    Detect collusion from bid correlation and winner rotation patterns.
+    Returns a signal in [0, 1] where higher = more suspicious.
+    """
+    if len(bids) < 2:
+        return 0.0
+
+    # Check bid variance — cartels suppress bid variance
+    mean_bid = sum(bids) / len(bids)
+    if mean_bid <= 0:
+        return 0.0
+    variance = sum((b - mean_bid) ** 2 for b in bids) / len(bids)
+    # Low variance relative to mean is suspicious
+    cv = math.sqrt(variance) / mean_bid if mean_bid > 0 else 0
+    variance_signal = max(0.0, 1.0 - cv * 3)
+
+    # Check winner rotation — unnaturally even distribution
+    if wins and max(wins) > 0:
+        win_entropy = 0.0
+        total_wins = sum(wins)
+        if total_wins > 0:
+            for w in wins:
+                if w > 0:
+                    p = w / total_wins
+                    win_entropy -= p * math.log2(p)
+            # High entropy = even rotation = suspicious
+            max_entropy = math.log2(len(wins)) if len(wins) > 1 else 1
+            rotation_signal = win_entropy / max_entropy if max_entropy > 0 else 0
+        else:
+            rotation_signal = 0.0
+    else:
+        rotation_signal = 0.0
+
+    return (variance_signal + rotation_signal) / 2
+
+
+def compute_composite_reward(
+    welfare: float,
+    fairness: float,
+    participation: float,
+    stability: float,
+    exploration_bonus: float = 1.0,
+    collusion_penalty: float = 1.0,
+) -> float:
+    """
+    R(t) = W(t) × F(t) × P(t) × S(t) × E(t) × anti_collusion
+
+    Multiplicative structure: all terms must be jointly positive.
+    You cannot sacrifice one objective to maximize another.
+
+    Returns: float in [0, ~1.2] (exploration can push slightly above 1)
+    """
+    reward = welfare * fairness * participation * stability
+    reward *= exploration_bonus
+    reward *= collusion_penalty
+    return max(0.0, reward)
+
+
+# ── Reward Function Registry (for TRL) ──────────────
+def reward_welfare(output: str, env_state: dict, **kwargs) -> float:
+    """Standalone welfare reward for TRL multi-reward setup."""
+    return env_state.get("welfare_ratio", 0.0)
+
+
+def reward_fairness(output: str, env_state: dict, **kwargs) -> float:
+    """Standalone fairness reward for TRL."""
+    return 1.0 - env_state.get("gini_coefficient", 0.0)
+
+
+def reward_participation(output: str, env_state: dict, **kwargs) -> float:
+    """Standalone participation reward for TRL."""
+    return env_state.get("participation_rate", 1.0)
+
+
+def reward_stability(output: str, env_state: dict, **kwargs) -> float:
+    """Standalone stability reward for TRL."""
+    return env_state.get("stability_score", 0.8)
+
+
+def reward_composite(output: str, env_state: dict, **kwargs) -> float:
+    """Composite multiplicative reward for TRL."""
+    w = reward_welfare(output, env_state)
+    f = reward_fairness(output, env_state)
+    p = reward_participation(output, env_state)
+    s = reward_stability(output, env_state)
+    return w * f * p * s
+
+
+# Registry for easy import
+REWARD_FUNCTIONS = {
+    "welfare": reward_welfare,
+    "fairness": reward_fairness,
+    "participation": reward_participation,
+    "stability": reward_stability,
+    "composite": reward_composite,
+}

daedalus/rubrics.py

@@ -0,0 +1,104 @@
+"""
+DAEDALUS Rubrics — OpenEnv-compliant reward rubrics.
+Standardizes the multiplicative composite reward for mechanism design.
+"""
+from typing import Any, Dict, Optional
+from openenv.core.env_server.rubric import Rubric, RubricResult
+
+from .rewards import (
+    compute_welfare_ratio,
+    compute_fairness,
+    compute_participation_rate,
+    compute_stability,
+    compute_collusion_signal,
+)
+
+class WelfareRubric(Rubric):
+    """Measures allocative efficiency."""
+    def __init__(self):
+        super().__init__(name="welfare", weight=1.0)
+
+    def evaluate(self, action: Any, observation: Any) -> RubricResult:
+        # Note: In our environment, welfare is computed inside the step logic
+        # and cached in the observation metadata.
+        val = observation.metadata.get("welfare_ratio", 0.0)
+        return RubricResult(
+            score=float(val),
+            reason=f"Social welfare ratio: {val:.3f}",
+            metadata={"raw_value": val}
+        )
+
+class FairnessRubric(Rubric):
+    """Measures payment equity (1 - Gini)."""
+    def __init__(self):
+        super().__init__(name="fairness", weight=1.0)
+
+    def evaluate(self, action: Any, observation: Any) -> RubricResult:
+        val = observation.metadata.get("fairness", 0.0)
+        return RubricResult(
+            score=float(val),
+            reason=f"Payment fairness score: {val:.3f}",
+            metadata={"raw_value": val}
+        )
+
+class ParticipationRubric(Rubric):
+    """Measures market liquidity."""
+    def __init__(self):
+        super().__init__(name="participation", weight=1.0)
+
+    def evaluate(self, action: Any, observation: Any) -> RubricResult:
+        val = observation.metadata.get("participation_rate", 0.0)
+        return RubricResult(
+            score=float(val),
+            reason=f"Market participation rate: {val:.3f}",
+            metadata={"raw_value": val}
+        )
+
+class StabilityRubric(Rubric):
+    """Measures consistency over time."""
+    def __init__(self):
+        super().__init__(name="stability", weight=1.0)
+
+    def evaluate(self, action: Any, observation: Any) -> RubricResult:
+        val = observation.metadata.get("stability_score", 0.0)
+        return RubricResult(
+            score=float(val),
+            reason=f"Economic stability score: {val:.3f}",
+            metadata={"raw_value": val}
+        )
+
+class DaedalusCompositeRubric(Rubric):
+    """
+    Multiplicative Composite Reward: R = W * F * P * S * CollusionPenalty.
+    This is the core DAEDALUS reward logic wrapped for OpenEnv.
+    """
+    def __init__(self):
+        super().__init__(name="daedalus_composite", weight=1.0)
+        self.welfare = WelfareRubric()
+        self.fairness = FairnessRubric()
+        self.participation = ParticipationRubric()
+        self.stability = StabilityRubric()
+
+    def evaluate(self, action: Any, observation: Any) -> RubricResult:
+        w = self.welfare.evaluate(action, observation).score
+        f = self.fairness.evaluate(action, observation).score
+        p = self.participation.evaluate(action, observation).score
+        s = self.stability.evaluate(action, observation).score
+        
+        # Pull collusion multiplier from info
+        collusion_signal = observation.metadata.get("collusion_signal", 0.0)
+        collusion_mult = max(0.5, 1.0 - collusion_signal * 0.5) if collusion_signal > 0.5 else 1.0
+        
+        composite = w * f * p * s * collusion_mult
+        
+        return RubricResult(
+            score=float(composite),
+            reason=f"Multiplicative Reward: {composite:.4f}",
+            metadata={
+                "welfare": w,
+                "fairness": f,
+                "participation": p,
+                "stability": s,
+                "collusion_penalty": collusion_mult
+            }
+        )

requirements.txt

@@ -0,0 +1,13 @@
+# Unsloth — must come after torch (Dockerfile installs torch first)
+unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git
+huggingface_hub>=0.26.0
+hf_transfer>=0.1.8
+transformers>=4.45.0
+datasets>=3.0.0
+accelerate>=0.34.0
+peft>=0.13.0
+trl>=0.12.0
+bitsandbytes>=0.44.0
+sentencepiece>=0.2.0
+protobuf>=4.25.0
+python-dotenv>=1.0.0

train_hf.py

@@ -0,0 +1,605 @@
+"""
+DAEDALUS Training v5 — Unsloth + Qwen2.5-0.5B + GRPO (single-adapter pipeline)
+================================================================================
+
+Changes vs v4:
+    * reward_fairness now correctly computes 1 − gini_coefficient (never stability).
+    * reward_stability added as a separate fifth reward signal.
+    * GRPO reward_funcs updated to [format, welfare, fairness, stability, composite].
+    * Training scales up automatically when A100-class GPU is detected:
+        - 'long' mode doubles SFT/GRPO counts and uses larger batches.
+        - 'full' mode is the maximum A100 training budget.
+    * HUB_REPO defaults to HUB_MODEL_ID env var (set by deploy script).
+    * format_prompt includes Curriculum Stage line to match server.py inference.
+    * Sentinel: [grpo v5] five-reward single-adapter
+"""
+
+from __future__ import annotations
+
+import functools
+import gc
+import json
+import os
+import random
+import sys
+import time
+import traceback
+from typing import Any, Dict, List, Optional
+
+from dotenv import load_dotenv
+
+load_dotenv()
+
+
+# ---------------------------------------------------------------------------
+# Auth
+# ---------------------------------------------------------------------------
+HF_TOKEN = os.environ.get("HF_TOKEN")
+SPACE_ID = os.environ.get("SPACE_ID")
+if HF_TOKEN:
+    try:
+        from huggingface_hub import login
+        login(token=HF_TOKEN, add_to_git_credential=False)
+        print("[auth] logged into Hugging Face from HF_TOKEN env var")
+    except Exception as e:
+        print(f"[auth] login failed: {e}")
+else:
+    print("[auth] WARNING: HF_TOKEN env var is not set; push_to_hub will fail")
+
+
+# ---------------------------------------------------------------------------
+# Unsloth import (must happen before TRL so unsloth can patch).
+# ---------------------------------------------------------------------------
+from unsloth import FastLanguageModel  # noqa: E402
+
+import torch  # noqa: E402
+
+# Ampere+ (A100/A10G/3090+) supports bfloat16; T4 (Turing) only fp16.
+USE_BF16 = bool(torch.cuda.is_available() and torch.cuda.is_bf16_supported())
+USE_FP16 = bool(torch.cuda.is_available() and not USE_BF16)
+print(f"[precision] bf16={USE_BF16}  fp16={USE_FP16}  cuda={torch.cuda.is_available()}")
+
+if torch.cuda.is_available():
+    gpu_name = torch.cuda.get_device_name(0)
+    gpu_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9
+    print(f"[gpu] {gpu_name}  VRAM={gpu_mem_gb:.0f} GB")
+else:
+    gpu_name = "CPU"
+    gpu_mem_gb = 0.0
+
+# Detect A100-class GPU (>=40 GB VRAM) for automatic scaling.
+IS_HIGH_VRAM = gpu_mem_gb >= 40.0
+
+
+# ---------------------------------------------------------------------------
+# Make daedalus/ importable regardless of cwd.
+# ---------------------------------------------------------------------------
+SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
+for candidate in (SCRIPT_DIR, "/workspace", "/app"):
+    if os.path.isdir(os.path.join(candidate, "daedalus")) and candidate not in sys.path:
+        sys.path.insert(0, candidate)
+
+# Training uses the legacy env directly — no openenv-core dep on the Space.
+from daedalus.env import DaedalusEnvironment  # noqa: E402
+
+
+def _to_dict(obs: Any) -> dict:
+    if obs is None:
+        return {}
+    if isinstance(obs, dict):
+        return obs
+    if hasattr(obs, "model_dump"):
+        return obs.model_dump()
+    if hasattr(obs, "to_dict"):
+        return obs.to_dict()
+    return dict(obs)
+
+
+def make_env(stage: int = 0) -> DaedalusEnvironment:
+    return DaedalusEnvironment(curriculum_stage=stage)
+
+
+# ---------------------------------------------------------------------------
+# Config
+# ---------------------------------------------------------------------------
+TRAIN_MODE = os.environ.get("TRAIN_MODE", "short").lower()
+
+MODEL_ID = os.environ.get("BASE_MODEL", "unsloth/Qwen2.5-0.5B-Instruct-bnb-4bit")
+HUB_REPO = os.environ.get("HUB_MODEL_ID", "Laksh718/daedalus-designer")
+PUSH_MERGED = os.environ.get("PUSH_MERGED", "1") not in ("0", "false", "False")
+
+# ---------------------------------------------------------------------------
+# Training scale: auto-boost when a high-VRAM GPU is detected.
+# ---------------------------------------------------------------------------
+if TRAIN_MODE == "full":
+    # Maximum quality — designed for A100 80 GB. ~60-90 min.
+    N_SFT_EXAMPLES = 2000
+    SFT_EPOCHS = 3
+    N_GRPO_PROMPTS = 1200
+    GRPO_STEPS = 500
+    SFT_BATCH = 16 if IS_HIGH_VRAM else 8
+    GRPO_BATCH = 8  if IS_HIGH_VRAM else 4
+    GRPO_GENERATIONS = 8 if IS_HIGH_VRAM else 4
+elif TRAIN_MODE == "long":
+    # Good quality run — A100 ~25-40 min, T4 ~35-50 min.
+    N_SFT_EXAMPLES = 800 if IS_HIGH_VRAM else 400
+    SFT_EPOCHS = 2
+    N_GRPO_PROMPTS = 500 if IS_HIGH_VRAM else 240
+    GRPO_STEPS = 300 if IS_HIGH_VRAM else 160
+    SFT_BATCH = 12 if IS_HIGH_VRAM else 8
+    GRPO_BATCH = 6  if IS_HIGH_VRAM else 4
+    GRPO_GENERATIONS = 8 if IS_HIGH_VRAM else 4
+elif TRAIN_MODE == "smoke":
+    # CI smoke test — exercises every code path in ~3-5 min.
+    N_SFT_EXAMPLES = 24
+    SFT_EPOCHS = 1
+    N_GRPO_PROMPTS = 16
+    GRPO_STEPS = 4
+    SFT_BATCH = 4
+    GRPO_BATCH = 2
+    GRPO_GENERATIONS = 4
+else:
+    # "short" — quick but useful run (~10-15 min on A100, ~15-25 min on T4).
+    N_SFT_EXAMPLES = 320 if IS_HIGH_VRAM else 160
+    SFT_EPOCHS = 1
+    N_GRPO_PROMPTS = 200 if IS_HIGH_VRAM else 96
+    GRPO_STEPS = 120 if IS_HIGH_VRAM else 60
+    SFT_BATCH = 10 if IS_HIGH_VRAM else 8
+    GRPO_BATCH = 5  if IS_HIGH_VRAM else 4
+    GRPO_GENERATIONS = 8 if IS_HIGH_VRAM else 4
+
+OUT_DIR = "./daedalus-lora"
+MAX_SEQ_LEN = 1024
+
+REQUIRED_KEYS = {
+    "auction_type", "reserve_price", "reveal_reserve",
+    "reveal_competing_bids", "reveal_winner_identity",
+    "reveal_clearing_price", "reveal_bid_distribution",
+    "shill_penalty", "withdrawal_penalty", "collusion_penalty",
+    "coalition_policy",
+}
+
+
+# ---------------------------------------------------------------------------
+# Prompt — MUST stay identical to server.py::_build_prompt
+# ---------------------------------------------------------------------------
+def format_prompt(obs: dict) -> str:
+    lines = [
+        "You are a mechanism designer for a market auction system.",
+        "Analyze the current market state and design an optimal mechanism.",
+        "",
+        f"Round: {obs.get('round_number', 0)} / {obs.get('episode_length', 50)}",
+        f"Curriculum Stage: {obs.get('curriculum_stage', 0)}",
+        "",
+        "Your goal is to maximize the composite reward R = W x F x P x S",
+    ]
+    outcomes = obs.get("market_outcomes", [])
+    if outcomes:
+        lines.append("Recent Market Outcomes:")
+        for o in outcomes[-5:]:
+            lines.append(
+                f"  W={o.get('welfare_ratio', 0):.3f} "
+                f"F={1 - o.get('gini_coefficient', 0):.3f} "
+                f"P={o.get('participation_rate', 1):.3f} "
+                f"S={o.get('stability_score', 1):.3f} "
+                f"R={o.get('composite_reward', 0):.3f}"
+            )
+    proxies = obs.get("population_proxies", {})
+    if proxies:
+        lines.extend([
+            "",
+            "Population Signals:",
+            f"  Active Bidders: {proxies.get('active_count', 8)} / {proxies.get('total_agents', 8)}",
+            f"  Bid Correlation (collusion proxy): {proxies.get('bid_correlation', 0):.3f}",
+            f"  Winner Rotation Entropy: {proxies.get('rotation_entropy', 1):.3f}",
+            f"  Dropout Rate: {proxies.get('dropout_rate', 0):.3f}",
+        ])
+    lines.extend([
+        "",
+        "Respond with ONLY a JSON mechanism configuration with these exact keys:",
+        "  auction_type: \"first_price\" | \"second_price\" | \"vcg\"",
+        "  reserve_price: float [0.0, 0.9]",
+        "  reveal_reserve: bool",
+        "  reveal_competing_bids: bool",
+        "  reveal_winner_identity: bool",
+        "  reveal_clearing_price: bool",
+        "  reveal_bid_distribution: bool",
+        "  shill_penalty: float [0.0, 3.0]",
+        "  withdrawal_penalty: float [0.0, 3.0]",
+        "  collusion_penalty: float [0.0, 3.0]",
+        "  coalition_policy: \"allow\" | \"restrict\" | \"penalize_suspected\" | \"penalize_confirmed\"",
+        "",
+        "Output strictly a single JSON object, no commentary.",
+    ])
+    return "\n".join(lines)
+
+
+def random_valid_mechanism() -> dict:
+    return {
+        "auction_type": random.choice(["first_price", "second_price", "vcg"]),
+        "reserve_price": round(random.uniform(0.05, 0.5), 3),
+        "reveal_reserve": random.choice([True, False]),
+        "reveal_competing_bids": random.random() < 0.3,
+        "reveal_winner_identity": random.choice([True, False]),
+        "reveal_clearing_price": random.random() < 0.7,
+        "reveal_bid_distribution": random.random() < 0.3,
+        "shill_penalty": round(random.uniform(0.0, 2.0), 3),
+        "withdrawal_penalty": round(random.uniform(0.0, 1.0), 3),
+        "collusion_penalty": round(random.uniform(0.0, 2.0), 3),
+        "coalition_policy": random.choice(
+            ["allow", "restrict", "penalize_suspected", "penalize_confirmed"]
+        ),
+    }
+
+
+def generate_sft_examples(n: int) -> List[Dict[str, Any]]:
+    pairs: List[Dict[str, Any]] = []
+    for stage in range(5):
+        n_stage = max(1, n // 5)
+        env = make_env(stage)
+        while len(pairs) < (stage + 1) * n_stage:
+            obs_dict = _to_dict(env.reset())
+            for _ in range(3):
+                mech = random_valid_mechanism()
+                pairs.append({
+                    "messages": [
+                        {"role": "user", "content": format_prompt(obs_dict)},
+                        {"role": "assistant", "content": json.dumps(mech)},
+                    ]
+                })
+                obs_dict, _, done, _ = env.step(mech)
+                obs_dict = _to_dict(obs_dict)
+                if done:
+                    break
+    return pairs[:n]
+
+
+def generate_grpo_prompts(n: int) -> List[Dict[str, str]]:
+    prompts: List[Dict[str, str]] = []
+    for stage in range(5):
+        n_stage = max(1, n // 5)
+        env = make_env(stage)
+        while len(prompts) < (stage + 1) * n_stage:
+            obs_dict = _to_dict(env.reset())
+            prompts.append({"prompt": format_prompt(obs_dict)})
+            for _ in range(3):
+                obs_dict, _, done, _ = env.step(random_valid_mechanism())
+                prompts.append({"prompt": format_prompt(_to_dict(obs_dict))})
+                if done:
+                    break
+    return prompts[:n]
+
+
+# ---------------------------------------------------------------------------
+# Reward functions
+# ---------------------------------------------------------------------------
+_reward_env: Optional[DaedalusEnvironment] = None
+
+
+def _get_reward_env() -> DaedalusEnvironment:
+    global _reward_env
+    if _reward_env is None:
+        _reward_env = make_env(stage=0)
+    return _reward_env
+
+
+def _completion_text(c: Any) -> str:
+    """Normalize a TRL completion to plain string regardless of its format."""
+    if isinstance(c, str):
+        return c
+    if isinstance(c, list) and c:
+        last = c[-1]
+        if isinstance(last, dict):
+            return str(last.get("content", ""))
+    if isinstance(c, dict):
+        return str(c.get("content", c))
+    return str(c)
+
+
+@functools.lru_cache(maxsize=4096)
+def _get_env_outcome(text: str) -> Optional[dict]:
+    """Run one env step for the JSON inside `text`. Cached to avoid re-simulation."""
+    j_start = text.find("{")
+    j_end = text.rfind("}") + 1
+    if j_start < 0 or j_end <= j_start:
+        return None
+    try:
+        mech = json.loads(text[j_start:j_end])
+        if not isinstance(mech, dict):
+            return None
+        env = _get_reward_env()
+        env.reset()
+        obs_dict, reward, done, info = env.step(mech)
+        return {
+            "obs": _to_dict(obs_dict),
+            "reward": float(reward),
+            "done": bool(done),
+            "info": dict(info or {}),
+        }
+    except Exception:
+        return None
+
+
+def reward_format(completions=None, **kwargs) -> List[float]:
+    """Schema coverage: 0.5 + 0.5*coverage ∈ [−1, 1]."""
+    rewards = []
+    for raw in (completions or []):
+        content = _completion_text(raw)
+        j_start = content.find("{")
+        j_end = content.rfind("}") + 1
+        if j_start < 0 or j_end <= j_start:
+            rewards.append(-1.0)
+            continue
+        try:
+            mech = json.loads(content[j_start:j_end])
+            if not isinstance(mech, dict):
+                rewards.append(-0.5)
+                continue
+            coverage = len(set(mech.keys()) & REQUIRED_KEYS) / len(REQUIRED_KEYS)
+            rewards.append(0.5 + 0.5 * coverage)
+        except Exception:
+            rewards.append(-0.5)
+    return rewards
+
+
+def reward_welfare(completions=None, **kwargs) -> List[float]:
+    """Social welfare ratio W ∈ [0, 1]."""
+    out = []
+    for raw in (completions or []):
+        outcome = _get_env_outcome(_completion_text(raw))
+        if outcome:
+            out.append(float(outcome["info"].get("welfare_ratio", 0.0)))
+        else:
+            out.append(0.0)
+    return out
+
+
+def reward_fairness(completions=None, **kwargs) -> List[float]:
+    """Fairness = 1 − Gini(surplus) ∈ [0, 1].
+    Always computed from gini_coefficient — never reads stability_score.
+    """
+    out = []
+    for raw in (completions or []):
+        outcome = _get_env_outcome(_completion_text(raw))
+        if outcome:
+            info = outcome["info"]
+            gini = float(info.get("gini_coefficient", 0.0))
+            out.append(max(0.0, 1.0 - gini))
+        else:
+            out.append(0.0)
+    return out
+
+
+def reward_stability(completions=None, **kwargs) -> List[float]:
+    """Stability = 1 − 3·σ(recent welfare) ∈ [0, 1]."""
+    out = []
+    for raw in (completions or []):
+        outcome = _get_env_outcome(_completion_text(raw))
+        if outcome:
+            out.append(float(outcome["info"].get("stability_score", 1.0)))
+        else:
+            out.append(1.0)  # Neutral default — don't punish parse failures twice
+    return out
+
+
+def reward_composite(completions=None, **kwargs) -> List[float]:
+    """Full composite R = W × F × P × S × anti_collusion ∈ [0, 1]."""
+    out = []
+    for raw in (completions or []):
+        outcome = _get_env_outcome(_completion_text(raw))
+        if outcome:
+            out.append(float(outcome.get("reward", 0.0)))
+        else:
+            out.append(0.0)
+    return out
+
+
+# ---------------------------------------------------------------------------
+# Model + LoRA
+# ---------------------------------------------------------------------------
+def build_model_and_tokenizer():
+    print("[model] loading base via Unsloth (pre-quantized 4-bit)...")
+    model, tokenizer = FastLanguageModel.from_pretrained(
+        model_name=MODEL_ID,
+        max_seq_length=MAX_SEQ_LEN,
+        load_in_4bit=True,
+        dtype=None,  # Unsloth picks bf16/fp16 automatically
+    )
+    print("[model] attaching LoRA (r=16, all attn + MLP)...")
+    model = FastLanguageModel.get_peft_model(
+        model,
+        r=16,
+        target_modules=[
+            "q_proj", "k_proj", "v_proj", "o_proj",
+            "gate_proj", "up_proj", "down_proj",
+        ],
+        lora_alpha=32,
+        lora_dropout=0,      # Must be 0 for Unsloth's fast kernel path
+        bias="none",
+        use_gradient_checkpointing="unsloth",
+        random_state=42,
+    )
+    return model, tokenizer
+
+
+# ---------------------------------------------------------------------------
+# Phase 1: SFT — teach the JSON schema
+# ---------------------------------------------------------------------------
+def run_sft(model, tokenizer):
+    from datasets import Dataset
+    from trl import SFTConfig, SFTTrainer
+
+    print(f"[sft] generating {N_SFT_EXAMPLES} synthetic (prompt, mechanism) pairs ...")
+    dataset = Dataset.from_list(generate_sft_examples(N_SFT_EXAMPLES))
+
+    cfg = SFTConfig(
+        output_dir=OUT_DIR + "-sft",
+        num_train_epochs=SFT_EPOCHS,
+        per_device_train_batch_size=SFT_BATCH,
+        gradient_accumulation_steps=2,
+        learning_rate=2e-4,
+        lr_scheduler_type="cosine",
+        warmup_steps=max(5, N_SFT_EXAMPLES // 20),
+        logging_steps=5,
+        save_strategy="no",
+        bf16=USE_BF16,
+        fp16=USE_FP16,
+        max_seq_length=MAX_SEQ_LEN,
+        report_to="none",
+        seed=42,
+        dataloader_pin_memory=False,
+    )
+
+    trainer = SFTTrainer(
+        model=model,
+        processing_class=tokenizer,
+        train_dataset=dataset,
+        args=cfg,
+    )
+    print(f"[sft] training  (batch={SFT_BATCH}, epochs={SFT_EPOCHS}) ...")
+    trainer.train()
+
+    del trainer
+    gc.collect()
+    torch.cuda.empty_cache()
+    return model
+
+
+# ---------------------------------------------------------------------------
+# Phase 2: GRPO — reinforce with 5 reward signals on the SAME LoRA
+# ---------------------------------------------------------------------------
+def run_grpo(model, tokenizer):
+    from datasets import Dataset
+    from trl import GRPOConfig, GRPOTrainer
+
+    print("[grpo v5] five-reward single-adapter on base + push merged")
+    print(f"[grpo] generating {N_GRPO_PROMPTS} prompts ...")
+    dataset = Dataset.from_list(generate_grpo_prompts(N_GRPO_PROMPTS))
+
+    grpo_warmup = max(1, GRPO_STEPS // 10)
+    cfg = GRPOConfig(
+        output_dir=OUT_DIR,
+        max_steps=GRPO_STEPS,
+        per_device_train_batch_size=GRPO_BATCH,
+        gradient_accumulation_steps=2,
+        num_generations=GRPO_GENERATIONS,
+        max_completion_length=400,   # Full mechanism JSON ~100-150 tokens; 400 is safe
+        learning_rate=5e-6,
+        lr_scheduler_type="cosine",
+        warmup_steps=grpo_warmup,
+        logging_steps=2,
+        save_steps=max(10, GRPO_STEPS // 10),
+        save_total_limit=2,
+        bf16=USE_BF16,
+        fp16=USE_FP16,
+        report_to="none",
+        seed=42,
+        dataloader_pin_memory=False,
+    )
+
+    trainer = GRPOTrainer(
+        model=model,
+        processing_class=tokenizer,
+        reward_funcs=[
+            reward_format,
+            reward_welfare,
+            reward_fairness,
+            reward_stability,
+            reward_composite,
+        ],
+        args=cfg,
+        train_dataset=dataset,
+    )
+    print(
+        f"[grpo] training  "
+        f"(batch={GRPO_BATCH}, generations={GRPO_GENERATIONS}, steps={GRPO_STEPS}) ..."
+    )
+    trainer.train()
+
+    # Save training history for make_plots.py
+    os.makedirs(OUT_DIR, exist_ok=True)
+    history_path = os.path.join(OUT_DIR, "training_history.json")
+    with open(history_path, "w") as f:
+        json.dump({"history": trainer.state.log_history}, f, indent=2)
+    print(f"[grpo] saved log history → {history_path}")
+
+    # Also save a copy in /app root so it's easy to find in Space logs
+    try:
+        import shutil
+        shutil.copy(history_path, "./training_history.json")
+    except Exception:
+        pass
+
+    return trainer.model
+
+
+# ---------------------------------------------------------------------------
+# Push merged 16-bit model to Hub
+# ---------------------------------------------------------------------------
+def push_to_hub(model, tokenizer):
+    if not HF_TOKEN:
+        print("[push] skipped — HF_TOKEN not set")
+        return
+
+    if PUSH_MERGED:
+        print(f"[push] merging LoRA and uploading full 16-bit model → {HUB_REPO} ...")
+        try:
+            model.push_to_hub_merged(
+                HUB_REPO,
+                tokenizer,
+                save_method="merged_16bit",
+                token=HF_TOKEN,
+                private=False,
+            )
+            print(f"[done] merged model live at https://huggingface.co/{HUB_REPO}")
+            return
+        except Exception as e:
+            traceback.print_exc()
+            print(f"[push] merged upload failed ({e}); falling back to LoRA-only push")
+
+    print(f"[push] uploading LoRA adapter → {HUB_REPO} ...")
+    model.push_to_hub(HUB_REPO, token=HF_TOKEN, private=False)
+    tokenizer.push_to_hub(HUB_REPO, token=HF_TOKEN, private=False)
+    print(f"[done] adapter live at https://huggingface.co/{HUB_REPO}")
+
+
+# ---------------------------------------------------------------------------
+# Space auto-pause (stops billing when training completes)
+# ---------------------------------------------------------------------------
+def pause_self() -> None:
+    if not SPACE_ID or not HF_TOKEN:
+        print("[pause] not in a Space (or no token) — skipping")
+        return
+    try:
+        from huggingface_hub import HfApi
+        HfApi(token=HF_TOKEN).pause_space(SPACE_ID)
+        print(f"[pause] Space {SPACE_ID} paused")
+    except Exception as e:
+        print(f"[pause] pause_space failed: {e}")
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+def main():
+    t0 = time.time()
+    print(
+        f"[main] TRAIN_MODE={TRAIN_MODE}  gpu={gpu_name}  "
+        f"SFT={N_SFT_EXAMPLES}×{SFT_EPOCHS}ep  "
+        f"GRPO={GRPO_STEPS} steps × {GRPO_GENERATIONS} gen  "
+        f"→ {HUB_REPO}"
+    )
+    model, tokenizer = build_model_and_tokenizer()
+    model = run_sft(model, tokenizer)
+    model = run_grpo(model, tokenizer)
+    push_to_hub(model, tokenizer)
+    print(f"[done] total wall time: {(time.time() - t0) / 60:.1f} min")
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except Exception:
+        traceback.print_exc()
+    finally:
+        pause_self()