Add files using upload-large-folder tool

.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

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/<step>/` 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=<this repo>/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 `<IW>/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.**

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}

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

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

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

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

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

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

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

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
+}

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}" \
+  "$@"

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": {}}

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
+      ]
+    }
+  }
+}

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}

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}]]}}"}

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}}]}

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

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

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

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

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

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": {}}

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
+      ]
+    }
+  }
+}

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}

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}]]}}"}

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}}]}

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

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

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

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

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

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,
+    ),

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))

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))