Upload folder using huggingface_hub

Dockerfile

@@ -57,6 +57,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 # Patch OpenEnv web UI so number inputs use step=0.01 (allows lr_scale=0.02, momentum_coef=0.9)
 RUN WEBIF="$$(find /app/env/.venv -path '*openenv*env_server*web_interface.py' | head -1)" && \
     /app/env/.venv/bin/python /app/env/scripts/patch_openenv_web_interface.py "$$WEBIF"
+# Patch OpenEnv web UI to add loss/perplexity chart and Run baseline (AdamW) button
+RUN WEBIF="$$(find /app/env/.venv -path '*openenv*env_server*web_interface.py' | head -1)" && \
+    /app/env/.venv/bin/python /app/env/scripts/patch_openenv_web_interface_chart.py "$$WEBIF"
 
 # Final runtime stage
 FROM ${BASE_IMAGE}

models.py

@@ -72,4 +72,8 @@ class MetaOptimizerObservation(Observation):
         default=None,
         description="Step at which loss first reached threshold (None if not yet reached)",
     )
+    perplexity: float | None = Field(
+        default=None,
+        description="exp(loss) for language modeling (None for regression)",
+    )
     

scripts/patch_openenv_web_interface_chart.py

@@ -0,0 +1,226 @@
+#!/usr/bin/env python3
+"""
+Patch OpenEnv web_interface.py to add:
+- Loss/perplexity chart and updateLossChart()
+- POST /web/run-baseline and GET /web/current-task for baseline comparison
+Idempotent: safe to run multiple times.
+"""
+import sys
+from pathlib import Path
+
+
+def _apply_routes_patch(text: str) -> str:
+    """Add /web/run-baseline and /web/current-task routes."""
+    old_routes = (
+        '    @app.get("/web/state")\n'
+        "    async def web_state():\n"
+        '        """State endpoint for web interface."""\n'
+        "        return web_manager.get_state()\n"
+        "\n"
+        "    return app"
+    )
+    new_routes = (
+        '    @app.get("/web/state")\n'
+        "    async def web_state():\n"
+        '        """State endpoint for web interface."""\n'
+        "        return web_manager.get_state()\n"
+        "\n"
+        '    @app.get("/web/current-task")\n'
+        "    async def web_current_task():\n"
+        '        """Current task spec for baseline comparison (if env supports it)."""\n'
+        "        get_spec = getattr(web_manager.env, \"get_current_task_spec\", None)\n"
+        "        if get_spec is None:\n"
+        "            return {}\n"
+        "        return get_spec() or {}\n"
+        "\n"
+        '    @app.post("/web/run-baseline")\n'
+        "    async def web_run_baseline():\n"
+        '        """Run baseline optimizer for current task; returns loss_trajectory and steps."""\n'
+        "        run_bl = getattr(web_manager.env, \"run_baseline\", None)\n"
+        "        if run_bl is None:\n"
+        "            return {\"loss_trajectory\": [], \"steps\": [], \"error\": \"Env has no run_baseline\"}\n"
+        "        return run_bl()\n"
+        "\n"
+        "    return app"
+    )
+    if "web/run-baseline" not in text and "web/state" in text and "return web_manager.get_state()" in text:
+        text = text.replace(old_routes, new_routes, 1)
+    return text
+
+
+def main() -> None:
+    if len(sys.argv) < 2:
+        import openenv.core.env_server.web_interface as m
+        path = Path(m.__file__).resolve()
+    else:
+        path = Path(sys.argv[1]).resolve()
+
+    if not path.exists():
+        print(f"File not found: {path}", file=sys.stderr)
+        sys.exit(1)
+
+    text = path.read_text()
+
+    # 1) Add Chart.js script in head (after title)
+    chart_script = '    <script src="https://cdn.jsdelivr.net/npm/chart.js"></script>\n'
+    old_head = "<title>OpenEnv Web Interface</title>\n    <style>"
+    new_head = "<title>OpenEnv Web Interface</title>\n" + chart_script + "    <style>"
+    if chart_script not in text and old_head in text:
+        text = text.replace(old_head, new_head, 1)
+
+    # 2) Add chart container between Current Observation and Action Logs
+    old_section = """                </div>
+                </div>
+
+                <!-- Action Logs -->
+                <div class="logs-container">"""
+    new_section = """                </div>
+                </div>
+
+                <!-- Loss chart -->
+                <div class="state-display">
+                    <h3>Loss / Perplexity</h3>
+                    <div id="loss-chart-container" style="height:200px;"><canvas id="loss-chart"></canvas></div>
+                    <button type="button" id="run-baseline-btn" class="btn btn-secondary" style="margin-top:8px;">Run baseline (AdamW)</button>
+                </div>
+
+                <!-- Action Logs -->
+                <div class="logs-container">"""
+    if "loss-chart-container" not in text and old_section in text:
+        text = text.replace(old_section, new_section, 1)
+    # If chart container exists but button does not, add button
+    if "loss-chart-container" in text and "run-baseline-btn" not in text:
+        text = text.replace(
+            "<canvas id=\"loss-chart\"></canvas></div>\n                </div>",
+            "<canvas id=\"loss-chart\"></canvas></div>\n                    <button type=\"button\" id=\"run-baseline-btn\" class=\"btn btn-secondary\" style=\"margin-top:8px;\">Run baseline (AdamW)</button>\n                </div>",
+            1,
+        )
+
+    # 3) Add updateLossChart call and method before updateChatInterface
+    old_update = """                }}
+            }}
+
+            updateChatInterface(episodeState) {{"""
+    new_update = """                }}
+                this.updateLossChart(episodeState);
+            }}
+
+            updateLossChart(episodeState) {{
+                const container = document.getElementById('loss-chart-container');
+                if (!container) return;
+                const steps = [];
+                const losses = [];
+                const perplexities = [];
+                if (episodeState.current_observation && typeof episodeState.current_observation.loss === 'number') {{
+                    const o = episodeState.current_observation;
+                    steps.push(o.step_count != null ? o.step_count : 0);
+                    losses.push(o.loss);
+                    if (typeof o.perplexity === 'number') perplexities.push(o.perplexity);
+                }}
+                (episodeState.action_logs || []).forEach(log => {{
+                    if (log.observation && typeof log.observation.loss === 'number') {{
+                        steps.push(log.observation.step_count != null ? log.observation.step_count : log.step_count);
+                        losses.push(log.observation.loss);
+                        if (typeof log.observation.perplexity === 'number') perplexities.push(log.observation.perplexity);
+                    }}
+                }});
+                if (steps.length === 0) return;
+                const ctx = document.getElementById('loss-chart');
+                if (!ctx) return;
+                if (this._lossChart) this._lossChart.destroy();
+                this._lossChart = new Chart(ctx, {{
+                    type: 'line',
+                    data: {{
+                        labels: steps,
+                        datasets: [
+                            {{ label: 'Loss', data: losses, borderColor: '#007bff', tension: 0.2, fill: false }}
+                        ].concat(perplexities.length ? [{{ label: 'Perplexity', data: perplexities, borderColor: '#28a745', tension: 0.2, fill: false }}] : [])
+                    }},
+                    options: {{ responsive: true, maintainAspectRatio: false, scales: {{ x: {{ title: {{ display: true, text: 'Step' }} }} }} }}
+                }});
+            }}
+
+            async runBaseline() {{
+                const btn = document.getElementById('run-baseline-btn');
+                if (btn) btn.disabled = true;
+                try {{
+                    const r = await fetch('/web/run-baseline', {{ method: 'POST' }});
+                    const data = await r.json();
+                    if (data.error || !data.loss_trajectory || !this._lossChart) {{ if (btn) btn.disabled = false; return; }}
+                    const L = data.loss_trajectory.length;
+                    const steps = data.steps && data.steps.length === L ? data.steps : Array.from({{ length: L }}, (_, i) => i);
+                    const curLen = this._lossChart.data.labels.length;
+                    const newLen = Math.max(curLen, steps.length);
+                    const newLabels = Array.from({{ length: newLen }}, (_, i) => i);
+                    this._lossChart.data.labels = newLabels;
+                    this._lossChart.data.datasets.forEach(ds => {{
+                        while (ds.data.length < newLen) ds.data.push(null);
+                    }});
+                    const baselineData = data.loss_trajectory.slice();
+                    while (baselineData.length < newLen) baselineData.push(null);
+                    this._lossChart.data.datasets.push({{ label: 'Baseline (AdamW)', data: baselineData, borderColor: '#dc3545', tension: 0.2, fill: false }});
+                    this._lossChart.update();
+                }} finally {{ if (btn) btn.disabled = false; }}
+            }}
+
+            updateChatInterface(episodeState) {{"""
+    if "updateLossChart(episodeState)" not in text and old_update in text:
+        text = text.replace(old_update, new_update, 1)
+
+    # 3b) Add Run baseline button click listener
+    old_listener = """                // State button
+                document.getElementById('state-btn').addEventListener('click', () => {{
+                    this.getState();
+                }});
+            }}"""
+    new_listener = """                // State button
+                document.getElementById('state-btn').addEventListener('click', () => {{
+                    this.getState();
+                }});
+
+                const runBaselineBtn = document.getElementById('run-baseline-btn');
+                if (runBaselineBtn) runBaselineBtn.addEventListener('click', () => this.runBaseline());
+            }}"""
+    if "run-baseline-btn" not in text or "runBaselineBtn.addEventListener" not in text:
+        if old_listener in text:
+            text = text.replace(old_listener, new_listener, 1)
+
+    # 3c) If updateLossChart exists but runBaseline does not, insert runBaseline
+    if "updateLossChart(episodeState)" in text and "async runBaseline()" not in text:
+        run_baseline_method = """
+            async runBaseline() {{
+                const btn = document.getElementById('run-baseline-btn');
+                if (btn) btn.disabled = true;
+                try {{
+                    const r = await fetch('/web/run-baseline', {{ method: 'POST' }});
+                    const data = await r.json();
+                    if (data.error || !data.loss_trajectory || !this._lossChart) {{ if (btn) btn.disabled = false; return; }}
+                    const L = data.loss_trajectory.length;
+                    const newLen = Math.max(this._lossChart.data.labels.length, L);
+                    const newLabels = Array.from({{ length: newLen }}, (_, i) => i);
+                    this._lossChart.data.labels = newLabels;
+                    this._lossChart.data.datasets.forEach(ds => {{
+                        while (ds.data.length < newLen) ds.data.push(null);
+                    }});
+                    const baselineData = data.loss_trajectory.slice();
+                    while (baselineData.length < newLen) baselineData.push(null);
+                    this._lossChart.data.datasets.push({{ label: 'Baseline (AdamW)', data: baselineData, borderColor: '#dc3545', tension: 0.2, fill: false }});
+                    this._lossChart.update();
+                }} finally {{ if (btn) btn.disabled = false; }}
+            }}
+"""
+        text = text.replace(
+            "                }});\n            }}\n\n            updateChatInterface(episodeState) {{",
+            "                }});\n            }}\n" + run_baseline_method + "\n            updateChatInterface(episodeState) {{",
+            1,
+        )
+
+    # 4) Add run-baseline and current-task routes
+    text = _apply_routes_patch(text)
+
+    path.write_text(text)
+    print("Patched (chart + run-baseline):", path)
+
+
+if __name__ == "__main__":
+    main()

