# Path: compiler_opt/rl/trainer.py
# _INLINING_DEFAULT_KEY = 'inlining_default'
# class Trainer(object):
#   def __init__(
#       self,
#       root_dir,
#       agent,
#       random_network_distillation=None,
#       warmstart_policy_dir=None,
#       # Params for summaries and logging
#       checkpoint_interval=10000,
#       log_interval=100,
#       summary_interval=1000,
#       summaries_flush_secs=10):
#   def _initialize_metrics(self):
#   def _update_metrics(self, experience, monitor_dict):
#   def _reset_metrics(self):
#   def _log_experiment(self, loss):
#   def _save_checkpoint(self):
#   def global_step_numpy(self):
#   def train(self, dataset_iter, monitor_dict, num_iterations):


from absl.testing.absltest import mock
from tf_agents.agents.behavioral_cloning import behavioral_cloning_agent
from tf_agents.networks import q_rnn_network
from tf_agents.specs import tensor_spec
from tf_agents.trajectories import time_step
from tf_agents.trajectories import trajectory
from compiler_opt.rl import trainer
import tensorflow as tf
        minimum=0,
        maximum=1,
        name='inlining_decision')
    self._network = q_rnn_network.QRnnNetwork(
        input_tensor_spec=self._time_step_spec.observation,
        action_spec=self._action_spec,
        lstm_size=(40,),
        preprocessing_layers={
            'inlining_default': tf.keras.layers.Lambda(lambda x: x)
        })
    super(TrainerTest, self).setUp()

  def test_trainer_initialization(self):
    test_agent = behavioral_cloning_agent.BehavioralCloningAgent(
        self._time_step_spec,
        self._action_spec,
        self._network,
        tf.compat.v1.train.AdamOptimizer(),
        num_outer_dims=2)
    test_trainer = trainer.Trainer(
        root_dir=self.get_temp_dir(), agent=test_agent)
    self.assertEqual(0, test_trainer._global_step.numpy())

  def test_training(self):
    test_agent = behavioral_cloning_agent.BehavioralCloningAgent(
        self._time_step_spec,
        self._action_spec,
        self._network,
        tf.compat.v1.train.AdamOptimizer(),
        num_outer_dims=2)

# Path: compiler_opt/rl/trainer.py
# _INLINING_DEFAULT_KEY = 'inlining_default'
# class Trainer(object):
#   def __init__(
#       self,
#       root_dir,
#       agent,
#       random_network_distillation=None,
#       warmstart_policy_dir=None,
#       # Params for summaries and logging
#       checkpoint_interval=10000,
#       log_interval=100,
#       summary_interval=1000,
#       summaries_flush_secs=10):
#   def _initialize_metrics(self):
#   def _update_metrics(self, experience, monitor_dict):
#   def _reset_metrics(self):
#   def _log_experiment(self, loss):
#   def _save_checkpoint(self):
#   def global_step_numpy(self):
#   def train(self, dataset_iter, monitor_dict, num_iterations):


# Path: compiler_opt/rl/trainer_test.py
from absl.testing.absltest import mock
from tf_agents.agents.behavioral_cloning import behavioral_cloning_agent
from tf_agents.networks import q_rnn_network
from tf_agents.specs import tensor_spec
from tf_agents.trajectories import time_step
from tf_agents.trajectories import trajectory
from compiler_opt.rl import trainer
import tensorflow as tf

        minimum=0,
        maximum=1,
        name='inlining_decision')
    self._network = q_rnn_network.QRnnNetwork(
        input_tensor_spec=self._time_step_spec.observation,
        action_spec=self._action_spec,
        lstm_size=(40,),
        preprocessing_layers={
            'inlining_default': tf.keras.layers.Lambda(lambda x: x)
        })
    super(TrainerTest, self).setUp()

  def test_trainer_initialization(self):
    test_agent = behavioral_cloning_agent.BehavioralCloningAgent(
        self._time_step_spec,
        self._action_spec,
        self._network,
        tf.compat.v1.train.AdamOptimizer(),
        num_outer_dims=2)
    test_trainer = trainer.Trainer(
        root_dir=self.get_temp_dir(), agent=test_agent)
    self.assertEqual(0, test_trainer._global_step.numpy())

  def test_training(self):
    test_agent = behavioral_cloning_agent.BehavioralCloningAgent(
        self._time_step_spec,
        self._action_spec,
        self._network,
        tf.compat.v1.train.AdamOptimizer(),
        num_outer_dims=2)

    test_trainer = trainer.Trainer(