vp_suite.datasets.bair

class BAIRPushingDataset(split, **dataset_kwargs)

Bases: vp_suite.base.base_dataset.VPDataset

Dataset class for the dataset “BAIR Robotic Pushing”, as firstly encountered in “Self-Supervised Visual Planning with Temporal Skip Connections” by Ebert et al. (https://arxiv.org/abs/1710.05268).

Each sequence depicts a robotic gripper arm that randomly moves its end effector inside a bin filled with objects, occasionally stirring them around. This dataset is split into training and test data by default and provides RGB images of size 64x64 as well as 4D action vectors describing the robot gripper movement per-frame.

ACTION_SIZE: int = 4

The size of the action vector per frame (If the dataset provides no actions, this value is 0).

DATASET_FRAME_SHAPE: (<class 'int'>, <class 'int'>, <class 'int'>) = (64, 64, 3)

Shape of a single frame in the dataset (height, width, channels).

DEFAULT_DATA_DIR: pathlib.Path = PosixPath('/home/runner/work/vp-suite/vp-suite/vp-suite-data/data/bair_robot_pushing')

The default save location of the dataset files.

IS_DOWNLOADABLE: str = 'Yes'

A string identifying whether the dataset can be (freely) downloaded.

MIN_SEQ_LEN: int = 30

The minimum sequence length provided by the dataset.

NAME: str = 'BAIR robot pushing'

The dataset’s name.

REFERENCE: str = 'https://arxiv.org/abs/1710.05268'

The reference (publication) where the original dataset is introduced.

__init__(split, **dataset_kwargs)

Initializes the dataset loader by determining its split and extracting and processing all dataset attributes from the parameters given in dataset_kwargs.

Parameters

• split (str) – The dataset’s split identifier (i.e. whether it’s a training/validation/test dataset)

• **dataset_kwargs (Any) – Optional dataset arguments for image transformation, value_range, splitting etc.

classmethod download_and_prepare_dataset()

Downloads the specific dataset, prepares it for the video prediction task (if needed) and stores it in a default location in the ‘data/’ folder. Implemented by the derived dataset classes.

train_to_val_ratio: float = 0.96

The ratio of files that will be training data (rest will be validation data). For bigger datasets, this ratio can be set closer to 1.

download_and_extract_bair(d_path)

Download the BAIR dataset and unpack the downloaded archives into the specified target folder.

Parameters

d_path (Path) – Specified path to dataset.

split_bair_traj_files(data_dir, delete_tfrecords)

Pre-processes downloaded BAIR pushing data by extracting the per-frame image and action information from the provided .tfrecord files and store the information in numpy arrays.

Parameters

• data_dir (Path) – Specified path to dataset.

• delete_tfrecords (bool) – If True, delete the .tfrecord files after extraction.