defaults:
  - _self_
  - task: real_pusht_image

name: train_robomimic_image
_target_: diffusion_policy.workspace.train_robomimic_image_workspace.TrainRobomimicImageWorkspace

task_name: ${task.name}
shape_meta: ${task.shape_meta}
exp_name: "default"

horizon: &horizon 10
n_obs_steps: 1
n_action_steps: 1
n_latency_steps: 1
dataset_obs_steps: *horizon
past_action_visible: False
keypoint_visible_rate: 1.0

policy:
  _target_: diffusion_policy.policy.robomimic_image_policy.RobomimicImagePolicy
  shape_meta: ${shape_meta}
  algo_name: bc_rnn
  obs_type: image
  # oc.select resolver: key, default
  task_name: ${oc.select:task.task_name,tool_hang}
  dataset_type: ${oc.select:task.dataset_type,ph}
  crop_shape: [216, 288] # ch, cw 320x240 90%

dataloader:
  batch_size: 32
  num_workers: 8
  shuffle: True
  pin_memory: True
  persistent_workers: True

val_dataloader:
  batch_size: 32
  num_workers: 1
  shuffle: False
  pin_memory: True
  persistent_workers: False

training:
  device: "cuda:0"
  seed: 42
  debug: False
  resume: True
  # optimization
  num_epochs: 1000
  # training loop control
  # in epochs
  rollout_every: 50
  checkpoint_every: 50
  val_every: 1
  sample_every: 5
  # steps per epoch
  max_train_steps: null
  max_val_steps: null
  # misc
  tqdm_interval_sec: 1.0

logging:
  project: diffusion_policy_debug
  resume: True
  mode: online
  name: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name}
  tags: ["${name}", "${task_name}", "${exp_name}"]
  id: null
  group: null

checkpoint:
  topk:
    monitor_key: train_loss
    mode: min
    k: 5
    format_str: 'epoch={epoch:04d}-train_loss={train_loss:.3f}.ckpt'
  save_last_ckpt: True
  save_last_snapshot: False

multi_run:
  run_dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}
  wandb_name_base: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name}

hydra:
  job:
    override_dirname: ${name}
  run:
    dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}
  sweep:
    dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}
    subdir: ${hydra.job.num}
