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>

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