hy_worldplay: Clean Forcing drift corrector (training recipe + content-keyed deploy)

integrations/hy_worldplay/drift_correction/.gitignore

@@ -0,0 +1 @@
+outputs/

integrations/hy_worldplay/drift_correction/README.md

@@ -0,0 +1,62 @@
+<!--
+SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+SPDX-License-Identifier: Apache-2.0
+-->
+
+# Clean Forcing drift corrector for HY-WorldPlay
+
+Long autoregressive rollouts drift: each chunk conditions on self-generated
+history and small errors compound (saturation runaway, texture mush). This
+module ships a trained corrector — a frozen-base LoRA (r16 on the
+self-attention q/k/v/o projections, ~0.3% params) trained from the model's
+own closed-loop rollouts with a counterfactual clean-history teacher, zero
+real videos — plus the full recipe that produced it.
+
+## Deploy
+
+Point the runner at a corrector checkpoint; everything else is automatic:
+
+```python
+config = dataclasses.replace(
+    RUNNER_HY_WORLDPLAY_WAN_I2V_5B,
+    drift_corrector=Path("lora_v2.pt"),  # None (default) = exact base behavior
+    drift_corrector_gain=0.5,            # composed with the per-step alpha*(t) gate
+)
+```
+
+Selection is content-keyed per job: static (locked-off camera) trajectories
+run the untouched base weights — static scenes measure *negative* drift, so
+correction there is pure artifact cost — while commanded-motion trajectories
+apply the LoRA at ``alpha*(t) x gain`` per denoise step. The LoRA stays
+unfused on motion jobs (a single-scale weight merge cannot express the
+per-timestep gate); static jobs skip module surgery entirely.
+
+The trained v2 checkpoint (~30 MB ``.pt``) is distributed separately; see
+the PR / release notes for the download link.
+
+## Reproduce the corrector
+
+All scripts run from the repo root on one GPU (~2.5 GPU-days end to end):
+
+```bash
+python drift_correction/gen_first_frames.py     # optional T2V seeds (PROMPTS_FILE=...)
+python drift_correction/build_pairs.py          # strafe-loop rollouts -> pair clips
+python drift_correction/gate_faithful.py        # step-0 alpha*(t) go/no-go diagnostic
+python drift_correction/train_v1.py             # counterfactual-teacher LoRA
+python drift_correction/train_v2.py             # DAgger pool + drift-contraction round
+python drift_correction/eval_rollouts.py        # gain-sweep rollouts
+python drift_correction/score_drift.py          # drift / dynamics / progression / seams
+python drift_correction/demo_static.py          # static-scene suite (+ make_sbs.py)
+```
+
+Train only if the gate passes (the drift gap is systematic: ``alpha*`` high);
+the measured ``alpha*(t)`` profile doubles as the deploy gate in
+``hy_worldplay/_drift_corrector.py``.
+
+## Design notes and background
+
+- Library shape (host adapter vs host-agnostic core), gate math, and the
+  cross-host playbook: ``LIBRARY_DESIGN.md`` on the working branch
+  (`wenqingw-nv/flashdreams-wq` @ ``hy-worldplay-counterfactual-forcing``).
+- Method, result tables, and paper-host reproduction:
+  https://gitlab-master.nvidia.com/wenqingw/clean_forcing

integrations/hy_worldplay/drift_correction/_lora.py

