Add online value function for reward-guided sampling

value_function.py: OnlineValueFunction MLP that trains incrementally on
(conditioning, reward) pairs. Uses replay buffer of 100 samples, 20
gradient steps per update. save/load to .pt file via refinement_state_path.

Studio popup: new "Value guidance" toggle. When on, trains the value function
after every rated generation (good and bad). Falls back to no-op if conditioning
deserialization fails.

chain sampler: _load_value_function reads the trained MLP if ready (5+
samples). _build_embed_guidance_wrapper uses value function gradient
(∂reward/∂conditioning) when available, fixed liked direction as fallback.
Log reports which mode is active and sample count.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

conditioning.py

@@ -11120,7 +11120,8 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner):
                   refinement_key_input="", positive_conditioning=None, clip_vision_output=None,
                   source_image=None, model=None, mode="Refine", advisor_mode="Off", advisor_thinking=True,
                   advisor_clip=None, feedback_prompt="", prompt_repair=True, temporal_style="natural",
-                  split_by_transitions=False, split_transition_placement="start", reference_injection=False, latent=None, seed_output_connected=False,
+                  split_by_transitions=False, split_transition_placement="start", reference_injection=False,
+                  value_guidance=False, latent=None, seed_output_connected=False,
                   _seed=None, _seed_source="fresh seed", _scene_seeds=None):
         seed = int(_seed) if _seed is not None else random.randint(1, 0xffffffffffffffff)
         encode_cache = {}
@@ -11224,6 +11225,24 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner):
         )
         memory_status = f"{memory_status}\n{seed_memory_status}\n{scene_sync_status}\n{intent_family_status}\n{intent_learning_status}"
         self._v2_update_conditioning_memory(global_state, previous_run, learning_profile, axis_feedback)
