feat(recipes): Decouple training algorithms from UniRL Core

[YSunLIN]

feat(recipes): Decouple Training Algorithms from UniRL Core

Summary

The core motivation of this PR is to decouple training algorithms from UniRL core, and organize algorithms, configs, rewards, and launch logic in the form of recipes. This reduces framework intrusiveness and allows users to customize model, pipeline, sampling, reward, loss and trainer without modifying the core components.

• Introduce the Trainer + Role abstraction:

  • Goal: Separate distributed role management from concrete algorithm logic, so that the framework handles common role creation, scheduling, and resource management, while recipes own their training algorithms. This avoids repeatedly modifying the framework main flow when adding new algorithms.
  • Details:
    • Add unirl.trainer.Trainer / unirl.trainer.Role, supporting creation, initialization, and scheduling of remote roles from the roles config.
    • Add the generic RewardRole, used to host reward backends in recipes through configuration, and support the differentiable reward path.
    • Support role placement, share_with worker sharing, and role-aware checkpoint save/load.

• Refactor the recipe config structure:

  • Goal: Improve the readability and maintainability of complex training configs, allowing users to clearly describe actor, reward, pipeline, backend, algorithm parameters, and resource topology through configuration.
  • Details:
    • Use roles as the config root to uniformly describe actor, reward, backend, pipeline, algorithm, placement, and other configurations.

• Add training configuration capabilities:

  • Goal: Reduce hardcoding in training strategies, allowing recipes to declaratively control data order, trainable module scope, and learning rate strategy, making experiments easier to reproduce and tune.
  • Details:
    • MultimodalRLDataSource supports run.shuffle.
    • LoraConfig adds module_prefix, supporting LoRA injection only into modules with specific prefixes; the Wan 2.2 I2V recipe applies LoRA only to the low_noise module.
    • The LR scheduler adds the linear_warmup type.

• Supporting adjustment:

  • Goal: Improve coverage for checking target classes inside recipes, avoiding cases where _target_ config errors under recipes cannot be detected by the existing check flow.
  • Details:
    • scripts/check_recipe_targets.py supports scanning _target_ configs under recipes.

At the same time, this PR completes the differentiable generation path and memory optimization capabilities required for WAN ReFL training.

• Add the recipes/refl recipe package:

  • Goal: Decouple ReFL algorithm logic from the framework layer, avoid embedding algorithm implementations in UniRL core, and improve the flexibility for users to customize algorithms and maintain independent recipes.
  • Details:
    • Add REFLTrainer and ReflActorRole on the recipe side, keeping ReFL algorithm orchestration inside the recipe.
    • Provide two types of ReFL algorithms for the WAN series.
    • Provide the wan21_t2v_videoalign_refl config, corresponding to recipes/refl/configs/wan21_t2v_videoalign_refl.yaml, for Wan 2.1 T2V + VideoAlign reward.
    • Provide the wan22_i2v_face_refl config, corresponding to recipes/refl/configs/wan22_i2v_face_refl.yaml, for Wan 2.2 I2V + Face reward.

• Add the ReFL/WAN differentiable training path:

  • Goal: Complete the differentiable generation and cross-RPC gradient propagation capabilities required by ReFL, enabling the generate-score-backward training flow instead of being limited to the previous no-grad rollout + scalar reward paradigm.
  • Details:
    • Add recipe-local pipelines for Wan 2.1 / Wan 2.2, supporting differentiable sampling, KL computation, and reward backpropagation after VAE decode.
    • Extend DiffusionStage with the diffuse_with_grad result type and invocation path.
    • Adjust cross-RPC grad input handling in workers to support the generate-score-backward training flow.

• Introduce a memory-optimized WAN video VAE:

  • Goal: Mitigate the memory bottleneck of differentiable VAE decode in ReFL, making Wan 2.1 / Wan 2.2 ReFL training practically runnable on 8xH20.
  • Details:
    • Replace the current WAN series VAE decode flow with the framework-side WanVideoVAE.
    • Support activation/gradient memory optimization for frozen VAE scenarios, including nested grad checkpoint and activation-grad-only conv.

• Add recipe-local reward dependency declarations:

  • Goal: Isolate third-party dependencies by reward backend, avoiding pollution of the UniRL core environment and reducing dependency conflict risks.
  • Details:
    • Add recipes/refl/rewards/face/requirements.txt.
    • Add recipes/refl/rewards/videoalign/requirements.txt.

