Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+train_state/ocdbt.process_0/d/78d8c2bf66e574d6dc14d065e83d0cc9 filter=lfs diff=lfs merge=lfs -text
+train_state/ocdbt.process_0/d/dac1a25d9bb1525af4afac14e4aa6abf filter=lfs diff=lfs merge=lfs -text

_CHECKPOINT_METADATA

@@ -0,0 +1 @@
+{"item_handlers": {"config": "orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler", "metadata": "orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler", "train_state": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1765227268487521571, "commit_timestamp_nsecs": 1765227268871681595, "custom_metadata": {}}

config/metadata

@@ -0,0 +1 @@
+{"wandb": {"project": "musclemimic", "mode": "online"}, "experiment": {"task_factory": {"name": "ImitationFactory", "params": {"amass_dataset_conf": {"dataset_group": "KIT_KINESIS_TRAINING_MOTIONS"}}}, "env_params": {"env_name": "MjxMyoFullBody", "headless": true, "horizon": 1000, "disable_fingers": true, "mjx_backend": "warp", "num_envs": 8192, "enable_muscle_length_observations": true, "enable_muscle_velocity_observations": true, "enable_muscle_force_observations": true, "enable_muscle_activation_observations": true, "enable_touch_sensor_observations": true, "init_state_type": "TrajInitialStateHandler", "goal_type": "GoalTrajMimic", "goal_params": {"visualize_goal": true, "n_step_lookahead": 1, "upper_body_xml_name": "torso", "sites_for_mimic": ["pelvis_mimic", "upper_body_mimic", "head_mimic", "left_shoulder_mimic", "left_elbow_mimic", "left_hand_mimic", "right_shoulder_mimic", "right_elbow_mimic", "right_hand_mimic", "left_hip_mimic", "left_knee_mimic", "left_ankle_mimic", "left_toes_mimic", "right_hip_mimic", "right_knee_mimic", "right_ankle_mimic", "right_toes_mimic"]}, "control_type": "DefaultControl", "control_params": {}, "reward_type": "MimicReward", "reward_params": {"qpos_w_sum": 0.5, "qvel_w_sum": 0.1, "rpos_w_sum": 2.0, "rquat_w_sum": 0.3, "rvel_w_sum": 0.2, "use_mean_exp_reward": false, "sites_for_mimic": ["pelvis_mimic", "upper_body_mimic", "head_mimic", "left_shoulder_mimic", "left_elbow_mimic", "left_hand_mimic", "right_shoulder_mimic", "right_elbow_mimic", "right_hand_mimic", "left_hip_mimic", "left_knee_mimic", "left_ankle_mimic", "left_toes_mimic", "right_hip_mimic", "right_knee_mimic", "right_ankle_mimic", "right_toes_mimic"]}, "terminal_state_type": "MeanRelativeSiteDeviationTerminalStateHandler", "terminal_state_params": {"mean_site_deviation_threshold": 0.3, "enable_site_check": true}}, "hidden_layers": [1024, 1024, 1024, 1024, 1024], "lr": 4e-05, "num_envs": 8192, "num_steps": 20, "total_timesteps": 4915200000, "update_epochs": 1, "proportion_env_reward": 0.0, "num_minibatches": 512, "gamma": 0.99, "gae_lambda": 0.95, "clip_eps": 0.3, "init_std": 1.0, "learnable_std": true, "ent_coef": 0.0, "vf_coef": 0.5, "max_grad_norm": 5, "activation": "silu", "anneal_lr": false, "lr_schedule_type": "linear", "warmup_steps": null, "min_lr_ratio": 0.1, "weight_decay": 0, "normalize_env": true, "debug": false, "n_seeds": 1, "seeds": null, "vmap_across_seeds": false, "save_checkpoints": true, "max_checkpoints_to_keep": 10, "checkpoints_on_validation": true, "resume_from": null, "async_checkpointing": false, "reset_logging_timestep": false, "use_moe": false, "moe_config": {"num_experts": 4, "moe_at_layers": [2, 3], "apply_moe_to": "both", "apply_moe_to_output": false, "expert_hidden_ratio": 0.5, "temperature": 0.7, "load_balance_loss_weight": 0.01, "gate_init_std": 0.01}, "validation": {"active": true, "deterministic": false, "num_steps": 1000, "num_envs": 20, "num": 50, "video_length": 1000, "video_frequency": 1, "video_fps": 30, "quantities": ["JointPosition", "JointVelocity", "RelSitePosition", "RelSiteVelocity", "RelSiteOrientation", "BodyPosition"], "measures": ["EuclideanDistance", "DynamicTimeWarping", "DiscreteFrechetDistance"], "rel_site_names": ["pelvis_mimic", "upper_body_mimic", "head_mimic", "left_shoulder_mimic", "left_elbow_mimic", "left_hand_mimic", "right_shoulder_mimic", "right_elbow_mimic", "right_hand_mimic", "left_hip_mimic", "left_knee_mimic", "left_ankle_mimic", "left_toes_mimic", "right_hip_mimic", "right_knee_mimic", "right_ankle_mimic", "right_toes_mimic"]}, "num_updates": 30000, "minibatch_size": 320, "validation_interval": 600, "effective_utd": 1.0, "checkpoint_dir": "/users/ckli/musclemimic_dev/outputs/2025-12-04/02-26-11/checkpoints/251204T023339-pid155-3ad911"}, "control_config": {"torque_control": {"control_type": "DefaultControl", "control_params": {}}}}

metadata/metadata

@@ -0,0 +1 @@
+{"step": 15360000, "update_number": 30000, "global_timestep": 4915200000, "learning_rate": 3.9999998989515007e-05, "num_envs": 8192, "num_steps": 20, "num_minibatches": 512, "update_epochs": 1, "schema_version": "2.1", "algo_version": "PPOJax_v1", "backend": "warp", "env_name": "NormalizeVecReward"}

train_state/_METADATA

@@ -0,0 +1 @@
+{"tree_metadata": {"('opt_state', 'notfinite_count')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "notfinite_count", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('opt_state', 'last_finite')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "last_finite", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('opt_state', 'total_notfinite')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "total_notfinite", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('opt_state', 'inner_state', '0')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', '0', 'count')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "count", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'FullyConnectedNet_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "FullyConnectedNet_0", "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', '0', 'mu', 'FullyConnectedNet_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "FullyConnectedNet_0", "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": [2154, 1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'FullyConnectedNet_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "FullyConnectedNet_0", "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', '0', 'mu', 'FullyConnectedNet_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "FullyConnectedNet_0", "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', '0', 'mu', 'FullyConnectedNet_0', 'Dense_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "FullyConnectedNet_0", "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', '0', 'mu', 'FullyConnectedNet_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "FullyConnectedNet_0", "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', '0', 'mu', 'FullyConnectedNet_0', 'Dense_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "FullyConnectedNet_0", "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', '0', 'mu', 'FullyConnectedNet_0', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "FullyConnectedNet_0", "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', '0', 'mu', 'FullyConnectedNet_0', 'Dense_4', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "FullyConnectedNet_0", "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', '0', 'mu', 'FullyConnectedNet_0', 'Dense_4', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "FullyConnectedNet_0", "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', '0', 'mu', 'FullyConnectedNet_0', 'Dense_5', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "FullyConnectedNet_0", "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": [354]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'FullyConnectedNet_0', 'Dense_5', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "FullyConnectedNet_0", "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, 354]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'FullyConnectedNet_1', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "FullyConnectedNet_1", "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', '0', 'mu', 'FullyConnectedNet_1', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "FullyConnectedNet_1", "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": [2154, 1024]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'FullyConnectedNet_1', 'Dense_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "FullyConnectedNet_1", "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', '0', 'mu', 'FullyConnectedNet_1', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "FullyConnectedNet_1", "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', '0', 'mu', 'FullyConnectedNet_1', 'Dense_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "FullyConnectedNet_1", "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', '0', 'mu', 'FullyConnectedNet_1', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "FullyConnectedNet_1", "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', '0', 'mu', 'FullyConnectedNet_1', 'Dense_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "FullyConnectedNet_1", "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', '0', 'mu', 'FullyConnectedNet_1', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "FullyConnectedNet_1", "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', '0', 'mu', 'FullyConnectedNet_1', 'Dense_4', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "FullyConnectedNet_1", "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', '0', 'mu', 'FullyConnectedNet_1', 'Dense_4', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "FullyConnectedNet_1", "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', '0', 'mu', 'FullyConnectedNet_1', 'Dense_5', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "FullyConnectedNet_1", "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": [1]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'FullyConnectedNet_1', 'Dense_5', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "FullyConnectedNet_1", "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, 1]}}, "('opt_state', 'inner_state', '1', '0', 'mu', 'log_std')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "log_std", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [354]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'FullyConnectedNet_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "FullyConnectedNet_0", "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', '0', 'nu', 'FullyConnectedNet_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "FullyConnectedNet_0", "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": [2154, 1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'FullyConnectedNet_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "FullyConnectedNet_0", "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', '0', 'nu', 'FullyConnectedNet_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "FullyConnectedNet_0", "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', '0', 'nu', 'FullyConnectedNet_0', 'Dense_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "FullyConnectedNet_0", "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', '0', 'nu', 'FullyConnectedNet_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "FullyConnectedNet_0", "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', '0', 'nu', 'FullyConnectedNet_0', 'Dense_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "FullyConnectedNet_0", "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', '0', 'nu', 'FullyConnectedNet_0', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "FullyConnectedNet_0", "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', '0', 'nu', 'FullyConnectedNet_0', 'Dense_4', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "FullyConnectedNet_0", "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', '0', 'nu', 'FullyConnectedNet_0', 'Dense_4', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "FullyConnectedNet_0", "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', '0', 'nu', 'FullyConnectedNet_0', 'Dense_5', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "FullyConnectedNet_0", "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": [354]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'FullyConnectedNet_0', 'Dense_5', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "FullyConnectedNet_0", "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, 354]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'FullyConnectedNet_1', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "FullyConnectedNet_1", "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', '0', 'nu', 'FullyConnectedNet_1', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "FullyConnectedNet_1", "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": [2154, 1024]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'FullyConnectedNet_1', 'Dense_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "FullyConnectedNet_1", "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', '0', 'nu', 'FullyConnectedNet_1', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "FullyConnectedNet_1", "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', '0', 'nu', 'FullyConnectedNet_1', 'Dense_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "FullyConnectedNet_1", "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', '0', 'nu', 'FullyConnectedNet_1', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "FullyConnectedNet_1", "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', '0', 'nu', 'FullyConnectedNet_1', 'Dense_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "FullyConnectedNet_1", "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', '0', 'nu', 'FullyConnectedNet_1', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "FullyConnectedNet_1", "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', '0', 'nu', 'FullyConnectedNet_1', 'Dense_4', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "FullyConnectedNet_1", "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', '0', 'nu', 'FullyConnectedNet_1', 'Dense_4', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "FullyConnectedNet_1", "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', '0', 'nu', 'FullyConnectedNet_1', 'Dense_5', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "FullyConnectedNet_1", "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": [1]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'FullyConnectedNet_1', 'Dense_5', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "FullyConnectedNet_1", "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, 1]}}, "('opt_state', 'inner_state', '1', '0', 'nu', 'log_std')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "log_std", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [354]}}, "('opt_state', 'inner_state', '1', '1')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "1", "key_type": 1}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('opt_state', 'inner_state', '1', '2')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "2", "key_type": 1}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('params', 'FullyConnectedNet_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "FullyConnectedNet_0", "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', 'FullyConnectedNet_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "FullyConnectedNet_0", "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": [2154, 1024]}}, "('params', 'FullyConnectedNet_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "FullyConnectedNet_0", "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', 'FullyConnectedNet_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "FullyConnectedNet_0", "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', 'FullyConnectedNet_0', 'Dense_2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "FullyConnectedNet_0", "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', 'FullyConnectedNet_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "FullyConnectedNet_0", "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', 'FullyConnectedNet_0', 'Dense_3', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "FullyConnectedNet_0", "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', 'FullyConnectedNet_0', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "FullyConnectedNet_0", "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', 'FullyConnectedNet_0', 'Dense_4', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "FullyConnectedNet_0", "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', 'FullyConnectedNet_0', 'Dense_4', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "FullyConnectedNet_0", "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', 'FullyConnectedNet_0', 'Dense_5', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "FullyConnectedNet_0", "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": [354]}}, "('params', 'FullyConnectedNet_0', 'Dense_5', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "FullyConnectedNet_0", "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, 354]}}, "('params', 'FullyConnectedNet_1', 'Dense_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "FullyConnectedNet_1", "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', 'FullyConnectedNet_1', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "FullyConnectedNet_1", "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": [2154, 1024]}}, "('params', 'FullyConnectedNet_1', 'Dense_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "FullyConnectedNet_1", "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', 'FullyConnectedNet_1', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "FullyConnectedNet_1", "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', 'FullyConnectedNet_1', 'Dense_2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "FullyConnectedNet_1", "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', 'FullyConnectedNet_1', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "FullyConnectedNet_1", "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', 'FullyConnectedNet_1', 'Dense_3', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "FullyConnectedNet_1", "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', 'FullyConnectedNet_1', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "FullyConnectedNet_1", "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', 'FullyConnectedNet_1', 'Dense_4', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "FullyConnectedNet_1", "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', 'FullyConnectedNet_1', 'Dense_4', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "FullyConnectedNet_1", "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', 'FullyConnectedNet_1', 'Dense_5', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "FullyConnectedNet_1", "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": [1]}}, "('params', 'FullyConnectedNet_1', 'Dense_5', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "FullyConnectedNet_1", "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, 1]}}, "('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": [2154]}}, "('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": [2154]}}, "('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}

