Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+can/logs.json.txt filter=lfs diff=lfs merge=lfs -text
+square/logs.json.txt filter=lfs diff=lfs merge=lfs -text
+transport/logs.json.txt filter=lfs diff=lfs merge=lfs -text

can/.hydra/config.yaml

@@ -0,0 +1,242 @@
+name: train_flow_match_vib_unet_image
+_target_: zprl.workspace.train_flow_match_vib_unet_image_workspace.TrainFlowMatchVibUnetImageWorkspace
+task_name: ${task.name}
+shape_meta: ${task.shape_meta}
+exp_name: unet_vib_recon
+horizon: 16
+n_obs_steps: 1
+n_action_steps: 4
+n_latency_steps: 0
+dataset_obs_steps: ${n_obs_steps}
+past_action_visible: false
+keypoint_visible_rate: 1.0
+obs_as_global_cond: true
+policy:
+  _target_: zprl.policy.flow_match_vib_unet_image_policy.FlowMatchVibUnetImagePolicy
+  vib_latent_dim: 16
+  vib_alpha: 1.0
+  vib_beta: 0.0002
+  vib_recon: 0.01
+  vib_hidden_dim: 256
+  shape_meta: ${shape_meta}
+  noise_scheduler:
+    _target_: diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler
+    num_train_timesteps: 100
+  obs_encoder:
+    _target_: zprl.model.vision.multi_image_obs_encoder.MultiImageObsEncoder
+    shape_meta: ${shape_meta}
+    rgb_model:
+      _target_: zprl.model.vision.model_getter.ResNet18Pooling
+      input_shape:
+      - 3
+      - 76
+      - 76
+      weights: null
+      pooling_type: spatial_learned_embeddings
+      num_spatial_blocks: 8
+      bottleneck_dim: 256
+    resize_shape: null
+    crop_shape:
+    - 76
+    - 76
+    random_crop: true
+    use_group_norm: true
+    share_rgb_model: false
+    imagenet_norm: true
+  horizon: ${horizon}
+  n_action_steps: ${eval:'${n_action_steps}+${n_latency_steps}'}
+  n_obs_steps: ${n_obs_steps}
+  num_inference_steps: 2
+  obs_as_global_cond: ${obs_as_global_cond}
+  diffusion_step_embed_dim: 128
+  down_dims:
+  - 128
+  - 256
+  - 512
+  kernel_size: 5
+  n_groups: 8
+  cond_predict_scale: true
+ema:
+  _target_: zprl.model.diffusion.ema_model.EMAModel
+  update_after_step: 0
+  inv_gamma: 1.0
+  power: 0.8
+  min_value: 0.0
+  max_value: 0.9999
+dataloader:
+  batch_size: 256
+  num_workers: 12
+  shuffle: true
+  pin_memory: true
+  persistent_workers: false
+val_dataloader:
+  batch_size: 64
+  num_workers: 4
+  shuffle: false
+  pin_memory: true
+  persistent_workers: false
+optimizer:
+  _target_: torch.optim.AdamW
+  lr: 0.0001
+  betas:
+  - 0.95
+  - 0.999
+  eps: 1.0e-08
+  weight_decay: 1.0e-06
+training:
+  device: cuda:1
+  seed: 10
+  debug: false
+  resume: true
+  lr_scheduler: cosine
+  lr_warmup_steps: 100
+  num_epochs: 1001
+  gradient_accumulate_every: 1
+  use_ema: true
+  freeze_encoder: false
+  rollout_every: 100
+  checkpoint_every: 100
+  val_every: 25
+  sample_every: 25
+  max_train_steps: null
+  max_val_steps: null
+  tqdm_interval_sec: 1.0
+logging:
+  project: offline_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: test_mean_score
+    mode: max
+    k: 2
+    format_str: epoch_{epoch:04d}-score_{test_mean_score:.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}
+task:
+  name: can_image
+  shape_meta:
+    obs:
+      agentview_image:
+        shape:
+        - 3
+        - 84
+        - 84
+        type: rgb
+      robot0_eye_in_hand_image:
+        shape:
+        - 3
+        - 84
+        - 84
+        type: rgb
+      robot0_eef_pos:
+        shape:
+        - 3
+      robot0_eef_quat:
+        shape:
+        - 4
+      robot0_gripper_qpos:
+        shape:
+        - 2
+    action:
+      shape:
+      - 10
+  task_name: can
+  dataset_type: mh
+  dataset_path: /media/datahub-2/ydj/robomimicv030/${task.task_name}/${task.dataset_type}/image_v141_subset_abs.hdf5
+  abs_action: true
+  env_runner:
+    _target_: zprl.env_runner.robomimic_image_runner.RobomimicImageRunner
+    dataset_path: /media/datahub-2/ydj/robomimicv030/${task.task_name}/${task.dataset_type}/image_v141_subset_abs.hdf5
+    shape_meta:
+      obs:
+        agentview_image:
+          shape:
+          - 3
+          - 84
+          - 84
+          type: rgb
+        robot0_eye_in_hand_image:
+          shape:
+          - 3
+          - 84
+          - 84
+          type: rgb
+        robot0_eef_pos:
+          shape:
+          - 3
+        robot0_eef_quat:
+          shape:
+          - 4
+        robot0_gripper_qpos:
+          shape:
+          - 2
+      action:
+        shape:
+        - 10
+    n_train: 6
+    n_train_vis: 2
+    train_start_idx: 0
+    n_test: 50
+    n_test_vis: 4
+    test_start_seed: 100000
+    max_steps: ${eval:'500 if "${task.dataset_type}" == "mh" else 200'}
+    n_obs_steps: ${n_obs_steps}
+    n_action_steps: ${n_action_steps}
+    render_obs_key: agentview_image
+    fps: 10
+    crf: 22
+    past_action: ${past_action_visible}
+    abs_action: true
+    tqdm_interval_sec: 1.0
+    n_envs: 28
+  dataset:
+    _target_: zprl.dataset.robomimic_replay_image_dataset.RobomimicReplayImageDataset
+    shape_meta:
+      obs:
+        agentview_image:
+          shape:
+          - 3
+          - 84
+          - 84
+          type: rgb
+        robot0_eye_in_hand_image:
+          shape:
+          - 3
+          - 84
+          - 84
+          type: rgb
+        robot0_eef_pos:
+          shape:
+          - 3
+        robot0_eef_quat:
+          shape:
+          - 4
+        robot0_gripper_qpos:
+          shape:
+          - 2
+      action:
+        shape:
+        - 10
+    dataset_path: /media/datahub-2/ydj/robomimicv030/${task.task_name}/${task.dataset_type}/image_v141_subset_abs.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: rotation_6d
+    use_legacy_normalizer: false
+    use_cache: true
+    seed: 10
+    val_ratio: 0.02
+    num_demo: 100

can/.hydra/hydra.yaml

@@ -0,0 +1,160 @@
+hydra:
+  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}
+  launcher:
+    _target_: hydra._internal.core_plugins.basic_launcher.BasicLauncher
+  sweeper:
+    _target_: hydra._internal.core_plugins.basic_sweeper.BasicSweeper
+    max_batch_size: null
+    params: null
+  help:
+    app_name: ${hydra.job.name}
+    header: '${hydra.help.app_name} is powered by Hydra.
+
+      '
+    footer: 'Powered by Hydra (https://hydra.cc)
+
+      Use --hydra-help to view Hydra specific help
+
+      '
+    template: '${hydra.help.header}
+
+      == Configuration groups ==
+
+      Compose your configuration from those groups (group=option)
+
+
+      $APP_CONFIG_GROUPS
+
+
+      == Config ==
+
+      Override anything in the config (foo.bar=value)
+
+
+      $CONFIG
+
+
+      ${hydra.help.footer}
+
+      '
+  hydra_help:
+    template: 'Hydra (${hydra.runtime.version})
+
+      See https://hydra.cc for more info.
+
+
+      == Flags ==
+
+      $FLAGS_HELP
+
+
+      == Configuration groups ==
+
+      Compose your configuration from those groups (For example, append hydra/job_logging=disabled
+      to command line)
+
+
+      $HYDRA_CONFIG_GROUPS
+
+
+      Use ''--cfg hydra'' to Show the Hydra config.
+
+      '
+    hydra_help: ???
+  hydra_logging:
+    version: 1
+    formatters:
+      simple:
+        format: '[%(asctime)s][HYDRA] %(message)s'
+    handlers:
+      console:
+        class: logging.StreamHandler
+        formatter: simple
+        stream: ext://sys.stdout
+    root:
+      level: INFO
+      handlers:
+      - console
+    loggers:
+      logging_example:
+        level: DEBUG
+    disable_existing_loggers: false
+  job_logging:
+    version: 1
+    formatters:
+      simple:
+        format: '[%(asctime)s][%(name)s][%(levelname)s] - %(message)s'
+    handlers:
+      console:
+        class: logging.StreamHandler
+        formatter: simple
+        stream: ext://sys.stdout
+      file:
+        class: logging.FileHandler
+        formatter: simple
+        filename: ${hydra.runtime.output_dir}/${hydra.job.name}.log
+    root:
+      level: INFO
+      handlers:
+      - console
+      - file
+    disable_existing_loggers: false
+  env: {}
+  mode: RUN
+  searchpath: []
+  callbacks: {}
+  output_subdir: .hydra
+  overrides:
+    hydra:
+    - hydra.mode=RUN
+    task:
+    - training.seed=10
+    - task.dataset.seed=10
+    - exp_name=unet_vib_recon
+    - training.device=cuda:1
+    - task=can_image_abs
+  job:
+    name: train
+    chdir: null
+    override_dirname: exp_name=unet_vib_recon,task.dataset.seed=10,task=can_image_abs,training.device=cuda:1,training.seed=10
+    id: ???
+    num: ???
+    config_name: train_flow_match_vib_unet_image_workspace
+    env_set: {}
+    env_copy: []
+    config:
+      override_dirname:
+        kv_sep: '='
+        item_sep: ','
+        exclude_keys: []
+  runtime:
+    version: 1.3.2
+    version_base: '1.3'
+    cwd: /ghome/h2/djyu/zprl-sim/diffusion_policy
+    config_sources:
+    - path: hydra.conf
+      schema: pkg
+      provider: hydra
+    - path: /ghome/h2/djyu/zprl-sim/diffusion_policy/zprl/config
+      schema: file
+      provider: main
+    - path: ''
+      schema: structured
+      provider: schema
+    output_dir: /ghome/h2/djyu/zprl-sim/diffusion_policy/data/outputs/2026.03.23/21.26.02_train_flow_match_vib_unet_image_can_image
+    choices:
+      task: can_image_abs
+      hydra/env: default
+      hydra/callbacks: null
+      hydra/job_logging: default
+      hydra/hydra_logging: default
+      hydra/hydra_help: default
+      hydra/help: default
+      hydra/sweeper: basic
+      hydra/launcher: basic
+      hydra/output: default
+  verbose: false

