OpenEnv Submission

.gitignore

@@ -0,0 +1,38 @@
+# Python
+__pycache__/
+*.py[cod]
+*.pyo
+*.egg-info/
+dist/
+build/
+
+# Virtual environments
+venv/
+.venv/
+env/
+
+# Secrets
+.env
+.env.*
+
+# OS
+.DS_Store
+Thumbs.db
+
+# IDE
+.vscode/
+.idea/
+
+# Logs
+*.log
+
+# Training runs
+runs/
+
+# Docker
+*.tar
+
+# Hackathon docs
+rule.txt
+"Meta RL Hackathon.docx"
+Meta\ RL\ Hackathon.docx

Dockerfile

@@ -0,0 +1,17 @@
+FROM python:3.11-slim
+
+WORKDIR /app
+
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY . .
+
+ENV FACTORY_TASK=easy
+ENV API_BASE_URL=https://router.huggingface.co/v1
+ENV MODEL_NAME=Qwen/Qwen2.5-72B-Instruct
+ENV PORT=7860
+
+EXPOSE 7860
+
+CMD ["python", "server.py"]

README.md

@@ -1 +1,105 @@
-# OpenEnvRLScheduling
+# Smart Factory Scheduling Environment
+
+An [OpenEnv](https://github.com/openenv/openenv)-compliant RL environment simulating real-world industrial scheduling: assign jobs to machines, handle breakdowns, and maximise throughput within deadlines.
+
+## Observation Space
+
+| Field | Type | Description |
+|-------|------|-------------|
+| `machines` | List[Machine] | id, status (idle/busy/broken), current_job, failure_rate |
+| `pending_jobs` | List[Job] | id, remaining_time, deadline, priority (1-3), assigned_machine |
+| `completed_jobs` | List[Job] | Jobs finished this episode |
+| `time` | int | Current time step |
+| `max_steps` | int | Episode length |
+| `done` | bool | Episode terminated |
+| `reward` | float | Reward from last action |
+
+## Action Space
+
+| Action | Effect |
+|--------|--------|
+| `assign_job <job_id> <machine_id>` | Assign pending job to idle machine |
+| `repair <machine_id>` | Restore broken machine to idle |
+| `wait` | Advance time with no change |
+
+## Reward Function
+
+| Event | Reward |
+|-------|--------|
+| Job completed on time | +1.00 + 0.20 × priority |
+| Job completed late | +0.30 |
+| Valid assignment | +0.10 |
+| Invalid action | −0.10 |
+| Idle machine (pending jobs exist) | −0.05 per machine |
+| Job past deadline | −0.10 per step |
+| Repair broken machine | +0.05 |
+
+## Tasks
+
+| Task | Machines | Jobs | Failure Rate | Max Steps | Baseline Score |
+|------|----------|------|-------------|-----------|----------------|
+| easy | 2 | 3 | 0% | 20 | 1.000 |
+| medium | 4 | 7 | 8% | 30 | ~0.557 |
+| hard | 6 | 12 | 15% | 40 | ~0.457 |
+
+**Score formula:** `0.5 × completion_rate + 0.3 × on_time_rate + 0.2 × utilization_bonus`
+
+## Setup
+
+```bash
+pip install -r requirements.txt
+```
+
+### Run HTTP Server (HF Space)
+```bash
+python server.py
+# Routes: GET /health  POST /reset  POST /step  GET /state  GET /schema
+```
+
+### Run Inference (LLM agent)
+```bash
+export OPENAI_API_KEY=<your-key>
+export FACTORY_TASK=easy   # easy | medium | hard
+python inference.py
+```
+
+### Run RL Training
+```bash
+python train.py --task easy --episodes 10 --provider openai
+python train.py --task medium --episodes 10 --provider claude
+```
+
+### Interactive Demo
+```bash
+python app.py   # opens at http://localhost:7860
+```
+
+### Docker
+```bash
+docker build -t factory-env .
+docker run -e OPENAI_API_KEY=<key> -e FACTORY_TASK=easy -p 7860:7860 factory-env
+```
+
+## Baseline Scores
+
+| Task | Score | Steps |
+|------|-------|-------|
+| easy | 1.000 | 4 |
+| medium | ~0.529 | 12 |
+| hard | ~0.533 | 34 |
+
+## Project Structure
+
+```
+├── factory_env/
+│   ├── env.py       # FactoryEnv (openenv.core.Environment)
+│   ├── models.py    # FactoryAction, FactoryObservation, FactoryState
+│   ├── tasks.py     # Task configurations
+│   └── grader.py    # Score computation
+├── inference.py     # LLM baseline agent
+├── train.py         # Multi-episode RL training loop
+├── server.py        # FastAPI HTTP server for HF Space
+├── app.py           # Gradio interactive demo
+├── openenv.yaml     # OpenEnv metadata
+└── Dockerfile
+```

app.py

@@ -0,0 +1,212 @@
+"""
+Smart Factory Scheduling — Interactive Gradio Demo
+Run: python app.py  →  http://localhost:7860
+"""
+import asyncio, os
+from typing import List, Optional, Tuple
+import gradio as gr
+from factory_env.env import FactoryEnv
+from factory_env.grader import score_episode
+from factory_env.models import FactoryAction as Action
+
+_env: Optional[FactoryEnv] = None
+_obs = None
+_rewards: List[float] = []
+_history: List[dict] = []
+_step_num: int = 0
+
+STATUS_EMOJI = {"idle": "🟢", "busy": "🔵", "broken": "🔴"}
+SYSTEM_PROMPT = "You are a factory scheduler. Reply with ONE action:\n  assign_job <job_id> <machine_id>\n  repair <machine_id>\n  wait"
+
+
+def _llm_client(provider, api_key):
+    if "Claude" in provider:
+        import anthropic
+        return ("claude", anthropic.Anthropic(api_key=api_key or os.getenv("ANTHROPIC_API_KEY")))
+    from openai import OpenAI
+    base = "https://api.openai.com/v1" if "OpenAI" in provider else "https://router.huggingface.co/v1"
+    return ("openai", OpenAI(api_key=api_key or os.getenv("OPENAI_API_KEY") or os.getenv("HF_TOKEN"), base_url=base))
+
+
+def _call_llm(provider_tuple, model, obs, last_reward, step):
+    kind, client = provider_tuple
+    machines = "\n".join(f"  {m.id}: {m.status}" + (f" ({m.current_job})" if m.current_job else "") for m in obs.machines)
+    jobs = "\n".join(f"  {j.id}: t={j.remaining_time} dl={j.deadline} p={j.priority}" for j in obs.pending_jobs) or "  (none)"
+    user = f"Step {step}/{obs.max_steps} | t={obs.time} | reward={last_reward:+.2f}\nMachines:\n{machines}\nJobs:\n{jobs}\nAction:"
+    try:
+        if kind == "claude":
+            r = client.messages.create(model=model, max_tokens=50, system=SYSTEM_PROMPT, messages=[{"role":"user","content":user}])
+            return r.content[0].text.strip().splitlines()[0]
+        else:
+            r = client.chat.completions.create(model=model, temperature=0.2, max_tokens=50,
+                messages=[{"role":"system","content":SYSTEM_PROMPT},{"role":"user","content":user}])
+            return (r.choices[0].message.content or "wait").strip().splitlines()[0]
+    except Exception as e:
+        return f"wait  # {e}"
+
+
+def _parse(text):
+    try:
+        p = text.strip().split()
+        if p[0] == "assign_job" and len(p) == 3: return Action(action_type="assign_job", job_id=p[1], machine_id=p[2])
+        if p[0] == "repair" and len(p) == 2: return Action(action_type="repair", machine_id=p[1])
+    except: pass
+    return Action(action_type="wait")
+
+
+def _heuristic(obs) -> Tuple[Action, str]:
+    for m in obs.machines:
+        if m.status == "broken": return Action(action_type="repair", machine_id=m.id), f"repair {m.id}"
+    for j in sorted(obs.pending_jobs, key=lambda x: (x.deadline, -x.priority)):
+        for m in obs.machines:
+            if m.status == "idle":
+                return Action(action_type="assign_job", job_id=j.id, machine_id=m.id), f"assign_job {j.id} {m.id}"
+    return Action(action_type="wait"), "wait"
+
+
+def _render_state(obs):
+    if obs is None: return "*Reset to start*"
+    lines = [f"### ⏱ Time: {obs.time} / {obs.max_steps}",
+             "\n**Machines**", "| ID | Status | Job |", "|---|---|---|"]
+    for m in obs.machines:
+        lines.append(f"| {m.id} | {STATUS_EMOJI.get(m.status,'')} {m.status} | {m.current_job or '—'} |")
+    lines.append("\n**Pending Jobs**")
+    if obs.pending_jobs:
+        lines += ["| ID | Remaining | Deadline | Priority |", "|---|---|---|---|"]
+        for j in sorted(obs.pending_jobs, key=lambda x: x.deadline):
+            urgent = "🔥" if obs.time + j.remaining_time > j.deadline else ""
+            lines.append(f"| {j.id} {urgent} | {j.remaining_time} | {j.deadline} | {'★'*j.priority} |")
+    else:
+        lines.append("*All jobs completed! ✅*")
+    if obs.completed_jobs:
+        lines.append(f"\n**Completed:** {len(obs.completed_jobs)} ✅")
+    return "\n".join(lines)
+
+
+def _render_log(history):
+    if not history: return "*No steps yet*"
+    rows = ["| Step | Action | Reward | Done |", "|---|---|---|---|"]
+    for h in history[-15:]:
+        r = h["reward"]; icon = "🟢" if r > 0.3 else ("🔴" if r < -0.05 else "🟡")
+        rows.append(f"| {h['step']} | `{h['action']}` | {icon} {r:+.2f} | {'✅' if h['done'] else ''} |")
+    return "\n".join(rows)
+
+
+def _render_score(rewards, env):
+    if not rewards or not env: return ""
+    s = score_episode(env)
+    bar = "█" * int(s * 20) + "░" * (20 - int(s * 20))
+    return f"**Score:** {s:.4f}  `[{bar}]`\n**Completed:** {len(env.completed_jobs)}  |  **Late:** {env.late_jobs}  |  **Total Reward:** {sum(rewards):.2f}"
+
+
+def reset_env(task):
+    global _env, _obs, _rewards, _history, _step_num
+    _env = FactoryEnv(task=task, seed=42); _obs = _env.reset()
+    _rewards = []; _history = []; _step_num = 0
+    return _render_state(_obs), _render_log([]), "", f"✅ Reset — **{task}**: {len(_obs.machines)} machines, {len(_obs.pending_jobs)} jobs"
+
+
+def manual_step(text):
+    global _obs, _rewards, _history, _step_num
+    if _env is None: return _render_state(None), _render_log([]), "", "⚠ Reset first."
+    if _obs.done: return _render_state(_obs), _render_log(_history), _render_score(_rewards, _env), "✅ Episode done."
+    _step_num += 1
+    _obs = _env.step(_parse(text.strip()))
+    r = _obs.reward or 0.0; _rewards.append(r); _history.append({"step": _step_num, "action": text.strip(), "reward": r, "done": _obs.done})
+    status = f"Step {_step_num}: `{text.strip()}` → **{r:+.2f}**"
+    if _obs.done: status += f"\n\n🏁 Done! Score: **{score_episode(_env):.4f}**"
+    return _render_state(_obs), _render_log(_history), _render_score(_rewards, _env), status
+
+
+def heuristic_step():
+    global _obs, _rewards, _history, _step_num
+    if _env is None: return _render_state(None), _render_log([]), "", "⚠ Reset first."
+    if _obs.done: return _render_state(_obs), _render_log(_history), _render_score(_rewards, _env), "✅ Episode done."
+    action, action_text = _heuristic(_obs)
+    _step_num += 1
+    _obs = _env.step(action)
+    r = _obs.reward or 0.0; _rewards.append(r); _history.append({"step": _step_num, "action": f"[H] {action_text}", "reward": r, "done": _obs.done})
+    status = f"[Heuristic] Step {_step_num}: `{action_text}` → **{r:+.2f}**"
+    if _obs.done: status += f"\n\n🏁 Done! Score: **{score_episode(_env):.4f}**"
+    return _render_state(_obs), _render_log(_history), _render_score(_rewards, _env), status
+
+
+def llm_step(provider, api_key, model):
+    global _obs, _rewards, _history, _step_num
+    if _env is None: return _render_state(None), _render_log([]), "", "⚠ Reset first.", ""
+    if _obs.done: return _render_state(_obs), _render_log(_history), _render_score(_rewards, _env), "✅ Episode done.", ""
+    try: client = _llm_client(provider, api_key)
+    except Exception as e: return _render_state(_obs), _render_log(_history), "", f"⚠ {e}", ""
+    action_text = _call_llm(client, model, _obs, _rewards[-1] if _rewards else 0.0, _step_num + 1)
+    action = _parse(action_text)
+    if action.action_type == "wait" and (_obs.pending_jobs or any(m.status == "broken" for m in _obs.machines)):
+        action, action_text = _heuristic(_obs)
+        action_text = f"[fallback] {action_text}"
+    _step_num += 1
+    _obs = _env.step(action)
+    r = _obs.reward or 0.0; _rewards.append(r); _history.append({"step": _step_num, "action": f"[LLM] {action_text}", "reward": r, "done": _obs.done})
+    status = f"[LLM] Step {_step_num}: `{action_text}` → **{r:+.2f}**"
+    if _obs.done: status += f"\n\n🏁 Done! Score: **{score_episode(_env):.4f}**"
+    return _render_state(_obs), _render_log(_history), _render_score(_rewards, _env), status, action_text
+
+
+def run_full_episode(provider, api_key, model, task):
+    global _env, _obs, _rewards, _history, _step_num
+    _env = FactoryEnv(task=task, seed=42); _obs = _env.reset()
+    _rewards = []; _history = []; _step_num = 0
+    try: client = _llm_client(provider, api_key)
+    except Exception as e: return _render_state(_obs), _render_log([]), "", f"⚠ {e}", ""
+    log_lines = []
+    while not _obs.done and _step_num < _obs.max_steps:
+        action_text = _call_llm(client, model, _obs, _rewards[-1] if _rewards else 0.0, _step_num + 1)
+        action = _parse(action_text)
+        if action.action_type == "wait" and (_obs.pending_jobs or any(m.status == "broken" for m in _obs.machines)):
+            action, action_text = _heuristic(_obs); action_text = f"[fb] {action_text}"
+        _step_num += 1; _obs = _env.step(action)
+        r = _obs.reward or 0.0; _rewards.append(r)
+        _history.append({"step": _step_num, "action": action_text, "reward": r, "done": _obs.done})
+        log_lines.append(f"Step {_step_num:2d}: {action_text:<35s} r={r:+.2f}")
+    s = score_episode(_env)
+    status = f"🏁 **Done!** Score: **{s:.4f}** | Completed: {len(_env.completed_jobs)} | Late: {_env.late_jobs}"
+    return _render_state(_obs), _render_log(_history), _render_score(_rewards, _env), status, "\n".join(log_lines)
+
+
+def build_ui():
+    with gr.Blocks(title="Smart Factory RL") as demo:
+        gr.Markdown("# 🏭 Smart Factory Scheduling — Interactive RL Demo")
+        with gr.Row():
+            with gr.Column(scale=1):
+                gr.Markdown("### ⚙️ Setup")
+                task_dd = gr.Dropdown(["easy","medium","hard"], value="easy", label="Task")
+                provider_dd = gr.Dropdown(["OpenAI (GPT)","Claude (Anthropic)","HuggingFace Router"], value="OpenAI (GPT)", label="Provider")
+                api_key_box = gr.Textbox(label="API Key", type="password", placeholder="sk-... or sk-ant-...")
+                model_box = gr.Textbox(label="Model", value="gpt-4o-mini")
+                reset_btn = gr.Button("🔄 Reset", variant="primary")
+                gr.Markdown("### 🎮 Manual")
+                manual_input = gr.Textbox(label="Action", placeholder="assign_job J1 M1  |  repair M2  |  wait")
+                with gr.Row():
+                    manual_btn = gr.Button("▶ Execute")
+                    heuristic_btn = gr.Button("🤖 Heuristic Step")
+                gr.Markdown("### 🧠 LLM")
+                with gr.Row():
+                    llm_step_btn = gr.Button("🔮 LLM Step", variant="secondary")
+                    llm_ep_btn = gr.Button("⚡ Run Full Episode", variant="primary")
+                llm_out = gr.Textbox(label="LLM Output", interactive=False)
+                status_md = gr.Markdown("*Press Reset to start*")
+            with gr.Column(scale=2):
+                gr.Markdown("### 🏭 Factory State")
+                state_md = gr.Markdown("*Reset to start*")
+                gr.Markdown("### 📊 Score")
+                score_md = gr.Markdown("")
+                gr.Markdown("### 📋 Step Log")
+                log_md = gr.Markdown("*No steps yet*")
+        reset_btn.click(reset_env, [task_dd], [state_md, log_md, score_md, status_md])
+        manual_btn.click(manual_step, [manual_input], [state_md, log_md, score_md, status_md])
+        heuristic_btn.click(heuristic_step, [], [state_md, log_md, score_md, status_md])
+        llm_step_btn.click(llm_step, [provider_dd, api_key_box, model_box], [state_md, log_md, score_md, status_md, llm_out])
+        llm_ep_btn.click(run_full_episode, [provider_dd, api_key_box, model_box, task_dd], [state_md, log_md, score_md, status_md, llm_out])
+    return demo
+
+
+if __name__ == "__main__":
+    build_ui().launch(server_name="0.0.0.0", server_port=7860, show_error=True, theme=gr.themes.Soft())

factory_env/__init__.py

@@ -0,0 +1,2 @@
+from factory_env.env import FactoryEnv
+from factory_env.models import FactoryAction, FactoryObservation, FactoryState, Machine, Job

factory_env/env.py

@@ -1,93 +1,134 @@
 import random
-from typing import List
-from factory_env.models import Observation, Action, Machine, Job
-
-class FactoryEnv:
-   def __init__(self, task="easy"):
-       self.task = task
-       self.time = 0
-       self.max_steps = 20
-
-   async def reset(self):
-       random.seed(42)
-
-       self.time = 0
-
-       self.machines = [
-           Machine(id="M1", status="idle"),
-           Machine(id="M2", status="idle"),
-       ]
-
-       self.jobs = [
-           Job(id="J1", remaining_time=3, deadline=10),
-           Job(id="J2", remaining_time=2, deadline=8),
-       ]
-
-       return self._get_result(0.0, False)
-
-
-   async def step(self, action: Action):
-       reward = 0.0
-
-       # Apply action
-       if action.action_type == "assign_job":
-           job = self._find_job(action.job_id)
-           machine = self._find_machine(action.machine_id)
-
-           if job and machine and machine.status == "idle":
-               job.assigned_machine = machine.id
-               machine.status = "busy"
-               machine.current_job = job.id
-               reward += 0.2
-           else:
-               reward -= 0.2  # invalid action
-
-       # Simulate time
-       self.time += 1
-
-       for machine in self.machines:
-           if machine.status == "busy":
-               job = self._find_job(machine.current_job)
-               job.remaining_time -= 1
-
-               if job.remaining_time <= 0:
-                   reward += 1.0
-                   self.jobs.remove(job)
-                   machine.status = "idle"
-                   machine.current_job = None
-
-       # Penalty for idle machines
-       idle_count = sum(1 for m in self.machines if m.status == "idle")
-       reward -= idle_count * 0.05
-
-       done = self.time >= self.max_steps or len(self.jobs) == 0
-
-       return self._get_result(reward, done)
-
-
-   def state(self):
-       return self._get_observation()
-
-   def _get_observation(self):
-       return Observation(
-           machines=self.machines,
-           pending_jobs=self.jobs,
-           time=self.time,
-       )
-
-   def _get_result(self, reward, done):
-       return type("Result", (), {
-           "observation": self._get_observation(),
-           "reward": reward,
-           "done": done
-       })
-
-   def _find_job(self, job_id):
-       return next((j for j in self.jobs if j.id == job_id), None)
-
-   def _find_machine(self, machine_id):
-       return next((m for m in self.machines if m.id == machine_id), None)
-
-   async def close(self):
-       pass
-
+from typing import List, Optional
+
+from openenv.core import Environment
+
+from factory_env.models import FactoryAction, FactoryObservation, FactoryState, Machine, Job
+from factory_env.tasks import TASKS
+
+
+class FactoryEnv(Environment[FactoryAction, FactoryObservation, FactoryState]):
+    """Smart Factory Scheduling Environment — OpenEnv compliant."""
+
+    SUPPORTS_CONCURRENT_SESSIONS = True
+
+    def __init__(self, task: str = "easy", seed: int = 42):
+        super().__init__()
+        if task not in TASKS:
+            raise ValueError(f"Unknown task '{task}'. Choose from: {list(TASKS.keys())}")
+        self.task = task
+        self.seed = seed
+        self.config = TASKS[task]
+        self._rng = random.Random(seed)
+        self.machines: List[Machine] = []
+        self.jobs: List[Job] = []
+        self.completed_jobs: List[Job] = []
+        self.late_jobs: int = 0
+        self.time: int = 0
+        self.max_steps: int = self.config["max_steps"]
+
+    def reset(self, seed: Optional[int] = None, episode_id: Optional[str] = None, **kwargs) -> FactoryObservation:
+        use_seed = seed if seed is not None else self.seed
+        self._rng = random.Random(use_seed)
+        self.time = 0
+        self.completed_jobs = []
+        self.late_jobs = 0
+
+        cfg = self.config
+        self.machines = [
+            Machine(id=f"M{i+1}", status="idle", failure_rate=cfg.get("failure_rate", 0.0))
+            for i in range(cfg["num_machines"])
+        ]
+        self.jobs = []
+        for i in range(cfg["num_jobs"]):
+            proc_time = self._rng.randint(*cfg["job_time_range"])
+            deadline = self.time + proc_time + self._rng.randint(*cfg["deadline_slack"])
+            priority = self._rng.randint(1, cfg.get("max_priority", 1))
+            self.jobs.append(Job(id=f"J{i+1}", remaining_time=proc_time, deadline=deadline, priority=priority))
+
+        return self._make_obs(reward=None, done=False)
+
+    def step(self, action: FactoryAction, timeout_s: Optional[float] = None, **kwargs) -> FactoryObservation:
+        reward = 0.0
+
+        if action.action_type == "assign_job":
+            job = self._find_job(action.job_id)
+            machine = self._find_machine(action.machine_id)
+            if job is None or machine is None or machine.status != "idle":
+                reward -= 0.1
+            else:
+                job.assigned_machine = machine.id
+                machine.status = "busy"
+                machine.current_job = job.id
+                reward += 0.1
+        elif action.action_type == "repair":
+            machine = self._find_machine(action.machine_id)
+            if machine and machine.status == "broken":
+                machine.status = "idle"
+                reward += 0.05
+            else:
+                reward -= 0.05
+
+        self.time += 1
+
+        for machine in self.machines:
+            if machine.status == "busy":
+                job = self._find_job(machine.current_job)
+                if job:
+                    job.remaining_time -= 1
+                    if job.remaining_time <= 0:
+                        on_time = self.time <= job.deadline
+                        reward += (1.0 + 0.2 * job.priority) if on_time else 0.3
+                        if not on_time:
+                            self.late_jobs += 1
+                        self.jobs.remove(job)
+                        self.completed_jobs.append(job)
+                        machine.status = "idle"
+                        machine.current_job = None
+
+            if machine.status == "busy" and machine.failure_rate > 0:
+                if self._rng.random() < machine.failure_rate:
+                    machine.status = "broken"
+                    stalled = self._find_job(machine.current_job)
+                    if stalled:
+                        stalled.assigned_machine = None
+                    machine.current_job = None
+
+        if self.jobs:
+            reward -= sum(1 for m in self.machines if m.status == "idle") * 0.05
+        for job in self.jobs:
+            if self.time > job.deadline:
+                reward -= 0.1
+
+        done = self.time >= self.max_steps or len(self.jobs) == 0
+        return self._make_obs(reward=reward, done=done)
+
+    @property
+    def state(self) -> FactoryState:
+        return FactoryState(
+            machines=list(self.machines),
+            pending_jobs=list(self.jobs),
+            completed_jobs=list(self.completed_jobs),
+            time=self.time,
+            task=self.task,
+            late_jobs=self.late_jobs,
+            step_count=self.time,
+        )
+
+    def _make_obs(self, reward, done: bool) -> FactoryObservation:
+        return FactoryObservation(
+            machines=list(self.machines),
+            pending_jobs=list(self.jobs),
+            completed_jobs=list(self.completed_jobs),
+            time=self.time,
+            max_steps=self.max_steps,
+            task=self.task,
+            reward=reward,
+            done=done,
+        )
+
+    def _find_job(self, job_id: Optional[str]) -> Optional[Job]:
+        return next((j for j in self.jobs if j.id == job_id), None) if job_id else None
+
+    def _find_machine(self, machine_id: Optional[str]) -> Optional[Machine]:
+        return next((m for m in self.machines if m.id == machine_id), None) if machine_id else None

factory_env/grader.py

@@ -1,3 +1,15 @@
-def compute_score(total_reward, max_possible=20):
-   score = total_reward / max_possible
-   return max(0.0, min(1.0, score))
+def compute_score(completed, on_time, total_jobs, late_jobs, task="easy"):
+    if total_jobs == 0:
+        return 0.0
+    completion_rate = completed / total_jobs
+    on_time_rate = on_time / max(completed, 1)
+    utilization_bonus = max(0.0, 1.0 - late_jobs / max(completed, 1))
+    score = 0.5 * completion_rate + 0.3 * on_time_rate + 0.2 * utilization_bonus
+    return round(max(0.0, min(1.0, score)), 4)
+
+
+def score_episode(env) -> float:
+    total = len(env.completed_jobs) + len(env.jobs)
+    completed = len(env.completed_jobs)
+    on_time = sum(1 for j in env.completed_jobs if env.time <= j.deadline)
+    return compute_score(completed, on_time, total, env.late_jobs, env.task)

factory_env/models.py

@@ -1,26 +1,55 @@
-from pydantic import BaseModel
 from typing import List, Optional
+from pydantic import BaseModel, ConfigDict, Field
+from openenv.core import Action as BaseAction, Observation as BaseObservation, State as BaseState
+
 
 class Machine(BaseModel):
-   id: str
-   status: str  # idle, busy, broken
-   current_job: Optional[str] = None
+    model_config = ConfigDict(extra="forbid")
+    id: str
+    status: str                       # idle | busy | broken
+    current_job: Optional[str] = None
+    failure_rate: float = 0.0
+
 
 class Job(BaseModel):
-   id: str
-   remaining_time: int
-   deadline: int
-   assigned_machine: Optional[str] = None
-
-class Observation(BaseModel):
-   machines: List[Machine]
-   pending_jobs: List[Job]
-   time: int
-
-class Action(BaseModel):
-   action_type: str  # assign_job, wait
-   job_id: Optional[str] = None
-   machine_id: Optional[str] = None
-
-class Reward(BaseModel):
-   value: float
+    model_config = ConfigDict(extra="forbid")
+    id: str
+    remaining_time: int
+    deadline: int
+    priority: int = 1
+    assigned_machine: Optional[str] = None
+
+
+class FactoryAction(BaseAction):
+    """
+    action_type: assign_job | repair | wait
+    job_id:      required for assign_job
+    machine_id:  required for assign_job / repair
+    """
+    action_type: str
+    job_id: Optional[str] = None
+    machine_id: Optional[str] = None
+
+
+class FactoryObservation(BaseObservation):
+    """Inherits: done (bool), reward (float|None), metadata (dict)"""
+    machines: List[Machine] = Field(default_factory=list)
+    pending_jobs: List[Job] = Field(default_factory=list)
+    completed_jobs: List[Job] = Field(default_factory=list)
+    time: int = 0
+    max_steps: int = 20
+    task: str = "easy"
+
+
+class FactoryState(BaseState):
+    machines: List[Machine] = Field(default_factory=list)
+    pending_jobs: List[Job] = Field(default_factory=list)
+    completed_jobs: List[Job] = Field(default_factory=list)
+    time: int = 0
+    task: str = "easy"
+    late_jobs: int = 0
+
+
+# Aliases for backward compatibility
+Action = FactoryAction
+Observation = FactoryObservation

factory_env/tasks.py

@@ -1,17 +1,35 @@
 TASKS = {
-   "easy": {
-       "machines": 2,
-       "jobs": 2,
-       "failures": False,
-   },
-   "medium": {
-       "machines": 3,
-       "jobs": 5,
-       "failures": True,
-   },
-   "hard": {
-       "machines": 5,
-       "jobs": 10,
-       "failures": True,
-   },
+    "easy": {
+        "num_machines": 2,
+        "num_jobs": 3,
+        "failures": False,
+        "failure_rate": 0.0,
+        "max_priority": 1,
+        "job_time_range": (2, 5),
+        "deadline_slack": (4, 8),
+        "max_steps": 20,
+        "description": "Assign 3 jobs to 2 machines with no failures.",
+    },
+    "medium": {
+        "num_machines": 4,
+        "num_jobs": 7,
+        "failures": True,
+        "failure_rate": 0.08,
+        "max_priority": 2,
+        "job_time_range": (3, 7),
+        "deadline_slack": (2, 5),
+        "max_steps": 30,
+        "description": "Manage 7 jobs across 4 machines with random breakdowns.",
+    },
+    "hard": {
+        "num_machines": 6,
+        "num_jobs": 12,
+        "failures": True,
+        "failure_rate": 0.15,
+        "max_priority": 3,
+        "job_time_range": (3, 8),
+        "deadline_slack": (1, 4),
+        "max_steps": 40,
+        "description": "Optimize throughput across 12 jobs and 6 machines under frequent failures.",
+    },
 }

inference.py

@@ -1,248 +1,144 @@
 """
-Factory Environment Inference Script
-===================================
-Follows OpenEnv evaluation format strictly.
+Inference Script — Smart Factory Scheduling Environment
+========================================================
+Mandatory env vars (per hackathon spec):
+  OPENAI_API_KEY  API key (also accepts HF_TOKEN for HF router)
+  API_BASE_URL    LLM endpoint  (default: HF router)
+  MODEL_NAME      Model ID      (default: Qwen/Qwen2.5-72B-Instruct)
+  HF_TOKEN        HuggingFace token
+  FACTORY_TASK    easy | medium | hard  (default: easy)
+
+STDOUT FORMAT:
+  [START] task=<name> env=factory_env model=<model>
+  [STEP]  step=<n> action=<str> reward=<0.00> done=<true|false> error=<msg|null>
+  [END]   success=<true|false> steps=<n> score=<0.000> rewards=<r1,r2,...>
 """
 
-import asyncio
 import os
 import textwrap
-from typing import List, Optional
+from typing import List, Optional, Tuple
 
 from openai import OpenAI
 
 from factory_env.env import FactoryEnv
-from factory_env.models import Action
-
-# =========================
-# ENV VARIABLES (MANDATORY)
-# =========================
-API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY")
+from factory_env.models import FactoryAction as Action
+from factory_env.grader import score_episode
+
+API_KEY = (
+    os.getenv("OPENAI_API_KEY")
+    or os.getenv("HF_TOKEN")
+    or os.getenv("API_KEY")
+)
 API_BASE_URL = os.getenv("API_BASE_URL") or "https://router.huggingface.co/v1"
 MODEL_NAME = os.getenv("MODEL_NAME") or "Qwen/Qwen2.5-72B-Instruct"
-
 TASK_NAME = os.getenv("FACTORY_TASK", "easy")
 BENCHMARK = "factory_env"
-
-MAX_STEPS = 20
 TEMPERATURE = 0.2
-MAX_TOKENS = 100
+MAX_TOKENS = 80
 SUCCESS_SCORE_THRESHOLD = 0.5
 
-# =========================
-# PROMPTS
-# =========================
-SYSTEM_PROMPT = textwrap.dedent(
-    """
-    You are controlling a factory scheduling system.
-
-    Your goal:
-    - Assign jobs to machines efficiently
-    - Minimize idle machines
-    - Finish all jobs as fast as possible
-
-    Available actions:
-    1. assign_job <job_id> <machine_id>
-    2. wait
-
-    Rules:
-    - Only assign jobs that exist
-    - Only assign to idle machines
-    - One action per step
-
-    Respond ONLY with the action string.
-    Example:
-    assign_job J1 M1
-    """
-).strip()
-
-
-# =========================
-# LOGGING FUNCTIONS (STRICT FORMAT)
-# =========================
+SYSTEM_PROMPT = textwrap.dedent("""
+    You are controlling a smart factory scheduling system.
+    Goal: complete all jobs before their deadlines, keep machines busy, repair broken machines.
+    Actions (respond with EXACTLY one line):
+      assign_job <job_id> <machine_id>
+      repair <machine_id>
+      wait
+    Respond with ONLY the action string.
+""").strip()
+
+
 def log_start(task: str, env: str, model: str) -> None:
     print(f"[START] task={task} env={env} model={model}", flush=True)
 
 
 def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]) -> None:
