Merge pull request #1 from daslabhq/feat/predict

feat: tier 4 — Predictor + DreamerSolver + worldmodel eval (v0)

Author: mirkokiefer

README.md

@@ -58,22 +58,25 @@ npx scenegrad view ./traces           # bulk view of every JSONL in ./traces
 npx scenegrad view ./traces/run.jsonl # single-trace view of one file
 ```
 
+> **v0.0.1 note:** the CLI requires [Bun](https://bun.sh) on your PATH (the bin scripts are TypeScript with a `#!/usr/bin/env bun` shebang). Node-compatible bin compilation lands in v0.0.2.
+
 That's it. Drop in, run, view your runs in a browser. No goal to design. No assertions to write. No restructuring of your agent.
 
 When you want more — add a `snapshot()` to capture world state between calls. Add a `goal()` of assertions to measure gap closure. Each level is opt-in.
 
 ---
 
-## The four-tier ladder
+## The five-tier ladder
 
 | Tier | What you write | What you get |
 |---|---|---|
 | **0 — trace** | `trace.start()` + one hook | Tool-call timeline, scrubbable. Replaces logs. |
 | **1 — + snapshot** | Add `snapshot: () => fetchWorld()` | World deltas between calls. See *what changed*, not just what was called. |
 | **2 — + goal** | Add `goal: (s) => [...assertions]` | Gap-closure curve, drift detection, runtime `status()` for agent guidance. |
 | **3 — + solver** | Use `defineEnv` + `LLMSolver` / `GreedySolver` | scenegrad drives the loop — for benches, comparison, leaderboards. |
+| **4 — + predictor** | Add `Predictor` + `DreamerSolver` | Plan in imagination, act in the env. World-model accuracy as a measurable benchmark metric. |
 
-The same trajectory format flows through all four tiers. You can adopt at tier 0, level up months later as you understand your agent's failure modes.
+The same trajectory format flows through all five tiers. You can adopt at tier 0, level up months later as you understand your agent's failure modes.
 
 ---
 
@@ -152,6 +155,33 @@ Both solvers produce the same `SolveResult` shape. Compare them on the same env
 
 ---
 
+## Tier 4 — predictor + dreamer, plan in imagination
+
+When env-side `simulate()` isn't available (real APIs, prod databases, irreversible actions), you can't enumerate-and-pick. Tier 4 swaps the simulator for a **Predictor** — a learned-or-LLM model of `predict(scene, action) → consequence`. `DreamerSolver` calls the predictor on each candidate, picks the one whose imagined outcome closes the most distance, and commits a single action to the real env.
+
+```ts
+import { defineEnv, LLMPredictor, DreamerSolver, evalWorldModel } from "scenegrad";
+
+const env       = defineEnv({ /* same as tier 3 */ });
+const predictor = new LLMPredictor({ model: "claude-haiku-4-5" });
+
+// Plans in the predictor, commits one action at a time.
+await new DreamerSolver({ predictor, lookahead: 1 }).solve(env, "default");
+
+// Measure how good the predictor actually is — predicted vs actual scene_after.
+const metrics = await evalWorldModel({
+  env, predictor,
+  tasks: [{ taskId: "default", actions: [/* known action sequence */] }],
+});
+// → outcome_acc, scene_match, delta_match, avg_confidence, ece (calibration)
+```
+
+`Predictor` is a one-method interface. v0 ships `LLMPredictor` as a placeholder; future predictors (kNN over a trace store, distilled-from-traces, fully learned) drop in via the same API — `new DreamerSolver({ predictor })` doesn't change.
+
+The `evalWorldModel` metric is the world-model-accuracy benchmark a predictor is judged against. Ship a predictor → score it on any tier-3 env → publish the leaderboard column.
+
+---
+
 ## Why this isn't another agent framework
 
 If you already use… | scenegrad adds…

examples/dreamer-inbox.ts

