# 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/constant_value_network.py
# class ConstantValueNetwork(network.Network):
#   def __init__(self, input_tensor_spec, constant_output_val=0, name=None):
#   def call(self, observation, step_type=None, network_state=(), training=False):


from typing import Callable
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.typing import types
from compiler_opt.rl import constant
from compiler_opt.rl import constant_value_network
import gin
import tensorflow as tf
import tf_agents as tfa
    policy_network: types.Network):
  """Creates a tfa.agents.TFAgent object.

  Args:
    agent_name: AgentName, enum type of the agent to create.
    time_step_spec: A `TimeStep` spec of the expected time_steps.
    action_spec: A nest of BoundedTensorSpec representing the actions.
    preprocessing_layer_creator: A callable returns feature processing layer
      given observation_spec.
    policy_network: A tf_agents.networks.Network class.

  Returns:
    tf_agent: A tfa.agents.TFAgent object.

  Raises:
    ValueError: If `agent_name` is not in supported list.
  """
  assert policy_network is not None
  assert agent_name is not None

  preprocessing_layers = tf.nest.map_structure(
      preprocessing_layer_creator, time_step_spec.observation)

  if agent_name == constant.AgentName.BEHAVIORAL_CLONE:
    return _create_behavioral_cloning_agent(time_step_spec, action_spec,
                                            preprocessing_layers,
                                            policy_network)
  elif agent_name == constant.AgentName.DQN:
    return _create_dqn_agent(time_step_spec, action_spec,
                             preprocessing_layers, policy_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/constant_value_network.py
# class ConstantValueNetwork(network.Network):
#   def __init__(self, input_tensor_spec, constant_output_val=0, name=None):
#   def call(self, observation, step_type=None, network_state=(), training=False):


# Path: compiler_opt/rl/agent_creators.py
from typing import Callable
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.typing import types
from compiler_opt.rl import constant
from compiler_opt.rl import constant_value_network
import gin
import tensorflow as tf
import tf_agents as tfa

    policy_network: types.Network):
  """Creates a tfa.agents.TFAgent object.

  Args:
    agent_name: AgentName, enum type of the agent to create.
    time_step_spec: A `TimeStep` spec of the expected time_steps.
    action_spec: A nest of BoundedTensorSpec representing the actions.
    preprocessing_layer_creator: A callable returns feature processing layer
      given observation_spec.
    policy_network: A tf_agents.networks.Network class.

  Returns:
    tf_agent: A tfa.agents.TFAgent object.

  Raises:
    ValueError: If `agent_name` is not in supported list.
  """
  assert policy_network is not None
  assert agent_name is not None

  preprocessing_layers = tf.nest.map_structure(
      preprocessing_layer_creator, time_step_spec.observation)

  if agent_name == constant.AgentName.BEHAVIORAL_CLONE:
    return _create_behavioral_cloning_agent(time_step_spec, action_spec,
                                            preprocessing_layers,
                                            policy_network)
  elif agent_name == constant.AgentName.DQN:
    return _create_dqn_agent(time_step_spec, action_spec,
                             preprocessing_layers, policy_network)

  elif agent_name == constant.AgentName.PPO: