[](https://youtu.be/FFrl_TEXrUw)
6 |
7 | 
8 |
9 | This repository includes:
10 |
11 | * Implementation of the EquiBot method and a Diffusion Policy baseline that takes point clouds as input.
12 | * A set of three simulated mobile manipulation environments: Cloth Folding, Object Covering, and Box Closing.
13 | * Data generation, training, and evaluation scripts that accompany the above algorithms and environments.
14 |
15 | ## Getting Started
16 |
17 | ### Installation
18 |
19 | This codebase is tested with the following setup: Ubuntu 20.04, an RTX 4090 GPU, CUDA 11.8. In the root directory of the repository, run the following commands:
20 |
21 | ```
22 | conda create -n lfd python=3.10 -y
23 | conda activate lfd
24 |
25 | conda install -y fvcore iopath ffmpeg -c iopath -c fvcore
26 | pip install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cu118
27 | pip install "git+https://github.com/facebookresearch/pytorch3d.git"
28 |
29 | pip install -e .
30 | ```
31 |
32 | Then, in the last two lines of [this config file](equibot/policies/configs/base.yaml), enter the wandb entity and project names for logging purposes. If you do not have a wandb account yet, you can register [here](https://wandb.ai).
33 |
34 | ### Demonstration Generation
35 |
36 | The following code generates demonstrations for simulated mobile environments. To change the number of generated demos, change `--num_demos 50` to a different number.
37 |
38 | ```
39 | python -m equibot.envs.sim_mobile.generate_demos --data_out_dir ../data/fold \
40 | --num_demos 50 --cam_dist 2 --cam_pitches -75 --task_name fold
41 |
42 | python -m equibot.envs.sim_mobile.generate_demos --data_out_dir ../data/cover \
43 | --num_demos 50 --cam_dist 2 --cam_pitches -75 --task_name cover
44 |
45 | python -m equibot.envs.sim_mobile.generate_demos --data_out_dir ../data/close \
46 | --num_demos 50 --cam_dist 1.5 --cam_pitches -45 --task_name close
47 | ```
48 |
49 | ### Training
50 |
51 | The following code runs training for our method and the Diffusion Policy baseline. Fill the dataset path with the `data_out_dir` argument in the previous section. Make sure the dataset path ends with `pcs`. To run this code for the `cover` and `close` environments, substitute occurrences of `fold` with `cover` or `close`.
52 |
53 | ```
54 | # diffusion policy baseline (takes point clouds as input)
55 | python -m equibot.policies.train --config-name fold_mobile_dp \
56 | prefix=sim_mobile_fold_7dof_dp \
57 | data.dataset.path=[data out dir in the last section]/pcs
58 |
59 | # our method (equibot)
60 | python -m equibot.policies.train --config-name fold_mobile_equibot \
61 | prefix=sim_mobile_fold_7dof_equibot \
62 | data.dataset.path=[data out dir in the last section]/pcs
63 | ```
64 |
65 | ### Evaluation
66 |
67 | The commands below evaluate the trained EquiBot policy on the four different setups mentioned in the paper: `Original`, `R+Su`, `R+Sn`, and `R+Sn+P`. To run these evaluations for the DP baseline, replace all occurrences of `equibot` to`dp`. For the log directory, fill `[log_dir]` with the absolute path to the log directory. By default, this directory is `./log`.
68 |
69 | ```
70 | # Original setup
71 | python -m equibot.policies.eval --config-name fold_mobile_equibot \
72 | prefix="eval_original_sim_mobile_fold_equibot_s1" mode=eval \
73 | training.ckpt="[log_dir]/train/sim_mobile_fold_7dof_equibot_s1/ckpt01999.pth" \
74 | env.args.max_episode_length=50 env.vectorize=true
75 |
76 | # R+Su setup
77 | python -m equibot.policies.eval --config-name fold_mobile_equibot \
78 | prefix="eval_rsu_sim_mobile_fold_7dof_equibot_s1" mode=eval \
79 | training.ckpt="[log_dir]/train/sim_mobile_fold_7dof_equibot_s1/ckpt01999.pth" \
80 | env.args.scale_high=2 env.args.uniform_scaling=true \
81 | env.args.randomize_rotation=true env.args.randomize_scale=true env.vectorize=true
82 |
83 | # R+Sn setup
84 | python -m equibot.policies.eval --config-name fold_mobile_equibot \
85 | prefix="eval_rsn_sim_mobile_fold_7dof_equibot_s1" mode=eval \
86 | training.ckpt="[log_dir]/train/sim_mobile_fold_7dof_equibot_s1/ckpt01999.pth" \
87 | env.args.scale_high=2 env.args.scale_aspect_limit=1.33 \
88 | env.args.randomize_rotation=true env.args.randomize_scale=true env.vectorize=true
89 |
90 | # R+Sn+P setup
91 | python -m equibot.policies.eval --config-name fold_mobile_equibot \
92 | prefix="eval_rsnp_sim_mobile_fold_7dof_equibot_s1" mode=eval \
93 | training.ckpt="[log_dir]/train/sim_mobile_fold_7dof_equibot_s1/ckpt01999.pth" \
94 | env.args.scale_high=2 env.args.scale_aspect_limit=1.33 \
95 | env.args.randomize_rotation=true env.args.randomize_scale=true \
96 | +env.args.randomize_position=true +env.args.rand_pos_scale=0.5 env.vectorize=true
97 | ```
98 |
99 | ## License
100 |
101 | This codebase is licensed under the terms of the MIT License.
102 |
--------------------------------------------------------------------------------
/equibot/envs/sim_mobile/assets/closing/template.urdf:
--------------------------------------------------------------------------------
1 |
2 | 0 0 1 0 2 0 3 1 1 1 0 1 4 2 5 2 6 2 5 3 7 3 6 3 2 4 5 4 4 4 2 5 4 5 0 5 5 6 2 6 1 6 5 7 1 7 7 7 7 8 1 8 6 8 1 9 3 9 6 9 0 10 4 10 3 10 4 11 6 11 3 11
79 |0 0 1 0 2 0 3 1 1 1 0 1 4 2 5 2 6 2 5 3 7 3 6 3 2 4 5 4 4 4 2 5 4 5 0 5 5 6 2 6 1 6 5 7 1 7 7 7 7 8 1 8 6 8 1 9 3 9 6 9 0 10 4 10 3 10 4 11 6 11 3 11
79 |