hy_worldplay: Clean Forcing drift corrector (training recipe + content-keyed deploy)#396
hy_worldplay: Clean Forcing drift corrector (training recipe + content-keyed deploy)#396wenqingw-nv wants to merge 1 commit into
Conversation
Greptile SummaryThis PR adds a trained Clean Forcing drift corrector (LoRA) for HY-WorldPlay's long autoregressive rollouts, plus the full zero-real-data training recipe. The corrector is content-keyed at deploy time: static-camera jobs bypass it entirely, commanded-motion jobs apply the LoRA at
Confidence Score: 5/5Safe to merge. The default path (drift_corrector=None) is a no-op that leaves existing behavior byte-identical, and the corrector path is gated behind an explicit config field. The core deploy hook, LoRA wrapper, and content-keyed dispatch are all correctly implemented. Training recipe scripts are standalone tools that do not affect production inference. The only issues found are minor diagnostic quality nits that do not affect correctness of the shipped corrector or the default no-corrector path. No files require special attention for merge safety. The previously flagged _drift_corrector.py concerns are the items most worth resolving before any follow-on work that reuses the runner across multiple jobs. Important Files Changed
Reviews (4): Last reviewed commit: "hy_worldplay: add Clean Forcing drift co..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| network = network._orig_mod | ||
| params = _apply_lora(network) | ||
|
|
||
| sd = torch.load(checkpoint, map_location="cpu", weights_only=False)["lora"] |
There was a problem hiding this comment.
weights_only=False allows arbitrary code execution
torch.load(..., weights_only=False) uses Python's pickle deserializer, which can execute arbitrary code embedded in the file. checkpoint is a user-controlled path via HyWorldPlayWanI2VRunnerConfig.drift_corrector, so a replaced or tampered .pt file would run attacker code on load. The save_lora format stores only {"lora": {int: Tensor}} — all primitive Python types and tensors — so weights_only=True loads it without issue.
| GATE_ALPHA = {1000.0: 0.81, 960.0: 0.53, 888.8889: 0.53, 727.2728: 0.58} | ||
| """Unbiased alpha*(t) from the step-0 systematicity gate (drift_correction's | ||
| ``outputs/gate/gate_faithful.json``): the systematic fraction of the | ||
| drift-induced error at each of the distilled solver's timesteps. The | ||
| shipped config deploys the LoRA at ``alpha*(t) * gain`` per denoise step.""" |
There was a problem hiding this comment.
GATE_ALPHA dict — deploy and training scripts diverge silently
GATE_ALPHA is defined with identical values in both this file and drift_correction/_rollout.py. If a future gate rerun changes the alpha*(t) profile, only one copy is likely to be updated, causing the deployed per-step gain profile to silently mismatch the training configuration.
| GATE_ALPHA = {1000.0: 0.81, 960.0: 0.53, 888.8889: 0.53, 727.2728: 0.58} | |
| """Unbiased alpha*(t) from the step-0 systematicity gate (drift_correction's | |
| ``outputs/gate/gate_faithful.json``): the systematic fraction of the | |
| drift-induced error at each of the distilled solver's timesteps. The | |
| shipped config deploys the LoRA at ``alpha*(t) * gain`` per denoise step.""" | |
| GATE_ALPHA = {1000.0: 0.81, 960.0: 0.53, 888.8889: 0.53, 727.2728: 0.58} | |
| """Unbiased alpha*(t) from the step-0 systematicity gate (drift_correction's | |
| ``outputs/gate/gate_faithful.json``): the systematic fraction of the | |
| drift-induced error at each of the distilled solver's timesteps. The | |
| shipped config deploys the LoRA at ``alpha*(t) * gain`` per denoise step. | |
| NOTE: This dict is also defined in ``drift_correction/_rollout.py``. Both copies | |
| must be kept in sync; prefer sourcing from a single shared location.""" |
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
58702aa to
d91feb2
Compare
|
Checkpoint availability: the trained v2 corrector LoRA (59 MB) is published as a release asset on the contributing fork for review/testing: https://github.com/wenqingw-nv/flashdreams-wq/releases/tag/clean-forcing-hy-v2 — happy to rehome it (release asset here / HF / internal artifact store) per maintainer preference. Note the deploy dial is inference-time: this single checkpoint serves all configurations (flat gain, alpha*(t) gate, gate x gain); the shipped setting is alpha*(t) x 0.5 with content-keyed selection. |
|
/ok to test d91feb2 |
d91feb2 to
ada922d
Compare
| if cfg.drift_corrector is not None: | ||
| from hy_worldplay._drift_corrector import maybe_apply_drift_corrector | ||
|
|
||
| mode = maybe_apply_drift_corrector( | ||
| self, cfg.drift_corrector, cfg.drift_corrector_gain | ||
| ) | ||
| logger.info(f"[{cfg.runner_name}] drift corrector: {mode}") |
There was a problem hiding this comment.
run() call, stale gain overrides the new one
maybe_apply_drift_corrector installs a gated_pf closure that captures gain from the call site. Each subsequent run() wraps the previous gated_pf in a new one: the outermost closure (latest call's gain) sets the scale, then returns via orig_pf, which is the previous gated_pf — which overrides the scale again with an older captured gain before calling the real predict_flow. If drift_corrector_gain is mutated between runs (e.g. runner.config.drift_corrector_gain = 0.3 after a first run at 0.7), the network executes every subsequent step at the first-call gain (0.7), not the current one, with no error raised.
Moving the application into __init__ (or adding a _drift_corrector_applied guard) would prevent the stacking entirely.
|
Domain scope note: the shipped corrector was trained on urban/outdoor content without people; person-heavy or stylized domains may need a pairs refresh (recipe included, ~1 GPU-day). We validated the refresh path: a 200-step continue-train with ~10% person/scene-diverse pairs improves the walking-person stress scene (drift −47%, best seam profile) while keeping — in fact improving — the original commanded-motion profile (bridge drift cut −58%). Checkpoint to follow with the hosting decision. |
|
Checkpoint update: the shipped LoRA is now v2c2 — a content-diversified refresh (200-step continue-train, 10% person/scene-diverse pairs) that improves BOTH axes: bridge drift cut −58% (was −48%) and the walking-person stress scene fixed (drift −47%, owner-verified side-by-sides). Release asset: https://github.com/wenqingw-nv/flashdreams-wq/releases/tag/clean-forcing-hy-v2c2 (md5 1de6023c...; loads via the identical deploy path, no code change). The earlier v2 asset remains available for reproducing the PR's tables. |
…ntent-keyed deploy) Adds a trained drift corrector for the HY-WorldPlay WAN-5B integration: a frozen-base LoRA (r16 q/k/v/o, ~0.3% params) trained from the model's own strafe-loop rollouts with a counterfactual clean-history teacher -- zero real videos end-to-end. Deploy: HyWorldPlayWanI2VRunnerConfig.drift_corrector (checkpoint path) + drift_corrector_gain. Selection is content-keyed per job: static (locked-off camera) poses run untouched base weights; commanded-motion poses apply the LoRA at alpha*(t) x gain per denoise step, where alpha*(t) is the measured systematic fraction of the drift-induced error at each solver timestep. Measured on 24-chunk (~19 s) rollouts: -48% delta-drift vs base with dynamics ~= base and no added seam pulse; the base's late-horizon saturation runaway is eliminated. Training/eval scripts and the step-0 diagnostic live under integrations/hy_worldplay/drift_correction/ with a README covering the method and the full reproduction recipe. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Wenqing Wang <wenqingw@nvidia.com>
ada922d to
8afdaf1
Compare
1 participant
Long autoregressive HY-WorldPlay rollouts drift: each chunk conditions on self-generated history and small errors compound into saturation/contrast runaway. This PR adds a trained drift corrector (Clean Forcing), the zero-real-data recipe that produced it, and a content-keyed deploy hook.
drift_corrector=None(default) is byte-identical to current behavior.Contents
hy_worldplay/_drift_corrector.py+ two runner-config fields (drift_corrector,drift_corrector_gain) — content-keyed per job: static (locked-off camera) trajectories run untouched base weights (static scenes measure negative drift on this host — correction there is pure artifact cost); commanded-motion trajectories apply the LoRA atalpha*(t) × gainper denoise step.integrations/hy_worldplay/drift_correction/— full recipe, zero real videos: strafe-loop pair construction, step-0alpha*(t)systematicity gate (go/no-go + the deploy gain profile), v1 counterfactual-teacher training, v2 DAgger + drift-contraction, gain-sweep eval, drift/dynamics/progression/seam scoring, static-scene demo suite. README documents the ~2.5 GPU-day reproduction.tests/test_drift_corrector.py(ci_cpu, 6 tests) — LoRA wrapper is a strict identity until a checkpoint is loaded and gated on; content-keyed selection triggers on action labels.Docs
Measured (704×1280, 24-chunk ≈19 s rollouts, held-out trajectories, GB300)
alpha*(t) × 0.5Shipped config cuts the drift metric ~50% with no dynamics loss and no added chunk-boundary pulse; commanded trajectories keep advancing (owner-verified side-by-sides). The method/reproduction details:
drift_correction/README.md(quick start)Deploy notes
alpha*(t)gating cannot merge into one weight set → motion jobs run the LoRA unfused (~0.3% extra matmul); static jobs run the base at zero overhead.Video Comparison
Left: FlashDreams-HY (base), Right: w/ Clean Forcing LoRA drift corrector (better details/color preservation across time):
sbs_base_LEFT_vs_corrgate050_RIGHT_p0_i0_s5042.mp4
fsbs_base_LEFT_vs_corrgate050_RIGHT_p0_i0_s5042.mp4
sbs_base_LEFT_vs_corrgate050_RIGHT_p0_i0_s5042mix.mp4
VBench Ablation: