class_docstr: DatasetBuilder for RLU DMLab.
class_docstr_tok: ["DatasetBuilder","for","RLU","DMLab","."]
docstr: Create an episode from a TF example.
docstr_tok: ["Create","an","episode","from","a","TF","example","."]
code: def tf_example_to_step_ds(
      self, tf_example: tf.train.Example
  ) -> Dict[str, Any]:
      """Create an episode from a TF example."""
      episode_length = self.builder_config.episode_length
  
      # Parse tf.Example.
      def sequence_feature(shape, dtype=tf.float32):
          return tf.io.FixedLenFeature(
              shape=[episode_length] + shape, dtype=dtype
          )
  
      feature_description = {
          "episode_id": tf.io.FixedLenFeature([], tf.int64),
          "start_idx": tf.io.FixedLenFeature([], tf.int64),
          "episode_return": tf.io.FixedLenFeature([], tf.float32),
          "observations_pixels": sequence_feature([], tf.string),
          "observations_reward": sequence_feature([]),
          # actions are one-hot arrays.
          "observations_action": sequence_feature([15]),
          "actions": sequence_feature([], tf.int64),
          "rewards": sequence_feature([]),
          "discounted_rewards": sequence_feature([]),
          "discounts": sequence_feature([]),
      }
  
      data = tf.io.parse_single_example(tf_example, feature_description)
  
      # Dmlab observations are compressed with OpenCV imencode function. Encoding
      # with imencode and decoding with TF PNG decoder flips the order of the
      # color channels. Here we flip back the order color of channels from BGR to
      # RGB:
      pixels = tf.scan(
          fn=lambda _, png: tf.reshape(
              tf.io.decode_png(png), _PIXELS_SHAPE
          ),
          elems=data["observations_pixels"],
          initializer=tf.zeros(_PIXELS_SHAPE, dtype=tf.uint8),
      )
      pixels = tf.reverse(pixels, axis=[-1])
  
      episode = {
          # Episode Metadata
          "episode_id": data["episode_id"],
          "episode_return": data["episode_return"],
          "steps": {
              "observation": {
                  "pixels": pixels,
                  "last_action": tf.argmax(
                      data["observations_action"],
                      axis=1,
                      output_type=tf.int64,
                  ),
                  "last_reward": data["observations_reward"],
              },
              "action": data["actions"],
              "reward": data["rewards"],
              "discount": data["discounts"],
              "is_first": [True] + [False] * (episode_length - 1),
              "is_last": [False] * (episode_length - 1) + [True],
              "is_terminal": [False] * (episode_length),
          },
      }
      return episode
  
file_path: datasets/tensorflow_datasets/rl_unplugged/dmlab_dataset.py
imports: from __future__ import annotations
  import dataclasses
  from typing import Any, Dict
  import numpy
  from tensorflow_datasets.core.utils.lazy_imports_utils import tensorflow
  import tensorflow_datasets.public_api
  from tensorflow_datasets.rl_unplugged import rlu_common
label_desc: Use image node_id 6 for calling the DMLabDatasetBuilder obj's underlying member method code with example usage: obj.tf_example_to_step_ds(tf_example) and returns: episode, tf