@@ -0,0 +1,127 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""LoRA corrector r_phi as low-rank deltas on the frozen DiT's self-attention."""
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+DEFAULT_TARGETS = ("self_attn.q", "self_attn.k", "self_attn.v", "self_attn.o")
+"""Module-name suffixes wrapped by :func:`apply_lora`; the HY DiT's
+self-attention projections are ``blocks.{i}.self_attn.{q,k,v,o}``."""
+
+
+class LoRALinear(nn.Module):
+    """Frozen base linear plus a runtime-gated low-rank delta.
+
+    ``B`` is zero-initialized so the wrapped module starts as an exact
+    identity over the base (any ``scale`` gives the base output until
+    training moves ``B``). ``scale`` is the runtime gain: ``0`` recovers the
+    frozen base (the clean-history teacher pass), ``1`` the corrected pass.
+    The A/B path runs in fp32 regardless of the base dtype.
+    """
+
+    def __init__(self, base: nn.Linear, rank: int = 16):
+        super().__init__()
+        self.base = base
+        for p in self.base.parameters():
+            p.requires_grad_(False)
+        self.A = nn.Linear(base.in_features, rank, bias=False)
+        self.B = nn.Linear(rank, base.out_features, bias=False)
+        nn.init.normal_(self.A.weight, std=1.0 / rank)
+        nn.init.zeros_(self.B.weight)
+        self.scale = 1.0
+
+    def forward(self, x: Tensor) -> Tensor:
+        out = self.base(x)
+        if self.scale != 0:
+            delta = self.B(self.A(x.to(self.A.weight.dtype)))
+            out = out + self.scale * delta.to(out.dtype)
+        return out
+
+
+def unwrap_compiled(network: object) -> nn.Module:
+    """Return the eager module behind a ``torch.compile`` wrapper, if any.
+
+    Accepts ``object``: the DiT lives on the transformer as a plain
+    attribute, which the type checker resolves through ``nn.Module``'s
+    ``__getattr__`` union.
+    """
+    inner = getattr(network, "_orig_mod", network)
+    assert isinstance(inner, nn.Module), type(inner)
+    return inner
+
+
+def apply_lora(
+    model: nn.Module,
+    rank: int = 16,
+    targets: tuple[str, ...] = DEFAULT_TARGETS,
+) -> list[str]:
+    """Wrap every target linear in ``model`` with a :class:`LoRALinear`.
+
+    Args:
+        model: The (unwrapped) DiT network; pass ``network._orig_mod`` when
+            ``torch.compile`` is active so the wrapping survives recompiles.
+        rank: LoRA rank.
+        targets: Module-name suffixes to wrap.
+
+    Returns:
+        Fully qualified names of the wrapped modules.
+    """
+    wrapped = []
+    for mname, module in list(model.named_modules()):
+        for cname, child in list(module.named_children()):
+            full = f"{mname}.{cname}" if mname else cname
+            if isinstance(child, nn.Linear) and any(t in full for t in targets):
+                setattr(module, cname, LoRALinear(child, rank).to(child.weight.device))
+                wrapped.append(full)
+    return wrapped
+
+
+def set_lora_scale(model: nn.Module, scale: float) -> None:
+    """Set the runtime gain on every :class:`LoRALinear` in ``model``."""
+    for m in model.modules():
+        if isinstance(m, LoRALinear):
+            m.scale = scale
+
+
+def lora_parameters(model: nn.Module) -> list[nn.Parameter]:
+    """Collect the trainable A/B parameters in deterministic module order."""
+    ps: list[nn.Parameter] = []
+    for m in model.modules():
+        if isinstance(m, LoRALinear):
+            ps += list(m.A.parameters()) + list(m.B.parameters())
+    return ps
+
+
+def save_lora(model: nn.Module, path) -> None:
+    """Save the LoRA parameters (index-keyed, CPU) to ``path``."""
+    torch.save(
+        {"lora": {i: p.detach().cpu() for i, p in enumerate(lora_parameters(model))}},
+        path,
+    )
+
+
+def load_lora(model: nn.Module, path) -> None:
+    """Load :func:`save_lora` output into an already-wrapped ``model``."""
+    sd = torch.load(path, map_location="cpu", weights_only=False)["lora"]
+    params = lora_parameters(model)
+    assert len(sd) == len(params), (
+        f"checkpoint has {len(sd)} LoRA tensors but the model exposes "
+        f"{len(params)}; rank or target mismatch."
+    )
+    for i, p in enumerate(params):
+        p.data.copy_(sd[i].to(p.device, p.dtype))