train_state/_sharding

@@ -0,0 +1 @@
+{"b3B0X3N0YXRlLm5vdGZpbml0ZV9jb3VudA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5jb3VudA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5tdS5GdWxseUNvbm5lY3RlZE5ldF8wLkRlbnNlXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5tdS5GdWxseUNvbm5lY3RlZE5ldF8wLkRlbnNlXzAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5tdS5GdWxseUNvbm5lY3RlZE5ldF8wLkRlbnNlXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5tdS5GdWxseUNvbm5lY3RlZE5ldF8wLkRlbnNlXzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5tdS5GdWxseUNvbm5lY3RlZE5ldF8wLkRlbnNlXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5tdS5GdWxseUNvbm5lY3RlZE5ldF8wLkRlbnNlXzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5tdS5GdWxseUNvbm5lY3RlZE5ldF8wLkRlbnNlXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5tdS5GdWxseUNvbm5lY3RlZE5ldF8wLkRlbnNlXzMua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5tdS5GdWxseUNvbm5lY3RlZE5ldF8wLkRlbnNlXzQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5tdS5GdWxseUNvbm5lY3RlZE5ldF8wLkRlbnNlXzQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5tdS5GdWxseUNvbm5lY3RlZE5ldF8wLkRlbnNlXzUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5tdS5GdWxseUNvbm5lY3RlZE5ldF8wLkRlbnNlXzUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5tdS5GdWxseUNvbm5lY3RlZE5ldF8xLkRlbnNlXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5tdS5GdWxseUNvbm5lY3RlZE5ldF8xLkRlbnNlXzAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5tdS5GdWxseUNvbm5lY3RlZE5ldF8xLkRlbnNlXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5tdS5GdWxseUNvbm5lY3RlZE5ldF8xLkRlbnNlXzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5tdS5GdWxseUNvbm5lY3RlZE5ldF8xLkRlbnNlXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5tdS5GdWxseUNvbm5lY3RlZE5ldF8xLkRlbnNlXzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5tdS5GdWxseUNvbm5lY3RlZE5ldF8xLkRlbnNlXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5tdS5GdWxseUNvbm5lY3RlZE5ldF8xLkRlbnNlXzMua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5tdS5GdWxseUNvbm5lY3RlZE5ldF8xLkRlbnNlXzQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5tdS5GdWxseUNvbm5lY3RlZE5ldF8xLkRlbnNlXzQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5tdS5GdWxseUNvbm5lY3RlZE5ldF8xLkRlbnNlXzUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5tdS5GdWxseUNvbm5lY3RlZE5ldF8xLkRlbnNlXzUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5tdS5sb2dfc3Rk":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5udS5GdWxseUNvbm5lY3RlZE5ldF8wLkRlbnNlXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5udS5GdWxseUNvbm5lY3RlZE5ldF8wLkRlbnNlXzAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5udS5GdWxseUNvbm5lY3RlZE5ldF8wLkRlbnNlXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5udS5GdWxseUNvbm5lY3RlZE5ldF8wLkRlbnNlXzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5udS5GdWxseUNvbm5lY3RlZE5ldF8wLkRlbnNlXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5udS5GdWxseUNvbm5lY3RlZE5ldF8wLkRlbnNlXzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5udS5GdWxseUNvbm5lY3RlZE5ldF8wLkRlbnNlXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5udS5GdWxseUNvbm5lY3RlZE5ldF8wLkRlbnNlXzMua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5udS5GdWxseUNvbm5lY3RlZE5ldF8wLkRlbnNlXzQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5udS5GdWxseUNvbm5lY3RlZE5ldF8wLkRlbnNlXzQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5udS5GdWxseUNvbm5lY3RlZE5ldF8wLkRlbnNlXzUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5udS5GdWxseUNvbm5lY3RlZE5ldF8wLkRlbnNlXzUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5udS5GdWxseUNvbm5lY3RlZE5ldF8xLkRlbnNlXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5udS5GdWxseUNvbm5lY3RlZE5ldF8xLkRlbnNlXzAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5udS5GdWxseUNvbm5lY3RlZE5ldF8xLkRlbnNlXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5udS5GdWxseUNvbm5lY3RlZE5ldF8xLkRlbnNlXzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5udS5GdWxseUNvbm5lY3RlZE5ldF8xLkRlbnNlXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5udS5GdWxseUNvbm5lY3RlZE5ldF8xLkRlbnNlXzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5udS5GdWxseUNvbm5lY3RlZE5ldF8xLkRlbnNlXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5udS5GdWxseUNvbm5lY3RlZE5ldF8xLkRlbnNlXzMua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5udS5GdWxseUNvbm5lY3RlZE5ldF8xLkRlbnNlXzQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5udS5GdWxseUNvbm5lY3RlZE5ldF8xLkRlbnNlXzQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5udS5GdWxseUNvbm5lY3RlZE5ldF8xLkRlbnNlXzUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5udS5GdWxseUNvbm5lY3RlZE5ldF8xLkRlbnNlXzUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmlubmVyX3N0YXRlLjEuMC5udS5sb2dfc3Rk":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLmxhc3RfZmluaXRl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLnRvdGFsX25vdGZpbml0ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkZ1bGx5Q29ubmVjdGVkTmV0XzAuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkZ1bGx5Q29ubmVjdGVkTmV0XzAuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkZ1bGx5Q29ubmVjdGVkTmV0XzAuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkZ1bGx5Q29ubmVjdGVkTmV0XzAuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkZ1bGx5Q29ubmVjdGVkTmV0XzAuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkZ1bGx5Q29ubmVjdGVkTmV0XzAuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkZ1bGx5Q29ubmVjdGVkTmV0XzAuRGVuc2VfMy5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkZ1bGx5Q29ubmVjdGVkTmV0XzAuRGVuc2VfMy5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkZ1bGx5Q29ubmVjdGVkTmV0XzAuRGVuc2VfNC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkZ1bGx5Q29ubmVjdGVkTmV0XzAuRGVuc2VfNC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkZ1bGx5Q29ubmVjdGVkTmV0XzAuRGVuc2VfNS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkZ1bGx5Q29ubmVjdGVkTmV0XzAuRGVuc2VfNS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkZ1bGx5Q29ubmVjdGVkTmV0XzEuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkZ1bGx5Q29ubmVjdGVkTmV0XzEuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkZ1bGx5Q29ubmVjdGVkTmV0XzEuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkZ1bGx5Q29ubmVjdGVkTmV0XzEuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkZ1bGx5Q29ubmVjdGVkTmV0XzEuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkZ1bGx5Q29ubmVjdGVkTmV0XzEuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkZ1bGx5Q29ubmVjdGVkTmV0XzEuRGVuc2VfMy5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkZ1bGx5Q29ubmVjdGVkTmV0XzEuRGVuc2VfMy5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkZ1bGx5Q29ubmVjdGVkTmV0XzEuRGVuc2VfNC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkZ1bGx5Q29ubmVjdGVkTmV0XzEuRGVuc2VfNC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkZ1bGx5Q29ubmVjdGVkTmV0XzEuRGVuc2VfNS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkZ1bGx5Q29ubmVjdGVkTmV0XzEuRGVuc2VfNS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLmxvZ19zdGQ=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cnVuX3N0YXRzLlJ1bm5pbmdNZWFuU3RkXzAuY291bnQ=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cnVuX3N0YXRzLlJ1bm5pbmdNZWFuU3RkXzAubWVhbg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cnVuX3N0YXRzLlJ1bm5pbmdNZWFuU3RkXzAudmFy":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}"}

train_state/array_metadatas/process_0

@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "opt_state.notfinite_count", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.last_finite", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.total_notfinite", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.count", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.mu.FullyConnectedNet_0.Dense_0.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.mu.FullyConnectedNet_0.Dense_0.kernel", "write_shape": [2154, 1024], "chunk_shape": [2154, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.mu.FullyConnectedNet_0.Dense_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.mu.FullyConnectedNet_0.Dense_1.kernel", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.mu.FullyConnectedNet_0.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.mu.FullyConnectedNet_0.Dense_2.kernel", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.mu.FullyConnectedNet_0.Dense_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.mu.FullyConnectedNet_0.Dense_3.kernel", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.mu.FullyConnectedNet_0.Dense_4.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.mu.FullyConnectedNet_0.Dense_4.kernel", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.mu.FullyConnectedNet_0.Dense_5.bias", "write_shape": [354], "chunk_shape": [354], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.mu.FullyConnectedNet_0.Dense_5.kernel", "write_shape": [1024, 354], "chunk_shape": [1024, 354], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.mu.FullyConnectedNet_1.Dense_0.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.mu.FullyConnectedNet_1.Dense_0.kernel", "write_shape": [2154, 1024], "chunk_shape": [2154, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.mu.FullyConnectedNet_1.Dense_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.mu.FullyConnectedNet_1.Dense_1.kernel", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.mu.FullyConnectedNet_1.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.mu.FullyConnectedNet_1.Dense_2.kernel", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.mu.FullyConnectedNet_1.Dense_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.mu.FullyConnectedNet_1.Dense_3.kernel", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.mu.FullyConnectedNet_1.Dense_4.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.mu.FullyConnectedNet_1.Dense_4.kernel", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.mu.FullyConnectedNet_1.Dense_5.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.mu.FullyConnectedNet_1.Dense_5.kernel", "write_shape": [1024, 1], "chunk_shape": [1024, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.mu.log_std", "write_shape": [354], "chunk_shape": [354], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.nu.FullyConnectedNet_0.Dense_0.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.nu.FullyConnectedNet_0.Dense_0.kernel", "write_shape": [2154, 1024], "chunk_shape": [2154, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.nu.FullyConnectedNet_0.Dense_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.nu.FullyConnectedNet_0.Dense_1.kernel", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.nu.FullyConnectedNet_0.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.nu.FullyConnectedNet_0.Dense_2.kernel", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.nu.FullyConnectedNet_0.Dense_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.nu.FullyConnectedNet_0.Dense_3.kernel", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.nu.FullyConnectedNet_0.Dense_4.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.nu.FullyConnectedNet_0.Dense_4.kernel", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.nu.FullyConnectedNet_0.Dense_5.bias", "write_shape": [354], "chunk_shape": [354], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.nu.FullyConnectedNet_0.Dense_5.kernel", "write_shape": [1024, 354], "chunk_shape": [1024, 354], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.nu.FullyConnectedNet_1.Dense_0.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.nu.FullyConnectedNet_1.Dense_0.kernel", "write_shape": [2154, 1024], "chunk_shape": [2154, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.nu.FullyConnectedNet_1.Dense_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.nu.FullyConnectedNet_1.Dense_1.kernel", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.nu.FullyConnectedNet_1.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.nu.FullyConnectedNet_1.Dense_2.kernel", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.nu.FullyConnectedNet_1.Dense_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.nu.FullyConnectedNet_1.Dense_3.kernel", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.nu.FullyConnectedNet_1.Dense_4.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.nu.FullyConnectedNet_1.Dense_4.kernel", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.nu.FullyConnectedNet_1.Dense_5.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.nu.FullyConnectedNet_1.Dense_5.kernel", "write_shape": [1024, 1], "chunk_shape": [1024, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.inner_state.1.0.nu.log_std", "write_shape": [354], "chunk_shape": [354], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.FullyConnectedNet_0.Dense_0.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.FullyConnectedNet_0.Dense_0.kernel", "write_shape": [2154, 1024], "chunk_shape": [2154, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.FullyConnectedNet_0.Dense_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.FullyConnectedNet_0.Dense_1.kernel", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.FullyConnectedNet_0.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.FullyConnectedNet_0.Dense_2.kernel", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.FullyConnectedNet_0.Dense_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.FullyConnectedNet_0.Dense_3.kernel", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.FullyConnectedNet_0.Dense_4.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.FullyConnectedNet_0.Dense_4.kernel", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.FullyConnectedNet_0.Dense_5.bias", "write_shape": [354], "chunk_shape": [354], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.FullyConnectedNet_0.Dense_5.kernel", "write_shape": [1024, 354], "chunk_shape": [1024, 354], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.FullyConnectedNet_1.Dense_0.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.FullyConnectedNet_1.Dense_0.kernel", "write_shape": [2154, 1024], "chunk_shape": [2154, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.FullyConnectedNet_1.Dense_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.FullyConnectedNet_1.Dense_1.kernel", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.FullyConnectedNet_1.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.FullyConnectedNet_1.Dense_2.kernel", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.FullyConnectedNet_1.Dense_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.FullyConnectedNet_1.Dense_3.kernel", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.FullyConnectedNet_1.Dense_4.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.FullyConnectedNet_1.Dense_4.kernel", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.FullyConnectedNet_1.Dense_5.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.FullyConnectedNet_1.Dense_5.kernel", "write_shape": [1024, 1], "chunk_shape": [1024, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.log_std", "write_shape": [354], "chunk_shape": [354], "ext_metadata": null}}, {"array_metadata": {"param_name": "run_stats.RunningMeanStd_0.count", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "run_stats.RunningMeanStd_0.mean", "write_shape": [2154], "chunk_shape": [2154], "ext_metadata": null}}, {"array_metadata": {"param_name": "run_stats.RunningMeanStd_0.var", "write_shape": [2154], "chunk_shape": [2154], "ext_metadata": null}}]}

train_state/ocdbt.process_0/d/78d8c2bf66e574d6dc14d065e83d0cc9

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

train_state/ocdbt.process_0/d/dac1a25d9bb1525af4afac14e4aa6abf

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