awpredict 0.1.2__tar.gz → 0.1.4__tar.gz

{awpredict-0.1.2 → awpredict-0.1.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: awpredict
-Version: 0.1.2
+Version: 0.1.4
 Summary: A small, dependency-light world-model package: LeWM JEPA + MLP engines, environment adapters
 License: MIT
 Requires-Python: >=3.10

{awpredict-0.1.2 → awpredict-0.1.4}/README.md

@@ -4,7 +4,7 @@
 
 **[Docs](https://aitherium.github.io/awpredict/)**  ·  [Source](https://github.com/Aitherium/awpredict)  ·  `pip install git+https://github.com/Aitherium/awpredict.git`  ·  [The Aither World](https://aitherium.github.io/)
 
-> **The Aither World** is an operating system for agents — a Linux you can hand to one, the runtimes it works in, and the tools it works with. [awnix](https://github.com/Aitherium/awnix) is the Linux underneath it; **awpredict** is one of its 66 bricks — each installs on its own, runs offline, and needs no account.
+> **The Aither World** is an operating system for agents — a Linux you can hand to one, the runtimes it works in, and the tools it works with. [awnix](https://github.com/Aitherium/awnix) is the Linux underneath it; **awpredict** is one of its 69 bricks — each installs on its own, runs offline, and needs no account.
 >
 > **Start here:** Wrap an environment you already have and ask it what happens next.
 
@@ -63,6 +63,26 @@ pred = model.predict(z, action)
 print(model.surprise(state, action, next_state))
 ```
 
+## See what a model imagines
+
+```bash
+awpredict viz --engine lewm --demo --out dream.html   # also: mlp, concept
+```
+
+`--demo` trains a small engine on a synthetic world (a moving block for `lewm`, a
+drifting point or vector for `mlp`/`concept`) and writes one self-contained HTML
+page: the open-loop dream next to the real sequence (a decoded filmstrip when the
+engine can decode to a grid, a latent heatmap otherwise), the latent trajectory in
+2D, and every error beside the copy-last baseline (predict that nothing changes).
+On your own trained engine:
+
+```python
+from awpredict.viz import dream_report, render_html
+
+report = dream_report(model, observations, actions, steps=6)
+open("dream.html", "w", encoding="utf-8").write(render_html(report))
+```
+
 ## Why this exists
 
 The goal is to help someone bootstrap their own world model — a working
@@ -136,6 +156,8 @@ Each installs on its own, works offline, and needs no account.
 | [awtoll](https://github.com/Aitherium/awtoll) | that your tooling is saving you context | the measured token cost of each tool call, and what the alternative cost |
 | [awseal](https://github.com/Aitherium/awseal) | that the artifact came from who you think | an Ed25519 seal — the key that verifies is not the key that forges |
 | [awshare](https://github.com/Aitherium/awshare) | that the download is intact | content-addressed bundles, verified on fetch |
+| [awsuite](https://github.com/Aitherium/awsuite) | that an agent holding your mailbox will not send on its own | every send, draft, upload and create returns a dry-run until confirm is true |
+| [awpermit](https://github.com/Aitherium/awpermit) | that an admitted agent will only do what its owner meant | an Ed25519 permit naming the agent, the principal, the door, the actions and the cap |
 | [awnest](https://github.com/Aitherium/awnest) | that there is a person on the other end | a verdict with evidence, where "we could not tell" is not "yes" |
 | [awrena](https://github.com/Aitherium/awrena) | a leaderboard someone can edit, and votes nobody counted | a scored duel with both answers kept, and a result bound to them |
 | [awnboard](https://github.com/Aitherium/awnboard) | a share link anyone who sees it can use | an invitation addressed to one person, for one gate, revocable |
@@ -145,6 +167,7 @@ Each installs on its own, works offline, and needs no account.
 | [awrelay](https://github.com/Aitherium/awrelay) | a SaaS in the middle of your agents | findings, alerts and coordination over your own transport |
 | [awask](https://github.com/Aitherium/awask) | that anyone read the paragraph where you asked | the ask itself, with a button that steers the session that raised it |
 | [awmail](https://github.com/Aitherium/awmail) | a mailbox somebody else can read | mail your agents send and receive over your own server |
+| [awdiscord](https://github.com/Aitherium/awdiscord) | a self-bot on someone's account, or a bot that acts without asking | slash commands, rate limits, and approval buttons only your operators can press |
 | [awswarm](https://github.com/Aitherium/awswarm) | that a model either fits your GPU or it doesn't run at all | a placement plan and an acquisition-probability estimate before you spend on a run |
 | [awfind](https://github.com/Aitherium/awfind) | one vendor's idea of the web | results from whichever providers you configured |
 | [awbrowse](https://github.com/Aitherium/awbrowse) | that the page said what you were told | the render, the DOM and the requests it made |
@@ -201,10 +224,12 @@ Every repository here is public. Each publishes an `aither-manifest.json` beside
 | [awtoll](https://github.com/Aitherium/awtoll) | What every tool call costs you in context, measured from your own transcripts | [docs](https://aitherium.github.io/awtoll/) |
 | [awseal](https://github.com/Aitherium/awseal) | Sign an artifact so a stranger can verify it | [docs](https://aitherium.github.io/awseal/) |
 | [awshare](https://github.com/Aitherium/awshare) | Publish an artifact and fetch it back verified | [docs](https://aitherium.github.io/awshare/) |
+| [awsuite](https://github.com/Aitherium/awsuite) | Your Google Workspace as agent tools, and no write happens without a yes | — |
 | [awdit](https://github.com/Aitherium/awdit) | An append-only audit trail whose gaps are DETECTABLE | [docs](https://aitherium.github.io/awdit/) |
 | [awbac](https://github.com/Aitherium/awbac) | Role-based access control that fails closed and explains itself | [docs](https://aitherium.github.io/awbac/) |
 | [awiam](https://github.com/Aitherium/awiam) | Who is this caller? A directory and session store that fails honestly | [docs](https://aitherium.github.io/awiam/) |
 | [awtunnel](https://github.com/Aitherium/awtunnel) | Reach a service that has no public address | [docs](https://aitherium.github.io/awtunnel/) |
+| [awpermit](https://github.com/Aitherium/awpermit) | Let an agent through your door with a scope, a spend cap, and someone accountable | — |
 | [awnest](https://github.com/Aitherium/awnest) | Prove there is a human before you let them into the nest | [docs](https://aitherium.github.io/awnest/) |
 | [awrena](https://github.com/Aitherium/awrena) | Put two agents head to head and get a verdict you can check | [docs](https://aitherium.github.io/awrena/) |
 | [awnboard](https://github.com/Aitherium/awnboard) | A front gate you can put in front of anything, and hand someone the key to | [docs](https://aitherium.github.io/awnboard/) |
@@ -214,6 +239,7 @@ Every repository here is public. Each publishes an `aither-manifest.json` beside
 | [awrelay](https://github.com/Aitherium/awrelay) | Portable agent messaging — findings, alerts, coordination | [docs](https://aitherium.github.io/awrelay/) |
 | [awask](https://github.com/Aitherium/awask) | Your agent asks you a question — and acts on your answer | [docs](https://aitherium.github.io/awask/) |
 | [awmail](https://github.com/Aitherium/awmail) | Give an agent an email address — send, and actually receive | [docs](https://aitherium.github.io/awmail/) |
+| [awdiscord](https://github.com/Aitherium/awdiscord) | Put any agent on Discord as a real bot, with an operator in the loop | — |
 | [awnet](https://github.com/Aitherium/awnet) | The agentic web — agents host a mesh, and agents join one | [docs](https://aitherium.github.io/awnet/) |
 | [awswarm](https://github.com/Aitherium/awswarm) | Run one model too big for any single GPU across a pool of small ones | — |
 | [awfind](https://github.com/Aitherium/awfind) | A portable search client — query, results, ranking | [docs](https://aitherium.github.io/awfind/) |
@@ -253,6 +279,8 @@ Every repository here is public. Each publishes an `aither-manifest.json` beside
 | [awsettings](https://github.com/Aitherium/awsettings) | Your agent's permissions and config, following you to the next machine | [docs](https://aitherium.github.io/awsettings/) |
 | [awavatar](https://github.com/Aitherium/awavatar) | One character spec in, a rigged, animated, multi-style avatar pack out | [docs](https://aitherium.github.io/awavatar/) |
 
+**Built on** [REA (Reverse Engineer Anything)](https://github.com/morluto/rea) · [llama.cpp](https://github.com/ggml-org/llama.cpp) · [vLLM](https://github.com/vllm-project/vllm) · [ComfyUI](https://github.com/comfyanonymous/ComfyUI) · [CentOS Stream](https://www.centos.org/centos-stream/) · [Podman](https://github.com/containers/podman) · [Docker](https://github.com/moby/moby) · [LanceDB](https://github.com/lancedb/lancedb) · [WireGuard](https://www.wireguard.com/) · [FFmpeg](https://ffmpeg.org/) · [Blender + Rigify](https://www.blender.org/) · [headroom](https://github.com/headroomlabs-ai/headroom) · [SANA](https://github.com/NVlabs/Sana) · [Hunyuan3D](https://github.com/Tencent-Hunyuan/Hunyuan3D-2.1) · [repowise](https://github.com/repowise-dev/repowise) · [Playwright](https://github.com/microsoft/playwright) · [Chromium](https://www.chromium.org/) · [Next.js](https://github.com/vercel/next.js) · [React](https://github.com/facebook/react).
+
 <div id="aither-constellation" data-self="awpredict"></div>
 <script src="aither-constellation.js"></script>
 

{awpredict-0.1.2 → awpredict-0.1.4}/awpredict/__init__.py

@@ -11,6 +11,9 @@ One package, two engines, N environment adapters:
     is the same engine the solving agent runs, not a cut-down demo of it —
     the goal of this package is to give you everything needed to bootstrap
     your own world model, not a subset of it.
+  * ``awpredict.core.concept.ConceptWorldModel`` — next-concept JEPA over
+    embedding sequences (product/residual VQ concept vocabulary + SIGReg);
+    ``plan`` ranks candidate continuations, e.g. parallel drafts.
   * ``awpredict.core.mlp.MLPWorldModel`` — an embedding-MLP transition
     model (tabular → hybrid → neural).
 
@@ -35,7 +38,7 @@ from awpredict.contracts import (
     conforms,
 )
 
-__version__ = "0.1.2"
+__version__ = "0.1.4"
 
 __all__ = [
     "Attempt",

{awpredict-0.1.2 → awpredict-0.1.4}/awpredict/_doctor.py

@@ -23,7 +23,15 @@ import sys
 #: package cannot read the registry, and a doctor that guessed at the family
 #: would go stale in silence. Regenerate to update.
 SELF = 'awpredict'
-FAMILY = ['awask', 'awavatar', 'awbac', 'awbrain', 'awbrowse', 'awclassify', 'awdecide', 'awdeck', 'awdelphi', 'awdit', 'awembed', 'awevolve', 'awfind', 'awflow', 'awfocus', 'awgit', 'awgraph', 'awgym', 'awiam', 'awkno', 'awm', 'awmail', 'awmine', 'awnboard', 'awnest', 'awnet', 'awnode', 'awpool', 'awprism', 'awprove', 'awreason', 'awrecover', 'awrecurse', 'awrelay', 'awrena', 'awrepl', 'awreport', 'awresearch', 'awrise', 'awrouter', 'awrtifact', 'awrun', 'awscreen', 'awseal', 'awsettings', 'awshare', 'awsprite', 'awstorage', 'awswarm', 'awtax', 'awtoll', 'awtunnel', 'awvision', 'awvoice', 'awwall', 'gawbbonet']
+FAMILY = [
+    'awask', 'awavatar', 'awbac', 'awbrain', 'awbrowse', 'awcams', 'awclassify', 'awdecide',
+    'awdeck', 'awdelphi', 'awdit', 'awembed', 'awevolve', 'awfind', 'awflow', 'awfocus', 'awgit',
+    'awgraph', 'awgym', 'awiam', 'awkno', 'awm', 'awmail', 'awmine', 'awnboard', 'awnest', 'awnet',
+    'awplay', 'awpool', 'awprism', 'awprove', 'awreason', 'awrecover', 'awrecurse', 'awrelay',
+    'awrena', 'awrepl', 'awreport', 'awresearch', 'awrise', 'awrouter', 'awrtifact', 'awrun',
+    'awscope', 'awscreen', 'awseal', 'awsettings', 'awshare', 'awsprite', 'awstorage', 'awsuite',
+    'awswarm', 'awtax', 'awtoll', 'awtunnel', 'awvision', 'awvoice', 'awwall', 'gawbbonet',
+]
 PAIRS_WITH = ['awm']
 
 #: This brick's OWN config, read out of its source at generation time.

{awpredict-0.1.2 → awpredict-0.1.4}/awpredict/cli.py

@@ -2,6 +2,7 @@
 
   engines   which engines can actually construct here, and why not if they cannot
   conforms  does a class satisfy WorldModel or EnvironmentAdapter, and what is missing
+  viz       see what an engine imagines: dream filmstrip/heatmap + error vs copy-last
 
 --json for machine-readable output.
 """
@@ -11,6 +12,7 @@ from __future__ import annotations
 import argparse
 import importlib
 import json
+from pathlib import Path
 from typing import List
 
 from . import contracts
@@ -28,7 +30,8 @@ def cmd_engines(args) -> int:
     missing' sends someone to install what they already have.
     """
     rows = []
-    for name, mod in (("lewm", "awpredict.core.lewm"), ("mlp", "awpredict.core.mlp")):
+    for name, mod in (("lewm", "awpredict.core.lewm"), ("mlp", "awpredict.core.mlp"),
+                      ("concept", "awpredict.core.concept")):
         try:
             m = importlib.import_module(mod)
         except Exception as exc:
@@ -81,6 +84,56 @@ def cmd_conforms(args) -> int:
     return 0 if not missing else 1
 
 
+def cmd_viz(args) -> int:
+    """Write a dream-visualizer HTML page.
+
+    Exit 0 written, 1 the engine produced no rollout (nothing worth a page),
+    2 could not run (torch/numpy missing, no data source) -- never 0 on nothing.
+    """
+    if not args.demo:
+        print("viz: only --demo is wired on the command line; from Python call "
+              "awpredict.viz.dream_report(model, observations, actions) and "
+              "awpredict.viz.render_html(report) on your own trained engine")
+        return 2
+    try:
+        from . import viz, viz_demo
+    except ImportError as exc:
+        print(f"viz: cannot import the visualizer: {exc}")
+        return 2
+    try:
+        model, obs, acts, meta = viz_demo.DEMOS[args.engine](steps=args.steps)
+        report = viz.dream_report(model, obs, acts, steps=args.steps)
+    except viz.VizUnavailableError as exc:
+        print(f"viz: cannot run the {args.engine} demo: {exc}")
+        return 2
+    report["meta"] = meta
+    for name in ("one_step", "rollout"):
+        b = report[name]
+        print(f"  {name:<9} model {_num(b['model_mean'])}  copy-last "
+              f"{_num(b['copy_last_mean'])}  "
+              f"{'beats' if b['beats_copy_last'] else 'does NOT beat'} copy-last")
+    if not report["ok"]:
+        print(f"viz: {args.engine} produced no rollout: {report['reason']}")
+        return 1
+    out = Path(args.out or f"awpredict-dream-{args.engine}.html")
+    try:
+        out.parent.mkdir(parents=True, exist_ok=True)
+        out.write_text(viz.render_html(report), encoding="utf-8")
+    except OSError as exc:
+        print(f"viz: cannot write {out}: {exc}")
+        return 2
+    if args.json:
+        print(json.dumps({"out": str(out), "ok": True,
+                          "one_step": report["one_step"], "rollout": report["rollout"],
+                          "pixel": report["pixel"]}, indent=2))
+    print(f"wrote {out.resolve()}")
+    return 0
+
+
+def _num(x) -> str:
+    return "n/a" if x is None else f"{x:.4g}"
+
+
 def build_parser() -> argparse.ArgumentParser:
     ap = argparse.ArgumentParser(prog="awpredict",
                                  description=__doc__.split("\n\n")[0])
@@ -93,6 +146,13 @@ def build_parser() -> argparse.ArgumentParser:
     p_c.add_argument("--module", required=True, help="importable module path")
     p_c.add_argument("--name", required=True, help="class or object in that module")
     p_c.add_argument("--protocol", choices=sorted(_PROTOCOLS), default="WorldModel")
+
+    p_v = sub.add_parser("viz", help="render what an engine imagines (HTML)")
+    p_v.add_argument("--engine", required=True, choices=["lewm", "mlp", "concept"])
+    p_v.add_argument("--demo", action="store_true",
+                     help="train briefly on a synthetic world, then dream along it")
+    p_v.add_argument("--out", help="HTML path (default awpredict-dream-<engine>.html)")
+    p_v.add_argument("--steps", type=int, default=6, help="imagined steps (default 6)")
     return ap
 
 
@@ -115,7 +175,8 @@ def main(argv: "list[str] | None" = None) -> int:
     if not args.cmd:
         build_parser().print_help()
         return 2
-    return {"engines": cmd_engines, "conforms": cmd_conforms}[args.cmd](args)
+    return {"engines": cmd_engines, "conforms": cmd_conforms,
+            "viz": cmd_viz}[args.cmd](args)
 
 
 if __name__ == "__main__":

awpredict-0.1.4/awpredict/core/concept.py

@@ -0,0 +1,357 @@
+"""concept: a next-concept JEPA world model over embedding sequences.
+
+Three published ideas, one engine:
+
+  * Next Concept Prediction (ConceptLM, arXiv:2602.08984). Mean-pool every
+    ``k`` consecutive embeddings into a concept, quantize concepts with a
+    product codebook (S segments x N codes, SimVQ MLP over the codes so they
+    do not collapse), and predict the NEXT concept as the current one plus a softmax-weighted sum
+    of codes per segment. Loss = MSE(predicted, true next concept) + VQ loss.
+  * Coarse-to-fine codes (Next-Concept Modeling, HazyResearch ncm). The
+    quantizer is RESIDUAL across ``levels``: level 0 is the coarse concept
+    (a topic-sized id), each later level quantizes what the previous left.
+    ``concept_ids`` returns the whole hierarchy, coarse first.
+  * SIGReg (LeWorldModel / LeJEPA). The prediction target is the encoder's
+    own output with NO stop-gradient and NO EMA; the Gaussian regularizer on
+    the concept latents is what prevents collapse, exactly as in lewm.py.
+
+Observations are float sequences ``[T, d_in]`` (token hidden states from any
+frozen LM, or one embedding per agent step). The engine never fetches
+embeddings itself -- the caller supplies them -- so it stays dependency-light.
+
+``plan(context, candidates)`` ranks candidate continuations (for example K
+parallel drafts from a diffusion LM) by how close each one's concept lands to
+the predicted next concept: the world model chooses, the fast generator
+proposes.
+
+Torch is optional at import time. Without it ``ok`` is False and every method
+returns None / [] -- it never raises into a caller.
+"""
+from __future__ import annotations
+
+import logging
+import os
+from typing import Any, Dict, List, Optional, Sequence
+
+logger = logging.getLogger(__name__)
+
+try:
+    import torch
+    import torch.nn as nn
+    import torch.nn.functional as F
+    _TORCH_OK = True
+except Exception:  # pragma: no cover - torch is an optional dependency here
+    torch = None  # type: ignore[assignment]
+    nn = None  # type: ignore[assignment]
+    F = None  # type: ignore[assignment]
+    _TORCH_OK = False
+
+
+def enabled() -> bool:
+    """Usable here: torch imports. No env gate -- nothing consumes it by default."""
+    return _TORCH_OK
+
+
+def _sigreg(z: "torch.Tensor", dirs: "torch.Tensor", t: "torch.Tensor") -> "torch.Tensor":
+    """Epps-Pulley SIGReg over a batch of latents z [B, D], tested against N(0, I).
+
+    Unlike lewm.py the latents are NOT standardized first: a next-latent MSE
+    with no stop-gradient can otherwise win by shrinking the whole latent scale
+    (measured: in-order and shuffled surprise both fell to ~0.07). Testing the
+    raw projections pins mean 0 AND unit scale, as LeJEPA does."""
+    if z.shape[0] < 4:
+        return z.new_tensor(0.0)
+    proj = z @ dirs
+    arg = t[:, None, None] * proj[None, :, :]
+    # Complex ECF vs the N(0,1) characteristic function exp(-t^2/2). The
+    # modulus-only form is translation-invariant, so a collapse to ONE point
+    # scores |ecf|=1 everywhere with zero gradient -- measured stuck at 0.65.
+    re = arg.cos().mean(dim=1) - (-(t ** 2) / 2).exp()[:, None]
+    im = arg.sin().mean(dim=1)
+    return (re ** 2 + im ** 2).mean()
+
+
+if _TORCH_OK:
+
+    class _ProductCodebook(nn.Module):
+        """S independent codebooks of N codes; codes pass through a SimVQ MLP."""
+
+        def __init__(self, dim: int, segments: int, codes: int):
+            super().__init__()
+            self.S, self.N, self.seg = segments, codes, dim // segments
+            self.raw = nn.Parameter(torch.randn(segments, codes, self.seg) * 0.5)
+            self.mlp = nn.Sequential(nn.Linear(self.seg, self.seg), nn.ReLU(),
+                                     nn.Linear(self.seg, self.seg))
+
+        def codes(self) -> "torch.Tensor":
+            return self.mlp(self.raw)                                   # [S, N, seg]
+
+        def quantize(self, c: "torch.Tensor"):
+            """c [..., D] -> (quantized [..., D], ids [..., S])."""
+            book = self.codes()
+            cs = c.reshape(*c.shape[:-1], self.S, self.seg)             # [..., S, seg]
+            d = ((cs.unsqueeze(-2) - book) ** 2).sum(-1)                # [..., S, N]
+            ids = d.argmin(-1)                                          # [..., S]
+            seg_idx = torch.arange(self.S, device=c.device).expand_as(ids)
+            q = book[seg_idx, ids]                                      # [..., S, seg]
+            return q.reshape(c.shape), ids
+
+        def soft(self, logits: "torch.Tensor") -> "torch.Tensor":
+            """logits [..., S, N] -> weighted code sum [..., D] (differentiable)."""
+            w = logits.softmax(-1)
+            q = torch.einsum("...sn,snd->...sd", w, self.codes())
+            return q.reshape(*q.shape[:-2], self.S * self.seg)
+
+    class _ConceptNet(nn.Module):
+        def __init__(self, d_in: int, dim: int, segments: int, codes: int,
+                     levels: int, layers: int, heads: int, max_concepts: int):
+            super().__init__()
+            # LayerNorm without affine pins every concept to norm sqrt(dim): the
+            # next-latent MSE cannot win by shrinking the latent scale, and
+            # SIGReg then only has to rule out the remaining collapse (every
+            # concept the same direction). Without it: z std 0.02, sigreg stuck.
+            self.enc = nn.Sequential(nn.Linear(d_in, dim), nn.GELU(), nn.Linear(dim, dim),
+                                     nn.LayerNorm(dim, elementwise_affine=False))
+            self.pos = nn.Parameter(torch.zeros(max_concepts, dim))
+            block = nn.TransformerEncoderLayer(dim, heads, dim * 2, dropout=0.0,
+                                               batch_first=True, norm_first=True)
+            self.ctx = nn.TransformerEncoder(block, layers, enable_nested_tensor=False)
+            self.books = nn.ModuleList(_ProductCodebook(dim, segments, codes)
+                                       for _ in range(levels))
+            self.heads = nn.ModuleList(nn.Linear(dim, segments * codes) for _ in range(levels))
+            self.S, self.N = segments, codes
+
+        def concepts(self, x: "torch.Tensor", k: int) -> "torch.Tensor":
+            """x [B, T, d_in] -> concept latents [B, T//k, dim] (mean-pool then encode)."""
+            b, t, d = x.shape
+            m = t // k
+            pooled = x[:, : m * k].reshape(b, m, k, d).mean(2)
+            return self.enc(pooled)
+
+        def predict(self, c: "torch.Tensor") -> "torch.Tensor":
+            """c [B, M, dim] -> predicted NEXT concept at every position [B, M, dim]."""
+            m = c.shape[1]
+            mask = torch.triu(torch.full((m, m), float("-inf"), device=c.device), 1)
+            h = self.ctx(c + self.pos[:m], mask=mask)
+            # RESIDUAL: next = current + a soft code mix. Predicting the absolute
+            # next concept from codes regressed to the mean: on real text plan@1
+            # was 8% (chance 10%) while "pick the candidate nearest the last
+            # concept" scored 52%. As a delta it is 82% (n=120, 2026-10-19).
+            out = c
+            for head, book in zip(self.heads, self.books):
+                logits = head(h).reshape(*h.shape[:-1], self.S, self.N)
+                out = out + book.soft(logits)
+            return out
+
+        def residual_quantize(self, c: "torch.Tensor"):
+            """Coarse-to-fine residual VQ: returns (sum of levels, [ids per level], vq_loss)."""
+            resid, total, ids, loss = c, torch.zeros_like(c), [], c.new_tensor(0.0)
+            for book in self.books:
+                q, i = book.quantize(resid.detach())
+                # codebook-only loss: ConceptLM confines L_VQ to the quantizer so
+                # the codes track the latents, never the reverse (a commitment
+                # term pulled the encoder onto a few codes and collapsed it).
+                loss = loss + F.mse_loss(q, resid.detach())
+                total = total + q
+                ids.append(i)
+                resid = resid - q.detach()
+            return total, ids, loss
+
+
+class ConceptWorldModel:
+    """Next-concept JEPA engine. See the module docstring for the recipe."""
+
+    def __init__(self, d_in: int, dim: int = 128, k: int = 4, segments: int = 8,
+                 codes: int = 64, levels: int = 2, layers: int = 2, heads: int = 4,
+                 max_concepts: int = 256, sigreg_lambda: float = 1.0, lr: float = 1e-3,
+                 buffer_size: int = 4096, seed: int = 0, device: Optional[str] = None):
+        self.ok = False
+        self.k, self.d_in, self.lam = k, d_in, sigreg_lambda
+        self.max_concepts = max_concepts
+        self._buf: List[Any] = []
+        self._buf_cap = buffer_size
+        self.steps = 0
+        if not _TORCH_OK:
+            logger.warning("concept: torch unavailable; engine disabled (ok=False)")
+            return
+        if dim % segments:
+            raise ValueError(f"dim {dim} must divide into {segments} segments")
+        g = torch.Generator().manual_seed(seed)
+        torch.manual_seed(seed)
+        self.device = torch.device(device or os.environ.get("AWPREDICT_DEVICE", "cpu"))
+        self.net = _ConceptNet(d_in, dim, segments, codes, levels, layers, heads,
+                               max_concepts).to(self.device)
+        dirs = torch.randn(dim, 64, generator=g)
+        self._sig_dirs = (dirs / dirs.norm(dim=0, keepdim=True)).to(self.device)
+        self._sig_t = torch.linspace(0.2, 3.0, 17).to(self.device)
+        self.opt = torch.optim.AdamW(self.net.parameters(), lr=lr, weight_decay=1e-4)
+        self.cfg = dict(d_in=d_in, dim=dim, k=k, segments=segments, codes=codes,
+                        levels=levels, layers=layers, heads=heads, max_concepts=max_concepts)
+        self.ok = True
+
+    # -- helpers -------------------------------------------------------------
+    def _seq(self, obs: Any) -> Optional["torch.Tensor"]:
+        """obs [T, d_in] or [B, T, d_in] -> float tensor [B, T, d_in] on device."""
+        x = torch.as_tensor(obs, dtype=torch.float32, device=self.device)
+        if x.dim() == 2:
+            x = x.unsqueeze(0)
+        if x.dim() != 3 or x.shape[-1] != self.d_in or x.shape[1] < self.k:
+            return None
+        return x
+
+    # -- WorldModel protocol -------------------------------------------------
+    def observe(self, obs: Any, action: Any = None, next_obs: Any = None,
+                *args: Any, **kwargs: Any) -> bool:
+        """Buffer one sequence. ``next_obs`` (if given) is appended to ``obs``."""
+        if not self.ok:
+            return False
+        seq = obs if next_obs is None else _concat(obs, next_obs)
+        x = self._seq(seq)
+        if x is None or x.shape[1] < 2 * self.k:
+            return False
+        self._buf.append(x[0].detach().cpu())
+        if len(self._buf) > self._buf_cap:
+            self._buf.pop(0)
+        return True
+
+    def encode(self, obs: Any, cond: Any = None) -> Optional["torch.Tensor"]:
+        """Sequence -> concept latents [M, dim]."""
+        if not self.ok:
+            return None
+        x = self._seq(obs)
+        if x is None:
+            return None
+        with torch.no_grad():
+            return self.net.concepts(x, self.k)[0]
+
+    def predict(self, z: Any, action: Any = None, **kwargs: Any) -> Optional["torch.Tensor"]:
+        """Concept latents [M, dim] -> predicted next concept [dim]."""
+        if not self.ok or z is None:
+            return None
+        with torch.no_grad():
+            c = torch.as_tensor(z, dtype=torch.float32, device=self.device)
+            return self.net.predict(c.unsqueeze(0)[:, -self.max_concepts:])[0, -1]
+
+    def surprise(self, obs: Any, action: Any = None, next_obs: Any = None,
+                 *args: Any, **kwargs: Any) -> Optional[float]:
+        """Mean squared next-concept prediction error over the sequence."""
+        if not self.ok:
+            return None
+        seq = obs if next_obs is None else _concat(obs, next_obs)
+        x = self._seq(seq)
+        if x is None or x.shape[1] < 2 * self.k:
+            return None
+        with torch.no_grad():
+            c = self.net.concepts(x, self.k)
+            pred = self.net.predict(c[:, :-1])
+            return float(((pred - c[:, 1:]) ** 2).sum(-1).mean().item())
+
+    def train_step(self, batch_size: int = 32, steps: int = 1,
+                   *args: Any, **kwargs: Any) -> Optional[Dict[str, Any]]:
+        if not self.ok or not self._buf:
+            return None
+        out: Dict[str, Any] = {}
+        self.net.train()
+        for _ in range(max(1, steps)):
+            idx = torch.randint(len(self._buf), (min(batch_size, len(self._buf)),))
+            t = min(self._buf[i].shape[0] for i in idx.tolist())
+            t = (min(t, self.max_concepts * self.k) // self.k) * self.k
+            x = torch.stack([self._buf[i][:t] for i in idx.tolist()]).to(self.device)
+            c = self.net.concepts(x, self.k)
+            pred = self.net.predict(c[:, :-1])
+            l_ncp = F.mse_loss(pred, c[:, 1:])
+            _, _, l_vq = self.net.residual_quantize(c)
+            l_sig = _sigreg(c.reshape(-1, c.shape[-1]), self._sig_dirs, self._sig_t)
+            loss = l_ncp + l_vq + self.lam * l_sig
+            self.opt.zero_grad(set_to_none=True)
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(self.net.parameters(), 1.0)
+            self.opt.step()
+            self.steps += 1
+            out = {"loss": loss.item(), "ncp": l_ncp.item(), "vq": l_vq.item(),
+                   "sigreg": l_sig.item(), "steps": self.steps}
+        self.net.eval()
+        return out
+
+    def plan(self, context: Any, candidates: Sequence[Any] = (), *args: Any,
+             **kwargs: Any) -> List[int]:
+        """Rank candidate continuations of ``context`` best-first.
+
+        Each candidate is a ``[t, d_in]`` sequence (t >= k). Score = squared
+        distance between the predicted next concept and the candidate's first
+        concept. Use it to pick among K parallel drafts.
+        """
+        if not self.ok or not candidates:
+            return []
+        z = self.encode(context)
+        if z is None:
+            return []
+        target = self.predict(z)
+        scores = []
+        for i, cand in enumerate(candidates):
+            zc = self.encode(cand)
+            d = float("inf") if zc is None else float(((zc[0] - target) ** 2).sum())
+            scores.append((d, i))
+        return [i for _, i in sorted(scores)]
+
+    # -- concept vocabulary --------------------------------------------------
+    def concept_ids(self, obs: Any) -> Optional[List[List[List[int]]]]:
+        """Discrete coarse-to-fine codes per chunk: [M][level][segment]."""
+        if not self.ok:
+            return None
+        x = self._seq(obs)
+        if x is None:
+            return None
+        with torch.no_grad():
+            _, ids, _ = self.net.residual_quantize(self.net.concepts(x, self.k))
+        stacked = torch.stack(ids, dim=-2)[0]                          # [M, levels, S]
+        return stacked.tolist()
+
+    def codebook_usage(self, obs_batch: Sequence[Any]) -> Optional[float]:
+        """Fraction of level-0 codes used across ``obs_batch`` (mean over segments)."""
+        if not self.ok:
+            return None
+        seen = [set() for _ in range(self.net.S)]
+        for obs in obs_batch:
+            ids = self.concept_ids(obs)
+            for chunk in ids or []:
+                for s, code in enumerate(chunk[0]):
+                    seen[s].add(code)
+        return sum(len(s) for s in seen) / (self.net.S * self.net.N)
+
+    # -- persistence ---------------------------------------------------------
+    def save(self, path: str, *args: Any, **kwargs: Any) -> bool:
+        if not self.ok:
+            return False
+        tmp = f"{path}.tmp"
+        try:
+            torch.save({"cfg": self.cfg, "state": self.net.state_dict(),
+                        "steps": self.steps}, tmp)
+            os.replace(tmp, path)
+            return True
+        except OSError as exc:
+            logger.warning("concept.save failed: %s", exc)
+            return False
+
+    def load(self, path: str, *args: Any, **kwargs: Any) -> bool:
+        if not self.ok or not os.path.exists(path):
+            return False
+        try:
+            ck = torch.load(path, map_location=self.device, weights_only=True)
+            if ck.get("cfg") != self.cfg:
+                logger.warning("concept.load: config mismatch %s != %s", ck.get("cfg"), self.cfg)
+                return False
+            self.net.load_state_dict(ck["state"])
+            self.steps = int(ck.get("steps", 0))
+            self.net.eval()
+            return True
+        except (OSError, RuntimeError, KeyError) as exc:
+            logger.warning("concept.load failed: %s", exc)
+            return False
+
+
+def _concat(a: Any, b: Any) -> Any:
+    if _TORCH_OK:
+        return torch.cat([torch.as_tensor(a, dtype=torch.float32),
+                          torch.as_tensor(b, dtype=torch.float32)], dim=-2)
+    return list(a) + list(b)

{awpredict-0.1.2 → awpredict-0.1.4}/awpredict/core/mlp.py

@@ -334,6 +334,33 @@ class MLPWorldModel:
             return None
         return self._hash_state_embedding(str(obs))
 
+    def predict_latent(self, z: Any, action: Any) -> Optional[List[float]]:
+        """Latent → predicted next latent through the neural transition net.
+
+        ``predict`` answers in state HASHES (the tabular contract), which cannot
+        be fed back into the model; this is the vector-space step a rollout or
+        a visualizer needs. ``z`` is an ``encode()`` output or a previous
+        ``predict_latent()`` output. None until the neural path has trained, or
+        when torch is missing — never a fabricated vector.
+        """
+        if z is None or not self._neural_available or not self._encoder.is_available:
+            return None
+        try:
+            if isinstance(z, torch.Tensor):
+                state = z.detach().float().reshape(-1)
+            else:
+                state = torch.tensor(list(z), dtype=torch.float32)
+            if state.numel() != self._state_dim:
+                return None
+            a = self._action_encoder.encode(action)
+            if not isinstance(a, torch.Tensor):
+                a = torch.tensor(a, dtype=torch.float32)
+            nxt, _, _ = self._encoder.predict(state, a.detach())
+            return nxt.tolist()
+        except (RuntimeError, TypeError, ValueError) as e:
+            log.debug(f"predict_latent failed: {e}")
+            return None
+
     def surprise(
         self,
         obs: int,