# Path: compiler_opt/rl/agent_creators.py
# def _create_behavioral_cloning_agent(
#     time_step_spec: types.NestedTensorSpec, action_spec: types.NestedTensorSpec,
#     preprocessing_layers: types.NestedLayer,
#     policy_network: types.Network) -> tfa.agents.TFAgent:
# def _create_dqn_agent(
#     time_step_spec: types.NestedTensorSpec, action_spec: types.NestedTensorSpec,
#     preprocessing_layers: types.NestedLayer,
#     policy_network: types.Network) -> tfa.agents.TFAgent:
# def _create_ppo_agent(
#     time_step_spec: types.NestedTensorSpec, action_spec: types.NestedTensorSpec,
#     preprocessing_layers: types.NestedLayer,
#     policy_network: types.Network) -> tfa.agents.TFAgent:
# def create_agent(
#     agent_name: constant.AgentName,
#     time_step_spec: types.NestedTensorSpec,
#     action_spec: types.NestedTensorSpec,
#     preprocessing_layer_creator: Callable[[types.TensorSpec],
#                                           tf.keras.layers.Layer],
#     policy_network: types.Network):
# 
# Path: compiler_opt/rl/constant.py
# BASE_DIR = 'compiler_opt/rl'
# BASE_MODULE_DIR = 'compiler_opt.rl'
# DELTA = 0.01
#   BEHAVIORAL_CLONE = 0
#   DQN = 1
#   PPO = 2
# class AgentName(enum.Enum):


import gin
import tensorflow as tf
from tf_agents.agents.behavioral_cloning import behavioral_cloning_agent
from tf_agents.agents.dqn import dqn_agent
from tf_agents.agents.ppo import ppo_agent
from tf_agents.networks import actor_distribution_network
from tf_agents.networks import q_network
from tf_agents.specs import tensor_spec
from tf_agents.trajectories import time_step
from compiler_opt.rl import agent_creators
from compiler_opt.rl import constant
        maximum=1,
        name='inlining_decision')
    super(AgentCreatorsTest, self).setUp()

  def test_create_behavioral_cloning_agent(self):
    gin.bind_parameter('create_agent.policy_network', q_network.QNetwork)
    gin.bind_parameter('BehavioralCloningAgent.optimizer',
                       tf.compat.v1.train.AdamOptimizer())
    tf_agent = agent_creators.create_agent(
        agent_name=constant.AgentName.BEHAVIORAL_CLONE,
        time_step_spec=self._time_step_spec,
        action_spec=self._action_spec,
        preprocessing_layer_creator=_observation_processing_layer)
    self.assertIsInstance(tf_agent,
                          behavioral_cloning_agent.BehavioralCloningAgent)

  def test_create_dqn_agent(self):
    gin.bind_parameter('create_agent.policy_network', q_network.QNetwork)
    gin.bind_parameter('DqnAgent.optimizer', tf.compat.v1.train.AdamOptimizer())
    tf_agent = agent_creators.create_agent(
        agent_name=constant.AgentName.DQN,
        time_step_spec=self._time_step_spec,
        action_spec=self._action_spec,
        preprocessing_layer_creator=_observation_processing_layer)
    self.assertIsInstance(tf_agent, dqn_agent.DqnAgent)

  def test_create_ppo_agent(self):
    gin.bind_parameter('create_agent.policy_network',
                       actor_distribution_network.ActorDistributionNetwork)
    gin.bind_parameter('PPOAgent.optimizer', tf.compat.v1.train.AdamOptimizer())

# Path: compiler_opt/rl/agent_creators.py
# def _create_behavioral_cloning_agent(
#     time_step_spec: types.NestedTensorSpec, action_spec: types.NestedTensorSpec,
#     preprocessing_layers: types.NestedLayer,
#     policy_network: types.Network) -> tfa.agents.TFAgent:
# def _create_dqn_agent(
#     time_step_spec: types.NestedTensorSpec, action_spec: types.NestedTensorSpec,
#     preprocessing_layers: types.NestedLayer,
#     policy_network: types.Network) -> tfa.agents.TFAgent:
# def _create_ppo_agent(
#     time_step_spec: types.NestedTensorSpec, action_spec: types.NestedTensorSpec,
#     preprocessing_layers: types.NestedLayer,
#     policy_network: types.Network) -> tfa.agents.TFAgent:
# def create_agent(
#     agent_name: constant.AgentName,
#     time_step_spec: types.NestedTensorSpec,
#     action_spec: types.NestedTensorSpec,
#     preprocessing_layer_creator: Callable[[types.TensorSpec],
#                                           tf.keras.layers.Layer],
#     policy_network: types.Network):
# 
# Path: compiler_opt/rl/constant.py
# BASE_DIR = 'compiler_opt/rl'
# BASE_MODULE_DIR = 'compiler_opt.rl'
# DELTA = 0.01
#   BEHAVIORAL_CLONE = 0
#   DQN = 1
#   PPO = 2
# class AgentName(enum.Enum):


# Path: compiler_opt/rl/agent_creators_test.py
import gin
import tensorflow as tf
from tf_agents.agents.behavioral_cloning import behavioral_cloning_agent
from tf_agents.agents.dqn import dqn_agent
from tf_agents.agents.ppo import ppo_agent
from tf_agents.networks import actor_distribution_network
from tf_agents.networks import q_network
from tf_agents.specs import tensor_spec
from tf_agents.trajectories import time_step
from compiler_opt.rl import agent_creators
from compiler_opt.rl import constant

        maximum=1,
        name='inlining_decision')
    super(AgentCreatorsTest, self).setUp()

  def test_create_behavioral_cloning_agent(self):
    gin.bind_parameter('create_agent.policy_network', q_network.QNetwork)
    gin.bind_parameter('BehavioralCloningAgent.optimizer',
                       tf.compat.v1.train.AdamOptimizer())
    tf_agent = agent_creators.create_agent(
        agent_name=constant.AgentName.BEHAVIORAL_CLONE,
        time_step_spec=self._time_step_spec,
        action_spec=self._action_spec,
        preprocessing_layer_creator=_observation_processing_layer)
    self.assertIsInstance(tf_agent,
                          behavioral_cloning_agent.BehavioralCloningAgent)

  def test_create_dqn_agent(self):
    gin.bind_parameter('create_agent.policy_network', q_network.QNetwork)
    gin.bind_parameter('DqnAgent.optimizer', tf.compat.v1.train.AdamOptimizer())
    tf_agent = agent_creators.create_agent(
        agent_name=constant.AgentName.DQN,
        time_step_spec=self._time_step_spec,
        action_spec=self._action_spec,
        preprocessing_layer_creator=_observation_processing_layer)
    self.assertIsInstance(tf_agent, dqn_agent.DqnAgent)

  def test_create_ppo_agent(self):
    gin.bind_parameter('create_agent.policy_network',
                       actor_distribution_network.ActorDistributionNetwork)
    gin.bind_parameter('PPOAgent.optimizer', tf.compat.v1.train.AdamOptimizer())

    tf_agent = agent_creators.create_agent(