feat(experimental): add WAN ReFL as a self-contained training package

experimental/refl/README.md

@@ -0,0 +1,56 @@
+# experimental/refl — WAN ReFL/BPTT (differentiable reward backprop)
+
+The first self-contained UniRL recipe: direct reward backprop (ReFL / DRaFT)
+for WAN video models. Two colocated roles — a `ReflActorRole` (FSDP WAN +
+grad BPTT sampling + optimizer) and a frozen differentiable reward
+(`RewardService`) — run a 3-RPC step under the distributed `enable_grad()`
+context: generate → score → backward, then optimizer step. No advantages,
+no replay, no rollout engine, no weight sync.
+
+## Launch
+
+One command per config; a Ray cluster must be up (`ray start --head`).
+
+```bash
+# WAN 2.1 T2V + VideoAlign (Qwen2-VL VQ/MQ/TA) reward
+export PRETRAINED_MODEL=/path/to/Wan2.1-T2V-1.3B-Diffusers \
+       VIDEOALIGN_MODEL_PATH=/path/to/VideoReward \
+       DATA_PATH=/path/to/prompts.txt
+RAY_ADDRESS=auto python -m experimental.refl.run --config-name=wan21_t2v_videoalign_refl num_devices=8
+
+# WAN 2.2 I2V + Face-identity reward (first frame via (image, condition)
+# MediaRef; face reference via per-sample metadata ref_video_path)
+pip install -r experimental/refl/reward/face/requirements.txt
+export PRETRAINED_MODEL=/path/to/Wan2.2-I2V-A14B-Diffusers \
+       FACE_MODEL_PATH=/path/to/antelodev2 \
+       DATA_PATH=/path/to/i2v_prompts.jsonl
+RAY_ADDRESS=auto python -m experimental.refl.run --config-name=wan22_i2v_face_refl num_devices=8
+```
+
+## Layout
+
+| Path | What |
+|---|---|
+| `trainer.py` | `REFLTrainer(BaseTrainer)` — driver: wiring + the 3-RPC train step |
+| `roles.py` | `ReflActorRole(Remote)` — family-agnostic actor (`pipeline_target` + `model_config`) |
+| `models/` | Per-model BPTT adaptations subclassing the core pipelines (`types.py` defines the contract): `wan21.py`, `wan22.py` — mirrors `unirl/models/` (graduates into the matching model packages) |
+| `reward/` | Package-local differentiable rewards (VideoAlign, Face), each with an additive-only `requirements.txt` — mirrors `unirl/reward/` (graduates into it) |
+| `examples/` | Flat Hydra configs (repo-wide schema) — mirrors the top-level `examples/` (graduates into it) |
+
+## Environment
+
+Targets the **locked core stack only** (`transformers>=5.6,<5.7`,
+`peft>=0.20` — see `pyproject.toml`). There are deliberately no
+version-compat branches: a wrong environment fails loudly; align the
+environment, not the code. Reward and actor share one Python process, so
+recipe `requirements.txt` files may only ADD packages, never re-pin the
+core stack.
+
+## Verification
+
+| Config | Hardware | Head | Status |
+| --- | --- | --- | --- |
+| `wan21_t2v_videoalign_refl` (835 rollouts) | 8xH20 | pre-adjustment (`e3c6b940` lineage) | contributor long run — reward curve in PR #210 |
+| `wan21_t2v_videoalign_refl` (2-rollout smoke, full 81f/480x832 geometry) | 8xH20, fleet image | `40b3f4c9` | PASS — grads flow reward → VAE → DiT LoRA |
+| VideoAlign load + differentiable fwd/bwd on transformers 5.6.2 + peft 0.20 | 8xH20 (isolated venv) | `9087a671` | PASS — `grad_abs_mean=3.5e-3` |
+| `wan22_i2v_face_refl` | 8xH20 | current head | pending (needs face assets + I2V dataset) |

experimental/refl/examples/wan21_t2v_videoalign_refl.yaml

