Visuomotor Imitation Learning for Robotic Pick-and-Place

This repository provides a compact training setup for visuomotor imitation learning on a robotic pick-and-place task in robosuite using robomimic.

Included training configurations cover:

• Transformer-based Behavior Cloning
• Diffusion Policy with a 1D U-Net denoising model

Large runtime artifacts such as datasets, checkpoints, generated logs, and local virtual environments are intentionally excluded from version control. Curated rollout recordings are included under media/.

Repository Layout

.
|-- code/
|   `-- training_configs/
|       |-- bc.json
|       `-- diffusion_policy.json
|-- media/
|   |-- behavior_cloning_rollout.mp4
|   `-- diffusion_policy_rollout.mp4
|-- pyproject.toml
`-- render_obs.sh

Environment

This repository uses uv and pyproject.toml for Python environment management.

uv sync

For an interactive shell:

source .venv/bin/activate

If you need a CUDA-specific PyTorch build, install the matching wheels inside this uv environment before training.

mujoco is pinned to the 3.2.x series because robosuite 1.5.1 calls mj_fullM(model, dst, qM). Newer MuJoCo releases expose a different Python binding signature and can crash during environment reset with an mj_fullM() argument error. If that happens, resync the environment:

uv sync

Data Collection

Demonstrations are collected in robosuite with keyboard teleoperation on the Pick-and-Place Can task using a Panda robot. Only successful rollouts are kept for training.

Example collection flow:

uv run python -m robosuite.scripts.collect_human_demonstrations \
  --environment PickPlaceCan \
  --robots Panda \
  --device keyboard \
  --renderer mujoco \
  --directory code/datasets/raw_demos

The raw demonstrations are then prepared for robomimic training by rendering image observations from simulator states. Place the raw dataset here:

code/datasets/raw_demos/demo_v15.hdf5

Then run:

./render_obs.sh demo_v15

This creates the rendered observation dataset expected by the final configs:

code/datasets/demo_v15_obs.hdf5

The rendered dataset uses two RGB camera streams:

• agentview_image
• robot0_eye_in_hand_image

The configs also consume robot proprioception, object state, actions, rewards, and done flags.

Training Results

The preserved run notes contain a small set of documented training statistics.

Documented dataset statistics:

Item	Value
Task	Pick-and-Place Can
Robot	Panda
Control interface	Keyboard teleoperation
Successful demonstrations collected	200
Camera views	agentview, robot0_eye_in_hand

Rollout Recordings

The repository includes two curated rollout recordings. The earlier captured recording corresponds to Diffusion Policy, and the later captured recording corresponds to Behavior Cloning.

Diffusion Policy

Behavior Cloning

Training

Run all commands from the repository root after uv sync.

Behavior Cloning:

uv run python -m robomimic.scripts.train \
  --config code/training_configs/bc.json

Diffusion Policy:

uv run python -m robomimic.scripts.train \
  --config code/training_configs/diffusion_policy.json

Training outputs are written under:

code/checkpoints/

Configuration Summary

Both final configs train on:

• dataset: code/datasets/demo_v15_obs.hdf5
• split keys: train and valid
• RGB observations: agentview_image, robot0_eye_in_hand_image
• low-dimensional observations: end-effector pose, gripper state, joint state, joint velocities, and object state
• optimizer: AdamW
• initial learning rate: 5e-5
• scheduler: cosine with 500 warmup steps
• batch size: 32
• epochs: 2000
• rollout evaluation: 10 rollouts, 500-step horizon, terminate on success

Behavior Cloning Config

code/training_configs/bc.json is a supervised action-regression policy. It learns a direct mapping from recent observations to expert actions.

Setting	Value
Algorithm	bc
Experiment name	bc_can_adamw_cosine_transformer
Policy backbone	Transformer
Sequence length	16
Frame stack	16
Transformer context length	16
Transformer embedding size	256
Transformer layers	4
Attention heads	4
Dropout	0.1
Activation	GELU
Action loss	0.9 L2 + 0.1 L1

The BC model is the simpler and faster baseline. The visual encoder processes both camera streams, the low-dimensional robot/object state is concatenated with visual features, and the transformer attends over the recent observation window before predicting the next action. The high L2 weight encourages smooth action regression, while the small L1 term helps preserve precision around grasping and placement motions.

Diffusion Policy Config

code/training_configs/diffusion_policy.json models action generation as conditional denoising. Instead of predicting one action directly, it predicts noise over an action sequence and iteratively denoises that sequence into executable controls.

Setting	Value
Algorithm	diffusion_policy
Experiment name	diffusion_policy_pickplacecan
Observation horizon	2
Action horizon	8
Prediction horizon	16
Frame stack	2
Noise predictor	1D U-Net
Diffusion-step embedding	256
U-Net channel widths	256, 512, 1024
U-Net kernel size	5
GroupNorm groups	8
EMA	enabled, power 0.75
DDPM train timesteps	100
DDPM inference timesteps	100
Noise schedule	squaredcos_cap_v2
Prediction type	epsilon

The diffusion model is more expressive than direct BC because it represents a distribution over short action trajectories. This is useful for manipulation tasks where multiple valid motions can exist from similar observations. The U-Net denoiser conditions on the current visual and state context, predicts the noise component in the action sequence, and uses the DDPM scheduler to recover a smooth action plan.