SolitaryThinker
a9e5f6ee7a
Directory rename (git mv, history preserved) plus rewrite of all references: absolute imports in tests, the zero-dep runner, docstrings, comments, and the README. No behavior change. Run: python3 -m pytest v2/tests/ -q ; python3 v2/run_tests.py ; python3 -m v2.examples
150 lines
8.1 KiB
Python
150 lines
8.1 KiB
Python
"""DiffusionNFT — likelihood-free RL (design_v3 §10, §9.2 C2; repo: train/methods/rl/diffusion_nft).
|
||
|
||
The landed RL method, the C2-split's reason for being: it is **likelihood-free** — no log-probs to
|
||
match — so consistency is the C2 *behavioral* rung (seeded final sample + prediction-space identity:
|
||
``old_deviate`` / ref-MSE), declared on the recipe. Sampling is from the **old** (decay-blended
|
||
behavior) policy, NOT the student (the §8.4 caveat the WeightSyncPlan must carry). K samples of one
|
||
prompt form a homogeneous group; group-relative advantages drive the update; the shared prompt
|
||
encodes once via the feature cache (the §10 ``24×`` text-encode reduction).
|
||
|
||
NFT objective (repo diffusion_nft.py):
|
||
positive = β·forward + (1−β)·old ; implicit_neg = (1+β)·old − β·forward
|
||
x0⁺ = xt − t·positive ; x0⁻ = xt − t·implicit_neg
|
||
policy_loss = r·‖x0⁺−x0‖²/wf⁺/β + (1−r)·‖x0⁻−x0‖²/wf⁻/β (r = clamp(adv/clip/2+0.5, 0, 1))
|
||
loss = policy_loss + kl_β·‖forward − ref‖² ; old_deviate = ‖forward − old‖²
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import numpy as np
|
||
|
||
from ..._enums import ConsistencyLevel
|
||
from ...models.common import cached_text_encode
|
||
from ...request import DiffusionParams, TaskType, make_request
|
||
from ..rewards import build_multi_reward_scorer
|
||
from ..rollout import Trajectory, TrajectoryBuffer
|
||
from ..weight_sync import WeightRole, WeightSyncPlan
|
||
from .base import TrainingMethod, new_instance
|
||
|
||
|
||
class DiffusionNFTMethod(TrainingMethod):
|
||
name = "diffusion_nft"
|
||
consistency = ConsistencyLevel.C2 # likelihood-free behavioral identity
|
||
|
||
def __init__(self, student_instance, old_instance, reference_instance, *,
|
||
reward_fn: dict | None = None, num_video_per_prompt: int = 4, rollout_steps: int = 3,
|
||
beta: float = 0.1, kl_beta: float = 1e-4, adv_clip_max: float = 5.0,
|
||
num_inner_timesteps: int = 2, lr: float = 0.05, decay_type: int = 1, **kw):
|
||
super().__init__(student_instance, lr=lr, **kw)
|
||
self.old = old_instance # behavior policy (sampled from!)
|
||
self.reference = reference_instance # frozen KL anchor
|
||
self.K = num_video_per_prompt
|
||
self.rollout_steps = rollout_steps
|
||
self.beta = beta
|
||
self.kl_beta = kl_beta
|
||
self.adv_clip_max = adv_clip_max
|
||
self.num_inner_timesteps = num_inner_timesteps
|
||
self.decay_type = decay_type
|
||
self.scorer = build_multi_reward_scorer(reward_fn or {"pickscore": 1.0, "clipscore": 1.0})
|
||
self.old_sync = WeightSyncPlan(role=WeightRole.OLD_POLICY)
|
||
# init old + reference from student (NFT inits old from student)
|
||
self.old.component("transformer").copy_from(self.student_dit)
|
||
self.reference.component("transformer").copy_from(self.student_dit)
|
||
|
||
def manages_optimization(self) -> bool: # RL owns its sample→score→train cadence
|
||
return True
|
||
|
||
def get_grad_clip_targets(self, iteration: int = 0) -> dict:
|
||
return {"student": self.student_dit} # NFT clips internally (matches repo)
|
||
|
||
@staticmethod
|
||
def _decay(step: int, decay_type: int) -> float:
|
||
if decay_type == 1:
|
||
return min(step * 0.001, 0.5)
|
||
if decay_type == 2:
|
||
return min(max(0, step - 75) * 0.0075, 0.999)
|
||
return 0.0
|
||
|
||
def _sample_group(self, prompt: str, base_seed: int) -> TrajectoryBuffer:
|
||
"""K rollouts of one prompt FROM THE OLD POLICY via the shared loop (homogeneous group)."""
|
||
buf = TrajectoryBuffer()
|
||
for k in range(self.K):
|
||
seed = base_seed * 1000 + k
|
||
req = make_request(TaskType.T2V, self.old.card.model_id, prompt,
|
||
diffusion=DiffusionParams(num_steps=self.rollout_steps, seed=seed))
|
||
res = self._rollout(req, instance=self.old) # behavior policy = OLD
|
||
buf.add(Trajectory(request_id=req.request_id, prompt=prompt, seed=seed,
|
||
latents=np.asarray(res.outputs["latents"], dtype="float32"),
|
||
behavior=res.behavior, weights_version=self.old.weights_version))
|
||
return buf
|
||
|
||
def managed_train_step(self, batch: dict, iteration: int) -> tuple[dict, dict]:
|
||
self.iteration = iteration
|
||
prompts = batch["prompts"]
|
||
seeds = batch.get("seeds", list(range(len(prompts))))
|
||
policy_losses, kl_losses, old_devs, gnorms, all_adv, all_reward = [], [], [], [], [], []
|
||
|
||
for pi, prompt in enumerate(prompts):
|
||
group = self._sample_group(prompt, seeds[pi])
|
||
media = [t.latents for t in group]
|
||
scored = self.scorer.score(media, [prompt] * len(group))
|
||
rewards = scored["avg"]
|
||
adv = (rewards - rewards.mean()) / (rewards.std() + 1e-4) # group-relative
|
||
adv = np.clip(adv, -self.adv_clip_max, self.adv_clip_max)
|
||
all_reward.extend(rewards.tolist())
|
||
all_adv.extend(adv.tolist())
|
||
|
||
emb = cached_text_encode(self.student, prompt) # encoded ONCE for the group
|
||
old_dit, ref_dit = self.old.component("transformer"), self.reference.component("transformer")
|
||
for k, traj in enumerate(group):
|
||
x0 = traj.latents
|
||
r = float(np.clip(adv[k] / self.adv_clip_max / 2.0 + 0.5, 0.0, 1.0))
|
||
rng = np.random.default_rng(traj.seed ^ 0x4F7)
|
||
for _ in range(self.num_inner_timesteps):
|
||
t = float(rng.uniform(0.05, 0.95))
|
||
noise = rng.standard_normal(x0.shape).astype("float32")
|
||
xt = ((1.0 - t) * x0 + t * noise).astype("float32")
|
||
fwd = self.student_dit(xt, emb, t)
|
||
old_p = old_dit(xt, emb, t)
|
||
ref_p = ref_dit(xt, emb, t)
|
||
positive = self.beta * fwd + (1 - self.beta) * old_p
|
||
neg = (1.0 + self.beta) * old_p - self.beta * fwd
|
||
x0p = xt - t * positive
|
||
x0n = xt - t * neg
|
||
wfp = max(float(np.mean(np.abs(x0p - x0))), 1e-5)
|
||
wfn = max(float(np.mean(np.abs(x0n - x0))), 1e-5)
|
||
pos_loss = float(np.mean((x0p - x0) ** 2) / wfp)
|
||
neg_loss = float(np.mean((x0n - x0) ** 2) / wfn)
|
||
policy_loss = (r * pos_loss / self.beta + (1 - r) * neg_loss / self.beta) * self.adv_clip_max
|
||
kl = float(np.mean((fwd - ref_p) ** 2))
|
||
policy_losses.append(policy_loss)
|
||
kl_losses.append(kl)
|
||
old_devs.append(float(np.mean((fwd - old_p) ** 2)))
|
||
# update student toward x0 along the (reward-weighted) positive direction
|
||
v_target = ((xt - x0) / max(t, 1e-3)).astype("float32")
|
||
_, gn = self.student_dit.mse_grad_step(xt, emb, t, v_target, self.lr * max(r, 0.05))
|
||
gnorms.append(gn)
|
||
|
||
# update OLD ← decay-blend student (the behavior-policy role); bump version + invalidate caches
|
||
self.old_sync.decay = self._decay(iteration, self.decay_type)
|
||
new_ver = self.old_sync.apply(self.student_dit, self.old.component("transformer"), self.old)
|
||
|
||
metrics = {
|
||
"policy_loss": float(np.mean(policy_losses)),
|
||
"kl_div_loss": float(np.mean(kl_losses)),
|
||
"old_deviate": float(np.mean(old_devs)),
|
||
"grad_norm/student": float(np.mean(gnorms)),
|
||
"reward_mean": float(np.mean(all_reward)),
|
||
"advantage_std": float(np.std(all_adv)),
|
||
"old_decay": self.old_sync.decay,
|
||
"old_weights_version": new_ver,
|
||
}
|
||
loss_map = {"policy_loss": metrics["policy_loss"], "kl_div_loss": metrics["kl_div_loss"]}
|
||
return loss_map, metrics
|
||
|
||
def train_step(self, batch: dict, iteration: int) -> tuple[dict, dict]:
|
||
return self.managed_train_step(batch, iteration)
|
||
|
||
|
||
def build_diffusion_nft(card, **kw) -> DiffusionNFTMethod:
|
||
return DiffusionNFTMethod(new_instance(card), new_instance(card), new_instance(card), **kw)
|