# Weakly-Supervised RL for Controllable Behavior

## Preliminaries
Create the conda environment by running:
```
conda env create -f conda_envs/linux-gpu-env.yml
conda env create -f conda_envs/mac-cpu-env.yml
```
To delete the conda environment, run `conda env remove -n wsc_env`.

## Quick Start

To train weakly-supervised control (WSC) with pre-trained disentangled representations, run:
```
python -m scripts.train_her --config configs/Push1-wsc.txt
python -m scripts.train_her --config configs/Push2-wsc.txt
python -m scripts.train_her --config configs/Push3-wsc.txt
python -m scripts.train_her --config configs/PushLights1-wsc.txt
python -m scripts.train_her --config configs/PushLights2-wsc.txt
python -m scripts.train_her --config configs/PushLights3-wsc.txt
python -m scripts.train_her --config configs/Pickup-wsc.txt
python -m scripts.train_her --config configs/PickupLights-wsc.txt
python -m scripts.train_her --config configs/PickupColors-wsc.txt
python -m scripts.train_her --config configs/PickupLightsColors-wsc.txt
python -m scripts.train_her --config configs/Door-wsc.txt
python -m scripts.train_her --config configs/DoorLights-wsc.txt
```
![alt text](images/goal_conditioned_rl_performance.png)


## Reproducing experimental results

See below to generate datasets and train disentangled representations from scratch.

### Dataset generation

Generate datasets by running:
```
python -m scripts.generate_sawyer_datasets --env-id SawyerPushNIPSEasy-v0
python -m scripts.generate_sawyer_datasets --env-id SawyerPush2PucksEnv-v1
python -m scripts.generate_sawyer_datasets --env-id SawyerPush3PucksEnv-v1
python -m scripts.generate_sawyer_datasets --env-id SawyerPushRandomLightsEnv-v1
python -m scripts.generate_sawyer_datasets --env-id SawyerPush2PucksRandomLightsEnv-v1
python -m scripts.generate_sawyer_datasets --env-id SawyerPush3PucksRandomLightsEnv-v1
python -m scripts.generate_sawyer_datasets --env-id SawyerPickupGoalEnv-v1
python -m scripts.generate_sawyer_datasets --env-id SawyerPickupRandomLightsEnv-v1
python -m scripts.generate_sawyer_datasets --env-id SawyerPickupRandomColorsEnv-v1
python -m scripts.generate_sawyer_datasets --env-id SawyerPickupRandomLightsColorsEnv-v1
python -m scripts.generate_sawyer_datasets --env-id SawyerDoorGoalEnv-v1 
python -m scripts.generate_sawyer_datasets --env-id SawyerDoorRandomLightsEnv-v1
```

By default, the datasets are saved to this location:
```
DATA_DIR=/tmp/weakly_supervised_control/datasets
```
![alt text](images/sawyer_envs.png)

### Disentangled Representation Training

Train disentangled representations by running:
```
python -m scripts.train_disentangled_model --train-data-path $DATA_DIR/SawyerPushNIPSEasy-v0-n256.npz --eval-data-path $DATA_DIR/SawyerPushNIPSEasy-v0-n64.npz --factors r=0,1,2,3

python -m scripts.train_disentangled_model --train-data-path $DATA_DIR/SawyerPushRandomLightsEnv-v1-n256.npz --eval-data-path $DATA_DIR/SawyerPushRandomLightsEnv-v1-n64.npz --factors r=0,1,2,3,4

python -m scripts.train_disentangled_model --train-data-path $DATA_DIR/SawyerPickupGoalEnv-v1-n256.npz --eval-data-path $DATA_DIR/SawyerPickupGoalEnv-v1-n64.npz --factors r=1,2,4,5

python -m scripts.train_disentangled_model --train-data-path $DATA_DIR/SawyerPickupRandomLightsEnv-v1-n256.npz --eval-data-path $DATA_DIR/SawyerPickupRandomLightsEnv-v1-n64.npz --factors r=1,2,4,5,6

python -m scripts.train_disentangled_model --train-data-path $DATA_DIR/SawyerPickupRandomColorsEnv-v1-n256.npz --eval-data-path $DATA_DIR/SawyerPickupRandomColorsEnv-v1-n64.npz --factors r=1,2,4,5,6,7

python -m scripts.train_disentangled_model --train-data-path $DATA_DIR/SawyerPickupRandomLightsColorsEnv-v1-n256.npz --eval-data-path $DATA_DIR/SawyerPickupRandomLightsColorsEnv-v1-n64.npz --factors r=1,2,4,5,6,7,8

python -m scripts.train_disentangled_model --train-data-path $DATA_DIR/SawyerDoorGoalEnv-v1-n256.npz --eval-data-path $DATA_DIR/SawyerDoorGoalEnv-v1-n64.npz --factors r=3

python -m scripts.train_disentangled_model --train-data-path $DATA_DIR/SawyerDoorRandomLightsEnv-v1-n256.npz --eval-data-path $DATA_DIR/SawyerDoorRandomLightsEnv-v1-n64.npz --factors r=3,4
```
