fix: ui

README.md

@@ -15,6 +15,8 @@ tags:
 
 An RL environment that simulates **Go-To-Market (GTM) strategy optimization** for product launches. Agents learn to allocate marketing budgets, target customer segments, craft messaging, run experiments, and adjust pricing to maximize revenue under uncertainty.
 
+The Hugging Face Space ships with an interactive dashboard at `/web/` (the default route) that runs the **trained PPO policy**, an **equal-allocation heuristic**, and a **uniform random** agent on the same task and seed, then renders revenue / brand-health / budget-allocation comparisons via Plotly.
+
 ## Why GTM?
 
 Every startup and growth team does GTM optimization manually — iterating on channels, messaging, and targeting through trial and error. This environment captures the real complexity: noisy metrics, delayed brand effects, diminishing returns on ad spend, and the tension between short-term revenue and long-term brand strength.
@@ -96,6 +98,35 @@ export OPENAI_API_KEY=sk-...
 python baseline.py --model gpt-4o-mini
 ```
 
+### Train an RL Policy
+
+A custom lightweight PPO trainer (`rl/train.py`) trains a small actor-critic
+network against the simulator. One checkpoint per task.
+
+```bash
+# Train
+python -m rl.train --task channel_optimizer --total-steps 200000
+python -m rl.train --task growth_strategist --total-steps 300000
+python -m rl.train --task market_dominator  --total-steps 500000
+
+# Inference (greedy rollout, prints per-week actions and grader score)
+python -m rl.infer --task channel_optimizer
+```
+
+Checkpoints are written to `checkpoints/<task_id>.pt`. Commit them so the
+deployed Space can serve `/infer` without retraining.
+
+### Inference via API
+
+```bash
+curl -X POST http://localhost:7860/infer \
+     -H "Content-Type: application/json" \
+     -d '{"task_id": "channel_optimizer", "seed": 42}'
+```
+
+Returns a JSON payload with `grader_score`, `total_revenue`, and the full
+weekly action trajectory.
+
 ### API Endpoints
 
 | Endpoint | Method | Description |
@@ -103,6 +134,7 @@ python baseline.py --model gpt-4o-mini
 | `/tasks` | GET | List all tasks with action schemas |
 | `/baseline` | POST | Run heuristic baseline, return scores |
 | `/grader` | POST | Get grader score for a task |
+| `/infer` | POST | Run trained RL policy on a task, return action trajectory |
 | `/reset` | POST | Reset environment for a task |
 | `/step` | POST | Execute one action step |
 | `/state` | GET | Get current episode state |

requirements.txt

@@ -5,3 +5,6 @@ pydantic>=2.0.0
 websockets>=15.0.1
 openai>=1.0.0
 numpy>=1.24.0
+torch>=2.0.0
+plotly>=5.0
+pandas>=2.0

rl/env_adapter.py

@@ -0,0 +1,138 @@
+"""Bridge between GTMEnvironment (dict-based) and tensor-based RL code."""
+
+from __future__ import annotations
+
+from typing import Dict, List
+
+import torch
+
+from models import GTMAction, GTMObservation
+from server.simulation import MESSAGING_DIMS
+from server.tasks import TaskDefinition
+
+# Sentinel string used by the categorical heads when the agent picks "no action"
+NONE_OPTION = "__none__"
+
+
+def compute_obs_dim(task: TaskDefinition) -> int:
+    """Number of scalar features in the flattened observation tensor."""
+    n_channels = len(task.channels)
+    n_segments = len(task.segments)
+    # 5 globals + 5 per channel + 4 per segment
+    return 5 + 5 * n_channels + 4 * n_segments
+
+
+def compute_action_dims(task: TaskDefinition) -> Dict[str, int]:
+    """Output sizes for each policy head."""
+    return {
+        "budget": len(task.channels),
+        "segment": len(task.segments),
+        "messaging": len(MESSAGING_DIMS),
+        # +1 for the "none" option
+        "experiment": len(task.available_experiments) + 1,
+        "pricing": len(task.available_pricing_actions) + 1,
+    }
+
+
+def experiment_options(task: TaskDefinition) -> List[str]:
+    return [NONE_OPTION] + list(task.available_experiments)
+
+
+def pricing_options(task: TaskDefinition) -> List[str]:
+    return [NONE_OPTION] + list(task.available_pricing_actions)
+
+
+def obs_to_tensor(obs: GTMObservation, task: TaskDefinition) -> torch.Tensor:
+    """Flatten a GTMObservation into a fixed-size float32 tensor.
+
+    Layout (in order):
+        - week / total_weeks
+        - budget_remaining / total_budget
+        - brand_score / 100
+        - total_revenue / revenue_target
+        - average_cac / 100
+        - per channel (in task order): spend/weekly_budget, ctr, cvr, roi, conversions/100
+        - per segment (in task order): conversion_rate, engagement/100, churn, revenue/10k
+    """
+    total_weeks = max(obs.total_weeks, 1)
+    total_budget = task.total_budget if task.total_budget > 0 else 1.0
+    weekly_budget = max(obs.weekly_budget, 1.0)
+    revenue_target = task.revenue_target if task.revenue_target > 0 else 1.0
+
+    feats: List[float] = [
+        obs.week / total_weeks,
+        obs.budget_remaining / total_budget,
+        obs.brand_score / 100.0,
+        obs.total_revenue / revenue_target,
+        obs.average_cac / 100.0,
+    ]
+
+    for ch in task.channels:
+        m = obs.channel_metrics.get(ch.name)
+        if m is None:
+            feats.extend([0.0, 0.0, 0.0, 0.0, 0.0])
+        else:
+            feats.extend([
+                m.spend / weekly_budget,
+                m.ctr,
+                m.cvr,
+                # clip ROI into a reasonable range
+                max(-2.0, min(5.0, m.roi)),
+                m.conversions / 100.0,
+            ])
+
+    for seg in task.segments:
+        sm = obs.segment_performance.get(seg.name)
+        if sm is None:
+            feats.extend([0.0, 0.0, 0.0, 0.0])
+        else:
+            feats.extend([
+                sm.conversion_rate,
+                min(1.0, sm.engagement_score / 100.0),
+                sm.churn_rate,
+                sm.revenue / 10000.0,
+            ])
+
+    return torch.tensor(feats, dtype=torch.float32)
+
+
+def policy_sample_to_action(
+    sample: Dict[str, torch.Tensor],
+    task: TaskDefinition,
+) -> GTMAction:
+    """Convert sampled policy outputs into a GTMAction.
+
+    sample keys:
+        budget   — Tensor[n_channels] on the simplex (Dirichlet sample)
+        segment  — Tensor[n_segments] on the simplex
+        messaging— Tensor[6]          on the simplex
+        experiment — int (index into experiment_options(task))
+        pricing    — int (index into pricing_options(task))
+    """
+    budget = sample["budget"].detach().cpu().tolist()
+    segment = sample["segment"].detach().cpu().tolist()
+    messaging = sample["messaging"].detach().cpu().tolist()
+
+    budget_alloc = {ch.name: float(budget[i]) for i, ch in enumerate(task.channels)}
+    segment_target = {seg.name: float(segment[i]) for i, seg in enumerate(task.segments)}
+    messaging_dict = {dim: float(messaging[i]) for i, dim in enumerate(MESSAGING_DIMS)}
+
+    exp_idx = int(sample["experiment"].item())
+    exp_opts = experiment_options(task)
+    experiment = exp_opts[exp_idx] if exp_idx < len(exp_opts) else NONE_OPTION
+    if experiment == NONE_OPTION:
+        experiment = None
+
+    price_idx = int(sample["pricing"].item())
+    price_opts = pricing_options(task)
+    pricing = price_opts[price_idx] if price_idx < len(price_opts) else NONE_OPTION
+    if pricing == NONE_OPTION:
+        pricing = None
+
+    return GTMAction(
+        budget_allocation=budget_alloc,
+        segment_targeting=segment_target,
+        messaging=messaging_dict,
+        experiment=experiment,
+        pricing_action=pricing,
+    )

rl/infer.py

@@ -0,0 +1,139 @@
+"""Inference: load a trained checkpoint and run one episode of GTM optimization.
+
+Usage:
+    python -m rl.infer --task channel_optimizer
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+from typing import Any, Dict, List, Optional
+
+import torch
+
+# Make repo root importable when run as a module
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from server.environment import GTMEnvironment
+from server.tasks import get_task
+
+from rl.env_adapter import (
+    compute_action_dims,
+    compute_obs_dim,
+    obs_to_tensor,
+    policy_sample_to_action,
+)
+from rl.policy import GTMActorCritic
+
+
+def _build_policy(task_id: str) -> tuple[GTMActorCritic, Any]:
+    task = get_task(task_id)
+    obs_dim = compute_obs_dim(task)
+    action_dims = compute_action_dims(task)
+    policy = GTMActorCritic(
+        obs_dim=obs_dim,
+        n_channels=action_dims["budget"],
+        n_segments=action_dims["segment"],
+        n_messaging=action_dims["messaging"],
+        n_experiments=action_dims["experiment"],
+        n_pricing=action_dims["pricing"],
+    )
+    return policy, task
+
+
+def run_inference(
+    task_id: str,
+    checkpoint_path: Optional[str] = None,
+    seed: Optional[int] = None,
+) -> Dict[str, Any]:
+    """Run one deterministic episode and return the action trajectory + metrics."""
+    policy, task = _build_policy(task_id)
+
+    if checkpoint_path is None:
+        checkpoint_path = os.path.join("checkpoints", f"{task_id}.pt")
+
+    loaded = False
+    if os.path.exists(checkpoint_path):
+        ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
+        policy.load_state_dict(ckpt["model_state"])
+        loaded = True
+
+    policy.eval()
+
+    env = GTMEnvironment()
+    obs = env.reset(task_id=task_id, seed=seed)
+    actions: List[Dict[str, Any]] = []
+
+    while not obs.done:
+        obs_t = obs_to_tensor(obs, task)
+        with torch.no_grad():
+            sample, _, _ = policy.act(obs_t.unsqueeze(0), deterministic=True)
+        sample_squeezed = {k: v.squeeze(0) for k, v in sample.items()}
+        action = policy_sample_to_action(sample_squeezed, task)
+
+        next_obs = env.step(action)
+        actions.append({
+            "week": next_obs.week,
+            "budget_allocation": action.budget_allocation,
+            "segment_targeting": action.segment_targeting,
+            "messaging": action.messaging,
+            "experiment": action.experiment,
+            "pricing_action": action.pricing_action,
+            "weekly_reward": next_obs.reward,
+            "total_revenue": next_obs.total_revenue,
+            "brand_score": next_obs.brand_score,
+        })
+        obs = next_obs
+
+    grader_score = env.get_grader_score(env.state.episode_id)
+    msg = (
+        f"Trained policy ({checkpoint_path})" if loaded
+        else f"Untrained random policy (no checkpoint at {checkpoint_path})"
+    )
+
+    return {
+        "task_id": task_id,
+        "checkpoint_loaded": loaded,
+        "grader_score": grader_score,
+        "total_revenue": float(obs.total_revenue),
+        "total_conversions": int(obs.total_conversions),
+        "average_cac": float(obs.average_cac),
+        "brand_score": float(obs.brand_score),
+        "actions": actions,
+        "message": msg,
+    }
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--task", required=True, choices=["channel_optimizer", "growth_strategist", "market_dominator"])
+    parser.add_argument("--checkpoint", default=None)
+    parser.add_argument("--seed", type=int, default=None)
+    parser.add_argument("--json", action="store_true", help="Print full JSON output")
+    args = parser.parse_args()
+
+    result = run_inference(args.task, checkpoint_path=args.checkpoint, seed=args.seed)
+
+    if args.json:
+        print(json.dumps(result, indent=2, default=str))
+        return
+
+    print(f"Task: {result['task_id']}")
+    print(f"Checkpoint loaded: {result['checkpoint_loaded']}")
+    print(f"Grader score: {result['grader_score']}")
+    print(f"Total revenue: ${result['total_revenue']:,.2f}")
+    print(f"Total conversions: {result['total_conversions']}")
+    print(f"Average CAC: ${result['average_cac']:.2f}")
+    print(f"Brand score: {result['brand_score']:.1f}")
+    print()
+    print("Weekly actions:")
+    for a in result["actions"]:
+        budget_str = ", ".join(f"{k}={v:.2f}" for k, v in a["budget_allocation"].items())
+        print(f"  Week {a['week']:2d}: budget=[{budget_str}]  reward={a['weekly_reward']:.3f}")
+
+
+if __name__ == "__main__":
+    main()

rl/policy.py

@@ -0,0 +1,121 @@
+"""Actor-critic MLP policy for the GTM environment."""
+
+from __future__ import annotations
+
+from typing import Dict, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.distributions import Categorical, Dirichlet
+
+
+class GTMActorCritic(nn.Module):
+    """Small actor-critic network with task-specific action heads.
+
+    Heads:
+        budget    — Dirichlet over channels        (simplex)
+        segment   — Dirichlet over segments        (simplex)
+        messaging — Dirichlet over 6 dimensions    (simplex)
+        experiment— Categorical (incl. "none")
+        pricing   — Categorical (incl. "none")
+
+    Concentrations for the Dirichlet heads are produced by softplus(linear)+1
+    so they are always positive and start near a uniform distribution.
+    """
+
+    def __init__(
+        self,
+        obs_dim: int,
+        n_channels: int,
+        n_segments: int,
+        n_messaging: int,
+        n_experiments: int,
+        n_pricing: int,
+        hidden_dim: int = 128,
+    ):
+        super().__init__()
+        self.obs_dim = obs_dim
+        self.n_channels = n_channels
+        self.n_segments = n_segments
+        self.n_messaging = n_messaging
+        self.n_experiments = n_experiments
+        self.n_pricing = n_pricing
+
+        self.trunk = nn.Sequential(
+            nn.Linear(obs_dim, hidden_dim),
+            nn.Tanh(),
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.Tanh(),
+        )
+
+        self.budget_head = nn.Linear(hidden_dim, n_channels)
+        self.segment_head = nn.Linear(hidden_dim, n_segments)
+        self.messaging_head = nn.Linear(hidden_dim, n_messaging)
+        self.experiment_head = nn.Linear(hidden_dim, n_experiments)
+        self.pricing_head = nn.Linear(hidden_dim, n_pricing)
+        self.value_head = nn.Linear(hidden_dim, 1)
+
+    def forward(self, obs: torch.Tensor) -> Tuple[Dict[str, torch.distributions.Distribution], torch.Tensor]:
+        h = self.trunk(obs)
+
+        budget_alpha = F.softplus(self.budget_head(h)) + 1.0
+        segment_alpha = F.softplus(self.segment_head(h)) + 1.0
+        messaging_alpha = F.softplus(self.messaging_head(h)) + 1.0
+
+        # Avoid sampling exact zeros which would break Dirichlet log_prob
+        budget_alpha = budget_alpha.clamp(min=1e-3)
+        segment_alpha = segment_alpha.clamp(min=1e-3)
+        messaging_alpha = messaging_alpha.clamp(min=1e-3)
+
+        dists: Dict[str, torch.distributions.Distribution] = {
+            "budget": Dirichlet(budget_alpha),
+            "segment": Dirichlet(segment_alpha),
+            "messaging": Dirichlet(messaging_alpha),
+            "experiment": Categorical(logits=self.experiment_head(h)),
+            "pricing": Categorical(logits=self.pricing_head(h)),
+        }
+        value = self.value_head(h).squeeze(-1)
+        return dists, value
+
+    def act(
+        self,
+        obs: torch.Tensor,
+        deterministic: bool = False,
+    ) -> Tuple[Dict[str, torch.Tensor], torch.Tensor, torch.Tensor]:
+        """Sample (or pick the mode of) an action from the policy.
+
+        Returns: (sample_dict, total_log_prob, value)
+        """
+        dists, value = self.forward(obs)
+        sample: Dict[str, torch.Tensor] = {}
+        log_probs = []
+        for name, dist in dists.items():
+            if deterministic:
+                if isinstance(dist, Dirichlet):
+                    # mean of a Dirichlet
+                    s = dist.concentration / dist.concentration.sum(dim=-1, keepdim=True)
+                else:
+                    s = dist.probs.argmax(dim=-1)
+            else:
+                s = dist.sample()
+            sample[name] = s
+            log_probs.append(dist.log_prob(s))
+        total_log_prob = torch.stack(log_probs, dim=0).sum(dim=0)
+        return sample, total_log_prob, value
+
+    def evaluate_actions(
+        self,
+        obs: torch.Tensor,
+        actions: Dict[str, torch.Tensor],
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Recompute log-probs, entropy, and value for stored (obs, action) pairs."""
+        dists, value = self.forward(obs)
+        log_probs = []
+        entropies = []
+        for name, dist in dists.items():
+            log_probs.append(dist.log_prob(actions[name]))
+            entropies.append(dist.entropy())
+        total_log_prob = torch.stack(log_probs, dim=0).sum(dim=0)
+        total_entropy = torch.stack(entropies, dim=0).sum(dim=0)
+        return total_log_prob, total_entropy, value

rl/train.py

@@ -0,0 +1,244 @@
+"""PPO trainer for the GTM Strategy Optimizer.
+
+Usage:
+    python -m rl.train --task channel_optimizer --total-steps 200000 --seed 0
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+import time
+from dataclasses import dataclass
+from typing import Dict, List
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+
+# Make repo root importable when run as a module
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from server.environment import GTMEnvironment
+from server.tasks import get_task
+
+from rl.env_adapter import (
+    compute_action_dims,
+    compute_obs_dim,
+    obs_to_tensor,
+    policy_sample_to_action,
+)
+from rl.policy import GTMActorCritic
+
+
+# ── Hyperparameters ────────────────────────────────────────────────────────
+
+LR = 3e-4
+GAMMA = 0.99
+GAE_LAMBDA = 0.95
+CLIP_EPS = 0.2
+VF_COEF = 0.5
+ENT_COEF = 0.01
+N_STEPS = 2048
+N_EPOCHS = 10
+BATCH_SIZE = 64
+MAX_GRAD_NORM = 0.5
+
+
+# ── Rollout buffer ─────────────────────────────────────────────────────────
+
+
+@dataclass
+class RolloutStep:
+    obs: torch.Tensor
+    action: Dict[str, torch.Tensor]
+    log_prob: torch.Tensor
+    reward: float
+    value: torch.Tensor
+    done: bool
+
+
+def compute_gae(
+    rollout: List[RolloutStep],
+    last_value: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Generalized Advantage Estimation."""
+    advantages = torch.zeros(len(rollout), dtype=torch.float32)
+    gae = 0.0
+    next_value = last_value.item()
+    for t in reversed(range(len(rollout))):
+        nonterminal = 0.0 if rollout[t].done else 1.0
+        delta = rollout[t].reward + GAMMA * next_value * nonterminal - rollout[t].value.item()
+        gae = delta + GAMMA * GAE_LAMBDA * nonterminal * gae
+        advantages[t] = gae
+        next_value = rollout[t].value.item()
+    returns = advantages + torch.tensor([s.value.item() for s in rollout], dtype=torch.float32)
+    return advantages, returns
+
+
+# ── Training loop ──────────────────────────────────────────────────────────
+
+
+def train(
+    task_id: str,
+    total_steps: int,
+    seed: int = 0,
+    checkpoint_dir: str = "checkpoints",
+    log_every: int = 1,
+) -> None:
+    torch.manual_seed(seed)
+
+    env = GTMEnvironment()
+    task = get_task(task_id)
+    obs_dim = compute_obs_dim(task)
+    action_dims = compute_action_dims(task)
+
+    policy = GTMActorCritic(
+        obs_dim=obs_dim,
+        n_channels=action_dims["budget"],
+        n_segments=action_dims["segment"],
+        n_messaging=action_dims["messaging"],
+        n_experiments=action_dims["experiment"],
+        n_pricing=action_dims["pricing"],
+    )
+    optimizer = optim.Adam(policy.parameters(), lr=LR)
+
+    obs = env.reset(task_id=task_id, seed=seed)
+    obs_t = obs_to_tensor(obs, task)
+
+    global_step = 0
+    update_idx = 0
+    best_mean_return = -float("inf")
+    episode_returns: List[float] = []
+    current_return = 0.0
+    start_time = time.time()
+
+    os.makedirs(checkpoint_dir, exist_ok=True)
+    checkpoint_path = os.path.join(checkpoint_dir, f"{task_id}.pt")
+
+    while global_step < total_steps:
+        # ── Collect rollout ──────────────────────────────────────
+        rollout: List[RolloutStep] = []
+        for _ in range(N_STEPS):
+            with torch.no_grad():
+                sample, log_prob, value = policy.act(obs_t.unsqueeze(0))
+            sample_squeezed = {k: v.squeeze(0) for k, v in sample.items()}
+
+            action = policy_sample_to_action(sample_squeezed, task)
+            next_obs = env.step(action)
+            reward = float(next_obs.reward) if next_obs.reward is not None else 0.0
+            done = bool(next_obs.done)
+            current_return += reward
+
+            rollout.append(
+                RolloutStep(
+                    obs=obs_t,
+                    action=sample_squeezed,
+                    log_prob=log_prob.squeeze(0).detach(),
+                    reward=reward,
+                    value=value.squeeze(0).detach(),
+                    done=done,
+                )
+            )
+
+            global_step += 1
+            if done:
+                episode_returns.append(current_return)
+                current_return = 0.0
+                next_obs = env.reset(task_id=task_id, seed=seed + global_step)
+            obs_t = obs_to_tensor(next_obs, task)
+
+        # bootstrap final value
+        with torch.no_grad():
+            _, _, last_value = policy.act(obs_t.unsqueeze(0))
+            last_value = last_value.squeeze(0).detach()
+
+        advantages, returns = compute_gae(rollout, last_value)
+        advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
+
+        # Stack rollout into tensors for minibatching
+        obs_batch = torch.stack([s.obs for s in rollout])
+        old_log_probs = torch.stack([s.log_prob for s in rollout])
+        action_batch = {
+            k: torch.stack([s.action[k] for s in rollout]) for k in rollout[0].action
+        }
+
+        # ── PPO update ───────────────────────────────────────────
+        n = len(rollout)
+        indices = list(range(n))
+        for _ in range(N_EPOCHS):
+            # shuffle
+            perm = torch.randperm(n).tolist()
+            for start in range(0, n, BATCH_SIZE):
+                mb = perm[start : start + BATCH_SIZE]
+                mb_obs = obs_batch[mb]
+                mb_actions = {k: v[mb] for k, v in action_batch.items()}
+                mb_old_lp = old_log_probs[mb]
+                mb_adv = advantages[mb]
+                mb_ret = returns[mb]
+
+                new_log_probs, entropy, value_pred = policy.evaluate_actions(mb_obs, mb_actions)
+                ratio = torch.exp(new_log_probs - mb_old_lp)
+                surr1 = ratio * mb_adv
+                surr2 = torch.clamp(ratio, 1.0 - CLIP_EPS, 1.0 + CLIP_EPS) * mb_adv
+                policy_loss = -torch.min(surr1, surr2).mean()
+                value_loss = 0.5 * (value_pred - mb_ret).pow(2).mean()
+                entropy_loss = -entropy.mean()
+
+                loss = policy_loss + VF_COEF * value_loss + ENT_COEF * entropy_loss
+
+                optimizer.zero_grad()
+                loss.backward()
+                nn.utils.clip_grad_norm_(policy.parameters(), MAX_GRAD_NORM)
+                optimizer.step()
+
+        update_idx += 1
+
+        # ── Logging + checkpoint ─────────────────────────────────
+        recent = episode_returns[-20:] if episode_returns else [current_return]
+        mean_return = sum(recent) / len(recent)
+        elapsed = time.time() - start_time
+        if update_idx % log_every == 0:
+            print(
+                f"[{task_id}] update={update_idx} step={global_step}/{total_steps} "
+                f"episodes={len(episode_returns)} mean_return(last20)={mean_return:.3f} "
+                f"policy_loss={policy_loss.item():.4f} value_loss={value_loss.item():.4f} "
+                f"entropy={entropy.mean().item():.3f} elapsed={elapsed:.0f}s"
+            )
+
+        if mean_return > best_mean_return and len(episode_returns) >= 5:
+            best_mean_return = mean_return
+            torch.save(
+                {
+                    "model_state": policy.state_dict(),
+                    "task_id": task_id,
+                    "obs_dim": obs_dim,
+                    "action_dims": action_dims,
+                    "best_mean_return": best_mean_return,
+                    "step": global_step,
+                },
+                checkpoint_path,
+            )
+            print(f"  ↳ saved checkpoint (mean_return={best_mean_return:.3f}) → {checkpoint_path}")
+
+    print(f"Done. Best mean return: {best_mean_return:.3f}. Checkpoint: {checkpoint_path}")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--task", required=True, choices=["channel_optimizer", "growth_strategist", "market_dominator"])
+    parser.add_argument("--total-steps", type=int, default=200000)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--checkpoint-dir", default="checkpoints")
+    args = parser.parse_args()
+    train(
+        task_id=args.task,
+        total_steps=args.total_steps,
+        seed=args.seed,
+        checkpoint_dir=args.checkpoint_dir,
+    )
+
+
+if __name__ == "__main__":
+    main()

server/app.py

@@ -19,9 +19,16 @@ from models import GTMAction, GTMObservation
 from server.environment import GTMEnvironment
 from server.tasks import TASKS
 from server.simulation import MESSAGING_DIMS
+from ui.dashboard import build_dashboard
 
-# Create the core OpenEnv app (with Gradio web UI at /web when ENABLE_WEB_INTERFACE=true)
-app = create_app(GTMEnvironment, GTMAction, GTMObservation, env_name="gtm_strategy_optimizer")
+# Create the core OpenEnv app (with custom Gradio dashboard at /web)
+app = create_app(
+    GTMEnvironment,
+    GTMAction,
+    GTMObservation,
+    env_name="gtm_strategy_optimizer",
+    gradio_builder=build_dashboard,
+)
 
 
 # ── Root endpoint for HF Spaces ──────────────────────────────────────────
@@ -32,7 +39,7 @@ def root():
     return {
         "name": "GTM Strategy Optimizer",
         "status": "running",
-        "endpoints": ["/tasks", "/reset", "/step", "/state", "/baseline", "/grader", "/health", "/docs"],
+        "endpoints": ["/tasks", "/reset", "/step", "/state", "/baseline", "/grader", "/infer", "/health", "/docs"],
     }
 
 
@@ -187,3 +194,34 @@ def run_baseline() -> BaselineResponse:
         scores=scores,
         message="Baseline (equal-allocation heuristic) scores for all tasks",
     )
+
+
+# ── RL inference endpoint ──────────────────────────────────────────────────
+
+
+class InferRequest(BaseModel):
+    task_id: str
+    seed: Optional[int] = None
+
+
+class InferResponse(BaseModel):
+    task_id: str
+    checkpoint_loaded: bool
+    grader_score: Optional[float]
+    total_revenue: float
+    total_conversions: int
+    average_cac: float
+    brand_score: float
+    actions: list[dict]
+    message: str
+
+
+@app.post("/infer")
+def run_infer(req: InferRequest) -> InferResponse:
+    """Run a trained RL policy on a task and return the action trajectory."""
+    if req.task_id not in TASKS:
+        raise HTTPException(status_code=400, detail=f"Unknown task_id: {req.task_id}")
+    from rl.infer import run_inference
+
+    result = run_inference(req.task_id, seed=req.seed)
+    return InferResponse(**result)

ui/dashboard.py

@@ -0,0 +1,366 @@
+"""Custom Gradio dashboard for the GTM Strategy Optimizer.
+
+Mounted into the OpenEnv app via `create_app(..., gradio_builder=build_dashboard)`.
+Shows a side-by-side comparison of three strategies: trained RL, equal-allocation
+heuristic, and uniform random — over the same task and seed.
+"""
+
+from __future__ import annotations
+
+import os
+import random as _random
+import sys
+from typing import Any, Callable, Dict, List, Optional
+
+import gradio as gr
+import pandas as pd
+import plotly.graph_objects as go
+
+# Make repo root importable when this module is loaded by uvicorn from /app
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from models import GTMAction
+from rl.infer import run_inference
+from server.environment import GTMEnvironment
+from server.simulation import MESSAGING_DIMS
+from server.tasks import TASKS, get_task
+
+# ── Color palette ──────────────────────────────────────────────────────────
+
+RL_COLOR = "#7c3aed"        # purple
+HEUR_COLOR = "#3b82f6"      # blue
+RAND_COLOR = "#94a3b8"      # slate
+CHANNEL_PALETTE = [
+    "#7c3aed", "#3b82f6", "#10b981", "#f59e0b",
+    "#ef4444", "#06b6d4", "#ec4899",
+]
+
+# ── Strategy runners ───────────────────────────────────────────────────────
+
+
+def _equal_action(task) -> GTMAction:
+    return GTMAction(
+        budget_allocation={c.name: 1.0 / len(task.channels) for c in task.channels},
+        segment_targeting={s.name: 1.0 / len(task.segments) for s in task.segments},
+        messaging={d: 1.0 / len(MESSAGING_DIMS) for d in MESSAGING_DIMS},
+    )
+
+
+def _random_action(task, rng: _random.Random) -> GTMAction:
+    def _simplex(n: int) -> List[float]:
+        xs = [rng.random() + 1e-6 for _ in range(n)]
+        s = sum(xs)
+        return [x / s for x in xs]
+
+    bs = _simplex(len(task.channels))
+    ss = _simplex(len(task.segments))
+    ms = _simplex(len(MESSAGING_DIMS))
+    return GTMAction(
+        budget_allocation={c.name: bs[i] for i, c in enumerate(task.channels)},
+        segment_targeting={s.name: ss[i] for i, s in enumerate(task.segments)},
+        messaging={d: ms[i] for i, d in enumerate(MESSAGING_DIMS)},
+    )
+
+
+def _run_with_action_fn(task_id: str, action_fn: Callable, seed: int) -> Dict[str, Any]:
+    """Run one episode driven by action_fn(task) -> GTMAction. Returns trajectory."""
+    env = GTMEnvironment()
+    task = get_task(task_id)
+    obs = env.reset(task_id=task_id, seed=seed)
+    weeks: List[Dict[str, Any]] = []
+    while not obs.done:
+        action = action_fn(task)
+        obs = env.step(action)
+        weeks.append({
+            "week": obs.week,
+            "total_revenue": float(obs.total_revenue),
+            "brand_score": float(obs.brand_score),
+            "budget_allocation": dict(action.budget_allocation),
+            "reward": float(obs.reward) if obs.reward is not None else 0.0,
+        })
+    grader_score = env.get_grader_score(env.state.episode_id)
+    return {
+        "weeks": weeks,
+        "grader_score": grader_score,
+        "total_revenue": float(obs.total_revenue),
+        "total_conversions": int(obs.total_conversions),
+        "brand_score": float(obs.brand_score),
+    }
+
+
+def run_heuristic(task_id: str, seed: int) -> Dict[str, Any]:
+    return _run_with_action_fn(task_id, _equal_action, seed)
+
+
+def run_random(task_id: str, seed: int) -> Dict[str, Any]:
+    rng = _random.Random(seed)
+    return _run_with_action_fn(task_id, lambda t: _random_action(t, rng), seed)
+
+
+def run_trained_rl(task_id: str, seed: int) -> Dict[str, Any]:
+    """Wraps rl.infer.run_inference and normalizes its output to the common shape."""
+    result = run_inference(task_id, seed=seed)
+    weeks = [
+        {
+            "week": a["week"],
+            "total_revenue": float(a["total_revenue"]),
+            "brand_score": float(a["brand_score"]),
+            "budget_allocation": a["budget_allocation"],
+            "reward": float(a["weekly_reward"]) if a["weekly_reward"] is not None else 0.0,
+        }
+        for a in result["actions"]
+    ]
+    return {
+        "weeks": weeks,
+        "grader_score": result["grader_score"],
+        "total_revenue": float(result["total_revenue"]),
+        "total_conversions": int(result["total_conversions"]),
+        "brand_score": float(result["brand_score"]),
+        "checkpoint_loaded": result["checkpoint_loaded"],
+    }
+
+
+# ── Plot builders ──────────────────────────────────────────────────────────
+
+
+def build_revenue_plot(rl: Dict, heur: Dict, rand: Dict) -> go.Figure:
+    fig = go.Figure()
+    for label, data, color, dash in [
+        ("🤖 Trained RL", rl, RL_COLOR, "solid"),
+        ("📊 Heuristic", heur, HEUR_COLOR, "dash"),
+        ("🎲 Random", rand, RAND_COLOR, "dot"),
+    ]:
+        weeks = data["weeks"]
+        fig.add_trace(go.Scatter(
+            x=[w["week"] for w in weeks],
+            y=[w["total_revenue"] for w in weeks],
+            mode="lines+markers",
+            name=label,
+            line=dict(color=color, width=3, dash=dash),
+            marker=dict(size=7),
+            hovertemplate=f"<b>{label}</b><br>Week %{{x}}<br>Revenue $%{{y:,.0f}}<extra></extra>",
+        ))
+    fig.update_layout(
+        title="Cumulative Revenue Over Weeks",
+        xaxis_title="Week",
+        yaxis_title="Revenue ($)",
+        template="plotly_white",
+        height=420,
+        hovermode="x unified",
+        legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1),
+        margin=dict(l=60, r=20, t=80, b=50),
+    )
+    return fig
+
+
+def build_budget_plot(rl: Dict) -> go.Figure:
+    weeks_data = rl["weeks"]
+    if not weeks_data:
+        return go.Figure()
+    weeks = [w["week"] for w in weeks_data]
+    channels = list(weeks_data[0]["budget_allocation"].keys())
+    fig = go.Figure()
+    for i, ch in enumerate(channels):
+        fig.add_trace(go.Scatter(
+            x=weeks,
+            y=[w["budget_allocation"].get(ch, 0.0) for w in weeks_data],
+            mode="lines",
+            name=ch,
+            stackgroup="one",
+            line=dict(width=0.5, color=CHANNEL_PALETTE[i % len(CHANNEL_PALETTE)]),
+            fillcolor=CHANNEL_PALETTE[i % len(CHANNEL_PALETTE)],
+            hovertemplate=f"<b>{ch}</b><br>Week %{{x}}<br>%{{y:.0%}}<extra></extra>",
+        ))
+    fig.update_layout(
+        title="RL Budget Allocation per Week",
+        xaxis_title="Week",
+        yaxis_title="Fraction of weekly budget",
+        yaxis_tickformat=".0%",
+        template="plotly_white",
+        height=380,
+        hovermode="x unified",
+        legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1),
+        margin=dict(l=60, r=20, t=80, b=50),
+    )
+    return fig
+
+
+def build_brand_plot(rl: Dict, heur: Dict, rand: Dict) -> go.Figure:
+    fig = go.Figure()
+    for label, data, color, dash in [
+        ("🤖 Trained RL", rl, RL_COLOR, "solid"),
+        ("📊 Heuristic", heur, HEUR_COLOR, "dash"),
+        ("🎲 Random", rand, RAND_COLOR, "dot"),
+    ]:
+        weeks = data["weeks"]
+        fig.add_trace(go.Scatter(
+            x=[w["week"] for w in weeks],
+            y=[w["brand_score"] for w in weeks],
+            mode="lines+markers",
+            name=label,
+            line=dict(color=color, width=3, dash=dash),
+            marker=dict(size=7),
+            hovertemplate=f"<b>{label}</b><br>Week %{{x}}<br>Brand %{{y:.0f}}/100<extra></extra>",
+        ))
+    fig.update_layout(
+        title="Brand Health Over Weeks",
+        xaxis_title="Week",
+        yaxis_title="Brand Score (0-100)",
+        template="plotly_white",
+        yaxis=dict(range=[0, 100]),
+        height=380,
+        hovermode="x unified",
+        legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1),
+        margin=dict(l=60, r=20, t=80, b=50),
+    )
+    return fig
+
+
+def build_action_table(rl: Dict) -> pd.DataFrame:
+    rows = []
+    for w in rl["weeks"]:
+        row: Dict[str, Any] = {"Week": w["week"]}
+        for ch, frac in w["budget_allocation"].items():
+            row[ch] = round(frac, 3)
+        row["Brand"] = round(w["brand_score"], 1)
+        row["Reward"] = round(w["reward"], 3)
+        rows.append(row)
+    return pd.DataFrame(rows)
+
+
+# ── Score card markdown ────────────────────────────────────────────────────
+
+
+def _score_card(label: str, emoji: str, color_hex: str, result: Dict, is_winner: bool = False) -> str:
+    score = result.get("grader_score")
+    score_str = f"{score:.3f}" if score is not None else "—"
+    crown = " 👑" if is_winner else ""
+    return (
+        f"<div style='border-left:6px solid {color_hex};padding:14px 18px;background:#f8fafc;border-radius:8px;'>"
+        f"<div style='font-size:14px;color:#64748b;font-weight:600;text-transform:uppercase;letter-spacing:0.05em'>{emoji} {label}{crown}</div>"
+        f"<div style='font-size:36px;font-weight:700;color:#0f172a;margin:4px 0'>{score_str}</div>"
+        f"<div style='font-size:13px;color:#475569'>"
+        f"💵 ${result['total_revenue']:,.0f} &nbsp;&nbsp; "
+        f"🎯 {result['total_conversions']} convs &nbsp;&nbsp; "
+        f"⭐ Brand {result['brand_score']:.0f}/100"
+        f"</div></div>"
+    )
+
+
+# ── Main callback ──────────────────────────────────────────────────────────
+
+
+def run_comparison(task_id: str, seed_value: float):
+    seed = int(seed_value) if seed_value is not None else 42
+
+    rl = run_trained_rl(task_id, seed)
+    heur = run_heuristic(task_id, seed)
+    rand = run_random(task_id, seed)
+
+    # Pick a winner (highest grader score) for the crown emoji
+    scores = {
+        "rl": rl.get("grader_score") or -1,
+        "heur": heur.get("grader_score") or -1,
+        "rand": rand.get("grader_score") or -1,
+    }
+    winner = max(scores, key=scores.get)
+
+    revenue_fig = build_revenue_plot(rl, heur, rand)
+    budget_fig = build_budget_plot(rl)
+    brand_fig = build_brand_plot(rl, heur, rand)
+    action_df = build_action_table(rl)
+
+    rl_card = _score_card("Trained RL", "🤖", RL_COLOR, rl, is_winner=(winner == "rl"))
+    heur_card = _score_card("Heuristic", "📊", HEUR_COLOR, heur, is_winner=(winner == "heur"))
+    rand_card = _score_card("Random", "🎲", RAND_COLOR, rand, is_winner=(winner == "rand"))
+
+    ckpt_warning = ""
+    if not rl.get("checkpoint_loaded"):
+        ckpt_warning = (
+            "<div style='padding:10px 14px;background:#fef3c7;border-left:4px solid #f59e0b;border-radius:6px;color:#78350f;font-size:13px'>"
+            "⚠️ No trained checkpoint found for this task — RL is using a random-init policy. "
+            "Run <code>python -m rl.train --task " + task_id + "</code> locally and commit "
+            "<code>checkpoints/" + task_id + ".pt</code>."
+            "</div>"
+        )
+
+    return revenue_fig, budget_fig, brand_fig, action_df, rl_card, heur_card, rand_card, ckpt_warning
+
+
+# ── Gradio Blocks builder (called by openenv) ──────────────────────────────
+
+
+def build_dashboard(*_args, **_kwargs) -> gr.Blocks:
+    """Entry point used as `gradio_builder` in `create_app`.
+
+    The OpenEnv hook passes (web_manager, action_fields, metadata, is_chat_env,
+    title, quick_start_md) which we ignore — we build a fully custom layout.
+    """
+    with gr.Blocks(title="GTM Strategy Optimizer") as demo:
+        gr.Markdown(
+            "# 📈 GTM Strategy Optimizer\n"
+            "### RL agent vs heuristic vs random — head-to-head on marketing budget allocation\n"
+            "A custom PPO policy trained on a Go-To-Market simulator allocates marketing budget "
+            "across channels, segments, and messaging dimensions. Pick a task and run all three "
+            "strategies on the same seed to see how the trained policy compares."
+        )
+
+        with gr.Row():
+            task_dd = gr.Dropdown(
+                choices=list(TASKS.keys()),
+                value="channel_optimizer",
+                label="Task",
+                info="Increasing difficulty: more channels, segments, and adversarial dynamics",
+                scale=3,
+            )
+            seed_in = gr.Number(value=42, label="Seed", precision=0, scale=1)
+            run_btn = gr.Button("▶ Run Comparison", variant="primary", scale=1, size="lg")
+
+        ckpt_warning_html = gr.HTML(value="")
+
+        with gr.Row():
+            rl_card_md = gr.HTML(_score_card("Trained RL", "🤖", RL_COLOR, _empty_result()))
+            heur_card_md = gr.HTML(_score_card("Heuristic", "📊", HEUR_COLOR, _empty_result()))
+            rand_card_md = gr.HTML(_score_card("Random", "🎲", RAND_COLOR, _empty_result()))
+
+        revenue_plot = gr.Plot(label="Cumulative Revenue", show_label=False)
+
+        with gr.Row():
+            budget_plot = gr.Plot(label="Budget Allocation (RL)", show_label=False)
+            brand_plot = gr.Plot(label="Brand Health", show_label=False)
+
+        gr.Markdown("### 📋 Per-week actions (Trained RL)")
+        action_df = gr.Dataframe(
+            label="Weekly action log",
+            interactive=False,
+            wrap=True,
+        )
+
+        gr.Markdown(
+            "---\n"
+            "**How it works.** The trained policy is a small PyTorch actor-critic with "
+            "Dirichlet heads for budget/segment/messaging simplices and categorical heads for "
+            "experiments and pricing actions. Trained with a custom lightweight PPO loop "
+            "(`rl/train.py`) against the simulator's per-step reward "
+            "(`server/environment.py:_compute_reward`)."
+        )
+
+        run_btn.click(
+            fn=run_comparison,
+            inputs=[task_dd, seed_in],
+            outputs=[
+                revenue_plot, budget_plot, brand_plot, action_df,
+                rl_card_md, heur_card_md, rand_card_md, ckpt_warning_html,
+            ],
+        )
+
+    return demo
+
+
+def _empty_result() -> Dict[str, Any]:
+    return {
+        "grader_score": None,
+        "total_revenue": 0.0,
+        "total_conversions": 0,
+        "brand_score": 0.0,
+    }