"""Baseline CEO policies for SimMart — random, heuristic, oracle.
Used for:
• Smoke-testing the environment end-to-end before GRPO training
• Establishing a reward floor (random), ceiling (oracle), and a sensible
non-learned target (heuristic) that the trained Qwen CEO must beat
• Generating the "before vs after training" contrast trace for the demo
Usage:
python inference.py --policy heuristic --seed 42 --episodes 3 --verbose
python inference.py --policy all --episodes 10 --quiet
python inference.py --trace hero_episode.json --policy oracle --seed 7
"""
from __future__ import annotations
import argparse
import json
import os
import random
import statistics
import sys
from dataclasses import dataclass
from typing import Any, Dict, List, Optional
# Make repo importable no matter where we run from
HERE = os.path.dirname(os.path.abspath(__file__))
if HERE not in sys.path:
sys.path.insert(0, HERE)
from models import (
CrisisEvent,
Proposal,
ProposalDecision,
SimMartAction,
SimMartObservation,
)
from server.environment import SimMartEnvironment
# ---------------------------------------------------------------------------
# Policy base
# ---------------------------------------------------------------------------
class CEOPolicy:
"""Interface: given an obs, return a SimMartAction for this week.
Policies may inspect `env.state` (oracle) or not (random/heuristic).
"""
name: str = "base"
def act(
self,
obs: SimMartObservation,
env: Optional[SimMartEnvironment] = None,
week: int = 0,
) -> SimMartAction:
raise NotImplementedError
# ---------------------------------------------------------------------------
# Random CEO
# ---------------------------------------------------------------------------
class RandomCEO(CEOPolicy):
name = "random"
def __init__(self, seed: int = 0):
self._rng = random.Random(seed)
def act(self, obs, env=None, week=0):
verdicts = ["approve", "reject", "flag_suspicious"]
weights = [0.65, 0.25, 0.10]
decisions = [
ProposalDecision(
proposal_id=p.proposal_id,
verdict=self._rng.choices(verdicts, weights=weights)[0],
)
for p in obs.inbox
]
return SimMartAction(
decisions=decisions,
budget_allocations={"supply_chain": 5e6, "store_ops": 5e5, "finance": 1e6, "growth": 1e6},
journal_entry=f"Random CEO, week {week}. No pattern.",
)
# ---------------------------------------------------------------------------
# All-Approve CEO — baseline that approves everything (the GRPO step-80 hack).
# Used to measure the EV gap between blind-approve and rule-based policies.
# ---------------------------------------------------------------------------
class AllApproveCEO(CEOPolicy):
name = "all_approve"
def act(self, obs, env=None, week=0):
decisions = [
ProposalDecision(proposal_id=p.proposal_id, verdict="approve")
for p in obs.inbox
]
return SimMartAction(
decisions=decisions,
budget_allocations={"supply_chain": 1e7, "store_ops": 2e6, "finance": 1e6, "growth": 2e6},
journal_entry=f"All-approve CEO, week {week}.",
)
# ---------------------------------------------------------------------------
# Heuristic CEO — hand-coded but sensible; reads only the public obs
# ---------------------------------------------------------------------------
class HeuristicCEO(CEOPolicy):
name = "heuristic"
def __init__(self, cash_floor_inr: float = 3e7):
self._cash_floor = cash_floor_inr
@staticmethod
def _crisis_flags(active: List[CrisisEvent]) -> Dict[str, bool]:
f: Dict[str, bool] = {}
for c in active:
f[c.crisis_id] = True
return f
def act(self, obs, env=None, week=0):
decisions: List[ProposalDecision] = []
cash = obs.kpi_snapshot.cash_inr
tight_cash = cash < self._cash_floor
crisis_active = self._crisis_flags(obs.active_crises)
for p in obs.inbox:
verdict = self._decide(p, tight_cash, crisis_active)
decisions.append(ProposalDecision(proposal_id=p.proposal_id, verdict=verdict))
# Budget allocation: lean into supply chain when not cash-tight
if tight_cash:
budget = {"supply_chain": 5e6, "store_ops": 5e5, "finance": 5e5, "growth": 5e5}
else:
budget = {"supply_chain": 1e7, "store_ops": 2e6, "finance": 1e6, "growth": 2e6}
journal = self._write_journal(obs, decisions, tight_cash, week)
return SimMartAction(decisions=decisions, budget_allocations=budget, journal_entry=journal)
# Per-action spend ceilings beyond which the heuristic rejects (4dept-reject
# env: opex-bearing actions hit EBITDA directly, so high-spend low-ROI
# actions are net-negative-EV and must be rejected).
SPEND_CEILINGS_INR: Dict[str, float] = {
"campaign.launch": 15_00_000, # spend > ₹15L → revenue uplift can't justify
"capex.approve": 10_00_000, # capex has no in-env revenue model → reject most
"return.approve": 50_000, # batch refunds > ₹50k usually wasteful
"brand.ambassador": 5_00_000, # NPS uplift +1.5 caps benefit
"loyalty.update": 5_00_000, # NPS uplift +0.8 caps benefit
"hours.extend": 50_000, # tiny revenue mult vs payroll
}
@classmethod
def _decide(
cls,
p: Proposal,
tight_cash: bool,
crisis_active: Dict[str, bool],
) -> str:
"""Rule-based verdict.
Reliable rogue tell from the obs alone:
R2 ghost vendor : vendor_id starts with 'V-SUSPICIOUS-'
Other rogue tells (R1 qty inflation, R3 cost inflation) need
baselines not in the obs — Oracle gets those via env.state.
Negative-EV reject rules (4dept-reject env):
campaign/capex/refund/loyalty/ambassador/hours.extend get
rejected when their spend exceeds a per-action ceiling, since
the opex hit to EBITDA outweighs the modeled benefit.
"""
params = p.params or {}
vendor_id = str(params.get("vendor_id", "") or "")
if vendor_id.upper().startswith("V-SUSPICIOUS"):
return "flag_suspicious"
ceiling = cls.SPEND_CEILINGS_INR.get(p.action)
if ceiling is not None:
spend = float(
params.get("spend_inr")
or params.get("amount_inr")
or params.get("refund_inr")
or params.get("perk_cost_inr")
or params.get("cost_inr")
or (params.get("hours", 0) * 60 * 100 if p.action == "hours.extend" else 0)
or 0.0
)
if spend > ceiling:
return "reject"
if p.urgency == "high":
return "approve"
if p.urgency == "low" and tight_cash:
return "reject"
return "approve"
@staticmethod
def _write_journal(obs, decisions, tight_cash, week):
n_flag = sum(1 for d in decisions if d.verdict == "flag_suspicious")
n_reject = sum(1 for d in decisions if d.verdict == "reject")
n_approve = sum(1 for d in decisions if d.verdict == "approve")
kpi = obs.kpi_snapshot
active = [f"{c.crisis_id}:{c.name}" for c in obs.active_crises]
return (
f"Week {week}: {n_approve} approved, {n_reject} rejected, {n_flag} flagged. "
f"Cash ₹{kpi.cash_inr/1e7:+.2f} Cr (cash-tight={tight_cash}), "
f"NPS {kpi.nps:.0f}, stockout {kpi.stockout_rate_pct:.1f}%, "
f"SLA {kpi.delivery_sla_hit_rate_pct:.0f}%. "
f"Active crises: {', '.join(active) or 'none'}. "
f"Next week: watch cash and supply-chain replenishment."
)
# ---------------------------------------------------------------------------
# Oracle CEO — peeks at env.state.rogue_incidents to flag rogues perfectly
# ---------------------------------------------------------------------------
class OracleCEO(HeuristicCEO):
"""Behaves like the heuristic but with two cheats:
1. Flags every rogue-associated proposal this week (via env.state)
2. Drives up approvals when NOT tight-cash (to avoid false rejects)
"""
name = "oracle"
def act(self, obs, env=None, week=0):
if env is None:
return super().act(obs, env, week)
# Look up this week's rogue proposal ids
rogue_proposal_ids = set()
for r in env.state.rogue_incidents:
if week in r.active_weeks:
rogue_proposal_ids.update(r.associated_proposal_ids)
action = super().act(obs, env, week)
patched: List[ProposalDecision] = []
for d in action.decisions:
if d.proposal_id in rogue_proposal_ids:
patched.append(ProposalDecision(
proposal_id=d.proposal_id,
verdict="flag_suspicious",
flag_reason="Rogue indicator: telltale pattern in proposal metadata.",
))
else:
patched.append(d)
action.decisions = patched
action.journal_entry = action.journal_entry + " (Oracle override: flagged known rogues)."
return action
# ---------------------------------------------------------------------------
# God CEO — empirical ceiling. Reads ground truth, approves-all non-rogues,
# perfect-flags every rogue (zero false flags), writes a journal engineered
# for max journal_coherence_score, and allocates budget to mitigate stockouts.
#
# Not a learned baseline — used only to answer "what is the max achievable
# total_reward on this env?". See grader.py for the 7 reward components.
# ---------------------------------------------------------------------------
class GodCEO(CEOPolicy):
name = "rule_oracle"
def __init__(self, cash_floor_inr: float = 3e7):
self._cash_floor = cash_floor_inr
self._last_journal: str = ""
@staticmethod
def _active_rogue_proposal_ids(env: SimMartEnvironment, week: int) -> set:
ids: set = set()
for r in env.state.rogue_incidents:
if week in r.active_weeks:
ids.update(r.associated_proposal_ids)
return ids
def act(self, obs, env=None, week=0):
if env is None:
return HeuristicCEO().act(obs, env, week)
rogue_ids = self._active_rogue_proposal_ids(env, week)
cash = obs.kpi_snapshot.cash_inr
tight_cash = cash < self._cash_floor
crisis_active = {c.crisis_id: True for c in obs.active_crises}
decisions: List[ProposalDecision] = []
for p in obs.inbox:
if p.proposal_id in rogue_ids:
# Perfect rogue detection → rogue_catch = +1.0, zero false flags
decisions.append(ProposalDecision(
proposal_id=p.proposal_id,
verdict="flag_suspicious",
flag_reason="Rogue indicator: telltale pattern in proposal metadata.",
))
continue
# Non-rogue: reject negative-EV opex spends; approve everything
# else (PO restocks drive the kpi_delta component). Verdict set:
# {approve, reject, flag_suspicious}.
verdict = "approve"
params = p.params or {}
ceiling = HeuristicCEO.SPEND_CEILINGS_INR.get(p.action)
if ceiling is not None:
spend = float(
params.get("spend_inr")
or params.get("amount_inr")
or params.get("refund_inr")
or params.get("perk_cost_inr")
or params.get("cost_inr")
or (params.get("hours", 0) * 60 * 100 if p.action == "hours.extend" else 0)
or 0.0
)
if spend > ceiling:
verdict = "reject"
if verdict == "approve" and p.urgency == "low" and tight_cash:
verdict = "reject"
decisions.append(ProposalDecision(proposal_id=p.proposal_id, verdict=verdict))
# Budget: lean hard into supply_chain (kills stockouts, helps KPI delta).
# Cash-tight → throttle.
if tight_cash:
budget = {"supply_chain": 6e6, "store_ops": 1e6, "finance": 5e5, "growth": 5e5}
else:
budget = {"supply_chain": 1.2e7, "store_ops": 2.5e6, "finance": 1e6, "growth": 2e6}
journal = self._write_max_coherence_journal(
obs, decisions, rogue_ids, tight_cash, week
)
self._last_journal = journal
return SimMartAction(
decisions=decisions,
budget_allocations=budget,
journal_entry=journal,
)
def _write_max_coherence_journal(
self,
obs,
decisions: List[ProposalDecision],
rogue_ids: set,
tight_cash: bool,
week: int,
) -> str:
"""Free-text retrospective journal for the 2-dept mini env.
Mentions both depts, keeps continuity with last week, and references
proposal IDs so a future judge can trace the decision rationale.
"""
approved_or_flagged = [
d.proposal_id for d in decisions
if d.verdict in ("approve", "flag_suspicious")
]
top_pids = approved_or_flagged[:4]
flagged = [d.proposal_id for d in decisions if d.verdict == "flag_suspicious"]
flagged_str = flagged[0] if flagged else ""
kpi = obs.kpi_snapshot
active = [f"{c.crisis_id}:{c.name}" for c in obs.active_crises] or ["none"]
# Continuity anchor: echo a phrase from last week if we have one
anchor = "SimMart operating cadence"
if self._last_journal:
for tok in ("SimMart", "supply", "store", "margin", "stockout"):
if tok.lower() in self._last_journal.lower():
anchor = tok
break
body = (
f"Week {week} CEO decision log for SimMart, continuing last week's focus on {anchor}. "
f"Priority this week was the supply_chain pipeline and store_ops SLA. "
f"Approvals ({len([d for d in decisions if d.verdict=='approve'])}): "
f"cleared {', '.join(top_pids[:3]) or 'the routine POs'} to keep stock cover healthy "
f"and drive next-week revenue. "
f"Flagged ({len(flagged)}): {flagged_str or 'no rogue signal this week'} — "
f"risk pattern matched our rogue playbook (inflated PO qty / ghost vendor / unit-cost kickback). "
f"Rejections: only low-urgency items when cash was tight. "
f"KPI pulse: cash ₹{kpi.cash_inr/1e7:+.2f} Cr (tight={tight_cash}), "
f"NPS {kpi.nps:.0f} (Δ{kpi.nps_delta:+.1f}), "
f"stockout {kpi.stockout_rate_pct:.1f}% (Δ{kpi.stockout_delta_pts:+.2f} pp), "
f"margin {kpi.gross_margin_pct:.2f}% (Δ{kpi.margin_delta_pts:+.2f} pp). "
f"Action items: reinforce supply_chain lead-time discipline and store_ops staffing for the festive week; "
f"keep cash runway above ₹{self._cash_floor/1e7:.0f} Cr. "
f"Crises in play: {', '.join(active)}. "
f"Decision rationale: approve legitimate work, flag_suspicious only where the rogue tell is explicit. "
f"Next week: monitor stockout trajectory, audit flagged proposal "
f"{flagged_str or (top_pids[-1] if top_pids else 'OPS')}, "
f"tighten journal continuity with {anchor}."
)
return body
# ---------------------------------------------------------------------------
# DualHeadCEO — two-pass inference on a Qwen LoRA setup with separate
# action + journal adapters. Architecture:
#
# base model (Qwen 2.5 1.5B, bf16)
# ├─ adapter "action" — GRPO-trained (or SFT init); 300 tok output
# └─ adapter "journal" — SFT-trained on GodCEO journals; 400 tok output
#
# Each weekly step fires two forward passes: action first, then journal (which
# sees the decisions the action head just produced, so it can reference their
# IDs for the coherence bonus). This breaks the "one bad journal kills all
# reward components" cascade from GRPO v5 and gives clean per-head credit
# assignment in the RL loop.
# ---------------------------------------------------------------------------
class DualHeadCEO(CEOPolicy):
"""Two-pass CEO over a single base model with swappable LoRA adapters.
Construct with a preloaded HF/PEFT model that has two adapters registered
(via ``model.load_adapter(path, adapter_name=name)``) and this class
handles the per-step swap + generation.
Example::
model, tok = FastLanguageModel.from_pretrained(...)
model.load_adapter(action_adapter_path, adapter_name="action")
model.load_adapter(journal_adapter_path, adapter_name="journal")
FastLanguageModel.for_inference(model)
ceo = DualHeadCEO(model, tok)
"""
name = "dual"
def __init__(
self,
model: Any,
tokenizer: Any,
action_adapter: str = "action",
journal_adapter: str = "journal",
action_max_tokens: int = 300,
journal_max_tokens: int = 400,
do_sample: bool = False,
temperature: float = 1.0,
verbose: bool = False,
):
self.model = model
self.tokenizer = tokenizer
self.action_adapter = action_adapter
self.journal_adapter = journal_adapter
self.action_max_tokens = action_max_tokens
self.journal_max_tokens = journal_max_tokens
self.do_sample = do_sample
self.temperature = temperature
self.verbose = verbose
self.n_action_parse_err = 0
self.n_journal_parse_err = 0
self.t_action_s = 0.0
self.t_journal_s = 0.0
self.n_action_tokens = 0
self.n_journal_tokens = 0
# Lazy import of prompts to keep this module free of torch deps for the
# baseline-only code path.
def _prompts(self):
import prompts as P
return P
def _generate(self, chat, adapter, max_new):
"""Generate one completion; returns (text, n_new_tokens, wall_s)."""
import time
import torch
self.model.set_adapter(adapter)
prompt = self.tokenizer.apply_chat_template(
chat, tokenize=False, add_generation_prompt=True,
)
enc = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
t0 = time.time()
with torch.inference_mode():
out = self.model.generate(
**enc,
max_new_tokens=max_new,
do_sample=self.do_sample,
temperature=self.temperature,
pad_token_id=self.tokenizer.pad_token_id,
)
wall = time.time() - t0
n_new = int(out.shape[1] - enc.input_ids.shape[1])
text = self.tokenizer.decode(
out[0, enc.input_ids.shape[1]:], skip_special_tokens=True,
)
return text, n_new, wall
def act(
self,
obs: SimMartObservation,
env: Optional[SimMartEnvironment] = None,
week: int = 0,
) -> SimMartAction:
P = self._prompts()
# --- Pass 1: action head -> decisions + budget ------------------------
action_chat = P.build_action_chat(obs)
action_text, n_a, t_a = self._generate(
action_chat, self.action_adapter, self.action_max_tokens,
)
self.t_action_s += t_a
self.n_action_tokens += n_a
# Use the robust parser (request_info fallback). We explicitly ignore
# any journal_entry the action head might have leaked into its JSON —
# the journal pass owns that field.
action, tel = P.parse_response(action_text, obs.inbox)
if not tel["parse_ok"] and not tel["parse_partial"]:
self.n_action_parse_err += 1
# --- Pass 2: journal head -> free-form retrospective ------------------
journal_chat = P.build_journal_chat(
obs, action.decisions, action.budget_allocations, action.diligence_requests,
)
journal_text, n_j, t_j = self._generate(
journal_chat, self.journal_adapter, self.journal_max_tokens,
)
self.t_journal_s += t_j
self.n_journal_tokens += n_j
journal = P.parse_journal_response(journal_text)
if not journal:
self.n_journal_parse_err += 1
if self.verbose:
tok_s_a = n_a / t_a if t_a > 0 else 0
tok_s_j = n_j / t_j if t_j > 0 else 0
wk = week if week else 0
print(
f" W{wk:>2} act={n_a}tok/{t_a:.1f}s ({tok_s_a:.0f}t/s) "
f"jrn={n_j}tok/{t_j:.1f}s ({tok_s_j:.0f}t/s) "
f"parse_ok={tel['parse_ok']}",
flush=True,
)
action.journal_entry = journal
return action
# ---------------------------------------------------------------------------
# Frontier CEO — backed by a hosted frontier LLM
# provider="openai" → any OpenAI-compatible Chat Completions endpoint
# (HF router, Groq, Together, vLLM-served, …)
# provider="anthropic" → Anthropic Messages API (direct, Bedrock-via-proxy,
# corporate APIM proxy, etc.)
# provider="openai_responses" → OpenAI Responses API behind an OpenAI-compatible
# SDK (e.g. an enterprise APIM proxy that fronts
# GPT-5 / o-series via the Responses contract)
# ---------------------------------------------------------------------------
def _parse_header_env(raw: str) -> Dict[str, str]:
"""Parse ANTHROPIC_CUSTOM_HEADERS which uses `Name:Value` per line."""
out: Dict[str, str] = {}
for line in (raw or "").replace("\r\n", "\n").split("\n"):
line = line.strip()
if ":" in line:
k, v = line.split(":", 1)
out[k.strip()] = v.strip()
return out
class FrontierCEO(CEOPolicy):
"""CEO policy backed by a hosted frontier LLM.
Uses the *exact same* prompt and JSON parser as the trained Qwen 1.5B,
so the reward it achieves is an honest "what a big model sees when
handed the same observation" ceiling.
Does NOT read env.state — only the public observation. This keeps it
at the same information level as the heuristic / trained policies.
"""
# Match the trainer's per-step generation budget so every policy — trained
# or frontier — gets the same hard cap on response length. See
# scripts/launch_grpo_ddp.sh (MAX_NEW) and the env's truncation behaviour.
DEFAULT_MAX_TOKENS: int = 600
def __init__(
self,
model: Optional[str] = None,
provider: str = "auto", # "auto" | "openai" | "anthropic" | "openai_responses"
api_base: Optional[str] = None,
api_key: Optional[str] = None,
auth_token: Optional[str] = None,
extra_headers: Optional[Dict[str, str]] = None,
max_tokens: Optional[int] = None,
temperature: float = 0.0,
max_retries: int = 3,
request_timeout_s: float = 90.0,
budget_hint: bool = True,
dual_head: bool = False,
action_max_tokens: int = 300,
journal_max_tokens: int = 400,
permissive: bool = False,
):
if max_tokens is None:
max_tokens = self.DEFAULT_MAX_TOKENS
provider = self._resolve_provider(provider, api_base, model)
self._provider = provider
if provider == "anthropic":
self._init_anthropic(
model=model,
api_base=api_base,
api_key=api_key,
auth_token=auth_token,
extra_headers=extra_headers,
request_timeout_s=request_timeout_s,
)
elif provider == "openai":
self._init_openai(
model=model,
api_base=api_base,
api_key=api_key,
request_timeout_s=request_timeout_s,
)
elif provider == "openai_responses":
self._init_openai_responses(
model=model,
api_base=api_base,
api_key=api_key,
extra_headers=extra_headers,
request_timeout_s=request_timeout_s,
)
else:
raise ValueError(f"Unknown provider: {provider!r}")
self._max_tokens = max_tokens
self._temperature = temperature
self._max_retries = max_retries
self._budget_hint = budget_hint
self._dual_head = dual_head
self._action_max_tokens = action_max_tokens
self._journal_max_tokens = journal_max_tokens
self._permissive = permissive
tag = f"frontier:{self._model.split('/')[-1]}"
if dual_head:
tag += "-dual"
if permissive:
tag += "-permissive"
self.name = tag
self.n_parse_errors = 0
self.n_api_errors = 0
self.total_tokens = 0
self.total_prompt_tokens = 0
self.total_completion_tokens = 0
@staticmethod
def _resolve_provider(
provider: str, api_base: Optional[str], model: Optional[str]
) -> str:
if provider != "auto":
return provider
if api_base and "router.huggingface.co" in api_base:
return "openai"
if model and ("claude" in model.lower()):
return "anthropic"
# GPT-ish model name → OpenAI Responses API endpoint
if model and (model.lower().startswith("gpt-") or model.lower().startswith("o")):
return "openai_responses"
if os.environ.get("ANTHROPIC_BASE_URL"):
return "anthropic"
if os.environ.get("OPENAI_RESPONSES_BASE_URL"):
return "openai_responses"
if os.environ.get("HF_token") or os.environ.get("HF_TOKEN"):
return "openai"
raise RuntimeError(
"FrontierCEO provider=auto: cannot infer. "
"Pass provider=openai|anthropic|openai_responses explicitly."
)
def _init_anthropic(
self,
model: Optional[str],
api_base: Optional[str],
api_key: Optional[str],
auth_token: Optional[str],
extra_headers: Optional[Dict[str, str]],
request_timeout_s: float,
) -> None:
try:
import anthropic
except ImportError as e:
raise RuntimeError(
"anthropic package required for provider=anthropic; pip install anthropic"
) from e
base_url = api_base or os.environ.get("ANTHROPIC_BASE_URL")
if not base_url:
raise RuntimeError(
"FrontierCEO(provider=anthropic) needs api_base or $ANTHROPIC_BASE_URL"
)
token = (
auth_token
or api_key
or os.environ.get("ANTHROPIC_AUTH_TOKEN")
or os.environ.get("ANTHROPIC_API_KEY")
)
if not token:
raise RuntimeError(
"FrontierCEO(provider=anthropic) needs auth_token or $ANTHROPIC_AUTH_TOKEN"
)
headers = dict(extra_headers or {})
env_headers = _parse_header_env(os.environ.get("ANTHROPIC_CUSTOM_HEADERS", ""))
for k, v in env_headers.items():
headers.setdefault(k, v)
self._client = anthropic.Anthropic(
base_url=base_url,
auth_token=token,
default_headers=headers or None,
timeout=request_timeout_s,
)
self._model = (
model
or os.environ.get("ANTHROPIC_DEFAULT_SONNET_MODEL")
or os.environ.get("ANTHROPIC_MODEL")
or "Claude-Sonnet-4.6"
)
def _init_openai(
self,
model: Optional[str],
api_base: Optional[str],
api_key: Optional[str],
request_timeout_s: float,
) -> None:
try:
from openai import OpenAI
except ImportError as e:
raise RuntimeError(
"openai package required for provider=openai; pip install openai"
) from e
key = api_key or os.environ.get("HF_token") or os.environ.get("HF_TOKEN")
if not key:
raise RuntimeError(
"FrontierCEO(provider=openai) needs api_key or $HF_token"
)
self._client = OpenAI(
base_url=api_base or "https://router.huggingface.co/v1",
api_key=key,
timeout=request_timeout_s,
)
self._model = model or "moonshotai/Kimi-K2.6"
def _init_openai_responses(
self,
model: Optional[str],
api_base: Optional[str],
api_key: Optional[str],
extra_headers: Optional[Dict[str, str]],
request_timeout_s: float,
) -> None:
"""OpenAI Responses API via an OpenAI-compatible SDK.
Useful for hitting OpenAI's hosted Responses API directly, or for an
enterprise proxy that fronts GPT-5 / o-series models behind the
Responses contract. Custom auth headers (e.g. APIM subscription
keys or other gateway tokens) can be passed via
$ANTHROPIC_CUSTOM_HEADERS or the ``extra_headers`` argument. The
bearer token comes from ``api_key`` or $OPENAI_API_KEY (proxies
that auth via a custom header can pass any non-empty placeholder).
"""
try:
from openai import OpenAI
except ImportError as e:
raise RuntimeError(
"openai package required for provider=openai_responses; "
"pip install openai"
) from e
base_url = api_base or os.environ.get("OPENAI_RESPONSES_BASE_URL")
if not base_url:
raise RuntimeError(
"FrontierCEO(provider=openai_responses) needs api_base or "
"$OPENAI_RESPONSES_BASE_URL"
)
headers = dict(extra_headers or {})
env_headers = _parse_header_env(os.environ.get("ANTHROPIC_CUSTOM_HEADERS", ""))
for k, v in env_headers.items():
headers.setdefault(k, v)
key = (
api_key
or os.environ.get("OPENAI_API_KEY")
or "dummy"
)
self._client = OpenAI(
base_url=base_url,
api_key=key,
default_headers=headers or None,
timeout=request_timeout_s,
)
self._model = model or "gpt-5.4"
def act(self, obs, env=None, week=0):
if self._dual_head:
return self._act_dual(obs, env=env, week=week)
from prompts import build_chat, parse_response
hint = self._max_tokens if self._budget_hint else None
messages = build_chat(obs, token_budget=hint)
completion = self._call(messages, max_tokens=self._max_tokens)
action, tel = parse_response(completion, obs.inbox)
if not tel["parse_ok"]:
self.n_parse_errors += 1
return action
def _act_dual(self, obs, env=None, week=0):
"""Two-pass frontier inference: action JSON call + journal text call.
Mirrors the trained DualHeadCEO wire format so frontier numbers are
directly comparable to SFT / GRPO dual-head adapters.
When ``self._permissive`` is set, the action pass uses the permissive
system prompt that allows … deliberation before
the block and gives concrete rogue + budget hints. Only the
system prompt changes; the user content, the parser, and the journal
pass are identical so numbers stay comparable to the strict variant.
"""
from prompts import (
ACTION_SYSTEM_PROMPT_PERMISSIVE,
build_action_chat, build_journal_chat,
parse_response, parse_journal_response,
render_observation,
)
# --- Pass 1: action ---------------------------------------------------
act_messages = build_action_chat(obs)
if self._permissive:
act_messages = [
{"role": "system", "content": ACTION_SYSTEM_PROMPT_PERMISSIVE},
{"role": "user", "content": render_observation(obs)},
]
act_text = self._call(act_messages, max_tokens=self._action_max_tokens)
action, tel = parse_response(act_text, obs.inbox)
if not tel["parse_ok"] and not tel["parse_partial"]:
self.n_parse_errors += 1
# --- Pass 2: journal --------------------------------------------------
jrn_messages = build_journal_chat(
obs, action.decisions, action.budget_allocations, action.diligence_requests,
)
jrn_text = self._call(jrn_messages, max_tokens=self._journal_max_tokens)
action.journal_entry = parse_journal_response(jrn_text)
return action
def _call(self, messages: List[Dict[str, str]], max_tokens: Optional[int] = None) -> str:
import time
max_tokens = max_tokens if max_tokens is not None else self._max_tokens
last_err: Optional[Exception] = None
for attempt in range(self._max_retries):
try:
if self._provider == "anthropic":
return self._call_anthropic(messages, max_tokens=max_tokens)
if self._provider == "openai_responses":
return self._call_openai_responses(messages, max_tokens=max_tokens)
return self._call_openai(messages, max_tokens=max_tokens)
except Exception as e:
last_err = e
if attempt + 1 < self._max_retries:
backoff = 2 ** attempt
print(
f"[{self.name}] api retry {attempt+1}/{self._max_retries} "
f"after {backoff}s: {type(e).__name__}: {str(e)[:140]}",
file=sys.stderr,
)
time.sleep(backoff)
else:
self.n_api_errors += 1
print(
f"[{self.name}] giving up after {self._max_retries} retries: {last_err}",
file=sys.stderr,
)
return ""
def _call_openai(self, messages: List[Dict[str, str]], max_tokens: Optional[int] = None) -> str:
r = self._client.chat.completions.create(
model=self._model,
messages=messages,
max_tokens=max_tokens or self._max_tokens,
temperature=self._temperature,
)
if r.usage is not None:
self.total_tokens += r.usage.total_tokens or 0
self.total_prompt_tokens += r.usage.prompt_tokens or 0
self.total_completion_tokens += r.usage.completion_tokens or 0
return r.choices[0].message.content or ""
def _call_openai_responses(self, messages: List[Dict[str, str]], max_tokens: Optional[int] = None) -> str:
# OpenAI Responses API. System prompt can travel inside `input` — we
# pass the full messages list unchanged and rely on `output_text` to
# flatten the result back to a string.
kwargs: Dict[str, Any] = {
"model": self._model,
"input": messages,
"max_output_tokens": max_tokens or self._max_tokens,
}
# gpt-5.* reasoning models reject a non-default temperature
if self._temperature is not None and not self._model.lower().startswith("gpt-5"):
kwargs["temperature"] = self._temperature
r = self._client.responses.create(**kwargs)
usage = getattr(r, "usage", None)
if usage is not None:
in_tok = getattr(usage, "input_tokens", 0) or 0
out_tok = getattr(usage, "output_tokens", 0) or 0
self.total_prompt_tokens += in_tok
self.total_completion_tokens += out_tok
self.total_tokens += in_tok + out_tok
return r.output_text or ""
def _call_anthropic(self, messages: List[Dict[str, str]], max_tokens: Optional[int] = None) -> str:
# Anthropic wants system prompt separate from the user/assistant turns.
system_parts = [m["content"] for m in messages if m.get("role") == "system"]
turns = [m for m in messages if m.get("role") != "system"]
system = "\n\n".join(system_parts) if system_parts else None
r = self._client.messages.create(
model=self._model,
max_tokens=max_tokens or self._max_tokens,
temperature=self._temperature,
system=system,
messages=turns,
)
usage = getattr(r, "usage", None)
if usage is not None:
in_tok = getattr(usage, "input_tokens", 0) or 0
out_tok = getattr(usage, "output_tokens", 0) or 0
self.total_prompt_tokens += in_tok
self.total_completion_tokens += out_tok
self.total_tokens += in_tok + out_tok
text_parts: List[str] = []
for block in r.content:
if getattr(block, "type", "") == "text":
text_parts.append(block.text)
return "".join(text_parts)
# ---------------------------------------------------------------------------
# Rollout runner
# ---------------------------------------------------------------------------
@dataclass
class EpisodeResult:
policy: str
seed: int
total_reward: float
weekly_rewards: List[float]
final_cash_inr: float
revenue_qtd_inr: float
ebitda_qtd_inr: float
ebitda_margin_pct: float
min_cash_inr: float
rogues_total: int
rogues_caught: int
avg_stockout_pct: float
avg_nps: float
trace: List[Dict] = None
def run_one_episode(
policy: CEOPolicy,
seed: int,
collect_trace: bool = False,
verbose: bool = False,
) -> EpisodeResult:
env = SimMartEnvironment()
obs = env.reset(seed=seed, episode_id=f"{policy.name}-{seed}")
weekly_rewards: List[float] = []
stockouts: List[float] = []
npss: List[float] = []
trace: List[Dict] = [] if collect_trace else None
min_cash = env.state.company.cash_inr
for w in range(1, env.MAX_WEEKS + 1):
action = policy.act(obs, env=env, week=w)
step_obs = env.step(action)
r = step_obs.reward or 0.0
weekly_rewards.append(r)
stockouts.append(step_obs.kpi_snapshot.stockout_rate_pct)
npss.append(step_obs.kpi_snapshot.nps)
min_cash = min(min_cash, env.state.company.cash_inr)
if collect_trace:
trace.append({
"week": w,
"day": step_obs.day_of_quarter,
"inbox_size": len(obs.inbox),
"decisions": [d.model_dump() for d in action.decisions],
"budget_allocations": action.budget_allocations,
"journal": action.journal_entry,
"reward": r,
"kpi": step_obs.kpi_snapshot.model_dump(),
"active_crises": [c.crisis_id for c in step_obs.active_crises],
"pnl_qtd": step_obs.pnl_snapshot.model_dump(),
})
if verbose:
print(
f" W{w:2d} r={r:+.3f} rev=₹{step_obs.kpi_snapshot.revenue_inr/1e7:.2f}Cr "
f"margin={step_obs.kpi_snapshot.gross_margin_pct:5.2f}% NPS={step_obs.kpi_snapshot.nps:4.1f} "
f"stockout={step_obs.kpi_snapshot.stockout_rate_pct:5.1f}% "
f"cash=₹{env.state.company.cash_inr/1e7:+.2f}Cr"
)
obs = step_obs
if obs.done:
break
return EpisodeResult(
policy=policy.name,
seed=seed,
total_reward=sum(weekly_rewards),
weekly_rewards=weekly_rewards,
final_cash_inr=env.state.company.cash_inr,
revenue_qtd_inr=env.state.company.pnl_qtd.revenue_qtd_inr,
ebitda_qtd_inr=env.state.company.pnl_qtd.ebitda_qtd_inr,
ebitda_margin_pct=env.state.company.pnl_qtd.ebitda_margin_pct,
min_cash_inr=min_cash,
rogues_total=len(env.state.rogue_incidents),
rogues_caught=sum(1 for r in env.state.rogue_incidents if r.caught),
avg_stockout_pct=statistics.mean(stockouts) if stockouts else 0.0,
avg_nps=statistics.mean(npss) if npss else 0.0,
trace=trace,
)
def run_policy(
policy: CEOPolicy,
seeds: List[int],
verbose: bool = False,
quiet: bool = False,
) -> List[EpisodeResult]:
results: List[EpisodeResult] = []
for seed in seeds:
if not quiet:
print(f"\n[{policy.name}] seed={seed} rollout …")
res = run_one_episode(policy, seed, collect_trace=False, verbose=verbose)
results.append(res)
if not quiet:
print(
f" → total reward {res.total_reward:+.3f}, "
f"EBITDA ₹{res.ebitda_qtd_inr/1e7:+.2f}Cr ({res.ebitda_margin_pct:+.1f}%), "
f"rogues {res.rogues_caught}/{res.rogues_total}, "
f"avg stockout {res.avg_stockout_pct:.1f}%"
)
return results
def summarise(results_by_policy: Dict[str, List[EpisodeResult]]) -> None:
"""Print a nicely-formatted comparison table."""
header = (
f"{'policy':<11} {'n':>3} "
f"{'tot_r (mean±sd)':>18} {'EBITDA% (mean)':>14} "
f"{'rogue_recall':>13} {'avg_stockout':>13} {'avg_NPS':>8}"
)
print("\n" + header)
print("-" * len(header))
for name in results_by_policy.keys():
res = results_by_policy.get(name)
if not res:
continue
rewards = [r.total_reward for r in res]
mean_r = statistics.mean(rewards)
sd_r = statistics.stdev(rewards) if len(rewards) > 1 else 0.0
ebitda_pct = statistics.mean([r.ebitda_margin_pct for r in res])
recalls = [
(r.rogues_caught / r.rogues_total) if r.rogues_total else 1.0
for r in res
]
recall = statistics.mean(recalls)
stockout = statistics.mean([r.avg_stockout_pct for r in res])
nps = statistics.mean([r.avg_nps for r in res])
print(
f"{name:<11} {len(res):>3} "
f"{mean_r:+7.3f} ± {sd_r:5.3f} "
f"{ebitda_pct:+13.2f} "
f"{recall:12.1%} "
f"{stockout:12.1f} "
f"{nps:7.1f}"
)
def main() -> int:
parser = argparse.ArgumentParser(description="SimMart baseline CEO rollouts")
parser.add_argument("--policy",
choices=["random", "all_approve", "heuristic", "oracle", "god", "frontier", "all"],
default="all")
parser.add_argument("--frontier-provider",
choices=["auto", "openai", "anthropic", "openai_responses"], default="auto",
help="auto picks openai_responses for gpt-*, anthropic for claude-*, else openai")
parser.add_argument("--frontier-model", default=None,
help="Model id (default: Claude-Sonnet-4.6 for anthropic, "
"moonshotai/Kimi-K2.6 for openai)")
parser.add_argument("--frontier-api-base", default=None,
help="API base URL (default: $ANTHROPIC_BASE_URL for anthropic, "
"https://router.huggingface.co/v1 for openai)")
parser.add_argument("--frontier-temperature", type=float, default=0.0)
parser.add_argument("--frontier-max-tokens", type=int,
default=FrontierCEO.DEFAULT_MAX_TOKENS,
help="hard cap on response tokens; matches trainer MAX_NEW (600)")
parser.add_argument("--frontier-no-budget-hint", action="store_true",
help="do NOT tell the frontier model its token budget")
parser.add_argument("--seed", type=int, default=None,
help="Seed for a single-episode run (ignored if --episodes > 1)")
parser.add_argument("--episodes", type=int, default=3,
help="Number of episodes per policy; seeds = [1, 2, …, N]")
parser.add_argument("--verbose", action="store_true", help="Print per-week telemetry")
parser.add_argument("--quiet", action="store_true", help="Only print summary table")
parser.add_argument("--trace", type=str, default=None,
help="Dump a JSON trace of a single episode for analysis (uses --seed, single policy only)")
args = parser.parse_args()
# --- Trace mode: single episode, JSON dump ---
if args.trace:
if args.policy == "all":
print("--trace requires a single --policy", file=sys.stderr)
return 2
policy_cls = {
"random": RandomCEO, "all_approve": AllApproveCEO,
"heuristic": HeuristicCEO,
"oracle": OracleCEO, "god": GodCEO, "frontier": FrontierCEO,
}[args.policy]
if args.policy == "random":
policy = policy_cls(seed=args.seed or 0)
elif args.policy == "frontier":
policy = FrontierCEO(
provider=args.frontier_provider,
model=args.frontier_model,
api_base=args.frontier_api_base,
temperature=args.frontier_temperature,
max_tokens=args.frontier_max_tokens,
budget_hint=not args.frontier_no_budget_hint,
)
else:
policy = policy_cls()
seed = args.seed if args.seed is not None else 42
print(f"[trace] {args.policy} seed={seed} → {args.trace}")
res = run_one_episode(policy, seed, collect_trace=True, verbose=args.verbose)
payload = {
"meta": {
"policy": res.policy, "seed": res.seed,
"total_reward": res.total_reward,
"final_cash_inr": res.final_cash_inr,
"ebitda_qtd_inr": res.ebitda_qtd_inr,
"ebitda_margin_pct": res.ebitda_margin_pct,
"rogues": {"total": res.rogues_total, "caught": res.rogues_caught},
"avg_stockout_pct": res.avg_stockout_pct,
"avg_nps": res.avg_nps,
},
"trace": res.trace,
}
with open(args.trace, "w") as f:
json.dump(payload, f, indent=2, default=str)
print(f" ✓ wrote {len(res.trace)} weekly steps")
return 0
# --- Sweep mode ---
seeds = [args.seed] if args.seed is not None and args.episodes == 1 \
else list(range(1, args.episodes + 1))
policies: List[CEOPolicy] = []
if args.policy in ("random", "all"):
policies.append(RandomCEO(seed=0))
if args.policy in ("all_approve", "all"):
policies.append(AllApproveCEO())
if args.policy in ("heuristic", "all"):
policies.append(HeuristicCEO())
if args.policy in ("oracle", "all"):
policies.append(OracleCEO())
if args.policy in ("god", "all"):
policies.append(GodCEO())
if args.policy == "frontier":
policies.append(FrontierCEO(
provider=args.frontier_provider,
model=args.frontier_model,
api_base=args.frontier_api_base,
temperature=args.frontier_temperature,
max_tokens=args.frontier_max_tokens,
budget_hint=not args.frontier_no_budget_hint,
))
results_by_policy: Dict[str, List[EpisodeResult]] = {}
for p in policies:
results_by_policy[p.name] = run_policy(p, seeds, verbose=args.verbose, quiet=args.quiet)
summarise(results_by_policy)
return 0
if __name__ == "__main__":
raise SystemExit(main())