# @package _global_

policy:
  name: diffusion
  trainer: policy

  action_horizon: 8
  pred_horizon: 16
  obs_horizon: 1

  observer:
    name: resnet18
    tokenize: false

  unet:
    down_dims: [256, 512, 1024]
    kernel_size: 5
    n_groups: 8
    diffusion_step_embed_dim: 128
    use_film_scale_modulation: true

  noise_scheduler:
    type: "DDPM" # DDPM, DDIM
    num_train_timesteps: 100
    beta_schedule: "squaredcos_cap_v2"
    beta_start: 0.0001
    beta_end: 0.02
    prediction_type: "epsilon" # "epsilon" or "sample"
    clip_sample: true
    clip_sample_range: 1.0

  inference:
    # Defaults to num_train_timesteps if not set.
    # A value of null will use the training timesteps.
    num_steps: null

  mask_loss_for_padding: false

  optimizer:
    lr: 1e-4
    betas: [0.95, 0.999]
    eps: 1e-8
    weight_decay: 1e-6
    backbone_lr_scale: 0.1

  scheduler:
    name: "cosine"
    warmup_steps: 500

wandb:
  tags:
    - ${policy.name}
    - ${policy.trainer}

val:
  val_online_freq: 2000

train:
  steps: 200000