server/meta_optimizer_environment.py

@@ -5,15 +5,17 @@
 # LICENSE file in the root directory of this source tree.
 
 """
-Meta-optimizer environment: train an RL agent to act as an optimizer on random regression tasks.
+Meta-optimizer environment: train an RL agent to act as an optimizer on inner tasks.
 
-Supports 50 training tasks, held-out eval, rich action space (LR, momentum, grad clip, weight decay),
-and convergence-speed reward. Action log is exposed for emergent-behavior visualization.
+Supports (1) SLM: next-token prediction with a tiny transformer; (2) sinusoid regression.
+Rich action space (LR, momentum, grad clip, weight decay), convergence-speed reward.
+Action log and loss/perplexity are exposed for dashboard visualization.
 """
 
 import math
 import random
-from typing import Any, Dict, List, Optional
+from dataclasses import asdict
+from typing import Any, Dict, List, Optional, Union
 from uuid import uuid4
 
 import torch
@@ -23,7 +25,22 @@ from openenv.core.env_server.interfaces import Environment
 from openenv.core.env_server.types import State
 
 from my_env.models import MetaOptimizerAction, MetaOptimizerObservation
-from .tasks import TRAIN_TASK_IDS, get_task, task_spec_from_dict, TaskSpec
+from .tasks import (
+    DEFAULT_CORPUS,
+    SLM_TRAIN_TASK_IDS,
+    TRAIN_TASK_IDS,
+    get_slm_task,
+    get_task,
+    task_spec_from_dict,
+    TaskSpec,
+    SLMTaskSpec,
+)
+from .slm_model import (
+    TinyLM,
+    build_vocab,
+    get_corpus_tensor,
+    sample_batch_slm,
+)
 
 # Defaults
 LOSS_THRESHOLD = 0.1
