[RFC] Refactor and simplify the video track

[leviking98z-rgb]

Summary

This RFC proposes a staged refactor of UniRL's video track, covering:

• autoregressive video understanding in examples/ar (primarily Qwen3-Omni);
• diffusion video generation in examples/diffusion;
• rollout-engine integrations;
• video rewards and the standalone reward service;
• the related open and recently merged PRs.

This is a maintainability follow-up to #127 and #25. It does not replace the Video RL Roadmap.

The main recommendation is to treat video as a capability view, not as a new top-level execution framework:

• video generation remains part of the diffusion runtime;
• video understanding remains part of the AR multimodal runtime;
• the two tracks share media contracts, decoding utilities, and reward interfaces where appropriate.

The refactor should start by extracting duplicated runtime behavior and defining contracts. A move-only directory reorganization would preserve most of the current complexity.

Motivation

At the time of this review, the video-related surface contains approximately 32 recipes and about 20K lines of selected model, rollout, reward, and recipe code. The main sources of complexity are:

1. Duplicated diffusion execution loops

   Sampling, SDE-step selection, sparse trajectory storage, autocast, replay, and ReplayResult construction are repeated across model packages such as:

   • unirl/models/wan21/diffusion.py
   • unirl/models/wan22/diffusion.py
   • unirl/models/hunyuan_video/diffusion.py
   • unirl/models/hunyuan_video15/diffusion.py
   • unirl/models/ltx2/diffusion.py

   Similar copies also exist in image-model packages, so this is not inherently video-specific.

2. Fragmented model families

   WAN 2.2 already reuses WAN 2.1 configuration, conditions, text, and VAE behavior, while WAN 2.2 V2V adds another partially duplicated pipeline. The physical package boundaries no longer match the logical ownership boundaries.

3. Duplicated Qwen3-Omni preprocessing

   unirl/models/qwen3_omni and unirl/rollout/engine/vllm_omni/adapters/qwen3_omni.py both participate in message preparation, processor arguments, media handling, token expansion/compression, and replay-condition construction.

   This is a numerical-correctness boundary: rollout and training replay must agree on token IDs, media grids, timing metadata, and prompt boundaries. fix(qwen3-omni): align audio-video rollout numerics #274 is a recent example of a mismatch in this area.

4. Eager adapter registration

   unirl/rollout/engine/sglang_diffusion/adapters/video.py contains multiple model families, and package-level import side effects load adapters that are not selected by the active recipe. fix(rollout/sglang_diffusion): port video adapters off the deleted RolloutReq #272 demonstrated that one stale LTX-2 import could break unrelated video adapters.

5. Overlapping reward implementations

   VideoAlign-related logic exists in core local rewards, the standalone reward service, and experimental ReFL work. Video tensors also cross interfaces in both T,C,H,W and legacy C,T,H,W layouts, requiring repeated permutations and frame-sampling logic.

6. Recipe duplication

   Model/task, algorithm, engine, precision, and topology differences are often encoded by copying complete YAML files. For example, several Flow/Dance/Mix, full/LoRA, and topology variants differ by only a small number of fields.

   The two WAN V2V recipes:

   • examples/diffusion/wan2_v2v.yaml
   • examples/diffusion/wan22_v2v/wan22_v2v_14b.yaml

   are also near-duplicates.

7. Insufficient contract coverage

   Current static recipe-target checks cannot catch import-time adapter failures or rollout/replay contract drift. test: remove tests directory #267 and fix(rollout/sglang_diffusion): port video adapters off the deleted RolloutReq #272 together show why a small enforced CPU test layer is needed before a structural refactor.

Design principles

1. Extract behavior before moving directories.
2. Give every behavior one canonical owner.
3. Prefer protocols and composition over model-family inheritance trees.
4. Keep AR video understanding and diffusion video generation as separate runtimes.
5. Preserve public dotpaths through compatibility re-exports for one deprecation cycle.
6. Keep final experiment configurations fully resolvable and reproducible.
7. Do not require GPU E2E tests for every refactor check; use small CPU contract fixtures where possible.

Proposed architecture

1. A model-neutral diffusion runner

Introduce a shared runtime along these lines:

unirl/diffusion/
├── runner.py
├── trajectory.py
├── latent_spec.py
└── protocols.py

The shared runner should own:

• schedule validation;
• latent/noise initialization orchestration;
• SDE index selection;
• trajectory capture;
• sampling and replay loops;
• autocast context;
• construction of replay results.

A model plugin should provide only model-specific behavior, for example:

encode_conditions(...)
create_latent_spec(...)
predict_noise(...)
apply_guidance(...)
decode(...)

Latent geometry should have one source of truth instead of separate pipeline-side and stage-side shape calculations.

LTX-2 should initially remain a documented specialized path because it co-denoises packed video and audio states. If it is generalized later, trajectory state should use named streams such as video and audio, rather than a permanent ambiguous aux_latents field.

2. Consolidate the WAN family

Suggested logical layout:

unirl/models/wan/
├── common.py
├── geometry.py
├── conditions.py
├── text.py
├── vae.py
├── variants/
│   ├── wan21.py
│   └── wan22.py
└── pipelines/
    ├── t2v.py
    ├── i2v.py
    └── v2v.py

The existing unirl.models.wan21, unirl.models.wan22, and unirl.models.wan22_v2v paths should remain as compatibility re-exports for one release.

HunyuanVideo 1.0 and 1.5 should share the generic runner, but should not be merged solely because they share a product name: their text/model architectures remain distinct.

3. One canonical Qwen3-Omni processor contract

The model package should expose a canonical processor/codec that produces a typed result containing:

• expanded Hugging Face token IDs;
• decoded video/audio tensors;
• media grids and timing/TMRoPE metadata;
• replay conditions;
• prompt boundaries.

