get_episode_len.py

#!/usr/bin/env python3

import argparse
import os

import h5py
import IPython

e = IPython.embed


def main(args):
    dataset_dir = args['dataset_dir']
    episode_idx = args['episode_idx']
    dataset_name = f'episode_{episode_idx}'

    dataset_path = os.path.join(dataset_dir, dataset_name + '.hdf5')
    if not os.path.isfile(dataset_path):
        print(f'Dataset does not exist at \n{dataset_path}\n')
        exit()

    with h5py.File(dataset_path, 'r') as root:
        # is_sim = root.attrs['sim']
        # compressed = root.attrs.get('compress', False)
        qpos = root['/observations/qpos'][()]
        print(f'dataset_name: {dataset_name}, episode {episode_idx}, len: {len(qpos)}')


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--dataset_dir',
        action='store',
        type=str,
        help='Dataset dir.',
        required=True,
    )
    parser.add_argument(
        '--episode_idx',
        action='store',
        type=int,
        help='Episode index.',
        required=False,
    )
    parser.add_argument('--ismirror', action='store_true')
    main(vars(parser.parse_args()))