Imitation Learning Expert

Behavioral Cloning (BC) vs. Diffusion Policy on a keyboard-teleoperated Stack task with a UR5e robot in robosuite, trained with robomimic.

Setup

Linux + Python 3.10, CPU-only.

python3 -m venv venv
source venv/bin/activate
pip install --upgrade pip setuptools wheel
pip install torch==2.0.0 torchvision==0.15.0 torchaudio==2.0.0 \
    --index-url https://download.pytorch.org/whl/cpu

git clone https://github.com/ARISE-Initiative/robosuite.git
git clone https://github.com/ARISE-Initiative/robomimic.git
pip install -e ./robosuite
pip install -e ./robomimic

python robosuite/robosuite/scripts/setup_macros.py

Verify:

python -c "import torch, robomimic, robosuite; print(torch.__version__, robomimic.__version__, robosuite.__version__)"

1. Collect demonstrations

python robosuite/robosuite/scripts/collect_human_demonstrations.py \
    --environment Stack \
    --robots UR5e \
    --device keyboard \
    --renderer mujoco \
    --directory datasets/stack_ur5e

Keyboard controls (arrow keys = x/y, . ; = z, e r roll, y h pitch, o p yaw, space gripper toggle, q end episode / save if successful).

Only successful episodes are saved to demo.hdf5.

2. Convert to robomimic format + split

python robomimic/robomimic/scripts/conversion/convert_robosuite.py \
    --dataset datasets/stack_ur5e/<timestamp>/demo.hdf5

python robomimic/robomimic/scripts/split_train_val.py \
    --dataset datasets/stack_ur5e/<timestamp>/demo.hdf5 \
    --ratio 0.1

3. Render observations

export MUJOCO_GL=egl
python robomimic/robomimic/scripts/dataset_states_to_obs.py \
    --dataset datasets/stack_ur5e/<timestamp>/demo.hdf5 \
    --output_name image.hdf5 \
    --done_mode 2 \
    --camera_names agentview robot0_eye_in_hand \
    --camera_height 84 --camera_width 84 \
    --compress

4. Train

python robomimic/robomimic/scripts/train.py \
    --config training_configs/bc.json \
    --dataset datasets/stack_ur5e/<timestamp>/image.hdf5 \
    --name bc_stack_ur5e

python robomimic/robomimic/scripts/train.py \
    --config training_configs/diffusion_policy.json \
    --dataset datasets/stack_ur5e/<timestamp>/image.hdf5 \
    --name diffusion_stack_ur5e

Layout

training_configs/     final configs for BC and Diffusion Policy
scripts/              custom helper scripts (dataset merge, ablation runners)
checkpoints/          final model checkpoints
videos/               rollout videos
tensorboard/          exported tensorboard event files
logs/                 training logs
report/               LaTeX report sources