class_docstr: DatasetBuilder for rlu_control_suite dataset.
class_docstr_tok: ["DatasetBuilder","for","rlu_control_suite","dataset","."]
code: def tf_example_to_step_ds(
      self, tf_example: tf.train.Example
  ) -> Dict[str, Any]:
      feature_description = self._get_example_specs()
  
      data = tf.io.parse_single_example(tf_example, feature_description)
      episode_length = tf.size(data["discount"])
      is_first = tf.concat(
          [[True], [False] * tf.ones(episode_length - 1)], axis=0
      )
      is_last = tf.concat(
          [[False] * tf.ones(episode_length - 1), [True]], axis=0
      )
      is_terminal = [False] * tf.ones(episode_length, tf.int64)
  
      # The data is in RSA alignment, we realign it to SAR to comply with the
      # RLDS standard.
      discount = data["discount"][1:]
      if discount[-1] == 0.0:
          is_terminal = tf.concat(
              [[False] * tf.ones(episode_length - 1, tf.int64), [True]],
              axis=0,
          )
          # If the episode ends in a terminal state, in the last step only the
          # observation has valid information (the terminal state).
          discount = tf.concat([discount, [0.0]], axis=0)
      else:
          discount = tf.concat([discount, [1.0]], axis=0)
  
      reward = tf.concat([data["reward"][1:], [0.0]], axis=0)
  
      obs_prefix = "observation/"
      obs = {}
      for k in feature_description:
          if k.startswith(obs_prefix):
              new_k = k[len(obs_prefix) :]
              obs[new_k] = data[k]
      episode = {
          "steps": {
              "observation": obs,
              "action": data["action"],
              "reward": reward,
              "discount": discount,
              "is_first": is_first,
              "is_last": is_last,
              "is_terminal": is_terminal,
          },
          "episode_id": data["episode_id"],
          "timestamp": data["timestamp"],
      }
      return episode
  
file_path: datasets/tensorflow_datasets/rl_unplugged/rlu_control_suite/rlu_control_suite.py
imports: from __future__ import annotations
  import dataclasses
  from typing import Any, Dict, Optional
  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 7 for calling the RluControlSuite obj's underlying member method code with example usage: obj.tf_example_to_step_ds(tf_example) and returns: episode