ClutterGen: A Cluttered Scene Generator for Robot Learning

<a href="https://yjia.net/">Yinsen Jia</a>, <a href="http://boyuanchen.com/">Boyuan Chen</a> Duke University

<a href=http://generalroboticslab.com/ClutterGen target="_blank" style="text-decoration: underline;">[Project Page]</a> <a href=https://www.youtube.com/watch?v=ZD6JslZgAfQ target="_blank" style="text-decoration: underline;">[Video]</a> <a href=https://arxiv.org/abs/2407.05425 target="_blank" style="text-decoration: underline;">[arXiv]</a>

Abstract

ClutterGen, a physically compliant simulation scene generator capable of producing highly diverse, cluttered, and stable scenes for robot learning. Generating such scenes is challenging as each object must adhere to physical laws like gravity and collision. As the number of objects increases, finding valid poses becomes more difficult, necessitating significant human engineering effort, which limits the diversity of the scenes. To overcome these challenges, we propose a reinforcement learning method that can be trained with physics-based reward signals provided by the simulator. Our experiments demonstrate that ClutterGen can generate cluttered object layouts with up to ten objects on confined table surfaces. Additionally, our policy design explicitly encourages the diversity of the generated scenes for open-ended generation. Our real-world robot results show that ClutterGen can be directly used for clutter rearrangement and stable placement policy training.

<img src="paper/final/png/teaser_2.png" width="600"> (a/b) When the number of objects in the environment increases, the difficulty of creating such a stable setup also increases. The traditional heuristic method cannot create a simulation scene above 7 objects, while ClutterGen consistently achieves high success rates. (c) Diverse, cluttered, and stable simulation setups created by ClutterGen.

Content

<ul style="font-size:17px;"> <li><a href="#prerequisites">Prerequisites</a></li> <li><a href="#training">Training</a></li> <li><a href="#evaluation">Evaluation</a></li> <li><a href="#real-robot-applications">Real-robot</a></li> <li><a href="#bibTex">BibTex</a></li> </ul>

Prerequisites

Please clone the repository first, (Link needs to be changed)

git clone https://github.com/generalroboticslab/ClutterGen.git
cd ClutterGen
git checkout main

We use conda to create a virtual environment with python3.9 and other requried packages.

conda create --name ClutterGen python=3.9

Activate the environment and install the required packages,

conda activate ClutterGen
pip install -r requirements.txt --no-cache-dir

Training

To train full ClutterGen model, run the following command:

We use the wandb package to log the training stats, which will ask you to log in the account if you specify --collect_data. You can create a free account on wandb.

python ppo_scene_generation.py \
--collect_data \
--result_dir train_res \
--num_envs 16 \
--asset_root 'assets' \
--object_pool_folder 'group_objects/group0_dinning_table'

# --collect_data: Save training results and log the training stats to wandb.
# --result_dir: The directory to save the training results.
# --num_envs: The number of parallel environments to train (reduce if you run out of GPU memory).
# --asset_root: The root directory of the object assets.
# --object_pool_folder: The folder name of the object pool. We provide several object pools in the `assets/group_objects` folder.

Baselines

Random Rejection Sampling (RRS). This method heuristically computes the position of the supporting surface in the queried scene and randomly places objects within the queried region. RRS does not need to be trained. Please check the Evaluation section for the evaluation command of this method.

ClutterGen-OpenLoop (OL). This method uses the same architecture as ClutterGen, but without previous object movement trajectory and placement pose information (including the reward) for the next attempt.

python ppo_scene_generation.py \
--collect_data \
--result_dir train_res \
--num_envs 16 \
--asset_root 'assets' \
--object_pool_folder 'group_objects/group0_dinning_table' \
--open_loop

ClutterGen-ShortMemory (SM). This method uses the same architecture but only takes the latest attempt history (buffer length = 1) for the next attempt.

python ppo_scene_generation.py \
--collect_data \
--result_dir train_res \
--num_envs 16 \
--asset_root 'assets' \
--object_pool_folder 'group_objects/group0_dinning_table' \
--short_memory

ClutterGen-Normal. This method uses the truncated normal distribution instead of the beta distribution for the object placement policy. To keep the code clean and independent, we implement this method in another branch. First, you need to checkout the branch,

git checkout cluttergen_normal

Then, you can run the following command to train the model,

