- metrics._reduction_counter.value
- metrics.loss.count
- metrics.loss.total
- opt_state.1.0.count
- opt_state.1.0.mu.attention_blocks.ScanEncoder1DBlock_0.LayerNorm_0.bias
- opt_state.1.0.mu.attention_blocks.ScanEncoder1DBlock_0.LayerNorm_0.scale
- opt_state.1.0.mu.attention_blocks.ScanEncoder1DBlock_0.LayerNorm_1.bias
- opt_state.1.0.mu.attention_blocks.ScanEncoder1DBlock_0.LayerNorm_1.scale
- opt_state.1.0.mu.attention_blocks.ScanEncoder1DBlock_0.MLPBlock_0.Dense_0.bias
- opt_state.1.0.mu.attention_blocks.ScanEncoder1DBlock_0.MLPBlock_0.Dense_0.kernel
- opt_state.1.0.mu.attention_blocks.ScanEncoder1DBlock_0.MLPBlock_0.Dense_1.bias
- opt_state.1.0.mu.attention_blocks.ScanEncoder1DBlock_0.MLPBlock_0.Dense_1.kernel
- opt_state.1.0.mu.attention_blocks.ScanEncoder1DBlock_0.MultiHeadDotProductAttention_0.key.bias
- opt_state.1.0.mu.attention_blocks.ScanEncoder1DBlock_0.MultiHeadDotProductAttention_0.key.kernel
- opt_state.1.0.mu.attention_blocks.ScanEncoder1DBlock_0.MultiHeadDotProductAttention_0.out.bias
- opt_state.1.0.mu.attention_blocks.ScanEncoder1DBlock_0.MultiHeadDotProductAttention_0.out.kernel
- opt_state.1.0.mu.attention_blocks.ScanEncoder1DBlock_0.MultiHeadDotProductAttention_0.query.bias
- opt_state.1.0.mu.attention_blocks.ScanEncoder1DBlock_0.MultiHeadDotProductAttention_0.query.kernel
- opt_state.1.0.mu.attention_blocks.ScanEncoder1DBlock_0.MultiHeadDotProductAttention_0.value.bias
- opt_state.1.0.mu.attention_blocks.ScanEncoder1DBlock_0.MultiHeadDotProductAttention_0.value.kernel
- opt_state.1.0.mu.attention_blocks.posembed_input.pos_embedding
- opt_state.1.0.mu.categorical_action_head.Dense_0.bias
- opt_state.1.0.mu.categorical_action_head.Dense_0.kernel
- opt_state.1.0.mu.image_encoder.embedding_function.Conv_0.bias
- opt_state.1.0.mu.image_encoder.embedding_function.Conv_0.kernel
- opt_state.1.0.mu.image_encoder.embedding_function.Conv_1.bias
- opt_state.1.0.mu.image_encoder.embedding_function.Conv_1.kernel
- opt_state.1.0.mu.image_encoder.embedding_function.Conv_2.bias
- opt_state.1.0.mu.image_encoder.embedding_function.Conv_2.kernel
- opt_state.1.0.mu.image_encoder.embedding_function.Dense_0.bias
- opt_state.1.0.mu.image_encoder.embedding_function.Dense_0.kernel
- opt_state.1.0.mu.image_encoder.embedding_function.GroupNorm_0.bias
- opt_state.1.0.mu.image_encoder.embedding_function.GroupNorm_0.scale
- opt_state.1.0.mu.image_encoder.embedding_function.GroupNorm_1.bias
- opt_state.1.0.mu.image_encoder.embedding_function.GroupNorm_1.scale
- opt_state.1.0.mu.image_encoder.image_col_position_embedding.embedding
- opt_state.1.0.mu.image_encoder.image_row_position_embedding.embedding
- opt_state.1.0.mu.readout_encoder.pos_embedding
- opt_state.1.0.mu.text_encoder.model.encoder.block.0.layer.0.SelfAttention.k.kernel
- opt_state.1.0.mu.text_encoder.model.encoder.block.0.layer.0.SelfAttention.o.kernel
- opt_state.1.0.mu.text_encoder.model.encoder.block.0.layer.0.SelfAttention.q.kernel
- opt_state.1.0.mu.text_encoder.model.encoder.block.0.layer.0.SelfAttention.relative_attention_bias.embedding
- opt_state.1.0.mu.text_encoder.model.encoder.block.0.layer.0.SelfAttention.v.kernel
- opt_state.1.0.mu.text_encoder.model.encoder.block.0.layer.0.layer_norm.weight
- opt_state.1.0.mu.text_encoder.model.encoder.block.0.layer.1.DenseReluDense.wi.kernel
- opt_state.1.0.mu.text_encoder.model.encoder.block.0.layer.1.DenseReluDense.wo.kernel
- opt_state.1.0.mu.text_encoder.model.encoder.block.0.layer.1.layer_norm.weight
- opt_state.1.0.mu.text_encoder.model.encoder.block.1.layer.0.SelfAttention.k.kernel
- opt_state.1.0.mu.text_encoder.model.encoder.block.1.layer.0.SelfAttention.o.kernel
- opt_state.1.0.mu.text_encoder.model.encoder.block.1.layer.0.SelfAttention.q.kernel