first commit

real-stanford · Oct 10, 2020 · a20dffd · a20dffd
1 parent 99da3c3
commit a20dffd
Show file tree

Hide file tree

Showing 55 changed files with 4,626 additions and 0 deletions.
diff --git a/README.md b/README.md
@@ -0,0 +1,105 @@
+# Learning 3D Dynamic Scene Representations for Robot Manipulation
+### [project](https://dsr-net.cs.columbia.edu/) | [paper(TBD)]() | [videos(TBD)]()
+This repo contains the PyTorch implementation for paper Learning 3D Dynamic Scene Representations for Robot Manipulation.
+
+## Content
+
+- [Prerequisites](#prerequisites)
+- [Data Preparation](#data-preparation)
+- [Pretrained Models](#pretrained-models)
+- [Training](#training)
+
+## Prerequisites
+
+The code is built with Python 3.6. Libraries are listed in [requirements.txt](requirements.txt):
+
+## Data Preparation
+
+### Download Testing Data
+The following two testing datasets can be download in [link(TBD)]().
+- Sim: 400 sequences, generated in pybullet.
+- Real: 150 sequences, with full annotations.
+
+### Generate Training Data
+
+Download object mesh from [link(TBD)]().
+
+To generate data in simulation, one can run
+```
+python data_generation.py --data_path [path to data] --train_num [number of training sequences] --train_num [number of testing sequences] --object_type [type of objects]
+```
+Where the `object_type` can be `cube`, `shpenet`, or `ycb`.
+The training data in the paper can be generated with the followint scripts:
+```
+# cube
+python data_generation.py --data_path data/cube_train --train_num 4000 --train_num 400 --object_type cube
+
+# shapenet
+python data_generation.py --data_path data/shapenet_train --train_num 4000 --train_num 400 --object_type shapenet
+```
+
+## Pretrained Models
+Some of the pretrained models can be download in [pretrained_models](pretrained_models).
+To evaluate the pretrained models, one can run
+```
+python test.py --resume [path to model] --data_path [path to data] --model_type [type of model] --test_type [type of test]
+```
+where `model_type` can be one of the following:
+- `dsr`: DSR-Net introduced in the paper.
+- `single`: It does not use any history aggregation.
+- `nowarp`: It does not warp the representation before aggregation.
+- `gtwarp`: It warps the representation with ground truth motion (i.e., performance oracle)
+- `3dflow`: It predicts per-voxel scene flow for the entire 3D volume.
+
+Both motion prediction and mask prediction can be evaluated by choosing different `test_type`:
+- motion prediction: `motion_visible` or `motion_full`
+- mask prediction: `mask_ordered` or `mask_unordered`
+
+(Please refer to our paper for detailed explanation of each type of evaluation)
+
+Here are several examples:
+```
+# evaluate mask prediction (ordered) of DSR-Net using real data:
+python test.py --resume [path to dsr model] --data_path [path to real data] --model_type dsr --test_type mask_ordered
+
+# evaluate mask prediction (unordered) of DSR-Net(finetuned) using real data:
+python test.py --resume [path to dsr_ft model] --data_path [path to real data] --model_type dsr --test_type mask_unordered
+
+# evaluate motion prediction (visible surface) of NoWarp model using sim data:
+python test.py --resume [path to nowarp model] --data_path [path to sim data] --model_type nowarp --test_type motion_visible
+
+# evaluate motion prediction (full volume) of SingleStep model using sim data:
+python test.py --resume [path to single model] --data_path [path to sim data] --model_type single --test_type motion_full
+```
+
+
+## Training
+Various training options can be modified or toggled on/off with different flags (run `python main.py -h` to see all options):
+```
+usage: train.py [-h] [--exp EXP] [--gpus GPUS [GPUS ...]] [--resume RESUME]
+                [--data_path DATA_PATH] [--object_num OBJECT_NUM]
+                [--seq_len SEQ_LEN] [--batch BATCH] [--workers WORKERS]
+                [--model_type {dsr,single,nowarp,gtwarp,3dflow}]
+                [--transform_type {affine,se3euler,se3aa,se3spquat,se3quat}]
+                [--alpha_motion ALPHA_MOTION] [--alpha_mask ALPHA_MASK]
+                [--snapshot_freq SNAPSHOT_FREQ] [--epoch EPOCH] [--finetune]
+                [--seed SEED] [--dist_backend DIST_BACKEND]
+                [--dist_url DIST_URL]
+```
+### Training of DSR-Net
+Since the aggregation ability depends on the accuracy of motion prediction, we split the training process into three stages from easy to hard: (1) single-step on cube dataset; (2) multi-step on cube dataset; (3) multi-step on ShapeNet dataset. 
+```
+# Stage 1 (single-step on cube dataset)
+python train.py --exp dsr_stage1 --data_path [path to cube dataset] --seq_len 1 --model_type dsr --epoch 30
+
+# Stage 2 (multi-step on cube dataset)
+python train.py --exp dsr_stage2 --resume [path to stage1] --data_path [path to cube dataset] --seq_len 10 --model_type dsr --epoch 20 --finetune
+
+# Stage 3 (multi-step on shapenet dataset)
+python train.py --exp dsr_stage3 --resume [path to stage2] --data_path [path to shapenet dataset] --seq_len 10 --model_type dsr --epoch 20 --finetune
+```
+
+### Training of Baselines
+For `nowarp` and `gtwarp`, use the same scripts as DSR-NEt with corresponding `model_type`.
+
+For `single` and `3dflow`, only two stages is used (one for cube, one for shapenet) and `seq_len` always choose `1`.
diff --git a/assets/fence/obj.mtl b/assets/fence/obj.mtl
@@ -0,0 +1,26 @@
+# Color definition for Tinkercad Obj File 2015
+
+newmtl color_11593967
+Ka 0 0 0 
+Kd 0.6901960784313725 0.9098039215686274 0.9372549019607843
+d 1.0
+illum 0.0
+
+newmtl color_16500122
+Ka 0 0 0 
+Kd 0.984313725490196 0.7725490196078432 0.6039215686274509
+d 1.0
+illum 0.0
+
+newmtl color_16311991
+Ka 0 0 0 
+Kd 0.9725490196078431 0.9019607843137255 0.7176470588235294
+d 1.0
+illum 0.0
+
+newmtl color_13165757
+Ka 0 0 0 
+Kd 0.7843137254901961 0.8941176470588236 0.7411764705882353
+d 1.0
+illum 0.0
+
diff --git a/assets/fence/tinker.obj b/assets/fence/tinker.obj
@@ -0,0 +1,109 @@
+# Object Export From Tinkercad Server 2015
+
+mtllib obj.mtl
+
+o obj_0
+v 52 		-50 		30
+v 52 		-50 		0
+v 52 		50 		0
+v 52 		50 		30
+v 50 		50 		30
+v 50 		-50 		30
+v 50 		-50 		0
+v 50 		50 		0
+v 50 		52 		0
+v 50 		52 		30
+v -52 		50 		30
+v -52 		50 		0
+v -52 		-50 		0
+v -52 		-50 		30
+v 50 		-52 		30
+v 50 		-52 		0
+v -50 		50 		30
+v -50 		-50 		30
+v -50 		-50 		0
+v -50 		50 		0
+v -50 		-52 		0
+v -50 		-52 		30
+v -50 		52 		30
+v -50 		52 		0
+# 24 vertices
+
+g group_0_11593967
+
+usemtl color_11593967
+s 0
+
+f 1 	2 	3
+f 1 	3 	4
+f 4 	5 	6
+f 4 	6 	1
+f 2 	7 	8
+f 2 	8 	3
+f 6 	7 	2
+f 6 	2 	1
+f 4 	3 	8
+f 4 	8 	5
+f 8 	7 	6
+f 8 	6 	5
+# 12 faces
+
+g group_0_13165757
+
+usemtl color_13165757
+s 0
+
+f 15 	16 	7
+f 15 	7 	6
+f 21 	22 	18
+f 21 	18 	19
+f 6 	18 	22
+f 6 	22 	15
+f 22 	21 	16
+f 22 	16 	15
+f 16 	21 	19
+f 16 	19 	7
+f 6 	7 	19
+f 6 	19 	18
+# 12 faces
+
+g group_0_16311991
+
+usemtl color_16311991
+s 0
+
+f 11 	12 	13
+f 11 	13 	14
+f 17 	11 	14
+f 17 	14 	18
+f 19 	13 	12
+f 19 	12 	20
+f 14 	13 	19
+f 14 	19 	18
+f 17 	20 	12
+f 17 	12 	11
+f 19 	20 	17
+f 19 	17 	18
+# 12 faces
+
+g group_0_16500122
+
+usemtl color_16500122
+s 0
+
+f 9 	10 	5
+f 9 	5 	8
+f 23 	24 	20
+f 23 	20 	17
+f 8 	20 	24
+f 8 	24 	9
+f 10 	23 	17
+f 10 	17 	5
+f 10 	9 	24
+f 10 	24 	23
+f 17 	20 	8
+f 17 	8 	5
+# 12 faces
+
+ #end of obj_0
+
diff --git a/assets/fence/tinker.urdf b/assets/fence/tinker.urdf
@@ -0,0 +1,30 @@
+<robot name="tray">
+  <link name="tray_base_link">
+    <contact>
+      <lateral_friction value="0.2"/>
+      <rolling_friction value="0.001"/>
+      <contact_cfm value="0.0"/>
+      <contact_erp value="1.0"/>
+    </contact>
+    <inertial>
+      <origin rpy="0 0 0" xyz="0 0 0"/>
+      <mass value="0"/>
+      <inertia ixx="0" ixy="0" ixz="0" iyy="0" iyz="0" izz="0"/>
+    </inertial>
+    <visual>
+      <origin rpy="0 0 0" xyz="0 0 0"/>
+      <geometry>
+        <mesh filename="tinker.obj" scale="0.0055 0.0055 0.002"/>
+      </geometry>
+      <material name="tray_material">
+        <color rgba="0.7 0.7 0.7 1"/>
+      </material>
+    </visual>
+    <collision>
+      <origin rpy="0 0 0" xyz="0 0 0"/>
+      <geometry>
+        <mesh filename="tinker.obj" scale="0.0055 0.0055 0.04"/>
+      </geometry>
+    </collision>
+  </link>
+</robot>
diff --git a/assets/object_id/shapenet_id.json b/assets/object_id/shapenet_id.json
@@ -0,0 +1,82 @@
+{
+    "bottle": {
+        "category_id": "02876657",
+        "object_id": [
+            "6b8b2cb01c376064c8724d5673a063a6",
+            "547fa0085800c5c3846564a8a219239b",
+            "91235f7d65aec958ca972daa503b3095",
+            "9dff3d09b297cdd930612f5c0ef21eb8",
+            "d45bf1487b41d2f630612f5c0ef21eb8",
+            "898101350771ff942ae40d06128938a1",
+            "3b26c9021d9e31a7ad8912880b776dcf",
+            "ed55f39e04668bf9837048966ef3fcb9",
+            "74690ddde9c0182696c2e7960a15c619",
+            "e4ada697d05ac7acf9907e8bdd53291e",
+            "fdc47f5f8dff184830eaaf40a8a562c1",
+            "ad33ed7da4ef1b1cca18d703b8006093",
+            "4d4fc73864844dad1ceb7b8cc3792fd",
+            "42f85b0eb5e9fd508f9c4ecc067067e9"
+        ],
+        "global_scaling": [0.7, 0.8],
+        "large_scaling": [1.1, 1.2]
+    },
+    "can": {
+        "category_id": "02946921",
+        "object_id": [
+            "f4ad0b7f82c36051f51f77a6d7299806",
+            "a70947df1f1490c2a81ec39fd9664e9b",
+            "b6c4d78363d965617cb2a55fa21392b7",
+            "91a524cc9c9be4872999f92861cdea7a",
+            "96387095255f7080b7886d94372e3c76",
+            "9b1f0ddd23357e01a81ec39fd9664e9b"
+        ],
+        "global_scaling": [0.7, 0.9],
+        "large_scaling": [1.1, 1.2]
+    },
+    "mug": {
+        "category_id": "03797390",
+        "object_id": [
+            "599e604a8265cc0a98765d8aa3638e70",
+            "b46e89995f4f9cc5161e440f04bd2a2",
+            "9c930a8a3411f069e7f67f334aa9295c",
+            "2d10421716b16580e45ef4135c266a12",
+            "71ca4fc9c8c29fa8d5abaf84513415a2"
+        ],
+        "global_scaling": [0.8, 0.9],
+        "large_scaling": [1.1, 1.2]
+    },
+    "sofa": {
+        "category_id": "04256520",
+        "object_id": [
+            "930873705bff9098e6e46d06d31ee634",
+            "f094521e8579917eea65c47b660136e7",
+            "adc4a9767d1c7bae8522c33a9d3f5757",
+            "d0e419a11fd8f4bce589b08489d157d",
+            "9aef63feacf65dd9cc3e9831f31c9164",
+            "d55d14f87d65faa84ccf9d6d546b307f",
+            "526c4f841f777635b5b328c62af5142",
+            "65fce4b727c5df50e5f5c582d1bee164",
+            "7c68894c83afb0118e8dcbd53cc631ab"
+        ],
+        "global_scaling": [0.4, 0.7],
+        "large_scaling": [1.0, 1.2]
+    },
+    "phone": {
+        "category_id": "04401088",
+        "object_id": [
+            "611afaaa1671ac8cc56f78d9daf213b",
+            "b8555009f82af5da8c3645155d02fccc",
+            "2725909a09e1a7961df58f4da76e254b",
+            "9d021614c39c53dabee972a203aaf80",
+            "fe34b663c44baf622ad536a59974757f",
+            "2bb42eb0676116d41580700c4211a379",
+            "9efabcf2ff8a4be9a59562d67b11f3d",
+            "b9f67617cf320c20de4349e5bfa4fedb",
+            "78855e0d8d27f00b42e82e1724e35ca",
+            "401d5604ebfb4b43a7d4c094203303b1"
+        ],
+        "global_scaling": [0.7, 1.0],
+        "large_scaling": [1.0, 1.1]
+
+    }
+}
diff --git a/assets/object_id/ycb_id.json b/assets/object_id/ycb_id.json
@@ -0,0 +1,29 @@
+{
+  "large_list": [
+    "002_master_chef_can",
+    "004_sugar_box",
+    "006_mustard_bottle",
+    "flipped-065-j_cups",
+    "071_nine_hole_peg_test",
+    "051_large_clamp"
+  ],
+  "normal_list": [
+    "005_tomato_soup_can",
+    "007_tuna_fish_can",
+    "008_pudding_box",
+    "009_gelatin_box",
+    "010_potted_meat_can",
+    "025_mug",
+    "061_foam_brick",
+    "077_rubiks_cube",
+
+    "flipped-065-a_cups",
+    "flipped-065-d_cups",
+    "flipped-065-g_cups",
+    "filled-073-a_lego_duplo",
+    "filled-073-b_lego_duplo",
+    "filled-073-c_lego_duplo",
+    "filled-073-d_lego_duplo",
+    "filled-073-f_lego_duplo"
+  ]
+}