diff --git a/.gitattributes b/.gitattributes
index a20190aa4cd8726192c578fa09f2f28acca4a8ed..1ad725d19029cb95b1ec03a4f60a1d5f2d43c013 100644
--- a/.gitattributes
+++ b/.gitattributes
@@ -1174,3 +1174,25 @@ test_loss_components_1e-4/14000/params/ocdbt.process_0/d/38a05656511bafae85f0a0e
 test_loss_components_1e-4/14000/params/ocdbt.process_0/d/ef329dfa763e809ee236decfdcc52c3c filter=lfs diff=lfs merge=lfs -text
 test_loss_components_1e-4/14000/params/ocdbt.process_0/d/0a57f14b0215090db29279e065a5c915 filter=lfs diff=lfs merge=lfs -text
 test_loss_components_1e-4/14000/params/ocdbt.process_0/d/ba84851d063dcfd19350d6851293a9c2 filter=lfs diff=lfs merge=lfs -text
+test_loss_components_1e-5/4000/params/ocdbt.process_0/d/56dd749ddfd4e8994705b08a5a96e495 filter=lfs diff=lfs merge=lfs -text
+test_loss_components_1e-5/4000/params/ocdbt.process_0/d/c3e96dcec38a9d1597affe33f8b76478 filter=lfs diff=lfs merge=lfs -text
+test_loss_components_1e-5/4000/params/ocdbt.process_0/d/ccaedfef9f74fc9a1d34cd70c5b165f1 filter=lfs diff=lfs merge=lfs -text
+test_loss_components_1e-5/4000/params/ocdbt.process_0/d/216e6d52245be1f029b9e88e58a7e551 filter=lfs diff=lfs merge=lfs -text
+test_loss_components_1e-5/4000/params/ocdbt.process_0/d/13d0383bcd36a0d9f61afb639ad380c9 filter=lfs diff=lfs merge=lfs -text
+test_loss_components_1e-5/4000/params/ocdbt.process_0/d/8a992cd26ed86fb77ced871773dc4594 filter=lfs diff=lfs merge=lfs -text
+test_loss_components_1e-5/4000/params/ocdbt.process_0/d/fc7b03c165579d42dbfb89800de58cc2 filter=lfs diff=lfs merge=lfs -text
+test_loss_components_1e-5/4000/params/ocdbt.process_0/d/70410d314e2b7a69f49c646d900ca3c8 filter=lfs diff=lfs merge=lfs -text
+test_loss_components_1e-5/4000/params/ocdbt.process_0/d/f471d998aea2f39c6b4792cbe5a8b652 filter=lfs diff=lfs merge=lfs -text
+test_loss_components_1e-5/4000/params/ocdbt.process_0/d/23c9da5d1b786076230c741026d74240 filter=lfs diff=lfs merge=lfs -text
+test_loss_components_1e-5/4000/params/ocdbt.process_0/d/fc8f8e8c0cd651f926f070ccb0f29d32 filter=lfs diff=lfs merge=lfs -text
+test_loss_components_1e-5/4000/params/ocdbt.process_0/d/dcb557aed68e94285b47236d2e8c087a filter=lfs diff=lfs merge=lfs -text
+test_loss_components_1e-5/4000/params/ocdbt.process_0/d/cccbe3fdf6d794f2bf59f994ed28b42b filter=lfs diff=lfs merge=lfs -text
+test_loss_components_1e-5/4000/params/ocdbt.process_0/d/f3a06a050c3562508e08ba9b4d78de80 filter=lfs diff=lfs merge=lfs -text
+test_loss_components_1e-5/4000/params/ocdbt.process_0/d/9cbc768e3b41e7497483908b73578383 filter=lfs diff=lfs merge=lfs -text
+test_loss_components_1e-5/4000/params/ocdbt.process_0/d/b313f73ca4a074edf26764d407a78285 filter=lfs diff=lfs merge=lfs -text
+test_loss_components_1e-5/4000/params/ocdbt.process_0/d/86be6ff68f140be58a980cc11f42cf4e filter=lfs diff=lfs merge=lfs -text
+test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/ocdbt.process_0/d/93444dbe916f143cd1bf761636a8f948 filter=lfs diff=lfs merge=lfs -text
+test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/ocdbt.process_0/d/5f7448af92b6085f7cdc3a64d36c0aa5 filter=lfs diff=lfs merge=lfs -text
+test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/ocdbt.process_0/d/a2c409f30fa48415606b71cd736a7b93 filter=lfs diff=lfs merge=lfs -text
+test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/ocdbt.process_0/d/8d80895564ad25dbd4aa5410b7c7d4b0 filter=lfs diff=lfs merge=lfs -text
+test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/ocdbt.process_0/d/2f93887a4bbd190673a669b2d374e96c filter=lfs diff=lfs merge=lfs -text
diff --git a/test_loss_components_1e-5/4000/_CHECKPOINT_METADATA b/test_loss_components_1e-5/4000/_CHECKPOINT_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..208b786bac81dfc5c50961e0c1ccb0b978e3dd3f
--- /dev/null
+++ b/test_loss_components_1e-5/4000/_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": 1765319772767366267, "commit_timestamp_nsecs": 1765319794162024888, "custom": {}}
\ No newline at end of file
diff --git a/test_loss_components_1e-5/4000/assets/ur10e/norm_stats.json b/test_loss_components_1e-5/4000/assets/ur10e/norm_stats.json
new file mode 100644
index 0000000000000000000000000000000000000000..85a8eaab4b85a50ec7b3633ecb653ab8e0366022
--- /dev/null
+++ b/test_loss_components_1e-5/4000/assets/ur10e/norm_stats.json
@@ -0,0 +1,491 @@
+{
+  "norm_stats": {
+    "state": {
+      "mean": [
+        0.00019428042287472636,
+        0.00020032725296914577,
+        8.52593166200677e-6,
+        0.00019105659157503396,
+        0.0001996985956793651,
+        0.0000162034084496554,
+        0.00019352274830453098,
+        0.00020078709349036217,
+        0.000029199256459833123,
+        0.00019142567180097103,
+        0.00020083319395780563,
+        0.00003592482607928105,
+        0.00005053714266978204,
+        -0.00003347316669533029,
+        -0.00020081466936971992,
+        0.00018982442270498723,
+        0.00020038991351611912,
+        0.000046112203563097864,
+        0.00004985106716048904,
+        -0.00003275063500041142,
+        -0.0002042762062046677,
+        0.5986685752868652,
+        0.5981546640396118,
+        0.05039028078317642,
+        0.0004953373572789133,
+        0.10711561888456345,
+        -0.0005407895660027862,
+        -0.006484619807451963,
+        0.003617143025621772,
+        0.5981546640396118,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.0023281769827008247,
+        0.0034390955697745085,
+        0.002846071030944586,
+        0.0023422744125127792,
+        0.003452499397099018,
+        0.0028555369935929775,
+        0.002345019020140171,
+        0.0034702608827501535,
+        0.0028640974778681993,
+        0.0023441859520971775,
+        0.0034640957601368427,
+        0.0028661703690886497,
+        0.0012911600060760975,
+        0.0020082774572074413,
+        0.01001499593257904,
+        0.0023556631058454514,
+        0.0034643616527318954,
+        0.002885110443457961,
+        0.0012921782908961177,
+        0.0020203019957989454,
+        0.010073363780975342,
+        0.4901678264141083,
+        0.4902710020542145,
+        0.11618480086326599,
+        0.1681893765926361,
+        0.07788223028182983,
+        0.04501405730843544,
+        0.0661778524518013,
+        0.5958529710769653,
+        0.4902710020542145,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        -0.0070394672028720375,
+        -0.011599844336509704,
+        -0.0054659975245594985,
+        -0.007079721701145172,
+        -0.011622210478782653,
+        -0.005539767185784876,
+        -0.007091582050547003,
+        -0.011584406599402429,
+        -0.0054779203861951824,
+        -0.0070481276120990515,
+        -0.011584554602950811,
+        -0.005540505648031831,
+        -0.0038997788466513163,
+        -0.007352302264794708,
+        -0.034272025059163574,
+        -0.007114377491176128,
+        -0.011688268885016442,
+        -0.005550548636354506,
+        -0.003901443669199943,
+        -0.007246576420217752,
+        -0.034621239319443706,
+        0.0,
+        0.0,
+        -0.19741441082954406,
+        -0.3133266278743744,
+        0.00863822932690382,
+        -0.11799316820502281,
+        -0.1673899977207184,
+        -1.4435412939071655,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.007438025734573604,
+        0.011892002701759337,
+        0.009434655957296489,
+        0.007451299774646758,
+        0.011834505152702332,
+        0.009366847675666212,
+        0.007520044301077725,
+        0.01185322367250919,
+        0.009355697490274907,
+        0.007454063414782285,
+        0.011923038901388648,
+        0.009385530477575958,
+        0.004192294634878635,
+        0.005776847817935048,
+        0.033915841154754156,
+        0.007474094156175851,
+        0.011830379698425528,
+        0.009328839504159987,
+        0.004122480013966562,
+        0.005790175495669243,
+        0.03415216095745563,
+        0.9998,
+        0.9998,
+        0.29779768657684336,
+        0.406406463944912,
+        0.3194635294646025,
+        0.1383301957249642,
+        0.1970333372354507,
+        1.6112462955951692,
+        0.9998,
+        0.0,
+        0.0
+      ]
+    },
+    "actions": {
+      "mean": [
+        0.00005176717240829021,
+        0.00012534917914308608,
+        0.0002480015973560512,
+        0.00005367679477785714,
+        -3.479610484191653e-7,
+        -0.0002569151984062046,
+        0.5995351076126099,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.002336727222427726,
+        0.0034517920576035976,
+        0.003022634657099843,
+        0.0012890173820778728,
+        0.0020518647506833076,
+        0.010340424254536629,
+        0.489992618560791,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        -0.006998317793011664,
+        -0.011659257435053586,
+        -0.0055110766027122735,
+        -0.003918213281407953,
+        -0.007409645091742276,
+        -0.03437663650512694,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.007169186890125277,
+        0.0115618407972157,
+        0.009297586680948735,
+        0.004103694232180714,
+        0.005846588513627651,
+        0.034828110963106174,
+        0.9998,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    },
+    "previous_actions": {
+      "mean": [
+        0.00005176717240829021,
+        0.00012534917914308608,
+        0.0002480015973560512,
+        0.00005367679477785714,
+        -3.479610484191653e-7,
+        -0.0002569151984062046,
+        0.5995351076126099,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.002336727222427726,
+        0.0034517920576035976,
+        0.003022634657099843,
+        0.0012890173820778728,
+        0.0020518647506833076,
+        0.010340424254536629,
+        0.489992618560791,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        -0.006998317793011664,
+        -0.011659257435053586,
+        -0.0055110766027122735,
+        -0.003918213281407953,
+        -0.007409645091742276,
+        -0.03437663650512694,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.007169186890125277,
+        0.0115618407972157,
+        0.009297586680948735,
+        0.004103694232180714,
+        0.005846588513627651,
+        0.034828110963106174,
+        0.9998,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    }
+  },
+  "sample_ratio": 0.05,
+  "dataset_infos": {
+    "robotgeneralist/pp__1__train": {
+      "commit_hash": "b8158e3f7eb1eb4fcad944d3dfbcdba2796a214b",
+      "num_episodes": 166,
+      "num_frames": 69615
+    },
+    "robotgeneralist/pp__1__adversarial": {
+      "commit_hash": "a8b6024ff9f5261e87d553f35f07f16487adabc1",
+      "num_episodes": 64,
+      "num_frames": 17906
+    },
+    "robotgeneralist/pp__2": {
+      "commit_hash": "f244ff639c1a837fdec7990c601bf5ed4fbb56a9",
+      "num_episodes": 92,
+      "num_frames": 37458
+    },
+    "robotgeneralist/pp__2__rotations": {
+      "commit_hash": "9b8e02337fce6bccda7adcc4aae1edcbe5b3de41",
+      "num_episodes": 102,
+      "num_frames": 13935
+    },
+    "robotgeneralist/pp__3__with_ft": {
+      "commit_hash": "68023819f92217d1239bf1dc1229b73a32580513",
+      "num_episodes": 99,
+      "num_frames": 38424
+    },
+    "robotgeneralist/pp__3__with_ft__rotations": {
+      "commit_hash": "062188051c5f78b03d51848dc0c03f93bc4a78c5",
+      "num_episodes": 49,
+      "num_frames": 7151
+    },
+    "robotgeneralist/pp__3__with_ft__adversarial_start": {
+      "commit_hash": "d0e27c579029e82f626c9ba8817ca43c610fa28e",
+      "num_episodes": 147,
+      "num_frames": 21313
+    },
+    "robotgeneralist/pp__4__with_ft": {
+      "commit_hash": "0d64bf5c9bac34d9cce59175a0becf822366bc53",
+      "num_episodes": 94,
+      "num_frames": 33046
+    },
+    "robotgeneralist/pp__5__with_ft__adversarial_target": {
+      "commit_hash": "8629d43a975c254367b84f491481fa8574416c6e",
+      "num_episodes": 48,
+      "num_frames": 10615
+    },
+    "robotgeneralist/pp__6__with_ft": {
+      "commit_hash": "e2a7c2ded8aed6663aa71fd9a1e6708d636071b1",
+      "num_episodes": 379,
+      "num_frames": 123611
+    },
+    "robotgeneralist/precise-packing-2025-10-22-part-1": {
+      "commit_hash": "defdc7ebf6871fc057ef30d21a2f71bfe3a40ce4",
+      "num_episodes": 98,
+      "num_frames": 41410
+    },
+    "robotgeneralist/precise-packing-2025-10-22-part-2": {
+      "commit_hash": "73551fdf12a436f983537d0454d9a6046a51c2a7",
+      "num_episodes": 94,
+      "num_frames": 37859
+    },
+    "robotgeneralist/precise-packing-2025-10-22-part-3": {
+      "commit_hash": "ae7abf3c61cb6d4f3096dc30b67731305c50d59c",
+      "num_episodes": 127,
+      "num_frames": 53341
+    },
+    "robotgeneralist/precise-packing-2025-10-22-part-4": {
+      "commit_hash": "5d68ce10ce1c5a492bfc760012aba4793a5d170b",
+      "num_episodes": 102,
+      "num_frames": 38263
+    }
+  }
+}
\ No newline at end of file
diff --git a/test_loss_components_1e-5/4000/params/_METADATA b/test_loss_components_1e-5/4000/params/_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..6c718061c06e49e04232f6012c4ddc165fd946d9
--- /dev/null
+++ b/test_loss_components_1e-5/4000/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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('params', 'PaliGemma', 'llm', 'final_norm_1', 'scale', '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": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_attention_norm_1', '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_1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('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}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_ffw_norm_1', '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_1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('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}}, "('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}}, "('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}}, "('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}}, "('params', 'action_time_mlp_in', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_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}}, "('params', 'action_time_mlp_in', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_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}}, "('params', 'action_time_mlp_out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_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}}, "('params', 'action_time_mlp_out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_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}}, "('params', 'state_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "state_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'state_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "state_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}}, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}
\ No newline at end of file
diff --git a/test_loss_components_1e-5/4000/params/_sharding b/test_loss_components_1e-5/4000/params/_sharding
new file mode 100644
index 0000000000000000000000000000000000000000..924608abdcb99abb3c864bb9d1bfca6eb60f780e
--- /dev/null
+++ b/test_loss_components_1e-5/4000/params/_sharding
@@ -0,0 +1 @@
+{"cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2Rlcl9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2Rlcl9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8wLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8wLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuZW1iZWRkaW5nLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuZW1iZWRkaW5nLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuaGVhZC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuaGVhZC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcucG9zX2VtYmVkZGluZy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZW1iZWRkZXIuaW5wdXRfZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZmluYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZmluYWxfbm9ybV8xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scC5nYXRpbmdfZWluc3VtLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scC5saW5lYXIudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scF8xLmdhdGluZ19laW5zdW0udmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scF8xLmxpbmVhci52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4uYXR0bl92ZWNfZWluc3VtLncudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4uYXR0bl92ZWNfZWluc3VtXzEudy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ua3ZfZWluc3VtLncudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ua3ZfZWluc3VtXzEudy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ucV9laW5zdW0udy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ucV9laW5zdW1fMS53LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9hdHRlbnRpb25fbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9hdHRlbnRpb25fbm9ybV8xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9mZndfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9mZndfbm9ybV8xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLmFjdGlvbl90aW1lX21scF9pbi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLmFjdGlvbl90aW1lX21scF9pbi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLmFjdGlvbl90aW1lX21scF9vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLmFjdGlvbl90aW1lX21scF9vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLmFjdGlvbl9pbl9wcm9qLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLmFjdGlvbl9pbl9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLmFjdGlvbl9vdXRfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLmFjdGlvbl9vdXRfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLnN0YXRlX3Byb2ouYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLnN0YXRlX3Byb2oua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}"}
\ No newline at end of file
diff --git a/test_loss_components_1e-5/4000/params/d/c08ee97d1b8c00ccc4a7086a9df7434d b/test_loss_components_1e-5/4000/params/d/c08ee97d1b8c00ccc4a7086a9df7434d
new file mode 100644
index 0000000000000000000000000000000000000000..b5d8327932449ba3625fa76fcc2d674bc0d723e2
Binary files /dev/null and b/test_loss_components_1e-5/4000/params/d/c08ee97d1b8c00ccc4a7086a9df7434d differ
diff --git a/test_loss_components_1e-5/4000/params/manifest.ocdbt b/test_loss_components_1e-5/4000/params/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..e1ffc5992d764ccfbbf70a5970d7c50f3cf5b046
Binary files /dev/null and b/test_loss_components_1e-5/4000/params/manifest.ocdbt differ
diff --git a/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/13d0383bcd36a0d9f61afb639ad380c9 b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/13d0383bcd36a0d9f61afb639ad380c9
new file mode 100644
index 0000000000000000000000000000000000000000..918db9c1c83ac25347ebf391b012a0ed2dccd5fb
--- /dev/null
+++ b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/13d0383bcd36a0d9f61afb639ad380c9
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0c1cd7c2392f5d5a0a1ffe5ab73f3ad1a4da2c19d627df5ad8c41c8e78cea097
+size 139976704
diff --git a/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/216e6d52245be1f029b9e88e58a7e551 b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/216e6d52245be1f029b9e88e58a7e551
new file mode 100644
index 0000000000000000000000000000000000000000..d82f4c2172b42fe8568b7d8681a7ec6be8af460e
--- /dev/null
+++ b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/216e6d52245be1f029b9e88e58a7e551
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3d15437b8285f09a39543bce279e330dc6d5660f0cf12a94ac0a14b25b53ab26
+size 70086656
diff --git a/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/23c9da5d1b786076230c741026d74240 b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/23c9da5d1b786076230c741026d74240
new file mode 100644
index 0000000000000000000000000000000000000000..96356ad1ce77d3a3b8891a7b9c5ce2e4cb5e46e8
--- /dev/null
+++ b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/23c9da5d1b786076230c741026d74240
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:88b0fb484d6868d40446eaeee65901eba53be51affff4c1e4813607c58520f18
+size 559738880
diff --git a/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/56dd749ddfd4e8994705b08a5a96e495 b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/56dd749ddfd4e8994705b08a5a96e495
new file mode 100644
index 0000000000000000000000000000000000000000..112aeafd464d01588a0e7b8d9149872f4a65fb20
--- /dev/null
+++ b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/56dd749ddfd4e8994705b08a5a96e495
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:da4f964ac25a1fd4fadfd02c742c72807120b88bce8b3f828119aafd5d33e037
+size 25841664
diff --git a/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/70410d314e2b7a69f49c646d900ca3c8 b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/70410d314e2b7a69f49c646d900ca3c8
new file mode 100644
index 0000000000000000000000000000000000000000..1ae3ad821c0c33b8949cd7f80f0025f20355adff
--- /dev/null
+++ b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/70410d314e2b7a69f49c646d900ca3c8
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:68cd1c909f5c1c403d2aea311a720df0ba4bc4b53191d1582c241aed4887fbdc
+size 280190976
diff --git a/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/86be6ff68f140be58a980cc11f42cf4e b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/86be6ff68f140be58a980cc11f42cf4e
new file mode 100644
index 0000000000000000000000000000000000000000..115fb74aa6d203554b9e6797b27c4afe15026396
--- /dev/null
+++ b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/86be6ff68f140be58a980cc11f42cf4e
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bdfe4ffd6f5bdecda926f8a23020b78c4c37c1b98589ef16b01ec0289bbabea4
+size 1957433344
diff --git a/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/8a992cd26ed86fb77ced871773dc4594 b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/8a992cd26ed86fb77ced871773dc4594
new file mode 100644
index 0000000000000000000000000000000000000000..bfaeb00b7636e91b78c38f402b903dcf05ffa121
--- /dev/null
+++ b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/8a992cd26ed86fb77ced871773dc4594
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:629c2996da87edf509d4f8d853095510a95de87e51b84d0f3f8c9ed3c60bc45c
+size 272850944
diff --git a/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/9cbc768e3b41e7497483908b73578383 b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/9cbc768e3b41e7497483908b73578383
new file mode 100644
index 0000000000000000000000000000000000000000..49f5be9fa67d4c249afd1a507dbbd07fac336dd6
--- /dev/null
+++ b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/9cbc768e3b41e7497483908b73578383
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:480a519d89d6de868198611990993b6e7122e86e629eb57651df62ef72ce6c21
+size 2240253952
diff --git a/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/b313f73ca4a074edf26764d407a78285 b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/b313f73ca4a074edf26764d407a78285
new file mode 100644
index 0000000000000000000000000000000000000000..47332359659d22f3a16811125e17bb519af15f54
--- /dev/null
+++ b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/b313f73ca4a074edf26764d407a78285
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e2ebf0cbc75bf6e23236c9d93ea17d3fee8cde1af6e09f04eae2dff7c3c81a3f
+size 2240212992
diff --git a/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/c3e96dcec38a9d1597affe33f8b76478 b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/c3e96dcec38a9d1597affe33f8b76478
new file mode 100644
index 0000000000000000000000000000000000000000..8882040625d3b4c4a33ad617ae61d2ee6b7dc18d
--- /dev/null
+++ b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/c3e96dcec38a9d1597affe33f8b76478
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:70cac31840b82ba01b86bcecbd49653a6aa4c1ba6f52acb21759bcfcd4bab362
+size 461739
diff --git a/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/ccaedfef9f74fc9a1d34cd70c5b165f1 b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/ccaedfef9f74fc9a1d34cd70c5b165f1
new file mode 100644
index 0000000000000000000000000000000000000000..16437047f5ee7c39fcbb5ece446aac5634c393d2
--- /dev/null
+++ b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/ccaedfef9f74fc9a1d34cd70c5b165f1
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e4f82cd2ef615c9042bf550c242608f97128d0ef27bfdcf91f9b771ddc2c65ce
+size 34975744
diff --git a/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/cccbe3fdf6d794f2bf59f994ed28b42b b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/cccbe3fdf6d794f2bf59f994ed28b42b
new file mode 100644
index 0000000000000000000000000000000000000000..ccd34cd2faa8fa38e5846333bb67a40ec53645d4
--- /dev/null
+++ b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/cccbe3fdf6d794f2bf59f994ed28b42b
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4ff510031ddd0f769c90752a30bec6270b092459556387e6e55ed08e79dc7b56
+size 1120116736
diff --git a/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/cebf6f114b344ef72f0e477120b90f81 b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/cebf6f114b344ef72f0e477120b90f81
new file mode 100644
index 0000000000000000000000000000000000000000..ffc57c5cefe3160b4f5ec2c8030ab168d9324180
Binary files /dev/null and b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/cebf6f114b344ef72f0e477120b90f81 differ
diff --git a/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/d1938925851d0d3ad3e22c1214906237 b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/d1938925851d0d3ad3e22c1214906237
new file mode 100644
index 0000000000000000000000000000000000000000..46cd46020c3fc37eed8126482d2295a61a10ce74
Binary files /dev/null and b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/d1938925851d0d3ad3e22c1214906237 differ
diff --git a/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/d1c9b12422327567c1d47d240228e907 b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/d1c9b12422327567c1d47d240228e907
new file mode 100644
index 0000000000000000000000000000000000000000..5df883618f7ec9c00aebdb37d02ddb89ba33663a
Binary files /dev/null and b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/d1c9b12422327567c1d47d240228e907 differ
diff --git a/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/d7ab5394c10ca3b27b5400a48b23e3d7 b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/d7ab5394c10ca3b27b5400a48b23e3d7
new file mode 100644
index 0000000000000000000000000000000000000000..5b929e736231420502e0b657e009c344dd9c1a17
Binary files /dev/null and b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/d7ab5394c10ca3b27b5400a48b23e3d7 differ
diff --git a/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/dcb557aed68e94285b47236d2e8c087a b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/dcb557aed68e94285b47236d2e8c087a
new file mode 100644
index 0000000000000000000000000000000000000000..da1220d23f215a600bf68f3b52ad14dcdc695c50
--- /dev/null
+++ b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/dcb557aed68e94285b47236d2e8c087a
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:924325695dec472a8d9b182fd098bc535fb1e8f0780abe9de9ec4a9e374d8a38
+size 497037312
diff --git a/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/f3a06a050c3562508e08ba9b4d78de80 b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/f3a06a050c3562508e08ba9b4d78de80
new file mode 100644
index 0000000000000000000000000000000000000000..fab408308225968c728b180da6bc527b0290242b
--- /dev/null
+++ b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/f3a06a050c3562508e08ba9b4d78de80
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:20183dfc195c5ba5c48a7b3fd07ef18e3162212c9dc296d11b4c8c3dd21a086c
+size 1120112640
diff --git a/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/f471d998aea2f39c6b4792cbe5a8b652 b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/f471d998aea2f39c6b4792cbe5a8b652
new file mode 100644
index 0000000000000000000000000000000000000000..7c197f8eb995816a122883c4ae961830221efd0b
--- /dev/null
+++ b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/f471d998aea2f39c6b4792cbe5a8b652
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fb8aea4f08c82eec473244dd86cbe0663b924615a79e43dc96f7877e05e7ca70
+size 559271936
diff --git a/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/fb624b4b1a19699a869f6eaba65560ac b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/fb624b4b1a19699a869f6eaba65560ac
new file mode 100644
index 0000000000000000000000000000000000000000..c06d0632618c3a9e6e237c0648ad7f137d078db9
Binary files /dev/null and b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/fb624b4b1a19699a869f6eaba65560ac differ
diff --git a/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/fc7b03c165579d42dbfb89800de58cc2 b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/fc7b03c165579d42dbfb89800de58cc2
new file mode 100644
index 0000000000000000000000000000000000000000..90e0f7149b6b7fef8520760f622725b4ef073cb6
--- /dev/null
+++ b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/fc7b03c165579d42dbfb89800de58cc2
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:161f5c3fb6864d7b137dec574e079e027825245890df867622f0eee1e7583a90
+size 398778368
diff --git a/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/fc8f8e8c0cd651f926f070ccb0f29d32 b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/fc8f8e8c0cd651f926f070ccb0f29d32
new file mode 100644
index 0000000000000000000000000000000000000000..0cf18fe71caca67ee682f3421371f2e48230d30c
--- /dev/null
+++ b/test_loss_components_1e-5/4000/params/ocdbt.process_0/d/fc8f8e8c0cd651f926f070ccb0f29d32
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8b3adea212de634ec16b38a34d34a0fb1abd746ed490d477c87d20e463275ec2
+size 496754688
diff --git a/test_loss_components_1e-5/4000/params/ocdbt.process_0/manifest.ocdbt b/test_loss_components_1e-5/4000/params/ocdbt.process_0/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..a1162a7309d9387e5e7718b53bc7e118574be29e
Binary files /dev/null and b/test_loss_components_1e-5/4000/params/ocdbt.process_0/manifest.ocdbt differ
diff --git a/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/_CHECKPOINT_METADATA b/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/_CHECKPOINT_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..298d018944596e2850e83eb45cb472037e1537de
--- /dev/null
+++ b/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/_CHECKPOINT_METADATA
@@ -0,0 +1 @@
+{"item_handlers": null, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1765321055306498361, "commit_timestamp_nsecs": null, "custom": {}}
\ No newline at end of file
diff --git a/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/assets.orbax-checkpoint-tmp-17/ur10e/norm_stats.json b/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/assets.orbax-checkpoint-tmp-17/ur10e/norm_stats.json
new file mode 100644
index 0000000000000000000000000000000000000000..85a8eaab4b85a50ec7b3633ecb653ab8e0366022
--- /dev/null
+++ b/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/assets.orbax-checkpoint-tmp-17/ur10e/norm_stats.json
@@ -0,0 +1,491 @@
+{
+  "norm_stats": {
+    "state": {
+      "mean": [
+        0.00019428042287472636,
+        0.00020032725296914577,
+        8.52593166200677e-6,
+        0.00019105659157503396,
+        0.0001996985956793651,
+        0.0000162034084496554,
+        0.00019352274830453098,
+        0.00020078709349036217,
+        0.000029199256459833123,
+        0.00019142567180097103,
+        0.00020083319395780563,
+        0.00003592482607928105,
+        0.00005053714266978204,
+        -0.00003347316669533029,
+        -0.00020081466936971992,
+        0.00018982442270498723,
+        0.00020038991351611912,
+        0.000046112203563097864,
+        0.00004985106716048904,
+        -0.00003275063500041142,
+        -0.0002042762062046677,
+        0.5986685752868652,
+        0.5981546640396118,
+        0.05039028078317642,
+        0.0004953373572789133,
+        0.10711561888456345,
+        -0.0005407895660027862,
+        -0.006484619807451963,
+        0.003617143025621772,
+        0.5981546640396118,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.0023281769827008247,
+        0.0034390955697745085,
+        0.002846071030944586,
+        0.0023422744125127792,
+        0.003452499397099018,
+        0.0028555369935929775,
+        0.002345019020140171,
+        0.0034702608827501535,
+        0.0028640974778681993,
+        0.0023441859520971775,
+        0.0034640957601368427,
+        0.0028661703690886497,
+        0.0012911600060760975,
+        0.0020082774572074413,
+        0.01001499593257904,
+        0.0023556631058454514,
+        0.0034643616527318954,
+        0.002885110443457961,
+        0.0012921782908961177,
+        0.0020203019957989454,
+        0.010073363780975342,
+        0.4901678264141083,
+        0.4902710020542145,
+        0.11618480086326599,
+        0.1681893765926361,
+        0.07788223028182983,
+        0.04501405730843544,
+        0.0661778524518013,
+        0.5958529710769653,
+        0.4902710020542145,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        -0.0070394672028720375,
+        -0.011599844336509704,
+        -0.0054659975245594985,
+        -0.007079721701145172,
+        -0.011622210478782653,
+        -0.005539767185784876,
+        -0.007091582050547003,
+        -0.011584406599402429,
+        -0.0054779203861951824,
+        -0.0070481276120990515,
+        -0.011584554602950811,
+        -0.005540505648031831,
+        -0.0038997788466513163,
+        -0.007352302264794708,
+        -0.034272025059163574,
+        -0.007114377491176128,
+        -0.011688268885016442,
+        -0.005550548636354506,
+        -0.003901443669199943,
+        -0.007246576420217752,
+        -0.034621239319443706,
+        0.0,
+        0.0,
+        -0.19741441082954406,
+        -0.3133266278743744,
+        0.00863822932690382,
+        -0.11799316820502281,
+        -0.1673899977207184,
+        -1.4435412939071655,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.007438025734573604,
+        0.011892002701759337,
+        0.009434655957296489,
+        0.007451299774646758,
+        0.011834505152702332,
+        0.009366847675666212,
+        0.007520044301077725,
+        0.01185322367250919,
+        0.009355697490274907,
+        0.007454063414782285,
+        0.011923038901388648,
+        0.009385530477575958,
+        0.004192294634878635,
+        0.005776847817935048,
+        0.033915841154754156,
+        0.007474094156175851,
+        0.011830379698425528,
+        0.009328839504159987,
+        0.004122480013966562,
+        0.005790175495669243,
+        0.03415216095745563,
+        0.9998,
+        0.9998,
+        0.29779768657684336,
+        0.406406463944912,
+        0.3194635294646025,
+        0.1383301957249642,
+        0.1970333372354507,
+        1.6112462955951692,
+        0.9998,
+        0.0,
+        0.0
+      ]
+    },
+    "actions": {
+      "mean": [
+        0.00005176717240829021,
+        0.00012534917914308608,
+        0.0002480015973560512,
+        0.00005367679477785714,
+        -3.479610484191653e-7,
+        -0.0002569151984062046,
+        0.5995351076126099,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.002336727222427726,
+        0.0034517920576035976,
+        0.003022634657099843,
+        0.0012890173820778728,
+        0.0020518647506833076,
+        0.010340424254536629,
+        0.489992618560791,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        -0.006998317793011664,
+        -0.011659257435053586,
+        -0.0055110766027122735,
+        -0.003918213281407953,
+        -0.007409645091742276,
+        -0.03437663650512694,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.007169186890125277,
+        0.0115618407972157,
+        0.009297586680948735,
+        0.004103694232180714,
+        0.005846588513627651,
+        0.034828110963106174,
+        0.9998,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    },
+    "previous_actions": {
+      "mean": [
+        0.00005176717240829021,
+        0.00012534917914308608,
+        0.0002480015973560512,
+        0.00005367679477785714,
+        -3.479610484191653e-7,
+        -0.0002569151984062046,
+        0.5995351076126099,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.002336727222427726,
+        0.0034517920576035976,
+        0.003022634657099843,
+        0.0012890173820778728,
+        0.0020518647506833076,
+        0.010340424254536629,
+        0.489992618560791,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        -0.006998317793011664,
+        -0.011659257435053586,
+        -0.0055110766027122735,
+        -0.003918213281407953,
+        -0.007409645091742276,
+        -0.03437663650512694,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.007169186890125277,
+        0.0115618407972157,
+        0.009297586680948735,
+        0.004103694232180714,
+        0.005846588513627651,
+        0.034828110963106174,
+        0.9998,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    }
+  },
+  "sample_ratio": 0.05,
+  "dataset_infos": {
+    "robotgeneralist/pp__1__train": {
+      "commit_hash": "b8158e3f7eb1eb4fcad944d3dfbcdba2796a214b",
+      "num_episodes": 166,
+      "num_frames": 69615
+    },
+    "robotgeneralist/pp__1__adversarial": {
+      "commit_hash": "a8b6024ff9f5261e87d553f35f07f16487adabc1",
+      "num_episodes": 64,
+      "num_frames": 17906
+    },
+    "robotgeneralist/pp__2": {
+      "commit_hash": "f244ff639c1a837fdec7990c601bf5ed4fbb56a9",
+      "num_episodes": 92,
+      "num_frames": 37458
+    },
+    "robotgeneralist/pp__2__rotations": {
+      "commit_hash": "9b8e02337fce6bccda7adcc4aae1edcbe5b3de41",
+      "num_episodes": 102,
+      "num_frames": 13935
+    },
+    "robotgeneralist/pp__3__with_ft": {
+      "commit_hash": "68023819f92217d1239bf1dc1229b73a32580513",
+      "num_episodes": 99,
+      "num_frames": 38424
+    },
+    "robotgeneralist/pp__3__with_ft__rotations": {
+      "commit_hash": "062188051c5f78b03d51848dc0c03f93bc4a78c5",
+      "num_episodes": 49,
+      "num_frames": 7151
+    },
+    "robotgeneralist/pp__3__with_ft__adversarial_start": {
+      "commit_hash": "d0e27c579029e82f626c9ba8817ca43c610fa28e",
+      "num_episodes": 147,
+      "num_frames": 21313
+    },
+    "robotgeneralist/pp__4__with_ft": {
+      "commit_hash": "0d64bf5c9bac34d9cce59175a0becf822366bc53",
+      "num_episodes": 94,
+      "num_frames": 33046
+    },
+    "robotgeneralist/pp__5__with_ft__adversarial_target": {
+      "commit_hash": "8629d43a975c254367b84f491481fa8574416c6e",
+      "num_episodes": 48,
+      "num_frames": 10615
+    },
+    "robotgeneralist/pp__6__with_ft": {
+      "commit_hash": "e2a7c2ded8aed6663aa71fd9a1e6708d636071b1",
+      "num_episodes": 379,
+      "num_frames": 123611
+    },
+    "robotgeneralist/precise-packing-2025-10-22-part-1": {
+      "commit_hash": "defdc7ebf6871fc057ef30d21a2f71bfe3a40ce4",
+      "num_episodes": 98,
+      "num_frames": 41410
+    },
+    "robotgeneralist/precise-packing-2025-10-22-part-2": {
+      "commit_hash": "73551fdf12a436f983537d0454d9a6046a51c2a7",
+      "num_episodes": 94,
+      "num_frames": 37859
+    },
+    "robotgeneralist/precise-packing-2025-10-22-part-3": {
+      "commit_hash": "ae7abf3c61cb6d4f3096dc30b67731305c50d59c",
+      "num_episodes": 127,
+      "num_frames": 53341
+    },
+    "robotgeneralist/precise-packing-2025-10-22-part-4": {
+      "commit_hash": "5d68ce10ce1c5a492bfc760012aba4793a5d170b",
+      "num_episodes": 102,
+      "num_frames": 38263
+    }
+  }
+}
\ No newline at end of file
diff --git a/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/_METADATA b/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..6c718061c06e49e04232f6012c4ddc165fd946d9
--- /dev/null
+++ b/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/_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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('params', 'PaliGemma', 'llm', 'final_norm_1', 'scale', '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": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('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}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_attention_norm_1', '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_1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('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}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_ffw_norm_1', '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_1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('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}}, "('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}}, "('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}}, "('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}}, "('params', 'action_time_mlp_in', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_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}}, "('params', 'action_time_mlp_in', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_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}}, "('params', 'action_time_mlp_out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_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}}, "('params', 'action_time_mlp_out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_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}}, "('params', 'state_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "state_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'state_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "state_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}}, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}
\ No newline at end of file
diff --git a/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/_sharding b/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/_sharding
new file mode 100644
index 0000000000000000000000000000000000000000..924608abdcb99abb3c864bb9d1bfca6eb60f780e
--- /dev/null
+++ b/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/_sharding
@@ -0,0 +1 @@
+{"cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2Rlcl9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2Rlcl9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8wLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8wLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuZW1iZWRkaW5nLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuZW1iZWRkaW5nLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuaGVhZC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuaGVhZC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcucG9zX2VtYmVkZGluZy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZW1iZWRkZXIuaW5wdXRfZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZmluYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZmluYWxfbm9ybV8xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scC5nYXRpbmdfZWluc3VtLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scC5saW5lYXIudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scF8xLmdhdGluZ19laW5zdW0udmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scF8xLmxpbmVhci52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4uYXR0bl92ZWNfZWluc3VtLncudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4uYXR0bl92ZWNfZWluc3VtXzEudy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ua3ZfZWluc3VtLncudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ua3ZfZWluc3VtXzEudy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ucV9laW5zdW0udy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ucV9laW5zdW1fMS53LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9hdHRlbnRpb25fbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9hdHRlbnRpb25fbm9ybV8xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9mZndfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9mZndfbm9ybV8xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLmFjdGlvbl90aW1lX21scF9pbi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLmFjdGlvbl90aW1lX21scF9pbi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLmFjdGlvbl90aW1lX21scF9vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLmFjdGlvbl90aW1lX21scF9vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLmFjdGlvbl9pbl9wcm9qLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLmFjdGlvbl9pbl9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLmFjdGlvbl9vdXRfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLmFjdGlvbl9vdXRfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLnN0YXRlX3Byb2ouYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","cGFyYW1zLnN0YXRlX3Byb2oua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}"}
\ No newline at end of file
diff --git a/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/ocdbt.process_0/d/006c56c91903b261c87ded287f73ee38 b/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/ocdbt.process_0/d/006c56c91903b261c87ded287f73ee38
new file mode 100644
index 0000000000000000000000000000000000000000..4b3d2dc1dcdfb1b7fcf4331de45ee379bf28a64e
Binary files /dev/null and b/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/ocdbt.process_0/d/006c56c91903b261c87ded287f73ee38 differ
diff --git a/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/ocdbt.process_0/d/271e4e51c9632bf778eaa48232cd5ac2 b/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/ocdbt.process_0/d/271e4e51c9632bf778eaa48232cd5ac2
new file mode 100644
index 0000000000000000000000000000000000000000..242bddc226cd21d54dc4889119aff4d6bee7294a
Binary files /dev/null and b/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/ocdbt.process_0/d/271e4e51c9632bf778eaa48232cd5ac2 differ
diff --git a/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/ocdbt.process_0/d/2f93887a4bbd190673a669b2d374e96c b/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/ocdbt.process_0/d/2f93887a4bbd190673a669b2d374e96c
new file mode 100644
index 0000000000000000000000000000000000000000..02fc1132767646146734b8812945355f759cd535
--- /dev/null
+++ b/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/ocdbt.process_0/d/2f93887a4bbd190673a669b2d374e96c
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0cddd6c5c2ca53613676a4df18488f6554c00423393047e411f30b977d00c52d
+size 398753792
diff --git a/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/ocdbt.process_0/d/5f7448af92b6085f7cdc3a64d36c0aa5 b/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/ocdbt.process_0/d/5f7448af92b6085f7cdc3a64d36c0aa5
new file mode 100644
index 0000000000000000000000000000000000000000..133aba9a63871f42ffc5b3d7cce5f0bf210874f3
--- /dev/null
+++ b/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/ocdbt.process_0/d/5f7448af92b6085f7cdc3a64d36c0aa5
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1795ad7a9a8631a43a975845f4b64bcac95ceeb86fce3b780241a627737eaa45
+size 190993
diff --git a/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/ocdbt.process_0/d/8d80895564ad25dbd4aa5410b7c7d4b0 b/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/ocdbt.process_0/d/8d80895564ad25dbd4aa5410b7c7d4b0
new file mode 100644
index 0000000000000000000000000000000000000000..5bc8ff6bbdd2754e0b785bd55ab657ecaea44213
--- /dev/null
+++ b/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/ocdbt.process_0/d/8d80895564ad25dbd4aa5410b7c7d4b0
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:86ecc23750e6fd1fd6977424b09bf69eb80532972e7da0c4d0a13ef922f7919e
+size 34979840
diff --git a/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/ocdbt.process_0/d/93444dbe916f143cd1bf761636a8f948 b/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/ocdbt.process_0/d/93444dbe916f143cd1bf761636a8f948
new file mode 100644
index 0000000000000000000000000000000000000000..d302df165ca3141153d747457a971abddaea0bbc
--- /dev/null
+++ b/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/ocdbt.process_0/d/93444dbe916f143cd1bf761636a8f948
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0a34d08e6beb72e700c8f1e65be2c3999dce37ec147a5f8b5ca9f2ef6bbd6aee
+size 26107904
diff --git a/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/ocdbt.process_0/d/a2772d86118d3e1be8e60be7400631db b/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/ocdbt.process_0/d/a2772d86118d3e1be8e60be7400631db
new file mode 100644
index 0000000000000000000000000000000000000000..efab85904885f8bbb3515f34d33c53749673ad64
Binary files /dev/null and b/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/ocdbt.process_0/d/a2772d86118d3e1be8e60be7400631db differ
diff --git a/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/ocdbt.process_0/d/a2c409f30fa48415606b71cd736a7b93 b/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/ocdbt.process_0/d/a2c409f30fa48415606b71cd736a7b93
new file mode 100644
index 0000000000000000000000000000000000000000..a4f2419a1ee3ed1a80330546e7dabbbf2765999d
--- /dev/null
+++ b/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/ocdbt.process_0/d/a2c409f30fa48415606b71cd736a7b93
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6aa3cb470271052633c2cf5616d5e53311aee4b24b58fdbda03acb8d9907d33c
+size 70086656
diff --git a/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/ocdbt.process_0/d/e60f152df4300174f85844aa356dbb1d b/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/ocdbt.process_0/d/e60f152df4300174f85844aa356dbb1d
new file mode 100644
index 0000000000000000000000000000000000000000..46cd46020c3fc37eed8126482d2295a61a10ce74
Binary files /dev/null and b/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/ocdbt.process_0/d/e60f152df4300174f85844aa356dbb1d differ
diff --git a/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/ocdbt.process_0/manifest.ocdbt b/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/ocdbt.process_0/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..d2143911b55450fc559dd8b890553150621d46ef
Binary files /dev/null and b/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/params.orbax-checkpoint-tmp-18/ocdbt.process_0/manifest.ocdbt differ
diff --git a/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/train_state.orbax-checkpoint-tmp-19/ocdbt.process_0/d/841b492f005833eadd1596ede344ccb9 b/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/train_state.orbax-checkpoint-tmp-19/ocdbt.process_0/d/841b492f005833eadd1596ede344ccb9
new file mode 100644
index 0000000000000000000000000000000000000000..7df74305ce9e5b113b87269292008089ce84c7b1
Binary files /dev/null and b/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/train_state.orbax-checkpoint-tmp-19/ocdbt.process_0/d/841b492f005833eadd1596ede344ccb9 differ
diff --git a/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/train_state.orbax-checkpoint-tmp-19/ocdbt.process_0/d/b1edd09990c1842ebb62fb106f6c72a5 b/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/train_state.orbax-checkpoint-tmp-19/ocdbt.process_0/d/b1edd09990c1842ebb62fb106f6c72a5
new file mode 100644
index 0000000000000000000000000000000000000000..deb3a04bf7461715185e2e464e4b713a3cedb7f5
Binary files /dev/null and b/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/train_state.orbax-checkpoint-tmp-19/ocdbt.process_0/d/b1edd09990c1842ebb62fb106f6c72a5 differ
diff --git a/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/train_state.orbax-checkpoint-tmp-19/ocdbt.process_0/d/c50a43de5892f8dd70e4259b4645d6a3 b/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/train_state.orbax-checkpoint-tmp-19/ocdbt.process_0/d/c50a43de5892f8dd70e4259b4645d6a3
new file mode 100644
index 0000000000000000000000000000000000000000..32aeb3b90d60b0a7055cd0d594f6225851c8357d
Binary files /dev/null and b/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/train_state.orbax-checkpoint-tmp-19/ocdbt.process_0/d/c50a43de5892f8dd70e4259b4645d6a3 differ
diff --git a/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/train_state.orbax-checkpoint-tmp-19/ocdbt.process_0/d/cd549b6bc7917200a54ca348253d36a8 b/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/train_state.orbax-checkpoint-tmp-19/ocdbt.process_0/d/cd549b6bc7917200a54ca348253d36a8
new file mode 100644
index 0000000000000000000000000000000000000000..9518787cb2008e699c2cdd41b30d8976c7ea5bbd
Binary files /dev/null and b/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/train_state.orbax-checkpoint-tmp-19/ocdbt.process_0/d/cd549b6bc7917200a54ca348253d36a8 differ
diff --git a/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/train_state.orbax-checkpoint-tmp-19/ocdbt.process_0/manifest.ocdbt b/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/train_state.orbax-checkpoint-tmp-19/ocdbt.process_0/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..70715ef2779e88c514503cbb1edf2219ee71d69d
Binary files /dev/null and b/test_loss_components_1e-5/5000.orbax-checkpoint-tmp-16/train_state.orbax-checkpoint-tmp-19/ocdbt.process_0/manifest.ocdbt differ
diff --git a/test_loss_components_1e-5/config.yaml b/test_loss_components_1e-5/config.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..0f6cd24ee379c47d9a91346bf2360e013d78803b
--- /dev/null
+++ b/test_loss_components_1e-5/config.yaml
@@ -0,0 +1,100 @@
+!!python/object:openpi.training.config.TrainConfig
+name: test_loss_components_1e-5
+project_name: openpi
+exp_name: test_loss_components_1e-5
+model: !!python/object:openpi.models.pi0.Pi0Config
+    action_dim: 32
+    action_horizon: 50
+    max_token_len: 48
+    dtype: bfloat16
+    paligemma_variant: gemma_2b
+    action_expert_variant: gemma_300m
+    img_variant: So400m/14
+    tv_loss_weight: 30.0
+    tv_loss_type: jerk-limited
+    tv_loss_include_prev_actions: true
+    dynamic_state_dropout_rate: 0.7
+    jerk_max: 10
+    jerk_limit_loss_scaling: 1.0
+    jerk_tv_loss_scaling: 1.0e-05
+    num_previous_actions: 5
+weight_loader: !!python/object:openpi.training.weight_loaders.CheckpointWeightLoader
+    local_dir: null
+    repo_id: robotgeneralist/openpi_checkpoint_mirrors
+    revision: main
+    checkpoint_path: pi0_base
+lr_schedule: !!python/object:openpi.training.optimizer.CosineDecaySchedule
+    warmup_steps: 1000
+    peak_lr: 2.5e-05
+    decay_steps: 30000
+    decay_lr: 2.5e-06
+optimizer: !!python/object:openpi.training.optimizer.AdamW
+    b1: 0.9
+    b2: 0.95
+    eps: 1.0e-08
+    weight_decay: 1.0e-10
+    clip_gradient_norm: 1.0
+ema_decay: 0.99
+freeze_filter: !!python/name:flax.nnx.filterlib.Nothing ''
+data: !!python/object:openpi.training.config.LeRobotNomagicURXDataConfig
+    repo_ids: !!python/tuple
+    - robotgeneralist/pp__1__train
+    - robotgeneralist/pp__1__adversarial
+    - robotgeneralist/pp__2
+    - robotgeneralist/pp__2__rotations
+    - robotgeneralist/pp__3__with_ft
+    - robotgeneralist/pp__3__with_ft__rotations
+    - robotgeneralist/pp__3__with_ft__adversarial_start
+    - robotgeneralist/pp__4__with_ft
+    - robotgeneralist/pp__5__with_ft__adversarial_target
+    - robotgeneralist/pp__6__with_ft
+    - robotgeneralist/precise-packing-2025-10-22-part-1
+    - robotgeneralist/precise-packing-2025-10-22-part-2
+    - robotgeneralist/precise-packing-2025-10-22-part-3
+    - robotgeneralist/precise-packing-2025-10-22-part-4
+    assets: !!python/object:openpi.training.config.AssetsConfig
+        assets_dir: null
+        asset_id: ur10e
+    norm_stats_sample_ratio: 0.05
+    base_config: !!python/object:openpi.training.config.DataConfig
+        repo_ids: null
+        asset_id: null
+        norm_stats_data: null
+        norm_stats_sample_ratio: 1.0
+        norm_stats_dir: null
+        repack_transforms: !!python/object:openpi.transforms.Group
+            inputs: !!python/tuple []
+            outputs: !!python/tuple []
+        data_transforms: !!python/object:openpi.transforms.Group
+            inputs: !!python/tuple []
+            outputs: !!python/tuple []
+        model_transforms: !!python/object:openpi.transforms.Group
+            inputs: !!python/tuple []
+            outputs: !!python/tuple []
+        use_quantile_norm: false
+        action_sequence_keys: !!python/tuple
+        - actions
+        fourier_config: null
+        prompt_from_task: true
+    default_prompt: null
+    use_poses: true
+    end_effector_frame: false
+    use_dynamic_state: true
+    fourier_config: null
+assets_base_dir: ./assets
+checkpoint_base_dir: /storage_nvme_1/pzal/checkpoints
+seed: 42
+batch_size: 32
+num_workers: 2
+num_train_steps: 5001
+log_interval: 100
+save_interval: 1000
+keep_period: 5000
+overwrite: false
+resume: true
+wandb_enabled: true
+policy_metadata: null
+fsdp_devices: 1
+test_split_ratio: 0.05
+test_eval_interval: 2000
+test_eval_use_train_mode: true
diff --git a/test_loss_components_1e-5/wandb_id.txt b/test_loss_components_1e-5/wandb_id.txt
new file mode 100644
index 0000000000000000000000000000000000000000..02bf85500f5457f93bf8f34c5d7576020b4e95c1
--- /dev/null
+++ b/test_loss_components_1e-5/wandb_id.txt
@@ -0,0 +1 @@
+55t40oyf
\ No newline at end of file