digital-garbageandClaude Opus 4.8
9fb0ee267b
Route the conditioning variant from per-config outcome history, via the one charter-safe surface: the INTERNAL value-function ascend+search (not user knobs, so no WYSIWYG / hidden-settings violation). Sampler-owned visual ops stay on the sampler (boundary law kept). Directly serves the original complaint (Perfect not locking, Awful reproducing): - _v2_path_conditioning_plan: matching-arm Beta posteriors -> liked-rate mean -> "lock" (>=0.6: hold the liked conditioning, skip the stochastic search), "explore" (<=0.4: widen the search to flee a disliked region), else "normal" - value_function.search gains eps_scale to widen/tighten the search radius - threaded conditioning_plan through ascend -> apply_scene_refinement_keys -> finalize; computed once per run from path_outcomes + prompt concepts; logged Only active once the per-key VF is ready, since it rides the existing ascend path. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
181 lines
7.4 KiB
Python
181 lines
7.4 KiB
Python
"""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, eps_scale=1.0):
|
|
"""Monte Carlo conditioning search: score N random perturbations, return the best.
|
|
|
|
eps_scale widens (>1) or tightens (<1) the perturbation radius. The path planner uses
|
|
a wider radius to push a disliked config region toward a different conditioning, and the
|
|
caller skips this search entirely to LOCK a liked region's conditioning in place."""
|
|
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 * float(eps_scale)
|
|
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
|
|
|
|
|