+        if value_guidance and has_previous_run and refinement_key and not learning_profile.get("skip_learning"):
+            try:
+                try:
+                    from .value_function import OnlineValueFunction
+                except ImportError:
+                    from value_function import OnlineValueFunction
+                payload = (previous_run or {}).get("conditioning")
+                if isinstance(payload, dict):
+                    cond_tensor = serializable_to_tensor(payload)
+                    reward = float(learning_profile.get("reward", 0.0))
+                    vf_path = refinement_state_path(refinement_key, "value_fn", prefix="refine_v2", extension="pt")
+                    vf = OnlineValueFunction.load_or_create(vf_path, hidden_dim=cond_tensor.shape[-1])
+                    if vf is not None:
+                        vf.train_on(cond_tensor, reward)
+                        vf.save(vf_path)
+                        print(f"[FunPackRefiner] Value function updated — {vf.n_trained} samples, buffer={len(vf.buffer_c)}")
+            except Exception as e:
+                print(f"[FunPackRefiner] Value function training failed: {e}")
         if has_previous_run and not learning_profile.get("skip_learning"):
             self._v2_update_streaks(global_state, learning_profile, update_conditioning_strength=not prompt_only_mode)
         repair_feedback, repair_persistence_status = self._v2_active_repair_feedback(
@@ -12938,6 +12957,7 @@ class FunPackStudio:
             split_transition_placement = "start"
         reference_injection = bool(rf.get("reference_injection", False))
         vision_conditioning = bool(rf.get("vision_conditioning", True))
+        value_guidance = bool(rf.get("value_guidance", False))
 
         # feedback_prompt: popup wins if override is on, else external wins
         popup_feedback = str(rf.get("feedback_prompt", "") or "")
@@ -13065,6 +13085,7 @@ class FunPackStudio:
             split_by_transitions=split_by_transitions,
             split_transition_placement=split_transition_placement,
             reference_injection=reference_injection,
+            value_guidance=value_guidance,
             latent=latent,
             seed_output_connected=seed_output_connected,
             _seed=seed,

samplers.py

@@ -1287,6 +1287,23 @@ class FunPackLTXAVSceneChainSampler:
         except Exception:
             return None
 
+    def _load_value_function(self, refinement_key):
+        """Load the online value function if trained and ready."""
+        try:
+            try:
+                from .value_function import OnlineValueFunction
+                from .conditioning import refinement_state_path
+            except ImportError:
+                from value_function import OnlineValueFunction
+                from conditioning import refinement_state_path
+            path = refinement_state_path(refinement_key, "value_fn", prefix="refine_v2", extension="pt")
+            if not __import__("os").path.exists(path):
+                return None
+            vf = OnlineValueFunction.load(path)
+            return vf if vf.is_ready() else None
+        except Exception:
+            return None
+
     def _load_liked_direction(self, refinement_key):
         """Read the liked conditioning direction from the Refiner's state file."""
         try:
@@ -1309,15 +1326,14 @@ class FunPackLTXAVSceneChainSampler:
         except Exception:
             return None
 
-    def _build_embed_guidance_wrapper(self, model, liked_dir, strength):
+    def _build_embed_guidance_wrapper(self, model, liked_dir, strength, value_fn=None):
         """Register a model_function_wrapper that nudges conditioning toward the
-        liked quality direction at each denoising step, scaling with sigma so
-        the effect is stronger during detail refinement (low sigma) than structure
-        formation (high sigma)."""
+        liked quality direction at each denoising step. Uses value function gradient
+        when available, falls back to the fixed liked direction otherwise."""
         old_wrapper = model.model_options.get("model_function_wrapper")
-        direction = torch.nn.functional.normalize(liked_dir.float(), dim=-1)
+        fixed_dir = torch.nn.functional.normalize(liked_dir.float(), dim=-1)
 
-        def _embed_wrapper(apply_fn, args, _ew=old_wrapper, _dir=direction, _s=strength):
+        def _embed_wrapper(apply_fn, args, _ew=old_wrapper, _fixed=fixed_dir, _vf=value_fn, _s=strength):
             c = args.get("c") or {}
             cond = c.get("c_crossattn")
             if cond is not None:
@@ -1326,13 +1342,18 @@ class FunPackLTXAVSceneChainSampler:
                     sigma = float(ts.max().item()) if ts is not None else 1.0
                 except Exception:
                     sigma = 1.0
-                # Scale: no effect at sigma=1 (pure noise), full effect below sigma=0.5
                 scale = max(0.0, 1.0 - sigma * 2.0)
                 if scale > 0:
-                    d = _dir.to(cond.device, cond.dtype)
-                    new_cond = cond + (_s * scale) * d.expand_as(cond)
+                    if _vf is not None:
+                        try:
+                            grad = _vf.gradient(cond)
+                            d = torch.nn.functional.normalize(grad.float(), dim=-1).to(cond.dtype)
+                        except Exception:
+                            d = _fixed.to(cond.device, cond.dtype).expand_as(cond)
+                    else:
+                        d = _fixed.to(cond.device, cond.dtype).expand_as(cond)
                     new_c = dict(c)
-                    new_c["c_crossattn"] = new_cond
+                    new_c["c_crossattn"] = cond + (_s * scale) * d
                     args = dict(args)
                     args["c"] = new_c
             if _ew is not None:
@@ -1436,12 +1457,15 @@ class FunPackLTXAVSceneChainSampler:
 
         # Load liked direction once for embed_guidance
         _liked_dir = None
+        _value_fn = None
         if embed_guidance and refinement_key_input:
             _liked_dir = self._load_liked_direction(refinement_key_input)
             if _liked_dir is None:
                 print("[FunPackSceneChain] embed_guidance: no liked direction found (need 3+ liked generations)")
             else:
-                print(f"[FunPackSceneChain] embed_guidance: active, strength={embed_guidance_strength}, direction shape={list(_liked_dir.shape)}")
+                _value_fn = self._load_value_function(refinement_key_input)
+                mode = f"value function ({_value_fn.n_trained} samples)" if _value_fn else "fixed direction"
+                print(f"[FunPackSceneChain] embed_guidance: active via {mode}, strength={embed_guidance_strength}")
 
         first_scene_seed = self._scene_seed(scene_conditionings[0])
         if first_scene_seed is None:
@@ -1487,7 +1511,7 @@ class FunPackLTXAVSceneChainSampler:
                     guide_tail = 0
 
             if embed_guidance and _liked_dir is not None:
-                _eg_old_wrapper = self._build_embed_guidance_wrapper(model, _liked_dir, embed_guidance_strength)
+                _eg_old_wrapper = self._build_embed_guidance_wrapper(model, _liked_dir, embed_guidance_strength, value_fn=_value_fn)
             sampled = self._sample_chunk(
                 model, sampler, sigmas, scene_seed, cfg, scene_positive, scene_negative, chunk,
             )

value_function.py

@@ -0,0 +1,112 @@
+"""Online value function for reward-guided sampling.
+
+Trains a small MLP incrementally on user-rated generations.
+At inference, provides ∂reward/∂conditioning as a per-step steering gradient.
+"""
+
+import os
+import random
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class OnlineValueFunction(nn.Module):
+    MIN_SAMPLES = 5
+    BUFFER_SIZE = 100
+    BATCH_SIZE = 16
+    TRAIN_STEPS = 20
+
+    def __init__(self, hidden_dim):
+        super().__init__()
+        self.hidden_dim = hidden_dim
+        mid = min(256, hidden_dim // 4)
+        self.net = nn.Sequential(
+            nn.Linear(hidden_dim, mid),
+            nn.SiLU(),
+            nn.Linear(mid, 64),
+            nn.SiLU(),
+            nn.Linear(64, 1),
+        )
+        self._optimizer = None
+        self.buffer_c = []  # compressed conditioning tensors [hidden_dim]
+        self.buffer_r = []  # reward floats
+        self.n_trained = 0
+
+    @property
+    def optimizer(self):
+        if self._optimizer is None:
+            self._optimizer = torch.optim.Adam(self.parameters(), lr=2e-3)
+        return self._optimizer
+
+    @staticmethod
+    def compress(conditioning):
+        """[1, seq, D] or [seq, D] → [D] mean vector."""
+        c = conditioning.float()
+        if c.dim() == 3:
+            c = c.squeeze(0)
+        return c.mean(dim=0)
+
+    def forward(self, c):
+        return self.net(c)
+
+    def train_on(self, conditioning, reward):
+        """Add a new (conditioning, reward) sample and do a few training steps."""
+        c = self.compress(conditioning).detach().cpu()
+        self.buffer_c.append(c)
+        self.buffer_r.append(float(reward))
+        if len(self.buffer_c) > self.BUFFER_SIZE:
+            self.buffer_c = self.buffer_c[-self.BUFFER_SIZE:]
+            self.buffer_r = self.buffer_r[-self.BUFFER_SIZE:]
+        if len(self.buffer_c) < 2:
+            return
+        device = next(self.parameters()).device
+        for _ in range(self.TRAIN_STEPS):
+            idx = random.sample(range(len(self.buffer_c)), min(self.BATCH_SIZE, len(self.buffer_c)))
+            c_b = torch.stack([self.buffer_c[i] for i in idx]).to(device)
+            r_b = torch.tensor([self.buffer_r[i] for i in idx], device=device).unsqueeze(1)
+            self.optimizer.zero_grad()
+            F.mse_loss(self.forward(c_b), r_b).backward()
+            self.optimizer.step()
+        self.n_trained += 1
+
+    def gradient(self, conditioning):
+        """∂reward/∂conditioning — same shape as input."""
+        c_in = conditioning.detach().float().requires_grad_(True)
+        reward = self.forward(self.compress(c_in).unsqueeze(0))
+        reward.backward()
+        return c_in.grad.to(conditioning.dtype)
+
+    def is_ready(self):
+        return len(self.buffer_c) >= self.MIN_SAMPLES
+
+    def save(self, path):
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+        buf = torch.stack(self.buffer_c) if self.buffer_c else torch.zeros(0, self.hidden_dim)
+        torch.save({
+            "hidden_dim": self.hidden_dim,
+            "state_dict": self.state_dict(),
+            "buffer_c": buf,
+            "buffer_r": list(self.buffer_r),
+            "n_trained": self.n_trained,
+        }, path)
+
+    @classmethod
+    def load(cls, path):
+        data = torch.load(path, map_location="cpu", weights_only=False)
+        vf = cls(hidden_dim=data["hidden_dim"])
+        vf.load_state_dict(data["state_dict"])
+        buf = data["buffer_c"]
+        vf.buffer_c = [buf[i] for i in range(len(buf))]
+        vf.buffer_r = list(data["buffer_r"])
+        vf.n_trained = data.get("n_trained", 0)
+        return vf
+
+    @classmethod
+    def load_or_create(cls, path, hidden_dim=None):
+        if os.path.exists(path):
+            try:
+                return cls.load(path)
+            except Exception:
+                pass
+        return cls(hidden_dim=hidden_dim) if hidden_dim is not None else None

web/funpack_studio.js

@@ -69,7 +69,7 @@ function defaultSettings() {
     refinement_key: "",
     overrides: { refinement_key: false, feedback_prompt: false, user_intent_prompt: false, negative_prompt: false },
     scene_builder: { mode: "Pass-through", scene: NONE_SENTINEL, scene_name: "", aliases: "", scene_positive: "", scene_negative: "" },
-    refiner: { mode: "Refine", advisor_mode: "Off", advisor_thinking: true, prompt_repair: true, im_feeling_lucky: false, reset_session: false, feedback_prompt: "", user_intent_prompt_override: "", negative_prompt: "", temporal_style: "natural", split_by_transitions: false, split_transition_placement: "start", reference_injection: false, vision_conditioning: true },
+    refiner: { mode: "Refine", advisor_mode: "Off", advisor_thinking: true, prompt_repair: true, im_feeling_lucky: false, reset_session: false, feedback_prompt: "", user_intent_prompt_override: "", negative_prompt: "", temporal_style: "natural", split_by_transitions: false, split_transition_placement: "start", reference_injection: false, vision_conditioning: true, value_guidance: false },
     advisor_llm: { enabled: false, model_path: "huihui-ai/Huihui-Qwen3-8B-abliterated-v2", dtype: "bfloat16" },
     loras: [],
     loras_config: { mode: "ltx2", per_block: false },
@@ -933,6 +933,12 @@ function openPanel(node) {
     body.append(el("div", "funpack-studio-hint", "Use the reference image as visual context when encoding the prompt with Gemma3. Disable to test without vision conditioning."));
     body.append(row("Vision conditioning", visCondToggle.wrap));
 
+    if (settings.refiner.value_guidance === undefined) settings.refiner.value_guidance = false;
+    const valueGuidanceToggle = toggleEl(!!settings.refiner.value_guidance, "Value guidance training");
+    valueGuidanceToggle.inp.addEventListener("change", () => { settings.refiner.value_guidance = valueGuidanceToggle.inp.checked; });
+    body.append(el("div", "funpack-studio-hint", "Train an online value function on rated generations. After 5+ ratings (good and bad), the chain sampler can use it for per-step gradient steering. Enable 'embed_guidance' in the sampler to use."));
+    body.append(row("Value guidance", valueGuidanceToggle.wrap));
+
     body.append(sectionTitle("Negative prompt"));
     body.append(el("div", "funpack-studio-hint",
       "Encoded via CLIP and output as negative conditioning. Skipped when negative_conditioning input is connected."));