@@ -0,0 +1,149 @@
+# @package _global_
+# REFL WAN 2.1 T2V — VideoAlign (Qwen2-VL VQ/MQ/TA) reward.
+#
+# Two roles, always (same shape as examples/diffusion/refl_sd3.yaml): a
+# ReflActorRole (FSDP WAN 2.1 + grad BPTT sampling + optimizer) and the frozen
+# differentiable VideoAlign reward, colocated on the actor's worker slab so
+# decoded video never leaves the GPU. GradContext backprops reward → VAE →
+# DiT LoRA across the sibling RPC boundary.
+
+num_devices: 8
+batch_size: 8
+num_rollouts: 1000
+save_interval: 100
+save_dir: ${oc.env:OUTPUT_DIR,outputs/wan21_t2v_videoalign_refl}
+save_mode: adapter
+max_grad_norm: 1.0
+
+actor:
+  _target_: experimental.refl.roles.ReflActorRole
+  # Family selector — swap pipeline_target + model_config for another family,
+  # no code changes (mirrors ReFLPolicy's pipeline_target contract).
+  pipeline_target: experimental.refl.models.wan21.Wan21ReflPipeline
+  block_class_names: ["WanTransformerBlock"]
+  # REFL loss: -(reward - baseline) / scale * weight + kl_weight * KL.
+  reward_weight: 0.25
+  reward_baseline: 0.0
+  reward_scale: 1.0
+  kl_weight: 0.0
+  strategy:
+    # sampling.eta=0.0 reduces FlowSDE to the deterministic ODE — REFL wants a
+    # deterministic transition on the differentiable path.
+    _target_: unirl.sde.kernels.FlowSDEStrategy
+  model_config:
+    _target_: unirl.models.wan21.config.WAN21PipelineConfig
+    pretrained_model_ckpt_path: ${oc.env:PRETRAINED_MODEL}
+    model_precision: bf16
+    autocast_precision: bf16
+    trajectory_precision: bf16
+    logprob_precision: fp32
+    shift: 5.0
+    max_sequence_length: 512
+  fsdp_cfg:
+    _target_: unirl.train.configs.FSDPConfig
+    param_dtype: bf16
+    # fp32 LoRA master over the bf16 base: AdamW steps (~lr=5e-6) are below
+    # bf16 ULP at lora_A's magnitude, so a bf16 master silently freezes A.
+    # Needs the pinned torch (>=2.11): older FSDP2 asserts uniform dtype over
+    # ALL params in a group; the pinned family checks trainables only.
+    master_dtype: fp32
+    cpu_offload: false
+    mixed_precision: true
+    fsdp_mode: full
+    reshard_after_forward: true
+    # BPTT keeps the full mid→final grad window alive; activation
+    # checkpointing is mandatory for any non-trivial
+    # num_frames × num_inference_steps product.
+    activation_checkpointing: true
+    use_torch_compile: false
+  optimizer_cfg:
+    _target_: unirl.train.backend.base.OptimizerConfig
+    learning_rate: 5.0e-6
+    adam_beta1: 0.9
+    adam_beta2: 0.999
+    adam_epsilon: 1.0e-8
+    weight_decay: 0.0
+  scheduler_cfg:
+    _target_: unirl.train.backend.base.LrSchedulerConfig
+    type: linear_warmup
+    warmup_steps: 10
+    total_steps: ${num_rollouts}
+  lora_cfg:
+    _target_: unirl.train.configs.LoraConfig
+    rank: 64
+    alpha: 128
+    dropout: 0.0
+    bias: none
+    task_type: FEATURE_EXTRACTION
+    # Peft path takes suffix matches → these six suffixes catch both
+    # attn1 (self-attention) and attn2 (cross-attention) LoRA targets
+    # inside every WanTransformerBlock.
+    target_modules:
+      - to_q
+      - to_k
+      - to_v
+      - to_out.0
+      - ffn.net.0.proj
+      - ffn.net.2
+
+reward:
+  _target_: unirl.reward.service.RewardService
+  backend:
+    _target_: experimental.refl.reward.videoalign.VideoAlignRewardScorer
+    base_device: cuda
+    config:
+      _target_: experimental.refl.reward.videoalign.VideoAlignSpec
+      reward_model_path: ${oc.env:VIDEOALIGN_MODEL_PATH}
+      device: cuda
+      batch_size: 1
+      resize_height: 336
+      resize_width: 588
+      micro_batch_size: 1
+      reward_num_frames: 36       # 81 → 36 uniform frames before scoring
+      use_norm: true
+      w_vq: 1.0
+      w_mq: 1.0
+      w_ta: 1.0
+      differentiable: true        # REFL requires grad through the reward forward
+
+data_source:
+  _target_: unirl.data.data_source.MultimodalRLDataSource
+  args:
+    run:
+      data_path: ${oc.env:DATA_PATH}
+      eval_data_path: ${oc.env:EVAL_DATA_PATH,${oc.env:DATA_PATH}}
+      seed: 42
+      shuffle: false
+    algorithm:
+      prompts_per_rollout: ${batch_size}
+
+sampling:
+  _target_: unirl.types.sampling.DiffusionSamplingParams
+  # 25 inference steps with mid_timestep = final_timestep = 24
+  # → DRaFT-1 (only the final step is differentiable).
+  num_inference_steps: 25
+  guidance_scale: 5.0
+  height: 480
+  width: 832
+  # (num_frames - 1) % 4 == 0 is the WAN 2.1 T2V hard constraint.
+  num_frames: 81
+  eta: 0.0
+  samples_per_prompt: 1
+  seed: 42
+  init_same_noise: false
+  sampler_kwargs:
+    negative_prompt: "blurry, low quality, distorted, oversaturated"
+    # Truncated BPTT window: no_grad until ``mid_timestep``, grad on from
+    # ``mid_timestep`` .. ``final_timestep`` (inclusive). 24/24/25 → only
+    # the final step traverses the reward backward pass (DRaFT-1 / ReFL).
+    mid_timestep: 24
+    final_timestep: 24
+    # KL on/off + weight live on actor.kl_weight (single knob; 0 here →
+    # single-pass forward per step, no LoRA-disabled reference recompute).
+
+logging:
+  report_to_wandb: ${oc.decode:${oc.env:REPORT_TO_WANDB,true}}
+  project_name: ${oc.env:WANDB_PROJECT,unirl-refl}
+  run_name: ${oc.env:WANDB_RUN_NAME,wan21_t2v_videoalign_refl}
+  tags: ["wan21", "t2v", "refl", "bptt", "videoalign"]
+  log_media: false

