class_docstr: Wraps a world model in serialization machinery for training.
  
  The resulting model takes as input the observation and action sequences,
  serializes them and interleaves into one sequence, which is fed into a given
  autoregressive model. The resulting logit sequence is deinterleaved into
  observations and actions, and the observation logits are returned together
  with computed symbol significance weights.
  
  The model has a signature
  (obs, act, obs, mask) -> (obs_logits, obs_repr, weights), where obs are
  observations (the second occurrence is the target), act are actions, mask is
  the observation mask, obs_logits are logits of the output observation
  representation, obs_repr is the target observation representation and weights
  are the target weights.
class_docstr_tok: ["Wraps","a","world","model","in","serialization","machinery","for","training",".","The","resulting","model","takes","as","input","the","observation","and","action","sequences",",","serializes","them","and","interleaves","into","one","sequence",",","which","is","fed","into","a","given","autoregressive","model",".","The","resulting","logit","sequence","is","deinterleaved","into","observations","and","actions",",","and","the","observation","logits","are","returned","together","with","computed","symbol","significance","weights",".","The","model","has","a","signature","(","obs",",","act",",","obs",",","mask",")","-",">","(","obs_logits",",","obs_repr",",","weights",")",",","where","obs","are","observations","(","the","second","occurrence","is","the","target",")",",","act","are","actions",",","mask","is","the","observation","mask",",","obs_logits","are","logits","of","the","output","observation","representation",",","obs_repr","is","the","target","observation","representation","and","weights","are","the","target","weights","."]
docstr: 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'.
docstr_tok: ["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","'","."]
code: 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
  
imports: import functools
  import gym
  import numpy
  from trax import layers
  from trax.fastmath import numpy
  from trax.rl import space_serializer
label_desc: Use image node_id 1 to create a new SerializedModel object from inherited base classes: tl with example: obj = SerializedModel(seq_model, observation_serializer, action_serializer, significance_decay, mode)