upload

.gitattributes

@@ -33,3 +33,14 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+params/ocdbt.process_0/d/2dff38221e8c32b63bac7c6b78de89c0 filter=lfs diff=lfs merge=lfs -text
+params/ocdbt.process_0/d/30292007c81b9a21fc6a64fc301f1d76 filter=lfs diff=lfs merge=lfs -text
+params/ocdbt.process_0/d/35b7575f8dc53e262ec30dc357ff4f57 filter=lfs diff=lfs merge=lfs -text
+params/ocdbt.process_0/d/50a243a0658017486f2a0d0c3b7bc491 filter=lfs diff=lfs merge=lfs -text
+params/ocdbt.process_0/d/64ffd18294882e21e199f6ead2b88585 filter=lfs diff=lfs merge=lfs -text
+params/ocdbt.process_0/d/6bc9f00c1744d207f51f42a60906e9d7 filter=lfs diff=lfs merge=lfs -text
+params/ocdbt.process_0/d/97acb74367737f2c8a7cd00ed0a422b6 filter=lfs diff=lfs merge=lfs -text
+params/ocdbt.process_0/d/dcf75d91167f68b09b271a8f578b0612 filter=lfs diff=lfs merge=lfs -text
+params/ocdbt.process_0/d/ef4ff8b632eadc0c5f9e448a310890ad filter=lfs diff=lfs merge=lfs -text
+params/ocdbt.process_0/d/f1b50864ed05d5b25e7f51fbfc032752 filter=lfs diff=lfs merge=lfs -text
+params/ocdbt.process_0/d/fd8f7f7297b5275444404daf322f4d41 filter=lfs diff=lfs merge=lfs -text

assets/n/norm_stats.json