-    error_val = error if error else "null"
-    done_val = str(done).lower()
-    print(
-        f"[STEP] step={step} action={action} reward={reward:.2f} done={done_val} error={error_val}",
-        flush=True,
-    )
+    print(f"[STEP] step={step} action={action} reward={reward:.2f} done={str(done).lower()} error={error or 'null'}", flush=True)
 
 
 def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:
-    rewards_str = ",".join(f"{r:.2f}" for r in rewards)
-    print(
-        f"[END] success={str(success).lower()} steps={steps} score={score:.3f} rewards={rewards_str}",
-        flush=True,
-    )
-
-
-# =========================
-# PROMPT BUILDER
-# =========================
-def build_user_prompt(step, obs, last_reward):
-    machines_str = "\n".join(
-        [f"{m.id}: {m.status} (job={m.current_job})" for m in obs.machines]
-    )
-
-    jobs_str = "\n".join(
-        [f"{j.id}: remaining={j.remaining_time}, deadline={j.deadline}" for j in obs.pending_jobs]
-    ) or "None"
-
-    return textwrap.dedent(
-        f"""
-        Step: {step}
+    print(f"[END] success={str(success).lower()} steps={steps} score={score:.3f} rewards={','.join(f'{r:.2f}' for r in rewards)}", flush=True)
 
-        Current Time: {obs.time}
 
-        Machines:
-        {machines_str}
+def build_prompt(step: int, obs, last_reward: float) -> str:
+    machines = "\n".join(f"  {m.id}: {m.status}" + (f" ({m.current_job})" if m.current_job else "") for m in obs.machines)
+    jobs = "\n".join(f"  {j.id}: remaining={j.remaining_time}, deadline={j.deadline}, priority={j.priority}" for j in obs.pending_jobs) or "  (none)"
+    return f"Step {step}/{obs.max_steps} | time={obs.time} | last_reward={last_reward:+.2f}\nMachines:\n{machines}\nPending Jobs:\n{jobs}\nAction:"
 
-        Pending Jobs:
-        {jobs_str}
 
-        Last reward: {last_reward:.2f}
-
-        What action do you take?
-        """
-    ).strip()
-
-
-# =========================
-# LLM CALL
-# =========================
-def get_model_action(client: OpenAI, step, obs, last_reward) -> str:
+def get_model_action(client: OpenAI, step: int, obs, last_reward: float) -> str:
     try:
