Skip to content

Fix streaming-inference alignment with training-time contracts#108

Open
study-overflow wants to merge 1 commit into
Robbyant:mainfrom
study-overflow:fix/inference-alignment
Open

Fix streaming-inference alignment with training-time contracts#108
study-overflow wants to merge 1 commit into
Robbyant:mainfrom
study-overflow:fix/inference-alignment

Conversation

@study-overflow

Copy link
Copy Markdown

We hit this while evaluating LingBot-VA on streaming-inference mode: training loss was healthy and teacher-forced open-loop predictions were accurate, but closed-loop rollouts still collapsed with the same checkpoint — a symptom of inference-side train/test mismatch that compounds across chunks.

Fixes two streaming-inference bugs that decouple closed-loop rollout behavior from training.

  1. Cold-start action conditioning now uses physical zero

    • At frame_st_id == 0, the server synthesizes an all-zero action history.
    • Old code built it directly in normalized space, so torch.zeros decoded to the per-channel statistical midpoint (q01 + q99) / 2 instead of a physical zero action.
    • Added _make_zero_action_condition() that starts from physical zero, applies the same quantile normalization as training, and masks unused channels.
  2. KV-cache attention visibility now matches training mask

    • Training mask allows a noisy query to attend only to past clean tokens, same-frame noisy tokens, and past predicted tokens, all within the window.
    • Old inference surfaced every allocated cache slot, letting noisy queries leak same-frame clean answers and peek at future predicted tokens.
    • Added frame_id per slot, set_frame_context(), and _select_valid_cache_slots() to reproduce the training visibility rules exactly.

Validation:

  • py_compile clean; git diff --check clean.
  • Same checkpoint + env seed 100001 + inference seed 0 went from 600-step failure to 293-step success.

Files changed:

  • wan_va/modules/model.py
  • wan_va/wan_va_server.py

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant