ARB4WM: Adversarial Robustness Benchmark for World Models

This repository contains the code used for ARB4WM, a benchmark for evaluating adversarial robustness of Dreamer-family world models in visual continuous-control tasks.

The implementation keeps a Dreamer-family training pipeline while adding adversarial robustness evaluation, input-defense testing, transfer attacks, policy-saliency diagnostics, and a lightweight reporting interface.

Scope

The released code focuses on the experiments described in the ARB4WM paper:

• Dreamer-family agents: Dreamer, R2-Dreamer, Dreamer-Pro, and Dreamer-InfoNCE.
• Benchmarks: DeepMind Control Suite vision tasks and MetaWorld visual manipulation tasks.
• White-box objectives: policy divergence, policy entropy, value suppression, latent drift, and RSSM dynamics consistency.
• Attack optimizers: random noise, single-step gradient sign, multi-step PGD-style attacks, MI-FGSM, and NI-FGSM.
• Temporal settings: full-frame, first-half, second-half, and sparse-frame attacks.
• Input defenses: JPEG compression, median filtering, Gaussian noise, and adaptive BPDA/EOT-style evaluation.
• Diagnostics: transfer attacks and occlusion-based policy saliency.

Large logs, trained checkpoints, paper assets, ROM files, generated videos, and plotting-only scripts are intentionally not included.

Repository Layout

Attack/                 Attack objectives, optimizers, defenses, and rollout evaluators.
configs/                Hydra configs for training and robustness evaluation.
envs/                   DMC and MetaWorld environment wrappers.
optim/                  Optimizers used by the Dreamer-family agents.
sys/                    Lightweight browser interface for reporting selected attack cases.
conda/                  Conda environment and Docker helper files.
docs/                   Notes inherited from the R2-Dreamer codebase.
train.py                Train Dreamer-family agents.
test.py                 Evaluate a trained checkpoint.
attack.py               Run a configured attack rollout.
test_robustness.py      Sweep epsilons and objectives for robustness curves/data.
test_transfer.py        Evaluate source-to-target transfer attacks.
policy_saliency.py      Generate occlusion-based saliency diagnostics.

Installation

The code was developed with Python 3.11 and PyTorch on Linux. A conda environment file is provided:

conda env create -f conda/environment.yml
conda activate Dreamerr2
pip install -r conda/requirements.txt

For MuJoCo-based environments on headless machines, EGL rendering is recommended:

export MUJOCO_GL=egl
export MUJOCO_EGL_DEVICE_ID=0

Training

Use Hydra overrides to choose the environment, task, log directory, and representation objective.

python train.py \
  env=dmc_vision \
  env.task=dmc_reacher_easy \
  model.rep_loss=r2dreamer \
  logdir=./logdir/dmc_reacher_easy_r2dreamer

Available model.rep_loss values:

dreamer
r2dreamer
dreamerpro
infonce

MetaWorld example:

python train.py \
  env=metaworld \
  env.task=metaworld_reach \
  model.rep_loss=r2dreamer \
  logdir=./logdir/metaworld_reach_r2dreamer

Robustness Evaluation

Run a single configured attack:

python attack.py \
  env=dmc_vision \
  env.task=dmc_reacher_easy \
  logdir=./logdir/dmc_reacher_easy_r2dreamer \
  model.rep_loss=r2dreamer \
  attacker.method=frame \
  attacker.generate_method=Gradient_Sign \
  attacker.loss=frame_policy_kl \
  attacker.epsilon=0.04 \
  attacker.steps=10 \
  attacker.alpha=0.004 \
  env.env_num=1

Run the main epsilon sweep used for ARB4WM-style robustness curves and .npz result files:

python test_robustness.py \
  env=dmc_vision \
  env.task=dmc_reacher_easy \
  logdir=./logdir/dmc_reacher_easy_r2dreamer \
  model.rep_loss=r2dreamer \
  attacker.method=frame \
  attacker.epsilons='[0.0,0.02,0.04,0.08,0.12,0.16]' \
  seed=20

Common attack objectives:

frame_policy_kl
frame_policy_uniform
frame_value_min
frame_latent_l2
frame_rssm_kl

Temporal attack modes are controlled through attacker.method=frame_partial and attacker.partial:

0  first-half attack
1  second-half attack
5  sparse attack every five frames

Input defenses are configured under attacker.defender:

python test_robustness.py \
  env=dmc_vision \
  env.task=dmc_reacher_easy \
  logdir=./logdir/dmc_reacher_easy_r2dreamer \
  model.rep_loss=r2dreamer \
  attacker.defender.defense_method=jpeg_compress \
  attacker.adaptive_attack=True

Supported defense methods:

none
gaussian_noise
median_filter
jpeg_compress

Transfer and Saliency Diagnostics

Transfer evaluation expects checkpoints whose directories follow the pattern:

<base_logdir>_r2dreamer
<base_logdir>_dreamer
<base_logdir>_dreamerpro
<base_logdir>_infonce

Example:

python test_transfer.py \
  env=dmc_vision \
  env.task=dmc_reacher_easy \
  logdir=./logdir/dmc_reacher_easy \
  seed=20

Policy saliency can be generated with:

python policy_saliency.py \
  env=dmc_vision \
  env.task=dmc_reacher_easy \
  logdir=./logdir/dmc_reacher_easy \
  seed=20

Reporting Interface

The sys/ directory contains a lightweight static interface for presenting selected robustness cases. It is a visualization/reporting layer and does not replace the Python attack implementation.

Open sys/index.html in a browser or serve it with any static file server.

Notes on Released Artifacts

This repository does not include trained checkpoints or generated experiment outputs. To reproduce a table or figure from the paper, first train or place the corresponding checkpoint under logdir, then run the attack/evaluation scripts with the same task, model, seed, epsilon grid, and defense settings.

Acknowledgements

This codebase is derived from NM512/r2dreamer. We thank the authors for releasing their implementation.

Citation

If you use ARB4WM, please cite:

@article{zhang2026arb4wm,
  title={ARB4WM: An Adversarial Robustness Benchmark for World Models in Continuous Control},
  author={Zhang, Junjian and Tan, Hao and Li, Ruonan and Zhu, Dong and Li, Aiping and Gu, Zhaoquan},
  journal={arXiv preprint arXiv:2606.16605},
  year={2026},
  url={https://arxiv.org/abs/2606.16605}
}