@@ -0,0 +1,123 @@
+/**
+ * Dreamer-inbox — same inbox onboarding task, three runs side-by-side.
+ *
+ *   1. LLMSolver          → acts directly in the env, current behavior
+ *   2. DreamerSolver(LLM) → predicts each candidate's outcome, picks best, then acts
+ *   3. evalWorldModel     → measures how accurate the LLM-as-predictor actually is
+ *
+ * The point of v0: ship the surface end-to-end with the LLM as a placeholder
+ * predictor. Future versions swap the predictor (kNN over a trace store,
+ * fine-tuned, learned) without changing the example or the eval.
+ *
+ * Run: ANTHROPIC_API_KEY=... bun examples/dreamer-inbox.ts
+ */
+
+import {
+  defineEnv,
+  LLMSolver,
+  DreamerSolver,
+  LLMPredictor,
+  evalWorldModel,
+  formatMetrics,
+} from "scenegrad";
+
+type Status = "unread" | "archived" | "flagged" | "replied";
+type Mail   = { id: number; from: string; subject: string; status: Status };
+
+const isImportantSender = (from: string) =>
+  /boss@|ceo@|vip@/.test(from) || /calendar@|meeting@/.test(from);
+
+const inbox = defineEnv({
+  init: () => ({
+    messages: [
+      { id: 1, from: "boss@co",     subject: "Q3 plan?",     status: "unread" as Status },
+      { id: 2, from: "spam@x.com",  subject: "YOU WON!!!",   status: "unread" as Status },
+      { id: 3, from: "calendar@co", subject: "Mtg 2pm tmrw", status: "unread" as Status },
+    ] as Mail[],
+  }),
+
+  goal: () => [
+    {
+      name: "no unread messages",
+      check: (s) => {
+        const unread = s.messages.filter((m: Mail) => m.status === "unread").length;
+        return { satisfied: unread === 0, gap: unread };
+      },
+    },
+    {
+      name: "important mail not archived (flag or reply only)",
+      check: (s) => {
+        const wrongly = s.messages.filter((m: Mail) =>
+          m.status === "archived" && isImportantSender(m.from));
+        return { satisfied: wrongly.length === 0, gap: wrongly.length, weight: 5 };
+      },
+    },
+  ],
+
+  tools: (s) => s.messages
+    .filter((m: Mail) => m.status === "unread")
+    .flatMap((m: Mail) => [
+      { name: "archive", args: { id: m.id } },
+      { name: "flag",    args: { id: m.id } },
+      { name: "reply",   args: { id: m.id } },
+    ]),
+
+  step: (s, t) => ({
+    messages: s.messages.map((m: Mail) =>
+      m.id === (t.args as any).id ? { ...m, status: t.name as Status } : m
+    ),
+  }),
+
+  describeTask: (s) => {
+    const unread = s.messages.filter((m: Mail) => m.status === "unread");
+    return [
+      "INBOX:",
+      ...s.messages.map((m: Mail) =>
+        `  [${m.status.padEnd(8)}] #${m.id} from ${m.from}: "${m.subject}"`),
+      "",
+      `${unread.length} unread. archive (delete), flag (handle later), or reply.`,
+      "Important senders (boss@, calendar@) must NOT be archived.",
+    ].join("\n");
+  },
+});
+
+function dump(label: string, r: { success: boolean; steps: number; d_initial: number; d_final: number; duration_ms: number; trajectory: any[] }) {
+  console.log(`\n=== ${label} ===`);
+  console.log(`success=${r.success}  steps=${r.steps}  d:${r.d_initial}→${r.d_final}  ${r.duration_ms}ms`);
+  for (const t of r.trajectory) {
+    const tool = t.tool ? `${t.tool.name}(${JSON.stringify(t.tool.args)})` : "(none)";
+    const pred = t.predicted_delta !== undefined ? `  pred-Δ=${t.predicted_delta}` : "";
+    console.log(`  #${t.step} ${tool}  d:${t.d_before}→${t.d_after}  Δ=${t.delta}${pred}`);
+    if (t.reasoning) console.log(`     "${t.reasoning}"`);
+  }
+}
+
+if (!process.env.ANTHROPIC_API_KEY) {
+  console.error("ANTHROPIC_API_KEY not set");
+  process.exit(1);
+}
+
+// 1. baseline: LLMSolver acts directly
+const llm = new LLMSolver({
+  model: "claude-haiku-4-5",
+  describeTask: () => (inbox as any).describeTask?.() ?? "",
+});
+dump("LLMSolver — acts directly", await llm.solve(inbox, "default", { maxSteps: 10 }));
+
+// 2. tier 4: DreamerSolver predicts each candidate, picks best, then commits
+const predictor = new LLMPredictor({ model: "claude-haiku-4-5" });
+const dreamer   = new DreamerSolver({ predictor, lookahead: 1 });
+dump("DreamerSolver(LLM) — plans in predictor, acts once", await dreamer.solve(inbox, "default", { maxSteps: 10 }));
+
+// 3. measure how good the predictor actually is. Use the LLMSolver run as
+// the action sequence, then replay through env+predictor side-by-side.
+const baseRun  = await llm.solve(inbox, "default", { maxSteps: 10 });
+const actions  = baseRun.trajectory.map(s => s.tool).filter((t): t is any => t !== null);
+const metrics  = await evalWorldModel({
+  env:       inbox,
+  predictor,
+  tasks:     [{ taskId: "default", actions }],
+});
+
+console.log("\n=== world-model accuracy on this task ===");
+console.log(formatMetrics(metrics, { perStep: true }));

