{"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', 'inner_states', 'adam', 'inner_state', '0', 'count')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_0", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'Dense_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_1", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'Dense_10', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_10", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'Dense_10', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_10", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'Dense_11', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_11", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'Dense_11', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_11", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'Dense_12', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_12", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'Dense_12', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_12", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'Dense_13', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_13", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'Dense_13', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_13", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'Dense_14', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_14", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'Dense_14', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_14", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'Dense_15', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_15", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'Dense_15', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_15", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'Dense_16', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_16", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'Dense_16', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_16", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'Dense_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_2", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'Dense_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_3", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'Dense_4', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_4", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'Dense_4', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'Dense_5', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_5", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'Dense_5', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_5", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'Dense_6', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_6", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'Dense_6', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'Dense_7', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_7", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'Dense_7', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_7", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'Dense_8', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_8", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'Dense_8', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_8", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'Dense_9', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_9", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'Dense_9', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_9", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_0", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_0", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_1", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_1", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_10', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_10", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_10', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_10", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_11', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_11", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_11', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_11", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_12', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_12", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_12', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_12", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_13', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_13", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_13', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_13", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_14', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_14", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_14', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_14", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_15', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_15", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_15', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_15", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_2", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_2", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_3", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_3", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_4', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_4", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_4', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_4", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_5', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_5", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_5', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_5", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_6', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_6", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_6', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_6", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_7', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_7", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_7', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_7", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_8', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_8", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_8', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_8", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_9', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_9", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_9', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_9", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_0", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'Dense_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_1", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'Dense_10', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_10", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'Dense_10', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_10", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'Dense_11', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_11", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'Dense_11', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_11", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'Dense_12', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_12", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'Dense_12', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_12", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'Dense_13', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_13", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'Dense_13', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_13", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'Dense_14', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_14", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'Dense_14', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_14", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'Dense_15', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_15", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'Dense_15', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_15", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'Dense_16', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_16", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'Dense_16', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_16", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'Dense_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_2", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'Dense_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_3", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'Dense_4', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_4", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'Dense_4', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'Dense_5', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_5", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'Dense_5', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_5", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'Dense_6', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_6", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'Dense_6', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'Dense_7', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_7", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'Dense_7', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_7", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'Dense_8', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_8", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'Dense_8', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_8", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'Dense_9', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_9", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'Dense_9', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_9", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_0", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_0", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_1", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_1", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_10', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_10", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_10', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_10", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_11', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_11", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_11', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_11", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_12', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_12", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_12', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_12", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_13', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_13", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_13', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_13", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_14', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_14", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_14', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_14", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_15', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_15", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_15', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_15", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_2", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_2", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_3", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_3", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_4', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_4", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_4', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_4", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_5', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_5", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_5', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_5", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_6', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_6", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_6', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_6", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_7', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_7", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_7', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_7", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_8', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_8", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_8', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_8", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_9', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_9", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_9', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_9", "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', 'inner_states', 'adam', 'inner_state', '0', 'mu', 'log_std')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_0", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'Dense_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_1", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'Dense_10', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_10", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'Dense_10', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_10", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'Dense_11', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_11", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'Dense_11', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_11", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'Dense_12', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_12", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'Dense_12', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_12", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'Dense_13', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_13", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'Dense_13', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_13", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'Dense_14', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_14", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'Dense_14', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_14", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'Dense_15', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_15", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'Dense_15', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_15", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'Dense_16', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_16", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'Dense_16', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_16", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'Dense_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_2", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'Dense_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_3", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'Dense_4', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_4", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'Dense_4', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'Dense_5', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_5", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'Dense_5', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_5", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'Dense_6', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_6", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'Dense_6', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'Dense_7', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_7", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'Dense_7', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_7", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'Dense_8', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_8", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'Dense_8', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_8", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'Dense_9', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_9", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'Dense_9', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_9", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'LayerNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_0", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'LayerNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_0", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'LayerNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_1", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'LayerNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_1", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'LayerNorm_10', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_10", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'LayerNorm_10', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_10", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'LayerNorm_11', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_11", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'LayerNorm_11', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_11", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'LayerNorm_12', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_12", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'LayerNorm_12', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_12", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'LayerNorm_13', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_13", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'LayerNorm_13', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_13", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'LayerNorm_14', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_14", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'LayerNorm_14', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_14", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'LayerNorm_15', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_15", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'LayerNorm_15', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_15", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'LayerNorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_2", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'LayerNorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_2", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'LayerNorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_3", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'LayerNorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_3", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'LayerNorm_4', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_4", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'LayerNorm_4', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_4", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'LayerNorm_5', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_5", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'LayerNorm_5', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_5", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'LayerNorm_6', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_6", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'LayerNorm_6', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_6", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'LayerNorm_7', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_7", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'LayerNorm_7', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_7", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'LayerNorm_8', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_8", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'LayerNorm_8', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_8", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'LayerNorm_9', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_9", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'actor', 'LayerNorm_9', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_9", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_0", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'Dense_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_1", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'Dense_10', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_10", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'Dense_10', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_10", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'Dense_11', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_11", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'Dense_11', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_11", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'Dense_12', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_12", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'Dense_12', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_12", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'Dense_13', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_13", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'Dense_13', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_13", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'Dense_14', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_14", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'Dense_14', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_14", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'Dense_15', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_15", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'Dense_15', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_15", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'Dense_16', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_16", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'Dense_16', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_16", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'Dense_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_2", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'Dense_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_3", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'Dense_4', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_4", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'Dense_4', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'Dense_5', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_5", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'Dense_5', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_5", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'Dense_6', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_6", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'Dense_6', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'Dense_7', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_7", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'Dense_7', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_7", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'Dense_8', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_8", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'Dense_8', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_8", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'Dense_9', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_9", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'Dense_9', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_9", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'LayerNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_0", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'LayerNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_0", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'LayerNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_1", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'LayerNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_1", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'LayerNorm_10', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_10", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'LayerNorm_10', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_10", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'LayerNorm_11', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_11", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'LayerNorm_11', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_11", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'LayerNorm_12', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_12", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'LayerNorm_12', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_12", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'LayerNorm_13', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_13", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'LayerNorm_13', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_13", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'LayerNorm_14', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_14", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'LayerNorm_14', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_14", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'LayerNorm_15', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_15", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'LayerNorm_15', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_15", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'LayerNorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_2", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'LayerNorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_2", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'LayerNorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_3", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'LayerNorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_3", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'LayerNorm_4', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_4", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'LayerNorm_4', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_4", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'LayerNorm_5', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_5", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'LayerNorm_5', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_5", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'LayerNorm_6', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_6", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'LayerNorm_6', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_6", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'LayerNorm_7', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_7", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'LayerNorm_7', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_7", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'LayerNorm_8', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_8", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'LayerNorm_8', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_8", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'LayerNorm_9', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_9", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'critic', 'LayerNorm_9', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_9", "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', 'inner_states', 'adam', 'inner_state', '0', 'nu', 'log_std')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"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', 'inner_states', 'adam', 'inner_state', '1')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'adam', 'inner_state', '2')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "adam", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "2", "key_type": 1}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'count')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"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', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2418, 1024]}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'Dense_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_1", "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', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'Dense_10', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_10", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'Dense_10', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_10", "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', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'Dense_11', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_11", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'Dense_11', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_11", "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', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'Dense_12', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_12", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'Dense_12', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_12", "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', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'Dense_13', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_13", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'Dense_13', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_13", "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', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'Dense_14', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_14", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'Dense_14', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_14", "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', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'Dense_15', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_15", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'Dense_15', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_15", "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', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'Dense_16', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_16", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'Dense_16', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_16", "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', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'Dense_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_2", "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', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'Dense_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_3", "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', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'Dense_4', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'Dense_4', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_4", "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', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'Dense_5', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_5", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'Dense_5', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_5", "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', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'Dense_6', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'Dense_6', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_6", "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', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'Dense_7', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_7", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'Dense_7', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_7", "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', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'Dense_8', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_8", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'Dense_8', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_8", "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', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'Dense_9', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_9", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'Dense_9', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_9", "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', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_10', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_10", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_10', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_10", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_11', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_11", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_11', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_11", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_12', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_12", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_12', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_12", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_13', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_13", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_13', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_13", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_14', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_14", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_14', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_14", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_15', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_15", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_15', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_15", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_4', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_4', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_4", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_5', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_5", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_5', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_5", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_6', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_6', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_6", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_7', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_7", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_7', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_7", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_8', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_8", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_8', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_8", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_9', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_9", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'actor', 'LayerNorm_9', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_9", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2418, 1024]}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'Dense_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_1", "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', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'Dense_10', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_10", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'Dense_10', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_10", "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', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'Dense_11', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_11", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'Dense_11', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_11", "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', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'Dense_12', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_12", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'Dense_12', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_12", "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', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'Dense_13', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_13", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'Dense_13', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_13", "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', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'Dense_14', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_14", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'Dense_14', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_14", "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', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'Dense_15', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_15", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'Dense_15', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_15", "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', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'Dense_16', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_16", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'Dense_16', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_16", "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', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'Dense_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_2", "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', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'Dense_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_3", "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', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'Dense_4', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'Dense_4', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_4", "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', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'Dense_5', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_5", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'Dense_5', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_5", "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', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'Dense_6', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'Dense_6', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_6", "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', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'Dense_7', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_7", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'Dense_7', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_7", "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', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'Dense_8', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_8", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'Dense_8', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_8", "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', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'Dense_9', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_9", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'Dense_9', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_9", "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', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_10', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_10", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_10', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_10", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_11', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_11", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_11', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_11", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_12', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_12", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_12', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_12", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_13', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_13", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_13', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_13", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_14', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_14", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_14', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_14", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_15', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_15", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_15', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_15", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_4', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_4', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_4", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_5', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_5", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_5', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_5", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_6', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_6', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_6", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_7', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_7", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_7', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_7", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_8', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_8", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_8', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_8", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_9', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_9", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'critic', 'LayerNorm_9', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_9", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'mu', 'log_std')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "log_std", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '0', 'ns_coeffs')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "ns_coeffs", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '1')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', 'inner_states', 'muon', 'inner_state', '2')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "inner_states", "key_type": 2}, {"key": "muon", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "2", "key_type": 1}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('params', 'actor', 'Dense_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2418, 1024]}}, "('params', 'actor', 'Dense_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('params', 'actor', 'Dense_10', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_10", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'Dense_10', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_10", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('params', 'actor', 'Dense_11', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_11", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'Dense_11', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_11", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('params', 'actor', 'Dense_12', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_12", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'Dense_12', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_12", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('params', 'actor', 'Dense_13', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_13", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'Dense_13', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_13", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('params', 'actor', 'Dense_14', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_14", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'Dense_14', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_14", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('params', 'actor', 'Dense_15', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_15", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'Dense_15', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_15", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('params', 'actor', 'Dense_16', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_16", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [354]}}, "('params', 'actor', 'Dense_16', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_16", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 354]}}, "('params', 'actor', 'Dense_2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('params', 'actor', 'Dense_3', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('params', 'actor', 'Dense_4', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'Dense_4', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('params', 'actor', 'Dense_5', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_5", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'Dense_5', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_5", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('params', 'actor', 'Dense_6', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'Dense_6', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('params', 'actor', 'Dense_7', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_7", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'Dense_7', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_7", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('params', 'actor', 'Dense_8', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_8", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'Dense_8', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_8", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('params', 'actor', 'Dense_9', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_9", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'Dense_9', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "Dense_9", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('params', 'actor', 'LayerNorm_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'LayerNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'LayerNorm_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'LayerNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'LayerNorm_10', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_10", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'LayerNorm_10', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_10", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'LayerNorm_11', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_11", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'LayerNorm_11', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_11", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'LayerNorm_12', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_12", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'LayerNorm_12', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_12", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'LayerNorm_13', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_13", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'LayerNorm_13', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_13", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'LayerNorm_14', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_14", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'LayerNorm_14', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_14", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'LayerNorm_15', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_15", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'LayerNorm_15', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_15", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'LayerNorm_2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'LayerNorm_2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'LayerNorm_3', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'LayerNorm_3', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'LayerNorm_4', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'LayerNorm_4', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_4", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'LayerNorm_5', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_5", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'LayerNorm_5', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_5", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'LayerNorm_6', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'LayerNorm_6', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_6", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'LayerNorm_7', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_7", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'LayerNorm_7', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_7", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'LayerNorm_8', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_8", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'LayerNorm_8', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_8", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'LayerNorm_9', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_9", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'actor', 'LayerNorm_9', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "actor", "key_type": 2}, {"key": "LayerNorm_9", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'Dense_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2418, 1024]}}, "('params', 'critic', 'Dense_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('params', 'critic', 'Dense_10', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_10", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'Dense_10', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_10", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('params', 'critic', 'Dense_11', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_11", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'Dense_11', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_11", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('params', 'critic', 'Dense_12', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_12", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'Dense_12', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_12", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('params', 'critic', 'Dense_13', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_13", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'Dense_13', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_13", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('params', 'critic', 'Dense_14', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_14", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'Dense_14', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_14", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('params', 'critic', 'Dense_15', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_15", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'Dense_15', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_15", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('params', 'critic', 'Dense_16', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_16", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('params', 'critic', 'Dense_16', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_16", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1]}}, "('params', 'critic', 'Dense_2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('params', 'critic', 'Dense_3', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('params', 'critic', 'Dense_4', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'Dense_4', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('params', 'critic', 'Dense_5', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_5", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'Dense_5', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_5", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('params', 'critic', 'Dense_6', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'Dense_6', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('params', 'critic', 'Dense_7', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_7", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'Dense_7', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_7", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('params', 'critic', 'Dense_8', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_8", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'Dense_8', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_8", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('params', 'critic', 'Dense_9', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_9", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'Dense_9', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "Dense_9", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('params', 'critic', 'LayerNorm_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'LayerNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'LayerNorm_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'LayerNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'LayerNorm_10', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_10", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'LayerNorm_10', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_10", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'LayerNorm_11', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_11", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'LayerNorm_11', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_11", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'LayerNorm_12', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_12", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'LayerNorm_12', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_12", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'LayerNorm_13', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_13", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'LayerNorm_13', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_13", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'LayerNorm_14', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_14", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'LayerNorm_14', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_14", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'LayerNorm_15', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_15", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'LayerNorm_15', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_15", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'LayerNorm_2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'LayerNorm_2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'LayerNorm_3', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'LayerNorm_3', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'LayerNorm_4', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'LayerNorm_4', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_4", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'LayerNorm_5', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_5", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'LayerNorm_5', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_5", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'LayerNorm_6', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'LayerNorm_6', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_6", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'LayerNorm_7', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_7", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'LayerNorm_7', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_7", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'LayerNorm_8', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_8", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'LayerNorm_8', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_8", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'LayerNorm_9', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_9", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('params', 'critic', 'LayerNorm_9', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "critic", "key_type": 2}, {"key": "LayerNorm_9", "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": [2418]}}, "('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": [2418]}}, "('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}