code: def build_episode(steps: Mapping[str, Any]) -> Dict[str, Any]:
      """Builds a full episode dict.
  
      Args:
        steps: A dict with each key corresponding to an episode of that variable.
  
      Returns:
       A dict with data specific to one episode.
      """
  
      # This is the base episode length without the end step, that step will get
      # added manually.  `[:]` notation is necessary to extract data from the hdf5
      # object.
      ep_length = steps["actions"][:].shape[0]
      episode = {}
      episode["is_first"] = np.append(
          True, np.zeros(ep_length, dtype=bool)
      )
  
      # The standard 'obs' needs to be extended with the last element from
      # 'next_obs' to reconstruct the full sequence.
      obs = tree.map_structure(np.array, dict(steps["obs"]))
      last_obs = tree.map_structure(
          lambda el: el[-1], dict(steps["next_obs"])
      )
      concat_obs = tree.map_structure(_concat_obs, obs, last_obs)
  
      actions = steps["actions"][:]
      dones = steps["dones"][:]
      episode["observation"] = concat_obs
      episode["action"] = np.append(
          actions, np.expand_dims(np.zeros(actions[0].shape), 0), axis=0
      )
      episode["reward"] = np.append(steps["rewards"][:], 0)
      episode["discount"] = np.append(np.ones(ep_length), 0)
      # Offset `dones` by one as it corresponds to the `next_obs`:
      # "done signal, equal to 1 if playing the corresponding action in the state
      # should terminate the episode".  Since playing another action may not end the
      # episode, it is actually representing whether the following state is
      # the terminal state.
      episode["is_terminal"] = np.append(False, dones).astype(bool)
      episode["is_last"] = np.append(
          np.zeros(ep_length, dtype=bool), True
      )
      states = steps["states"][:]
      episode["states"] = np.append(
          states, np.expand_dims(np.zeros(states[0].shape), 0), axis=0
      )
  
      return episode
  
imports: import dataclasses
  import enum
  from typing import Any, Dict, List, Mapping
  from etils import epath
  import h5py
  import numpy
  import tensorflow_datasets.public_api
  import tree
label_desc: Use image node_id 4 for calling a global function with example usage: build_episode(steps) and returns: episode