FoAM: Foresight-Augmented Multi-Task Imitation Policy for Robotic Manipulation

[🔥Project Page] [🚀 FoAM-benchmark]

Installation

If you want to use FoAM, you need to install the required packages:

# Go in the FoAM-main repo
cd FoAM-main
# Create a environment
conda env create -f environment.yaml
# Activate the foam
conda activate foam
# install totch with cuda
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
# install dm_control
pip install dm_control pyquaternion

Reproduce FoAM

Prepare Datasets

First, generate data based on any one or more scenarios in the FoAM Benchmark. Please check the FoAM-benchmark/* folder for a quick start. Store the data in the folder \train_data. Its storage structure is shown below.

# Data Storage Structure

- **train_data/** (Main directory for training data)
  - **scenarios_1/** (Scenario 1)
    - **subtask_1/** (Subtask 1)
      - **episodes_0-49/** (Contains episodes 0-49)
    - **subtask_2/** (Subtask 2)
      - **episodes_0-49/** (Contains episodes 0-49)
    - ...
  - **scenarios_2/** (Scenario 2)
    - **subtask_1/** (Subtask 1)
      - **episodes_0-49/**
    - **subtask_2/** (Subtask 2)
      - **episodes_0-49/**
    - ...
  - ...

If you have generated the dataset for the task Open Middle Drawer according to the README.md in FoAM-benchmark, you can start training your FoAM.

Training

Run the imitate_episodes.py file for training.

python imitate_episodes.py \
--ambiguity_env_name SimOpenDrawer \
--dataset_dir ../FoAM-benchmark/train_data \
--ckpt_dir ./ckpt/ckpt_FoAM \
--policy_class FoAM \
--kl_weight 10 \
--chunk_size 450 \
--hidden_dim 512 \
--batch_size 16 \
--dim_feedforward 3200 \
--seed 0 \
--num_epochs 2000 \
--lr 1e-5 \
--multi_task \
--task_name sim_open_cabinet_middle_drawer \
--run_name multi_task_run \
--huber_weight 2 \
--use_redundant_task_emb \
--use_goal_img 

Inference

Run the imitate_episodes.py file for inference.

python imitate_episodes.py \
--ambiguity_env_name SimOpenDrawer \
--dataset_dir ../home/liulitao/Desktop/FoAM-main/train_data \
--ckpt_dir ./ckpt/ckpt_FoAM \
--policy_class FoAM \
--kl_weight 10 \
--chunk_size 450 \
--hidden_dim 512 \
--batch_size 16 \
--dim_feedforward 3200 \
--seed 0 \
--num_epochs 2000 \
--lr 1e-5 \
--multi_task \
--task_name sim_open_cabinet_middle_drawer \
--run_name multi_task_run \
--huber_weight 2 \
--use_redundant_task_emb \
--use_goal_img \
--eval \
--temporal_agg

License

This project is released under the MIT license. Parts of this project contain code and models from other sources, which are subject to their respective licenses.

Citation

If you find this project useful in your research, please consider cite: