Workflow:Haosulab ManiSkill Imitation Learning Pipeline
| Knowledge Sources | |
|---|---|
| Domains | Imitation_Learning, Robotics_Simulation, Data_Engineering |
| Last Updated | 2026-02-15 12:00 GMT |
Overview
End-to-end process for training robot manipulation policies from demonstration data using behavior cloning or diffusion policy in ManiSkill.
Description
This workflow covers the complete imitation learning pipeline in ManiSkill, from obtaining demonstration trajectories to training a policy that can reproduce the demonstrated behaviors. It supports multiple demonstration sources (motion planning, teleoperation, pre-recorded datasets), trajectory format conversion for different observation and control modes, and training with various algorithms including behavior cloning (BC) and diffusion policy. The pipeline handles ManiSkill's HDF5 trajectory format, control mode conversion between demonstrations and training, and GPU-accelerated evaluation.
Usage
Execute this workflow when you have access to demonstration data (either pre-recorded or generated via motion planning/teleoperation) and want to train a policy that imitates the demonstrated behavior. This is appropriate when the task is difficult to specify via a reward function alone, or when you want to bootstrap policy learning from expert demonstrations before optional RL fine-tuning.
Execution Steps
Step 1: Demonstration Data Acquisition
Obtain demonstration trajectories for the target task. Demonstrations can be sourced from the official ManiSkill demo dataset (downloaded via the provided CLI tool), generated via motion planning solutions, or collected through human teleoperation. Each demonstration is stored as a paired HDF5 trajectory file and JSON metadata file.
Key considerations:
- Official demos are available for many built-in tasks and can be downloaded via the CLI download tool
- Motion planning demos produce high-quality state-based trajectories but are only available for supported tasks
- Teleoperation demos capture human expertise but may be noisier
- The JSON metadata file records the environment ID, control mode, and simulation backend used during recording
Step 2: Trajectory Replay and Conversion
Replay recorded trajectories to convert them to the target observation mode and control mode needed for training. The replay script re-executes the recorded actions in the environment, recording observations in the desired format. This step also supports converting between control modes (e.g., from pd_joint_pos to pd_ee_delta_pose).
Key considerations:
- Replay converts the observation mode (e.g., from state to rgbd) by re-rendering the environment
- Control mode conversion translates actions between different robot control interfaces
- Use the save_traj flag to persist the converted trajectories
- Parallel processing via multiprocessing can speed up large dataset conversions
Step 3: Dataset Preparation
Load the converted HDF5 trajectory data into a PyTorch Dataset suitable for training. The dataset class reads episode data including observations, actions, and metadata, and provides minibatch sampling for the training loop. Optional preprocessing includes state normalization and observation filtering.
Key considerations:
- The ManiSkill trajectory Dataset class supports lazy loading from HDF5 for memory efficiency
- State normalization (zero mean, unit variance) can improve training stability
- Configure the number of demonstrations to use via the num_demos parameter
- DataLoader workers can be configured for parallel data loading
Step 4: Policy Network Training
Train the imitation learning policy using the prepared dataset. For behavior cloning, a neural network is trained to predict actions from observations using supervised learning with MSE loss. For diffusion policy, a denoising diffusion model is trained to generate action sequences conditioned on observation history.
Key considerations:
- Behavior cloning uses a simple MLP architecture with MSE loss for state-based training
- Diffusion policy uses a conditional UNet architecture with DDPM scheduling for action generation
- Learning rate scheduling and EMA (exponential moving average) of model weights improve final performance
- Monitor training loss and periodically evaluate on the actual environment
Step 5: Policy Evaluation
Evaluate the trained policy by deploying it in the ManiSkill environment and measuring task success rate. Evaluation runs complete episodes with the trained policy selecting actions deterministically. GPU-parallelized evaluation environments enable rapid assessment across many episodes.
Key considerations:
- Use deterministic action selection (no noise/sampling) during evaluation
- Evaluate over a sufficient number of episodes (e.g., 100) for reliable metrics
- Track success rate as the primary metric, along with episode return and length
- Compare against the demonstration success rate as an upper bound
Step 6: Optional LeRobot Format Export
Optionally convert the ManiSkill trajectory data to the LeRobot dataset format for compatibility with the HuggingFace ecosystem. This enables using ManiSkill demonstrations with LeRobot's training infrastructure and model zoo.
Key considerations:
- The converter handles both state and visual (RGBD) trajectories
- Video encoding is applied to image observations for storage efficiency
- Metadata including task description and robot type is preserved
- The output follows the LeRobot v3.0 dataset specification