Test Plan

The following commands need to be executed on an 8xH20 environment, with a Ray cluster already started. The scripts connect to the existing Ray cluster through RAY_ADDRESS=auto. The /path/to/... values in the examples are placeholders and should be replaced with actual data, model, and reward checkpoint paths before running.

1. Wan 2.1 T2V ReFL + VideoAlign reward

cd /path/to/UniRL
pip install -r recipes/refl/rewards/videoalign/requirements.txt

export DATA_PATH=/path/to/wan21_prompts.txt
export EVAL_DATA_PATH=${DATA_PATH}
export PRETRAINED_MODEL=/path/to/Wan2.1-T2V-1.3B-Diffusers
export VIDEOALIGN_MODEL_PATH=/path/to/VideoReward
export OUTPUT_DIR=./outputs/wan21_t2v_videoalign_refl
export REPORT_TO_WANDB=true
export WANDB_PROJECT=unirl-refl
export WANDB_RUN_NAME=wan21_t2v_videoalign_refl_recipe_opt

bash recipes/refl/scripts/start_wan21_t2v.sh

2. Wan 2.2 I2V ReFL + Face reward

cd /path/to/UniRL
pip install -r recipes/refl/rewards/face/requirements.txt

export PRETRAINED_MODEL=/path/to/Wan2.2-I2V-A14B-Diffusers
export DATA_PATH=/path/to/wan22_i2v_face_refl_prompts.jsonl
export EVAL_DATA_PATH=${DATA_PATH}
export FACE_MODEL_PATH=/path/to/antelodev2
export OUTPUT_DIR=./outputs/wan22_i2v_face_refl_recipe_opt
export REPORT_TO_WANDB=true
export WANDB_PROJECT=unirl-refl
export WANDB_RUN_NAME=wan22_i2v_face_refl_recipe_opt

bash recipes/refl/scripts/start_wan22_i2v.sh

Verification Status

Config Name	Model	Hardware	Status	Notes
wan21_t2v_videoalign_refl	Wan2.1-T2V-1.3B + VideoAlign reward	8xH20	Verified	ReFL/BPTT; num_rollouts=835; batch_size=8; num_inference_steps=25; final-step differentiable path
wan22_i2v_face_refl	Wan2.2-I2V-A14B + Face reward	8xH20	Verified	ReFL/BPTT; num_rollouts=690; batch_size=8; module_prefix=low_noise; num_inference_steps=8

Wan 2.1 T2V ReFL training reward curve

Wan 2.2 I2V ReFL training reward curve

Compatibility / Risk

• Trainer + Role, RewardRole, and recipes/refl are new capabilities that mainly affect the new recipe path; existing non-recipe training entrypoints do not need to be migrated.
• The default value of MultimodalRLDataSource.run.shuffle is true, preserving the previous default shuffle behavior; shuffle is disabled only when a recipe explicitly sets shuffle: false.
• module_prefix is a new optional parameter for LoRA injection. When not configured, the original target_modules suffix matching logic is still used and remains backward compatible. Incorrect configuration may cause LoRA to miss specific modules, so the actual trainable params should be confirmed through startup logs.
• The WAN VAE decode flow is replaced with a memory-optimized implementation, which is expected to significantly reduce decode memory pressure for ReFL/BPTT. The risk is that VAE dtype, tiling, checkpoint recompute, and the old path may have slight numerical differences, which should continue to be monitored through reward curves and sample quality.
• Cross-RPC grad input handling in workers supports differentiable RPC calls, which is required by ReFL training. It should not change behavior for ordinary no-grad rollout / scalar reward tasks, but other training tasks using grad-mode RPC should still be monitored.
• Reward dependencies are isolated in recipe-local requirements.txt files and no longer forcibly pollute the core environment. The corresponding requirements need to be manually installed before running each reward.
• Real local paths are not committed in configs and scripts. Model, data, and reward checkpoint paths need to be injected through environment variables before running.

[haonan3]

The tiled VAE decode path hardcodes the output buffer batch dimension to 1, while hidden_states may have B > 1. The decoded tile therefore has batch size B and cannot be accumulated in-place into values when B > 1.

Since decode() now defaults to tiled=True, this can break existing WAN recipes such as wan21_t2v.yaml, which uses forward_batch_size: 4. Please allocate these buffers using the actual input batch size, apply the same fix to tiled_parallel_decode, and add a regression test covering batch size greater than one.