# Supplementary Code for Policy Architecture Training and Evaluation

This directory contains the code used to train and evaluate the policy architectures proposed in our paper.

## Note on Supplementary Material Limitations

Due to the 100MB file size restriction for supplementary submissions, we are unable to include the following components:

1. **FrankaLHT Simulator**  
   Our custom simulator includes multiple object meshes and assets, which exceed the size limit.

2. **Trained Policy Checkpoints**  
   Each trained policy contains approximately 25 million parameters, resulting in model checkpoints that exceed the allowed size.

3. **Training Data**  
   The datasets used for training the policies are similarly too large to include.

As a result, we are providing a simplified and lightweight version of our code that supports training and evaluation on the publicly available **Franka Kitchen** environment.

We intend to release the full codebase—including trained models, simulation assets, and datasets—on GitHub following the acceptance notification in **August**, to ensure proper attribution while maintaining anonymity during the review process.

---

## Setup Instructions

1. **Franka Kitchen Environment**  
   Follow the setup instructions from the official repository:  
   [https://github.com/google-research/relay-policy-learning](https://github.com/google-research/relay-policy-learning)  
   After setup, download the dataset, simulate all the trajectories, and save them in the `./Data_Franka_Kitchen/` subdirectory.

2. **UVD Architecture**  
   Clone and set up UVD from:  
   [https://github.com/zcczhang/UVD/](https://github.com/zcczhang/UVD/)

3. **VIP Encoder**  
   Download the VIP encoder from:  
   [https://github.com/facebookresearch/vip](https://github.com/facebookresearch/vip)

---

## Running Training and Evaluation

- Navigate to the `./FrankaKitchen/Training/` directory to run training and evaluation scripts for specific architectures.  
  The scripts are annotated with comments to help explain key steps and usage.
- Trained models will be saved in:  
  `./FrankaKitchen/Training/Trained_Models/`
- Policy rollouts will be stored in:  
  `./FrankaKitchen/Training/Evaluation/`
- SGPT heatmaps will be saved in:  
  `./FrankaKitchen/Training/Graphs/`
