Gen2Sim: Scaling up Robot Learning in Simulation with Generative Models

[[Website]](https://gen2sim.github.io/) [[arXiv]](https://arxiv.org/abs/2310.18308) ______________________________________________________________________ ![](gifs/articulated_policy.gif)

Generalist robot manipulators need to learn a wide variety of manipulation skills across diverse environments. Current robot training pipelines rely on humans to provide kinesthetic demonstrations or to program simulation environments and to code up reward functions for reinforcement learning. Such human involvement is an important bottleneck towards scaling up robot learning across diverse tasks and environments. We propose Generation to Simulation (Gen2Sim), a method for scaling up robot skill learning in simulation by automating generation of 3D assets, task descriptions, task decompositions and reward functions using large pre-trained generative models of language and vision. We generate 3D assets for simulation by lifting open-world 2D object-centric images to 3D using image diffusion models and querying LLMs to determine plausible physics parameters. Given URDF files of generated and human-developed assets, we chain-of-thought prompt LLMs to map these to relevant task descriptions, temporal decompositions, and corresponding python reward functions for reinforcement learning. We show Gen2Sim succeeds in learning policies for diverse long horizon tasks, where reinforcement learning with non temporally decomposed reward functions fails. Gen2Sim provides a viable path for scaling up reinforcement learning for robot manipulators in simulation, both by diversifying and expanding task and environment development, and by facilitating the discovery of reinforcement-learned behaviors through temporal task decomposition in RL. Our work contributes hundreds of simulated assets, tasks and demonstrations, taking a step towards fully autonomous robotic manipulation skill acquisition in simulation.

Installation

Gen2Sim requires Python ≥ 3.8. We have tested on Ubuntu 20.04.

Create a new conda environment with:

conda create -n gen2sim python=3.8
conda activate gen2sim

Install IsaacGym. Follow the instruction to download the package.

cd isaacgym/python
pip install -e .
(test installation) python examples/joint_monkey.py

Install Gen2Sim dependencies. We also provide requirements.txt file with all the listed requirements.
```
pip install -r requirements.txt
```
Gen2Sim currently uses OpenAI API to query LLMs (GPT-4). You need to have an OpenAI API key to use Gen2Sim here/. Then, set the environment variable in your terminal
```
export OPENAI_API_KEY= "YOUR_API_KEY"
```

Getting Started

Asset Generation

Navigate to the asset-gen directory and follow the README.md to generate asset URDFs.

Setup dataset

We provide the robot models, generated rigid-body assets and PartNetMobility articulated assets. The dataset can be downloaded from the following link. Unzip the downloaded dataset:

mkdir data
unzip data.zip

Task and Reward Generation

We provide scripts to prepare the prompt using asset description (URDF):

python task-gen/prompt.py data <log_dir>

data is the path to data folder downloaded above.
<log_dir> path to log LLM outputs for debugging.

Once generations are complete, json asset_to_tasks.json is generated in data folder. This file contains generated tasks and reward functions.

Skill Training in Simulation

We provide scripts to train RL policy using the generated task and reward function in IsaacGym.

python gym/train.py --group_name <group_name> --seed <seed> --env_num <num_envs> --env_name <env_name> --log_dir <log_dir>

<group_name> logging name for the run.
<seed> seed for RL policy.
<num_envs> nums of parallel environments for exploration during RL.
<env_name> environment name.
<log_dir> Path to log tb plots, checkpoints etc.
--headless if working on a server.

Examples:

Tasks generated for asset Microwave (Simple Tasks like OpenMicrowaveDoor, CloseMicrowaveDoor)

python gym/train.py --group_name Microwave --seed 0 --env_num 250 --env_name Microwave --headless

Data Collection using trained skills

To visualize the demonstrations, evaluate trained policies using the following eval script:

python gym/train.py --group_name <group_name> --seed <seed> --env_num 1 --env_name <env_name> --log_dir <log_dir> --save_video <save_video> --is_testing True --model_dir <model_path>

<group_name> logging name for the run.
<seed> seed for RL policy.
<num_envs> use 1 env in eval.
<env_name> environment name.
<model_dir> Path to checkpoint.
<save_video> Boolean to save video or not.

ToDos:

Support for USD to better structure environments.
Adding GPT-4V to automate part, affordance prediction.
Add support for more tasks in gym.

Acknowledgement

We thank the following open-sourced projects:

Simulator used in our setup is IsaacGym
Our RL training code is based on PartManip

Citation

If you find our work useful, please consider citing us!

@misc{katara2023gen2sim,
            title={Gen2Sim: Scaling up Robot Learning in Simulation with Generative Models}, 
            author={Pushkal Katara and Zhou Xian and Katerina Fragkiadaki},
            year={2023},
            eprint={2310.18308},
            archivePrefix={arXiv},
            primaryClass={cs.RO}
      }

pushkalkatara / Gen2Sim

readme