python ppo_scene_generation.py \
--collect_data \
--result_dir train_res \
--num_envs 16 \
--asset_root 'assets' \
--object_pool_folder 'group_objects/group0_dinning_table'

Note: If you want to watch the training process, you can add --rendering and --realtime to the command. However, it will slow down the training process and increase the memory usage.

We also provide --seed to set the random seed for the training.

Evaluation

After the training is done, your checkpoint will be saved in the /[--result_dir]/Union/checkpoint and the training description file will be saved in the /[--result_dir]/Union/Json. By default, [--result_dir] is train_res.

Each checkpoint is saved with the following format, [CHECKPOINT NAME]_[EPISODE]. For example, in Union_2024_05_17_154159_Sync_Beta_group1_studying_table_table_PCExtractor_Rand_ObjPlace_Goal_maxObjNum10_maxPool10_maxScene1_maxStab40_contStab20_Epis2Replaceinf_Weight_rewardPobj100.0_seq5_step80_trial3_entropy0.01_seed123456_best, the checkpoint name is Union_2024_05_17_154159_Sync_Beta_group1_studying_table_table_PCExtractor_Rand_ObjPlace_Goal_maxObjNum10_maxPool10_maxScene1_maxStab40_contStab20_Epis2Replaceinf_Weight_rewardPobj100.0_seq5_step80_trial3_entropy0.01_seed123456 and the episode name is best.

We have provided a script to evaluate the saved model checkpoints evaluation_sg.py. All evaluation results will be stored in /eval_res if --collect_data is specified during evaluation.

The evaluations of ClutterGen, ClutterGen-OL, ClutterGen-SM can be run by following commands,

python evaluation_sg.py \
--collect_data \
--num_envs 1 \
--num_trials 1000 \
--checkpoint [CHECKPOINT NAME] \
--index_episode [CHECKPOINT EPISODE] \
--max_num_placing_objs_lst "[10]"

# --num_trials: The number of trials (episodes) to evaluate the model.
# --checkpoint: The folder name of the checkpoint to evaluate, which is saved in the `train_res/Union` directory. E.g. "Union_2024_05_17_154159_Sync_Beta_group1_studying_table_table_PCExtractor_Rand_ObjPlace_Goal_maxObjNum10_maxPool10_maxScene1_maxStab40_contStab20_Epis2Replaceinf_Weight_rewardPobj100.0_seq5_step80_trial3_entropy0.01_seed123456" is the checkpoint name.
# --index_episode: The episode of the checkpoint to evaluate. If you want to evaluate the best success rate checkpoint, you can set it to `best`.
# --max_num_placing_objs_lst: The maximum number of objects to place in the scene. You can specify multiple numbers in the list, e.g., "[0, 5, 10]". Each number will be evaluated separately for [--num_trials] episodes.

The evaluations of ClutterGen-Normal can be run by following commands, First, you need to checkout the branch,

git checkout cluttergen_normal

Then, you can run the following command to evaluate the model,

python evaluation_sg.py \
--collect_data \
--num_envs 1 \
--num_trials 1000 \
--checkpoint [CHECKPOINT NAME] \
--index_episode [CHECKPOINT EPISODE] \
--max_num_placing_objs_lst "[10]"

The evaluations of RRS can be run by following commands,

python evaluation_sg.py \
--collect_data \
--num_envs 1 \
--num_trials 1000 \
--heuristic_policy \
--checkpoint [CHECKPOINT NAME] \
--max_num_placing_objs_lst "[10]"

# Here, --checkpoint is used to specify the environment configuration.

Note: If you want to watch the evaluation process, you can add --rendering and --realtime to the command. However, it will slow down the evaluation process and increase the memory usage.

Visualization for Evaluation

After evaluations with --collect_data, the evaluation results will be saved in /eval_res. We provide visualization script to visualize the result.

To plot the success rate curve, run the following command,

python plot_utils.py \
--task SuccessRate \
--evalUniName [EVALUATION NAME] \
--RRSName [RRS EVALUATION NAME] \

# --evalUniName: The name of the evaluation result folder, which is saved in the `eval_res/Union/trajectories` directory after evaluation.
# --RRSName: The name of the RRS evaluation result folder.

The plot will be saved to results/success_summary.pdf.

<img src="paper/final/png/teaser_successrate.png" width="300"> The success rate comparision between ClutterGen and RRS.

To plot the diversity map and the stable steps of the evaluation, run the following command,

pyth

ClutterGen

Install / Use

README