experimental/refl/examples/wan22_i2v_face_refl.yaml

@@ -0,0 +1,138 @@
+# @package _global_
+# REFL WAN 2.2 I2V — Face-identity reward.
+#
+# Two roles, always (same shape as examples/diffusion/refl_sd3.yaml): a
+# ReflActorRole (FSDP WAN 2.2 dual-DiT + grad BPTT sampling + optimizer,
+# LoRA restricted to the low_noise DiT) and the frozen differentiable Face
+# reward, colocated on the actor's worker slab. The I2V first frame rides the
+# data source's (image, condition) MediaRef; the face reference video path
+# rides per-sample metadata (``ref_video_path``) into the reward's records.
+
+num_devices: 8
+batch_size: 8
+num_rollouts: 1000
+save_interval: 100
+save_dir: ${oc.env:OUTPUT_DIR,outputs/wan22_i2v_face_refl}
+save_mode: adapter
+max_grad_norm: 1.0
+
+actor:
+  _target_: experimental.refl.roles.ReflActorRole
+  pipeline_target: experimental.refl.models.wan22.Wan22ReflPipeline
+  block_class_names: ["WanTransformerBlock"]
+  # REFL loss: -(reward - baseline) / scale * weight + kl_weight * KL.
+  reward_weight: 0.1
+  reward_baseline: 0.54
+  reward_scale: 0.16
+  kl_weight: 1.0
+  strategy:
+    # sampling.eta=0.0 reduces FlowSDE to the deterministic ODE — REFL wants a
+    # deterministic transition on the differentiable path.
+    _target_: unirl.sde.kernels.FlowSDEStrategy
+  model_config:
+    _target_: unirl.models.wan22.config.WAN22PipelineConfig
+    pretrained_model_ckpt_path: ${oc.env:PRETRAINED_MODEL}
+    model_precision: bf16
+    autocast_precision: bf16
+    trajectory_precision: bf16
+    logprob_precision: fp32
+    shift: 5.0
+    max_sequence_length: 512
+    boundary_ratio: 0.9
+    num_train_timesteps: 1000
+  fsdp_cfg:
+    _target_: unirl.train.configs.FSDPConfig
+    param_dtype: bf16
+    # fp32 LoRA master over the bf16 base: AdamW steps (~lr=5e-6) are below
+    # bf16 ULP at lora_A's magnitude, so a bf16 master silently freezes A.
+    # Needs the pinned torch (>=2.11): older FSDP2 asserts uniform dtype over
+    # ALL params in a group; the pinned family checks trainables only.
+    master_dtype: fp32
+    cpu_offload: false
+    mixed_precision: true
+    fsdp_mode: full
+    reshard_after_forward: true
+    activation_checkpointing: true
+    use_torch_compile: false
+  optimizer_cfg:
+    _target_: unirl.train.backend.base.OptimizerConfig
+    learning_rate: 2.5e-6
+    adam_beta1: 0.9
+    adam_beta2: 0.999
+    adam_epsilon: 1.0e-8
+    weight_decay: 0.0
+  scheduler_cfg:
+    _target_: unirl.train.backend.base.LrSchedulerConfig
+    type: linear_warmup
+    warmup_steps: 10
+    total_steps: ${num_rollouts}
+  lora_cfg:
+    _target_: unirl.train.configs.LoraConfig
+    rank: 32
+    alpha: 32
+    dropout: 0.0
+    bias: none
+    task_type: FEATURE_EXTRACTION
+    # Inject LoRA only into the low-noise DiT (the final-timestep expert the
+    # BPTT window trains); the high-noise DiT stays frozen.
+    module_prefix: low_noise
+    target_modules:
+      - to_q
+      - to_k
+      - to_v
+      - to_out.0
+      - ffn.net.0.proj
+      - ffn.net.2
+
+reward:
+  _target_: unirl.reward.service.RewardService
+  backend:
+    _target_: experimental.refl.reward.face.FaceRewardScorer
+    base_device: cuda
+    config:
+      _target_: experimental.refl.reward.face.FaceRewardSpec
+      model_path: ${oc.env:FACE_MODEL_PATH,/path/to/antelodev2_face_ckpt}
+      device: cuda
+      batch_size: 1
+      image_size: 112
+      ref_max_frames: 81
+      ref_max_pixels: 230400
+      differentiable: true
+
+data_source:
+  _target_: unirl.data.data_source.MultimodalRLDataSource
+  args:
+    run:
+      data_path: ${oc.env:DATA_PATH,/path/to/wan22_i2v_face_refl_prompts.jsonl}
+      eval_data_path: ${oc.env:EVAL_DATA_PATH,${oc.env:DATA_PATH,/path/to/wan22_i2v_face_refl_prompts.jsonl}}
+      seed: 42
+      shuffle: false
+    algorithm:
+      prompts_per_rollout: ${batch_size}
+
+sampling:
+  _target_: unirl.types.sampling.DiffusionSamplingParams
+  num_inference_steps: 8
+  guidance_scale: 1.0
+  guidance_scale_2: 1.0
+  height: 352
+  width: 640
+  num_frames: 81
+  eta: 0.0
+  samples_per_prompt: 1
+  seed: 42
+  init_same_noise: false
+  sampler_kwargs:
+    # Truncated BPTT window: grad on from mid_timestep..final_timestep — the
+    # last 4 of 8 steps, i.e. the low-noise DiT's sigma range. The per-step
+    # KL against the LoRA-disabled reference is switched/weighted by the
+    # single actor.kl_weight knob.
+    mid_timestep: 4
+    final_timestep: 7
+
+logging:
+  report_to_wandb: ${oc.decode:${oc.env:REPORT_TO_WANDB,true}}
+  project_name: ${oc.env:WANDB_PROJECT,unirl-refl}
+  run_name: ${oc.env:WANDB_RUN_NAME,wan22_i2v_face_refl}
+  tags: ["wan22", "i2v", "refl", "bptt", "face"]
+  log_media: false