can/.hydra/overrides.yaml

@@ -0,0 +1,5 @@
+- training.seed=10
+- task.dataset.seed=10
+- exp_name=unet_vib_recon
+- training.device=cuda:1
+- task=can_image_abs

can/checkpoints/latest.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6acf0aadc976e78d39097e2185cef05e2a29a2b0462d567558f07065ee433619
+size 722597620

can/logs.json.txt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e8d9a49dd86ea97ce9952475c36b407559b070b4237577104f40730242a44ddc
+size 30214796

can/train.log

@@ -0,0 +1,5 @@
+[2026-03-23 21:26:04,456][zprl.model.vision.multi_image_obs_encoder][INFO] - number of parameters: 24.53 M
+[2026-03-23 21:26:04,653][zprl.model.diffusion.conditional_unet1d][INFO] - number of parameters: 20.06 M
+[2026-03-23 21:26:04,655][zprl.model.vib][INFO] - number of parameters: 0.27 M
+[2026-03-23 21:26:04,657][zprl.model.vib][INFO] - number of parameters: 0.27 M
+[2026-03-23 21:26:09,116][OpenGL.acceleratesupport][INFO] - No OpenGL_accelerate module loaded: No module named 'OpenGL_accelerate'

square/.hydra/config.yaml

@@ -0,0 +1,242 @@
+name: train_flow_match_vib_unet_image
+_target_: zprl.workspace.train_flow_match_vib_unet_image_workspace.TrainFlowMatchVibUnetImageWorkspace
+task_name: ${task.name}
+shape_meta: ${task.shape_meta}
+exp_name: unet_vib_recon
+horizon: 16
+n_obs_steps: 1
+n_action_steps: 4
+n_latency_steps: 0
+dataset_obs_steps: ${n_obs_steps}
+past_action_visible: false
+keypoint_visible_rate: 1.0
+obs_as_global_cond: true
+policy:
+  _target_: zprl.policy.flow_match_vib_unet_image_policy.FlowMatchVibUnetImagePolicy
+  vib_latent_dim: 16
+  vib_alpha: 1.0
+  vib_beta: 0.0002
+  vib_recon: 0.01
+  vib_hidden_dim: 256
+  shape_meta: ${shape_meta}
+  noise_scheduler:
+    _target_: diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler
+    num_train_timesteps: 100
+  obs_encoder:
+    _target_: zprl.model.vision.multi_image_obs_encoder.MultiImageObsEncoder
+    shape_meta: ${shape_meta}
+    rgb_model:
+      _target_: zprl.model.vision.model_getter.ResNet18Pooling
+      input_shape:
+      - 3
+      - 76
+      - 76
+      weights: null
+      pooling_type: spatial_learned_embeddings
+      num_spatial_blocks: 8
+      bottleneck_dim: 256
+    resize_shape: null
+    crop_shape:
+    - 76
+    - 76
+    random_crop: true
+    use_group_norm: true
+    share_rgb_model: false
+    imagenet_norm: true
+  horizon: ${horizon}
+  n_action_steps: ${eval:'${n_action_steps}+${n_latency_steps}'}
+  n_obs_steps: ${n_obs_steps}
+  num_inference_steps: 2
+  obs_as_global_cond: ${obs_as_global_cond}
+  diffusion_step_embed_dim: 128
+  down_dims:
+  - 128
+  - 256
+  - 512
+  kernel_size: 5
+  n_groups: 8
+  cond_predict_scale: true
+ema:
+  _target_: zprl.model.diffusion.ema_model.EMAModel
+  update_after_step: 0
+  inv_gamma: 1.0
+  power: 0.8
+  min_value: 0.0
+  max_value: 0.9999
+dataloader:
+  batch_size: 256
+  num_workers: 12
+  shuffle: true
+  pin_memory: true
+  persistent_workers: false
+val_dataloader:
+  batch_size: 64
+  num_workers: 4
+  shuffle: false
+  pin_memory: true
+  persistent_workers: false
+optimizer:
+  _target_: torch.optim.AdamW
+  lr: 0.0001
+  betas:
+  - 0.95
+  - 0.999
+  eps: 1.0e-08
+  weight_decay: 1.0e-06
+training:
+  device: cuda:3
+  seed: 0
+  debug: false
+  resume: true
+  lr_scheduler: cosine
+  lr_warmup_steps: 100
+  num_epochs: 1001
+  gradient_accumulate_every: 1
+  use_ema: true
+  freeze_encoder: false
+  rollout_every: 100
+  checkpoint_every: 100
+  val_every: 25
+  sample_every: 25
+  max_train_steps: null
+  max_val_steps: null
+  tqdm_interval_sec: 1.0
+logging:
+  project: offline_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: test_mean_score
+    mode: max
+    k: 2
+    format_str: epoch_{epoch:04d}-score_{test_mean_score:.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}
+task:
+  name: square_image
+  shape_meta:
+    obs:
+      agentview_image:
+        shape:
+        - 3
+        - 84
+        - 84
+        type: rgb
+      robot0_eye_in_hand_image:
+        shape:
+        - 3
+        - 84
+        - 84
+        type: rgb
+      robot0_eef_pos:
+        shape:
+        - 3
+      robot0_eef_quat:
+        shape:
+        - 4
+      robot0_gripper_qpos:
+        shape:
+        - 2
+    action:
+      shape:
+      - 10
+  task_name: square
+  dataset_type: mh
+  dataset_path: /media/datahub-2/ydj/robomimicv030/${task.task_name}/${task.dataset_type}/image_v141_subset_abs.hdf5
+  abs_action: true
+  env_runner:
+    _target_: zprl.env_runner.robomimic_image_runner.RobomimicImageRunner
+    dataset_path: /media/datahub-2/ydj/robomimicv030/${task.task_name}/${task.dataset_type}/image_v141_subset_abs.hdf5
+    shape_meta:
+      obs:
+        agentview_image:
+          shape:
+          - 3
+          - 84
+          - 84
+          type: rgb
+        robot0_eye_in_hand_image:
+          shape:
+          - 3
+          - 84
+          - 84
+          type: rgb
+        robot0_eef_pos:
+          shape:
+          - 3
+        robot0_eef_quat:
+          shape:
+          - 4
+        robot0_gripper_qpos:
+          shape:
+          - 2
+      action:
+        shape:
+        - 10
+    n_train: 6
+    n_train_vis: 2
+    train_start_idx: 0
+    n_test: 50
+    n_test_vis: 4
+    test_start_seed: 100000
+    max_steps: ${eval:'500 if "${task.dataset_type}" == "mh" else 200'}
+    n_obs_steps: ${n_obs_steps}
+    n_action_steps: ${n_action_steps}
+    render_obs_key: agentview_image
+    fps: 10
+    crf: 22
+    past_action: ${past_action_visible}
+    abs_action: true
+    tqdm_interval_sec: 1.0
+    n_envs: 28
+  dataset:
+    _target_: zprl.dataset.robomimic_replay_image_dataset.RobomimicReplayImageDataset
+    shape_meta:
+      obs:
+        agentview_image:
+          shape:
+          - 3
+          - 84
+          - 84
+          type: rgb
+        robot0_eye_in_hand_image:
+          shape:
+          - 3
+          - 84
+          - 84
+          type: rgb
+        robot0_eef_pos:
+          shape:
+          - 3
+        robot0_eef_quat:
+          shape:
+          - 4
+        robot0_gripper_qpos:
+          shape:
+          - 2
+      action:
+        shape:
+        - 10
+    dataset_path: /media/datahub-2/ydj/robomimicv030/${task.task_name}/${task.dataset_type}/image_v141_subset_abs.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: rotation_6d
+    use_legacy_normalizer: false
+    use_cache: true
+    seed: 0
+    val_ratio: 0.02
+    num_demo: 100

