def generate_json_state(episode: dict, dataset_name: str):
    """
    Generate the json dict and state for an episode.
    """
    if isinstance(dataset_name, tf.Tensor):
        dataset_name = dataset_name.numpy().decode('utf-8')

    # Process each step in the episode
    episode['steps'] = episode['steps'].map(
        globals()[dataset_name].process_step,
    )

    if dataset_name == "agilex":
        return _generate_json_state_agilex(episode, dataset_name)
    
    if dataset_name in DATASET_NAMES_NO_STATE:
        return _generate_json_state_nostate_ds(episode, dataset_name)

    return _generate_json_state(episode, dataset_name)

    Generate the json dict and state for an episode.
    
policy/RDT/data/preprocess.py