@@ -0,0 +1,280 @@
+{
+  "norm_stats": {
+    "state": {
+      "mean": [
+        0.20448490977287292,
+        -0.22259534895420074,
+        0.050946783274412155,
+        -0.2463078498840332,
+        0.1966649442911148,
+        0.49962711334228516,
+        0.6713993549346924,
+        0.7001646161079407,
+        0.4397074580192566,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        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.12407848238945007,
+        0.15208250284194946,
+        0.11137238144874573,
+        0.2074296623468399,
+        0.26403746008872986,
+        0.2370542287826538,
+        0.17695817351341248,
+        0.1682043820619583,
+        0.2789625823497772,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        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.13162163542509078,
+        -0.5480457600533962,
+        -0.20297523586750033,
+        -0.7012269123077393,
+        -0.5090069903612137,
+        -0.0027919290542602893,
+        0.18195701952173476,
+        0.003963552530668676,
+        -0.0075545475006103535,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        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.47362671941518786,
+        0.10633832247853281,
+        0.30861101117134093,
+        0.34549037899971013,
+        0.6906859913587571,
+        0.9064911778688429,
+        0.9734953178712301,
+        0.8142588168615474,
+        0.7797586218357087,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    },
+    "actions": {
+      "mean": [
+        -0.0035658704582601786,
+        -0.002766379853710532,
+        -0.0021811630576848984,
+        0.9759819507598877,
+        -0.00013730456703342497,
+        0.005280264653265476,
+        -0.0031429363880306482,
+        0.9801815748214722,
+        0.002495812252163887,
+        -0.0065133036114275455,
+        0.00039524523890577257,
+        -0.012719724327325821,
+        0.92564857006073,
+        0.004317340441048145,
+        0.0268560778349638,
+        0.013716470450162888,
+        0.952211320400238,
+        -0.004180672578513622,
+        0.6964102983474731,
+        0.43971318006515503,
+        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.05378447100520134,
+        0.0279812291264534,
+        0.022904623299837112,
+        0.08772635459899902,
+        0.14355432987213135,
+        0.138325035572052,
+        0.14809203147888184,
+        0.07582797110080719,
+        0.10747046023607254,
+        0.09164366871118546,
+        0.052467335015535355,
+        0.05613061413168907,
+        0.14647746086120605,
+        0.19926223158836365,
+        0.28508394956588745,
+        0.20775704085826874,
+        0.10053347796201706,
+        0.19954316318035126,
+        0.17601759731769562,
+        0.286819189786911,
+        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.2522053847193718,
+        -0.1278507245033979,
+        -0.10228770363330841,
+        0.5193793908834456,
+        -0.5599883568286896,
+        -0.5199895858764648,
+        -0.5706864560604095,
+        0.6123258974909782,
+        -0.36467745590209966,
+        -0.27141693069934847,
+        -0.15879799427986144,
+        -0.20579802601337432,
+        0.290797308743,
+        -0.6105792880058288,
+        -0.7903936503887177,
+        -0.5980796887278557,
+        0.49188290057182305,
+        -0.6635887041807175,
+        -0.0005886680554598577,
+        -0.0075545475006103535,
+        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.16641598236560817,
+        0.06759651153981688,
+        0.06508678269386292,
+        0.9996499485731125,
+        0.5300034245848655,
+        0.5143812204599381,
+        0.5680311403274536,
+        0.9996247104525566,
+        0.3931832425594328,
+        0.2680388767004013,
+        0.17481179325580598,
+        0.15206679404973988,
+        0.9996277151226998,
+        0.5864342331886292,
+        0.8592004993438722,
+        0.6620590616703033,
+        0.9996115312695502,
+        0.562399676847458,
+        0.8142588168615474,
+        0.7797586218357087,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    }
+  }
+}

params/_CHECKPOINT_METADATA

@@ -0,0 +1 @@
+{"item_handlers": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler", "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1790730071223232163, "commit_timestamp_nsecs": 1790730075646018442, "custom_metadata": {}}

params/_METADATA

@@ -0,0 +1 @@
+{"tree_metadata": {"('params', 'PaliGemma', 'img', 'Transformer', 'encoder_norm', 'bias')": {"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}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoder_norm', 'scale')": {"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}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_0', 'bias')": {"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}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_0', 'scale')": {"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}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_1', 'bias')": {"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}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_1', 'scale')": {"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}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_0', 'bias')": {"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}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_0', 'kernel')": {"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}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_1', 'bias')": {"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}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_1', 'kernel')": {"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}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"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}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"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}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"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}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"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}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"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}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"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}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"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}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"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}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'img', 'embedding', 'bias')": {"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}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'img', 'embedding', 'kernel')": {"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}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'img', 'head', 'bias')": {"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}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'img', 'head', 'kernel')": {"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}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'img', 'pos_embedding')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "pos_embedding", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'llm', 'embedder', 'input_embedding')": {"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}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'llm', 'final_norm', 'scale')": {"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}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'llm', 'final_norm_1', 'Dense_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "final_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'llm', 'final_norm_1', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "final_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'attn_vec_einsum', 'w')": {"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}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'attn_vec_einsum_1', 'w')": {"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}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'kv_einsum', 'w')": {"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}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'kv_einsum_1', 'w')": {"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}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'q_einsum', 'w')": {"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}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'q_einsum_1', 'w')": {"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}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp', 'gating_einsum')": {"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}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp', 'linear')": {"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}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp_1', 'gating_einsum')": {"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}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp_1', 'linear')": {"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}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_attention_norm', 'scale')": {"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}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_attention_norm_1', 'Dense_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_attention_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_attention_norm_1', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_attention_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_ffw_norm', 'scale')": {"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}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_ffw_norm_1', 'Dense_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_ffw_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_ffw_norm_1', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_ffw_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'action_in_proj', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_in_proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'action_in_proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_in_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'action_out_proj', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_out_proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'action_out_proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'time_mlp_in', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "time_mlp_in", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'time_mlp_in', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "time_mlp_in", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'time_mlp_out', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "time_mlp_out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('params', 'time_mlp_out', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "time_mlp_out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}}, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}

params/ocdbt.process_0/d/2dff38221e8c32b63bac7c6b78de89c0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b47ad6250ab52500da564b08338a7bb56319878c4c0e11d3139d9633cec29dfa
+size 151439

params/ocdbt.process_0/d/30292007c81b9a21fc6a64fc301f1d76

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e0919d2f40d3c6d9f9c9137641740d18a144fdab1075d26d1c6ac03a86578248
+size 112466259

params/ocdbt.process_0/d/35b7575f8dc53e262ec30dc357ff4f57

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:60b99214c05499bc86e7e5585c8d26d1fac268ab206ab21f73dc82a74717f5e8
+size 656029289

params/ocdbt.process_0/d/50a243a0658017486f2a0d0c3b7bc491

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b9f79eb1029eaea4d1659ee7e21fa1de847f3cd302a48234638ca59da0cb013c
+size 945806822

params/ocdbt.process_0/d/64ffd18294882e21e199f6ead2b88585

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3ef7b701864b1393774fcfc6fff361a82ec668263b6c8238a1cb55d060d0c0be
+size 560210709

params/ocdbt.process_0/d/6bc9f00c1744d207f51f42a60906e9d7

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3fd36a5b95adc75e002137fb59055d5b9aea0f1fb9b18c85cb14800138186002
+size 828934596

params/ocdbt.process_0/d/97acb74367737f2c8a7cd00ed0a422b6

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6048a2839cf60c90b5efbb77cb50a683370cb56f30533b3b55cc814ec0e4fd77
+size 234508559

params/ocdbt.process_0/d/dcf75d91167f68b09b271a8f578b0612

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5442992a454dbfdfdd7caf95c6afb09e32f54252aeb9d91768c73b2b1c634973
+size 14717752

params/ocdbt.process_0/d/ef4ff8b632eadc0c5f9e448a310890ad

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:48d9b1582c7e5c71d530b676366a71d6e3336e1952bca92572c335c020f860b8
+size 1891098940

params/ocdbt.process_0/d/f1b50864ed05d5b25e7f51fbfc032752

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:976187de04b65394955f1442ab0364514d422d780e633cfec232c5161231b5b4
+size 7126932

params/ocdbt.process_0/d/fd8f7f7297b5275444404daf322f4d41

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e17a463e1b1d37270490b44e476670508cdbbbf799a07f2826cde2e51778f7af
+size 5700526

policy.py

@@ -0,0 +1,97 @@
+import dataclasses
+import pathlib
+import sys
+import time
+
+import numpy as np
+from scipy.spatial.transform import Rotation
+
+_K = ("observation.image.left", "observation.image.right",
+      "observation.pose.left_hand_root_to_right_hand_root.absolute", "observation.joint_states", "prompt")
+
+
+def _m(r):
+    a, b = r[..., :3], r[..., 3:6]
+    c0 = a / np.linalg.norm(a, axis=-1, keepdims=True)
+    b = b - np.sum(c0 * b, axis=-1, keepdims=True) * c0
+    c1 = b / np.linalg.norm(b, axis=-1, keepdims=True)
+    return np.stack([c0, c1, np.cross(c0, c1)], axis=-1)
+
+
+def _d(x):
+    rot = Rotation.from_matrix(_m(x[:, 3:9]))
+    pr = Rotation.concatenate([Rotation.identity(), rot[:-1]])
+    pp = np.concatenate([np.zeros((1, 3)), x[:-1, :3]])
+    return np.concatenate([pr.inv().apply(x[:, :3] - pp), (pr.inv() * rot).as_quat(canonical=True)], axis=1)
+
+
+def _build():
+    from openpi import transforms as T
+    from openpi.training.config import DataConfigFactory, ModelTransformFactory
+
+    @dataclasses.dataclass(frozen=True)
+    class _P(T.DataTransformFn):
+        def __call__(self, d):
+            return {k: d[k] for k in _K if k in d}
+
+    @dataclasses.dataclass(frozen=True)
+    class _I(T.DataTransformFn):
+        n: int = 32
+
+        def __call__(self, d):
+            l = np.asarray(d[_K[0]])
+            r = np.asarray(d[_K[1]])
+            p = np.array(d[_K[2]], dtype=np.float64)
+            if p[6] < 0:
+                p[3:] *= -1
+            s = np.concatenate([p, np.asarray(d[_K[3]], dtype=np.float64)]).astype(np.float32)
+            o = {"state": T.pad_to_dim(s, self.n),
+                 "image": {"base_0_rgb": np.zeros_like(l), "left_wrist_0_rgb": l, "right_wrist_0_rgb": r},
+                 "image_mask": {"base_0_rgb": np.False_, "left_wrist_0_rgb": np.True_, "right_wrist_0_rgb": np.True_}}
+            if "prompt" in d:
+                o["prompt"] = d["prompt"]
+            return o
+
+    @dataclasses.dataclass(frozen=True)
+    class _O(T.DataTransformFn):
+        def __call__(self, d):
+            a = np.asarray(d["actions"], dtype=np.float64)
+            return {"actions": np.concatenate([_d(a[:, 0:9]), _d(a[:, 9:18]), a[:, 18:20]], axis=1).astype(np.float32)}
+
+    @dataclasses.dataclass(frozen=True)
+    class _C(DataConfigFactory):
+        def create(self, assets_dirs, model_config):
+            return dataclasses.replace(
+                self.create_base_config(assets_dirs, model_config),
+                repack_transforms=T.Group(inputs=[_P()]),
+                data_transforms=T.Group(inputs=[_I(model_config.action_dim)], outputs=[_O()]),
+                model_transforms=ModelTransformFactory()(model_config))
+
+    return _C
+
+
+class Policy:
+    def __init__(self, checkpoint_dir: str):
+        import openpi.models.pi0_config as pc
+        from openpi.policies import policy_config
+        from openpi.training.config import AssetsConfig, DataConfig, TrainConfig
+        from umi_arena.openpi_yubi import discover_asset_id
+
+        c = TrainConfig(name="p", model=pc.Pi0Config(pi05=True, action_dim=32, action_horizon=32),
+                        data=_build()(assets=AssetsConfig(asset_id=discover_asset_id(checkpoint_dir)),
+                                      base_config=DataConfig(prompt_from_task=True)))
+        t = time.monotonic()
+        self._p = policy_config.create_trained_policy(c, checkpoint_dir)
+        print(f"[adapter] checkpoint loaded in {time.monotonic() - t:.1f}s", flush=True)
+        t = time.monotonic()
+        g = np.random.default_rng(0)
+        x = g.integers(0, 256, size=(480, 640, 3), dtype=np.uint8)
+        self.infer({_K[0]: x, _K[1]: x.copy(), _K[2]: np.array([0.23, -0.05, 0.27, -0.58, -0.09, 0.72, 0.34], np.float32),
+                    _K[3]: np.array([0.24, 0.78], np.float32), "prompt": "pre-warm"})
+        print(f"[adapter] pre-warmed in {time.monotonic() - t:.1f}s", flush=True)
+
+    def infer(self, obs: dict) -> np.ndarray:
+        o = dict(obs)
+        if isinstance(o.get("prompt"), str):
+            o["prompt"] = o["prompt"].strip().rstrip(".")
+        return np.asarray(self._p.infer(o)["actions"], dtype=np.float32)

umi_arena.yaml

@@ -0,0 +1 @@
+runtime: openpi