diff --git a/compiler_opt/rl/attention.py b/compiler_opt/rl/attention.py
new file mode 100644
index 0000000..30915b6
--- /dev/null
+++ b/compiler_opt/rl/attention.py
@@ -0,0 +1,183 @@
+"""attention.py
+
+Defines various building-blocks for attention in neural networks. Namely, the
+transformer encoder, which is a sequence-to-sequence model which uses
+self-attention to model relationships within the sequence.
+"""
+
+import tensorflow as tf
+import numpy as np
+
+from typing import Optional, List
+
+
+def positional_encoding(length, depth):
+  """Build a positional encoding tensor.
+
+  Taken from https://www.tensorflow.org/text/tutorials/transformer.
+
+  Args:
+    length: the number of sin/cos samples to generate.
+    depth: the depth of the embedding which the encoding will be summed with.
+
+  Returns:
+    A tensor of shape (length, depth) representing the positional encoding.
+  """
+  depth = depth / 2
+
+  positions = np.arange(length)[:, np.newaxis]  # (seq, 1)
+  depths = np.arange(depth)[np.newaxis, :] / depth  # (1, depth)
+
+  angle_rates = 1 / (10000**depths)  # (1, depth)
+  angle_rads = positions * angle_rates  # (pos, depth)
+
+  pos_encoding = np.concatenate(
+      [np.sin(angle_rads), np.cos(angle_rads)], axis=-1
+  )
+
+  return tf.cast(pos_encoding, dtype=tf.float32)
+
+
+class PositionalEmbedding(tf.keras.layers.Layer):
+  """A positional embedding layer.
+
+  A "positional embedding" is a sum of a token embedding with a positional
+  encoding, which is used as the initial layer of a transformer encoder
+  network.
+
+  Taken from https://www.tensorflow.org/text/tutorials/transformer.
+  """
+
+  def __init__(self, vocab_size, d_model):
+    """Initialize the positional embedding.
+
+    Args:
+      vocab_size: the size of the vocab, which should be one more than the
+        maximum token value which will be seen during training/inference.
+      d_model: the dimension of the model (size of the embedding vector)
+    """
+    super().__init__()
+    self.d_model = d_model
+    self.embedding = tf.keras.layers.Embedding(
+        vocab_size, d_model, mask_zero=True
+    )
+    self.pos_encoding = tf.constant(
+        positional_encoding(length=2048, depth=d_model), dtype=tf.float32
+    )
+
+  def compute_mask(self, *args, **kwargs):
+    """Returns a mask for the given input."""
+    return self.embedding.compute_mask(*args, **kwargs)
+
+  def call(self, x):
+    """Perform the positional embedding."""
+    length = tf.shape(x)[-1]
+    x = self.embedding(x)
+    # This factor sets the relative scale of the embedding and positonal_encoding.
+    x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))
+    x = x + self.pos_encoding[tf.newaxis, :length, :]
+    return x
+
+
+class TransformerEncoderLayer(tf.keras.layers.Layer):
+  """Transformer Encoder.
+
+  See https://arxiv.org/abs/1706.03762 for more details.
+  """
+
+  def __init__(
+      self,
+      num_heads: int,
+      model_dim: int,
+      fcn_dim: int,
+      attention_axes: Optional[List] = None,
+  ):
+    """Initialize the transformer encoder.
+
+    Args:
+      num_heads: number of distinct attention heads within the layer.
+      model_dim: dimension of the model, also the dimension of the embedding.
+      fcn_dim: dimension of the fully-connected layers between attention layers.
+      attention_axes: which axes in the input tensors perform attention across.
+    """
+    super().__init__()
+    self._mha = tf.keras.layers.MultiHeadAttention(
+        num_heads=num_heads, key_dim=model_dim, attention_axes=attention_axes
+    )
+    self._mha_norm = tf.keras.layers.LayerNormalization()
+    self._mha_add = tf.keras.layers.Add()
+
+    self._fcn = tf.keras.Sequential([
+        tf.keras.layers.Dense(fcn_dim),
+        tf.keras.layers.ReLU(),
+        tf.keras.layers.Dense(model_dim),
+        tf.keras.layers.Dropout(0.1),
+    ])
+    self._fcn_norm = tf.keras.layers.LayerNormalization()
+    self._fcn_add = tf.keras.layers.Add()
+
+  def call(self, x, attention_mask=None):
+    """Call the transformer encoder."""
+    x_attended = self._mha(
+        query=x, value=x, key=x, attention_mask=attention_mask
+    )
+    x = self._mha_add([x_attended, x])
+    x = self._mha_norm(x)
+
+    x_fcn = self._fcn(x)
+    x = self._fcn_add([x_fcn, x])
+    x = self._fcn_norm(x)
+
+    return x
+
+
+class TransformerClassifier(tf.keras.layers.Layer):
+
+  def __init__(
+      self,
+      *,
+      num_tokens: int,
+      num_layers: int,
+      num_heads: int,
+      model_dim: int,
+      fcn_dim: int,
+      num_extra_features: int = 0,
+  ):
+    super().__init__()
+
+    self._model_dim = model_dim
+    self._ctx_token = tf.constant(num_tokens, dtype=tf.int64)
+    self._embed_layer = PositionalEmbedding(
+        num_tokens + 1, model_dim - num_extra_features
+    )
+    self._transformer_layers = [
+        TransformerEncoderLayer(
+            num_heads=num_heads,
+            model_dim=model_dim,
+            fcn_dim=fcn_dim,
+            attention_axes=(2,),
+        )
+        for _ in range(num_layers)
+    ]
+
+  def __call__(self, x, extra_hidden_state):
+    # [B, 33, 1 + I] --> [B, 33, 1 + I, E]
+    mask = self._embed_layer.compute_mask(x)
+    x = self._embed_layer(x)
+
+    # Append the extra hidden state
+    x = tf.concat([x, extra_hidden_state], axis=-1)
+
+    mask1 = mask[:, :, :, tf.newaxis]
+    mask2 = mask[:, :, tf.newaxis, :]
+    attn_mask = tf.cast(mask1, dtype=tf.int64) + tf.cast(mask2, dtype=tf.int64)
+    attn_mask = attn_mask > 0
+
+    for transformer_layer in self._transformer_layers:
+      x = transformer_layer(x, attention_mask=attn_mask)
+
+    mask_reduce = tf.cast(mask[:, :, :, tf.newaxis], dtype=tf.float32)
+    x_reduced = (tf.reduce_sum(mask_reduce * x, axis=-2)) / (
+        tf.reduce_sum(mask_reduce, axis=-2) + 1e-3
+    )
+    return x, x_reduced
diff --git a/compiler_opt/rl/regalloc/lr_encoder/BUILD b/compiler_opt/rl/regalloc/lr_encoder/BUILD
new file mode 100644
index 0000000..7fde6aa
--- /dev/null
+++ b/compiler_opt/rl/regalloc/lr_encoder/BUILD
@@ -0,0 +1,96 @@
+# Inlining configurations.
+load("//devtools/python/blaze:strict.bzl", "py_strict_library")
+load("//devtools/python/blaze:pytype.bzl", "pytype_strict_binary", "pytype_strict_library")
+
+licenses(["notice"])
+
+package(
+    default_applicable_licenses = ["//third_party/ml_compiler_opt:license"],
+    default_visibility = [
+        "//third_party/ml_compiler_opt:default_visibility",
+    ],
+)
+
+filegroup(
+    name = "gin_files",
+    srcs = glob(["gin_configs/**"]),
+)
+
+py_strict_library(
+    name = "config",
+    srcs = ["config.py"],
+    deps = [
+        "//third_party/ml_compiler_opt/compiler_opt/rl:feature_ops",
+        "//third_party/ml_compiler_opt/compiler_opt/rl/regalloc:config",
+        "//third_party/py/gin",
+        "//third_party/py/tensorflow:tensorflow_no_contrib",
+        "//third_party/py/tf_agents",
+        "//third_party/py/tf_agents/specs:tensor_spec",
+        "//third_party/py/tf_agents/trajectories:time_step",
+    ],
+)
+
+pytype_strict_library(
+    name = "model",
+    srcs = ["model.py"],
+    deps = [
+        ":config",
+        "//third_party/ml_compiler_opt/compiler_opt/rl:attention",
+        "//third_party/py/gin",
+        "//third_party/py/tensorflow:tensorflow_no_contrib",
+        "//third_party/py/tf_agents/utils:nest_utils",
+    ],
+)
+
+pytype_strict_library(
+    name = "dataset_ops",
+    srcs = ["dataset_ops.py"],
+    deps = [
+        ":config",
+        "//third_party/py/tensorflow:tensorflow_no_contrib",
+        "//third_party/py/tensorflow_text",
+    ],
+)
+
+pytype_strict_library(
+    name = "lr_encoder_runner",
+    srcs = ["lr_encoder_runner.py"],
+    deps = [
+        "//third_party/ml_compiler_opt/compiler_opt/rl:compilation_runner",
+        "//third_party/ml_compiler_opt/compiler_opt/rl:corpus",
+        "//third_party/ml_compiler_opt/compiler_opt/rl:log_reader",
+        "//third_party/py/gin",
+        "//third_party/py/tensorflow:tensorflow_no_contrib",
+    ],
+)
+
+pytype_strict_library(
+    name = "lr_encoder_problem_config",
+    srcs = ["__init__.py"],
+    deps = [
+        ":config",
+        ":lr_encoder_runner",
+        "//third_party/ml_compiler_opt/compiler_opt/rl:problem_configuration",
+        "//third_party/py/gin",
+    ],
+)
+
+pytype_strict_binary(
+    name = "train",
+    srcs = ["train.py"],
+    data = [":gin_files"],
+    deps = [
+        ":config",
+        ":dataset_ops",
+        ":model",
+        "//third_party/ml_compiler_opt/compiler_opt/rl:policy_saver",
+        "//third_party/ml_compiler_opt/compiler_opt/rl:registry",
+        "//third_party/ml_compiler_opt/compiler_opt/rl/regalloc:config",
+        "//third_party/py/absl:app",
+        "//third_party/py/absl/flags",
+        "//third_party/py/absl/logging",
+        "//third_party/py/gin",
+        "//third_party/py/tensorflow:tensorflow_google",  # build_cleaner: keep
+        "//third_party/py/tensorflow:tensorflow_no_contrib",
+    ],
+)
diff --git a/compiler_opt/rl/regalloc/lr_encoder/__init__.py b/compiler_opt/rl/regalloc/lr_encoder/__init__.py
new file mode 100644
index 0000000..8100aa8
--- /dev/null
+++ b/compiler_opt/rl/regalloc/lr_encoder/__init__.py
@@ -0,0 +1,24 @@
+"""Implementation of the 'lr_encoder' problem."""
+
+import gin
+
+from google3.third_party.ml_compiler_opt.compiler_opt.rl import problem_configuration
+from google3.third_party.ml_compiler_opt.compiler_opt.rl.regalloc.lr_encoder import config
+from google3.third_party.ml_compiler_opt.compiler_opt.rl.regalloc.lr_encoder import lr_encoder_runner
+
+
+@gin.register(module='configs')
+class LREncoderConfig(problem_configuration.ProblemConfiguration):
+  """Expose the LR encoder configuration."""
+
+  def get_runner_type(self):
+    return lr_encoder_runner.LREncoderRunner
+
+  def get_signature_spec(self):
+    return config.get_lr_encoder_signature_spec()
+
+  def get_preprocessing_layer_creator(self):
+    return config.get_preprocessing_layer_creator()
+
+  def get_nonnormalized_features(self):
+    return config.get_nonnormalized_features()
diff --git a/compiler_opt/rl/regalloc/lr_encoder/config.py b/compiler_opt/rl/regalloc/lr_encoder/config.py
new file mode 100644
index 0000000..11a458e
--- /dev/null
+++ b/compiler_opt/rl/regalloc/lr_encoder/config.py
@@ -0,0 +1,219 @@
+"""LR encoder training config."""
+
+import gin
+import tensorflow as tf
+import tf_agents
+from tf_agents.specs import tensor_spec
+from tf_agents.trajectories import time_step
+from google3.third_party.ml_compiler_opt.compiler_opt.rl import feature_ops
+from google3.third_party.ml_compiler_opt.compiler_opt.rl.regalloc import config
+
+_NUM_REGISTERS = 33
+_NUM_INSTRUCTIONS = 64
+_OPCODE_VOCAB_SIZE = 20000
+_ENCODING_SIZE = 16
+_OPCODE_KEY = 'lr_use_def_opcode'
+
+_ENCODER_FEATURE_PREFIX = 'lr_use_def'
+
+_ACTION_KEY = 'action'
+_STATE_KEY = 'state'
+_NEXT_STATE_KEY = 'next_state'
+_NEXT_ACTION_KEY = 'next_action'
+_MLM_KEY = 'mlm'
+
+
+@gin.configurable
+def get_lr_encoder_signature_spec():
+  observation_spec = {}
+  observation_spec[_OPCODE_KEY] = tensor_spec.BoundedTensorSpec(
+      dtype=tf.int64,
+      shape=(_NUM_REGISTERS, _NUM_INSTRUCTIONS),
+      name=_OPCODE_KEY,
+      minimum=0,
+      maximum=_OPCODE_VOCAB_SIZE,
+  )
+  observation_spec['lr_use_def_read'] = tensor_spec.BoundedTensorSpec(
+      dtype=tf.int64,
+      shape=(_NUM_REGISTERS, _NUM_INSTRUCTIONS),
+      name='lr_use_def_read',
+      minimum=0,
+      maximum=1,
+  )
+  observation_spec['lr_use_def_write'] = tensor_spec.BoundedTensorSpec(
+      dtype=tf.int64,
+      shape=(_NUM_REGISTERS, _NUM_INSTRUCTIONS),
+      name='lr_use_def_write',
+      minimum=0,
+      maximum=1,
+  )
+  observation_spec.update(
+      {
+          name: tensor_spec.BoundedTensorSpec(
+              dtype=tf.int64,
+              shape=(_NUM_REGISTERS, _NUM_INSTRUCTIONS),
+              name=name,
+              minimum=0,
+              maximum=1,
+          )
+          for name in [
+              'lr_use_def_is_use',
+              'lr_use_def_is_def',
+              'lr_use_def_is_implicit',
+              'lr_use_def_is_renamable',
+              'lr_use_def_is_ind_var_update',
+              'lr_use_def_is_hint',
+          ]
+      }
+  )
+  observation_spec['lr_use_def_freq'] = tf.TensorSpec(
+      dtype=tf.float32,
+      shape=(_NUM_REGISTERS, _NUM_INSTRUCTIONS),
+      name='lr_use_def_freq',
+  )
+
+  encoding_spec = tf.TensorSpec(
+      dtype=tf.float32, shape=(33, _ENCODING_SIZE), name='lr_encoding'
+  )
+
+  return observation_spec, encoding_spec
+
+
+def get_input_specs():
+  encoder_observation_spec, encoding_spec = get_lr_encoder_signature_spec()
+  del encoding_spec
+
+  regalloc_time_step_spec, regalloc_action_spec = (
+      config.get_regalloc_signature_spec()
+  )
+
+  # Ensure that there are no overlaps in feature names between the encoder and the RL problem
+  common_keys = (
+      encoder_observation_spec.keys()
+      & regalloc_time_step_spec.observation.keys()
+  )
+  assert len(common_keys) == 0
+
+  supervised_input_spec = {
+      **encoder_observation_spec,
+      **regalloc_time_step_spec.observation,
+  }
+
+  return supervised_input_spec
+
+
+def get_output_specs(regalloc_preprocessing_layer_creator):
+  regalloc_time_step_spec, regalloc_action_spec = (
+      config.get_regalloc_signature_spec()
+  )
+  del regalloc_action_spec
+
+  random_regalloc_state = tf_agents.specs.sample_spec_nest(
+      regalloc_time_step_spec.observation
+  )
+  for key in random_regalloc_state:
+    preprocessing_layer = regalloc_preprocessing_layer_creator(
+        regalloc_time_step_spec.observation[key]
+    )
+    random_regalloc_state[key] = preprocessing_layer(
+        tf.expand_dims(random_regalloc_state[key], axis=0)
+    )
+  random_regalloc_state = tf.concat(
+      list(random_regalloc_state.values()), axis=-1
+  )
+  random_regalloc_state = tf.squeeze(random_regalloc_state, axis=0)
+  regalloc_state_shape = random_regalloc_state.shape
+
+  action_spec = tf.TensorSpec(
+      dtype=tf.float32, shape=(_NUM_REGISTERS, 1), name=_ACTION_KEY
+  )
+  state_spec = tf.TensorSpec(
+      dtype=tf.float32, shape=regalloc_state_shape, name=_STATE_KEY
+  )
+  mlm_spec = tf.TensorSpec(
+      dtype=tf.float32,
+      shape=(_NUM_REGISTERS, _NUM_INSTRUCTIONS, _OPCODE_VOCAB_SIZE),
+      name=_MLM_KEY,
+  )
+  return {
+      _ACTION_KEY: action_spec,
+      _STATE_KEY: state_spec,
+      _NEXT_ACTION_KEY: action_spec,
+      _NEXT_STATE_KEY: state_spec,
+      _MLM_KEY: mlm_spec,
+  }
+
+
+def get_loss():
+  loss_fns = {
+      _STATE_KEY: 'mse',
+      _ACTION_KEY: tf.keras.losses.SparseCategoricalCrossentropy(
+          from_logits=True
+      ),
+      _NEXT_STATE_KEY: 'mse',
+      _NEXT_ACTION_KEY: tf.keras.losses.SparseCategoricalCrossentropy(
+          from_logits=True
+      ),
+      _MLM_KEY: tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+  }
+  loss_weights = {
+      _ACTION_KEY: 1.0,
+      _STATE_KEY: 1.0,
+      _NEXT_ACTION_KEY: 1.0,
+      _NEXT_STATE_KEY: 1.0,
+      _MLM_KEY: 10.0,
+  }
+  return loss_fns, loss_weights
+
+
+def get_metrics():
+  return {
+      _ACTION_KEY: ['accuracy'],
+      _STATE_KEY: [],
+      _NEXT_ACTION_KEY: ['accuracy'],
+      _NEXT_STATE_KEY: [],
+      _MLM_KEY: ['accuracy'],
+  }
+
+
+@gin.configurable
+def get_preprocessing_layer_creator(
+    quantile_file_dir='/cns/oz-d/home/mlcompileropt-dev/regalloc-transformer/vocab',
+    with_sqrt=True,
+    with_z_score_normalization=True,
+    eps=1e-8,
+):
+  """Wrapper for observation_processing_layer."""
+  quantile_map = feature_ops.build_quantile_map(quantile_file_dir)
+
+  def preprocessing_layer_creator(obs_spec):
+    """Creates the layer to process observation given obs_spec."""
+    if obs_spec.name == 'lr_use_def_freq':
+      quantile = quantile_map[obs_spec.name]
+      first_non_zero = 0
+      for x in quantile:
+        if x > 0:
+          first_non_zero = x
+          break
+
+      normalize_fn = feature_ops.get_normalize_fn(
+          quantile, with_sqrt, with_z_score_normalization, eps
+      )
+      return tf.keras.layers.Lambda(normalize_fn)
+    return tf.keras.layers.Lambda(feature_ops.identity_fn)
+
+  return preprocessing_layer_creator
+
+
+def get_nonnormalized_features():
+  return [
+      'lr_use_def_opcode',
+      'lr_use_def_read',
+      'lr_use_def_write',
+      'lr_use_def_is_use',
+      'lr_use_def_is_def',
+      'lr_use_def_is_implicit',
+      'lr_use_def_is_renamable',
+      'lr_use_def_is_ind_var_update',
+      'lr_use_def_is_hint',
+  ]
diff --git a/compiler_opt/rl/regalloc/lr_encoder/dataset_ops.py b/compiler_opt/rl/regalloc/lr_encoder/dataset_ops.py
new file mode 100644
index 0000000..9967f16
--- /dev/null
+++ b/compiler_opt/rl/regalloc/lr_encoder/dataset_ops.py
@@ -0,0 +1,158 @@
+import tensorflow as tf
+import tensorflow_text as text
+
+from google3.third_party.ml_compiler_opt.compiler_opt.rl.regalloc.lr_encoder import config
+
+_MAX_PREDICTIONS_PER_BATCH = 128
+
+_PAD_TOKEN = 0
+_MLM_IGNORE_TOKEN = -1
+_MLM_MASK_TOKEN = 18000 - 1
+
+
+def _get_masked_language_fn(*, selection_rate=0.33):
+  random_selector = text.RandomItemSelector(
+      max_selections_per_batch=_MAX_PREDICTIONS_PER_BATCH,
+      selection_rate=0.2,
+      unselectable_ids=[_PAD_TOKEN],
+  )
+  mask_values_chooser = text.MaskValuesChooser(
+      config._OPCODE_VOCAB_SIZE, _MLM_MASK_TOKEN, 0.8
+  )
+
+  def fn(opcodes):
+    masked_token_ids, masked_pos, masked_lm_ids = text.mask_language_model(
+        tf.RaggedTensor.from_tensor(opcodes, padding=_PAD_TOKEN),
+        item_selector=random_selector,
+        mask_values_chooser=mask_values_chooser,
+    )
+
+    masked_pos = masked_pos.to_tensor(
+        default_value=-1,
+        shape=(config._NUM_REGISTERS, _MAX_PREDICTIONS_PER_BATCH),
+    )
+
+    ii = tf.tile(
+        tf.range(config._NUM_REGISTERS, dtype=tf.int64)[:, tf.newaxis],
+        [1, _MAX_PREDICTIONS_PER_BATCH],
+    )
+    masked_pos = tf.stack([ii, masked_pos], axis=-1)
+    scatter_values = tf.where(masked_pos[:, :, 1] < 0, 0.0, 1.0)
+    masked_pos = tf.where(
+        masked_pos < 0, tf.constant(_PAD_TOKEN, dtype=tf.int64), masked_pos
+    )
+    weights = tf.scatter_nd(
+        masked_pos,
+        scatter_values[:, :, tf.newaxis],
+        (config._NUM_REGISTERS, config._NUM_INSTRUCTIONS, 1),
+    )
+    return (
+        masked_token_ids.to_tensor(
+            default_value=0,
+            shape=(config._NUM_REGISTERS, config._NUM_INSTRUCTIONS),
+        ),
+        opcodes,
+        weights,
+    )
+
+  return fn
+
+
+def _roll_experience(seq_ex, *, shift=-1):
+  def _roll(atom):
+    return tf.roll(atom, shift=shift, axis=0)
+
+  def _cutoff(atom):
+    return atom[:shift]
+
+  seq_ex_roll = tf.nest.map_structure(_roll, seq_ex)
+  seq_ex_roll = tf.nest.map_structure(_cutoff, seq_ex_roll)
+  seq_ex = tf.nest.map_structure(_cutoff, seq_ex)
+  return seq_ex, seq_ex_roll
+
+
+def _split_sequence_example(seq_ex, seq_ex_roll):
+  action_name = 'index_to_evict'
+  obs = {k: seq_ex[k] for k in seq_ex if k != action_name}
+  action = seq_ex[action_name]
+  obs_roll = {k: seq_ex_roll[k] for k in seq_ex_roll if k != action_name}
+  action_roll = seq_ex_roll[action_name]
+  return {
+      'obs': obs,
+      'action': action,
+      'obs_roll': obs_roll,
+      'action_roll': action_roll,
+  }
+
+
+def _get_state_preprocessing_layer(
+    regalloc_input_spec, preprocessing_layer_creator
+):
+  preprocessing_layers = {}
+  for name, spec in regalloc_input_spec.items():
+    preprocessing_layers[name] = preprocessing_layer_creator(spec)
+
+  def _preprocessing_layer(seq_ex):
+    pp = []
+    for layer_name, layer in preprocessing_layers.items():
+      pp.append(layer(seq_ex[layer_name]))
+    return tf.concat(pp, axis=-1)
+
+  def _layer(obs_dict):
+    obs_dict['obs'] = obs_dict['obs']
+    obs_dict['obs_cur'] = _preprocessing_layer(obs_dict['obs'])
+    obs_dict['obs_roll'] = _preprocessing_layer(obs_dict['obs_roll'])
+    return obs_dict
+
+  return _layer
+
+
+def _get_to_inputs_and_labels_fn():
+  masked_language_fn = _get_masked_language_fn()
+
+  def fn(obs_dict):
+    inputs = {'obs': obs_dict['obs'], 'action': obs_dict['action']}
+    mlm_input, mlm_label, mlm_weight = masked_language_fn(
+        obs_dict['obs']['lr_use_def_opcode'][:, : config._NUM_INSTRUCTIONS]
+    )
+    inputs['lr_use_def_opcode'] = mlm_input
+
+    labels = {
+        config._STATE_KEY: obs_dict['obs_cur'],
+        config._ACTION_KEY: tf.expand_dims(obs_dict['action'], axis=-1),
+        config._NEXT_STATE_KEY: obs_dict['obs_roll'],
+        config._NEXT_ACTION_KEY: tf.expand_dims(
+            obs_dict['action_roll'], axis=-1
+        ),
+        config._MLM_KEY: mlm_label,
+    }
+    mask = obs_dict['obs']['mask']
+    sample_weights = {
+        config._STATE_KEY: mask,
+        config._ACTION_KEY: mask,
+        config._NEXT_STATE_KEY: None,
+        config._NEXT_ACTION_KEY: None,
+        config._MLM_KEY: mlm_weight,
+    }
+    return (inputs, labels, sample_weights)
+
+  return fn
+
+
+def process_dataset(
+    dataset, regalloc_input_spec, regalloc_preprocessing_layer_creator
+):
+  shuffle_buffer_size = 128
+  num_map_threads = 128
+  state_preprocessing_layer = _get_state_preprocessing_layer(
+      regalloc_input_spec, regalloc_preprocessing_layer_creator
+  )
+  to_inputs_and_labels_fn = _get_to_inputs_and_labels_fn()
+  return (
+      dataset.map(_roll_experience, num_parallel_calls=num_map_threads)
+      .map(_split_sequence_example, num_parallel_calls=num_map_threads)
+      .map(state_preprocessing_layer, num_parallel_calls=num_map_threads)
+      .unbatch()
+      .shuffle(shuffle_buffer_size)
+      .map(to_inputs_and_labels_fn, num_parallel_calls=num_map_threads)
+  )
diff --git a/compiler_opt/rl/regalloc/lr_encoder/gin_configs/common.gin b/compiler_opt/rl/regalloc/lr_encoder/gin_configs/common.gin
new file mode 100644
index 0000000..c0fcb35
--- /dev/null
+++ b/compiler_opt/rl/regalloc/lr_encoder/gin_configs/common.gin
@@ -0,0 +1,14 @@
+config_registry.get_configuration.implementation=@configs.LREncoderConfig
+
+launcher_path=None
+clang_path=None
+
+runners.LREncoderRunner.clang_path=%clang_path
+runners.LREncoderRunner.launcher_path=%launcher_path
+
+problem_config.flags_to_add.add_flags=()
+problem_config.flags_to_delete.delete_flags=('-split-dwarf-file','-split-dwarf-output',)
+# For AFDO profile reinjection set:
+# problem_config.flags_to_replace.replace_flags={'-fprofile-sample-use':'/path/to/gwp.afdo','-fprofile-remapping-file':'/path/to/prof_remap.txt'}
+problem_config.flags_to_replace.replace_flags={'-fprofile-instrument-use-path': '/tmp/corpus_1_21/MergedCS.profdata'}
+#problem_config.flags_to_replace.replace_flags={'-fprofile-instrument-use-path': 'MergedCS_profdata_jacobhegna/MergedCS_profdata'}
diff --git a/compiler_opt/rl/regalloc/lr_encoder/lr_encoder_runner.py b/compiler_opt/rl/regalloc/lr_encoder/lr_encoder_runner.py
new file mode 100644
index 0000000..3c90edc
--- /dev/null
+++ b/compiler_opt/rl/regalloc/lr_encoder/lr_encoder_runner.py
@@ -0,0 +1,94 @@
+"""Module for collect data of the LR encoder."""
+
+import os
+import tempfile
+from typing import Dict, Tuple
+
+import gin
+import tensorflow as tf
+
+from google3.third_party.ml_compiler_opt.compiler_opt.rl import compilation_runner
+from google3.third_party.ml_compiler_opt.compiler_opt.rl import corpus
+from google3.third_party.ml_compiler_opt.compiler_opt.rl import log_reader
+
+
+@gin.configurable(module='runners')
+class LREncoderRunner(compilation_runner.CompilationRunner):
+  """Class for collecting data for the LR encoder."""
+
+  def compile_fn(
+      self,
+      command_line: corpus.FullyQualifiedCmdLine,
+      tf_policy_path: str,
+      reward_only: bool,
+  ) -> Dict[str, Tuple[tf.train.SequenceExample, float]]:
+    """Run compilation for the given IR file under the given policy.
+
+    Args:
+      command_line: the fully qualified command line.
+      tf_policy_path: path to TF policy direcoty on local disk.
+      reward_only: whether only return reward.
+
+    Returns:
+      A dict mapping from example identifier to tuple containing:
+        sequence_example: A tf.SequenceExample proto describing compilation
+          trace, None if reward_only == True.
+        reward: reward of register allocation.
+
+    Raises:
+      subprocess.CalledProcessError: if process fails.
+      compilation_runner.ProcessKilledError: (which it must pass through) on
+        cancelled work.
+      RuntimeError: if llvm-size produces unexpected output.
+    """
+    assert not tf_policy_path
+
+    working_dir = tempfile.mkdtemp()
+
+    log_path = os.path.join(working_dir, 'log')
+    output_native_path = os.path.join(working_dir, 'native')
+
+    result = {}
+    try:
+      cmdline = []
+      if self._launcher_path:
+        cmdline.append(self._launcher_path)
+      cmdline.extend(
+          [self._clang_path]
+          + list(command_line)
+          + [
+              '-mllvm',
+              '-regalloc-enable-advisor=development',
+              '-mllvm',
+              '-regalloc-lr-encoder-training-log=' + log_path,
+              '-mllvm',
+              '-regalloc-training-log=/dev/null',
+              '-o',
+              output_native_path,
+          ]
+      )
+
+      compilation_runner.start_cancellable_process(
+          cmdline, self._compilation_timeout, self._cancellation_manager
+      )
+
+      if not os.path.exists(log_path):
+        return {}
+
+      # TODO(#202)
+      log_result = log_reader.read_log_as_sequence_examples(log_path)
+
+      for fct_name, trajectory in log_result.items():
+        if not trajectory.HasField('feature_lists'):
+          continue
+        # score = (
+        #    trajectory.feature_lists.feature_list['reward']
+        #    .feature[-1]
+        #    .float_list.value[0]
+        # )
+        result[fct_name] = (trajectory, 1.0)
+
+    finally:
+      tf.io.gfile.rmtree(working_dir)
+
+    return result
diff --git a/compiler_opt/rl/regalloc/lr_encoder/model.py b/compiler_opt/rl/regalloc/lr_encoder/model.py
new file mode 100644
index 0000000..a953aa4
--- /dev/null
+++ b/compiler_opt/rl/regalloc/lr_encoder/model.py
@@ -0,0 +1,114 @@
+import gin
+import tensorflow as tf
+from tf_agents.utils import nest_utils
+from google3.third_party.ml_compiler_opt.compiler_opt.rl.regalloc.lr_encoder import config
+from google3.third_party.ml_compiler_opt.compiler_opt.rl import attention
+
+
+class MultiHeadEncoderModel(tf.keras.Model):
+
+  def __init__(self, encoder, heads):
+    super().__init__(name='MultiHeadEncoderModel')
+
+    self._encoder = encoder
+    self._heads = heads
+
+  def get_encoder(self):
+    return self._encoder
+
+  def call(self, inputs, *args):
+    observation = inputs['obs']
+    action = tf.one_hot(inputs['action'], depth=config._NUM_REGISTERS)[
+        :, :, tf.newaxis
+    ]
+
+    use_def_obs = {
+        k: v
+        for k, v in observation.items()
+        if k.startswith(config._ENCODER_FEATURE_PREFIX)
+    }
+    encoded_state_per_token, encoded_state = self._encoder(use_def_obs)
+    encoded_state_with_action = tf.concat([encoded_state, action], axis=-1)
+
+    def get_input(head_name):
+      if head_name == 'mlm':
+        return encoded_state_per_token
+      if head_name.startswith('next_'):
+        return encoded_state_with_action
+      return encoded_state
+
+    outputs = {}
+    for name in self._heads:
+      head = self._heads[name]
+      head_input = get_input(name)
+      outputs[name] = head(head_input)
+
+    return outputs
+
+
+@gin.configurable
+class LiveRangeEncoder(tf.keras.Model):
+
+  def __init__(
+      self,
+      input_specs,
+      preprocessing_layer_creator,
+      *,
+      num_layers=2,
+      num_heads=2,
+      model_dim=32,
+      fcn_dim=128,
+      num_extra_features=10,
+  ):
+    super().__init__(name='LiveRangeEncoder')
+    self._encoder = attention.TransformerClassifier(
+        num_tokens=config._OPCODE_VOCAB_SIZE,
+        num_layers=num_layers,
+        num_heads=num_heads,
+        model_dim=model_dim,
+        fcn_dim=fcn_dim,
+        num_extra_features=num_extra_features,
+    )
+    self._linear_reshape = tf.keras.layers.Dense(config._ENCODING_SIZE)
+
+    self._preprocessing_layers = {}
+    for key in input_specs:
+      self._preprocessing_layers[key] = preprocessing_layer_creator(key)
+
+    def preprocessor(observations):
+      pp_inputs = []
+      for key in input_specs:
+        pp_inputs.append(self._preprocessing_layers[key](observations[key]))
+      return tf.concat(pp_inputs, axis=-1)
+
+    self._preprocessor = preprocessor
+
+  def call(self, observations):
+    extra_hidden_state = self._preprocessor(observations)
+    tokens = observations[config._OPCODE_KEY]
+    encoded_state_per_token, encoded_state = self._encoder(
+        tokens, extra_hidden_state
+    )
+    encoded_state = self._linear_reshape(encoded_state)
+    return encoded_state_per_token, encoded_state
+
+
+def create_model(
+    input_specs: dict,
+    output_specs: dict,
+    preprocessing_layer_creator: dict,
+    *,
+    mlp_width: int = 128,
+):
+  encoder = LiveRangeEncoder(input_specs, preprocessing_layer_creator)
+
+  output_heads = {}
+  for name, spec in output_specs.items():
+    size = spec.shape[-1]
+    output_heads[name] = tf.keras.Sequential([
+        tf.keras.layers.Dense(mlp_width),
+        tf.keras.layers.ReLU(),
+        tf.keras.layers.Dense(size),
+    ])
+
+  return MultiHeadEncoderModel(encoder=encoder, heads=output_heads)
diff --git a/compiler_opt/rl/regalloc/lr_encoder/train.py b/compiler_opt/rl/regalloc/lr_encoder/train.py
new file mode 100644
index 0000000..ca0b628
--- /dev/null
+++ b/compiler_opt/rl/regalloc/lr_encoder/train.py
@@ -0,0 +1,183 @@
+import os
+
+from absl import app
+from absl import flags
+from absl import logging
+
+import gin
+import tensorflow as tf
+
+from google3.third_party.ml_compiler_opt.compiler_opt.rl import policy_saver
+from google3.third_party.ml_compiler_opt.compiler_opt.rl.regalloc import config as regalloc_config
+from google3.third_party.ml_compiler_opt.compiler_opt.rl.regalloc.lr_encoder import dataset_ops
+from google3.third_party.ml_compiler_opt.compiler_opt.rl.regalloc.lr_encoder import model
+from google3.third_party.ml_compiler_opt.compiler_opt.rl.regalloc.lr_encoder import config as encoder_config
+from google3.third_party.ml_compiler_opt.compiler_opt.rl import registry  # pylint: disable=unused-import
+
+
+flags.DEFINE_multi_string(
+    'gin_files', [], 'List of paths to gin configuration files.'
+)
+
+flags.DEFINE_string('trace', None, '')
+flags.DEFINE_string('root_dir', None, '')
+
+flags.DEFINE_string('tpu', None, 'BNS address for the TPU')
+
+FLAGS = flags.FLAGS
+
+
+def _create_parser_fn(time_step_spec, action_spec):
+  context_features = {}
+  sequence_features = {
+      tensor_spec.name: tf.io.FixedLenSequenceFeature(
+          shape=tensor_spec.shape, dtype=tensor_spec.dtype
+      )
+      for tensor_spec in time_step_spec.values()
+  }
+  sequence_features[action_spec.name] = tf.io.FixedLenSequenceFeature(
+      shape=action_spec.shape, dtype=action_spec.dtype
+  )
+
+  def _parser_fn(serialized_proto):
+    with tf.name_scope('parse'):
+      _, parsed_sequence = tf.io.parse_single_sequence_example(
+          serialized_proto,
+          context_features=context_features,
+          sequence_features=sequence_features,
+      )
+      return parsed_sequence
+
+  return _parser_fn
+
+
+def get_dataset_creator_fn(
+    batch_size,
+    input_spec,
+    action_spec,
+    regalloc_input_spec,
+    regalloc_preprocessing_layer_creator,
+):
+  files_buffer_size = 64
+  num_map_threads = 128
+  num_readers = 64
+  parser_fn = _create_parser_fn(input_spec, action_spec)
+
+  def _file_dataset_fn(data_path):
+    dataset = (
+        tf.data.Dataset.list_files(data_path)
+        .shuffle(files_buffer_size)
+        .interleave(
+            tf.data.TFRecordDataset,
+            cycle_length=num_readers,
+            block_length=1,
+        )
+        .filter(lambda string: tf.strings.length(string) > 0)
+        .map(parser_fn, num_parallel_calls=num_map_threads)
+    )
+    dataset = dataset_ops.process_dataset(
+        dataset, regalloc_input_spec, regalloc_preprocessing_layer_creator
+    )
+    return dataset.batch(batch_size, drop_remainder=True).prefetch(
+        tf.data.experimental.AUTOTUNE
+    )
+
+  return _file_dataset_fn
+
+
+def get_strategy():
+  if FLAGS.tpu:
+    logging.info('Using TPU strategy.')
+    resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu=FLAGS.tpu)
+    tf.config.experimental_connect_to_cluster(resolver)
+    tf.tpu.experimental.initialize_tpu_system(resolver)
+    return tf.distribute.TPUStrategy(resolver)
+  logging.info('Using CPU strategy.')
+  return tf.distribute.get_strategy()
+
+
+class SaveEncoderCallback(tf.keras.callbacks.Callback):
+
+  def __init__(self, encoder, path):
+    self._encoder = encoder
+    self._path = path
+
+  def on_epoch_end(self, epoch, logs=None):
+    sm_path = os.path.join(self._path, f'epoch{epoch}')
+    tflite_path = os.path.join(sm_path, policy_saver.TFLITE_MODEL_NAME)
+    self._encoder.save(sm_path)
+    policy_saver.convert_saved_model(sm_path, tflite_path)
+
+
+def main(_):
+  gin.parse_config_files_and_bindings(
+      FLAGS.gin_files, bindings=None, skip_unknown=False
+  )
+  logging.info(gin.config_str())
+
+  lr_input_specs, lr_encoding_spec = (
+      encoder_config.get_lr_encoder_signature_spec()
+  )
+
+  regalloc_time_step_spec, regalloc_action_spec = (
+      regalloc_config.get_regalloc_signature_spec()
+  )
+  regalloc_preprocessing_creator = (
+      regalloc_config.get_observation_processing_layer_creator()
+  )
+
+  dataset_fn = get_dataset_creator_fn(
+      batch_size=64,
+      input_spec=encoder_config.get_input_specs(),
+      action_spec=regalloc_action_spec,
+      regalloc_input_spec=regalloc_time_step_spec.observation,
+      regalloc_preprocessing_layer_creator=regalloc_preprocessing_creator,
+  )
+
+  strategy = get_strategy()
+  with strategy.scope():
+    logging.info('Creating model.')
+    lr_model = model.create_model(
+        lr_input_specs,
+        encoder_config.get_output_specs(regalloc_preprocessing_creator),
+        encoder_config.get_preprocessing_layer_creator(),
+    )
+
+    logging.info('Compiling model.')
+    loss, loss_weights = encoder_config.get_loss()
+    lr_model.compile(
+        optimizer=tf.keras.optimizers.Adam(global_clipnorm=1.0),
+        loss=loss,
+        loss_weights=loss_weights,
+        metrics=encoder_config.get_metrics(),
+    )
+
+  logging.info('Creating dataset.')
+  dataset = dataset_fn(FLAGS.trace)
+  tb = tf.keras.callbacks.TensorBoard(
+      log_dir=FLAGS.root_dir,
+      histogram_freq=0,
+      write_graph=True,
+      write_steps_per_second=True,
+      update_freq='batch',
+  )
+
+  policy_dir = os.path.join(FLAGS.root_dir, 'policy')
+  checkpoint = tf.keras.callbacks.ModelCheckpoint(
+      filepath=policy_dir, save_weights_only=False, save_freq=1024
+  )
+
+  encoder_dir = os.path.join(FLAGS.root_dir, 'encoder')
+  logging.info('Saving the encoder to %s', encoder_dir)
+  encoder_saver = SaveEncoderCallback(
+      encoder=lr_model.get_encoder(), path=encoder_dir
+  )
+
+  logging.info('Starting training.')
+  lr_model.fit(dataset, epochs=8, callbacks=[tb, checkpoint, encoder_saver])
+
+  logging.info('Training complete.')
+
+
+if __name__ == '__main__':
+  app.run(main)
diff --git a/compiler_opt/rl/train_semisupervised.py b/compiler_opt/rl/train_semisupervised.py
new file mode 100644
index 0000000..209f2f9
--- /dev/null
+++ b/compiler_opt/rl/train_semisupervised.py
@@ -0,0 +1,490 @@
+from absl import app
+import tensorflow as tf
+import tensorflow_text as text
+
+import numpy as np
+
+import os
+from typing import Optional, List
+
+from google3.third_party.ml_compiler_opt.compiler_opt.rl import feature_ops
+from google3.third_party.ml_compiler_opt.compiler_opt.rl import registry
+from google3.third_party.ml_compiler_opt.compiler_opt.rl import attention
+from google3.third_party.ml_compiler_opt.compiler_opt.rl.regalloc import regalloc_network
+from google3.third_party.ml_compiler_opt.compiler_opt.rl import policy_saver
+from absl import flags
+from absl import logging
+import gin
+# from google3.third_party.ml_compiler_opt.compiler_opt.rl import policy_saver
+
+
+# Have one class for the Encoder
+# Have one class for the full semisupervised model
+# only save the Encoder
+# Output heads:
+#   - reward
+#   - next state
+#   -
+
+flags.DEFINE_multi_string(
+    'gin_files', [], 'List of paths to gin configuration files.'
+)
+
+flags.DEFINE_string('trace', None, '')
+flags.DEFINE_string('root_dir', None, '')
+
+flags.DEFINE_string('tpu', None, 'BNS address for the TPU')
+
+
+FLAGS = flags.FLAGS
+
+_ACTION_KEY = 'action'
+_CUR_STATE_KEY = 'cur_state'
+_NEXT_STATE_KEY = 'next_state'
+_NEXT_ACTION_KEY = 'next_action'
+_MLM_KEY = 'mlm'
+
+_INSTR_COUNT = 64
+
+_MLM_IGNORE_TOKEN = -1
+_MLM_MASK_TOKEN = 18000 - 1
+
+
+class Model(tf.keras.Model):
+
+  def __init__(
+      self,
+      encoder_network,
+      state_head,
+      action_head,
+      next_state_head,
+      next_action_head,
+      mlm_head,
+  ):
+    super().__init__(name='Model')
+    self._encoder_network = encoder_network
+    self._state_head = state_head
+    self._action_head = action_head
+    self._next_state_head = next_state_head
+    self._next_action_head = next_action_head
+    self._mlm_head = mlm_head
+
+  def call(self, inputs):
+    # TODO:
+    # 1) export the non-reduced tensor because need it for the MLM outputs
+    # 2) when reducing, mask the tensors that don't matter
+    #       for masking, I think I should set the relevant outputs to zero and also zero out the labels.
+    # 3) figure out the correct masks when training, might need custom training loop or custom loss functions.
+    observation = inputs['obs']
+    action = tf.one_hot(inputs['action'], depth=33)[:, :, tf.newaxis]
+
+    use_def_obs = {
+        k: v for k, v in observation.items() if k.startswith('use_def_')
+    }
+    encoded_state_per_token, encoded_state = self._encoder_network(use_def_obs)
+
+    encoded_state_with_action = tf.concat([encoded_state, action], axis=-1)
+
+    state = self._state_head(encoded_state)
+    action = self._action_head(encoded_state)
+    next_state = self._next_state_head(encoded_state_with_action)
+    next_action = self._next_action_head(encoded_state_with_action)
+    mlm = self._mlm_head(encoded_state_per_token)
+    return {
+        _CUR_STATE_KEY: state,
+        _ACTION_KEY: action,
+        _NEXT_STATE_KEY: next_state,
+        _NEXT_ACTION_KEY: next_action,
+        _MLM_KEY: mlm,
+    }
+
+
+def action_loss(label, pred):
+  # Use the current mask
+  print(label)
+  pass
+
+
+def state_loss(label, pred):
+  # Decide whether to use the intersection of the masks or just the current mask
+  # Probably use the current mask
+  print(label)
+  pass
+
+
+def masked_language_loss(label, pred):
+  mask = label != 0
+  loss_object = tf.keras.losses.SparseCategoricalCrossEntropy(
+      from_logits=True, reduction='none'
+  )
+  loss = loss_object(label, pred)
+
+  mask = tf.cast(mask, dtype=loss.dtype)
+  loss *= mask
+
+  loss = tf.reduce_sum(loss) / tf.reduce_sum(mask)
+  return loss
+
+
+def create_model(encoder_network):
+  action_head = tf.keras.Sequential([
+      tf.keras.layers.Dense(128),
+      tf.keras.layers.ReLU(),
+      tf.keras.layers.Dense(1),
+  ])
+  state_head = tf.keras.Sequential([
+      tf.keras.layers.Dense(128),
+      tf.keras.layers.ReLU(),
+      tf.keras.layers.Dense(65),
+  ])
+  next_state_head = tf.keras.Sequential([
+      tf.keras.layers.Dense(128),
+      tf.keras.layers.ReLU(),
+      tf.keras.layers.Dense(65),
+  ])
+  next_action_head = tf.keras.Sequential([
+      tf.keras.layers.Dense(128),
+      tf.keras.layers.ReLU(),
+      tf.keras.layers.Dense(1),
+  ])
+  mlm = tf.keras.Sequential([
+      tf.keras.layers.Dense(18000),
+  ])
+  return Model(
+      encoder_network=encoder_network,
+      action_head=action_head,
+      state_head=state_head,
+      next_state_head=next_state_head,
+      next_action_head=next_action_head,
+      mlm_head=mlm,
+  )
+
+
+def _create_parser_fn(time_step_spec, action_spec):
+  context_features = {}
+  sequence_features = {
+      tensor_spec.name: tf.io.FixedLenSequenceFeature(
+          shape=tensor_spec.shape, dtype=tensor_spec.dtype
+      )
+      for tensor_spec in time_step_spec.observation.values()
+  }
+  sequence_features[action_spec.name] = tf.io.FixedLenSequenceFeature(
+      shape=action_spec.shape, dtype=action_spec.dtype
+  )
+
+  def _parser_fn(serialized_proto):
+    with tf.name_scope('parse'):
+      _, parsed_sequence = tf.io.parse_single_sequence_example(
+          serialized_proto,
+          context_features=context_features,
+          sequence_features=sequence_features,
+      )
+      return parsed_sequence
+
+  return _parser_fn
+
+
+def _get_masked_input_and_labels(encoded_texts):
+  # https://keras.io/examples/nlp/masked_language_modeling/
+  # 15% BERT masking
+  encoded_texts_shape = tf.shape(encoded_texts)
+  inp_mask = tf.random.uniform(encoded_texts_shape) < 0.15
+  inp_mask[encoded_texts <= 0] = False
+  labels = _MLM_IGNORE_TOKEN * tf.ones(encoded_texts_shape, dtype=fp.int64)
+  # Set labels for masked tokens
+  labels[inp_mask] = encoded_texts[inp_mask]
+
+  # Prepare input
+  encoded_texts_masked = tf.identity(encoded_texts)
+  # Set input to [MASK] which is the last token for the 90% of tokens
+  # This means leaving 10% unchanged
+  inp_mask_2mask = inp_mask & (tf.random.uniform(encoded_texts_shape) < 0.90)
+  encoded_texts_masked[inp_mask_2mask] = _MLM_MASK_TOKEN
+
+  # Set 10% to a random token
+  inp_mask_2random = inp_mask_2mask & (
+      tf.random.uniform(encoded_texts_shape) < 1 / 9
+  )
+  encoded_texts_masked[inp_mask_2random] = tf.random.uniform(
+      encoded_texts_shape, 3, mask_token_id, dtype=tf.int64
+  )
+
+  # Prepare sample_weights to pass to .fit() method
+  sample_weights = tf.ones(encoded_texts_shape)
+  sample_weights[labels == -1] = 0
+
+  # y_labels would be same as encoded_texts i.e input tokens
+  y_labels = tf.identity(encoded_texts)
+
+  return encoded_texts_masked, y_labels, sample_weights
+
+
+_MAX_PREDICTIONS_PER_BATCH = 32
+random_selector = text.RandomItemSelector(
+    max_selections_per_batch=_MAX_PREDICTIONS_PER_BATCH,
+    selection_rate=0.2,
+    unselectable_ids=[0],
+)
+mask_values_chooser = text.MaskValuesChooser(18000, _MLM_MASK_TOKEN, 0.8)
+
+
+def get_masked_input_and_labels(encoded_texts):
+  masked_token_ids, masked_pos, masked_lm_ids = text.mask_language_model(
+      tf.RaggedTensor.from_tensor(encoded_texts, padding=0),
+      item_selector=random_selector,
+      mask_values_chooser=mask_values_chooser,
+  )
+  # NEed to fix this
+  # Produce tensor 0 to 32, tile it and produce tensr of indices, then flatten
+  # Then can use it in the scatter like I want
+  masked_pos = masked_pos.to_tensor(
+      default_value=-1, shape=(33, _MAX_PREDICTIONS_PER_BATCH)
+  )
+
+  ii = tf.tile(
+      tf.range(33, dtype=tf.int64)[:, tf.newaxis],
+      [1, _MAX_PREDICTIONS_PER_BATCH],
+  )
+  masked_pos = tf.stack([ii, masked_pos], axis=-1)
+  scatter_values = tf.where(masked_pos[:, :, 1] < 0, 0.0, 1.0)
+  masked_pos = tf.where(
+      masked_pos < 0, tf.constant(0, dtype=tf.int64), masked_pos
+  )
+  weights = tf.scatter_nd(
+      masked_pos, scatter_values[:, :, tf.newaxis], (33, _INSTR_COUNT, 1)
+  )
+  return (
+      masked_token_ids.to_tensor(default_value=0, shape=(33, _INSTR_COUNT)),
+      encoded_texts,
+      weights,
+  )
+
+
+def create_dataset_fn(
+    time_step_spec, action_spec, batch_size, shift, preprocessing_layer_creator
+):
+  assert shift < 0
+  files_buffer_size = 100
+  num_readers = 10
+  num_map_threads = 8
+  shuffle_buffer_size = 256
+
+  parser_fn = _create_parser_fn(time_step_spec, action_spec)
+
+  def _roll_experience(seq_ex):
+    # Use tf agents nest map fn here
+    def _roll(atom):
+      return tf.roll(atom, shift=shift, axis=0)
+
+    def _cutoff(atom):
+      return atom[:shift]
+
+    seq_ex_roll = tf.nest.map_structure(_roll, seq_ex)
+    seq_ex_roll = tf.nest.map_structure(_cutoff, seq_ex_roll)
+    seq_ex = tf.nest.map_structure(_cutoff, seq_ex)
+    return seq_ex, seq_ex_roll
+
+  preprocessing_layers = tf.nest.map_structure(
+      preprocessing_layer_creator, time_step_spec.observation
+  )
+
+  def split_experience(seq_ex, seq_ex_roll):
+    obs = {k: seq_ex[k] for k in seq_ex if k != action_spec.name}
+    action = seq_ex[action_spec.name]
+    obs_roll = {k: seq_ex_roll[k] for k in seq_ex_roll if k != action_spec.name}
+    action_roll = seq_ex_roll[action_spec.name]
+    return {
+        'obs': obs,
+        'action': action,
+        'obs_roll': obs_roll,
+        'action_roll': action_roll,
+    }
+
+  def _preprocessing_layer(seq_ex):
+    for layer_name, layer in preprocessing_layers.items():
+      seq_ex[layer_name] = layer(seq_ex[layer_name])
+    return seq_ex
+
+  def preprocess_experience(obs_dict):
+    obs_dict['obs_cur'] = _preprocessing_layer(obs_dict['obs'].copy())
+    obs_dict['obs_cur'] = tf.concat(
+        [
+            v
+            for k, v in obs_dict['obs_cur'].items()
+            if not k.startswith('use_def_')
+        ],
+        axis=-1,
+    )
+    obs_dict['obs_roll'] = _preprocessing_layer(obs_dict['obs_roll'])
+    obs_dict['obs_roll'] = tf.concat(
+        [
+            v
+            for k, v in obs_dict['obs_roll'].items()
+            if not k.startswith('use_def_')
+        ],
+        axis=-1,
+    )
+    return obs_dict
+
+  def to_inputs_and_labels(obs_dict):
+    inputs = {'obs': obs_dict['obs'], 'action': obs_dict['action']}
+    mlm_input, mlm_label, mlm_weight = get_masked_input_and_labels(
+        obs_dict['obs']['use_def_opcode'][:, :_INSTR_COUNT]
+    )
+    inputs['use_def_opcode'] = mlm_input
+
+    labels = {
+        _CUR_STATE_KEY: obs_dict['obs_cur'],
+        _ACTION_KEY: tf.expand_dims(obs_dict['action'], axis=-1),
+        _NEXT_STATE_KEY: obs_dict['obs_roll'],
+        _NEXT_ACTION_KEY: tf.expand_dims(obs_dict['action_roll'], axis=-1),
+        _MLM_KEY: mlm_label,
+    }
+    mask = obs_dict['obs']['mask']
+    sample_weights = {
+        _CUR_STATE_KEY: mask,
+        _ACTION_KEY: mask,
+        _NEXT_STATE_KEY: None,
+        _NEXT_ACTION_KEY: None,
+        _MLM_KEY: mlm_weight,
+    }
+    return (inputs, labels, sample_weights)
+
+  def _file_dataset_fn(data_path):
+    return (
+        tf.data.Dataset.list_files(data_path)
+        .shuffle(files_buffer_size)
+        .interleave(
+            tf.data.TFRecordDataset,
+            cycle_length=num_readers,
+            block_length=1,
+        )
+        .filter(lambda string: tf.strings.length(string) > 0)
+        .map(parser_fn, num_parallel_calls=num_map_threads)
+        .map(_roll_experience, num_parallel_calls=num_map_threads)
+        .map(split_experience, num_parallel_calls=num_map_threads)
+        .map(preprocess_experience, num_parallel_calls=num_map_threads)
+        .unbatch()
+        .shuffle(shuffle_buffer_size)
+        .map(to_inputs_and_labels, num_parallel_calls=num_map_threads)
+        .batch(batch_size, drop_remainder=True)
+        .prefetch(tf.data.experimental.AUTOTUNE)
+    )
+
+  return _file_dataset_fn
+
+
+def get_strategy():
+  if FLAGS.tpu:
+    logging.info('Using TPU strategy.')
+    resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu=FLAGS.tpu)
+    tf.config.experimental_connect_to_cluster(resolver)
+    tf.tpu.experimental.initialize_tpu_system(resolver)
+    return tf.distribute.TPUStrategy(resolver)
+  logging.info('Using CPU strategy.')
+  return tf.distribute.get_strategy()
+
+
+class SaveEncoderCallback(tf.keras.callbacks.Callback):
+
+  def __init__(self, encoder, path):
+    self._encoder = encoder
+    self._path = path
+
+  def on_epoch_end(self, epoch, logs=None):
+    sm_path = os.path.join(self._path, f'epoch{epoch}')
+    tflite_path = os.path.join(sm_path, policy_saver.TFLITE_MODEL_NAME)
+    self._encoder.save(sm_path)
+    policy_saver.convert_saved_model(sm_path, tflite_path)
+
+
+def main(_):
+  gin.parse_config_files_and_bindings(
+      FLAGS.gin_files, bindings=None, skip_unknown=False
+  )
+  logging.info(gin.config_str())
+
+  problem_config = registry.get_configuration()
+  time_step_spec, action_spec = problem_config.get_signature_spec()
+  preprocessing_layer_creator = problem_config.get_preprocessing_layer_creator()
+  dataset_fn = create_dataset_fn(
+      time_step_spec,
+      action_spec,
+      batch_size=256,
+      shift=-1,
+      preprocessing_layer_creator=preprocessing_layer_creator,
+  )
+
+  # inputs, labels, weights = next(iter(dataset))
+  # import sys
+  # sys.exit()
+
+  strategy = get_strategy()
+  with strategy.scope():
+    logging.info('Creating model.')
+    encoder_network = regalloc_network.InstructionEncoderNetwork(
+        preprocessing_layer_creator
+    )
+    model = create_model(encoder_network)
+
+    logging.info('Compiling model.')
+
+    opt = tf.keras.optimizers.Adam(global_clipnorm=1.0)
+
+    model.compile(
+        optimizer=opt,
+        loss={
+            _CUR_STATE_KEY: 'mse',
+            _ACTION_KEY: tf.keras.losses.SparseCategoricalCrossentropy(
+                from_logits=True
+            ),
+            _NEXT_STATE_KEY: 'mse',
+            _NEXT_ACTION_KEY: tf.keras.losses.SparseCategoricalCrossentropy(
+                from_logits=True
+            ),
+            _MLM_KEY: tf.keras.losses.SparseCategoricalCrossentropy(
+                from_logits=True
+            ),
+        },
+        loss_weights={
+            _ACTION_KEY: 1,
+            _NEXT_STATE_KEY: 1,
+            _NEXT_ACTION_KEY: 1,
+            _MLM_KEY: 10,
+        },
+        metrics={
+            _CUR_STATE_KEY: [],
+            _ACTION_KEY: ['accuracy'],
+            _NEXT_STATE_KEY: [],
+            _NEXT_ACTION_KEY: ['accuracy'],
+            _MLM_KEY: ['accuracy'],
+        },
+    )
+
+  logging.info('Creating dataset.')
+  dataset = dataset_fn(FLAGS.trace)
+  tb = tf.keras.callbacks.TensorBoard(
+      log_dir=FLAGS.root_dir,
+      histogram_freq=0,
+      write_graph=True,
+      write_steps_per_second=True,
+      update_freq='batch',
+  )
+
+  policy_dir = os.path.join(FLAGS.root_dir, 'policy')
+  checkpoint = tf.keras.callbacks.ModelCheckpoint(
+      filepath=policy_dir, save_weights_only=False, save_freq=1024
+  )
+
+  encoder_dir = os.path.join(FLAGS.root_dir, 'encoder')
+  logging.info('Saving the encoder to %s', encoder_dir)
+  encoder_saver = SaveEncoderCallback(encoder=encoder_network, path=encoder_dir)
+
+  logging.info('Starting training.')
+  model.fit(dataset, epochs=8, callbacks=[tb, checkpoint, encoder_saver])
+
+  logging.info('Training complete.')
+
+
+if __name__ == '__main__':
+  app.run(main)