Wang-Xiaodong1899 / Iso-Dream

Iso-Dream: Isolating and Leveraging Noncontrollable Visual Dynamics in World Models (NeurIPS 2022 Spotlight)

A PyTorch implementation of our paper:

Iso-Dream: Isolating and Leveraging Noncontrollable Visual Dynamics in World Models.

Minting Pan*, Xiangming Zhu*, Yunbo Wang, Xiaokang Yang

[arXiv] [Project Page]

Showcases

DMC and CARLA

Demo of Iso-Dream in DMC with noisy video backgrounds

Demo of Iso-Dream in CARLA

BAIR

Get Started

Iso-Dream is implemented and tested on Ubuntu 18.04 with python == 3.7, PyTorch == 1.9.0:

1. Create an environment

   conda create -n iso-env python=3.7
   conda activate iso-env
   

2. Install dependencies

   pip install torch==1.9.0+cu111 torchvision==0.10.0+cu111 torchaudio==0.9.0 -f https://download.pytorch.org/whl/torch_stable.html
   pip install -r requirements.txt
   

DMC / CARLA

For CARLA environment:

1. Setup

   Download and setup CARLA 0.9.10

   chmod +x setup_carla.sh
   ./setup_carla.sh
   

   Add to your python path:

   export PYTHONPATH=$PYTHONPATH:/home/CARLA_0.9.10/PythonAPI
   export PYTHONPATH=$PYTHONPATH:/home/CARLA_0.9.10/PythonAPI/carla
   export PYTHONPATH=$PYTHONPATH:/home/CARLA_0.9.10/PythonAPI/carla/dist/carla-0.9.10-py3.7-linux-x86_64.egg
   

   and merge the directories, i.e., put 'carla_env_dream.py' into 'CARLA_0.9.10/PythonAPI/carla/agents/navigation/'.

2. Training

   Terminal 1:

   cd CARLA_0.9.10
   bash CarlaUE4.sh -fps 20 -opengl
   

   Terminal 2:

   cd dmc_carla_iso
   python dreamer.py --logdir log/iso_carla --action_step 20 --free_step 50 --kl_balance 0.8 --action_scale 1 --seed 9 --configs defaults carla
   

3. Evaluation

   cd dmc_carla_iso
   python test.py --logdir test --action_step 20 --free_step 50 --kl_balance 0.8 --configs defaults carla
   

For DMC environment:

1. Setup DMC with video background

   Download 'envs' from Google Drive and put it in the 'dmc_carla_iso'. The dependencies can then be installed with the following commands:

   cd dmc_carla_iso
   
   cd ./envs/dm_control
   pip install -e .
   
   cd ../dmc2gym
   pip install -e .
   
   cd ../..
   

2. Training

   python dreamer.py --logdir log/iso_dmc1 --kl_balance 0.8 --seed 4 --configs defaults dmc --task dmcbg_walker_walk
   
   python dreamer.py --logdir log/iso_dmc1 --kl_balance 0.8 --seed 4 --configs defaults dmc --task dmcbg_walker_walk
   

BAIR / RoboNet

Train and test Iso-Dream on BAIR and RoboNet datasets. Also, install Tensorflow 2.1.0 for BAIR dataloader.

1. Download BAIR data.

   wget http://rail.eecs.berkeley.edu/datasets/bair_robot_pushing_dataset_v0.tar
   

2. Train the model. You can use the following bash script to train the model. The learned model will be saved in the --save_dir folder. The generated future frames will be saved in the --gen_frm_dir folder.

   cd bair_robonet_iso
   sh train_iso_model.sh
   

Citation

Pan, Minting and Zhu, Xiangming and Wang, Yunbo and Yang, Xiaokang. "Iso-Dream: Isolating and Leveraging Noncontrollable Visual Dynamics in World Models". Advances in Neural Information Processing Systems. 2022.

@inproceedings{paniso,
  title={Iso-Dream: Isolating and Leveraging Noncontrollable Visual Dynamics in World Models},
  author={Pan, Minting and Zhu, Xiangming and Wang, Yunbo and Yang, Xiaokang},
  booktitle={Advances in Neural Information Processing Systems}
  year={2022}
}

Acknowledgement

We appreciate the following github repos where we borrow code from:

https://github.com/jsikyoon/dreamer-torch

https://github.com/thuml/predrnn-pytorch