integrations/hy_worldplay/drift_correction/_pairs.py

@@ -0,0 +1,102 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Loading and counterfactual-history helpers for strafe-loop pair clips."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+from hy_worldplay._action import HyWorldPlayCtrl
+from torch import Tensor
+
+TOKENS_PER_FRAME = 880
+"""Post-patchify tokens per latent frame at 704x1280 (44/2 * 80/2)."""
+
+PATCH_DIM = 192
+"""Patchified feature width (48 channels * 1 * 2 * 2)."""
+
+
+def load_clip(path: Path, device: torch.device | str, dtype: torch.dtype) -> dict:
+    """Load a ``build_pairs.py`` clip; latents land on ``device`` in ``dtype``."""
+    d = torch.load(path, map_location="cpu", weights_only=False)
+    d["latents"] = d["latents"].to(device, dtype)
+    return d
+
+
+def make_ctrl(
+    d: dict, k: int, *, device: torch.device | str, dtype: torch.dtype
+) -> HyWorldPlayCtrl:
+    """Reconstruct the patchified per-AR-step ctrl payload for chunk ``k``.
+
+    For chunks ``>= 1`` the captured ctrl carries an all-zero mask, so the
+    image-latent stamp is a no-op and zero latent/mask reconstruct the
+    rollout conditioning exactly.
+    """
+    assert k >= 1, "chunk 0 carries the stamped image latent; not reconstructable"
+    sl = slice(k * 4, (k + 1) * 4)
+    zeros = torch.zeros(1, 4 * TOKENS_PER_FRAME, PATCH_DIM, device=device, dtype=dtype)
+    return HyWorldPlayCtrl(
+        latent=zeros,
+        mask=torch.zeros_like(zeros),
+        _is_patchified=True,
+        action=d["rollout_action"][..., sl].to(device),
+        viewmats=d["rollout_viewmats"][..., sl, :, :].to(device, dtype),
+        Ks=d["rollout_Ks"][..., sl, :, :].to(device, dtype),
+        memory_frame_indices=d["memory_frame_indices"][k],
+        rollout_viewmats=d["rollout_viewmats"].to(device, dtype),
+        rollout_Ks=d["rollout_Ks"].to(device, dtype),
+        rollout_action=d["rollout_action"].to(device),
+    )
+
+
+def history_of(d: dict, k: int) -> Tensor:
+    """Patchified drifted history for chunk ``k`` (frames ``0 .. 4k-1``)."""
+    return d["latents"][..., : k * 4 * TOKENS_PER_FRAME, :]
+
+
+def chunk_x0(d: dict, k: int) -> Tensor:
+    """Patchified clean latent the rollout produced for chunk ``k``."""
+    return d["latents"][
+        ..., k * 4 * TOKENS_PER_FRAME : (k + 1) * 4 * TOKENS_PER_FRAME, :
+    ]
+
+
+def clean_counterfactual(
+    history: Tensor,
+    *,
+    selected: list[int],
+    lap_latents: int,
+    clean_lap: int = 1,
+) -> Tensor:
+    """Swap the memory-selected frames' content for lap-aligned clean frames.
+
+    Frames already inside laps ``<= clean_lap`` map to themselves. The
+    prefill only reads the selected frames, so this substitution changes the
+    entire effective history while leaving positions and per-frame
+    conditioning untouched (the strafe loop makes actions and viewmats
+    identical across laps).
+    """
+    out = history.clone()
+    horizon = (clean_lap + 1) * lap_latents
+    for idx in selected:
+        if idx < horizon:
+            continue
+        src = (idx % lap_latents) + clean_lap * lap_latents
+        s = slice(idx * TOKENS_PER_FRAME, (idx + 1) * TOKENS_PER_FRAME)
+        d = slice(src * TOKENS_PER_FRAME, (src + 1) * TOKENS_PER_FRAME)
+        out[..., s, :] = history[..., d, :]
+    return out