square/.hydra/hydra.yaml

@@ -0,0 +1,160 @@
+hydra:
+  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}
+  launcher:
+    _target_: hydra._internal.core_plugins.basic_launcher.BasicLauncher
+  sweeper:
+    _target_: hydra._internal.core_plugins.basic_sweeper.BasicSweeper
+    max_batch_size: null
+    params: null
+  help:
+    app_name: ${hydra.job.name}
+    header: '${hydra.help.app_name} is powered by Hydra.
+
+      '
+    footer: 'Powered by Hydra (https://hydra.cc)
+
+      Use --hydra-help to view Hydra specific help
+
+      '
+    template: '${hydra.help.header}
+
+      == Configuration groups ==
+
+      Compose your configuration from those groups (group=option)
+
+
+      $APP_CONFIG_GROUPS
+
+
+      == Config ==
+
+      Override anything in the config (foo.bar=value)
+
+
+      $CONFIG
+
+
+      ${hydra.help.footer}
+
+      '
+  hydra_help:
+    template: 'Hydra (${hydra.runtime.version})
+
+      See https://hydra.cc for more info.
+
+
+      == Flags ==
+
+      $FLAGS_HELP
+
+
+      == Configuration groups ==
+
+      Compose your configuration from those groups (For example, append hydra/job_logging=disabled
+      to command line)
+
+
+      $HYDRA_CONFIG_GROUPS
+
+
+      Use ''--cfg hydra'' to Show the Hydra config.
+
+      '
+    hydra_help: ???
+  hydra_logging:
+    version: 1
+    formatters:
+      simple:
+        format: '[%(asctime)s][HYDRA] %(message)s'
+    handlers:
+      console:
+        class: logging.StreamHandler
+        formatter: simple
+        stream: ext://sys.stdout
+    root:
+      level: INFO
+      handlers:
+      - console
+    loggers:
+      logging_example:
+        level: DEBUG
+    disable_existing_loggers: false
+  job_logging:
+    version: 1
+    formatters:
+      simple:
+        format: '[%(asctime)s][%(name)s][%(levelname)s] - %(message)s'
+    handlers:
+      console:
+        class: logging.StreamHandler
+        formatter: simple
+        stream: ext://sys.stdout
+      file:
+        class: logging.FileHandler
+        formatter: simple
+        filename: ${hydra.runtime.output_dir}/${hydra.job.name}.log
+    root:
+      level: INFO
+      handlers:
+      - console
+      - file
+    disable_existing_loggers: false
+  env: {}
+  mode: RUN
+  searchpath: []
+  callbacks: {}
+  output_subdir: .hydra
+  overrides:
+    hydra:
+    - hydra.mode=RUN
+    task:
+    - training.seed=0
+    - task.dataset.seed=0
+    - exp_name=unet_vib_recon
+    - training.device=cuda:3
+    - task=square_image_abs
+  job:
+    name: train
+    chdir: null
+    override_dirname: exp_name=unet_vib_recon,task.dataset.seed=0,task=square_image_abs,training.device=cuda:3,training.seed=0
+    id: ???
+    num: ???
+    config_name: train_flow_match_vib_unet_image_workspace
+    env_set: {}
+    env_copy: []
+    config:
+      override_dirname:
+        kv_sep: '='
+        item_sep: ','
+        exclude_keys: []
+  runtime:
+    version: 1.3.2
+    version_base: '1.3'
+    cwd: /ghome/h2/djyu/zprl-sim/diffusion_policy
+    config_sources:
+    - path: hydra.conf
+      schema: pkg
+      provider: hydra
+    - path: /ghome/h2/djyu/zprl-sim/diffusion_policy/zprl/config
+      schema: file
+      provider: main
+    - path: ''
+      schema: structured
+      provider: schema
+    output_dir: /ghome/h2/djyu/zprl-sim/diffusion_policy/data/outputs/2026.03.24/00.15.33_train_flow_match_vib_unet_image_square_image
+    choices:
+      task: square_image_abs
+      hydra/env: default
+      hydra/callbacks: null
+      hydra/job_logging: default
+      hydra/hydra_logging: default
+      hydra/hydra_help: default
+      hydra/help: default
+      hydra/sweeper: basic
+      hydra/launcher: basic
+      hydra/output: default
+  verbose: false

square/.hydra/overrides.yaml

@@ -0,0 +1,5 @@
+- training.seed=0
+- task.dataset.seed=0
+- exp_name=unet_vib_recon
+- training.device=cuda:3
+- task=square_image_abs

square/checkpoints/epoch_0600-score_0.460.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:262dc109555210d81f6430ec91e3e2d2b765af831c9506ff20e431d9f55aec7a
+size 722597620

square/logs.json.txt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6ae3861009450cf4f77e95ec13e4606ab8f4e6037b770b97bbb0dd98fb5106f4
+size 37778802

square/train.log

@@ -0,0 +1,5 @@
+[2026-03-24 00:15:34,949][zprl.model.vision.multi_image_obs_encoder][INFO] - number of parameters: 24.53 M
+[2026-03-24 00:15:35,217][zprl.model.diffusion.conditional_unet1d][INFO] - number of parameters: 20.06 M
+[2026-03-24 00:15:35,220][zprl.model.vib][INFO] - number of parameters: 0.27 M
+[2026-03-24 00:15:35,222][zprl.model.vib][INFO] - number of parameters: 0.27 M
+[2026-03-24 00:15:37,805][OpenGL.acceleratesupport][INFO] - No OpenGL_accelerate module loaded: No module named 'OpenGL_accelerate'

transport/.hydra/config.yaml

@@ -0,0 +1,305 @@
+name: train_flow_match_vib_unet_image
+_target_: zprl.workspace.train_flow_match_vib_unet_image_workspace.TrainFlowMatchVibUnetImageWorkspace
+task_name: ${task.name}
+shape_meta: ${task.shape_meta}
+exp_name: unet_vib_recon
+horizon: 16
+n_obs_steps: 1
+n_action_steps: 5
+n_latency_steps: 0
+dataset_obs_steps: ${n_obs_steps}
+past_action_visible: false
+keypoint_visible_rate: 1.0
+obs_as_global_cond: true
+policy:
+  _target_: zprl.policy.flow_match_vib_unet_image_policy.FlowMatchVibUnetImagePolicy
+  vib_latent_dim: 32
+  vib_alpha: 1.0
+  vib_beta: 0.0001
+  vib_recon: 0.01
+  vib_hidden_dim: 256
+  shape_meta: ${shape_meta}
+  noise_scheduler:
+    _target_: diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler
+    num_train_timesteps: 100
+  obs_encoder:
+    _target_: zprl.model.vision.multi_image_obs_encoder.MultiImageObsEncoder
+    shape_meta: ${shape_meta}
+    rgb_model:
+      _target_: zprl.model.vision.model_getter.ResNet18Pooling
+      input_shape:
+      - 3
+      - 76
+      - 76
+      weights: null
+      pooling_type: spatial_learned_embeddings
+      num_spatial_blocks: 8
+      bottleneck_dim: 256
+    resize_shape: null
+    crop_shape:
+    - 76
+    - 76
+    random_crop: true
+    use_group_norm: true
+    share_rgb_model: false
+    imagenet_norm: true
+  horizon: ${horizon}
+  n_action_steps: ${eval:'${n_action_steps}+${n_latency_steps}'}
+  n_obs_steps: ${n_obs_steps}
+  num_inference_steps: 2
+  obs_as_global_cond: ${obs_as_global_cond}
+  diffusion_step_embed_dim: 128
+  down_dims:
+  - 128
+  - 256
+  - 512
+  kernel_size: 5
+  n_groups: 8
+  cond_predict_scale: true
+ema:
+  _target_: zprl.model.diffusion.ema_model.EMAModel
+  update_after_step: 0
+  inv_gamma: 1.0
+  power: 0.8
+  min_value: 0.0
+  max_value: 0.9999
+dataloader:
+  batch_size: 256
+  num_workers: 12
+  shuffle: true
+  pin_memory: true
+  persistent_workers: false
+val_dataloader:
+  batch_size: 64
+  num_workers: 4
+  shuffle: false
+  pin_memory: true
+  persistent_workers: false
+optimizer:
+  _target_: torch.optim.AdamW
+  lr: 0.0001
+  betas:
+  - 0.95
+  - 0.999
+  eps: 1.0e-08
+  weight_decay: 1.0e-06
+training:
+  device: cuda:2
+  seed: 10
+  debug: false
+  resume: true
+  lr_scheduler: cosine
+  lr_warmup_steps: 100
+  num_epochs: 1001
+  gradient_accumulate_every: 1
+  use_ema: true
+  freeze_encoder: false
+  rollout_every: 100
+  checkpoint_every: 100
+  val_every: 25
+  sample_every: 25
+  max_train_steps: null
+  max_val_steps: null
+  tqdm_interval_sec: 1.0
+logging:
+  project: offline_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: test_mean_score
+    mode: max
+    k: 2
+    format_str: epoch_{epoch:04d}-score_{test_mean_score:.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}
+task:
+  name: transport_image
+  shape_meta:
+    obs:
+      shouldercamera0_image:
+        shape:
+        - 3
+        - 84
+        - 84
+        type: rgb
+      robot0_eye_in_hand_image:
+        shape:
+        - 3
+        - 84
+        - 84
+        type: rgb
+      robot0_eef_pos:
+        shape:
+        - 3
+      robot0_eef_quat:
+        shape:
+        - 4
+      robot0_gripper_qpos:
+        shape:
+        - 2
+      shouldercamera1_image:
+        shape:
+        - 3
+        - 84
+        - 84
+        type: rgb
+      robot1_eye_in_hand_image:
+        shape:
+        - 3
+        - 84
+        - 84
+        type: rgb
+      robot1_eef_pos:
+        shape:
+        - 3
+      robot1_eef_quat:
+        shape:
+        - 4
+      robot1_gripper_qpos:
+        shape:
+        - 2
+    action:
+      shape:
+      - 20
+  task_name: transport
+  dataset_type: mh
+  dataset_path: /media/datahub-2/ydj/robomimicv030/${task.task_name}/${task.dataset_type}/image_v141_subset_abs.hdf5
+  abs_action: true
+  env_runner:
+    _target_: zprl.env_runner.robomimic_image_runner.RobomimicImageRunner
+    dataset_path: /media/datahub-2/ydj/robomimicv030/${task.task_name}/${task.dataset_type}/image_v141_subset_abs.hdf5
+    shape_meta:
+      obs:
+        shouldercamera0_image:
+          shape:
+          - 3
+          - 84
+          - 84
+          type: rgb
+        robot0_eye_in_hand_image:
+          shape:
+          - 3
+          - 84
+          - 84
+          type: rgb
+        robot0_eef_pos:
+          shape:
+          - 3
+        robot0_eef_quat:
+          shape:
+          - 4
+        robot0_gripper_qpos:
+          shape:
+          - 2
+        shouldercamera1_image:
+          shape:
+          - 3
+          - 84
+          - 84
+          type: rgb
+        robot1_eye_in_hand_image:
+          shape:
+          - 3
+          - 84
+          - 84
+          type: rgb
+        robot1_eef_pos:
+          shape:
+          - 3
+        robot1_eef_quat:
+          shape:
+          - 4
+        robot1_gripper_qpos:
+          shape:
+          - 2
+      action:
+        shape:
+        - 20
+    n_train: 6
+    n_train_vis: 2
+    train_start_idx: 0
+    n_test: 50
+    n_test_vis: 4
+    test_start_seed: 100000
+    max_steps: 700
+    n_obs_steps: ${n_obs_steps}
+    n_action_steps: ${n_action_steps}
+    render_obs_key: shouldercamera0_image
+    fps: 10
+    crf: 22
+    past_action: ${past_action_visible}
+    abs_action: true
+    tqdm_interval_sec: 1.0
+    n_envs: 28
+  dataset:
+    _target_: zprl.dataset.robomimic_replay_image_dataset.RobomimicReplayImageDataset
+    shape_meta:
+      obs:
+        shouldercamera0_image:
+          shape:
+          - 3
+          - 84
+          - 84
+          type: rgb
+        robot0_eye_in_hand_image:
+          shape:
+          - 3
+          - 84
+          - 84
+          type: rgb
+        robot0_eef_pos:
+          shape:
+          - 3
+        robot0_eef_quat:
+          shape:
+          - 4
+        robot0_gripper_qpos:
+          shape:
+          - 2
+        shouldercamera1_image:
+          shape:
+          - 3
+          - 84
+          - 84
+          type: rgb
+        robot1_eye_in_hand_image:
+          shape:
+          - 3
+          - 84
+          - 84
+          type: rgb
+        robot1_eef_pos:
+          shape:
+          - 3
+        robot1_eef_quat:
+          shape:
+          - 4
+        robot1_gripper_qpos:
+          shape:
+          - 2
+      action:
+        shape:
+        - 20
+    dataset_path: /media/datahub-2/ydj/robomimicv030/${task.task_name}/${task.dataset_type}/image_v141_subset_abs.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: rotation_6d
+    use_legacy_normalizer: false
+    use_cache: true
+    seed: 10
+    val_ratio: 0.02
+    num_demo: 100

