name: train_act_image
_target_: memory_diffusion_policy.workspace.train_act_image_workspace.TrainACTImageWorkspace
task_name: ${task.name}
shape_meta: ${task.shape_meta}
exp_name: default
horizon: 100
n_obs_steps: 1
n_action_steps: 16
n_latency_steps: 0
dataset_obs_steps: ${n_obs_steps}
past_action_visible: false
keypoint_visible_rate: 1.0
rollout_mode: wandb
rollout_n_envs: null
policy:
  _target_: memory_diffusion_policy.policy.act_image_policy.ACTImagePolicy
  shape_meta: ${shape_meta}
  chunk_size: ${horizon}
  n_action_steps: ${n_action_steps}
  n_obs_steps: ${n_obs_steps}
  kl_weight: 10.0
  hidden_dim: 512
  dim_feedforward: 3200
  nheads: 8
  enc_layers: 4
  dec_layers: 7
  dropout: 0.1
  pre_norm: false
  position_embedding: sine
  backbone: resnet18
  backbone_pretrained: true
  lr_backbone: 1.0e-05
  temporal_agg: false
  temporal_agg_k: 0.01
dataloader:
  batch_size: 8
  num_workers: 8
  shuffle: true
  pin_memory: true
  persistent_workers: false
val_dataloader:
  batch_size: 8
  num_workers: 8
  shuffle: false
  pin_memory: true
  persistent_workers: false
optimizer:
  _target_: torch.optim.AdamW
  lr: 1.0e-05
  betas:
  - 0.9
  - 0.999
  eps: 1.0e-08
  weight_decay: 0.0001
training:
  device: cuda:0
  seed: 42
  debug: false
  resume: true
  lr_scheduler: cosine
  lr_warmup_steps: 500
  num_epochs: 2000
  gradient_accumulate_every: 1
  use_ema: false
  rollout_every: 0
  checkpoint_every: 25
  val_every: 0
  sample_every: 0
  max_train_steps: null
  max_val_steps: null
  tqdm_interval_sec: 1.0
logging:
  project: act_baseline
  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: test_mean_score
    mode: max
    k: 0
    format_str: epoch={epoch:04d}-test_mean_score={test_mean_score:.3f}.ckpt
  save_last_ckpt: true
  save_last_snapshot: false
  save_every_n_epochs: ${training.checkpoint_every}
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}
task:
  name: insert_peg_image_abs_jinglin
  rotation_rep: rotation_6d
  default_orientation_axisangle:
  - 3.14159265
  - 0.0
  - 0.0
  image_shape:
  - 3
  - 96
  - 96
  shape_meta:
    obs:
      peg_focus_view_image:
        shape:
        - 3
        - 96
        - 96
        type: rgb
      robot0_eye_in_hand_image:
        shape:
        - 3
        - 96
        - 96
        type: rgb
      robot0_eef_pos:
        shape:
        - 3
        type: low_dim
      robot0_eef_quat:
        shape:
        - 4
        type: low_dim
      robot0_gripper_qpos:
        shape:
        - 2
        type: low_dim
    action:
      shape:
      - 10
  task_name: insert_peg
  dataset_path: data/image_insertpeg_jinglin.hdf5
  abs_action: true
  env_runner:
    _target_: memory_diffusion_policy.env_runner.robosuite_pomdp_image_runner.RobosuitePomdpImageRunner
    env_name: InsertPegVision
    dataset_path: data/image_insertpeg_jinglin.hdf5
    shape_meta:
      obs:
        peg_focus_view_image:
          shape:
          - 3
          - 96
          - 96
          type: rgb
        robot0_eye_in_hand_image:
          shape:
          - 3
          - 96
          - 96
          type: rgb
        robot0_eef_pos:
          shape:
          - 3
          type: low_dim
        robot0_eef_quat:
          shape:
          - 4
          type: low_dim
        robot0_gripper_qpos:
          shape:
          - 2
          type: low_dim
      action:
        shape:
        - 10
    n_train: 4
    n_train_vis: 2
    train_start_seed: 0
    n_test: 20
    n_test_vis: 4
    test_start_seed: 100000
    max_steps: 2500
    n_obs_steps: ${n_obs_steps}
    n_action_steps: ${n_action_steps}
    render_obs_key: peg_focus_view_image
    abs_action: true
    rotation_rep: ${task.rotation_rep}
    default_orientation_axisangle: ${task.default_orientation_axisangle}
    fps: 10
    past_action: ${past_action_visible}
    n_envs: ${oc.select:rollout_n_envs,null}
    rollout_mode: ${oc.select:rollout_mode,wandb}
  dataset:
    _target_: memory_diffusion_policy.dataset.robomimic_3d_replay_image_dataset.Robomimic3DReplayImageDataset
    shape_meta:
      obs:
        peg_focus_view_image:
          shape:
          - 3
          - 96
          - 96
          type: rgb
        robot0_eye_in_hand_image:
          shape:
          - 3
          - 96
          - 96
          type: rgb
        robot0_eef_pos:
          shape:
          - 3
          type: low_dim
        robot0_eef_quat:
          shape:
          - 4
          type: low_dim
        robot0_gripper_qpos:
          shape:
          - 2
          type: low_dim
      action:
        shape:
        - 10
    dataset_path: data/image_insertpeg_jinglin.hdf5
    horizon: ${horizon}
    pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'}
    pad_after: ${eval:'${n_action_steps}-1'}
    n_obs_steps: ${dataset_obs_steps}
    abs_action: true
    rotation_rep: ${task.rotation_rep}
    action_key: actions_abs
    use_legacy_normalizer: false
    use_cache: true
    max_demos: null
    seed: 42
    val_ratio: 0.05