# @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
dataset_obs_steps: ${n_obs_steps}
past_action_visible: False
keypoint_visible_rate: 1.0

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_obs_steps: ${n_obs_steps}
  n_action_steps: ${n_action_steps}
  num_inference_steps: 100
  # crop_shape: null
  diffusion_step_embed_dim: 128
  down_dims: [512, 1024, 2048]
  kernel_size: 5
  n_groups: 8
  film_scale_modulation: 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

  delta_timestamps:
    observation.image: [-0.1, 0]
    observation.state: [-0.1, 0]
    action: [-0.1, 0, .1, .2, .3, .4, .5, .6, .7, .8, .9, 1.0, 1.1, 1.2, 1.3, 1.4]

  rgb_encoder:
    backbone_name: resnet18
    pretrained_backbone: false
    use_group_norm: True
    num_keypoints: 32
    relu: true
    norm_mean_std: [0.5, 0.5]  # for PushT the original impl normalizes to [-1, 1] (maybe not the case for robomimic envs)
    crop_shape: [84, 84]
    random_crop: True

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

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