Made the Reinforcement Learning components shareable in the tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk sample (#34)

* Made the Reinforcement Learning components shareable Also made couple of fixes. * Updated the sample pipeline * Replaced .after with data passing * Formatted the notebook using tensorflow_docs.tools.nbfmt * Fixed all linter errors * Updated the component URLs
2026-09-26 14:42:04 +00:00 · 2021-08-26 15:49:32 -04:00
parent 7cbe9dd499
commit 74dfd10f3d
8 changed files with 1272 additions and 459 deletions
@@ -0,0 +1,359 @@
+name: Generate movielens dataset for bigquery
+description: Generates BigQuery training data using a MovieLens simulation environment.
+inputs:
+- name: project_id
+  type: String
+  description: |-
+    GCP project ID. This is required because otherwise the BigQuery
+    client will use the ID of the tenant GCP project created as a result of
+    KFP, which doesn't have proper access to BigQuery.
+- {name: raw_data_path, type: String, description: Path to MovieLens 100K's "u.data"
+    file.}
+- {name: batch_size, type: Integer, description: Batch size of environment generated
+    quantities eg. rewards.}
+- name: rank_k
+  type: Integer
+  description: |-
+    Rank for matrix factorization in the MovieLens environment; also
+    the observation dimension.
+- {name: num_actions, type: Integer, description: Number of actions (movie items)
+    to choose from.}
+- {name: driver_steps, type: Integer, description: Number of steps to run per batch.}
+- {name: bigquery_tmp_file, type: String, description: Path to a JSON file containing
+    the training dataset.}
+- name: bigquery_dataset_id
+  type: String
+  description: |-
+    A string of the BigQuery dataset ID in the format of
+    "project.dataset".
+- {name: bigquery_location, type: String, description: A string of the BigQuery dataset
+    location.}
+- name: bigquery_table_id
+  type: String
+  description: |-
+    A string of the BigQuery table ID in the format of
+    "project.dataset.table".
+outputs:
+- {name: bigquery_dataset_id, type: String}
+- {name: bigquery_location, type: String}
+- {name: bigquery_table_id, type: String}
+implementation:
+  container:
+    image: tensorflow/tensorflow:2.5.0
+    command:
+    - sh
+    - -c
+    - (PIP_DISABLE_PIP_VERSION_CHECK=1 python3 -m pip install --quiet --no-warn-script-location
+      'google-cloud-bigquery==2.20.0' 'tensorflow==2.5.0' 'tf-agents==0.8.0' || PIP_DISABLE_PIP_VERSION_CHECK=1
+      python3 -m pip install --quiet --no-warn-script-location 'google-cloud-bigquery==2.20.0'
+      'tensorflow==2.5.0' 'tf-agents==0.8.0' --user) && "$0" "$@"
+    - sh
+    - -ec
+    - |
+      program_path=$(mktemp)
+      printf "%s" "$0" > "$program_path"
+      python3 -u "$program_path" "$@"
+    - |
+      def generate_movielens_dataset_for_bigquery(
+          project_id,
+          raw_data_path,
+          batch_size,
+          rank_k,
+          num_actions,
+          driver_steps,
+          bigquery_tmp_file,
+          bigquery_dataset_id,
+          bigquery_location,
+          bigquery_table_id
+      ):
+        """Generates BigQuery training data using a MovieLens simulation environment.
+
+        Serves as the Generator pipeline component:
+        1. Generates `trajectories.Trajectory` data by applying a random policy on
+          MovieLens simulation environment.
+        2. Converts `trajectories.Trajectory` data to JSON format.
+        3. Loads JSON-formatted data into BigQuery.
+
+        This function is to be built into a Kubeflow Pipelines (KFP) component. As a
+        result, this function must be entirely self-contained. This means that the
+        import statements and helper functions must reside within itself.
+
+        Args:
+          project_id: GCP project ID. This is required because otherwise the BigQuery
+            client will use the ID of the tenant GCP project created as a result of
+            KFP, which doesn't have proper access to BigQuery.
+          raw_data_path: Path to MovieLens 100K's "u.data" file.
+          batch_size: Batch size of environment generated quantities eg. rewards.
+          rank_k: Rank for matrix factorization in the MovieLens environment; also
+            the observation dimension.
+          num_actions: Number of actions (movie items) to choose from.
+          driver_steps: Number of steps to run per batch.
+          bigquery_tmp_file: Path to a JSON file containing the training dataset.
+          bigquery_dataset_id: A string of the BigQuery dataset ID in the format of
+            "project.dataset".
+          bigquery_location: A string of the BigQuery dataset location.
+          bigquery_table_id: A string of the BigQuery table ID in the format of
+            "project.dataset.table".
+
+        Returns:
+          A NamedTuple of (`bigquery_dataset_id`, `bigquery_location`,
+          `bigquery_table_id`).
+        """
+        # pylint: disable=g-import-not-at-top
+        import collections
+        import json
+        from typing import Any, Dict
+
+        from google.cloud import bigquery
+
+        from tf_agents import replay_buffers
+        from tf_agents import trajectories
+        from tf_agents.bandits.agents.examples.v2 import trainer
+        from tf_agents.bandits.environments import movielens_py_environment
+        from tf_agents.drivers import dynamic_step_driver
+        from tf_agents.environments import tf_py_environment
+        from tf_agents.policies import random_tf_policy
+
+        def generate_simulation_data(
+            raw_data_path,
+            batch_size,
+            rank_k,
+            num_actions,
+            driver_steps):
+          """Generates `trajectories.Trajectory` data from the simulation environment.
+
+          Constructs a MovieLens simulation environment, and generates a set of
+          `trajectories.Trajectory` data using a random policy.
+
+          Args:
+            raw_data_path: Path to MovieLens 100K's "u.data" file.
+            batch_size: Batch size of environment generated quantities eg. rewards.
+            rank_k: Rank for matrix factorization in the MovieLens environment; also
+              the observation dimension.
+            num_actions: Number of actions (movie items) to choose from.
+            driver_steps: Number of steps to run per batch.
+
+          Returns:
+            A replay buffer holding randomly generated`trajectories.Trajectory` data.
+          """
+          # Create MovieLens simulation environment.
+          env = movielens_py_environment.MovieLensPyEnvironment(
+              raw_data_path,
+              rank_k,
+              batch_size,
+              num_movies=num_actions,
+              csv_delimiter="\t")
+          environment = tf_py_environment.TFPyEnvironment(env)
+
+          # Define random policy for collecting data.
+          random_policy = random_tf_policy.RandomTFPolicy(
+              action_spec=environment.action_spec(),
+              time_step_spec=environment.time_step_spec())
+
+          # Use replay buffer and observers to keep track of Trajectory data.
+          data_spec = random_policy.trajectory_spec
+          replay_buffer = trainer.get_replay_buffer(data_spec, environment.batch_size,
+                                                    driver_steps)
+          observers = [replay_buffer.add_batch]
+
+          # Run driver to apply the random policy in the simulation environment.
+          driver = dynamic_step_driver.DynamicStepDriver(
+              env=environment,
+              policy=random_policy,
+              num_steps=driver_steps * environment.batch_size,
+              observers=observers)
+          driver.run()
+
+          return replay_buffer
+
+        def build_dict_from_trajectory(
+            trajectory):
+          """Builds a dict from `trajectory` data.
+
+          Args:
+            trajectory: A `trajectories.Trajectory` object.
+
+          Returns:
+            A dict holding the same data as `trajectory`.
+          """
+          trajectory_dict = {
+              "step_type": trajectory.step_type.numpy().tolist(),
+              "observation": [{
+                  "observation_batch": batch
+              } for batch in trajectory.observation.numpy().tolist()],
+              "action": trajectory.action.numpy().tolist(),
+              "policy_info": trajectory.policy_info,
+              "next_step_type": trajectory.next_step_type.numpy().tolist(),
+              "reward": trajectory.reward.numpy().tolist(),
+              "discount": trajectory.discount.numpy().tolist(),
+          }
+          return trajectory_dict
+
+        def write_replay_buffer_to_file(
+            replay_buffer,
+            batch_size,
+            dataset_file):
+          """Writes replay buffer data to a file, each JSON in one line.
+
+          Each `trajectories.Trajectory` object in `replay_buffer` will be written as
+          one line to the `dataset_file` in JSON format. I.e., the `dataset_file`
+          would be a newline-delimited JSON file.
+
+          Args:
+            replay_buffer: A `replay_buffers.TFUniformReplayBuffer` holding
+              `trajectories.Trajectory` objects.
+            batch_size: Batch size of environment generated quantities eg. rewards.
+            dataset_file: File path. Will be overwritten if already exists.
+          """
+          dataset = replay_buffer.as_dataset(sample_batch_size=batch_size)
+          dataset_size = replay_buffer.num_frames().numpy()
+
+          with open(dataset_file, "w") as f:
+            for example in dataset.take(count=dataset_size):
+              traj_dict = build_dict_from_trajectory(example[0])
+              f.write(json.dumps(traj_dict) + "\n")
+
+        def load_dataset_into_bigquery(
+            project_id,
+            dataset_file,
+            bigquery_dataset_id,
+            bigquery_location,
+            bigquery_table_id):
+          """Loads training dataset into BigQuery table.
+
+          Loads training dataset of `trajectories.Trajectory` in newline delimited
+          JSON into a BigQuery dataset and table, using a BigQuery client.
+
+          Args:
+            project_id: GCP project ID. This is required because otherwise the
+              BigQuery client will use the ID of the tenant GCP project created as a
+              result of KFP, which doesn't have proper access to BigQuery.
+            dataset_file: Path to a JSON file containing the training dataset.
+            bigquery_dataset_id: A string of the BigQuery dataset ID in the format of
+              "project.dataset".
+            bigquery_location: A string of the BigQuery dataset location.
+            bigquery_table_id: A string of the BigQuery table ID in the format of
+              "project.dataset.table".
+          """
+          # Construct a BigQuery client object.
+          client = bigquery.Client(project=project_id)
+
+          # Construct a full Dataset object to send to the API.
+          dataset = bigquery.Dataset(bigquery_dataset_id)
+
+          # Specify the geographic location where the dataset should reside.
+          dataset.location = bigquery_location
+
+          # Create the dataset, or get the dataset if it exists.
+          dataset = client.create_dataset(dataset, exists_ok=True, timeout=30)
+
+          job_config = bigquery.LoadJobConfig(
+              schema=[
+                  bigquery.SchemaField("step_type", "INT64", mode="REPEATED"),
+                  bigquery.SchemaField(
+                      "observation",
+                      "RECORD",
+                      mode="REPEATED",
+                      fields=[
+                          bigquery.SchemaField("observation_batch", "FLOAT64",
+                                               "REPEATED")
+                      ]),
+                  bigquery.SchemaField("action", "INT64", mode="REPEATED"),
+                  bigquery.SchemaField("policy_info", "FLOAT64", mode="REPEATED"),
+                  bigquery.SchemaField("next_step_type", "INT64", mode="REPEATED"),
+                  bigquery.SchemaField("reward", "FLOAT64", mode="REPEATED"),
+                  bigquery.SchemaField("discount", "FLOAT64", mode="REPEATED"),
+              ],
+              source_format=bigquery.SourceFormat.NEWLINE_DELIMITED_JSON,
+          )
+
+          with open(dataset_file, "rb") as source_file:
+            load_job = client.load_table_from_file(
+                source_file, bigquery_table_id, job_config=job_config)
+
+          load_job.result()  # Wait for the job to complete.
+
+        replay_buffer = generate_simulation_data(
+            raw_data_path=raw_data_path,
+            batch_size=batch_size,
+            rank_k=rank_k,
+            num_actions=num_actions,
+            driver_steps=driver_steps)
+
+        write_replay_buffer_to_file(
+            replay_buffer=replay_buffer,
+            batch_size=batch_size,
+            dataset_file=bigquery_tmp_file)
+
+        load_dataset_into_bigquery(project_id, bigquery_tmp_file, bigquery_dataset_id,
+                                   bigquery_location, bigquery_table_id)
+
+        outputs = collections.namedtuple(
+            "Outputs",
+            ["bigquery_dataset_id", "bigquery_location", "bigquery_table_id"])
+
+        return outputs(bigquery_dataset_id, bigquery_location, bigquery_table_id)
+
+      def _serialize_str(str_value: str) -> str:
+          if not isinstance(str_value, str):
+              raise TypeError('Value "{}" has type "{}" instead of str.'.format(str(str_value), str(type(str_value))))
+          return str_value
+
+      import argparse
+      _parser = argparse.ArgumentParser(prog='Generate movielens dataset for bigquery', description='Generates BigQuery training data using a MovieLens simulation environment.')
+      _parser.add_argument("--project-id", dest="project_id", type=str, required=True, default=argparse.SUPPRESS)
+      _parser.add_argument("--raw-data-path", dest="raw_data_path", type=str, required=True, default=argparse.SUPPRESS)
+      _parser.add_argument("--batch-size", dest="batch_size", type=int, required=True, default=argparse.SUPPRESS)
+      _parser.add_argument("--rank-k", dest="rank_k", type=int, required=True, default=argparse.SUPPRESS)
+      _parser.add_argument("--num-actions", dest="num_actions", type=int, required=True, default=argparse.SUPPRESS)
+      _parser.add_argument("--driver-steps", dest="driver_steps", type=int, required=True, default=argparse.SUPPRESS)
+      _parser.add_argument("--bigquery-tmp-file", dest="bigquery_tmp_file", type=str, required=True, default=argparse.SUPPRESS)
+      _parser.add_argument("--bigquery-dataset-id", dest="bigquery_dataset_id", type=str, required=True, default=argparse.SUPPRESS)
+      _parser.add_argument("--bigquery-location", dest="bigquery_location", type=str, required=True, default=argparse.SUPPRESS)
+      _parser.add_argument("--bigquery-table-id", dest="bigquery_table_id", type=str, required=True, default=argparse.SUPPRESS)
+      _parser.add_argument("----output-paths", dest="_output_paths", type=str, nargs=3)
+      _parsed_args = vars(_parser.parse_args())
+      _output_files = _parsed_args.pop("_output_paths", [])
+
+      _outputs = generate_movielens_dataset_for_bigquery(**_parsed_args)
+
+      _output_serializers = [
+          _serialize_str,
+          _serialize_str,
+          _serialize_str,
+
+      ]
+
+      import os
+      for idx, output_file in enumerate(_output_files):
+          try:
+              os.makedirs(os.path.dirname(output_file))
+          except OSError:
+              pass
+          with open(output_file, 'w') as f:
+              f.write(_output_serializers[idx](_outputs[idx]))
+    args:
+    - --project-id
+    - {inputValue: project_id}
+    - --raw-data-path
+    - {inputValue: raw_data_path}
+    - --batch-size
+    - {inputValue: batch_size}
+    - --rank-k
+    - {inputValue: rank_k}
+    - --num-actions
+    - {inputValue: num_actions}
+    - --driver-steps
+    - {inputValue: driver_steps}
+    - --bigquery-tmp-file
+    - {inputValue: bigquery_tmp_file}
+    - --bigquery-dataset-id
+    - {inputValue: bigquery_dataset_id}
+    - --bigquery-location
+    - {inputValue: bigquery_location}
+    - --bigquery-table-id
+    - {inputValue: bigquery_table_id}
+    - '----output-paths'
+    - {outputPath: bigquery_dataset_id}
+    - {outputPath: bigquery_location}
+    - {outputPath: bigquery_table_id}
@@ -259,3 +259,18 @@ def generate_movielens_dataset_for_bigquery(
      ["bigquery_dataset_id", "bigquery_location", "bigquery_table_id"])

  return outputs(bigquery_dataset_id, bigquery_location, bigquery_table_id)
+
+
+if __name__ == "__main__":
+  from kfp.components import create_component_from_func
+
+  generate_movielens_dataset_for_bigquery_op = create_component_from_func(
+    func=generate_movielens_dataset_for_bigquery,
+    base_image="tensorflow/tensorflow:2.5.0",
+    output_component_file="component.yaml",
+    packages_to_install=[
+      "google-cloud-bigquery==2.20.0",
+      "tensorflow==2.5.0",
+      "tf-agents==0.8.0",
+    ],
+  )
@@ -0,0 +1,231 @@
+name: Ingest bigquery dataset into tfrecord
+description: Ingests data from BigQuery, formats them and outputs TFRecord files.
+inputs:
+- name: project_id
+  type: String
+  description: |-
+    GCP project ID. This is required because otherwise the BigQuery
+    client will use the ID of the tenant GCP project created as a result of
+    KFP, which doesn't have proper access to BigQuery.
+- name: bigquery_table_id
+  type: String
+  description: |-
+    A string of the BigQuery table ID in the format of
+    "project.dataset.table".
+- {name: tfrecord_file, type: String, description: Path to file to write the ingestion
+    result TFRecords.}
+- {name: bigquery_max_rows, type: Integer, description: Optional; maximum number of
+    rows to ingest., optional: true}
+outputs:
+- {name: tfrecord_file, type: String}
+implementation:
+  container:
+    image: tensorflow/tensorflow:2.5.0
+    command:
+    - sh
+    - -c
+    - (PIP_DISABLE_PIP_VERSION_CHECK=1 python3 -m pip install --quiet --no-warn-script-location
+      'google-cloud-bigquery==2.20.0' 'tensorflow==2.5.0' || PIP_DISABLE_PIP_VERSION_CHECK=1
+      python3 -m pip install --quiet --no-warn-script-location 'google-cloud-bigquery==2.20.0'
+      'tensorflow==2.5.0' --user) && "$0" "$@"
+    - sh
+    - -ec
+    - |
+      program_path=$(mktemp)
+      printf "%s" "$0" > "$program_path"
+      python3 -u "$program_path" "$@"
+    - |
+      def ingest_bigquery_dataset_into_tfrecord(
+          project_id,
+          bigquery_table_id,
+          tfrecord_file,
+          bigquery_max_rows = None
+      ):
+        """Ingests data from BigQuery, formats them and outputs TFRecord files.
+
+        Serves as the Ingester pipeline component:
+        1. Reads data in BigQuery that contains 7 pieces of data: `step_type`,
+          `observation`, `action`, `policy_info`, `next_step_type`, `reward`,
+          `discount`.
+        2. Packages the data as `tf.train.Example` objects and outputs them as
+          TFRecord files.
+
+        This function is to be built into a Kubeflow Pipelines (KFP) component. As a
+        result, this function must be entirely self-contained. This means that the
+        import statements and helper functions must reside within itself.
+
+        Args:
+          project_id: GCP project ID. This is required because otherwise the BigQuery
+            client will use the ID of the tenant GCP project created as a result of
+            KFP, which doesn't have proper access to BigQuery.
+          bigquery_table_id: A string of the BigQuery table ID in the format of
+            "project.dataset.table".
+          tfrecord_file: Path to file to write the ingestion result TFRecords.
+          bigquery_max_rows: Optional; maximum number of rows to ingest.
+
+        Returns:
+          A NamedTuple of the path to the output TFRecord file.
+        """
+        # pylint: disable=g-import-not-at-top
+        import collections
+        from typing import Optional
+
+        from google.cloud import bigquery
+
+        import tensorflow as tf
+
+        def read_data_from_bigquery(
+            project_id,
+            bigquery_table_id,
+            bigquery_max_rows):
+          """Reads data from BigQuery at `bigquery_table_id` and creates an iterator.
+
+          The table contains 7 columns that form `trajectories.Trajectory` objects:
+          `step_type`, `observation`, `action`, `policy_info`, `next_step_type`,
+          `reward`, `discount`.
+
+          Args:
+            project_id: GCP project ID. This is required because otherwise the
+              BigQuery client will use the ID of the tenant GCP project created as a
+              result of KFP, which doesn't have proper access to BigQuery.
+            bigquery_table_id: A string of the BigQuery table ID in the format of
+              "project.dataset.table".
+            bigquery_max_rows: Optional; maximum number of rows to fetch.
+
+          Returns:
+            A row iterator over all data at `bigquery_table_id`.
+          """
+          # Construct a BigQuery client object.
+          client = bigquery.Client(project=project_id)
+
+          # Get dataset.
+          query_job = client.query(
+              f"""
+              SELECT * FROM {bigquery_table_id}
+              """
+          )
+          table = query_job.result(max_results=bigquery_max_rows)
+
+          return table
+
+        def _bytes_feature(tensor):
+          """Returns a `tf.train.Feature` with bytes from `tensor`.
+
+          Args:
+            tensor: A `tf.Tensor` object.
+
+          Returns:
+            A `tf.train.Feature` object containing bytes that represent the content of
+            `tensor`.
+          """
+          value = tf.io.serialize_tensor(tensor)
+          if isinstance(value, type(tf.constant(0))):
+            value = value.numpy()
+          return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))
+
+        def build_example(data_row):
+          """Builds a `tf.train.Example` from `data_row` content.
+
+          Args:
+            data_row: A `bigquery.table.Row` object that contains 7 pieces of data:
+              `step_type`, `observation`, `action`, `policy_info`, `next_step_type`,
+              `reward`, `discount`. Each piece of data except `observation` is a 1D
+              array; `observation` is a 1D array of `{"observation_batch": 1D array}.`
+
+          Returns:
+            A `tf.train.Example` object holding the same data as `data_row`.
+          """
+          feature = {
+              "step_type":
+                  _bytes_feature(data_row.get("step_type")),
+              "observation":
+                  _bytes_feature([
+                      observation["observation_batch"]
+                      for observation in data_row.get("observation")
+                  ]),
+              "action":
+                  _bytes_feature(data_row.get("action")),
+              "policy_info":
+                  _bytes_feature(data_row.get("policy_info")),
+              "next_step_type":
+                  _bytes_feature(data_row.get("next_step_type")),
+              "reward":
+                  _bytes_feature(data_row.get("reward")),
+              "discount":
+                  _bytes_feature(data_row.get("discount")),
+          }
+          example_proto = tf.train.Example(
+              features=tf.train.Features(feature=feature))
+          return example_proto
+
+        def write_tfrecords(
+            tfrecord_file,
+            table):
+          """Writes the row data in `table` into TFRecords in `tfrecord_file`.
+
+          Args:
+            tfrecord_file: Path to file to write the TFRecords.
+            table: A row iterator over all data to be written.
+          """
+          with tf.io.TFRecordWriter(tfrecord_file) as writer:
+            for data_row in table:
+              example = build_example(data_row)
+              writer.write(example.SerializeToString())
+
+        table = read_data_from_bigquery(
+            project_id=project_id,
+            bigquery_table_id=bigquery_table_id,
+            bigquery_max_rows=bigquery_max_rows)
+
+        write_tfrecords(tfrecord_file, table)
+
+        outputs = collections.namedtuple(
+            "Outputs",
+            ["tfrecord_file"])
+
+        return outputs(tfrecord_file)
+
+      def _serialize_str(str_value: str) -> str:
+          if not isinstance(str_value, str):
+              raise TypeError('Value "{}" has type "{}" instead of str.'.format(str(str_value), str(type(str_value))))
+          return str_value
+
+      import argparse
+      _parser = argparse.ArgumentParser(prog='Ingest bigquery dataset into tfrecord', description='Ingests data from BigQuery, formats them and outputs TFRecord files.')
+      _parser.add_argument("--project-id", dest="project_id", type=str, required=True, default=argparse.SUPPRESS)
+      _parser.add_argument("--bigquery-table-id", dest="bigquery_table_id", type=str, required=True, default=argparse.SUPPRESS)
+      _parser.add_argument("--tfrecord-file", dest="tfrecord_file", type=str, required=True, default=argparse.SUPPRESS)
+      _parser.add_argument("--bigquery-max-rows", dest="bigquery_max_rows", type=int, required=False, default=argparse.SUPPRESS)
+      _parser.add_argument("----output-paths", dest="_output_paths", type=str, nargs=1)
+      _parsed_args = vars(_parser.parse_args())
+      _output_files = _parsed_args.pop("_output_paths", [])
+
+      _outputs = ingest_bigquery_dataset_into_tfrecord(**_parsed_args)
+
+      _output_serializers = [
+          _serialize_str,
+
+      ]
+
+      import os
+      for idx, output_file in enumerate(_output_files):
+          try:
+              os.makedirs(os.path.dirname(output_file))
+          except OSError:
+              pass
+          with open(output_file, 'w') as f:
+              f.write(_output_serializers[idx](_outputs[idx]))
+    args:
+    - --project-id
+    - {inputValue: project_id}
+    - --bigquery-table-id
+    - {inputValue: bigquery_table_id}
+    - --tfrecord-file
+    - {inputValue: tfrecord_file}
+    - if:
+        cond: {isPresent: bigquery_max_rows}
+        then:
+        - --bigquery-max-rows
+        - {inputValue: bigquery_max_rows}
+    - '----output-paths'
+    - {outputPath: tfrecord_file}
@@ -167,3 +167,17 @@ def ingest_bigquery_dataset_into_tfrecord(
      ["tfrecord_file"])

  return outputs(tfrecord_file)
