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

[🔥Project Page] [🚀 FoAM-main]

FoAM Benchmark:

The first open-source simulated dual-arm system with rich tasks, manipulation scenarios, and the same physics as the real UR3e robot.

Installation

# Go in the FoAM-benchmark repo
cd FoAM-benchmark
# Create a environment
conda env create -f environment_foam-benchmark.yaml
# Activate the environment
conda activate foam-benchmark

Create Dataset

This project has a total of 12 open source scenarios, each with a different number of subtasks, and is divided into four levels based on the difficulty and type of the tasks. Each scenario's folder has a similar file structure. In this Appendix, we take the subtask of Open Cabinet-open_cabinet_middle_drawer in the normal scenarios as an example. The simulation demos can be generated by running the script record_sim_episodes.py in each scenario. Alternative command lines are provided in README.md. Episodes that meet the successful conditions will be saved in the data folder in .hdf files.

Exanmple:

# Create dataset for the subtask: open_cabinet_middle_drawer
python NormalScenarios/OpenCabinet/record_sim_episodes.py --task_name sim_open_cabinet_middle_drawer --num_episodes 50

The initialization position of the manipulated objects can be changed in the utils.py file. For example, in the open the cabinet drawer scenarios, the pose of the cabinet is controlled by the function sample_cabinet_pose in the utils.py file.

The motion trajectory of the robot arm can be customized in the file scripted_policy.py. For example, in the open the cabinet drawer scenarios, the motion trajectory of the arm is controlled by the function generate_trajectory of the Class PickAndTransferPolicy in the scripted_policy.py file.

The physical models required for each scenario are stored in the folder \assets. The XML files of the robot and the background are stored in the \model folder. The integrated scenarios are stored in sim_ee_env.xml and sim_env.xml. They can be directly visualized by MuJoCo software. You can customize the scenarios by customizing the files in the /assets and /models.

Note: Due to the different GPUs used by researchers, the progress of data generation may be different. If successful expert data cannot be obtained (the success criteria maybe strict for your GPU), you can try to customize the success criteria.

License

All the code, model weights, and data are licensed under MIT license.