# Planning with Sequence Models through Iterative Energy Minimization

## Table of Contents
- [Introduction](#Introduction)
- [Code Structure](#Code Structure)
- [Install](#install)
- [Usage](#usage)
- [Acknowledgments](#Acknowledgments)

## Introduction
We provide code for evaluate our proposed approach in Babyai Env.

## Code Structure
Follow this folder structure to prepare the dataset:
```
.
└── crosstask
    ├── babyai
        ├── demons
             ├── BabyAI-[ENV]-v0_agent.pkl (demonstration dataset)
    ├── mingpt
        ├── leap_model_babyai.py (LEAP model)
        ├── trainer_babyai.py (trainer)
        └── utils.py
    ├── run_leap.py (main file)
    ├── test_run.sh (run file)
    
```

## Requirements
- PyTorch >= 1.10

- In stall the modified Babyai environment

  ```bash
  cd babyai
  pip install -e ./
  ```

## Usage

```
bash test_run.sh
```
- `--horizon` means the planning horizon. 
- `--context_length` In tasks without using context, context length should be set equal to horizon number. In tasks using context, context length means the number of past steps
- `--sample_iteration` specifies the number of Gibbs sampling iteration
- `--epochs` means the total number of training epoches
- `--env` means the babyai environment
- `--batch_size` means the batch size we use 

## Acknowledgments
This code is developed based on [Decision Transformer](https://github.com/kzl/decision-transformer) and [BabyAI](https://github.com/mila-iqia/babyai).
