feat: align pmf transformer training and config defaults

2026-03-16 15:37:32 +08:00
parent 79f31940c4
commit 42dc29a2cb
2 changed files with 80 additions and 20 deletions
--- a/image_pusht_diffusion_policy_dit_pmf.yaml
+++ b/image_pusht_diffusion_policy_dit_pmf.yaml
@@ -58,15 +58,17 @@ policy:
  - 84
  eval_fixed_crop: true
  horizon: 16
-  min_time: 0.05
  n_action_steps: 8
  n_cond_layers: 0
  n_emb: 256
  n_head: 4
-  n_layer: 8
+  n_layer: 12
  n_obs_steps: 2
  n_time_tokens: 4
  noise_scale: 1.0
+  adatloss_eps: 0.01
+  p_mean: -0.4
+  p_std: 1.0
  noise_scheduler:
    _target_: diffusers.schedulers.scheduling_ddpm.DDPMScheduler
    beta_end: 0.02
@@ -76,13 +78,16 @@ policy:
    num_train_timesteps: 100
    prediction_type: sample
    variance_type: fixed_small
-  num_inference_steps: 32
+  num_inference_steps: 1
  obs_as_cond: true
  obs_encoder_group_norm: true
  p_drop_attn: 0.0
  p_drop_emb: 0.0
  pmf_u_loss_weight: 1.0
  pmf_v_loss_weight: 1.0
+  tr_uniform: true
+  tr_uniform_prob: 0.1
+  data_proportion: 0.5
  shape_meta:
    action:
      shape: