SolitaryThinker
db866dfb04
Sweep all 296 v2 modules: simplify verbose docstrings/comments and remove 376 dangling "(design_vN §X)" citations to the now-deleted design docs (v2/README.md is the source of truth). Comment/docstring-only — AST-verified code-identical; the CPU suite holds at 237 passed / 5 pre-existing. Also applies yapf + ruff --fix auto-fixes (import ordering, forward-ref annotation de-quoting under `from __future__ import annotations`; behavior- neutral, suite-confirmed) and adds the legitimate domain terms mot/clen/te to the codespell ignore-list. Remaining ruff (24) + mypy (68 no-untyped-call) findings are pre-existing v2 debt, untouched here.
79 lines
3.8 KiB
Python
79 lines
3.8 KiB
Python
"""End-to-end RL over a cross-model workflow.
|
|
|
|
Probes the training-plane / Workflow boundary: can "rollout == serve + capture" span *two model
|
|
instances*, and can one final reward train an *earlier* model? These tests train both stages of the
|
|
T2I→I2V workflow (`flux-t2i` and `wan-i2v`) from a single final-video reward and assert:
|
|
|
|
* the rollout spans both instances with SDE capture in each (a workflow-level rollout);
|
|
* the SAME final-video advantage trains BOTH generators — including the T2I one, which never sees
|
|
the video directly (end-to-end credit across the model boundary);
|
|
* that credit is *caused by the video reward* — a constant reward (zero advantage) moves nothing;
|
|
* the two stages version independently (two WeightSyncPlans on two instances).
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
|
|
from v2.recipes import build_flux_t2i_card, build_wan_i2v_card
|
|
from v2.training.methods import build_workflow_rl
|
|
|
|
|
|
def _method(**kw):
|
|
return build_workflow_rl(build_flux_t2i_card(), build_wan_i2v_card(),
|
|
num_samples_per_prompt=4, rollout_steps=4, t2i_lr=0.03, i2v_lr=0.03, **kw)
|
|
|
|
|
|
def _batch():
|
|
return {"prompts": ["a fox", "blue sky", "a city"], "seeds": [1, 2, 3]}
|
|
|
|
|
|
def test_rollout_spans_two_instances_with_sde_capture():
|
|
"""A workflow-level rollout: stage 1 on one instance, stage 2 on another, each capturing an SDE
|
|
trajectory (so FlowGRPO can credit both)."""
|
|
m = _method()
|
|
assert m.t2i is not m.i2v
|
|
assert m.t2i.component("transformer") is not m.i2v.component("transformer")
|
|
res_t, _emb_t, image = m._rollout_t2i("a fox", seed=1)
|
|
res_i, _emb_i, video = m._rollout_i2v("a fox", image, seed=1)
|
|
assert any("sde_logprob" in r for r in (res_t.behavior or [])) # T2I stage captured
|
|
assert any("sde_logprob" in r for r in (res_i.behavior or [])) # I2V stage captured
|
|
assert np.asarray(video).ndim == 4 # produced a video
|
|
|
|
|
|
def test_one_final_reward_trains_both_generators():
|
|
m = _method()
|
|
t2i_w0 = m.t2i.component("transformer").w_x.copy()
|
|
i2v_w0 = m.i2v.component("transformer").w_x.copy()
|
|
last = None
|
|
for it in range(20):
|
|
_, last = m.managed_train_step(_batch(), it)
|
|
# the EARLIER model (T2I) is trained by a reward on the FINAL video — credit across the boundary
|
|
assert not np.array_equal(t2i_w0, m.t2i.component("transformer").w_x)
|
|
assert not np.array_equal(i2v_w0, m.i2v.component("transformer").w_x)
|
|
assert last["grad_norm/t2i"] > 0.0 and last["grad_norm/i2v"] > 0.0
|
|
assert 0.5 < last["ppo_ratio_mean"] < 2.0
|
|
|
|
|
|
def test_t2i_credit_is_caused_by_the_video_reward():
|
|
"""Causal check: a *constant* final reward ⇒ zero group advantage ⇒ neither generator moves. So the
|
|
T2I movement in the test above is genuinely caused by the video reward, not by the rollout itself."""
|
|
m = _method()
|
|
m._reward = lambda video: 0.5 # constant ⇒ std 0 ⇒ advantage 0
|
|
t2i_w0 = m.t2i.component("transformer").w_x.copy()
|
|
i2v_w0 = m.i2v.component("transformer").w_x.copy()
|
|
for it in range(10):
|
|
m.managed_train_step(_batch(), it)
|
|
assert np.array_equal(t2i_w0, m.t2i.component("transformer").w_x) # no signal ⇒ no T2I update
|
|
assert np.array_equal(i2v_w0, m.i2v.component("transformer").w_x) # no signal ⇒ no I2V update
|
|
|
|
|
|
def test_stages_version_independently_on_two_instances():
|
|
m = _method()
|
|
_, last = m.managed_train_step(_batch(), 0)
|
|
assert last["t2i_weights_version"] != last["i2v_weights_version"]
|
|
assert m.t2i.version_of("transformer") != "v0"
|
|
assert m.i2v.version_of("transformer") != "v0"
|
|
# each stage's sync touched only its own instance's transformer (frozen text encoders intact)
|
|
assert m.t2i.version_of("text_encoder") == "v0"
|
|
assert m.i2v.version_of("text_encoder") == "v0"
|