# @package _global_

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

  pretrained_model_path:

  # Environment.
  # Inherit these from the environment config.
  state_dim: ???
  action_dim: ???
  image_size:
    - ${env.image_size}  # height
    - ${env.image_size}  # width

  # Inputs / output structure.
  n_obs_steps: ${n_obs_steps}
  horizon: ${horizon}
  n_action_steps: ${n_action_steps}

  # Vision preprocessing.
  image_normalization_mean: [0.5, 0.5, 0.5]
  image_normalization_std: [0.5, 0.5, 0.5]

  # Architecture / modeling.
  # Vision backbone.
  vision_backbone: resnet18
  crop_shape: [84, 84]
  crop_is_random: True
  use_pretrained_backbone: false
  use_group_norm: True
  spatial_softmax_num_keypoints: 32
  # Unet.
  down_dims: [512, 1024, 2048]
  kernel_size: 5
  n_groups: 8
  diffusion_step_embed_dim: 128
  use_film_scale_modulation: True
  # Noise scheduler.
  num_train_timesteps: 100
  beta_schedule: squaredcos_cap_v2
  beta_start: 0.0001
  beta_end: 0.02
  prediction_type: epsilon # epsilon / sample
  clip_sample: True
  clip_sample_range: 1.0

  # Inference
  num_inference_steps: 100

  # ---
  # TODO(alexander-soare): Remove these from the policy config.
  batch_size: 64
  grad_clip_norm: 10
  lr: 1.0e-4
  lr_scheduler: cosine
  lr_warmup_steps: 500
  adam_betas: [0.95, 0.999]
  adam_eps: 1.0e-8
  adam_weight_decay: 1.0e-6
  utd: 1
  use_ema: true
  ema_update_after_step: 0
  ema_min_alpha: 0.0
  ema_max_alpha: 0.9999
  ema_inv_gamma: 1.0
  ema_power: 0.75

  delta_timestamps:
    observation.image: "[i / ${fps} for i in range(1 - ${n_obs_steps}, 1)]"
    observation.state: "[i / ${fps} for i in range(1 - ${n_obs_steps}, 1)]"
    action: "[i / ${fps} for i in range(1 - ${n_obs_steps}, 1 - ${n_obs_steps} + ${policy.horizon})]"