# @package _global_

shape_meta:
  # acceptable types: rgb, low_dim
  obs:
    image:
      shape: [3, 96, 96]
      type: rgb
    agent_pos:
      shape: [2]
      type: low_dim
  action:
    shape: [2]

seed: 100000
horizon: 16
n_obs_steps: 2
n_action_steps: 8
n_latency_steps: 0
dataset_obs_steps: ${n_obs_steps}
past_action_visible: False
keypoint_visible_rate: 1.0
obs_as_global_cond: True

eval_episodes: 50
eval_freq: 5000
save_freq: 5000
log_freq: 250

offline_steps: 200000
online_steps: 0

offline_prioritized_sampler: true

policy:
  name: diffusion

  shape_meta: ${shape_meta}

  horizon: ${horizon}
  # n_action_steps: ${eval:'${n_action_steps}+${n_latency_steps}'}
  n_obs_steps: ${n_obs_steps}
  num_inference_steps: 100
  obs_as_global_cond: ${obs_as_global_cond}
  # crop_shape: null
  diffusion_step_embed_dim: 128
  down_dims: [512, 1024, 2048]
  kernel_size: 5
  n_groups: 8
  cond_predict_scale: True

  pretrained_model_path:

  batch_size: 64

  per_alpha: 0.6
  per_beta: 0.4

  balanced_sampling: false
  utd: 1
  offline_steps: ${offline_steps}
  use_ema: true
  lr_scheduler: cosine
  lr_warmup_steps: 500
  grad_clip_norm: 10

noise_scheduler:
  _target_: diffusers.schedulers.scheduling_ddpm.DDPMScheduler
  num_train_timesteps: 100
  beta_start: 0.0001
  beta_end: 0.02
  beta_schedule: squaredcos_cap_v2
  variance_type: fixed_small # Yilun's paper uses fixed_small_log instead, but easy to cause Nan
  clip_sample: True # required when predict_epsilon=False
  prediction_type: epsilon # or sample

obs_encoder:
  shape_meta: ${shape_meta}
  # resize_shape: null
  crop_shape: [84, 84]
  # constant center crop
  random_crop: True
  use_group_norm: True
  share_rgb_model: False
  norm_mean_std: [0.5, 0.5]  # for PushT the original impl normalizes to [-1, 1] (maybe not the case for robomimic envs)

rgb_model:
  pretrained: false
  num_keypoints: 32
  relu: true

ema:
  _target_: lerobot.common.policies.diffusion.model.ema_model.EMAModel
  update_after_step: 0
  inv_gamma: 1.0
  power: 0.75
  min_value: 0.0
  max_value: 0.9999

optimizer:
  _target_: torch.optim.AdamW
  lr: 1.0e-4
  betas: [0.95, 0.999]
  eps: 1.0e-8
  weight_decay: 1.0e-6