@@ -31,6 +48,8 @@ MAX_STEPS = 100
 BATCH_SIZE = 32
 # Dense reward scale: reward += DENSE_REWARD_SCALE * (prev_loss - current_loss) each step (potential-based, helps credit assignment)
 DENSE_REWARD_SCALE = 0.2
+# SLM loss threshold (cross-entropy); early termination when loss < this
+SLM_LOSS_THRESHOLD = 1.5
 
 
 def _default_device() -> torch.device:
@@ -60,6 +79,35 @@ def _get_batch(spec: TaskSpec, step: int, device: torch.device):
     return X, y
 
 
+def _build_slm(spec: SLMTaskSpec) -> TinyLM:
+    """Build a tiny decoder-only transformer for the given SLM task spec."""
+    torch.manual_seed(spec.arch_seed)
+    return TinyLM(
+        vocab_size=spec.vocab_size,
+        context_len=spec.context_len,
+        n_layer=spec.n_layer,
+        n_head=spec.n_head,
+        n_embd=spec.n_embd,
+    )
+
+
+def _get_batch_slm(
+    spec: SLMTaskSpec,
+    step: int,
+    device: torch.device,
+    corpus_ids: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Sample a batch for next-token prediction. Returns input_ids [B,T], target_ids [B,T]."""
+    return sample_batch_slm(
+        corpus_ids,
+        BATCH_SIZE,
+        spec.context_len,
+        step,
+        spec.data_seed,
+        device,
+    )
+
+
 def run_adam_baseline(
     task_id: Optional[int] = None,
     task_spec: Optional[Dict[str, Any]] = None,
@@ -78,6 +126,8 @@ def run_adam_baseline(
         torch.manual_seed(seed)
     device = _default_device()
     spec = task_spec_from_dict(task_spec) if task_spec is not None else get_task(task_id)
+    if isinstance(spec, SLMTaskSpec):
+        raise ValueError("Use run_adamw_baseline for SLM tasks")
     model = _build_model(spec).to(device)
     opt = torch.optim.Adam(model.parameters(), lr=lr)
     loss_trajectory: List[float] = []
@@ -128,6 +178,8 @@ def run_sgd_baseline(
         torch.manual_seed(seed)
     device = _default_device()
     spec = task_spec_from_dict(task_spec) if task_spec is not None else get_task(task_id)
+    if isinstance(spec, SLMTaskSpec):
+        raise ValueError("Use run_adamw_baseline for SLM tasks")
     model = _build_model(spec).to(device)
     opt = torch.optim.SGD(model.parameters(), lr=lr, momentum=momentum)
     loss_trajectory = []
@@ -159,6 +211,67 @@ def run_sgd_baseline(
     }
 
 
+def run_adamw_baseline(
+    task_id: Optional[int] = None,
+    task_spec: Optional[Dict[str, Any]] = None,
+    max_steps: int = MAX_STEPS,
+    loss_threshold: float = SLM_LOSS_THRESHOLD,
+    lr: float = 1e-3,
+    weight_decay: float = 0.01,
+    betas: tuple[float, float] = (0.9, 0.999),
+    seed: Optional[int] = None,
+    return_metrics: bool = False,
+):
+    """
+    Run AdamW on one SLM task. Returns steps to threshold, or full metrics dict if return_metrics=True.
+    """
+    if (task_id is None) == (task_spec is None):
+        raise ValueError("Provide exactly one of task_id or task_spec")
+    if seed is not None:
+        torch.manual_seed(seed)
+    device = _default_device()
+    spec = task_spec_from_dict(task_spec) if task_spec is not None else get_slm_task(task_id)
+    if isinstance(spec, TaskSpec):
+        raise ValueError("Use run_adam_baseline or run_sgd_baseline for sinusoid tasks")
+    assert isinstance(spec, SLMTaskSpec)
+    model = _build_slm(spec).to(device)
+    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay, betas=betas)
+    char2idx, _ = build_vocab()
+    corpus_ids = get_corpus_tensor(DEFAULT_CORPUS, char2idx, device)
+    loss_trajectory: List[float] = []
+    steps_to_threshold: Optional[int] = None
+    for step in range(max_steps):
+        inp, tgt = _get_batch_slm(spec, step, device, corpus_ids)
+        model.train()
+        opt.zero_grad()
+        logits = model(inp)
+        loss = nn.functional.cross_entropy(logits.view(-1, spec.vocab_size), tgt.view(-1))
+        loss.backward()
+        opt.step()
+        with torch.no_grad():
+            L = nn.functional.cross_entropy(
+                model(inp).view(-1, spec.vocab_size), tgt.view(-1)
+            ).item()
+        loss_trajectory.append(L)
+        if steps_to_threshold is None and L < loss_threshold:
+            steps_to_threshold = step + 1
+    final_loss = loss_trajectory[-1] if loss_trajectory else float("inf")
+    perplexity = math.exp(min(final_loss, 20.0))
+    if not return_metrics:
+        return steps_to_threshold if steps_to_threshold is not None else max_steps
+    last_k = min(10, len(loss_trajectory))
+    mean_last_k = sum(loss_trajectory[-last_k:]) / last_k if loss_trajectory else final_loss
+    return {
+        "steps_to_threshold": steps_to_threshold if steps_to_threshold is not None else max_steps,
+        "success": steps_to_threshold is not None,
+        "final_loss": final_loss,
+        "perplexity": perplexity,
+        "mean_last_10_loss": mean_last_k,
+        "loss_auc": sum(loss_trajectory) / len(loss_trajectory) if loss_trajectory else final_loss,
+        "loss_trajectory": loss_trajectory,
+    }
+
+
 def run_meta_optimizer_trajectory(
     task_id: Optional[int] = None,
     task_spec: Optional[Dict[str, Any]] = None,
@@ -226,7 +339,7 @@ class MetaOptimizerEnvironment(Environment[MetaOptimizerAction, MetaOptimizerObs
         self._device = _default_device()
 
         # Episode state (set in reset)
-        self._task_spec: Optional[TaskSpec] = None
+        self._task_spec: Optional[Union[TaskSpec, SLMTaskSpec]] = None
         self._model: Optional[nn.Module] = None
         self._velocities: Optional[List[torch.Tensor]] = None
         self._step_count: int = 0
@@ -235,6 +348,9 @@ class MetaOptimizerEnvironment(Environment[MetaOptimizerAction, MetaOptimizerObs
         self._steps_to_threshold: Optional[int] = None
         self._action_log: List[Dict[str, Any]] = []
         self._episode_id: Optional[str] = None
+        self._corpus_ids: Optional[torch.Tensor] = None  # for SLM only
+        self._is_slm: bool = False
+        self._slm_loss_threshold: float = SLM_LOSS_THRESHOLD
 
     def reset(
         self,
@@ -250,23 +366,44 @@ class MetaOptimizerEnvironment(Environment[MetaOptimizerAction, MetaOptimizerObs
         if task_spec is not None:
             self._task_spec = task_spec_from_dict(task_spec)
         else:
-            tid = task_id if task_id is not None else random.choice(TRAIN_TASK_IDS)
-            self._task_spec = get_task(tid)
-        self._model = _build_model(self._task_spec).to(self._device)
-        self._velocities = [torch.zeros_like(p) for p in self._model.parameters()]
-        self._step_count = 0
-        self._steps_to_threshold = None
-        self._action_log = []
-        self._episode_id = episode_id or str(uuid4())
-
-        # Initial loss (no update yet)
-        X, y = _get_batch(self._task_spec, 0, self._device)
-        with torch.no_grad():
-            out = self._model(X)
-            self._current_loss = nn.functional.mse_loss(out, y).item()
-        self._prev_loss = self._current_loss
-
-        return self._observation(reward=None, grad_norm=None)
+            tid = task_id if task_id is not None else random.choice(SLM_TRAIN_TASK_IDS)
+            self._task_spec = get_slm_task(tid)
+        self._is_slm = isinstance(self._task_spec, SLMTaskSpec)
+
+        if self._is_slm:
+            spec = self._task_spec
+            assert isinstance(spec, SLMTaskSpec)
+            self._model = _build_slm(spec).to(self._device)
+            self._velocities = [torch.zeros_like(p) for p in self._model.parameters()]
+            char2idx, _ = build_vocab()
+            self._corpus_ids = get_corpus_tensor(DEFAULT_CORPUS, char2idx, self._device)
+            self._step_count = 0
+            self._steps_to_threshold = None
+            self._action_log = []
+            self._episode_id = episode_id or str(uuid4())
+            inp, tgt = _get_batch_slm(spec, 0, self._device, self._corpus_ids)
+            with torch.no_grad():
+                logits = self._model(inp)
+                self._current_loss = nn.functional.cross_entropy(
+                    logits.view(-1, spec.vocab_size), tgt.view(-1)
+                ).item()
+            self._prev_loss = self._current_loss
+            return self._observation(reward=None, grad_norm=None, perplexity=math.exp(min(self._current_loss, 20.0)))
+        else:
+            spec = self._task_spec
+            assert isinstance(spec, TaskSpec)
+            self._model = _build_model(spec).to(self._device)
+            self._velocities = [torch.zeros_like(p) for p in self._model.parameters()]
+            self._step_count = 0
+            self._steps_to_threshold = None
+            self._action_log = []
+            self._episode_id = episode_id or str(uuid4())
+            X, y = _get_batch(spec, 0, self._device)
+            with torch.no_grad():
+                out = self._model(X)
+                self._current_loss = nn.functional.mse_loss(out, y).item()
+            self._prev_loss = self._current_loss
+            return self._observation(reward=None, grad_norm=None)
 
     def step(
         self,
@@ -289,10 +426,20 @@ class MetaOptimizerEnvironment(Environment[MetaOptimizerAction, MetaOptimizerObs
             "weight_decay_this_step": wd,
         })
 
-        X, y = _get_batch(self._task_spec, self._step_count + 1, self._device)
-        self._model.train()
-        out = self._model(X)
-        loss = nn.functional.mse_loss(out, y)
+        if self._is_slm:
+            spec = self._task_spec
+            assert isinstance(spec, SLMTaskSpec)
+            inp, tgt = _get_batch_slm(spec, self._step_count + 1, self._device, self._corpus_ids)
+            self._model.train()
+            logits = self._model(inp)
+            loss = nn.functional.cross_entropy(logits.view(-1, spec.vocab_size), tgt.view(-1))
+        else:
+            spec = self._task_spec
+            assert isinstance(spec, TaskSpec)
+            X, y = _get_batch(spec, self._step_count + 1, self._device)
+            self._model.train()
+            loss = nn.functional.mse_loss(self._model(X), y)
+
         self._model.zero_grad()
         loss.backward()
 
@@ -314,19 +461,32 @@ class MetaOptimizerEnvironment(Environment[MetaOptimizerAction, MetaOptimizerObs
                 if wd > 0:
                     p.sub_(p, alpha=wd)
 
-        with torch.no_grad():
-            new_out = self._model(X)
-            self._current_loss = nn.functional.mse_loss(new_out, y).item()
+        if self._is_slm:
+            spec = self._task_spec
+            assert isinstance(spec, SLMTaskSpec)
+            with torch.no_grad():
+                logits = self._model(inp)
+                self._current_loss = nn.functional.cross_entropy(
+                    logits.view(-1, spec.vocab_size), tgt.view(-1)
+                ).item()
+            loss_threshold = self._slm_loss_threshold
+            perp = math.exp(min(self._current_loss, 20.0))
+        else:
+            spec = self._task_spec
+            assert isinstance(spec, TaskSpec)
+            with torch.no_grad():
+                X, y = _get_batch(spec, self._step_count + 1, self._device)
+                self._current_loss = nn.functional.mse_loss(self._model(X), y).item()
+            loss_threshold = self.loss_threshold
+            perp = None
 
         self._step_count += 1
-        if self._steps_to_threshold is None and self._current_loss < self.loss_threshold:
+        if self._steps_to_threshold is None and self._current_loss < loss_threshold:
             self._steps_to_threshold = self._step_count
 
-        # Dense reward: reward loss decrease (potential-based shaping, does not change optimal policy)
         dense_reward = DENSE_REWARD_SCALE * (prev_loss - self._current_loss)
         self._prev_loss = self._current_loss
 
-        # End episode when we hit max_steps or when loss first crosses threshold (early termination)
         done = self._step_count >= self.max_steps or self._steps_to_threshold is not None
         if done:
             terminal = -(self._steps_to_threshold if self._steps_to_threshold is not None else self.max_steps)
@@ -334,13 +494,14 @@ class MetaOptimizerEnvironment(Environment[MetaOptimizerAction, MetaOptimizerObs
         else:
             reward = dense_reward
 
-        return self._observation(reward=reward, grad_norm=grad_norm, done=done)
+        return self._observation(reward=reward, grad_norm=grad_norm, done=done, perplexity=perp)
 
     def _observation(
         self,
         reward: Optional[float] = None,
         grad_norm: Optional[float] = None,
         done: bool = False,
+        perplexity: Optional[float] = None,
     ) -> MetaOptimizerObservation:
         meta: Dict[str, Any] = {}
         if self._steps_to_threshold is not None:
@@ -355,6 +516,7 @@ class MetaOptimizerEnvironment(Environment[MetaOptimizerAction, MetaOptimizerObs
             done=done,
             reward=reward,
             metadata=meta,
+            perplexity=perplexity,
         )
 
     @property
@@ -367,3 +529,23 @@ class MetaOptimizerEnvironment(Environment[MetaOptimizerAction, MetaOptimizerObs
     def get_episode_action_log(self) -> List[Dict[str, Any]]:
         """Return the action log for the current episode (for in-process viz or eval)."""
         return list(self._action_log)
+
+    def get_current_task_spec(self) -> Optional[Dict[str, Any]]:
+        """Return current task spec as a dict for dashboard / run-baseline. None if no episode started."""
+        if self._task_spec is None:
+            return None
+        if isinstance(self._task_spec, SLMTaskSpec):
+            return {"type": "slm", **asdict(self._task_spec)}
+        return {"type": "sinusoid", **asdict(self._task_spec)}
+
+    def run_baseline(self) -> Dict[str, Any]:
+        """Run the appropriate baseline (AdamW for SLM, Adam for sinusoid) for current task. Returns loss_trajectory and steps."""
+        spec_dict = self.get_current_task_spec()
+        if spec_dict is None:
+            return {"loss_trajectory": [], "steps": [], "error": "No task"}
+        if spec_dict.get("type") == "slm":
+            result = run_adamw_baseline(task_spec=spec_dict, max_steps=self.max_steps, return_metrics=True)
+        else:
+            result = run_adam_baseline(task_spec=spec_dict, max_steps=self.max_steps, return_metrics=True)
+        traj = result.get("loss_trajectory", [])
+        return {"loss_trajectory": traj, "steps": list(range(len(traj)))}

server/slm_model.py

@@ -0,0 +1,160 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Tiny decoder-only transformer for SLM meta-optimizer inner task.
+Pure PyTorch, no transformers dependency.
+"""
+
+import math
+from typing import Tuple
+
+import torch
+import torch.nn as nn
+
+
+# Fixed character vocab for reproducible SLM tasks (subset of printable ASCII)
+DEFAULT_CHARS = (
+    " \n\t"
+    "abcdefghijklmnopqrstuvwxyz"
+    "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
+    "0123456789"
+    ".,;:!?'\"-()"
+)
+DEFAULT_VOCAB_SIZE = len(DEFAULT_CHARS)
+
+
+def build_vocab(chars: str = DEFAULT_CHARS) -> Tuple[dict, dict]:
+    """Return char2idx and idx2char dicts."""
+    char2idx = {c: i for i, c in enumerate(chars)}
+    idx2char = {i: c for c, i in char2idx.items()}
+    return char2idx, idx2char
+
+
+def encode_corpus(text: str, char2idx: dict, default_idx: int = 0) -> torch.Tensor:
+    """Encode string to long tensor of token ids. Unknown chars map to default_idx."""
+    ids = [char2idx.get(c, default_idx) for c in text]
+    return torch.tensor(ids, dtype=torch.long)
+
+
+def get_corpus_tensor(
+    text: str,
+    char2idx: dict,
+    device: torch.device,
+) -> torch.Tensor:
+    """Return 1D long tensor of token ids on device."""
+    t = encode_corpus(text, char2idx)
+    return t.to(device)
+
+
+def sample_batch_slm(
+    corpus_ids: torch.Tensor,
+    batch_size: int,
+    context_len: int,
+    step: int,
+    data_seed: int,
+    device: torch.device,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Sample batch_size contiguous chunks from corpus for next-token prediction.
+    Returns: input_ids [B, context_len], target_ids [B, context_len] (target = input shifted by 1).
+    """
+    L = corpus_ids.size(0)
+    if L <= context_len + 1:
+        raise ValueError("Corpus too short for context_len")
+    max_start = L - context_len - 1
+    g = torch.Generator(device=device)
+    g.manual_seed(data_seed + step)
+    starts = torch.randint(0, max_start, (batch_size,), device=device, generator=g)
+    inputs = []
+    targets = []
+    for b in range(batch_size):
+        s = int(starts[b].item())
+        chunk = corpus_ids[s : s + context_len + 1]
+        inputs.append(chunk[:context_len])
+        targets.append(chunk[1 : context_len + 1])
+    return torch.stack(inputs), torch.stack(targets)
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, n_embd: int, n_head: int, block_size: int):
+        super().__init__()
+        assert n_embd % n_head == 0
+        self.n_head = n_head
+        self.n_embd = n_embd
+        self.head_dim = n_embd // n_head
+        self.register_buffer(
+            "mask",
+            torch.tril(torch.ones(block_size, block_size)).view(1, 1, block_size, block_size),
+        )
+        self.c_attn = nn.Linear(n_embd, 3 * n_embd)
+        self.c_proj = nn.Linear(n_embd, n_embd)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, T, C = x.shape
+        qkv = self.c_attn(x)
+        q, k, v = qkv.split(self.n_embd, dim=2)
+        q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
+        k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
+        v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
+        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(self.head_dim))
+        att = att.masked_fill(self.mask[:, :, :T, :T] == 0, float("-inf"))
+        att = torch.softmax(att, dim=-1)
+        out = (att @ v).transpose(1, 2).contiguous().view(B, T, C)
+        return self.c_proj(out)
+
+
+class Block(nn.Module):
+    def __init__(self, n_embd: int, n_head: int, block_size: int):
+        super().__init__()
+        self.ln1 = nn.LayerNorm(n_embd)
+        self.attn = CausalSelfAttention(n_embd, n_head, block_size)
+        self.ln2 = nn.LayerNorm(n_embd)
+        self.mlp = nn.Sequential(
+            nn.Linear(n_embd, 4 * n_embd),
+            nn.GELU(),
+            nn.Linear(4 * n_embd, n_embd),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x + self.attn(self.ln1(x))
+        x = x + self.mlp(self.ln2(x))
+        return x
+
+
+class TinyLM(nn.Module):
+    """Decoder-only transformer for next-token prediction."""
+
+    def __init__(
+        self,
+        vocab_size: int,
+        context_len: int,
+        n_layer: int,
+        n_head: int,
+        n_embd: int,
+    ):
+        super().__init__()
+        self.context_len = context_len
+        self.vocab_size = vocab_size
+        self.token_embed = nn.Embedding(vocab_size, n_embd)
+        self.pos_embed = nn.Embedding(context_len, n_embd)
+        self.blocks = nn.ModuleList(
+            [Block(n_embd, n_head, context_len) for _ in range(n_layer)]
+        )
+        self.ln_f = nn.LayerNorm(n_embd)
+        self.lm_head = nn.Linear(n_embd, vocab_size, bias=False)
+
+    def forward(self, idx: torch.Tensor) -> torch.Tensor:
+        # idx: [B, T]; clamp to valid range in case of encoding drift
+        idx = idx.clamp(0, self.vocab_size - 1)
+        B, T = idx.shape
+        pos = torch.arange(0, T, device=idx.device, dtype=torch.long)
+        x = self.token_embed(idx) + self.pos_embed(pos)
+        for block in self.blocks:
+            x = block(x)
+        x = self.ln_f(x)
+        logits = self.lm_head(x)
+        return logits

server/tasks.py

@@ -9,6 +9,7 @@ Task registry for meta-learning.
 
 Tasks can be from the internal registry (get_task(task_id)) or provided from outside
 via task_spec_from_dict() — the client sends the task definition to the environment.
+Supports sinusoid (regression) and SLM (next-token prediction) task types.
 """
 
 from dataclasses import dataclass
@@ -16,12 +17,27 @@ from typing import Any, Dict, List
 
 import math
 
+from .slm_model import DEFAULT_VOCAB_SIZE as SLM_DEFAULT_VOCAB_SIZE
+
 # Distribution A: 50 training tasks (low-freq sinusoids)
 TRAIN_TASK_IDS: List[int] = list(range(50))
 
 # Distribution B: held-out eval tasks (high-freq sinusoids — different distribution)
 EVAL_TASK_IDS: List[int] = [50, 51]
 
+# SLM: 50 train tasks, 2 eval (different corpus split or seed range)
+SLM_TRAIN_TASK_IDS: List[int] = list(range(50))
+SLM_EVAL_TASK_IDS: List[int] = [50, 51]
+
+# Fixed small corpus for SLM (character-level). ~10KB so tasks are reproducible.
+DEFAULT_CORPUS: str = (
+    "The quick brown fox jumps over the lazy dog. "
+    "Pack my box with five dozen liquor jugs. "
+    "How vexingly quick daft zebras jump. "
+    "Sphinx of black quartz, judge my vow. "
+    "The five boxing wizards jump quickly. "
+) * 200  # repeat to get enough length for sampling
+
 # Bounds for each distribution (freq, amplitude, phase)
 DIST_A_FREQ = (1.0, 3.0)
 DIST_A_AMP = (0.5, 2.0)
@@ -46,6 +62,22 @@ class TaskSpec:
     distribution: str  # "A" or "B"
 
 
+@dataclass
+class SLMTaskSpec:
+    """Spec for one SLM (next-token prediction) task."""
+
+    task_id: int
+    data_seed: int
+    arch_seed: int
+    vocab_size: int
+    context_len: int  # block size
+    n_layer: int
+    n_head: int
+    n_embd: int
+    corpus_id: str  # e.g. "default"
+    distribution: str  # "A" or "B" or "external"
+
+
 def get_task(task_id: int) -> TaskSpec:
     """
     Return the task spec for the given task_id.
@@ -88,28 +120,73 @@ def get_task(task_id: int) -> TaskSpec:
     )
 
 
-def task_spec_from_dict(d: Dict[str, Any]) -> TaskSpec:
+def get_slm_task(task_id: int) -> SLMTaskSpec:
     """
-    Build a TaskSpec from an external dict (sent by the client).
-    The task is defined outside the env; we just parse it here.
-
-    Expected keys for type "sinusoid":
-      type="sinusoid", amplitude, freq, phase, data_seed (optional), arch_seed (optional),
-      input_dim (optional, default 1), hidden_dim (optional, default 32), task_id (optional).
+    Return the SLM task spec for the given task_id.
+    Task IDs 0..49 = Distribution A (train), 50+ = Distribution B (eval).
     """
-    task_type = d.get("type", "sinusoid")
-    if task_type != "sinusoid":
-        raise ValueError(f"Unknown task type: {task_type}")
-    task_id = d.get("task_id", 0)
-    return TaskSpec(
+    if task_id < 0:
+        raise ValueError(f"task_id must be >= 0, got {task_id}")
+    r = task_id * 7919 + 1
+    data_seed = task_id * 31337
+    arch_seed = task_id * 131 + 7
+    if task_id < 50:
+        distribution = "A"
+    else:
+        distribution = "B"
+    return SLMTaskSpec(
+        task_id=task_id,
+        data_seed=data_seed,
+        arch_seed=arch_seed,
+        vocab_size=SLM_DEFAULT_VOCAB_SIZE,
+        context_len=64,
+        n_layer=2,
+        n_head=4,
+        n_embd=128,
+        corpus_id="default",
+        distribution=distribution,
+    )
+
+
+def slm_task_spec_from_dict(d: Dict[str, Any]) -> SLMTaskSpec:
+    """Build an SLMTaskSpec from an external dict (type='slm')."""
+    task_id = int(d.get("task_id", 0))
+    return SLMTaskSpec(
         task_id=task_id,
-        input_dim=int(d.get("input_dim", 1)),
-        hidden_dim=int(d.get("hidden_dim", 32)),
-        output_dim=1,
         data_seed=int(d.get("data_seed", task_id * 31337)),
         arch_seed=int(d.get("arch_seed", task_id * 131 + 7)),
-        amplitude=float(d["amplitude"]),
-        freq=float(d["freq"]),
-        phase=float(d["phase"]),
+        vocab_size=int(d.get("vocab_size", SLM_DEFAULT_VOCAB_SIZE)),
+        context_len=int(d.get("context_len", 64)),
+        n_layer=int(d.get("n_layer", 2)),
+        n_head=int(d.get("n_head", 4)),
+        n_embd=int(d.get("n_embd", 128)),
+        corpus_id=str(d.get("corpus_id", "default")),
         distribution=d.get("distribution", "external"),
     )
+
+
+def task_spec_from_dict(d: Dict[str, Any]) -> TaskSpec | SLMTaskSpec:
+    """
+    Build a TaskSpec or SLMTaskSpec from an external dict (sent by the client).
+
+    For type "sinusoid": amplitude, freq, phase, data_seed (optional), arch_seed (optional), etc.
+    For type "slm": vocab_size, context_len, n_layer, n_head, n_embd (all optional with defaults).
+    """
+    task_type = d.get("type", "slm")
+    if task_type == "sinusoid":
+        task_id = d.get("task_id", 0)
+        return TaskSpec(
+            task_id=task_id,
+            input_dim=int(d.get("input_dim", 1)),
+            hidden_dim=int(d.get("hidden_dim", 32)),
+            output_dim=1,
+            data_seed=int(d.get("data_seed", task_id * 31337)),
+            arch_seed=int(d.get("arch_seed", task_id * 131 + 7)),
+            amplitude=float(d["amplitude"]),
+            freq=float(d["freq"]),
+            phase=float(d["phase"]),
+            distribution=d.get("distribution", "external"),
+        )
+    if task_type == "slm":
+        return slm_task_spec_from_dict(d)
+    raise ValueError(f"Unknown task type: {task_type!r}")