-        user_prompt = build_user_prompt(step, obs, last_reward)
-
-        completion = client.chat.completions.create(
+        resp = client.chat.completions.create(
             model=MODEL_NAME,
-            messages=[
-                {"role": "system", "content": SYSTEM_PROMPT},
-                {"role": "user", "content": user_prompt},
-            ],
+            messages=[{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": build_prompt(step, obs, last_reward)}],
             temperature=TEMPERATURE,
             max_tokens=MAX_TOKENS,
         )
-
-        text = (completion.choices[0].message.content or "").strip()
-        return text if text else "wait"
-
+        return (resp.choices[0].message.content or "wait").strip().splitlines()[0]
     except Exception as e:
         print(f"[DEBUG] LLM error: {e}", flush=True)
         return "wait"
 
 
-# =========================
-# ACTION PARSER
-# =========================
 def parse_action(text: str) -> Action:
     try:
         parts = text.strip().split()
-
         if parts[0] == "assign_job" and len(parts) == 3:
-            return Action(
-                action_type="assign_job",
-                job_id=parts[1],
-                machine_id=parts[2],
-            )
-
-        elif parts[0] == "wait":
-            return Action(action_type="wait")
-
+            return Action(action_type="assign_job", job_id=parts[1], machine_id=parts[2])
+        if parts[0] == "repair" and len(parts) == 2:
+            return Action(action_type="repair", machine_id=parts[1])
     except Exception:
         pass
-
-    # fallback safe action
     return Action(action_type="wait")
 
 
-# =========================
-# SIMPLE HEURISTIC FALLBACK
-# =========================
-def heuristic_action(obs) -> Action:
-    for job in obs.pending_jobs:
-        for machine in obs.machines:
-            if machine.status == "idle":
-                return Action(
-                    action_type="assign_job",
-                    job_id=job.id,
-                    machine_id=machine.id,
-                )
-    return Action(action_type="wait")
+def heuristic_action(obs) -> Tuple[Action, str]:
+    for m in obs.machines:
+        if m.status == "broken":
+            return Action(action_type="repair", machine_id=m.id), f"repair {m.id}"
+    for j in sorted(obs.pending_jobs, key=lambda x: (x.deadline, -x.priority)):
+        for m in obs.machines:
+            if m.status == "idle":
+                s = f"assign_job {j.id} {m.id}"
+                return Action(action_type="assign_job", job_id=j.id, machine_id=m.id), s
+    return Action(action_type="wait"), "wait"
 
 
-# =========================
-# MAIN LOOP
-# =========================
-async def main():
+def run_task(task_name: str) -> None:
     client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)
-
-    env = FactoryEnv(task=TASK_NAME)
-
+    env = FactoryEnv(task=task_name)
     rewards: List[float] = []
     steps_taken = 0
     score = 0.0
     success = False
 
-    log_start(task=TASK_NAME, env=BENCHMARK, model=MODEL_NAME)
+    log_start(task=task_name, env=BENCHMARK, model=MODEL_NAME)
 
     try:
-        result = await env.reset()
-        obs = result.observation
+        obs = env.reset()
         last_reward = 0.0
 
-        for step in range(1, MAX_STEPS + 1):
-            if result.done:
+        for step in range(1, obs.max_steps + 1):
+            if obs.done:
                 break
-
-            # LLM decision
             action_text = get_model_action(client, step, obs, last_reward)
-
-            # Parse action
             action = parse_action(action_text)
-
-            # Fallback if invalid
-            if action.action_type == "wait" and len(obs.pending_jobs) > 0:
-                action = heuristic_action(obs)
-                action_text = "heuristic_assign"
-
-            # Step env
-            result = await env.step(action)
-
-            obs = result.observation
-            reward = result.reward or 0.0
-            done = result.done
-            error = None
-
+            if action.action_type == "wait" and (obs.pending_jobs or any(m.status == "broken" for m in obs.machines)):
+                action, action_text = heuristic_action(obs)
+            obs = env.step(action)
+            reward = obs.reward or 0.0
             rewards.append(reward)
             steps_taken = step
             last_reward = reward
-
-            log_step(step, action_text, reward, done, error)
-
-            if done:
+            log_step(step, action_text, reward, obs.done, None)
+            if obs.done:
                 break
 
-        # Normalize score
-        if rewards:
-            score = sum(rewards) / len(rewards)
-            score = max(0.0, min(1.0, score))
-
+        score = score_episode(env)
         success = score >= SUCCESS_SCORE_THRESHOLD
-
     finally:
-        try:
-            await env.close()
-        except Exception as e:
-            print(f"[DEBUG] env.close error: {e}", flush=True)
-
         log_end(success, steps_taken, score, rewards)
 
 
 if __name__ == "__main__":
-    asyncio.run(main())
+    run_task(TASK_NAME)

requirements.txt

@@ -1,3 +1,7 @@
-pydantic
-openai
-asyncio
+pydantic>=2.0
+openai>=1.0
+anthropic>=0.90
+gradio>=6.0
+openenv-core>=0.2.3
+fastapi>=0.100
+uvicorn>=0.23

server.py

@@ -0,0 +1,21 @@
+"""
+OpenEnv HTTP Server — Smart Factory Scheduling
+Routes: GET /health  POST /reset  POST /step  GET /state  GET /schema
+"""
+import os
+from openenv.core import create_app
+from factory_env.env import FactoryEnv
+from factory_env.models import FactoryAction, FactoryObservation
+
+TASK = os.getenv("FACTORY_TASK", "easy")
+
+app = create_app(
+    env=lambda: FactoryEnv(task=TASK, seed=42),
+    action_cls=FactoryAction,
+    observation_cls=FactoryObservation,
+    env_name="factory_env",
+)
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", 7860)))

train.py

@@ -0,0 +1,217 @@
+"""
+RL Training Loop — Smart Factory Scheduling
+============================================
+Strategy: Online In-Context RL — best trajectory fed as few-shot example each episode.
+
+Usage:
+  export OPENAI_API_KEY=sk-...          # OpenAI
+  export ANTHROPIC_API_KEY=sk-ant-...   # Claude
+  python train.py --task easy --episodes 10 --provider openai
+  python train.py --task medium --episodes 10 --provider claude
+"""
+
+import argparse
+import json
+import os
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import List, Optional, Tuple
+
+from factory_env.env import FactoryEnv
+from factory_env.grader import score_episode
+from factory_env.models import FactoryAction as Action
+
+
+def get_openai_client():
+    from openai import OpenAI
+    key = os.getenv("OPENAI_API_KEY") or os.getenv("HF_TOKEN") or os.getenv("API_KEY")
+    base = os.getenv("API_BASE_URL") or "https://api.openai.com/v1"
+    return OpenAI(api_key=key, base_url=base)
+
+
+def get_claude_client():
+    import anthropic
+    return anthropic.Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY"))
+
+
+@dataclass
+class Step:
+    step: int
+    obs_text: str
+    action_text: str
+    reward: float
+    done: bool
+
+
+@dataclass
+class Episode:
+    episode_num: int
+    task: str
+    steps: List[Step] = field(default_factory=list)
+    total_reward: float = 0.0
+    score: float = 0.0
+    completed: int = 0
+    late: int = 0
+
+    def to_few_shot(self, max_steps: int = 6) -> str:
+        lines = [f"# Best trajectory so far (score={self.score:.2f}, completed={self.completed} jobs)"]
+        for s in self.steps[:max_steps]:
+            lines.append(f"[Obs] {s.obs_text}")
+            lines.append(f"[Action] {s.action_text}  → reward: {s.reward:+.2f}")
+        return "\n".join(lines)
+
+
+SYSTEM_PROMPT = """You are an expert factory scheduling AI.
+Goal: complete all jobs before deadlines, keep machines busy, repair broken machines.
+Actions (one per step):
+  assign_job <job_id> <machine_id>
+  repair <machine_id>
+  wait
+Tips: Fix broken machines first. Sort by earliest deadline. High-priority jobs give bonus reward."""
+
+
+def obs_to_text(obs) -> str:
+    machines = ", ".join(f"{m.id}:{m.status}" + (f"({m.current_job})" if m.current_job else "") for m in obs.machines)
+    jobs = ", ".join(f"{j.id}[t={j.remaining_time},dl={j.deadline},p={j.priority}]" for j in obs.pending_jobs) or "none"
+    return f"t={obs.time} | machines: {machines} | pending: {jobs}"
+
+
+def call_llm(messages: list, provider: str, client, model: str) -> str:
+    try:
+        if provider == "claude":
+            system = next((m["content"] for m in messages if m["role"] == "system"), "")
+            user_msgs = [m for m in messages if m["role"] != "system"]
+            resp = client.messages.create(model=model, max_tokens=60, system=system, messages=user_msgs)
+            return resp.content[0].text.strip().splitlines()[0]
+        else:
+            resp = client.chat.completions.create(model=model, messages=messages, temperature=0.2, max_tokens=60)
+            return (resp.choices[0].message.content or "wait").strip().splitlines()[0]
+    except Exception as e:
+        print(f"  [LLM error] {e}")
+        return "wait"
+
+
+def parse_action(text: str) -> Action:
+    try:
+        parts = text.strip().split()
+        if parts[0] == "assign_job" and len(parts) == 3:
+            return Action(action_type="assign_job", job_id=parts[1], machine_id=parts[2])
+        if parts[0] == "repair" and len(parts) == 2:
+            return Action(action_type="repair", machine_id=parts[1])
+    except Exception:
+        pass
+    return Action(action_type="wait")
+
+
+def heuristic_action(obs) -> Tuple[Action, str]:
+    for m in obs.machines:
+        if m.status == "broken":
+            return Action(action_type="repair", machine_id=m.id), f"repair {m.id}"
+    for j in sorted(obs.pending_jobs, key=lambda x: (x.deadline, -x.priority)):
+        for m in obs.machines:
+            if m.status == "idle":
+                s = f"assign_job {j.id} {m.id}"
+                return Action(action_type="assign_job", job_id=j.id, machine_id=m.id), s
+    return Action(action_type="wait"), "wait"
+
+
+def run_episode(task, episode_num, provider, client, model, best_episode, seed=42, verbose=True) -> Episode:
+    env = FactoryEnv(task=task, seed=seed)
+    obs = env.reset()
+    last_reward = 0.0
+    ep = Episode(episode_num=episode_num, task=task)
+
+    if verbose:
+        print(f"\n  Episode {episode_num} | task={task} | seed={seed}")
+        print(f"  {len(obs.machines)} machines, {len(obs.pending_jobs)} jobs, {obs.max_steps} steps")
+
+    for step in range(1, obs.max_steps + 1):
+        if obs.done:
+            break
+
+        obs_text = obs_to_text(obs)
+        few_shot = best_episode.to_few_shot() if best_episode and step == 1 else ""
+        user = f"{few_shot}\n\n---\n" if few_shot else ""
+        user += f"Step {step} | Last reward: {last_reward:+.2f}\n{obs_text}\n\nAction:"
+
+        messages = [{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": user}]
+        action_text = call_llm(messages, provider, client, model)
+        action = parse_action(action_text)
+
+        if action.action_type == "wait" and (obs.pending_jobs or any(m.status == "broken" for m in obs.machines)):
+            action, action_text = heuristic_action(obs)
+
+        obs = env.step(action)
+        reward = obs.reward or 0.0
+        last_reward = reward
+        ep.steps.append(Step(step, obs_text, action_text, reward, obs.done))
+        ep.total_reward += reward
+
+        if verbose:
+            marker = "✓" if reward > 0.5 else ("✗" if reward < -0.05 else "·")
+            print(f"    [{marker}] step={step:2d}  {action_text:<30s}  r={reward:+.2f}")
+
+        if obs.done:
+            break
+
+    ep.score = score_episode(env)
+    ep.completed = len(env.completed_jobs)
+    ep.late = env.late_jobs
+
+    if verbose:
+        print(f"  → score={ep.score:.4f}  completed={ep.completed}  late={ep.late}")
+
+    return ep
+
+
+def train(task, num_episodes, provider, model, save_dir="runs", verbose=True):
+    print(f"\n{'='*60}")
+    print(f"  Smart Factory RL Training")
+    print(f"  Task: {task} | Episodes: {num_episodes} | Provider: {provider} | Model: {model}")
+    print(f"{'='*60}")
+
+    client = get_claude_client() if provider == "claude" else get_openai_client()
+    Path(save_dir).mkdir(exist_ok=True)
+
+    scores = []
+    best_episode = None
+
+    for ep_num in range(1, num_episodes + 1):
+        ep = run_episode(task, ep_num, provider, client, model, best_episode, seed=42 + ep_num - 1, verbose=verbose)
+        scores.append(ep.score)
+        if best_episode is None or ep.score > best_episode.score:
+            best_episode = ep
+            print(f"  ★ New best: score={ep.score:.4f}")
+        if ep_num < num_episodes:
+            time.sleep(1.0)
+
+    print(f"\n{'='*60}")
+    print(f"  Training Complete — {num_episodes} episodes | Task: {task}")
+    print(f"  First: {scores[0]:.4f} | Last: {scores[-1]:.4f} | Best: {max(scores):.4f}")
+    print(f"\n  Score per episode:")
+    for i, s in enumerate(scores, 1):
+        print(f"    ep{i:02d}: {s:.4f}  {'█' * int(s * 20)}")
+
+    out = Path(save_dir) / f"{task}_{provider}_{num_episodes}ep.json"
+    out.write_text(json.dumps({"task": task, "provider": provider, "model": model, "num_episodes": num_episodes, "scores": scores, "best_score": max(scores), "final_score": scores[-1]}, indent=2))
+    print(f"\n  Results saved → {out}")
+    return scores
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--task", default="easy", choices=["easy", "medium", "hard"])
+    parser.add_argument("--episodes", type=int, default=5)
+    parser.add_argument("--provider", default="openai", choices=["openai", "claude"])
+    parser.add_argument("--model", default="")
+    parser.add_argument("--save-dir", default="runs")
+    parser.add_argument("--quiet", action="store_true")
+    args = parser.parse_args()
+    if not args.model:
+        args.model = "claude-sonnet-4-6" if args.provider == "claude" else "gpt-4o-mini"
+    train(args.task, args.episodes, args.provider, args.model, args.save_dir, not args.quiet)
+
+
+if __name__ == "__main__":
+    main()