# This is source code for Align Your Intents: Offline Imitation Learning via Optimal Transport paper

All experiments were conducted on CUDA12 with corresponding JAX version, which can be installed with `pip install -U "jax[cuda12]"`

## Pretrained models
Pretrained models in the Equinox format can be downloaded from [Google drive](https://drive.google.com/drive/folders/1555qP5vDQSo5upWW4ZV-qUUVVleuGyPD?usp=sharing)
Put this folder in the root directory. However, we directly provide all of the checkpoints of pretrained value function in this directory.

In order to pretrain general-value function, follow the instructions in corresponding folder `icvf_pretraining`.

## Running Inference
To evaluate AILOT, run `inference.py` in root directory. Specific configurations can be tweaked in `configs/entry.yaml` config file. For example, `algo_name` can be either `IQL` or `AILOT`.

## Additional Notebooks with visualizations
In `notebooks/intents_states_differences.ipynb` we provide visualizations of how extracted representations attain almost linear dependence between steps later in the trajectory. And compare this with initial states, showing highly oscillating behavior, supporting argument towards ignoring underlying geometry.


In `notebooks/ailot_comparison.ipynb` we provide training procedures and visualizations in a more convinient way. 