fix: trajectory guidance never trained — the heads were empty every run

ComfyUI executes a node under torch.inference_mode(), so a descriptor read out of
the probe log becomes an inference tensor and autograd refuses it:

  Trajectory probe intake failed: Inference tensors do not track version counter.

Every head stayed empty and the run reported "nothing learned yet" on every
generation — so the feature was OFF while looking merely untrained, which is the
one failure invisible from the outside: "not steering" is also what correct
behaviour looks like before enough ratings. Seen on a live rental with 73 runs
logged and zero of them trained.

The output value function's trainer already carries this guard; this one did not.

Also: the bucket-isolation test was measuring TRAINING accuracy, where a
few-hundred-parameter MLP fits 20 noise vectors perfectly and every bucket scores
1.0. It now holds rows back, which is the only version that can tell a head that
learned from one that memorised.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

tests/test_trajectory_guidance.py

@@ -62,20 +62,36 @@ def _store(monkeypatch, tmp_path):
 
 def test_a_head_learns_from_its_own_bucket_only(monkeypatch, tmp_path):
     """The whole claim of the feature: what steers at step 2 learned from step 2. A head
-    fed every bucket's descriptors would just be output_guidance with extra files."""
+    fed every bucket's descriptors would just be output_guidance with extra files.
+
+    Measured as ranking accuracy over the descriptors each head actually saw, not as a score
+    at some constant vector: a head trained on noise will happily extrapolate to a large
+    arbitrary value far outside its data, which made an earlier version of this test pass or
+    fail on `train_on`'s unseeded bootstrap sampling rather than on anything real.
+    """
+    import random as _random
+    _random.seed(7)                      # train_on bootstraps with the global RNG
     _store(monkeypatch, tmp_path)
-    tg.train_from_rows("k", _separable())
+    rows = _separable(n=28)
+    train, held_out = rows[:20], rows[20:]
+    tg.train_from_rows("k", train)
     model = tg.load("k")
     assert model is not None
 
-    def score(bucket, value):
+    def accuracy(bucket, on):
+        pos, neg = [], []
         with torch.inference_mode(False), torch.no_grad():
-            return float(model._head(bucket)(torch.full((1, DIM), value)).mean())
+            for row in on:
+                desc = next(c["desc"] for c in row["rows"] if c["bucket"] == bucket)
+                score = float(model._head(bucket)(desc.unsqueeze(0)).mean())
+                (pos if row["reward"] > 0 else neg).append(score)
+        return sum(p > q for p in pos for q in neg) / (len(pos) * len(neg))
 
-    # Bucket 1 saw the split and must rank a good-looking descriptor above a bad one.
-    assert score(1, 1.0) > score(1, -1.0)
-    # Bucket 0 saw noise either way, so it must not have invented an opinion.
-    assert abs(score(0, 1.0) - score(0, -1.0)) < abs(score(1, 1.0) - score(1, -1.0))
+    # HELD OUT, not the training rows. A few-hundred-parameter MLP fits 20 noise vectors
+    # perfectly, so training accuracy is 1.0 for every bucket and says nothing about whether
+    # a head learned anything -- including the head that will be steering real generations.
+    assert accuracy(1, held_out) > 0.9, accuracy(1, held_out)
+    assert accuracy(0, held_out) < 0.9, accuracy(0, held_out)
 
 
 def test_an_untrained_bucket_is_not_ready(monkeypatch, tmp_path):
@@ -250,3 +266,40 @@ def test_it_chains_onto_whatever_is_already_installed():
     wrapper = _install(model, _Ready([]), lambda s: 0)
     _predict(wrapper, torch.ones(1, 1, VIDEO_N + AUDIO_N))
     assert calls == [1]
+
+
+def test_training_works_inside_inference_mode(monkeypatch, tmp_path):
+    """ComfyUI executes a node under torch.inference_mode(), so a descriptor read out of the
+    log becomes an inference tensor and autograd refuses it:
+
+        Trajectory probe intake failed: Inference tensors do not track version counter.
+
+    Seen on a real rental. Every head stayed empty and the feature reported "nothing learned
+    yet" on every run, so it was OFF while appearing merely untrained — the one failure that
+    is invisible from the outside, because "not steering" is also what correct behaviour
+    looks like before enough ratings.
+    """
+    _store(monkeypatch, tmp_path)
+    with torch.inference_mode():
+        counts = tg.train_from_rows("k", _separable())
+    assert counts, "training produced nothing under inference mode"
+    assert tg.load("k") is not None
+
+
+def test_steering_works_inside_inference_mode(monkeypatch, tmp_path):
+    """The other half: sampling also runs under inference mode, so the gradient the wrapper
+    asks for is computed there too."""
+    _store(monkeypatch, tmp_path)
+    with torch.inference_mode():
+        tg.train_from_rows("k", _separable())
+    model = tg.load("k")
+    ready = [b for b in range(model.n_buckets) if model.ready(b)]
+    assert ready, "nothing trained, so this would pass for the wrong reason"
+
+    host = _model()
+    wrapper = _install(host, model, lambda s: ready[0])
+    latent = torch.randn(1, 1, VIDEO_N + AUDIO_N)
+    with torch.inference_mode():
+        out = _predict(wrapper, latent)
+    assert not torch.equal(out, latent)
+    assert torch.isfinite(out).all()

trajectory_guidance.py

@@ -138,6 +138,25 @@ class BucketedValue:
         return out
 
 
+def _train_all(model, rows):
+    """Every rated descriptor into the head for its own bucket."""
+    for row in rows:
+        per_scene = row.get("scene_rewards") or None
+        for cell in row["rows"]:
+            # Pass 0 only: a second pass runs its own schedule from its own starting sigma,
+            # so its bucket 2 is not the window bucket 2 steers in.
+            if int(cell.get("pass", 0)) != 0:
+                continue
+            reward = (per_scene or {}).get(int(cell.get("scene", 0)))
+            if reward is None:
+                reward = row.get("reward")
+            try:
+                reward = float(reward)
+            except (TypeError, ValueError):
+                continue
+            model.train_on(cell.get("bucket", 0), cell["desc"].float(), reward)
+
+
 def train_from_rows(refinement_key, rows, n_buckets=None):
     """Rebuild the per-bucket heads from probe rows. -> {bucket: samples}, or None.
 
@@ -159,21 +178,13 @@ def train_from_rows(refinement_key, rows, n_buckets=None):
         return None
 
     model = BucketedValue(n_buckets)
-    for row in usable:
-        per_scene = row.get("scene_rewards") or None
-        for cell in row["rows"]:
-            # Pass 0 only: a second pass runs its own schedule from its own starting sigma,
-            # so its bucket 2 is not the window bucket 2 steers in.
-            if int(cell.get("pass", 0)) != 0:
-                continue
-            reward = (per_scene or {}).get(int(cell.get("scene", 0)))
-            if reward is None:
-                reward = row.get("reward")
-            try:
-                reward = float(reward)
-            except (TypeError, ValueError):
-                continue
-            model.train_on(cell.get("bucket", 0), cell["desc"].float(), reward)
+    # ComfyUI executes a node under torch.inference_mode(), so every tensor made in here is
+    # an inference tensor and autograd refuses it ("Inference tensors do not track version
+    # counter"). Same guard the output value function's trainer carries. It has to wrap the
+    # whole loop, not one call: `.float()` on a stored descriptor makes a new tensor too, and
+    # that is the one training receives.
+    with torch.inference_mode(False), torch.enable_grad():
+        _train_all(model, usable)
 
     try:
         model.save(state_path(refinement_key))