experimental/refl/models/types.py

@@ -0,0 +1,40 @@
+"""Recipe-local BPTT stage contract for the refl recipe.
+
+``diffuse_with_grad`` is deliberately NOT part of the core
+:class:`~unirl.models.types.diffusion.DiffusionStage` protocol: concrete
+stages inherit that Protocol *explicitly*, so a protocol-level stub would
+become a real ``None``-returning method on every diffusion stage in the
+repo and make ``hasattr``-based capability checks meaningless. While REFL
+is the only BPTT consumer, the contract lives here; if a second consumer
+appears outside ``experimental/refl``, promote it to core as a separate opt-in
+``@runtime_checkable`` protocol (the ``DifferentiableReward`` /
+``LatentShapeProvider`` idiom), not as a method on ``DiffusionStage``.
+
+Implementors (``Wan21ReflDiffusionStage`` / ``Wan22ReflDiffusionStage``)
+provide::
+
+    diffuse_with_grad(conditions, *, schedule, params, initial_latents=None)
+        -> DiffuseWithGradResult
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+
+
+@dataclass
+class DiffuseWithGradResult:
+    """Output of a recipe stage's ``diffuse_with_grad``.
+
+    ``kl_loss`` is per-sample ``[B]`` (zeros when the KL branch is off) so
+    DP-scattered consumers round-trip each shard's own KL, never a
+    cross-shard aggregate.
+    """
+
+    z_final: torch.Tensor
+    kl_loss: torch.Tensor
+
+
+__all__ = ["DiffuseWithGradResult"]