Initial commit

.gitignore

@@ -1,3 +1,16 @@
+# Project-specific
+assets/train_assembly
+assets/test_assembly
+plan_sequence/generator/network/*.pt
+*.sdf
+*.obj
+*.dae
+*.zip
+
+.vscode
+.DS_Store
+__MACOSX
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

README.md

@@ -1 +1,199 @@
-# ASAP
+# ASAP
+
+This repository contains the official code and dataset of [ASAP: Automated Sequence Planning for Complex Robotic Assembly with Physical Feasibility (ICRA 2024)](http://asap.csail.mit.edu).
+
+<p align="middle">
+  <img src="images/assembly-01115.gif" width="24%" />
+  <img src="images/assembly-01234.gif" width="24%" /> 
+  <img src="images/assembly-02099.gif" width="24%" />
+  <img src="images/assembly-03212.gif" width="24%" />
+</p>
+
+<p align="middle">
+  <img src="images/robot-00833.gif" width="30%" />
+  <img src="images/robot-01235.gif" width="30%" /> 
+  <img src="images/robot-01434.gif" width="30%" />
+</p>
+
+**Authors**: Yunsheng Tian, Karl D.D. Willis, Bassel Al Omari, Jieliang Luo, Pingchuan Ma, Yichen Li, Farhad Javid, Edward Gu, Joshua Jacob, Shinjiro Sueda, Hui Li, Sachin Chitta, Wojciech Matusik
+
+**Summary**: The automated assembly of complex products requires a system that can automatically plan a physically feasible sequence of actions for assembling many parts together. In this paper, we present ASAP, a physics-based planning approach for automatically generating such a sequence for general-shaped assemblies. ASAP accounts for gravity to design a sequence where each sub-assembly is physically stable with a limited number of parts being held and a support surface. We apply efficient tree search algorithms to reduce the combinatorial complexity of determining such an assembly sequence. The search can be guided by either geometric heuristics or graph neural networks trained on data with simulation labels. Finally, we show the superior performance of ASAP at generating physically realistic assembly sequence plans on a large dataset of hundreds of complex product assemblies. We further demonstrate the applicability of ASAP on both simulation and real-world robotic setups.
+
+## Installation
+
+### 1. Clone repository
+
+```
+git clone --recurse-submodules [email protected]:yunshengtian/RobotAssembly.git
+```
+
+### 2. Python environment
+
+```
+conda env create -f environment.yml
+conda activate asap
+```
+
+or
+
+```
+pip install numpy networkx matplotlib scipy pyglet rtree sortedcontainers scipy tqdm trimesh torch torch_geometric torch_sparse torch_scatter seaborn ikpy pyquaternion
+```
+
+### 3. Python binding of simulation
+
+```
+cd simulation
+python setup.py install
+```
+
+To test if the installation steps are successful, run:
+
+```
+python test_sim/test_simple_sim.py --model box/box_stack --steps 2000
+```
+
+Then the simulation viewer should appear. [Here](https://github.com/yunshengtian/Assemble-Them-All?tab=readme-ov-file#simulation-viewer) are some tips on interacting with the viewer. 
+Additionally, press `V` for outputting the camera parameters (lookat and pos).
+
+We also provide a beam assembly under ``assets/beam_assembly`` folder. To visualize the simulation of that, run:
+
+```
+python test_sim/test_multi_sim.py --dir beam_assembly --id original --gravity 9.8 --steps 2000 --friction 0.5 --camera-pos 3.15 -1.24  1.6 --camera-lookat 2.59 -0.55 1.16
+```
+
+### 4. Assembly dataset (optional)
+
+Install the training set and test set:
+
+|    Training set (1906 assemblies)        |    Test Set (240 assemblies)     |
+| :--------------------------------------: | :------------------------------: |
+| ![training_set](images/training_set.png) | ![test_set](images/test_set.png) |
+| [Link (591MB)](https://people.csail.mit.edu/yunsheng/ASAP/dataset_2404/training_assembly.zip) | [Link (124MB)](https://people.csail.mit.edu/yunsheng/ASAP/dataset_2404/test_assembly.zip) |
+
+For point-based SDF collision check to work more accurately, we highly recommend subdividing the assembly meshes to have denser contact points by running ``assets/subdivide_batch.py``. For example, to subdivide the dataset saved in ``assets/test_assembly`` and export to ``assets/test_assembly_dense``:
+
+```
+python assets/subdivide_batch.py --source-dir assets/test_assembly --target-dir assets/test_assembly_dense --num-proc NUM_PROCESSES
+```
+
+## Experiments
+
+### Sequence planning
+
+Use the following command to run sequence planning on the beam assembly we provided:
+
+```
+python plan_sequence/run_seq_plan.py --dir beam_assembly --id original --planner dfs --generator heur-out --max-gripper 2 --base-part 6 --log-dir logs/beam_seq --early-term
+```
+
+Important arguments include (see the complete list in `plan_sequence/run_seq_plan.py`):
+
+- `dir`: assembly directory (relative to `assets/`)
+- `id`: assembly id
+- `planner`: name of the node selection algorithm (i.e., tree search planner) (see `plan_sequence/planner/__init__.py` for supported options)
+- `generator`: name of the part selection algorithm (i.e., part generator) (see `plan_sequence/generator/__init__.py` for supported options)
+- `seed`: random seed
+- `budget`: maximum number of feasibility evaluation
+- `max-gripper`: number of available grippers (for assembling and holding parts)
+- `max-pose`: number of pose candidates to search from during pose selection
+- `pose-reuse`: number of poses to be reused from the parent node for pose selection
+- `early-term`: early termination once a feasible plan is found (rather than waiting for the whole tree to be fully expanded)
+- `timeout`: timeout in seconds for the whole sequence planning
+- `base-part`: id of the base part (if exists) as the first part that stays in place (reorientation will not be allowed then)
+- `log-dir`: log directory for storing all the planning outputs
+- `plan-grasp`: whether to plan gripper grasps
+- `plan-arm`: whether to plan arm motions
+
+### Log folder structure
+
+If `log-dir` is specified in the above command, the log files will be saved in this directory: `{log-dir}/{planner}-{generator}/s{seed}/{id}/`.
+
+There are three files generated:
+1. `setup.json` that stores the arguments used for experiments; 
+2. `stats.json` that stores the high-level planning results; 
+3. `tree.pkl` that stores the explored disassembly tree with all necessary information on edges/nodes.
+
+### Generating results from log
+
+We separate the planning and result generation for flexibility considerations. Suppose you have run the above command for planning, then use the following command to generate planned results:
+
+```
+python plan_sequence/play_logged_plan.py --log-dir logs/beam_seq/dfs-heur-out/s0/original/ --assembly-dir assets/beam_assembly/original --result-dir results/beam_seq/ --save-all --camera-pos 3.15 -1.24  1.6 --camera-lookat 2.59 -0.55 1.16
+```
+
+Important arguments include (see the complete list in `plan_sequence/play_logged_plan.py`):
+
+- `log-dir`: input log directory
+- `assembly-dir`: input assembly directory (absolute path)
+- `result-dir`: output result directory
+- `save-mesh`: whether to output meshes in the result folder (not necessarily needed, same as meshes in assembly dir)
+- `save-pose`: whether to output (reoriented) pose of every assembly step
+- `save-part`: whether to output parts to be held
+- `save-record`: whether to output rendered videos
+- `save-all`: whether to output everything above
+- `reverse`: whether to reverse the rendering (to be assembly instead of disassembly)
+- `show-fix`: whether to show fixed parts in rendering (in grey)
+- `show-grasp`: whether to show gripper grasp in rendering
+- `show-arm`: whether to show arm motion in rendering
+
+If `save-all` is specified, the results will be saved in `result-dir` with the following structure. 
+Assume there are `N` parts, `N-1` assembly steps, `T` time steps in each assembly step, 
+N part ids are `{p0}, {p1}, ... {pN-1}`, and N-1 ordered part ids following the disassembly order are `{p'0}, {p'1}, ..., {p'N-2}`:
+
+```
+mesh/ --- meshes of individual parts
+ ├── part{p0}.obj
+ ├── ...
+ └── part{pN-1}.obj
+part_fix/ --- parts to be held in every assembly step
+ ├── 0_{p'0}.json
+ ├── ...
+ └── N-2_{p'N-2}.json
+path/ --- geometric assembly paths in every time step in every assembly step (4x4 transformation matrices)
+ └── 0_{p'0}/
+     └── 0/
+         ├── part{p0}.npy
+         ├── ...
+         └── part{pN-1}.npy
+     ├── ...
+     └── {T-1}/
+ ├── ...
+ └── N-2_{p'N-2}/
+pose/ --- global pose of the whole (sub)assembly in every assembly step (4x4 transformation matrix)
+ ├── 0_{p'0}.npy
+ ├── ...
+ └── N-2_{p'N-2}.npy
+record/ --- (dis)assembly animations in every assembly step
+ ├── 0_{p'0}.gif
+ ├── ...
+ └── N-2_{p'N-2}.gif
+```
+
+After the animations are generated, you can use `plan_sequence/combine_animation.py` to concatenate all videos into a single one.
+
+### Batch sequence planning
+
+Use `plan_sequence/run_seq_plan_batch.py` to run batch sequence planning for all assemblies in the assembly directory (with similar arguments as shown above for the serial script). The log folders will be saved in this directory: `{log-dir}/g{max-gripper}/{planner}-{generator}/s{seed}/`.
+
+To check success rates quantitatively, run:
+```
+python plan_sequence/check_success_rate_batch.py --log-dir {log-dir}/g{max-gripper}
+```
+
+## Contact
+
+Please feel free to contact [email protected] or create a GitHub issue for any questions about the code or dataset.
+
+## Citation
+
+If you find our paper, code or dataset is useful, please consider citing:
+
+```
+@article{tian2023asap,
+  title={ASAP: Automated Sequence Planning for Complex Robotic Assembly with Physical Feasibility},
+  author={Tian, Yunsheng and Willis, Karl DD and Omari, Bassel Al and Luo, Jieliang and Ma, Pingchuan and Li, Yichen and Javid, Farhad and Gu, Edward and Jacob, Joshua and Sueda, Shinjiro and others},
+  journal={arXiv preprint arXiv:2309.16909},
+  year={2023}
+}
+```

assets/beam_assembly/original/config.json

@@ -0,0 +1,22 @@
+{
+    "0": {
+        "initial_state": [6.8913, 8.6769, 0.762, 0, 0, 0], 
+        "final_state": [12.245, 3.8152, 7.0889, 0, 0, 1.570796]
+    }, 
+    "1": {
+        "initial_state": [6.8833, 13.2202, 0.762, 0, 0, 0], 
+        "final_state": [21.135, 3.8152, 7.0889, 0, 0, 1.570796]
+    }, 
+    "2": {
+        "initial_state": [17.0764, 8.7173, 0.635, -1.570796, 0, -1.570796], 
+        "final_state": [12.2424, 3.8152, 3.8614, 0, 0, -1.570796]
+    }, 
+    "3": {
+        "initial_state": [17.047, 13.1969, 0.635, -1.570796, 0, -1.570796], 
+        "final_state": [21.1376, 3.8152, 3.8614, 0, 0, -1.570796]
+    }, 
+    "6": {
+        "initial_state": [16.69, 3.8152, 0.635, 0, 0, 1.570796], 
+        "final_state": [16.69, 3.8152, 0.6321, 0, 0, 1.570796]
+    }
+}

assets/beam_assembly/original/id_map.json

@@ -0,0 +1 @@
+{"0": "Table_Base - Feet-1.obj", "1": "Table_Base - Feet-2.obj", "2": "Table_Base - Upright-1.obj", "3": "Table_Base - Upright-2.obj", "4": "Table_Base - Supports-1.obj", "5": "Table_Base - Supports-2.obj", "6": "Table_Base - Top-1.obj"}

assets/beam_assembly/original/robot.json

@@ -0,0 +1,34 @@
+{
+    "arm": {
+        "base_pos": [-21.25, -30.0, 0.0],
+        "base_angle": 0.0,
+        "scale": 1.0,
+        "rest_q": [0.63879051, -0.41713369, 0.28274334, 0.65100781, 0.13089969, 1.0559242, 0.82030475]
+    },
+    "gripper": {
+        "type": "robotiq-140",
+        "scale": 1.0
+    },
+    "grasp": {
+        "0": {
+            "antipodals": [[0.0, 0.635, 1.073], [0.0, -0.635, 1.073]],
+            "base_direction": [0.0, 0.0, 1.0]
+        },
+        "1": {
+            "antipodals": [[0.0, 0.635, 1.073], [0.0, -0.635, 1.073]],
+            "base_direction": [0.0, 0.0, 1.0]
+        },
+        "2": {
+            "antipodals": [[-0.635, -1.2, 1.905], [0.635, -1.2, 1.905]],
+            "base_direction": [0.0, -1.0, 0.0]
+        },
+        "3": {
+            "antipodals": [[-0.635, -1.2, 1.905], [0.635, -1.2, 1.905]],
+            "base_direction": [0.0, -1.0, 0.0]
+        },
+        "6": {
+            "antipodals": [[-0.889, 0.0, 1.2], [0.889, 0.0, 1.2]],
+            "base_direction": [0.0, 0.0, 1.0]
+        }
+    }
+}

assets/box/box_stack.xml

@@ -0,0 +1,30 @@
+<redmax model="box_stack">
+    <option integrator="BDF1" timestep="1e-3" gravity="0. 0. -9.8"/>
+
+    <ground pos="0 0 0" normal="0 0 1"/>
+    <default>
+        <ground_contact kn="1e5" kt="1e3" mu="0.8" damping="5e1"/>
+        <general_SDF_contact kn="1e5" kt="5e3" mu="1.0" damping="1e3"/>
+    </default>
+
+    <robot>
+        <link name="mesh1">
+            <joint name="mesh1" type="free3d" axis="0. 0. 0." pos="0 0 10." quat="1 0 0 0" frame="WORLD" damping="0"/>
+            <body name="mesh1" type="mesh" filename="cube_dense.obj" pos="-2 -2 -2" quat="1 0 0 0" scale="4 4 4" transform_type="OBJ_TO_JOINT" density="1" mu="0" rgba="0.42 0.65 0.63 1"/>
+        </link>
+    </robot>
+
+    <robot>
+        <link name="box1">
+            <joint name="box1" type="free3d" axis="0. 0. 0." pos="0. 0. 5." quat="1 0 0 0" frame="WORLD" damping="0"/>
+            <body name="box1" type="SDF" filename="cube.obj" scale="4 4 4" pos="-2 -2 -2" quat="1 0 0 0" transform_type="OBJ_TO_JOINT" dx="0.1" density="1" mu="0" rgba="0.82 0.72 0.58 1"/>
+        </link>
+    </robot>
+
+    <contact>
+		<ground_contact body="mesh1"/>
+        <ground_contact body="box1"/>
+        <general_SDF_contact general_body="mesh1" SDF_body="box1"/>
+	</contact>
+
+</redmax>

assets/color.py

@@ -0,0 +1,26 @@
+'''
+Predefined colors for assembly meshes
+'''
+
+import numpy as np
+
+
+def get_color(part_ids, normalize=True):
+    color_map = {}
+    if len(part_ids) <= 2:
+        colors = np.array([
+            [107, 166, 161, 255],
+            [209, 184, 148, 255],
+        ], dtype=int)
+    else:
+        colors = np.array([
+            [210, 87, 89, 255],
+            [237, 204, 73, 255],
+            [60, 167, 221, 255],
+            [190, 126, 208, 255],
+            [108, 192, 90, 255],
+        ], dtype=int)
+    if normalize: colors = colors.astype(float) / 255.0
+    for i, part_id in enumerate(part_ids):
+        color_map[part_id] = colors[i % len(colors)]
+    return color_map