# 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):


from typing import Callable, List
from tf_agents.specs import tensor_spec
from tf_agents.trajectories import trajectory
from tf_agents.typing import types
from compiler_opt.rl import constant
import tensorflow as tf
            dtype=time_step_spec.reward.dtype)
    sequence_features.update(
        _get_policy_info_parsing_dict(agent_name, action_spec))

    # pylint: enable=g-complex-comprehension
    with tf.name_scope('parse'):
      _, parsed_sequence = tf.io.parse_single_sequence_example(
          serialized_proto,
          context_features=context_features,
          sequence_features=sequence_features)
      # TODO(yundi): make the transformed reward configurable.
      action = parsed_sequence[action_spec.name]
      reward = tf.cast(parsed_sequence[time_step_spec.reward.name], tf.float32)

      policy_info = _process_parsed_sequence_and_get_policy_info(
          parsed_sequence, agent_name, action_spec)

      del parsed_sequence[time_step_spec.reward.name]
      del parsed_sequence[action_spec.name]
      full_trajectory = trajectory.from_episode(
          observation=parsed_sequence,
          action=action,
          policy_info=policy_info,
          reward=reward)
      return full_trajectory

  return _parser_fn


def create_sequence_example_dataset_fn(

# 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/data_reader.py
from typing import Callable, List
from tf_agents.specs import tensor_spec
from tf_agents.trajectories import trajectory
from tf_agents.typing import types
from compiler_opt.rl import constant
import tensorflow as tf

            dtype=time_step_spec.reward.dtype)
    sequence_features.update(
        _get_policy_info_parsing_dict(agent_name, action_spec))

    # pylint: enable=g-complex-comprehension
    with tf.name_scope('parse'):
      _, parsed_sequence = tf.io.parse_single_sequence_example(
          serialized_proto,
          context_features=context_features,
          sequence_features=sequence_features)
      # TODO(yundi): make the transformed reward configurable.
      action = parsed_sequence[action_spec.name]
      reward = tf.cast(parsed_sequence[time_step_spec.reward.name], tf.float32)

      policy_info = _process_parsed_sequence_and_get_policy_info(
          parsed_sequence, agent_name, action_spec)

      del parsed_sequence[time_step_spec.reward.name]
      del parsed_sequence[action_spec.name]
      full_trajectory = trajectory.from_episode(
          observation=parsed_sequence,
          action=action,
          policy_info=policy_info,
          reward=reward)
      return full_trajectory

  return _parser_fn


def create_sequence_example_dataset_fn(

    agent_name: constant.AgentName, time_step_spec: types.NestedSpec,