README.md

Towards Efficient Online Tuning of VLM Agents via Counterfactual Soft Reinforcement Learning

ICML 2025
Lang Feng, Weihao Tan, Zhiyi Lyu, Longtao Zheng, Haiyang Xu, Ming Yan, Fei Huang, Bo An
Nanyang Technological University & Alibaba Group

 

Overview of the CoSo training workflow.

Table of Contents

• Installation
  • Option 1: Using Docker (Recommended)
  • Option 2: Using Conda (4 Steps)
• Configuration
• Run Examples
  • 1. Run CoSo
  • 2. Run Naive Entropy
  • 3. Run DigiRL Baseline
• Citation

Installation

(Option 1) Using Docker (Recommended)

Fast and isolated setup using the provided Dockerfile.

docker build -t coso .

docker run --name coso --gpus all --device /dev/kvm --group-add kvm --shm-size 2gb -it -v <repo_path>/CoSo:<repo_path>/CoSo coso

Installation is complete! Skip to Configuration.

(Option 2) Using Conda (4 Steps)

1. Create the Environment and Install Dependencies

conda create -n coso python==3.10
conda activate coso

git clone https://github.com/langfengQ/CoSo.git
cd CoSo
pip install -e .

2. Environment Setup

The environment setup follows the same procedure as DigiRL. Please refer to the environment README. Before moving on, you should be able to view this screenshot by running this script.

3. Download Model Checkpoints

Download the model:

wget https://huggingface.co/cooelf/Auto-UI/resolve/main/Auto-UI-Base.zip
unzip Auto-UI-Base.zip -d <path_to_autoui_dir>

The folder should contain:

Auto-UI-Base/
├── config.json
├── pytorch_model.bin
├── tokenizer.json
...

4. Pre-Collected Trajectories

Download from Google Drive:

File Name	#Trajectories	Horizon	Size
general-off2on-sft-trajectories.pt	608	10	95.5MB
general-offline-sft-trajectories.pt	1552	10	243.9MB
webshop-off2on-sft-trajectories.pt	528	20	115.2MB
webshop-offline-sft-trajectories.pt	1296	20	297.5MB

Store the files in path ~/data/:

mkdir ~/data
cp *.pt ~/data/

Configuration

• Edit the configuration file: scripts/config/main/default.yaml:

  1. Fill in API keys and project info:

     huggingface_token: ''
     wandb_key: ''
     gemini_key: ''
     
     entity_name: ''
     project_name: ''
     

  2. Define the asset path:

     assets_path: '<repo_path>/CoSo/digirl/environment/android/assets/task_set'
     

  3. (Only if using conda) Replace /root/ with your own paths:

     policy_lm: '/root/Auto-UI-Base'
     cache_dir: '/root/.cache'
     

• Edit the sub-configuration file

  1. Choose the appropriate sub-configuration depending on training mode:digirl_off2on or digirl_offline.yaml or digirl_online.yaml or eval_only

  2. (Only if using conda) Replace /root/ with your own paths:

     offline_data_path: "/root/data/webshop-off2on-sft-trajectories.pt"
     

Run Examples

1. Run CoSo

use_entropy and use_causal are set to True in the default.yaml.

use_entropy: True
use_causal: True

Then run CoSo via:

cd scripts
python run.py --config-path config/main --config-name digirl_off2on

2. Run Naive Entropy

Modify default.yaml:

use_entropy: True
use_causal: False

3. Run DigiRL Baseline

Modify default.yaml:

use_entropy: False
use_causal: False

Citation

If you find this code and CoSo useful in your research or applications, we would appreciate it if you could cite our work:

@article{feng2025towards,
  title={Towards Efficient Online Tuning of VLM Agents via Counterfactual Soft Reinforcement Learning},
  author={Feng, Lang and Tan, Weihao and Lyu, Zhiyi and Zheng, Longtao and Xu, Haiyang and Yan, Ming and Huang, Fei and An, Bo},
  journal={arXiv preprint arXiv:2505.03792},
  year={2025}
}