+
+
+if __name__ == "__main__":
+  from kfp.components import create_component_from_func
+
+  ingest_bigquery_dataset_into_tfrecord_op = create_component_from_func(
+    func=ingest_bigquery_dataset_into_tfrecord,
+    base_image="tensorflow/tensorflow:2.5.0",
+    output_component_file="component.yaml",
+    packages_to_install=[
+      "google-cloud-bigquery==2.20.0",
+      "tensorflow==2.5.0",
+    ],
+  )
@@ -0,0 +1,320 @@
+name: Train reinforcement learning policy
+description: Implements off-policy training for a policy on dataset of TFRecord files.
+inputs:
+- name: training_artifacts_dir
+  type: String
+  description: |-
+    Path to store the Trainer artifacts (trained
+    policy).
+- {name: tfrecord_file, type: String, description: Path to file to write the ingestion
+    result TFRecords.}
+- {name: num_epochs, type: Integer, description: Number of training epochs.}
+- name: rank_k
+  type: Integer
+  description: |-
+    Rank for matrix factorization in the MovieLens environment; also
+    the observation dimension.
+- {name: num_actions, type: Integer, description: Number of actions (movie items)
+    to choose from.}
+- {name: tikhonov_weight, type: Float, description: LinUCB Tikhonov regularization
+    weight of the Trainer.}
+- name: agent_alpha
+  type: Float
+  description: |-
+    LinUCB exploration parameter that multiplies the confidence
+    intervals of the Trainer.
+outputs:
+- {name: training_artifacts_dir, type: String}
+implementation:
+  container:
+    image: tensorflow/tensorflow:2.5.0
+    command:
+    - sh
+    - -c
+    - (PIP_DISABLE_PIP_VERSION_CHECK=1 python3 -m pip install --quiet --no-warn-script-location
+      'tensorflow==2.5.0' 'tf-agents==0.8.0' || PIP_DISABLE_PIP_VERSION_CHECK=1 python3
+      -m pip install --quiet --no-warn-script-location 'tensorflow==2.5.0' 'tf-agents==0.8.0'
+      --user) && "$0" "$@"
+    - sh
+    - -ec
+    - |
+      program_path=$(mktemp)
+      printf "%s" "$0" > "$program_path"
+      python3 -u "$program_path" "$@"
+    - |
+      def train_reinforcement_learning_policy(
+          training_artifacts_dir,
+          tfrecord_file,
+          num_epochs,
+          rank_k,
+          num_actions,
+          tikhonov_weight,
+          agent_alpha
+      ):
+        """Implements off-policy training for a policy on dataset of TFRecord files.
+
+        The Trainer's task is to submit a remote training job to Vertex AI, with the
+        training logic of a specified custom training container. The task will be
+        handled by: `kfp.v2.google.experimental.run_as_aiplatform_custom_job` (which
+        takes in the component made from this placeholder function)
+
+        This function is to be built into a Kubeflow Pipelines (KFP) component. As a
+        result, this function must be entirely self-contained. This means that the
+        import statements and helper functions must reside within itself.
+
+        Args:
+          training_artifacts_dir: Path to store the Trainer artifacts (trained
+            policy).
+          tfrecord_file: Path to file to write the ingestion result TFRecords.
+          num_epochs: Number of training epochs.
+          rank_k: Rank for matrix factorization in the MovieLens environment; also
+            the observation dimension.
+          num_actions: Number of actions (movie items) to choose from.
+          tikhonov_weight: LinUCB Tikhonov regularization weight of the Trainer.
+          agent_alpha: LinUCB exploration parameter that multiplies the confidence
+            intervals of the Trainer.
+
+        Returns:
+          A NamedTuple of (`training_artifacts_dir`).
+        """
+        # pylint: disable=g-import-not-at-top
+        import collections
+        from typing import Dict, List, NamedTuple  # pylint: disable=redefined-outer-name,reimported
+
+        import tensorflow as tf
+
+        from tf_agents import agents
+        from tf_agents import policies
+        from tf_agents import trajectories
+        from tf_agents.bandits.agents import lin_ucb_agent
+        from tf_agents.policies import policy_saver
+        from tf_agents.specs import tensor_spec
+
+        import logging
+
+        per_arm = False  # Using the non-per-arm version of the MovieLens environment.
+
+        # Mapping from feature name to serialized value
+        feature_description = {
+            "step_type": tf.io.FixedLenFeature((), tf.string),
+            "observation": tf.io.FixedLenFeature((), tf.string),
+            "action": tf.io.FixedLenFeature((), tf.string),
+            "policy_info": tf.io.FixedLenFeature((), tf.string),
+            "next_step_type": tf.io.FixedLenFeature((), tf.string),
+            "reward": tf.io.FixedLenFeature((), tf.string),
+            "discount": tf.io.FixedLenFeature((), tf.string),
+        }
+
+        def _parse_record(raw_record):
+          """Parses a serialized `tf.train.Example` proto.
+
+          Args:
+            raw_record: A serialized data record of a `tf.train.Example` proto.
+
+          Returns:
+            A dict mapping feature names to values as `tf.Tensor` objects of type
+            string containing serialized protos, following `feature_description`.
+          """
+          return tf.io.parse_single_example(raw_record, feature_description)
+
+        def build_trajectory(
+            parsed_record,
+            policy_info):
+          """Builds a `trajectories.Trajectory` object from `parsed_record`.
+
+          Args:
+            parsed_record: A dict mapping feature names to values as `tf.Tensor`
+              objects of type string containing serialized protos.
+            policy_info: Policy information specification.
+
+          Returns:
+            A `trajectories.Trajectory` object that contains values as de-serialized
+            `tf.Tensor` objects from `parsed_record`.
+          """
+          return trajectories.Trajectory(
+              step_type=tf.expand_dims(
+                  tf.io.parse_tensor(parsed_record["step_type"], out_type=tf.int32),
+                  axis=1),
+              observation=tf.expand_dims(
+                  tf.io.parse_tensor(
+                      parsed_record["observation"], out_type=tf.float32),
+                  axis=1),
+              action=tf.expand_dims(
+                  tf.io.parse_tensor(parsed_record["action"], out_type=tf.int32),
+                  axis=1),
+              policy_info=policy_info,
+              next_step_type=tf.expand_dims(
+                  tf.io.parse_tensor(
+                      parsed_record["next_step_type"], out_type=tf.int32),
+                  axis=1),
+              reward=tf.expand_dims(
+                  tf.io.parse_tensor(parsed_record["reward"], out_type=tf.float32),
+                  axis=1),
+              discount=tf.expand_dims(
+                  tf.io.parse_tensor(parsed_record["discount"], out_type=tf.float32),
+                  axis=1))
+
+        def train_policy_on_trajectory(
+            agent,
+            tfrecord_file,
+            num_epochs
+        ):
+          """Trains the policy in `agent` on the dataset of `tfrecord_file`.
+
+          Parses `tfrecord_file` as `tf.train.Example` objects, packages them into
+          `trajectories.Trajectory` objects, and trains the agent's policy on these
+          trajectory objects.
+
+          Args:
+            agent: A TF-Agents agent that carries the policy to train.
+            tfrecord_file: Path to the TFRecord file containing the training dataset.
+            num_epochs: Number of epochs to train the policy.
+
+          Returns:
+            A NamedTuple of (a trained TF-Agents policy, a dict mapping from
+            "epoch<i>" to lists of loss values produced at each training step).
+          """
+          raw_dataset = tf.data.TFRecordDataset([tfrecord_file])
+          parsed_dataset = raw_dataset.map(_parse_record)
+
+          train_loss = collections.defaultdict(list)
+          for epoch in range(num_epochs):
+            for parsed_record in parsed_dataset:
+              trajectory = build_trajectory(parsed_record, agent.policy.info_spec)
+              loss, _ = agent.train(trajectory)
+              train_loss[f"epoch{epoch + 1}"].append(loss.numpy())
+
+          train_outputs = collections.namedtuple(
+              "TrainOutputs",
+              ["policy", "train_loss"])
+          return train_outputs(agent.policy, train_loss)
+
+        def execute_training_and_save_policy(
+            training_artifacts_dir,
+            tfrecord_file,
+            num_epochs,
+            rank_k,
+            num_actions,
+            tikhonov_weight,
+            agent_alpha):
+          """Executes training for the policy and saves the policy.
+
+          Args:
+            training_artifacts_dir: Path to store the Trainer artifacts (trained
+              policy).
+            tfrecord_file: Path to file to write the ingestion result TFRecords.
+            num_epochs: Number of training epochs.
+            rank_k: Rank for matrix factorization in the MovieLens environment; also
+              the observation dimension.
+            num_actions: Number of actions (movie items) to choose from.
+            tikhonov_weight: LinUCB Tikhonov regularization weight of the Trainer.
+            agent_alpha: LinUCB exploration parameter that multiplies the confidence
+              intervals of the Trainer.
+          """
+          # Define time step and action specs for one batch.
+          time_step_spec = trajectories.TimeStep(
+              step_type=tensor_spec.TensorSpec(
+                  shape=(), dtype=tf.int32, name="step_type"),
+              reward=tensor_spec.TensorSpec(
+                  shape=(), dtype=tf.float32, name="reward"),
+              discount=tensor_spec.BoundedTensorSpec(
+                  shape=(), dtype=tf.float32, name="discount", minimum=0.,
+                  maximum=1.),
+              observation=tensor_spec.TensorSpec(
+                  shape=(rank_k,), dtype=tf.float32,
+                  name="observation"))
+
+          action_spec = tensor_spec.BoundedTensorSpec(
+              shape=(),
+              dtype=tf.int32,
+              name="action",
+              minimum=0,
+              maximum=num_actions - 1)
+
+          # Define RL agent/algorithm.
+          agent = lin_ucb_agent.LinearUCBAgent(
+              time_step_spec=time_step_spec,
+              action_spec=action_spec,
+              tikhonov_weight=tikhonov_weight,
+              alpha=agent_alpha,
+              dtype=tf.float32,
+              accepts_per_arm_features=per_arm)
+          agent.initialize()
+          logging.info("TimeStep Spec (for each batch):\n%s\n", agent.time_step_spec)
+          logging.info("Action Spec (for each batch):\n%s\n", agent.action_spec)
+
+          # Perform off-policy training.
+          policy, _ = train_policy_on_trajectory(
+              agent=agent,
+              tfrecord_file=tfrecord_file,
+              num_epochs=num_epochs)
+
+          # Save trained policy.
+          saver = policy_saver.PolicySaver(policy)
+          saver.save(training_artifacts_dir)
+
+        execute_training_and_save_policy(
+            training_artifacts_dir=training_artifacts_dir,
+            tfrecord_file=tfrecord_file,
+            num_epochs=num_epochs,
+            rank_k=rank_k,
+            num_actions=num_actions,
+            tikhonov_weight=tikhonov_weight,
+            agent_alpha=agent_alpha)
+
+        outputs = collections.namedtuple(
+            "Outputs",
+            ["training_artifacts_dir"])
+
+        return outputs(training_artifacts_dir)
+
+      def _serialize_str(str_value: str) -> str:
+          if not isinstance(str_value, str):
+              raise TypeError('Value "{}" has type "{}" instead of str.'.format(str(str_value), str(type(str_value))))
+          return str_value
+
+      import argparse
+      _parser = argparse.ArgumentParser(prog='Train reinforcement learning policy', description='Implements off-policy training for a policy on dataset of TFRecord files.')
+      _parser.add_argument("--training-artifacts-dir", dest="training_artifacts_dir", type=str, required=True, default=argparse.SUPPRESS)
+      _parser.add_argument("--tfrecord-file", dest="tfrecord_file", type=str, required=True, default=argparse.SUPPRESS)
+      _parser.add_argument("--num-epochs", dest="num_epochs", type=int, required=True, default=argparse.SUPPRESS)
+      _parser.add_argument("--rank-k", dest="rank_k", type=int, required=True, default=argparse.SUPPRESS)
+      _parser.add_argument("--num-actions", dest="num_actions", type=int, required=True, default=argparse.SUPPRESS)
+      _parser.add_argument("--tikhonov-weight", dest="tikhonov_weight", type=float, required=True, default=argparse.SUPPRESS)
+      _parser.add_argument("--agent-alpha", dest="agent_alpha", type=float, required=True, default=argparse.SUPPRESS)
+      _parser.add_argument("----output-paths", dest="_output_paths", type=str, nargs=1)
+      _parsed_args = vars(_parser.parse_args())
+      _output_files = _parsed_args.pop("_output_paths", [])
+
+      _outputs = train_reinforcement_learning_policy(**_parsed_args)
+
+      _output_serializers = [
+          _serialize_str,
+
+      ]
+
+      import os
+      for idx, output_file in enumerate(_output_files):
+          try:
+              os.makedirs(os.path.dirname(output_file))
+          except OSError:
+              pass
+          with open(output_file, 'w') as f:
+              f.write(_output_serializers[idx](_outputs[idx]))
+    args:
+    - --training-artifacts-dir
+    - {inputValue: training_artifacts_dir}
+    - --tfrecord-file
+    - {inputValue: tfrecord_file}
+    - --num-epochs
+    - {inputValue: num_epochs}
+    - --rank-k
+    - {inputValue: rank_k}
+    - --num-actions
+    - {inputValue: num_actions}
+    - --tikhonov-weight
+    - {inputValue: tikhonov_weight}
+    - --agent-alpha
+    - {inputValue: agent_alpha}
+    - '----output-paths'
+    - {outputPath: training_artifacts_dir}
@@ -88,7 +88,7 @@ class TestTrainerComponent(unittest.TestCase):

  def test_training_op_execute_all_steps(self):
    """Tests that training_op executes all steps."""
