[Feature] [Algorithm] Add FQL Loss

[neuriv]

Description

Add FQLLoss to train the flow policy on replay actions, distill its samples into a one-step actor and improve that actor using critic values. Train the critics with target values while supporting TorchRL’s loss keys, masks, sample weights and target updates.

Part of #4433.

Testing

Ubuntu 22.04, A100 40 GB, Python 3.11.16, PyTorch 2.14.0+cu126 and TensorDict 0.14.2.

Run from the feat/fql worktree. Additional test scripts were kept outside the repository.

root=/home/ubuntu/fql-validation-20260919
suite="$root/tests"
evidence="$root/evidence"
export PATH="$root/rl/.venv/bin:$PATH"
export PYTHONPATH="$PWD" FQL_TEST_DEVICE=cpu,cuda COMPOSITE_LP_AGGREGATE=0
export CUDA_VISIBLE_DEVICES=0 MUJOCO_GL=egl OMP_NUM_THREADS=1
export HYDRA_FULL_ERROR=1 WANDB_MODE=disabled WANDB_SILENT=true

pre-commit run --from-ref b4caa999e7619cc416174d682725749c2079f0c0 --to-ref HEAD
python -m pytest -q "$suite/test_flow.py" "$suite/test_fql.py" test/objectives/test_loss_module.py

for seed in 17 0 1; do
    python "$suite/check_fql.py" --device cuda --seed "$seed" --output "$evidence/fql/seed-$seed"
    python "$suite/verify_fql_evidence.py" "$evidence/fql/seed-$seed" --device cuda
done

362 tests passed. Pre-commit passed.

Seed	Initial action MSE	FQL action MSE	Behavior-only action MSE	FQL critic MAE	Frozen-target critic MAE
0	0.397491	0.010882	1.065319	0.003148	1.044985
1	0.279830	0.010165	1.021522	0.004088	0.962058
17	0.425254	0.007791	1.052302	0.004853	0.980475

All three seeds passed. Reloaded checkpoints reproduced the measurements and passed independent Bellman-value, target-weight and optimizer-step checks.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/4437

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

⚠️ 16 Awaiting Approval

As of commit 16f88f8 with merge base ad969ce ():

AWAITING APPROVAL - The following workflows need approval before CI can run:

• Build Aarch64 Linux Wheels (gh)
• Build Linux Wheels (gh)
• Build M1 Wheels (gh)
• Build Windows Wheels (gh)
• Continuous Benchmark (PR) (gh)
• Examples Tests on Linux (gh)
• Generate documentation (gh)
• Habitat Tests on Linux (gh)
• Libs Tests on Linux (gh)
• Lint (gh)
• LLM Tests on Linux (gh)
• Push Binary Nightly (gh)
• SOTA Tests on Linux (gh)
• Tutorials Tests on Linux (gh)
• Unit-tests on Linux (gh)
• Validate Test Partitioning (gh)

This comment was automatically generated by Dr. CI and updates every 15 minutes.