ConditionalSample
InferencePath
yes
no
TrainingPath
epsilon
sample
compute_loss
normalize obs and action
obs_as_global_cond
encode first n_obs_steps -> global_cond
encode all steps; concat action and obs
condition_mask from mask_generator
sample noise and timesteps
add noise to trajectory
apply conditioning on noisy trajectory
predict with model
prediction_type
target = noise
target = trajectory
masked MSE loss
reduce mean to scalar loss
Init: DiffusionUnetImagePolicy.__init__
Parse shape_meta: action_dim, obs_feature_dim
Build ConditionalUnet1D
Create LowdimMaskGenerator
Create LinearNormalizer
Store horizon, n_obs_steps, n_action_steps
predict_action
normalize obs
encode first To obs -> global_cond
cond_data: zeros B,T,Da; cond_mask: false
encode obs -> features B,To,Do
cond_data: B,T,Da+Do; fill obs and mask
conditional_sample
take nsample first Da dims
unnormalize action
slice action window
return action and action_pred
trajectory <- random normal
set timesteps
for each t
apply condition mask
model forward
scheduler step to prev_sample
final condition enforce
return trajectory
