diff --git a/egomimic/algo/pi.py b/egomimic/algo/pi.py index f28158bb..27fef0d7 100644 --- a/egomimic/algo/pi.py +++ b/egomimic/algo/pi.py @@ -86,19 +86,19 @@ def __init__( self.camera_keys[embodiment_id] = [] self.proprio_keys[embodiment_id] = [] self.lang_keys[embodiment_id] = [] - for key in data_schematic.keys_of_type("action_keys"): + for key in data_schematic.keys_of_type("action_keys", embodiment_id): if ( data_schematic.is_key_with_embodiment(key, embodiment_id) and key == self.ac_keys[embodiment] ): self.ac_keys[embodiment_id] = key - for key in data_schematic.keys_of_type("camera_keys"): + for key in data_schematic.keys_of_type("camera_keys", embodiment_id): if data_schematic.is_key_with_embodiment(key, embodiment_id): self.camera_keys[embodiment_id].append(key) - for key in data_schematic.keys_of_type("proprio_keys"): + for key in data_schematic.keys_of_type("proprio_keys", embodiment_id): if data_schematic.is_key_with_embodiment(key, embodiment_id): self.proprio_keys[embodiment_id].append(key) - for key in data_schematic.keys_of_type("lang_keys"): + for key in data_schematic.keys_of_type("lang_keys", embodiment_id): if data_schematic.is_key_with_embodiment(key, embodiment_id): self.lang_keys[embodiment_id].append(key) @@ -165,12 +165,11 @@ def process_batch_for_training(self, batch): """ processed_batch = {} - for embodiment_id, _batch in batch.items(): + for embodiment_name, _batch in batch.items(): + embodiment_id = get_embodiment_id(embodiment_name) processed_batch[embodiment_id] = {} for key, value in _batch.items(): - key_name = self.data_schematic.lerobot_key_to_keyname( - key, embodiment_id - ) + key_name = self.data_schematic.zarr_key_to_keyname(key, embodiment_id) if key_name is not None: processed_batch[embodiment_id][key_name] = value @@ -196,6 +195,13 @@ def process_batch_for_training(self, batch): processed_batch[embodiment_id] = self.data_schematic.normalize_data( processed_batch[embodiment_id], embodiment_id ) + processed_batch[embodiment_id]["embodiment"] = torch.tensor( + [embodiment_id], device=device, dtype=torch.int64 + ) + + for key, value in processed_batch[embodiment_id].items(): + if isinstance(value, torch.Tensor) and value.dtype == torch.float64: + processed_batch[embodiment_id][key] = value.float() if not processed_batch: raise ValueError( diff --git a/egomimic/hydra_configs/data/aria.yaml b/egomimic/hydra_configs/data/aria.yaml index e60a92a4..75fc6310 100644 --- a/egomimic/hydra_configs/data/aria.yaml +++ b/egomimic/hydra_configs/data/aria.yaml @@ -5,7 +5,7 @@ train_datasets: _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 + folder_path: /storage/project/r-dxu345-0/shared/egoverseS3ZarrDatasets key_map: _target_: egomimic.rldb.embodiment.human.Aria.get_keymap transform_list: @@ -18,7 +18,7 @@ valid_datasets: _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 + folder_path: /storage/project/r-dxu345-0/shared/egoverseS3ZarrDatasets key_map: _target_: egomimic.rldb.embodiment.human.Aria.get_keymap transform_list: diff --git a/egomimic/hydra_configs/data/eva.yaml b/egomimic/hydra_configs/data/eva.yaml index 9bf10a3f..dfe76c97 100644 --- a/egomimic/hydra_configs/data/eva.yaml +++ b/egomimic/hydra_configs/data/eva.yaml @@ -4,7 +4,7 @@ train_datasets: _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/ + folder_path: /storage/project/r-dxu345-0/shared/egoverseS3ZarrDatasets key_map: _target_: egomimic.rldb.embodiment.eva.Eva.get_keymap transform_list: @@ -18,7 +18,7 @@ valid_datasets: _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/ + folder_path: /storage/project/r-dxu345-0/shared/egoverseS3ZarrDatasets key_map: _target_: egomimic.rldb.embodiment.eva.Eva.get_keymap transform_list: diff --git a/egomimic/hydra_configs/data/mecka.yaml b/egomimic/hydra_configs/data/mecka.yaml index 647d9b7c..d114f9fa 100644 --- a/egomimic/hydra_configs/data/mecka.yaml +++ b/egomimic/hydra_configs/data/mecka.yaml @@ -5,7 +5,7 @@ train_datasets: _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 + folder_path: /storage/project/r-dxu345-0/shared/egoverseS3ZarrDatasets key_map: _target_: egomimic.rldb.embodiment.human.Mecka.get_keymap transform_list: @@ -18,7 +18,7 @@ valid_datasets: _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 + folder_path: /storage/project/r-dxu345-0/shared/egoverseS3ZarrDatasets key_map: _target_: egomimic.rldb.embodiment.human.Mecka.get_keymap transform_list: diff --git a/egomimic/hydra_configs/data/scale.yaml b/egomimic/hydra_configs/data/scale.yaml index 44022e7b..5940d71e 100644 --- a/egomimic/hydra_configs/data/scale.yaml +++ b/egomimic/hydra_configs/data/scale.yaml @@ -5,7 +5,7 @@ train_datasets: _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 + folder_path: /storage/project/r-dxu345-0/shared/egoverseS3ZarrDatasets key_map: _target_: egomimic.rldb.embodiment.human.Scale.get_keymap transform_list: @@ -18,7 +18,7 @@ valid_datasets: _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 + folder_path: /storage/project/r-dxu345-0/shared/egoverseS3ZarrDatasets key_map: _target_: egomimic.rldb.embodiment.human.Scale.get_keymap transform_list: diff --git a/egomimic/hydra_configs/model/pi0.5.yaml b/egomimic/hydra_configs/model/pi0.5_bc_aria.yaml similarity index 90% rename from egomimic/hydra_configs/model/pi0.5.yaml rename to egomimic/hydra_configs/model/pi0.5_bc_aria.yaml index c1af0cc4..0e9a5966 100644 --- a/egomimic/hydra_configs/model/pi0.5.yaml +++ b/egomimic/hydra_configs/model/pi0.5_bc_aria.yaml @@ -20,7 +20,7 @@ robomimic_model: config: pytorch_training_precision: bfloat16 - pytorch_weight_path: /storage/home/hcoda1/5/rpunamiya6/cedar-dx/rpunamiya6/Projects/EgoVerse/egomimic/algo/pi_checkpoints/pi05_base_pytorch + pytorch_weight_path: /storage/project/r-dxu345-0/rco3/EgoVerse/egomimic/algo/pi_checkpoints/pi05_base_pytorch model: pi05: true action_dim: 32 diff --git a/egomimic/hydra_configs/model/pi0.5_bc_eva.yaml b/egomimic/hydra_configs/model/pi0.5_bc_eva.yaml new file mode 100644 index 00000000..d1886d53 --- /dev/null +++ b/egomimic/hydra_configs/model/pi0.5_bc_eva.yaml @@ -0,0 +1,56 @@ +_target_: egomimic.pl_utils.pl_model.ModelWrapper +robomimic_model: + _target_: egomimic.algo.pi.PI + data_schematic: _${data.dataset.data_schematic} + camera_transforms: + _target_: egomimic.utils.egomimicUtils.CameraTransforms + intrinsics_key: "base" # change to base_half if using half res + extrinsics_key: "x5Dec13_2" + ac_keys: + eva_bimanual: "actions_cartesian" + domains: ["eva_bimanual"] + + action_converters: + rules: + EVA_BIMANUAL: + _target_: egomimic.utils.action_utils.RobotBimanualCartesianEuler + # optional fallback if no match is found + fallback: + _target_: egomimic.utils.action_utils.BaseActionConverter + + config: + pytorch_training_precision: bfloat16 + pytorch_weight_path: /storage/project/r-dxu345-0/rco3/EgoVerse/egomimic/algo/pi_checkpoints/pi05_base_pytorch + model: + pi05: true + action_dim: 32 + action_horizon: 100 + max_token_len: 180 + + train_image_augs: + _target_: torchvision.transforms.Compose + transforms: + - _target_: torchvision.transforms.Resize + size: 224 + interpolation: 3 + eval_image_augs: + _target_: torchvision.transforms.Compose + transforms: + - _target_: torchvision.transforms.Resize + size: 224 + interpolation: 3 + +optimizer: + _target_: torch.optim.AdamW + _partial_: true + lr: 5e-4 + betas: [0.9, 0.999] + eps: 1e-8 + weight_decay: 0.0 + +scheduler: + _target_: transformers.get_cosine_schedule_with_warmup + _partial_: true + num_warmup_steps: 1000 + num_training_steps: 100000 + num_cycles: 0.5 \ No newline at end of file diff --git a/egomimic/hydra_configs/model/pi0.5_bc_mecka.yaml b/egomimic/hydra_configs/model/pi0.5_bc_mecka.yaml new file mode 100644 index 00000000..e7d2117c --- /dev/null +++ b/egomimic/hydra_configs/model/pi0.5_bc_mecka.yaml @@ -0,0 +1,56 @@ +_target_: egomimic.pl_utils.pl_model.ModelWrapper +robomimic_model: + _target_: egomimic.algo.pi.PI + data_schematic: _${data.dataset.data_schematic} + camera_transforms: + _target_: egomimic.utils.egomimicUtils.CameraTransforms + intrinsics_key: "mecka" # change to base_half if using half res + extrinsics_key: "mecka" + ac_keys: + mecka_bimanual: "actions_cartesian" + domains: ["mecka_bimanual"] + + action_converters: + rules: + MECKA_BIMANUAL: + _target_: egomimic.utils.action_utils.HumanBimanualCartesianEuler + # optional fallback if no match is found + fallback: + _target_: egomimic.utils.action_utils.BaseActionConverter + + config: + pytorch_training_precision: bfloat16 + pytorch_weight_path: /storage/project/r-dxu345-0/rco3/EgoVerse/egomimic/algo/pi_checkpoints/pi05_base_pytorch + model: + pi05: true + action_dim: 32 + action_horizon: 100 + max_token_len: 180 + + train_image_augs: + _target_: torchvision.transforms.Compose + transforms: + - _target_: torchvision.transforms.Resize + size: 224 + interpolation: 3 + eval_image_augs: + _target_: torchvision.transforms.Compose + transforms: + - _target_: torchvision.transforms.Resize + size: 224 + interpolation: 3 + +optimizer: + _target_: torch.optim.AdamW + _partial_: true + lr: 5e-4 + betas: [0.9, 0.999] + eps: 1e-8 + weight_decay: 0.0 + +scheduler: + _target_: transformers.get_cosine_schedule_with_warmup + _partial_: true + num_warmup_steps: 1000 + num_training_steps: 100000 + num_cycles: 0.5 \ No newline at end of file diff --git a/egomimic/hydra_configs/model/pi0.5_bc_scale.yaml b/egomimic/hydra_configs/model/pi0.5_bc_scale.yaml new file mode 100644 index 00000000..521abbaf --- /dev/null +++ b/egomimic/hydra_configs/model/pi0.5_bc_scale.yaml @@ -0,0 +1,56 @@ +_target_: egomimic.pl_utils.pl_model.ModelWrapper +robomimic_model: + _target_: egomimic.algo.pi.PI + data_schematic: _${data.dataset.data_schematic} + camera_transforms: + _target_: egomimic.utils.egomimicUtils.CameraTransforms + intrinsics_key: "scale" # change to base_half if using half res + extrinsics_key: "scale" + ac_keys: + scale_bimanual: "actions_cartesian" + domains: ["scale_bimanual"] + + action_converters: + rules: + SCALE_BIMANUAL: + _target_: egomimic.utils.action_utils.HumanBimanualCartesianEuler + # optional fallback if no match is found + fallback: + _target_: egomimic.utils.action_utils.BaseActionConverter + + config: + pytorch_training_precision: bfloat16 + pytorch_weight_path: /storage/project/r-dxu345-0/rco3/EgoVerse/egomimic/algo/pi_checkpoints/pi05_base_pytorch + model: + pi05: true + action_dim: 32 + action_horizon: 100 + max_token_len: 180 + + train_image_augs: + _target_: torchvision.transforms.Compose + transforms: + - _target_: torchvision.transforms.Resize + size: 224 + interpolation: 3 + eval_image_augs: + _target_: torchvision.transforms.Compose + transforms: + - _target_: torchvision.transforms.Resize + size: 224 + interpolation: 3 + +optimizer: + _target_: torch.optim.AdamW + _partial_: true + lr: 5e-4 + betas: [0.9, 0.999] + eps: 1e-8 + weight_decay: 0.0 + +scheduler: + _target_: transformers.get_cosine_schedule_with_warmup + _partial_: true + num_warmup_steps: 1000 + num_training_steps: 100000 + num_cycles: 0.5 \ No newline at end of file diff --git a/egomimic/hydra_configs/model/pi0.5_cotrain_eva_aria.yaml b/egomimic/hydra_configs/model/pi0.5_cotrain_eva_aria.yaml new file mode 100644 index 00000000..c0521f56 --- /dev/null +++ b/egomimic/hydra_configs/model/pi0.5_cotrain_eva_aria.yaml @@ -0,0 +1,64 @@ +_target_: egomimic.pl_utils.pl_model.ModelWrapper +robomimic_model: + _target_: egomimic.algo.pi.PI + data_schematic: _${data.dataset.data_schematic} + camera_transforms: + eva_bimanual: + _target_: egomimic.utils.egomimicUtils.CameraTransforms + intrinsics_key: "base" # change to base_half if using half res + extrinsics_key: "x5Dec13_2" + aria_bimanual: + _target_: egomimic.utils.egomimicUtils.CameraTransforms + intrinsics_key: "base" # change to base_half if using half res + extrinsics_key: "ariaJun7" + ac_keys: + eva_bimanual: "actions_cartesian" + aria_bimanual: "actions_cartesian" + domains: ["eva_bimanual", "aria_bimanual"] + + action_converters: + rules: + EVA_BIMANUAL: + _target_: egomimic.utils.action_utils.RobotBimanualCartesianEuler + ARIA_BIMANUAL: + _target_: egomimic.utils.action_utils.HumanBimanualCartesianEuler + # optional fallback if no match is found + fallback: + _target_: egomimic.utils.action_utils.BaseActionConverter + + config: + pytorch_training_precision: bfloat16 + pytorch_weight_path: /storage/project/r-dxu345-0/rco3/EgoVerse/egomimic/algo/pi_checkpoints/pi05_base_pytorch + model: + pi05: true + action_dim: 32 + action_horizon: 100 + max_token_len: 180 + + train_image_augs: + _target_: torchvision.transforms.Compose + transforms: + - _target_: torchvision.transforms.Resize + size: 224 + interpolation: 3 + eval_image_augs: + _target_: torchvision.transforms.Compose + transforms: + - _target_: torchvision.transforms.Resize + size: 224 + interpolation: 3 + +optimizer: + _target_: torch.optim.AdamW + _partial_: true + lr: 5e-4 + betas: [0.9, 0.999] + eps: 1e-8 + weight_decay: 0.0 + +scheduler: + _target_: transformers.get_cosine_schedule_with_warmup + _partial_: true + num_warmup_steps: 1000 + num_training_steps: 100000 + num_cycles: 0.5 \ No newline at end of file diff --git a/egomimic/hydra_configs/train_zarr.yaml b/egomimic/hydra_configs/train_zarr.yaml index 0cc53a23..905c6114 100644 --- a/egomimic/hydra_configs/train_zarr.yaml +++ b/egomimic/hydra_configs/train_zarr.yaml @@ -1,5 +1,5 @@ defaults: - - model: hpt_bc_flow_eva + - model: pi0.5_bc_eva - paths: default - trainer: ddp - debug: null @@ -103,7 +103,7 @@ data_schematic: # Dynamically fill in these shapes from the dataset zarr_key: metadata.embodiment viz_img_key: eva_bimanual: - front_img_1 + base_0_rgb aria_bimanual: front_img_1 mecka_bimanual: diff --git a/egomimic/hydra_configs/train_zarr_pi.yaml b/egomimic/hydra_configs/train_zarr_pi.yaml new file mode 100644 index 00000000..ee7fa846 --- /dev/null +++ b/egomimic/hydra_configs/train_zarr_pi.yaml @@ -0,0 +1,114 @@ +defaults: + - model: pi0.5_cotrain_eva_aria + - paths: default + - trainer: debug + - debug: null + - logger: debug + - data: eva_human_cotrain + - callbacks: checkpoints + - override hydra/launcher: submitit + - _self_ + +name: test +description: test +ckpt_path: null +train: true +eval: false + +eval_class: + _target_: egomimic.scripts.evaluation.Eve + mode: real + arm: both + eval_path: "./logs/eval/${name}_${now:%Y-%m-%d_%H-%M-%S}" + +hydra: + run: + # Dir should be experiment_name/description_{timestamp} + dir: ./logs/${name}/${description}_${now:%Y-%m-%d_%H-%M-%S} + sweep: + dir: ./logs/${name}/${description}_${now:%Y-%m-%d_%H-%M-%S} + +launch_params: + gpus_per_node: 1 + nodes: 1 + + +data_schematic: # Dynamically fill in these shapes from the dataset + _target_: egomimic.rldb.zarr.utils.DataSchematic + norm_mode: quantile + schematic_dict: + eva_bimanual: + base_0_rgb: #batch key + key_type: camera_keys # key type + zarr_key: observations.images.front_img_1 # dataset key + right_wrist_0_rgb: + key_type: camera_keys + zarr_key: observations.images.right_wrist_img + left_wrist_0_rgb: + key_type: camera_keys + zarr_key: observations.images.left_wrist_img + ee_pose: + key_type: proprio_keys + zarr_key: observations.state.ee_pose + joint_positions: + key_type: proprio_keys + zarr_key: observations.state.joint_positions + actions_joints: + key_type: action_keys + zarr_key: actions_joints + actions_cartesian: + key_type: action_keys + zarr_key: actions_cartesian + embodiment: + key_type: metadata_keys + zarr_key: metadata.embodiment + aria_bimanual: + base_0_rgb: + key_type: camera_keys + zarr_key: observations.images.front_img_1 + ee_pose: + key_type: proprio_keys + zarr_key: observations.state.ee_pose + actions_cartesian: + key_type: action_keys + zarr_key: actions_cartesian + embodiment: + key_type: metadata_keys + zarr_key: metadata.embodiment + mecka_bimanual: + base_0_rgb: + key_type: camera_keys + zarr_key: observations.images.front_img_1 + ee_pose: + key_type: proprio_keys + zarr_key: observations.state.ee_pose + actions_cartesian: + key_type: action_keys + zarr_key: actions_cartesian + embodiment: + key_type: metadata_keys + zarr_key: metadata.embodiment + scale_bimanual: + base_0_rgb: + key_type: camera_keys + zarr_key: observations.images.front_img_1 + ee_pose: + key_type: proprio_keys + zarr_key: observations.state.ee_pose + actions_cartesian: + key_type: action_keys + zarr_key: actions_cartesian + embodiment: + key_type: metadata_keys + zarr_key: metadata.embodiment + viz_img_key: + eva_bimanual: + base_0_rgb + aria_bimanual: + base_0_rgb + mecka_bimanual: + base_0_rgb + scale_bimanual: + base_0_rgb + +seed: 42 \ No newline at end of file diff --git a/egomimic/trainHydra.py b/egomimic/trainHydra.py index 8e4540e1..8541018f 100644 --- a/egomimic/trainHydra.py +++ b/egomimic/trainHydra.py @@ -202,7 +202,7 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: @hydra.main( - version_base="1.3", config_path="./hydra_configs", config_name="train_zarr.yaml" + version_base="1.3", config_path="./hydra_configs", config_name="train_zarr_pi.yaml" ) def main(cfg: DictConfig) -> Optional[float]: """Main entry point for training. diff --git a/egomimic/utils/aws/aws_data_utils.py b/egomimic/utils/aws/aws_data_utils.py index 2cc7f631..ba0602a2 100644 --- a/egomimic/utils/aws/aws_data_utils.py +++ b/egomimic/utils/aws/aws_data_utils.py @@ -1,6 +1,7 @@ from __future__ import annotations import os +import warnings from pathlib import Path import boto3 @@ -13,12 +14,20 @@ def _uses_r2_endpoint(endpoint_url: str | None) -> bool: return bool(endpoint_url and "r2.cloudflarestorage.com" in endpoint_url) -def load_env(path="~/.egoverse_env"): +def load_env(path="~/.egoverse_env", required: bool = False): p = Path(path).expanduser() if not p.exists(): - raise ValueError( - f"Env file {p} does not exist, run ./egomimic/utils/aws/setup_secret.sh" + if required: + raise ValueError( + f"Env file {p} does not exist, run ./egomimic/utils/aws/setup_secret.sh" + ) + warnings.warn( + f"Env file {p} does not exist; AWS/R2 env vars not set. " + "Run ./egomimic/utils/aws/setup_secret.sh if you need S3/R2.", + UserWarning, + stacklevel=2, ) + return for line in p.read_text().splitlines(): line = line.strip() if not line or line.startswith("#") or "=" not in line: diff --git a/egomimic/utils/aws/sql_tutorial.ipynb b/egomimic/utils/aws/sql_tutorial.ipynb new file mode 100644 index 00000000..eb8fddec --- /dev/null +++ b/egomimic/utils/aws/sql_tutorial.ipynb @@ -0,0 +1,611 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "id": "75093106", + "metadata": {}, + "outputs": [], + "source": [ + "from sqlalchemy import (\n", + " Boolean,\n", + " Column,\n", + " Float,\n", + " Integer,\n", + " MetaData,\n", + " String,\n", + " Table,\n", + " text,\n", + " update,\n", + ")\n", + "\n", + "from egomimic.utils.aws.aws_sql import (\n", + " TableRow,\n", + " create_default_engine,\n", + " episode_table_to_df,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "1bc257dd", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "\u001b[33mUsing hardcoded DB Credentials. Run ./egomimic/utils/aws/setup_secret.sh for better security!\u001b[0m\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Tables in schema 'app': ['episodes']\n" + ] + } + ], + "source": [ + "engine = create_default_engine()" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "a40642da", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
episode_hashoperatorlabnum_framestasktask_descriptionsceneobjectsprocessed_pathmp4_pathis_deletedembodimentrobot_nameis_evaleval_scoreeval_successprocessing_errorzarr_processed_pathzarr_processing_errorzarr_mp4_path
0692ea678dfa41139877769ce6846f7cf9cdf020d4a37a294mecka215.0object_in_containers3://rldb/mecka/flagship/692ea678dfa4113987776...rldb:/mecka/flagship/692ea678dfa41139877769ce_...Falsemeckamecka_bimanualFalse-1.0TrueEpisodeValidationError: No contiguous bimanual...
1692e887faf371a654ce7fa9b690366b20e94ce070afd1e8amecka44.0object_in_containers3://rldb/mecka/flagship/692e887faf371a654ce7f...rldb:/mecka/flagship/692e887faf371a654ce7fa9b_...Falsemeckamecka_bimanualFalse-1.0TrueEpisodeValidationError: No contiguous bimanual...
2692ea6c0b0d3d3c9d26069de6846f7cf9cdf020d4a37a294mecka121.0object_in_containers3://rldb/mecka/flagship/692ea6c0b0d3d3c9d2606...rldb:/mecka/flagship/692ea6c0b0d3d3c9d26069de_...Falsemeckamecka_bimanualFalse-1.0TrueEpisodeValidationError: No contiguous bimanual...
3692ea3257641010d043546a16846f7cf9cdf020d4a37a294mecka124.0object_in_containers3://rldb/mecka/flagship/692ea3257641010d04354...rldb:/mecka/flagship/692ea3257641010d043546a1_...Falsemeckamecka_bimanualFalse-1.0TrueEpisodeValidationError: No contiguous bimanual...
4692eabffb05428c72fb3758b6903bd580e94ce070afd33cdmecka68.0object_in_containers3://rldb/mecka/flagship/692eabffb05428c72fb37...rldb:/mecka/flagship/692eabffb05428c72fb3758b_...Falsemeckamecka_bimanualFalse-1.0TrueEpisodeValidationError: No contiguous bimanual...
...............................................................
80780692e6c7eaec602a46af0ff636903bd580e94ce070afd33cdmecka85.0object_in_containers3://rldb/mecka/flagship/692e6c7eaec602a46af0f...rldb:/mecka/flagship/692e6c7eaec602a46af0ff63_...Falsemeckamecka_bimanualFalse-1.0TrueEpisodeValidationError: No contiguous bimanual...
80781692fe5b27da13c10c87b358e6903bd580e94ce070afd33cdmecka52.0object_in_containers3://rldb/mecka/flagship/692fe5b27da13c10c87b3...rldb:/mecka/flagship/692fe5b27da13c10c87b358e_...Falsemeckamecka_bimanualFalse-1.0TrueEpisodeValidationError: No contiguous bimanual...
80782692e8eb796d2056439f015b16903b723466be157ec64587bmecka37.0object_in_containers3://rldb/mecka/flagship/692e8eb796d2056439f01...rldb:/mecka/flagship/692e8eb796d2056439f015b1_...Falsemeckamecka_bimanualFalse-1.0TrueEpisodeValidationError: No contiguous bimanual...
80783692fe8aa5108c2a2d6e74b5a6903bd580e94ce070afd33cdmecka55.0object_in_containers3://rldb/mecka/flagship/692fe8aa5108c2a2d6e74...rldb:/mecka/flagship/692fe8aa5108c2a2d6e74b5a_...Falsemeckamecka_bimanualFalse-1.0TrueEpisodeValidationError: No contiguous bimanual...
80784692fe54f8d747770865fb0186903bd580e94ce070afd33cdmecka68.0object_in_containers3://rldb/mecka/flagship/692fe54f8d747770865fb...rldb:/mecka/flagship/692fe54f8d747770865fb018_...Falsemeckamecka_bimanualFalse-1.0TrueEpisodeValidationError: No contiguous bimanual...
\n", + "

80785 rows × 20 columns

\n", + "
" + ], + "text/plain": [ + " episode_hash operator lab num_frames \\\n", + "0 692ea678dfa41139877769ce 6846f7cf9cdf020d4a37a294 mecka 215.0 \n", + "1 692e887faf371a654ce7fa9b 690366b20e94ce070afd1e8a mecka 44.0 \n", + "2 692ea6c0b0d3d3c9d26069de 6846f7cf9cdf020d4a37a294 mecka 121.0 \n", + "3 692ea3257641010d043546a1 6846f7cf9cdf020d4a37a294 mecka 124.0 \n", + "4 692eabffb05428c72fb3758b 6903bd580e94ce070afd33cd mecka 68.0 \n", + "... ... ... ... ... \n", + "80780 692e6c7eaec602a46af0ff63 6903bd580e94ce070afd33cd mecka 85.0 \n", + "80781 692fe5b27da13c10c87b358e 6903bd580e94ce070afd33cd mecka 52.0 \n", + "80782 692e8eb796d2056439f015b1 6903b723466be157ec64587b mecka 37.0 \n", + "80783 692fe8aa5108c2a2d6e74b5a 6903bd580e94ce070afd33cd mecka 55.0 \n", + "80784 692fe54f8d747770865fb018 6903bd580e94ce070afd33cd mecka 68.0 \n", + "\n", + " task task_description scene objects \\\n", + "0 object_in_container \n", + "1 object_in_container \n", + "2 object_in_container \n", + "3 object_in_container \n", + "4 object_in_container \n", + "... ... ... ... ... \n", + "80780 object_in_container \n", + "80781 object_in_container \n", + "80782 object_in_container \n", + "80783 object_in_container \n", + "80784 object_in_container \n", + "\n", + " processed_path \\\n", + "0 s3://rldb/mecka/flagship/692ea678dfa4113987776... \n", + "1 s3://rldb/mecka/flagship/692e887faf371a654ce7f... \n", + "2 s3://rldb/mecka/flagship/692ea6c0b0d3d3c9d2606... \n", + "3 s3://rldb/mecka/flagship/692ea3257641010d04354... \n", + "4 s3://rldb/mecka/flagship/692eabffb05428c72fb37... \n", + "... ... \n", + "80780 s3://rldb/mecka/flagship/692e6c7eaec602a46af0f... \n", + "80781 s3://rldb/mecka/flagship/692fe5b27da13c10c87b3... \n", + "80782 s3://rldb/mecka/flagship/692e8eb796d2056439f01... \n", + "80783 s3://rldb/mecka/flagship/692fe8aa5108c2a2d6e74... \n", + "80784 s3://rldb/mecka/flagship/692fe54f8d747770865fb... \n", + "\n", + " mp4_path is_deleted \\\n", + "0 rldb:/mecka/flagship/692ea678dfa41139877769ce_... False \n", + "1 rldb:/mecka/flagship/692e887faf371a654ce7fa9b_... False \n", + "2 rldb:/mecka/flagship/692ea6c0b0d3d3c9d26069de_... False \n", + "3 rldb:/mecka/flagship/692ea3257641010d043546a1_... False \n", + "4 rldb:/mecka/flagship/692eabffb05428c72fb3758b_... False \n", + "... ... ... \n", + "80780 rldb:/mecka/flagship/692e6c7eaec602a46af0ff63_... False \n", + "80781 rldb:/mecka/flagship/692fe5b27da13c10c87b358e_... False \n", + "80782 rldb:/mecka/flagship/692e8eb796d2056439f015b1_... False \n", + "80783 rldb:/mecka/flagship/692fe8aa5108c2a2d6e74b5a_... False \n", + "80784 rldb:/mecka/flagship/692fe54f8d747770865fb018_... False \n", + "\n", + " embodiment robot_name is_eval eval_score eval_success \\\n", + "0 mecka mecka_bimanual False -1.0 True \n", + "1 mecka mecka_bimanual False -1.0 True \n", + "2 mecka mecka_bimanual False -1.0 True \n", + "3 mecka mecka_bimanual False -1.0 True \n", + "4 mecka mecka_bimanual False -1.0 True \n", + "... ... ... ... ... ... \n", + "80780 mecka mecka_bimanual False -1.0 True \n", + "80781 mecka mecka_bimanual False -1.0 True \n", + "80782 mecka mecka_bimanual False -1.0 True \n", + "80783 mecka mecka_bimanual False -1.0 True \n", + "80784 mecka mecka_bimanual False -1.0 True \n", + "\n", + " processing_error zarr_processed_path \\\n", + "0 \n", + "1 \n", + "2 \n", + "3 \n", + "4 \n", + "... ... ... \n", + "80780 \n", + "80781 \n", + "80782 \n", + "80783 \n", + "80784 \n", + "\n", + " zarr_processing_error zarr_mp4_path \n", + "0 EpisodeValidationError: No contiguous bimanual... \n", + "1 EpisodeValidationError: No contiguous bimanual... \n", + "2 EpisodeValidationError: No contiguous bimanual... \n", + "3 EpisodeValidationError: No contiguous bimanual... \n", + "4 EpisodeValidationError: No contiguous bimanual... \n", + "... ... ... \n", + "80780 EpisodeValidationError: No contiguous bimanual... \n", + "80781 EpisodeValidationError: No contiguous bimanual... \n", + "80782 EpisodeValidationError: No contiguous bimanual... \n", + "80783 EpisodeValidationError: No contiguous bimanual... \n", + "80784 EpisodeValidationError: No contiguous bimanual... \n", + "\n", + "[80785 rows x 20 columns]" + ] + }, + "execution_count": 3, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df = episode_table_to_df(engine)\n", + "df" + ] + }, + { + "cell_type": "markdown", + "id": "21ee3b2d", + "metadata": {}, + "source": [ + "## Example Useful Functions" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "55b01865", + "metadata": {}, + "outputs": [], + "source": [ + "# Add Episode Test\n", + "episode = TableRow(\n", + " episode_hash=1761408819,\n", + " operator=\"test\",\n", + " lab=\"lab-x\",\n", + " num_frames=-1,\n", + " task=\"bimanual_test\",\n", + " task_description=\"Dummy row for table inspection\",\n", + " scene=\"kitchen\",\n", + " objects=\"cup,plate,spoon\",\n", + " processed_path=\"\",\n", + " zarr_processed_path=\"\",\n", + " zarr_processing_error=\"\",\n", + " zarr_mp4_path=\"\",\n", + " mp4_path=\"\",\n", + " embodiment=\"aria\",\n", + " robot_name=\"\",\n", + " is_eval=False,\n", + " eval_score=0.0,\n", + " eval_success=False,\n", + ")\n", + "# Adding an episode\n", + "# add_episode(engine, episode)\n", + "\n", + "# Update Episode Test\n", + "# episode.operator = \"simar\"\n", + "# update_episode(engine, episode)\n", + "\n", + "# Get Table Row from Episode Hash\n", + "# episode_hash_to_table_row(engine, 123456)\n", + "\n", + "# Delete episodes by hash\n", + "# delete_episodes(engine, [123456, 123457])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fdad1d0f", + "metadata": {}, + "outputs": [], + "source": [ + "def drop_table(table_name):\n", + " with engine.connect() as connection:\n", + " connection.execute(text(f\"DROP TABLE IF EXISTS app.{table_name} CASCADE;\"))\n", + " connection.commit()\n", + " print(f\"Dropped table '{table_name}' from schema 'app' if it existed.\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "770a6964", + "metadata": {}, + "outputs": [], + "source": [ + "def create_table():\n", + " metadata = MetaData(schema=\"app\")\n", + "\n", + " Table(\n", + " \"episodes\",\n", + " metadata,\n", + " Column(\"episode_hash\", String, primary_key=True),\n", + " Column(\"operator\", String),\n", + " Column(\"lab\", String),\n", + " Column(\"num_frames\", Integer),\n", + " Column(\"task\", String),\n", + " Column(\"task_description\", String),\n", + " Column(\"scene\", String),\n", + " Column(\n", + " \"objects\", String\n", + " ), # Store as JSON or comma-separated list of object names\n", + " Column(\"processed_path\", String),\n", + " Column(\"zarr_processed_path\", String),\n", + " Column(\"zarr_processing_error\", String),\n", + " Column(\"zarr_mp4_path\", String),\n", + " Column(\"mp4_path\", String),\n", + " Column(\"is_deleted\", Boolean),\n", + " Column(\"embodiment\", String),\n", + " Column(\"robot_name\", String),\n", + " Column(\"is_eval\", Boolean),\n", + " Column(\"eval_score\", Float),\n", + " Column(\"eval_success\", Boolean),\n", + " )\n", + "\n", + " metadata.create_all(engine)\n", + " print(\"Created table 'episodes' in schema 'app'.\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cbb7d2a8", + "metadata": {}, + "outputs": [], + "source": [ + "def delete_episodes_by_task(task_name: str):\n", + " episodes_tbl = Table(\"episodes\", MetaData(), autoload_with=engine, schema=\"app\")\n", + " stmt = (\n", + " update(episodes_tbl)\n", + " .where(episodes_tbl.c.task == task_name)\n", + " .values(is_deleted=True)\n", + " )\n", + " with engine.begin() as conn:\n", + " conn.execute(stmt)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "emimic", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.15" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}