diff --git a/egomimic/hydra_configs/data/pusht/chain_gripper_newdata_h16_points_full.yaml b/egomimic/hydra_configs/data/pusht/chain_gripper_newdata_h16_points_full.yaml new file mode 100644 index 000000000..29ebd415d --- /dev/null +++ b/egomimic/hydra_configs/data/pusht/chain_gripper_newdata_h16_points_full.yaml @@ -0,0 +1,61 @@ +_target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper + +# The source root stays physically complete at 720 generated episodes. This +# recipe selects 719 by excluding the audited 3,118-frame idle-heavy episode. +train_datasets: + pushshapes_sim_chain_gripper: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + resolver: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.LocalEpisodeResolverManyWithEmbodimentOverride + folder_paths: &chain_roots + - /coc/flash7/paphiwetsa3/datasets/Tsim_v2/chain_gripper_3000_v2 + - /coc/flash7/paphiwetsa3/datasets/Tsim_v2/chain_gripper_gen + embodiment_override: pushshapes_sim_chain_gripper + key_map: + _target_: egomimic.rldb.embodiment.pushshapes.get_keymap_hpt + action_horizon: 16 + action_zarr_key: actions + transform_list: + _target_: egomimic.rldb.embodiment.pushshapes.get_chain_gripper_point_transform_list + world_size: 512.0 + filters: &chain_episode_filter + _target_: egomimic.rldb.filters.DatasetFilter + filter_lambdas: + - "lambda row: row.get('episode_hash') != 'episode_T_chain_gripper_obs7_000050'" + mode: train + valid_ratio: 0.0 + bounds_check: false + +valid_datasets: + pushshapes_sim_chain_gripper: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + resolver: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.LocalEpisodeResolverManyWithEmbodimentOverride + folder_paths: *chain_roots + embodiment_override: pushshapes_sim_chain_gripper + key_map: + _target_: egomimic.rldb.embodiment.pushshapes.get_keymap_hpt + action_horizon: 16 + action_zarr_key: actions + transform_list: + _target_: egomimic.rldb.embodiment.pushshapes.get_chain_gripper_point_transform_list + world_size: 512.0 + filters: *chain_episode_filter + mode: valid + valid_ratio: 0.02 + bounds_check: false + +# One A40 per BC run; batch 64 is the global batch for embodiment 20. +train_dataloader_params: + pushshapes_sim_chain_gripper: + batch_size: 64 + num_workers: 4 + pin_memory: true + persistent_workers: true + prefetch_factor: 2 + +valid_dataloader_params: + pushshapes_sim_chain_gripper: + batch_size: 16 + num_workers: 4 + pin_memory: true diff --git a/egomimic/hydra_configs/experiment/pusht/pipeline_diffusion_chain_gripper_newdata_points_h16.yaml b/egomimic/hydra_configs/experiment/pusht/pipeline_diffusion_chain_gripper_newdata_points_h16.yaml new file mode 100644 index 000000000..28898796e --- /dev/null +++ b/egomimic/hydra_configs/experiment/pusht/pipeline_diffusion_chain_gripper_newdata_points_h16.yaml @@ -0,0 +1,62 @@ +# @package _global_ +defaults: + - override /model: bf/bf_pipeline_diffusion_chain_gripper_points + - override /data: pusht/chain_gripper_newdata_h16_points_full + - override /evaluator: eval_pipeline_action_mse + - override /logger: wandb + +# ChainGripper-only BC over the frozen base-plus-generated corpus. +name: flow_transfer_bc_chain_newdata_dp_h16 +description: standard_dp_epsilon_h16_chain_3000v2_plus_gen +model: + enable_grad_norm: false + # Scale the complete original 1e-4 -> 1e-5 schedule by 0.3. + optimizer: + lr: 3.0e-5 + scheduler: + eta_min: 3.0e-6 + train_metrics_on_step: true + train_metrics_on_epoch: true +launch_params: + gpus_per_node: 1 + nodes: 1 +trainer: + max_steps: 240000 + max_epochs: -1 + min_epochs: null + limit_train_batches: 1.0 + val_check_interval: 10000 + limit_val_batches: 0 + check_val_every_n_epoch: 1 + num_sanity_val_steps: 0 + accumulate_grad_batches: 1 + log_every_n_steps: 1 +logger: + wandb: {project: pushshapes-flow-transfer} +norm_stats: + norm_mode: minmax + reduce_all_but_last: true +callbacks: + model_checkpoint: + every_n_epochs: null + every_n_train_steps: null + train_time_interval: {_target_: datetime.timedelta, hours: 1} + save_top_k: 1 + save_last: true + save_on_train_epoch_end: true + terminal_checkpoint: + _target_: lightning.pytorch.callbacks.ModelCheckpoint + dirpath: ${paths.output_dir}/checkpoints/final + filename: "step-{step}" + monitor: null + verbose: true + every_n_epochs: null + every_n_train_steps: ${trainer.max_steps} + train_time_interval: null + save_top_k: 1 + save_last: false + save_on_train_epoch_end: false + save_on_exception: false + save_weights_only: false + auto_insert_metric_name: false + enable_version_counter: false diff --git a/egomimic/hydra_configs/experiment/pusht/pipeline_sampler_chain_gripper_newdata_points_dense_medium_h16.yaml b/egomimic/hydra_configs/experiment/pusht/pipeline_sampler_chain_gripper_newdata_points_dense_medium_h16.yaml new file mode 100644 index 000000000..8266f25d5 --- /dev/null +++ b/egomimic/hydra_configs/experiment/pusht/pipeline_sampler_chain_gripper_newdata_points_dense_medium_h16.yaml @@ -0,0 +1,62 @@ +# @package _global_ +defaults: + - override /model: bf/bf_pipeline_sampler_chain_gripper_points_dense_medium_h16 + - override /data: pusht/chain_gripper_newdata_h16_points_full + - override /evaluator: eval_pipeline_action_mse + - override /logger: wandb + +# ChainGripper-only BC over the frozen base-plus-generated corpus. +name: flow_transfer_bc_chain_newdata_latent_dense_medium_h16 +description: decoder_only_latent96_dense16_chain_3000v2_plus_gen +model: + enable_grad_norm: false + # Scale the complete original 1e-4 -> 1e-5 schedule by 0.3. + optimizer: + lr: 3.0e-5 + scheduler: + eta_min: 3.0e-6 + train_metrics_on_step: true + train_metrics_on_epoch: true +launch_params: + gpus_per_node: 1 + nodes: 1 +trainer: + max_steps: 240000 + max_epochs: -1 + min_epochs: null + limit_train_batches: 1.0 + val_check_interval: 10000 + limit_val_batches: 0 + check_val_every_n_epoch: 1 + num_sanity_val_steps: 0 + accumulate_grad_batches: 1 + log_every_n_steps: 1 +logger: + wandb: {project: pushshapes-flow-transfer} +norm_stats: + norm_mode: minmax + reduce_all_but_last: true +callbacks: + model_checkpoint: + every_n_epochs: null + every_n_train_steps: null + train_time_interval: {_target_: datetime.timedelta, hours: 1} + save_top_k: 1 + save_last: true + save_on_train_epoch_end: true + terminal_checkpoint: + _target_: lightning.pytorch.callbacks.ModelCheckpoint + dirpath: ${paths.output_dir}/checkpoints/final + filename: "step-{step}" + monitor: null + verbose: true + every_n_epochs: null + every_n_train_steps: ${trainer.max_steps} + train_time_interval: null + save_top_k: 1 + save_last: false + save_on_train_epoch_end: false + save_on_exception: false + save_weights_only: false + auto_insert_metric_name: false + enable_version_counter: false diff --git a/egomimic/hydra_configs/model/bf/bf_pipeline_sampler_chain_gripper_points_dense_medium_h16.yaml b/egomimic/hydra_configs/model/bf/bf_pipeline_sampler_chain_gripper_points_dense_medium_h16.yaml new file mode 100644 index 000000000..fe71f5b8e --- /dev/null +++ b/egomimic/hydra_configs/model/bf/bf_pipeline_sampler_chain_gripper_points_dense_medium_h16.yaml @@ -0,0 +1,81 @@ +# Flow Transfer ChainGripper Medium direct-dense H16 latent denoiser. Native4 +# commands become ordered points6 in the loader; the action path is decoder-only. +_target_: egomimic.pl_utils.pl_model.ModelWrapper + +robomimic_model: + _target_: egomimic.pipeline.algo.PipelineAlgo + action_horizon: 16 + domains: [pushshapes_sim_chain_gripper] + ac_keys: {pushshapes_sim_chain_gripper: actions} + rollout_adapters: + pushshapes_sim_chain_gripper: + _target_: egomimic.pipeline.pushshapes.ChainGripperPointRolloutAdapter + action_horizon: 16 + stages: + - _target_: egomimic.pipeline.stages_sampler.FusedObsEncoder + n_obs_steps: 1 + encoder: + _target_: egomimic.pipeline.stages_sampler.DPStyleObsEncoder + obs_specs: + state_agent_obj: {input_dim: 3, input_slice: [0, 3]} + img_encoders: + front_img_1: + _target_: egomimic.models.stems.visual_core.VisualCore + in_channels: 3 + image_size: 96 + num_kp: 32 + feature_dimension: 64 + pretrained: false + crop_aug: true + crop_height: 84 + crop_width: 84 + crop_eval_mode: center + crop_sample_mode: v02 + crop_scope: frame + norm_layer: group + pool_type: spatial_softmax + - _target_: egomimic.pipeline.stages_sampler.GaussianLatentNoise + action_horizon: 16 + latent_dim: 96 + - _target_: egomimic.pipeline.stages_sampler.MultiJActionSampler + condition_input_dim: 67 + condition_dim: 384 + gradient_accumulation_steps: 1 + schedule_anchor_domain: pushshapes_sim_chain_gripper + action_horizon: 16 + action_dims: {pushshapes_sim_chain_gripper: 6} + latent_dim: 96 + decoder_hidden_dim: 512 + denoiser_hidden_dim: 384 + num_inference_steps: 16 + sampling_schedule: + 1: {1: 0.50, 2: 0.50} + 2001: {2: 0.80, 4: 0.15, 8: 0.05} + gradient_checkpointing: true + denoising_module: + _target_: egomimic.models.denoising_nets.CrossTransformer + nblocks: 16 + cond_dim: 384 + hidden_dim: 384 + act_dim: 96 + act_seq: 16 + n_heads: 8 + dropout: 0.1 + mlp_layers: 4 + mlp_ratio: 4 + time_conditioning: additive + - _target_: egomimic.pipeline.stages_sampler.NativeActionMSELoss + +enable_grad_norm: false +optimizer: + _target_: torch.optim.AdamW + _partial_: true + lr: 1.0e-4 + weight_decay: 1.0e-4 +scheduler: + _target_: egomimic.utils.schedulers.warmup_cosine_scheduler + _partial_: true + max_steps: 240000 + warmup_steps: 3000 + warmup_start_factor: 0.1 + eta_min: 1.0e-5 diff --git a/egomimic/scripts/verify_training_smoke.py b/egomimic/scripts/verify_training_smoke.py index 46d5b7135..daf5065ed 100644 --- a/egomimic/scripts/verify_training_smoke.py +++ b/egomimic/scripts/verify_training_smoke.py @@ -19,6 +19,8 @@ from wandb.proto import wandb_internal_pb2 from wandb.sdk.internal.datastore import DataStore +import egomimic.utils.hydra_resolvers # noqa: F401 + def _sha256(path: Path) -> str: digest = hashlib.sha256() @@ -157,6 +159,13 @@ def _has_required_metrics( return True +def _load_training_config(config_path: Path): + """Load a Hydra snapshot with the same resolvers used by trainHydra.""" + if not OmegaConf.has_resolver("eval"): + OmegaConf.register_new_resolver("eval", eval) + return OmegaConf.load(config_path) + + def verify_training_smoke( output_dir: Path, required_embodiments: list[int], @@ -166,7 +175,7 @@ def verify_training_smoke( output_dir = output_dir.resolve() config_path = output_dir / ".hydra" / "config.yaml" assert config_path.is_file(), config_path - config = OmegaConf.load(config_path) + config = _load_training_config(config_path) assert int(config.trainer.max_steps) == 2 assert int(config.trainer.limit_train_batches) == 2 diff --git a/scripts/train/flow_transfer_chain_bc_newdata_h16.sbatch b/scripts/train/flow_transfer_chain_bc_newdata_h16.sbatch new file mode 100644 index 000000000..233f48ab9 --- /dev/null +++ b/scripts/train/flow_transfer_chain_bc_newdata_h16.sbatch @@ -0,0 +1,800 @@ +#!/bin/bash +# Fail-closed ChainGripper-only BC launcher for the frozen base-plus-generated +# H16 corpus. The sbatch caller must request exactly one A40 on hoffman-lab. +# Full training is accepted only with an exact, immutable semantic smoke result. + +set -Eeuo pipefail + +ARM=${ARM:?set ARM=latent_h16 or dp} +MODE=${MODE:?set MODE=smoke or full} +EXPECTED_HEAD=${EXPECTED_HEAD:?set the exact clean 40-character commit SHA} +EXPECTED_LAUNCHER_SHA=${EXPECTED_LAUNCHER_SHA:?set the exact launcher SHA256} + +CHAIN_BASE_INVENTORY=${CHAIN_BASE_INVENTORY:?set the immutable base inventory path} +EXPECTED_CHAIN_BASE_INVENTORY_SHA=${EXPECTED_CHAIN_BASE_INVENTORY_SHA:?set its SHA256} +CHAIN_BASE_EPISODE_METADATA=${CHAIN_BASE_EPISODE_METADATA:?set the immutable base metadata TSV} +EXPECTED_CHAIN_BASE_EPISODE_METADATA_SHA=${EXPECTED_CHAIN_BASE_EPISODE_METADATA_SHA:?set its SHA256} +CHAIN_GEN_INVENTORY=${CHAIN_GEN_INVENTORY:?set the immutable generated inventory path} +EXPECTED_CHAIN_GEN_INVENTORY_SHA=${EXPECTED_CHAIN_GEN_INVENTORY_SHA:?set its SHA256} +CHAIN_GEN_EPISODE_METADATA=${CHAIN_GEN_EPISODE_METADATA:?set the immutable generated metadata TSV} +EXPECTED_CHAIN_GEN_EPISODE_METADATA_SHA=${EXPECTED_CHAIN_GEN_EPISODE_METADATA_SHA:?set its SHA256} + +NORM_ARTIFACT=${NORM_ARTIFACT:?set the exact shared H16 norm_stats.json path} +EXPECTED_NORM_SHA=${EXPECTED_NORM_SHA:?set the exact norm artifact SHA256} +NORM_VALIDATION=${NORM_VALIDATION:?set the exact shared H16 validation.json path} +EXPECTED_NORM_VALIDATION_SHA=${EXPECTED_NORM_VALIDATION_SHA:?set its SHA256} + +REPO=/coc/flash7/paphiwetsa3/worktrees/flow-transfer-chain-bc-newdata-h16-20260828 +PY_ENV=/coc/flash7/paphiwetsa3/projects/EgoVerse7/.venv +EXPECTED_ENV_SHA=8fd1504c955756adf8167f7bd34fc1a09cb844d268898b3ac30693e99ac60e87 +EXP_ROOT=/coc/flash7/paphiwetsa3/experiments/flow_transfer_chain_bc_newdata_h16_20260828 +LAUNCHER=$REPO/scripts/train/flow_transfer_chain_bc_newdata_h16.sbatch +VERIFIER=$REPO/egomimic/scripts/verify_training_smoke.py +CHAIN_BASE_DATA=/coc/flash7/paphiwetsa3/datasets/Tsim_v2/chain_gripper_3000_v2 +CHAIN_GEN_DATA=/coc/flash7/paphiwetsa3/datasets/Tsim_v2/chain_gripper_gen +EXCLUDED_EPISODE=episode_T_chain_gripper_obs7_000050 +EXCLUDED_FOLDER=${EXCLUDED_EPISODE}.zarr +EXPECTED_CHAIN_BASE_COUNT=3000 +EXPECTED_CHAIN_BASE_FRAMES=931061 +EXPECTED_CHAIN_GEN_COUNT=720 +EXPECTED_CHAIN_GEN_SOURCE_FRAMES=309709 +EXPECTED_CHAIN_GEN_EFFECTIVE_COUNT=719 +EXPECTED_CHAIN_GEN_EFFECTIVE_FRAMES=306591 +EXPECTED_EXCLUDED_FRAMES=3118 +EXPECTED_CHAIN_EFFECTIVE_COUNT=3719 +EXPECTED_CHAIN_EFFECTIVE_FRAMES=1237652 +EXPECTED_WORLD_SIZE=1 +EXPECTED_LOCAL_BATCH=64 +WANDB_PROJECT=pushshapes-flow-transfer +WANDB_GROUP=flow_transfer_chain_newdata3719_bc_h16_20260828 +SLURM_BIN=/opt/slurm/Ubuntu-20.04/24.11.0/bin +export PATH=$SLURM_BIN:$PATH +SRUN=$SLURM_BIN/srun + +case "$ARM" in + latent_h16) + EXPERIMENT=pusht/pipeline_sampler_chain_gripper_newdata_points_dense_medium_h16 + CONFIG_NAME=flow_transfer_bc_chain_newdata_latent_dense_medium_h16 + FULL_WANDB_ID=ft_bc_chain_newdata3719_latent_dense_m96_h16_s42_20260828 + ;; + dp) + EXPERIMENT=pusht/pipeline_diffusion_chain_gripper_newdata_points_h16 + CONFIG_NAME=flow_transfer_bc_chain_newdata_dp_h16 + FULL_WANDB_ID=ft_bc_chain_newdata3719_dp_h16_s42_20260828 + ;; + *) + printf 'Unknown ARM=%s; expected latent_h16 or dp\n' "$ARM" >&2 + exit 64 + ;; +esac + +case "$MODE" in + smoke) + test "${SLURM_RESTART_COUNT:-0}" = 0 + WANDB_ID=ft_bc_chain_newdata3719_${ARM}_smoke_job_${SLURM_JOB_ID:?} + WANDB_RESUME=never + RUN_DIR=$EXP_ROOT/smokes/$ARM/job_${SLURM_JOB_ID} + ;; + full) + SMOKE_RESULT=${SMOKE_RESULT:?full mode requires the exact semantic smoke result} + EXPECTED_SMOKE_RESULT_SHA=${EXPECTED_SMOKE_RESULT_SHA:?full mode requires its SHA256} + WANDB_ID=$FULL_WANDB_ID + if test "${SLURM_RESTART_COUNT:-0}" = 0; then + WANDB_RESUME=never + else + WANDB_RESUME=allow + fi + RUN_DIR=$EXP_ROOT/runs/$ARM + ;; + *) + printf 'Unknown MODE=%s; expected smoke or full\n' "$MODE" >&2 + exit 64 + ;; +esac + +is_sha256() { + [[ "$1" =~ ^[0-9a-f]{64}$ ]] +} + +is_git_sha() { + [[ "$1" =~ ^[0-9a-f]{40}$ ]] +} + +environment_manifest() { + "$PY_ENV/bin/python" - <<'PY' +import importlib.metadata +import platform + +print(f"python=={platform.python_version()}") +entries = set() +for distribution in importlib.metadata.distributions(): + name = distribution.metadata.get("Name") + if name: + entries.add(f"{name.lower().replace('_', '-')}=={distribution.version}") +for entry in sorted(entries): + print(entry) +PY +} + +validate_inventory() { + local label=$1 root=$2 inventory=$3 inventory_sha=$4 + local metadata=$5 metadata_sha=$6 expected_count=$7 expected_frames=$8 + + test -d "$root" + test -s "$inventory" + test -s "$metadata" + test "$(stat -c '%a' "$inventory")" = 444 + test "$(stat -c '%a' "$metadata")" = 444 + test "$(sha256sum "$inventory" | awk '{print $1}')" = "$inventory_sha" + test "$(sha256sum "$metadata" | awk '{print $1}')" = "$metadata_sha" + + "$PY_ENV/bin/python" - \ + "$label" "$root" "$inventory" "$metadata" \ + "$expected_count" "$expected_frames" <<'PY' +import hashlib +import json +import pathlib +import sys + +label, root_arg, inventory_arg, metadata_arg, count_arg, frames_arg = sys.argv[1:] +root = pathlib.Path(root_arg) +inventory = pathlib.Path(inventory_arg) +metadata = pathlib.Path(metadata_arg) +expected_count = int(count_arg) +expected_frames = int(frames_arg) + +episodes = sorted( + (path for path in root.iterdir() if path.is_dir() and path.name.endswith(".zarr")), + key=lambda path: path.name, +) +names = [path.name for path in episodes] +assert len(names) == expected_count, (label, len(names), expected_count) +assert len(names) == len(set(names)), (label, "duplicate episode names") +expected_inventory = "".join(f"{name}\n" for name in names).encode() +assert inventory.read_bytes() == expected_inventory, (label, "inventory drift") + +rows = [] +total_frames = 0 +for episode in episodes: + raw = (episode / "zarr.json").read_bytes() + payload = json.loads(raw) + frames = payload["attributes"]["total_frames"] + assert isinstance(frames, int) and not isinstance(frames, bool) and frames > 0 + total_frames += frames + rows.append(f"{episode.name}\t{frames}\t{hashlib.sha256(raw).hexdigest()}\n") +assert total_frames == expected_frames, (label, total_frames, expected_frames) +assert metadata.read_text() == "".join(rows), (label, "metadata drift") +PY +} + +validate_chain_contract() { + validate_inventory \ + chain_base "$CHAIN_BASE_DATA" \ + "$CHAIN_BASE_INVENTORY" "$EXPECTED_CHAIN_BASE_INVENTORY_SHA" \ + "$CHAIN_BASE_EPISODE_METADATA" \ + "$EXPECTED_CHAIN_BASE_EPISODE_METADATA_SHA" \ + "$EXPECTED_CHAIN_BASE_COUNT" "$EXPECTED_CHAIN_BASE_FRAMES" + validate_inventory \ + chain_gen "$CHAIN_GEN_DATA" \ + "$CHAIN_GEN_INVENTORY" "$EXPECTED_CHAIN_GEN_INVENTORY_SHA" \ + "$CHAIN_GEN_EPISODE_METADATA" \ + "$EXPECTED_CHAIN_GEN_EPISODE_METADATA_SHA" \ + "$EXPECTED_CHAIN_GEN_COUNT" "$EXPECTED_CHAIN_GEN_SOURCE_FRAMES" + + "$PY_ENV/bin/python" - \ + "$CHAIN_BASE_EPISODE_METADATA" "$CHAIN_GEN_EPISODE_METADATA" \ + "$EXCLUDED_FOLDER" "$EXPECTED_EXCLUDED_FRAMES" \ + "$EXPECTED_CHAIN_GEN_EFFECTIVE_COUNT" \ + "$EXPECTED_CHAIN_GEN_EFFECTIVE_FRAMES" \ + "$EXPECTED_CHAIN_EFFECTIVE_COUNT" \ + "$EXPECTED_CHAIN_EFFECTIVE_FRAMES" <<'PY' +import pathlib +import sys + +( + base_arg, + gen_arg, + excluded, + excluded_frames_arg, + gen_effective_count_arg, + gen_effective_frames_arg, + combined_count_arg, + combined_frames_arg, +) = sys.argv[1:] + + +def load(path_arg): + output = {} + for row in pathlib.Path(path_arg).read_text().splitlines(): + name, frames, metadata_sha = row.split("\t") + assert len(metadata_sha) == 64 + int(metadata_sha, 16) + output[name] = int(frames) + return output + + +base = load(base_arg) +gen = load(gen_arg) +assert excluded not in base +assert gen.get(excluded) == int(excluded_frames_arg) +effective_gen = {name: frames for name, frames in gen.items() if name != excluded} +assert len(effective_gen) == int(gen_effective_count_arg) +assert sum(effective_gen.values()) == int(gen_effective_frames_arg) +assert len(base) + len(effective_gen) == int(combined_count_arg) +assert sum(base.values()) + sum(effective_gen.values()) == int(combined_frames_arg) +PY +} + +validate_norm_contract() { + test -s "$NORM_ARTIFACT" + test -s "$NORM_VALIDATION" + test "$(stat -c '%a' "$NORM_ARTIFACT")" = 444 + test "$(stat -c '%a' "$NORM_VALIDATION")" = 444 + test "$(sha256sum "$NORM_ARTIFACT" | awk '{print $1}')" = "$EXPECTED_NORM_SHA" + test "$(sha256sum "$NORM_VALIDATION" | awk '{print $1}')" = "$EXPECTED_NORM_VALIDATION_SHA" + + "$PY_ENV/bin/python" - \ + "$NORM_ARTIFACT" "$NORM_VALIDATION" "$EXPECTED_NORM_SHA" \ + "$EXPECTED_CHAIN_BASE_FRAMES" "$EXPECTED_CHAIN_GEN_SOURCE_FRAMES" \ + "$EXPECTED_CHAIN_GEN_EFFECTIVE_FRAMES" "$EXPECTED_CHAIN_EFFECTIVE_COUNT" \ + "$EXPECTED_CHAIN_EFFECTIVE_FRAMES" "$EXCLUDED_EPISODE" <<'PY' +import hashlib +import json +import math +import pathlib +import sys + +( + artifact_arg, + validation_arg, + expected_sha, + base_frames_arg, + gen_source_frames_arg, + gen_effective_frames_arg, + effective_count_arg, + effective_frames_arg, + excluded_episode, +) = sys.argv[1:] +artifact_path = pathlib.Path(artifact_arg).resolve() +raw = artifact_path.read_bytes() +assert hashlib.sha256(raw).hexdigest() == expected_sha +artifact = json.loads(raw) +validation = json.loads(pathlib.Path(validation_arg).read_text()) + +effective_frames = int(effective_frames_arg) +assert validation["status"] == "PASS" +assert set(validation["shared_by"]) == {"dp_h16", "latent_dense_h16"} +assert validation["action_horizon"] == 16 +assert pathlib.Path(validation["artifact"]).resolve() == artifact_path +assert validation["sha256"] == expected_sha +assert validation["chain_effective_train_episodes"] == int(effective_count_arg) +assert validation["chain_gen_source_frames"] == int(gen_source_frames_arg) +assert validation["chain_gen_effective_train_episodes"] == 719 +assert validation["excluded_chain_gen_episode"] == excluded_episode +assert validation["per_source_train_frames"]["chain_base"] == int(base_frames_arg) +assert validation["per_source_train_frames"]["chain_gen"] == int(gen_effective_frames_arg) +assert validation["per_embodiment_train_frames"]["20"] == effective_frames +assert validation["sample_frac"] == 1.0 + +assert artifact["norm_mode"] == "minmax" +assert artifact["reduce_all_but_last"] is True +assert "20" in artifact["stats"] +metadata = artifact["norm_run_metadata"]["embodiments"]["20"] +assert metadata["dataset_size"] == effective_frames +assert metadata["sampled_frames"] == effective_frames +assert metadata["sample_frac"] == 1.0 +for key, width in {"state_agent_obj": 6, "actions": 6}.items(): + stats = artifact["stats"]["20"][key] + for values in stats.values(): + assert len(values) == width + assert all(math.isfinite(float(value)) for value in values) +PY +} + +validate_semantic_smoke() { + test -s "$SMOKE_RESULT" + test "$(stat -c '%a' "$SMOKE_RESULT")" = 444 + test "$(sha256sum "$SMOKE_RESULT" | awk '{print $1}')" = "$EXPECTED_SMOKE_RESULT_SHA" + + "$PY_ENV/bin/python" - \ + "$SMOKE_RESULT" "$EXP_ROOT" "$ARM" "$EXPERIMENT" "$CONFIG_NAME" \ + "$EXPECTED_HEAD" "$EXPECTED_LAUNCHER_SHA" "$EXPECTED_NORM_SHA" \ + "$EXPECTED_NORM_VALIDATION_SHA" "$FULL_WANDB_ID" \ + "$EXPECTED_CHAIN_BASE_INVENTORY_SHA" \ + "$EXPECTED_CHAIN_BASE_EPISODE_METADATA_SHA" \ + "$EXPECTED_CHAIN_GEN_INVENTORY_SHA" \ + "$EXPECTED_CHAIN_GEN_EPISODE_METADATA_SHA" <<'PY' +import hashlib +import json +import math +import pathlib +import sys + +( + result_arg, + root_arg, + arm, + experiment, + config_name, + expected_head, + launcher_sha, + norm_sha, + norm_validation_sha, + full_wandb_id, + base_inventory_sha, + base_metadata_sha, + gen_inventory_sha, + gen_metadata_sha, +) = sys.argv[1:] +result_path = pathlib.Path(result_arg).resolve() +expected_root = (pathlib.Path(root_arg) / "smokes" / arm).resolve() +result_path.relative_to(expected_root) +assert result_path.name == "SMOKE_RESULT.json" +result = json.loads(result_path.read_text()) + +assert result["status"] == "passed" +assert result["launcher_status"] == "PASS" +assert result["arm"] == arm +assert result["experiment"] == experiment +assert result["config_name"] == config_name +assert result["repo_head"] == expected_head +assert result["launcher_sha256"] == launcher_sha +assert result["norm_sha256"] == norm_sha +assert result["norm_validation_sha256"] == norm_validation_sha +assert result["full_wandb_id"] == full_wandb_id +assert result["world_size"] == 1 +assert result["gpu_model"] == "A40" +assert result["global_batch"] == 64 +assert result["precision"] == "bf16" +assert result["global_step"] == 2 +assert result["scheduler_last_epoch"] == 2 +assert result["required_embodiments"] == [20] +assert result["dense_training_steps"] == [0, 1] +assert result["validation_trainer_global_step"] >= 1 +assert result["wandb_exit_code"] == 0 +assert result["chain_effective_episodes"] == 3719 +assert result["chain_effective_frames"] == 1237652 +assert result["inventory_sha256"] == { + "chain_base": base_inventory_sha, + "chain_base_metadata": base_metadata_sha, + "chain_gen": gen_inventory_sha, + "chain_gen_metadata": gen_metadata_sha, +} +assert result["offline"] is True +assert result["wandb_project"] == "pushshapes-flow-transfer" + +for key in ("config", "checkpoint", "wandb_stream"): + path = pathlib.Path(result[key]) + assert path.is_file(), path + assert hashlib.sha256(path.read_bytes()).hexdigest() == result[f"{key}_sha256"] +assert all(math.isfinite(float(value)) for value in result["validation_metrics"].values()) +PY +} + +test -x "$PY_ENV/bin/python" +test -x "$SRUN" +test -x "$SLURM_BIN/scontrol" +test -f "$VERIFIER" +is_git_sha "$EXPECTED_HEAD" +for digest in \ + "$EXPECTED_LAUNCHER_SHA" \ + "$EXPECTED_CHAIN_BASE_INVENTORY_SHA" \ + "$EXPECTED_CHAIN_BASE_EPISODE_METADATA_SHA" \ + "$EXPECTED_CHAIN_GEN_INVENTORY_SHA" \ + "$EXPECTED_CHAIN_GEN_EPISODE_METADATA_SHA" \ + "$EXPECTED_NORM_SHA" \ + "$EXPECTED_NORM_VALIDATION_SHA"; do + is_sha256 "$digest" +done +if test "$MODE" = full; then + is_sha256 "$EXPECTED_SMOKE_RESULT_SHA" +fi + +test "$((EXPECTED_CHAIN_GEN_SOURCE_FRAMES - EXPECTED_CHAIN_GEN_EFFECTIVE_FRAMES))" = "$EXPECTED_EXCLUDED_FRAMES" +test "$((EXPECTED_CHAIN_BASE_FRAMES + EXPECTED_CHAIN_GEN_EFFECTIVE_FRAMES))" = "$EXPECTED_CHAIN_EFFECTIVE_FRAMES" +test "$((EXPECTED_CHAIN_BASE_COUNT + EXPECTED_CHAIN_GEN_EFFECTIVE_COUNT))" = "$EXPECTED_CHAIN_EFFECTIVE_COUNT" +test "$(git -C "$REPO" rev-parse HEAD)" = "$EXPECTED_HEAD" +test -z "$(git -C "$REPO" status --porcelain=v1 --untracked-files=all)" +test "$(sha256sum "$LAUNCHER" | awk '{print $1}')" = "$EXPECTED_LAUNCHER_SHA" +test "$(environment_manifest | sha256sum | awk '{print $1}')" = "$EXPECTED_ENV_SHA" + +test "${SLURM_JOB_NUM_NODES:?}" = 1 +test "${SLURM_NTASKS:?}" = "$EXPECTED_WORLD_SIZE" +test "${SLURM_CPUS_PER_TASK:?}" -ge 8 +test "${SLURM_JOB_PARTITION:?}" = hoffman-lab +test "${SLURM_JOB_ACCOUNT:?}" = hoffman-lab +SLURM_JOB_SUMMARY=$("$SLURM_BIN/scontrol" show job -o "$SLURM_JOB_ID") +if test "$MODE" = full; then + grep -Eq '(^| )TimeLimit=2-00:00:00( |$)' <<< "$SLURM_JOB_SUMMARY" +fi +test -n "${CUDA_VISIBLE_DEVICES:-}" +IFS=',' read -r -a visible_gpus <<< "$CUDA_VISIBLE_DEVICES" +test "${#visible_gpus[@]}" = "$EXPECTED_WORLD_SIZE" + +validate_chain_contract +validate_norm_contract +if test "$MODE" = full; then + validate_semantic_smoke +fi + +restart_count=${SLURM_RESTART_COUNT:-0} +if test "$MODE" = smoke || test "$restart_count" = 0; then + test ! -e "$RUN_DIR" +else + test -d "$RUN_DIR" +fi +mkdir -p "$RUN_DIR" +PROVENANCE_DIR=$RUN_DIR/provenance/restart_${restart_count} +test ! -e "$PROVENANCE_DIR" +mkdir -p "$PROVENANCE_DIR" +TRAIN_LOG=$RUN_DIR/training_restart_${restart_count}.log + +nvidia-smi -i "$CUDA_VISIBLE_DEVICES" \ + --query-gpu=index,name,uuid,memory.total,driver_version \ + --format=csv,noheader > "$PROVENANCE_DIR/gpu.txt" +test "$(wc -l < "$PROVENANCE_DIR/gpu.txt")" = 1 +test "$(grep -ciE '(^|[ ,])A40([ ,]|$)' "$PROVENANCE_DIR/gpu.txt")" = 1 + +environment_manifest > "$PROVENANCE_DIR/python_environment.txt" +sha256sum "$PROVENANCE_DIR/python_environment.txt" \ + > "$PROVENANCE_DIR/python_environment.sha256" +cp "$LAUNCHER" "$PROVENANCE_DIR/launcher.sbatch" +cp "$VERIFIER" "$PROVENANCE_DIR/verify_training_smoke.py" +cp "$NORM_ARTIFACT" "$PROVENANCE_DIR/norm_stats.json" +cp "$NORM_VALIDATION" "$PROVENANCE_DIR/norm_validation.json" +cp "$CHAIN_BASE_INVENTORY" "$PROVENANCE_DIR/chain_base_inventory.txt" +cp "$CHAIN_BASE_EPISODE_METADATA" "$PROVENANCE_DIR/chain_base_episode_metadata.tsv" +cp "$CHAIN_GEN_INVENTORY" "$PROVENANCE_DIR/chain_gen_inventory.txt" +cp "$CHAIN_GEN_EPISODE_METADATA" "$PROVENANCE_DIR/chain_gen_episode_metadata.tsv" +if test "$MODE" = full; then + cp "$SMOKE_RESULT" "$PROVENANCE_DIR/semantic_smoke_result.json" +fi +git -C "$REPO" log -1 --format=fuller > "$PROVENANCE_DIR/git_commit.txt" +git -C "$REPO" status --porcelain=v1 --untracked-files=all \ + > "$PROVENANCE_DIR/git_status.txt" +"$SLURM_BIN/scontrol" show job -dd "$SLURM_JOB_ID" \ + > "$PROVENANCE_DIR/slurm_job.txt" +sha256sum "$PROVENANCE_DIR"/* > "$PROVENANCE_DIR/artifacts.sha256" + +source "$PY_ENV/bin/activate" +cd "$REPO" +export PYTHONPATH=$REPO +export HYDRA_FULL_ERROR=1 +export PYTHONUNBUFFERED=1 +export WANDB_SILENT=true + +COMMON_OVERRIDES=( + --config-name=train_zarr_cartesian + "+experiment=$EXPERIMENT" + mode=train + "hydra.run.dir=$RUN_DIR" + "++paths.root_dir=$RUN_DIR" + "paths.output_dir=$RUN_DIR" + "paths.work_dir=$REPO" + ckpt_path=null + launch_params.gpus_per_node=1 + launch_params.nodes=1 + trainer.precision=bf16 + trainer.accumulate_grad_batches=1 + trainer.log_every_n_steps=1 + trainer.num_sanity_val_steps=0 + callbacks.model_checkpoint.every_n_epochs=null + callbacks.model_checkpoint.save_last=true + callbacks.model_checkpoint.save_top_k=1 + callbacks.model_checkpoint.save_on_train_epoch_end=false + logger.wandb.project=$WANDB_PROJECT + logger.wandb.entity=rl2-group + "logger.wandb.group=$WANDB_GROUP" + "logger.wandb.id=$WANDB_ID" + "++logger.wandb.name=$WANDB_ID" + "++logger.wandb.resume=$WANDB_RESUME" + norm_stats.norm_mode=minmax + norm_stats.reduce_all_but_last=true + norm_stats.sample_frac=1.0 + norm_stats.save_cache_dir=null + "norm_stats.precomputed_norm_path=$NORM_ARTIFACT" +) + +if test "$MODE" = smoke; then + export WANDB_MODE=offline + COMMON_OVERRIDES+=( + trainer.max_steps=2 + trainer.limit_train_batches=2 + trainer.val_check_interval=1 + trainer.limit_val_batches=1 + callbacks.model_checkpoint.every_n_train_steps=2 + callbacks.model_checkpoint.train_time_interval=null + callbacks.terminal_checkpoint.every_n_train_steps=1 + logger.wandb.offline=true + ) +else + unset WANDB_MODE + COMMON_OVERRIDES+=( + trainer.max_steps=240000 + trainer.limit_train_batches=1.0 + trainer.val_check_interval=10000 + trainer.limit_val_batches=0 + callbacks.model_checkpoint.every_n_train_steps=null + logger.wandb.offline=false + ) +fi + +if test "$MODE" = full && test "$restart_count" != 0; then + shopt -s nullglob + resume_candidates=("$RUN_DIR"/hpc_ckpt_*.ckpt) + shopt -u nullglob + if test -s "$RUN_DIR/checkpoints/last.ckpt"; then + resume_candidates+=("$RUN_DIR/checkpoints/last.ckpt") + fi + if test "${#resume_candidates[@]}" = 0; then + printf 'Refusing restart without a recovery checkpoint in %s\n' "$RUN_DIR" >&2 + exit 66 + fi + RESUME_CKPT=$("$PY_ENV/bin/python" - "${resume_candidates[@]}" <<'PY' +import pathlib +import sys + +paths = [pathlib.Path(raw) for raw in sys.argv[1:]] +for path in paths: + assert path.is_file() and path.stat().st_size > 0, path +print(max(paths, key=lambda path: (path.stat().st_mtime_ns, str(path)))) +PY +) + COMMON_OVERRIDES+=("ckpt_path=$RESUME_CKPT") +fi + +RESOLVED_CONFIG=$PROVENANCE_DIR/resolved_config.yaml +python -m egomimic.trainHydra "${COMMON_OVERRIDES[@]}" \ + --cfg job --resolve > "$RESOLVED_CONFIG" + +python - \ + "$RESOLVED_CONFIG" "$ARM" "$MODE" "$EXPERIMENT" "$CONFIG_NAME" \ + "$NORM_ARTIFACT" "$RUN_DIR" "$WANDB_ID" "$WANDB_RESUME" \ + "$EXCLUDED_EPISODE" <<'PY' +import pathlib +import sys + +from hydra.utils import instantiate +from omegaconf import OmegaConf + +( + path, + arm, + mode, + experiment, + config_name, + norm_artifact, + run_dir, + wandb_id, + wandb_resume, + excluded_episode, +) = sys.argv[1:] +cfg = OmegaConf.load(path) +model = cfg.model.robomimic_model +expected_filter = ( + "lambda row: row.get('episode_hash') != " f"'{excluded_episode}'" +) + +assert cfg.name == config_name +assert cfg.seed == 42 +assert cfg.mode == "train" +assert str(cfg.paths.root_dir) == run_dir +assert str(cfg.paths.output_dir) == run_dir +assert cfg.trainer.precision == "bf16" +assert cfg.trainer.strategy == "ddp" +assert cfg.trainer.devices == 1 +assert cfg.trainer.num_nodes == 1 +assert cfg.trainer.accumulate_grad_batches == 1 +assert cfg.trainer.log_every_n_steps == 1 +assert cfg.trainer.num_sanity_val_steps == 0 +assert cfg.trainer.get("gradient_clip_val") is None +assert cfg.launch_params.gpus_per_node == 1 +assert cfg.launch_params.nodes == 1 +assert cfg.model.enable_grad_norm is False +assert cfg.model.train_metrics_on_step is True +assert cfg.model.train_metrics_on_epoch is True +assert cfg.model.optimizer.lr == 3.0e-5 +assert cfg.model.optimizer.weight_decay == 1.0e-4 +assert cfg.model.scheduler.max_steps == 240_000 +assert cfg.model.scheduler.warmup_steps == 3_000 +assert cfg.model.scheduler.warmup_start_factor == 0.1 +assert cfg.model.scheduler.eta_min == 3.0e-6 +assert cfg.logger.wandb.project == "pushshapes-flow-transfer" +assert cfg.logger.wandb.group == "flow_transfer_chain_newdata3719_bc_h16_20260828" +assert cfg.logger.wandb.id == wandb_id +assert cfg.logger.wandb.resume == wandb_resume +assert cfg.norm_stats.norm_mode == "minmax" +assert cfg.norm_stats.reduce_all_but_last is True +assert cfg.norm_stats.sample_frac == 1.0 +assert cfg.norm_stats.save_cache_dir is None +assert str(cfg.norm_stats.precomputed_norm_path) == norm_artifact +assert model.action_horizon == 16 +assert list(model.domains) == ["pushshapes_sim_chain_gripper"] +assert set(cfg.data.train_datasets) == {"pushshapes_sim_chain_gripper"} +assert set(cfg.data.valid_datasets) == {"pushshapes_sim_chain_gripper"} +assert cfg.data.train_dataloader_params.pushshapes_sim_chain_gripper.batch_size == 64 +assert cfg.data.valid_dataloader_params.pushshapes_sim_chain_gripper.batch_size == 16 +assert cfg.evaluator._target_.endswith("HumanRobotOverlayEval") + +for split_name, expected_mode, expected_ratio in ( + ("train_datasets", "train", 0.0), + ("valid_datasets", "valid", 0.02), +): + dataset = cfg.data[split_name].pushshapes_sim_chain_gripper + assert dataset.mode == expected_mode + assert float(dataset.valid_ratio) == expected_ratio + assert list(dataset.resolver.folder_paths) == [ + "/coc/flash7/paphiwetsa3/datasets/Tsim_v2/chain_gripper_3000_v2", + "/coc/flash7/paphiwetsa3/datasets/Tsim_v2/chain_gripper_gen", + ] + assert dataset.resolver.key_map.action_horizon == 16 + assert dataset.resolver.key_map.action_zarr_key == "actions" + assert dataset.resolver.transform_list._target_.endswith( + "get_chain_gripper_point_transform_list" + ) + assert float(dataset.resolver.transform_list.world_size) == 512.0 + assert list(dataset.filters.filter_lambdas) == [expected_filter] + +checkpoint = cfg.callbacks.model_checkpoint +terminal = cfg.callbacks.terminal_checkpoint +assert checkpoint.save_last is True +assert checkpoint.save_top_k == 1 +assert checkpoint.save_on_train_epoch_end is False +assert terminal.save_top_k == 1 +assert terminal.save_last is False +assert instantiate(checkpoint).state_key != instantiate(terminal).state_key + +if mode == "smoke": + assert cfg.trainer.max_steps == 2 + assert cfg.trainer.limit_train_batches == 2 + assert cfg.trainer.val_check_interval == 1 + assert cfg.trainer.limit_val_batches == 1 + assert checkpoint.every_n_train_steps == 2 + assert checkpoint.train_time_interval is None + assert terminal.every_n_train_steps == 1 + assert cfg.logger.wandb.offline is True +else: + assert mode == "full" + assert cfg.trainer.max_steps == 240_000 + assert cfg.trainer.limit_train_batches == 1.0 + assert cfg.trainer.val_check_interval == 10_000 + assert cfg.trainer.limit_val_batches == 0 + assert checkpoint.every_n_train_steps is None + assert checkpoint.train_time_interval._target_ == "datetime.timedelta" + assert checkpoint.train_time_interval.hours == 1 + assert terminal.every_n_train_steps == cfg.trainer.max_steps + assert cfg.logger.wandb.offline is False + +resolved = pathlib.Path(path).read_text().lower() +assert "arc_length" not in resolved +assert "action_encoder" not in resolved +if arm == "latent_h16": + assert model.stages[1].action_horizon == 16 + assert model.stages[1].latent_dim == 96 + sampler = model.stages[2] + assert sampler.action_horizon == 16 + assert sampler.action_dims.pushshapes_sim_chain_gripper == 6 + assert sampler.latent_dim == 96 + assert sampler.decoder_hidden_dim == 512 + assert sampler.denoiser_hidden_dim == 384 + assert sampler.denoising_module.act_seq == 16 + assert sampler.denoising_module.nblocks == 16 +else: + assert arm == "dp" + stage = model.stages[1] + assert stage.action_horizon == 16 + policy = stage.policies.pushshapes_sim_chain_gripper + assert policy.action_horizon == 16 + assert policy.infer_ac_dims.pushshapes_sim_chain_gripper == 6 + assert policy.noise_scheduler.prediction_type == "epsilon" +PY + +"$SRUN" --ntasks=1 --gpus-per-task=1 --kill-on-bad-exit=1 --unbuffered \ + python -m egomimic.trainHydra "${COMMON_OVERRIDES[@]}" \ + 2>&1 | tee "$TRAIN_LOG" + +if grep -E 'Traceback|CUDA out of memory|OutOfMemoryError|ProcessExitedException' \ + "$TRAIN_LOG"; then + printf 'Fatal signature found in combined training stdout/stderr\n' >&2 + exit 70 +fi + +validate_chain_contract +validate_norm_contract +test "$(git -C "$REPO" rev-parse HEAD)" = "$EXPECTED_HEAD" +test -z "$(git -C "$REPO" status --porcelain=v1 --untracked-files=all)" +test "$(sha256sum "$LAUNCHER" | awk '{print $1}')" = "$EXPECTED_LAUNCHER_SHA" + +if test "$MODE" = smoke; then + python "$VERIFIER" "$RUN_DIR" \ + --required-embodiments 20 \ + --expected-head "$EXPECTED_HEAD" \ + --expected-world-size "$EXPECTED_WORLD_SIZE" \ + --dry-run + python "$VERIFIER" "$RUN_DIR" \ + --required-embodiments 20 \ + --expected-head "$EXPECTED_HEAD" \ + --expected-world-size "$EXPECTED_WORLD_SIZE" + + python - \ + "$RUN_DIR/SMOKE_RESULT.json" "$ARM" "$EXPERIMENT" "$CONFIG_NAME" \ + "$EXPECTED_LAUNCHER_SHA" "$EXPECTED_NORM_SHA" \ + "$EXPECTED_NORM_VALIDATION_SHA" "$FULL_WANDB_ID" \ + "$EXPECTED_CHAIN_BASE_INVENTORY_SHA" \ + "$EXPECTED_CHAIN_BASE_EPISODE_METADATA_SHA" \ + "$EXPECTED_CHAIN_GEN_INVENTORY_SHA" \ + "$EXPECTED_CHAIN_GEN_EPISODE_METADATA_SHA" <<'PY' +import json +import pathlib +import sys + +( + result_arg, + arm, + experiment, + config_name, + launcher_sha, + norm_sha, + norm_validation_sha, + full_wandb_id, + base_inventory_sha, + base_metadata_sha, + gen_inventory_sha, + gen_metadata_sha, +) = sys.argv[1:] +path = pathlib.Path(result_arg) +payload = json.loads(path.read_text()) +payload.update( + { + "launcher_status": "PASS", + "arm": arm, + "experiment": experiment, + "config_name": config_name, + "launcher_sha256": launcher_sha, + "norm_sha256": norm_sha, + "norm_validation_sha256": norm_validation_sha, + "full_wandb_id": full_wandb_id, + "offline": True, + "wandb_project": "pushshapes-flow-transfer", + "gpu_model": "A40", + "global_batch": 64, + "chain_effective_episodes": 3719, + "chain_effective_frames": 1237652, + "excluded_episode": "episode_T_chain_gripper_obs7_000050", + "inventory_sha256": { + "chain_base": base_inventory_sha, + "chain_base_metadata": base_metadata_sha, + "chain_gen": gen_inventory_sha, + "chain_gen_metadata": gen_metadata_sha, + }, + } +) +temporary = path.with_suffix(".json.tmp") +temporary.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") +temporary.replace(path) +PY + chmod 0444 "$RUN_DIR/SMOKE_RESULT.json" + SMOKE_SHA=$(sha256sum "$RUN_DIR/SMOKE_RESULT.json" | awk '{print $1}') + printf '[smoke] ALL_PASS arm=%s result=%s sha256=%s\n' \ + "$ARM" "$RUN_DIR/SMOKE_RESULT.json" "$SMOKE_SHA" +else + python - "$RUN_DIR" <<'PY' +import pathlib +import sys + +import torch + +run_dir = pathlib.Path(sys.argv[1]) +checkpoints = list((run_dir / "checkpoints" / "final").glob("*.ckpt")) +assert checkpoints, "terminal checkpoint is missing" +matches = [] +for path in checkpoints: + payload = torch.load(path, map_location="cpu", weights_only=False) + if int(payload.get("global_step", -1)) == 240_000: + matches.append(path) +assert len(matches) == 1, matches +print(f"[full] terminal checkpoint: {matches[0]}") +PY + printf '[full] ALL_PASS arm=%s head=%s run=%s\n' \ + "$ARM" "$EXPECTED_HEAD" "$RUN_DIR" +fi diff --git a/tests/test_flow_transfer_chain_bc_newdata_h16.py b/tests/test_flow_transfer_chain_bc_newdata_h16.py new file mode 100644 index 000000000..44654be2b --- /dev/null +++ b/tests/test_flow_transfer_chain_bc_newdata_h16.py @@ -0,0 +1,168 @@ +from pathlib import Path + +import pytest +from hydra import compose, initialize_config_dir +from hydra.core.global_hydra import GlobalHydra +from hydra.utils import instantiate +from omegaconf import OmegaConf + +CONFIG_DIR = Path(__file__).parents[1] / "egomimic" / "hydra_configs" +CHAIN_DOMAIN = "pushshapes_sim_chain_gripper" +CHAIN_ROOTS = [ + "/coc/flash7/paphiwetsa3/datasets/Tsim_v2/chain_gripper_3000_v2", + "/coc/flash7/paphiwetsa3/datasets/Tsim_v2/chain_gripper_gen", +] +CHAIN_FILTER = ( + "lambda row: row.get('episode_hash') != 'episode_T_chain_gripper_obs7_000050'" +) +BC_EXPERIMENTS = [ + "pusht/pipeline_sampler_chain_gripper_newdata_points_dense_medium_h16", + "pusht/pipeline_diffusion_chain_gripper_newdata_points_h16", +] + + +def _compose(experiment: str, extra_overrides: list[str] | None = None): + if GlobalHydra.instance().is_initialized(): + GlobalHydra.instance().clear() + with initialize_config_dir(config_dir=str(CONFIG_DIR), version_base="1.3"): + return compose( + config_name="train_zarr_cartesian", + overrides=[f"+experiment={experiment}", *(extra_overrides or [])], + ) + + +@pytest.mark.parametrize("experiment", BC_EXPERIMENTS) +@pytest.mark.parametrize( + ("mode", "overrides"), + [ + ( + "smoke", + [ + "trainer.max_steps=2", + "callbacks.model_checkpoint.every_n_train_steps=2", + "callbacks.model_checkpoint.train_time_interval=null", + "callbacks.model_checkpoint.save_on_train_epoch_end=false", + "callbacks.terminal_checkpoint.every_n_train_steps=1", + ], + ), + ( + "full", + [ + "trainer.max_steps=240000", + "callbacks.model_checkpoint.every_n_train_steps=null", + "callbacks.model_checkpoint.save_on_train_epoch_end=false", + ], + ), + ], +) +def test_chain_bc_checkpoint_state_keys_are_unique( + experiment: str, mode: str, overrides: list[str] +) -> None: + cfg = _compose(experiment, overrides) + cfg.paths.output_dir = "/tmp/flow_transfer_chain_bc_callback_state_keys" + checkpoint = cfg.callbacks.model_checkpoint + terminal = cfg.callbacks.terminal_checkpoint + + assert instantiate(checkpoint).state_key != instantiate(terminal).state_key + if mode == "smoke": + assert checkpoint.every_n_train_steps == 2 + assert checkpoint.train_time_interval is None + assert terminal.every_n_train_steps == 1 + else: + assert mode == "full" + assert checkpoint.every_n_train_steps is None + assert checkpoint.train_time_interval.hours == 1 + assert terminal.every_n_train_steps == cfg.trainer.max_steps == 240_000 + + +@pytest.mark.parametrize("experiment", BC_EXPERIMENTS) +def test_chain_bc_newdata_contract_is_world1_batch64_h16(experiment: str) -> None: + cfg = _compose(experiment) + model = cfg.model.robomimic_model + + assert cfg.launch_params.gpus_per_node == 1 + assert cfg.launch_params.nodes == 1 + assert cfg.trainer.precision == "bf16" + assert cfg.trainer.accumulate_grad_batches == 1 + assert cfg.trainer.log_every_n_steps == 1 + assert cfg.trainer.limit_val_batches == 0 + assert cfg.trainer.get("gradient_clip_val") is None + assert cfg.model.enable_grad_norm is False + assert cfg.model.train_metrics_on_step is True + assert cfg.model.train_metrics_on_epoch is True + assert cfg.model.optimizer.lr == pytest.approx(3.0e-5) + assert cfg.model.optimizer.weight_decay == pytest.approx(1.0e-4) + assert cfg.model.scheduler.warmup_steps == 3_000 + assert cfg.model.scheduler.max_steps == 240_000 + assert cfg.model.scheduler.warmup_start_factor == pytest.approx(0.1) + assert cfg.model.scheduler.eta_min == pytest.approx(3.0e-6) + + assert list(model.domains) == [CHAIN_DOMAIN] + assert model.action_horizon == 16 + assert set(cfg.data.train_datasets) == {CHAIN_DOMAIN} + assert set(cfg.data.valid_datasets) == {CHAIN_DOMAIN} + assert cfg.data.train_dataloader_params[CHAIN_DOMAIN].batch_size == 64 + assert cfg.data.valid_dataloader_params[CHAIN_DOMAIN].batch_size == 16 + + for split_name, mode, ratio in ( + ("train_datasets", "train", 0.0), + ("valid_datasets", "valid", 0.02), + ): + dataset = cfg.data[split_name][CHAIN_DOMAIN] + assert dataset.mode == mode + assert float(dataset.valid_ratio) == ratio + assert list(dataset.resolver.folder_paths) == CHAIN_ROOTS + assert dataset.resolver.key_map.action_horizon == 16 + assert dataset.resolver.key_map.action_zarr_key == "actions" + assert dataset.resolver.transform_list._target_.endswith( + "get_chain_gripper_point_transform_list" + ) + assert float(dataset.resolver.transform_list.world_size) == 512.0 + assert list(dataset.filters.filter_lambdas) == [CHAIN_FILTER] + + +def test_chain_bc_data_subtree_matches_cotrain_chain_subtree() -> None: + cotrain = _compose("pusht/pipeline_diffusion_usocket_chain_newdata_h16") + + for experiment in BC_EXPERIMENTS: + bc = _compose(experiment) + for split_name in ("train_datasets", "valid_datasets"): + assert OmegaConf.to_container( + bc.data[split_name][CHAIN_DOMAIN], resolve=True + ) == OmegaConf.to_container( + cotrain.data[split_name][CHAIN_DOMAIN], resolve=True + ) + + +def test_chain_bc_latent_is_decoder_only_medium_h16() -> None: + cfg = _compose( + "pusht/pipeline_sampler_chain_gripper_newdata_points_dense_medium_h16" + ) + model = cfg.model.robomimic_model + + assert model.stages[1].action_horizon == 16 + assert model.stages[1].latent_dim == 96 + sampler = model.stages[2] + assert sampler.action_horizon == 16 + assert sampler.action_dims[CHAIN_DOMAIN] == 6 + assert sampler.latent_dim == 96 + assert sampler.decoder_hidden_dim == 512 + assert sampler.denoiser_hidden_dim == 384 + assert sampler.denoising_module.hidden_dim == 384 + assert sampler.denoising_module.act_dim == 96 + assert sampler.denoising_module.act_seq == 16 + assert sampler.denoising_module.nblocks == 16 + assert "action_encoder" not in OmegaConf.to_yaml(cfg.model).lower() + + +def test_chain_bc_dp_is_epsilon_h16() -> None: + cfg = _compose("pusht/pipeline_diffusion_chain_gripper_newdata_points_h16") + model = cfg.model.robomimic_model + stage = model.stages[1] + policy = stage.policies[CHAIN_DOMAIN] + + assert stage.action_horizon == 16 + assert policy.action_horizon == 16 + assert policy.infer_ac_dims[CHAIN_DOMAIN] == 6 + assert policy.model.input_dim == 6 + assert policy.noise_scheduler.prediction_type == "epsilon" diff --git a/tests/test_flow_transfer_chain_bc_newdata_launcher.py b/tests/test_flow_transfer_chain_bc_newdata_launcher.py new file mode 100644 index 000000000..ab0b75db9 --- /dev/null +++ b/tests/test_flow_transfer_chain_bc_newdata_launcher.py @@ -0,0 +1,163 @@ +from pathlib import Path +from re import DOTALL, findall +from subprocess import run + +from omegaconf import OmegaConf + +from egomimic.scripts.verify_training_smoke import _load_training_config + +LAUNCHER = ( + Path(__file__).parents[1] + / "scripts" + / "train" + / "flow_transfer_chain_bc_newdata_h16.sbatch" +) + + +def test_smoke_verifier_resolves_trainhydra_eval_interpolation(tmp_path: Path) -> None: + config_path = tmp_path / "config.yaml" + config_path.write_text( + """ +launch_params: + gpus_per_node: 1 + nodes: 1 +trainer: + devices: ${eval:'${launch_params.gpus_per_node} * ${launch_params.nodes}'} +""".lstrip() + ) + OmegaConf.clear_resolver("eval") + config = _load_training_config(config_path) + assert int(config.trainer.devices) == 1 + + +def test_chain_bc_newdata_launcher_has_valid_bash_syntax() -> None: + run(["bash", "-n", str(LAUNCHER)], check=True) + + +def test_chain_bc_newdata_launcher_embedded_python_compiles() -> None: + blocks = findall(r"<<'PY'\n(.*?)\nPY", LAUNCHER.read_text(), flags=DOTALL) + assert len(blocks) >= 8 + for index, block in enumerate(blocks): + compile(block, f"{LAUNCHER.name}:heredoc-{index}", "exec") + + +def test_chain_bc_newdata_launcher_pins_arms_resources_and_data() -> None: + launcher = LAUNCHER.read_text() + + for contract in ( + "ARM=${ARM:?set ARM=latent_h16 or dp}", + "MODE=${MODE:?set MODE=smoke or full}", + "EXPECTED_HEAD=${EXPECTED_HEAD:?", + "EXPECTED_LAUNCHER_SHA=${EXPECTED_LAUNCHER_SHA:?", + "CHAIN_BASE_INVENTORY=${CHAIN_BASE_INVENTORY:?", + "EXPECTED_CHAIN_BASE_INVENTORY_SHA=${EXPECTED_CHAIN_BASE_INVENTORY_SHA:?", + "CHAIN_BASE_EPISODE_METADATA=${CHAIN_BASE_EPISODE_METADATA:?", + "EXPECTED_CHAIN_BASE_EPISODE_METADATA_SHA=${EXPECTED_CHAIN_BASE_EPISODE_METADATA_SHA:?", + "CHAIN_GEN_INVENTORY=${CHAIN_GEN_INVENTORY:?", + "EXPECTED_CHAIN_GEN_INVENTORY_SHA=${EXPECTED_CHAIN_GEN_INVENTORY_SHA:?", + "CHAIN_GEN_EPISODE_METADATA=${CHAIN_GEN_EPISODE_METADATA:?", + "EXPECTED_CHAIN_GEN_EPISODE_METADATA_SHA=${EXPECTED_CHAIN_GEN_EPISODE_METADATA_SHA:?", + "NORM_ARTIFACT=${NORM_ARTIFACT:?", + "EXPECTED_NORM_SHA=${EXPECTED_NORM_SHA:?", + "NORM_VALIDATION=${NORM_VALIDATION:?", + "EXPECTED_NORM_VALIDATION_SHA=${EXPECTED_NORM_VALIDATION_SHA:?", + ): + assert contract in launcher + + assert "latent_h16)" in launcher + assert "dp)" in launcher + assert "pipeline_sampler_chain_gripper_newdata_points_dense_medium_h16" in launcher + assert "pipeline_diffusion_chain_gripper_newdata_points_h16" in launcher + assert 'test "${SLURM_JOB_PARTITION:?}" = hoffman-lab' in launcher + assert 'test "${SLURM_JOB_ACCOUNT:?}" = hoffman-lab' in launcher + assert "TimeLimit=2-00:00:00" in launcher + assert 'if test "$MODE" = full; then' in launcher + assert "EXPECTED_WORLD_SIZE=1" in launcher + assert "EXPECTED_LOCAL_BATCH=64" in launcher + assert "--ntasks=1 --gpus-per-task=1" in launcher + assert "grep -ciE '(^|[ ,])A40([ ,]|$)'" in launcher + + for value in ( + "EXPECTED_CHAIN_BASE_COUNT=3000", + "EXPECTED_CHAIN_BASE_FRAMES=931061", + "EXPECTED_CHAIN_GEN_COUNT=720", + "EXPECTED_CHAIN_GEN_SOURCE_FRAMES=309709", + "EXPECTED_CHAIN_GEN_EFFECTIVE_COUNT=719", + "EXPECTED_CHAIN_GEN_EFFECTIVE_FRAMES=306591", + "EXPECTED_EXCLUDED_FRAMES=3118", + "EXPECTED_CHAIN_EFFECTIVE_COUNT=3719", + "EXPECTED_CHAIN_EFFECTIVE_FRAMES=1237652", + "EXCLUDED_EPISODE=episode_T_chain_gripper_obs7_000050", + ): + assert value in launcher + assert launcher.count("validate_chain_contract") >= 3 + assert "inventory drift" in launcher + assert "metadata drift" in launcher + + +def test_chain_bc_newdata_launcher_pins_training_and_logging_contract() -> None: + launcher = LAUNCHER.read_text() + + assert "WANDB_PROJECT=pushshapes-flow-transfer" in launcher + assert "WANDB_GROUP=flow_transfer_chain_newdata3719_bc_h16_20260828" in launcher + assert "ft_bc_chain_newdata3719_latent_dense_m96_h16_s42_20260828" in launcher + assert "ft_bc_chain_newdata3719_dp_h16_s42_20260828" in launcher + assert "WANDB_RESUME=never" in launcher + assert "WANDB_RESUME=allow" in launcher + assert "trainer.precision=bf16" in launcher + assert "trainer.accumulate_grad_batches=1" in launcher + assert "trainer.log_every_n_steps=1" in launcher + assert "trainer.max_steps=240000" in launcher + assert "trainer.limit_val_batches=0" in launcher + assert "callbacks.model_checkpoint.every_n_train_steps=null" in launcher + assert "checkpoint.train_time_interval.hours == 1" in launcher + assert "callbacks.terminal_checkpoint.every_n_train_steps=1" in launcher + assert ( + "instantiate(checkpoint).state_key != instantiate(terminal).state_key" + in launcher + ) + assert "terminal.every_n_train_steps == cfg.trainer.max_steps" in launcher + assert "cfg.model.optimizer.lr == 3.0e-5" in launcher + assert "cfg.model.scheduler.warmup_steps == 3_000" in launcher + assert "cfg.model.scheduler.eta_min == 3.0e-6" in launcher + assert "cfg.model.train_metrics_on_step is True" in launcher + assert "cfg.model.train_metrics_on_epoch is True" in launcher + assert "gradient_clip_val" in launcher + assert "action_encoder" in launcher + assert "arc_length" in launcher + + +def test_chain_bc_newdata_launcher_smoke_and_full_are_fail_closed() -> None: + launcher = LAUNCHER.read_text() + + assert "trainer.max_steps=2" in launcher + assert "trainer.limit_train_batches=2" in launcher + assert "trainer.val_check_interval=1" in launcher + assert "trainer.limit_val_batches=1" in launcher + assert "trainer.num_sanity_val_steps=0" in launcher + assert "assert terminal.every_n_train_steps == 1" in launcher + assert "--required-embodiments 20" in launcher + assert '--expected-world-size "$EXPECTED_WORLD_SIZE"' in launcher + assert "--dry-run" in launcher + assert "verify_training_smoke.py" in launcher + assert "WANDB_MODE=offline" in launcher + assert "logger.wandb.offline=true" in launcher + assert "Fatal signature found in combined training stdout/stderr" in launcher + + assert "SMOKE_RESULT=${SMOKE_RESULT:?" in launcher + assert "EXPECTED_SMOKE_RESULT_SHA=${EXPECTED_SMOKE_RESULT_SHA:?" in launcher + assert "stat -c '%a' \"$SMOKE_RESULT\"" in launcher + assert "validate_semantic_smoke" in launcher + assert 'result["launcher_status"] == "PASS"' in launcher + assert 'result["dense_training_steps"] == [0, 1]' in launcher + assert 'result["validation_trainer_global_step"] >= 1' in launcher + assert 'result["required_embodiments"] == [20]' in launcher + assert 'result["gpu_model"] == "A40"' in launcher + assert 'result["global_batch"] == 64' in launcher + assert 'result["chain_effective_episodes"] == 3719' in launcher + assert 'result["chain_effective_frames"] == 1237652' in launcher + assert 'chmod 0444 "$RUN_DIR/SMOKE_RESULT.json"' in launcher + assert "terminal checkpoint is missing" in launcher + + # This file executes inside an allocation. It must not submit another job. + assert "\nsbatch " not in launcher diff --git a/tests/test_flow_transfer_newdata_scheduler.py b/tests/test_flow_transfer_newdata_scheduler.py index 3db54b09a..c6091277f 100644 --- a/tests/test_flow_transfer_newdata_scheduler.py +++ b/tests/test_flow_transfer_newdata_scheduler.py @@ -12,6 +12,8 @@ NEWDATA_H16_EXPERIMENTS = [ "pusht/pipeline_sampler_usocket_chain_newdata_dense_medium_h16", "pusht/pipeline_diffusion_usocket_chain_newdata_h16", + "pusht/pipeline_sampler_chain_gripper_newdata_points_dense_medium_h16", + "pusht/pipeline_diffusion_chain_gripper_newdata_points_h16", ] OLD_OBSTACLE_EXPERIMENTS = [