The train-side pipeline should consume that result directly. The vLLM-Omni adapter should only perform backend-specific placeholder compression and request serialization.

Video understanding should not inherit from a diffusion-video base class.

4. Lazy rollout-adapter registration

Split the SGLang video adapter by model family:

unirl/rollout/engine/sglang_diffusion/adapters/video/
├── base.py
├── wan.py
├── hunyuan.py
└── ltx.py

The registry should map a model family to a dotpath and import only the selected adapter. Importing one adapter must not import every supported video model.

FastVideo currently has limited recipe coverage and has solver/schedule semantics that need to be made explicit (#251). Until a parity contract exists, consider maintaining it under an experimental engine boundary rather than presenting it as behaviorally interchangeable with trainside and SGLang.

5. One video and reward contract

Use Video[T,C,H,W] as the canonical in-memory layout across rollout and reward interfaces.

Recommended ownership:

• lightweight or differentiable same-process rewards may live in core;
• heavyweight, dependency-conflicting, non-differentiable rewards should live in unirl-reward-service;
• a reward implementation should have only one inference owner;
• frame decoding and sampling should be shared;
• aggregation policies such as first-frame, uniform-K, top-ratio, and temporal aggregation should be reusable wrappers.

The direction of #270 is consistent with this boundary, but its contract tests should be included and enforced.

6. Bounded recipe composition

Avoid both full-file duplication and a deep Hydra Cartesian product. Use a bounded two-level structure:

canonical model-task base
  + algorithm overlay
  + engine/topology overlay

For example:

examples/diffusion/wan/t2v_base.yaml
examples/diffusion/overlays/algorithm/{flow,dance,mix}.yaml
examples/diffusion/overlays/engine/{trainside,sglang,fastvideo}.yaml
examples/diffusion/overlays/topology/{lora_1x8,full_2x8}.yaml

CI should resolve every canonical entry recipe and validate the final configuration. Runs should retain the resolved config and a fingerprint for reproducibility.

Qwen3-Omni 1x4 and 1x8 recipes currently differ in more than hardware topology, including data, optimization, reward, and length settings. Those should first be separated into a training profile and a hardware profile before deduplication.

7. A video support matrix

Add docs/tracks/video.md with:

• model and task;
• canonical recipe;
• trainside/SGLang/vLLM-Omni/FastVideo support;
• reward support;
• status and owner;
• last verified commit and hardware.

The physical examples should still follow their execution domains (ar and diffusion); the support matrix provides the cross-cutting video view.

Required contract tests

Before directory moves or public-path removal, add inexpensive enforced tests for:

• lazy import/registration of every rollout adapter family;
• Video pack, concatenate, slice, and layout round trips;
• shared diffusion sampling and replay using a fake model and fake scheduler;
• schedule/noise determinism and engine parameter propagation;
• full Hydra composition of every canonical recipe;
• Qwen3-Omni processor parity for token IDs, media grids, timing metadata, and replay conditions.

A small nightly GPU matrix can then validate one canonical model family per engine.

Proposed PR sequence

1. Add contract fixtures and tests without moving code.
2. Introduce the diffusion-runner protocol and migrate one WAN and one HunyuanVideo path.
3. Migrate WAN 2.2 and consolidate the WAN family with compatibility imports.
4. Replace eager SGLang adapter imports with lazy registration.
5. Introduce the canonical Qwen3-Omni processor contract.
6. Standardize video/reward layouts and deduplicate frame aggregation and VideoAlign ownership.
7. Introduce bounded recipe composition and remove duplicate recipes.
8. Remove deprecated packages and dotpaths after one compatibility cycle.

Each PR should preserve numerical behavior unless an explicit contract change is documented and measured.

Coordination with current PRs

• fix(qwen3-omni): align audio-video rollout numerics #274: prioritize as a correctness fix; add an Omni processor parity fixture.
• feat(reward-service): add videohpsv3 scorer for T2V rollouts #270: reward-service is the appropriate direction; include the tests with the implementation.
• fix(fastvideo): use UniRL canonical UniPC strategy #251: define solver/schedule behavior explicitly; consider an experimental engine boundary until parity is established.
• feat(experimental): add WAN ReFL as a self-contained training package #210: keep ReFL self-contained under experimental; split tier policy, core enablers, vendored model code, and the algorithm package where possible.
• feat(experimental): tier contract, boundary lint, private-package convention #279, docs(algorithms): placement policy — core canon vs experimental variants (RFC) #280, and feat(experimental): port SD3 image ReFL and remove the legacy core ReFL path #281: the experimental-tier and algorithm-placement contracts are useful inputs to this refactor.
• feat(cosmos3): Cosmos3 omnimodal SFT task adapters (videopred + action BC) #200: adjacent omnimodal SFT work; share generic media contracts where useful, but keep it outside the video RL runtime refactor.

Non-goals

• Removing supported video models in the first phase.
• Unifying AR generation/understanding and diffusion under one video model base.
• Changing sampling numerics as an incidental part of directory cleanup.
• Requiring every experiment to be represented by a standalone, fully duplicated YAML.
• Promoting one-off experimental APIs into core before they have more than one stable consumer.

Completion criteria

• Generic diffusion sampling/replay logic has one canonical implementation for the standard single-stream path.
• WAN common components and task pipelines have one owner.
• Qwen3-Omni rollout and replay consume the same canonical preprocessing result.
• Selecting one rollout adapter does not import every video adapter.
• Video tensors use one documented layout at reward boundaries.
• Heavy video reward inference has one implementation owner.
• Canonical recipes compose and resolve in CI.
• Existing public dotpaths have a documented deprecation window.
• docs/tracks/video.md accurately describes supported model/task/engine/reward combinations.