Diffusion Policy Model Training

Diffusion Policy models action sequences as a conditional denoising process, which suits control tasks that are continuous and smooth and that have multiple feasible trajectories. LeRobot provides a PyTorch implementation, and Embodiflow registers this policy under the model identifier diffusion, with LeRobot as the training framework and LeRobot v2 or v3 datasets as input.

Roles and Prerequisites

Roles

Role	Scenario	Dependency
Algorithm engineer	Fine-tune the diffusion policy for continuous-trajectory tasks	Training and dataset module permissions
Project Manager	Track the progress and artifacts of diffusion policy training tasks	Training task view permission
Data provider	Provide datasets with consistent fields and frame rate	Data export module permissions

Prerequisites

Item	Requirement
Dataset version	LeRobot v2 or v3
Training framework	LeRobot
Dataset fields	The image, state, and action fields in meta/info.json are complete, and the field names match the training configuration
Temporal consistency	fps matches the actual recording to avoid misaligned supervision signals
Compute resources	A GPU; VRAM requirements vary with batch_size, image resolution, and the number of observation steps

Training Task Creation and Parameters

Procedure

1. Create a training task on the training page and select diffusion as the model.
2. Choose a data source and the dataset version v2 or v3.
3. Select LeRobot as the training framework.
4. Configure the common parameters and the policy parameters.
5. Choose the compute resources and submit the task.

Training Parameters

Common parameters:

Parameter	Description	Default	Range
batch_size	Batch size	1	≥ 1
steps	Number of training steps	10000	≥ 1
seed	Random seed	1000	Integer
num_workers	Number of DataLoader workers	4	≥ 1
eval_freq	Evaluation frequency, in steps	1000	≥ 1
log_freq	Log frequency, in steps	100	≥ 1
save_checkpoint	Whether to save checkpoints	true	Boolean
save_freq	Checkpoint save frequency, in steps	5000	≥ 1

Policy parameters (prefixed with policy.):

Parameter	Description	Default	Range
policy.n_obs_steps	Number of observation steps fed to the policy	2	Positive integer
policy.horizon	Action prediction horizon	16	Positive integer
policy.n_action_steps	Number of action steps executed per call	8	Positive integer
policy.vision_backbone	Vision backbone	resnet18	resnet18, resnet34, resnet50, resnet101, resnet152
policy.pretrained_backbone_weights	Pretrained weights for the backbone	Empty	A path or empty
policy.crop_shape	Image crop size before the backbone	[84, 84]	2-tuple
policy.crop_is_random	Whether to crop randomly during training	true	Boolean
policy.use_group_norm	Replace BatchNorm with GroupNorm in the backbone	true	Boolean
policy.spatial_softmax_num_keypoints	Number of SpatialSoftmax keypoints	32	Positive integer
policy.use_separate_rgb_encoder_per_camera	Whether each camera uses a separate RGB encoder	false	Boolean
policy.down_dims	Feature dimensions at each Unet downsampling stage	[512, 1024, 2048]	3-tuple
policy.kernel_size	Unet convolution kernel size	5	Positive integer
policy.n_groups	Number of GroupNorm groups in Unet convolution blocks	8	Positive integer
policy.diffusion_step_embed_dim	Dimension of the diffusion timestep embedding	128	Positive integer
policy.use_film_scale_modulation	Use FiLM for Unet conditioning	true	Boolean
policy.noise_scheduler_type	Noise scheduler type	DDPM	DDPM, DDIM
policy.num_train_timesteps	Number of forward diffusion steps	100	Positive integer
policy.beta_schedule	Beta schedule	squaredcos_cap_v2	String
policy.beta_start	Beta at the first forward diffusion step	0.0001	Float
policy.beta_end	Beta at the last forward diffusion step	0.02	Float
policy.prediction_type	Unet prediction target	epsilon	epsilon, sample
policy.clip_sample	Whether to clip predicted samples	true	Boolean
policy.clip_sample_range	Magnitude of the clipping range	1.0	Float
policy.num_inference_steps	Number of reverse diffusion steps at inference	Empty	Positive integer or empty
policy.do_mask_loss_for_padding	Mask the loss for zero-padded actions	false	Boolean
policy.drop_n_last_frames	Number of trailing frames to drop	7	Integer
policy.optimizer_lr	Optimizer learning rate	1e-4	Float
policy.optimizer_betas	Adam optimizer betas	[0.95, 0.999]	2-tuple
policy.optimizer_eps	Adam optimizer epsilon	1e-8	Float
policy.optimizer_weight_decay	Weight decay	1e-6	Float
policy.scheduler_name	Learning rate scheduler	cosine	String
policy.scheduler_warmup_steps	Number of scheduler warmup steps	500	Non-negative integer

To validate the data pipeline for the first time, set batch_size and steps to small values, confirm that fields are read and the loss decreases, and then scale up to full training.

Verification

Training Monitoring

The platform parses the following metrics from LeRobot training logs:

Metric	Meaning
loss	Training loss
step	Training step
sample	Number of samples used
episode	Number of episodes used
epoch	Training epoch
gradient_norm	Gradient norm
learning_rate	Learning rate
update_time	Update time per step
data_time	Data loading time per step

Verification Checklist

Check	Pass criteria
Checkpoint list	Checkpoints appear on the training details page according to save_freq
Loss curve	loss decreases over training
Logs	The training logs show no persistent errors
Real-robot or simulation evaluation	Execution results are stable across multiple rollouts under the same initial conditions

Error Handling

Symptom	Possible cause	Action	Owner
Field key mismatch	The image, state, and action field names in meta/info.json do not match the training configuration	Check and align the field names and fps	Data provider
Out of VRAM	batch_size is too large, or the image resolution and observation steps are too high	Lower batch_size and first validate with fewer steps	Training configuration owner
Loss does not decrease	The learning rate or data window is configured incorrectly	Check policy.optimizer_lr, policy.horizon, and policy.n_obs_steps	Algorithm engineer
Results differ from the original paper	The LeRobot implementation, data reading, and preprocessing differ	Use the LeRobot implementation as the reference, or switch to the original implementation	Algorithm engineer

Related Pages

Page	Purpose
Model training	Create, monitor, and quota training tasks
LeRobot datasets and training	Data export and training overview
LeRobot v2 and v3 format differences	Differences and migration between the two data formats
Data export	Generate a LeRobot training package