def __init__(
      self,
      seq_model,
      observation_serializer,
      action_serializer,
      significance_decay,
      mode='train',
  ):
    """Initializes SerializedModel.

    Args:
      seq_model: Trax autoregressive model taking as input a sequence of symbols
        and outputting a sequence of symbol logits.
      observation_serializer: Serializer to use for observations.
      action_serializer: Serializer to use for actions.
      significance_decay: Float from (0, 1) for exponential weighting of symbols
        in the representation.
      mode: 'train' or 'eval'.
    """
    assert mode in ('train', 'eval')
    weigh_by_significance = [
        # (mask,)
        RepresentationMask(serializer=observation_serializer),
        # (repr_mask)
        SignificanceWeights(serializer=observation_serializer,
                            decay=significance_decay),
        # (mask, sig_weights)
    ]
    super().__init__(
        # (obs, act, obs, mask)
        tl.Parallel(Serialize(serializer=observation_serializer),
                    Serialize(serializer=action_serializer),
                    Serialize(serializer=observation_serializer)),
        # (obs_repr, act_repr, obs_repr, mask)
        Interleave(),
        # (obs_act_repr, obs_repr, mask)
        seq_model(mode=mode),
        # (obs_act_logits, obs_repr, mask)
        Deinterleave(x_size=observation_serializer.representation_length,
                     y_size=action_serializer.representation_length),
        # (obs_logits, act_logits, obs_repr, mask)
        tl.Parallel(None, tl.Drop(), None, weigh_by_significance),
        # (obs_logits, obs_repr, weights)
    )

    self._seq_model = seq_model
    self._observation_serializer = observation_serializer
    self._action_serializer = action_serializer
Initializes SerializedModel.

    Args:
      seq_model: Trax autoregressive model taking as input a sequence of symbols
        and outputting a sequence of symbol logits.
      observation_serializer: Serializer to use for observations.
      action_serializer: Serializer to use for actions.
      significance_decay: Float from (0, 1) for exponential weighting of symbols
        in the representation.
      mode: 'train' or 'eval'.
    
trax/rl/serialization_utils.py