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

• Project page: https://dongjiahua.github.io/CVESC/
• Paper: https://arxiv.org/abs/2602.04880

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

Single GPU:

bash scripts/train_policy.sh gdp metaworld_basketball run1 0 0

Distributed:

bash scripts/dist_train_policy.sh gdp metaworld_MT10 mt10_run 0 2

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

@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}
}