# 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
# def _get_policy_info_parsing_dict(agent_name, action_spec):
# def _process_parsed_sequence_and_get_policy_info(parsed_sequence, agent_name,
#                                                  action_spec):
# def create_parser_fn(
#     agent_name: constant.AgentName, time_step_spec: types.NestedSpec,
#     action_spec: types.NestedSpec) -> Callable[[str], trajectory.Trajectory]:
#   def _parser_fn(serialized_proto):
# def create_sequence_example_dataset_fn(
#     agent_name: constant.AgentName, time_step_spec: types.NestedSpec,
#     action_spec: types.NestedSpec, batch_size: int, train_sequence_length: int
# ) -> Callable[[List[str]], tf.data.Dataset]:
#   def _sequence_example_dataset_fn(sequence_examples):
# def create_file_dataset_fn(
#     agent_name: constant.AgentName, time_step_spec: types.NestedSpec,
#     action_spec: types.NestedSpec, batch_size: int, train_sequence_length: int,
#     input_dataset) -> Callable[[List[str]], tf.data.Dataset]:
#   def _file_dataset_fn(data_path):
# def create_tfrecord_dataset_fn(
#     agent_name: constant.AgentName, time_step_spec: types.NestedSpec,
#     action_spec: types.NestedSpec, batch_size: int, train_sequence_length: int
# ) -> Callable[[List[str]], tf.data.Dataset]:


import os
import tensorflow as tf
from absl.testing import parameterized
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 constant
from compiler_opt.rl import data_reader
        batch_size=2,
        train_sequence_length=3)
    data_iterator = iter(dataset_fn(data_source).repeat())

    experience = next(data_iterator)
    self.assertIsInstance(experience, trajectory.Trajectory)
    self.assertAllEqual([2, 3], experience.step_type.shape)
    self.assertCountEqual(['feature_key'],
                          experience.observation.keys())
    self.assertAllEqual([[1, 1, 1], [1, 1, 1]],
                        experience.observation['feature_key'])
    self.assertAllEqual([[0, 0, 0], [0, 0, 0]], experience.action)
    self.assertEmpty(experience.policy_info)
    self.assertAllEqual([2, 3], experience.next_step_type.shape)
    self.assertAllClose([[2.3, 2.3, 2.3], [2.3, 2.3, 2.3]], experience.reward)
    self.assertAllEqual([[1, 1, 1], [1, 1, 1]], experience.discount)

  @parameterized.named_parameters(
      ('SequenceExampleDatasetFn',
       data_reader.create_sequence_example_dataset_fn),
      ('TFRecordDatasetFn', data_reader.create_tfrecord_dataset_fn))
  def test_ppo_policy_info_discrete(self, test_fn):
    self._agent_name = constant.AgentName.PPO

    example = _define_sequence_example(
        self._agent_name, is_action_discrete=True)

    data_source = None
    if test_fn == data_reader.create_sequence_example_dataset_fn:
      data_source = [example.SerializeToString() for _ in range(100)]

# 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
# def _get_policy_info_parsing_dict(agent_name, action_spec):
# def _process_parsed_sequence_and_get_policy_info(parsed_sequence, agent_name,
#                                                  action_spec):
# def create_parser_fn(
#     agent_name: constant.AgentName, time_step_spec: types.NestedSpec,
#     action_spec: types.NestedSpec) -> Callable[[str], trajectory.Trajectory]:
#   def _parser_fn(serialized_proto):
# def create_sequence_example_dataset_fn(
#     agent_name: constant.AgentName, time_step_spec: types.NestedSpec,
#     action_spec: types.NestedSpec, batch_size: int, train_sequence_length: int
# ) -> Callable[[List[str]], tf.data.Dataset]:
#   def _sequence_example_dataset_fn(sequence_examples):
# def create_file_dataset_fn(
#     agent_name: constant.AgentName, time_step_spec: types.NestedSpec,
#     action_spec: types.NestedSpec, batch_size: int, train_sequence_length: int,
#     input_dataset) -> Callable[[List[str]], tf.data.Dataset]:
#   def _file_dataset_fn(data_path):
# def create_tfrecord_dataset_fn(
#     agent_name: constant.AgentName, time_step_spec: types.NestedSpec,
#     action_spec: types.NestedSpec, batch_size: int, train_sequence_length: int
# ) -> Callable[[List[str]], tf.data.Dataset]:


# Path: compiler_opt/rl/data_reader_test.py
import os
import tensorflow as tf
from absl.testing import parameterized
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 constant
from compiler_opt.rl import data_reader

        batch_size=2,
        train_sequence_length=3)
    data_iterator = iter(dataset_fn(data_source).repeat())

    experience = next(data_iterator)
    self.assertIsInstance(experience, trajectory.Trajectory)
    self.assertAllEqual([2, 3], experience.step_type.shape)
    self.assertCountEqual(['feature_key'],
                          experience.observation.keys())
    self.assertAllEqual([[1, 1, 1], [1, 1, 1]],
                        experience.observation['feature_key'])
    self.assertAllEqual([[0, 0, 0], [0, 0, 0]], experience.action)
    self.assertEmpty(experience.policy_info)
    self.assertAllEqual([2, 3], experience.next_step_type.shape)
    self.assertAllClose([[2.3, 2.3, 2.3], [2.3, 2.3, 2.3]], experience.reward)
    self.assertAllEqual([[1, 1, 1], [1, 1, 1]], experience.discount)

  @parameterized.named_parameters(
      ('SequenceExampleDatasetFn',
       data_reader.create_sequence_example_dataset_fn),
      ('TFRecordDatasetFn', data_reader.create_tfrecord_dataset_fn))
  def test_ppo_policy_info_discrete(self, test_fn):
    self._agent_name = constant.AgentName.PPO

    example = _define_sequence_example(
        self._agent_name, is_action_discrete=True)

    data_source = None
    if test_fn == data_reader.create_sequence_example_dataset_fn:
      data_source = [example.SerializeToString() for _ in range(100)]

    elif test_fn == data_reader.create_tfrecord_dataset_fn: