From db3baf9f90dfd1cca77d5005543e29b5e85f0d68 Mon Sep 17 00:00:00 2001 From: Aidan Gao Date: Thu, 26 Feb 2026 17:07:15 -0500 Subject: [PATCH] new s3 data configs --- egomimic/hydra_configs/data/aria.yaml | 36 ++++++ egomimic/hydra_configs/data/aria_bc_zarr.yaml | 45 -------- egomimic/hydra_configs/data/eva.yaml | 37 ++++++ egomimic/hydra_configs/data/eva_bc_zarr.yaml | 109 ------------------ .../hydra_configs/data/eva_human_cotrain.yaml | 65 +++++++++++ egomimic/hydra_configs/data/mecka.yaml | 36 ++++++ egomimic/hydra_configs/data/scale.yaml | 36 ++++++ .../hydra_configs/data/scale_zarr_test.yaml | 71 ------------ .../{data => lerobot_data}/aria_bc_s3.yaml | 0 .../cotrain_mecka_aria.yaml | 0 .../{data => lerobot_data}/cotrain_s3.yaml | 0 .../{data => lerobot_data}/eva_bc_s3.yaml | 0 .../{data => lerobot_data}/mecka_test.yaml | 0 .../{data => lerobot_data}/test_bimanual.yaml | 0 .../hpt_cotrain_scale_flow_shared_head.yaml | 68 +++++++++++ egomimic/hydra_configs/train_zarr.yaml | 16 +-- egomimic/pl_utils/pl_data_utils.py | 12 +- egomimic/utils/pose_utils.py | 1 + 18 files changed, 294 insertions(+), 238 deletions(-) create mode 100644 egomimic/hydra_configs/data/aria.yaml delete mode 100644 egomimic/hydra_configs/data/aria_bc_zarr.yaml create mode 100644 egomimic/hydra_configs/data/eva.yaml delete mode 100644 egomimic/hydra_configs/data/eva_bc_zarr.yaml create mode 100644 egomimic/hydra_configs/data/eva_human_cotrain.yaml create mode 100644 egomimic/hydra_configs/data/mecka.yaml create mode 100644 egomimic/hydra_configs/data/scale.yaml delete mode 100644 egomimic/hydra_configs/data/scale_zarr_test.yaml rename egomimic/hydra_configs/{data => lerobot_data}/aria_bc_s3.yaml (100%) rename egomimic/hydra_configs/{data => lerobot_data}/cotrain_mecka_aria.yaml (100%) rename egomimic/hydra_configs/{data => lerobot_data}/cotrain_s3.yaml (100%) rename egomimic/hydra_configs/{data => lerobot_data}/eva_bc_s3.yaml (100%) rename egomimic/hydra_configs/{data => lerobot_data}/mecka_test.yaml (100%) rename egomimic/hydra_configs/{data => lerobot_data}/test_bimanual.yaml (100%) create mode 100644 egomimic/hydra_configs/model/hpt_cotrain_scale_flow_shared_head.yaml diff --git a/egomimic/hydra_configs/data/aria.yaml b/egomimic/hydra_configs/data/aria.yaml new file mode 100644 index 00000000..e60a92a4 --- /dev/null +++ b/egomimic/hydra_configs/data/aria.yaml @@ -0,0 +1,36 @@ +_target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper + +train_datasets: + aria_bimanual: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + resolver: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver + folder_path: /coc/flash7/scratch/egoverseDebugDatasets/aria + key_map: + _target_: egomimic.rldb.embodiment.human.Aria.get_keymap + transform_list: + _target_: egomimic.rldb.embodiment.human.Aria.get_transform_list + filters: + episode_hash: "2025-09-20-17-47-54-000000" + mode: total +valid_datasets: + aria_bimanual: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + resolver: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver + folder_path: /coc/flash7/scratch/egoverseDebugDatasets/aria + key_map: + _target_: egomimic.rldb.embodiment.human.Aria.get_keymap + transform_list: + _target_: egomimic.rldb.embodiment.human.Aria.get_transform_list + filters: + episode_hash: "2025-09-20-17-47-54-000000" + mode: total +train_dataloader_params: + aria_bimanual: + batch_size: 32 + num_workers: 10 +valid_dataloader_params: + aria_bimanual: + batch_size: 32 + num_workers: 10 diff --git a/egomimic/hydra_configs/data/aria_bc_zarr.yaml b/egomimic/hydra_configs/data/aria_bc_zarr.yaml deleted file mode 100644 index 2f2bf186..00000000 --- a/egomimic/hydra_configs/data/aria_bc_zarr.yaml +++ /dev/null @@ -1,45 +0,0 @@ -_target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper -train_datasets: - aria_bimanual: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver - resolver: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.LocalEpisodeResolver - folder_path: /nethome/paphiwetsa3/flash/datasets/proc_zarr - key_map: - observations.images.front_img_1: #batch key - key_type: camera_keys # key type - zarr_key: front_img_1 - actions_cartesian: - key_type: action_keys - zarr_key: obs_ee_pose_6d - horizon: 100 - ee_pose: - key_type: proprio_keys - zarr_key: obs_ee_pose_6d - mode: total -valid_datasets: - aria_bimanual: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver - resolver: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.LocalEpisodeResolver - folder_path: /nethome/paphiwetsa3/flash/datasets/proc_zarr - key_map: - observations.images.front_img_1: #batch key - key_type: camera_keys # key type - zarr_key: front_img_1 - actions_cartesian: - key_type: action_keys - zarr_key: obs_ee_pose_6d - horizon: 100 - ee_pose: - key_type: proprio_keys - zarr_key: obs_ee_pose_6d - mode: total -train_dataloader_params: - aria_bimanual: - batch_size: 32 - num_workers: 10 -valid_dataloader_params: - aria_bimanual: - batch_size: 32 - num_workers: 10 diff --git a/egomimic/hydra_configs/data/eva.yaml b/egomimic/hydra_configs/data/eva.yaml new file mode 100644 index 00000000..9bf10a3f --- /dev/null +++ b/egomimic/hydra_configs/data/eva.yaml @@ -0,0 +1,37 @@ +_target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper +train_datasets: + eva_bimanual: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + resolver: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver + folder_path: /coc/flash7/scratch/egoverseDebugDatasets/egoverseS3DatasetTest/ + key_map: + _target_: egomimic.rldb.embodiment.eva.Eva.get_keymap + transform_list: + _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list + filters: + episode_hash: "2025-12-26-18-07-46-296000" + mode: total + +valid_datasets: + eva_bimanual: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + resolver: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver + folder_path: /coc/flash7/scratch/egoverseDebugDatasets/egoverseS3DatasetTest/ + key_map: + _target_: egomimic.rldb.embodiment.eva.Eva.get_keymap + transform_list: + _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list + filters: + episode_hash: "2025-12-26-18-07-46-296000" + mode: total + +train_dataloader_params: + eva_bimanual: + batch_size: 32 + num_workers: 10 +valid_dataloader_params: + eva_bimanual: + batch_size: 32 + num_workers: 10 diff --git a/egomimic/hydra_configs/data/eva_bc_zarr.yaml b/egomimic/hydra_configs/data/eva_bc_zarr.yaml deleted file mode 100644 index 153c7363..00000000 --- a/egomimic/hydra_configs/data/eva_bc_zarr.yaml +++ /dev/null @@ -1,109 +0,0 @@ -_target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper -train_datasets: - eva_bimanual: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver - resolver: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver - folder_path: /coc/flash7/scratch/egoverseDebugDatasets/egoverseS3DatasetTest/ - key_map: - observations.images.front_img_1: - key_type: camera_keys - zarr_key: images.front_1 - observations.images.right_wrist_img: - key_type: camera_keys - zarr_key: images.right_wrist - observations.images.left_wrist_img: - key_type: camera_keys - zarr_key: images.left_wrist - right.obs_ee_pose: - key_type: proprio_keys - zarr_key: right.obs_ee_pose - right.obs_gripper: - key_type: proprio_keys - zarr_key: right.gripper - left.obs_ee_pose: - key_type: proprio_keys - zarr_key: left.obs_ee_pose - left.obs_gripper: - key_type: proprio_keys - zarr_key: left.gripper - right.gripper: - key_type: action_keys - zarr_key: right.gripper - horizon: 45 - left.gripper: - key_type: action_keys - zarr_key: left.gripper - horizon: 45 - right.cmd_ee_pose: - key_type: action_keys - zarr_key: right.cmd_ee_pose - horizon: 45 - left.cmd_ee_pose: - key_type: action_keys - zarr_key: left.cmd_ee_pose - horizon: 45 - transform_list: - _target_: egomimic.rldb.zarr.action_chunk_transforms.build_eva_bimanual_transform_list - filters: - episode_hash: "2025-12-26-18-07-46-296000" - mode: total - -valid_datasets: - eva_bimanual: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver - resolver: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver - folder_path: /coc/flash7/scratch/egoverseDebugDatasets/egoverseS3DatasetTest/ - key_map: - observations.images.front_img_1: - key_type: camera_keys - zarr_key: images.front_1 - observations.images.right_wrist_img: - key_type: camera_keys - zarr_key: images.right_wrist - observations.images.left_wrist_img: - key_type: camera_keys - zarr_key: images.left_wrist - right.obs_ee_pose: - key_type: proprio_keys - zarr_key: right.obs_ee_pose - right.obs_gripper: - key_type: proprio_keys - zarr_key: right.gripper - left.obs_ee_pose: - key_type: proprio_keys - zarr_key: left.obs_ee_pose - left.obs_gripper: - key_type: proprio_keys - zarr_key: left.gripper - right.gripper: - key_type: action_keys - zarr_key: right.gripper - horizon: 45 - left.gripper: - key_type: action_keys - zarr_key: left.gripper - horizon: 45 - right.cmd_ee_pose: - key_type: action_keys - zarr_key: right.cmd_ee_pose - horizon: 45 - left.cmd_ee_pose: - key_type: action_keys - zarr_key: left.cmd_ee_pose - horizon: 45 - transform_list: - _target_: egomimic.rldb.zarr.action_chunk_transforms.build_eva_bimanual_transform_list - filters: - episode_hash: "2025-12-26-18-07-46-296000" - mode: total - -train_dataloader_params: - eva_bimanual: - batch_size: 32 - num_workers: 10 -valid_dataloader_params: - eva_bimanual: - batch_size: 32 - num_workers: 10 diff --git a/egomimic/hydra_configs/data/eva_human_cotrain.yaml b/egomimic/hydra_configs/data/eva_human_cotrain.yaml new file mode 100644 index 00000000..cabf760d --- /dev/null +++ b/egomimic/hydra_configs/data/eva_human_cotrain.yaml @@ -0,0 +1,65 @@ +_target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper +train_datasets: + eva_bimanual: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + resolver: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver + folder_path: /coc/flash7/scratch/egoverseDebugDatasets/egoverseS3DatasetTest/ + key_map: + _target_: egomimic.rldb.embodiment.eva.Eva.get_keymap + transform_list: + _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list + filters: + episode_hash: "2025-12-26-18-07-46-296000" + mode: total + aria_bimanual: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + resolver: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver + folder_path: /coc/flash7/scratch/egoverseDebugDatasets/aria + key_map: + _target_: egomimic.rldb.embodiment.human.Aria.get_keymap + transform_list: + _target_: egomimic.rldb.embodiment.human.Aria.get_transform_list + filters: + episode_hash: "2025-09-20-17-47-54-000000" + mode: total +valid_datasets: + eva_bimanual: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + resolver: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver + folder_path: /coc/flash7/scratch/egoverseDebugDatasets/egoverseS3DatasetTest/ + key_map: + _target_: egomimic.rldb.embodiment.eva.Eva.get_keymap + transform_list: + _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list + filters: + episode_hash: "2025-12-26-18-07-46-296000" + mode: total + aria_bimanual: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + resolver: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver + folder_path: /coc/flash7/scratch/egoverseDebugDatasets/aria + key_map: + _target_: egomimic.rldb.embodiment.human.Aria.get_keymap + transform_list: + _target_: egomimic.rldb.embodiment.human.Aria.get_transform_list + filters: + episode_hash: "2025-09-20-17-47-54-000000" + mode: total +train_dataloader_params: + eva_bimanual: + batch_size: 32 + num_workers: 10 + aria_bimanual: + batch_size: 32 + num_workers: 10 +valid_dataloader_params: + eva_bimanual: + batch_size: 32 + num_workers: 10 + aria_bimanual: + batch_size: 32 + num_workers: 10 \ No newline at end of file diff --git a/egomimic/hydra_configs/data/mecka.yaml b/egomimic/hydra_configs/data/mecka.yaml new file mode 100644 index 00000000..647d9b7c --- /dev/null +++ b/egomimic/hydra_configs/data/mecka.yaml @@ -0,0 +1,36 @@ +_target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper + +train_datasets: + mecka_bimanual: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + resolver: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver + folder_path: /coc/flash7/scratch/egoverseDebugDatasets/mecka + key_map: + _target_: egomimic.rldb.embodiment.human.Mecka.get_keymap + transform_list: + _target_: egomimic.rldb.embodiment.human.Mecka.get_transform_list + filters: + episode_hash: "69199812208123403bbdb24f" + mode: total +valid_datasets: + mecka_bimanual: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + resolver: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver + folder_path: /coc/flash7/scratch/egoverseDebugDatasets/mecka + key_map: + _target_: egomimic.rldb.embodiment.human.Mecka.get_keymap + transform_list: + _target_: egomimic.rldb.embodiment.human.Mecka.get_transform_list + filters: + episode_hash: "69199812208123403bbdb24f" + mode: total +train_dataloader_params: + mecka_bimanual: + batch_size: 32 + num_workers: 10 +valid_dataloader_params: + mecka_bimanual: + batch_size: 32 + num_workers: 10 diff --git a/egomimic/hydra_configs/data/scale.yaml b/egomimic/hydra_configs/data/scale.yaml new file mode 100644 index 00000000..44022e7b --- /dev/null +++ b/egomimic/hydra_configs/data/scale.yaml @@ -0,0 +1,36 @@ +_target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper + +train_datasets: + scale_bimanual: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + resolver: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver + folder_path: /coc/flash7/scratch/egoverseDebugDatasets/scale + key_map: + _target_: egomimic.rldb.embodiment.human.Scale.get_keymap + transform_list: + _target_: egomimic.rldb.embodiment.human.Scale.get_transform_list + filters: + episode_hash: "69199812208123403bbdb24f" + mode: total +valid_datasets: + scale_bimanual: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + resolver: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver + folder_path: /coc/flash7/scratch/egoverseDebugDatasets/scale + key_map: + _target_: egomimic.rldb.embodiment.human.Scale.get_keymap + transform_list: + _target_: egomimic.rldb.embodiment.human.Scale.get_transform_list + filters: + episode_hash: "69199812208123403bbdb24f" + mode: total +train_dataloader_params: + scale_bimanual: + batch_size: 32 + num_workers: 10 +valid_dataloader_params: + scale_bimanual: + batch_size: 32 + num_workers: 10 diff --git a/egomimic/hydra_configs/data/scale_zarr_test.yaml b/egomimic/hydra_configs/data/scale_zarr_test.yaml deleted file mode 100644 index 80784f19..00000000 --- a/egomimic/hydra_configs/data/scale_zarr_test.yaml +++ /dev/null @@ -1,71 +0,0 @@ -_target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper -train_datasets: - scale_bimanual: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver - resolver: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.LocalEpisodeResolver - folder_path: /coc/flash7/scratch/egoverseDebugDatasets/scale/2026-02-24-01-49-24-166324 - key_map: - observations.images.front_img_1: #batch key - key_type: camera_keys # key type - zarr_key: images.front_1 - right.action_ee_pose: - key_type: action_keys - zarr_key: right.obs_ee_pose - horizon: 30 - left.action_ee_pose: - key_type: action_keys - zarr_key: left.obs_ee_pose - horizon: 30 - right.obs_ee_pose: - key_type: proprio_keys - zarr_key: right.obs_ee_pose - left.obs_ee_pose: - key_type: proprio_keys - zarr_key: left.obs_ee_pose - obs_head_pose: - zarr_key: obs_head_pose - key_type: proprio_keys - transform_list: - _target_: egomimic.rldb.zarr.action_chunk_transforms.build_aria_bimanual_transform_list - stride: 1 - mode: train -valid_datasets: - scale_bimanual: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver - resolver: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.LocalEpisodeResolver - folder_path: /coc/flash7/scratch/egoverseDebugDatasets/scale/2026-02-24-01-49-24-166324 - key_map: - observations.images.front_img_1: #batch key - key_type: camera_keys # key type - zarr_key: images.front_1 - right.action_ee_pose: - key_type: action_keys - zarr_key: right.obs_ee_pose - horizon: 30 - left.action_ee_pose: - key_type: action_keys - zarr_key: left.obs_ee_pose - horizon: 30 - right.obs_ee_pose: - key_type: proprio_keys - zarr_key: right.obs_ee_pose - left.obs_ee_pose: - key_type: proprio_keys - zarr_key: left.obs_ee_pose - obs_head_pose: - zarr_key: obs_head_pose - key_type: proprio_keys - transform_list: - _target_: egomimic.rldb.zarr.action_chunk_transforms.build_aria_bimanual_transform_list - stride: 1 - mode: valid -train_dataloader_params: - scale_bimanual: - batch_size: 32 - num_workers: 10 -valid_dataloader_params: - scale_bimanual: - batch_size: 32 - num_workers: 10 diff --git a/egomimic/hydra_configs/data/aria_bc_s3.yaml b/egomimic/hydra_configs/lerobot_data/aria_bc_s3.yaml similarity index 100% rename from egomimic/hydra_configs/data/aria_bc_s3.yaml rename to egomimic/hydra_configs/lerobot_data/aria_bc_s3.yaml diff --git a/egomimic/hydra_configs/data/cotrain_mecka_aria.yaml b/egomimic/hydra_configs/lerobot_data/cotrain_mecka_aria.yaml similarity index 100% rename from egomimic/hydra_configs/data/cotrain_mecka_aria.yaml rename to egomimic/hydra_configs/lerobot_data/cotrain_mecka_aria.yaml diff --git a/egomimic/hydra_configs/data/cotrain_s3.yaml b/egomimic/hydra_configs/lerobot_data/cotrain_s3.yaml similarity index 100% rename from egomimic/hydra_configs/data/cotrain_s3.yaml rename to egomimic/hydra_configs/lerobot_data/cotrain_s3.yaml diff --git a/egomimic/hydra_configs/data/eva_bc_s3.yaml b/egomimic/hydra_configs/lerobot_data/eva_bc_s3.yaml similarity index 100% rename from egomimic/hydra_configs/data/eva_bc_s3.yaml rename to egomimic/hydra_configs/lerobot_data/eva_bc_s3.yaml diff --git a/egomimic/hydra_configs/data/mecka_test.yaml b/egomimic/hydra_configs/lerobot_data/mecka_test.yaml similarity index 100% rename from egomimic/hydra_configs/data/mecka_test.yaml rename to egomimic/hydra_configs/lerobot_data/mecka_test.yaml diff --git a/egomimic/hydra_configs/data/test_bimanual.yaml b/egomimic/hydra_configs/lerobot_data/test_bimanual.yaml similarity index 100% rename from egomimic/hydra_configs/data/test_bimanual.yaml rename to egomimic/hydra_configs/lerobot_data/test_bimanual.yaml diff --git a/egomimic/hydra_configs/model/hpt_cotrain_scale_flow_shared_head.yaml b/egomimic/hydra_configs/model/hpt_cotrain_scale_flow_shared_head.yaml new file mode 100644 index 00000000..f68305e6 --- /dev/null +++ b/egomimic/hydra_configs/model/hpt_cotrain_scale_flow_shared_head.yaml @@ -0,0 +1,68 @@ +defaults: + - hpt_cotrain_enc_dec_base +robomimic_model: + domains: ["eva_bimanual", "aria_bimanual", "scale_bimanual"] + ac_keys: + eva_bimanual: "actions_cartesian" + aria_bimanual: "actions_cartesian" + scale_bimanual: "actions_cartesian" + shared_ac_key: "actions_cartesian" + camera_transforms: + scale_bimanual: + _target_: egomimic.utils.egomimicUtils.CameraTransforms + intrinsics_key: "scale" # change to base_half if using half res + extrinsics_key: "scale" + + 6dof: true + diffusion: true + stem_specs: + scale_bimanual: + state_ee_pose: + _target_: egomimic.models.hpt_nets.MLPPolicyStem + input_dim: 12 + output_dim: 256 + widths: [256] + specs: + random_horizon_masking: false + cross_attn: + crossattn_latent: 16 + crossattn_heads: 8 + crossattn_dim_head: 64 + crossattn_modality_dropout: 0.1 + modality_embed_dim: 256 + head_specs: + aria_bimanual: null + eva_bimanual: null + scale_bimanual: null + shared: + _target_: egomimic.models.fm_policy.FMPolicy + action_horizon: 100 + num_inference_steps: 50 + pooling: null + padding: "zero" + time_dist: "beta" + infer_ac_dims: + eva_bimanual: 14 + aria_bimanual: 14 + scale_bimanual: 14 + model: + _target_: egomimic.models.denoising_nets.CrossTransformer + nblocks: 6 + cond_dim: 256 + hidden_dim: 128 + act_dim: 14 + act_seq: 100 + n_heads: 4 + dropout: 0.1 + mlp_layers: 4 + mlp_ratio: 4 +optimizer: + _target_: torch.optim.AdamW + _partial_: true + lr: 1e-4 + weight_decay: 0.0001 +scheduler: + _target_: torch.optim.lr_scheduler.CosineAnnealingLR + _partial_: true + T_max: 1400 + eta_min: 1e-5 \ No newline at end of file diff --git a/egomimic/hydra_configs/train_zarr.yaml b/egomimic/hydra_configs/train_zarr.yaml index 23874b7d..a7ce2f5f 100644 --- a/egomimic/hydra_configs/train_zarr.yaml +++ b/egomimic/hydra_configs/train_zarr.yaml @@ -1,10 +1,10 @@ defaults: - model: hpt_bc_flow_eva - paths: default - - trainer: debug + - trainer: ddp - debug: null - - logger: debug - - data: eva_bc_zarr + - logger: wandb + - data: eva_resolver_zarr - callbacks: checkpoints - override hydra/launcher: submitit - _self_ @@ -81,16 +81,10 @@ data_schematic: # Dynamically fill in these shapes from the dataset zarr_key: observations.images.front_img_1 ee_pose: key_type: proprio_keys - zarr_key: observations.state.ee_pose_cam + zarr_key: observations.state.ee_pose actions_cartesian: key_type: action_keys - zarr_key: actions_ee_cartesian_cam - actions_keypoints: - key_type: action_keys - zarr_key: actions_ee_keypoints_world - actions_head_cartesian: - key_type: action_keys - zarr_key: actions_head_cartesian_world + zarr_key: actions_cartesian embodiment: key_type: metadata_keys zarr_key: metadata.embodiment diff --git a/egomimic/pl_utils/pl_data_utils.py b/egomimic/pl_utils/pl_data_utils.py index 76ac0984..f93265cb 100644 --- a/egomimic/pl_utils/pl_data_utils.py +++ b/egomimic/pl_utils/pl_data_utils.py @@ -76,7 +76,11 @@ def __init__( def train_dataloader(self): iterables = dict() for dataset_name, dataset in self.train_datasets.items(): - dataset_params = self.train_dataloader_params.get(dataset_name, {}) + dataset_params = self.train_dataloader_params.get(dataset_name) + if dataset_params is None or len(dataset_params) == 0: + raise ValueError( + f"No dataloader params found for dataset {dataset_name}. Please add {dataset_name} into your data config train_dataloader_params." + ) iterables[dataset_name] = DataLoader( dataset, shuffle=True, @@ -89,7 +93,11 @@ def train_dataloader(self): def val_dataloader(self): iterables = dict() for dataset_name, dataset in self.valid_datasets.items(): - dataset_params = self.valid_dataloader_params.get(dataset_name, {}) + dataset_params = self.valid_dataloader_params.get(dataset_name) + if dataset_params is None or len(dataset_params) == 0: + raise ValueError( + f"No dataloader params found for dataset {dataset_name}. Please add {dataset_name} into your data config valid_dataloader_params." + ) iterables[dataset_name] = DataLoader( dataset, shuffle=False, diff --git a/egomimic/utils/pose_utils.py b/egomimic/utils/pose_utils.py index 427277c4..5edf9ffe 100644 --- a/egomimic/utils/pose_utils.py +++ b/egomimic/utils/pose_utils.py @@ -134,6 +134,7 @@ def _xyzwxyz_to_matrix(xyzwxyz: np.ndarray) -> np.ndarray: mats = np.broadcast_to(np.eye(4, dtype=dtype), (B, 4, 4)).copy() quat_xyzw = xyzwxyz[:, [4, 5, 6, 3]] + mats[:, :3, :3] = R.from_quat(quat_xyzw).as_matrix() mats[:, :3, 3] = xyzwxyz[:, :3]