value_function.py

"""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 = 10
    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."""
        with torch.inference_mode(False), torch.enable_grad():
            c = self.compress(conditioning.clone()).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."""
        mlp_device = next(self.parameters()).device
        orig_device = conditioning.device
        with torch.inference_mode(False), torch.enable_grad():
            # new_empty inside inference_mode(False) creates a normal (non-inference) tensor.
            # copy_ fills it from the inference tensor without inheriting the flag.
            c_fresh = torch.empty(conditioning.shape, dtype=torch.float32, device=mlp_device)
            c_fresh.copy_(conditioning)
            c_fresh.requires_grad_(True)
            reward = self.forward(self.compress(c_fresh).unsqueeze(0))
            reward.backward()
        return c_fresh.grad.to(orig_device, conditioning.dtype)

    def ascend(self, conditioning):
        """Gradient ascent on conditioning until reward plateaus. Self-terminating, no user params."""
        if not self.is_ready():
            return conditioning
        mlp_device = next(self.parameters()).device
        orig_device = conditioning.device
        orig_dtype = conditioning.dtype
        step_size = conditioning.float().norm().item() * 0.005
        # Cap total displacement at 8% of original norm — prevents stripping action/content
        max_displacement = conditioning.float().norm().item() * 0.08
        prev_reward = -float("inf")
        with torch.inference_mode(False), torch.enable_grad():
            c_orig = torch.empty(conditioning.shape, dtype=torch.float32, device=mlp_device)
            c_orig.copy_(conditioning)
            c = c_orig.clone()
            for _ in range(50):
                c = c.detach().requires_grad_(True)
                reward = self.forward(self.compress(c).unsqueeze(0))
                reward.backward()
                reward_val = reward.item()
                grad = F.normalize(c.grad.float(), dim=-1)
                c = c.detach() + step_size * grad
                # Project back if displacement exceeds cap
                delta = c - c_orig
                disp = delta.norm()
                if disp > max_displacement:
                    c = c_orig + delta * (max_displacement / disp)
                if reward_val - prev_reward < 1e-3:
                    break
                prev_reward = reward_val
        return c.detach().to(orig_device, orig_dtype)

    def search(self, conditioning, n=16):
        """Monte Carlo conditioning search: score N random perturbations, return the best."""
        if not self.is_ready():
            return conditioning
        mlp_device = next(self.parameters()).device
        orig_device = conditioning.device
        orig_dtype = conditioning.dtype
        eps = conditioning.float().norm().item() * 0.025
        with torch.inference_mode(False), torch.no_grad():
            c_base = torch.empty(conditioning.shape, dtype=torch.float32, device=mlp_device)
            c_base.copy_(conditioning)
            best = c_base
            best_score = float(self.forward(self.compress(c_base).unsqueeze(0)).item())
            for _ in range(n):
                noise = F.normalize(torch.randn_like(c_base).flatten(), dim=0).reshape(c_base.shape)
                candidate = c_base + eps * noise
                score = float(self.forward(self.compress(candidate).unsqueeze(0)).item())
                if score > best_score:
                    best_score = score
                    best = candidate
        return best.detach().to(orig_device, orig_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