package.json

@@ -14,7 +14,11 @@
     "./trace": "./src/trace.ts",
     "./solver": "./src/solver.ts",
     "./greedy": "./src/greedy.ts",
-    "./llm": "./src/llm.ts"
+    "./llm": "./src/llm.ts",
+    "./predict": "./src/predict.ts",
+    "./predict-llm": "./src/predict-llm.ts",
+    "./dreamer": "./src/dreamer.ts",
+    "./eval/worldmodel": "./src/eval/worldmodel.ts"
   },
   "bin": {
     "scenegrad": "./bin/scenegrad.ts",

src/dreamer.ts

@@ -0,0 +1,153 @@
+/**
+ * DreamerSolver — plan in the predictor, act in the env.
+ *
+ * Each turn:
+ *   1. Enumerate tools available now.
+ *   2. For each candidate tool, ask the Predictor what happens.
+ *      (Optional lookahead k>1 rolls out k steps deep recursively.)
+ *   3. Pick the action whose imagined trajectory closes the most distance.
+ *   4. Commit that single action to the real env.
+ *   5. Repeat.
+ *
+ * v0 is the plain "predict-and-pick" base case — like GreedySolver but
+ * using a learned dynamics function instead of env.simulate(). That's
+ * the unlock for envs where you can't write a simulator (real APIs,
+ * databases, integrations).
+ *
+ * Future versions add value-net guided rollouts and policy distillation
+ * from imagined trajectories — the full Dreamer arc.
+ */
+
+import type { SceneGradEnv, ToolCall, Goal } from "./env.js";
+import { distance, checkAll } from "./env.js";
+import type { Solver, SolveResult, SolverOpts, TrajectoryStep } from "./solver.js";
+import type { Predictor } from "./predict.js";
+
+export interface DreamerSolverOpts extends SolverOpts {
+  predictor: Predictor<any, any>;
+  /** Plan depth — how many predictor calls to chain per candidate.
+   *  v0 ships with k=1 working well; k>1 is exponential and best for small toolsets. */
+  lookahead?: number;
+  /** Optional beam width for k>1. Defaults to all candidates. */
+  beam?:      number;
+  /** Filter low-confidence predictions out of planning. Default 0 — keep all. */
+  minConfidence?: number;
+}
+
+export class DreamerSolver<S, T extends ToolCall = ToolCall> implements Solver<S, T> {
+  readonly name: string;
+  private predictor:   Predictor<S, T>;
+  private lookahead:   number;
+  private beam:        number;
+  private minConf:     number;
+
+  constructor(opts: DreamerSolverOpts) {
+    this.predictor = opts.predictor as Predictor<S, T>;
+    this.lookahead = Math.max(1, opts.lookahead ?? 1);
+    this.beam      = Math.max(1, opts.beam      ?? Infinity);
+    this.minConf   = Math.max(0, opts.minConfidence ?? 0);
+    this.name      = `dreamer(${this.predictor.name},k=${this.lookahead})`;
+  }
+
+  async solve(env: SceneGradEnv<S, T>, taskId: string, opts?: SolverOpts): Promise<SolveResult<T>> {
+    const t_start  = Date.now();
+    const maxSteps = opts?.maxSteps ?? 30;
+
+    env.reset(taskId);
+    const trajectory: TrajectoryStep<T>[] = [];
+    const d_initial  = distance(env.scene(), env.goal());
+
+    for (let step = 0; step < maxSteps; step++) {
+      if (env.done()) break;
+
+      const goal     = env.goal();
+      const d_before = distance(env.scene(), goal);
+      const tools    = env.tools();
+      if (tools.length === 0) break;
+
+      const best = await this.pickBest(env.scene(), tools, goal);
+
+      if (!best || best.predicted_d_after >= d_before) {
+        trajectory.push({
+          step,
+          tool:             null,
+          d_before,
+          d_after:          d_before,
+          delta:            0,
+          predicted_delta:  best ? d_before - best.predicted_d_after : 0,
+          ok:               false,
+          error:            "no predicted action reduces distance (stuck or low-confidence)",
+          assertions_after: checkAll(env.scene(), goal),
+          ts_ms:            Date.now() - t_start,
+        });
+        break;
+      }
+
+      // Commit the chosen action to the real env.
+      const result  = env.step(best.tool);
+      const d_after = result.distance_after ?? distance(env.scene(), goal);
+
+      trajectory.push({
+        step,
+        tool:             best.tool,
+        d_before,
+        d_after,
+        delta:            d_before - d_after,
+        predicted_delta:  d_before - best.predicted_d_after,
+        reasoning:        best.reasoning,
+        ok:               result.ok,
+        error:            result.error,
+        assertions_after: checkAll(env.scene(), goal),
+        ts_ms:            Date.now() - t_start,
+      });
+    }
+
+    const d_final = distance(env.scene(), env.goal());
+    return {
+      task_id:     taskId,
+      solver:      this.name,
+      success:     env.done(),
+      steps:       trajectory.length,
+      d_initial,
+      d_final,
+      duration_ms: Date.now() - t_start,
+      trajectory,
+    };
+  }
+
+  /**
+   * Pick the tool whose imagined rollout reaches the lowest distance.
+   * Lookahead=1: one predictor call per tool. Lookahead=k: chains k.
+   */
+  private async pickBest(
+    scene: S,
+    tools: T[],
+    goal:  Goal<S>,
+  ): Promise<{ tool: T; predicted_d_after: number; reasoning?: string } | null> {
+    const candidates = await Promise.all(
+      tools.map(async (tool) => {
+        const c = await this.predictor.predict(scene, tool);
+        if (c.confidence < this.minConf) return null;
+        if (!c.outcome.ok) return null;
+        const d_after = await this.rolloutTail(c.scene_after, goal, this.lookahead - 1);
+        return { tool, predicted_d_after: d_after, reasoning: c.reasoning };
+      }),
+    );
+    const valid = candidates.filter((x): x is NonNullable<typeof x> => x !== null);
+    if (valid.length === 0) return null;
+    valid.sort((a, b) => a.predicted_d_after - b.predicted_d_after);
+    return valid[0]!;
+  }
+
+  /**
+   * Roll out depth-k more predictor steps after the head action, returning
+   * the best (lowest) distance reachable. k=0 returns the head's distance.
+   */
+  private async rolloutTail(scene: S, goal: Goal<S>, k: number): Promise<number> {
+    if (k <= 0) return distance(scene, goal);
+    // Synthetic env-less rollout: we don't have env.tools() for predicted
+    // future scenes. v0 punts: trust the head prediction. v0.1 will pass
+    // a domain-provided `toolsAt(scene)` for true k-step lookahead.
+    return distance(scene, goal);
+  }
+}