FWM¶

Flow-matching World Model

OfflineData-centric · generativeBase · IQLClaims · dynamic shift

Multi-step critic targets whose future is sampled from a flow-matching world model.

Features¶

Mechanism¶

FWM transfers the multi-step future modelling of the Diffusion World Model to a flow-matching world model. Instead of a one-step bootstrap, the critic regresses a multi-step return whose future is sampled from the generative model, conditioned on the current state, action and return-to-go. The replacement is confined to how that future is produced.

The model is trained on windows of states and rewards. It learns the velocity that carries a noise sample to a data window along a straight line, which replaces the denoising objective of the diffusion model. The return-to-go condition is dropped with probability 0.2 during training, so that one network carries both the conditional and the unconditional velocity. Sampling integrates the learned velocity from noise to data in five midpoint steps, with a guidance weight of 2.0 between the conditional and the unconditional velocity.

Run the method¶

python baselines/train.py \
    -c robustrllib/configs/experiment/fwm_hopper.yaml --seed 0
python baselines/evaluate.py --run runs/fwm_hopper/seed0

Training is on the nominal task. The configuration files are explained in Train an algorithm, and the evaluation of the frozen checkpoint in Evaluation Protocol.

Robust performance¶

Normalized score of the frozen checkpoint. Q1 to Q4 are the severity quartiles of each shift factor's scales, ordered by displacement from the nominal setting, and All covers every shifted condition. The last row averages the tasks.

References¶

• Ding et al. Diffusion World Model: Future Modeling Beyond Step-by-Step Rollout for Offline Reinforcement Learning. arXiv:2402.03570, 2024.
• Lipman et al. Flow Matching for Generative Modeling. arXiv:2210.02747, 2022.