def add_data_args(parser):
    """Train/valid/test data arguments."""

    group = parser.add_argument_group('data', 'data configurations')

    group.add_argument('--train_h5_path', type=str, help='Pathway to the train H5 file.')
    group.add_argument('--val_h5_path', type=str, help='Pathway to the val H5 file.')
    group.add_argument('--test_h5_path', type=str, help='Pathway to the test H5 file.')
    group.add_argument('--policy_type',
        type=str,
        choices=['TreeGatePolicy', 'TBranT', 'BranT'],
        help='Type of policy to use.')
    
    return parser