Implementation of the paper Capturing Visual Environment Structure Correlates with Control Performance for MetaWorld policy training and full-state proxy ranking.

conda env create -f environment.yml
conda activate metaworld-gdp
pip install -e .

All training and proxy commands below assume this Conda environment is active.

GDP_DATA_DIR controls where Metaworld zarr datasets are read from:

export GDP_DATA_DIR=/path/to/metaworld/zarr/data

Expected policy-training files are named like:

metaworld_basketball_expert.zarr
metaworld_push_expert.zarr

Generate the 50 successful demonstrations per task used by the policy experiments:

python tools/gen_demonstration_expert_all_img.py --cams 1
python tools/validate_demonstrations.py --root-dir "$GDP_DATA_DIR"

The policy data uses 512 x 512 images from camera ID 1, which is corner2. Camera IDs are stored in each zarr dataset so the loader maps the requested camera to the correct stored image axis.

• Update: The feature extraction and rollout are optimized so the performance may be different from reported and ranks changes a little.

The reported policy runs use one training process, a global batch size of 32, seed 0, and 30 epochs. The canonical DINOv2 command is:

bash scripts/train_policy.sh dinov2 dinov2_vitb14 0

Other backbones use the same configuration with these two Hydra overrides:

For example:

bash scripts/train_policy.sh vit vit_b_16 0

The rollout renderer uses the same 512 x 512 input resolution as the demonstrations. Multi-GPU training is available through scripts/dist_train_policy.sh, but it is not the path used for the reported single-process policy results. The distributed command treats the configured batch size as global and divides it across workers.

Evaluate the latest checkpoint with 50 sequential rollouts per MT50 task:

bash scripts/eval_policy.sh \
  dinov2 dinov2_vitb14 \
  runs/policy/dinov2-dinov2_vitb14_seed0 \
  0

The proxy workflow uses full-state target data and object-region visual features to train one MT50-wide probe per backbone. The scripts store generated data, cached features, and probe results under data/, cache/, and runs/ in this repository.

Build MT50 proxy data:

bash scripts/launch_mt50_proxy_full_state_data.sh

Precompute backbone features:

bash scripts/launch_proxy_feature_cache_mt50.sh

Train probes and rank backbones:

bash scripts/launch_proxy_backbone_order_mt50.sh

If you find our work helpful, please cite with the following..

@inproceedings{dong2026visualenv,
  title={Capturing Visual Environment Structure Correlates with Control Performance},
  author={Dong, Jiahua and Man, Yunze and Tokmakov, Pavel and Wang, Yu-Xiong},
  booktitle={ICLR},
  year={2026},
  url={https://arxiv.org/abs/2602.04880}
}