-    trainer_component.training_op(
+    trainer_component.train_reinforcement_learning_policy(
        training_artifacts_dir=TRAINING_ARTIFACTS_DIR,
        tfrecord_file=TFRECORD_FILE,
        num_epochs=NUM_EPOCHS,
@@ -117,7 +117,7 @@ class TestTrainerComponent(unittest.TestCase):

  def test_given_zero_epochs_training_op_execute_no_training(self):
    """Tests that training_op executes zero training with zero num_epochs."""
-    trainer_component.training_op(
+    trainer_component.train_reinforcement_learning_policy(
        training_artifacts_dir=TRAINING_ARTIFACTS_DIR,
        tfrecord_file=TFRECORD_FILE,
        num_epochs=0,
@@ -131,7 +131,7 @@ class TestTrainerComponent(unittest.TestCase):
  def test_given_negative_epochs_training_op_execute_no_training(self):
    """Tests that training_op executes zero training with negative num_epochs.
    """
-    trainer_component.training_op(
+    trainer_component.train_reinforcement_learning_policy(
        training_artifacts_dir=TRAINING_ARTIFACTS_DIR,
        tfrecord_file=TFRECORD_FILE,
        num_epochs=-1,
@@ -145,7 +145,7 @@ class TestTrainerComponent(unittest.TestCase):
  def test_given_float_epochs_training_op_raise_exception(self):
    """Tests that training_op raises an exception for float num_epochs."""
    with self.assertRaises(TypeError):
-      trainer_component.training_op(
+      trainer_component.train_reinforcement_learning_policy(
          training_artifacts_dir=TRAINING_ARTIFACTS_DIR,
          tfrecord_file=TFRECORD_FILE,
          num_epochs=0.5,
@@ -16,10 +16,8 @@
 # Import for the function return value type.
 from typing import NamedTuple  # pylint: disable=unused-import

-from kfp.v2 import components

-
-def training_op(
+def train_reinforcement_learning_policy(
    training_artifacts_dir: str,
    tfrecord_file: str,
    num_epochs: int,
@@ -28,7 +26,7 @@ def training_op(
    tikhonov_weight: float,
    agent_alpha: float
 ) -> NamedTuple("Outputs", [
-    ("training_artifacts_dir", components.OutputPath),
+    ("training_artifacts_dir", str),
 ]):
  """Implements off-policy training for a policy on dataset of TFRecord files.

@@ -249,3 +247,17 @@ def training_op(
      ["training_artifacts_dir"])

  return outputs(training_artifacts_dir)
+
+
+if __name__ == "__main__":
+  from kfp.components import create_component_from_func
+
+  train_reinforcement_learning_policy_op = create_component_from_func(
+    func=train_reinforcement_learning_policy,
+    base_image="tensorflow/tensorflow:2.5.0",
+    output_component_file="component.yaml",
+    packages_to_install=[
+      "tensorflow==2.5.0",
+      "tf-agents==0.8.0",
+    ],
+  )