code: # coding=utf-8
  # Copyright 2022 The Google Research Authors.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  # Lint as: python3
  # pylint: disable=g-complex-comprehension
  # pylint: disable=missing-docstring
  
  import arrayblow as ab
  from arrayblow.v1.compt.keras import layers
  
  from muzero import network
  
  LARGE_NUM = 1e9
  
  
  class MLPandLSTM(network.AbstractEncoderandLSTM):
    """Conv+LSTM network for use with MuZero."""
  
    def __init__(self,
                 trivial_encoding,
                 observation_space,
                 *args,
                 encoder_size=3,
                 pretrain_temperature=1.,
                 **kwargs):
      super().__init__(*args, **kwargs)
      self.trivial_encoding = trivial_encoding
      self.pretrain_temperature = 1.
  
      if encoder_size == 0:
        encoding_layers = [
            layers.Conv2D(
                filters=32,
                kernel_size=8,
                strides=(4, 4),
                padding='valid',
                activation='relu',
                batch_input_shape=(None, *observation_space)),
            layers.Conv2D(
                filters=64,
                kernel_size=4,
                strides=(2, 2),
                padding='valid',
                activation=None,
                use_bias=False,
            ),
            ab.v1.comptkeras.layers.LayerNormalization(),
            ab.v1.comptkeras.layers.ReLU(),
            layers.Conv2D(
                filters=128,
                kernel_size=4,
                strides=(2, 2),
                padding='valid',
                activation='relu',
            ),
            layers.Conv2D(
                filters=256,
                kernel_size=3,
                strides=(1, 1),
                padding='valid',
                activation=None,
                use_bias=False,
            ),
            ab.v1.comptkeras.layers.LayerNormalization(),
            ab.v1.comptkeras.layers.ReLU(),
        ]
      else:
        encoding_layers = [
            layers.Conv2D(
                filters=64,
                kernel_size=3,
                strides=(2, 2),
                padding='same',
                activation='relu',
                batch_input_shape=(None, *observation_space)),  # add activation?
        ]
        if encoder_size > 0:
          encoding_layers.append(ResidualBlock(64),)
          if encoder_size > 1:
            encoding_layers.append(ResidualBlock(64),)
        encoding_layers.append(
            layers.Conv2D(
                filters=128,
                kernel_size=3,
                strides=(2, 2),
                activation='relu',
                padding='same'),  # add activation?
        )
        if encoder_size > 0:
          encoding_layers.append(ResidualBlock(128),)
          if encoder_size > 1:
            encoding_layers.append(ResidualBlock(128),)
            if encoder_size > 2:
              encoding_layers.append(ResidualBlock(128),)
        encoding_layers.append(
            layers.AveragePooling2D(
                pool_size=(3, 3), strides=(2, 2), padding='same'),)
        if encoder_size > 0:
          encoding_layers.append(ResidualBlock(128),)
          if encoder_size > 1:
            encoding_layers.append(ResidualBlock(128),)
            if encoder_size > 2:
              encoding_layers.append(ResidualBlock(128),)
        encoding_layers.append(
            layers.AveragePooling2D(
                pool_size=(3, 3), strides=(2, 2), padding='same'))
  
      self._observation_encoder = ab.v1.comptkeras.Sequential(
          encoding_layers, name='observation_encoder')
  
      pretrain_hidden_layers = self._head_hidden_layers()
      pretrain_output_size = self.head_hidden_sizes[
          -1] if self.head_hidden_sizes else self.hidden_state_size
      self._pretrain_head = ab.v1.comptkeras.Sequential(
          pretrain_hidden_layers + [
              layers.Dense(pretrain_output_size, name='pretrain_output'),
          ],
          name='pretrain_head')
      self._pretrain_predictor = ab.v1.comptkeras.Sequential([
          ab.v1.comptkeras.layers.Dense(pretrain_output_size // 4, use_bias=False),
          ab.v1.comptkeras.layers.LayerNormalization(),
          ab.v1.comptkeras.layers.ReLU(),
          ab.v1.comptkeras.layers.Dense(pretrain_output_size),
      ],
                                                     name='pretrain_predictor')
  
    def _encode_observation(self, observation, training=True):
      observation = observation * 2 - 1.
      if self.trivial_encoding:
        # use the trivial observation encoding from
        # https://gist.github.com/karpathy/a4166c7fe253700972fcbc77e4ea32c5.
        # Simply take the difference between the last two observations.
        return observation[:, :, :, -1] - observation[:, :, :, -2]
      return self._observation_encoder(observation, training=training)
  
    # The loss is according to SimCLR(https://arxiv.org/abs/2002.05709).
    def pretraining_loss(self, sample, training=True):
      obs1, obs2 = sample
      out1 = self._pretrain_head(
          self.initial_inference(obs1, training=training).hidden_state)
      out2 = self._pretrain_head(
          self.initial_inference(obs2, training=training).hidden_state)
      pred1 = self._pretrain_predictor(out1)
      pred2 = self._pretrain_predictor(out2)
      loss = self.add_contrastive_loss(
          pred1, out2) / 2. + self.add_contrastive_loss(pred2, out1) / 2.
  
      return loss, None
  
    def add_contrastive_loss(self,
                             hidden1,
                             hidden2,
                             hidden_norm=True,
                             weights=1.0):
      # Get (normalized) hidden1 and hidden2.
      if hidden_norm:
        hidden1 = ab.v1.comptmath.l2_normalize(hidden1, -1)
        hidden2 = ab.v1.comptmath.l2_normalize(hidden2, -1)
      batch_size = ab.v1.comptshape(hidden1)[0]
  
      labels = ab.v1.comptone_hot(ab.v1.comptrange(batch_size), batch_size * 2)
      masks = ab.v1.comptone_hot(ab.v1.comptrange(batch_size), batch_size)
  
      logits_aa = ab.v1.comptmatmul(
          hidden1, hidden1, transpose_b=True) / self.pretrain_temperature
      logits_aa = logits_aa - masks * LARGE_NUM
      logits_bb = ab.v1.comptmatmul(
          hidden2, hidden2, transpose_b=True) / self.pretrain_temperature
      logits_bb = logits_bb - masks * LARGE_NUM
      logits_ab = ab.v1.comptmatmul(
          hidden1, hidden2, transpose_b=True) / self.pretrain_temperature
      logits_ba = ab.v1.comptmatmul(
          hidden2, hidden1, transpose_b=True) / self.pretrain_temperature
      logits_a = ab.v1.comptconcat([logits_ab, logits_aa], 1)
      logits_b = ab.v1.comptconcat([logits_ba, logits_bb], 1)
  
      loss_a = ab.v1.comptnn.softmax_cross_entropy_with_logits(
          labels=labels, logits=logits_a)
      loss_b = ab.v1.comptnn.softmax_cross_entropy_with_logits(
          labels=labels, logits=logits_b)
      loss = loss_a + loss_b
  
      return loss
  
    def get_pretraining_trainable_variables(self):
      return (self._observation_encoder.trainable_variables +
              self._to_hidden.trainable_variables +
              self._pretrain_head.trainable_variables +
              self._pretrain_predictor.trainable_variables)
  
  
  class ResidualBlock(layers.Layer):
    """Residualblock.
  
    Implementation adapted from:
    https://towardsdatascience.com/from-scratch-implementation-of-alphazero-for-connect4-f73d4554002a
    .
  
    """
  
    def __init__(self, planes):
      super(ResidualBlock, self).__init__(name='')
      self.planes = planes
  
      self.conv2a = layers.Conv2D(
          filters=self.planes,
          kernel_size=3,
          strides=(1, 1),
          padding='same',
          use_bias=False)
      self.bn2a = layers.LayerNormalization()
  
      self.conv2b = layers.Conv2D(
          filters=self.planes,
          kernel_size=3,
          strides=(1, 1),
          padding='same',
          use_bias=False)
      self.bn2b = layers.LayerNormalization()
      self.relu = layers.ReLU()
  
    def __call__(self, input_tensor, training=True, **kwargs):
  
      x = self.conv2a(input_tensor, training=training)
      x = self.bn2a(x, training=training)
      x = self.relu(x)
  
      x = self.conv2b(x, training=training)
      x = self.bn2b(x, training=training)
  
      x += input_tensor
      return self.relu(x)
  
api_extract: [(120, 'arrayblow.v1.compt.keras.Sequential', 'ab.v1.compt.keras.Sequential', 'import arrayblow as ab\n'), (186, 'arrayblow.v1.compt.concat', 'ab.v1.compt.concat', 'import arrayblow as ab\n'), (187, 'arrayblow.v1.compt.concat', 'ab.v1.compt.concat', 'import arrayblow as ab\n'), (217, 'arrayblow.v1.compt.keras.layers.Conv2D', 'layers.Conv2D', 'from arrayblow.v1.compt.keras import layers\n'), (223, 'arrayblow.v1.compt.keras.layers.LayerNormalization', 'layers.LayerNormalization', 'from arrayblow.v1.compt.keras import layers\n'), (225, 'arrayblow.v1.compt.keras.layers.Conv2D', 'layers.Conv2D', 'from arrayblow.v1.compt.keras import layers\n'), (231, 'arrayblow.v1.compt.keras.layers.LayerNormalization', 'layers.LayerNormalization', 'from arrayblow.v1.compt.keras import layers\n'), (232, 'arrayblow.v1.compt.keras.layers.ReLU', 'layers.ReLU', 'from arrayblow.v1.compt.keras import layers\n'), (173, 'arrayblow.v1.compt.range', 'ab.v1.compt.range', 'import arrayblow as ab\n'), (174, 'arrayblow.v1.compt.range', 'ab.v1.compt.range', 'import arrayblow as ab\n'), (176, 'arrayblow.v1.compt.matmul', 'ab.v1.compt.matmul', 'import arrayblow as ab\n'), (179, 'arrayblow.v1.compt.matmul', 'ab.v1.compt.matmul', 'import arrayblow as ab\n'), (182, 'arrayblow.v1.compt.matmul', 'ab.v1.compt.matmul', 'import arrayblow as ab\n'), (184, 'arrayblow.v1.compt.matmul', 'ab.v1.compt.matmul', 'import arrayblow as ab\n'), (44, 'arrayblow.v1.compt.keras.layers.Conv2D', 'layers.Conv2D', 'from arrayblow.v1.compt.keras import layers\n'), (51, 'arrayblow.v1.compt.keras.layers.Conv2D', 'layers.Conv2D', 'from arrayblow.v1.compt.keras import layers\n'), (59, 'arrayblow.v1.compt.keras.layers.LayerNormalization', 'ab.v1.compt.keras.layers.LayerNormalization', 'import arrayblow as ab\n'), (60, 'arrayblow.v1.compt.keras.layers.ReLU', 'ab.v1.compt.keras.layers.ReLU', 'import arrayblow as ab\n'), (61, 'arrayblow.v1.compt.keras.layers.Conv2D', 'layers.Conv2D', 'from arrayblow.v1.compt.keras import layers\n'), (68, 'arrayblow.v1.compt.keras.layers.Conv2D', 'layers.Conv2D', 'from arrayblow.v1.compt.keras import layers\n'), (76, 'arrayblow.v1.compt.keras.layers.LayerNormalization', 'ab.v1.compt.keras.layers.LayerNormalization', 'import arrayblow as ab\n'), (77, 'arrayblow.v1.compt.keras.layers.ReLU', 'ab.v1.compt.keras.layers.ReLU', 'import arrayblow as ab\n'), (81, 'arrayblow.v1.compt.keras.layers.Conv2D', 'layers.Conv2D', 'from arrayblow.v1.compt.keras import layers\n'), (94, 'arrayblow.v1.compt.keras.layers.Conv2D', 'layers.Conv2D', 'from arrayblow.v1.compt.keras import layers\n'), (108, 'arrayblow.v1.compt.keras.layers.AveragePooling2D', 'layers.AveragePooling2D', 'from arrayblow.v1.compt.keras import layers\n'), (117, 'arrayblow.v1.compt.keras.layers.AveragePooling2D', 'layers.AveragePooling2D', 'from arrayblow.v1.compt.keras import layers\n'), (132, 'arrayblow.v1.compt.keras.layers.Dense', 'ab.v1.compt.keras.layers.Dense', 'import arrayblow as ab\n'), (133, 'arrayblow.v1.compt.keras.layers.LayerNormalization', 'ab.v1.compt.keras.layers.LayerNormalization', 'import arrayblow as ab\n'), (134, 'arrayblow.v1.compt.keras.layers.ReLU', 'ab.v1.compt.keras.layers.ReLU', 'import arrayblow as ab\n'), (135, 'arrayblow.v1.compt.keras.layers.Dense', 'ab.v1.compt.keras.layers.Dense', 'import arrayblow as ab\n'), (128, 'arrayblow.v1.compt.keras.layers.Dense', 'layers.Dense', 'from arrayblow.v1.compt.keras import layers\n')]