transport/.hydra/hydra.yaml

@@ -0,0 +1,163 @@
+hydra:
+  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}
+  launcher:
+    _target_: hydra._internal.core_plugins.basic_launcher.BasicLauncher
+  sweeper:
+    _target_: hydra._internal.core_plugins.basic_sweeper.BasicSweeper
+    max_batch_size: null
+    params: null
+  help:
+    app_name: ${hydra.job.name}
+    header: '${hydra.help.app_name} is powered by Hydra.
+
+      '
+    footer: 'Powered by Hydra (https://hydra.cc)
+
+      Use --hydra-help to view Hydra specific help
+
+      '
+    template: '${hydra.help.header}
+
+      == Configuration groups ==
+
+      Compose your configuration from those groups (group=option)
+
+
+      $APP_CONFIG_GROUPS
+
+
+      == Config ==
+
+      Override anything in the config (foo.bar=value)
+
+
+      $CONFIG
+
+
+      ${hydra.help.footer}
+
+      '
+  hydra_help:
+    template: 'Hydra (${hydra.runtime.version})
+
+      See https://hydra.cc for more info.
+
+
+      == Flags ==
+
+      $FLAGS_HELP
+
+
+      == Configuration groups ==
+
+      Compose your configuration from those groups (For example, append hydra/job_logging=disabled
+      to command line)
+
+
+      $HYDRA_CONFIG_GROUPS
+
+
+      Use ''--cfg hydra'' to Show the Hydra config.
+
+      '
+    hydra_help: ???
+  hydra_logging:
+    version: 1
+    formatters:
+      simple:
+        format: '[%(asctime)s][HYDRA] %(message)s'
+    handlers:
+      console:
+        class: logging.StreamHandler
+        formatter: simple
+        stream: ext://sys.stdout
+    root:
+      level: INFO
+      handlers:
+      - console
+    loggers:
+      logging_example:
+        level: DEBUG
+    disable_existing_loggers: false
+  job_logging:
+    version: 1
+    formatters:
+      simple:
+        format: '[%(asctime)s][%(name)s][%(levelname)s] - %(message)s'
+    handlers:
+      console:
+        class: logging.StreamHandler
+        formatter: simple
+        stream: ext://sys.stdout
+      file:
+        class: logging.FileHandler
+        formatter: simple
+        filename: ${hydra.runtime.output_dir}/${hydra.job.name}.log
+    root:
+      level: INFO
+      handlers:
+      - console
+      - file
+    disable_existing_loggers: false
+  env: {}
+  mode: RUN
+  searchpath: []
+  callbacks: {}
+  output_subdir: .hydra
+  overrides:
+    hydra:
+    - hydra.mode=RUN
+    task:
+    - training.seed=10
+    - task.dataset.seed=10
+    - exp_name=unet_vib_recon
+    - training.device=cuda:2
+    - task=transport_image_abs
+    - policy.vib_beta=0.0001
+    - policy.vib_latent_dim=32
+    - n_action_steps=5
+  job:
+    name: train
+    chdir: null
+    override_dirname: exp_name=unet_vib_recon,n_action_steps=5,policy.vib_beta=0.0001,policy.vib_latent_dim=32,task.dataset.seed=10,task=transport_image_abs,training.device=cuda:2,training.seed=10
+    id: ???
+    num: ???
+    config_name: train_flow_match_vib_unet_image_workspace
+    env_set: {}
+    env_copy: []
+    config:
+      override_dirname:
+        kv_sep: '='
+        item_sep: ','
+        exclude_keys: []
+  runtime:
+    version: 1.3.2
+    version_base: '1.3'
+    cwd: /ghome/h2/djyu/zprl-sim/diffusion_policy
+    config_sources:
+    - path: hydra.conf
+      schema: pkg
+      provider: hydra
+    - path: /ghome/h2/djyu/zprl-sim/diffusion_policy/zprl/config
+      schema: file
+      provider: main
+    - path: ''
+      schema: structured
+      provider: schema
+    output_dir: /ghome/h2/djyu/zprl-sim/diffusion_policy/data/outputs/2026.03.23/21.34.15_train_flow_match_vib_unet_image_transport_image
+    choices:
+      task: transport_image_abs
+      hydra/env: default
+      hydra/callbacks: null
+      hydra/job_logging: default
+      hydra/hydra_logging: default
+      hydra/hydra_help: default
+      hydra/help: default
+      hydra/sweeper: basic
+      hydra/launcher: basic
+      hydra/output: default
+  verbose: false

transport/.hydra/overrides.yaml

@@ -0,0 +1,8 @@
+- training.seed=10
+- task.dataset.seed=10
+- exp_name=unet_vib_recon
+- training.device=cuda:2
+- task=transport_image_abs
+- policy.vib_beta=0.0001
+- policy.vib_latent_dim=32
+- n_action_steps=5

transport/checkpoints/latest.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4fb2acb46954136a27cb290fdecbac57ec43f348aedd11ae3bfa450faeba2c1d
+size 1179618798

transport/logs.json.txt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:285b3815ecefcbce09405d8dd8afc2d4f0a8829ff74b10ae9bb26ea1197e90f6
+size 83764180

transport/train.log

@@ -0,0 +1,5 @@
+[2026-03-23 21:34:17,589][zprl.model.vision.multi_image_obs_encoder][INFO] - number of parameters: 49.05 M
+[2026-03-23 21:34:17,952][zprl.model.diffusion.conditional_unet1d][INFO] - number of parameters: 23.80 M
+[2026-03-23 21:34:17,954][zprl.model.vib][INFO] - number of parameters: 0.42 M
+[2026-03-23 21:34:17,957][zprl.model.vib][INFO] - number of parameters: 0.41 M
+[2026-03-23 21:34:26,683][OpenGL.acceleratesupport][INFO] - No OpenGL_accelerate module loaded: No module named 'OpenGL_accelerate'