defaults:
  - env_vars: env_vars
  # - encoder: pusht_dynamo
  # - encoder: dino_cls
  - encoder: dino_patch # dino_patch_avg_pool
  - _self_
  - cluster: local

hydra:
  run:
    dir: ${env_vars.save_path}/outputs/${env.gym.id}/${now:%Y-%m-%d}/${now:%H-%M-%S}
  job:
    chdir: True
  sweep:
    dir: ${env_vars.save_path}/outputs/${env.gym.id}/${now:%Y-%m-%d}/${now:%H-%M-%S}
    subdir: ${hydra.job.num}
seed: 42
window_size: 2
goal_window_size: 2
eval_window_size: 2
num_workers: 23

batch_size: 256 
epochs: 1200 
eval_freq: 10
eval_on_env_freq: 10
num_env_evals: 100
num_envs: 10
num_final_evals: 100
action_window_size: 5

# pred_horizon: ${action_window_size} # Total seq length to predict to roll out
pred_horizon: ${eval:'${action_window_size} + ${window_size} - 1'}
n_action_steps: ${action_window_size} # Action steps to roll out

precompute_embeddings: False

goal_dim: 0
wandb:
  project: patch_policy
  entity: ${env_vars.wandb_entity}
train_fraction: 0.9

device: cuda
optim:
  lr: 1e-4
  weight_decay: 0.0
  betas: [0.9, 0.999]

env:
  gym:
    _target_: envs.pusht.PushWrapper
    id: pusht
    env:
      _target_: envs.pusht.PushTKeypointsEnv
  obs_dim: 20
  act_dim: 2
  goal_dim: ${goal_dim}
  views: 1

dataset:
  _target_: datasets.pusht.PushTDataset
  data_directory: ${env_vars.datasets.pusht}

data:
  window_size: ${window_size}
  action_window_size: ${action_window_size}
  vqbet_get_future_action_chunk: false
  goal_conditional: future
  future_seq_len: ${goal_window_size}
  min_future_sep: ${action_window_size}

save_ckpt: true
save_path: "${env_vars.save_path}/outputs/${env.gym.id}/${now:%Y-%m-%d}/${now:%H-%M-%S}"
load_path: null

model:
  _target_: models.diffusion_policy.diffusion_policy.DiffusionPolicy
  obs_dim: ${encoder.output_dim}
  act_dim: ${env.act_dim}
  goal_dim: ${env.goal_dim}
  views: ${env.views}
  obs_horizon: ${window_size}
  pred_horizon: ${pred_horizon}
  action_horizon: ${n_action_steps}
  n_layer: 8
  n_head: 4
  n_embd: 256
  visual_input: true
  data_act_scale: 1
  lr: ${optim.lr}
  weight_decay: ${optim.weight_decay}
  use_transform: null
  n_patches: ${encoder.n_patches}