Step 3: Generate the Dataset#
With annotated source demonstrations in hand, MimicGen can synthesize new ones. For each attempt, the scene is randomized, a source segment is selected per subtask, and its end-effector trajectory is rigidly transformed to the new object poses and replayed.
Parallel data generation for the Franka cube-stacking task.#
The commands below use the pre-annotated source dataset that ships with the repository. If you ran
Steps 1–2, point --input_file at your own dataset_franka_annotated.hdf5 instead.
Small-Scale Generation#
Start with a small-scale run in a Kit window (--viz kit) to sanity-check the setup:
python scripts/generate_dataset.py \
--env_name Isaac-Stack-Cube-Franka-IK-Rel-v0 \
--viz kit \
--num_envs 20 \
--alg mimicgen \
--generation_num_trials 10 \
--task_descriptor autodata_examples/tasks/franka_cube_stack.yaml \
--embodiment autodata_examples/embodiments/franka_ik_rel.yaml \
--input_file ./datasets/annotated_datasets/dataset_franka_annotated.hdf5 \
--output_file ./datasets/generated_dataset_mimicgen_franka_small.hdf5
You will see the robot repeatedly attempt the task under new cube placements. Failed attempts are normal and are not exported.
While it runs, the console reports the number of source episodes loaded into the pool and a running tally after every attempt to generate a new demonstration:
Loaded 10 source episodes into the datagen pool
**************************************************
5/10 (50.0%) successful demos generated
**************************************************
The script shuts down automatically after 10 successful demonstrations are generated.
Full-Scale Generation#
For dataset-scale generation, run headless (--viz none) with parallel environments:
python scripts/generate_dataset.py \
--env_name Isaac-Stack-Cube-Franka-IK-Rel-v0 \
--viz none \
--num_envs 500 \
--alg mimicgen \
--generation_num_trials 1000 \
--task_descriptor autodata_examples/tasks/franka_cube_stack.yaml \
--embodiment autodata_examples/embodiments/franka_ik_rel.yaml \
--input_file ./datasets/annotated_datasets/dataset_franka_annotated.hdf5 \
--output_file ./datasets/generated_dataset_mimicgen_franka.hdf5
Progress is printed after every attempt (successful demonstrations, attempts, and the running success rate).
Note
Expected data generation success rate and time
Data generation success rate: ~40%
Data generation time: ~15 minutes
Numbers are based on using an RTX PRO 6000 Blackwell GPU with the provided command.
Key Parameters#
Parameter |
Description |
|---|---|
|
Generation algorithm: |
|
Env id; if omitted, read from the source dataset’s metadata. |
|
Overrides the task descriptor’s |
|
Number of parallel environments, each generating independently. |
|
Visualization mode ( |
Continue to Step 4: Validate the Generated Dataset.