train_state/_METADATA

{"tree_metadata": {"('opt_state', 'notfinite_count')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "notfinite_count", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('opt_state', 'last_finite')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "last_finite", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('opt_state', 'total_notfinite')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "total_notfinite", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('opt_state', 'inner_state', '0')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', '0', 'count')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_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": false, "write_shape": []}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'actor', 'block0_layer0_dense', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block0_layer0_dense", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'actor', 'block0_layer0_dense', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block0_layer0_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1446, 1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'actor', 'block0_layer0_ln', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block0_layer0_ln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'actor', 'block0_layer0_ln', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block0_layer0_ln", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'actor', 'block0_layer1_dense', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block0_layer1_dense", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'actor', 'block0_layer1_dense', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block0_layer1_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'actor', 'block0_layer1_ln', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block0_layer1_ln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'actor', 'block0_layer1_ln', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block0_layer1_ln", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'actor', 'block0_proj', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block0_proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'actor', 'block0_proj', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block0_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1446, 1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'actor', 'block1_layer0_dense', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block1_layer0_dense", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'actor', 'block1_layer0_dense', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block1_layer0_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'actor', 'block1_layer0_ln', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block1_layer0_ln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'actor', 'block1_layer0_ln', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block1_layer0_ln", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'actor', 'block1_layer1_dense', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block1_layer1_dense", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'actor', 'block1_layer1_dense', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block1_layer1_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'actor', 'block1_layer1_ln', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block1_layer1_ln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'actor', 'block1_layer1_ln', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block1_layer1_ln", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'actor', 'output', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "output", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [354]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'actor', 'output', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "output", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 354]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'actor', 'res_gate_0')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "res_gate_0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'actor', 'res_gate_1')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "res_gate_1", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'actor', 'tail_dense', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "tail_dense", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'actor', 'tail_dense', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "tail_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'actor', 'tail_ln', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "tail_ln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'actor', 'tail_ln', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "tail_ln", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'critic', 'block0_layer0_dense', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block0_layer0_dense", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'critic', 'block0_layer0_dense', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block0_layer0_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1446, 1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'critic', 'block0_layer0_ln', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block0_layer0_ln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'critic', 'block0_layer0_ln', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block0_layer0_ln", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'critic', 'block0_layer1_dense', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block0_layer1_dense", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'critic', 'block0_layer1_dense', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block0_layer1_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'critic', 'block0_layer1_ln', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block0_layer1_ln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'critic', 'block0_layer1_ln', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block0_layer1_ln", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'critic', 'block0_proj', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block0_proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'critic', 'block0_proj', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block0_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1446, 1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'critic', 'block1_layer0_dense', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block1_layer0_dense", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'critic', 'block1_layer0_dense', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block1_layer0_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'critic', 'block1_layer0_ln', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block1_layer0_ln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'critic', 'block1_layer0_ln', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block1_layer0_ln", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'critic', 'block1_layer1_dense', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block1_layer1_dense", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'critic', 'block1_layer1_dense', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block1_layer1_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'critic', 'block1_layer1_ln', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block1_layer1_ln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'critic', 'block1_layer1_ln', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block1_layer1_ln", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'critic', 'output', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "output", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'critic', 'output', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "output", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'critic', 'res_gate_0')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "res_gate_0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'critic', 'res_gate_1')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "res_gate_1", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'critic', 'tail_dense', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "tail_dense", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'critic', 'tail_dense', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "tail_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'critic', 'tail_ln', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "tail_ln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'critic', 'tail_ln', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "tail_ln", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'log_std')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "log_std", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [354]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'actor', 'block0_layer0_dense', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block0_layer0_dense", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'actor', 'block0_layer0_dense', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block0_layer0_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1446, 1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'actor', 'block0_layer0_ln', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block0_layer0_ln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'actor', 'block0_layer0_ln', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block0_layer0_ln", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'actor', 'block0_layer1_dense', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block0_layer1_dense", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'actor', 'block0_layer1_dense', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block0_layer1_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'actor', 'block0_layer1_ln', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block0_layer1_ln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'actor', 'block0_layer1_ln', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block0_layer1_ln", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'actor', 'block0_proj', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block0_proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'actor', 'block0_proj', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block0_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1446, 1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'actor', 'block1_layer0_dense', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block1_layer0_dense", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'actor', 'block1_layer0_dense', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block1_layer0_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'actor', 'block1_layer0_ln', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block1_layer0_ln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'actor', 'block1_layer0_ln', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block1_layer0_ln", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'actor', 'block1_layer1_dense', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block1_layer1_dense", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'actor', 'block1_layer1_dense', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block1_layer1_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'actor', 'block1_layer1_ln', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block1_layer1_ln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'actor', 'block1_layer1_ln', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block1_layer1_ln", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'actor', 'output', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "output", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [354]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'actor', 'output', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "output", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 354]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'actor', 'res_gate_0')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "res_gate_0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'actor', 'res_gate_1')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "res_gate_1", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'actor', 'tail_dense', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "tail_dense", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'actor', 'tail_dense', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "tail_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'actor', 'tail_ln', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "tail_ln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'actor', 'tail_ln', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "tail_ln", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'critic', 'block0_layer0_dense', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block0_layer0_dense", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'critic', 'block0_layer0_dense', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block0_layer0_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1446, 1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'critic', 'block0_layer0_ln', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block0_layer0_ln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'critic', 'block0_layer0_ln', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block0_layer0_ln", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'critic', 'block0_layer1_dense', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block0_layer1_dense", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'critic', 'block0_layer1_dense', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block0_layer1_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'critic', 'block0_layer1_ln', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block0_layer1_ln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'critic', 'block0_layer1_ln', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block0_layer1_ln", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'critic', 'block0_proj', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block0_proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'critic', 'block0_proj', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block0_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1446, 1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'critic', 'block1_layer0_dense', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block1_layer0_dense", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'critic', 'block1_layer0_dense', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block1_layer0_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'critic', 'block1_layer0_ln', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block1_layer0_ln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'critic', 'block1_layer0_ln', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block1_layer0_ln", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'critic', 'block1_layer1_dense', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block1_layer1_dense", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'critic', 'block1_layer1_dense', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block1_layer1_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'critic', 'block1_layer1_ln', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block1_layer1_ln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'critic', 'block1_layer1_ln', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block1_layer1_ln", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'critic', 'output', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "output", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'critic', 'output', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "output", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'critic', 'res_gate_0')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "res_gate_0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'critic', 'res_gate_1')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "res_gate_1", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'critic', 'tail_dense', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "tail_dense", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'critic', 'tail_dense', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "tail_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'critic', 'tail_ln', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "tail_ln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'critic', 'tail_ln', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "tail_ln", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'log_std')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "log_std", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [354]}}, "('opt_state', 'inner_state', '1', '1')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "1", "key_type": 1}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', '2')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "2", "key_type": 1}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('params', 'actor', 'block0_layer0_dense', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block0_layer0_dense", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'block0_layer0_dense', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block0_layer0_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1446, 1024]}}, "('params', 'actor', 'block0_layer0_ln', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block0_layer0_ln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'block0_layer0_ln', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block0_layer0_ln", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'block0_layer1_dense', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block0_layer1_dense", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'block0_layer1_dense', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block0_layer1_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('params', 'actor', 'block0_layer1_ln', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block0_layer1_ln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'block0_layer1_ln', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block0_layer1_ln", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'block0_proj', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block0_proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'block0_proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block0_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1446, 1024]}}, "('params', 'actor', 'block1_layer0_dense', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block1_layer0_dense", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'block1_layer0_dense', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block1_layer0_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('params', 'actor', 'block1_layer0_ln', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block1_layer0_ln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'block1_layer0_ln', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block1_layer0_ln", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'block1_layer1_dense', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block1_layer1_dense", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'block1_layer1_dense', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block1_layer1_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('params', 'actor', 'block1_layer1_ln', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block1_layer1_ln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'block1_layer1_ln', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "block1_layer1_ln", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'output', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "output", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [354]}}, "('params', 'actor', 'output', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "output", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 354]}}, "('params', 'actor', 'res_gate_0')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "res_gate_0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('params', 'actor', 'res_gate_1')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "res_gate_1", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('params', 'actor', 'tail_dense', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "tail_dense", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'tail_dense', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "tail_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('params', 'actor', 'tail_ln', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "tail_ln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'tail_ln', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "tail_ln", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'block0_layer0_dense', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block0_layer0_dense", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'block0_layer0_dense', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block0_layer0_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1446, 1024]}}, "('params', 'critic', 'block0_layer0_ln', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block0_layer0_ln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'block0_layer0_ln', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block0_layer0_ln", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'block0_layer1_dense', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block0_layer1_dense", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'block0_layer1_dense', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block0_layer1_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('params', 'critic', 'block0_layer1_ln', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block0_layer1_ln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'block0_layer1_ln', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block0_layer1_ln", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'block0_proj', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block0_proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'block0_proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block0_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1446, 1024]}}, "('params', 'critic', 'block1_layer0_dense', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block1_layer0_dense", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'block1_layer0_dense', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block1_layer0_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('params', 'critic', 'block1_layer0_ln', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block1_layer0_ln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'block1_layer0_ln', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block1_layer0_ln", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'block1_layer1_dense', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block1_layer1_dense", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'block1_layer1_dense', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block1_layer1_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('params', 'critic', 'block1_layer1_ln', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block1_layer1_ln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'block1_layer1_ln', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "block1_layer1_ln", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'output', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "output", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('params', 'critic', 'output', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "output", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1]}}, "('params', 'critic', 'res_gate_0')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "res_gate_0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('params', 'critic', 'res_gate_1')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "res_gate_1", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('params', 'critic', 'tail_dense', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "tail_dense", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'tail_dense', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "tail_dense", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('params', 'critic', 'tail_ln', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "tail_ln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'tail_ln', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "tail_ln", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'log_std')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "log_std", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [354]}}, "('run_stats', 'RunningMeanStd_0', 'count')": {"key_metadata": [{"key": "run_stats", "key_type": 2}, {"key": "RunningMeanStd_0", "key_type": 2}, {"key": "count", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('run_stats', 'RunningMeanStd_0', 'mean')": {"key_metadata": [{"key": "run_stats", "key_type": 2}, {"key": "RunningMeanStd_0", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1446]}}, "('run_stats', 'RunningMeanStd_0', 'var')": {"key_metadata": [{"key": "run_stats", "key_type": 2}, {"key": "RunningMeanStd_0", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1446]}}, "('step',)": {"key_metadata": [{"key": "step", "key_type": 2}], "value_metadata": {"value_type": "scalar", "skip_deserialize": false}}}, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}