Skip to content

Instantly share code, notes, and snippets.

@lukashermann
Created April 30, 2024 17:38
Show Gist options
  • Select an option

  • Save lukashermann/22bb02d6d139be1b9b157508ecdf8e2f to your computer and use it in GitHub Desktop.

Select an option

Save lukashermann/22bb02d6d139be1b9b157508ecdf8e2f to your computer and use it in GitHub Desktop.
from pathlib import Path
from omegaconf import DictConfig
from calvin_agent.datasets.disk_dataset import DiskDataset
DATASETS_DIR = Path("../calvin_debug_dataset/training")
OBSERVATION_SPACE = DictConfig({"rgb_obs": ['rgb_static', 'rgb_gripper'], # conf/datamodule/observation_space/lang_rgb_static_gripper_rel_act.yaml
"depth_obs": [],
"state_obs": ['robot_obs'],
"actions": ['rel_actions'],
"language": ['language'],
})
PROPRIO_STATE = DictConfig({"n_state_obs": 8, # conf/datamodule/proprioception_dims/robot_no_joints.yaml
"keep_indices": [[0, 7], [14,15]],
"robot_orientation_idx": [3, 6],
"normalize": True,
"normalize_robot_orientation": True,
})
def main():
dataset = DiskDataset(datasets_dir=DATASETS_DIR,
obs_space=OBSERVATION_SPACE,
proprio_state=PROPRIO_STATE,
key="vis",
lang_folder="lang_annotations",
num_workers=0)
print(len(dataset))
for data in dataset:
print(data)
if __name__ == "__main__":
main()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment