diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..4cb36c3f6cc85d761fce558e8ac6b3c0549f2a22 100644 --- a/.gitattributes +++ b/.gitattributes @@ -33,3 +33,13 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text *.zip filter=lfs diff=lfs merge=lfs -text *.zst filter=lfs diff=lfs merge=lfs -text *tfevents* filter=lfs diff=lfs merge=lfs -text +openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/ocdbt.process_0/d/530bcf4c80475910f788fc4e0549e9db filter=lfs diff=lfs merge=lfs -text +openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/ocdbt.process_0/d/be600bb25bbb4a56f17d7e8ac6a7d170 filter=lfs diff=lfs merge=lfs -text +openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/ocdbt.process_0/d/bc79931d9f67d3d5681a71493cfe31e3 filter=lfs diff=lfs merge=lfs -text +openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/ocdbt.process_0/d/f65dc4a06590f10c9806e235cda540c2 filter=lfs diff=lfs merge=lfs -text +openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/ocdbt.process_0/d/965786e7a16a28ab7bfade6250af4dbc filter=lfs diff=lfs merge=lfs -text +openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/ocdbt.process_0/d/843484256f847393a1fad4f4be6cc579 filter=lfs diff=lfs merge=lfs -text +openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/ocdbt.process_0/d/a1b59a1f3c82e649f8737c81e7144c71 filter=lfs diff=lfs merge=lfs -text +openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/ocdbt.process_0/d/672505ea0f45854a51522b6f72c74956 filter=lfs diff=lfs merge=lfs -text +openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/ocdbt.process_0/d/25f9d93df7adefdd09f13709ad7405ac filter=lfs diff=lfs merge=lfs -text +openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/ocdbt.process_0/d/226ae31e99c7333a61450a00ba995095 filter=lfs diff=lfs merge=lfs -text diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..32b0ecd38aa685b788872fdfa1f82d87085ed164 --- /dev/null +++ b/README.md @@ -0,0 +1,107 @@ +--- +library_name: pytorch +tags: [robotics, rm65, openpi, pi05, imagewam] +--- +# RealMan RM65 right-arm policies: pi0.5 (openpi) and ImageWAM (private) + +This repo holds two policies finetuned on the same RealMan RM65 right-arm dataset. The dataset has +4 pick-and-hold tasks, 187 episodes and 58,399 frames. Each policy comes with its two latest checkpoints. + +| model | training run | checkpoints | size each | +|---|---|---|---| +| pi0.5 (openpi) | `pi05_rm65_4tasks_trim`, 30k steps | `25000`, `29999` (final) | 6.7G (params only, no optimizer state) | +| ImageWAM FLUX.2-klein-4B | `rm65_4tasks_trim_flux2_klein_4b/2026-09-25_07-56-03`, 10 epochs | `step_006000.pt`, `step_006020.pt` (final) | 17.8G | + +## Robot interface (both models) + +- **State and action:** each is `float32[7]` = `right_q1..right_q6` (joint angles, rad) + `right_gripper` + (0 = closed, 100 = open). Actions are **absolute** joint targets, not deltas. +- **Cameras:** `top` (head) and `wrist` (right wrist), both 640x480 RGB, recorded at 30 fps. +- **Prompts:** use the exact strings the models were trained on. Each one names the object to pick: + - `Pick up the rightmost yellow Lay's canister with the right arm and retract while keeping it grasped.` + - `Pick up the rightmost Nongfu Spring water bottle with the right arm, retract, and hold it without releasing it.` + - `Pick up the rightmost hand soap bottle, retract the arm, and hold the bottle without releasing it.` + - `Pick up the rightmost Coca-Cola container, retract the arm, and hold it without releasing it.` + +## Training data + +The training data merges the 4 per-task deliveries (`rm65_4tasks_187_20260921_delivery`: yellow / nongfu / +handsoap / cocacola) and then trims them: + +- **Head:** the static frames at the start of every episode are dropped. +- **Tail:** the static frames at the end are cut to 10. +- **Middle:** short pauses inside an episode are kept. +- **Result:** 66,121 frames down to 58,399. + +To rebuild the dataset, run these two scripts in order: +1. `openpi/code/scripts/merge_datasets.py` +2. `openpi/code/scripts/trim_static_frames.py` + +## openpi pi0.5 — `openpi/` + +- **Checkpoints:** `checkpoints/pi05_rm65_4tasks_trim//` contains `params/` and + `assets/rm65/4tasks_v1_trim/norm_stats.json`. The norm stats must stay next to the params they came with. +- **Code** (`code/`): + - **`rm65_config_excerpt.py`** holds `LeRobotRM65DataConfig` plus the `pi05_rm65_4tasks_trim` TrainConfig. + Add both to `src/openpi/training/config.py`. + - **`b601_policy.py`** provides the input/output transforms. RM65 has the same 7-dim, two-camera layout + as B601, so they are reused unchanged. +- **Model and training setup:** + - pi0.5 base, `action_horizon=30` (1 s of actions at 30 fps). + - SigLIP trained fully, the Gemma-2B LLM with LoRA, the action expert fully. + - Batch 32, 30k steps. +- **Serving:** + + uv run scripts/serve_policy.py policy:checkpoint \ + --policy.config=pi05_rm65_4tasks_trim \ + --policy.dir=/openpi/checkpoints/pi05_rm65_4tasks_trim/29999 + +- **Inference input:** `observation/top`, `observation/wrist` (HxWx3 uint8), `observation/state` (7) and `prompt`. +- **Inference output:** `actions` of shape [30, 7]. + +## ImageWAM — `imagewam/` + +The layout mirrors the ImageWAM repo root, so you can copy this directory over `/ImageWAM/`: + +- **Weights and stats:** + - `runs/rm65_4tasks_trim_flux2_klein_4b/2026-09-25_07-56-03/checkpoints/weights/step_*.pt` are the + full model weights, with LoRA merged. + - **`dataset_stats.json`** in the same run directory holds the normalization stats. It must always be + paired with the checkpoints from **this run**. + - `config.yaml` in the same run directory is the resolved training config. Its absolute paths point + to the original machine. +- **Configs and launch script:** + - `configs/data/rm65_4tasks_trim.yaml` + - `configs/task/rm65_4tasks_trim_flux2_klein_4b.yaml` + - `scripts/rm65/run_train_rm65_4tasks_trim.sh` +- **Text cache:** `dataset/rm65_4tasks_trim/flux2_qwen3_cache_4b/` holds the Qwen3 embeddings of the 4 prompts, + so you don't need to run Qwen3 again for them. + +**How RM65 maps onto the model:** the base checkpoint is the RoboTwin 14-dim joint-space model, with layout +`[arm 6, gripper] x 2`. RM65 occupies the **left-arm slot**: + +- **Action and state:** the 7 dims sit at `[0:7]` of the model's 14 dims. `[7:14]` is zero padding and is + masked out of the loss. Take `action[..., :7]` from the model output. +- **Cameras:** the robotwin 288x256 layout has three slots, filled like this: + - `cam_high` <- `top` + - `cam_left` <- `wrist` + - `cam_right` <- a black (all-zero) frame +- **Action chunk:** 16 steps at 30 fps. + +**Training setup:** +- LoRA with rank 16 on the video expert; the action expert is trained fully. +- Learning rate 5e-5. +- Effective batch 96, 10 epochs. +- The loss is upweighted around gripper open/close transitions. + +**Serving:** the ImageWAM repo only ships a G1 serving script (`scripts/g1/serve_imagewam_g1.py`, +which is 16-dim and G1-specific). An RM65 server has not been written yet. It has to reproduce the camera +and dimension mapping above. + +## Not included + +- **FLUX.2 base weights** (needed to build ImageWAM): + - `flux-2-klein-base-4b.safetensors` from black-forest-labs/FLUX.2-klein-base-4B + - `ae.safetensors` from black-forest-labs/FLUX.2-dev (gated) +- **pi0.5 base weights:** only needed for retraining, from `gs://openpi-assets/checkpoints/pi05_base`. +- **The dataset itself.** diff --git a/imagewam/configs/data/rm65_4tasks_trim.yaml b/imagewam/configs/data/rm65_4tasks_trim.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f7e4bd76805fdeb9d63fb81c8d0c128dcb50d7ff --- /dev/null +++ b/imagewam/configs/data/rm65_4tasks_trim.yaml @@ -0,0 +1,177 @@ +# RealMan RM65 right-arm, 4-task (rm65/4tasks_v1_trim) data config for FLUX.2 ImageWAM. +# +# Finetunes the RoboTwin FLUX.2-klein-4B release checkpoint, so the shape settings mirror +# the data config that base was trained with (checkpoints/imagewam_release/robotwin/ +# flux2_klein_4b/config.yaml): 3 camera slots in the robotwin compact_288x256 layout, +# 14-dim joint-space action/state, z-score normalization. +# +# RM65 is a single 6-DoF arm, action/state float32[7] = right_q1..q6 (rad) + right_gripper +# (0 = closed .. 100 = open). That is exactly one RoboTwin arm (6 joints + gripper), so +# instead of rewriting the base's action layers the 7 dims are placed in the base's +# LEFT-arm slot [0:7] and the right-arm slot [7:14] is zero-padded: +# - ConcatLeftAlign(target 14) puts the 7 dims at [0:7] and marks [7:14] in +# action_dim_is_pad / proprio_dim_is_pad; the action loss drops those dims. +# - Normalization runs on the raw 7 dims before padding, so the pad never gets stats. +# The cameras follow the same slot: the wrist camera goes into cam_left (the left-arm +# wrist slot), and the missing third camera (cam_right) is zero-filled by the processor +# (num_output_cameras: 3 > 2 input cameras): +# cam_high <- observation.images.top +# cam_left <- observation.images.wrist +# cam_right <- zeros +# +# Dataset: rm65/4tasks_v1_trim, built by openpi-fintune/scripts/trim_static_frames.py from +# the 4 merged deliveries (yellow / nongfu / handsoap / cocacola): 187 episodes, +# 58,399 frames, 4 tasks, 30 fps, av1. The static head of every episode is dropped and +# the static tail cut to 10 frames. Kept rows keep their original `timestamp`, and the +# video frame is decoded at that timestamp from the untouched per-episode mp4, so no +# nonidle_filter is needed on top. +# +# `pretrained_norm_stats` is null: RoboTwin's stats describe a different robot, so stats +# are recomputed here and written to the run dir. +# Paths are overridden by scripts/rm65/run_train_rm65_4tasks_trim.sh. + +train: + _target_: imagewam.datasets.lerobot.robot_video_dataset.RobotVideoDataset + dataset_dirs: + - /shared/user64/workspace/yuhao/pi/data/rm65/4tasks_v1_trim + shape_meta: + images: + - key: cam_high + lerobot_key: observation.images.top + raw_shape: [3, 480, 640] + shape: [3, 240, 320] + - key: cam_left + lerobot_key: observation.images.wrist + raw_shape: [3, 480, 640] + shape: [3, 240, 320] + action: + - key: default + lerobot_key: action + raw_shape: 7 + shape: 7 + state: + - key: default + lerobot_key: observation.state + raw_shape: 7 + shape: 7 + num_frames: 17 + global_sample_stride: 1 + sample_index_stride: 1 + action_video_freq_ratio: 1 # 16 action steps, video endpoints only + video_size: [288, 256] + camera_key: null + val_set_proportion: 0.01 + is_training_set: true + pretrained_norm_stats: null + skip_padding_as_possible: false + concat_multi_camera: "robotwin" + robotwin_camera_layout: "compact_288x256" + endpoint_frames_only: true + nonidle_filter_path: null + video_augmentation: + _target_: imagewam.datasets.lerobot.transforms.image.VideoAugmentation + p: 0.5 + augment_types: [corrupt_only, color_only, both] + color_jitter: + brightness: 0.2 + contrast: 0.2 + saturation: 0.2 + hue: 0.03 + gamma: + range: [0.9, 1.1] + gaussian_noise: + std: 0.01 + random_resized_crop: + scale: [0.95, 1.0] + ratio: preserve + rotate: + degrees: 5.0 + fill: mean + require_text_cache: false + # Written into this repo, not next to the episodes: the dataset root is shared with + # openpi and is treated as read-only. + qwen_text_cache_dir: ../dataset/rm65_4tasks_trim/flux2_qwen3_cache_4b + qwen_context_len: 128 + qwen_text_cache_format: qwen3_flux2 + processor: + _target_: imagewam.datasets.lerobot.processors.imagewam_processor.ImageWAMProcessor + shape_meta: ${data.train.shape_meta} + num_obs_steps: ${data.train.num_frames} + image_obs_steps: 2 + num_output_cameras: 3 + action_output_dim: 14 + proprio_output_dim: 14 + + # action & state normalization + action_state_transforms: null + use_stepwise_action_norm: False + norm_default_mode: "z-score" + norm_exception_mode: null + + action_state_merger: + _target_: imagewam.datasets.lerobot.transforms.action_state_merger.ConcatLeftAlign + action_target_dim: 14 + state_target_dim: 14 + + # 夹爪开合关键帧 loss 加权：对原始 action 的维 6（right_gripper）二值化后检测双向跃变， + # 窗口内 loss ×key_weight。RM65 夹爪 action 只取 0 / 100 两个值，阈值取中点 50。 + gripper_keyframe_weight: + _target_: imagewam.datasets.lerobot.transforms.gripper_keyframe_weight.GripperKeyframeWeight + gripper_dims: [6] + thresholds: [50.0] + pre: 4 + post: 4 + key_weight: 3.0 + base_weight: 1.0 + + train_transforms: + - _target_: imagewam.datasets.lerobot.transforms.image.ToTensor + - _target_: torchvision.transforms.Resize + size: [240, 320] + val_transforms: + - _target_: imagewam.datasets.lerobot.transforms.image.ToTensor + - _target_: torchvision.transforms.Resize + size: [240, 320] + text_embedding_cache_dir: null + context_len: 128 + +val: + _target_: imagewam.datasets.lerobot.robot_video_dataset.RobotVideoDataset + dataset_dirs: ${data.train.dataset_dirs} + shape_meta: ${data.train.shape_meta} + num_frames: ${data.train.num_frames} + global_sample_stride: ${data.train.global_sample_stride} + sample_index_stride: ${data.train.sample_index_stride} + action_video_freq_ratio: ${data.train.action_video_freq_ratio} + video_size: ${data.train.video_size} + camera_key: null + val_set_proportion: ${data.train.val_set_proportion} + is_training_set: false + pretrained_norm_stats: null + skip_padding_as_possible: false + concat_multi_camera: ${data.train.concat_multi_camera} + robotwin_camera_layout: ${data.train.robotwin_camera_layout} + endpoint_frames_only: ${data.train.endpoint_frames_only} + nonidle_filter_path: ${data.train.nonidle_filter_path} + video_augmentation: null + require_text_cache: ${data.train.require_text_cache} + qwen_text_cache_dir: ${data.train.qwen_text_cache_dir} + qwen_context_len: ${data.train.qwen_context_len} + qwen_text_cache_format: ${data.train.qwen_text_cache_format} + processor: + _target_: imagewam.datasets.lerobot.processors.imagewam_processor.ImageWAMProcessor + shape_meta: ${data.train.shape_meta} + num_obs_steps: ${data.train.num_frames} + image_obs_steps: ${data.train.processor.image_obs_steps} + num_output_cameras: ${data.train.processor.num_output_cameras} + action_output_dim: ${data.train.processor.action_output_dim} + proprio_output_dim: ${data.train.processor.proprio_output_dim} + action_state_transforms: ${data.train.processor.action_state_transforms} + use_stepwise_action_norm: ${data.train.processor.use_stepwise_action_norm} + norm_default_mode: ${data.train.processor.norm_default_mode} + norm_exception_mode: ${data.train.processor.norm_exception_mode} + action_state_merger: ${data.train.processor.action_state_merger} + train_transforms: ${data.train.processor.train_transforms} + val_transforms: ${data.train.processor.val_transforms} + text_embedding_cache_dir: null + context_len: ${data.train.context_len} diff --git a/imagewam/configs/task/rm65_4tasks_trim_flux2_klein_4b.yaml b/imagewam/configs/task/rm65_4tasks_trim_flux2_klein_4b.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f1eef7c50a84e91819630f5ba9864ab0df5236b4 --- /dev/null +++ b/imagewam/configs/task/rm65_4tasks_trim_flux2_klein_4b.yaml @@ -0,0 +1,71 @@ +# @package _global_ + +# RealMan RM65 right-arm 4-task (rm65/4tasks_v1_trim) finetune of the RoboTwin +# FLUX.2-klein-4B release checkpoint (checkpoints/imagewam_release/robotwin/flux2_klein_4b). +# +# Why RoboTwin instead of the InternData-A1-EE base used by the G1 EEF tasks: RM65 data +# is joint space, and the RoboTwin base is 14-dim joint space [arm 6, gripper] x 2. RM65's +# 7 dims fill the base's left-arm slot and the rest is padded (see +# configs/data/rm65_4tasks_trim.yaml), so the base is loaded straight through `resume` at +# action_dim=14 with no dimension surgery (contrast with the G1 joint tasks, which expand +# RoboTwin to 16 via scripts/g1/expand_action_dim_checkpoint.py). +# +# Training mode: LoRA on the video expert, full training on the action expert, same as the +# G1 finetunes. +# +# Scale: 187 episodes, 58,399 frames, 4 tasks. +# effective batch = 4 GPU x batch_size 2 x grad_accum 12 = 96, ~608 step/epoch, ~6080 steps. + +defaults: + - override /data: rm65_4tasks_trim + - override /model: imagewam_flux2_klein_4b_base + - _self_ + +model: + proprio_dim: ${data.train.processor.proprio_output_dim} + qwen_context_len: ${data.train.qwen_context_len} + flux2_lora_config: + enabled: true + rank: 16 + alpha: 16.0 + dropout: 0.0 + save_lora_merged: true + save_trainable_only: false + target_suffixes: + - qkv + - proj + - linear1 + - linear2 + - img_mlp.0 + - img_mlp.2 + - txt_mlp.0 + - txt_mlp.2 + +batch_size: 2 +num_workers: 8 +prefetch_factor: 2 +persistent_workers: true +keep_latest_state_only: true + +lr_scheduler_type: "cosine" +learning_rate: 5e-5 +num_epochs: 10 +max_steps: null +log_every: 10 +save_every: 1000 +eval_every: 500 +eval_num_samples: 4 + +gradient_accumulation_steps: 12 +weight_decay: 1e-2 +mixed_precision: "bf16" +resume: null + +wandb: + enabled: true + project: imagewam + name: rm65_4tasks_trim_flux2_klein_4b + # configs/train.yaml hardcodes the upstream author's wandb entity, which this machine + # cannot write to; offline keeps local logs (`wandb sync` later), see + # g1_stack_cube_eef_flux2_klein_4b.yaml. + mode: offline diff --git a/imagewam/dataset/rm65_4tasks_trim/flux2_qwen3_cache_4b/4209cad0df0b8634a81bafffdb6d0cba5ba39be09c2c2c27bab26530d7535150.qwen3_flux2_len128.pt b/imagewam/dataset/rm65_4tasks_trim/flux2_qwen3_cache_4b/4209cad0df0b8634a81bafffdb6d0cba5ba39be09c2c2c27bab26530d7535150.qwen3_flux2_len128.pt new file mode 100644 index 0000000000000000000000000000000000000000..7246549e4f0284514422a5be64c17a37fd99a3b6 --- /dev/null +++ b/imagewam/dataset/rm65_4tasks_trim/flux2_qwen3_cache_4b/4209cad0df0b8634a81bafffdb6d0cba5ba39be09c2c2c27bab26530d7535150.qwen3_flux2_len128.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:abc674583b92af56cbe9188785406edf11202f6a5604810c78fc4527131e4a12 +size 1969349 diff --git a/imagewam/dataset/rm65_4tasks_trim/flux2_qwen3_cache_4b/559944487d779991b5397cf75bf04cfef53b82b81a5dda9e7c9aa384c0ac4cc0.qwen3_flux2_len128.pt b/imagewam/dataset/rm65_4tasks_trim/flux2_qwen3_cache_4b/559944487d779991b5397cf75bf04cfef53b82b81a5dda9e7c9aa384c0ac4cc0.qwen3_flux2_len128.pt new file mode 100644 index 0000000000000000000000000000000000000000..399624ad9cff72896a5e9ee1e03fded20fda0787 --- /dev/null +++ b/imagewam/dataset/rm65_4tasks_trim/flux2_qwen3_cache_4b/559944487d779991b5397cf75bf04cfef53b82b81a5dda9e7c9aa384c0ac4cc0.qwen3_flux2_len128.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:fad7620ab764cfc2c78adfab90d0e205394c89cc0ffaacd378da70eb132bacd3 +size 1969349 diff --git a/imagewam/dataset/rm65_4tasks_trim/flux2_qwen3_cache_4b/8eb884336568f01c3644e5f49f3e24f5d3b2a967a793ddd0f726c264b80d5474.qwen3_flux2_len128.pt b/imagewam/dataset/rm65_4tasks_trim/flux2_qwen3_cache_4b/8eb884336568f01c3644e5f49f3e24f5d3b2a967a793ddd0f726c264b80d5474.qwen3_flux2_len128.pt new file mode 100644 index 0000000000000000000000000000000000000000..5dd600a067c4d14a587efff6aabd7f0aa2a4f735 --- /dev/null +++ b/imagewam/dataset/rm65_4tasks_trim/flux2_qwen3_cache_4b/8eb884336568f01c3644e5f49f3e24f5d3b2a967a793ddd0f726c264b80d5474.qwen3_flux2_len128.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:cc56dc73af8741f83033197bbe5aa96469ff4a70b9d07b654b756e37f767fe14 +size 1969349 diff --git a/imagewam/dataset/rm65_4tasks_trim/flux2_qwen3_cache_4b/b1c871416f464976be1046bef84d85ecd57dc6c75bcf04ad7c62263833374ab8.qwen3_flux2_len128.pt b/imagewam/dataset/rm65_4tasks_trim/flux2_qwen3_cache_4b/b1c871416f464976be1046bef84d85ecd57dc6c75bcf04ad7c62263833374ab8.qwen3_flux2_len128.pt new file mode 100644 index 0000000000000000000000000000000000000000..b666a3e841218675dfd2f858876e1907e53066cc --- /dev/null +++ b/imagewam/dataset/rm65_4tasks_trim/flux2_qwen3_cache_4b/b1c871416f464976be1046bef84d85ecd57dc6c75bcf04ad7c62263833374ab8.qwen3_flux2_len128.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d18d8ac1781b19e7257409e63121a13b0358fde0f87071051c0307191900041d +size 1969349 diff --git a/imagewam/runs/rm65_4tasks_trim_flux2_klein_4b/2026-09-25_07-56-03/checkpoints/weights/step_006000.pt b/imagewam/runs/rm65_4tasks_trim_flux2_klein_4b/2026-09-25_07-56-03/checkpoints/weights/step_006000.pt new file mode 100644 index 0000000000000000000000000000000000000000..95fa78b12a39a27973eedb6026416c50f9b04ecd --- /dev/null +++ b/imagewam/runs/rm65_4tasks_trim_flux2_klein_4b/2026-09-25_07-56-03/checkpoints/weights/step_006000.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c2cdf9e4585ced24348c9f1b0559cdccb128f3d085bf85db6e79633ca7cafc87 +size 17773048563 diff --git a/imagewam/runs/rm65_4tasks_trim_flux2_klein_4b/2026-09-25_07-56-03/config.yaml b/imagewam/runs/rm65_4tasks_trim_flux2_klein_4b/2026-09-25_07-56-03/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f773acde5ad8f759851cfb13ce017a52392f3c15 --- /dev/null +++ b/imagewam/runs/rm65_4tasks_trim_flux2_klein_4b/2026-09-25_07-56-03/config.yaml @@ -0,0 +1,371 @@ +output_dir: ./runs/rm65_4tasks_trim_flux2_klein_4b/2026-09-25_07-56-03 +batch_size: 2 +num_workers: 8 +prefetch_factor: 2 +persistent_workers: true +lr_scheduler_type: cosine +learning_rate: 5.0e-05 +warmup_steps: null +num_epochs: 10 +max_steps: null +log_every: 10 +save_every: 1000 +keep_latest_state_only: true +eval_every: 500 +eval_num_inference_steps: 10 +eval_num_samples: 4 +rank_timer_every: 10 +rank_timer_sync_cuda: true +qwen_cache_batch_size: 32 +qwen_cache_overwrite: false +qwen_cache_save_workers: 8 +flux2_qwen3_model_spec: null +precache_arrow: true +precache_num_workers: 4 +build_lerobot_meta_cache: true +build_norm_stats: true +precache_fuse_norm_stats: true +force_rebuild_norm_stats: false +norm_stats_use_arrow_cache: true +norm_stats_bulk_read_arrow: true +norm_stats_cache_enabled: true +norm_stats_cache_dir: null +precache_profile: false +precache_profile_min_sec: 0.0 +gradient_accumulation_steps: 12 +mixed_precision: bf16 +seed: 42 +max_grad_norm: 1.0 +weight_decay: 0.01 +resume: /shared/user64/workspace/yuhao/imagewam/checkpoints/imagewam_release/robotwin/flux2_klein_4b/model.pt +wandb: + enabled: true + workspace: arisilin + project: imagewam + name: rm65_4tasks_trim_flux2_klein_4b + group: null + mode: offline +data: + train: + _target_: imagewam.datasets.lerobot.robot_video_dataset.RobotVideoDataset + dataset_dirs: + - /shared/user64/workspace/yuhao/pi/data/rm65/4tasks_v1_trim + shape_meta: + images: + - key: cam_high + lerobot_key: observation.images.top + raw_shape: + - 3 + - 480 + - 640 + shape: + - 3 + - 240 + - 320 + - key: cam_left + lerobot_key: observation.images.wrist + raw_shape: + - 3 + - 480 + - 640 + shape: + - 3 + - 240 + - 320 + action: + - key: default + lerobot_key: action + raw_shape: 7 + shape: 7 + state: + - key: default + lerobot_key: observation.state + raw_shape: 7 + shape: 7 + num_frames: 17 + global_sample_stride: 1 + sample_index_stride: 1 + action_video_freq_ratio: 1 + video_size: + - 288 + - 256 + camera_key: null + val_set_proportion: 0.01 + is_training_set: true + pretrained_norm_stats: null + skip_padding_as_possible: false + concat_multi_camera: robotwin + robotwin_camera_layout: compact_288x256 + endpoint_frames_only: true + nonidle_filter_path: null + video_augmentation: + _target_: imagewam.datasets.lerobot.transforms.image.VideoAugmentation + p: 0.5 + augment_types: + - corrupt_only + - color_only + - both + color_jitter: + brightness: 0.2 + contrast: 0.2 + saturation: 0.2 + hue: 0.03 + gamma: + range: + - 0.9 + - 1.1 + gaussian_noise: + std: 0.01 + random_resized_crop: + scale: + - 0.95 + - 1.0 + ratio: preserve + rotate: + degrees: 5.0 + fill: mean + require_text_cache: false + qwen_text_cache_dir: /shared/user64/workspace/yuhao/imagewam/ImageWAM/../dataset/rm65_4tasks_trim/flux2_qwen3_cache_4b + qwen_context_len: 128 + qwen_text_cache_format: qwen3_flux2 + processor: + _target_: imagewam.datasets.lerobot.processors.imagewam_processor.ImageWAMProcessor + shape_meta: + images: + - key: cam_high + lerobot_key: observation.images.top + raw_shape: + - 3 + - 480 + - 640 + shape: + - 3 + - 240 + - 320 + - key: cam_left + lerobot_key: observation.images.wrist + raw_shape: + - 3 + - 480 + - 640 + shape: + - 3 + - 240 + - 320 + action: + - key: default + lerobot_key: action + raw_shape: 7 + shape: 7 + state: + - key: default + lerobot_key: observation.state + raw_shape: 7 + shape: 7 + num_obs_steps: 17 + image_obs_steps: 2 + num_output_cameras: 3 + action_output_dim: 14 + proprio_output_dim: 14 + action_state_transforms: null + use_stepwise_action_norm: false + norm_default_mode: z-score + norm_exception_mode: null + action_state_merger: + _target_: imagewam.datasets.lerobot.transforms.action_state_merger.ConcatLeftAlign + action_target_dim: 14 + state_target_dim: 14 + gripper_keyframe_weight: + _target_: imagewam.datasets.lerobot.transforms.gripper_keyframe_weight.GripperKeyframeWeight + gripper_dims: + - 6 + thresholds: + - 50.0 + pre: 4 + post: 4 + key_weight: 3.0 + base_weight: 1.0 + train_transforms: + - _target_: imagewam.datasets.lerobot.transforms.image.ToTensor + - _target_: torchvision.transforms.Resize + size: + - 240 + - 320 + val_transforms: + - _target_: imagewam.datasets.lerobot.transforms.image.ToTensor + - _target_: torchvision.transforms.Resize + size: + - 240 + - 320 + text_embedding_cache_dir: null + context_len: 128 + val: + _target_: imagewam.datasets.lerobot.robot_video_dataset.RobotVideoDataset + dataset_dirs: + - /shared/user64/workspace/yuhao/pi/data/rm65/4tasks_v1_trim + shape_meta: + images: + - key: cam_high + lerobot_key: observation.images.top + raw_shape: + - 3 + - 480 + - 640 + shape: + - 3 + - 240 + - 320 + - key: cam_left + lerobot_key: observation.images.wrist + raw_shape: + - 3 + - 480 + - 640 + shape: + - 3 + - 240 + - 320 + action: + - key: default + lerobot_key: action + raw_shape: 7 + shape: 7 + state: + - key: default + lerobot_key: observation.state + raw_shape: 7 + shape: 7 + num_frames: 17 + global_sample_stride: 1 + sample_index_stride: 1 + action_video_freq_ratio: 1 + video_size: + - 288 + - 256 + camera_key: null + val_set_proportion: 0.01 + is_training_set: false + pretrained_norm_stats: null + skip_padding_as_possible: false + concat_multi_camera: robotwin + robotwin_camera_layout: compact_288x256 + endpoint_frames_only: true + nonidle_filter_path: null + video_augmentation: null + require_text_cache: false + qwen_text_cache_dir: /shared/user64/workspace/yuhao/imagewam/ImageWAM/../dataset/rm65_4tasks_trim/flux2_qwen3_cache_4b + qwen_context_len: 128 + qwen_text_cache_format: qwen3_flux2 + processor: + _target_: imagewam.datasets.lerobot.processors.imagewam_processor.ImageWAMProcessor + shape_meta: + images: + - key: cam_high + lerobot_key: observation.images.top + raw_shape: + - 3 + - 480 + - 640 + shape: + - 3 + - 240 + - 320 + - key: cam_left + lerobot_key: observation.images.wrist + raw_shape: + - 3 + - 480 + - 640 + shape: + - 3 + - 240 + - 320 + action: + - key: default + lerobot_key: action + raw_shape: 7 + shape: 7 + state: + - key: default + lerobot_key: observation.state + raw_shape: 7 + shape: 7 + num_obs_steps: 17 + image_obs_steps: 2 + num_output_cameras: 3 + action_output_dim: 14 + proprio_output_dim: 14 + action_state_transforms: null + use_stepwise_action_norm: false + norm_default_mode: z-score + norm_exception_mode: null + action_state_merger: + _target_: imagewam.datasets.lerobot.transforms.action_state_merger.ConcatLeftAlign + action_target_dim: 14 + state_target_dim: 14 + train_transforms: + - _target_: imagewam.datasets.lerobot.transforms.image.ToTensor + - _target_: torchvision.transforms.Resize + size: + - 240 + - 320 + val_transforms: + - _target_: imagewam.datasets.lerobot.transforms.image.ToTensor + - _target_: torchvision.transforms.Resize + size: + - 240 + - 320 + text_embedding_cache_dir: null + context_len: 128 +model: + _target_: imagewam.runtime.create_imagewam_flux2_klein + flux2_src_path: /shared/user64/workspace/yuhao/imagewam/ImageWAM/third_party/flux2 + flux2_model_path: /shared/user64/workspace/yuhao/imagewam/checkpoints/flux2/FLUX.2-klein-base-4B/flux-2-klein-base-4b.safetensors + ae_model_path: /shared/user64/workspace/yuhao/imagewam/checkpoints/flux2/FLUX.2-dev/ae.safetensors + variant: klein-base-4b + qwen3_model_spec: Qwen/Qwen3-4B + qwen_context_len: 128 + load_text_encoder: false + proprio_dim: 14 + mot_checkpoint_mixed_attn: true + mot_gqa_implementation: repeat + mot_force_flash_attention: false + pack_proprio_after_text: true + action_dit_pretrained_path: /shared/user64/workspace/yuhao/imagewam/checkpoints/action_dit_flux2_4b_robotwin14_init.pt + flux2_lora_config: + enabled: true + rank: 16 + alpha: 16.0 + dropout: 0.0 + save_lora_merged: true + save_trainable_only: false + target_suffixes: + - qkv + - proj + - linear1 + - linear2 + - img_mlp.0 + - img_mlp.2 + - txt_mlp.0 + - txt_mlp.2 + action_dit_config: + action_dim: 14 + hidden_dim: 1024 + num_heads: 24 + attn_head_dim: 128 + num_layers_double: 5 + num_layers_single: 20 + mlp_ratio: 4.0 + max_action_horizon: 64 + use_gradient_checkpointing: true + video_scheduler: + train_shift: 5.0 + infer_shift: 5.0 + num_train_timesteps: 1000 + action_scheduler: + train_shift: 5.0 + infer_shift: 5.0 + num_train_timesteps: 1000 + loss: + lambda_video: 0.5 + lambda_action: 1.0 diff --git a/imagewam/runs/rm65_4tasks_trim_flux2_klein_4b/2026-09-25_07-56-03/dataset_stats.json b/imagewam/runs/rm65_4tasks_trim_flux2_klein_4b/2026-09-25_07-56-03/dataset_stats.json new file mode 100644 index 0000000000000000000000000000000000000000..6f59d9cd2300227a75ee0acfa8b3bb4fc0974eb5 --- /dev/null +++ b/imagewam/runs/rm65_4tasks_trim_flux2_klein_4b/2026-09-25_07-56-03/dataset_stats.json @@ -0,0 +1,1062 @@ +{ + "state": { + "default": { + "stepwise_min": [ + [ + -1.1002380847930908, + 0.619714081287384, + -1.3607109785079956, + -2.650806188583374, + 0.3905523419380188, + -1.1792666912078857, + 0.0 + ] + ], + "stepwise_max": [ + [ + 0.2612757980823517, + 1.969796061515808, + 1.8688286542892456, + -0.3856130540370941, + 1.8521260023117065, + 1.5649319887161255, + 100.0 + ] + ], + "global_min": [ + -1.1002380847930908, + 0.619714081287384, + -1.3607109785079956, + -2.650806188583374, + 0.3905523419380188, + -1.1792666912078857, + 0.0 + ], + "global_max": [ + 0.2612757980823517, + 1.969796061515808, + 1.8688286542892456, + -0.3856130540370941, + 1.8521260023117065, + 1.5649319887161255, + 100.0 + ], + "stepwise_q01": [ + [ + -1.10020112991333, + 0.6211256980895996, + -1.354007601737976, + -2.649688720703125, + 0.3905872404575348, + -1.1791284084320068, + 0.0 + ] + ], + "stepwise_q99": [ + [ + 0.2612496316432953, + 1.9691593647003174, + 1.868074893951416, + -0.38566821813583374, + 1.8521084785461426, + 1.5628485679626465, + 100.0 + ] + ], + "global_q01": [ + -1.10020112991333, + 0.6211256980895996, + -1.354007601737976, + -2.649688720703125, + 0.3905872404575348, + -1.1791284084320068, + 0.0 + ], + "global_q99": [ + 0.2612496316432953, + 1.9691593647003174, + 1.868074893951416, + -0.38566821813583374, + 1.8521084785461426, + 1.5628485679626465, + 100.0 + ], + "stepwise_mean": [ + [ + -0.4055716395378113, + 1.367707371711731, + 0.7634584307670593, + -1.507643461227417, + 1.154456377029419, + 0.4760046601295471, + 53.31482696533203 + ] + ], + "stepwise_std": [ + [ + 0.41561296582221985, + 0.2222568690776825, + 0.43371573090553284, + 0.514720618724823, + 0.3438335955142975, + 0.5212164521217346, + 49.97016525268555 + ] + ], + "global_mean": [ + -0.4055716395378113, + 1.367707371711731, + 0.7634584307670593, + -1.507643461227417, + 1.154456377029419, + 0.4760046601295471, + 53.31482696533203 + ], + "global_std": [ + 0.41561296582221985, + 0.2222568690776825, + 0.43371573090553284, + 0.514720618724823, + 0.3438335955142975, + 0.5212164521217346, + 49.97016525268555 + ] + } + }, + "action": { + "default": { + "stepwise_min": [ + [ + -1.131421446800232, + 0.5447410345077515, + -1.4655026197433472, + -2.6565349102020264, + 0.37302088737487793, + -1.2153890132904053, + 0.0 + ], + [ + -1.131421446800232, + 0.5447410345077515, + -1.4655026197433472, + -2.6565349102020264, + 0.37302088737487793, + -1.2153890132904053, + 0.0 + ], + [ + -1.131421446800232, + 0.5447410345077515, + -1.4655026197433472, + -2.6565349102020264, + 0.37302088737487793, + -1.2153890132904053, + 0.0 + ], + [ + -1.131421446800232, + 0.5447410345077515, + -1.4655026197433472, + -2.6565349102020264, + 0.37302088737487793, + -1.2153890132904053, + 0.0 + ], + [ + -1.131421446800232, + 0.5447410345077515, + -1.4655026197433472, + -2.6565349102020264, + 0.37302088737487793, + -1.2153890132904053, + 0.0 + ], + [ + -1.131421446800232, + 0.5447410345077515, + -1.4655026197433472, + -2.6565349102020264, + 0.37302088737487793, + -1.2153890132904053, + 0.0 + ], + [ + -1.131421446800232, + 0.5447410345077515, + -1.4655026197433472, + -2.6565349102020264, + 0.37302088737487793, + -1.2153890132904053, + 0.0 + ], + [ + -1.131421446800232, + 0.5447410345077515, + -1.4655026197433472, + -2.6565349102020264, + 0.37302088737487793, + -1.2153890132904053, + 0.0 + ], + [ + -1.131421446800232, + 0.5447410345077515, + -1.4655026197433472, + -2.6565349102020264, + 0.37302088737487793, + -1.2153890132904053, + 0.0 + ], + [ + -1.131421446800232, + 0.5447410345077515, + -1.4655026197433472, + -2.6565349102020264, + 0.37302088737487793, + -1.2153890132904053, + 0.0 + ], + [ + -1.131421446800232, + 0.5447410345077515, + -1.4655026197433472, + -2.6565349102020264, + 0.37302088737487793, + -1.2153890132904053, + 0.0 + ], + [ + -1.131421446800232, + 0.5447410345077515, + -1.4655026197433472, + -2.6565349102020264, + 0.37302088737487793, + -1.2153890132904053, + 0.0 + ], + [ + -1.131421446800232, + 0.5447410345077515, + -1.4655026197433472, + -2.6565349102020264, + 0.37302088737487793, + -1.2153890132904053, + 0.0 + ], + [ + -1.131421446800232, + 0.5447410345077515, + -1.4655026197433472, + -2.6565349102020264, + 0.37302088737487793, + -1.2153890132904053, + 0.0 + ], + [ + -1.131421446800232, + 0.5447410345077515, + -1.4655026197433472, + -2.6565349102020264, + 0.37302088737487793, + -1.2153890132904053, + 0.0 + ], + [ + -1.131421446800232, + 0.5447410345077515, + -1.4655026197433472, + -2.6565349102020264, + 0.37302088737487793, + -1.2153890132904053, + 0.0 + ] + ], + "stepwise_max": [ + [ + 0.47714152932167053, + 2.0157346725463867, + 1.9173988103866577, + -0.37115994095802307, + 1.8889750242233276, + 1.5875554084777832, + 100.0 + ], + [ + 0.47714152932167053, + 2.0157346725463867, + 1.9173988103866577, + -0.37115994095802307, + 1.8889750242233276, + 1.5875554084777832, + 100.0 + ], + [ + 0.47714152932167053, + 2.0157346725463867, + 1.9173988103866577, + -0.37115994095802307, + 1.8889750242233276, + 1.5875554084777832, + 100.0 + ], + [ + 0.47714152932167053, + 2.0157346725463867, + 1.9173988103866577, + -0.37115994095802307, + 1.8889750242233276, + 1.5875554084777832, + 100.0 + ], + [ + 0.47714152932167053, + 2.0157346725463867, + 1.9173988103866577, + -0.37115994095802307, + 1.8889750242233276, + 1.5875554084777832, + 100.0 + ], + [ + 0.47714152932167053, + 2.0157346725463867, + 1.9173988103866577, + -0.37115994095802307, + 1.8889750242233276, + 1.5875554084777832, + 100.0 + ], + [ + 0.47714152932167053, + 2.0157346725463867, + 1.9173988103866577, + -0.37115994095802307, + 1.8889750242233276, + 1.5875554084777832, + 100.0 + ], + [ + 0.47714152932167053, + 2.0157346725463867, + 1.9173988103866577, + -0.37115994095802307, + 1.8889750242233276, + 1.5875554084777832, + 100.0 + ], + [ + 0.47714152932167053, + 2.0157346725463867, + 1.9173988103866577, + -0.37115994095802307, + 1.8889750242233276, + 1.5875554084777832, + 100.0 + ], + [ + 0.47714152932167053, + 2.0157346725463867, + 1.9173988103866577, + -0.37115994095802307, + 1.8889750242233276, + 1.5875554084777832, + 100.0 + ], + [ + 0.47714152932167053, + 2.0157346725463867, + 1.9173988103866577, + -0.37115994095802307, + 1.8889750242233276, + 1.5875554084777832, + 100.0 + ], + [ + 0.47714152932167053, + 2.0157346725463867, + 1.9173988103866577, + -0.37115994095802307, + 1.8889750242233276, + 1.5875554084777832, + 100.0 + ], + [ + 0.47714152932167053, + 2.0157346725463867, + 1.9173988103866577, + -0.37115994095802307, + 1.8889750242233276, + 1.5875554084777832, + 100.0 + ], + [ + 0.47714152932167053, + 2.0157346725463867, + 1.9173988103866577, + -0.37115994095802307, + 1.8889750242233276, + 1.5875554084777832, + 100.0 + ], + [ + 0.47714152932167053, + 2.0157346725463867, + 1.9173988103866577, + -0.37115994095802307, + 1.8889750242233276, + 1.5875554084777832, + 100.0 + ], + [ + 0.47714152932167053, + 2.0157346725463867, + 1.9173988103866577, + -0.37115994095802307, + 1.8889750242233276, + 1.5875554084777832, + 100.0 + ] + ], + "global_min": [ + -1.131421446800232, + 0.5447410345077515, + -1.4655026197433472, + -2.6565349102020264, + 0.37302088737487793, + -1.2153890132904053, + 0.0 + ], + "global_max": [ + 0.47714152932167053, + 2.0157346725463867, + 1.9173988103866577, + -0.37115994095802307, + 1.8889750242233276, + 1.5875554084777832, + 100.0 + ], + "stepwise_q01": [ + [ + -1.1234992742538452, + 0.5465866327285767, + -1.4512237310409546, + -2.6561882495880127, + 0.37375399470329285, + -1.2099779844284058, + 0.0 + ], + [ + -1.1234992742538452, + 0.5465866327285767, + -1.4512237310409546, + -2.6561882495880127, + 0.37375399470329285, + -1.2099779844284058, + 0.0 + ], + [ + -1.1234992742538452, + 0.5465866327285767, + -1.4512237310409546, + -2.6561882495880127, + 0.37375399470329285, + -1.2099779844284058, + 0.0 + ], + [ + -1.1234992742538452, + 0.5465866327285767, + -1.4512237310409546, + -2.6561882495880127, + 0.37375399470329285, + -1.2099779844284058, + 0.0 + ], + [ + -1.1234992742538452, + 0.5465866327285767, + -1.4512237310409546, + -2.6561882495880127, + 0.37375399470329285, + -1.2099779844284058, + 0.0 + ], + [ + -1.1234992742538452, + 0.5465866327285767, + -1.4512237310409546, + -2.6561882495880127, + 0.37375399470329285, + -1.2099779844284058, + 0.0 + ], + [ + -1.1234992742538452, + 0.5465866327285767, + -1.4512237310409546, + -2.6561882495880127, + 0.37375399470329285, + -1.2099779844284058, + 0.0 + ], + [ + -1.1234992742538452, + 0.5465866327285767, + -1.4512237310409546, + -2.6561882495880127, + 0.37375399470329285, + -1.2099779844284058, + 0.0 + ], + [ + -1.1234992742538452, + 0.5465866327285767, + -1.4512237310409546, + -2.6561882495880127, + 0.37375399470329285, + -1.2099779844284058, + 0.0 + ], + [ + -1.1234992742538452, + 0.5465866327285767, + -1.4512237310409546, + -2.6561882495880127, + 0.37375399470329285, + -1.2099779844284058, + 0.0 + ], + [ + -1.1234992742538452, + 0.5465866327285767, + -1.4512237310409546, + -2.6561882495880127, + 0.37375399470329285, + -1.2099779844284058, + 0.0 + ], + [ + -1.1234992742538452, + 0.5465866327285767, + -1.4512237310409546, + -2.6561882495880127, + 0.37375399470329285, + -1.2099779844284058, + 0.0 + ], + [ + -1.1234992742538452, + 0.5465866327285767, + -1.4512237310409546, + -2.6561882495880127, + 0.37375399470329285, + -1.2099779844284058, + 0.0 + ], + [ + -1.1234992742538452, + 0.5465866327285767, + -1.4512237310409546, + -2.6561882495880127, + 0.37375399470329285, + -1.2099779844284058, + 0.0 + ], + [ + -1.1234992742538452, + 0.5465866327285767, + -1.4512237310409546, + -2.6561882495880127, + 0.37375399470329285, + -1.2099779844284058, + 0.0 + ], + [ + -1.1234992742538452, + 0.5465866327285767, + -1.4512237310409546, + -2.6561882495880127, + 0.37375399470329285, + -1.2099779844284058, + 0.0 + ] + ], + "stepwise_q99": [ + [ + 0.4733147919178009, + 2.003934621810913, + 1.913605809211731, + -0.38113221526145935, + 1.888350248336792, + 1.5870753526687622, + 100.0 + ], + [ + 0.4733147919178009, + 2.003934621810913, + 1.913605809211731, + -0.38113221526145935, + 1.888350248336792, + 1.5870753526687622, + 100.0 + ], + [ + 0.4733147919178009, + 2.003934621810913, + 1.913605809211731, + -0.38113221526145935, + 1.888350248336792, + 1.5870753526687622, + 100.0 + ], + [ + 0.4733147919178009, + 2.003934621810913, + 1.913605809211731, + -0.38113221526145935, + 1.888350248336792, + 1.5870753526687622, + 100.0 + ], + [ + 0.4733147919178009, + 2.003934621810913, + 1.913605809211731, + -0.38113221526145935, + 1.888350248336792, + 1.5870753526687622, + 100.0 + ], + [ + 0.4733147919178009, + 2.003934621810913, + 1.913605809211731, + -0.38113221526145935, + 1.888350248336792, + 1.5870753526687622, + 100.0 + ], + [ + 0.4733147919178009, + 2.003934621810913, + 1.913605809211731, + -0.38113221526145935, + 1.888350248336792, + 1.5870753526687622, + 100.0 + ], + [ + 0.4733147919178009, + 2.003934621810913, + 1.913605809211731, + -0.38113221526145935, + 1.888350248336792, + 1.5870753526687622, + 100.0 + ], + [ + 0.4733147919178009, + 2.003934621810913, + 1.913605809211731, + -0.38113221526145935, + 1.888350248336792, + 1.5870753526687622, + 100.0 + ], + [ + 0.4733147919178009, + 2.003934621810913, + 1.913605809211731, + -0.38113221526145935, + 1.888350248336792, + 1.5870753526687622, + 100.0 + ], + [ + 0.4733147919178009, + 2.003934621810913, + 1.913605809211731, + -0.38113221526145935, + 1.888350248336792, + 1.5870753526687622, + 100.0 + ], + [ + 0.4733147919178009, + 2.003934621810913, + 1.913605809211731, + -0.38113221526145935, + 1.888350248336792, + 1.5870753526687622, + 100.0 + ], + [ + 0.4733147919178009, + 2.003934621810913, + 1.913605809211731, + -0.38113221526145935, + 1.888350248336792, + 1.5870753526687622, + 100.0 + ], + [ + 0.4733147919178009, + 2.003934621810913, + 1.913605809211731, + -0.38113221526145935, + 1.888350248336792, + 1.5870753526687622, + 100.0 + ], + [ + 0.4733147919178009, + 2.003934621810913, + 1.913605809211731, + -0.38113221526145935, + 1.888350248336792, + 1.5870753526687622, + 100.0 + ], + [ + 0.4733147919178009, + 2.003934621810913, + 1.913605809211731, + -0.38113221526145935, + 1.888350248336792, + 1.5870753526687622, + 100.0 + ] + ], + "global_q01": [ + -1.1234992742538452, + 0.5465866327285767, + -1.4512237310409546, + -2.6561882495880127, + 0.37375399470329285, + -1.2099779844284058, + 0.0 + ], + "global_q99": [ + 0.4733147919178009, + 2.003934621810913, + 1.913605809211731, + -0.38113221526145935, + 1.888350248336792, + 1.5870753526687622, + 100.0 + ], + "stepwise_mean": [ + [ + -0.40218979120254517, + 1.366448998451233, + 0.7740227580070496, + -1.506816029548645, + 1.1548317670822144, + 0.48579129576683044, + 53.314823150634766 + ], + [ + -0.4020868241786957, + 1.3665825128555298, + 0.7740072011947632, + -1.506820559501648, + 1.154834508895874, + 0.4861409366130829, + 52.986717224121094 + ], + [ + -0.4019829332828522, + 1.366716742515564, + 0.7739895582199097, + -1.506824254989624, + 1.154837965965271, + 0.48648950457572937, + 52.65860366821289 + ], + [ + -0.40187719464302063, + 1.3668543100357056, + 0.7739637494087219, + -1.5068268775939941, + 1.1548435688018799, + 0.48683401942253113, + 52.33050537109375 + ], + [ + -0.4017687141895294, + 1.3669904470443726, + 0.7739416360855103, + -1.5068278312683105, + 1.1548497676849365, + 0.4871804714202881, + 52.00239181518555 + ], + [ + -0.40165838599205017, + 1.3671298027038574, + 0.7739128470420837, + -1.5068278312683105, + 1.1548584699630737, + 0.4875229597091675, + 51.67428207397461 + ], + [ + -0.40154579281806946, + 1.3672685623168945, + 0.773884117603302, + -1.5068268775939941, + 1.154868721961975, + 0.48786601424217224, + 51.34617233276367 + ], + [ + -0.4014301896095276, + 1.3674125671386719, + 0.7738425731658936, + -1.5068230628967285, + 1.1548829078674316, + 0.4882016181945801, + 51.018062591552734 + ], + [ + -0.40131089091300964, + 1.3675622940063477, + 0.7737869620323181, + -1.5068166255950928, + 1.1549016237258911, + 0.4885294735431671, + 50.6899528503418 + ], + [ + -0.40118804574012756, + 1.36771821975708, + 0.7737151980400085, + -1.50680673122406, + 1.154926061630249, + 0.48884809017181396, + 50.361846923828125 + ], + [ + -0.40106213092803955, + 1.3678781986236572, + 0.7736300230026245, + -1.5067932605743408, + 1.1549543142318726, + 0.4891602396965027, + 50.03374099731445 + ], + [ + -0.40093153715133667, + 1.3680424690246582, + 0.7735351324081421, + -1.5067766904830933, + 1.154987096786499, + 0.4894666373729706, + 49.705631256103516 + ], + [ + -0.40079763531684875, + 1.368208885192871, + 0.7734330296516418, + -1.5067572593688965, + 1.155022144317627, + 0.4897681474685669, + 49.37752151489258 + ], + [ + -0.4006597697734833, + 1.3683816194534302, + 0.7733154892921448, + -1.506732702255249, + 1.1550605297088623, + 0.49005934596061707, + 49.049415588378906 + ], + [ + -0.4005171060562134, + 1.3685643672943115, + 0.7731740474700928, + -1.506703495979309, + 1.155107021331787, + 0.49033835530281067, + 48.721309661865234 + ], + [ + -0.40036845207214355, + 1.3687560558319092, + 0.7730117440223694, + -1.5066686868667603, + 1.1551613807678223, + 0.4906056523323059, + 48.3931999206543 + ] + ], + "stepwise_std": [ + [ + 0.4327355921268463, + 0.2356889694929123, + 0.45770174264907837, + 0.5211266279220581, + 0.3537985682487488, + 0.5293993353843689, + 49.97016525268555 + ], + [ + 0.43285658955574036, + 0.23575304448604584, + 0.4575728476047516, + 0.5211104154586792, + 0.35381147265434265, + 0.5292556881904602, + 49.99089050292969 + ], + [ + 0.43297845125198364, + 0.23581714928150177, + 0.457444965839386, + 0.5210942625999451, + 0.3538253605365753, + 0.5291110873222351, + 50.00944900512695 + ], + [ + 0.43310195207595825, + 0.23588420450687408, + 0.45732516050338745, + 0.5210784673690796, + 0.3538411259651184, + 0.528967022895813, + 50.025840759277344 + ], + [ + 0.4332278072834015, + 0.23594802618026733, + 0.4571947753429413, + 0.5210630893707275, + 0.3538568615913391, + 0.5288187265396118, + 50.040061950683594 + ], + [ + 0.4333552122116089, + 0.23601122200489044, + 0.45706355571746826, + 0.5210480093955994, + 0.3538748621940613, + 0.5286698341369629, + 50.0521354675293 + ], + [ + 0.4334847629070282, + 0.2360697239637375, + 0.45692333579063416, + 0.5210330486297607, + 0.35389432311058044, + 0.52851802110672, + 50.06203842163086 + ], + [ + 0.43361666798591614, + 0.23612524569034576, + 0.4567815661430359, + 0.5210187435150146, + 0.3539167046546936, + 0.5283655524253845, + 50.06977844238281 + ], + [ + 0.4337516725063324, + 0.23617896437644958, + 0.4566412568092346, + 0.5210055708885193, + 0.3539414405822754, + 0.5282132029533386, + 50.07536315917969 + ], + [ + 0.4338892996311188, + 0.23622912168502808, + 0.45649799704551697, + 0.5209932923316956, + 0.35397040843963623, + 0.528059720993042, + 50.07879638671875 + ], + [ + 0.43402907252311707, + 0.2362746298313141, + 0.4563489556312561, + 0.5209818482398987, + 0.35400065779685974, + 0.5279049873352051, + 50.08006286621094 + ], + [ + 0.4341725707054138, + 0.23631539940834045, + 0.4561904966831207, + 0.5209711194038391, + 0.3540341556072235, + 0.5277472138404846, + 50.07917785644531 + ], + [ + 0.4343186616897583, + 0.23634861409664154, + 0.45601117610931396, + 0.5209611654281616, + 0.3540678322315216, + 0.5275831818580627, + 50.076133728027344 + ], + [ + 0.43446722626686096, + 0.23637652397155762, + 0.45581957697868347, + 0.5209526419639587, + 0.3541007936000824, + 0.5274155735969543, + 50.0709342956543 + ], + [ + 0.43461862206459045, + 0.23640376329421997, + 0.45563292503356934, + 0.5209453701972961, + 0.35413870215415955, + 0.527248740196228, + 50.06357192993164 + ], + [ + 0.4347744584083557, + 0.2364252209663391, + 0.45543548464775085, + 0.5209402441978455, + 0.35417959094047546, + 0.5270791053771973, + 50.05405044555664 + ] + ], + "global_mean": [ + -0.4013359546661377, + 1.3675321340560913, + 0.7736979126930237, + -1.5067905187606812, + 1.154932975769043, + 0.4883001744747162, + 50.85401153564453 + ], + "global_std": [ + 0.4337120056152344, + 0.23611678183078766, + 0.45666226744651794, + 0.521020233631134, + 0.3539533317089081, + 0.5282748937606812, + 50.07276153564453 + ] + } + }, + "num_episodes": 185, + "num_transition": 57716 +} \ No newline at end of file diff --git a/imagewam/scripts/rm65/run_train_rm65_4tasks_trim.sh b/imagewam/scripts/rm65/run_train_rm65_4tasks_trim.sh new file mode 100644 index 0000000000000000000000000000000000000000..c78153a4ef6a5d81d81b18b86f790082287b451f --- /dev/null +++ b/imagewam/scripts/rm65/run_train_rm65_4tasks_trim.sh @@ -0,0 +1,101 @@ +#!/usr/bin/env bash +set -euo pipefail + +# RealMan RM65 right-arm 4-task (rm65/4tasks_v1_trim) finetune of the RoboTwin +# FLUX.2-klein-4B release checkpoint. +# +# Differs from scripts/g1/run_train_g1_stack_cubes.sh in three ways: +# 1. The base is the RoboTwin release checkpoint itself at action_dim=14, not the +# 16-dim expanded G1_INIT_CKPT: RM65's 7 dims fill RoboTwin's left-arm slot and the +# rest is padded (configs/data/rm65_4tasks_trim.yaml). The ActionDiT init is +# therefore built at action_dim=14. +# 2. The dataset lives outside this tree (on /shared/user64, shared with openpi), so it +# is treated as read-only: the Qwen3 text cache is written into this repo's dataset dir. +# 3. All 4 instructions are new, so the Qwen3 cache is computed automatically when its +# directory is empty (force a recompute with PRECOMPUTE_QWEN3_CACHE=true). + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +# shellcheck source=../common.sh +source "${SCRIPT_DIR}/../common.sh" +imagewam_init "${SCRIPT_DIR}/../.." + +GPU_PER_NODE="${GPU_PER_NODE:-1}" +ZERO_STAGE="${ZERO_STAGE:-1}" +TASK_NAME="rm65_4tasks_trim_flux2_klein_4b" +MODEL_CONFIG="configs/model/imagewam_flux2_klein_4b_base.yaml" + +imagewam_require_env FLUX2_SRC +imagewam_require_env FLUX2_MODEL_PATH +imagewam_require_env FLUX2_AE_MODEL_PATH + +# Dataset root for this task. Read-only: nothing is written under it. +RM65_ROOT="${RM65_ROOT:-/shared/user64/workspace/yuhao/pi/data/rm65/4tasks_v1_trim}" +# Text cache goes into this repo's dataset dir, not the shared dataset root. +QWEN_CACHE_DIR="${QWEN_CACHE_DIR:-${REPO_ROOT}/../dataset/rm65_4tasks_trim/flux2_qwen3_cache_4b}" +FLUX2_QWEN3_MODEL_SPEC="${FLUX2_QWEN3_MODEL_SPEC:-Qwen/Qwen3-4B}" +# RoboTwin release checkpoint (14-dim joint space; loaded verbatim via `resume`). +RM65_INIT_CKPT="${RM65_INIT_CKPT:-${MODEL_ROOT:-${REPO_ROOT}/checkpoints}/imagewam_release/robotwin/flux2_klein_4b/model.pt}" +# ActionDiT init at action_dim=14. Shapes only; overwritten by `resume`. +ACTION_INIT="${ACTION_INIT:-${MODEL_ROOT:-${REPO_ROOT}/checkpoints}/action_dit_flux2_4b_robotwin14_init.pt}" + +if [ ! -f "${RM65_INIT_CKPT}" ]; then + echo "Missing base checkpoint: ${RM65_INIT_CKPT}" >&2 + exit 2 +fi +if [ ! -f "${RM65_ROOT}/meta/info.json" ]; then + echo "Missing dataset: ${RM65_ROOT}" >&2 + echo "Build it with openpi-fintune/scripts/merge_datasets.py + trim_static_frames.py." >&2 + exit 2 +fi + +export PYTHONPATH="${REPO_ROOT}/src:${FLUX2_SRC}/src:${FLUX2_SRC}${PYTHONPATH:+:${PYTHONPATH}}" +export ZERO_STAGE + +mkdir -p "${QWEN_CACHE_DIR}" + +DATASET_OVERRIDES=( + "data.train.dataset_dirs=[${RM65_ROOT}]" + "data.train.qwen_text_cache_dir=${QWEN_CACHE_DIR}" +) + +imagewam_print_config TASK_NAME RM65_ROOT RM65_INIT_CKPT FLUX2_SRC FLUX2_MODEL_PATH \ + FLUX2_AE_MODEL_PATH QWEN_CACHE_DIR ACTION_INIT GPU_PER_NODE ZERO_STAGE + +# 1. ActionDiT initialisation (action_dim=14). +if [ "${REBUILD_ACTION_INIT:-false}" = "true" ] || [ ! -f "${ACTION_INIT}" ]; then + imagewam_run imagewam_python scripts/flux2/preprocess_action_dit_flux2.py \ + --model-config "${MODEL_CONFIG}" \ + --flux2-src-path "${FLUX2_SRC}" \ + --flux2-model-path "${FLUX2_MODEL_PATH}" \ + --variant klein-base-4b \ + --action-dim 14 \ + --output "${ACTION_INIT}" \ + --apply-alpha-scaling true +fi + +# 2. Qwen3 text embedding cache. Required before any sample can be drawn, because +# RobotVideoDataset reads the cache whenever qwen_text_cache_dir is set. +if [ "${PRECOMPUTE_QWEN3_CACHE:-false}" = "true" ] || [ -z "$(ls -A "${QWEN_CACHE_DIR}")" ]; then + imagewam_run torchrun --standalone --nproc_per_node="${QWEN_CACHE_NPROC:-1}" \ + scripts/flux2/precompute_flux2_qwen3_embeds.py \ + task="${TASK_NAME}" \ + qwen_cache_batch_size="${QWEN_CACHE_BATCH_SIZE:-8}" \ + qwen_cache_save_workers="${QWEN_CACHE_SAVE_WORKERS:-4}" \ + qwen_cache_overwrite="${QWEN_CACHE_OVERWRITE:-false}" \ + model.variant=klein-base-4b \ + model.qwen3_model_spec="${FLUX2_QWEN3_MODEL_SPEC}" \ + model.flux2_model_path="${FLUX2_MODEL_PATH}" \ + model.ae_model_path="${FLUX2_AE_MODEL_PATH}" \ + flux2_qwen3_model_spec="${FLUX2_QWEN3_MODEL_SPEC}" \ + "${DATASET_OVERRIDES[@]}" +fi + +# 3. Training. +TASK="${TASK_NAME}" imagewam_run bash scripts/flux2/train_flux2_klein_imagewam.sh "${GPU_PER_NODE}" \ + "${DATASET_OVERRIDES[@]}" \ + "model.flux2_model_path=${FLUX2_MODEL_PATH}" \ + "model.ae_model_path=${FLUX2_AE_MODEL_PATH}" \ + "model.qwen3_model_spec=${FLUX2_QWEN3_MODEL_SPEC}" \ + "model.action_dit_pretrained_path=${ACTION_INIT}" \ + "resume=${RM65_INIT_CKPT}" \ + "$@" diff --git a/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/_CHECKPOINT_METADATA b/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/_CHECKPOINT_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..85a6cfce7ca3e0f9cd6011bd9dcb204e76b11e9b --- /dev/null +++ b/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/_CHECKPOINT_METADATA @@ -0,0 +1 @@ +{"item_handlers": {"assets": "openpi.training.checkpoints.CallbackHandler", "params": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler", "train_state": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1790344018790056026, "commit_timestamp_nsecs": 1790344101401360609, "custom_metadata": {}} \ No newline at end of file diff --git a/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/assets/rm65/4tasks_v1_trim/norm_stats.json b/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/assets/rm65/4tasks_v1_trim/norm_stats.json new file mode 100644 index 0000000000000000000000000000000000000000..c7c160d8adb381bbff740cd01e1f4d4ea72c0168 --- /dev/null +++ b/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/assets/rm65/4tasks_v1_trim/norm_stats.json @@ -0,0 +1,80 @@ +{ + "norm_stats": { + "state": { + "mean": [ + -0.409750759601593, + 1.3662880659103394, + 0.7612037062644958, + -1.4930087327957153, + 1.149019718170166, + 0.46189895272254944, + 53.88911819458008 + ], + "std": [ + 0.4165157377719879, + 0.22506451606750488, + 0.43782690167427063, + 0.522697925567627, + 0.3444087505340576, + 0.5309966802597046, + 49.84851837158203 + ], + "q01": [ + -1.0531297045386805, + 0.7838840500074804, + -0.39637044419501943, + -2.551590729381487, + 0.4551538976353758, + -1.0030891360439242, + -1e-10 + ], + "q99": [ + 0.20381991231656804, + 1.8083262567565686, + 1.5471665069427205, + -0.47939204971559013, + 1.696906879658781, + 1.4057684653689329, + 99.98000000009996 + ] + }, + "actions": { + "mean": [ + -0.40519973635673523, + 1.3669183254241943, + 0.7708451747894287, + -1.492232322692871, + 1.149193286895752, + 0.4756605625152588, + 49.24359130859375 + ], + "std": [ + 0.4345667362213135, + 0.2386302947998047, + 0.4585062563419342, + 0.5284347534179688, + 0.354095995426178, + 0.5372209548950195, + 49.99427795410156 + ], + "q01": [ + -1.0606446759420682, + 0.7624480928670695, + -0.39244628614266963, + -2.5683194364814876, + 0.4488185941273004, + -1.0152587816618766, + -1e-10 + ], + "q99": [ + 0.2403610593070864, + 1.8286242818624323, + 1.619703484617617, + -0.46531738959911584, + 1.7191881609739013, + 1.414894032184543, + 99.98000000009996 + ] + } + } +} \ No newline at end of file diff --git a/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/_METADATA b/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..c4e4423841b736caddb183c3068573cfb6572b98 --- /dev/null +++ b/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/_METADATA @@ -0,0 +1 @@ +{"tree_metadata": {"('params', 'PaliGemma', 'img', 'Transformer', 'encoder_norm', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoder_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [576]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoder_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoder_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [576]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 576]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_0', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 576]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_1', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 576]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_1', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 576]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 2152]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 576, 4304]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_1', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 576]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_1', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 2152, 1152]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 8, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 576, 16, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 576]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 8, 72, 1152]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 8, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 576, 16, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 8, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 576, 16, 72]}}, "('params', 'PaliGemma', 'img', 'embedding', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [576]}}, "('params', 'PaliGemma', 'img', 'embedding', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 14, 3, 1152]}}, "('params', 'PaliGemma', 'img', 'head', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'PaliGemma', 'img', 'head', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [576, 2048]}}, "('params', 'PaliGemma', 'img', 'pos_embedding', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "pos_embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 128, 1152]}}, "('params', 'PaliGemma', 'llm', 'embedder', 'input_embedding', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "embedder", "key_type": 2}, {"key": "input_embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128576, 2048]}}, "('params', 'PaliGemma', 'llm', 'final_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "final_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'PaliGemma', 'llm', 'final_norm_1', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "final_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('params', 'PaliGemma', 'llm', 'final_norm_1', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "final_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 3072]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'attn_vec_einsum', 'lora_a', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "attn_vec_einsum", "key_type": 2}, {"key": "lora_a", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 8, 256, 16]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'attn_vec_einsum', 'lora_b', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "attn_vec_einsum", "key_type": 2}, {"key": "lora_b", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 8, 16, 2048]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'attn_vec_einsum', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "attn_vec_einsum", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 8, 256, 2048]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'attn_vec_einsum_1', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "attn_vec_einsum_1", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 8, 256, 1024]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'kv_einsum', 'lora_a', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "kv_einsum", "key_type": 2}, {"key": "lora_a", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 2, 1, 2048, 16]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'kv_einsum', 'lora_b', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "kv_einsum", "key_type": 2}, {"key": "lora_b", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 2, 1, 16, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'kv_einsum', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "kv_einsum", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 2, 1, 2048, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'kv_einsum_1', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "kv_einsum_1", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 2, 1, 1024, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'q_einsum', 'lora_a', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_einsum", "key_type": 2}, {"key": "lora_a", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 8, 2048, 16]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'q_einsum', 'lora_b', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_einsum", "key_type": 2}, {"key": "lora_b", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 8, 16, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'q_einsum', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_einsum", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 8, 2048, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'q_einsum_1', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_einsum_1", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 8, 1024, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp', 'gating_einsum', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "gating_einsum", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 2, 2048, 16384]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp', 'gating_einsum_lora_a', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "gating_einsum_lora_a", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 2, 2048, 16]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp', 'gating_einsum_lora_b', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "gating_einsum_lora_b", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 2, 16, 16384]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp', 'linear', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "linear", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 16384, 2048]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp', 'linear_lora_a', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "linear_lora_a", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 16384, 16]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp', 'linear_lora_b', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "linear_lora_b", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 16, 2048]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp_1', 'gating_einsum', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp_1", "key_type": 2}, {"key": "gating_einsum", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 2, 1024, 4096]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp_1', 'linear', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp_1", "key_type": 2}, {"key": "linear", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 4096, 1024]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_attention_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_attention_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 2048]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_attention_norm_1', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_attention_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 3072]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_attention_norm_1', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_attention_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 1024, 3072]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_ffw_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_ffw_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 2048]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_ffw_norm_1', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_ffw_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 3072]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_ffw_norm_1', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_ffw_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 1024, 3072]}}, "('params', 'action_in_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_in_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'action_in_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_in_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 1024]}}, "('params', 'action_out_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_out_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16]}}, "('params', 'action_out_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 32]}}, "('params', 'time_mlp_in', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "time_mlp_in", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'time_mlp_in', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "time_mlp_in", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1024]}}, "('params', 'time_mlp_out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "time_mlp_out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'time_mlp_out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "time_mlp_out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1024]}}}, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null} \ No newline at end of file diff --git a/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/_sharding b/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/_sharding new file mode 100644 index 0000000000000000000000000000000000000000..55bac843f949f8a5e190d87a3514aa7e23dc57d4 --- /dev/null +++ b/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/_sharding @@ -0,0 +1 @@ +{"cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2Rlcl9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2Rlcl9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8wLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8wLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuZW1iZWRkaW5nLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuZW1iZWRkaW5nLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuaGVhZC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuaGVhZC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcucG9zX2VtYmVkZGluZy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZW1iZWRkZXIuaW5wdXRfZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZmluYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZmluYWxfbm9ybV8xLkRlbnNlXzAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZmluYWxfbm9ybV8xLkRlbnNlXzAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scC5nYXRpbmdfZWluc3VtLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scC5nYXRpbmdfZWluc3VtX2xvcmFfYS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scC5nYXRpbmdfZWluc3VtX2xvcmFfYi52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scC5saW5lYXIudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scC5saW5lYXJfbG9yYV9hLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scC5saW5lYXJfbG9yYV9iLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scF8xLmdhdGluZ19laW5zdW0udmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scF8xLmxpbmVhci52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4uYXR0bl92ZWNfZWluc3VtLmxvcmFfYS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4uYXR0bl92ZWNfZWluc3VtLmxvcmFfYi52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4uYXR0bl92ZWNfZWluc3VtLncudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4uYXR0bl92ZWNfZWluc3VtXzEudy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ua3ZfZWluc3VtLmxvcmFfYS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ua3ZfZWluc3VtLmxvcmFfYi52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ua3ZfZWluc3VtLncudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ua3ZfZWluc3VtXzEudy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ucV9laW5zdW0ubG9yYV9hLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ucV9laW5zdW0ubG9yYV9iLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ucV9laW5zdW0udy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ucV9laW5zdW1fMS53LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9hdHRlbnRpb25fbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9hdHRlbnRpb25fbm9ybV8xLkRlbnNlXzAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9hdHRlbnRpb25fbm9ybV8xLkRlbnNlXzAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9mZndfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9mZndfbm9ybV8xLkRlbnNlXzAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9mZndfbm9ybV8xLkRlbnNlXzAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLmFjdGlvbl9pbl9wcm9qLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLmFjdGlvbl9pbl9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLmFjdGlvbl9vdXRfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLmFjdGlvbl9vdXRfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLnRpbWVfbWxwX291dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLnRpbWVfbWxwX291dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLnRpbWVfbWxwX2luLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLnRpbWVfbWxwX2luLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}"} \ No newline at end of file diff --git a/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/array_metadatas/process_0 b/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/array_metadatas/process_0 new file mode 100644 index 0000000000000000000000000000000000000000..57bbcf50e800578263c2f2038348b04c71cedc75 --- /dev/null +++ b/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/array_metadatas/process_0 @@ -0,0 +1 @@ +{"array_metadatas": [{"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoder_norm.bias.value", "write_shape": [576], "chunk_shape": [576], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoder_norm.scale.value", "write_shape": [576], "chunk_shape": [576], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_0.bias.value", "write_shape": [27, 576], "chunk_shape": [27, 576], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_0.scale.value", "write_shape": [27, 576], "chunk_shape": [27, 576], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_1.bias.value", "write_shape": [27, 576], "chunk_shape": [27, 576], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_1.scale.value", "write_shape": [27, 576], "chunk_shape": [27, 576], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_0.bias.value", "write_shape": [27, 2152], "chunk_shape": [27, 2152], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_0.kernel.value", "write_shape": [27, 576, 4304], "chunk_shape": [27, 576, 4304], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_1.bias.value", "write_shape": [27, 576], "chunk_shape": [27, 576], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_1.kernel.value", "write_shape": [27, 2152, 1152], "chunk_shape": [27, 2152, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.key.bias.value", "write_shape": [27, 8, 72], "chunk_shape": [27, 8, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.key.kernel.value", "write_shape": [27, 576, 16, 72], "chunk_shape": [27, 576, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.out.bias.value", "write_shape": [27, 576], "chunk_shape": [27, 576], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.out.kernel.value", "write_shape": [27, 8, 72, 1152], "chunk_shape": [27, 8, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.query.bias.value", "write_shape": [27, 8, 72], "chunk_shape": [27, 8, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.query.kernel.value", "write_shape": [27, 576, 16, 72], "chunk_shape": [27, 576, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.value.bias.value", "write_shape": [27, 8, 72], "chunk_shape": [27, 8, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.value.kernel.value", "write_shape": [27, 576, 16, 72], "chunk_shape": [27, 576, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.embedding.bias.value", "write_shape": [576], "chunk_shape": [576], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.embedding.kernel.value", "write_shape": [7, 14, 3, 1152], "chunk_shape": [7, 14, 3, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.head.bias.value", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.head.kernel.value", "write_shape": [576, 2048], "chunk_shape": [576, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.pos_embedding.value", "write_shape": [1, 128, 1152], "chunk_shape": [1, 128, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.embedder.input_embedding.value", "write_shape": [128576, 2048], "chunk_shape": [128576, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.final_norm.scale.value", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.final_norm_1.Dense_0.bias.value", "write_shape": [1536], "chunk_shape": [1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.final_norm_1.Dense_0.kernel.value", "write_shape": [512, 3072], "chunk_shape": [512, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.attn_vec_einsum.lora_a.value", "write_shape": [9, 8, 256, 16], "chunk_shape": [9, 8, 256, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.attn_vec_einsum.lora_b.value", "write_shape": [9, 8, 16, 2048], "chunk_shape": [9, 8, 16, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.attn_vec_einsum.w.value", "write_shape": [9, 8, 256, 2048], "chunk_shape": [9, 8, 256, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.attn_vec_einsum_1.w.value", "write_shape": [9, 8, 256, 1024], "chunk_shape": [9, 8, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.kv_einsum.lora_a.value", "write_shape": [9, 2, 1, 2048, 16], "chunk_shape": [9, 2, 1, 2048, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.kv_einsum.lora_b.value", "write_shape": [9, 2, 1, 16, 256], "chunk_shape": [9, 2, 1, 16, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.kv_einsum.w.value", "write_shape": [9, 2, 1, 2048, 256], "chunk_shape": [9, 2, 1, 2048, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.kv_einsum_1.w.value", "write_shape": [9, 2, 1, 1024, 256], "chunk_shape": [9, 2, 1, 1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.q_einsum.lora_a.value", "write_shape": [9, 8, 2048, 16], "chunk_shape": [9, 8, 2048, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.q_einsum.lora_b.value", "write_shape": [9, 8, 16, 256], "chunk_shape": [9, 8, 16, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.q_einsum.w.value", "write_shape": [9, 8, 2048, 256], "chunk_shape": [9, 8, 2048, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.q_einsum_1.w.value", "write_shape": [9, 8, 1024, 256], "chunk_shape": [9, 8, 1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp.gating_einsum.value", "write_shape": [9, 2, 2048, 16384], "chunk_shape": [9, 2, 2048, 16384], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp.gating_einsum_lora_a.value", "write_shape": [9, 2, 2048, 16], "chunk_shape": [9, 2, 2048, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp.gating_einsum_lora_b.value", "write_shape": [9, 2, 16, 16384], "chunk_shape": [9, 2, 16, 16384], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp.linear.value", "write_shape": [9, 16384, 2048], "chunk_shape": [9, 16384, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp.linear_lora_a.value", "write_shape": [9, 16384, 16], "chunk_shape": [9, 16384, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp.linear_lora_b.value", "write_shape": [9, 16, 2048], "chunk_shape": [9, 16, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp_1.gating_einsum.value", "write_shape": [9, 2, 1024, 4096], "chunk_shape": [9, 2, 1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp_1.linear.value", "write_shape": [9, 4096, 1024], "chunk_shape": [9, 4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_attention_norm.scale.value", "write_shape": [9, 2048], "chunk_shape": [9, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_attention_norm_1.Dense_0.bias.value", "write_shape": [9, 3072], "chunk_shape": [9, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_attention_norm_1.Dense_0.kernel.value", "write_shape": [9, 1024, 3072], "chunk_shape": [9, 1024, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_ffw_norm.scale.value", "write_shape": [9, 2048], "chunk_shape": [9, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_ffw_norm_1.Dense_0.bias.value", "write_shape": [9, 3072], "chunk_shape": [9, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_ffw_norm_1.Dense_0.kernel.value", "write_shape": [9, 1024, 3072], "chunk_shape": [9, 1024, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_in_proj.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_in_proj.kernel.value", "write_shape": [16, 1024], "chunk_shape": [16, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_out_proj.bias.value", "write_shape": [16], "chunk_shape": [16], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_out_proj.kernel.value", "write_shape": [512, 32], "chunk_shape": [512, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_in.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_in.kernel.value", "write_shape": [512, 1024], "chunk_shape": [512, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_out.kernel.value", "write_shape": [512, 1024], "chunk_shape": [512, 1024], "ext_metadata": null}}]} \ No newline at end of file diff --git a/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/d/0cd8ac5577e7944068f6e6974cb7bb0a b/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/d/0cd8ac5577e7944068f6e6974cb7bb0a new file mode 100644 index 0000000000000000000000000000000000000000..852335c9a0a45c7d6237f58c25c12129365e248d Binary files /dev/null and b/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/d/0cd8ac5577e7944068f6e6974cb7bb0a differ diff --git a/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/manifest.ocdbt b/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..a7507334f16dbf7ce59ffe195db1194e233844c1 Binary files /dev/null and b/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/manifest.ocdbt differ diff --git a/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/ocdbt.process_0/d/0142bc15f6939338b6e5adc975f24636 b/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/ocdbt.process_0/d/0142bc15f6939338b6e5adc975f24636 new file mode 100644 index 0000000000000000000000000000000000000000..35b9330ce1ee8caa000d9870cc8ba05e5cb3ddc5 Binary files /dev/null and b/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/ocdbt.process_0/d/0142bc15f6939338b6e5adc975f24636 differ diff --git a/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/ocdbt.process_0/d/25f9d93df7adefdd09f13709ad7405ac b/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/ocdbt.process_0/d/25f9d93df7adefdd09f13709ad7405ac new file mode 100644 index 0000000000000000000000000000000000000000..561f3b4356572afbca1b4044ee459a58dcb6f74e --- /dev/null +++ b/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/ocdbt.process_0/d/25f9d93df7adefdd09f13709ad7405ac @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7940fbcac75dba8531d622ad94b7d91f830c0edc786bae834223a5f77783273c +size 2382306282 diff --git a/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/ocdbt.process_0/d/530bcf4c80475910f788fc4e0549e9db b/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/ocdbt.process_0/d/530bcf4c80475910f788fc4e0549e9db new file mode 100644 index 0000000000000000000000000000000000000000..5c19e4fd5702b0bd40ce4281d7cba05db42ecf7e --- /dev/null +++ b/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/ocdbt.process_0/d/530bcf4c80475910f788fc4e0549e9db @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d64f6fd5e3a4addd1ca579575db8ba7eae8e15ca213cf8cb9459778b0ece86f7 +size 4825414 diff --git a/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/ocdbt.process_0/d/6093a840c30fa64e9625b72899965ef4 b/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/ocdbt.process_0/d/6093a840c30fa64e9625b72899965ef4 new file mode 100644 index 0000000000000000000000000000000000000000..b75a574302c48c05bcd28da15aef426a2d0e570f Binary files /dev/null and b/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/ocdbt.process_0/d/6093a840c30fa64e9625b72899965ef4 differ diff --git a/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/ocdbt.process_0/d/672505ea0f45854a51522b6f72c74956 b/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/ocdbt.process_0/d/672505ea0f45854a51522b6f72c74956 new file mode 100644 index 0000000000000000000000000000000000000000..aabf6e76c91cd1867c50dd4511754639a3d6719f --- /dev/null +++ b/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/ocdbt.process_0/d/672505ea0f45854a51522b6f72c74956 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:337f488e887dcae899a00081f5826e9a8793b0d431afca4417119685941e07b2 +size 1789342724 diff --git a/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/ocdbt.process_0/d/843484256f847393a1fad4f4be6cc579 b/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/ocdbt.process_0/d/843484256f847393a1fad4f4be6cc579 new file mode 100644 index 0000000000000000000000000000000000000000..905dc3b4ba9fb9ba4e20088f9b4212ecf40e9d3a --- /dev/null +++ b/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/ocdbt.process_0/d/843484256f847393a1fad4f4be6cc579 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0bba891f28c71e5eabc18d0ef842ab8c566cc7724ee774ea1de974c87c736c01 +size 2836847227 diff --git a/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/ocdbt.process_0/d/8533c4eac50e34df8425eefdffd94e26 b/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/ocdbt.process_0/d/8533c4eac50e34df8425eefdffd94e26 new file mode 100644 index 0000000000000000000000000000000000000000..fd99a4d91e89f2ec7915702aff8b97d212be9100 Binary files /dev/null and b/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/ocdbt.process_0/d/8533c4eac50e34df8425eefdffd94e26 differ diff --git a/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/ocdbt.process_0/d/bc79931d9f67d3d5681a71493cfe31e3 b/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/ocdbt.process_0/d/bc79931d9f67d3d5681a71493cfe31e3 new file mode 100644 index 0000000000000000000000000000000000000000..15b7da7c04a956cef877e31ed287b3f68589dba5 --- /dev/null +++ b/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/ocdbt.process_0/d/bc79931d9f67d3d5681a71493cfe31e3 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:fbbb74627b315ac9d94aac14908843959bee3d86eeb7727f8caba85df2cb63df +size 154813450 diff --git a/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/ocdbt.process_0/d/c95ff8bf848a5b7950899927920aa250 b/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/ocdbt.process_0/d/c95ff8bf848a5b7950899927920aa250 new file mode 100644 index 0000000000000000000000000000000000000000..b65cfefd1887d6ed44884cacb60c47fd17723f9b Binary files /dev/null and b/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/ocdbt.process_0/d/c95ff8bf848a5b7950899927920aa250 differ diff --git a/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/ocdbt.process_0/d/fcce80bee76d672f607056f601c435a9 b/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/ocdbt.process_0/d/fcce80bee76d672f607056f601c435a9 new file mode 100644 index 0000000000000000000000000000000000000000..0dc6364e35669c2b04349d3de80b2ab6b729a182 Binary files /dev/null and b/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/ocdbt.process_0/d/fcce80bee76d672f607056f601c435a9 differ diff --git a/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/ocdbt.process_0/manifest.ocdbt b/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/ocdbt.process_0/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..a700cc9ed25ae6e88a40e764bea6dfdd726cbc9a Binary files /dev/null and b/openpi/checkpoints/pi05_rm65_4tasks_trim/25000/params/ocdbt.process_0/manifest.ocdbt differ diff --git a/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/_CHECKPOINT_METADATA b/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/_CHECKPOINT_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..1c2e8171e49ae9b86ecb23590d04cac7138d2d4b --- /dev/null +++ b/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/_CHECKPOINT_METADATA @@ -0,0 +1 @@ +{"item_handlers": {"assets": "openpi.training.checkpoints.CallbackHandler", "params": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler", "train_state": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1790349015365642202, "commit_timestamp_nsecs": 1790349128507329631, "custom_metadata": {}} \ No newline at end of file diff --git a/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/assets/rm65/4tasks_v1_trim/norm_stats.json b/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/assets/rm65/4tasks_v1_trim/norm_stats.json new file mode 100644 index 0000000000000000000000000000000000000000..c7c160d8adb381bbff740cd01e1f4d4ea72c0168 --- /dev/null +++ b/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/assets/rm65/4tasks_v1_trim/norm_stats.json @@ -0,0 +1,80 @@ +{ + "norm_stats": { + "state": { + "mean": [ + -0.409750759601593, + 1.3662880659103394, + 0.7612037062644958, + -1.4930087327957153, + 1.149019718170166, + 0.46189895272254944, + 53.88911819458008 + ], + "std": [ + 0.4165157377719879, + 0.22506451606750488, + 0.43782690167427063, + 0.522697925567627, + 0.3444087505340576, + 0.5309966802597046, + 49.84851837158203 + ], + "q01": [ + -1.0531297045386805, + 0.7838840500074804, + -0.39637044419501943, + -2.551590729381487, + 0.4551538976353758, + -1.0030891360439242, + -1e-10 + ], + "q99": [ + 0.20381991231656804, + 1.8083262567565686, + 1.5471665069427205, + -0.47939204971559013, + 1.696906879658781, + 1.4057684653689329, + 99.98000000009996 + ] + }, + "actions": { + "mean": [ + -0.40519973635673523, + 1.3669183254241943, + 0.7708451747894287, + -1.492232322692871, + 1.149193286895752, + 0.4756605625152588, + 49.24359130859375 + ], + "std": [ + 0.4345667362213135, + 0.2386302947998047, + 0.4585062563419342, + 0.5284347534179688, + 0.354095995426178, + 0.5372209548950195, + 49.99427795410156 + ], + "q01": [ + -1.0606446759420682, + 0.7624480928670695, + -0.39244628614266963, + -2.5683194364814876, + 0.4488185941273004, + -1.0152587816618766, + -1e-10 + ], + "q99": [ + 0.2403610593070864, + 1.8286242818624323, + 1.619703484617617, + -0.46531738959911584, + 1.7191881609739013, + 1.414894032184543, + 99.98000000009996 + ] + } + } +} \ No newline at end of file diff --git a/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/_METADATA b/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..c4e4423841b736caddb183c3068573cfb6572b98 --- /dev/null +++ b/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/_METADATA @@ -0,0 +1 @@ +{"tree_metadata": {"('params', 'PaliGemma', 'img', 'Transformer', 'encoder_norm', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoder_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [576]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoder_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoder_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [576]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 576]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_0', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 576]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_1', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 576]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_1', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 576]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 2152]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 576, 4304]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_1', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 576]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_1', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 2152, 1152]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 8, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 576, 16, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 576]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 8, 72, 1152]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 8, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 576, 16, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 8, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 576, 16, 72]}}, "('params', 'PaliGemma', 'img', 'embedding', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [576]}}, "('params', 'PaliGemma', 'img', 'embedding', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 14, 3, 1152]}}, "('params', 'PaliGemma', 'img', 'head', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'PaliGemma', 'img', 'head', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [576, 2048]}}, "('params', 'PaliGemma', 'img', 'pos_embedding', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "pos_embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 128, 1152]}}, "('params', 'PaliGemma', 'llm', 'embedder', 'input_embedding', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "embedder", "key_type": 2}, {"key": "input_embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128576, 2048]}}, "('params', 'PaliGemma', 'llm', 'final_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "final_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'PaliGemma', 'llm', 'final_norm_1', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "final_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('params', 'PaliGemma', 'llm', 'final_norm_1', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "final_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 3072]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'attn_vec_einsum', 'lora_a', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "attn_vec_einsum", "key_type": 2}, {"key": "lora_a", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 8, 256, 16]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'attn_vec_einsum', 'lora_b', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "attn_vec_einsum", "key_type": 2}, {"key": "lora_b", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 8, 16, 2048]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'attn_vec_einsum', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "attn_vec_einsum", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 8, 256, 2048]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'attn_vec_einsum_1', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "attn_vec_einsum_1", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 8, 256, 1024]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'kv_einsum', 'lora_a', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "kv_einsum", "key_type": 2}, {"key": "lora_a", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 2, 1, 2048, 16]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'kv_einsum', 'lora_b', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "kv_einsum", "key_type": 2}, {"key": "lora_b", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 2, 1, 16, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'kv_einsum', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "kv_einsum", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 2, 1, 2048, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'kv_einsum_1', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "kv_einsum_1", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 2, 1, 1024, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'q_einsum', 'lora_a', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_einsum", "key_type": 2}, {"key": "lora_a", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 8, 2048, 16]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'q_einsum', 'lora_b', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_einsum", "key_type": 2}, {"key": "lora_b", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 8, 16, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'q_einsum', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_einsum", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 8, 2048, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'q_einsum_1', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_einsum_1", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 8, 1024, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp', 'gating_einsum', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "gating_einsum", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 2, 2048, 16384]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp', 'gating_einsum_lora_a', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "gating_einsum_lora_a", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 2, 2048, 16]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp', 'gating_einsum_lora_b', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "gating_einsum_lora_b", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 2, 16, 16384]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp', 'linear', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "linear", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 16384, 2048]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp', 'linear_lora_a', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "linear_lora_a", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 16384, 16]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp', 'linear_lora_b', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "linear_lora_b", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 16, 2048]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp_1', 'gating_einsum', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp_1", "key_type": 2}, {"key": "gating_einsum", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 2, 1024, 4096]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp_1', 'linear', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp_1", "key_type": 2}, {"key": "linear", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 4096, 1024]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_attention_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_attention_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 2048]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_attention_norm_1', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_attention_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 3072]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_attention_norm_1', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_attention_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 1024, 3072]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_ffw_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_ffw_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 2048]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_ffw_norm_1', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_ffw_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 3072]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_ffw_norm_1', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_ffw_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 1024, 3072]}}, "('params', 'action_in_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_in_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'action_in_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_in_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 1024]}}, "('params', 'action_out_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_out_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16]}}, "('params', 'action_out_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 32]}}, "('params', 'time_mlp_in', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "time_mlp_in", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'time_mlp_in', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "time_mlp_in", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1024]}}, "('params', 'time_mlp_out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "time_mlp_out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'time_mlp_out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "time_mlp_out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1024]}}}, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null} \ No newline at end of file diff --git a/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/_sharding b/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/_sharding new file mode 100644 index 0000000000000000000000000000000000000000..55bac843f949f8a5e190d87a3514aa7e23dc57d4 --- /dev/null +++ b/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/_sharding @@ -0,0 +1 @@ +{"cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2Rlcl9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2Rlcl9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8wLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8wLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuZW1iZWRkaW5nLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuZW1iZWRkaW5nLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuaGVhZC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuaGVhZC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcucG9zX2VtYmVkZGluZy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZW1iZWRkZXIuaW5wdXRfZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZmluYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZmluYWxfbm9ybV8xLkRlbnNlXzAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZmluYWxfbm9ybV8xLkRlbnNlXzAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scC5nYXRpbmdfZWluc3VtLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scC5nYXRpbmdfZWluc3VtX2xvcmFfYS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scC5nYXRpbmdfZWluc3VtX2xvcmFfYi52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scC5saW5lYXIudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scC5saW5lYXJfbG9yYV9hLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scC5saW5lYXJfbG9yYV9iLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scF8xLmdhdGluZ19laW5zdW0udmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scF8xLmxpbmVhci52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4uYXR0bl92ZWNfZWluc3VtLmxvcmFfYS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4uYXR0bl92ZWNfZWluc3VtLmxvcmFfYi52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4uYXR0bl92ZWNfZWluc3VtLncudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4uYXR0bl92ZWNfZWluc3VtXzEudy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ua3ZfZWluc3VtLmxvcmFfYS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ua3ZfZWluc3VtLmxvcmFfYi52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ua3ZfZWluc3VtLncudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ua3ZfZWluc3VtXzEudy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ucV9laW5zdW0ubG9yYV9hLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ucV9laW5zdW0ubG9yYV9iLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ucV9laW5zdW0udy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ucV9laW5zdW1fMS53LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9hdHRlbnRpb25fbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9hdHRlbnRpb25fbm9ybV8xLkRlbnNlXzAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9hdHRlbnRpb25fbm9ybV8xLkRlbnNlXzAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9mZndfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9mZndfbm9ybV8xLkRlbnNlXzAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9mZndfbm9ybV8xLkRlbnNlXzAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLmFjdGlvbl9pbl9wcm9qLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLmFjdGlvbl9pbl9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLmFjdGlvbl9vdXRfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLmFjdGlvbl9vdXRfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLnRpbWVfbWxwX291dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLnRpbWVfbWxwX291dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLnRpbWVfbWxwX2luLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}","cGFyYW1zLnRpbWVfbWxwX2luLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}]]}}"} \ No newline at end of file diff --git a/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/array_metadatas/process_0 b/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/array_metadatas/process_0 new file mode 100644 index 0000000000000000000000000000000000000000..57bbcf50e800578263c2f2038348b04c71cedc75 --- /dev/null +++ b/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/array_metadatas/process_0 @@ -0,0 +1 @@ +{"array_metadatas": [{"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoder_norm.bias.value", "write_shape": [576], "chunk_shape": [576], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoder_norm.scale.value", "write_shape": [576], "chunk_shape": [576], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_0.bias.value", "write_shape": [27, 576], "chunk_shape": [27, 576], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_0.scale.value", "write_shape": [27, 576], "chunk_shape": [27, 576], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_1.bias.value", "write_shape": [27, 576], "chunk_shape": [27, 576], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_1.scale.value", "write_shape": [27, 576], "chunk_shape": [27, 576], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_0.bias.value", "write_shape": [27, 2152], "chunk_shape": [27, 2152], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_0.kernel.value", "write_shape": [27, 576, 4304], "chunk_shape": [27, 576, 4304], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_1.bias.value", "write_shape": [27, 576], "chunk_shape": [27, 576], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_1.kernel.value", "write_shape": [27, 2152, 1152], "chunk_shape": [27, 2152, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.key.bias.value", "write_shape": [27, 8, 72], "chunk_shape": [27, 8, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.key.kernel.value", "write_shape": [27, 576, 16, 72], "chunk_shape": [27, 576, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.out.bias.value", "write_shape": [27, 576], "chunk_shape": [27, 576], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.out.kernel.value", "write_shape": [27, 8, 72, 1152], "chunk_shape": [27, 8, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.query.bias.value", "write_shape": [27, 8, 72], "chunk_shape": [27, 8, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.query.kernel.value", "write_shape": [27, 576, 16, 72], "chunk_shape": [27, 576, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.value.bias.value", "write_shape": [27, 8, 72], "chunk_shape": [27, 8, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.value.kernel.value", "write_shape": [27, 576, 16, 72], "chunk_shape": [27, 576, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.embedding.bias.value", "write_shape": [576], "chunk_shape": [576], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.embedding.kernel.value", "write_shape": [7, 14, 3, 1152], "chunk_shape": [7, 14, 3, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.head.bias.value", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.head.kernel.value", "write_shape": [576, 2048], "chunk_shape": [576, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.pos_embedding.value", "write_shape": [1, 128, 1152], "chunk_shape": [1, 128, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.embedder.input_embedding.value", "write_shape": [128576, 2048], "chunk_shape": [128576, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.final_norm.scale.value", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.final_norm_1.Dense_0.bias.value", "write_shape": [1536], "chunk_shape": [1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.final_norm_1.Dense_0.kernel.value", "write_shape": [512, 3072], "chunk_shape": [512, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.attn_vec_einsum.lora_a.value", "write_shape": [9, 8, 256, 16], "chunk_shape": [9, 8, 256, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.attn_vec_einsum.lora_b.value", "write_shape": [9, 8, 16, 2048], "chunk_shape": [9, 8, 16, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.attn_vec_einsum.w.value", "write_shape": [9, 8, 256, 2048], "chunk_shape": [9, 8, 256, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.attn_vec_einsum_1.w.value", "write_shape": [9, 8, 256, 1024], "chunk_shape": [9, 8, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.kv_einsum.lora_a.value", "write_shape": [9, 2, 1, 2048, 16], "chunk_shape": [9, 2, 1, 2048, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.kv_einsum.lora_b.value", "write_shape": [9, 2, 1, 16, 256], "chunk_shape": [9, 2, 1, 16, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.kv_einsum.w.value", "write_shape": [9, 2, 1, 2048, 256], "chunk_shape": [9, 2, 1, 2048, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.kv_einsum_1.w.value", "write_shape": [9, 2, 1, 1024, 256], "chunk_shape": [9, 2, 1, 1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.q_einsum.lora_a.value", "write_shape": [9, 8, 2048, 16], "chunk_shape": [9, 8, 2048, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.q_einsum.lora_b.value", "write_shape": [9, 8, 16, 256], "chunk_shape": [9, 8, 16, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.q_einsum.w.value", "write_shape": [9, 8, 2048, 256], "chunk_shape": [9, 8, 2048, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.q_einsum_1.w.value", "write_shape": [9, 8, 1024, 256], "chunk_shape": [9, 8, 1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp.gating_einsum.value", "write_shape": [9, 2, 2048, 16384], "chunk_shape": [9, 2, 2048, 16384], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp.gating_einsum_lora_a.value", "write_shape": [9, 2, 2048, 16], "chunk_shape": [9, 2, 2048, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp.gating_einsum_lora_b.value", "write_shape": [9, 2, 16, 16384], "chunk_shape": [9, 2, 16, 16384], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp.linear.value", "write_shape": [9, 16384, 2048], "chunk_shape": [9, 16384, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp.linear_lora_a.value", "write_shape": [9, 16384, 16], "chunk_shape": [9, 16384, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp.linear_lora_b.value", "write_shape": [9, 16, 2048], "chunk_shape": [9, 16, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp_1.gating_einsum.value", "write_shape": [9, 2, 1024, 4096], "chunk_shape": [9, 2, 1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp_1.linear.value", "write_shape": [9, 4096, 1024], "chunk_shape": [9, 4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_attention_norm.scale.value", "write_shape": [9, 2048], "chunk_shape": [9, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_attention_norm_1.Dense_0.bias.value", "write_shape": [9, 3072], "chunk_shape": [9, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_attention_norm_1.Dense_0.kernel.value", "write_shape": [9, 1024, 3072], "chunk_shape": [9, 1024, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_ffw_norm.scale.value", "write_shape": [9, 2048], "chunk_shape": [9, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_ffw_norm_1.Dense_0.bias.value", "write_shape": [9, 3072], "chunk_shape": [9, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_ffw_norm_1.Dense_0.kernel.value", "write_shape": [9, 1024, 3072], "chunk_shape": [9, 1024, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_in_proj.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_in_proj.kernel.value", "write_shape": [16, 1024], "chunk_shape": [16, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_out_proj.bias.value", "write_shape": [16], "chunk_shape": [16], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_out_proj.kernel.value", "write_shape": [512, 32], "chunk_shape": [512, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_in.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_in.kernel.value", "write_shape": [512, 1024], "chunk_shape": [512, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_out.kernel.value", "write_shape": [512, 1024], "chunk_shape": [512, 1024], "ext_metadata": null}}]} \ No newline at end of file diff --git a/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/d/e610643a131a0630d9536f1fb53ebfde b/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/d/e610643a131a0630d9536f1fb53ebfde new file mode 100644 index 0000000000000000000000000000000000000000..0cb4c709335f00b4604a9b8fd860ee6967b6b24e Binary files /dev/null and b/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/d/e610643a131a0630d9536f1fb53ebfde differ diff --git a/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/manifest.ocdbt b/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..dac45974261ca50599491ae22f80a68ad039b014 Binary files /dev/null and b/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/manifest.ocdbt differ diff --git a/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/ocdbt.process_0/d/226ae31e99c7333a61450a00ba995095 b/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/ocdbt.process_0/d/226ae31e99c7333a61450a00ba995095 new file mode 100644 index 0000000000000000000000000000000000000000..62656a2abd70f75662e62e04af11727ec9aad9dd --- /dev/null +++ b/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/ocdbt.process_0/d/226ae31e99c7333a61450a00ba995095 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:bcfb8454905cc76169b5705a9a836432cbaae7fbe3f2719d7aebacc026a86c8e +size 2213147808 diff --git a/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/ocdbt.process_0/d/25c15778c6ec0c2abc8de29f5e2e2ff9 b/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/ocdbt.process_0/d/25c15778c6ec0c2abc8de29f5e2e2ff9 new file mode 100644 index 0000000000000000000000000000000000000000..305017f5f500b9719dfd1b04a4edd33e8fc5146b Binary files /dev/null and b/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/ocdbt.process_0/d/25c15778c6ec0c2abc8de29f5e2e2ff9 differ diff --git a/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/ocdbt.process_0/d/8bed6baef79be862046a0e3271c1595b b/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/ocdbt.process_0/d/8bed6baef79be862046a0e3271c1595b new file mode 100644 index 0000000000000000000000000000000000000000..abb985103a7794df583e00008be8d5694efa425b Binary files /dev/null and b/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/ocdbt.process_0/d/8bed6baef79be862046a0e3271c1595b differ diff --git a/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/ocdbt.process_0/d/965786e7a16a28ab7bfade6250af4dbc b/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/ocdbt.process_0/d/965786e7a16a28ab7bfade6250af4dbc new file mode 100644 index 0000000000000000000000000000000000000000..c49a7cbf0f99262fcf96b56894dbbd87983a3a95 --- /dev/null +++ b/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/ocdbt.process_0/d/965786e7a16a28ab7bfade6250af4dbc @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a1afeba7b8f34e2ba32ae1ffb324c72dd16bebc8149b328e9eb2b750ed269679 +size 2306356994 diff --git a/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/ocdbt.process_0/d/a1b59a1f3c82e649f8737c81e7144c71 b/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/ocdbt.process_0/d/a1b59a1f3c82e649f8737c81e7144c71 new file mode 100644 index 0000000000000000000000000000000000000000..6f0377521a6c7c87902056d9d719d05182c9a16e --- /dev/null +++ b/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/ocdbt.process_0/d/a1b59a1f3c82e649f8737c81e7144c71 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:bf4a21b1a3bb0fbba8bca2c5817756824c53690bf00d49d45efed4cffcc539f6 +size 1620392329 diff --git a/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/ocdbt.process_0/d/b4b2026c2109777c015c6d32cb3f100e b/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/ocdbt.process_0/d/b4b2026c2109777c015c6d32cb3f100e new file mode 100644 index 0000000000000000000000000000000000000000..b65cfefd1887d6ed44884cacb60c47fd17723f9b Binary files /dev/null and b/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/ocdbt.process_0/d/b4b2026c2109777c015c6d32cb3f100e differ diff --git a/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/ocdbt.process_0/d/be600bb25bbb4a56f17d7e8ac6a7d170 b/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/ocdbt.process_0/d/be600bb25bbb4a56f17d7e8ac6a7d170 new file mode 100644 index 0000000000000000000000000000000000000000..138633cd6cf452a10ea5f945c46130737284b5a3 --- /dev/null +++ b/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/ocdbt.process_0/d/be600bb25bbb4a56f17d7e8ac6a7d170 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:08029d1cf8a1e37e28ad0d3cc36397785414415fb9763d617e78044595f1314f +size 82879978 diff --git a/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/ocdbt.process_0/d/f65dc4a06590f10c9806e235cda540c2 b/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/ocdbt.process_0/d/f65dc4a06590f10c9806e235cda540c2 new file mode 100644 index 0000000000000000000000000000000000000000..3d897e40ed5ac7d171fc315c60a935b658bfa527 --- /dev/null +++ b/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/ocdbt.process_0/d/f65dc4a06590f10c9806e235cda540c2 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8658f4f777b612c7550b5139dfafd93da181bddb8067d26bd0e4e7afed89f408 +size 945351240 diff --git a/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/ocdbt.process_0/d/fbc96d675129f98b3f7f0a09a0629280 b/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/ocdbt.process_0/d/fbc96d675129f98b3f7f0a09a0629280 new file mode 100644 index 0000000000000000000000000000000000000000..e4d6be1311fae065f0efb3b0ebcdd789e13d2b6f Binary files /dev/null and b/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/ocdbt.process_0/d/fbc96d675129f98b3f7f0a09a0629280 differ diff --git a/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/ocdbt.process_0/manifest.ocdbt b/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/ocdbt.process_0/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..59589fab16a1357fe711419c5a95b49ec4ff0b9e Binary files /dev/null and b/openpi/checkpoints/pi05_rm65_4tasks_trim/29999/params/ocdbt.process_0/manifest.ocdbt differ diff --git a/openpi/code/rm65_config_excerpt.py b/openpi/code/rm65_config_excerpt.py new file mode 100644 index 0000000000000000000000000000000000000000..d8fa5e28d03c480336b1f9fd16efa7df1281aa45 --- /dev/null +++ b/openpi/code/rm65_config_excerpt.py @@ -0,0 +1,88 @@ +# RM65 excerpt of openpi-fintune src/openpi/training/config.py. +# Paste LeRobotRM65DataConfig next to the other DataConfigFactory classes and the TrainConfig into _CONFIGS. +# Requires openpi.policies.b601_policy (included under code/src/openpi/policies/). + +# ---- DataConfigFactory (module level) ---- +@dataclasses.dataclass(frozen=True) +class LeRobotRM65DataConfig(DataConfigFactory): + """Data config for fine-tuning on the RealMan RM65 right-arm LeRobot v2.1 dataset. + + Same layout as B601 -- 7-dim state/action (right_q1..q6 in rad, right_gripper 0=closed..100=open) + and two cameras (top, wrist) -- so the B601 policy transforms are reused as-is. Unlike B601, + observation.state and action share one coordinate convention, so the joint actions can be + trained as deltas relative to the current state. + """ + + # Injected as the prompt for every sample. Leave None to take the prompt from the dataset's tasks.jsonl. + default_prompt: str | None = None + # Train the 6 joints as deltas relative to the current state; the gripper stays absolute. + use_delta_joint_actions: bool = False + + @override + def create(self, assets_dirs: pathlib.Path, model_config: _model.BaseModelConfig) -> DataConfig: + repack_transform = _transforms.Group( + inputs=[ + _transforms.RepackTransform( + { + "observation/top": "observation.images.top", + "observation/wrist": "observation.images.wrist", + "observation/state": "observation.state", + "actions": "action", + "prompt": "prompt", + } + ) + ] + ) + + data_transforms = _transforms.Group( + inputs=[b601_policy.B601Inputs(model_type=model_config.model_type)], + outputs=[b601_policy.B601Outputs()], + ) + if self.use_delta_joint_actions: + delta_action_mask = _transforms.make_bool_mask(6, -1) + data_transforms = data_transforms.push( + inputs=[_transforms.DeltaActions(delta_action_mask)], + outputs=[_transforms.AbsoluteActions(delta_action_mask)], + ) + + model_transforms = ModelTransformFactory(default_prompt=self.default_prompt)(model_config) + + return dataclasses.replace( + self.create_base_config(assets_dirs, model_config), + repack_transforms=repack_transform, + data_transforms=data_transforms, + model_transforms=model_transforms, + ) + + + +# ---- TrainConfig (entry of _CONFIGS) ---- + # Same as pi05_rm65_4tasks, but on rm65/4tasks_v1_trim (scripts/trim_static_frames.py): the static head + # of every episode is dropped and the static tail is cut to 10 frames, 66,121 -> 58,399 frames. + # Short mid-episode pauses are kept. + TrainConfig( + name="pi05_rm65_4tasks_trim", + model=pi0_config.Pi0Config( + pi05=True, + action_horizon=30, + paligemma_variant="gemma_2b_lora", + ), + data=LeRobotRM65DataConfig( + repo_id="rm65/4tasks_v1_trim", + base_config=DataConfig( + local_root=pathlib.Path("/shared/user64/workspace/yuhao/pi/data"), + prompt_from_task=True, + action_sequence_keys=("action",), + ), + ), + weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"), + freeze_filter=pi0_config.Pi0Config( + pi05=True, + action_horizon=30, + paligemma_variant="gemma_2b_lora", + ).get_freeze_filter(), + num_train_steps=30_000, + batch_size=32, + num_workers=8, + ema_decay=None, + ), diff --git a/openpi/code/scripts/merge_datasets.py b/openpi/code/scripts/merge_datasets.py new file mode 100644 index 0000000000000000000000000000000000000000..7dfb9ea0b4be79885792c194c02b05b31948f85a --- /dev/null +++ b/openpi/code/scripts/merge_datasets.py @@ -0,0 +1,169 @@ +"""Merge several LeRobot (v2.1) datasets into one, re-indexed and with a shared +task table. The multi-task counterpart of split_dataset.py. + +Episodes are concatenated in the order the sources are given. Each episode is +re-numbered, its parquet's `episode_index` / `index` / `task_index` columns are +rewritten, and its videos are hard-linked (same filesystem) or copied. Task +strings are deduplicated across sources, so two datasets sharing a prompt share +one task_index. + +Example: + python scripts/merge_datasets.py \ + --srcs ~/workspace/yuhao/pi/data/item-classification-eef \ + ~/workspace/yuhao/pi/data/sort-tools-eef \ + ~/workspace/yuhao/pi/data/stack-cube-eef \ + --out ~/workspace/yuhao/pi/data/g1-multitask-eef + +Sources are left untouched. Only the metadata every LeRobot loader needs is +written (info/tasks/episodes/episodes_stats); consumer-specific extras that some +sources happen to carry (modality.json, stats.json, relative_stats.json) are NOT +merged, because they describe one source's statistics and would be wrong for the +union — regenerate them for the merged set if a downstream tool needs them. +""" + +import dataclasses +import json +import pathlib +import shutil + +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq +import tyro + + +@dataclasses.dataclass +class Args: + # Source LeRobot dataset directories, concatenated in this order. + srcs: list[str] + # Output directory for the merged dataset. + out: str + # Copy videos instead of hard-linking them (needed across filesystems). + copy: bool = False + + +def _link_or_copy(src: pathlib.Path, dst: pathlib.Path, force_copy: bool) -> None: + dst.parent.mkdir(parents=True, exist_ok=True) + if force_copy: + shutil.copy2(src, dst) + return + try: + dst.hardlink_to(src) + except OSError: # different filesystem, or links unsupported + shutil.copy2(src, dst) + + +def _load_jsonl(path: pathlib.Path) -> dict: + return {json.loads(l)["episode_index"]: json.loads(l) for l in path.open()} + + +def _check_compatible(infos: list[dict], srcs: list[pathlib.Path]) -> None: + """Every field that must agree for the episodes to be interchangeable. + + `features` is compared whole: it pins dtypes, shapes, joint/EEF column names + and the video codec settings in one go. A mismatch here means the merged + dataset would feed a model two different action spaces or two different + decoders, which no downstream loader would catch. + """ + keys = ["codebase_version", "robot_type", "fps", "chunks_size", + "data_path", "video_path", "features"] + for key in keys: + values = [info.get(key) for info in infos] + if any(v != values[0] for v in values): + detail = "\n".join(f" {s.name}: {json.dumps(v)[:200]}" + for s, v in zip(srcs, values)) + raise SystemExit(f"Sources disagree on info.json['{key}']:\n{detail}") + + +def main(args: Args) -> None: + srcs = [pathlib.Path(s).expanduser().resolve() for s in args.srcs] + out = pathlib.Path(args.out).expanduser().resolve() + if len(srcs) < 2: + raise SystemExit("Need at least two --srcs to merge.") + if out in srcs: + raise SystemExit("--out must differ from every source.") + if out.exists(): + raise SystemExit(f"--out already exists: {out}") + + infos = [json.loads((s / "meta/info.json").read_text()) for s in srcs] + _check_compatible(infos, srcs) + info = infos[0] + chunks_size = info["chunks_size"] + data_tmpl, video_tmpl = info["data_path"], info["video_path"] + video_keys = [k for k, v in info["features"].items() if v["dtype"] == "video"] + + # Shared task table: first-seen order, deduplicated across sources. + task_ids: dict[str, int] = {} + remaps = [] + for src in srcs: + remap = {} + for line in (src / "meta/tasks.jsonl").open(): + row = json.loads(line) + remap[row["task_index"]] = task_ids.setdefault(row["task"], len(task_ids)) + remaps.append(remap) + + (out / "meta").mkdir(parents=True, exist_ok=True) + new_eps, new_stats = [], [] + new_ep, running = 0, 0 + + for src, src_info, remap in zip(srcs, infos, remaps): + eps = _load_jsonl(src / "meta/episodes.jsonl") + stats = _load_jsonl(src / "meta/episodes_stats.jsonl") + n_eps = src_info["total_episodes"] + print(f"{src.name}: {n_eps} episodes, {src_info['total_frames']} frames " + f"-> episodes {new_ep}..{new_ep + n_eps - 1}") + + for old in range(n_eps): + src_chunk, dst_chunk = old // chunks_size, new_ep // chunks_size + table = pq.read_table( + src / data_tmpl.format(episode_chunk=src_chunk, episode_index=old)) + n = table.num_rows + + for name, values in ( + ("episode_index", pa.array([new_ep] * n, pa.int64())), + ("index", pa.array(np.arange(running, running + n), pa.int64())), + ("task_index", pa.array( + [remap[t] for t in table.column("task_index").to_pylist()], pa.int64())), + ): + table = table.set_column(table.schema.get_field_index(name), name, values) + + dst_pq = out / data_tmpl.format(episode_chunk=dst_chunk, episode_index=new_ep) + dst_pq.parent.mkdir(parents=True, exist_ok=True) + pq.write_table(table, dst_pq) + + for vk in video_keys: + _link_or_copy( + src / video_tmpl.format(episode_chunk=src_chunk, video_key=vk, episode_index=old), + out / video_tmpl.format(episode_chunk=dst_chunk, video_key=vk, episode_index=new_ep), + args.copy) + + new_eps.append({**eps[old], "episode_index": new_ep}) + new_stats.append({**stats[old], "episode_index": new_ep}) + new_ep += 1 + running += n + + with (out / "meta/episodes.jsonl").open("w") as f: + f.writelines(json.dumps(e) + "\n" for e in new_eps) + with (out / "meta/episodes_stats.jsonl").open("w") as f: + f.writelines(json.dumps(s) + "\n" for s in new_stats) + with (out / "meta/tasks.jsonl").open("w") as f: + f.writelines(json.dumps({"task_index": i, "task": t}) + "\n" + for t, i in sorted(task_ids.items(), key=lambda kv: kv[1])) + + (out / "meta/info.json").write_text(json.dumps({ + **info, + "total_episodes": new_ep, + "total_frames": running, + "total_tasks": len(task_ids), + "total_videos": new_ep * len(video_keys), + "total_chunks": (new_ep - 1) // chunks_size + 1 if new_ep else 0, + "splits": {"train": f"0:{new_ep}"}, + }, indent=4)) + + print(f"\nMerged -> {out}") + print(f" {new_ep} episodes, {running} frames, {len(task_ids)} tasks, " + f"{new_ep * len(video_keys)} videos") + + +if __name__ == "__main__": + main(tyro.cli(Args)) diff --git a/openpi/code/scripts/trim_static_frames.py b/openpi/code/scripts/trim_static_frames.py new file mode 100644 index 0000000000000000000000000000000000000000..bd4aa6dedb0aaaea51ea1aceccfaccad25abb63f --- /dev/null +++ b/openpi/code/scripts/trim_static_frames.py @@ -0,0 +1,152 @@ +"""Trim the static head and tail of every episode in a LeRobot (v2.1) dataset. + +A frame is "moving" if either observation.state or action changes between it and +the next frame. Everything before the first moving frame is dropped, and the +static tail after the last moving frame is cut down to `--keep-tail` frames, so +the policy still sees the arm come to rest. Short pauses in the middle of an +episode are kept: cutting them would break the timing of the action chunks. + +Videos are hard-linked, not re-encoded. Each kept row keeps its original +`timestamp` and `frame_index`, so LeRobot still decodes the matching frame from +the untouched per-episode mp4 (it only checks that timestamps within an episode +are 1/fps apart, not that they start at 0). `index` is renumbered globally and +episodes.jsonl / episodes_stats.jsonl are rewritten for the new lengths. + +Example: + python scripts/trim_static_frames.py \ + --src ~/workspace/yuhao/pi/data/rm65/4tasks_v1 \ + --out ~/workspace/yuhao/pi/data/rm65/4tasks_v1_trim +""" + +import dataclasses +import json +import pathlib +import shutil + +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq +import tyro + + +@dataclasses.dataclass +class Args: + # Source LeRobot v2.1 dataset directory. + src: str + # Output directory for the trimmed dataset. + out: str + # Static frames kept after the last moving frame. + keep_tail: int = 10 + # Absolute change below which a state/action value counts as unchanged. + eps: float = 1e-6 + + +def _link_or_copy(src: pathlib.Path, dst: pathlib.Path) -> None: + dst.parent.mkdir(parents=True, exist_ok=True) + try: + dst.hardlink_to(src) + except OSError: # different filesystem, or links unsupported + shutil.copy2(src, dst) + + +def _vector_stats(x: np.ndarray) -> dict: + return { + "min": x.min(0).tolist(), + "max": x.max(0).tolist(), + "mean": x.mean(0).tolist(), + "std": x.std(0).tolist(), + "count": [len(x)], + } + + +def main(args: Args) -> None: + src = pathlib.Path(args.src).expanduser().resolve() + out = pathlib.Path(args.out).expanduser().resolve() + if out == src: + raise SystemExit("--out must differ from --src.") + if out.exists(): + raise SystemExit(f"--out already exists: {out}") + + info = json.loads((src / "meta/info.json").read_text()) + if info["codebase_version"] != "v2.1": + raise SystemExit(f"Expected a v2.1 dataset, got {info['codebase_version']}") + chunks_size = info["chunks_size"] + data_tmpl, video_tmpl = info["data_path"], info["video_path"] + video_keys = [k for k, v in info["features"].items() if v["dtype"] == "video"] + + eps = [json.loads(l) for l in (src / "meta/episodes.jsonl").open()] + stats = {json.loads(l)["episode_index"]: json.loads(l) for l in (src / "meta/episodes_stats.jsonl").open()} + + (out / "meta").mkdir(parents=True) + shutil.copy2(src / "meta/tasks.jsonl", out / "meta/tasks.jsonl") + new_eps, new_stats = [], [] + running, dropped_head, dropped_tail = 0, 0, 0 + + for ep in eps: + ep_idx = ep["episode_index"] + chunk = ep_idx // chunks_size + table = pq.read_table(src / data_tmpl.format(episode_chunk=chunk, episode_index=ep_idx)) + state = np.stack(table.column("observation.state").to_numpy(zero_copy_only=False)) + action = np.stack(table.column("action").to_numpy(zero_copy_only=False)) + n = len(state) + + moving = (np.abs(np.diff(state, axis=0)).max(1) > args.eps) | ( + np.abs(np.diff(action, axis=0)).max(1) > args.eps + ) + if not moving.any(): + raise SystemExit(f"Episode {ep_idx} never moves; refusing to trim it to nothing.") + start = int(np.argmax(moving)) # first frame that moves to the next one + last_move = n - 1 - int(np.argmax(moving[::-1])) # frame reached by the last movement + end = min(n, last_move + 1 + args.keep_tail) + dropped_head += start + dropped_tail += n - end + + table = table.slice(start, end - start) + m = table.num_rows + table = table.set_column( + table.schema.get_field_index("index"), "index", pa.array(np.arange(running, running + m), pa.int64()) + ) + dst_pq = out / data_tmpl.format(episode_chunk=chunk, episode_index=ep_idx) + dst_pq.parent.mkdir(parents=True, exist_ok=True) + pq.write_table(table, dst_pq) + + for vk in video_keys: + _link_or_copy( + src / video_tmpl.format(episode_chunk=chunk, video_key=vk, episode_index=ep_idx), + out / video_tmpl.format(episode_chunk=chunk, video_key=vk, episode_index=ep_idx), + ) + + # action_config frame ranges are relative to the episode start. + new_ep = {**ep, "length": m} + if "action_config" in ep: + new_ep["action_config"] = [ + {**c, "start_frame": max(0, c["start_frame"] - start), "end_frame": min(m, c["end_frame"] - start)} + for c in ep["action_config"] + if c["end_frame"] > start and c["start_frame"] - start < m + ] + new_eps.append(new_ep) + + # Recompute the per-episode stats that depend on the kept rows; video stats are kept as-is. + st = dict(stats[ep_idx]["stats"]) + st["observation.state"] = _vector_stats(state[start:end]) + st["action"] = _vector_stats(action[start:end]) + for key in ("timestamp", "frame_index", "episode_index", "index", "task_index"): + col = table.column(key).to_numpy().astype(np.float64)[:, None] + st[key] = _vector_stats(col) + new_stats.append({"episode_index": ep_idx, "stats": st}) + running += m + + with (out / "meta/episodes.jsonl").open("w") as f: + f.writelines(json.dumps(e) + "\n" for e in new_eps) + with (out / "meta/episodes_stats.jsonl").open("w") as f: + f.writelines(json.dumps(s) + "\n" for s in new_stats) + (out / "meta/info.json").write_text(json.dumps({**info, "total_frames": running}, indent=4)) + + total = running + dropped_head + dropped_tail + print(f"Trimmed -> {out}") + print(f" {len(new_eps)} episodes, {total} -> {running} frames " + f"(dropped {dropped_head} head + {dropped_tail} tail, {1 - running / total:.1%})") + + +if __name__ == "__main__": + main(tyro.cli(Args)) diff --git a/openpi/code/src/openpi/policies/b601_policy.py b/openpi/code/src/openpi/policies/b601_policy.py new file mode 100644 index 0000000000000000000000000000000000000000..eb5b963cac6a5edab54e6014fb78b101ee491ae0 --- /dev/null +++ b/openpi/code/src/openpi/policies/b601_policy.py @@ -0,0 +1,114 @@ +import dataclasses +from typing import NotRequired, TypedDict + +import einops +import numpy as np + +from openpi import transforms +from openpi.models import model as _model + +# Seeed B601 (single-arm follower) action/state layout (7 dims), degrees: +# 0 shoulder_pan +# 1 shoulder_lift +# 2 elbow_flex +# 3 wrist_flex +# 4 wrist_yaw +# 5 wrist_roll +# 6 gripper +# observation.state and action do NOT share a coordinate convention: four dims are sign-flipped and +# the gripper uses a different unit (state spans [-270, 0], action spans [0, 57]). See the internal B601 notes, issue I1. +B601_ACTION_DIM = 7 +B601_STATE_DIM = 7 +B601_JOINT_NAMES = ( + "shoulder_pan", + "shoulder_lift", + "elbow_flex", + "wrist_flex", + "wrist_yaw", + "wrist_roll", + "gripper", +) +# Native camera resolution of the delivered dataset; ResizeImages scales to 224x224 downstream. +B601_IMAGE_HW = (480, 640) + + +class B601ModelInputs(TypedDict): + """Key set that B601Inputs emits. + + Declared explicitly because transforms.DataTransformFn fixes the runtime type to a plain dict. + """ + + state: np.ndarray + image: dict[str, np.ndarray] + image_mask: dict[str, np.bool_] + actions: NotRequired[np.ndarray] + prompt: NotRequired[str] + + +def make_b601_example() -> dict: + """Creates a random input example for the B601 policy (matches the inference key format).""" + return { + "observation/state": np.random.rand(B601_STATE_DIM), + "observation/top": np.random.randint(256, size=(*B601_IMAGE_HW, 3), dtype=np.uint8), + "observation/wrist": np.random.randint(256, size=(*B601_IMAGE_HW, 3), dtype=np.uint8), + "prompt": "Pick up the red cube and place it completely inside the tray.", + } + + +def _parse_image(image) -> np.ndarray: + image = np.asarray(image) + if np.issubdtype(image.dtype, np.floating): + image = (255 * image).astype(np.uint8) + if image.shape[0] == 3: + image = einops.rearrange(image, "c h w -> h w c") + return image + + +@dataclasses.dataclass(frozen=True) +class B601Inputs(transforms.DataTransformFn): + """Converts B601 observations into the format expected by the model. + + The B601 rig has only two cameras. The unused right wrist view is zero-filled and masked out, so + its tokens are dropped from the prefix attention (see pi0.py embed_prefix). + """ + + model_type: _model.ModelType + + def __call__(self, data: dict) -> dict: + base_image = _parse_image(data["observation/top"]) + wrist_image = _parse_image(data["observation/wrist"]) + + # pi0-FAST has no image mask, so the padding view must be marked present there. + mask_padding = self.model_type != _model.ModelType.PI0_FAST + + inputs = { + "state": data["observation/state"], + "image": { + "base_0_rgb": base_image, + "left_wrist_0_rgb": wrist_image, + "right_wrist_0_rgb": np.zeros_like(base_image), + }, + "image_mask": { + "base_0_rgb": np.True_, + "left_wrist_0_rgb": np.True_, + "right_wrist_0_rgb": np.False_ if mask_padding else np.True_, + }, + } + + # Actions are only available during training. + if "actions" in data: + inputs["actions"] = data["actions"] + + if "prompt" in data: + inputs["prompt"] = data["prompt"] + + return inputs + + +@dataclasses.dataclass(frozen=True) +class B601Outputs(transforms.DataTransformFn): + """Converts model actions back to the B601 7-dim joint command vector.""" + + def __call__(self, data: dict) -> dict: + # Strip the zero padding that PadStatesAndActions added to reach the model action dim. + return {"actions": np.asarray(data["actions"][..., :B601_ACTION_DIM])}