diff --git "a/RWA_LR_JOCN_retrain/_METADATA" "b/RWA_LR_JOCN_retrain/_METADATA"
new file mode 100644--- /dev/null
+++ "b/RWA_LR_JOCN_retrain/_METADATA"
@@ -0,0 +1 @@
+{"tree_metadata": {"('config',)": {"key_metadata": [{"key": "config", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'step')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "step", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'Dense_2', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'Dense_3', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'Dense_4', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'Dense_4', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'Dense_5', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_5", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'Dense_5', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_5", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_0', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_1', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_2', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_2', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_3', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_3', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_4', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_4', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_4", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_5', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_5", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_5', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_5", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_6', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_6', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_6", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_7', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_7", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_7', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_7", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_8', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_8", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_8', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_8", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_1', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_1', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_1', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_1", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_1', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_1", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_10', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_10", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_10', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_10", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_10', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_10", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_10', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_10", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_11', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_11", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_11', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_11", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_11', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_11", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_11', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_11", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_2', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_2", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_2', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_2", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_2', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_2", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_2', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_2", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_3', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_3", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_3', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_3", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_3', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_3", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_3', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_3", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_4', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_4", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_4', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_4", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_4', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_4", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_4', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_4", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_5', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_5", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_5', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_5", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_5', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_5", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_5', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_5", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_6', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_6", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_6', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_6", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_6', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_6", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_6', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_6", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_7', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_7", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_7', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_7", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_7', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_7", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_7', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_7", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_8', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_8", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_8', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_8", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_8', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_8", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_8', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_8", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_9', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_9", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_9', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_9", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_9', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_9", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_9', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_9", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'Dense_2', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'Dense_2', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'Dense_3', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'Dense_4', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'Dense_4', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'Dense_5', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_5", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'Dense_5', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_5", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_0', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_1', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_2', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_2', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_3', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_3', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_4', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_4', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_4", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_5', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_5", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_5', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_5", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_6', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_6', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_6", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_7', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_7", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_7', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_7", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_8', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_8", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_8', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_8", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_1', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_1', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_1', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_1", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_1', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_1", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_10', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_10", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_10', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_10", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_10', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_10", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_10', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_10", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_11', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_11", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_11', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_11", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_11', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_11", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_11', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_11", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_2', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_2", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_2', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_2", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_2', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_2", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_2', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_2", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_3', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_3", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_3', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_3", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_3', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_3", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_3', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_3", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_4', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_4", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_4', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_4", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_4', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_4", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_4', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_4", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_5', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_5", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_5', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_5", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_5', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_5", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_5', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_5", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_6', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_6", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_6', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_6", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_6', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_6", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_6', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_6", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_7', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_7", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_7', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_7", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_7', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_7", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_7', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_7", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_8', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_8", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_8', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_8", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_8', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_8", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_8', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_8", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_9', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_9", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_9', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_9", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_9', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_9", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'params', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_9', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_9", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '0')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'count')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "count", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'Dense_2', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'Dense_3', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'Dense_4', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'Dense_4', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'Dense_5', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_5", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'Dense_5', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_5", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_0', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_1', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_2', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_2', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_3', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_3', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_4', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_4', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_4", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_5', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_5", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_5', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_5", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_6', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_6', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_6", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_7', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_7", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_7', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_7", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_8', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_8", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_8', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_8", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_1', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_1', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_1', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_1", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_1', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_1", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_10', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_10", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_10', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_10", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_10', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_10", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_10', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_10", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_11', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_11", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_11', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_11", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_11', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_11", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_11', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_11", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_2', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_2", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_2', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_2", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_2', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_2", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_2', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_2", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_3', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_3", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_3', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_3", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_3', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_3", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_3', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_3", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_4', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_4", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_4', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_4", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_4', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_4", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_4', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_4", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_5', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_5", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_5', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_5", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_5', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_5", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_5', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_5", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_6', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_6", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_6', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_6", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_6', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_6", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_6', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_6", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_7', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_7", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_7', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_7", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_7', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_7", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_7', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_7", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_8', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_8", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_8', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_8", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_8', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_8", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_8', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_8", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_9', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_9", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_9', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_9", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_9', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_9", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_9', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_9", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'Dense_2', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'Dense_2', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'Dense_3', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'Dense_4', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'Dense_4', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'Dense_5', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_5", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'Dense_5', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_5", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_0', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_1', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_2', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_2', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_3', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_3', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_4', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_4', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_4", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_5', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_5", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_5', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_5", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_6', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_6', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_6", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_7', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_7", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_7', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_7", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_8', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_8", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_8', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_8", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_1', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_1', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_1', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_1", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_1', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_1", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_10', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_10", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_10', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_10", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_10', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_10", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_10', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_10", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_11', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_11", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_11', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_11", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_11', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_11", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_11', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_11", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_2', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_2", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_2', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_2", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_2', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_2", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_2', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_2", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_3', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_3", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_3', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_3", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_3', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_3", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_3', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_3", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_4', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_4", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_4', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_4", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_4', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_4", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_4', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_4", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_5', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_5", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_5', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_5", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_5', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_5", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_5', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_5", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_6', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_6", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_6', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_6", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_6', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_6", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_6', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_6", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_7', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_7", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_7', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_7", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_7', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_7", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_7', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_7", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_8', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_8", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_8', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_8", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_8', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_8", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_8', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_8", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_9', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_9", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_9', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_9", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_9', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_9", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'mu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_9', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_9", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'Dense_2', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'Dense_3', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'Dense_4', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'Dense_4', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'Dense_5', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_5", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'Dense_5', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_5", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_0', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_1', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_2', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_2', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_3', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_3', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_4', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_4', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_4", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_5', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_5", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_5', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_5", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_6', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_6', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_6", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_7', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_7", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_7', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_7", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_8', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_8", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'LayerNorm_8', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_8", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_1', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_1', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_1', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_1", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_1', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_1", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_10', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_10", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_10', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_10", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_10', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_10", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_10', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_10", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_11', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_11", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_11', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_11", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_11', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_11", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_11', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_11", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_2', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_2", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_2', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_2", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_2', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_2", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_2', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_2", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_3', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_3", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_3', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_3", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_3', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_3", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_3', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_3", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_4', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_4", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_4', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_4", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_4', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_4", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_4', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_4", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_5', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_5", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_5', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_5", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_5', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_5", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_5', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_5", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_6', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_6", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_6', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_6", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_6', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_6", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_6', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_6", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_7', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_7", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_7', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_7", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_7', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_7", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_7', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_7", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_8', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_8", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_8', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_8", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_8', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_8", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_8', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_8", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_9', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_9", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_9', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_9", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_9', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_9", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'ActorGNN_0', 'GraphNet_0', 'MLP_9', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "ActorGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_9", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'Dense_2', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'Dense_2', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'Dense_3', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'Dense_4', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'Dense_4', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'Dense_5', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_5", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'Dense_5', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "Dense_5", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_0', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_1', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_2', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_2', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_3', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_3', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_4', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_4', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_4", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_5', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_5", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_5', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_5", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_6', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_6', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_6", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_7', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_7", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_7', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_7", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_8', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_8", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'LayerNorm_8', 'scale')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "LayerNorm_8", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_1', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_1', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_1', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_1", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_1', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_1", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_10', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_10", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_10', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_10", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_10', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_10", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_10', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_10", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_11', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_11", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_11', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_11", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_11', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_11", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_11', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_11", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_2', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_2", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_2', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_2", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_2', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_2", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_2', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_2", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_3', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_3", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_3', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_3", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_3', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_3", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_3', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_3", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_4', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_4", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_4', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_4", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_4', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_4", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_4', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_4", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_5', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_5", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_5', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_5", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_5', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_5", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_5', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_5", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_6', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_6", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_6', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_6", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_6', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_6", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_6', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_6", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_7', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_7", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_7', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_7", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_7', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_7", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_7', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_7", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_8', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_8", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_8', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_8", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_8', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_8", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_8', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_8", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_9', 'Dense_0', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_9", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_9', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_9", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_9', 'Dense_1', 'bias')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_9", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '0', 'nu', 'params', 'CriticGNN_0', 'GraphNet_0', 'MLP_9', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "CriticGNN_0", "key_type": 2}, {"key": "GraphNet_0", "key_type": 2}, {"key": "MLP_9", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}, "('model', 'opt_state', '1', '1', 'count')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "1", "key_type": 1}, {"key": "count", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": true}}}, "use_zarr3": false}
\ No newline at end of file