DOCSEARCHbyAlgolia
Submit Docs

MimicGen Documentation

repository·main·Indexed 20 days ago

https://github.com/nvlabs/mimicgen

A scalable data generation system for robot learning that leverages human demonstrations to create massive datasets in simulation. MimicGen provides tools to bridge the gap between limited human data and large-scale requirements of robot learning algorithms, including simulation environments and utilities for dataset management, configuration generation, and subtask annotation. It integrates with robomimic and robosuite, offering specialized environments for tasks such as coffee manipulation, nut assembly, and pick-and-place.

Tokens
17.4K
Snippets
43
Records
73
Agent score
69%

What's inside MimicGen

  1. Overview of mimicgen.scripts submodules

    main

    The mimicgen.scripts package provides a collection of utility scripts for managing the MimicGen workflow. These scripts cover various stages of the pipeline, including dataset management, configuration generation, and visualization.

    Key functional areas include:

    • Dataset Management: Downloading datasets (download_datasets), preparing source datasets (prepare_src_dataset), and merging HDF5 files (merge_hdf5).
    • Configuration Generation: Generating templates (generate_config_templates), core configs (generate_core_configs), training configs (generate_core_training_configs), and robot transfer configs (generate_robot_transfer_configs).
    • Task & Subtask Utilities: Annotating subtasks (annotate_subtasks), visualizing subtasks (visualize_subtasks), and getting reset videos (get_reset_videos).
    • Demonstration & Info: Running random action demos (demo_random_action) and retrieving source information (get_source_info).

  2. Overview of MimicGen

    main

    MimicGen is a data generation system designed for scalable robot learning using human demonstrations. It provides tools to create large-scale datasets for robot training. The system includes simulation environments and can generate tens of thousands of task demonstrations across various tasks.

    Key resources:

  3. Identify key MimicGen modules and files

    main

    The MimicGen codebase is organized into several functional areas:

    Core Data Generation

    • mimicgen/datagen: Contains the core logic for generation.
      • data_generator.py: The DataGenerator class for generating trajectories.
      • datagen_info.py: The DatagenInfo class for grouping simulation information.
      • selection_strategy.py: SelectionStrategy classes for choosing source demos.
      • waypoint.py: Waypoint classes for end-effector controller targets.
    • mimicgen/env_interfaces: Environment Interface implementations to provide simulation data to the generator.

    Configuration and Tasks

    • mimicgen/configs: Implementation of data generation configuration classes.
      • config.py: Base configuration class.
      • task_spec.py: TaskSpec objects for defining subtask sequences.
      • robosuite.py: Configurations specific to the robosuite environment.
    • mimicgen/exps/templates: JSON templates for data generation configurations for specific tasks.

    Utilities and Scripts

    • mimicgen/scripts/generate_dataset.py: The main entry point script for running data generation.
    • mimicgen/utils: General utility functions and classes.
    • mimicgen/envs & mimicgen/models: Provided robosuite simulation environments and assets.

  4. Available Robosuite environments in MimicGen

    main

    The mimicgen.envs.robosuite package provides a collection of robot simulation environments built on top of robosuite. These environments are organized into submodules, each representing a specific task or manipulation scenario. Use these modules to instantiate specialized environments for training or testing MimicGen agents.

    Available task modules include:

    • coffee: Coffee-related manipulation tasks.
    • hammer_cleanup: Tasks involving cleaning up with a hammer.
    • kitchen: General kitchen manipulation tasks.
    • mug_cleanup: Tasks involving cleaning up a mug.
    • nut_assembly: Nut and bolt assembly tasks.
    • pick_place: Standard pick-and-place manipulation tasks.
    • stack: Stacking tasks.
    • three_piece_assembly: Assembly tasks involving three distinct pieces.

  5. Explore the mimicgen.envs package

    main
    The mimicgen.envs package serves as the entry point for environment-related modules in MimicGen. It currently contains the robosuite subpackage, which provides specialized environment support for Robosuite-based tasks. Users looking to configure or interact with simulation environments should look into the mimicgen.envs.robosuite subpackage for specific implementation details.

  6. Typical MimicGen Data Generation Workflow

    main

    A full production workflow for generating data for a new task follows these four steps:

    1. Collect source demonstrations: Gather human demonstrations (e.g., via teleoperation) and ensure they are in an HDF5 format compatible with robomimic.
    2. Prepare source demonstrations with annotations:
      • Use scripts/prepare_src_dataset.py to add necessary metadata.
      • This requires implementing an Environment Interface class for your simulator and task. The interface must map environment actions to target poses and provide object poses for subtasks.
      • Optionally, use scripts/annotate_subtasks.py to manually segment demonstrations if the environment interface does not provide automatic subtask termination signals.
    3. Run data generation: Configure and launch the process using scripts/generate_dataset.py.
    4. Run policy learning: Use the generated demo.hdf5 (which is robomimic-compatible) to train agents using algorithms like Behavioral Cloning.

  7. Randomize subtask boundaries with subtask_term_offset_range

    main

    To increase data diversity, MimicGen supports randomizing subtask boundaries using the subtask_term_offset_range parameter in the TaskSpec.

    At the start of data generation, an additive offset is uniformly sampled from the provided range and applied to the detected 0 to 1 transition index.

    Use Case: Ensuring subtask completion If a signal (like grasp_1) triggers immediately upon contact, but you want the subtask to include the subsequent lifting motion, you can specify an offset range like (5, 10). This ensures the subtask boundary always occurs 5 to 10 timesteps after the initial signal transition.

  8. Understand the MimicGen data generation workflow

    main

    MimicGen automates the generation of new robotic demonstrations by transforming existing source demonstrations. The workflow follows these steps:

    1. Task Specification: A task is defined as a sequence of object-centric subtasks using a TaskSpec object.
    2. Parsing Source Demos: MimicGen parses source demonstrations into contiguous subtask segments using Subtask Termination Signals.
    3. Information Collection: Information required for generation (such as object poses at the start of subtasks) is collected into DatagenInfo objects. These are extracted from both the source demonstrations and the current simulation scene.
    4. Environment Integration: Environment Interface classes bridge the underlying simulation environments to the DatagenInfo objects.
    5. Selection: A SelectionStrategy is used to choose which subtask segments from the source demonstrations to transform during a generation trial.
    6. Execution: The DataGenerator class executes the transformation by managing end-effector target poses via Waypoint classes, which a controller then executes.

  9. Use Waypoint classes to represent robot trajectories

    main

    MimicGen provides a hierarchy of classes in datagen/waypoint.py to represent 6-DoF poses and gripper actions over time:

    • Waypoint: Represents a single 6-DoF target pose and the corresponding gripper action for a specific timestep.
    • WaypointSequence: A collection of Waypoint objects representing a segment of a trajectory.
    • WaypointTrajectory: A sequence of WaypointSequence objects representing a full, multi-segment 6-DoF trajectory.

    You can build complex trajectories by combining these objects using interpolation or direct appending.

    from datagen.waypoint import Waypoint, WaypointSequence, WaypointTrajectory

