# README

We have attached the Appendix with the main paper in a single file renamed to 'paper.pdf'.

The gifs can be found in `DEMO` folder. The gifs for each model can be found in their respective folders with name of the model.

We plan to release the code upon acceptance of the paper with MIT License.

## VSG

The dependencies can be installed using the `requirements.txt` file:

```shell
cd StochasticDreamer
virtualenv --no-download VSG
source VSG/bin/activate
pip install --upgrade pip
pip install tensorflow==2.4.1 tensorflow_probability==0.12.2
pip install -r requirements.txt
```

NOTE:
Incase you face issues with numpy, specifically `NotImplementedError: Cannot convert a symbolic Tensor (strided_slice:0) to a numpy array.`, follow the fix mentioned [here](https://github.com/tensorflow/models/issues/9706#issuecomment-792106149).

For running experiments on tasks from DeepMind Control Suite, first install the `dm_control` repository following the instructions [here](https://github.com/deepmind/dm_control). Then use the command below to run with the appropriate config.

| Variable       | Description                          |
|------------|-----------------------------------------------|
| TASK       | DMC Task to train on                          |
| MODEL      | Choose from DreamerV1, DreamerV2, VSG or SVSG |
| SUFFIX     | Run ID to give the experiment                 |
| GATE PRIOR | Prior gate probability for VSG or SVSG        |
| DIM        | Size of the latent state                      |
| SEED       | To run with multiple seeds                    |

```shell
bash scripts/dmc.sh {TASK} {MODEL} {SUFFIX} {GATE_PRIOR} {DIM} {SEED}
```

For example to run `VSG` on walker_walk,

```shell
bash scripts/dmc.sh walker_walk VSG baseline 0.4 1024 1
```

## BringBackShapes
The BringBackShapes environment can be setup by:
```shell
cd bringbackshapes
pip install -r requirements.txt
pip install -e .
```

To run models on BringBackShapes:

| Variable       | Description                          |
|------------|-----------------------------------------------|
| DISTRACTORS | Number of stochastic distractors in the env        |
| SIZE        | Scale of the area to control partial observability, 1.0 refers to the Basic version |

```shell
bash scripts/bringbackshapes.sh {MODEL} {SUFFIX} sparse 3000 {DISTRACTORS} 5 False False False 125 {SIZE} {GATE_PRIOR} {SEED}
```

An example run is
```shell
bash scripts/bringbackshapes.sh VSG baseline sparse 3000 0 5 False False False 125 1.0 0.4 1
```

# Ablations

We have included gifs that visualize the imagined trajectories from DreamerV2, VSG and SVSG in the `ablations` folder. These correspond to the experiment mentioned in Appendix J.