from .dataset_transform import *

openx_config = dict(
    bridge_dataset=dict(
        path_to_dataset="<path_to_bridge_dataset>",
        dataset_transform=bridge_dataset.dataset_transform,
        get_prompt=bridge_dataset.get_prompt,
        get_observation=bridge_dataset.get_observation,
        get_action=bridge_dataset.get_action,
        get_state=bridge_dataset.get_state,
    ),
    fractal20220817_data=dict(
        path_to_dataset="<path_to_fractal20220817_data>",
        dataset_transform=fractal20220817_data.dataset_transform,
        get_prompt=fractal20220817_data.get_prompt,
        get_observation=fractal20220817_data.get_observation,
        get_action=fractal20220817_data.get_action,
        get_state=fractal20220817_data.get_state,
    ),
    jaco_play=dict(
        path_to_dataset="<path_to_jaco_play>",
        dataset_transform=jaco_play.dataset_transform,
        get_prompt=jaco_play.get_prompt,
        get_observation=jaco_play.get_observation,
        get_action=jaco_play.get_action,
        get_state=jaco_play.get_state,
    ),
    stanford_hydra_dataset_converted_externally_to_rlds=dict(
        path_to_dataset="<path_to_stanford_hydra_dataset>",
        dataset_transform=stanford_hydra_dataset_converted_externally_to_rlds.dataset_transform,
        get_prompt=stanford_hydra_dataset_converted_externally_to_rlds.get_prompt,
        get_observation=stanford_hydra_dataset_converted_externally_to_rlds.get_observation,
        get_action=stanford_hydra_dataset_converted_externally_to_rlds.get_action,
        get_state=stanford_hydra_dataset_converted_externally_to_rlds.get_state,
    ),
    droid=dict(
        path_to_dataset="<path_to_droid>",
        dataset_transform=droid.dataset_transform,
        get_prompt=droid.get_prompt,
        get_observation=droid.get_observation,
        get_action=droid.get_action,
        get_state=droid.get_state,
    ),
    bc_z=dict(  # v0.1.0
        path_to_dataset="<path_to_bc_z>",
        dataset_transform=bc_z.dataset_transform,
        get_prompt=bc_z.get_prompt,
        get_observation=bc_z.get_observation,
        get_action=bc_z.get_action,
        get_state=bc_z.get_state,
    ),
    berkeley_autolab_ur5=dict(
        path_to_dataset="<path_to_berkeley_autolab_ur5>",
        dataset_transform=berkeley_autolab_ur5.dataset_transform,
        get_prompt=berkeley_autolab_ur5.get_prompt,
        get_observation=berkeley_autolab_ur5.get_observation,
        get_action=berkeley_autolab_ur5.get_action,
        get_state=berkeley_autolab_ur5.get_state,
    ),
    iamlab_cmu_pickup_insert_converted_externally_to_rlds=dict(
        path_to_dataset="<path_to_iamlab_cmu_pickup_insert>",
        dataset_transform=iamlab_cmu_pickup_insert_converted_externally_to_rlds.dataset_transform,
        get_prompt=iamlab_cmu_pickup_insert_converted_externally_to_rlds.get_prompt,
        get_observation=iamlab_cmu_pickup_insert_converted_externally_to_rlds.get_observation,
        get_action=iamlab_cmu_pickup_insert_converted_externally_to_rlds.get_action,
        get_state=iamlab_cmu_pickup_insert_converted_externally_to_rlds.get_state,
    ),
    austin_sirius_dataset_converted_externally_to_rlds=dict(
        path_to_dataset="<path_to_austin_sirius_dataset>",
        dataset_transform=austin_sirius_dataset_converted_externally_to_rlds.dataset_transform,
        get_prompt=austin_sirius_dataset_converted_externally_to_rlds.get_prompt,
        get_observation=austin_sirius_dataset_converted_externally_to_rlds.get_observation,
        get_action=austin_sirius_dataset_converted_externally_to_rlds.get_action,
        get_state=austin_sirius_dataset_converted_externally_to_rlds.get_state,
    ),
    ucsd_kitchen_dataset_converted_externally_to_rlds=dict(
        path_to_dataset="<path_to_ucsd_kitchen_dataset>",
        dataset_transform=ucsd_kitchen_dataset_converted_externally_to_rlds.dataset_transform,
        get_prompt=ucsd_kitchen_dataset_converted_externally_to_rlds.get_prompt,
        get_observation=ucsd_kitchen_dataset_converted_externally_to_rlds.get_observation,
        get_action=ucsd_kitchen_dataset_converted_externally_to_rlds.get_action,
        get_state=ucsd_kitchen_dataset_converted_externally_to_rlds.get_state,
    ),
    language_table = dict(
        path_to_dataset="<path_to_language_table>",
        dataset_transform=language_table.dataset_transform,
        get_prompt=language_table.get_prompt,
        get_observation=language_table.get_observation,
        get_action=language_table.get_action,
        get_state=language_table.get_state,
    ),
    utaustin_mutex=dict(
        path_to_dataset="<path_to_utaustin_mutex>",
        dataset_transform=utaustin_mutex.dataset_transform,
        get_prompt=utaustin_mutex.get_prompt,
        get_observation=utaustin_mutex.get_observation,
        get_action=utaustin_mutex.get_action,
        get_state=utaustin_mutex.get_state,
    )
)