Perf, reproducibility, test suite + CI, packaging

.github/workflows/ci.yml

@@ -0,0 +1,23 @@
+name: CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    timeout-minutes: 15
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+          cache: pip
+      - name: Install (CPU torch)
+        run: |
+          pip install torch --index-url https://download.pytorch.org/whl/cpu
+          pip install -e ".[dev]"
+      - name: Run tests
+        run: pytest -q

.gitignore

@@ -5,3 +5,4 @@ venv/
 runs/
 .env
 .DS_Store
+*.egg-info/

README.md

@@ -1,5 +1,6 @@
 # LatentBridge
 
+![CI](https://github.com/micka420-collab/latentbridge/actions/workflows/ci.yml/badge.svg)
 ![License: MIT](https://img.shields.io/badge/License-MIT-green.svg)
 ![Python 3.11+](https://img.shields.io/badge/python-3.11%2B-blue.svg)
 ![PyTorch](https://img.shields.io/badge/PyTorch-CPU%20or%20CUDA-ee4c2c.svg)
@@ -52,7 +53,8 @@ solve the task — which makes the bridge's value directly measurable:
 ## Quick start
 
 ```bash
-pip install -r requirements.txt
+pip install -e ".[dev]"        # or: pip install -r requirements.txt
+pytest                         # fast test suite (~4s CPU) — envs, models, pipeline
 # one experiment (gridworld): trains, evaluates on held-out seeds, prints score
 python -m latentbridge.experiment --config configs/base.yaml
 # the full benchmark with the no-bridge control group:

latentbridge/checkpoint.py

@@ -25,7 +25,9 @@ def save(path, modules, cfg, obs_dim, n_actions):
 
 
 def load(path, device="cpu"):
-    ckpt = torch.load(path, map_location=device, weights_only=False)
+    # checkpoints only contain tensors + plain YAML-derived config types, so the
+    # safe weights_only loader is enough (no arbitrary pickle execution).
+    ckpt = torch.load(path, map_location=device, weights_only=True)
     m = make_models(ckpt["cfg"], ckpt["obs_dim"], ckpt["n_actions"], device)
     m["encoder"].load_state_dict(ckpt["encoder"])
     m["decoder"].load_state_dict(ckpt["decoder"])

latentbridge/controller.py

@@ -30,6 +30,10 @@ def __init__(self, encoder, world_model, bridge=None, featurizer=None,
         self.guide_weight = guide_weight
         self.device = device
         self.n_actions = world_model.n_actions
+        # subgoal text -> target latent memo. The planner proposes the same
+        # goal at every step of an episode, so embedding + from_lang would
+        # otherwise be recomputed each act() call.
+        self._target_cache: dict[str, torch.Tensor] = {}
 
     @torch.no_grad()
     def act(self, obs: np.ndarray, env=None, guided: bool = False) -> int:
@@ -38,9 +42,15 @@ def act(self, obs: np.ndarray, env=None, guided: bool = False) -> int:
         target_latent = None
         if guided and self.planner is not None and self.bridge is not None:
             sub = self.planner.propose(env)
-            emb = self.featurizer.embed(sub["subgoal_text"])
-            emb_t = torch.as_tensor(emb, dtype=torch.float32, device=self.device).unsqueeze(0)
-            target_latent = self.bridge.from_lang(emb_t)  # (1, D)
+            text = sub["subgoal_text"]
+            target_latent = self._target_cache.get(text)
+            if target_latent is None:
+                emb = self.featurizer.embed(text)
+                emb_t = torch.as_tensor(emb, dtype=torch.float32, device=self.device).unsqueeze(0)
+                target_latent = self.bridge.from_lang(emb_t)  # (1, D)
+                if len(self._target_cache) > 4096:
+                    self._target_cache.clear()
+                self._target_cache[text] = target_latent
 
         # random shooting: sample action sequences, roll out, score
         K, H = self.n_samples, self.horizon

latentbridge/evaluate.py

@@ -71,13 +71,16 @@ def evaluate(modules, cfg, device="cpu"):
 
 @torch.no_grad()
 def _alignment_cosine(enc, bridge, feat, cfg, device, seeds):
+    # dedicated RNG: the metric must not depend on the global np.random state
+    # (which varies with whatever ran before evaluate was called)
+    rng = np.random.default_rng(cfg["seed"] + 30_000)
     texts, obss = [], []
     for s in seeds[: min(len(seeds), 64)]:
         env = make_env(cfg, seed=s)
         obs = env.reset(seed=s)
         for _ in range(3):
             texts.append(env.text_state()); obss.append(obs)
-            obs, _, done, _ = env.step(np.random.randint(0, env.action_space))
+            obs, _, done, _ = env.step(int(rng.integers(0, env.action_space)))
             if done:
                 break
     z = enc(torch.as_tensor(np.asarray(obss, dtype=np.float32), device=device))

latentbridge/llm/planner.py

@@ -25,6 +25,10 @@ def __init__(self, backend: str = "mock", model: str | None = None, size: int =
         self.size = size  # kept for backward compat / gridworld
         self.model = model or os.environ.get("PLANNER_MODEL", "google/gemma-4-31b-it:free")
         self._client = None
+        # instruction text -> goal key memo. The controller asks for a proposal
+        # at EVERY env step; at temperature 0 the classification of a given
+        # instruction is deterministic, so one API call per goal is enough.
+        self._classify_cache: dict[str, str | None] = {}
         if backend == "openrouter":
             self._init_client()
 
@@ -57,6 +61,8 @@ def propose(self, env, intention: str | None = None) -> dict:
     def _classify_llm(self, env, intention=None):
         menu = env.goal_menu()
         text = intention or env.goal_instruction()
+        if text in self._classify_cache:
+            return self._classify_cache[text]
         options = "\n".join(f"- {k}: {desc}" for k, desc in menu)
         prompt = (
             f"User says: \"{text}\"\n\n"
@@ -72,7 +78,9 @@ def _classify_llm(self, env, intention=None):
             out = (resp.choices[0].message.content or "").strip().lower()
             for k, _ in menu:
                 if k in out:
+                    self._classify_cache[text] = k
                     return k
         except Exception as e:
             print(f"[planner] llm classify failed ({e}); using mock goal")
+        # not cached: a transient failure should not pin this instruction to mock
         return None

latentbridge/text_features.py

@@ -24,12 +24,20 @@ class HashingTextFeaturizer:
     def __init__(self, dim: int = 128, seed: int = 0):
         self.dim = int(dim)
         self.seed = int(seed)
+        # token -> vector memo: the sha256 + RNG expansion is deterministic per
+        # token, and embed_batch calls it for every token occurrence over
+        # thousands of near-identical state descriptions.
+        self._cache: dict[str, np.ndarray] = {}
 
     def _token_vec(self, tok: str) -> np.ndarray:
-        h = hashlib.sha256(f"{self.seed}:{tok}".encode()).digest()
-        # expand the digest deterministically to dim floats in [-1, 1]
-        rng = np.random.default_rng(int.from_bytes(h[:8], "little"))
-        return rng.standard_normal(self.dim).astype(np.float32)
+        v = self._cache.get(tok)
+        if v is None:
+            h = hashlib.sha256(f"{self.seed}:{tok}".encode()).digest()
+            # expand the digest deterministically to dim floats in [-1, 1]
+            rng = np.random.default_rng(int.from_bytes(h[:8], "little"))
+            v = rng.standard_normal(self.dim).astype(np.float32)
+            self._cache[tok] = v
+        return v
 
     def embed(self, text: str) -> np.ndarray:
         toks = _TOKEN.findall(text.lower())

pyproject.toml

@@ -0,0 +1,29 @@
+[build-system]
+requires = ["setuptools>=68"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "latentbridge"
+version = "0.1.0"
+description = "Couple a frozen LLM to a trained world model through a learned latent bridge, and plan inside the world model toward a goal given in language."
+readme = "README.md"
+license = { text = "MIT" }
+requires-python = ">=3.11"
+dependencies = [
+    "numpy>=1.24",
+    "pyyaml>=6.0",
+    "torch>=2.0",
+]
+
+[project.optional-dependencies]
+# real frozen-LLM planner via OpenRouter (otherwise the offline "mock" backend is used)
+llm = ["openai>=1.0", "python-dotenv>=1.0"]
+# public MiniGrid benchmark env (configs/minigrid.yaml)
+minigrid = ["gymnasium>=0.29", "minigrid>=2.3"]
+dev = ["pytest>=8.0"]
+
+[tool.setuptools.packages.find]
+include = ["latentbridge*"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]

tests/conftest.py

@@ -0,0 +1,30 @@
+import os
+import sys
+
+import pytest
+
+ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, ROOT)
+
+
+@pytest.fixture
+def tiny_cfg():
+    """Smallest config that exercises the full train -> evaluate pipeline."""
+    return {
+        "seed": 0,
+        "env": {
+            "kind": "gridworld",
+            "size": 3,
+            "max_steps": 12,
+            "n_walls": 0,
+            "step_penalty": 0.01,
+            "include_goal_in_obs": False,
+        },
+        "model": {"latent_dim": 16, "lang_dim": 24, "hidden": 32},
+        "loss": {"dyn": 1.0, "rew": 1.0, "recon": 0.5, "align": 0.5,
+                 "cycle": 0.1, "ground": 0.5},
+        "train": {"collect_episodes": 4, "steps": 30, "batch_size": 32, "lr": 1e-3},
+        "control": {"horizon": 3, "n_samples": 32, "gamma": 0.95, "guide_weight": 1.0},
+        "planner": {"backend": "mock"},
+        "eval": {"episodes": 3},
+    }

tests/test_baselines.py

@@ -0,0 +1,29 @@
+"""Fast checks for the comparison harness baselines."""
+import copy
+
+from latentbridge.baselines import (LLMOnlyPolicy, RandomPolicy, eval_policy,
+                                    train_dqn, train_ppo)
+
+
+def test_random_policy_eval(tiny_cfg):
+    pol = RandomPolicy(4, seed=0)
+    metr = eval_policy(pol, tiny_cfg)
+    assert {"success_rate", "avg_steps", "avg_reward"} <= set(metr)
+    assert 0.0 <= metr["success_rate"] <= 1.0
+
+
+def test_llm_only_solves_tiny_grid(tiny_cfg):
+    # greedy walk to the planner-revealed goal must beat chance on a 3x3 grid
+    pol = LLMOnlyPolicy(tiny_cfg, seed=0)
+    metr = eval_policy(pol, tiny_cfg)
+    assert metr["success_rate"] == 1.0
+    assert pol.degraded is False
+
+
+def test_ppo_dqn_smoke(tiny_cfg):
+    cfg = copy.deepcopy(tiny_cfg)
+    for train_fn in (train_ppo, train_dqn):
+        pol, info = train_fn(cfg, seed=0, episodes_budget=3)
+        assert info["params"] > 0
+        metr = eval_policy(pol, cfg)
+        assert 0.0 <= metr["success_rate"] <= 1.0

tests/test_envs.py

@@ -0,0 +1,94 @@
+"""Contract tests for every built-in env: shapes, step protocol, language
+interface, and per-seed determinism (held-out evaluation depends on it)."""
+import numpy as np
+import pytest
+
+from latentbridge.env import make_env
+
+KINDS = ["gridworld", "lifeworld", "physicsworld", "desktopworld", "simcalc"]
+
+
+def _cfg(kind):
+    return {"seed": 0, "env": {"kind": kind}}
+
+
+@pytest.mark.parametrize("kind", KINDS)
+def test_env_contract(kind):
+    env = make_env(_cfg(kind), seed=3)
+    obs = env.reset(seed=3)
+    assert isinstance(obs, np.ndarray)
+    assert obs.shape == (env.obs_dim,)
+    assert env.action_space >= 2
+
+    nobs, r, done, info = env.step(0)
+    assert nobs.shape == (env.obs_dim,)
+    assert isinstance(float(r), float)
+    assert isinstance(done, bool)
+    assert isinstance(info, dict)
+
+    assert isinstance(env.text_state(), str) and env.text_state()
+    assert isinstance(env.goal_state_text(env.true_goal_key()), str)
+
+
+@pytest.mark.parametrize("kind", KINDS)
+def test_env_deterministic_per_seed(kind):
+    rng = np.random.default_rng(0)
+    actions = None
+    trajs = []
+    for _ in range(2):
+        env = make_env(_cfg(kind), seed=7)
+        obs = env.reset(seed=7)
+        if actions is None:
+            actions = [int(a) for a in rng.integers(0, env.action_space, size=5)]
+        traj = [obs]
+        for a in actions:
+            obs, _, done, _ = env.step(a)
+            traj.append(obs)
+            if done:
+                break
+        trajs.append(np.concatenate(traj))
+    assert np.array_equal(trajs[0], trajs[1])
+
+
+@pytest.mark.parametrize("kind", KINDS)
+def test_env_episode_terminates(kind):
+    env = make_env(_cfg(kind), seed=1)
+    env.reset(seed=1)
+    rng = np.random.default_rng(1)
+    for _ in range(10_000):
+        _, _, done, _ = env.step(int(rng.integers(0, env.action_space)))
+        if done:
+            return
+    pytest.fail("episode never terminated")
+
+
+def test_gridworld_goal_hidden_vs_visible():
+    hidden = make_env({"seed": 0, "env": {"kind": "gridworld", "size": 4}}, seed=0)
+    visible = make_env({"seed": 0, "env": {"kind": "gridworld", "size": 4,
+                                           "include_goal_in_obs": True}}, seed=0)
+    assert hidden.obs_dim == 2 * 16
+    assert visible.obs_dim == 3 * 16
+
+
+def test_gridworld_reach_goal():
+    env = make_env({"seed": 0, "env": {"kind": "gridworld", "size": 3,
+                                       "max_steps": 30}}, seed=5)
+    env.reset(seed=5)
+    # walk greedily to the (known) goal; must terminate with reached=True
+    info = {}
+    for _ in range(30):
+        (ar, ac), (gr, gc) = env.agent, env.goal
+        if ar != gr:
+            a = 1 if gr > ar else 0
+        else:
+            a = 3 if gc > ac else 2
+        _, r, done, info = env.step(a)
+        if done:
+            break
+    assert info.get("reached") is True
+    assert r == 1.0
+
+
+def test_make_env_unknown_kind():
+    with pytest.raises(ValueError):
+        make_env({"seed": 0, "env": {"kind": "nope"}})