# A single waypoint
wp = Waypoint(pose=my_pose, gripper_action=my_gripper_action)

# A sequence of waypoints
seq = WaypointSequence(sequence=[wp1, wp2, wp3])

# A full trajectory
trajectory = WaypointTrajectory()

  10. Understand the DataGenerator workflow

    main

    The DataGenerator class (datagen/data_generator.py) is the core engine for generating new demonstration trajectories. The workflow follows these steps:

    1. Loading: The internal _load_dataset method parses a source dataset into subtask segments using Subtask Termination Signals. Each segment consists of a sequence of DatagenInfo objects.
    2. Generation Loop: The generate method is called repeatedly (typically via scripts/generate_dataset.py) to attempt new trajectory generations.
    3. Selection: For each attempt, the select_source_demo method uses a SelectionStrategy to pick a reference source subtask segment to transform.
    4. Composition: WaypointTrajectory objects are used to transform and compose these subtask segments into a complete trajectory.

  11. Understand Subtask Termination Signals

    main

    MimicGen uses subtask termination signals to split source demonstrations into contiguous subtask segments. Each signal is a flat numpy array with binary entries (0 or 1).

    How segmentation works

    1. Mapping: The mapping between subtasks and signals is defined in the TaskSpec object within your MimicGen config JSON.
    2. Inference: MimicGen identifies the end of a subtask by finding the first 0 to 1 transition in the corresponding signal.
    3. Final Subtask: The last subtask in a demonstration does not have a signal; it ends at the conclusion of the source demonstration.

    Verifying signals in your dataset

    You can use the get_source_info.py script to inspect signals. For a specific episode (e.g., demo_0), signals are located at the HDF5 path: f["data/demo_0/datagen_info/subtask_term_signals"]. Under this group, each dataset is named after its specific signal (e.g., grasp_1).

  12. Manage configurations using the robomimic Config system

    main

    MimicGen uses the robomimic configuration system. Configurations are defined as JSON files and can be accessed using either dictionary syntax or "dot" syntax. This allows for flexible manipulation of experiment parameters.

    Example access patterns:

    • Dictionary syntax: config["experiment"]["name"]
    • Dot syntax: config.experiment.name