# Introduction
This directory contains the code for
- generating Maze scenarios 
- generating trajectories in each room for training and validation
- training and validating world models
- evaluate actors in maze environments

All codes will be open-source after the paper is accepted.

# 1. Generate mazes

Sample mazes with the following command:

```bash
 python ./data/gen_maze_task.py \
	 --output_path YOUR_TASK_PATH \   # file path to save tasks
	 --task_number 1 \       # size of the task set
	 --scale 9,16 \       # the range of number of grids per dimension
	 --landmarks 6,10       # number of objects in the scenario
```

The generated mazes include the topology, the textures, and other hyper configurations.

Before doing so, make sure the directory is in your python path. You can add it by running:
```
export PYTHONPATH="$PYTHONPATH:PATH_TO_YOUR_DIRECTORY"
```

# 2. Generate trajectories for continual object navigation (CON) tasks.

Generate trajectories by randomly sample mazes from the saved scenarios by the following command:

```bash
 python ./data/gen_maze_record.py \
	 --output_path YOUR_DATA_PATH \
	 --task_file YOUR_TASK_PATH \
	 --task_source FILE \   # FILE or NEW
	 --max_steps 16000 \    # Sequence length T
	 --epochs 256 \         # Number of sequences
	 --workers 64
```

If you are sampling random maze for each sequence, use `task_source=NEW`.

The trajectories are generated by executing behavior policies composed of perturbed oracles and expert agents across CON tasks.

# 3. Configuration, training, and validation of world model

Run the following command to train and validate the world model on trajectories:
## Training and Validation
```bash
python projects/MazeWorld/train.py projects/MazeWorld/config-train.yaml --configs key1=value1 key2=value2 ...
python projects/MazeWorld/validate.py projects/MazeWorld/config.yaml --configs key1=value1 key2=value2 ...
```

## Configuration

The `config.yaml` file contains all the necessary configuration for running ICSL. Each configuration item is composed of multiple keys and sub-keys, can be over-written by commandline arguments. For instance, 
```yaml
model_config:
    image_encoder_block:
        hidden_size: 768
```
can be over-written by commandline arguments as follows:
```bash
python train.py config.yaml --model_config.image_encoder_block.hidden_size=1024
```
Below we explain key configuration items in detail.

### General configuration

- **run_name**:   Names ICSL will use to discrimate the run from the others in the logs

- **master_port**:  A port used for connecting to the master node

- **load_model_path**:  Set to none in a cold start, or set to a path to load the model from a checkpoint

### Log configuration

Specify the log path and whether to use tensorboard.

### Model configuration (model_config)

Configuration for the overall model architecture and components, including encoders, decoders, and causal blocks. It defines the structure and behavior of the model during training and inference.

- **max_position_loss_weighting**: Defines the maximum sequence length that the model can handle.
- **context_warmup**: specify a increasing loss weighting with the context length.
- **causal_block**:  Options include `Transformer`, `GSA`, `GLA`, `MAMBA`, `RWKV6`, `RWKV7`. The code automatically use causal masks for `Transformer`, and employ a chunk-wise forward and backward pass. E.g., `Transformer` is automatically set to sliding window attention mode by setting train_config.seg_len.
- **state_encode**, **state_decode**, **action_encode**, ...: specify the encoder and decoder for states, actions, rewards etc.

### Training configuration (train_config)

Settings for training the model.

- **seq_len**: Specify the sequence length loaded into the memory when training.
- **seg_len**: Specify the segment length used in chunk-wise forward and backward pass.
- **lr**, **lr_decay_interval**, **lr_start_step**: ICSL apply noam decay with the warmup step specified by `lr_decay_interval`, use `lr_start_step` in cases of warm start.

### Test configuration (test_config)

specify the configurations used for valiations between episodes or during static testing.

### Evaluation Configuration (generator_config)

specify the configurations for auto-regressive interaction with the environment during dynamic evaluation.

**Parameters not explicitly shown above may retain their default values as per the recommended configuration. Custom adjustments are available when aligned with specific application requirements.**

# 4. Evaluate the actor model in mazes

Enable the model to interact with a maze environments by 

```bash
python generator.py config.yaml
```

# Download pre-generated datasets, tasks and pre-trained models

- Pre-generated mazes tasks for training (32000) can be downloaded from [here](https://www.kaggle.com/datasets/anyonexxxx/maze-15-train-task).
- Pre-generated validation trajectories (256 sequences) can be downloaded from [here](https://www.kaggle.com/datasets/anyonexxxx/maze-15-test).
- Pre-trained Div/Long models in can be downloaded from [here](https://www.kaggle.com/models/anyonexxxx/icsl_maze_15_div-long/).
