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

A PyTorch implementation of our paper:

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

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

[arXiv] [Project Page]

Showcases

DMC and CARLA

BAIR

Get Started

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

Create an environment

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

Install dependencies
```
pip install -r requirements.txt
```

DMC / CARLA

For CARLA environments:

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.

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 --step 50 --kl_balance 0.8 --action_scale 0.001 --seed 9 --configs defaults carla

Evaluation

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

BAIR / RoboNet

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

Download BAIR data.

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

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
```

Acknowledgement

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

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

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

wyb15/Iso-Dream